From 1634dec042e85ac816c97780e53bf1b1f6ef5fa2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 15:13:36 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20eval=20v=5Frange=20width=20from=20Q-gap,?= =?UTF-8?q?=20baseline=20from=20Q-std=20=E2=80=94=20eliminates=20-40=20dip?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes: 1. Width: half = max(10*q_gap, 3*q_std, 0.1) instead of fixed 3*q_std+1.0. With q_gap=0.06 and 51 atoms, gives ~5 atoms of action resolution (was 1.5 atoms with the 1.0 floor). No more catastrophic -40 dips from atom underresolution. 2. Baseline: uses q_std_ema (proportional to natural Q oscillation) instead of fixed 0.01. Normal Q-oscillation (±q_std) no longer triggers aggressive alpha — only abnormal shifts do. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 26 ++++++++++--------- crates/ml/src/trainers/dqn/fused_training.rs | 19 ++++++++++---- .../src/trainers/dqn/trainer/training_loop.rs | 5 ++-- 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4ebc57fc9..d069682ed 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1016,29 +1016,31 @@ impl GpuDqnTrainer { /// Update eval v_range from observed Q-value statistics. /// Called at epoch boundary after compute_q_stats. - pub fn update_eval_v_range(&mut self, q_mean: f32, q_std: f32) { - // Adaptive-rate EMA: tracks fast when Q-values shift, smooths when stable. - // Alpha = prediction_error / (prediction_error + baseline). When the EMA - // is far from the current value (shift detected), alpha → 0.5 (fast). - // When close (stable), alpha → 0.01 (smooth). Self-calibrating. + pub fn update_eval_v_range(&mut self, q_mean: f32, q_std: f32, q_gap: f32) { + // Adaptive-rate EMA with baseline proportional to Q-std. + // Baseline = max(q_std_ema, 0.001) so normal Q-oscillation doesn't + // trigger aggressive tracking — only abnormal shifts do. if !self.eval_ema_initialized { - // Seed from config v_range center (0.0) and width — no discontinuous snap. - // The adaptive alpha will quickly track to actual Q-stats from here. self.eval_q_mean_ema = 0.0; - self.eval_q_std_ema = 1.0; // gives v_range [-4, 4] initially + self.eval_q_std_ema = q_std.max(0.01); self.eval_ema_initialized = true; } { + let baseline = self.eval_q_std_ema.max(0.001); let mean_err = (q_mean - self.eval_q_mean_ema).abs(); let std_err = (q_std - self.eval_q_std_ema).abs(); - let baseline = 0.01_f32; - let alpha_mean = (mean_err / (mean_err + baseline)).clamp(0.01, 0.5); - let alpha_std = (std_err / (std_err + baseline)).clamp(0.01, 0.5); + let alpha_mean = (mean_err / (mean_err + baseline)).clamp(0.01, 0.3); + let alpha_std = (std_err / (std_err + baseline)).clamp(0.01, 0.3); self.eval_q_mean_ema = (1.0 - alpha_mean) * self.eval_q_mean_ema + alpha_mean * q_mean; self.eval_q_std_ema = (1.0 - alpha_std) * self.eval_q_std_ema + alpha_std * q_std; } - let half = (3.0 * self.eval_q_std_ema + 1.0).max(1.0); + // Width from Q-gap (action differentiation range), not Q-std. + // Floor = max(10 * q_gap, 3 * q_std) — enough atoms to resolve actions. + // With 51 atoms, a width of 10*q_gap gives q_gap/delta_z ≈ 5 atoms of + // resolution between best and worst action. + let gap_width = (10.0 * q_gap).max(3.0 * self.eval_q_std_ema).max(0.1); + let half = gap_width; let v_min = self.eval_q_mean_ema - half; let v_max = self.eval_q_mean_ema + half; unsafe { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index dcc29f87a..2ff145c58 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2042,14 +2042,23 @@ impl FusedTrainingCtx { pub(crate) fn per_sample_support_ptr(&self) -> u64 { self.gpu_iql.per_sample_support_ptr() } /// Fixed wide v_range pointer for eval compute_expected_q (not C51 loss). pub(crate) fn eval_v_range_ptr(&self) -> u64 { self.trainer.eval_v_range_ptr() } - /// Return EMA-smoothed eval v_range as [v_min, v_max]. + /// Return current eval v_range as [v_min, v_max] via synchronous DtoH. + /// Reads the same buffer that update_eval_v_range wrote. pub(crate) fn eval_v_range(&self) -> [f32; 2] { - let half = (3.0 * self.trainer.eval_q_std_ema + 1.0_f32).max(1.0); - [self.trainer.eval_q_mean_ema - half, self.trainer.eval_q_mean_ema + half] + let mut vr = [0.0_f32; 2]; + unsafe { + cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); + cudarc::driver::sys::cuMemcpyDtoH_v2( + vr.as_mut_ptr().cast(), + self.trainer.eval_v_range_ptr(), + 2 * std::mem::size_of::(), + ); + } + vr } /// Update eval v_range from observed Q-value statistics. - pub(crate) fn update_eval_v_range(&mut self, q_mean: f32, q_std: f32) { - self.trainer.update_eval_v_range(q_mean, q_std); + pub(crate) fn update_eval_v_range(&mut self, q_mean: f32, q_std: f32, q_gap: f32) { + self.trainer.update_eval_v_range(q_mean, q_std, q_gap); } /// Per-sample epsilon from IQL expectile gap. pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index d783e0836..e9c2cc7f6 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1389,9 +1389,10 @@ impl DQNTrainer { self.epoch_atom_entropy = stats.atom_entropy; self.epoch_atom_utilization = stats.atom_utilization; - // Update eval v_range from observed Q-stats + // Update eval v_range from observed Q-stats + Q-gap let q_std = stats.q_variance.max(0.0).sqrt(); - fused.update_eval_v_range(stats.q_mean, q_std); + let q_gap = (stats.avg_max_q as f32 - stats.q_mean).max(0.0); + fused.update_eval_v_range(stats.q_mean, q_std, q_gap); } } train_step_count += 1;