# Per-Branch C51/MSE Loss Weighting — 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:** Make the magnitude branch MSE-dominant (4× MSE, 0.2× C51) to prevent C51's variance-biased cross-entropy from collapsing magnitude to Small.

**Architecture:** 2 CUDA kernel edits. Replace the existing `mag_amp = 1.5` with per-branch loss weighting: C51 grad scales magnitude by 0.2×, MSE grad scales magnitude by 4.0×.

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

---

### Task 1: Per-Branch Loss Weighting in Grad Kernels

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

- [ ] **Step 1: Update C51 grad kernel**

In `c51_grad_kernel.cu`, find the existing magnitude amplification block (around line 81):

```c
    /* Magnitude gradient amplification: scale d==1 (magnitude) by 1.5× to
     * compensate for Flat counterfactual discount (0.3× on ~58% of samples).
     * Effective gradient per magnitude sample ≈ 0.3*0.58 + 1.0*0.42 = 0.594.
     * Scale = 1.5 ≈ 1/0.594, conservative to avoid gradient explosion.
     * Gradient clipping after this kernel prevents any runaway. */
    float mag_amp = (d == 1) ? 1.5f : 1.0f;

    /* d_adv[b, d, a, j] = mag_amp * d_combined * (delta(a, a_d) - 1/A_d) */
    for (int a = 0; a < A_d; a++) {
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        float grad_val = mag_amp * d_combined * dueling_grad;
```

Replace with:

```c
    /* Per-branch loss weighting: magnitude (d==1) gets 0.2× C51 gradient.
     * C51 cross-entropy inherently prefers low-variance actions (Small positions
     * have tighter return distributions → lower cross-entropy). Reducing C51's
     * influence on magnitude to 20% lets MSE (variance-neutral) dominate.
     * Direction/order/urgency keep full C51 gradient for distributional risk awareness. */
    float branch_scale = (d == 1) ? 0.2f : 1.0f;

    /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */
    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;
```

- [ ] **Step 2: Update MSE grad kernel**

In `mse_grad_kernel.cu`, find the same block (around line 80):

```c
    float mag_amp = (d == 1) ? 1.5f : 1.0f;

    /* d_adv[b, d, a, j] = mag_amp * d_combined * (delta(a, a_d) - 1/A_d) */
    for (int a = 0; a < A_d; a++) {
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        float grad_val = mag_amp * d_combined * dueling_grad;
```

Replace with:

```c
    /* Per-branch loss weighting: magnitude (d==1) gets 4.0× MSE gradient.
     * MSE loss is variance-neutral (optimizes expected Q only, no distributional
     * shape bias). Amplifying MSE for magnitude makes it the dominant loss signal,
     * preventing C51's cross-entropy from collapsing magnitude to Small.
     * Combined: magnitude = 0.2×C51 + 4.0×MSE ≈ 80% MSE-driven. */
    float branch_scale = (d == 1) ? 4.0f : 1.0f;

    /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */
    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;
```

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

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

- [ ] **Step 4: Local smoke test**

```bash
FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::walk_forward --ignored --nocapture 2>&1 | grep -E 'Action div|Sharpe=|PF='
```

Expected: diversity ≥5/9 at epoch 10.

- [ ] **Step 5: Commit + push**

```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): per-branch loss weighting — 0.2× C51 + 4.0× MSE for magnitude (variance-neutral)"
git push origin main
```

---

### Task 2: H100 Validation

- [ ] **Step 1: Delete old fxcache + deploy**

```bash
# Delete cached fxcache on PVC (force rebuild with new binary)
kubectl run cleanup --image=ubuntu:24.04 --restart=Never -n foxhunt \
  --overrides='...' # (delete /feature-cache/*.fxcache)

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

- [ ] **Step 2: Monitor**

Watch for:
- Diversity ≥5/9 at epochs 5-10 (not 3/9)
- Sharpe > 0 sustained
- PF ≥ 1.0

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