gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(rl): PPO NaN detection every mini-batch, delete dead DQN clipping code

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-01 16:04:00 +01:00
parent 7744becdd5
commit a86ec4f6d1
2 changed files with 26 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
crates/ml/src/dqn/agent.rs
View File

@@ -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<String, Tensor>,
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::<f32>().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;

56
crates/ml/src/ppo/ppo.rs
View File

@@ -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