# Magnitude Branch Collapse Fix — 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:** Fix the 3/9 dir×mag diversity collapse where the magnitude branch converges to always-Small (0.25×max_position) due to insufficient exploration, gradient starvation from Flat zeroing, and uncompensated gradient magnitude.

**Architecture:** Three targeted changes in CUDA kernels + Rust config. (1) Give magnitude its own epsilon = direction epsilon. (2) On Flat samples, compute counterfactual magnitude gradient for the most likely non-Flat direction at 0.3× weight. (3) Scale magnitude gradient by `1/(1-flat_fraction)` in the grad kernel to compensate for the ~58% zero-gradient Flat samples. Additionally, scale the dense micro-reward by magnitude fraction so Full positions get 4× the signal of Small.

**Tech Stack:** CUDA (cubin via build.rs), Rust (cudarc), TOML config

**CRITICAL CONSTRAINT:** All CUDA kernels compile via build.rs (NOT NVRTC). Kernel arg order in Rust MUST match CUDA parameter order EXACTLY. No CPU path. No stubs.

---

### Task 1: Per-Branch Epsilon — Magnitude Gets Full Exploration

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

**Problem:** Line 722 computes `eps_other = eps_end + 0.1f * (epsilon - eps_end)` giving magnitude ~2.3% exploration. Direction gets the full cosine-scheduled epsilon (~5-20%). The model greedily picks Small because it minimizes loss variance, and never explores Medium/Full.

**Fix:** Give magnitude its own epsilon equal to direction's. Keep order/urgency at reduced epsilon (they genuinely are less important).

- [ ] **Step 1: Change epsilon computation**

In `experience_kernels.cu`, find the Gem 2 epsilon block at lines 719-722:

```c
    /* Gem 2: Per-branch epsilon — direction explores more (the hard problem).
     * Magnitude/order/urgency get reduced epsilon: 10% of epsilon range. */
    float eps_dir = epsilon;
    float eps_other = eps_end + 0.1f * (epsilon - eps_end);
```

Replace with:

```c
    /* Gem 2: Per-branch epsilon — direction AND magnitude explore fully.
     * Both are hard problems: direction chooses side, magnitude chooses size.
     * Order/urgency get reduced epsilon: 10% of epsilon range. */
    float eps_dir = epsilon;
    float eps_mag = epsilon;  /* magnitude needs full exploration — same as direction */
    float eps_other = eps_end + 0.1f * (epsilon - eps_end);
```

- [ ] **Step 2: Use eps_mag in magnitude branch selection**

Find the magnitude selection block at lines 796-802:

```c
        /* Branch 1: magnitude — Gem 2: uses eps_other (reduced epsilon) */
        if (lcg_random(&rng) < eps_other) {
```

Replace with:

```c
        /* Branch 1: magnitude — uses eps_mag (full exploration, same as direction) */
        if (lcg_random(&rng) < eps_mag) {
```

- [ ] **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/experience_kernels.cu
git commit -m "fix(magnitude): full epsilon exploration for magnitude branch — same as direction"
```

---

### Task 2: Counterfactual Magnitude Gradient on Flat Samples

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

**Problem:** Gem 3 zeros the magnitude gradient when direction=Flat (line 350 in c51_loss, line 298 in mse_loss). With 58% Flat actions, the magnitude branch trains on <42% of the batch. This is gradient starvation.

**Fix:** On Flat samples, instead of zeroing the magnitude loss entirely, compute what the magnitude Q distribution WOULD look like for the most likely non-Flat direction (argmax of Short vs Long Q-values from the direction branch). Use this counterfactual at 0.3× weight. This gives the magnitude branch gradient signal even on Flat samples.

**Design:** The counterfactual works by replacing `shmem_lp[j] = shmem_val[j]` (V-only, zero advantage) with `shmem_lp[j] = shmem_val[j] + 0.3 * (shmem_adv[a_d_cf * NA + j] - a_mean)` where `a_d_cf` is the magnitude action that the model WOULD have chosen if the best non-Flat direction were selected. Since the advantage values are already in shmem_adv (loaded for the magnitude branch), this requires zero extra memory or kernel launches.

The non-Flat direction argmax comes from the direction advantage pointer (on_adv_row[0]), which is already available. We read `on_adv_row[0]` to find which non-Flat direction has the highest advantage, then use that as the hypothetical direction for the counterfactual.

Note: We don't need the actual direction's action — we need the MAGNITUDE's advantage values (already in shmem_adv for d==1), but weighted at 0.3× instead of zeroed.

- [ ] **Step 1: Add counterfactual to C51 loss kernel**

In `c51_loss_kernel.cu`, find the Gem 3 block at lines 349-355:

```c
        /* Gem 3: Zero magnitude loss when direction=Flat */
        if (d == 1 && branch_action[0] == 1) {
            /* Direction was Flat — magnitude is irrelevant, zero its gradient */
            for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
                shmem_lp[j] = shmem_val[j];  /* Q = V only */
        }
        __syncthreads();
```

Replace with:

```c
        /* Gem 3+: Counterfactual magnitude gradient on Flat samples.
         * Instead of zeroing magnitude loss (which starves the branch of gradient),
         * compute what the magnitude Q distribution would look like for the best
         * non-Flat direction, weighted at 0.3×. This gives the magnitude branch
         * gradient signal on ALL samples, not just the 42% with directional actions. */
        if (d == 1 && branch_action[0] == 1) {
            /* Magnitude advantage is already in shmem_adv. Compute counterfactual
             * Q = V + 0.3 * (A_mag[a_d] - mean(A_mag)) instead of Q = V only.
             * The 0.3× weight prevents counterfactual from dominating real gradient. */
            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();
```

**Key insight:** `a_d` is the magnitude action that WAS selected (even though direction is Flat, the action encoding still encodes a magnitude choice — `mag_idx` is set to 1 by the Flat shortcut, so `a_d=1` always on Flat). The counterfactual computes loss against what the network outputs for magnitude=Half, weighted at 0.3×. This teaches the magnitude branch to maintain meaningful Q-value distributions even when Flat is selected, so when a directional trade IS taken, the magnitude branch already has learned preferences.

- [ ] **Step 2: Apply same fix to MSE loss kernel**

In `mse_loss_kernel.cu`, find the identical Gem 3 block at lines 297-303:

```c
        /* Gem 3: Zero magnitude loss when direction=Flat */
        if (d == 1 && branch_action[0] == 1) {
            /* Direction was Flat — magnitude is irrelevant, zero its gradient */
            for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
                shmem_lp[j] = shmem_val[j];  /* Q = V only */
        }
        __syncthreads();
```

Replace with the IDENTICAL counterfactual block from Step 1:

```c
        /* Gem 3+: Counterfactual magnitude gradient on Flat samples. */
        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();
```

- [ ] **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 \
       crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
git commit -m "feat(magnitude): counterfactual gradient on Flat samples — 0.3× weight"
```

---

### Task 3: Magnitude Gradient Amplification in Grad Kernels

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

**Problem:** Even with the counterfactual (Task 2), the magnitude branch's gradient is still weaker than other branches because:
1. Counterfactual uses 0.3× weight on Flat samples (not 1.0×)
2. The effective gradient magnitude for magnitude is ~0.3×58% + 1.0×42% = 59.4% of what direction gets

**Fix:** Scale the magnitude gradient (d==1) by a compensation factor in the grad kernels. The factor is `1.0 / (0.3 * flat_fraction + (1.0 - flat_fraction))`. With 58% Flat, this is `1.0 / (0.3*0.58 + 0.42) = 1.0 / 0.594 ≈ 1.68×`.

Rather than computing the exact flat fraction (which varies per epoch), use a fixed scale of `1.5×` for the magnitude branch. This is conservative enough to avoid gradient explosion but sufficient to compensate for the counterfactual discount.

**Design:** The gradient amplification goes in the grad kernel (not the loss kernel) because:
- The loss kernel computes per-sample loss (scalar) — scaling it would distort PER priorities
- The grad kernel computes per-element gradients — scaling the magnitude branch's elements is clean
- The gradient clipping AFTER the grad kernel will prevent any explosion

- [ ] **Step 1: Add magnitude gradient amplification to C51 grad kernel**

In `c51_grad_kernel.cu`, find the dueling gradient loop at lines 76-82:

```c
    /* d_adv[b, d, a, j] = 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 = d_combined * dueling_grad;
        int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], grad_val);
    }
```

Replace with:

```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;
        int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], grad_val);
    }
```

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

In `mse_grad_kernel.cu`, find the identical dueling gradient loop and apply the same `mag_amp` pattern. Search for `d_adv[b, d, a, j]` or `d_combined * dueling_grad`.

Read the file first, then apply the same change: add `float mag_amp = (d == 1) ? 1.5f : 1.0f;` before the loop, and multiply `grad_val` by `mag_amp`.

- [ ] **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): 1.5× gradient amplification to compensate Flat counterfactual discount"
```

---

### Task 4: Magnitude-Scaled Micro-Reward

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

**Problem:** The dense micro-reward (Gem 1) currently rewards directional accuracy: Short profits when price falls, Long profits when price rises. But it doesn't scale by magnitude — Short-Small and Short-Full get identical micro-reward. This means the model gets no dense signal about whether Full or Small was the right size choice.

**Fix:** Scale the micro-reward by the magnitude fraction: Small(0.25), Half(0.50), Full(1.00). This gives Full positions 4× the micro-reward of Small positions, creating a direct gradient signal for the magnitude branch through the reward channel.

- [ ] **Step 1: Add magnitude scaling to micro-reward**

In `experience_kernels.cu`, find the micro-reward computation at lines 1439-1454. After line 1453 (`reward += adaptive_scale * dir_reward / micro_vol_proxy;`), the reward is not magnitude-scaled.

Find:

```c
        reward += adaptive_scale * dir_reward / micro_vol_proxy;
```

Replace with:

```c
        /* Scale micro-reward by magnitude fraction: Full(1.0) gets 4× Small(0.25).
         * This gives the magnitude branch direct reward signal via the dense reward. */
        float mag_fraction;
        if (mag_idx == 0)      mag_fraction = 0.25f;
        else if (mag_idx == 1) mag_fraction = 0.50f;
        else                   mag_fraction = 1.00f;
        reward += adaptive_scale * dir_reward * mag_fraction / micro_vol_proxy;
```

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

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

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "feat(magnitude): scale micro-reward by magnitude fraction — Full gets 4× Small signal"
```

---

### Task 5: Unit Tests + Verification

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

- [ ] **Step 1: Add magnitude exploration test**

Add to `reward_v8.rs` after the existing `test_flat_shortcut`:

```rust
    #[test]
    fn test_magnitude_epsilon_equals_direction() {
        // Verify magnitude branch gets full epsilon, not 10%
        let eps_start = 0.20_f64;
        let eps_end = 0.02;
        // Old: eps_other = eps_end + 0.1 * (eps - eps_end) = 0.02 + 0.018 = 0.038
        let eps_old_mag = eps_end + 0.1 * (eps_start - eps_end);
        // New: eps_mag = eps (same as direction)
        let eps_new_mag = eps_start;
        // Magnitude should get 5× more exploration than before
        assert!(eps_new_mag > eps_old_mag * 4.0,
            "New magnitude epsilon ({eps_new_mag}) should be >4× old ({eps_old_mag})");
        assert!((eps_new_mag - eps_start).abs() < 1e-10,
            "Magnitude epsilon should equal direction epsilon");
    }

    #[test]
    fn test_counterfactual_weight() {
        // Verify counterfactual doesn't dominate real gradient
        let cf_weight = 0.3_f64;
        let flat_fraction = 0.58;
        // Effective gradient per sample: cf_weight × flat + 1.0 × non-flat
        let effective = cf_weight * flat_fraction + 1.0 * (1.0 - flat_fraction);
        assert!(effective > 0.5, "Effective gradient should be >50%: {effective}");
        assert!(effective < 1.0, "Counterfactual must reduce magnitude gradient: {effective}");
        // Amplification factor compensates: 1/effective ≈ 1.68
        let amp = 1.0 / effective;
        assert!((amp - 1.68).abs() < 0.1, "Expected ~1.68× amplification, got {amp}");
    }

    #[test]
    fn test_magnitude_micro_reward_scaling() {
        // Verify Full gets 4× the micro-reward of Small
        let dir_reward = 0.01_f64; // price moved in our direction
        let adaptive_scale = 1.0;
        let micro_vol = 0.005;
        let small_reward = adaptive_scale * dir_reward * 0.25 / micro_vol;
        let full_reward = adaptive_scale * dir_reward * 1.00 / micro_vol;
        assert!((full_reward / small_reward - 4.0).abs() < 1e-10,
            "Full/Small reward ratio should be 4.0, got {}", full_reward / small_reward);
    }
```

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

```bash
SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::reward_v8 --nocapture 2>&1 | tail -20
```
Expected: all tests pass including the 3 new ones.

- [ ] **Step 3: Run full test suite**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | tail -3
```
Expected: 906+ passed, 0 failed.

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs
git commit -m "test(magnitude): 3 unit tests — epsilon parity, counterfactual weight, micro-reward scaling"
```

---

### Task 6: Push + Deploy Hyperopt

- [ ] **Step 1: Push**

```bash
git push origin main
```

- [ ] **Step 2: Launch hyperopt**

```bash
./scripts/argo-train.sh dqn
```

- [ ] **Step 3: Monitor early epochs**

Watch for:
- dir×mag diversity should be >3/9 (target: 5-6/9 by epoch 10)
- Magnitude distribution: Small/Medium/Full should be more balanced
- PF trending toward 1.0+
- No `CUDA_ERROR` or SIGSEGV

```bash
argo logs <workflow-name> -n foxhunt --follow=false -c main 2>&1 | grep -E 'Action div|actions=\[|PF='
```
