# Advantage Standardization + Distribution Entropy — Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Prevent magnitude branch Q-value collapse by normalizing magnitude advantages to unit variance (self-adapting, no temperature tuning) and boosting distribution entropy for magnitude atoms (5× base entropy coefficient).

**Architecture:** 2 changes in 4 CUDA kernel files. No new parameters, no Rust changes, no buffer size changes. CUDA graph capture unaffected.

**Tech Stack:** CUDA (cubin via build.rs)

**CRITICAL CONSTRAINT:** All CUDA kernels compile via build.rs (NOT NVRTC). No CPU path. No new kernel parameters.

---

### Task 1: Advantage Standardization in C51 Loss Kernel

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu`

**Design:** For d==1 (magnitude branch), replace the standard dueling computation `Q[j] = V[j] + A[a_d,j] - mean(A)` with standardized advantages `Q[j] = V[j] + (A[a_d,j] - mean(A)) / (std(A) + 1e-6)`. The std is computed per-atom across the 3 magnitude actions. This also applies to the Gem 3+ counterfactual block.

- [ ] **Step 1: Add advantage standardization to the main dueling block**

In `c51_loss_kernel.cu`, find the dueling computation at lines ~338-347:

```c
        /* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j])
         * All 4 branches now have small action counts — mean subtraction is safe. */
        for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * num_atoms + j];
            a_mean /= (float)n_d;
            shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean;
        }
        __syncthreads();
```

Replace with:

```c
        /* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j])
         * For magnitude branch (d==1): standardize advantages to unit variance.
         * This prevents winner-take-all Q-value collapse where Small dominates
         * due to lower PnL variance causing tighter C51 distributions. */
        for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * num_atoms + j];
            a_mean /= (float)n_d;
            float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
            if (d == 1) {
                /* Advantage standardization: divide by std to bound scale.
                 * std = sqrt(mean(x^2) - mean(x)^2) where x = A - mean(A).
                 * With n_d=3, this is cheap (3 values per atom). */
                float sq_sum = 0.0f;
                for (int a = 0; a < n_d; a++) {
                    float diff = shmem_adv[a * num_atoms + j] - a_mean;
                    sq_sum += diff * diff;
                }
                float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
                centered /= (a_std + 1e-6f);
            }
            shmem_lp[j] = shmem_val[j] + centered;
        }
        __syncthreads();
```

- [ ] **Step 2: Apply same standardization to the Gem 3+ counterfactual block**

Find the counterfactual block at lines ~354-368:

```c
        if (d == 1 && branch_action[0] == 1) {
            float cf_weight = 0.3f;
            for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
                float a_mean = 0.0f;
                for (int a = 0; a < n_d; a++)
                    a_mean += shmem_adv[a * num_atoms + j];
                a_mean /= (float)n_d;
                shmem_lp[j] = shmem_val[j]
                    + cf_weight * (shmem_adv[a_d * num_atoms + j] - a_mean);
            }
        }
        __syncthreads();
```

Replace with:

```c
        if (d == 1 && branch_action[0] == 1) {
            /* Gem 3+ counterfactual with advantage standardization. */
            float cf_weight = 0.3f;
            for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
                float a_mean = 0.0f;
                for (int a = 0; a < n_d; a++)
                    a_mean += shmem_adv[a * num_atoms + j];
                a_mean /= (float)n_d;
                float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
                float sq_sum = 0.0f;
                for (int a = 0; a < n_d; a++) {
                    float diff = shmem_adv[a * num_atoms + j] - a_mean;
                    sq_sum += diff * diff;
                }
                float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
                centered /= (a_std + 1e-6f);
                shmem_lp[j] = shmem_val[j] + cf_weight * centered;
            }
        }
        __syncthreads();
```

- [ ] **Step 3: Verify build**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
git commit -m "feat(magnitude): advantage standardization in C51 loss — prevents Q-value scale collapse"
```

---

### Task 2: Advantage Standardization in MSE Loss Kernel

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu`

Apply the IDENTICAL changes from Task 1 to the MSE loss kernel. The code structure is identical.

- [ ] **Step 1: Add advantage standardization to the main dueling block**

In `mse_loss_kernel.cu`, find the dueling computation at lines ~286-295. Apply the identical replacement from Task 1 Step 1.

- [ ] **Step 2: Apply same standardization to the Gem 3+ counterfactual block**

Find the counterfactual block at lines ~302-316. Apply the identical replacement from Task 1 Step 2.

- [ ] **Step 3: Verify build**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
git commit -m "feat(magnitude): advantage standardization in MSE loss — mirrors C51 kernel"
```

---

### Task 3: Distribution Entropy Boost in Grad Kernels

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu`

**Design:** Multiply the entropy regularization gradient by 5× for d==1 (magnitude). This prevents the atom distribution from becoming too peaked for any magnitude action. The MSE grad kernel doesn't have entropy, so we add it.

- [ ] **Step 1: Boost entropy in C51 grad kernel**

In `c51_grad_kernel.cu`, find the entropy block at lines ~44-48:

```c
    /* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp) */
    if (entropy_coeff > 0.0f) {
        float lp_clamped = fmaxf(lp, -10.0f);
        d_combined += entropy_coeff * (1.0f + lp_clamped);
    }
```

Replace with:

```c
    /* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp)
     * Magnitude branch (d==1) gets 5× entropy boost to prevent distribution
     * collapse. The standardized advantages (from loss kernel) prevent Q-value
     * scale collapse; this prevents atom sharpness collapse. */
    if (entropy_coeff > 0.0f) {
        float ent_scale = (d == 1) ? 5.0f : 1.0f;
        float lp_clamped = fmaxf(lp, -10.0f);
        d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
    }
```

- [ ] **Step 2: Add entropy to MSE grad kernel**

In `mse_grad_kernel.cu`, there is NO entropy term. Add one after line ~46 (`float d_combined = isw * td_error * p_j * (z_j - e_q);`):

Find:

```c
    /* Gradient of MSE loss through softmax expectation (float arithmetic) */
    float d_combined = isw * td_error * p_j * (z_j - e_q);
```

Replace with:

```c
    /* Gradient of MSE loss through softmax expectation (float arithmetic) */
    float d_combined = isw * td_error * p_j * (z_j - e_q);

    /* Magnitude entropy boost: prevent atom distribution collapse for d==1.
     * MSE path doesn't have entropy by default — add for magnitude only.
     * p_j is the softmax probability; lp = log(p_j) for entropy gradient. */
    if (d == 1 && entropy_coeff > 0.0f) {
        float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
        d_combined += 5.0f * entropy_coeff * (1.0f + lp_approx);
    }
```

Wait — the MSE grad kernel doesn't receive `entropy_coeff` as a parameter. Check the kernel signature. If it doesn't have it, we need to add it.

Read the MSE grad kernel signature first. If `entropy_coeff` is not a parameter, do NOT add it (would change Rust launch code). Instead, use a hardcoded constant: `0.005f` (= 5 × 0.001 default). This avoids changing the kernel ABI.

```c
    /* Magnitude entropy boost: prevent atom distribution collapse for d==1.
     * Hardcoded 0.005 = 5× base entropy_coeff (0.001). No kernel param needed. */
    if (d == 1) {
        float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
        d_combined += 0.005f * (1.0f + lp_approx);
    }
```

- [ ] **Step 3: Verify build**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu \
       crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
git commit -m "feat(magnitude): 5× entropy boost in grad kernels — prevents atom distribution collapse"
```

---

### Task 4: Unit Tests + Verification

**Files:**
- Modify: `crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs`

- [ ] **Step 1: Add advantage standardization test**

Add to `reward_v8.rs`:

```rust
    #[test]
    fn test_advantage_standardization() {
        // Verify standardization bounds advantages to O(1)
        // Raw advantages: Small=10.0, Medium=1.0, Full=0.5 (collapse scenario)
        let advantages = [10.0_f64, 1.0, 0.5];
        let mean = advantages.iter().sum::<f64>() / 3.0;
        let centered: Vec<f64> = advantages.iter().map(|a| a - mean).collect();
        let variance = centered.iter().map(|c| c * c).sum::<f64>() / 3.0;
        let std = variance.sqrt();
        let standardized: Vec<f64> = centered.iter().map(|c| c / (std + 1e-6)).collect();
        // After standardization, values are O(1) — no 20× gap
        assert!(standardized[0] < 2.0, "Small should be <2.0 after std, got {}", standardized[0]);
        assert!(standardized[2] > -2.0, "Full should be >-2.0 after std, got {}", standardized[2]);
        // Relative ordering preserved
        assert!(standardized[0] > standardized[1], "Small > Medium ordering preserved");
        assert!(standardized[1] > standardized[2], "Medium > Full ordering preserved");
        // Gap is bounded
        let gap = standardized[0] - standardized[2];
        assert!(gap < 3.0, "Max gap should be <3.0, got {gap}");
    }
```

- [ ] **Step 2: Run tests**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::reward_v8 --nocapture 2>&1 | tail -20
```

- [ ] **Step 3: Full test suite**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | tail -3
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs
git commit -m "test(magnitude): advantage standardization unit test"
```

---

### Task 5: Push + Deploy + Monitor

- [ ] **Step 1: Push**

```bash
git push origin main
```

- [ ] **Step 2: Launch baseline 10 epochs**

```bash
./scripts/argo-train.sh dqn --baseline --epochs 10
```

- [ ] **Step 3: Monitor diversity**

Target: dir×mag diversity should hold at ≥5/9 through epoch 10 (not collapse to 3/9 at epoch 4).

```bash
argo logs <workflow> -n foxhunt --follow=false -c main 2>&1 | grep 'Action div'
```
