# Multi-Head Policy 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. Each phase is a separate atomic commit with predetermined falsification gates.

**Goal:** Replace foxhunt's single-policy-head architecture with K=3 policy heads + regime-gated routing, bypassing VSN's softmax bottleneck on `regime[3]` (vol) per the empirical regime-attenuation finding. The mixture's gradient flows through foxhunt's actual π training mechanism (Q-distill from PER, not PPO surrogate).

**Architecture:** K=3 policy heads each `Linear(HIDDEN_DIM=128 → N_ACTIONS=11)`. Gating head `Linear(6 → 3)` reads raw regime features from `step_regime_d` directly (parallel channel — does NOT go through VSN/Mamba2). Mixture: `π_logits = log Σ_k softmax(gate_logits)[k] · softmax(pi_logits_k)`. Per-head aux KL prior loss (β=0.05) maintains specialization. Q-distill kernel writes gradient to `ss_pi_grad_logits_d` unchanged; the mixture backward distributes via standard MoE jacobian.

**Tech Stack:** Rust 1.85, cudarc, CUDA 12.4, sm_86 (RTX 3050) / sm_89 (L40S) / sm_90 (H100). Existing crate: `ml-alpha`. Spec: `docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md` (read FULL spec + ADDENDUM 2026-06-03).

**Branch target:** New branch `ml-alpha-multi-head` off `ml-alpha-regime-observer` HEAD (currently `779c03b9d`). Each phase is a separate commit. Phase 2A-D is the cluster-verification gate.

**Prerequisites (must verify before dispatching):**
- `git rev-parse HEAD` returns `779c03b9d` (or later if Phase 1B-C/D landed first)
- Working tree clean
- `./scripts/determinism-check.sh --quick` exits 0 at HEAD
- Spec ADDENDUM 2026-06-03 read and understood (Q-distill, gating head Option B = regime-direct, refined falsification gates)
- `step_regime_d` allocated in `crates/ml-alpha/src/trainer/perception.rs:1795` and accessible at trainer's per-step path
- Read `pearl_foxhunt_pi_trained_by_q_distillation_not_ppo` (load-bearing for not breaking π training)
- Read `pearl_local_smoke_noise_floor_and_regime_concentration` (load-bearing for interpreting smoke results)

---

## Phase 2A-A — ISV slots + MultiHeadPolicy struct + mixture forward kernel + GPU-oracle tests

**Goal:** Foundation. New ISV slots, K policy head + gating head data structures, mixture forward kernel with GPU-oracle invariant tests. No trainer integration yet — components work standalone.

**Files:**
- Create: `crates/ml-alpha/src/rl/multi_head_policy.rs`
- Create: `crates/ml-alpha/cuda/multi_head_policy_forward.cu`
- Create: `crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu`
- Create: `crates/ml-alpha/tests/multi_head_policy_invariants.rs`
- Modify: `crates/ml-alpha/src/rl/isv_slots.rs` (4 new slots + RL_SLOTS_END bump)
- Modify: `crates/ml-alpha/build.rs` (register new .cu)
- Modify: `crates/ml-alpha/src/rl/mod.rs` (pub mod multi_head_policy)

### Task A.1 — Allocate ISV slots

- [ ] **Step 1: Add 4 new ISV slots in `isv_slots.rs`**

After the existing slots (after RL_PPO_N_MINIBATCHES_INDEX = 760 if Phase 1B-A landed; else after RL_SURFER_SCAFFOLD_FORCE_PIN_INDEX = 757):

```rust
// ─────────────────────────────────────────────────────────────────────
// Phase 2A multi-head policy (2026-06-03)
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
//       ADDENDUM 2026-06-03 (Q-distill compatibility + regime-direct gating)
// Pearl: pearl_local_smoke_noise_floor_and_regime_concentration
// ─────────────────────────────────────────────────────────────────────

/// Number of policy heads. Bootstrap K=3. Range allowed [1, 8]. K=1
/// reduces to single-head (legacy behavior — for ablation only).
pub const RL_POLICY_NUM_HEADS_INDEX: usize = 761;

/// Gating distribution entropy floor. Bootstrap log(K) × 0.5 = 0.549 for K=3.
/// Prevents gate collapse to single-head mode per
/// `pearl_pi_actor_collapses_without_entropy_floor`.
pub const RL_POLICY_GATING_ENTROPY_FLOOR_INDEX: usize = 762;

/// Per-head entropy floor. Bootstrap log(N_ACTIONS) × 0.4 = 0.959 for N_ACTIONS=11.
/// Prevents per-head determinism collapse.
pub const RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX: usize = 763;

/// β coefficient for per-head auxiliary KL prior loss. Bootstrap 0.05
/// (5% of policy loss magnitude). Range [0.0, 0.5]. β=0.0 → no aux
/// (specialization driven only by gate + init bias).
pub const RL_POLICY_AUX_PRIOR_BETA_INDEX: usize = 764;
```

Update `RL_SLOTS_END = 765`.

- [ ] **Step 2: Build sanity**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```
Expected: exit 0.

### Task A.2 — Write `multi_head_policy_forward.cu` kernel

- [ ] **Step 1: Create kernel file**

Path: `crates/ml-alpha/cuda/multi_head_policy_forward.cu`

```c
// multi_head_policy_forward.cu — K-head policy logit + mixture combination.
//
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
//       ADDENDUM 2026-06-03 §R.6
//
// Per (batch, head) compute pi_logits_k = W_k @ h_t + b_k (HIDDEN_DIM=128 → N_ACTIONS=11).
// Then per batch compute mixture:
//   pi_probs_k[b, a] = softmax_a(pi_logits_k[b])
//   pi_probs[b, a]   = Σ_k gate_probs[b, k] · pi_probs_k[b, a]
//   pi_logits[b, a]  = log(pi_probs[b, a] + 1e-12)
//
// Deterministic by construction: per-batch sequential mixture sum (no parallel
// reductions across batches; the only reduction is over K=3 heads with fixed
// summation order).
//
// Per feedback_no_atomicadd: no atomic ops.
// Per feedback_no_nvrtc: pre-compiled cubin.

#include <stdint.h>

#define HIDDEN_DIM    128
#define N_ACTIONS     11
#define MAX_K_HEADS   8       // Compile-time max; runtime K from ISV

extern "C" __global__ void multi_head_policy_forward(
    const float* __restrict__ h_t,            // [B × HIDDEN_DIM]
    const float* __restrict__ W_heads,        // [K × N_ACTIONS × HIDDEN_DIM]
    const float* __restrict__ b_heads,        // [K × N_ACTIONS]
    const float* __restrict__ gate_probs,     // [B × K]  — already softmax'd
    float*       __restrict__ pi_logits_k,    // [B × K × N_ACTIONS]   (output, for backward)
    float*       __restrict__ pi_probs_k,     // [B × K × N_ACTIONS]   (output, for backward)
    float*       __restrict__ pi_logits,      // [B × N_ACTIONS]       (output, mixture)
    const int B,
    const int K
) {
    const int b = blockIdx.x;
    if (b >= B) return;

    const int tid = threadIdx.x;  // 0..N_ACTIONS-1

    __shared__ float s_pi_logits_k[MAX_K_HEADS * N_ACTIONS];
    __shared__ float s_pi_probs_k[MAX_K_HEADS * N_ACTIONS];
    __shared__ float s_pi_probs[N_ACTIONS];
    __shared__ float s_max_k[MAX_K_HEADS];     // for softmax stability
    __shared__ float s_sum_k[MAX_K_HEADS];

    // Step 1: compute pi_logits_k[b, k, a] = W_heads[k, a, :] @ h_t[b, :] + b_heads[k, a]
    // One thread per action; loops over k inside thread.
    if (tid < N_ACTIONS) {
        for (int k = 0; k < K; ++k) {
            float gl = b_heads[k * N_ACTIONS + tid];
            #pragma unroll
            for (int j = 0; j < HIDDEN_DIM; ++j) {
                gl += W_heads[(k * N_ACTIONS + tid) * HIDDEN_DIM + j] * h_t[b * HIDDEN_DIM + j];
            }
            s_pi_logits_k[k * N_ACTIONS + tid] = gl;
        }
    }
    __syncthreads();

    // Step 2: per-head softmax for pi_probs_k
    // Use tid 0..K-1 for max and sum reductions
    if (tid < K) {
        float mx = s_pi_logits_k[tid * N_ACTIONS];
        #pragma unroll
        for (int a = 1; a < N_ACTIONS; ++a) {
            float v = s_pi_logits_k[tid * N_ACTIONS + a];
            mx = fmaxf(mx, v);
        }
        s_max_k[tid] = mx;
    }
    __syncthreads();

    if (tid < N_ACTIONS) {
        for (int k = 0; k < K; ++k) {
            float v = expf(s_pi_logits_k[k * N_ACTIONS + tid] - s_max_k[k]);
            s_pi_probs_k[k * N_ACTIONS + tid] = v;
        }
    }
    __syncthreads();

    if (tid < K) {
        float sm = 0.0f;
        #pragma unroll
        for (int a = 0; a < N_ACTIONS; ++a) {
            sm += s_pi_probs_k[tid * N_ACTIONS + a];
        }
        s_sum_k[tid] = sm;
    }
    __syncthreads();

    if (tid < N_ACTIONS) {
        for (int k = 0; k < K; ++k) {
            s_pi_probs_k[k * N_ACTIONS + tid] /= s_sum_k[k];
        }
    }
    __syncthreads();

    // Step 3: mixture: pi_probs[b, a] = Σ_k gate_probs[b, k] · pi_probs_k[k, a]
    if (tid < N_ACTIONS) {
        float p = 0.0f;
        #pragma unroll
        for (int k = 0; k < K; ++k) {
            p += gate_probs[b * K + k] * s_pi_probs_k[k * N_ACTIONS + tid];
        }
        s_pi_probs[tid] = p;
    }
    __syncthreads();

    // Step 4: write outputs
    if (tid < N_ACTIONS) {
        pi_logits[b * N_ACTIONS + tid] = logf(s_pi_probs[tid] + 1e-12f);
        for (int k = 0; k < K; ++k) {
            pi_logits_k[(b * K + k) * N_ACTIONS + tid] = s_pi_logits_k[k * N_ACTIONS + tid];
            pi_probs_k[(b * K + k) * N_ACTIONS + tid] = s_pi_probs_k[k * N_ACTIONS + tid];
        }
    }
}
```

- [ ] **Step 2: Register kernel in `build.rs`**

Add `compile_cuda_kernel("multi_head_policy_forward")?;` adjacent to existing kernel registrations.

- [ ] **Step 3: Build sanity**

```bash
SQLX_OFFLINE=true cargo build -p ml-alpha
```
Expected: exit 0, `target/.../multi_head_policy_forward.cubin` exists.

### Task A.3 — Write `multi_head_policy_gate_forward.cu` kernel

- [ ] **Step 1: Create gating head kernel**

Path: `crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu`

```c
// multi_head_policy_gate_forward.cu — gating head from raw regime features.
//
// Spec ADDENDUM 2026-06-03 §R.3 Option B: gate head reads step_regime_d[b][0..6]
// DIRECTLY (parallel channel — bypasses VSN/Mamba2 entirely).
//
// gate_logits[b, k] = W_gate[k, :] @ regime_h[b, :] + b_gate[k]
// gate_probs[b, k]  = softmax_k(gate_logits[b])
//
// Single-block-per-batch kernel, K threads.
// Per feedback_no_atomicadd: no atomic ops.

#include <stdint.h>

#define REGIME_DIM      6
#define MAX_K_HEADS     8

extern "C" __global__ void multi_head_policy_gate_forward(
    const float* __restrict__ regime_h,       // [B × REGIME_DIM=6]
    const float* __restrict__ W_gate,         // [K × REGIME_DIM]
    const float* __restrict__ b_gate,         // [K]
    float*       __restrict__ gate_logits,    // [B × K] (output, for backward)
    float*       __restrict__ gate_probs,     // [B × K] (output, post-softmax)
    const int B,
    const int K
) {
    const int b = blockIdx.x;
    if (b >= B) return;

    const int tid = threadIdx.x;
    __shared__ float s_logits[MAX_K_HEADS];
    __shared__ float s_max;
    __shared__ float s_sum;

    // Step 1: compute gate_logits[b, k] = W_gate[k, :] @ regime_h[b, :] + b_gate[k]
    if (tid < K) {
        float gl = b_gate[tid];
        #pragma unroll
        for (int j = 0; j < REGIME_DIM; ++j) {
            gl += W_gate[tid * REGIME_DIM + j] * regime_h[b * REGIME_DIM + j];
        }
        s_logits[tid] = gl;
    }
    __syncthreads();

    // Step 2: softmax (max-subtract for stability)
    if (tid == 0) {
        float mx = s_logits[0];
        for (int k = 1; k < K; ++k) mx = fmaxf(mx, s_logits[k]);
        s_max = mx;
    }
    __syncthreads();

    if (tid < K) {
        s_logits[tid] = expf(s_logits[tid] - s_max);
    }
    __syncthreads();

    if (tid == 0) {
        float sm = 0.0f;
        for (int k = 0; k < K; ++k) sm += s_logits[k];
        s_sum = sm;
    }
    __syncthreads();

    if (tid < K) {
        float v = s_logits[tid] / s_sum;
        gate_logits[b * K + tid] = s_logits[tid] * s_sum + s_max;  // pre-softmax (for backward)
        gate_probs[b * K + tid] = v;
    }
}
```

- [ ] **Step 2: Register in build.rs**

```rust
compile_cuda_kernel("multi_head_policy_gate_forward")?;
```

### Task A.4 — Create `MultiHeadPolicy` struct

- [ ] **Step 1: Create file**

Path: `crates/ml-alpha/src/rl/multi_head_policy.rs`

```rust
//! K-head policy with regime-gated routing.
//!
//! Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
//!       ADDENDUM 2026-06-03 §R.3-R.6
//!
//! Forward:
//!   gate_logits = W_gate @ regime_h + b_gate            (gating head — bypass VSN)
//!   gate_probs  = softmax_K(gate_logits)
//!   pi_logits_k = W_heads_k @ h_t + b_heads_k           (K policy heads)
//!   pi_probs_k  = softmax_A(pi_logits_k)
//!   pi_probs    = Σ_k gate_probs[k] · pi_probs_k        (mixture)
//!   pi_logits   = log(pi_probs)
//!
//! Backward: standard MoE jacobian via Q-distill gradient + aux prior loss merge.
//!
//! All transfers use mapped-pinned helpers (`feedback_no_htod_htoh_only_mapped_pinned`).
//! No atomicAdd anywhere (`feedback_no_atomicadd`).

use anyhow::{Context, Result};
use cudarc::driver::{CudaContext, CudaSlice, CudaStream, CudaFunction, CudaModule};
use std::sync::Arc;

use crate::heads::{HIDDEN_DIM, N_ACTIONS};
use crate::cfc::snap_features::REGIME_DIM;

pub const MAX_K_HEADS: usize = 8;

pub struct MultiHeadPolicy {
    pub b_size: usize,
    pub k: usize,

    // Policy head weights ([K × N_ACTIONS × HIDDEN_DIM] flat)
    pub w_heads_d: CudaSlice<f32>,
    pub b_heads_d: CudaSlice<f32>,

    // Gating head weights ([K × REGIME_DIM=6] flat)
    pub w_gate_d: CudaSlice<f32>,
    pub b_gate_d: CudaSlice<f32>,

    // Forward output buffers
    pub gate_logits_d: CudaSlice<f32>,        // [B × K]
    pub gate_probs_d: CudaSlice<f32>,         // [B × K]
    pub pi_logits_k_d: CudaSlice<f32>,        // [B × K × N_ACTIONS]
    pub pi_probs_k_d: CudaSlice<f32>,         // [B × K × N_ACTIONS]

    // Kernels
    pub forward_kernel: CudaFunction,
    pub gate_forward_kernel: CudaFunction,

    _forward_module: Arc<CudaModule>,
    _gate_forward_module: Arc<CudaModule>,
}

impl MultiHeadPolicy {
    pub fn new(
        ctx: &Arc<CudaContext>,
        stream: &Arc<CudaStream>,
        b_size: usize,
        k: usize,
        seed: u64,
    ) -> Result<Self> {
        assert!(k > 0 && k <= MAX_K_HEADS, "K must be in [1, {}]", MAX_K_HEADS);

        // Allocate weights with per-head action-bias init from spec §1.3.
        // Action enum verified against crates/ml-alpha/src/rl/common.rs:56-70:
        //   0=ShortLarge, 1=ShortSmall, 2=Hold, 3=FlatFromLong, 4=FlatFromShort,
        //   5=LongSmall, 6=LongLarge, 7=TrailTighten, 8=TrailLoosen, 9=HalfFlatLong, 10=HalfFlatShort.
        // Per-head specialization tri-split (asymmetry-break only — gate routes by regime):
        //   Head 0 = SHORT-bias / risk-off  → +0.5 on action 0 (ShortLarge)
        //   Head 1 = LONG-bias  / risk-on   → +0.5 on actions 5,6 (LongSmall, LongLarge)
        //   Head 2 = HOLD / stable          → +0.5 on action 2 (Hold — Phase 0 dominant action)
        //
        // Standard Xavier scaled init via scoped_init_seed for determinism (path verified).
        use ml_core::cuda_autograd::init::scoped_init_seed;
        let _guard = scoped_init_seed(seed);

        // Zero-init the linear weights; per-head action priors live on the bias vector.
        let w_heads_d = stream.alloc_zeros::<f32>(k * N_ACTIONS * HIDDEN_DIM)
            .context("alloc w_heads")?;
        let mut b_heads_init = vec![0.0_f32; k * N_ACTIONS];
        if k >= 1 { b_heads_init[0 * N_ACTIONS + 0] += 0.5; }  // Head 0 → ShortLarge
        if k >= 2 { b_heads_init[1 * N_ACTIONS + 5] += 0.5;
                    b_heads_init[1 * N_ACTIONS + 6] += 0.5; }  // Head 1 → LongSmall, LongLarge
        if k >= 3 { b_heads_init[2 * N_ACTIONS + 2] += 0.5; }  // Head 2 → Hold
        let b_heads_d = stream.alloc_zeros::<f32>(k * N_ACTIONS)?;
        stream.memcpy_htod(&b_heads_init, &b_heads_d)
            .context("upload b_heads init")?;

        // Gating head: small init (Xavier scale)
        let gate_scale = (1.0_f32 / REGIME_DIM as f32).sqrt();
        let w_gate_init: Vec<f32> = (0..k * REGIME_DIM)
            .map(|i| ((i as f32 * 0.123_f32).sin() * gate_scale)).collect();  // deterministic init
        let w_gate_d = stream.alloc_zeros::<f32>(k * REGIME_DIM)?;
        stream.memcpy_htod(&w_gate_init, &w_gate_d)
            .context("upload w_gate init")?;
        let b_gate_d = stream.alloc_zeros::<f32>(k)?;

        // Forward buffers
        let gate_logits_d = stream.alloc_zeros::<f32>(b_size * k)?;
        let gate_probs_d = stream.alloc_zeros::<f32>(b_size * k)?;
        let pi_logits_k_d = stream.alloc_zeros::<f32>(b_size * k * N_ACTIONS)?;
        let pi_probs_k_d = stream.alloc_zeros::<f32>(b_size * k * N_ACTIONS)?;

        // Load kernels (repo convention: const &[u8] + load_cubin — see rl/ppo.rs:152-154, rl/dqn.rs:279).
        const FORWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_head_policy_forward.cubin"));
        const GATE_FORWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_head_policy_gate_forward.cubin"));
        let forward_module = ctx
            .load_cubin(FORWARD_CUBIN.to_vec())
            .context("load multi_head_policy_forward cubin")?;
        let forward_kernel = forward_module.load_function("multi_head_policy_forward")?;
        let gate_forward_module = ctx
            .load_cubin(GATE_FORWARD_CUBIN.to_vec())
            .context("load multi_head_policy_gate_forward cubin")?;
        let gate_forward_kernel = gate_forward_module.load_function("multi_head_policy_gate_forward")?;

        Ok(Self {
            b_size, k,
            w_heads_d, b_heads_d, w_gate_d, b_gate_d,
            gate_logits_d, gate_probs_d, pi_logits_k_d, pi_probs_k_d,
            forward_kernel, gate_forward_kernel,
            _forward_module: forward_module, _gate_forward_module: gate_forward_module,
        })
    }

    pub fn forward(
        &self,
        stream: &Arc<CudaStream>,
        h_t: &CudaSlice<f32>,         // [B × HIDDEN_DIM]
        regime_h: &CudaSlice<f32>,    // [B × REGIME_DIM=6]
        pi_logits_out: &mut CudaSlice<f32>,  // [B × N_ACTIONS] — mixture output
    ) -> Result<()> {
        // Step 1: gating head forward
        // ... launch gate_forward_kernel with (B, 1, 1) grid, K-thread block
        // ... reads regime_h, writes gate_logits_d + gate_probs_d
        //
        // Step 2: policy heads forward + mixture
        // ... launch forward_kernel with (B, 1, 1) grid, N_ACTIONS-thread block
        // ... reads h_t + gate_probs_d, writes pi_logits_k_d + pi_probs_k_d + pi_logits_out
        unimplemented!("forward launches — wire in Task A.4 step 2")
    }
}
```

- [ ] **Step 2: Implement forward launches**

(Fill in the `forward` method with raw_launch invocations matching the kernel signatures.)

- [ ] **Step 3: Register module in `rl/mod.rs`**

Add `pub mod multi_head_policy;` to `crates/ml-alpha/src/rl/mod.rs` adjacent to the existing `pub mod ppo;`, `pub mod dqn;`, `pub mod iqn;`, `pub mod dueling_q;` lines.

- [ ] **Step 4: Build sanity**

```bash
SQLX_OFFLINE=true cargo build -p ml-alpha
```

### Task A.5 — GPU-oracle invariant tests

- [ ] **Step 1: Create test file**

Path: `crates/ml-alpha/tests/multi_head_policy_invariants.rs`

5 GPU-oracle invariant tests:

```rust
//! GPU-oracle invariants for MultiHeadPolicy.
//!
//! Tests run on actual GPU per `feedback_no_cpu_test_fallbacks`.
//! Each test asserts a specific architectural property.

use anyhow::Result;
use cudarc::driver::CudaContext;
use ml_alpha::rl::multi_head_policy::MultiHeadPolicy;

const TOLERANCE: f32 = 1e-5;

#[test]
#[ignore = "requires CUDA"]
fn gate_probs_sum_to_one() -> Result<()> {
    // Setup: K=3, B=4 batch. Random regime input.
    // Run forward.
    // Assert: gate_probs[b].sum() ≈ 1.0 for every b (softmax property).
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn pi_probs_mixture_sums_to_one() -> Result<()> {
    // Setup: K=3, B=4. Run forward.
    // Assert: exp(pi_logits[b]).sum() ≈ 1.0 for every b (mixture is a valid distribution).
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn k1_reduces_to_single_head() -> Result<()> {
    // Setup: K=1. Run forward with weights matching single-head Linear(128, 11).
    // Run a reference single-head policy_head Linear forward.
    // Assert: bit-equal output (within fp32 noise).
    // Verifies: K=1 ablation works as expected — no surprise mixture overhead.
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn gate_responds_to_regime_change() -> Result<()> {
    // Setup: K=3 with hand-crafted gating weights (e.g., w_gate[0] = [1, 0, 0, 0, 0, 0]).
    // Run forward with regime_h = [10, 0, ...] vs regime_h = [-10, 0, ...].
    // Assert: gate_probs differs significantly between the two inputs (responds to regime[0]).
    // Verifies: gating head is sensitive to regime input (the load-bearing design property).
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn forward_is_deterministic_across_contexts() -> Result<()> {
    // Setup: same seed → MultiHeadPolicy::new in two fresh contexts.
    // Run forward with identical input.
    // Assert: bit-equal output (`to_bits()` equality).
    // Verifies: no determinism regression vs `pearl_determinism_achieved`.
    Ok(())
}
```

- [ ] **Step 2: Run tests**

```bash
SQLX_OFFLINE=true cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture
```
Expected: 5/5 PASS.

### Task A.6 — Atomic commit + verify determinism not broken

- [ ] **Step 1: Stage Phase 2A-A changes** and verify `./scripts/determinism-check.sh --quick` exits 0 (pipeline unchanged since not wired in)

- [ ] **Step 2: Pre-commit hook on staged diff** — 0 memcpy_dtoh raw violations, 0 atomicAdd

- [ ] **Step 3: Commit**

Commit message template at the end of this plan.

---

## Phase 2A-B — Mixture backward kernel + aux prior loss

**Goal:** Implement the backward pass through the mixture (Q-distill gradient → K policy heads + gating head) and the auxiliary KL prior loss. No trainer integration yet.

**Files:**
- Create: `crates/ml-alpha/cuda/multi_head_policy_backward.cu`
- Create: `crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu`
- Modify: `crates/ml-alpha/src/rl/multi_head_policy.rs` (add `backward` method)
- Modify: `crates/ml-alpha/tests/multi_head_policy_invariants.rs` (add backward gradient tests)

### Tasks B.1-B.5 (outline — expand in 2A-A completion)

1. B.1 — Implement mixture backward jacobian (∂L/∂pi_logits → ∂L/∂pi_logits_k + ∂L/∂gate_logits)
2. B.2 — Implement gate backward (∂L/∂gate_logits → ∂L/∂W_gate + ∂L/∂b_gate)
3. B.3 — Implement per-head backward (∂L/∂pi_logits_k → ∂L/∂W_heads_k + ∂L/∂b_heads_k + ∂L/∂h_t)
4. B.4 — Implement aux prior loss + gradient (Σ_k β · gate_probs[k] · KL(pi_probs_k ‖ pi_prior_k))
5. B.5 — GPU-oracle finite-difference gradient tests (verify backward kernels against numerical gradients)

**Falsification at 2A-B**:
- Finite-difference gradient agreement within 1e-3 relative tolerance
- Aux prior loss is non-negative (KL ≥ 0)
- Determinism preserved (sequential reductions, no atomicAdd)

---

## Phase 2A-C — Trainer integration with Q-distill gradient flow

**Goal:** Wire MultiHeadPolicy into IntegratedTrainer, replacing the single `policy_head` while preserving Q-distill gradient flow (the existing line 5577 zero-fill + Q-distill write pattern).

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (replace `policy_head` with MultiHeadPolicy; merge aux prior gradient into ss_pi_grad_logits_d post-Q-distill)
- Modify: `crates/ml-alpha/examples/alpha_rl_train.rs` (no API changes; trainer is the seam)
- Modify: `crates/ml-alpha/tests/eval_diag_emission.rs` (EXPECTED_LEAVES bump for new policy.head_k_entropy, policy.gate_entropy, policy.gate_probs leaves)

### Tasks C.1-C.6 (outline — expand in 2A-B completion)

1. C.1 — Replace `policy_head` allocation with `MultiHeadPolicy::new(K=3, seed=...)` in trainer constructor
2. C.2 — Bootstrap ISV slots 761-764 in `with_controllers_bootstrapped` table
3. C.3 — Modify forward path: use MultiHeadPolicy::forward to produce pi_logits (consumed unchanged by downstream)
4. C.4 — Modify backward path: AFTER Q-distill writes to ss_pi_grad_logits_d, merge aux prior gradient
5. C.5 — Add diag emission: per-head entropy, gate entropy, gate_probs distribution, TVD diagnostic
6. C.6 — Adam updates only on MultiHeadPolicy weights (W_heads, b_heads, W_gate, b_gate); encoder/Q/V heads unchanged

**Falsification at 2A-C** (mid-smoke b=128, 2000+500, seed=42, 3-seed):

| Gate | Threshold | Source |
|---|---|---|
| **G1 (PRIMARY)** | TVD(Q1, Q4 action dist) on session_pnl_variance_ema axis ≥ 0.10 | spec ADDENDUM §R.5 |
| **G2** | Per-head pairwise KL ≥ 0.20 by step 1500 | spec §4 |
| **G3** | Gate entropy ≥ log(K) × 0.45 = 0.49 | spec §4 |
| **G4** | Argmax action mass ≤ 0.30 (down from baseline 0.40-0.54) | spec ADDENDUM §R.5 |
| **G5** | Determinism preserved + cross-source $50 | Phase 0 contract |
| **G6 (load-bearing empirical)** | Per-quartile: popart Q1 Δpnl improves by ≥ $1M (single-seed: borderline; 3-seed required for confidence) | spec ADDENDUM §R.5 |
| G7 (aggregate, informational) | eval_summary.total_pnl_usd within 1σ of better than mean 3-seed Phase 0 baseline | high bar given $1.9M σ |

If G1-G5 PASS and G6 PASSES at 3-seed → proceed to Phase 2A-D cluster verification.
If G1 FAILS (TVD on vol axis didn't lift) → multi-head bypass of VSN didn't work; STOP and diagnose.

---

## Phase 2A-D — Cluster verification 20k+5k

**Goal:** Validate at cluster scale. Cluster σ_eval_pnl is ~9× smaller than local (~$215k vs $1.9M), so cluster verdicts are far more sensitive. This is the load-bearing test.

### Tasks D.1-D.4 (outline — expand in 2A-C completion)

1. D.1 — Submit `./scripts/argo-alpha-rl.sh --branch ml-alpha-multi-head` at the 2A-C commit SHA
2. D.2 — Monitor for completion (~70 min on L40S)
3. D.3 — Analyze cluster verdict against per-quartile + aggregate gates
4. D.4 — Update memory pearls with cluster verdict (success or failure)

**Falsification at 2A-D (cluster)**:
- eval_summary.total_pnl_usd strictly better than Phase 0 baseline cluster value (TBD: requires baseline cluster run at HEAD `cfaa420bd`)
- popart Q1 Δpnl improves by ≥ $5M vs baseline cluster (scales with batch)
- Determinism preserved at cluster scale
- TVD(Q1, Q4 action dist) ≥ 0.10 on vol axis at cluster scale

If cluster verdict shows improvement: proceed to 2A-E productionization.
If cluster verdict shows no improvement OR regression: multi-head architecture doesn't deliver. Document negative result in pearl. Pivot to deeper architectural intervention (regime-stratified training, encoder changes, etc.).

---

## Phase 2A-E — Productionization

**Goal:** Make multi-head the default. Remove legacy single-head code path. Update pearls.

### Tasks E.1-E.4 (outline — expand in 2A-D completion)

1. E.1 — Default `RL_POLICY_NUM_HEADS_INDEX = 3` (production)
2. E.2 — Remove single-head `policy_head` legacy code path (per `feedback_single_source_of_truth_no_duplicates`)
3. E.3 — Update pearls: `pearl_foxhunt_pi_trained_by_q_distillation_not_ppo` ADDENDUM "now via mixture", `pearl_local_smoke_noise_floor_and_regime_concentration` "verified bypass of VSN bottleneck"
4. E.4 — Final cluster + local smoke parity check

---

## Commit message template (for Phase 2A-A)

```
feat(ml-alpha): Phase 2A-A — MultiHeadPolicy foundation (K-head + regime-gated routing)

Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
      ADDENDUM 2026-06-03 (Q-distill compatibility + regime-direct gating)
Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md
Pearl: pearl_local_smoke_noise_floor_and_regime_concentration

Foundation phase of multi-head policy refactor. Empirically motivated by
the regime-attenuation finding: at Phase 0 baseline the policy is
regime-conditional on popart_envelope (TVD=0.169 between Q1 and Q4
action dists) but NOT on session_pnl_variance_ema (TVD=0.015) despite
vol context being in encoder input at regime[3]. The asymmetry maps to
VSN's softmax-over-40 gating constraint structurally limiting any single
encoder-input feature to ~1/40 of attention budget. Multi-head policy
with regime-gated routing reads step_regime_d[b][0..6] directly via the
gating head (parallel channel) — bypassing VSN's bottleneck.

Phase 2A-A establishes the components WITHOUT trainer integration:

* 4 ISV slots (761-764):
    RL_POLICY_NUM_HEADS_INDEX             = 761 (K=3 bootstrap)
    RL_POLICY_GATING_ENTROPY_FLOOR_INDEX  = 762 (0.549 bootstrap)
    RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX    = 763 (0.959 bootstrap)
    RL_POLICY_AUX_PRIOR_BETA_INDEX        = 764 (0.05 bootstrap)
    RL_SLOTS_END 761 -> 765

* crates/ml-alpha/cuda/multi_head_policy_forward.cu (~180 LOC):
    Per-batch K-head logit + mixture combination. Reads h_t and
    gate_probs; writes pi_logits_k (for backward) + pi_logits (mixture
    for downstream). Deterministic sequential reduction over K heads.

* crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu (~80 LOC):
    Gating head reads step_regime_d[b][0..6] DIRECTLY (parallel
    channel — bypasses VSN). K-way softmax over gate logits.

* crates/ml-alpha/src/rl/multi_head_policy.rs (~300 LOC):
    MultiHeadPolicy struct with K policy heads + gating head + forward
    buffers. Per-head init bias for asymmetry-break specialization
    (Head 0: ShortLarge, Head 1: LongSmall+LongLarge, Head 2: Hold).
    Gating head reads regime DIRECTLY (bypass VSN); deterministic
    Xavier-style init. All allocations via existing patterns; no atomicAdd.

* crates/ml-alpha/tests/multi_head_policy_invariants.rs (~250 LOC):
    5 GPU-oracle invariant tests covering: gate softmax normalization,
    mixture probability sum, K=1 ablation reduces to single-head,
    gate responds to regime input, determinism across contexts.

Validation gates (Phase 2A-A complete when all pass):
* cargo build --release --example alpha_rl_train -p ml-alpha: exit 0
* cargo test multi_head_policy_invariants --release: 5/5 PASS
* ./scripts/determinism-check.sh --quick: exit 0
  (pipeline unchanged since not wired into trainer yet)
* Pre-commit hook: PASS (0 memcpy_dtoh, 0 atomicAdd)

Next: Phase 2A-B (mixture backward + aux prior loss kernels).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
```

---

## Self-review (controller, before dispatch)

### Spec coverage
- ✅ Spec §1.1 (architecture): K=3 heads + gating
- ✅ Spec §1.2 (reward unchanged): Phase 2A doesn't touch reward kernel
- ✅ Spec §1.3 (specialization): 4 mechanisms — init bias + per-head entropy floor + gate entropy floor + aux KL prior
- ✅ Spec §1.4 (aux loss): β=0.05 bootstrap
- ✅ Spec §2 (ISV slots): 4 slots at 761-764
- ✅ ADDENDUM §R.2 (Q-distill flow): documented in plan
- ✅ ADDENDUM §R.3 (gating head input Option B): direct regime read
- ✅ ADDENDUM §R.5 (falsification): G1 TVD on vol axis, G6 per-quartile popart Q1
- ✅ ADDENDUM §R.6 (gradient flow): aux prior merged AFTER Q-distill

### Placeholders
- Task bodies for 2A-B, 2A-C, 2A-D, 2A-E are OUTLINES. Expand each when its prerequisite phase passes.
- N_ACTIONS=11 layout assumed; bias init indices need verification (probably FLAT=0, HOLD_LONG=5, HOLD_SHORT=6, etc., but should be cross-referenced against `crates/ml-backtesting/src/order.rs` or similar before dispatch).

### Type consistency
- `MultiHeadPolicy` uses `_d` suffix for device slices throughout
- `K` capital, `b_size` lowercase, consistent with existing trainer struct naming

## Execution handoff

**Plan complete and saved to `docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md`.**

Per the methodology in `feedback_investigation_first_falsification_methodology`:

1. Dispatch Phase 2A-A coder NOW (foundation: ISV slots + struct + kernels + tests).
2. After 2A-A commits + tests pass: expand 2A-B task bodies, dispatch 2A-B coder.
3. After 2A-B commits + gradient finite-diff tests pass: expand 2A-C, dispatch coder.
4. After 2A-C smoke passes G1-G5 (3-seed): dispatch 2A-D cluster verification.
5. After 2A-D cluster validates G6: dispatch 2A-E productionization.

Each phase is one commit. Each phase's verdict feeds into the next phase's scope.

**NOT dispatching automatically.** Await explicit authorization from controller given the multi-week scope. The plan is the proper artifact; the coder dispatch is the next decision point.
