gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: gap-aware v_range with decaying reward-scale floor

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-12 20:34:16 +02:00
parent d42a8f4000
commit 9054b6aefb
3 changed files with 37 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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 {

1
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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); }

5
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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);