diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d63eeea6b..cf4324d33 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2060,6 +2060,38 @@ impl GpuDqnTrainer { &self.commit_lambda_buf } + /// Gentle weight decay for risk branch parameters (MVP training signal). + /// + /// Scales risk weights by `(1 - lr * 0.01)` each step, preventing R from + /// collapsing to 0 or 1. Full BPTT backward deferred — the risk branch + /// learns its initial representation from trunk gradients flowing through + /// shared weights. Uses scale_f32_kernel (no memcpy, all device-side). + pub(crate) fn step_risk_sgd(&mut self) -> Result<(), MLError> { + let param_sizes = compute_param_sizes(&self.config); + // Risk tensors are indices 64..67 (last 4 in flat param buffer). + // Total aligned element count across all 4 risk tensors: + let risk_aligned_count: usize = (64..NUM_WEIGHT_TENSORS) + .map(|i| align4(param_sizes[i])) + .sum(); + + let base_lr: f32 = 1e-4; + let lr = base_lr * self.new_component_lr_scale(); + let decay = 1.0_f32 - lr * 0.01; // very gentle decay + + let risk_ptr = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 64); + + unsafe { + self.stream.launch_builder(&self.scale_f32_kernel) + .arg(&risk_ptr) + .arg(&decay) + .arg(&(risk_aligned_count as i32)) + .launch(LaunchConfig::for_num_elems(risk_aligned_count as u32)) + .map_err(|e| MLError::ModelError(format!("risk_sgd decay: {e}")))?; + } + self.risk_adam_step += 1; + Ok(()) + } + /// Update regime gate atom utilization from latest Q-stats. pub fn set_regime_util(&mut self, util: f32) { unsafe { *self.regime_util_pinned = util; } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 281ddf014..4bc1f5fc0 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1027,6 +1027,11 @@ impl FusedTrainingCtx { // then apply standalone SGD update (separate from main Adam). self.mamba2_backward(self.batch_size)?; self.step_mamba2_adam()?; + + // ── Step 5a-risk: Risk branch gentle weight decay ─────────────── + // MVP training signal: decay toward initial values to prevent R + // collapse. Full BPTT backward deferred. + self.step_risk_sgd()?; self.increment_warmup(); // Multi-horizon value forward (outside graph, degenerate MVP) @@ -2073,6 +2078,12 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("step_mamba2_adam: {e}")) } + /// Gentle weight decay for risk branch parameters (MVP training signal). + pub(crate) fn step_risk_sgd(&mut self) -> Result<()> { + self.trainer.step_risk_sgd() + .map_err(|e| anyhow::anyhow!("step_risk_sgd: {e}")) + } + /// Increment new component LR warmup counter. pub(crate) fn increment_warmup(&mut self) { self.trainer.increment_warmup();