# Trainer Rollout-Buffer + GAE Refactor — Design Spec

**Date:** 2026-06-02
**Status:** Design — awaiting Phase 1A verdict before implementation
**Branch target:** `ml-alpha-regime-observer` (off `cfaa420bd` + Phase 1A)
**Estimated effort:** 1-2 weeks focused work + 2-3 cluster verification cycles

---

## 0. Why this refactor

### 0.1 Current state — single-step trainer

Per `pearl_foxhunt_trainer_is_genuinely_single_step`: foxhunt's PPO + distributional RL trainer has ALL per-step buffers `[B]`-sized only. `step_with_lobsim_gpu_body` (`crates/ml-alpha/src/trainer/integrated.rs:8807`) processes ONE env step per call, then runs forward → advantage compute → K-loop PER-replay training → PER push → returns to outer loop. `RL_N_ROLLOUT_STEPS_INDEX` (slot 404) is the K-loop training-intensification divisor, NOT a trajectory horizon. PER stores individual transitions; no contiguous trajectory reconstruction is possible.

### 0.2 What this architecture forecloses

True GAE backward sweep `A_t = δ_t + γλ·A_{t+1}` is structurally impossible. The recurrence operates over FUTURE TD errors `{δ_{t+1}, …, δ_{t+T}}` that haven't been computed yet at training time t.

`compute_advantage_return.cu:43` consequently uses done-gated 1-step TD: `advantages[b] = is_done ? (ret − vt) : 0.0f`. Effective horizon `T_eff = 1`. Per `pearl_advantage_kernel_is_done_gated_td_not_gae` CORRECTED COST ESTIMATE 2026-06-02d.

### 0.3 Mathematical case for GAE in foxhunt

Schulman, Moritz, Levine, Jordan, Abbeel 2015 GAE theorem: with `λ=0.95, γ=0.975`, T_eff = `1/(1 − γλ) ≈ 13.56` steps ≈ foxhunt's measured 19-step close interval.

**Improvement (1) — V regression quality.** Current done-gated TD effectively sets `returns_t ≈ V(s_t)` on non-done steps (since `A_t = 0 → returns_t = A_t + V_t = V_t`). V regression target on non-done is self-referential — V learns to predict V → zero gradient → V only learns from `V(s_close) ≈ r_close`. V quality on held states is poor.

With GAE: `returns_t = A_t^GAE + V_t` is the λ-return. V learns from every state with properly anchored target. **Load-bearing improvement** — bad V means bad advantages mean bad policy gradient.

**Improvement (2) — sample efficiency.** Non-zero gradient observations per epoch:
- Done-gated 1-step TD: `n ≈ b · n_steps · close_rate = 128 · 2000 · 0.053 ≈ 13,500`
- GAE: `n ≈ b · n_steps ≈ 256,000` — **~19× more advantage observations per epoch**

By Konda-Tsitsiklis 2003 O(1/√n) policy gradient convergence: policy noise reduces by `√19 ≈ 4.4×` per epoch. Equivalent to ~80,000 env steps of done-gated TD in current 2000-step mid-smoke budget.

**Improvement (3) — credit assignment.** A held action at step t contributing to a winning close at step t+10 receives credit `γ^10·λ^10 · A_close ≈ 0.46 · A_close`. Done-gated TD gives credit 0. Foxhunt's reward dynamics — where close-event rewards are sparse but held-state decisions matter — benefit substantially.

**Empirical case**: every production PPO impl since 2017 uses GAE (OpenAI Baselines, SB3, CleanRL, Acme). PPO without GAE is non-standard. Zhang-Zohren-Roberts 2020 + Bauman 2024 (the closest published trading-RL precedents) both use GAE.

### 0.4 What this refactor unlocks beyond GAE

The trainer rearchitecture from per-step to `[B × T]` rollout-based also enables:

1. **Multi-epoch PPO updates over the same rollout** (Schulman 2017 PPO paper: typically 4-10 epochs). Currently impossible — each env step gets one gradient update.
2. **λ-returns for V regression** — free byproduct of GAE.
3. **Truncated importance sampling correction** (V-trace, IMPALA-style off-policy).
4. **Clean separation of concerns**: rollout buffer serves on-policy PPO + V; PER continues serving off-policy distributional Q + aux heads. Currently both paths are intermixed in single-step semantics.

---

## 1. Architecture decisions

### 1.1 Rollout buffer

New device buffers, all `[B × T_rollout]` shape:

```
rollout_rewards_d:      f32  [B × T]     # r_t  (post-Phase-1A: pure realized_pnl_delta)
rollout_dones_d:        u8   [B × T]     # is_done[b][t]
rollout_v_t_d:          f32  [B × T]     # V(s_t) at collection time
rollout_v_tp1_d:        f32  [B × T]     # V(s_{t+1}) at collection time (or boot-strap V_T)
rollout_actions_d:      i32  [B × T]     # selected action
rollout_log_pi_old_d:   f32  [B × T]     # log π_old(a_t | s_t) — frozen at collection
rollout_iqn_q_values_d: f32  [B × T × N_TAU × N_ACTIONS]  # IQN snapshots for distill (if needed)
rollout_h_t_d:          f32  [B × T × HIDDEN_DIM]  # encoder output snapshot — needed for re-evaluation in multi-epoch updates
```

Memory at b=1024, T=256, HIDDEN_DIM=256: ~270MB for h_t buffer. Acceptable on L40S (48GB) and H100 (80GB). For RTX 3050 mid-smoke (b=128, T=256): ~34MB. Trivial.

**T_rollout choice**: 256 default. ISV-controlled via new slot `RL_PPO_ROLLOUT_HORIZON_INDEX` (deprecating the misnamed `RL_N_ROLLOUT_STEPS_INDEX = 404` which stays as K-loop divisor).

### 1.2 Rollout-collection loop

```
for t in 0..T_rollout:
    h_t = encoder(s_t)
    a_t, log_pi_old_t = sample_action_policy(h_t)
    v_t = V_head(h_t)
    (r_t, done_t, s_{t+1}) = env.step(a_t)
    rollout[b, t] = {h_t, a_t, log_pi_old_t, v_t, r_t, done_t}

# Bootstrap V at trajectory end (for GAE backward sweep boundary)
h_T = encoder(s_{T_rollout})
v_T = V_head(h_T)  # used as v_{T+1} in GAE recurrence
```

No training during collection. Pure data gathering. Forward passes are graph-capturable as a single mega-graph (subject to memory).

### 1.3 GAE backward sweep kernel

New CUDA kernel `gae_backward_sweep.cu`:

```c
// Input:  rewards[B][T], dones[B][T], v_t[B][T], v_T (bootstrap), γ, λ
// Output: advantages[B][T], returns[B][T]
//
// Per batch b, sequential backward sweep:
//   A_T   = 0
//   for t in (T-1)..0:
//     v_tp1 = (t == T-1) ? v_T : v_t[b][t+1]
//     non_terminal = 1.0f - dones[b][t]
//     δ_t = r[b][t] + γ * v_tp1 * non_terminal - v_t[b][t]
//     A_t = δ_t + γ * λ * A_next * non_terminal
//     A_next = A_t
//   returns[b][t] = A_t + v_t[b][t]
```

Single-thread-per-batch kernel (`Grid=(ceil(B/256))`, `Block=(256)`). Deterministic by construction (sequential backward sweep, no parallel reductions).

### 1.4 Training pass — multi-epoch over rollout minibatches

```
for epoch in 0..n_ppo_epochs:               # default 4
    shuffle_indices(seeded_rng)             # seed = epoch_seed
    for mb in minibatches(rollout):         # size = MB_SIZE flat samples
        h_t_mb   = gather(rollout_h_t, mb_indices)
        # Re-evaluate current policy + value heads on h_t (FROZEN encoder during multi-epoch)
        a_logits_new = policy_head(h_t_mb)
        v_new = V_head(h_t_mb)
        # PPO surrogate
        ratio = exp(log_pi_new(a_t) - log_pi_old_t)
        clipped = clip(ratio, 1-ε, 1+ε)
        L_pi = -E[min(ratio * A, clipped * A)]
        # V regression target
        L_v = E[(v_new - returns)²]
        # Backprop + Adam
```

`MB_SIZE` is a fraction of `B · T_rollout`. Default: 8 minibatches per epoch.

**Encoder frozen during multi-epoch** to preserve `h_t` snapshots' validity. Standard PPO practice; alternative is to re-encode from `s_t` (expensive).

### 1.5 Reconciliation with PER and distributional Q

**PER**: continues storing individual transitions for off-policy distributional Q + aux replay. Updated at END of each rollout collection (push T_rollout transitions per batch in one bulk operation). Sampling for distributional Q training remains as today.

**Distributional Q + aux heads**: train from PER samples as today. NOT from rollout buffer (those are on-policy for PPO + V only).

**Dueling Q**: trains from PER samples for distributional component, but can ALSO be evaluated on rollout buffer if needed for diagnostic comparison.

This separation is cleaner than the current intermixed semantics.

### 1.6 K-loop semantics

Current K-loop = K iterations of PER-replay training per env step. Reinterpret post-refactor:
- K_ppo = number of PPO epochs over the rollout buffer (default 4)
- K_per = number of PER replay batches per env step (continues existing behavior for distributional Q)

`RL_N_ROLLOUT_STEPS_INDEX` (slot 404) and `RL_K_LOOP_DIVISOR_INDEX` (slot 450) get re-documented but keep current semantics for PER replay. New slots:
- `RL_PPO_ROLLOUT_HORIZON_INDEX`: T_rollout, bootstrap 256
- `RL_PPO_N_EPOCHS_INDEX`: K_ppo, bootstrap 4
- `RL_PPO_MINIBATCH_FRACTION_INDEX`: minibatches per epoch, bootstrap 1/8

### 1.7 CUDA Graph capture

Each component graph-capturable independently:
- Rollout collection inner loop: ONE mega-graph capturing T env steps + forward + buffer writes. Per-step memcpy_dtod for diag readback (Phase 0 contract) stays outside the graph.
- GAE backward sweep: single kernel launch, trivially capturable.
- Training pass: K_ppo × n_minibatches launches; capturable as a graph per epoch.

Three graphs total, replayed in sequence. Each replay corresponds to one rollout cycle (T env steps + training).

### 1.8 Determinism preservation

All design elements determinism-compatible:
- Rollout collection: same as current per-step pipeline (already deterministic post-determinism foundation)
- GAE backward sweep: single-thread sequential per batch — trivially deterministic
- Multi-epoch shuffling: seeded RNG with deterministic seed derivation (e.g., `seed = base_seed XOR rollout_idx XOR epoch_idx`)
- Minibatch gradient accumulation: must use block-tree-reduce (no atomicAdd), per `feedback_no_atomicadd`

Phase 0 cross-source consistency contract ($50 tolerance between `diag.realized_pnl_usd_cum` and `eval_summary.total_pnl_usd`) preserved by construction — pnl_track.cu writes USD pnl unchanged; rollout buffer reads from the same pnl source.

---

## 2. Diagnostic emissions

New diag fields under `rollout`:

```json
"rollout": {
    "horizon": 256,                  // current T_rollout
    "n_epochs": 4,                   // current K_ppo
    "advantage_mean": 0.034,         // mean A_t over rollout
    "advantage_std": 1.27,           // std of A_t
    "returns_mean": 12.4,            // mean returns_t
    "returns_std": 18.9,             // std of returns_t
    "v_pred_mean": 8.2,
    "v_pred_target_pearson": 0.74,   // Pearson(V(s_t), returns_t) — V quality diagnostic
    "delta_t_mean": 0.0012,          // TD error mean
    "delta_t_var": 0.034,
    "fraction_done": 0.053,          // close events per step
    "T_eff_observed": 13.4           // measured 1/(1-γλ_effective)
}
```

`EXPECTED_LEAVES` bump: +10 (current 712-713 → ~723).

---

## 3. Falsification criteria (predetermined)

### G1 — Alignment preserved (Phase 1A carryover)
`Pearson(rewards.sum, Δrealized_pnl_usd_delta)` train last 500 ≥ 0.85. Must NOT regress from Phase 1A.

### G2 — V regression quality
`Pearson(v_pred_t, returns_t)` over train last 500 steps ≥ 0.5. Currently with done-gated TD this is unmeasurable (V on non-done states is self-referential). After GAE, V should track the bootstrapped returns.

### G3 — Sample efficiency
At same wallclock budget (12 min mid-smoke), Phase 1B agent achieves train win-rate ≥ 0.40 (currently 0.33 at HEAD). If GAE truly delivers 19× more gradient observations, equivalent improvement should be observable in wr.

### G4 — Determinism preserved
`./scripts/determinism-check.sh --quick` exits 0. No regression from `cfaa420bd` baseline.

### G5 — Cross-source consistency preserved
`diag.realized_pnl_usd_cum` vs `eval_summary.total_pnl_usd` within $50.

### G6 — Eval pnl direction
Eval `total_pnl_usd` improves vs Phase 1A baseline (less negative or positive). Magnitude target: ≥ 50% reduction in absolute eval loss vs Phase 1A baseline.

If any of G1, G4, G5 fail: refactor has introduced a bug. STOP and diagnose. Per `feedback_no_quickfixes`: do not tune hyperparameters trying to lift G2/G3/G6 — those are downstream consequences of correct G1/G4/G5.

---

## 4. File-level design

### 4.1 New files

| Path | Purpose | LOC est |
|---|---|---|
| `crates/ml-alpha/cuda/gae_backward_sweep.cu` | GAE kernel (single-thread per batch, sequential backward) | 150 |
| `crates/ml-alpha/src/trainer/rollout_buffer.rs` | RolloutBuffer struct: allocation, push, gather, GAE invocation | 400 |
| `crates/ml-alpha/src/trainer/rollout_collection.rs` | T env-step loop without training, mega-graph compatible | 250 |
| `crates/ml-alpha/src/trainer/ppo_minibatch.rs` | Multi-epoch PPO update over rollout minibatches | 350 |
| `crates/ml-alpha/tests/rollout_buffer_invariants.rs` | GPU-oracle tests for GAE correctness, sample-by-sample | 300 |
| `docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md` | Implementation plan (separate file) | — |

### 4.2 Modified files

| Path | Changes | LOC est |
|---|---|---|
| `crates/ml-alpha/src/rl/isv_slots.rs` | New slots: RL_PPO_ROLLOUT_HORIZON_INDEX, RL_PPO_N_EPOCHS_INDEX, RL_PPO_MINIBATCH_FRACTION_INDEX. Update RL_SLOTS_END. | 40 |
| `crates/ml-alpha/src/trainer/integrated.rs` | Replace `step_with_lobsim_gpu_body` per-step loop with rollout collection + GAE + multi-epoch training. ~12 ISV controller adaptations. New diag fields. | 1500 (refactor — net delta may be small if existing code gets consolidated) |
| `crates/ml-alpha/cuda/compute_advantage_return.cu` | DEPRECATE — kept as fallback under FOXHUNT_LEGACY_TD env flag for A/B comparison, removed in follow-up commit | 0 (no change initial) |
| `crates/ml-alpha/examples/alpha_rl_train.rs` | Adapt train + eval loop to rollout cadence | 300 |
| `crates/ml-alpha/tests/eval_diag_emission.rs` | EXPECTED_LEAVES bump, new rollout invariants | 100 |
| `scripts/tier1_5_verdict.py` | New signal: V regression Pearson (G2) | 50 |

Total: ~3400 LOC new + modified, single feature branch.

### 4.3 NOT modified (intentionally)

- `crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu` — reward Phase 1-7 unchanged; already producing pure pnl after Phase 1A
- `crates/ml-alpha/cuda/rl_surfer_scaffold_controller.cu` — pin remains, controller writes still gated
- `crates/ml-alpha/src/rl/{dqn,iqn,outcome_head}.rs` — distributional Q + aux heads continue from PER (unchanged)
- `crates/ml-alpha/cuda/bellman_target_projection.cu` — distributional Q targets from PER unchanged

---

## 5. Sequencing — phased delivery

This refactor cannot ship atomically given the LOC. Suggested phases (each a separate commit + cluster verification):

| Phase | Scope | Verification |
|---|---|---|
| 1B-A | New ISV slots + RolloutBuffer struct + `gae_backward_sweep.cu` kernel + GPU-oracle invariant tests | `rollout_buffer_invariants.rs` PASS |
| 1B-B | Rollout collection loop (parallel-existing-path, behind FOXHUNT_USE_ROLLOUT env flag) | Mid-smoke with flag=0 = current behavior unchanged |
| 1B-C | Multi-epoch PPO update over rollout buffer | Mid-smoke with flag=1: G1, G2, G4, G5 PASS; G3 measured |
| 1B-D | Cluster verification 20k + flag=1 | G6: eval pnl improvement vs Phase 1A baseline |
| 1B-E | Default flag=1; remove legacy `compute_advantage_return.cu` path; cleanup | Mid-smoke + cluster sanity |

Phase 1B-A through 1B-C are mostly local-test driven. 1B-D is the load-bearing cluster verification. 1B-E is cleanup.

---

## 6. Risks

### 6.1 Memory pressure at cluster b=1024
- `rollout_h_t_d` size: b · T · HIDDEN_DIM · 4 = 1024 · 256 · 256 · 4 = ~270MB
- Aggregate rollout buffer: ~350MB
- L40S 48GB, H100 80GB — well within capacity
- Risk: low

### 6.2 Controller starvation under new reward distribution
- Per `pearl_pure_pnl_mode_starves_b16_controllers`: b=16 starves; b=128 was safe in Phase 0/1A. Cluster b=1024 even safer.
- Watch popart EMA stability post-1B
- Mitigation: same adaptive controller framework, with extended observation window if needed

### 6.3 CUDA Graph re-capture invalidation
- Existing graphs (Graph A pre-step, Graph B post-step) must be rebuilt for new pipeline
- Mitigation: validate graph compatibility in 1B-C before cluster verification

### 6.4 PER off-policy timing skew
- PER receives transitions in bulk at end of each rollout (vs one-at-a-time currently)
- Skew = up to T_rollout = 256 steps
- For distributional Q (priority-based sampling), the staleness is bounded
- Risk: low; if priorities lag, can adjust PER alpha controller

### 6.5 The 12 ISV controllers tuned to per-step semantics
- popart_sigma, edge_decay PH, Kelly, V_blend_α, etc. — currently update per-step
- Post-1B: same controllers update per env step (same cadence). Should be transparent.
- Risk: low if rollout uses same env-step cadence; medium if rollout introduces a "rollout boundary" semantic that controllers don't handle

### 6.6 Multi-epoch encoder freeze
- Encoder NOT updated during 4 PPO epochs over the rollout
- Standard PPO impl; well-understood tradeoff
- Mitigation: if encoder underfits, can experiment with single-epoch or with encoder re-eval

---

## 7. Decision tree from Phase 1A verdict (which gates Phase 1B urgency)

| Phase 1A result | Phase 1B priority | Reasoning |
|---|---|---|
| G1 pass (Pearson ≥ 0.85) AND eval pnl improves significantly | Low — schedule for next sprint | Phase 1A solved most of the problem; GAE is quality upgrade |
| G1 pass AND eval still bad | High — start 1B-A immediately | Alignment fixed, but sample efficiency / V quality / credit assignment likely the next bottleneck |
| G1 fail (Pearson < 0.85) | Medium — diagnose Phase 1A root cause first | If Phase 1A doesn't lift alignment, GAE alone won't fix it either; need diagnosis before adding refactor complexity |

---

## 8. References

- Schulman, Moritz, Levine, Jordan, Abbeel 2015. *High-Dimensional Continuous Control Using Generalized Advantage Estimation*. arXiv 1506.02438.
- Schulman, Wolski, Dhariwal, Radford, Klimov 2017. *Proximal Policy Optimization Algorithms*. arXiv 1707.06347.
- Konda, Tsitsiklis 2003. *On Actor-Critic Algorithms*. SIAM Journal on Control and Optimization.
- Ng, Harada, Russell 1999. *Policy Invariance Under Reward Transformations*. ICML.
- Zhang, Zohren, Roberts 2020. *Deep Reinforcement Learning for Trading*. Oxford-Man Institute.
- Bauman 2024. *Active High Frequency Trading via DRL*. PMLR v233.
- foxhunt internal pearls:
  - `pearl_foxhunt_trainer_is_genuinely_single_step` — the structural finding driving this refactor
  - `pearl_advantage_kernel_is_done_gated_td_not_gae` — current state + corrected cost estimate
  - `pearl_reward_signal_anti_aligned_with_pnl` (all 4 ADDENDUM sections) — load-bearing problem chain
  - `pearl_diag_pnl_fields_are_shaped_reward_not_usd` — Phase 0 prerequisite, falsification gate
  - `pearl_determinism_achieved` — determinism contract preserved by refactor
- foxhunt internal specs (sibling):
  - `docs/superpowers/specs/2026-06-02-determinism-foundation.md`
  - `docs/superpowers/specs/2026-06-01-reward-policy-alignment-investigation.md` (v4 — supplanted by Phase 1A direct shaping-drop)
