From a86ec4f6d1206298a435f855cdf54ba4d503e378 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 1 Mar 2026 16:04:00 +0100 Subject: [PATCH] fix(rl): PPO NaN detection every mini-batch, delete dead DQN clipping code PPO was only checking for NaN losses every 10th epoch, allowing NaN to propagate for up to 9 epochs and corrupt model weights before detection. Now checks every mini-batch in both MLP and LSTM paths. Delete three dead gradient clipping methods from DQN agent: - compute_gradients_and_clip (never called, hardcoded max_norm=1.0) - clip_gradients (returns error, deprecated) - clip_gradients_map (computes clip factor but never applies it) Active DQN training uses AdamOptimizer::backward_step_with_monitoring which correctly delegates to gradient_utils::clip_grad_norm. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/dqn/agent.rs | 62 -------------------------------------- crates/ml/src/ppo/ppo.rs | 56 ++++++++++++++++------------------ 2 files changed, 26 insertions(+), 92 deletions(-) diff --git a/crates/ml/src/dqn/agent.rs b/crates/ml/src/dqn/agent.rs index 3859e7404..632525078 100644 --- a/crates/ml/src/dqn/agent.rs +++ b/crates/ml/src/dqn/agent.rs @@ -513,40 +513,6 @@ impl DQNAgent { Ok(x.detach()) } - /// Compute gradients and apply gradient clipping - /// - /// NOTE: This is a legacy method that is NOT used in the current training pipeline. - /// The correct workflow is: - /// 1. let mut grads = loss.backward()?; - /// 2. clip_grad_norm(&varmap, &mut grads, max_norm)?; - /// 3. optimizer.step(&grads)?; - fn compute_gradients_and_clip(&self, loss: &Tensor) -> Result<(), MLError> { - // Compute gradients via backward pass - let mut grads = loss.backward() - .map_err(|e| MLError::TrainingError(format!("Backward pass failed: {}", e)))?; - - // Apply gradient clipping to prevent exploding gradients - let vars = self.q_network.vars(); - let max_norm = 1.0; - let (_actual_norm, _clipped_norm) = crate::gradient_utils::clip_grad_norm(&vars.all_vars(), &mut grads, max_norm) - .map_err(|e| MLError::TrainingError(format!("Gradient clipping failed: {}", e)))?; - - // Note: In actual use, you would now call optimizer.step(&grads) - // This function doesn't do that, which is why it's marked as dead_code - - Ok(()) - } - - /// Legacy gradient clipping method - NOT USED - /// - /// Gradient clipping must be done with access to the GradStore from backward(). - /// This method signature is incorrect for the Candle v0.9.1 API. - fn clip_gradients(&self, _max_norm: f32) -> Result<(), MLError> { - Err(MLError::TrainingError( - "clip_gradients is deprecated - use clip_grad_norm with GradStore instead".to_owned() - )) - } - fn update_target_network_weights(&mut self) -> Result<(), MLError> { // Implement soft update of target network using Polyak averaging let tau = self.config.tau; // BUG #4 FIX: Use configurable tau from config @@ -783,34 +749,6 @@ impl DQNAgent { .unwrap_or(self.config.learning_rate) } - /// Apply gradient clipping to prevent exploding gradients - fn clip_gradients_map( - &self, - gradients: &mut HashMap, - max_norm: f32, - ) -> Result<(), MLError> { - let mut total_norm = 0.0_f32; - - // Calculate total gradient norm - for grad in gradients.values() { - let grad_norm = grad.powf(2.0)?.sum_all()?.to_scalar::().map_err(|e| { - MLError::TrainingError(format!("Failed to compute gradient norm: {}", e)) - })?; - total_norm += grad_norm; - } - - total_norm = total_norm.sqrt(); - - // Clip gradients if necessary - if total_norm > max_norm { - let _clip_factor = max_norm / total_norm; - // For simplicity, we'll skip gradient clipping for now - // In production, we would create proper scalar tensors for multiplication - } - - Ok(()) - } - /// Update reward statistics for metrics pub fn update_reward_stats(&mut self, episode_reward: f64, episode_won: bool) { self.metrics.total_episodes += 1; diff --git a/crates/ml/src/ppo/ppo.rs b/crates/ml/src/ppo/ppo.rs index 3975e0814..79ce2fa24 100644 --- a/crates/ml/src/ppo/ppo.rs +++ b/crates/ml/src/ppo/ppo.rs @@ -957,21 +957,19 @@ impl PPO { MLError::TrainingError(format!("Failed to extract value loss: {}", e)) })?; - // NaN detection every 10 epochs - if epoch % 10 == 0 { - if policy_loss_scalar.is_nan() { - return Err(MLError::TrainingError( - format!("NaN detected in policy loss at epoch {} - training unstable. \ - Consider reducing learning rate or increasing entropy coefficient.", epoch) - )); - } - if value_loss_scalar.is_nan() { - return Err(MLError::TrainingError(format!( - "NaN detected in value loss at epoch {} - training unstable. \ - Consider reducing learning rate.", - epoch - ))); - } + // NaN detection — check every mini-batch to prevent corruption + if policy_loss_scalar.is_nan() { + return Err(MLError::TrainingError( + format!("NaN detected in policy loss at epoch {} - training unstable. \ + Consider reducing learning rate or increasing entropy coefficient.", epoch) + )); + } + if value_loss_scalar.is_nan() { + return Err(MLError::TrainingError(format!( + "NaN detected in value loss at epoch {} - training unstable. \ + Consider reducing learning rate.", + epoch + ))); } // Accumulate policy gradients @@ -1338,21 +1336,19 @@ impl PPO { MLError::TrainingError(format!("Failed to extract value loss: {}", e)) })?; - // NaN detection every 10 epochs - if epoch % 10 == 0 { - if policy_loss_scalar.is_nan() { - return Err(MLError::TrainingError( - format!("NaN detected in policy loss at epoch {} - LSTM training unstable. \ - Consider reducing learning rate or sequence length.", epoch) - )); - } - if value_loss_scalar.is_nan() { - return Err(MLError::TrainingError(format!( - "NaN detected in value loss at epoch {} - LSTM training unstable. \ - Consider reducing learning rate.", - epoch - ))); - } + // NaN detection — check every mini-batch to prevent corruption + if policy_loss_scalar.is_nan() { + return Err(MLError::TrainingError( + format!("NaN detected in policy loss at epoch {} - LSTM training unstable. \ + Consider reducing learning rate or sequence length.", epoch) + )); + } + if value_loss_scalar.is_nan() { + return Err(MLError::TrainingError(format!( + "NaN detected in value loss at epoch {} - LSTM training unstable. \ + Consider reducing learning rate.", + epoch + ))); } // Update policy network with gradient clipping