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