diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2ff145c58..74df8bb36 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1001,6 +1001,13 @@ impl FusedTrainingCtx { } self.pending_vaccine_batch = None; + // ── Adaptive IQN lambda: read loss OUTSIDE graph capture ───────── + if let Some(ref iqn) = self.gpu_iqn { + if let Ok(loss) = iqn.read_total_loss() { + self.trainer.update_iqn_readiness(loss); + } + } + // ── Clip grad_buf L2 norm BEFORE Adam (outside graph) ─────────── // Root cause: backward weight gradient dW = activation^T × d_output. // Spectral norm bounds W, IS-weight clamp bounds d_logits (~2.0), @@ -1272,11 +1279,6 @@ impl FusedTrainingCtx { dqn_actions, dqn_rewards, dqn_dones, ) { Ok(_) => { - // Update adaptive IQN lambda from current loss - if let Ok(loss) = iqn.read_total_loss() { - self.trainer.update_iqn_readiness(loss); - } - let d_h_s2_ptr = iqn.d_h_s2_raw_ptr(); self.trainer.apply_iqn_trunk_gradient( d_h_s2_ptr, &mut self.online_dueling,