diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index bc16def12..5de0b01e8 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -477,8 +477,23 @@ impl DQNTrainer { let num_atoms = self.hyperparams.num_atoms; match collector.compute_reward_quantiles(count, num_atoms) { Ok(quantiles) => { + // Raw reward quantiles are unbounded — a single extreme + // PnL sample from the environment becomes `atom_positions[50]`, + // and the C51 expected-value read `Q = Σ prob × atom_pos` + // inherits that magnitude. Observed on L40S: epoch 1 Q in + // [0, 6], epoch 2 Q at ±333k — the tail of the sorted + // reward distribution from the first experience buffer. + // Clamp to the configured `[v_min, v_max]` so warm-start + // respects the C51 support bounds that every downstream + // path (loss projection, eval_v_range, IQL support) assumes. + let v_min = self.hyperparams.v_min as f32; + let v_max = self.hyperparams.v_max as f32; + let clamped_quantiles: Vec = quantiles + .iter() + .map(|&q| q.clamp(v_min, v_max)) + .collect(); if let Some(ref mut fused) = self.fused_ctx { - if let Err(e) = fused.warm_start_atom_positions(&quantiles) { + if let Err(e) = fused.warm_start_atom_positions(&clamped_quantiles) { warn!("C51 atom warm-start failed (non-fatal): {e}"); } }