diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4bd03122c..0dbeefe13 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1671,6 +1671,19 @@ pub struct GpuDqnTrainer { /// B3/G4: Last health-scaled Expected SARSA tau factor (1.0 at full health, 6.0 at collapse). /// Populated by `launch_c51_loss`. Initialized to 0.0. pub(crate) last_sarsa_tau_factor: f32, + + /// B4/G5: Last adaptive gradient budget for IQN (0.10 + 0.30×health). + /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. Default: health=1 → 0.40. + pub(crate) last_iqn_budget_eff: f32, + /// B4/G5: Last adaptive gradient budget for CQL (0.10×(1−regime)×health). + /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. Default: 0.00 (regime_stability=1). + pub(crate) last_cql_budget_eff: f32, + /// B4/G5: Last adaptive gradient budget for C51 (1−iqn−cql−ens). + /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. Default: health=1 → 0.55. + pub(crate) last_c51_budget_eff: f32, + /// B4/G5: Adaptive gradient budget for ensemble (constant 0.05). + /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. + pub(crate) last_ens_budget_eff: f32, } impl GpuDqnTrainer { @@ -4064,6 +4077,7 @@ impl GpuDqnTrainer { &mut self, iqn_d_h_s2_ptr: u64, _online_dueling: &mut DuelingWeightSet, + iqn_budget: f32, ) -> Result<(), MLError> { let b = self.config.batch_size; let _eg = EventTrackingGuard::new(self.stream.context()); @@ -4184,13 +4198,14 @@ impl GpuDqnTrainer { )?; } - // ── 7. Plain SAXPY: grad_buf[trunk] += iqn_lambda * scratch ── + // ── 7. Plain SAXPY: grad_buf[trunk] += iqn_lambda * iqn_budget * scratch ── // Adaptive lambda: scales with IQN loss readiness (0 when uncertain, 1 when converged). + // B4/G5: further scaled by iqn_budget (0.10+0.30×health). // Uses aux_child-specific handle — saxpy_f32_kernel is captured in forward_child. { let scratch_ptr = self.ptrs.iqn_trunk_m; let grad_ptr = self.ptrs.grad_buf; - let scale = self.config.iqn_lambda * self.iqn_readiness; + let scale = self.config.iqn_lambda * self.iqn_readiness * iqn_budget; let n_i32 = trunk_grad_total as i32; let blocks = ((trunk_grad_total + 255) / 256) as u32; unsafe { @@ -4771,14 +4786,15 @@ impl GpuDqnTrainer { /// /// Called after `apply_cql_gradient` populated `cql_grad_scratch`. /// No per-component clip — single global clip in Adam handles safety. - pub fn apply_cql_saxpy(&mut self) -> Result<(), MLError> { + pub fn apply_cql_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { let total = self.total_params as i32; let blocks = self.grad_norm_blocks as u32; let scratch_ptr = self.ptrs.cql_grad_scratch; let grad_ptr = self.ptrs.grad_buf; - let alpha = 1.0_f32; + // B4/G5: scale CQL gradient contribution by adaptive budget (0.10×(1−regime)×health). + let alpha = cql_budget; - // Plain SAXPY: grad_buf += 1.0 * cql_scratch + // Plain SAXPY: grad_buf += cql_budget * cql_scratch // Uses aux_child-specific handle — saxpy_f32_kernel is captured in forward_child. unsafe { self.stream @@ -7452,6 +7468,10 @@ impl GpuDqnTrainer { last_cql_alpha_eff: 0.0, last_tau_eff: 0.0, last_sarsa_tau_factor: 0.0, + last_iqn_budget_eff: 0.40, + last_cql_budget_eff: 0.00, + last_c51_budget_eff: 0.55, + last_ens_budget_eff: 0.05, }) } @@ -7633,6 +7653,18 @@ impl GpuDqnTrainer { } } + /// B4/G5: Read ISV health index and regime stability from pinned host memory. + /// Returns (health [0,1], regime_stability [0,1]). Falls back to (0.5, 0.5) if + /// pinned pointer is null. + pub(crate) fn read_isv_health_and_regime(&self) -> (f32, f32) { + if self.isv_signals_pinned.is_null() { return (0.5_f32, 0.5_f32); } + unsafe { + let h = (*self.isv_signals_pinned.add(LEARNING_HEALTH_INDEX)).clamp(0.0, 1.0); + let s = (*self.isv_signals_pinned.add(11)).clamp(0.0, 1.0); + (h, s) + } + } + /// B2/G3: Apply a health-coupled minimum to a scheduled tau value. /// Returns `max(tau_scheduled, 0.01 * (1.0 - health))`. During collapse /// (health≈0), the floor rises to 0.01 to accelerate target-network adaptation. diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 0ad25a206..e3bc4a8bd 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1157,11 +1157,13 @@ impl FusedTrainingCtx { ).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?; } - // NOTE: Per-component gradient budget clip REMOVED. - // The old code clipped C51+MSE gradient to a fraction of max_grad_norm, - // producing a CONSTANT gradient norm every step. Adam degenerates into - // SignSGD when it never sees true gradient magnitude variation. - // Safety is now handled by the single global clip inside Adam's kernel. + // NOTE: Per-component gradient budget clip REMOVED (old explicit clip → constant grad norm). + // B4/G5: Budgets are now applied as SCALE factors on individual SAXPY contributions, + // not as norm clips. This preserves gradient magnitude variation while steering + // component balance. Safety is still handled by the single global clip in Adam. + + // B4/G5: Compute adaptive budgets once per step from ISV health/regime signals. + let (_c51_budget, iqn_budget, cql_budget, _ens_budget) = self.compute_adaptive_budgets(); // DIAGNOSTIC: sync between each aux op to find hanging kernel @@ -1425,7 +1427,7 @@ impl FusedTrainingCtx { if let Some(ref mut iqn) = self.gpu_iqn { let d_h_s2_ptr = iqn.d_h_s2_raw_ptr(); self.trainer.apply_iqn_trunk_gradient( - d_h_s2_ptr, &mut self.online_dueling, + d_h_s2_ptr, &mut self.online_dueling, iqn_budget, ).map_err(|e| anyhow::anyhow!("IQN trunk gradient: {e}"))?; let dqn = agent.primary_dqn_mut(); @@ -1457,10 +1459,11 @@ impl FusedTrainingCtx { // CQL gradient + SAXPY (no per-component clip — Adam handles safety). + // B4/G5: cql_budget scales the SAXPY contribution (0.10×(1−regime)×health). if self.trainer.has_cql() { match self.trainer.apply_cql_gradient() { Ok(true) => { - self.trainer.apply_cql_saxpy() + self.trainer.apply_cql_saxpy(cql_budget) .map_err(|e| anyhow::anyhow!("CQL SAXPY: {e}"))?; } Ok(false) => {} @@ -2203,6 +2206,42 @@ impl FusedTrainingCtx { self.trainer.last_sarsa_tau_factor } + /// B4/G5: Compute adaptive gradient budgets from learning_health and regime_stability. + /// Returns (c51_budget, iqn_budget, cql_budget, ens_budget). Caches to trainer fields + /// for HEALTH_DIAG logging. + /// + /// Formula: + /// iqn_budget = 0.10 + 0.30 × health + /// cql_budget = 0.10 × (1 − regime_stability) × health + /// ens_budget = 0.05 (constant) + /// c51_budget = 1.0 − iqn_budget − cql_budget − ens_budget + /// + /// At health=1: iqn=0.40, cql variable (depends on regime_stability), ens=0.05, c51=rest. + /// At health=0: iqn=0.10, cql=0, ens=0.05, c51=0.85 (C51 takes over). + pub(crate) fn compute_adaptive_budgets(&mut self) -> (f32, f32, f32, f32) { + let (health, regime_stability) = self.trainer.read_isv_health_and_regime(); + let iqn_budget = 0.10_f32 + 0.30_f32 * health; + let cql_budget = 0.10_f32 * (1.0 - regime_stability) * health; + let ens_budget = 0.05_f32; + let c51_budget = (1.0_f32 - iqn_budget - cql_budget - ens_budget).max(0.0); + + self.trainer.last_iqn_budget_eff = iqn_budget; + self.trainer.last_cql_budget_eff = cql_budget; + self.trainer.last_c51_budget_eff = c51_budget; + self.trainer.last_ens_budget_eff = ens_budget; + + (c51_budget, iqn_budget, cql_budget, ens_budget) + } + + /// B4/G5: Last adaptive IQN gradient budget (0.10 + 0.30×health). + pub(crate) fn last_iqn_budget_eff(&self) -> f32 { self.trainer.last_iqn_budget_eff } + /// B4/G5: Last adaptive CQL gradient budget (0.10×(1−regime)×health). + pub(crate) fn last_cql_budget_eff(&self) -> f32 { self.trainer.last_cql_budget_eff } + /// B4/G5: Last adaptive C51 gradient budget (1−iqn−cql−ens). + pub(crate) fn last_c51_budget_eff(&self) -> f32 { self.trainer.last_c51_budget_eff } + /// B4/G5: Last ensemble gradient budget (constant 0.05). + pub(crate) fn last_ens_budget_eff(&self) -> f32 { self.trainer.last_ens_budget_eff } + /// Compute spectral gap on the last Q readback slice (coarse max/min ratio proxy). pub(crate) fn compute_q_spectral_gap(&self) -> f32 { self.trainer.compute_q_spectral_gap() diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 60ca387b6..50a453fff 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1873,6 +1873,13 @@ impl DQNTrainer { self.last_sarsa_tau_factor = Some(fused.last_sarsa_tau_factor()); } + // B4/G5: propagate adaptive gradient budgets for logging. + if let Some(ref fused) = self.fused_ctx { + self.last_iqn_budget_eff = Some(fused.last_iqn_budget_eff()); + self.last_cql_budget_eff = Some(fused.last_cql_budget_eff()); + self.last_c51_budget_eff = Some(fused.last_c51_budget_eff()); + } + // HEALTH_DIAG: components are [0, 1] normalized. effective = hyperparams after health-adaptation. novels = mechanism states. tracing::info!( "HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}]",