plan: ABI Complete — per-branch epsilon, IQN anti-collapse, hindsight relabeling, vaccine, ensemble, bottleneck (8 tasks)
Unified plan combining Sub-plans B+C + novel findings: 1. Fix c51_alpha ramp (eliminate 3 graph recaptures) 2. Per-branch epsilon (magnitude 2×, Boltzmann for order/urgency) 3. IQN quantile spread as anti-collapse bonus 4. Hindsight magnitude relabeling 5. Gradient vaccine with diverse batches 6. Ensemble variance at inference 7. Per-branch bottleneck width 8. Full verification + deploy Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
382
docs/superpowers/plans/2026-04-11-abi-complete.md
Normal file
382
docs/superpowers/plans/2026-04-11-abi-complete.md
Normal file
@@ -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<f32>,
|
||||
q_order: &CudaSlice<f32>,
|
||||
q_urgency: &CudaSlice<f32>,
|
||||
epsilon_dir: f32,
|
||||
epsilon_mag: f32,
|
||||
epsilon_ord: f32,
|
||||
epsilon_urg: f32,
|
||||
batch_size: usize,
|
||||
) -> Result<CudaSlice<u32>, 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)
|
||||
Reference in New Issue
Block a user