diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 3d4850d0b..6951daecb 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2148,10 +2148,26 @@ impl GpuDqnTrainer { // Floor = max(10 * q_gap, 3 * q_std) — enough atoms to resolve actions. // With 51 atoms, a width of 10*q_gap gives q_gap/delta_z ≈ 5 atoms of // resolution between best and worst action. + // + // Clamp the adaptive [v_min, v_max] to the theoretical bounds supplied + // at construction (config.v_{min,max}, derived from reward_scale/gamma). + // Without this clamp the window can drift arbitrarily far from 0: once + // the policy overestimates, eval_q_mean_ema tracks it up → v_max follows + // → C51 atoms span a wider range → even larger TD targets → loss + // explodes geometrically (Fold 1 smoke: final_loss=1.6e16, grad=62B). + // + // Two clamps are needed — one on half-width so the band can't exceed + // the full theoretical range, and one on the centre so the band always + // fits inside [config.v_min, config.v_max]: let gap_width = (10.0 * q_gap).max(3.0 * self.eval_q_std_ema).max(0.1); - let half = gap_width; - let v_min = self.eval_q_mean_ema - half; - let v_max = self.eval_q_mean_ema + half; + let abs_half = (self.config.v_max - self.config.v_min) * 0.5; + let half = gap_width.min(abs_half); + let center = self.eval_q_mean_ema.clamp( + self.config.v_min + half, + self.config.v_max - half, + ); + let v_min = center - half; + let v_max = center + half; // Pinned device-mapped: CPU writes directly, GPU reads via dev_ptr. No HtoD copy. unsafe { *self.eval_v_range_pinned = v_min;