gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(iqn-readiness): reset iqn_loss_initial + ema + readiness at fold boundary

Browse Source 
The IQN readiness gauge on `GpuDqnTrainer` is a streaming improvement
fraction `(iqn_loss_initial - iqn_loss_ema) / iqn_loss_initial`
(gpu_dqn_trainer.rs:5805-5818). `iqn_loss_initial` is captured once on
the first `update_iqn_readiness` call where it is still ~0 and never
resets — across folds it stays pinned to the fold-0 epoch-1 IQN loss.
The pinned device-mapped readiness slot is read by `c51_loss_kernel`
as the CVaR α and by IQN gradient-weight gating, so a stale anchor
either pins readiness near 1 (over-confident, full IQN gradient
weight on a still-recovering head) or near 0 (under-weighting a
converged head).

Fix:
- New `pub fn reset_iqn_readiness_state` on `GpuDqnTrainer` zeroes the
  three coupled scalars (iqn_loss_initial / iqn_loss_ema /
  iqn_readiness) and writes 0.0 through the device-mapped pinned host
  slot — no HtoD copy, the mapping propagates the host write directly.
- New `pub(crate) fn reset_iqn_readiness_state` wrapper on
  `FusedTrainingCtx` mirrors the existing `reset_eval_v_range_state`
  pattern.
- Three new `FoldReset` registry entries (`isv_iqn_loss_initial`,
  `isv_iqn_loss_ema`, `isv_iqn_readiness`) dispatch through one match
  arm in `reset_named_state` (training_loop.rs).

`update_iqn_readiness` already handles `iqn_loss_initial < 1e-12` as
the bootstrap branch, so the new fold's first batch is recaptured as
the anchor on the first call after reset, by design.

Per `feedback_no_partial_refactor.md` — same fold-boundary
state-migration contract, all three coupled scalars migrate together.

This commit completes the 3-fix sequence (Fix 1 IQN target sync,
Fix 2 aux Adam states, Fix 3 IQN readiness) closing the fold-boundary
contract gap surfaced by the audit. Bug signature being resolved:
  F0 ep5 Q=+0.79  (healthy)
  F1 ep1 Q=+0.82  (boundary fine; c51_alpha warmup blends IQN low)
  F1 ep2 Q=+2.23  (drift starts — IQN online↔target gap widens)
  F1 ep3 Q=+4.05
  F1 ep4 Q=+10.66 (geometric ~2.3×/epoch)
  F1 ep5 NaN at step 5 (fp32 overflow past atom support)

Fix 1 (IQN target sync) is the load-bearing fix for the geometric
drift itself; Fix 2 prevents first-epoch Adam-step overshoot in
seven aux optimizers; Fix 3 closes the gauge-staleness path.

cargo check -p ml --lib clean at 13 warnings (workspace baseline).
cargo test -p ml --lib --no-run clean. All three existing
`state_reset_registry::tests` pass without modification.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 19:17:47 +02:00
parent e1efa8994f
commit a894eca98f
5 changed files with 139 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -13152,6 +13152,38 @@ impl GpuDqnTrainer {
}
}
/// Reset IQN readiness gating state at fold boundary.
///
/// `update_iqn_readiness` (gpu_dqn_trainer.rs:5804) is a streaming
/// improvement gauge: `iqn_readiness = (initial - ema) / initial`.
/// `iqn_loss_initial` is captured once on the first call where it is
/// still ~0 and never resets — across folds it stays pinned to the
/// fold-0 epoch-1 IQN loss. After fold-boundary shrink-and-perturb +
/// IQN target sync the new fold's first IQN losses can land far
/// from the fold-0 baseline (different reward distribution post-
/// adversarial regime, fresh online↔target alignment), so the
/// readiness fraction is computed against a stale anchor — either
/// pinning readiness near 1 (over-confident, full IQN gradient
/// weight on a still-recovering head) or near 0 (under-weighting a
/// converged head).
///
/// **Fix**: zero `iqn_loss_initial` so the next `update_iqn_readiness`
/// call captures fold N+1's first batch as the new baseline; zero
/// the EMA + the readiness scalar; write 0.0 through the device-
/// mapped pinned slot so any in-flight kernel reads the reset value
/// without an HtoD copy. `update_iqn_readiness` already handles the
/// `iqn_loss_initial < 1e-12` case as the bootstrap branch.
pub fn reset_iqn_readiness_state(&mut self) {
self.iqn_loss_initial = 0.0;
self.iqn_loss_ema = 0.0;
self.iqn_readiness = 0.0;
if !self.iqn_readiness_pinned.is_null() {
// Pinned device-mapped: writing the host slot propagates
// immediately to GPU via the mapping; no HtoD copy issued.
unsafe { *self.iqn_readiness_pinned = 0.0; }
}
}
/// Set per_branch_q_gap_ema — HtoD for trajectory backtracking restore.
pub fn set_per_branch_q_gap_ema(&mut self, v: [f32; 4]) {
unsafe {

11
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -996,6 +996,17 @@ impl FusedTrainingCtx {
self.trainer.reset_eval_v_range_state();
}
/// Reset the IQN readiness gauge (`iqn_loss_initial`, `iqn_loss_ema`,
/// `iqn_readiness` scalar + pinned slot).
///
/// Called from DQNTrainer::reset_named_state for the
/// "isv_iqn_loss_initial" / "isv_iqn_loss_ema" / "isv_iqn_readiness"
/// FoldReset entries. The three names map to the same logical reset
/// (the streaming-improvement gauge's three coupled scalars).
pub(crate) fn reset_iqn_readiness_state(&mut self) {
self.trainer.reset_iqn_readiness_state();
}
/// Steps since last flat buffer sync.
pub(crate) fn steps_since_varmap_sync(&self) -> usize {
self.steps_since_varmap_sync

26
crates/ml/src/trainers/dqn/state_reset_registry.rs
View File

@@ -77,6 +77,32 @@ impl StateResetRegistry {
category: ResetCategory::FoldReset,
description: "ISV[13..22) — per-direction/magnitude Q-mean EMAs + |Q| refs",
},
// ───── IQN readiness gauge (Fix 3 of 3 fold-boundary audit) ─────
// Three coupled scalars on `GpuDqnTrainer`:
// `iqn_loss_initial`: anchor captured on first `update_iqn_readiness`
// `iqn_loss_ema`: streaming smoothed loss
// `iqn_readiness`: (initial - ema)/initial gauge ∈ [0,1]
// plus the pinned-host-mapped device slot read by c51_loss
// kernel as CVaR α + by gating consumers.
// Without fold-boundary reset the anchor stays pinned to the
// fold-0 epoch-1 IQN loss for the entire run, so the gauge in
// fold N+1 is computed against a stale baseline. All three names
// dispatch to the same `reset_iqn_readiness_state` helper.
RegistryEntry {
name: "isv_iqn_loss_initial",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.iqn_loss_initial — fold-anchor for IQN improvement gauge; producer is `update_iqn_readiness` (gpu_dqn_trainer.rs:5804); never reset before Fix 3 of 2026-04-28 fold-boundary audit",
},
RegistryEntry {
name: "isv_iqn_loss_ema",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.iqn_loss_ema — streaming smoothed IQN loss; producer is `update_iqn_readiness`; coupled to isv_iqn_loss_initial in the readiness fraction",
},
RegistryEntry {
name: "isv_iqn_readiness",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.iqn_readiness scalar + pinned device-mapped slot (read by c51_loss_kernel as CVaR α and by IQN gradient-weight gating); reset to 0.0 at fold boundary so the new fold's first `update_iqn_readiness` recaptures the bootstrap anchor",
},
// ───── Window-reset state ───────────────────────────────────
RegistryEntry {
name: "portfolio_state",

14
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -4581,6 +4581,20 @@ impl DQNTrainer {
fused.reset_eval_v_range_state();
}
}
"isv_iqn_loss_initial" | "isv_iqn_loss_ema" | "isv_iqn_readiness" => {
// Three coupled scalars on `GpuDqnTrainer` (iqn_loss_initial,
// iqn_loss_ema, iqn_readiness + pinned device-mapped slot)
// form a single streaming-improvement gauge. Listed
// separately in the registry so future tasks can migrate
// them independently; all three currently dispatch to the
// same `reset_iqn_readiness_state` helper. Closes Fix 3 of
// the 2026-04-28 fold-boundary audit — without this reset
// the readiness fraction is computed against a fold-0
// anchor for every subsequent fold.
if let Some(ref mut fused) = self.fused_ctx {
fused.reset_iqn_readiness_state();
}
}
"isv_learning_health" => {
if let Some(ref fused) = self.fused_ctx {
fused.trainer().write_isv_signal_at(

56
docs/dqn-wire-up-audit.md
View File

@@ -2,6 +2,62 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
Fold-boundary reset gap — IQN readiness gauge (2026-04-28, Fix 3 of
3): the IQN readiness gauge on `GpuDqnTrainer` is a streaming
improvement fraction `iqn_readiness = (iqn_loss_initial -
iqn_loss_ema) / iqn_loss_initial` (gpu_dqn_trainer.rs:5805-5818).
The `iqn_loss_initial` anchor is captured once on the first
`update_iqn_readiness` call where it is still ~0 and never resets,
so for every fold beyond fold 0 the gauge is computed against the
fold-0 epoch-1 IQN loss — a stale anchor under the fold-N+1
post-S&P + fresh online↔target alignment + adversarial regime
shift. The pinned device-mapped readiness slot is read by
`c51_loss_kernel` as the CVaR α and by IQN gradient-weight gating,
so a stale anchor either pins readiness near 1 (over-confident,
full IQN gradient weight on a still-recovering head) or near 0
(under-weighting a converged head). Both modes contribute to the
fold-1 Q-drift trajectory documented in Fix 1's audit entry.
**Fix**: add `pub fn reset_iqn_readiness_state` to `GpuDqnTrainer`
(zeroes `iqn_loss_initial`, `iqn_loss_ema`, `iqn_readiness` scalar,
and writes 0.0 through the device-mapped pinned host slot — no HtoD
copy issued, mapping propagates the host write directly to GPU).
Add `pub(crate) fn reset_iqn_readiness_state` wrapper on
`FusedTrainingCtx` mirroring the existing `reset_eval_v_range_state`
pattern. Add three new `FoldReset` registry entries
(`isv_iqn_loss_initial`, `isv_iqn_loss_ema`, `isv_iqn_readiness`)
that all dispatch to the same helper through one match arm in
`reset_named_state` (training_loop.rs) — the three names are
listed separately so future tasks can migrate them independently
but currently form a single logical reset. Per
`feedback_no_partial_refactor.md` — same fold-boundary contract,
all three coupled scalars migrate together. The bootstrap branch
of `update_iqn_readiness` (`iqn_loss_initial < 1e-12`) recaptures
the new fold's first batch as the anchor on the first call after
this reset, by design.
Touched: `gpu_dqn_trainer.rs` (new `reset_iqn_readiness_state` after
the existing `reset_eval_v_range_state`), `fused_training.rs` (new
wrapper after the existing `reset_eval_v_range_state` wrapper),
`state_reset_registry.rs` (three new FoldReset entries), `trainer/
training_loop.rs` (one new dispatch arm covering all three names).
cargo check -p ml --lib clean at 13 warnings (workspace baseline).
cargo test -p ml --lib --no-run clean. The three existing
`state_reset_registry::tests` (`registry_classifies_known_state`,
`registry_fold_reset_iterates_only_fold_reset_entries`,
`registry_unknown_name_returns_none`) all pass without modification.
This commit completes the 3-fix sequence (Fix 1 IQN target sync,
Fix 2 aux Adam states, Fix 3 IQN readiness) closing the
fold-boundary state-migration contract gap surfaced by the audit.
Bug signature resolved (geometric Q-drift through fold 1):
`F0 ep5 +0.79 → F1 ep1 +0.82 → ep2 +2.23 → ep3 +4.05 → ep4 +10.66
→ ep5 NaN @ step 5 (fp32 overflow past atom support)`. Fix 1 is
load-bearing for the geometric drift itself; Fix 2 prevents
fold-N+1 first-epoch Adam-step overshoot in seven aux optimizers;
Fix 3 closes the gauge-staleness path that gated IQN gradient
weight against a fold-0 anchor. L40S validation deferred to
post-merge.
Fold-boundary reset gap — auxiliary Adam state resets (2026-04-28,
Fix 2 of 3): the same fold-boundary contract that resets the main
DQN's `m_buf` / `v_buf` / `adam_step` (gpu_dqn_trainer.rs:3428) was