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

fix(magnitude): restore MSE gradient killed by alpha blend

Browse Source 
The alpha blend (grad = alpha*C51 + (1-alpha)*MSE) progressively killed
magnitude gradient: at alpha=1.0 post-warmup, magnitude got 0*C51 + 0*MSE
= zero gradient. The magnitude branch head stopped learning at epoch 5.

H100 confirmed: grad_norm stabilized at 0.44 (epochs 6-9) — other
components alive but magnitude frozen. Q-values plateaued at 2.03.

Fix: after the alpha blend, SAXPY d_adv[d==1] += alpha * d_adv_mse[d==1]
to restore the MSE portion removed by the (1-alpha) scaling. This gives
magnitude full MSE gradient at all alpha values.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 11:10:56 +02:00
parent 1540c0287d
commit 22b7bc0834
1 changed files with 20 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4573,6 +4573,26 @@ impl GpuDqnTrainer {
.arg(&adv_ptr).arg(&adv_mse_ptr)
.arg(&scale_mse).arg(&n_adv)
.launch(cfg_adv).map_err(|e| MLError::ModelError(format!("blend mse adv: {e}")))?;
// ── Restore full MSE gradient for magnitude branch (d==1) ────
// The alpha blend above gave magnitude: (1-alpha)*MSE_grad
// (because C51_grad is zeroed for d==1). At alpha=1.0, this is ZERO.
// Magnitude needs full MSE gradient always — it's the ONLY training
// signal for the magnitude branch head (C51 zeroed, IQN trunk-only).
// Fix: SAXPY d_adv[d==1] += alpha * d_adv_mse[d==1] to restore
// the alpha*MSE portion that was removed by the (1-alpha) scale.
if alpha > 1e-6 {
let mag_offset = (b * b0 * na) as u64 * 4; // bytes: b0 branch comes first
let n_mag = (b * b1 * na) as i32;
let mag_blocks = ((b * b1 * na + 255) / 256) as u32;
let cfg_mag = LaunchConfig { grid_dim: (mag_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
let mag_adv_ptr = adv_ptr + mag_offset;
let mag_mse_ptr = adv_mse_ptr + mag_offset;
self.stream.launch_builder(&self.saxpy_f32_kernel)
.arg(&mag_adv_ptr).arg(&mag_mse_ptr)
.arg(&alpha).arg(&n_mag)
.launch(cfg_mag).map_err(|e| MLError::ModelError(format!("restore mag mse: {e}")))?;
}
}
}