# Autonomous Branch Intelligence (ABI) — Design Spec

## Problem

The 4-branch factored DQN collapses to 82% Flat/Small exposures. The model learns (loss drops, gradients stable) but its policy is degenerate — it avoids risk entirely. Root causes:

1. **Degenerate action space**: 3 of 9 exposure levels map to position=0.0 (Flat×Small = Flat×Half = Flat×Full)
2. **C51 variance bias**: zero-variance Flat distributions get highest softmax Q-values
3. **Sparse reward**: Flat positions get zero reward per timestep (no signal to differentiate)
4. **Shared gradient**: all 4 branches get equal 1/D weight from a single scalar reward
5. **Weak exploration**: magnitude branch has same epsilon as direction despite needing more
6. **Missing connectors**: PopArt (disabled), Gradient Vaccine (starved), Ensemble inference (discarded)

## Solution: Autonomous Branch Intelligence

Each branch becomes an autonomous learner with its own reward signal, exploration rate, gradient flow, and curiosity — not 4 copies of the same Q-learning sharing one scalar.

---

## Component 1: Structural — Collapse Degenerate Action Space

**Change**: 9-level → 7-level exposure

Remove the 3 degenerate Flat variants. When direction=Flat, magnitude is meaningless (0.0 × anything = 0.0). Collapse to:

| Index | Direction | Magnitude | Exposure |
|-------|-----------|-----------|----------|
| 0 | Short | Small | -0.25 |
| 1 | Short | Half | -0.50 |
| 2 | Short | Full | -1.00 |
| 3 | Flat | — | 0.00 |
| 4 | Long | Small | +0.25 |
| 5 | Long | Half | +0.50 |
| 6 | Long | Full | +1.00 |

**Impact on branches**: Direction becomes 3 actions (Short/Flat/Long). When direction=Flat, magnitude branch output is MASKED (no gradient, no selection). When direction=Short or Long, magnitude branch selects from 3 levels.

**Files**: `crates/ml-core/src/action_space.rs`, `experience_kernels.cu` (action encoding), `gpu_action_selector.rs`

Remove the 30% Flat floor in `experience_kernels.cu:830-831` — the following mechanisms replace it.

---

## Component 2: Per-Branch Reward Decomposition

Replace single scalar reward with 4 branch-specific signals computed in `experience_kernels.cu`:

**Direction reward**: `R_dir = sign(position) × sign(Δprice) × |vol_normalized_return|`
- Did the agent bet the right direction? Pure directional signal.

**Magnitude reward**: `R_mag = |actual_PnL| / optimal_PnL(realized_vol)`
- Was position sized correctly for observed volatility?
- `optimal_PnL = kelly_fraction(realized_sharpe) × max_position × |Δprice|`
- Ratio > 1.0 means oversized, < 1.0 means undersized, = 1.0 is optimal

**Order reward**: `R_ord = (midprice - fill_price) / spread × sign(position)`
- Did the order type achieve good execution quality?
- Positive when fill beats midprice, negative when worse

**Urgency reward**: `R_urg = PnL_at_exit / max_unrealized_PnL`
- Did the agent exit at the right time?
- 1.0 = captured peak PnL, 0.5 = exited at half the peak, negative = held through reversal

**Loss computation**: Each branch optimizes its own reward via branch-specific C51 target distribution. The aggregate Q-value remains `V(s) + (1/4) Σ A_d(s, a_d)` for action selection, but gradients flow through branch-specific losses.

**Files**: `experience_kernels.cu` (reward kernel), `c51_loss_kernel.cu` (per-branch targets), `gpu_dqn_trainer.rs` (4-component reward buffer)

---

## Component 3: Hindsight Action Relabeling

After each trade completes, compute optimal magnitude in hindsight and create relabeled experiences:

**Mechanism**: For each completed trade with `(direction, magnitude, reward)`:
1. Compute what reward would have been at each magnitude level:
   - `R_small = R_actual × (0.25 / actual_mag_fraction)`
   - `R_half = R_actual × (0.50 / actual_mag_fraction)`
   - `R_full = R_actual × (1.00 / actual_mag_fraction)`
2. Create 2 additional experiences with same state but different magnitude + scaled reward
3. Insert into PER buffer with priority = |R_hindsight - R_actual| (high priority when hindsight differs)

**Why this works**: The magnitude branch gets 3× more training signal — every trade teaches it what ALL magnitudes would have produced. Breaks the sparse reward problem for magnitude.

**Files**: `gpu_her.rs` (extend relabeling), new kernel `her_action_relabel_kernel.cu`

---

## Component 4: Position-Aware Temporal Exploration

Epsilon varies by position state and episode phase:

```
epsilon(state, phase) = base_epsilon × position_mult × phase_mult × branch_mult

position_mult:
  flat (no position):     2.0   (cheap to explore)
  small position (≤0.25): 1.0   (moderate risk)
  large position (>0.5):  0.3   (expensive to explore)

phase_mult:
  early (0-25%):   1.5   (explore: find entries)
  middle (25-75%): 0.5   (exploit: manage position)
  late (75-100%):  1.2   (explore: find exits)

branch_mult:
  direction:  1.0
  magnitude:  2.0   (needs more exploration — collapsed to Small)
  order:      1.5   (needs more than current 5.5%)
  urgency:    1.5
```

**Implementation**: Pass `[4]` epsilon vector to GPU action selector instead of scalar. Compute per-step in experience kernel from portfolio state.

**Files**: `experience_kernels.cu` (epsilon computation), `gpu_action_selector.rs` (accept vector epsilon)

---

## Component 5: IQN Quantile Spread as Anti-Collapse

Use IQN's distributional uncertainty to PREVENT Flat/Small collapse:

**Mechanism**: Before Boltzmann action selection, add exploration bonus:
```
Q_select(s, a) = Q_c51(s, a) + alpha × IQR(s, a)
```
Where `IQR(s, a) = quantile_75(s, a) - quantile_25(s, a)` from IQN.

- Flat actions: low IQR (confident about mediocre return) → small bonus
- Large-exposure actions: high IQR (uncertain about outcome) → large bonus → explored more
- Self-calibrating: as model learns, IQR shrinks for well-understood actions, bonus fades

**alpha**: Start at 0.5, anneal to 0.1 over training (encourage early exploration, reduce late).

**Files**: `gpu_iqn_head.rs` (compute IQR per action), `gpu_action_selector.rs` (add bonus before Boltzmann)

---

## Component 6: Activate PopArt + Couple with Tau

**PopArt**: Already fully implemented (`gpu_dqn_trainer.rs:1882-1904`), just disabled. Enable it.

PopArt normalizes rewards to zero-mean unit-variance using Welford's online algorithm. This:
- Prevents C51 from favoring low-variance (Flat) actions
- Makes reward scale consistent across market regimes
- 100-step warmup before normalization kicks in

**Tau coupling**: When PopArt detects reward distribution shift (running variance changes >2× in one epoch):
- Reset tau to initial value (0.005)
- Prevents target network from tracking a non-stationary Q-landscape too quickly

**Files**: `fused_training.rs` (enable flag), `gpu_dqn_trainer.rs` (PopArt kernel), `training_loop.rs` (tau reset logic)

---

## Component 7: Feed the Gradient Vaccine

The gradient vaccine infrastructure (PCGrad-like orthogonal projection) is fully built (`gpu_dqn_trainer.rs:3650-3746`) but `pending_vaccine_batch` is always `None`.

**Fix**: Every 20 steps, create a DIVERSE vaccine batch:
1. Sample from PER buffer with STRATIFIED sampling — ensure all 3 magnitude levels are represented equally
2. Set `fused.pending_vaccine_batch = Some(diverse_batch)`
3. The existing vaccine mechanism projects training gradients to remove components that conflict with diverse-batch gradients

**Effect**: The model CANNOT converge to a non-diverse policy because the vaccine batch always contains magnitude diversity — gradients pushing toward "always Small" get projected out.

**Files**: `training_loop.rs` (populate vaccine batch), `replay_buffer_type.rs` (stratified sample method)

---

## Component 8: Per-Branch Bottleneck Width

Current bottleneck compresses 42 market features → 2 dimensions for ALL branches. This starves magnitude/order branches of information.

**Change**: Branch-specific bottleneck widths:
- Direction: bottleneck=4 (coarse regime + momentum sufficient)
- Magnitude: bottleneck=8 (needs volatility, spread, volume details)
- Order: bottleneck=6 (needs spread, depth, fill probability)
- Urgency: bottleneck=4 (needs momentum, time-of-day)

**Implementation**: Replace single bottleneck with 4 parallel projections. Total params increase: (42×4 + 42×8 + 42×6 + 42×4) = 924 params (negligible vs 336K total).

**Files**: `gpu_dqn_trainer.rs` (bottleneck allocation), `batched_forward.rs` (per-branch projection GEMMs)

---

## Component 9: Ensemble Variance at Inference

5 ensemble heads are trained with diversity loss but DISCARDED at inference (only head 0 used).

**Change**: At action selection, compute ensemble variance and add as exploration bonus:
```
Q_select(s, a) = Q_head0(s, a) + beta × std(Q_all_heads(s, a))
```

**Implementation**: At epoch boundary, sync all ensemble head Q-values into a single buffer. During action selection, compute std across heads per-action. Add to Q before Boltzmann.

**beta**: 0.3 (tunable via hyperopt). Higher = more exploration of uncertain actions.

**Files**: `fused_training.rs` (ensemble Q sync), `gpu_action_selector.rs` (add variance bonus)

---

## Component 10: Clear PER Between Folds

Currently walk-forward folds carry PER buffer across, creating implicit data leakage. Weights carry over (intentional transfer learning), but replay experiences from a different policy contaminate the new fold.

**Change**: At fold boundary, clear the PER buffer completely. Weights stay (transfer learning), but the replay buffer starts fresh. The first epoch of each fold fills the buffer with experiences from the current policy on the new data range.

**Files**: `trainer/mod.rs` (add `clear_replay_buffer()` call at fold boundary)

---

## Verification

### Local testing
```bash
SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget smoke_tests --nocapture
```

### H100 deployment
```bash
./scripts/argo-train.sh dqn --baseline --epochs 10
```

### Success criteria
1. **Action diversity**: No single exposure level >40% (currently S_25 is 34%)
2. **Magnitude usage**: All 3 magnitude levels >15% each (currently Small is 82%)
3. **Sharpe stability**: Epochs 2-5 Sharpe within ±0.5 of epoch 1 (currently drops by -0.89)
4. **Order diversity**: LimitMaker >10%, IoC >5% (currently <2%)
5. **GPU timing**: fwd_bwd ≤ 330ms (no regression from current 319ms)

### Metrics to monitor
- Per-branch action entropy (should be >0.8 normalized for each branch)
- Per-branch TD-error variance (should decrease over training for all branches)
- IQN IQR per action (should shrink for well-learned actions, stay high for uncertain)
- PopArt running variance (should stabilize after 100 steps)
- Vaccine projection frequency (how often gradients get corrected)

---

## Implementation Order

1. **Structural**: Collapse 7-level exposure + remove Flat floor (prerequisite for everything)
2. **PopArt**: Enable flag + tau coupling (quick win, high impact)
3. **Position-aware epsilon**: Per-branch epsilon vector (quick win)
4. **IQN anti-collapse**: IQR bonus at action selection (medium effort)
5. **Reward decomposition**: 4-component reward (high effort, high impact)
6. **Hindsight relabeling**: Magnitude counterfactual (medium effort)
7. **Gradient vaccine**: Stratified diverse batch source (medium effort)
8. **Per-branch bottleneck**: Branch-specific widths (medium effort)
9. **Ensemble inference**: Variance bonus (low effort)
10. **Clear PER between folds**: One-line fix (trivial)

## Key Files

| File | Components |
|------|-----------|
| `crates/ml-core/src/action_space.rs` | 1: 7-level exposure |
| `experience_kernels.cu` | 1: action encoding, 2: reward decomposition, 4: temporal epsilon |
| `gpu_action_selector.rs` | 4: vector epsilon, 5: IQR bonus, 9: ensemble variance |
| `gpu_iqn_head.rs` | 5: compute IQR per action |
| `gpu_her.rs` | 3: hindsight action relabeling |
| `gpu_dqn_trainer.rs` | 2: 4-component reward, 6: PopArt enable, 8: per-branch bottleneck |
| `fused_training.rs` | 6: tau coupling, 7: vaccine batch, 9: ensemble sync |
| `training_loop.rs` | 7: stratified sampling, 10: PER clear |
| `batched_forward.rs` | 8: per-branch projection GEMMs |
| `c51_loss_kernel.cu` | 2: per-branch C51 targets |
| `trainer/mod.rs` | 10: fold boundary cleanup |
