foxhunt/docs/superpowers/plans/2026-04-11-abi-complete.md

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:

self.trainer.c51_alpha = (epoch as f32 / warmup_epochs as f32).min(c51_alpha_max);

Replace with fixed value from config:

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):

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

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:

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:

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:

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:

// 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:

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

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:

// 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

pub iqr_exploration_alpha: f64,  // default 0.3, anneals to 0.05

• Step 5: Compile + test + commit

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

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:

let vaccine_batch = if train_step_count % 5 == 0 {

• Step 3: Compile + test + commit

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:

Q_select(s,a) = Q_head0(s,a) + beta * ensemble_std(s,a)

• Step 3: Add ensemble_exploration_beta config

pub ensemble_exploration_beta: f64,  // default 0.2

• Step 4: Compile + test + commit

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

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

git commit -m "feat: per-branch bottleneck — magnitude sees volatility, order sees spread"

---

Task 8: Full Verification + Deploy

• Step 1: Run all test suites

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

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

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)