fix: eval v_range width from Q-gap, baseline from Q-std — eliminates -40 dips

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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 15:13:36 +02:00
parent 60528390e1
commit 1634dec042
3 changed files with 31 additions and 19 deletions

View File

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

View File

@@ -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::<f32>(),
);
}
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() }

View File

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