diff --git a/docs/superpowers/plans/2026-03-23-9-exposure-dual-distributional.md b/docs/superpowers/plans/2026-03-23-9-exposure-dual-distributional.md new file mode 100644 index 000000000..29f968275 --- /dev/null +++ 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, // [batch, num_quantiles, num_actions] + actions: &CudaSlice, // [batch] selected exposure indices + batch_size: usize, +) -> Result, 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>` 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].