From d38a8cf9970da69c14580fb440bcb29781604729 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 23 Apr 2026 19:03:17 +0200 Subject: [PATCH] fix(dqn): clamp reward quantiles to v_min/v_max in C51 warm-start MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of the Q=±333k explosion at every run's epoch 2. `warm_start_atom_positions` writes quantiles from the raw environment reward distribution directly into `atom_positions_buf`. Raw rewards are unbounded PnL-scaled values — a single extreme sample in the first experience buffer becomes `atom_positions[num_atoms-1]`, and the C51 expected-value readback `Q = Σ prob × atom_pos` inherits that magnitude. Observed deterministically across train-7rgqd, train-5gzpn, and train-gj54m: epoch 1 Q in `[0, ~6]` (initial Xavier atoms), epoch 2 Q at exactly `±333406` once warm-start writes the sorted-reward-tail into the atom grid. Every downstream path — the C51 loss projection, eval_v_range EMA, IQL support, HEALTH_DIAG q_gap — assumes `atom_positions ∈ [v_min, v_max]`. The warm-start path was the only one bypassing that assumption. Clamp each quantile to the configured `[v_min, v_max]` before writing. This is a safety rail, not a tuning parameter: config.v_{min,max} are already derived from reward_scale (±15 default), which is the support range the rest of the system expects. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/trainers/dqn/trainer/training_loop.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) 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}"); } }