foxhunt/docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md

Multi-Head Policy with R-Multiple Reward (Intervention #4-clean)

Date: 2026-06-02 Status: Spec draft (post-sp-critical-reviewer BLOCK on prior v6 attempt) Prerequisite: Intervention #1 (R-multiple reward) has shipped and the eval verdict is in Linked specs: 2026-06-02-regime-invariance-four-interventions.md (parent) Linked pearls: pearl_surfer_baseline_was_train_only_never_eval, pearl_two_alpha_modes_surfer_vs_trend, pearl_reward_signal_anti_aligned_with_pnl, pearl_pi_actor_collapses_without_entropy_floor

§0. Empirical motivation — single-policy regime collapse is EMPIRICAL, not theoretical

The previous v6 attempt tried to motivate multi-head policy via a theoretical "average optimal action" argument. The sp-critical-reviewer correctly BLOCKED that as hand-waving. This spec grounds the motivation in cluster-run evidence:

Run SHA	Train wr	Eval wr	wr collapse	Per-pearl
dd049d9a4	0.56 (hold=30)	0.219	2.6× drop	pearl_surfer_baseline_was_train_only_never_eval
87a8259c6 (8gtk2)	wr_max 0.418 / final 0.343	0.340	minor train→eval but train wr peaked then declined within training itself	pearl_reward_signal_anti_aligned_with_pnl
63fc16f17 (grwwh, current)	0.312 final	TBD ~25min	TBD	project_adaptive_scaffold_session_2026_06_02

The dd049d9a4 evidence is decisive: a single policy network trained on Phase-5 reward shape achieved profitable train trade-quality and collapsed to barely-above-random on eval data. Per pearl_two_alpha_modes_surfer_vs_trend, this isn't a bug — it's "philosophy not bug" — but the philosophy was BAKED INTO weights by training. Without architectural support for multiple philosophies, the single policy can only embed one philosophy at a time.

This spec hypothesizes (NOT proves) that giving the policy K heads with different initial biases + a learned gating produces a policy that retains MULTIPLE philosophies in weights, allowing it to switch at inference time based on encoder context. The verification is empirical: cluster smoke + eval-wr stratified by regime markers.

§1. Design

§1.1 Architecture (verified against current codebase)

Current architecture (read from crates/ml-alpha/src/networks/policy_head.rs and crates/ml-alpha/src/trainer/integrated.rs):

• Encoder: mamba2 trunk → encoder_output of size HIDDEN_DIM
• Policy head: Linear(HIDDEN_DIM → N_ACTIONS=11), outputs pi_logits[b, a]
• Q head: dueling distributional → q_logits[b, a, z]
• V head: scalar critic for PPO

Proposed change (multi-head policy ONLY — Q and V heads unchanged):

• Encoder: UNCHANGED
• K policy heads (K=3 proposed): K × Linear(HIDDEN_DIM → N_ACTIONS), outputs pi_logits_k[b, a] for k ∈ {0, 1, 2}
• Gating head: small Linear(HIDDEN_DIM → K), outputs gate_logits[b, k]
• Final pi_logits: pi_logits[b, a] = log( Σ_k softmax(gate_logits[b])[k] × softmax(pi_logits_k[b])[a] )
• Equivalently in logit space, the mixture is taken on probabilities (not logits) then log'd back. This is the product-of-experts ↔ mixture-of-experts distinction; we use mixture (additive over probabilities).

§1.2 Reward signal — UNCHANGED from Intervention #1

This spec does NOT modify the reward kernel. R-multiple reward from Intervention #1 stays exactly as-is. The reward path is identical for all K heads — they all receive the SAME per-event reward signal.

This addresses BLOCKER 4 (Pearson alignment): since the reward path is unchanged, Pearson(reward, Δpnl) is identical to whatever #1 achieves. Multi-head policy adds POLICY expressiveness without touching the reward. No anti-aligned multiplicative shaping is introduced.

§1.3 Specialization mechanism (no new ISV signals)

This addresses BLOCKER 1: heads specialize WITHOUT requiring any new regime-observation ISV slot. Specialization arises from FOUR sources internal to the network:

1. Different head initialization: each head's output projection is initialized with a different action bias:
   • Head 0 (trend): bias toward HOLD_LONG / HOLD_SHORT actions (+0.5 logit on HOLD actions, 0 elsewhere)
   • Head 1 (reversion): bias toward FLAT_LONG / FLAT_SHORT actions (+0.5 logit on FLAT actions, 0 elsewhere)
   • Head 2 (flat): bias toward FLAT_FLAT action (+0.5 logit on FLAT only)
2. Per-head entropy floor: each head has an independent entropy lower bound preventing collapse to deterministic. Floor value: log(11) × 0.4 (~0.96).
3. Gating entropy floor: gating distribution forced to retain at least log(K) × 0.5 (~0.55 for K=3) entropy. Prevents gate-collapse-to-one-head mode (per pearl_pi_actor_collapses_without_entropy_floor).
4. Per-head auxiliary action-prior loss (§1.4 below) — continuous regularizer that prevents heads from converging to identical distributions.

The init bias alone (1) is FRAGILE: heads start different but see identical PPO gradient, so they would converge over training. Item 4 provides the CONTINUOUSLY-APPLIED regularizer that maintains specialization across training.

§1.4 Per-head auxiliary action-prior loss (CLEAN per-expert specialization)

For each head k, an auxiliary KL term pulls its action distribution toward a fixed prior p_prior_k:

aux_loss_k = β × KL( p_k(a|s) ∥ p_prior_k(a) )

Combined loss per head:

total_loss_k = ppo_loss_k + aux_loss_k

The mixture policy's total loss is Σ_k g_k × total_loss_k where g_k is the gating weight.

Priors (defined per head, FIXED — no learned drift):

• p_prior_0 (trend): 0.10 uniform over non-HOLD, 0.30 on each of the 3 HOLD_LONG / HOLD_SHORT / HOLD_FLAT actions (sums to 1.0)
• p_prior_1 (reversion): 0.10 uniform over non-FLAT-of-position, 0.40 on FLAT_LONG and FLAT_SHORT
• p_prior_2 (flat): 0.40 on FLAT_FLAT, 0.06 uniform over others

(The exact prior weights are tunable but should sum to 1.0 and concentrate ~30-40% of mass on the head's specialty action group.)

β coefficient: ISV-driven via new slot RL_POLICY_AUX_PRIOR_BETA_INDEX. Bootstrap 0.05 (5% of PPO loss magnitude). When β=0, the spec reduces to §1.3 init-bias-only — useful as ablation.

Why this addresses Pearson alignment (re-confirms BLOCKER 4 stays addressed): the aux_loss is added to the POLICY LOSS, not to the REWARD signal. Pearson(reward, Δpnl) is unchanged because the reward kernel is unmodified. The aux loss is a REGULARIZER on policy output; it does not alter the gradient direction of the PPO surrogate's REWARD term — only adds an orthogonal pull toward each head's prior.

Sign-preservation: KL divergence is always ≥ 0; aux_loss is always positive; subtracting it from the maximization objective always reduces the policy's deviation from prior. Cannot induce sign-flips in the reward direction.

Reference: standard mixture-of-experts auxiliary loss pattern. Shazeer et al. 2017 (sparsely-gated MoE), Lepikhin et al. 2020 (GShard) — well-studied technique, low novelty risk.

§2. ISV slot additions (4 slots, verified against current RL_SLOTS_END)

Current RL_SLOTS_END = 757 (post-grwwh / v5.2, will be 757 again after #1 R-multiple ships per the #1 plan). New slots:

RL_POLICY_NUM_HEADS_INDEX            usize = 757   // K (current 3 — could allow runtime override)
RL_POLICY_GATING_ENTROPY_FLOOR_INDEX usize = 758   // ~0.55 (log(K) × 0.5)
RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX   usize = 759   // ~0.96 (log(11) × 0.4)
RL_POLICY_AUX_PRIOR_BETA_INDEX       usize = 760   // 0.05 bootstrap; β for per-head action-prior KL

RL_SLOTS_END = 761 after this spec ships.

No new producer kernels needed. The entropy floors are read by the EXISTING entropy controller infrastructure (per crates/ml-alpha/cuda/rl_entropy_coef_controller.cu and similar).

§3. Implementation tasks (concrete, ordered)

§3.1 Network changes (Rust)

1. Modify PolicyHead struct in crates/ml-alpha/src/networks/policy_head.rs:
   • Add Vec<Linear> for K head projections instead of single Linear
   • Add Linear for gating projection
   • Adam optimizer states for each (use existing Adam per-tensor pattern)
2. Modify forward pass:
   • Compute K head logits in parallel (matmul against shared encoder_output)
   • Compute gate logits
   • Mixture via log Σ softmax(gate) × softmax(head)
3. Backward pass: gradient flows through gating × head expression. Standard PPO surrogate on the mixture pi_logits.

§3.2 Kernel changes (CUDA)

Multi-head policy is mostly Rust-side because the heads are linear projections. New kernels needed:

• rl_policy_mixture_forward.cu: takes K head logits + gate logits → produces mixture logits (LogSumExp over heads with gate weighting). Block-tree-reduce, no atomicAdd. ~80 LOC.
• rl_policy_mixture_backward.cu: backward through the mixture. ~120 LOC.
• rl_policy_aux_prior_loss.cu: computes Σ_k g_k × β × KL(p_k ∥ p_prior_k) and gradient ∂L_aux/∂head_k_logits. Reads RL_POLICY_AUX_PRIOR_BETA_INDEX for β. ~60 LOC.

Per feedback_no_nvrtc: pre-compiled cubins via build.rs.

§3.3 Trainer integration

• Allocate K head weights + gating weights in trainer constructor
• Initialize each head with the action-bias pattern from §1.3
• Wire entropy floors via existing controller infrastructure
• Per-head entropy diagnostic emitted to diag.jsonl (3 new leaves: policy.head_0_entropy, etc.) + gate entropy diagnostic (1 leaf)

§3.4 Diag emission

Add to 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 (mean across batch)

EXPECTED_LEAVES: 679 → 679 + (K + 1 + K) = 686 (K=3).

§4. Falsification — strict kill gates (addresses BLOCKER 5)

§4.1 Per-step kill criteria (per feedback_kill_runs_on_anomaly_quickly)

Kill the cluster smoke if ANY of:

• Pearson(reward, Δpnl) < 0.5 at step 2000 (sparse close events should align by then)
• gate_probs[k] < 0.05 for any k at step 5000 (mode collapse — gate ignored a head)
• gate_entropy < log(K) × 0.4 at step 5000 (gate determinism — should be > 0.45)
• action_entropy < log(11) × 0.5 at step 5000 (overall policy collapse)
• realized_pnl_cum < -$50M at step 10000 (catastrophic loss, abandon)

§4.2 Verdict gates (PASS / FAIL)

PASS criteria (ALL must hold):

• Train Pearson(reward, Δpnl) ≥ 0.6 at steps 1500-2000
• Train wr ≥ 0.35 by step 10k (above any reasonable random baseline; threshold above current grwwh peak of 0.31)
• Train action_entropy ∈ [log(11) × 0.5, log(11) × 0.85] (concentrated but not collapsed)
• Train gate_entropy ≥ log(K) × 0.4 (heads actively used)
• Eval pnl > 0 (the actual generalization test)
• Eval wr ≥ 0.30 (above the random-baseline floor for K=11 actions)
• Per-regime eval wr shows specialization: split eval rows by session_pnl_variance_ema quartiles; the spread max(wr_q) − min(wr_q) should exceed 0.05 (heads pulling agent in different directions based on regime).

FAIL criteria:

• Any per-step kill gate fires
• Eval pnl < -$2M (defines the fail boundary; in [-$2M, 0] is marginal → next intervention)
• Per-regime spread < 0.02 (heads didn't specialize → architectural gain near zero)

§4.3 Per-regime stratification methodology

The falsification needs regime markers that ALREADY EXIST in diag.jsonl. Verified existing signals:

• risk_stack.regime.session_pnl_variance_ema (slot 700) — proxy for vol regime
• risk_stack.regime.dead_zone_flag (slot 696) — proxy for chop regime
• risk_stack.regime.recovery_factor (slot 699) — proxy for drawdown regime

Stratify eval steps by quartiles of session_pnl_variance_ema (low-vol vs high-vol regimes); compute wr separately on each quartile. The hypothesis is that multi-head policy will produce different wr in different quartiles (heads specializing), while single-policy grwwh would produce uniform wr across quartiles.

§5. Risks and mitigations

§5.1 Mode collapse on gating

Risk: gating distribution collapses to single head (gate_probs ≈ [1, 0, 0]); other heads become dead weight.

Mitigation: RL_POLICY_GATING_ENTROPY_FLOOR_INDEX enforces hard floor on gate entropy. Per-step kill gate at step 5000 catches this.

Reference: pearl_pi_actor_collapses_without_entropy_floor — entropy floor canonically applied to action distribution; here applied to gating distribution at the same conceptual layer.

§5.2 Heads converge to identical behavior — MITIGATED by §1.4 aux prior loss

Risk: with shared encoder + shared reward, the heads might converge to identical action distributions, making multi-head equivalent to a single head with K× parameters wasted.

Primary mitigation: §1.4 per-head auxiliary action-prior KL loss. Continuously pulls each head toward its distinct prior, preventing convergence. Standard MoE specialization technique (Shazeer 2017).

Secondary mitigations: per-head init bias (§1.3 item 1) + per-head entropy floor (§1.3 item 2). The aux prior provides the lasting pressure; init + entropy provide the boundary conditions.

Diagnostic: emit pairwise KL between heads to diag (policy.head_pairwise_kl[i,j]). If KL drops below 0.1 across all pairs, heads have converged → β too low or aux prior loss disconnected.

β-tuning escape hatch: RL_POLICY_AUX_PRIOR_BETA_INDEX is ISV-driven. If empirically heads converge despite β=0.05 bootstrap, the controller can adapt β upward via an ISV controller (similar pattern to existing entropy-coef controller). Add this as a follow-up controller only if needed.

§5.3 Training instability from mixture

Risk: PPO surrogate on a mixture has higher variance than on a single Gaussian-ish policy. Could destabilize training.

Mitigation: existing PPO ratio clamp from B-7 / B-9 work handles this. Per-batch advantage normalization (Phase 4.5) helps.

§5.4 Eval doesn't show specialization (heads pulled to average)

Risk: even with all mitigations, eval might produce uniform wr across regime strata → no specialization actually achieved.

Mitigation: this IS the falsification gate. If specialization spread < 0.02, the architectural change failed. The conclusion is then NOT to ship #4.B; reroute to #3 (cross-fold training) or #4.C (full MoE).

§6. Open questions (to resolve before implementation, NOT during)

1. Initial K: 3 (trend/reversion/flat) vs 5 (add trend-reversal, vol-breakout)? Larger K is more expressive but harder to keep specialized. Recommend K=3 to start.
2. Loss decomposition: PPO loss on the mixture (single loss) vs per-head PPO loss weighted by gate (decomposed). Empirically the mixture-loss is simpler and standard in MoE literature. Recommend mixture-loss.
3. Eval boundary discipline: does the gating's running state need reset at eval boundary? Per pearl_adaptive_carryover_discipline, any stateful adaptive signal needs either reset or re-bootstrap. The gating is INSTANT (computed from encoder state per step) so no carry-over — but the per-head entropy controller IS stateful. Use neutral sentinel reset at eval boundary.
4. Distillation interaction: existing B-11-β-style Q→π distill could conflict with mixture policy. Verify the distill grad kernel handles K-head mixture output correctly before enabling distill.

§7. Out of scope

• Per-expert REWARD shaping (which the BLOCKED v6 attempt proposed via multiplicative r_trend × hold_bonus_factor) — explicitly rejected because it violates Pearson alignment per §6.2 of the parent spec. NOTE: per-expert POLICY-LOSS shaping IS in scope (§1.4 aux prior loss) — the distinction is reward-vs-loss; reward stays clean (Pearson-preserving), loss has per-head regularizer.
• Per-expert IQN τ (different risk preference per head) — clean alternative to §1.4 but interacts with the existing IQN risk-stack. Defer until §1.4 is validated. Add as §1.5 follow-up if needed.
• Per-expert encoder (full MoE / §4.C of parent) — defer to follow-up spec if multi-head with shared encoder is insufficient.
• Cross-fold training (#3) — orthogonal intervention; can be combined with multi-head as a follow-up.

§8. Spec/plan separation discipline (addresses BLOCKER 2)

This spec is INTENTIONALLY separate from 2026-06-02-regime-invariance-four-interventions.md and from any implementation plan. The implementation plan for this spec will be drafted in a separate document AFTER:

1. Intervention #1 (R-multiple) ships and eval verdict is collected
2. If eval verdict is negative, sp-critical-reviewer reviews this spec
3. If reviewer APPROVES, the implementation plan is drafted

Per feedback_no_partial_refactor: spec and plan must align atomically. The plan does not exist yet.

§9. Verified facts (audit trail for sp-critical-reviewer)

The following claims in this spec have been verified by direct code inspection:

Claim	Verification
RL_SLOTS_END = 757 post-grwwh	crates/ml-alpha/src/rl/isv_slots.rs:1771
regime_observer writes slots 696-701 (no trend_strength)	crates/ml-alpha/cuda/rl_regime_observer.cu:36-100 and crates/ml-alpha/src/rl/isv_slots.rs:1522-1534
Reward kernel is rl_fused_reward_pipeline.cu (current Phase 5 logic)	crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu
Per feedback_pi_actor_collapses_without_entropy_floor: entropy floor canonical pattern	pearl exists in MEMORY.md
Per feedback_kill_runs_on_anomaly_quickly: aggressive kill gates required	pearl exists in MEMORY.md
Per feedback_no_partial_refactor: spec/plan must align	pearl exists in MEMORY.md

Claims NOT verified that the reviewer should check:

• Exact PolicyHead struct layout in policy_head.rs (referenced but not opened)
• Exact EXPECTED_LEAVES value pre-this-spec (referenced as 679 from prior R-multiple invariant test)
• Whether mixture-of-policies has any existing precedent in foxhunt (search for prior MoE attempts)