# ABI Foundation 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 degenerate 9-level action space, enable PopArt reward normalization, clear PER between folds, and couple PopArt with tau reset — the structural prerequisites for Autonomous Branch Intelligence.

**Architecture:** Collapse 3 Flat variants to 1 (9→7 exposure levels), enable the already-implemented PopArt reward normalization, add tau reset on reward distribution shift, and ensure PER buffer doesn't leak across walk-forward folds.

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

---

## File Structure

| Action | File | Purpose |
|--------|------|---------|
| Modify | `crates/ml-core/src/action_space.rs` | 7-level ExposureLevel enum + updated from_index/target_exposure |
| Modify | `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Remove 30% Flat floor, update action encoding |
| Modify | `crates/ml/src/trainers/dqn/config.rs` | Enable PopArt by default |
| Modify | `crates/ml/src/trainers/dqn/fused_training.rs` | PopArt enable, tau coupling |
| Modify | `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Tau reset on PopArt variance shift |
| Modify | `crates/ml/src/trainers/dqn/trainer/mod.rs` | Verify PER clear at fold boundary |
| Modify | `crates/ml-dqn/src/evaluation/engine.rs` | Update tests for 7-level exposure |

---

### Task 1: Collapse ExposureLevel from 9 to 7 Levels

**Files:**
- Modify: `crates/ml-core/src/action_space.rs`

- [ ] **Step 1: Update ExposureLevel enum**

Replace the current 9-variant enum with 7 variants. The 3 Flat variants (Short25=dir1×mag0, Flat=dir1×mag1, Long25=dir1×mag2) all map to 0.0 exposure — collapse them into a single Flat.

In `crates/ml-core/src/action_space.rs`, replace the ExposureLevel enum (lines 68-84):

```rust
/// 7-level exposure: Short×{Small,Half,Full}, Flat, Long×{Small,Half,Full}
/// When direction=Flat, magnitude is masked (not selected).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[repr(u8)]
pub enum ExposureLevel {
    ShortSmall  = 0,  // dir=Short, mag=Small → -0.25
    ShortHalf   = 1,  // dir=Short, mag=Half  → -0.50
    ShortFull   = 2,  // dir=Short, mag=Full  → -1.00
    Flat        = 3,  // dir=Flat,  mag=N/A   →  0.00
    LongSmall   = 4,  // dir=Long,  mag=Small → +0.25
    LongHalf    = 5,  // dir=Long,  mag=Half  → +0.50
    LongFull    = 6,  // dir=Long,  mag=Full  → +1.00
}
```

- [ ] **Step 2: Update target_exposure()**

Replace the `target_exposure()` method:

```rust
impl ExposureLevel {
    pub fn target_exposure(&self) -> f64 {
        match self {
            Self::ShortSmall => -0.25,
            Self::ShortHalf  => -0.50,
            Self::ShortFull  => -1.00,
            Self::Flat       =>  0.00,
            Self::LongSmall  =>  0.25,
            Self::LongHalf   =>  0.50,
            Self::LongFull   =>  1.00,
        }
    }
}
```

- [ ] **Step 3: Update from_index()**

```rust
impl ExposureLevel {
    pub fn from_index(idx: usize) -> Result<Self, MLError> {
        match idx {
            0 => Ok(Self::ShortSmall),
            1 => Ok(Self::ShortHalf),
            2 => Ok(Self::ShortFull),
            3 => Ok(Self::Flat),
            4 => Ok(Self::LongSmall),
            5 => Ok(Self::LongHalf),
            6 => Ok(Self::LongFull),
            _ => Err(MLError::InvalidInput(format!("ExposureLevel index {idx} out of range 0..6"))),
        }
    }
}
```

- [ ] **Step 4: Update is_buy/is_sell/is_hold**

```rust
impl ExposureLevel {
    pub const fn is_buy(&self) -> bool {
        matches!(self, Self::LongSmall | Self::LongHalf | Self::LongFull)
    }
    pub const fn is_sell(&self) -> bool {
        matches!(self, Self::ShortSmall | Self::ShortHalf | Self::ShortFull)
    }
    pub const fn is_hold(&self) -> bool {
        matches!(self, Self::Flat)
    }
}
```

- [ ] **Step 5: Update direction() and magnitude() helpers**

```rust
impl ExposureLevel {
    /// Direction component: 0=Short, 1=Flat, 2=Long
    pub const fn direction(&self) -> u8 {
        match self {
            Self::ShortSmall | Self::ShortHalf | Self::ShortFull => 0,
            Self::Flat => 1,
            Self::LongSmall | Self::LongHalf | Self::LongFull => 2,
        }
    }

    /// Magnitude component: 0=Small(0.25), 1=Half(0.50), 2=Full(1.00), 3=N/A (Flat)
    pub const fn magnitude(&self) -> u8 {
        match self {
            Self::ShortSmall | Self::LongSmall => 0,
            Self::ShortHalf | Self::LongHalf => 1,
            Self::ShortFull | Self::LongFull => 2,
            Self::Flat => 3, // no magnitude for Flat
        }
    }

    /// Construct from direction + magnitude.
    /// When direction=Flat (1), magnitude is ignored → returns Flat.
    pub fn from_dir_mag(dir: usize, mag: usize) -> Result<Self, MLError> {
        match (dir, mag) {
            (0, 0) => Ok(Self::ShortSmall),
            (0, 1) => Ok(Self::ShortHalf),
            (0, 2) => Ok(Self::ShortFull),
            (1, _) => Ok(Self::Flat), // magnitude ignored for Flat
            (2, 0) => Ok(Self::LongSmall),
            (2, 1) => Ok(Self::LongHalf),
            (2, 2) => Ok(Self::LongFull),
            _ => Err(MLError::InvalidInput(format!("Invalid dir={dir} mag={mag}"))),
        }
    }
}
```

- [ ] **Step 6: Update FactoredAction encoding**

The FactoredAction `to_index()` and `from_index()` must account for 7 exposure levels:

```rust
impl FactoredAction {
    pub fn to_index(&self) -> usize {
        let exp_idx = self.exposure as usize;
        let ord_idx = self.order as usize;
        let urg_idx = self.urgency as usize;
        exp_idx * 9 + ord_idx * 3 + urg_idx  // 7 × 3 × 3 = 63 total
    }

    pub fn from_index(idx: usize) -> Result<Self, MLError> {
        let exp_idx = idx / 9;
        let remainder = idx % 9;
        let ord_idx = remainder / 3;
        let urg_idx = remainder % 3;
        Ok(Self {
            exposure: ExposureLevel::from_index(exp_idx)?,
            order: OrderType::from_index(ord_idx)?,
            urgency: Urgency::from_index(urg_idx)?,
        })
    }
}
```

- [ ] **Step 7: Fix all tests in action_space.rs**

Update test assertions for 7-level exposure. Key tests to fix:
- `test_target_exposure_values` — update expected values
- `test_exposure_level_indices` — 0..6 not 0..8
- `test_factored_action_roundtrip` — 63 total actions not 81

- [ ] **Step 8: Compile check**

Run: `SQLX_OFFLINE=true cargo check --workspace 2>&1 | grep 'error\[' | head -20`

This will show ALL downstream breakage from the enum change. Fix each error — they'll be in:
- `experience_kernels.cu` (action encoding constants)
- `gpu_action_selector.rs` (branch sizes)
- `dqn.rs` (epsilon-greedy action selection)
- `regime_conditional.rs` (action routing)
- `evaluation/engine.rs` (backtesting)

- [ ] **Step 9: Commit**

```bash
git add -A
git commit -m "refactor: collapse 9-level to 7-level ExposureLevel — eliminate degenerate Flat variants"
```

---

### Task 2: Remove 30% Flat Floor in Experience Collection

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

- [ ] **Step 1: Find and remove the 30% Flat floor**

In `experience_kernels.cu`, find the block around line 830:
```c
if (lcg_random(&rng) < 0.30f) {
    dir_idx = 1; /* Flat floor: 30% unconditional */
}
```

Remove this entire `if` block. The ABI exploration mechanisms (per-branch epsilon, IQN anti-collapse) replace it.

- [ ] **Step 2: Also remove the forced magnitude when Flat**

Find around line 854:
```c
if (dir_idx == 1) {
    mag_idx = 1;  /* Default to Half — doesn't affect position (0 × anything = 0) */
}
```

Replace with:
```c
if (dir_idx == 1) {
    mag_idx = 0;  /* Flat: magnitude is masked, set to 0 for deterministic encoding */
}
```

With the 7-level system, when direction=Flat, we encode exposure=3 (Flat), and magnitude is not part of the encoding.

- [ ] **Step 3: Update action encoding constants**

Update the branch size constants in the kernel to reflect 7-level exposure:
- `b0_size` (exposure/direction): stays 3 (Short/Flat/Long)
- `b1_size` (magnitude): stays 3 (Small/Half/Full) BUT is MASKED when direction=Flat
- Total factored actions: changes from 81 to 63 (or stays 81 with masking)

- [ ] **Step 4: Compile check**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | grep 'error\[' | head -10`

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "fix: remove 30% Flat floor + forced magnitude — ABI exploration replaces it"
```

---

### Task 3: Enable PopArt Reward Normalization

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

- [ ] **Step 1: Enable PopArt by default**

In `crates/ml/src/trainers/dqn/config.rs`, find line ~1358:
```rust
popart_enabled: false,
```
Change to:
```rust
popart_enabled: true,
```

- [ ] **Step 2: Also enable in the TOML configs**

In `config/training/dqn-production.toml` and `config/training/dqn-localdev.toml`, add:
```toml
popart_enabled = true
```

In `config/training/dqn-smoketest.toml`, add:
```toml
popart_enabled = true
```

- [ ] **Step 3: Verify PopArt is wired in fused_training.rs**

Read `crates/ml/src/trainers/dqn/fused_training.rs` line ~874. The PopArt code should already be:
```rust
if self.popart_enabled {
    let bs = self.batch_size;
    if let Err(e) = self.trainer.normalize_rewards_popart_inplace(bs) {
        tracing::warn!("PopArt normalize failed (non-fatal): {e}");
    }
}
```

This is inside `submit_forward_ops_main()` which is part of the mega-graph. PopArt normalization runs before C51 loss computation — normalizing rewards to zero-mean unit-variance prevents C51 from favoring low-variance (Flat) actions.

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

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

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/trainers/dqn/config.rs config/training/dqn-*.toml
git commit -m "feat: enable PopArt reward normalization — prevents C51 variance bias toward Flat"
```

---

### Task 4: Couple PopArt with Tau Reset

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

- [ ] **Step 1: Add PopArt variance tracking to FusedTrainingCtx**

In `crates/ml/src/trainers/dqn/fused_training.rs`, add field to `FusedTrainingCtx`:

```rust
/// PopArt running variance from previous epoch (for tau reset detection).
prev_popart_var: f32,
```

Initialize to `0.0` in the constructor.

- [ ] **Step 2: Add method to read PopArt variance**

Add to `FusedTrainingCtx`:

```rust
/// Read PopArt running variance from GPU (single scalar DtoH, cold path).
pub(crate) fn read_popart_variance(&self) -> f32 {
    if !self.popart_enabled { return 0.0; }
    let mut var = [0.0_f32];
    if self.trainer.read_popart_var(&mut var).is_ok() {
        var[0]
    } else {
        0.0
    }
}

/// Check if PopArt variance shifted significantly, suggesting tau reset.
pub(crate) fn should_reset_tau(&mut self) -> bool {
    let current_var = self.read_popart_variance();
    if self.prev_popart_var > 0.0 && current_var > 0.0 {
        let ratio = current_var / self.prev_popart_var;
        let shifted = ratio > 2.0 || ratio < 0.5; // >2× change in variance
        self.prev_popart_var = current_var;
        shifted
    } else {
        self.prev_popart_var = current_var;
        false
    }
}
```

- [ ] **Step 3: Add read_popart_var to GpuDqnTrainer**

In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`, add:

```rust
/// Read PopArt running variance from GPU (single scalar DtoH).
pub fn read_popart_var(&self, out: &mut [f32; 1]) -> Result<(), MLError> {
    self.stream.memcpy_dtoh(&self.popart_var, out)
        .map_err(|e| MLError::ModelError(format!("PopArt var readback: {e}")))?;
    Ok(())
}
```

- [ ] **Step 4: Add tau reset in training loop**

In `crates/ml/src/trainers/dqn/trainer/training_loop.rs`, in the epoch boundary section (after process_epoch_boundary), add:

```rust
// PopArt-coupled tau reset: if reward distribution shifts, reset tau to initial value
if let Some(ref mut fused) = self.fused_ctx {
    if fused.should_reset_tau() {
        self.adaptive_tau = self.hyperparams.tau;
        tracing::info!(
            tau = self.adaptive_tau,
            "PopArt variance shift detected — tau reset to initial value"
        );
    }
}
```

- [ ] **Step 5: Compile check**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib`

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat: couple PopArt variance with tau reset — stabilizes target network on reward shift"
```

---

### Task 5: Verify PER Clear Between Folds

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

- [ ] **Step 1: Verify PER clear exists in reset_for_fold**

Read `crates/ml/src/trainers/dqn/trainer/mod.rs` at the `reset_for_fold()` method. Verify it calls `clear_replay_buffer()`. The exploration found this IS already called — verify and add logging.

- [ ] **Step 2: Add explicit best_sharpe reset**

In `reset_for_fold()`, ensure `best_sharpe` is reset:

```rust
// Reset best Sharpe tracking for new fold
self.best_sharpe = f64::NEG_INFINITY;
self.best_epoch = 0;
self.best_val_loss = f64::INFINITY;
```

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

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

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/trainers/dqn/trainer/mod.rs
git commit -m "fix: ensure PER buffer + best_sharpe reset at fold boundary"
```

---

### Task 6: Full Smoke Test + Deploy

**Files:**
- No modifications — verification only

- [ ] **Step 1: Run ml-dqn tests**

Run: `SQLX_OFFLINE=true cargo test -p ml-dqn --lib -- --test-threads=1 2>&1 | tail -5`
Expected: All tests pass (may need to fix tests broken by 7-level enum change).

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

Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget --nocapture`
Expected: 7/7 pass.

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

Run: `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`
Expected: Pass with PopArt enabled, 7-level actions.

- [ ] **Step 4: Push and deploy to H100**

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

Watch for:
- PopArt normalization active in logs
- Action distribution: Flat should be <40% (was 82%)
- Magnitude diversity: all 3 levels >10% each
- Sharpe stability across epochs
