From 3286dc7dee0b9502d46e4c6370350434a9de5ab6 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 14 May 2026 02:26:39 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22-vnext):=20Phase=20C-1=20=E2=80=94=20K?= =?UTF-8?q?=3D3=20softmax=20=E2=86=92=20per-env=203-slot=20cache=20(produc?= =?UTF-8?q?er)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First half of Phase C. Lands the producer side of the K=3 trade- outcome aux head's state bridge: new kernel populates a per-env 3-slot cache from the K=3 softmax tile every rollout step. The consumer side (state gather reading from this cache → state slots [121..124)) lands in Phase C-2. Mirrors the K=2 head's existing aux_softmax_to_per_env_kernel exactly at K=3: - K=2: prev_aux_dir_prob[env] = 2*softmax[env, 1] - 1 (recentered) - K=3: prev_aux_outcome_probs[env, k] = softmax[env, k] for k in [0, 3) Changes: - state_layout.rs: 3 new constants AUX_OUTCOME_PROFIT_INDEX = 121, AUX_OUTCOME_STOP_INDEX = 122, AUX_OUTCOME_TIMEOUT_INDEX = 123. PROFIT_INDEX aliases AUX_DIR_PROB_INDEX (same value, different semantic). Phase C-2 flips slot 121's meaning from K=2's recentered p_up to K=3's p_Profit. - aux_outcome_softmax_to_per_env_kernel.cu: new kernel + cubin. - gpu_dqn_trainer.rs: new SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN embed. - gpu_experience_collector.rs: 2 new struct fields (cache buffer + kernel handle); cubin load + alloc in constructor; struct-init; per-step launch in rollout loop after K=3 forward. - build.rs: kernel registered. Encoding shift K=2 → K=3: K=2 used recentered [-1, +1] to match "no signal = 0" baseline of every other slot. K=3 keeps raw softmax probabilities [0, 1]. Cold-start sentinel 0.0 for all 3 slots = "no prediction yet" (mask). The 3-slot natural distribution is more informative than a scalar. Dead-code status: producer populates cache every step but experience_state_gather doesn't read from it yet — state slot 121 still receives K=2's prev_aux_dir_prob write. Phase C-2 swaps the state gather's source from K=2 cache to K=3 cache (3-slot write). Why split C into C-1 + C-2: experience_state_gather is a hot-path kernel with many consumers. Updating it touches training collector, eval-side backtest evaluator, Rust launcher arg list. C-2 lands that as an atomic state-semantic flip; C-1 lands the GPU-side scaffolding independently so the producer chain can be validated first. Verification: - cargo check -p ml clean. - cargo test -p ml --lib → 1016/0 green. Audit: docs/dqn-wire-up-audit.md Phase C-1 section. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-core/src/state_layout.rs | 22 +++++ crates/ml/build.rs | 9 ++ .../aux_outcome_softmax_to_per_env_kernel.cu | 98 +++++++++++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 8 ++ .../cuda_pipeline/gpu_experience_collector.rs | 91 +++++++++++++++++ docs/dqn-wire-up-audit.md | 27 +++++ 6 files changed, 255 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu diff --git a/crates/ml-core/src/state_layout.rs b/crates/ml-core/src/state_layout.rs index 77748d5f1..2970d0e4d 100644 --- a/crates/ml-core/src/state_layout.rs +++ b/crates/ml-core/src/state_layout.rs @@ -67,6 +67,28 @@ pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 121 // slot's "no signal = 0" baseline per `pearl_first_observation_bootstrap`). pub const AUX_DIR_PROB_INDEX: usize = PADDING_START; // 121 +// SP22 H6 vNext Phase C (2026-05-14) — trade-outcome aux head's K=3 +// softmax probabilities, raw [0, 1] range. Replaces the K=2 head's +// single-slot recentered `2*p_up - 1` at slot 121 (aliased — same index +// value, different semantic). The K=3 head's three-slot representation +// is more informative than a single directional scalar: +// - p_Profit: probability the trade closes via profit_target hit +// - p_Stop: probability the trade closes via stop_loss hit +// - p_Timeout:probability the trade closes via target_bars elapsed +// Softmax outputs are inherently in [0, 1] and sum to 1 per row, so no +// recentering needed for "no signal = 0" baseline (cold-start all 3 +// slots = 0.0 means "all-mask, no trade-close prediction yet" — also +// the FoldReset sentinel). +// +// Phase C (this commit): constants defined + producer kernel +// (`aux_outcome_softmax_to_per_env_kernel`) populates the per-env +// cache. Phase C-2 (follow-up): `experience_state_gather` reads from +// the K=3 cache and writes these 3 slots, deprecating the K=2 head's +// single-slot input to the policy. +pub const AUX_OUTCOME_PROFIT_INDEX: usize = PADDING_START; // 121 +pub const AUX_OUTCOME_STOP_INDEX: usize = PADDING_START + 1; // 122 +pub const AUX_OUTCOME_TIMEOUT_INDEX: usize = PADDING_START + 2; // 123 + const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM"); const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores"); const _: () = assert!(STATE_DIM_PADDED % 128 == 0, "STATE_DIM_PADDED must be 128-aligned for cuBLAS"); diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 5b1d6ba50..d915b48a4 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -841,6 +841,15 @@ fn main() { // this commit — Phase B5b allocates buffers + wires launches + // flips SH2_TOTAL=262 in forward/backward. "aux_to_input_concat_kernel.cu", + // SP22 H6 vNext Phase C (2026-05-14): per-env 3-slot softmax + // extractor for the K=3 trade-outcome aux head. Copies + // aux_outcome_softmax [n_envs, K=3] → prev_aux_outcome_probs + // [n_envs, 3]. Consumed at the NEXT step's state assembly via + // experience_state_gather (Phase C-2 follow-up). For Phase C + // (this commit) the producer fires every step but no consumer + // reads — dead-code scaffolding parallel to the K=2 head's + // aux_softmax_to_per_env_kernel. + "aux_outcome_softmax_to_per_env_kernel.cu", // SP22 H6 (2026-05-12): per-env p_up extractor that copies // `aux_softmax[env, 1]` → `prev_aux_dir_prob[env]` after each aux // forward. The cache buffer is read by `experience_state_gather` diff --git a/crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu b/crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu new file mode 100644 index 000000000..2b1fcbbda --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu @@ -0,0 +1,98 @@ +// crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu +// +// SP22 H6 vNext Phase C (2026-05-14): aux→policy state bridge — per-env +// trade-outcome probability extractor. +// +// Role +// ──── +// At rollout step t, after the K=3 trade-outcome aux head's softmax +// forward produces the per-env tile `aux_outcome_softmax [n_envs, K=3]`, +// this kernel copies all three probability columns into the per-env +// cache buffer `prev_aux_outcome_probs [n_envs, 3]`. The cache is read +// by `experience_state_gather` at step t+1 to populate state slots: +// +// state[AUX_OUTCOME_PROFIT_INDEX = 121] ← prev_aux_outcome_probs[env, 0] +// state[AUX_OUTCOME_STOP_INDEX = 122] ← prev_aux_outcome_probs[env, 1] +// state[AUX_OUTCOME_TIMEOUT_INDEX = 123] ← prev_aux_outcome_probs[env, 2] +// +// Replaces the K=2 head's single-slot bridge +// (`aux_softmax_to_per_env_kernel` → `prev_aux_dir_prob` → state[121] = +// `2*p_up - 1`). The K=3 head's 3-slot bridge supersedes that semantic +// at step t+1 once the state gather consumer is updated (Phase C-2, +// follow-up commit). For Phase C (this commit) the buffer is populated +// but no consumer reads it yet — dead-code scaffolding. +// +// Encoding +// ──────── +// Slots [121..124) hold raw softmax probabilities in [0, 1] (sum to 1 +// per row). Unlike the K=2 head's slot 121 which used a recentered +// `2*p - 1` encoding to land in [-1, +1], the K=3 head's three slots +// keep the natural [0, 1] range: +// - No structural bound violation (softmax outputs are inherently +// [0, 1] and sum to 1) +// - Three-slot probabilistic representation is more informative than +// a single "p_up − p_down" scalar +// - Cold-start sentinel: all three slots = 0.0. The policy reads +// "no signal" for all three classes, which is the conservative +// baseline (vs. the K=2 head's 0.0 = neutral encoding). After +// first rollout step the producer overwrites with real K=3 probs. +// +// Discipline +// ────────── +// - No atomicAdd (no reduction, just a strided gather) +// - GPU-only per `feedback_cpu_is_read_only.md` +// - Stream-ordered with the K=3 forward producer (same-stream barrier +// suffices — no explicit sync needed) +// +// Launch +// ────── +// grid = ((n_envs + 255) / 256), block = (256). One thread per env. +// Each thread writes 3 consecutive output elements (K=3 unrolled). +// +// Cost +// ──── +// Three coalesced 4-byte loads + three 4-byte stores per env. n_envs is +// O(1k); ~1-2µs total. Negligible vs the existing K=2 sibling. + +#include + +extern "C" __global__ void aux_outcome_softmax_to_per_env_kernel( + /* Per-env K=3 softmax tile [n_envs, K=3]. Written by + * `aux_trade_outcome_forward` during the captured rollout-step + * graph; valid after same-stream sync. Layout row-major; this + * kernel reads `aux_outcome_softmax[env * K + 0/1/2]` per env. */ + const float* __restrict__ aux_outcome_softmax, + /* Per-env cache buffer [n_envs, K=3] row-major. Written here, read + * by `experience_state_gather` (Phase C-2) at the next step's + * state assembly. Slot 0 = p_Profit, slot 1 = p_Stop, slot 2 = + * p_Timeout per the K=3 outcome label convention from + * `trade_outcome_label_kernel`. */ + float* __restrict__ prev_aux_outcome_probs, + /* Number of envs (rollout-time batch width = `alloc_episodes`). */ + int n_envs, + /* Softmax K (= AUX_OUTCOME_K = 3 in Phase A3+). Kernel reads slots + * [0, K). Passed runtime so the kernel signature stays stable if + * the head's K flips in a future redesign. */ + int K +) { + int env = blockIdx.x * blockDim.x + threadIdx.x; + if (env >= n_envs) return; + + /* Phase C (2026-05-14): write raw softmax probabilities [0, 1] + * into 3 per-env slots. The state gather consumer (Phase C-2) + * will copy these into state[121..124). Unlike the K=2 head's + * recentered encoding, the K=3 head's natural softmax range is + * already in [0, 1] and the three-slot representation preserves + * the full categorical distribution. + * + * The kernel writes K slots (passed runtime) so the loop tolerates + * future K flips; in production K=3 always so the loop trips + * exactly 3 times. Bounds: K can be any positive value, but the + * cache buffer must be sized [n_envs, K] by the caller. */ + const size_t src_base = (size_t)env * (size_t)K; + const size_t dst_base = (size_t)env * (size_t)K; + for (int k = 0; k < K; ++k) { + prev_aux_outcome_probs[dst_base + (size_t)k] = + aux_outcome_softmax[src_base + (size_t)k]; + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 9591b9a07..31882d1f1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -685,6 +685,14 @@ pub(crate) static SP14_Q_DISAGREEMENT_CUBIN: &[u8] = pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_softmax_to_per_env_kernel.cubin")); +/// SP22 H6 vNext Phase C (2026-05-14): per-env K=3 trade-outcome softmax +/// extractor cubin. Loaded by `GpuExperienceCollector` (and eventually +/// `GpuBacktestEvaluator` in the eval-side mirror). Producer of the 3- +/// slot per-env cache that Phase C-2 wires into state[121..124). +#[allow(dead_code)] +pub(crate) static SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_outcome_softmax_to_per_env_kernel.cubin")); + /// SP22 H6 Phase 3 α structural-prior init (atom-shift design 2026-05-13). /// One-time GPU init kernel that writes the per-action prior /// `[-0.5, 0.0, +0.5, 0.0]` for `[Short, Hold, Long, Flat]` into the diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 2bd7f2b40..2ea92ad64 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1297,6 +1297,32 @@ pub struct GpuExperienceCollector { /// consumer. exp_aux_to_label_per_sample: cudarc::driver::CudaSlice, + /// SP22 H6 vNext Phase C (2026-05-14) — per-env K=3 trade-outcome + /// softmax cache. Sized `[alloc_episodes × AUX_OUTCOME_K=3]` f32, + /// row-major: env-major then class. Written end-of-step by + /// `aux_outcome_softmax_to_per_env_kernel` (Phase C kernel) from + /// `exp_aux_to_softmax_buf`. Read start-of-next-step by + /// `experience_state_gather` (Phase C-2 wireup) to populate state + /// slots `[AUX_OUTCOME_PROFIT_INDEX..AUX_OUTCOME_TIMEOUT_INDEX+1)` + /// = `[121, 122, 123]`. + /// + /// Replaces the K=2 head's single-slot `prev_aux_dir_prob` semantic + /// at state slot 121. Both buffers coexist in this commit (K=2 head + /// remains active; its prev_aux_dir_prob still feeds state[121] via + /// the current state gather). Phase C-2 flips the state gather to + /// read from this 3-slot K=3 buffer instead. + /// + /// Cold-start sentinel = 0.0 across all 3 slots. The state gather + /// reads "no Profit/Stop/Timeout signal" on step 0, then this + /// kernel writes real K=3 softmax probs for step 1+ state assembly. + prev_aux_outcome_probs: cudarc::driver::CudaSlice, + /// SP22 H6 vNext Phase C (2026-05-14) — kernel handle for + /// `aux_outcome_softmax_to_per_env_kernel`. Loaded from the Phase + /// C cubin. Launched on the collector's stream right after the + /// K=3 `aux_trade_outcome_forward` populates + /// `exp_aux_to_softmax_buf`. + exp_aux_outcome_softmax_to_per_env_kernel: CudaFunction, + /// SP22 H6 Phase 2 (2026-05-12) — per-env aux directional probability /// cache, RECENTERED encoding. Sized `[alloc_episodes]` f32 (one slot /// per env). Written end-of-step by `exp_aux_softmax_to_per_env_kernel` @@ -2764,6 +2790,31 @@ impl GpuExperienceCollector { "sp22-vnext B4b: alloc exp_aux_to_label_per_sample ({} i32): {e}", alloc_episodes * alloc_timesteps )))?; + + // SP22 H6 vNext Phase C (2026-05-14): K=3 per-env softmax cache + // + producer kernel handle. Sized [alloc_episodes × AUX_OUTCOME_K=3] + // f32. Cold-start sentinel 0.0 across all 3 slots = "no + // Profit/Stop/Timeout signal" (the consumer state gather will + // read this baseline on step 0; subsequent steps populated + // from the K=3 forward's softmax). + let prev_aux_outcome_probs = stream + .alloc_zeros::(alloc_episodes * AUX_OUTCOME_K) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext C: alloc prev_aux_outcome_probs ({} f32): {e}", + alloc_episodes * AUX_OUTCOME_K + )))?; + let exp_aux_outcome_softmax_to_per_env_kernel = { + use super::gpu_dqn_trainer::SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN; + let m = stream.context() + .load_cubin(SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext C: aux_outcome_softmax_to_per_env cubin (collector): {e}" + )))?; + m.load_function("aux_outcome_softmax_to_per_env_kernel") + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext C: aux_outcome_softmax_to_per_env_kernel load (collector): {e}" + )))? + }; // SP22 H6 (2026-05-12): aux→policy state bridge. Load the per-env // softmax extractor (`aux_softmax_to_per_env_kernel`) + the shared // `fill_f32` from the epsilon_greedy cubin (for 0.5 sentinel init at @@ -3354,6 +3405,12 @@ impl GpuExperienceCollector { exp_aux_to_softmax_buf, exp_aux_to_label_buf, exp_aux_to_label_per_sample, + // SP22 H6 vNext Phase C (2026-05-14): K=3 trade-outcome + // per-env softmax cache + producer kernel handle (Phase C + // scaffolding — populates `prev_aux_outcome_probs` every + // step; state gather consumer wired in Phase C-2). + prev_aux_outcome_probs, + exp_aux_outcome_softmax_to_per_env_kernel, // SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up // cache + extractor kernel + fill_f32 handle for 0.5 sentinel. prev_aux_dir_prob, @@ -5887,6 +5944,40 @@ impl GpuExperienceCollector { self.exp_aux_to_logits_buf.raw_ptr(), self.exp_aux_to_softmax_buf.raw_ptr(), )?; + + // ── SP22 H6 vNext Phase C (2026-05-14): K=3 softmax → + // per-env 3-slot cache. Mirrors the K=2 sibling's + // `aux_softmax_to_per_env_kernel` launch (Phase 2) but + // writes 3 slots instead of 1. Reads + // `exp_aux_to_softmax_buf [n_envs, 3]` (just produced + // by the K=3 forward above on the same stream — stream- + // implicit ordering); writes `prev_aux_outcome_probs + // [n_envs, 3]`. Phase C-2 wires `experience_state + // _gather` to read this cache and populate state slots + // [121..124). + { + let n_envs_i32 = n as i32; + let k_i32 = AUX_OUTCOME_K as i32; + let blocks = ((n as u32).div_ceil(256)).max(1); + let softmax_ptr = self.exp_aux_to_softmax_buf.raw_ptr(); + let cache_ptr = self.prev_aux_outcome_probs.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.exp_aux_outcome_softmax_to_per_env_kernel) + .arg(&softmax_ptr) + .arg(&cache_ptr) + .arg(&n_envs_i32) + .arg(&k_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "aux_outcome_softmax_to_per_env_kernel t={t}: {e}" + )))?; + } + } } // ── 3. Convert logits → expected Q-values ─────────────────── diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 45a5fa98a..f2d089c1c 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -18010,3 +18010,30 @@ Without Phase C, validation runs would show "K=3 head trains and converges" but **Recommendation**: skip the full B5 for now, do Phase C next, then Phase D (atom-shift). Phase B5b (plan-conditioning) is a refinement we add IF Phase C/D's no-plan-params version shows promise but plateaus below the WR=0.55 target. Verification: `cargo check -p ml` clean. Lib test suite unchanged from B4b-2. + +#### Phase C-1 — K=3 softmax → per-env 3-slot cache (producer scaffold) (2026-05-14) + +First half of Phase C. Lands the producer side of the K=3 trade-outcome aux head's state bridge: new kernel `aux_outcome_softmax_to_per_env_kernel` populates a per-env 3-slot cache from the K=3 softmax tile every rollout step. The consumer side (state gather reading from this cache → state slots [121..124)) lands in Phase C-2. + +**Mirrors the K=2 head's existing `aux_softmax_to_per_env_kernel` exactly** at K=3 instead of K=2: +- K=2: writes `prev_aux_dir_prob[env] = 2*softmax[env, 1] - 1` (recentered single scalar) +- K=3: writes `prev_aux_outcome_probs[env, k] = softmax[env, k]` for k ∈ [0, 3) (raw probabilities) + +**Changes**: +- `state_layout.rs`: 3 new constants `AUX_OUTCOME_PROFIT_INDEX = 121`, `AUX_OUTCOME_STOP_INDEX = 122`, `AUX_OUTCOME_TIMEOUT_INDEX = 123`. Note: `AUX_OUTCOME_PROFIT_INDEX` aliases `AUX_DIR_PROB_INDEX` (same value 121, different semantic). The state gather consumer update in Phase C-2 will flip slot 121's meaning from K=2's recentered p_up to K=3's p_Profit. +- `aux_outcome_softmax_to_per_env_kernel.cu`: new kernel + cubin (~5 KB). Pure GPU strided gather, no atomicAdd. +- `gpu_dqn_trainer.rs`: new `SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN` embed. +- `gpu_experience_collector.rs`: 2 new struct fields (cache buffer `prev_aux_outcome_probs [alloc_episodes × 3]` + kernel handle); cubin load + alloc in constructor; struct-init entry; per-step launch in rollout loop immediately after the K=3 `aux_trade_outcome_forward` (stream-implicit producer→consumer ordering). +- `build.rs`: kernel registered. + +**Encoding shift K=2 → K=3**: +- K=2 used recentered `2*p_up - 1` ∈ [-1, +1] to match the "no signal = 0" baseline of every other state slot. +- K=3 keeps raw softmax probabilities in [0, 1]. Cold-start sentinel = 0.0 for all three slots = "no prediction yet" (mask). After first step the producer writes real K=3 probs. The 3-slot natural distribution is more informative than the K=2 scalar. + +**Dead-code status**: For Phase C-1 the producer kernel populates the cache every step but `experience_state_gather` doesn't read from it yet — state slot 121 still receives the K=2 head's `prev_aux_dir_prob` write. Phase C-2 swaps the state gather's source from the K=2 cache to the K=3 cache (3-slot write). + +**Why split C into C-1 + C-2**: the state gather (`experience_state_gather` kernel + Rust launchers) is a hot-path kernel with many consumers. Updating it touches the training collector, the eval-side backtest evaluator, and the Rust launcher's argument list. C-2 lands that as an atomic state-semantic flip; C-1 lands the GPU-side scaffolding independently so the producer chain can be validated first. + +**Verification**: `cargo check -p ml` clean. `cargo test -p ml --lib` → **1016/0 green**. The new kernel launch runs every rollout step but doesn't affect training behavior (no consumer reads the cache yet). + +**Next**: Phase C-2 (`experience_state_gather` update — read `prev_aux_outcome_probs[env, 0..3]` → state[121..124], replacing the K=2 single-slot write). This is the moment the K=3 head's predictions actually reach the policy.