plan: 9-exposure action space + dual C51/IQN distributional architecture

Phase A: expand exposure 5→9 (25% steps, 81 factored actions) Phase B: wire existing GpuIqnHead for CVaR position scaling — C51 picks direction, IQN sizes the risk. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 09:03:19 +01:00
parent 93dc7f811e
commit 98afee0a51
1 changed files with 470 additions and 0 deletions
--- a/docs/superpowers/plans/2026-03-23-9-exposure-dual-distributional.md
+++ b/docs/superpowers/plans/2026-03-23-9-exposure-dual-distributional.md
@@ -0,0 +1,470 @@
+# 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].