foxhunt/docs/superpowers/plans/2026-06-02-next-session-multi-head-implementation.md

Next Session — Multi-Head Policy Implementation Plan

Date: 2026-06-02 (handoff document) Status: Ready-to-execute plan for next session Linked specs:

• Primary: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
• Parent: docs/superpowers/specs/2026-06-02-regime-invariance-four-interventions.md
• Existing #1 plan: docs/superpowers/plans/2026-06-02-regime-invariance-implementation.md Branch at handoff: ml-alpha-regime-observer at SHA 63fc16f17 (v5.2) Memory snapshot: pearl_grwwh_eval_catastrophic_collapse (just saved)

§0. State at handoff

alpha-rl-grwwh verdict (just landed):

• train +$481M wr 0.312 hold 15.1 — strong train signal
• eval −$160.5M wr 0.264 pf 0.83 sharpe −1.70 — catastrophic generalization failure
• 3rd consecutive eval-collapse on single-policy + Phase 5 hold-bonus architecture

Empirically validated: the scaffold-fade approach (v4/v5/v5.1/v5.2) cannot fix the eval generalization gap because the gap is architectural — single policy network with Phase 5 reward shape over-fits to train regime.

Direction confirmed: ship multi-head policy + R-multiple reward together per the spec.

§1. Pre-flight before any code

§1.1 sp-critical-reviewer on the multi-head spec FIRST

The spec at 2026-06-02-multi-head-policy-with-r-multiple.md has been written with the FIVE BLOCKERs from the previous v6 spec review explicitly addressed. Before any implementation:

1. Dispatch sp-critical-reviewer agent on the multi-head spec
2. Reviewer should especially verify:
   • PolicyHead struct layout in crates/ml-alpha/src/networks/policy_head.rs (referenced but not personally verified in spec §9)
   • EXPECTED_LEAVES = 679 value referenced in §3.4 still matches current test
   • No prior MoE attempts in foxhunt (grep for existing multi-head code)
   • The aux-prior-loss math (§1.4) doesn't break PPO surrogate convexity
3. If reviewer BLOCKS, fix and re-review BEFORE implementation
4. If reviewer FIX-AND-PROCEED, apply edits then start §2

§1.2 Verify the per-step diag / eval_summary discrepancy bug

Per pearl_grwwh_eval_catastrophic_collapse: at eval step 4863 the per-step diag showed +$74M but eval_summary showed −$160M. This is a $234M discrepancy. Investigate BEFORE shipping a new run that relies on per-step diag for verdict tracking:

1. Pull diag.jsonl rows 22000-22826 from PVC for grwwh (/feature-cache/alpha-rl-runs/63fc16f17/fold1/diag.jsonl)
2. Check if the last ~137 eval steps show a catastrophic per-step pnl collapse — if YES, per-step is honest, the agent really did crash at the end
3. If NO collapse visible in per-step → the per-step diag's realized_pnl_cum_usd accumulator differs from eval_summary's total_pnl_usd calculation — file a separate bug/diagnostic spec

This investigation gates whether per-step diag can be trusted for kill-criteria in §3.1 falsification gates. The reviewer-mandated kill gates rely on per-step signal being authoritative.

§2. Decide bundling: ship #1 + multi-head together OR sequentially

Option A: Bundle (single commit, v6.0)

• R-multiple reward + multi-head policy in one commit
• Faster to verdict (one cluster smoke instead of two)
• Higher risk: if eval fails, can't isolate which intervention failed
• Aligns with the §0.5 finding that #1 alone is insufficient

Option B: Sequential (#1 first, then multi-head)

• Ship #1 R-multiple first, get cluster verdict
• IF #1 eval positive: stop (cheapest win)
• IF #1 eval negative: ship multi-head on top
• More cluster smokes burned
• Cleaner attribution

Recommendation: Option A (bundle)

The grwwh result rules out #1-alone as sufficient. v5.2 with Phase 5 still active failed eval; #1 just removes the contamination (R-multiple replaces hold-bonus). A clean single policy on R-multiple alone would likely still over-fit to fold-1 trend regime — same failure mode, different reward shape.

Multi-head is the architectural fix. Bundle.

§3. Implementation ordering (assuming Option A bundle)

Phase A: Foundation (sequential, atomic commit)

1. Slot retire + reallocate (crates/ml-alpha/src/rl/isv_slots.rs)

   • Delete: RL_SURFER_SCAFFOLD_WEIGHT_INDEX (753), RL_SURFER_BREAK_EVEN_WR_INDEX (754), RL_SURFER_K_SHARPNESS_INDEX (755), RL_SURFER_WARMUP_TRADES_INDEX (756)
   • Allocate (R-multiple, per existing #1 plan): RL_ATR_EMA_INDEX (753), RL_R_MULTIPLE_CLAMP (754), RL_ATR_EMA_ALPHA (755), RL_ATR_FLOOR (756)
   • Allocate (multi-head, per multi-head spec §2): RL_POLICY_NUM_HEADS_INDEX (757), RL_POLICY_GATING_ENTROPY_FLOOR_INDEX (758), RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX (759), RL_POLICY_AUX_PRIOR_BETA_INDEX (760)
   • RL_SLOTS_END = 761

2. Delete obsolete code

   • Delete: crates/ml-alpha/cuda/rl_surfer_scaffold_controller.cu
   • Delete: surfer-scaffold launch + module load in crates/ml-alpha/src/trainer/integrated.rs (6 touchpoints per coder analysis)
   • Delete: surfer-scaffold bootstrap entries in with_controllers_bootstrapped
   • Delete: surfer-scaffold registration in crates/ml-alpha/build.rs

3. New kernels (5 total)

   • rl_atr_ema_update.cu — ATR EMA producer (R-multiple input)
   • rl_policy_mixture_forward.cu — K head logits + gate → mixture logits via LogSumExp
   • rl_policy_mixture_backward.cu — backward through mixture
   • rl_policy_aux_prior_loss.cu — Σ_k g_k × β × KL(p_k ∥ p_prior_k) + gradient
   • Modify: rl_fused_reward_pipeline.cu Phase 5 — replace hold-bonus with r = clamp(Δpnl / max(atr_at_entry, floor), ±5)

Phase B: Rust integration

4. Network refactor (crates/ml-alpha/src/networks/policy_head.rs)

   • PolicyHead struct: single Linear → Vec for K heads + Linear for gating
   • Adam optimizer states per tensor
   • Forward pass through new mixture kernels
   • Backward pass with per-head + gate gradients

5. Trainer integration (crates/ml-alpha/src/trainer/integrated.rs)

   • Allocate K head weights at construction (init with action-bias per spec §1.3)
   • Allocate gating weights
   • Add unit_atr_at_entry_d per-batch buffer (R-multiple capture)
   • Wire 5 new kernel launches in step_with_lobsim
   • Bootstrap new ISV slots (β=0.05, entropy floors per spec)

Phase C: Diag + tests

6. Diag emission (build_diag_value)

   • policy.head_count = K
   • policy.head_k_entropy[k] for k in 0..K
   • policy.gate_entropy
   • policy.gate_probs[k] for k in 0..K
   • policy.head_pairwise_kl[i,j] for diagnostic
   • rewards.atr_ema_usd
   • rewards.r_multiple_last
   • EXPECTED_LEAVES: 679 → 679 + 7 (multi-head) + 2 (R-multiple) − 1 (drop pure_pnl_mode/surfer_scaffold_weight) = 687

7. Test migration

   • tests/reward_alignment_invariants.rs: I1-I3 → assert rewards.atr_ema_usd ≥ 0 and rewards.r_multiple_last ∈ [-5, +5]. Delete surfer_scaffold_weight assertions.
   • New test: tests/multi_head_policy_invariants.rs
     • I1: gate_probs[k] in [0, 1] for all k, Σ = 1
     • I2: gate_entropy ≥ entropy_floor (with epsilon)
     • I3: head_pairwise_kl > 0 (at least one pair distinct)
     • I4: each head's entropy ≥ head_entropy_floor

Phase D: Local validation + cluster

8. Local validation

   • cargo check + cargo build --release
   • cargo test reward_alignment_invariants multi_head_policy_invariants --ignored --nocapture
   • 200-step local smoke (matches existing test pattern)
   • GPU compute-sanitizer racecheck on the new kernels

9. Cluster smoke

   • Pin SHA via argo-alpha-rl.sh --sha
   • L40S, fold-1, b=1024, 20k train + 5k eval
   • Apply §4.1 per-step kill criteria from multi-head spec
   • Apply §4.2 verdict gates

§4. Failure-mode adaptive replanning

If multi-head + R-multiple cluster smoke ALSO produces negative eval:

Diagnostic step 1: per-regime eval stratification per spec §4.3. If specialization spread < 0.02 → heads didn't diverge → aux prior β too low. Boost β to 0.15 and retry.

Diagnostic step 2: if specialization shows but eval still negative → the heads' priors are wrong (e.g., trend+reversion+flat doesn't cover the actual eval-period regime). Try different prior set (e.g., add momentum-breakout head).

Diagnostic step 3: if priors are diverse and eval still negative → architectural ceiling hit. Pivot to:

• Cross-fold training (#3 from parent spec) — train on mixed folds 0/1/2
• Regime features wired into encoder (#2 from parent spec)
• Behavior cloning warm-start from a profitable rule-based bot

Termination: if 3 multi-head variants all fail eval, the conclusion is the architecture cannot extract alpha from foxhunt's MBP-10 dataset at b=1024 with 20k step budget. Different data scale or different model architecture required.

§5. Discipline reminders (DO NOT VIOLATE)

• feedback_no_feature_flags — surfer-scaffold deletion in SAME commit as new code, no toggles
• feedback_single_source_of_truth_no_duplicates — slot reuse 753-756 with rename, no v2_ suffixes
• feedback_local_smoke_before_cluster — Phase D step 8 BEFORE step 9, no exceptions
• feedback_kill_runs_on_anomaly_quickly — kill at first anomaly per spec §4.1, don't wait for verdict
• feedback_no_atomicadd — all new kernels use block-tree-reduce
• pearl_argo_branch_ref_trap_strikes_again — cluster command pins exact SHA, NOT branch tip
• feedback_going_in_circles_pattern — if 2 multi-head variants in a row fail with new failure modes, STOP iterating and reconsider foundation

§6. Session opening checklist (for the next session)

1. Read pearl_grwwh_eval_catastrophic_collapse (just saved)
2. Read docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
3. Read docs/superpowers/specs/2026-06-02-regime-invariance-four-interventions.md (parent)
4. Read this plan
5. Dispatch sp-critical-reviewer on multi-head spec (per §1.1)
6. Resolve reviewer findings
7. Investigate per-step / eval_summary discrepancy (per §1.2)
8. Decide bundling (per §2) — likely Option A
9. Begin Phase A implementation (per §3)

Total estimated session time: 4-6 hours for full bundle, cluster smoke runs overnight if needed.