diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index b39fc694b..7a47af43c 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -35,7 +35,19 @@ extern "C" __global__ void c51_grad_kernel( * (d==0). Slot 22 is SHARPE_EMA, consumed on-GPU by isv_signal_update to * drive slot 12 (learning_health). NULL → identity (safe: reduces to * pre-fix behaviour). */ - const float* __restrict__ isv_signals) + const float* __restrict__ isv_signals, + /* Backward-standardization fix (2026-04-23): the forward kernel + * `c51_loss_batched` applies per-atom advantage standardization + * `centered /= (a_std + 1e-6)` for the magnitude branch (d==1). + * The backward gradient must replicate the same normalisation factor + * or the gradient is not the partial derivative of the (scaled) loss. + * We approximate `a_std` as constant w.r.t. advantage values (same + * simplification used in batch-norm with `track_running_stats=False`), + * so the chain rule reduces to `d_adv[a,j] *= 1/(a_std[b,j] + 1e-6)`. + * `on_adv_logits_b1` layout matches the forward: f32 stride + * `[B, b1_size, num_atoms]` = `[B, b1_size * num_atoms]` per sample. + * NULL → standardization not applied (legacy behaviour). */ + const float* __restrict__ on_adv_logits_b1) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_elems = batch_size * num_atoms; @@ -235,9 +247,41 @@ extern "C" __global__ void c51_grad_kernel( int n_rest = A_d - 2; /* actions that are neither taken nor challenger */ float rest_weight = (n_rest > 0) ? (1.0f + 0.5f) / (float)n_rest : 0.0f; + /* Backward-standardization fix (2026-04-23): for d==1 (magnitude) + * the forward applies `centered /= (a_std + 1e-6f)` per (b, j). + * Recompute the same a_std here so the chain-rule factor lands + * on every d==1 grad write. `on_adv_logits_b1` layout is + * `[B, b1_size, num_atoms]`; we read the 3 advantage values at + * this (b, *, j) slice, compute mean + variance, then sqrt. + * Approximates a_std as constant w.r.t. advantage — the simpler + * chain rule used in practice (skips da_std/dadv terms). Guarded + * on `on_adv_logits_b1 != NULL` so the kernel is backward- + * compatible with smoke-test launchers that don't wire it. */ + float inv_a_std = 1.0f; + if (d == 1 && on_adv_logits_b1 != NULL) { + int A1 = A_d; + /* Only compute when b1_size looks sane (avoids OOB on + * degenerate configurations). */ + if (A1 > 0) { + float sum = 0.0f; + for (int a = 0; a < A1; a++) { + sum += on_adv_logits_b1[(long long)b * A1 * num_atoms + a * num_atoms + j]; + } + float mean = sum / (float)A1; + float sq_sum = 0.0f; + for (int a = 0; a < A1; a++) { + float v = on_adv_logits_b1[(long long)b * A1 * num_atoms + a * num_atoms + j] - mean; + sq_sum += v * v; + } + float a_std = sqrtf(sq_sum / (float)A1 + 1e-12f); + inv_a_std = 1.0f / (a_std + 1e-6f); + } + } + for (int a = 0; a < A_d; a++) { float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); float grad_val = branch_scale * d_combined * dueling_grad; + if (d == 1) grad_val *= inv_a_std; float spread_grad; if (a == a_d) { diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f0f37cf30..028aa657d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -13528,6 +13528,16 @@ impl GpuDqnTrainer { let liquid_mod_buf_ptr = self.liquid_mod_buf.raw_ptr(); let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); + // Backward-standardization fix (2026-04-23): thread the magnitude- + // branch online advantage logits pointer through to the backward + // kernel so it can recompute `a_std` at each (b, j) and apply the + // same `1 / (a_std + 1e-6)` factor the forward applied. The layout + // matches the forward: f32 stride `[B, b1_size, num_atoms]`. Offset + // base = on_b_logits_buf + B * b0_size * num_atoms * 4 (b0 slice). + let f32_sz = std::mem::size_of::(); + let on_b_base = self.on_b_logits_buf.raw_ptr(); + let on_adv_b1_ptr: u64 = on_b_base + (b * self.config.branch_0_size * na * f32_sz) as u64; + unsafe { self.stream .launch_builder(&self.c51_grad_kernel) @@ -13554,6 +13564,10 @@ impl GpuDqnTrainer { // slots [13] (q_mag_spread_ema) and [14] (q_dir_spread_ema) to // derive the magnitude-branch adaptive bin weight. .arg(&self.isv_signals_dev_ptr) + // Backward-standardization fix (2026-04-23): magnitude-branch + // advantage logits for per-(b, j) a_std recomputation in the + // d==1 dueling-grad block. Matches forward layout. + .arg(&on_adv_b1_ptr) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1),