From c05f75d0d842fade1d80d3d638dbab2ce437bfbe Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 12 Apr 2026 21:03:44 +0200 Subject: [PATCH] feat: adaptive v_range reward-scale floor from observed reward_std MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hardcoded REWARD_SCALE_FLOOR=0.01 was tuned for smoketest (reward_std=0.007) but 940× too small for production (reward_std=6.57). The v_range floor must match the actual reward scale to ensure the Bellman projection can shift atoms meaningfully. - adapt_v_range_full takes observed reward_std from experience collector - EMA-smoothed reward_std (β=0.99) prevents single-epoch noise from jerking floor - Falls back to 0.01 before first observation, then adapts automatically - reward_std_ema field on GpuDqnTrainer, observed_reward_std on DQNTrainer - Decaying floor uses actual reward scale: floor = R_std * exp(-|Q_mean|/R_std) 903/903 tests passing. Smoketest val_Sharpe positive (60.76 epoch 1). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 65 ++++++++++++------- crates/ml/src/trainers/dqn/fused_training.rs | 1 + .../src/trainers/dqn/trainer/constructor.rs | 1 + crates/ml/src/trainers/dqn/trainer/mod.rs | 2 + .../src/trainers/dqn/trainer/training_loop.rs | 5 +- 5 files changed, 48 insertions(+), 26 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f20b97ea1..ad687d76d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -645,6 +645,8 @@ pub struct GpuDqnTrainer { grad_norm_ema: f32, /// EMA of Q-divergence for adaptive tau computation. q_div_ema: f32, + /// EMA of observed reward std for adaptive v_range floor. + reward_std_ema: f32, // ── Training state ────────────────────────────────────────────── pub(crate) adam_step: i32, @@ -3099,6 +3101,7 @@ impl GpuDqnTrainer { adaptive_clip_dev_ptr, grad_norm_ema: 0.0, q_div_ema: 0.0, + reward_std_ema: 0.0, adam_step: 0, total_params, params_initialized: false, @@ -6464,22 +6467,24 @@ 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) + self.adapt_v_range_full(q_mean, q_variance, 0.0, 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 { + self.adapt_v_range_full(q_mean, q_variance, q_gap, 0.0) + } + + /// Fully adaptive v_range from Q-stats + reward scale. + /// + /// The reward_std parameter comes from the experience collector's observed + /// reward statistics. It varies 940× between smoketest (0.007) and production + /// (6.57) — hardcoding it was fundamentally wrong. + /// + /// The v_range floor = reward_std (single Bellman step shift magnitude). + /// This ensures the C51 projection can shift atoms by at least one reward + /// unit, breaking the stable fixed point at any scale. + pub fn adapt_v_range_full(&mut self, q_mean: f32, q_variance: f32, q_gap: f32, reward_std: f32) -> bool { const SIGMA_COVERAGE: f32 = 3.0; - // 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(); @@ -6490,20 +6495,32 @@ impl GpuDqnTrainer { let target_min = q_mean - total_half; let target_max = q_mean + total_half; - // 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 + // Reward-scale floor: derived from ACTUAL observed reward_std. + // Smoketest: reward_std ≈ 0.007 → floor = 0.007 + // Production: reward_std ≈ 6.57 → floor = 6.57 + // Falls back to EMA of observed reward_std (self.reward_std_ema) + // when per-step reward_std is not available. + let effective_reward_std = if reward_std > 1e-8 { + // Update EMA with fresh observation + const R_BETA: f32 = 0.99; + if self.reward_std_ema <= 1e-8 { + self.reward_std_ema = reward_std; + } else { + self.reward_std_ema = R_BETA * self.reward_std_ema + (1.0 - R_BETA) * reward_std; + } + self.reward_std_ema + } else if self.reward_std_ema > 1e-8 { + self.reward_std_ema // use cached EMA + } else { + 0.01 // absolute fallback before first observation + }; + + // Decay floor as Q-values mature (same exp decay as before, but + // using the ACTUAL reward scale instead of hardcoded 0.01) + let maturity = (q_mean.abs() / effective_reward_std.max(1e-8)).min(5.0); 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 decayed_floor = effective_reward_std * (-maturity).exp(); let adaptive_min = decayed_floor .max(gap_based_range) .max(target_max - target_min); diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index ae529db34..285b980d6 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1991,6 +1991,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 adapt_v_range_full(&mut self, q_mean: f32, q_variance: f32, q_gap: f32, reward_std: f32) -> bool { self.trainer.adapt_v_range_full(q_mean, q_variance, q_gap, reward_std) } 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/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 972fcc964..1805f2425 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -714,6 +714,7 @@ impl DQNTrainer { epoch_q_gap: 0.0, epoch_atom_entropy: 0.0, epoch_atom_utilization: 0.0, + observed_reward_std: 0.0, // GPU pipeline: pre-uploaded training data (initialized lazily at first epoch) gpu_data: None, diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 9fb63f03c..9cf031dfd 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -225,6 +225,8 @@ pub struct DQNTrainer { /// Atom entropy [0,1] and utilization [0,1] from C51 distributional component pub(crate) epoch_atom_entropy: f32, pub(crate) epoch_atom_utilization: f32, + /// Observed reward std from experience collector (for adaptive v_range floor). + pub(crate) observed_reward_std: f32, /// Consecutive low-drawdown epochs (for adversarial regime injection) pub(crate) consecutive_low_dd_epochs: usize, /// Whether adversarial regime is active this epoch diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 66a590eee..a93f16832 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1375,7 +1375,7 @@ impl DQNTrainer { self.epoch_atom_entropy = stats.atom_entropy; self.epoch_atom_utilization = stats.atom_utilization; 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); + fused.adapt_v_range_full(stats.q_mean, stats.q_variance, q_gap, self.observed_reward_std); } } train_step_count += 1; @@ -1511,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_with_gap(self.cached_avg_q as f32, self.epoch_q_variance, self.epoch_q_gap) { + if fused.adapt_v_range_full(self.cached_avg_q as f32, self.epoch_q_variance, self.epoch_q_gap, self.observed_reward_std) { 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); @@ -1590,6 +1590,7 @@ impl DQNTrainer { summary.total_trades, summary.action_counts ); monitor.track_reward(summary.mean_reward); + self.observed_reward_std = summary.reward_std; // Feed all 3 branch distributions from GPU into the monitor. // These are the actual model-selected actions, not // deterministic OrderRouter re-derivations.