# ABI Complete 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 magnitude collapse (Small=72%), order/urgency degeneration (<2%), and epoch-over-epoch Sharpe degradation through per-branch exploration, IQN uncertainty bonuses, hindsight action relabeling, gradient vaccine feeding, c51_alpha warmup fix, and ensemble inference variance.

**Architecture:** Each branch becomes an autonomous learner with its own epsilon, exploration bonus, and gradient signal. IQN quantile spread provides anti-collapse signal. Hindsight relabeling teaches magnitude what "bigger" would have earned. Gradient vaccine prevents convergence to non-diverse policy. Fixed c51_alpha eliminates mid-training loss landscape shifts.

**Tech Stack:** Rust, CUDA kernels (experience_kernels.cu, iqn_dual_head_kernel.cu), cudarc, ml-core

**Current state after ABI Foundation:** 7-level exposure, PopArt enabled, Flat floor removed, tau coupling active. Flat dropped from 82% → 12%. But Small still dominates at 72%, Half/Full at 4% each. Order diversity: Market >96%.

---

## File Structure

| Action | File | Purpose |
|--------|------|---------|
| Modify | `experience_kernels.cu` | Per-branch epsilon vector, magnitude Boltzmann fix |
| Modify | `gpu_action_selector.rs` | Accept per-branch epsilon, IQR bonus |
| Modify | `gpu_iqn_head.rs` | Compute IQR per action, export to action selector |
| Modify | `fused_training.rs` | Fixed c51_alpha, vaccine batch source, ensemble sync |
| Modify | `training_loop.rs` | Per-branch epsilon computation, vaccine sampling |
| Modify | `gpu_her.rs` | Hindsight magnitude relabeling kernel |
| Modify | `config.rs` | New hyperparams for per-branch epsilon multipliers |
| Modify | `gpu_dqn_trainer.rs` | Per-branch bottleneck, Boltzmann for order/urgency |

---

### Task 1: Fix c51_alpha — Start at Target Value, No Ramp

The c51_alpha ramp (0→0.2→0.4→0.5) changes the loss landscape every epoch, causing 3 CUDA graph recaptures and destabilizing Q-values.

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

- [ ] **Step 1: Find c51_alpha ramp logic**

In `fused_training.rs`, search for `c51_alpha` assignment. Find where it ramps:
```rust
self.trainer.c51_alpha = (epoch as f32 / warmup_epochs as f32).min(c51_alpha_max);
```
Replace with fixed value from config:
```rust
self.trainer.c51_alpha = self.hyperparams.c51_alpha_max as f32;
```

- [ ] **Step 2: Set c51_alpha_max = 0.5 from epoch 0**

In `config.rs`, find `c51_alpha_max` default. It should already be 0.5. Verify the `c51_warmup_epochs` is set to 0 (no warmup):

```rust
c51_warmup_epochs: 0,  // was: 5 (ramped over 5 epochs)
```

- [ ] **Step 3: Verify no graph recapture from c51_alpha changes**

The RECAPTURE_DIAG logs showed c51_alpha changing each epoch (0→0.2→0.4→0.5). With fixed alpha, the graph captures ONCE and never recaptures for alpha changes.

- [ ] **Step 4: Compile + test**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib && cargo test -p ml --lib -- gradient_budget --nocapture`

- [ ] **Step 5: Commit**

```bash
git add -A && git commit -m "fix: fixed c51_alpha=0.5 from epoch 0 — eliminates 3 graph recaptures per fold"
```

---

### Task 2: Per-Branch Epsilon in GPU Action Selection

The magnitude branch needs 2× more exploration than direction. Order/urgency need Boltzmann instead of argmax.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_action_selector.rs`
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
- Modify: `crates/ml/src/trainers/dqn/config.rs`

- [ ] **Step 1: Add per-branch epsilon config fields**

In `config.rs`, add to DQNHyperparameters:
```rust
pub epsilon_direction_mult: f64,   // default 1.0
pub epsilon_magnitude_mult: f64,   // default 2.0 (needs more exploration)
pub epsilon_order_mult: f64,       // default 1.5
pub epsilon_urgency_mult: f64,     // default 1.5
```

- [ ] **Step 2: Update GPU action selector to accept 4 epsilons**

In `gpu_action_selector.rs`, change `select_actions_branching`:
```rust
pub fn select_actions_branching(
    &mut self,
    q_exposure: &CudaSlice<f32>,
    q_order: &CudaSlice<f32>,
    q_urgency: &CudaSlice<f32>,
    epsilon_dir: f32,
    epsilon_mag: f32,
    epsilon_ord: f32,
    epsilon_urg: f32,
    batch_size: usize,
) -> Result<CudaSlice<u32>, MLError>
```

- [ ] **Step 3: Update experience kernel epsilon usage**

In `experience_kernels.cu`, find the epsilon usage for each branch. Change from:
```c
float eps_dir = epsilon;
float eps_mag = epsilon;
float eps_other = eps_end + 0.1f * (epsilon - eps_end);
```
To reading per-branch epsilons from kernel arguments.

- [ ] **Step 4: Add Boltzmann to order and urgency branches**

Currently order/urgency use direct argmax. Change to Boltzmann:
```c
// Order branch: Boltzmann instead of argmax
float tau_ord = fmaxf(q_max_ord - q_min_ord, 0.01f);
// ... softmax with temperature tau_ord ...
```

- [ ] **Step 5: Wire per-branch epsilon in training loop**

In `training_loop.rs` where epsilon is passed to action selection:
```rust
let eps_dir = epsilon * self.hyperparams.epsilon_direction_mult as f32;
let eps_mag = epsilon * self.hyperparams.epsilon_magnitude_mult as f32;
let eps_ord = epsilon * self.hyperparams.epsilon_order_mult as f32;
let eps_urg = epsilon * self.hyperparams.epsilon_urgency_mult as f32;
```

- [ ] **Step 6: Update all callers of select_actions_branching**

Search for `select_actions_branching` across the codebase. Update each call site to pass 4 epsilons.

- [ ] **Step 7: Compile + test + commit**

```bash
SQLX_OFFLINE=true cargo check -p ml --lib
git commit -m "feat: per-branch epsilon — magnitude gets 2× exploration, Boltzmann for order/urgency"
```

---

### Task 3: IQN Quantile Spread as Anti-Collapse Bonus

Add `IQR = Q75 - Q25` per action from IQN. High IQR = uncertain = explore more.

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

- [ ] **Step 1: Add IQR computation kernel to IQN**

In `gpu_iqn_head.rs`, after the forward kernel runs, add a new kernel that computes IQR:
- Input: `save_q_online [B, num_quantiles, total_branch_actions]`
- Output: `iqr_buf [B, total_branch_actions]` (new buffer)
- Kernel: for each (b, a), sort quantiles, compute Q75 - Q25

Allocate `iqr_buf` in GpuIqnHead constructor.

- [ ] **Step 2: Export IQR to action selector**

Add method: `pub fn iqr_buf_ptr(&self) -> u64` returning the raw pointer.

In `fused_training.rs`, after IQN training step, pass IQR to the action selector via a new field.

- [ ] **Step 3: Add IQR bonus in action selection**

In `gpu_action_selector.rs`, before Boltzmann selection:
```rust
// Q_select(s,a) = Q_c51(s,a) + alpha * IQR(s,a)
```
Accept `iqr_ptr: u64` and `iqr_alpha: f32` as parameters. Add the bonus in the CUDA kernel.

- [ ] **Step 4: Add iqr_alpha config parameter**

```rust
pub iqr_exploration_alpha: f64,  // default 0.3, anneals to 0.05
```

- [ ] **Step 5: Compile + test + commit**

```bash
git commit -m "feat: IQN quantile spread (IQR) as anti-collapse exploration bonus"
```

---

### Task 4: Hindsight Magnitude Relabeling

After each trade, create 2 extra experiences with different magnitudes and scaled rewards.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_her.rs`
- Create: `crates/ml/src/cuda_pipeline/her_magnitude_relabel_kernel.cu` (if needed inline)

- [ ] **Step 1: Add magnitude relabeling to HER**

After the existing HER goal relabeling, add a magnitude relabeling pass:

For each completed trade `(state, direction, magnitude, reward)`:
1. Compute rewards for other magnitudes: `R_k = R_actual × (mag_k / actual_mag)`
2. Create 2 relabeled experiences: same state, same direction, different magnitude
3. Insert into PER with priority = `|R_relabeled - R_actual|`

- [ ] **Step 2: Add kernel for magnitude relabeling**

The kernel takes:
- `actions [B]` — factored actions
- `rewards [B]` — trade rewards
- `segment_complete [B]` — which samples are trade completions
- Output: `relabeled_actions [2*B]`, `relabeled_rewards [2*B]`

For each complete trade: generate the two alternative magnitude actions.

- [ ] **Step 3: Integrate with PER insert**

After the relabeling kernel, insert the new experiences into the GPU PER buffer alongside the originals.

- [ ] **Step 4: Compile + test + commit**

```bash
git commit -m "feat: hindsight magnitude relabeling — teaches magnitude what bigger/smaller earns"
```

---

### Task 5: Feed the Gradient Vaccine with Diverse Batches

The gradient vaccine infrastructure exists but `pending_vaccine_batch` is poorly populated. Ensure it contains magnitude-diverse samples.

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

- [ ] **Step 1: Improve vaccine batch sampling**

Currently the vaccine batch is a random PER sample. Change to stratified:
- Sample N/3 experiences with magnitude=Small
- Sample N/3 with magnitude=Half
- Sample N/3 with magnitude=Full
- If any stratum has insufficient samples, fill from other strata

This requires the PER buffer to track action metadata. If that's too complex, simply use the existing random sample but VERIFY it contains all 3 magnitudes — if not, resample up to 3 times.

- [ ] **Step 2: Increase vaccine frequency**

Change from every 10th step to every 5th step:
```rust
let vaccine_batch = if train_step_count % 5 == 0 {
```

- [ ] **Step 3: Compile + test + commit**

```bash
git commit -m "feat: diverse vaccine batches — prevents convergence to magnitude-homogeneous policy"
```

---

### Task 6: Ensemble Variance as Exploration Bonus at Inference

The 5 ensemble heads are trained but discarded at inference.

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

- [ ] **Step 1: Add ensemble Q-value export**

In `fused_training.rs`, at epoch boundary after training completes, compute ensemble variance for the current batch and store in a buffer.

Or simpler: during the existing `run_ensemble_step()`, export `ensemble_var_q_buf` pointer.

- [ ] **Step 2: Add variance bonus in action selection**

In `gpu_action_selector.rs`:
```rust
Q_select(s,a) = Q_head0(s,a) + beta * ensemble_std(s,a)
```

- [ ] **Step 3: Add ensemble_exploration_beta config**

```rust
pub ensemble_exploration_beta: f64,  // default 0.2
```

- [ ] **Step 4: Compile + test + commit**

```bash
git commit -m "feat: ensemble variance as exploration bonus — high disagreement = explore"
```

---

### Task 7: Per-Branch Bottleneck Width

The bottleneck compresses 42 market features → 2 dims for ALL branches. Magnitude can't see volatility.

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

- [ ] **Step 1: Add per-branch bottleneck config**

```rust
pub bottleneck_dim_direction: usize,  // default 4
pub bottleneck_dim_magnitude: usize,  // default 8
pub bottleneck_dim_order: usize,      // default 6
pub bottleneck_dim_urgency: usize,    // default 4
```

- [ ] **Step 2: Allocate per-branch bottleneck buffers**

Replace single `bn_hidden_buf [B, 2]` with 4 separate buffers:
```
bn_hidden_dir  [B, 4]
bn_hidden_mag  [B, 8]
bn_hidden_ord  [B, 6]
bn_hidden_urg  [B, 4]
```

- [ ] **Step 3: Add per-branch projection GEMMs**

4 small GEMMs in the forward pass:
```
bn_dir = tanh(states[B, 42] @ W_bn_dir[42, 4])
bn_mag = tanh(states[B, 42] @ W_bn_mag[42, 8])
...
```

Each branch head receives its own compressed input concatenated with portfolio features.

- [ ] **Step 4: Update backward pass**

Chain rule through per-branch bottleneck projections.

- [ ] **Step 5: Compile + test + commit**

```bash
git commit -m "feat: per-branch bottleneck — magnitude sees volatility, order sees spread"
```

---

### Task 8: Full Verification + Deploy

- [ ] **Step 1: Run all test suites**

```bash
SQLX_OFFLINE=true cargo test -p ml-core --lib              # 302 tests
SQLX_OFFLINE=true cargo test -p ml-dqn --lib -- --test-threads=1  # 287 tests
SQLX_OFFLINE=true cargo test -p ml --lib                    # 903 tests
```

- [ ] **Step 2: Run smoke test**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
  cargo test -p ml --lib -- smoke_tests::regression::test_no_hang_single_epoch --ignored --nocapture
```

- [ ] **Step 3: Deploy to H100**

```bash
git push origin main
./scripts/argo-train.sh dqn --baseline --epochs 5 --watch
```

### Success Criteria
1. **Magnitude diversity**: Half >15%, Full >10% (currently 4% each)
2. **Order diversity**: LimitMaker >5%, IoC >3% (currently <2%)
3. **Sharpe stability**: Epochs 2-5 within ±0.3 of epoch 1
4. **GPU timing**: fwd_bwd ≤ 330ms
5. **No graph recapture** from c51_alpha changes (was 3 per fold)
