diff --git a/docs/superpowers/plans/2026-04-11-abi-complete.md b/docs/superpowers/plans/2026-04-11-abi-complete.md new file mode 100644 index 000000000..9ea915c72 --- /dev/null +++ b/docs/superpowers/plans/2026-04-11-abi-complete.md @@ -0,0 +1,382 @@ +# 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, + q_order: &CudaSlice, + q_urgency: &CudaSlice, + epsilon_dir: f32, + epsilon_mag: f32, + epsilon_ord: f32, + epsilon_urg: f32, + batch_size: usize, +) -> Result, 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)