fix(dqn): clamp reward quantiles to v_min/v_max in C51 warm-start

Root cause of the Q=±333k explosion at every run's epoch 2.

`warm_start_atom_positions` writes quantiles from the raw environment
reward distribution directly into `atom_positions_buf`. Raw rewards
are unbounded PnL-scaled values — a single extreme sample in the
first experience buffer becomes `atom_positions[num_atoms-1]`, and
the C51 expected-value readback `Q = Σ prob × atom_pos` inherits
that magnitude.

Observed deterministically across train-7rgqd, train-5gzpn, and
train-gj54m: epoch 1 Q in `[0, ~6]` (initial Xavier atoms), epoch 2
Q at exactly `±333406` once warm-start writes the sorted-reward-tail
into the atom grid. Every downstream path — the C51 loss projection,
eval_v_range EMA, IQL support, HEALTH_DIAG q_gap — assumes
`atom_positions ∈ [v_min, v_max]`. The warm-start path was the only
one bypassing that assumption.

Clamp each quantile to the configured `[v_min, v_max]` before writing.
This is a safety rail, not a tuning parameter: config.v_{min,max} are
already derived from reward_scale (±15 default), which is the support
range the rest of the system expects.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 19:03:17 +02:00
parent 9c3ddf8b37
commit d38a8cf997

View File

@@ -477,8 +477,23 @@ impl DQNTrainer {
let num_atoms = self.hyperparams.num_atoms;
match collector.compute_reward_quantiles(count, num_atoms) {
Ok(quantiles) => {
// Raw reward quantiles are unbounded — a single extreme
// PnL sample from the environment becomes `atom_positions[50]`,
// and the C51 expected-value read `Q = Σ prob × atom_pos`
// inherits that magnitude. Observed on L40S: epoch 1 Q in
// [0, 6], epoch 2 Q at ±333k — the tail of the sorted
// reward distribution from the first experience buffer.
// Clamp to the configured `[v_min, v_max]` so warm-start
// respects the C51 support bounds that every downstream
// path (loss projection, eval_v_range, IQL support) assumes.
let v_min = self.hyperparams.v_min as f32;
let v_max = self.hyperparams.v_max as f32;
let clamped_quantiles: Vec<f32> = quantiles
.iter()
.map(|&q| q.clamp(v_min, v_max))
.collect();
if let Some(ref mut fused) = self.fused_ctx {
if let Err(e) = fused.warm_start_atom_positions(&quantiles) {
if let Err(e) = fused.warm_start_atom_positions(&clamped_quantiles) {
warn!("C51 atom warm-start failed (non-fatal): {e}");
}
}