# Phase 2: Hierarchical 4-Branch DQN — Complete Architectural Redesign

**Goal:** Replace the 9-output exposure branch (i%3 degeneracy) with hierarchical direction(3) × magnitude(3), adding 3 novel gems. Eliminate all i%3-related workarounds. Production-ready, no stubs, no TODOs.

**Architecture:** 4-branch branching DQN: direction(3) × magnitude(3) × order(3) × urgency(3) = 81 actions. Each branch has 3 outputs with dueling mean subtraction (-1/2 gradient = 33%). Direction gets its own dense micro-reward, per-branch epsilon, and Flat detection shortcut.

**Tech Stack:** CUDA (nvcc → cubin), Rust (cudarc), cuBLAS SGEMM, TOML config

---

## 1. Network Architecture Change

### Current: 3 branches, 15 Q-outputs
```
Shared trunk: state[80] → h1[hidden] → h2[hidden]
Value head:   h2 → v1[value_h] → v2[num_atoms]
Branch 0:     h2 → b0_fc[adv_h] → b0_out[9 × num_atoms]  (exposure — BROKEN)
Branch 1:     h2 → b1_fc[adv_h] → b1_out[3 × num_atoms]  (order)
Branch 2:     h2 → b2_fc[adv_h] → b2_out[3 × num_atoms]  (urgency)

Q-values: [9+3+3 = 15] per sample
```

### New: 4 branches, 12 Q-outputs
```
Shared trunk: state[80] → h1[hidden] → h2[hidden]
Value head:   h2 → v1[value_h] → v2[num_atoms]
Branch 0:     h2 → b0_fc[adv_h] → b0_out[3 × num_atoms]  (direction: Short/Flat/Long)
Branch 1:     h2 → b1_fc[adv_h] → b1_out[3 × num_atoms]  (magnitude: Quarter/Half/Full)
Branch 2:     h2 → b2_fc[adv_h] → b2_out[3 × num_atoms]  (order: Market/LimitMaker/IoC)
Branch 3:     h2 → b3_fc[adv_h] → b3_out[3 × num_atoms]  (urgency: Patient/Normal/Aggressive)

Q-values: [3+3+3+3 = 12] per sample
```

### Weight tensor layout (flat params buffer)
```
[0]  w_s1       [1]  b_s1       (shared)
[2]  w_s2       [3]  b_s2       (shared)
[4]  w_v1       [5]  b_v1       (value)
[6]  w_v2       [7]  b_v2       (value)
[8]  w_b0fc     [9]  b_b0fc     (direction FC)
[10] w_b0out    [11] b_b0out    (direction output)
[12] w_b1fc     [13] b_b1fc     (magnitude FC)
[14] w_b1out    [15] b_b1out    (magnitude output)
[16] w_b2fc     [17] b_b2fc     (order FC)
[18] w_b2out    [19] b_b2out    (order output)
[20] w_b3fc     [21] b_b3fc     (urgency FC)
[22] w_b3out    [23] b_b3out    (urgency output)
[24] w_bn       [25] b_bn       (bottleneck)
```

`NUM_WEIGHT_TENSORS` = 26 (was 22). `compute_param_sizes` returns `[usize; 26]`.

---

## 2. Position Mapping

### New function (replaces `compute_target_position`)
```cuda
__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;
}
```

Position grid:
```
           Qtr(0.25)  Half(0.50)  Full(1.00)
Short(-1): S25        S50         S100
Flat(0):   0          0           0
Long(+1):  L25        L50         L100
```

7 unique positions. The 3 Flat duplicates are harmless — the model learns direction=Flat quickly.

---

## 3. Action Encoding/Decoding

### New 4-branch encoding
```cuda
// Encode: 4 branch indices → flat action (0-80)
int action = dir * 27 + mag * 9 + order * 3 + urgency;

// Decode: flat action → 4 branch indices
int dir = action / 27;
int mag = (action % 27) / 9;
int order = (action % 9) / 3;
int urgency = action % 3;
```

### Update in `trade_physics.cuh`
```cuda
__device__ __forceinline__ int decode_direction(int action, int b1_size, int b2_size, int b3_size) {
    return action / (b1_size * b2_size * b3_size);
}
__device__ __forceinline__ int decode_magnitude(int action, int b1_size, int b2_size, int b3_size) {
    return (action / (b2_size * b3_size)) % b1_size;
}
__device__ __forceinline__ int decode_order_type_4branch(int action, int b2_size, int b3_size) {
    return (action / b3_size) % b2_size;
}
__device__ __forceinline__ int decode_urgency_4branch(int action, int b3_size) {
    return action % b3_size;
}
```

---

## 4. Gem 1: Direction-Specific Dense Micro-Reward

In `experience_env_step`, the dense micro-reward becomes direction-aware:

```cuda
/* Dense micro-reward: direction-specific signal */
if (fabsf(position) > 0.001f && micro_reward_scale > 0.0f) {
    int dir_idx = decode_direction(action_idx, b1_size, b2_size, b3_size);
    float price_change = (raw_next - raw_close) / fmaxf(raw_close, 1.0f);

    /* Direction-specific: Short profits when price falls, Long when rises */
    float dir_reward = 0.0f;
    if (dir_idx == 0) dir_reward = -price_change;  /* Short */
    if (dir_idx == 2) dir_reward = price_change;   /* Long */
    /* dir_idx == 1 (Flat): no micro-reward — position should be 0 */

    float vol_ratio = micro_vol_proxy / 0.005f;
    float adaptive_scale = micro_reward_scale / fmaxf(sqrtf(vol_ratio), 0.5f);
    reward += adaptive_scale * dir_reward / micro_vol_proxy;
}
```

---

## 5. Gem 2: Per-Branch Epsilon (Direction-Specific Exploration)

In `experience_action_select`, compute different epsilon per branch:

```cuda
/* Per-branch epsilon: direction explores more (the hard problem) */
float eps_dir = epsilon;                                    /* Full epsilon for direction */
float eps_other = eps_end + 0.1f * (epsilon - eps_end);    /* 10% of epsilon range for others */

/* Direction: epsilon-greedy with eps_dir */
if (random < eps_dir) { dir_idx = random_choice(3); }
else { dir_idx = argmax(q_dir[0..3]); }

/* Magnitude/order/urgency: epsilon-greedy with eps_other */
if (random < eps_other) { mag_idx = random_choice(3); }
else { mag_idx = argmax(q_mag[0..3]); }
/* ... same for order and urgency ... */
```

---

## 6. Gem 3: Flat Detection Shortcut

In `experience_action_select`:

```cuda
/* Flat shortcut: when direction=Flat, magnitude is irrelevant */
if (dir_idx == 1) {
    mag_idx = 1;  /* Default to Half — doesn't affect position (0 × anything = 0) */
    /* Skip magnitude argmax — save compute and clarify gradient */
}
```

In C51/MSE loss kernels: when the taken direction is Flat, the magnitude branch receives NO gradient for this sample (the outcome doesn't depend on magnitude). This naturally focuses the magnitude branch on learning sizing for non-Flat positions.

```cuda
/* In C51 loss, branch d=1 (magnitude): */
if (branch_action[0] == 1) {  /* direction was Flat */
    /* Zero the gradient for magnitude — outcome was independent of magnitude */
    for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
        shmem_lp[j] = 0.0f;  /* No CE loss for magnitude when Flat */
}
```

---

## 7. Config Changes

### New fields
```rust
pub b3_size: usize,  // urgency branch size (was b2_size)
// b0_size = 3 (direction), b1_size = 3 (magnitude)
// b2_size = 3 (order), b3_size = 3 (urgency)
```

### Removed fields
```rust
// DELETE: exposure_aux_weight, exposure_aux_warmup_epochs
// DELETE: All aux_adam_* fields (14 fields)
// DELETE: best_exposure_out buffer
// DELETE: exposure_aux_grad_kernel, exposure_pretrain_step kernel
```

### TOML updates
```toml
b0_size = 3   # direction (was 9 — exposure)
b1_size = 3   # magnitude (was 3 — order)
b2_size = 3   # order (was 3 — urgency)
b3_size = 3   # urgency (NEW)
```

---

## 8. CUDA Kernel Changes

### `NUM_BRANCHES` constant
All kernels that reference `NUM_BRANCHES`: change from 3 to 4.
- `c51_loss_kernel.cu`
- `mse_loss_kernel.cu`
- `cql_grad_kernel.cu`
- `expected_q_kernel.cu`
- `experience_kernels.cu` (action_select + env_step)
- `backtest_env_kernel.cu`

### C51 loss kernel
- Loop `for d in 0..4` instead of `0..3`
- `branch_sizes = [3, 3, 3, 3]` instead of `[9, 3, 3]`
- `branch_action = [dir, mag, order, urgency]` decoded from flat action
- Gem 3: zero magnitude gradient when direction=Flat
- Restore mean subtraction for ALL branches (remove the d>0 skip — all 3-output branches work with dueling)

### MSE loss kernel
- Same changes as C51

### Experience action_select kernel
- 4 branch argmax instead of 3
- Gem 2: per-branch epsilon
- Gem 3: Flat detection shortcut
- Q-values stride: 12 instead of 15

### Experience env_step kernel
- `compute_target_position_4branch(dir_idx, mag_idx, max_pos)` replaces `compute_target_position(exposure_idx, b0_size, max_pos)`
- 4-branch action decoding
- Gem 1: direction-specific micro-reward
- Remove CEA counterfactual loop (no longer needed)
- Remove `out_best_exposure` write

### Backtest env_step kernel
- Same position mapping change

### Expected-Q kernel
- 4 branches in Q-value computation

---

## 9. Rust Changes

### `gpu_dqn_trainer.rs`
- `NUM_WEIGHT_TENSORS = 26` (was 22)
- `compute_param_sizes` returns 26-element array with 4 branch heads
- Add w_b3fc, b_b3fc, w_b3out, b_b3out to weight layout
- Add b3 branch weight sets (DuelingWeightSet → 4 BranchingWeightSets)
- DELETE: `aux_adam_*` (14 fields), `exposure_aux_*` (5 fields), `best_exposure_out`
- DELETE: `launch_exposure_aux_grad`, `set_exposure_aux_weight`, `run_pretrain_step`
- DELETE: `exposure_aux_grad_kernel`, `exposure_pretrain_kernel`, `popart_normalize_kernel` (not wired anyway)

### `batched_forward.rs`
- 4 branch GEMMs instead of 3
- Q-values output: [batch, 12] instead of [batch, 15]

### `batched_backward.rs`
- 4 branch backward GEMMs
- 4 branches in the d-loop

### `gpu_experience_collector.rs`
- `branch_sizes = [3, 3, 3, 3]`
- Q-values buffer: [alloc × 12] instead of [alloc × 15]
- DELETE: `best_exposure_out` buffer and accessor
- Config: `b3_size` field added

### `fused_training.rs`
- 4 branch weight sets (online + target)
- DELETE: exposure aux GEMM call from `submit_aux_ops`
- DELETE: `exposure_aux_weight`, `exposure_targets`, `pending_exposure_targets`
- Mega-graph capture with 4 branches

### `training_loop.rs`
- DELETE: pre-training phase (epoch 0)
- DELETE: aux target DtoD copy
- DELETE: aux weight decay schedule
- Config wiring for b3_size

### `config.rs`
- `b3_size: usize` (default 3)
- DELETE: `exposure_aux_weight`, `exposure_aux_warmup_epochs`
- Position mapping: `compute_max_position` unchanged (uses max_leverage)

### `dqn.rs` (ml-dqn crate)
- Network construction: 4 advantage heads instead of 3
- `BranchingQNetwork`: 4 linear layers for advantages
- Action selection: 4-branch argmax

### `gpu_backtest_evaluator.rs`
- 4-branch DqnBacktestConfig
- 4-branch forward pass

### `training_profile.rs`
- Parse b3_size from TOML

---

## 10. Dead Code Removal

| Component | Lines (~) | Why |
|-----------|-----------|-----|
| `exposure_aux_grad_kernel` (CUDA) | 60 | 3-output branches don't need aux |
| `exposure_pretrain_step` (CUDA) | 70 | Direction pre-training replaced by dense micro-reward |
| `launch_exposure_aux_grad` (Rust) | 90 | Dead with aux kernel |
| `set_exposure_aux_weight` (Rust) | 15 | Dead |
| `run_pretrain_step` (Rust) | 40 | Dead |
| `aux_adam_*` fields (14) + allocation | 60 | Separate optimizer no longer needed |
| `best_exposure_out` buffer + wiring | 30 | CEA loop removed |
| Pre-training phase (training_loop) | 40 | Replaced by Gem 1 dense reward |
| Aux target DtoD copy | 15 | Dead |
| CEA counterfactual loop (experience kernel) | 30 | Replaced by Gem 1 |
| Mean subtraction skip (d>0 guard) | 20 | All branches use normal dueling |
| `compute_target_position` linear function | 10 | Replaced by 4-branch version |
| **Total removed** | **~480 lines** | |

---

## 11. Testing

### Unit tests
- `test_4branch_action_encoding` — encode/decode roundtrip for all 81 actions
- `test_4branch_position_mapping` — direction × magnitude → correct position
- `test_flat_shortcut` — Flat × any_magnitude = 0
- `test_per_branch_epsilon` — direction gets higher epsilon than others
- `test_direction_micro_reward` — Short gets negative on price up, Long gets positive

### GPU smoke tests
- Walk-forward: Sharpe should be positive (features have r=0.37 predictive power)
- Q-values: NO i%3 pattern (each of 3 direction outputs is unique)
- Action diversity: 3/3 direction, 3/3 magnitude by epoch 3

### Regression tests
- All existing 900+ unit tests pass
- Existing GPU smoke tests pass (adjusted for new branch sizes)
- No SIGSEGV from kernel arg order changes

---

## 12. Performance Impact

| Metric | Before (3-branch) | After (4-branch) |
|--------|-------------------|-------------------|
| Q-network output | 15 × num_atoms | 12 × num_atoms (20% fewer) |
| Forward GEMMs | 3 branch heads | 4 branch heads (1 more, but smaller) |
| Backward GEMMs | 3 | 4 |
| C51 loss loop | 3 iterations | 4 iterations |
| Net FLOPS change | — | ~-10% (smaller output matrices) |
| Epoch time impact | — | ~neutral (fewer Q-values, more branches) |

---

## 13. Files Modified (comprehensive)

| File | Changes |
|------|---------|
| `trade_physics.cuh` | New decode functions, `compute_target_position_4branch` |
| `c51_loss_kernel.cu` | NUM_BRANCHES=4, restore mean subtraction, Gem 3 (Flat gradient zero) |
| `mse_loss_kernel.cu` | NUM_BRANCHES=4, restore mean subtraction |
| `cql_grad_kernel.cu` | NUM_BRANCHES=4 |
| `expected_q_kernel.cu` | 4-branch Q computation |
| `experience_kernels.cu` | 4-branch action select + env_step, Gem 1+2+3, delete CEA loop |
| `backtest_env_kernel.cu` | 4-branch position mapping |
| `backtest_gather_kernel.cu` | Portfolio feature computation (if branch-dependent) |
| `dqn_utility_kernels.cu` | DELETE exposure_aux_grad_kernel, exposure_pretrain_step |
| `gpu_dqn_trainer.rs` | 26 weight tensors, 4 branch sets, DELETE aux optimizer |
| `batched_forward.rs` | 4-branch GEMMs |
| `batched_backward.rs` | 4-branch backward |
| `gpu_experience_collector.rs` | 4-branch config, DELETE best_exposure |
| `fused_training.rs` | 4 branch weight sets, DELETE aux GEMM, mega-graph update |
| `training_loop.rs` | DELETE pre-training, DELETE aux targets, b3_size wiring |
| `config.rs` | b3_size, DELETE aux fields |
| `training_profile.rs` | b3_size parsing |
| `dqn.rs` (ml-dqn) | 4-head network construction |
| `gpu_backtest_evaluator.rs` | 4-branch config |
| `config/training/*.toml` (×4) | b0=3, b1=3, b2=3, b3=3 |
| `smoke_tests/reward_v8.rs` | Update tests for 4-branch |
