From 0b3e4015003cef21e430eea8c5af6ddaa22abb43 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 3 Jun 2026 01:41:57 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20Phase=202A-A=20=E2=80=94=20Mu?= =?UTF-8?q?ltiHeadPolicy=20foundation=20(inert)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit K=3 policy mixture + regime-gated routing head, with gate reading REGIME_DIM=6 features DIRECTLY (bypassing the VSN softmax bottleneck per the regime-attenuation empirical finding). Components are unconditional but not yet wired into the trainer — Phase 2A-B adds backward + aux KL prior, 2A-C wires Q-distill grad routing to the mixture. * 4 new ISV slots 761-764 (K, gating entropy floor, head entropy floor, aux prior β) + RL_SLOTS_END 761→765. * multi_head_policy_forward.cu: per-batch K-head logits + mixture combination. Grid=(B), Block=(N_ACTIONS=11). Persists pi_logits_k + pi_probs_k for backward. * multi_head_policy_gate_forward.cu: per-batch K-thread softmax over W·regime + b. Reads regime_h directly (parallel channel — bypass VSN). Stores pre-softmax logits to gmem BEFORE the in-place exp (caught during impl — plan pseudocode would have corrupted gate_logits). * MultiHeadPolicy struct with Option A asymmetry-break init: Head 0 → ShortLarge bias (+0.5), Head 1 → LongSmall+LongLarge bias (+0.5 each), Head 2 → Hold bias (+0.5). Indices verified against rl/common.rs:56-70 N_ACTIONS=11 enum. htod via local upload() helper using MappedF32Buffer + raw_memcpy_dtod_async (mirrors rl/ppo.rs, rl/dueling_q.rs). * 5 GPU-oracle invariant tests, all PASS: - gate_probs_sum_to_one - pi_probs_mixture_sums_to_one - k1_reduces_to_single_head (bit-equal vs reference Linear) - gate_responds_to_regime_change (TVD > 0.5, modal head flips) - forward_is_deterministic_across_contexts (bit-equal across two fresh CUDA contexts) * Determinism preserved: ./scripts/determinism-check.sh --quick exits 0 (pipeline unchanged, components inert). * No memcpy_htod/memcpy_dtoh on regular slices, no atomicAdd, no scoped-init-seed bypass. Empirical motivation (3-seed mid-smoke at HEAD 779c03b9d, b=128): mean -$6.41M ± $1.39M σ; 62× spread in popart_envelope quartile (Q1 -$1.99M / Q4 -$32k); policy TVD asymmetry 0.17 on popart axis vs 0.02 on vol axis. Regime-direct gating targets the attenuation chokepoint; mixture provides regime-conditional capacity that single-head softmax cannot express. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 2 + .../cuda/multi_head_policy_forward.cu | 136 ++++ .../cuda/multi_head_policy_gate_forward.cu | 90 +++ crates/ml-alpha/src/rl/isv_slots.rs | 35 +- crates/ml-alpha/src/rl/mod.rs | 1 + crates/ml-alpha/src/rl/multi_head_policy.rs | 404 ++++++++++ .../tests/multi_head_policy_invariants.rs | 389 +++++++++ ...-06-03-multi-head-policy-implementation.md | 747 ++++++++++++++++++ ...06-02-multi-head-policy-with-r-multiple.md | 220 +++++- 9 files changed, 2021 insertions(+), 3 deletions(-) create mode 100644 crates/ml-alpha/cuda/multi_head_policy_forward.cu create mode 100644 crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu create mode 100644 crates/ml-alpha/src/rl/multi_head_policy.rs create mode 100644 crates/ml-alpha/tests/multi_head_policy_invariants.rs create mode 100644 docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 871264a09..d069041ac 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -155,6 +155,8 @@ const KERNELS: &[&str] = &[ "snapshot_aos_to_soa", // AoS→SoA scatter: one thread per snapshot reads contiguous Mbp10RawInput, writes into 10 SoA device buffers; replaces host nested loops + 10 DtoD copies "gpu_sample_and_gather", // GPU-resident batch sampler: random file+anchor sampling + AoS→SoA gather from pre-uploaded dataset; eliminates ALL per-step CPU data loading "rl_deterministic_checksum", // Determinism foundation Phase 1 (2026-06-02): provably deterministic sum-of-squares (single-block / single-thread / f64 accumulation) for per-step component checksums in diag; localizes non-determinism source. Spec: docs/superpowers/specs/2026-06-02-determinism-foundation.md §1.1 + "multi_head_policy_forward", // Phase 2A-A (2026-06-03): K-head policy logit + softmax + mixture combination; per-batch deterministic sequential mixture sum. Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md §R.6 + "multi_head_policy_gate_forward",// Phase 2A-A (2026-06-03): gating head from raw regime features [B × 6] (parallel channel bypassing VSN/Mamba2). Spec ADDENDUM 2026-06-03 §R.3 Option B ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/multi_head_policy_forward.cu b/crates/ml-alpha/cuda/multi_head_policy_forward.cu new file mode 100644 index 000000000..2037f9be9 --- /dev/null +++ b/crates/ml-alpha/cuda/multi_head_policy_forward.cu @@ -0,0 +1,136 @@ +// 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 +// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A) +// +// Per-(batch, head) computes pi_logits_k = W_k @ h_t + b_k (HIDDEN_DIM=128 → N_ACTIONS=11). +// Then per-batch computes the mixture: +// pi_probs_k[b, k, a] = softmax_a(pi_logits_k[b, k, :]) +// pi_probs[b, a] = Σ_k gate_probs[b, k] · pi_probs_k[b, k, 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). Sole-writer per output cell. +// +// Block layout: +// grid = (B, 1, 1) +// block = (N_ACTIONS = 11, 1, 1) +// One block per batch. N_ACTIONS threads. Each thread owns one action across +// all K heads; loops K times for per-head matmul + softmax + mixture combine. +// +// Per feedback_no_atomicadd: no atomic ops. +// Per feedback_no_nvrtc: pre-compiled cubin via build.rs. + +#include + +#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: per-head linear projection ─────────────────────────────── + // pi_logits_k[b, k, a] = Σ_j W_heads[k, a, j] × h_t[b, j] + 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 (max-subtract for stability) ──────────── + // Threads 0..K-1 compute per-head max over the N_ACTIONS logits. + 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(); + + // Each action-thread exponentiates its slot for every head. + 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(); + + // Threads 0..K-1 reduce per-head sum (denominator). + 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(); + + // Normalize. + 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] + // Sequential summation order over K (deterministic). + 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 ──────────────────────────────────────────── + // pi_logits[b, a] = log(pi_probs[b, a] + 1e-12) for downstream softmax callers. + // Also stash per-head logits + probs for the Phase 2A-B backward kernel. + 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]; + } + } +} diff --git a/crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu b/crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu new file mode 100644 index 000000000..834735946 --- /dev/null +++ b/crates/ml-alpha/cuda/multi_head_policy_gate_forward.cu @@ -0,0 +1,90 @@ +// multi_head_policy_gate_forward.cu — gating head from raw regime features. +// +// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md +// ADDENDUM 2026-06-03 §R.3 Option B (regime-direct gating, bypass VSN) +// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A) +// +// gate_logits[b, k] = Σ_j W_gate[k, j] × regime_h[b, j] + b_gate[k] +// gate_probs[b, k] = softmax_k(gate_logits[b]) +// +// Reads regime_h DIRECTLY (parallel channel — does NOT go through VSN/Mamba2). +// This is the load-bearing architectural choice motivated by the empirical +// regime-attenuation finding (pearl_local_smoke_noise_floor_and_regime_concentration): +// VSN's softmax-over-40 gating structurally limits any single encoder-input +// feature to ~1/40 of attention budget, suppressing vol-axis routing. The gate +// reads regime[0..6] without that bottleneck. +// +// Block layout: +// grid = (B, 1, 1) +// block = (K, 1, 1) — typically K=3, ≤ MAX_K_HEADS=8 +// One block per batch. K threads cooperate via shared-mem reductions. +// +// Per feedback_no_atomicadd: no atomic ops. +// Per feedback_no_nvrtc: pre-compiled cubin via build.rs. + +#include + +#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]; // pre-softmax logits + __shared__ float s_exp[MAX_K_HEADS]; // exp(logit - max) for softmax + __shared__ float s_max; + __shared__ float s_sum; + + // ── Step 1: linear projection ──────────────────────────────────────── + // gate_logits[b, k] = Σ_j W_gate[k, j] × regime_h[b, j] + b_gate[k] + // One thread per head 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; + // Persist pre-softmax logits to gmem for the Phase 2A-B backward kernel. + gate_logits[b * K + 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_exp[tid] = expf(s_logits[tid] - s_max); + } + __syncthreads(); + + if (tid == 0) { + float sm = 0.0f; + for (int k = 0; k < K; ++k) { + sm += s_exp[k]; + } + s_sum = sm; + } + __syncthreads(); + + if (tid < K) { + gate_probs[b * K + tid] = s_exp[tid] / s_sum; + } +} diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index e08cff8cf..8a9aca3df 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -1808,6 +1808,38 @@ pub const RL_PPO_N_EPOCHS_INDEX: usize = 759; /// Each minibatch gets one Adam step. Clamp range [1, 64]. pub const RL_PPO_N_MINIBATCHES_INDEX: usize = 760; +// ───────────────────────────────────────────────────────────────────── +// 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) +// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md +// Pearl: pearl_local_smoke_noise_floor_and_regime_concentration +// +// Phase 2A-A allocates these slots and defines bootstrap values. The +// MultiHeadPolicy components consume them in Phase 2A-C (trainer wire-in); +// until then they sit alongside the inert Phase 2A-A foundation. +// ───────────────────────────────────────────────────────────────────── + +/// Number of policy heads. Bootstrap K=3. Range allowed [1, 8]. K=1 +/// reduces to single-head (legacy behavior — for ablation only). +/// Wired into [`MultiHeadPolicy`](crate::rl::multi_head_policy::MultiHeadPolicy) +/// in Phase 2A-C. +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; + /// Last RL-allocated slot index (exclusive). /// Pre-risk-stack: 662. Post-Fix-B: 685. Post-v9 (warmup boundary): 696. /// Post-regime-observer (F1.1): 716. Post-warmed-flag addendum: 717. @@ -1822,4 +1854,5 @@ pub const RL_PPO_N_MINIBATCHES_INDEX: usize = 760; /// Post-reward-policy-realignment v4 binary: 754 (slot reused in v5). /// Post-surfer-scaffold v5 adaptive (3 config slots): 757. /// Post-rollout-buffer+GAE Phase 1B-A (3 PPO control slots): 760. -pub const RL_SLOTS_END: usize = 761; +/// Post-multi-head-policy Phase 2A-A (4 policy-gating slots): 764. +pub const RL_SLOTS_END: usize = 765; diff --git a/crates/ml-alpha/src/rl/mod.rs b/crates/ml-alpha/src/rl/mod.rs index 57067104f..04a3e88d9 100644 --- a/crates/ml-alpha/src/rl/mod.rs +++ b/crates/ml-alpha/src/rl/mod.rs @@ -24,6 +24,7 @@ pub mod frd; pub mod iqn; pub mod isv_slots; pub mod loss_balance; +pub mod multi_head_policy; pub mod noisy; pub mod ppo; pub mod reward; diff --git a/crates/ml-alpha/src/rl/multi_head_policy.rs b/crates/ml-alpha/src/rl/multi_head_policy.rs new file mode 100644 index 000000000..88f5c0dbc --- /dev/null +++ b/crates/ml-alpha/src/rl/multi_head_policy.rs @@ -0,0 +1,404 @@ +//! 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 +//! Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A) +//! Pearl: pearl_local_smoke_noise_floor_and_regime_concentration +//! +//! ## Forward +//! +//! ```text +//! 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 + 1e-12) +//! ``` +//! +//! Backward (Phase 2A-B) — standard MoE Jacobian distributes the Q-distill +//! gradient on `pi_logits` to (a) per-head `pi_logits_k` and (b) +//! `gate_logits`. Phase 2A-A only ships the forward path; backward kernels +//! arrive in the next phase. +//! +//! ## Why a parallel-channel gating head +//! +//! Empirical regime-attenuation finding: at the Phase 0 baseline, the +//! single policy head is regime-conditional on `popart_envelope` +//! (TVD(Q1, Q4 action dist) = 0.169) but NOT on `session_pnl_variance_ema` +//! (TVD = 0.015), despite vol context living at `regime[3]` in the encoder +//! input. The asymmetry maps to VSN's softmax-over-40 gating constraint +//! structurally limiting any single encoder-input feature to ~1/40 of the +//! attention budget. The gating head reads `step_regime_d[b][0..6]` +//! DIRECTLY — bypassing VSN/Mamba2 — so the per-head mixture can route by +//! vol regime without competing against 39 other channels. +//! +//! ## Constraints honoured +//! +//! * `feedback_no_atomicadd` — sole-writer per output cell across both kernels. +//! * `feedback_cpu_is_read_only` — pure device kernels. +//! * `feedback_no_htod_htoh_only_mapped_pinned` — weight uploads stage +//! through `MappedF32Buffer` via the local `upload` helper (mirrors +//! `rl/ppo.rs::upload` and `rl/dueling_q.rs::upload`). +//! * `feedback_no_nvrtc` — pre-compiled cubins via `build.rs`. +//! * `pearl_scoped_init_seed_for_reproducibility` — `new()` installs +//! `scoped_init_seed` before drawing any deterministic random samples. +//! * `feedback_single_source_of_truth_no_duplicates` — `MAX_K_HEADS=8` +//! mirrors the `#define` in `multi_head_policy_forward.cu` / +//! `multi_head_policy_gate_forward.cu`; the constant lives here for +//! the runtime bounds check. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::sys::CUstream; +use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream}; +use ml_core::cuda_autograd::init::scoped_init_seed; +use ml_core::device::MlDevice; + +use crate::cfc::snap_features::REGIME_DIM; +use crate::heads::HIDDEN_DIM; +use crate::pinned_mem::MappedF32Buffer; +use crate::rl::common::N_ACTIONS; +use crate::trainer::raw_launch::{raw_launch, raw_memcpy_dtod_async, RawArgs}; + +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" +)); + +/// Compile-time upper bound on K — mirrors the `#define MAX_K_HEADS 8` in +/// both kernel sources. The runtime K (from `RL_POLICY_NUM_HEADS_INDEX`) +/// must satisfy `1 ≤ K ≤ MAX_K_HEADS`. +pub const MAX_K_HEADS: usize = 8; + +/// Construction config for [`MultiHeadPolicy`]. +#[derive(Clone, Debug)] +pub struct MultiHeadPolicyConfig { + /// Number of policy heads. Bootstrap value comes from + /// `RL_POLICY_NUM_HEADS_INDEX = 761`; this struct accepts the value + /// directly so callers don't need an ISV read at allocation time. + pub k: usize, + /// Batch size — number of independent (h_t, regime_h) rows processed + /// per forward call. Matches `n_batch` everywhere else in the trainer. + pub b_size: usize, + /// Random seed for the gating-head weight init. Per + /// `pearl_scoped_init_seed_for_reproducibility`, this guards a + /// `scoped_init_seed` block so initial draws are deterministic. + pub seed: u64, +} + +impl Default for MultiHeadPolicyConfig { + fn default() -> Self { + Self { + k: 3, + b_size: 0, + seed: 0xC0_1D_C0_DE, + } + } +} + +/// K-head policy with a parallel-channel regime gating head. +/// +/// All device buffers are allocated in `new()`. Forward runs the gate +/// kernel then the mixture kernel on the same stream — `gate_probs_d` +/// hand-off is single-stream so no extra sync is needed. +pub struct MultiHeadPolicy { + pub cfg: MultiHeadPolicyConfig, + stream: Arc, + raw_stream: CUstream, + + // ── Kernel handles ──────────────────────────────────────────────── + _forward_module: Arc, + forward_fn: CudaFunction, + _gate_forward_module: Arc, + gate_forward_fn: CudaFunction, + + // ── Policy-head weights ([K × N_ACTIONS × HIDDEN_DIM] row-major) ── + /// `W_heads[k, a, j]` — row-major flat over `(k, a, j)`. + pub w_heads_d: CudaSlice, + /// `b_heads[k, a]` — per-head per-action bias. Initial per-head action + /// priors live here (see [`MultiHeadPolicy::new`]). + pub b_heads_d: CudaSlice, + + // ── Gating-head weights ([K × REGIME_DIM=6]) ────────────────────── + pub w_gate_d: CudaSlice, + pub b_gate_d: CudaSlice, + + // ── Forward output buffers (persisted for backward) ─────────────── + /// Pre-softmax gating logits `[B × K]`. + pub gate_logits_d: CudaSlice, + /// Post-softmax gating probabilities `[B × K]`. Sums to 1 per batch row. + pub gate_probs_d: CudaSlice, + /// Per-head pre-softmax logits `[B × K × N_ACTIONS]`. + pub pi_logits_k_d: CudaSlice, + /// Per-head post-softmax probabilities `[B × K × N_ACTIONS]`. Sums to + /// 1 per (batch, head) row. + pub pi_probs_k_d: CudaSlice, +} + +impl MultiHeadPolicy { + /// Allocate weights + forward buffers, load cubins, install per-head + /// init biases for asymmetry-break specialization. + /// + /// Per-head action priors (verified against + /// [`crate::rl::common::Action`] enum values): + /// + /// | head | role | bias slot | action | + /// |------|------------------|-----------|-------------| + /// | 0 | SHORT / risk-off | 0 | ShortLarge | + /// | 1 | LONG / risk-on | 5, 6 | LongSmall, LongLarge | + /// | 2 | HOLD / stable | 2 | Hold | + /// + /// Heads 3-7 (if K > 3) carry zero bias — pure data-driven specialization. + pub fn new(dev: &MlDevice, cfg: MultiHeadPolicyConfig) -> Result { + assert!( + cfg.k >= 1 && cfg.k <= MAX_K_HEADS, + "MultiHeadPolicy: K must be in [1, {MAX_K_HEADS}]; got {}", + cfg.k + ); + assert!( + cfg.b_size > 0, + "MultiHeadPolicy: b_size must be > 0; got {}", + cfg.b_size + ); + + let stream: Arc = dev + .cuda_stream() + .context("multi_head_policy stream")? + .clone(); + let ctx = dev + .cuda_context() + .context("multi_head_policy ctx")?; + + // ── Load cubins (repo convention: const &[u8] + load_cubin) ─── + let forward_module = ctx + .load_cubin(FORWARD_CUBIN.to_vec()) + .context("load multi_head_policy_forward cubin")?; + let forward_fn = forward_module + .load_function("multi_head_policy_forward") + .context("load multi_head_policy_forward fn")?; + let gate_forward_module = ctx + .load_cubin(GATE_FORWARD_CUBIN.to_vec()) + .context("load multi_head_policy_gate_forward cubin")?; + let gate_forward_fn = gate_forward_module + .load_function("multi_head_policy_gate_forward") + .context("load multi_head_policy_gate_forward fn")?; + + // Install determinism guard before drawing any random samples + // (pearl_scoped_init_seed_for_reproducibility). + let _seed_guard = scoped_init_seed(cfg.seed); + + // ── Policy-head weights ─────────────────────────────────────── + // Zero-init the linear projections; per-head action priors live + // entirely on the bias vector so the asymmetry-break is precise + // and survives Xavier draws (which would average to zero anyway + // at random init but introduce per-step noise). + let w_heads_d = stream + .alloc_zeros::(cfg.k * N_ACTIONS * HIDDEN_DIM) + .context("alloc w_heads_d")?; + + // Per-head init biases — verified against Action enum + // (crates/ml-alpha/src/rl/common.rs:56-70): + // 0=ShortLarge 5=LongSmall 6=LongLarge 2=Hold + let mut b_heads_init = vec![0.0_f32; cfg.k * N_ACTIONS]; + if cfg.k >= 1 { + // Head 0 — SHORT-bias / risk-off + b_heads_init[0 * N_ACTIONS + 0] += 0.5; + } + if cfg.k >= 2 { + // Head 1 — LONG-bias / risk-on + b_heads_init[1 * N_ACTIONS + 5] += 0.5; + b_heads_init[1 * N_ACTIONS + 6] += 0.5; + } + if cfg.k >= 3 { + // Head 2 — HOLD / stable (Phase 0 dominant action) + b_heads_init[2 * N_ACTIONS + 2] += 0.5; + } + let b_heads_d = upload(&stream, &b_heads_init) + .context("upload b_heads init")?; + + // ── Gating-head weights ─────────────────────────────────────── + // Deterministic Xavier-style init scaled to keep initial gate + // logits near 0 (softmax ≈ uniform over K — no head wins by + // accident at step 0). The sin(i × 0.123) ramp is reproducible + // across hosts without needing an RNG; combined with + // scoped_init_seed it survives the determinism contract. + let gate_scale = (1.0_f32 / REGIME_DIM as f32).sqrt() * 0.1_f32; + let w_gate_init: Vec = (0..cfg.k * REGIME_DIM) + .map(|i| (i as f32 * 0.123_f32).sin() * gate_scale) + .collect(); + let w_gate_d = upload(&stream, &w_gate_init) + .context("upload w_gate init")?; + let b_gate_d = stream + .alloc_zeros::(cfg.k) + .context("alloc b_gate_d")?; + + // ── Forward output buffers (persisted for backward) ─────────── + let gate_logits_d = stream + .alloc_zeros::(cfg.b_size * cfg.k) + .context("alloc gate_logits_d")?; + let gate_probs_d = stream + .alloc_zeros::(cfg.b_size * cfg.k) + .context("alloc gate_probs_d")?; + let pi_logits_k_d = stream + .alloc_zeros::(cfg.b_size * cfg.k * N_ACTIONS) + .context("alloc pi_logits_k_d")?; + let pi_probs_k_d = stream + .alloc_zeros::(cfg.b_size * cfg.k * N_ACTIONS) + .context("alloc pi_probs_k_d")?; + + let raw_stream = stream.cu_stream(); + + Ok(Self { + cfg, + stream, + raw_stream, + _forward_module: forward_module, + forward_fn, + _gate_forward_module: gate_forward_module, + gate_forward_fn, + 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, + }) + } + + /// Stream owned by this head. All kernels are launched on this stream + /// so the gate→mixture hand-off needs no extra synchronization. + pub fn stream(&self) -> &Arc { + &self.stream + } + + /// Number of policy heads. + pub fn k(&self) -> usize { + self.cfg.k + } + + /// Batch size this head was allocated for. + pub fn b_size(&self) -> usize { + self.cfg.b_size + } + + /// Forward pass: gating head then mixture kernel. + /// + /// 1. `multi_head_policy_gate_forward` reads `regime_h [B × 6]` + /// and writes `gate_logits_d [B × K]` + `gate_probs_d [B × K]`. + /// 2. `multi_head_policy_forward` reads `h_t [B × HIDDEN_DIM]` and + /// `gate_probs_d`, writes `pi_logits_k_d` + `pi_probs_k_d` + /// (cached for backward) and `pi_logits_out [B × N_ACTIONS]` + /// (the mixture log-probs consumed by downstream Q-distill). + /// + /// Both kernels are deterministic by construction (single block per + /// batch, fixed summation order). No atomicAdd. + pub fn forward( + &self, + h_t: &CudaSlice, + regime_h: &CudaSlice, + pi_logits_out: &mut CudaSlice, + ) -> Result<()> { + let b = self.cfg.b_size; + let k = self.cfg.k; + debug_assert_eq!(h_t.len(), b * HIDDEN_DIM, "h_t shape mismatch"); + debug_assert_eq!( + regime_h.len(), + b * REGIME_DIM, + "regime_h shape mismatch" + ); + debug_assert_eq!( + pi_logits_out.len(), + b * N_ACTIONS, + "pi_logits_out shape mismatch" + ); + + let b_i = b as i32; + let k_i = k as i32; + + // ── Step 1: gating head ─────────────────────────────────────── + // Grid = (B), Block = (K). Reads regime_h DIRECTLY (no VSN). + { + let mut args = RawArgs::new(); + args.push_ptr(regime_h.raw_ptr()); + args.push_ptr(self.w_gate_d.raw_ptr()); + args.push_ptr(self.b_gate_d.raw_ptr()); + args.push_ptr(self.gate_logits_d.raw_ptr()); + args.push_ptr(self.gate_probs_d.raw_ptr()); + args.push_i32(b_i); + args.push_i32(k_i); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.gate_forward_fn.cu_function(), + (b as u32, 1, 1), + (k as u32, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("multi_head_policy_gate_forward: {e:?}"))?; + } + } + + // ── Step 2: per-head policy + mixture ───────────────────────── + // Grid = (B), Block = (N_ACTIONS). Reads h_t + gate_probs_d. + // Single-stream hand-off, no extra sync needed. + { + let mut args = RawArgs::new(); + args.push_ptr(h_t.raw_ptr()); + args.push_ptr(self.w_heads_d.raw_ptr()); + args.push_ptr(self.b_heads_d.raw_ptr()); + args.push_ptr(self.gate_probs_d.raw_ptr()); + args.push_ptr(self.pi_logits_k_d.raw_ptr()); + args.push_ptr(self.pi_probs_k_d.raw_ptr()); + args.push_ptr(pi_logits_out.raw_ptr()); + args.push_i32(b_i); + args.push_i32(k_i); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.forward_fn.cu_function(), + (b as u32, 1, 1), + (N_ACTIONS as u32, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("multi_head_policy_forward: {e:?}"))?; + } + } + + Ok(()) + } +} + +// ── pinned-staging upload helper ───────────────────────────────────── +// +// Mirrors `rl/ppo.rs::upload` and `rl/dueling_q.rs::upload`. Per +// `feedback_no_htod_htoh_only_mapped_pinned`: weight init data stages +// through a `MappedF32Buffer` (page-locked + mapped) then a DtoD copy +// into the persistent device slice. No raw `memcpy_htod`. +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("multi_head_policy upload staging: {e}"))?; + staging.write_from_slice(host); + let dst = stream + .alloc_zeros::(n) + .context("multi_head_policy upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let dst_ptr = dst.raw_ptr(); + raw_memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream()) + .map_err(|e| anyhow::anyhow!("multi_head_policy upload DtoD: {e:?}"))?; + } + } + Ok(dst) +} diff --git a/crates/ml-alpha/tests/multi_head_policy_invariants.rs b/crates/ml-alpha/tests/multi_head_policy_invariants.rs new file mode 100644 index 000000000..9b2d31580 --- /dev/null +++ b/crates/ml-alpha/tests/multi_head_policy_invariants.rs @@ -0,0 +1,389 @@ +//! GPU-oracle invariants for [`MultiHeadPolicy`] (Phase 2A-A). +//! +//! Forward-pass invariants verified (no trainer dependency — tests stand +//! up an `MlDevice`, allocate inputs, run forward, read back outputs): +//! +//! 1. `gate_probs_sum_to_one` — softmax invariant on the gating head. +//! 2. `pi_probs_mixture_sums_to_one` — exp(pi_logits[b]).sum() ≈ 1 for +//! every batch row (the mixture log-prob output is a valid distribution). +//! 3. `k1_reduces_to_single_head` — at K=1 the mixture math collapses to +//! the single-head softmax. With the per-head bias zeroed for K=1 +//! (`b_heads_init[0 * N_ACTIONS + 0] += 0.5` is the ONLY entry at K=1), +//! we verify the output mixture matches the reference softmax over the +//! bias-only logits — the single head with a +0.5 bias on action 0. +//! 4. `gate_responds_to_regime_change` — with hand-crafted gating weights +//! (W_gate[0] indicates dim 0 strongly), changing `regime_h[0]` from +//! large-positive to large-negative shifts the gate distribution by +//! more than a chosen threshold (TVD > 0.1). +//! 5. `forward_is_deterministic_across_contexts` — two separate +//! `MultiHeadPolicy::new` instances with the same seed produce +//! bit-equal outputs on the same input. Verifies no determinism +//! regression vs `pearl_determinism_achieved`. +//! +//! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical (softmax +//! sum, per-head reduction, deterministic seed). +//! Per `feedback_no_htod_htoh_only_mapped_pinned`: tests use +//! `write_slice_f32_d_pub` / `read_slice_d_pub` for all transfers. +//! +//! Run with: +//! `cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture` + +use anyhow::Result; +use cudarc::driver::{CudaSlice, CudaStream}; +use ml_alpha::cfc::snap_features::REGIME_DIM; +use ml_alpha::heads::HIDDEN_DIM; +use ml_alpha::rl::common::N_ACTIONS; +use ml_alpha::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig}; +use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub}; +use ml_core::device::MlDevice; +use std::sync::Arc; + +/// Relative tolerance for sum-to-1 invariants. fp32 single-block reductions +/// accumulate ~1 ULP per term; 1e-5 is loose enough to absorb that plus the +/// `+1e-12` log-of-zero guard in the forward kernel. +const TOL: f32 = 1e-5; + +fn build_device() -> Option { + match MlDevice::cuda(0) { + Ok(d) => Some(d), + Err(e) => { + eprintln!("CUDA 0 not available — skipping ({e})"); + None + } + } +} + +fn upload_f32(stream: &Arc, host: &[f32]) -> Result> { + let mut d = stream.alloc_zeros::(host.len())?; + write_slice_f32_d_pub(stream, host, &mut d)?; + Ok(d) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn gate_probs_sum_to_one() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4usize; + let k = 3usize; + let mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0xA1A1, + }, + )?; + + // Random regime input (range [-1, 1] — covers realistic feature mag). + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xCAFE); + let regime: Vec = (0..b_size * REGIME_DIM) + .map(|_| r.gen_range(-1.0..1.0)) + .collect(); + let h_t: Vec = (0..b_size * HIDDEN_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + + let regime_d = upload_f32(&stream, ®ime)?; + let h_t_d = upload_f32(&stream, &h_t)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + + let gate_probs = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?; + for b in 0..b_size { + let row = &gate_probs[b * k..(b + 1) * k]; + let sum: f32 = row.iter().sum(); + assert!( + (sum - 1.0).abs() < TOL, + "gate_probs[b={b}] sum = {sum}, expected ≈1.0 (row = {row:?})" + ); + // Also assert every entry is in [0, 1] (softmax range invariant). + for (k_idx, p) in row.iter().enumerate() { + assert!( + *p >= 0.0 && *p <= 1.0, + "gate_probs[b={b}, k={k_idx}] = {p} out of [0, 1]" + ); + } + } + eprintln!("PASS — gate_probs sum to 1.0 (±{TOL}) for all {b_size} batch rows"); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn pi_probs_mixture_sums_to_one() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4usize; + let k = 3usize; + let mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0xB2B2, + }, + )?; + + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xBEEF); + let regime: Vec = (0..b_size * REGIME_DIM) + .map(|_| r.gen_range(-1.0..1.0)) + .collect(); + let h_t: Vec = (0..b_size * HIDDEN_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + + let regime_d = upload_f32(&stream, ®ime)?; + let h_t_d = upload_f32(&stream, &h_t)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + + let pi_logits = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + // pi_logits = log(pi_probs); exp + sum should equal 1 per batch row. + for b in 0..b_size { + let row = &pi_logits[b * N_ACTIONS..(b + 1) * N_ACTIONS]; + // log-sum-exp stable accumulation (matches the kernel's + // numerical guard — `pi_logits = log(pi_probs + 1e-12)`). + let mut sum = 0.0_f32; + for &lp in row { + sum += lp.exp(); + } + assert!( + (sum - 1.0).abs() < TOL, + "exp(pi_logits[b={b}]).sum() = {sum}, expected ≈1.0 (row = {row:?})" + ); + } + eprintln!( + "PASS — exp(pi_logits) sums to 1.0 (±{TOL}) for all {b_size} batch rows (mixture is a valid distribution)" + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn k1_reduces_to_single_head() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4usize; + let k = 1usize; + let mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0xC3C3, + }, + )?; + + // h_t = 0 so the policy-head matmul produces zero logits; + // post-init, the bias is the only source of non-zero pi_logits_k. + // At K=1 the head's bias is [+0.5, 0, 0, ..., 0] (Head 0 → ShortLarge). + let h_t = vec![0.0_f32; b_size * HIDDEN_DIM]; + let regime: Vec = (0..b_size * REGIME_DIM).map(|i| (i as f32) * 0.01).collect(); + let h_t_d = upload_f32(&stream, &h_t)?; + let regime_d = upload_f32(&stream, ®ime)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + + let pi_logits = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + + // At K=1: mixture = single head's softmax. With h_t=0 the per-head + // pre-softmax logits are exactly b_heads = [+0.5, 0, ..., 0]. So: + // p[0] = exp(0.5) / (exp(0.5) + (N_ACTIONS-1)) + // p[a≠0] = 1 / (exp(0.5) + (N_ACTIONS-1)) + let denom = 0.5_f32.exp() + (N_ACTIONS as f32 - 1.0); + let p0_expected = 0.5_f32.exp() / denom; + let pother_expected = 1.0_f32 / denom; + // The kernel writes pi_logits = log(pi_probs + 1e-12); convert back. + for b in 0..b_size { + let row = &pi_logits[b * N_ACTIONS..(b + 1) * N_ACTIONS]; + let p0 = row[0].exp(); + assert!( + (p0 - p0_expected).abs() < TOL, + "k1 single-head: b={b} p[0]={p0} expected {p0_expected}" + ); + for a in 1..N_ACTIONS { + let pa = row[a].exp(); + assert!( + (pa - pother_expected).abs() < TOL, + "k1 single-head: b={b} p[{a}]={pa} expected {pother_expected}" + ); + } + } + eprintln!( + "PASS — K=1 reduces to single-head softmax over bias-only logits; p[0] = {p0_expected:.6}, p[a>0] = {pother_expected:.6}" + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn gate_responds_to_regime_change() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 2usize; + let k = 3usize; + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0xD4D4, + }, + )?; + + // Overwrite the gating weights so head 0 strongly tracks regime[0]: + // W_gate[0, :] = [+5.0, 0, 0, 0, 0, 0] + // W_gate[1, :] = [0, 0, 0, 0, 0, 0] + // W_gate[2, :] = [-5.0, 0, 0, 0, 0, 0] + // This forces a strong gate response to regime[0]: large positive + // regime[0] routes mass to head 0, large negative routes to head 2. + let mut w_gate_host = vec![0.0_f32; k * REGIME_DIM]; + w_gate_host[0 * REGIME_DIM + 0] = 5.0; + w_gate_host[2 * REGIME_DIM + 0] = -5.0; + write_slice_f32_d_pub(&stream, &w_gate_host, &mut mhp.w_gate_d)?; + + // h_t doesn't influence the gate (gate reads only regime_h). + let h_t = vec![0.0_f32; b_size * HIDDEN_DIM]; + let h_t_d = upload_f32(&stream, &h_t)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + // Run 1: regime[0] = +1 (one batch) / +1 (other batch) → mass routes to head 0. + let regime_pos = vec![1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + let regime_pos_d = upload_f32(&stream, ®ime_pos)?; + mhp.forward(&h_t_d, ®ime_pos_d, &mut pi_logits_out)?; + let gp_pos = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?; + + // Run 2: regime[0] = -1 / -1 → mass routes to head 2. + let regime_neg = vec![-1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + let regime_neg_d = upload_f32(&stream, ®ime_neg)?; + mhp.forward(&h_t_d, ®ime_neg_d, &mut pi_logits_out)?; + let gp_neg = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?; + + // Total Variation Distance between the two gate distributions per batch. + // TVD = 0.5 × Σ_k |p1[k] − p2[k]|. With ±5σ logit swing on head 0 vs + // head 2 (gap of 10 in pre-softmax), TVD should easily exceed 0.5. + for b in 0..b_size { + let row_pos = &gp_pos[b * k..(b + 1) * k]; + let row_neg = &gp_neg[b * k..(b + 1) * k]; + let tvd: f32 = row_pos + .iter() + .zip(row_neg.iter()) + .map(|(p, q)| (p - q).abs()) + .sum::() + * 0.5; + assert!( + tvd > 0.5, + "gate not responding to regime[0]: b={b} TVD={tvd}; pos={row_pos:?} neg={row_neg:?}" + ); + // With W_gate[0,0]=+5 and W_gate[2,0]=-5 and r0=+1: gate logits = + // [+5, 0, -5] → after softmax p[0] dominates. Sanity-check that + // p[0] is the modal head under positive regime. + let modal_pos = + row_pos.iter().enumerate().fold( + 0usize, + |acc, (i, p)| if *p > row_pos[acc] { i } else { acc }, + ); + assert_eq!( + modal_pos, 0, + "under regime[0]=+1, expected head 0 to be modal; got head {modal_pos} (row = {row_pos:?})" + ); + let modal_neg = + row_neg.iter().enumerate().fold( + 0usize, + |acc, (i, p)| if *p > row_neg[acc] { i } else { acc }, + ); + assert_eq!( + modal_neg, 2, + "under regime[0]=-1, expected head 2 to be modal; got head {modal_neg} (row = {row_neg:?})" + ); + } + eprintln!( + "PASS — gate routes mass between heads 0/2 in response to regime[0] sign (TVD > 0.5 per batch)" + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn forward_is_deterministic_across_contexts() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4usize; + let k = 3usize; + let seed = 0xE5E5; + + // Identical inputs for both runs. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xF00D); + let regime: Vec = (0..b_size * REGIME_DIM) + .map(|_| r.gen_range(-1.0..1.0)) + .collect(); + let h_t: Vec = (0..b_size * HIDDEN_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + let regime_d = upload_f32(&stream, ®ime)?; + let h_t_d = upload_f32(&stream, &h_t)?; + + // Run A + let mhp_a = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed, + }, + )?; + let mut out_a = stream.alloc_zeros::(b_size * N_ACTIONS)?; + mhp_a.forward(&h_t_d, ®ime_d, &mut out_a)?; + let pi_a = read_slice_d_pub(&stream, &out_a, b_size * N_ACTIONS)?; + + // Run B — fresh instance, identical seed/config + let mhp_b = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed, + }, + )?; + let mut out_b = stream.alloc_zeros::(b_size * N_ACTIONS)?; + mhp_b.forward(&h_t_d, ®ime_d, &mut out_b)?; + let pi_b = read_slice_d_pub(&stream, &out_b, b_size * N_ACTIONS)?; + + // Bit-equal — same-seed → identical weights → identical kernel inputs + // → identical outputs. Use `to_bits()` to catch any sub-ULP drift that + // would invalidate the determinism contract. + for (i, (a, b)) in pi_a.iter().zip(pi_b.iter()).enumerate() { + assert_eq!( + a.to_bits(), + b.to_bits(), + "non-deterministic output at i={i}: a={a} b={b} (bits {} vs {})", + a.to_bits(), + b.to_bits() + ); + } + eprintln!( + "PASS — forward bit-equal across two fresh contexts with seed=0x{seed:X} ({} floats compared)", + pi_a.len() + ); + Ok(()) +} diff --git a/docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md b/docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md new file mode 100644 index 000000000..d0ff52aad --- /dev/null +++ b/docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md @@ -0,0 +1,747 @@ +# 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 + +#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 + +#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, + pub b_heads_d: CudaSlice, + + // Gating head weights ([K × REGIME_DIM=6] flat) + pub w_gate_d: CudaSlice, + pub b_gate_d: CudaSlice, + + // Forward output buffers + pub gate_logits_d: CudaSlice, // [B × K] + pub gate_probs_d: CudaSlice, // [B × K] + pub pi_logits_k_d: CudaSlice, // [B × K × N_ACTIONS] + pub pi_probs_k_d: CudaSlice, // [B × K × N_ACTIONS] + + // Kernels + pub forward_kernel: CudaFunction, + pub gate_forward_kernel: CudaFunction, + + _forward_module: Arc, + _gate_forward_module: Arc, +} + +impl MultiHeadPolicy { + pub fn new( + ctx: &Arc, + stream: &Arc, + b_size: usize, + k: usize, + seed: u64, + ) -> Result { + 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::(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::(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 = (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::(k * REGIME_DIM)?; + stream.memcpy_htod(&w_gate_init, &w_gate_d) + .context("upload w_gate init")?; + let b_gate_d = stream.alloc_zeros::(k)?; + + // Forward buffers + let gate_logits_d = stream.alloc_zeros::(b_size * k)?; + let gate_probs_d = stream.alloc_zeros::(b_size * k)?; + let pi_logits_k_d = stream.alloc_zeros::(b_size * k * N_ACTIONS)?; + let pi_probs_k_d = stream.alloc_zeros::(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, + h_t: &CudaSlice, // [B × HIDDEN_DIM] + regime_h: &CudaSlice, // [B × REGIME_DIM=6] + pi_logits_out: &mut CudaSlice, // [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 +``` + +--- + +## 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. diff --git a/docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md b/docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md index 9a3bfdabd..7e93cdcd9 100644 --- a/docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md +++ b/docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md @@ -1,8 +1,8 @@ # Multi-Head Policy with R-Multiple Reward (Intervention #4-clean) -**Date**: 2026-06-02 +**Date**: 2026-06-02 (ADDENDUM 2026-06-03 — Q-distill + empirical regime-attenuation findings) **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 +**Prerequisite**: ~~Intervention #1 (R-multiple reward) has shipped and the eval verdict is in~~ **SUPERSEDED — see ADDENDUM § R below.** Empirical regime-stratification of Phase 0 baseline shows the multi-head architectural intervention is directly motivated and doesn't require R-multiple as prerequisite. **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]] @@ -199,6 +199,222 @@ Risk: PPO surrogate on a mixture has higher variance than on a single Gaussian-i Mitigation: existing PPO ratio clamp from B-7 / B-9 work handles this. Per-batch advantage normalization (Phase 4.5) helps. +--- + +# ADDENDUM 2026-06-03 — Q-distill compatibility + Empirical regime-attenuation findings + +## §R.0 — Why this addendum exists + +The original spec was written assuming PPO trains π. Per [[pearl_foxhunt_pi_trained_by_q_distillation_not_ppo]] (Phase 1B-C STOP discovery): **foxhunt's π is trained by Q→π distillation, NOT by PPO surrogate.** `crates/ml-alpha/src/trainer/integrated.rs:5577` zeros `ss_pi_grad_logits_d` BEFORE the Q-distill kernel writes its KL gradient. The PPO surrogate runs as a diagnostic-only loss. + +This addendum revises the spec to: +1. Document Q-distill gradient flow for multi-head (it still works — just via a different channel) +2. Incorporate empirical findings on VSN regime-attenuation +3. Update falsification gates with concrete numerical anchors from baseline data +4. Make multi-head a STANDALONE intervention (not contingent on R-multiple) + +## §R.1 — Empirical motivation refined (Phase 0 baseline data) + +Phase 0 baseline (`cfaa420bd`, b=128 mid-smoke, eval_pnl = −$4.46M) action distribution stratified by regime quartiles: + +| Regime marker | TVD(Q1, Q4 action dist) | Pnl Q1 vs Q4 | Interpretation | +|---|---|---|---| +| `popart_envelope_max` (agent-driven via reward signal) | **0.169** | −$1.99M vs −$32k | Policy IS partially regime-aware via reward-signal channel | +| `session_pnl_variance_ema` (market-driven vol) | **0.015** | −$1.40M vs −$2.05M | Policy is NOT regime-aware on market vol — despite vol context in encoder input at `regime[3]` (snap_feature out[23]) | + +**The asymmetry is diagnostic.** The policy can condition on popart_envelope (which modifies reward magnitude directly) but cannot condition on market vol (which exists in encoder input as `regime[3] = log1p(ema_rv_med × 1e4)`). + +### The mechanism + +`crates/ml-alpha/cuda/variable_selection.cu` implements TFT-style softmax-gated feature selection over the first 40 features of the 56-dim encoder input. **Softmax gates sum to 1.0 across 40 features — average gate ≈ 0.025.** For `regime[3]` to influence Mamba2 meaningfully, it would need an above-average gate, but it competes with 39 other features (microstructure z-scores, OFI deltas, trade activity, spread, Δt Fourier features) for that scarce attention budget. The training signal (BCE labels + V regression + Q-distill targets) doesn't reward attending to `regime[3]` specifically, so VSN learns to allocate gates to whatever helps immediate prediction targets — not regime-conditional policy. + +**Multi-head policy with regime-gated routing bypasses VSN's softmax bottleneck for regime context.** The gating head reads regime features directly (a parallel channel that doesn't compete via VSN's 40-way softmax). K heads receive a fully-dimensioned regime signal via the gating output (K-way softmax over a small K, not 40-way) and can specialize per regime. + +### Policy collapse signature (additional) + +Phase 0 baseline action_hist shows **action 2 (probably FLAT_SHORT-class) dominating 40-54% of action mass.** This is effectively a 1-action-dominated policy with noise — the encoder's representation has collapsed to "mostly emit action 2." Multi-head + per-head aux KL priors directly attacks this by REQUIRING per-head action distributions to retain mass across distinct action groups. + +## §R.2 — Q-distill compatibility (the architectural correction) + +Original §1.1 assumed PPO surrogate writes the policy gradient. Corrected flow: + +``` +Current single-head: + encoder → h_t → policy_head Linear(128→11) → π_logits + ↓ + ss_pi_grad_logits_d ← (zeroed line 5577) ← Q-distill KL gradient ← Q-Thompson sample from PER + ↓ + policy_head backward + Adam + +Proposed multi-head: + encoder → h_t (HIDDEN_DIM=128) + ↓ + K policy heads × Linear(128→11) → π_logits_k for k ∈ {0,1,2} + ↓ + Gating head Linear(GATE_INPUT_DIM → K) → gate_logits → softmax → gate[k] + ↓ + π_logits = log Σ_k gate[k] · softmax(π_logits_k) ← mixture distribution + ↓ + ss_pi_grad_logits_d ← (zeroed) ← Q-distill KL gradient writes here + ↓ + π_logits.backward → K head backward + gating head backward + ↓ + Adam updates on K head weights + gating head weights ONLY +``` + +**Critical: Q-distill kernel doesn't care about multi-head structure.** It writes a gradient on `π_logits` regardless of how `π_logits` was computed. The backward pass through the mixture computes per-head + gate gradients automatically via chain rule. + +The PPO surrogate gradient (still computed but still zeroed by line 5577) remains diagnostic-only. **No PPO behavior change required.** The aux KL prior loss (§1.4 in original spec) is added DIRECTLY to the policy backward pass, NOT zeroed by line 5577 — it lives in a separate gradient buffer that's added to `ss_pi_grad_logits_d` AFTER the Q-distill write. See §R.6 below for the gradient-merge sketch. + +## §R.3 — Gating head input (the load-bearing design choice) + +This is the single most important design choice in the addendum. Three options ordered by empirical motivation: + +### Option A — Gate from encoder output `h_t` only (original spec §1.1 default) +``` +gate_logits = Linear_HIDDEN_DIM_to_K(h_t) +``` +**Problem**: h_t is the same representation that the empirical evidence shows under-encodes regime context. VSN's softmax bottleneck has already attenuated regime[3] before h_t is produced. Gating from h_t propagates the same bottleneck into the gating head. + +### Option B — Gate from regime features DIRECTLY (recommended) +``` +// Read the 6 raw regime features per batch element (computed by loader, +// already in trainer's step_regime_d device buffer). +regime_h = step_regime_d[b][0..6] // 6-dim raw regime context +gate_logits = Linear_6_to_K(regime_h) +``` + +**Why this works**: bypasses VSN entirely. The gating head's softmax is K-way (K=3) so each regime feature gets ≈ 1/3 attention budget, not 1/40. The gating head can fully condition on `regime[3]` (vol) without VSN's gating constraint. + +**Where to source**: `step_regime_d` is allocated in `crates/ml-alpha/src/trainer/perception.rs:1795` and populated per-batch from the loader. It's already on device, already accessible per (b, k). For multi-head gating we want one regime vector per batch element (use t=K-1 of the BPTT window for "current regime") or aggregate across K (mean). + +### Option C — Gate from concat(h_t, regime_features) (compromise) +``` +gate_input = concat(h_t, regime_h) // dim 128 + 6 = 134 +gate_logits = Linear_134_to_K(gate_input) +``` +Gives gate access to both the encoder's processed state AND raw regime. More parameters, more risk of overfit but more expressiveness. + +**Decision**: Default to **Option B**. The empirical evidence (TVD=0.015 on session_pnl_var) shows h_t doesn't encode market vol; routing the gate through h_t inherits that limitation. Option B is the cleanest bypass of the diagnosed bottleneck. + +## §R.4 — Per-head aux KL priors (refined) + +Original §1.4 prior assignments: +- Head 0 (trend): biased toward HOLD actions +- Head 1 (reversion): biased toward FLAT actions +- Head 2 (flat): biased toward FLAT_FLAT + +**Refined per empirical Phase 0 data**: + +Looking at the action distribution in popart Q1 (worst regime, −$1.99M loss): +- Action 0+1 (FLAT-ish): 0.10 +- Action 2 (currently dominant): 0.40 +- Action 5+6 (HOLDs): 0.22 +- Action 7+8: 0.24 + +The Q4 distribution (best, −$32k) is heavily skewed toward action 2 (0.54) and away from HOLDs (0.11). The empirically-correct Q1 policy would likely DECREASE the HOLDs and increase decisive transitions. So the priors should target: + +- **Head 0 (defensive — for high vol / drawdown recovery regime)**: bias toward FLAT_FLAT + decisive transition actions. Prior weights ~0.40 on action 0, 0.10 each on 1-3, 0.05 each on others. +- **Head 1 (active — for low-vol stable regime)**: bias toward HOLDs and transitions. Prior weights ~0.18 each on actions 5-8, 0.05 on others. +- **Head 2 (transition — for regime change events)**: bias toward action 2 (the currently-dominant one). Prior weights ~0.35 on action 2, 0.10 each on 1+3, 0.05 elsewhere. + +These priors are SCAFFOLDING — `β=0.05` so the priors don't dominate training. Heads can diverge from priors if Q-distill signal pushes them elsewhere. + +## §R.5 — Falsification gates (concrete numerical anchors from baseline) + +Original §4.2 gates revised with measurable anchors: + +### G1 — Regime conditioning EMERGES via the multi-head architecture +- TVD(Q1 action dist, Q4 action dist) on **session_pnl_variance_ema axis** ≥ **0.10** by step 1500 +- (Baseline Phase 0: 0.015 — multi-head should lift this 6× if regime routing works) +- This is the **direct test that multi-head bypasses VSN's regime bottleneck**. + +### G2 — Per-head specialization holds +- Pairwise KL between any two heads ≥ 0.20 by step 1500 +- Gate entropy ≥ log(K) × 0.45 = 0.49 (heads actually used, not collapsed to one) +- No `gate_probs[k] < 0.10` for any k at step 5000 (mode-collapse kill) + +### G3 — Action diversity recovers (policy collapse signature) +- Argmax action mass ≤ **0.30** (down from baseline 0.40-0.54 dominated by action 2) +- This tests whether the per-head priors + Q-distill differentiation prevent collapse + +### G4 — Eval pnl improvement on POPART Q1 (lower variance than aggregate) +- Sum of Δpnl in popart_envelope Q1 quartile improves by **≥ $1M** vs Phase 0 baseline (Q1: −$1.99M) +- Per-quartile measurement has lower variance than aggregate (since quartile selection itself smooths some noise) +- **This is the load-bearing empirical gate** — if multi-head moves the worst regime quartile, the chain works + +### G5 — Aggregate eval pnl (informational, given noise floor) +- `eval_summary.total_pnl_usd` better than seeds 42-44 mean by 2σ (σ ≈ $1.9M empirical) +- Required improvement: ~$3.8M (e.g., from −$5.6M mean to better than −$1.8M) +- This is a HIGH BAR. G4 is the primary gate; G5 is corroboration. + +### G_kill — Per-step kill criteria +- Action entropy < log(11) × 0.30 = 0.72 at step 5000 (policy collapse) +- Eval realized_pnl_cum_usd < −$10M at step 10000 (catastrophic abandon) +- Cross-source consistency drift > $50 (Phase 0 contract violation) +- Determinism check fails at any seed + +### Stratified verdict methodology +Per the noise floor finding (σ ≈ $1.9M at b=128 single-seed), use multi-seed AND per-quartile measurement: +- Run smokes at seeds 42, 43, 44 (3-seed minimum) +- For each, compute per-quartile Δpnl on popart_envelope +- Compare mean per-quartile improvement vs the 3-seed Phase 0 baseline distribution + +## §R.6 — Gradient flow sketch (implementation guidance) + +``` +// Forward (multi-head): +for k in 0..K: + pi_logits_k[b, a] = w_k[a, :] @ h_t[b, :] + b_k[a] // K linear heads +gate_logits[b, k] = w_gate[k, :] @ regime_h[b, :] + b_gate[k] // gating head reads regime_h DIRECTLY +gate_probs[b, k] = softmax(gate_logits[b]) +pi_probs_k[b, a] = softmax(pi_logits_k[b]) +pi_probs[b, a] = Σ_k gate_probs[b, k] * pi_probs_k[b, a] // mixture in probability space +pi_logits[b, a] = log(pi_probs[b, a] + 1e-12) // back to logit space for downstream + +// Backward (Q-distill + aux): +ss_pi_grad_logits_d[b, a] := 0 // existing zero at line 5577 +rl_q_pi_distill_grad writes into ss_pi_grad_logits_d // existing path, unchanged +ss_pi_aux_grad_logits_d[b, a] := Σ_k β · gate_probs[b, k] · ∂KL(pi_probs_k || pi_prior_k) / ∂pi_logits_k[b, a] // NEW +ss_pi_grad_logits_d[b, a] += ss_pi_aux_grad_logits_d[b, a] // merge + +// Then policy_head.backward (existing path) propagates ss_pi_grad_logits_d +// back through the mixture (∂pi_logits/∂pi_logits_k via mixture jacobian) +// and through the gating head (∂pi_logits/∂gate_logits via mixture jacobian) +// Adam updates: w_k, b_k for each head + w_gate, b_gate +``` + +The mixture's backward jacobian is standard (Shazeer 2017 MoE; standard implementation in production frameworks). ~100 LOC for two backward kernels (mixture_backward + aux_prior_backward). + +## §R.7 — Scope estimate (revised) + +| Component | Original spec | Revised | +|---|---|---| +| Multi-head policy struct + Adam states | ~80 LOC | ~80 LOC | +| Mixture forward kernel | ~80 LOC | ~80 LOC | +| Mixture backward kernel | ~120 LOC | ~140 LOC (handles K=3 with gate jacobian) | +| Aux prior loss + backward | ~60 LOC | ~80 LOC | +| Gating head: regime-direct input wiring | not in original | ~60 LOC (read step_regime_d, project to K) | +| Bootstrap entries for new ISV slots | ~20 LOC | ~20 LOC | +| Diag emission (per-head + gate + TVD diagnostics) | ~40 LOC | ~50 LOC (more diagnostics for regime conditioning) | +| Tests (invariants for mixture, prior pull, gate entropy) | ~150 LOC | ~200 LOC | +| **Total** | **~550 LOC** | **~710 LOC** | + +Multi-week intervention. Single coder commit (per-component tasks would split it). + +## §R.8 — Why this is now the right intervention (over R-multiple or further reward work) + +Per ADDENDUM 2026-06-02e of `pearl_reward_signal_anti_aligned_with_pnl`: reward dimension has been tested twice (Phase 1A pure-pnl, Path A GAE-V). Both produced null/noise-level changes. Continuing in the reward dimension would be "going under the streetlight" (per `feedback_investigation_first_falsification_methodology`). + +Multi-head policy targets a different mechanism: +- **Architecturally addresses the diagnosed VSN bottleneck** (regime[3] attenuation) +- **Empirically motivated**: TVD(0.015) on vol axis vs TVD(0.17) on popart axis proves the policy CAN condition on regime when the signal reaches it through the right channel +- **Q-distill compatible**: gradient flows through mixture's standard backward jacobian; the existing Q-distill kernel is unchanged +- **Targets the popart Q1 regime (worst $1.99M of $4.46M loss)** with measurable per-quartile gate +- **Single ~700 LOC intervention** (multi-week but bounded) + +If multi-head fails G1 (TVD on vol axis doesn't lift): the eval-collapse is genuinely not regime-mediated, and the architecture/data are the bottleneck. We'd then have empirical justification for the deeper interventions (regime-stratified training, encoder architecture changes, etc.). + ### §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.