fix(dqn): C51 backward matches forward advantage-standardization for d==1

Root-cause bug from C-audit of train-5wb4n magnitude-branch Q-saturation
(task #92). Forward kernel `c51_loss_batched` (c51_loss_kernel.cu
lines 736-752) applies per-atom advantage standardization for the
magnitude branch:

    centered[j] = (adv[a_d,j] - a_mean[j]) / (a_std[j] + 1e-6)   // d==1 only
    combined[j] = value[j] + centered[j]
    lp[j]       = log_softmax(combined)

But the backward kernel `c51_grad_kernel` treated the forward as if
`combined[j] = value[j] + (adv[a_d,j] - a_mean[j])` — no a_std divisor.
That is chain-rule-inconsistent: the computed `d_combined / d_adv[a,j]`
is the gradient of a *different* loss than the one forward-pass'd.
Gradient points in the wrong direction relative to the actual loss
surface.

Fix: pass the magnitude-branch online advantage logits pointer
(`on_adv_logits_b1`) to the backward kernel. For each (b, j) in the
d==1 dueling-grad block, recompute a_std from the 3 advantage values
and multiply `grad_val *= 1/(a_std + 1e-6)`. Counterfactual magnitude
gradient (Hold samples) inherits the same factor automatically —
single correction point.

Approximation: treats a_std as constant w.r.t. advantage values (skips
`da_std/d_adv` chain-rule terms). This is the standard simplification
used in batch-norm with `track_running_stats=False`; produces the
dominant scale correction without the full Jacobian complexity.

Call site updates the single launch in `launch_c51_grad`. Kernel
signature guards on NULL for backward compatibility with any smoke-
test launcher that may not wire the argument; d==1 block falls back
to pre-fix behaviour (identity) in that case.

Part of task #92 fix triad. Previous commit d61aefe2b addressed the
balancer (safety-net fix). Next: IQL branch_scales floor for
direction-branch Hold/Flat per-sample starvation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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) {

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::<f32>();
+        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),