From e07c82976e591a73d126d169f9002f443da62e56 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 20 Apr 2026 20:02:43 +0200 Subject: [PATCH] =?UTF-8?q?feat(B3/G4):=20health-scaled=20Expected=20SARSA?= =?UTF-8?q?=20temperature=20=E2=80=94=20breaks=20collapse=20attractor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When health=1 (healthy): tau unchanged → sharp softmax → near-argmax target (deterministic). When health=0 (collapsed): tau scales 6× → wide softmax → stochastic sampling breaks Q-collapse attractor. - c51_loss_kernel.cu: add get_learning_health() helper, isv_signals as last param, tau_base → tau * factor - gpu_dqn_trainer.rs: add last_sarsa_tau_factor field + init, launch_c51_loss → &mut self, host-side mirror, append isv_signals_dev_ptr arg - fused_training.rs: add last_sarsa_tau_factor() accessor - training_loop.rs: propagate SARSA tau factor from fused ctx for HEALTH_DIAG logging Co-Authored-By: Claude Sonnet 4.6 --- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 20 +++++++++++++++++-- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 17 +++++++++++++++- crates/ml/src/trainers/dqn/fused_training.rs | 5 +++++ .../src/trainers/dqn/trainer/training_loop.rs | 5 +++++ 4 files changed, 44 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index c486a2d77..1c6b16804 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -161,6 +161,12 @@ __device__ void block_bellman_project_f( } } +/* B3/G4: Read learning_health from ISV buffer at index 12. + * Returns 0.5 if pointer is null (safe fallback). */ +__device__ __forceinline__ float get_learning_health(const float* __restrict__ isv_signals_ptr) { + return (isv_signals_ptr != NULL) ? isv_signals_ptr[12] : 0.5f; +} + /* ══════════════════════════════════════════════════════════════════════ * MAIN KERNEL: c51_loss_batched (float arithmetic, BF16 I/O) * ══════════════════════════════════════════════════════════════════════ */ @@ -232,7 +238,10 @@ extern "C" __global__ void c51_loss_batched( const float* __restrict__ atom_positions, /* [4, num_atoms] adaptive positions. NULL = linear. */ /* ── Per-sample CVaR alpha from learned risk branch ── */ - const float* __restrict__ cvar_alpha_buf /* [B] per-sample CVaR alpha from risk branch. NULL = use iqn_readiness. */ + const float* __restrict__ cvar_alpha_buf, /* [B] per-sample CVaR alpha from risk branch. NULL = use iqn_readiness. */ + + /* ── B3/G4: ISV signals for health-scaled Expected SARSA temperature ── */ + const float* __restrict__ isv_signals /* [ISV_DIM=13] pinned device-mapped. isv_signals[12] = learning_health. NULL = 0.5 fallback. */ ) { extern __shared__ float shmem_f[]; @@ -543,7 +552,14 @@ extern "C" __global__ void c51_loss_batched( * meaningful at any Q-value scale. 1% of |mean Q| or 1e-6. */ float mean_q = (max_eq + min_eq) * 0.5f; float tau_floor = fmaxf(fabsf(mean_q) * 0.01f, 1e-6f); - float tau = fmaxf(q_gap_local, tau_floor); + float tau_base = fmaxf(q_gap_local, tau_floor); + + // B3/G4: health-coupled temperature scaling. + // health=1 (healthy): factor=1.0 → sharp softmax (near argmax target, deterministic). + // health=0 (collapsed): factor=6.0 → wide softmax (stochastic sampling breaks attractor). + float health = get_learning_health(isv_signals); + float tau_health_factor = 1.0f + 5.0f * (1.0f - health); + float tau = tau_base * tau_health_factor; float sum_exp = 0.0f; for (int a = 0; a < n_d; a++) { action_weights[a] = expf((eq_per_action[a] - max_eq) / tau); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1b15d7c2a..4bd03122c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1667,6 +1667,10 @@ pub struct GpuDqnTrainer { /// Last effective tau value (cosine-annealed, potentially floored by 0.01×(1-health)). /// Populated by `apply_health_coupled_tau_floor`. Initialized to 0.0. pub(crate) last_tau_eff: f32, + + /// 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, } impl GpuDqnTrainer { @@ -7447,6 +7451,7 @@ impl GpuDqnTrainer { tau_anneal_steps, last_cql_alpha_eff: 0.0, last_tau_eff: 0.0, + last_sarsa_tau_factor: 0.0, }) } @@ -10703,7 +10708,7 @@ impl GpuDqnTrainer { /// Reads: on_v_logits, on_b_logits, tg_v_logits, tg_b_logits, /// on_next_v_logits, on_next_b_logits (Double DQN). /// Writes: per_sample_loss, td_errors, total_loss, save_current_lp, save_projected. - fn launch_c51_loss(&self) -> Result<(), MLError> { + fn launch_c51_loss(&mut self) -> Result<(), MLError> { let kernel = &self.c51_loss_kernel; let b = self.config.batch_size; @@ -10769,6 +10774,14 @@ impl GpuDqnTrainer { let shmem_floats = na + na + max_branch * na + na + na + na + 8; let shmem_bytes = (shmem_floats * std::mem::size_of::()) as u32; + // B3/G4: mirror the kernel's tau_health_factor on host for logging. + let current_health = if self.isv_signals_pinned.is_null() { + 0.5_f32 + } else { + unsafe { (*self.isv_signals_pinned.add(LEARNING_HEALTH_INDEX)).clamp(0.0, 1.0) } + }; + self.last_sarsa_tau_factor = 1.0 + 5.0 * (1.0 - current_health); + unsafe { self.stream .launch_builder(kernel) @@ -10832,6 +10845,8 @@ impl GpuDqnTrainer { .arg(&atom_positions_buf_ptr) // ── Per-sample CVaR alpha from learned risk branch ── .arg(&cvar_alpha_buf_ptr) + // ── B3/G4: ISV signals for health-scaled Expected SARSA temperature ── + .arg(&self.isv_signals_dev_ptr) .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (256, 1, 1), diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 94c82e158..0ad25a206 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2198,6 +2198,11 @@ impl FusedTrainingCtx { self.trainer.last_tau_eff } + /// B3/G4: Last health-scaled Expected SARSA tau factor (1.0 at full health, 6.0 at collapse). + pub(crate) fn last_sarsa_tau_factor(&self) -> f32 { + self.trainer.last_sarsa_tau_factor + } + /// 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 3326f3da2..6b8423449 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1868,6 +1868,11 @@ impl DQNTrainer { self.last_tau_eff = Some(fused.last_tau_eff()); } + // B3/G4: propagate SARSA tau factor for logging. + if let Some(ref fused) = self.fused_ctx { + self.last_sarsa_tau_factor = Some(fused.last_sarsa_tau_factor()); + } + // 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}]",