# 9-Exposure Action Space + Dual Distributional Architecture — 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:** Expand the exposure branch from 5 to 9 levels (25% steps) and wire the existing IQN dual-head for risk-sensitive position sizing alongside C51 action selection.

**Architecture:** Two phases. Phase A: expand exposure from 5→9 across CUDA kernels, Rust types, monitoring, and configs. Phase B: wire the existing `GpuIqnHead` (already built, shares DQN trunk) to provide CVaR-based position scaling — C51 picks direction, IQN sizes the position. The IQN head already trains in `FusedTrainingCtx::run_full_step()` — we just need to consume its quantile output in action selection.

**Tech Stack:** CUDA C (kernels), Rust (types, trainer, hyperopt), TOML (config)

**Spec:** `docs/superpowers/specs/2026-03-23-9-exposure-action-space-design.md`

---

## File Structure

### Phase A: 9-Exposure Action Space

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml/src/cuda_pipeline/common_device_functions.cuh` | Modify | `DQN_NUM_ACTIONS 5→9`, `DQN_TOTAL_ACTIONS 45→81`, `BRANCH_EXPOSURE_ACTIONS 5→9` |
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | `exposure_idx_to_fraction()` formula, `flat_idx=2→4` in Q-gap filter |
| `crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu` | Modify | Hardcoded `5`→parameterized, `flat_idx=2→4` in Q-gap filter |
| `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` | Modify | `MAX_BRANCH_SIZE 5→9`, `BRANCH_0_SIZE 5→9` |
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | `branch_sizes: [5,3,3]→[9,3,3]` |
| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Modify | `branch_0_size: 5→9` default |
| `crates/ml-core/src/common/action.rs` | Modify | Add 4 `ExposureLevel` variants, update `from_index()`, `target_exposure()` |
| `crates/ml/src/trainers/dqn/trainer/metrics.rs` | Modify | `DIRECTION_LUT: [f32; 5]→[f32; 9]`, exposure parsing |
| `config/training/dqn-production.toml` | Modify | `branch_0_size = 9` |
| `config/training/dqn-smoketest.toml` | Modify | `branch_0_size = 9` |

### Phase B: Dual Distributional (C51 + IQN CVaR)

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Modify | Add `cvar_action_scale()` method that returns per-action risk scaling [0.5, 1.0] |
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | Q-gap filter uses IQN CVaR to scale exposure: `target_pos *= cvar_scale` |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Modify | Pass IQN quantile output to experience collector |
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | Accept CVaR scale buffer, pass to env_step kernel |

---

### Task 1: CUDA Header — Update Action Space Constants

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/common_device_functions.cuh:47-60`

- [ ] **Step 1: Update macros**

```cuda
#define DQN_NUM_ACTIONS    9      /* was 5 */
#define DQN_TOTAL_ACTIONS  81     /* was 45: 9 * 3 * 3 */
#define BRANCH_EXPOSURE_ACTIONS 9 /* was 5 */
#define BRANCH_MAX_ACTIONS      BRANCH_EXPOSURE_ACTIONS
```

- [ ] **Step 2: Update `diversity_entropy()` max_entropy comment**

The function uses `DQN_NUM_ACTIONS` dynamically, so code auto-updates. Just fix the comment:
```cuda
float max_entropy = log2f((float)DQN_NUM_ACTIONS);  /* 3.17 for 9 actions */
```

- [ ] **Step 3: Update `curiosity_inference()` 3-class mapping**

```cuda
/* Map 9 exposure levels to 3 categories for curiosity: Short(0-2), Flat(3-5), Long(6-8) */
int category = exposure_idx / 3;
if (category >= 3) category = 2;
```

Old code had `exposure_idx / 2` for 5→3 mapping. With 9 actions: `exposure_idx / 3` gives {0,0,0,1,1,1,2,2,2}.

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

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

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/common_device_functions.cuh
git commit -m "feat: expand exposure action space 5→9 (CUDA header macros)"
```

---

### Task 2: Experience Kernel — Exposure Fraction Formula + Q-Gap Flat Index

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu:114-123`
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` (Q-gap filter, ~line 260)

- [ ] **Step 1: Replace switch-case with formula**

```cuda
__device__ __forceinline__ float exposure_idx_to_fraction(int idx) {
    /* 9 levels: -1.0, -0.75, -0.50, -0.25, 0.0, +0.25, +0.50, +0.75, +1.0 */
    return -1.0f + (float)idx * 0.25f;
}
```

- [ ] **Step 2: Generalize Q-gap flat_idx**

In `experience_action_select`, replace hardcoded `flat_idx = 2`:

```cuda
int flat_idx = b0_size / 2; /* 9/2 = 4 for 9-action, was hardcoded 2 */
```

Apply in BOTH branching and flat mode paths.

- [ ] **Step 3: Update `backtest_env_kernel.cu`**

Check `action_to_exposure()` in the backtest kernel — if it has a separate switch-case, update to match.

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

```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/backtest_env_kernel.cu
git commit -m "feat: 9-exposure fraction formula + generalized Q-gap flat_idx"
```

---

### Task 3: Epsilon-Greedy Kernel — Remove Hardcoded 5

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

- [ ] **Step 1: Update `branching_action_select`**

Replace all hardcoded `5` with `b0_size` parameter. The kernel already receives `b0_size` via batch_size position — add explicit parameter if needed.

Current: `exposure = (int)(gpu_random(&rng) * 5.0f); if (exposure >= 5) exposure = 4;`
New: `exposure = (int)(gpu_random(&rng) * (float)b0_size); if (exposure >= b0_size) exposure = b0_size - 1;`

- [ ] **Step 2: Update Q-gap flat_idx in all 3 kernels**

Replace `qe[2]` with `qe[b0_size / 2]` in conviction filter for:
- `epsilon_greedy_select`
- `epsilon_greedy_routed`
- `branching_action_select`

- [ ] **Step 3: Update Rust `gpu_action_selector.rs`**

If `branching_action_select` kernel signature changed (new `b0_size` param), update the launch args.

- [ ] **Step 4: Compile check + run GPU action selector test**

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

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu crates/ml/src/cuda_pipeline/gpu_action_selector.rs
git commit -m "feat: epsilon-greedy kernel parameterized for 9-action exposure"
```

---

### Task 4: C51 Loss Kernel — MAX_BRANCH_SIZE

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu:29,51`

- [ ] **Step 1: Update branch size defaults**

```cuda
#ifndef BRANCH_0_SIZE
#define BRANCH_0_SIZE 9    /* was 5 */
#endif
...
#define MAX_BRANCH_SIZE 9  /* was 5: max(BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE) */
```

Note: these are defaults. NVRTC injection already overrides via `gpu_dqn_trainer.rs` defines.

- [ ] **Step 2: Verify shared memory sizing**

`shmem_adv[MAX_BRANCH_SIZE * NUM_ATOMS]` = `9 * 101 = 909 floats = 3636 bytes`. Under 4KB — no occupancy impact.

`eq_per_action[MAX_BRANCH_SIZE]` = `9 floats = 36 bytes`. Fine.

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

```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
git commit -m "feat: C51 loss kernel MAX_BRANCH_SIZE 5→9"
```

---

### Task 5: Rust Core Types — ExposureLevel Enum

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

- [ ] **Step 1: Add 4 new ExposureLevel variants**

Add `Short75`, `Short25`, `Long25`, `Long75` to the enum. Update:
- `from_index()`: map indices 0-8 to the 9 variants
- `target_exposure()`: return the fraction (-1.0 to +1.0 in 0.25 steps)
- `to_trading_action()`: Short75/Short25 → Sell, Long25/Long75 → Buy
- `is_buy()`, `is_sell()`: include new variants

- [ ] **Step 2: Update FactoredAction bounds**

`from_index()` bound check: `45` → `81`
`to_index()` max value: `80` (was `44`)

Note: the stride formula `exposure * 9 + order * 3 + urgency` is UNCHANGED (order*urgency = 3*3 = 9 is coincidentally the same stride).

- [ ] **Step 3: Run core tests**

```bash
SQLX_OFFLINE=true cargo test -p ml-core --lib -- --test-threads=1
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml-core/src/common/action.rs
git commit -m "feat: ExposureLevel 5→9 variants (25% step increments)"
```

---

### Task 6: GPU Pipeline — branch_sizes + Config

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (NVRTC injection)
- Modify: `config/training/dqn-production.toml`
- Modify: `config/training/dqn-smoketest.toml`

- [ ] **Step 1: Update gpu_experience_collector.rs branch_sizes**

Find all `[5, 3, 3]` or `b0 = 5` and change to `[9, 3, 3]` or `b0 = 9`. These should derive from config, not hardcode.

- [ ] **Step 2: Update gpu_iqn_head.rs default**

```rust
branch_0_size: 9,  // was 5
```

- [ ] **Step 3: Update gpu_dqn_trainer.rs NVRTC injection**

The NVRTC `#define BRANCH_0_SIZE {b0}` should already use `config.branch_0_size` — verify it does and that config is set to 9.

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

`dqn-production.toml`:
```toml
[branching]
branch_0_size = 9
```

`dqn-smoketest.toml`:
```toml
branch_0_size = 9
```

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

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

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/cuda_pipeline/gpu_iqn_head.rs crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs config/training/
git commit -m "feat: branch_0_size 5→9 across GPU pipeline + TOML configs"
```

---

### Task 7: Monitoring + Metrics — Array Expansion

**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/metrics.rs`
- Modify: `crates/ml/src/trainers/dqn/monitoring.rs` (if it has fixed arrays)

- [ ] **Step 1: Update DIRECTION_LUT**

```rust
const DIRECTION_LUT: [f32; 9] = [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0];
```

- [ ] **Step 2: Update all fixed-size arrays**

Search for `[_; 5]`, `[usize; 5]`, `[f64; 5]` in monitoring.rs and metrics.rs. Change to `[_; 9]`.

- [ ] **Step 3: Update exposure label arrays**

```rust
const EXPOSURE_NAMES: [&str; 9] = [
    "S100", "S75", "S50", "S25", "Flat", "L25", "L50", "L75", "L100"
];
```

- [ ] **Step 4: Update all `45` → `81` assertions and bounds**

Search for `45` in test files: `dqn_action_collapse_fix_test.rs`, `ensemble_hyperopt_tests.rs`, `smoke_test_real_data.rs`.

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

```bash
SQLX_OFFLINE=true cargo test -p ml -p ml-core -p ml-dqn --lib -- --test-threads=1
```

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/trainers/dqn/trainer/metrics.rs crates/ml/src/trainers/dqn/monitoring.rs crates/ml/tests/
git commit -m "feat: monitoring arrays 5→9, action labels, test assertions 45→81"
```

---

### Task 8: Phase B — IQN CVaR Position Scaling

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

- [ ] **Step 1: Add `compute_cvar_scales()` to GpuIqnHead**

Given quantile Q-values for each action, compute CVaR at α=0.05 (5th percentile worst-case):

```rust
/// Compute CVaR-based position scaling from IQN quantiles.
/// Returns a GPU buffer [batch_size] with scaling factors in [0.25, 1.0].
/// Low CVaR (high tail risk) → small scale. High CVaR (safe) → full scale.
pub fn compute_cvar_scales(
    &self,
    quantile_q: &CudaSlice<f32>,  // [batch, num_quantiles, num_actions]
    actions: &CudaSlice<i32>,      // [batch] selected exposure indices
    batch_size: usize,
) -> Result<CudaSlice<f32>, MLError>
```

The CVaR at 5% = mean of the lowest 5% of quantile samples. Normalize to [0.25, 1.0]:
- CVaR ≥ 0 (expected profit even in worst case) → scale = 1.0
- CVaR < 0 (expected loss in worst case) → scale = max(0.25, 1.0 + CVaR)

- [ ] **Step 2: Write CUDA kernel for CVaR computation**

Add to `gpu_iqn_head.rs` inline kernel source:

```cuda
extern "C" __global__ void compute_cvar_scale(
    const float* quantile_q,  // [B, N_TAU, N_ACTIONS]
    const int* actions,       // [B]
    float* cvar_scales,       // [B] output
    int B, int N_TAU, int N_ACTIONS, float alpha
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= B) return;
    int a = actions[i];
    // Gather quantiles for selected action
    float sum = 0.0f;
    int count = max(1, (int)(alpha * N_TAU));
    // Quantiles are sorted (IQN property)
    for (int t = 0; t < count; t++) {
        sum += quantile_q[i * N_TAU * N_ACTIONS + t * N_ACTIONS + a];
    }
    float cvar = sum / (float)count;
    // Scale: [0.25, 1.0]
    float scale = (cvar >= 0.0f) ? 1.0f : fmaxf(0.25f, 1.0f + cvar * 0.5f);
    cvar_scales[i] = scale;
}
```

- [ ] **Step 3: Wire CVaR scales into experience collector**

In `gpu_experience_collector.rs`, accept an optional `cvar_scales: Option<&CudaSlice<f32>>` buffer. Pass to `experience_env_step` kernel.

In the env_step kernel, after computing `target_position`:
```cuda
/* Apply CVaR risk scaling: reduce position when tail risk is high */
if (cvar_scales != NULL) {
    target_position *= cvar_scales[i];
}
```

- [ ] **Step 4: Wire in fused_training.rs**

After the IQN forward pass in `run_full_step()`, call `compute_cvar_scales()` with the IQN quantile output and pass the result to the experience collector.

- [ ] **Step 5: Add `iqn_cvar_alpha` to DQNHyperparameters**

```rust
pub iqn_cvar_alpha: f64,  // CVaR percentile (default 0.05 = 5th percentile worst case)
```

Add to TOML configs and hyperopt search space as a searchable dimension.

- [ ] **Step 6: Compile + smoke test**

```bash
SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5
SQLX_OFFLINE=true cargo test -p ml --test dqn_training_smoke_test -- --nocapture --test-threads=1
```

- [ ] **Step 7: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_iqn_head.rs crates/ml/src/cuda_pipeline/experience_kernels.cu crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/config.rs config/training/
git commit -m "feat: IQN CVaR position scaling — C51 picks direction, IQN sizes risk"
```

---

### Task 9: Integration Test + H100 Validation

- [ ] **Step 1: Run all lib tests**

```bash
SQLX_OFFLINE=true cargo test -p ml -p ml-core -p ml-dqn --lib -- --test-threads=1
```

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

```bash
SQLX_OFFLINE=true cargo test -p ml --test dqn_action_collapse_fix_test --test ensemble_hyperopt_tests -- --test-threads=1
```

- [ ] **Step 3: Local hyperopt smoke test**

```bash
SQLX_OFFLINE=true timeout 120 ./target/release/examples/hyperopt_baseline_rl \
  --model dqn --phase fast --trials 3 --n-initial 2 --epochs 5 \
  --data-dir test_data/futures-baseline --symbol ES.FUT \
  --output /tmp/9action_dual_test.json --seed 42
```

Verify: 81 total actions in output, CVaR scales logged, epoch time ≤ 5s.

- [ ] **Step 4: Push and submit H100 run**

```bash
git push origin main
argo submit --from workflowtemplate/compile-and-train \
  -p commit-sha=HEAD -p hyperopt-trials=20 -p hyperopt-epochs=50 ...
```

- [ ] **Step 5: Monitor first trial**

Check: epoch time ~3-4s, Q-gap > 0 by epoch 10, 81 factored actions, CVaR scales in [0.25, 1.0].
