diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index dbd3bdfd2..ef9547ad0 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -858,14 +858,11 @@ impl FusedTrainingCtx { } } - // Reset the adaptive C51 v_range EMAs and the pinned range buffer. - // Without this, the new fold inherits the prior fold's tight atom - // support; if the new fold's Q-distribution falls outside, TD errors - // saturate the atom bins and gradients blow up (observed: 50-epoch - // smoke train-92xbj Fold 1 Epoch 12 -> grad_norm=inf -> NaN loss). - // After this reset, the first update_eval_v_range call in Fold N+1 - // reinitialises from the newly observed Q-statistics. - self.trainer.reset_eval_v_range_state(); + // FoldReset state (eval_q_mean_ema, eval_q_std_ema, isv_v_range_slots, + // isv_learning_health, isv_sharpe_ema, isv_q_means) is now handled by + // StateResetRegistry-driven dispatch in DQNTrainer::reset_for_fold + // (trainer/mod.rs). reset_eval_v_range_state is called via the + // FusedTrainingCtx::reset_eval_v_range_state wrapper from there. // Clear replay buffer — stale experiences from previous fold have wrong // reward distributions (different data window, different policy). // NOTE: the replay buffer clear is handled by DQNTrainer::reset_for_fold() @@ -882,6 +879,17 @@ impl FusedTrainingCtx { Ok(()) } + /// Reset the adaptive C51 v_range EMAs and the pinned range buffer. + /// + /// Called from DQNTrainer::reset_named_state for the "eval_q_mean_ema", + /// "eval_q_std_ema", and "isv_v_range_slots" FoldReset entries. + /// These three states form a single logical unit — the EMA state that + /// drives the ISV v-range slots. Separating them in the registry allows + /// future tasks to migrate them independently; for now one call resets all. + pub(crate) fn reset_eval_v_range_state(&mut self) { + self.trainer.reset_eval_v_range_state(); + } + /// Steps since last flat buffer sync. pub(crate) fn steps_since_varmap_sync(&self) -> usize { self.steps_since_varmap_sync diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index 918891bfc..9b4107411 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -47,11 +47,6 @@ impl StateResetRegistry { category: ResetCategory::FoldReset, description: "Per-window Kelly win/loss counters (ps[14..18])", }, - RegistryEntry { - name: "plan_state", - category: ResetCategory::FoldReset, - description: "Val plan_state_buf [N, 7] — plan entry snapshots", - }, RegistryEntry { name: "eval_q_mean_ema", category: ResetCategory::FoldReset, diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 5fc5619c5..4ffa766d4 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -36,6 +36,7 @@ use crate::labeling::triple_barrier::TripleBarrierEngine; use super::config::{DQNAgentType, DQNHyperparameters}; use super::statistics::{FeatureStatistics, QValueStats}; +use super::StateResetRegistry; mod action; mod metrics; @@ -1472,7 +1473,23 @@ impl DQNTrainer { self.best_epoch = 0; self.best_val_loss = f64::INFINITY; - // Reset fused training context state (keep graphs, reset step counters) + // Plan 1 Task 3: registry-driven fold-boundary reset. Replaces scattered + // ad-hoc reset calls. Adding a new FoldReset state requires a registry + // entry AND a dispatch arm in reset_named_state (Invariant 2 Wire-It-Up). + // SoftReset category is handled separately (Plan 2); the pre-existing + // ad-hoc handling for adaptive_gamma and grad_balance targets is + // preserved in the training loop and is not replaced here. + let registry = StateResetRegistry::new(); + for entry in registry.fold_reset_entries() { + self.reset_named_state(entry.name, 0u32) + .map_err(|e| anyhow::anyhow!("fold-reset '{}': {}", entry.name, e))?; + } + + // Reset fused training context state (keep graphs, reset step counters). + // This handles Adam state, IQN Adam, shrink-and-perturb, CUDA graph + // invalidation, and PopArt — none of which are FoldReset registry entries + // (Adam/network are TrainingPersist; Adam is reset explicitly here as a + // special-case stabilisation, not a lifecycle category). if let Some(ref mut fused) = self.fused_ctx { fused.reset_for_fold()?; } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 8e60d5f91..5a43aef7e 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -3980,5 +3980,67 @@ impl DQNTrainer { Ok(false) } + + /// Dispatch a named-state reset to its owning subsystem. + /// + /// Invariant: every entry in the StateResetRegistry's FoldReset category + /// has a match arm here. Unknown names are a bug (the registry entry + /// should have been added in the same commit as the state it describes). + pub(crate) fn reset_named_state(&mut self, name: &str, fold_idx: u32) -> Result<(), crate::MLError> { + match name { + "kelly_stats" => { + // Kelly stats are held in portfolio state ps[14..18] — cleared + // per-window by the env_step kernel's fold_start detection. + // No-op here; the per-window reset logic handles it. + } + "eval_q_mean_ema" | "eval_q_std_ema" | "isv_v_range_slots" => { + // All three names map to the same reset — the EMA slots and ISV + // v-range slots are a single logical unit. The registry lists them + // separately so future tasks can migrate them independently; + // for now one call covers all three. + if let Some(ref mut fused) = self.fused_ctx { + fused.reset_eval_v_range_state(); + } + } + "isv_learning_health" => { + if let Some(ref fused) = self.fused_ctx { + fused.trainer().write_isv_signal_at( + crate::cuda_pipeline::gpu_dqn_trainer::LEARNING_HEALTH_INDEX, + 0.5, + ); + } + } + "isv_sharpe_ema" => { + if let Some(ref fused) = self.fused_ctx { + fused.trainer().write_isv_signal_at( + crate::cuda_pipeline::gpu_dqn_trainer::SHARPE_EMA_INDEX, + 0.0, + ); + } + } + "isv_q_means" => { + if let Some(ref fused) = self.fused_ctx { + // Zero the 9 Q-mean EMA slots [13..22). + for idx in 13..22 { + fused.trainer().write_isv_signal_at(idx, 0.0); + } + } + } + _ => { + return Err(crate::MLError::ModelError(format!( + "StateResetRegistry fold-reset dispatch: unknown name '{}'. \ + Every FoldReset entry must have a dispatch arm in reset_named_state.", + name + ))); + } + } + tracing::debug!( + target: "state_reset_registry", + fold = fold_idx, + name = name, + "fold-reset applied" + ); + Ok(()) + } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 4a2c89933..90bec883b 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -14,3 +14,4 @@ |---|---|---|---| | (populated during A.5 audit) | | | | | `crates/ml/src/trainers/dqn/state_reset_registry.rs` | consumers in Plan 1 Task 3+ (reset_fold_state call) | Wired (consumer added in same plan) | A.1 primary implementation | +| `training_loop.rs::reset_named_state` | consumer of StateResetRegistry::fold_reset_entries | Wired | A.1 Task 3 dispatch |