# Phase 2: Hierarchical 4-Branch DQN — 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:** Replace the 9-output exposure branch (i%3 degeneracy) with hierarchical direction(3) × magnitude(3) × order(3) × urgency(3) = 81 actions. Eliminate all i%3-related workarounds. Add 3 gems (direction-specific micro-reward, per-branch epsilon, Flat detection shortcut). Delete ~480 lines of dead code.

**Architecture:** 10 tasks in dependency order. Tasks 1-2 (config + CUDA header) are foundation. Tasks 3-5 (CUDA kernels) modify 8 .cu files. Tasks 6-8 (Rust integration) wire 4 branches through the training pipeline. Task 9 removes dead code. Task 10 tests and deploys. Each task produces a compiling codebase.

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

**CRITICAL CONSTRAINT:** All CUDA kernels are compiled via build.rs at build time (NOT NVRTC). Kernel arg order in Rust MUST match CUDA parameter order EXACTLY. All branches use standard dueling mean subtraction. No CPU path, no stubs.

---

### Task 1: Config Foundation — b3_size + Delete Aux Fields

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs`
- Modify: `crates/ml/src/training_profile.rs`
- Modify: `config/training/dqn-production.toml`
- Modify: `config/training/dqn-hyperopt.toml`
- Modify: `config/training/dqn-smoketest.toml`
- Modify: `config/training/dqn-localdev.toml`

- [ ] **Step 1: Add b3_size to DQNHyperparameters and GpuDqnTrainConfig**

In `config.rs`, find `pub b2_size: usize`. Add after:
```rust
    /// Branch 3 size: urgency (Patient/Normal/Aggressive). Default 3.
    pub b3_size: usize,
```

Change defaults: `b0_size: 9` → `b0_size: 3` in BOTH `Default` and `conservative()`.

Add `b3_size: 3` defaults in BOTH constructors.

In `GpuDqnTrainConfig`, find `pub branch_2_size: usize`. Add:
```rust
    pub branch_3_size: usize,
```
Set default to 3.

- [ ] **Step 2: Delete aux-related config fields**

In `config.rs`, delete:
```rust
    pub exposure_aux_weight: f64,
    pub exposure_aux_warmup_epochs: usize,
```

Remove from `Default`, `conservative()`, and `apply_family_scaling`.

- [ ] **Step 3: Update all 4 TOML configs**

Set in each TOML:
```toml
b0_size = 3   # direction
b1_size = 3   # magnitude
b2_size = 3   # order
b3_size = 3   # urgency
```

Remove `exposure_aux_weight` and `exposure_aux_warmup_epochs` from all TOMLs.

- [ ] **Step 4: Update training_profile.rs**

Add parsing for `b3_size`. Search for where `b2_size` is parsed and add `b3_size` in the same pattern.

- [ ] **Step 5: Verify build + fix compilation errors**

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

Fix any compilation errors from removed fields (grep for `exposure_aux_weight` and `exposure_aux_warmup_epochs` in all .rs files and remove/update references).

- [ ] **Step 6: Commit**

```bash
git add -A
git commit -m "feat(4branch): config foundation — b0=3 direction, b1=3 magnitude, b2=3 order, b3=3 urgency

Delete exposure_aux_weight, exposure_aux_warmup_epochs.
All branch sizes now 3 (was 9,3,3). Total actions still 81 (3×3×3×3)."
```

---

### Task 2: CUDA Shared Header — Position Mapping + Action Decode

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/trade_physics.cuh`

- [ ] **Step 1: Add 4-branch position mapping function**

After `compute_target_position`, add:

```cuda
/* ── 4-branch hierarchical position mapping (Phase 2) ────────────────
 * direction(3) × magnitude(3) replaces flat exposure(9).
 * Direction: 0=Short(-1), 1=Flat(0), 2=Long(+1)
 * Magnitude: 0=Quarter(0.25), 1=Half(0.50), 2=Full(1.00)
 * Position = direction × magnitude × max_position */
__device__ __forceinline__ float compute_target_position_4branch(
    int dir_idx, int mag_idx, float max_position
) {
    float direction = (dir_idx == 0) ? -1.0f : (dir_idx == 2) ? 1.0f : 0.0f;
    float magnitude = (mag_idx == 0) ? 0.25f : (mag_idx == 1) ? 0.5f : 1.0f;
    return direction * magnitude * max_position;
}

/* ── 4-branch action decoding ────────────────────────────────────────
 * Flat action = dir * (b1*b2*b3) + mag * (b2*b3) + order * b3 + urgency
 * With all sizes=3: action = dir*27 + mag*9 + order*3 + urgency */
__device__ __forceinline__ int decode_direction_4b(int action, int b1, int b2, int b3) {
    return action / (b1 * b2 * b3);
}
__device__ __forceinline__ int decode_magnitude_4b(int action, int b1, int b2, int b3) {
    return (action / (b2 * b3)) % b1;
}
__device__ __forceinline__ int decode_order_4b(int action, int b2, int b3) {
    return (action / b3) % b2;
}
__device__ __forceinline__ int decode_urgency_4b(int action, int b3) {
    return action % b3;
}
```

- [ ] **Step 2: Commit**

```bash
git add crates/ml/src/cuda_pipeline/trade_physics.cuh
git commit -m "feat(4branch): add compute_target_position_4branch + 4-branch action decode functions"
```

---

### Task 3: CUDA Experience Kernels — Action Select + Env Step (Gems 1, 2, 3)

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

This is the most impactful kernel change — it determines how the agent selects actions and computes rewards.

- [ ] **Step 1: Update action_select kernel for 4 branches**

Find `experience_action_select`. The kernel currently selects actions for 3 branches. Update to 4 branches with per-branch epsilon (Gem 2) and Flat shortcut (Gem 3).

Key changes:
- Add `int b3_size` parameter
- Q-values stride: `b0_size + b1_size + b2_size + b3_size` (12 instead of 15)
- 4-branch argmax with per-branch epsilon
- When direction=Flat (dir_idx==1), set mag_idx=1 (Gem 3: skip magnitude)
- Factored action = `dir * b1*b2*b3 + mag * b2*b3 + order * b3 + urgency`

- [ ] **Step 2: Update env_step kernel for 4 branches**

Find `experience_env_step`. Key changes:
- Add `int b3_size` parameter
- Replace `compute_target_position(exposure_idx, b0_size, max_pos)` with `compute_target_position_4branch(dir_idx, mag_idx, max_pos)`
- 4-branch action decoding: `dir = decode_direction_4b(action, b1, b2, b3)`, etc.
- Replace the flat micro-reward with direction-specific (Gem 1)
- Remove the CEA counterfactual loop
- Remove `out_best_exposure` write
- Remove the `out_best_exposure` parameter from the kernel signature

- [ ] **Step 3: Wire b3_size in gpu_experience_collector.rs**

Add `b3_size` to kernel launch args for both action_select and env_step. Remove `out_best_exposure` arg.

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

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

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu \
       crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
git commit -m "feat(4branch): 4-branch action select + env step with Gems 1,2,3

Gem 1: direction-specific dense micro-reward
Gem 2: per-branch epsilon (direction explores more)
Gem 3: Flat detection shortcut (skip magnitude when Flat)"
```

---

### Task 4: CUDA Loss Kernels — C51 + MSE + CQL + Expected-Q + Grad

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/cql_grad_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/expected_q_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu`

All 6 kernels need the same changes:

- [ ] **Step 1: Update NUM_BRANCHES and branch arrays in ALL 6 kernels**

For EACH kernel file:
1. Change `#define NUM_BRANCHES 3` → `#define NUM_BRANCHES 4`
2. Add `int b3_size` parameter where `b2_size` is the last branch param
3. Update branch arrays:
```cuda
// BEFORE:
int a0 = action / (b1_size * b2_size);
int a1 = (action / b2_size) % b1_size;
int a2 = action % b2_size;
const int branch_action[3] = { a0, a1, a2 };
const int branch_sizes[3]  = { b0_size, b1_size, b2_size };

// AFTER:
int a0 = action / (b1_size * b2_size * b3_size);
int a1 = (action / (b2_size * b3_size)) % b1_size;
int a2 = (action / b3_size) % b2_size;
int a3 = action % b3_size;
const int branch_action[4] = { a0, a1, a2, a3 };
const int branch_sizes[4]  = { b0_size, b1_size, b2_size, b3_size };
```
4. Update advantage pointer arrays to include 4th branch
5. Restore mean subtraction for ALL branches (remove `if (d > 0)` guard)

- [ ] **Step 2: Add Gem 3 to C51 and MSE loss kernels — zero magnitude gradient when Flat**

In the C51 loss main loop, after computing the CE loss for branch d=1 (magnitude), add:

```cuda
        /* Gem 3: Zero magnitude gradient when direction=Flat */
        if (d == 1 && branch_action[0] == 1) {
            /* Direction was Flat — magnitude doesn't affect outcome */
            for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
                shmem_lp[j] = shmem_val[j];  /* Q = V only (no advantage) */
        }
```

Same in MSE loss.

- [ ] **Step 3: Wire b3_size in Rust kernel launch calls**

In `gpu_dqn_trainer.rs`, find ALL kernel launch calls for C51, MSE, CQL, expected_q, and grad kernels. Add `.arg(&b3)` where b2 is the last branch arg.

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

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

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_loss_kernel.cu \
       crates/ml/src/cuda_pipeline/mse_loss_kernel.cu \
       crates/ml/src/cuda_pipeline/cql_grad_kernel.cu \
       crates/ml/src/cuda_pipeline/expected_q_kernel.cu \
       crates/ml/src/cuda_pipeline/c51_grad_kernel.cu \
       crates/ml/src/cuda_pipeline/mse_grad_kernel.cu \
       crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat(4branch): NUM_BRANCHES=4 in all 6 loss/grad kernels + Gem 3 Flat gradient zero"
```

---

### Task 5: CUDA Backtest Kernel — 4-Branch Position Mapping

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

- [ ] **Step 1: Update backtest env_step for 4 branches**

Find `backtest_env_step_batch`. Add `int b3_size` parameter. Replace `compute_target_position` with `compute_target_position_4branch`. Update action decoding to 4-branch.

- [ ] **Step 2: Wire in gpu_backtest_evaluator.rs**

Add `.arg(&b3)` to the backtest kernel launch.

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/cuda_pipeline/backtest_env_kernel.cu \
       crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
git commit -m "feat(4branch): backtest env_step uses 4-branch position mapping"
```

---

### Task 6: Rust Weight Management — 26 Tensors, 4 Branch Heads

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs`
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs`

- [ ] **Step 1: Update compute_param_sizes to 26 tensors**

```rust
pub(crate) const NUM_WEIGHT_TENSORS: usize = 26;

pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT_TENSORS] {
    // ... existing [0]-[19] entries ...
    // ADD at indices [20]-[25]:
    //   [20] w_b3fc = adv_h * shared_h2
    //   [21] b_b3fc = adv_h
    //   [22] w_b3out = branch_3_size * num_atoms * adv_h
    //   [23] b_b3out = branch_3_size * num_atoms
    //   [24] w_bn (bottleneck)
    //   [25] b_bn
}
```

Note: the bottleneck tensors move from [20]-[21] to [24]-[25]. All spectral norm descriptors, weight flatten/unflatten, and GOFF offset computations must be updated.

- [ ] **Step 2: Add 4th branch to batched_forward.rs**

In `CublasForward`, add a 4th branch GEMM + bias. The Q-values output becomes `[batch, 12]` (was `[batch, 15]`).

The forward loop `for d in 0..3` becomes `for d in 0..4`. Add `w_b3fc`, `b_b3fc`, `w_b3out`, `b_b3out` weight pointers.

- [ ] **Step 3: Add 4th branch to batched_backward.rs**

In `backward_full`, add the 4th branch backward GEMM. The loop `for d in 0..3` becomes `for d in 0..4`.

Add GOFF offsets:
```rust
let goff_w_b3fc:  u64 = goff_b_b2out + b2 * na64 * f32;
let goff_b_b3fc:  u64 = goff_w_b3fc  + ah * sh2  * f32;
let goff_w_b3out: u64 = goff_b_b3fc  + ah         * f32;
let goff_b_b3out: u64 = goff_w_b3out + b3 * na64 * ah * f32;
```

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

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

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
       crates/ml/src/cuda_pipeline/batched_forward.rs \
       crates/ml/src/cuda_pipeline/batched_backward.rs
git commit -m "feat(4branch): 26 weight tensors, 4-branch forward + backward GEMMs"
```

---

### Task 7: Rust Integration — Fused Training + Experience Collector

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`

- [ ] **Step 1: Update FusedTrainingCtx for 4 branches**

Add 4th branch weight sets (online + target). Delete `exposure_aux_weight`, `exposure_targets`, `pending_exposure_targets`. Delete aux GEMM call from `submit_aux_ops`. Update mega-graph capture for 4 branches.

- [ ] **Step 2: Update GpuExperienceCollector for 4 branches**

- `branch_sizes = [3, 3, 3, 3]` (4 elements)
- Q-values buffer: `[alloc × 12]` instead of `[alloc × 15]`
- Delete `best_exposure_out` buffer and accessor
- Add `b3_size` to config
- Remove `best_exposure_buf()` accessor

- [ ] **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/trainers/dqn/fused_training.rs \
       crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
git commit -m "feat(4branch): fused training + collector wired for 4 branches, aux GEMM deleted"
```

---

### Task 8: Rust Training Loop + Network Construction

**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs`
- Modify: `crates/ml-dqn/src/dqn.rs`

- [ ] **Step 1: Update training_loop.rs**

- Delete pre-training phase (epoch 0)
- Delete aux target DtoD copy
- Delete aux weight decay schedule
- Wire `b3_size` in ExperienceCollectorConfig
- Update action diversity logging for 4 branches

- [ ] **Step 2: Update dqn.rs for 4-head network**

Find `BranchingQNetwork` or similar. Add 4th advantage head linear layer. Update action selection for 4-branch argmax. Update action encoding/decoding.

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

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

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/trainers/dqn/trainer/training_loop.rs \
       crates/ml-dqn/src/dqn.rs
git commit -m "feat(4branch): training loop + DQN network with 4 advantage heads"
```

---

### Task 9: Dead Code Removal (~480 lines)

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`

- [ ] **Step 1: Delete CUDA kernels from dqn_utility_kernels.cu**

Delete `exposure_aux_grad_kernel` (~60 lines) and `exposure_pretrain_step` (~70 lines).

- [ ] **Step 2: Delete aux optimizer from gpu_dqn_trainer.rs**

Delete ALL `aux_adam_*` fields (14 fields), allocations, struct literal entries, and the `launch_exposure_aux_grad` method (~90 lines). Delete `set_exposure_aux_weight` (~15 lines). Delete `run_pretrain_step` (~40 lines). Delete `exposure_aux_weight_gpu`, `exposure_aux_scratch`, kernel function handles.

Also delete `best_exposure_out` from the collector.

- [ ] **Step 3: Delete remaining references**

Search for `exposure_aux`, `best_exposure`, `pretrain_step`, `aux_adam` across ALL .rs files. Delete any remaining references.

- [ ] **Step 4: Delete old `compute_target_position` from trade_physics.cuh**

The old linear interpolation function is no longer used. Delete it (keep `compute_target_position_4branch`).

- [ ] **Step 5: Verify build + run tests**

```bash
SQLX_OFFLINE=true cargo check -p ml -p ml-dqn 2>&1 | tail -5
SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | grep "^test result:"
```

- [ ] **Step 6: Commit**

```bash
git add -A
git commit -m "cleanup(4branch): delete ~480 lines — aux optimizer, CEA, pre-training, old position mapping"
```

---

### Task 10: Tests + GPU Smoke Test + Final Verification + Push

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

- [ ] **Step 1: Add 4-branch unit tests**

Add to `reward_v8.rs`:

```rust
#[test]
fn test_4branch_action_encoding() {
    // All 81 actions encode/decode correctly
    for dir in 0..3 {
        for mag in 0..3 {
            for order in 0..3 {
                for urgency in 0..3 {
                    let action = dir * 27 + mag * 9 + order * 3 + urgency;
                    assert_eq!(action / 27, dir);
                    assert_eq!((action % 27) / 9, mag);
                    assert_eq!((action % 9) / 3, order);
                    assert_eq!(action % 3, urgency);
                    assert!(action < 81);
                }
            }
        }
    }
}

#[test]
fn test_4branch_position_mapping() {
    let max_pos = 2.0_f64;
    // Short × Full = -2.0
    assert!(((-1.0 * 1.0 * max_pos) - (-2.0)).abs() < 1e-10);
    // Flat × anything = 0
    assert!((0.0 * 0.25 * max_pos).abs() < 1e-10);
    assert!((0.0 * 1.0 * max_pos).abs() < 1e-10);
    // Long × Quarter = +0.5
    assert!((1.0 * 0.25 * max_pos - 0.5).abs() < 1e-10);
    // Long × Half = +1.0
    assert!((1.0 * 0.5 * max_pos - 1.0).abs() < 1e-10);
}

#[test]
fn test_flat_shortcut() {
    // Flat direction × any magnitude = 0
    for mag in 0..3 {
        let dir_fraction = 0.0_f64; // Flat
        let mag_scale = match mag { 0 => 0.25, 1 => 0.5, _ => 1.0 };
        assert!((dir_fraction * mag_scale).abs() < 1e-10);
    }
}
```

- [ ] **Step 2: Fix existing smoke test expectations**

Update any tests that reference `b0_size=9`, branch count=3, or the old 15-output Q-values.

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

```bash
SQLX_OFFLINE=true cargo test -p ml --lib -p ml-dqn --lib 2>&1 | grep "^test result:"
```
Expected: 900+ passed, 0 failed

- [ ] **Step 4: GPU 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 | tail -10
```

Expected: passes, NO i%3 pattern in Q-values, all 3 direction outputs are different.

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

```bash
git add -A
git commit -m "test(4branch): 3 new unit tests + smoke test verification — NO i%3 pattern"
git push origin main
./scripts/argo-train.sh dqn --epochs 40
```
