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:
@@ -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 {
|
||||
|
||||
@@ -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() }
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user