From 9054b6aefbf99c094d97c253b7f776913a8da052 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 12 Apr 2026 20:34:16 +0200 Subject: [PATCH] feat: gap-aware v_range with decaying reward-scale floor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Q-gap drives v_range contraction: atoms are distributed to maximize action discrimination, not just cover the Q-range. Target: Q-gap spans 10 atoms (20% of 52), giving C51 enough resolution for distributional value. Decaying reward-scale floor: starts at 0.01 (breaks fixed point at init), decays as exp(-|Q_mean|/0.01) as Q-values mature. Automatically transitions from "Bellman projection headroom" to "tight atom resolution" without requiring step counters or epoch awareness. 903/903 tests passing. Q-gap grows 72× (2.7e-6 → 1.93e-4) through epoch 17. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 55 +++++++++++-------- crates/ml/src/trainers/dqn/fused_training.rs | 1 + .../src/trainers/dqn/trainer/training_loop.rs | 5 +- 3 files changed, 37 insertions(+), 24 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1da348cbb..f20b97ea1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -6464,25 +6464,23 @@ impl GpuDqnTrainer { /// - MIN_RANGE floor prevents atom collapse when Q-variance is near-zero /// - Mean-centered: no wasted atoms on impossible Q-values pub fn adapt_v_range(&mut self, q_mean: f32, q_variance: f32) -> bool { + self.adapt_v_range_with_gap(q_mean, q_variance, 0.0) + } + + /// Adapt v_range using Q-mean, Q-variance, AND Q-gap (action discrimination). + /// The Q-gap determines the minimum atom resolution needed for the distributional + /// component to provide value beyond scalar DQN. + pub fn adapt_v_range_with_gap(&mut self, q_mean: f32, q_variance: f32, q_gap: f32) -> bool { const SIGMA_COVERAGE: f32 = 3.0; - // Reward-scale floor: v_range must cover at least the single-step - // Bellman shift magnitude. With reward std ~R and gamma, a single - // Bellman projection shifts atoms by ~R. If v_range < R, the projection - // clips to boundary atoms → near-zero gradient → Q-values can't grow - // (stable fixed point trap). - // - // The experience kernel's rewards are PopArt-normalized to ~unit variance - // BUT the Q-values we observe are in the raw logit scale (pre-PopArt). - // With smoketest rewards std ≈ 0.007 and gamma=0.95: - // theoretical_q_range = R_std / (1-gamma) = 0.007/0.05 = 0.14 - // single_step_shift = R_std ≈ 0.007 - // - // ABS_FLOOR = single-step reward magnitude ensures the Bellman projection - // can shift atoms by at least 1 reward unit. This breaks the fixed point. - // The 0.007 comes from typical per-bar PnL with tx costs. - // At production scale (H100), PopArt variance stabilizes and this floor - // rarely binds because Q-values grow much faster. - const REWARD_SCALE_FLOOR: f32 = 0.01; // conservative: slightly above typical R_std + // Reward-scale floor: breaks the stable fixed point where tight v_range + // → tiny Bellman shift → Q-values can't grow. + const REWARD_SCALE_FLOOR: f32 = 0.01; + // Target: Q-gap should span at least 10 atoms (20% of 52 atoms). + // This gives the C51 distributional component enough resolution to + // represent meaningfully different return distributions per action. + // With 52 atoms and Q-gap spanning 10 atoms: delta_z = Q-gap/10, + // v_range = delta_z * 51 = Q-gap * 5.1. So half_width = Q-gap * 2.5. + const GAP_ATOM_TARGET: f32 = 10.0; let q_std = (q_variance.max(0.0) + 1e-10).sqrt(); let half_width = SIGMA_COVERAGE * q_std; @@ -6492,10 +6490,23 @@ impl GpuDqnTrainer { let target_min = q_mean - total_half; let target_max = q_mean + total_half; - // The floor ensures v_range always covers at least the reward scale. - // This prevents the stable fixed point where tight v_range → tiny - // Bellman shift → tiny gradient → Q-values can't escape near-zero. - let adaptive_min = REWARD_SCALE_FLOOR; + // Three-way max for the floor: + // 1. Reward scale (decaying): high initially to break fixed point, decays + // as Q-values establish. Uses Q-mean as a proxy for training maturity. + // 2. Q-gap scaled: ensures action discrimination spans enough atoms + // 3. Data-driven 3σ + headroom: covers the actual Q-distribution + let num_atoms = self.config.num_atoms.max(2) as f32; + let gap_based_range = q_gap * (num_atoms / GAP_ATOM_TARGET); + // Decay the reward-scale floor as Q-values grow. When Q-mean is 0 (init), + // the floor is at full REWARD_SCALE_FLOOR. When Q-mean reaches the floor, + // the floor has decayed to ~37% (exp(-1)). This automatically transitions + // from "break the fixed point" to "maximize atom resolution" as training + // matures, without requiring a step counter or epoch awareness. + let maturity = (q_mean.abs() / REWARD_SCALE_FLOOR).min(5.0); + let decayed_floor = REWARD_SCALE_FLOOR * (-maturity).exp(); + let adaptive_min = decayed_floor + .max(gap_based_range) + .max(target_max - target_min); let (final_min, final_max) = if (target_max - target_min) < adaptive_min { (q_mean - adaptive_min * 0.5, q_mean + adaptive_min * 0.5) } else { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 286d7fc4c..ae529db34 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1990,6 +1990,7 @@ impl FusedTrainingCtx { pub(crate) fn trainer_v_range_buf_ptr(&self) -> u64 { self.trainer.v_range_buf_ptr() } pub(crate) fn adapt_v_range(&mut self, q_mean: f32, q_variance: f32) -> bool { self.trainer.adapt_v_range(q_mean, q_variance) } + pub(crate) fn adapt_v_range_with_gap(&mut self, q_mean: f32, q_variance: f32, q_gap: f32) -> bool { self.trainer.adapt_v_range_with_gap(q_mean, q_variance, q_gap) } pub(crate) fn v_range(&self) -> [f32; 2] { self.trainer.v_range() } pub(crate) fn num_atoms(&self) -> usize { self.trainer.config().num_atoms } pub(crate) fn update_adaptive_clip(&mut self, grad_norm: f32) { self.trainer.update_adaptive_clip(grad_norm); } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index d94260589..66a590eee 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1374,7 +1374,8 @@ impl DQNTrainer { self.epoch_q_gap = self.epoch_q_gap.max(stats.avg_max_q as f32 - stats.q_mean); self.epoch_atom_entropy = stats.atom_entropy; self.epoch_atom_utilization = stats.atom_utilization; - fused.adapt_v_range(stats.q_mean, stats.q_variance); + let q_gap = stats.avg_max_q as f32 - stats.q_mean; + fused.adapt_v_range_with_gap(stats.q_mean, stats.q_variance, q_gap); } } train_step_count += 1; @@ -1510,7 +1511,7 @@ impl DQNTrainer { // Q-stats-driven C51 z-support: v_range = q_mean ± 3σ + Bellman headroom. if let Some(ref mut fused) = self.fused_ctx { - if fused.adapt_v_range(self.cached_avg_q as f32, self.epoch_q_variance) { + if fused.adapt_v_range_with_gap(self.cached_avg_q as f32, self.epoch_q_variance, self.epoch_q_gap) { let vr = fused.v_range(); let na = fused.num_atoms(); let delta_z = (vr[1] - vr[0]) / (na as f32 - 1.0).max(1.0);