diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 6e6e0c49d..9604f6f8d 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -37,9 +37,11 @@ const KERNELS: &[&str] = &[ "ppo_clipped_surrogate", // RL Phase D: PPO clipped-surrogate + entropy bonus + value MSE fwd/bwd "rl_ppo_clip_controller", // RL Phase D: ISV controller emitting ε to ISV[RL_PPO_CLIP_INDEX=402] "rl_entropy_coef_controller", // RL Phase D: ISV controller emitting entropy bonus weight to ISV[RL_ENTROPY_COEF_INDEX=403] + "v_head_fwd_bwd", // RL Phase E.2: scalar V(s) head fwd + MSE bwd (linear layer; per-batch scratch + reduce_axis0) + "grad_h_accumulate", // RL Phase E.2: element-wise grad_h_encoder += λ × grad_h_head accumulator (one head at a time, serialised by stream) ]; -// Cache bust v21 (2026-05-22): RL Phase D — ppo_clipped_surrogate.cu (PPO L_π = -min(r·A, clip(r,1±ε)·A) + L_v = (V-R)² + L_ent = -coef·H(π) fwd/bwd over 9 actions; ε read from ISV[402], coef from ISV[403]), rl_ppo_clip_controller.cu (ε→ISV[402], anchored on KL(π_new‖π_old) EMA, target KL=0.01), rl_entropy_coef_controller.cu (coef→ISV[403], anchored on entropy deficit vs 0.7·ln(9)). Per pearl_controller_anchors_isv_driven, ε + coef are not constants — Wiener-α adaptive with first-observation bootstrap (ε=0.2, coef=0.01). +// Cache bust v22 (2026-05-22): RL Phase E.2 — v_head_fwd_bwd.cu (scalar V head fwd + MSE bwd, per-batch scratch for grad_w/grad_b — reduced by caller via reduce_axis0; no atomicAdd), grad_h_accumulate.cu (element-wise grad_h_encoder += λ × grad_h_head; trainer folds Q/π/V per-head grad_h_t contributions into encoder grad slot via scaled +=; deferred to Phase E.3 — encoder backward integration). ALSO extends dqn_distributional_q.cu with dqn_grad_w_b_h_t (mirrors aux_heads_bwd pattern: grad_logits → grad_w/grad_b/grad_h_t for the C51 head; per-batch scratch) and ppo_clipped_surrogate.cu with ppo_grad_w_b_h_t (same pattern for PPO π head). All new kernels honour feedback_no_atomicadd — per-batch grad scratch reduced by the existing reduce_axis0 path, never atomics. fn main() { println!("cargo:rerun-if-changed=build.rs"); diff --git a/crates/ml-alpha/cuda/dqn_distributional_q.cu b/crates/ml-alpha/cuda/dqn_distributional_q.cu index 977ffa2bf..c43b6f40d 100644 --- a/crates/ml-alpha/cuda/dqn_distributional_q.cu +++ b/crates/ml-alpha/cuda/dqn_distributional_q.cu @@ -218,3 +218,82 @@ extern "C" __global__ void dqn_distributional_q_bwd( atomicAdd(loss_out, ce); } } + + +// ───────────────────────────────────────────────────────────────────── +// dqn_grad_w_b_h_t (Phase E.2): +// Given per-batch grad_logits from `dqn_distributional_q_bwd`, propagate +// to weight / bias gradients (per-batch scratch — caller reduces via +// reduce_axis0) and to grad_h_t (per-(batch, c) sole writer; OVERWRITE). +// +// Mathematics (per batch b, hidden unit c): +// grad_w_per_batch[b, a*Q_N_ATOMS + z, c] +// = grad_logits[b, a, z] × h_t[b, c] +// grad_b_per_batch[b, a*Q_N_ATOMS + z] +// = grad_logits[b, a, z] +// grad_h_t[b, c] +// = Σ_{a, z} grad_logits[b, a, z] × w[a*Q_N_ATOMS + z, c] +// +// Block layout (mirrors aux_heads_bwd): +// grid = (B, 1, 1) +// block = (HIDDEN_DIM = 128, 1, 1) +// shared_mem_bytes = (N_ACTIONS * Q_N_ATOMS) * sizeof(float) +// — stages grad_logits row [a*Q_N_ATOMS + z] once +// +// Thread c is the sole writer of grad_w_per_batch[batch, k, c] (for all k +// in [0, N_ACTIONS * Q_N_ATOMS)) and grad_h_t[batch, c]. The first +// N_ACTIONS * Q_N_ATOMS threads (= 189 ≤ 128 → spills past block size) +// would need the bias write — instead, atom-style threads aren't +// available; we use a tile-wise loop where the first 128 threads +// cover the first 128 bias slots and a stride-loop covers the +// remaining 189 - 128 = 61. Per `feedback_no_atomicadd.md`: each k +// has a single thread writing it (one stride loop, single writer per k). +// +// No atomicAdd. No nvrtc. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void dqn_grad_w_b_h_t( + const float* __restrict__ w, // [N_ACTIONS * Q_N_ATOMS, HIDDEN_DIM] + const float* __restrict__ h_t, // [B * HIDDEN_DIM] + const float* __restrict__ grad_logits, // [B * N_ACTIONS * Q_N_ATOMS] + int B, + float* __restrict__ grad_w_per_batch, // [B * N_ACTIONS * Q_N_ATOMS * HIDDEN_DIM] + float* __restrict__ grad_b_per_batch, // [B * N_ACTIONS * Q_N_ATOMS] + float* __restrict__ grad_h_t // [B * HIDDEN_DIM] (OVERWRITE) +) { + // grad_logits row for this batch: [N_ACTIONS * Q_N_ATOMS] = 189 floats. + extern __shared__ float s_gl[]; // [N_ACTIONS * Q_N_ATOMS] + + const int batch = blockIdx.x; + const int c = threadIdx.x; + if (batch >= B) return; + + const int K_OUT = N_ACTIONS * Q_N_ATOMS; // 189 + + // Cooperative stage of grad_logits[batch, :] into shared. Stride-loop + // over K_OUT in steps of HIDDEN_DIM (block size). + for (int k = c; k < K_OUT; k += HIDDEN_DIM) { + s_gl[k] = grad_logits[batch * K_OUT + k]; + } + __syncthreads(); + + // ── grad_b_per_batch ─ stride-loop, sole writer per k ──────────── + for (int k = c; k < K_OUT; k += HIDDEN_DIM) { + grad_b_per_batch[batch * K_OUT + k] = s_gl[k]; + } + + if (c >= HIDDEN_DIM) return; + + // ── grad_w + grad_h_t (thread role c = hidden unit) ────────────── + const float h_bc = h_t[batch * HIDDEN_DIM + c]; + float gh_acc = 0.0f; + for (int k = 0; k < K_OUT; ++k) { + const float g_k = s_gl[k]; + // grad_w[batch, k, c] = g_k × h_bc — sole writer per (batch, k, c). + grad_w_per_batch[(long long) batch * K_OUT * HIDDEN_DIM + + (long long) k * HIDDEN_DIM + c] + = g_k * h_bc; + gh_acc += g_k * w[k * HIDDEN_DIM + c]; + } + // OVERWRITE — caller folds via grad_h_accumulate_scaled. + grad_h_t[batch * HIDDEN_DIM + c] = gh_acc; +} diff --git a/crates/ml-alpha/cuda/grad_h_accumulate.cu b/crates/ml-alpha/cuda/grad_h_accumulate.cu new file mode 100644 index 000000000..871d14f63 --- /dev/null +++ b/crates/ml-alpha/cuda/grad_h_accumulate.cu @@ -0,0 +1,38 @@ +// grad_h_accumulate.cu — element-wise scaled gradient accumulator +// (Phase E.2 of the integrated RL trainer). +// +// Combines per-head encoder-input gradients into the shared encoder +// grad slot via scaled +=: +// +// grad_h_encoder[i] += lambda × grad_h_head[i] +// +// Used by IntegratedTrainer to fold the Q / π / V heads' +// `grad_h_t [B × HIDDEN_DIM]` outputs into the encoder backward +// input, on top of the BCE + aux contribution already deposited by +// the PerceptionTrainer's own grad chain. +// +// One thread per element. Each thread is the sole writer of its +// destination slot (within a single launch) — no atomicAdd per +// `feedback_no_atomicadd.md`. Multiple launches in sequence on the +// same destination are safe because each launch reads-modifies-writes +// the same slot serially via stream ordering. +// +// Constraints honoured: +// * `feedback_no_atomicadd.md` — sole-writer per (i) +// * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs +// * `pearl_no_host_branches_in_captured_graph` — `n` is a captured +// scalar; `lambda` is a captured scalar (the trainer reads it from +// the host ISV mirror BEFORE the capture region). +// +// Launch: grid=(ceil(n/256), 1, 1); block=(256, 1, 1); shared=0. + +extern "C" __global__ void grad_h_accumulate_scaled( + const float* __restrict__ grad_h_head, // [n] — input head gradient + float lambda, // scalar weight (host-passed) + int n, // total elements + float* __restrict__ grad_h_encoder // [n] — encoder grad accumulator (additive) +) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + grad_h_encoder[i] += lambda * grad_h_head[i]; +} diff --git a/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu b/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu index 031b4edd0..f7ac35d28 100644 --- a/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu +++ b/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu @@ -58,6 +58,51 @@ #define RL_ENTROPY_COEF_INDEX 403 +// ───────────────────────────────────────────────────────────────────── +// ppo_policy_logits_fwd (Phase E.2): +// Standalone linear forward producing per-batch logits from h_t. The +// surrogate kernel takes these logits as input; previously the +// integrated trainer had no path to materialise them outside of +// the surrogate compute. +// +// logits[b, a] = b_pi[a] + Σ_c W_pi[a, c] × h_t[b, c] +// +// Block layout: +// grid = (B, N_ACTIONS, 1) +// block = (HIDDEN_DIM, 1, 1) +// shared_mem_bytes = 0 (each thread independently strides; thread 0 +// performs the per-output dot product. The +// block over HIDDEN_DIM is sized for backward +// compatibility with grad_w_b_h_t's reuse of +// the same h_t layout in shared.) +// +// Single output slot per block (b, a) → thread 0 reduces inline. Other +// threads idle (block size mirrors backward's HIDDEN_DIM for caller +// uniformity; perf is non-critical since N_ACTIONS=9 means 9×B blocks). +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void ppo_policy_logits_fwd( + const float* __restrict__ w, // [N_ACTIONS, HIDDEN_DIM] + const float* __restrict__ b, // [N_ACTIONS] + const float* __restrict__ h_t, // [B * HIDDEN_DIM] + int B, + float* __restrict__ logits // [B * N_ACTIONS] +) { + const int batch = blockIdx.x; + const int act = blockIdx.y; + const int c = threadIdx.x; + if (batch >= B) return; + if (act >= N_ACTIONS) return; + if (c != 0) return; // single-writer per output slot + + float acc = b[act]; + #pragma unroll + for (int i = 0; i < HIDDEN_DIM; ++i) { + acc += w[act * HIDDEN_DIM + i] * h_t[batch * HIDDEN_DIM + i]; + } + logits[batch * N_ACTIONS + act] = acc; +} + + // ───────────────────────────────────────────────────────────────────── // ppo_clipped_surrogate_fwd: // Computes softmax → log π_new(a_taken|s) → ratio → clipped surrogate + @@ -280,3 +325,66 @@ extern "C" __global__ void ppo_clipped_surrogate_bwd( } grad_logits[base + act] = pg_grad; } + + +// ───────────────────────────────────────────────────────────────────── +// ppo_grad_w_b_h_t (Phase E.2): +// Given per-batch grad_logits [B × N_ACTIONS] from +// `ppo_clipped_surrogate_bwd`, propagate to weight / bias gradients +// (per-batch scratch — caller reduces via reduce_axis0) and to grad_h_t +// (per-(batch, c) sole writer; OVERWRITE). +// +// Mathematics (per batch b, action a, hidden unit c): +// grad_w_per_batch[b, a, c] = grad_logits[b, a] × h_t[b, c] +// grad_b_per_batch[b, a] = grad_logits[b, a] +// grad_h_t[b, c] = Σ_a grad_logits[b, a] × w[a, c] +// +// Block layout (mirrors aux_heads_bwd): +// grid = (B, 1, 1) +// block = (HIDDEN_DIM = 128, 1, 1) +// shared_mem_bytes = N_ACTIONS * sizeof(float) (grad_logits stage) +// +// Thread c is the sole writer of grad_w_per_batch[batch, a, c] for all a, +// and of grad_h_t[batch, c]. The first N_ACTIONS threads additionally +// write grad_b_per_batch[batch, a]. Per `feedback_no_atomicadd.md`: no +// atomics; each output slot has exactly one writer. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void ppo_grad_w_b_h_t( + const float* __restrict__ w, // [N_ACTIONS, HIDDEN_DIM] + const float* __restrict__ h_t, // [B * HIDDEN_DIM] + const float* __restrict__ grad_logits, // [B * N_ACTIONS] + int B, + float* __restrict__ grad_w_per_batch, // [B * N_ACTIONS * HIDDEN_DIM] + float* __restrict__ grad_b_per_batch, // [B * N_ACTIONS] + float* __restrict__ grad_h_t // [B * HIDDEN_DIM] (OVERWRITE) +) { + __shared__ float s_gl[N_ACTIONS]; + + const int batch = blockIdx.x; + const int c = threadIdx.x; + if (batch >= B) return; + + // Stage grad_logits[batch, :] (9 floats) — first N_ACTIONS threads. + if (c < N_ACTIONS) { + s_gl[c] = grad_logits[batch * N_ACTIONS + c]; + // Sole writer of grad_b_per_batch[batch, c]. + grad_b_per_batch[batch * N_ACTIONS + c] = s_gl[c]; + } + __syncthreads(); + + if (c >= HIDDEN_DIM) return; + + // grad_w + grad_h_t — thread role c = hidden unit. + const float h_bc = h_t[batch * HIDDEN_DIM + c]; + float gh_acc = 0.0f; + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) { + const float g_a = s_gl[a]; + grad_w_per_batch[(long long) batch * N_ACTIONS * HIDDEN_DIM + + (long long) a * HIDDEN_DIM + c] + = g_a * h_bc; + gh_acc += g_a * w[a * HIDDEN_DIM + c]; + } + // OVERWRITE — caller folds via grad_h_accumulate_scaled. + grad_h_t[batch * HIDDEN_DIM + c] = gh_acc; +} diff --git a/crates/ml-alpha/cuda/v_head_fwd_bwd.cu b/crates/ml-alpha/cuda/v_head_fwd_bwd.cu new file mode 100644 index 000000000..a89ea58d4 --- /dev/null +++ b/crates/ml-alpha/cuda/v_head_fwd_bwd.cu @@ -0,0 +1,163 @@ +// v_head_fwd_bwd.cu — scalar V(s) head fwd + MSE bwd (Phase E.2). +// +// Per the integrated RL trainer plan +// (docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md) the value +// head is a single linear layer `V(s) = b_v + Σ_c w_v[c] × h_t[b, c]` +// supervised by MSE against the bootstrap return target. +// +// Forward (`v_head_fwd`): +// v_pred[b] = b_v + Σ_c w_v[c] × h_t[b, c] +// +// Backward (`v_head_bwd`): +// per-batch grad scalar s_b = 2 × (v_pred[b] - r_target[b]) / B +// (mean-MSE — the /B keeps the +// gradient scale stable as B varies, +// matching the BCE/aux convention) +// per-batch loss scratch loss_per_batch[b] = (v_pred[b] - r_target[b])² +// per-batch weight grad grad_w_per_batch[b, c] = s_b × h_t[b, c] +// per-batch bias grad grad_b_per_batch[b] = s_b +// per-batch h grad grad_h_t[b, c] = s_b × w_v[c] +// (OVERWRITE — the trainer's +// grad_h_accumulate kernel folds it into the +// encoder gradient via scaled +=) +// +// Cross-batch reduction of grad_w_per_batch / grad_b_per_batch is the +// caller's responsibility via the existing `reduce_axis0` kernel — same +// pattern as `aux_heads_bwd`. Per `feedback_no_atomicadd.md`: this file +// introduces no atomicAdd. +// +// Block layout: +// Forward: grid=(B, 1, 1); block=(HIDDEN_DIM, 1, 1). +// Thread c stages w_v[c] into shared, then thread 0 reduces +// the dot product. (B writes one scalar per block — no +// cross-thread reduction across HIDDEN_DIM needed for +// cosmetic perf at B≤256.) +// Backward: same layout. Thread c is sole writer of +// grad_w_per_batch[b, c] and grad_h_t[b, c]; thread 0 +// additionally writes loss_per_batch[b] and grad_b_per_batch[b]. +// +// Constraints honoured: +// * `feedback_no_atomicadd.md` — per-batch scratch + reduce_axis0 +// * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs +// * `feedback_nvidia_grade_perf_for_kernels.md` — warp-uniform; one +// block per batch; cooperative staging of the small w_v vector. +// * `pearl_no_host_branches_in_captured_graph` — no host scalars +// affect control flow (only B as a captured int). + +#define HIDDEN_DIM 128 + + +// ───────────────────────────────────────────────────────────────────── +// v_head_fwd: per-batch scalar V projection. +// +// Launch: +// grid = (B, 1, 1) +// block = (HIDDEN_DIM = 128, 1, 1) +// shared_mem_bytes = HIDDEN_DIM * sizeof(float) (w_v stage) +// +// Inputs: +// w_v [HIDDEN_DIM] — single-output linear weights +// b_v [1] — bias scalar +// h_t [B × HIDDEN_DIM] — encoder hidden state +// B — batch size +// Outputs: +// v_pred[B] — predicted scalar V(s) per batch sample +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void v_head_fwd( + const float* __restrict__ w_v, // [HIDDEN_DIM] + const float* __restrict__ b_v, // [1] + const float* __restrict__ h_t, // [B * HIDDEN_DIM] + int B, + float* __restrict__ v_pred // [B] +) { + extern __shared__ float s_w[]; // [HIDDEN_DIM] + + const int batch = blockIdx.x; + const int c = threadIdx.x; + if (batch >= B) return; + if (c >= HIDDEN_DIM) return; + + s_w[c] = w_v[c]; + __syncthreads(); + + if (c == 0) { + float acc = b_v[0]; + #pragma unroll + for (int i = 0; i < HIDDEN_DIM; ++i) { + acc += s_w[i] * h_t[batch * HIDDEN_DIM + i]; + } + v_pred[batch] = acc; + } +} + + +// ───────────────────────────────────────────────────────────────────── +// v_head_bwd: per-batch MSE backward. +// +// Per-batch grad scalar : s_b = 2 × (v_pred[b] - r_target[b]) / B +// (mean-MSE divides by B at the source so +// downstream reductions can sum without +// an explicit normalisation step). +// Outputs: +// loss_per_batch [B] — (v - r)² (mean +// applied at reduction +// by the caller — kernel +// emits the squared err) +// grad_w_per_batch [B × HIDDEN_DIM] — per-batch scratch +// (caller reduces via +// reduce_axis0) +// grad_b_per_batch [B] — per-batch scratch +// grad_h_t [B × HIDDEN_DIM] — per-(batch, c) writes +// (OVERWRITE — sole +// writer per (b, c)) +// +// Launch: +// grid = (B, 1, 1) +// block = (HIDDEN_DIM = 128, 1, 1) +// shared_mem_bytes = HIDDEN_DIM * sizeof(float) (w_v stage) +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void v_head_bwd( + const float* __restrict__ w_v, // [HIDDEN_DIM] + const float* __restrict__ h_t, // [B * HIDDEN_DIM] + const float* __restrict__ v_pred, // [B] + const float* __restrict__ r_target, // [B] + int B, + float* __restrict__ loss_per_batch, // [B] + float* __restrict__ grad_w_per_batch,// [B * HIDDEN_DIM] + float* __restrict__ grad_b_per_batch,// [B] + float* __restrict__ grad_h_t // [B * HIDDEN_DIM] (OVERWRITE) +) { + extern __shared__ float s_w[]; // [HIDDEN_DIM] + + const int batch = blockIdx.x; + const int c = threadIdx.x; + if (batch >= B) return; + if (c >= HIDDEN_DIM) return; + + s_w[c] = w_v[c]; + __syncthreads(); + + // ── Per-batch residual + loss + grad scalar (thread 0 only) ───── + __shared__ float s_grad; + if (c == 0) { + const float diff = v_pred[batch] - r_target[batch]; + // mean-MSE: divide grad by B at the source so per-batch + // scratches sum to mean(grad) after reduce_axis0. + const float inv_B = 1.0f / (float) B; + s_grad = 2.0f * diff * inv_B; + loss_per_batch[batch] = diff * diff; // raw squared err — caller divides by B + grad_b_per_batch[batch] = s_grad; + } + __syncthreads(); + + const float g = s_grad; + const float w_c = s_w[c]; + const float h_bc = h_t[batch * HIDDEN_DIM + c]; + + // Sole writer of grad_w_per_batch[batch, c] — no atomic. + grad_w_per_batch[batch * HIDDEN_DIM + c] = g * h_bc; + // Sole writer of grad_h_t[batch, c] — OVERWRITE. Caller's + // grad_h_accumulate combines this into the encoder gradient via + // scaled +=. + grad_h_t[batch * HIDDEN_DIM + c] = g * w_c; +} diff --git a/crates/ml-alpha/src/rl/dqn.rs b/crates/ml-alpha/src/rl/dqn.rs index 61cd1565e..e789d36d8 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -41,12 +41,15 @@ use std::sync::Arc; use anyhow::{Context, Result}; -use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut}; +use cudarc::driver::{ + CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg, +}; use ml_core::cuda_autograd::init::scoped_init_seed; use ml_core::device::MlDevice; use rand::{Rng, SeedableRng}; use rand_chacha::ChaCha8Rng; +use crate::heads::HIDDEN_DIM; use crate::pinned_mem::MappedF32Buffer; use crate::rl::common::{N_ACTIONS, Q_N_ATOMS}; @@ -92,6 +95,12 @@ pub struct DqnHead { /// Backward kernel: `dqn_distributional_q_bwd`. Computes categorical /// CE loss + `∂L/∂logits` given a pre-projected target distribution. pub bwd_fn: CudaFunction, + /// Phase E.2 backward chain: `dqn_grad_w_b_h_t`. Maps the `grad_logits` + /// output of `bwd_fn` into per-batch `grad_w` / `grad_b` scratch buffers + /// + an OVERWRITE `grad_h_t` slot for the integrated trainer to fold + /// into the encoder grad accumulator. Per-batch scratch — cross-batch + /// reduction is the caller's responsibility via `reduce_axis0`. + pub grad_w_b_h_t_fn: CudaFunction, /// Online-network weights `[N_ACTIONS × Q_N_ATOMS, HIDDEN_DIM]`, /// row-major. Each "row" indexes one `(action, atom)` output slot. @@ -123,6 +132,9 @@ impl DqnHead { let bwd_fn = module .load_function("dqn_distributional_q_bwd") .context("load dqn_distributional_q_bwd")?; + let grad_w_b_h_t_fn = module + .load_function("dqn_grad_w_b_h_t") + .context("load dqn_grad_w_b_h_t")?; // Per `pearl_scoped_init_seed_for_reproducibility`: install the // scoped seed guard BEFORE drawing any Xavier samples. @@ -155,6 +167,7 @@ impl DqnHead { _module: module, fwd_fn, bwd_fn, + grad_w_b_h_t_fn, w_d, b_d, w_target_d, @@ -162,6 +175,131 @@ impl DqnHead { }) } + /// Phase E.2 forward: launch `dqn_distributional_q_fwd` on `h_t`, + /// writing raw atom logits `[B × N_ACTIONS × Q_N_ATOMS]` into + /// `logits_out`. The trainer pre-allocates the output buffer so the + /// same memory is reused across step() calls. + pub fn forward( + &self, + h_t: &CudaSlice, + b_size: usize, + logits_out: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(logits_out.len(), b_size * N_ACTIONS * Q_N_ATOMS); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, N_ACTIONS as u32, 1), + block_dim: (Q_N_ATOMS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.fwd_fn); + launch + .arg(&self.w_d) + .arg(&self.b_d) + .arg(h_t) + .arg(&b_i) + .arg(logits_out); + unsafe { + launch.launch(cfg).context("dqn_distributional_q_fwd launch")?; + } + Ok(()) + } + + /// Phase E.2 backward chain stage 1 — categorical CE on the taken + /// action's atom distribution. Writes `loss_out [1]` and + /// `grad_logits [B × N_ACTIONS × Q_N_ATOMS]`. The `loss_out` scalar + /// is accumulated with atomicAdd (Phase C deferred fix — see kernel + /// header). `grad_logits` flows into `grad_w_b_h_t` next. + #[allow(clippy::too_many_arguments)] + pub fn backward_logits( + &self, + logits: &CudaSlice, + target_dist: &CudaSlice, + actions_taken: &CudaSlice, + b_size: usize, + loss_out: &mut CudaSlice, + grad_logits: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(logits.len(), b_size * N_ACTIONS * Q_N_ATOMS); + debug_assert_eq!(target_dist.len(), b_size * Q_N_ATOMS); + debug_assert_eq!(actions_taken.len(), b_size); + debug_assert_eq!(loss_out.len(), 1); + debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS * Q_N_ATOMS); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (Q_N_ATOMS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + launch + .arg(logits) + .arg(target_dist) + .arg(actions_taken) + .arg(&b_i) + .arg(loss_out) + .arg(grad_logits); + unsafe { + launch.launch(cfg).context("dqn_distributional_q_bwd launch")?; + } + Ok(()) + } + + /// Phase E.2 backward chain stage 2 — propagate `grad_logits` into + /// `grad_w_per_batch [B × N_ACTIONS × Q_N_ATOMS × HIDDEN_DIM]`, + /// `grad_b_per_batch [B × N_ACTIONS × Q_N_ATOMS]`, and + /// `grad_h_t [B × HIDDEN_DIM]` (OVERWRITE). The caller reduces + /// `grad_w_per_batch` / `grad_b_per_batch` along axis 0 via the + /// existing `reduce_axis0` kernel to get the final + /// `grad_w [N_ACTIONS × Q_N_ATOMS × HIDDEN_DIM]` / + /// `grad_b [N_ACTIONS × Q_N_ATOMS]`. + /// + /// `grad_h_t` is OVERWRITE: the trainer's `grad_h_accumulate` + /// kernel reads it and folds via scaled +=. + #[allow(clippy::too_many_arguments)] + pub fn backward_to_w_b_h( + &self, + h_t: &CudaSlice, + grad_logits: &CudaSlice, + b_size: usize, + grad_w_per_batch: &mut CudaSlice, + grad_b_per_batch: &mut CudaSlice, + grad_h_t: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS * Q_N_ATOMS); + debug_assert_eq!( + grad_w_per_batch.len(), + b_size * N_ACTIONS * Q_N_ATOMS * HIDDEN_DIM + ); + debug_assert_eq!(grad_b_per_batch.len(), b_size * N_ACTIONS * Q_N_ATOMS); + debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM); + + let k_out = (N_ACTIONS * Q_N_ATOMS) as usize; + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: (k_out * std::mem::size_of::()) as u32, + }; + let mut launch = self.stream.launch_builder(&self.grad_w_b_h_t_fn); + launch + .arg(&self.w_d) + .arg(h_t) + .arg(grad_logits) + .arg(&b_i) + .arg(grad_w_per_batch) + .arg(grad_b_per_batch) + .arg(grad_h_t); + unsafe { + launch.launch(cfg).context("dqn_grad_w_b_h_t launch")?; + } + Ok(()) + } + /// Stream used to launch all kernels owned by this head. Phase E's /// training loop reads this when scheduling the soft-update kernel. pub fn stream(&self) -> &Arc { diff --git a/crates/ml-alpha/src/rl/ppo.rs b/crates/ml-alpha/src/rl/ppo.rs index cbf5000f7..469dd7fe5 100644 --- a/crates/ml-alpha/src/rl/ppo.rs +++ b/crates/ml-alpha/src/rl/ppo.rs @@ -34,12 +34,15 @@ use std::sync::Arc; use anyhow::{Context, Result}; -use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut}; +use cudarc::driver::{ + CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg, +}; use ml_core::cuda_autograd::init::scoped_init_seed; use ml_core::device::MlDevice; use rand::{Rng, SeedableRng}; use rand_chacha::ChaCha8Rng; +use crate::heads::HIDDEN_DIM; use crate::pinned_mem::MappedF32Buffer; use crate::rl::common::N_ACTIONS; @@ -47,6 +50,8 @@ const PPO_SURR_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/ppo_clipped_surrogate.cubin" )); +const V_HEAD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/v_head_fwd_bwd.cubin")); /// Construction config for [`PolicyHead`] and [`ValueHead`]. #[derive(Clone, Debug)] @@ -82,6 +87,22 @@ pub struct PolicyHead { /// Backward kernel: `ppo_clipped_surrogate_bwd`. Computes per-logit /// gradient with the clip-mask applied (zero outside the clip band). pub surrogate_bwd_fn: CudaFunction, + /// Phase E.2 logits-projection kernel: `ppo_clipped_logits_fwd`. The + /// policy head uses the existing `ppo_clipped_surrogate_fwd` for the + /// loss path; for the standalone logits forward (used in the rollout + /// path AND as the input to the surrogate kernel) we run a separate + /// linear projection. Since PPO's surrogate kernel takes pre-computed + /// logits as input, the trainer first runs a small linear forward + /// (`policy_linear_fwd_fn`) that produces `[B × N_ACTIONS]` logits + /// from `h_t`. The surrogate kernel then consumes them. + pub policy_linear_fwd_fn: CudaFunction, + /// Phase E.2 backward chain stage 2: `ppo_grad_w_b_h_t`. Maps the + /// `grad_logits` output of `surrogate_bwd_fn` into per-batch + /// `grad_w` / `grad_b` scratch buffers + an OVERWRITE `grad_h_t` + /// slot for the integrated trainer to fold into the encoder grad + /// accumulator. Per-batch scratch — cross-batch reduction is the + /// caller's responsibility via `reduce_axis0`. + pub grad_w_b_h_t_fn: CudaFunction, /// Online weights `[N_ACTIONS, HIDDEN_DIM]`, row-major. Each row is /// one action's projection vector. @@ -110,6 +131,12 @@ impl PolicyHead { let surrogate_bwd_fn = module .load_function("ppo_clipped_surrogate_bwd") .context("load ppo_clipped_surrogate_bwd")?; + let policy_linear_fwd_fn = module + .load_function("ppo_policy_logits_fwd") + .context("load ppo_policy_logits_fwd")?; + let grad_w_b_h_t_fn = module + .load_function("ppo_grad_w_b_h_t") + .context("load ppo_grad_w_b_h_t")?; // Per pearl_scoped_init_seed_for_reproducibility: install the // scoped seed guard BEFORE drawing any Xavier samples. @@ -133,6 +160,8 @@ impl PolicyHead { _module: module, surrogate_fwd_fn, surrogate_bwd_fn, + policy_linear_fwd_fn, + grad_w_b_h_t_fn, w_d, b_d, }) @@ -150,18 +179,203 @@ impl PolicyHead { pub fn hidden_dim(&self) -> usize { self.cfg.hidden_dim } + + /// Phase E.2 forward — materialise per-batch action logits `[B × N_ACTIONS]` + /// via a small linear projection. Caller then feeds into the + /// surrogate kernel which fuses softmax + ratio + clip + loss. + pub fn forward_logits( + &self, + h_t: &CudaSlice, + b_size: usize, + logits_out: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(logits_out.len(), b_size * N_ACTIONS); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, N_ACTIONS as u32, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.policy_linear_fwd_fn); + launch + .arg(&self.w_d) + .arg(&self.b_d) + .arg(h_t) + .arg(&b_i) + .arg(logits_out); + unsafe { + launch.launch(cfg).context("ppo_policy_logits_fwd launch")?; + } + Ok(()) + } + + /// Phase E.2 surrogate forward — runs `ppo_clipped_surrogate_fwd` + /// using the trainer-provided logits + advantages + returns + V_pred. + /// Reads ε from `isv[RL_PPO_CLIP_INDEX]` and entropy coef from + /// `isv[RL_ENTROPY_COEF_INDEX]`. + #[allow(clippy::too_many_arguments)] + pub fn surrogate_forward( + &self, + logits: &CudaSlice, + log_pi_old: &CudaSlice, + actions: &CudaSlice, + advantages: &CudaSlice, + returns: &CudaSlice, + v_pred: &CudaSlice, + isv_d: &CudaSlice, + b_size: usize, + pi_log_prob: &mut CudaSlice, + entropy: &mut CudaSlice, + loss_pi: &mut CudaSlice, + loss_v: &mut CudaSlice, + loss_entropy: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(logits.len(), b_size * N_ACTIONS); + debug_assert_eq!(log_pi_old.len(), b_size); + debug_assert_eq!(actions.len(), b_size); + debug_assert_eq!(advantages.len(), b_size); + debug_assert_eq!(returns.len(), b_size); + debug_assert_eq!(v_pred.len(), b_size); + debug_assert_eq!(pi_log_prob.len(), b_size); + debug_assert_eq!(entropy.len(), b_size); + debug_assert_eq!(loss_pi.len(), 1); + debug_assert_eq!(loss_v.len(), 1); + debug_assert_eq!(loss_entropy.len(), 1); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (N_ACTIONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.surrogate_fwd_fn); + launch + .arg(logits) + .arg(log_pi_old) + .arg(actions) + .arg(advantages) + .arg(returns) + .arg(v_pred) + .arg(isv_d) + .arg(&b_i) + .arg(pi_log_prob) + .arg(entropy) + .arg(loss_pi) + .arg(loss_v) + .arg(loss_entropy); + unsafe { + launch + .launch(cfg) + .context("ppo_clipped_surrogate_fwd launch")?; + } + Ok(()) + } + + /// Phase E.2 surrogate backward — emits `grad_logits [B × N_ACTIONS]`. + /// Subsequent stage is `backward_to_w_b_h`. + #[allow(clippy::too_many_arguments)] + pub fn surrogate_backward_logits( + &self, + logits: &CudaSlice, + log_pi_old: &CudaSlice, + actions: &CudaSlice, + advantages: &CudaSlice, + isv_d: &CudaSlice, + b_size: usize, + grad_logits: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(logits.len(), b_size * N_ACTIONS); + debug_assert_eq!(log_pi_old.len(), b_size); + debug_assert_eq!(actions.len(), b_size); + debug_assert_eq!(advantages.len(), b_size); + debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (N_ACTIONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.surrogate_bwd_fn); + launch + .arg(logits) + .arg(log_pi_old) + .arg(actions) + .arg(advantages) + .arg(isv_d) + .arg(&b_i) + .arg(grad_logits); + unsafe { + launch + .launch(cfg) + .context("ppo_clipped_surrogate_bwd launch")?; + } + Ok(()) + } + + /// Phase E.2 backward chain stage 2 — propagate `grad_logits` into + /// `grad_w_per_batch [B × N_ACTIONS × HIDDEN_DIM]`, + /// `grad_b_per_batch [B × N_ACTIONS]`, and + /// `grad_h_t [B × HIDDEN_DIM]` (OVERWRITE). Caller reduces axis 0 via + /// `reduce_axis0`. + #[allow(clippy::too_many_arguments)] + pub fn backward_to_w_b_h( + &self, + h_t: &CudaSlice, + grad_logits: &CudaSlice, + b_size: usize, + grad_w_per_batch: &mut CudaSlice, + grad_b_per_batch: &mut CudaSlice, + grad_h_t: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS); + debug_assert_eq!(grad_w_per_batch.len(), b_size * N_ACTIONS * HIDDEN_DIM); + debug_assert_eq!(grad_b_per_batch.len(), b_size * N_ACTIONS); + debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: (N_ACTIONS * std::mem::size_of::()) as u32, + }; + let mut launch = self.stream.launch_builder(&self.grad_w_b_h_t_fn); + launch + .arg(&self.w_d) + .arg(h_t) + .arg(grad_logits) + .arg(&b_i) + .arg(grad_w_per_batch) + .arg(grad_b_per_batch) + .arg(grad_h_t); + unsafe { + launch.launch(cfg).context("ppo_grad_w_b_h_t launch")?; + } + Ok(()) + } } /// Value head: linear projection `h_t [B, HIDDEN_DIM] → scalar V(s) [B]`. -/// Trained by MSE between `V_pred` and the bootstrap target — Phase E -/// computes the target as `Q(s, a*) - A_t` (or equivalently the -/// done-aware discounted return); the surrogate kernel reports per-batch -/// squared error and Phase E accumulates the V-head gradient inline. +/// Trained by MSE between `V_pred` and the bootstrap target — Phase E.2 +/// computes the per-batch squared error + per-batch grad scratch via the +/// dedicated `v_head_fwd_bwd.cu` kernel pair. pub struct ValueHead { #[allow(dead_code)] cfg: PpoHeadsConfig, stream: Arc, + _module: Arc, + /// Forward kernel: `v_head_fwd`. Reads h_t and emits per-batch + /// V_pred [B]. + pub fwd_fn: CudaFunction, + /// Backward kernel: `v_head_bwd`. Emits per-batch loss + per-batch + /// grad_w / grad_b scratch + OVERWRITE grad_h_t. Caller reduces + /// grad_w / grad_b via `reduce_axis0`. + pub bwd_fn: CudaFunction, + /// Online weights `[1, HIDDEN_DIM]` (single output neuron). pub w_d: CudaSlice, /// Online bias `[1]`. @@ -177,6 +391,16 @@ impl ValueHead { .cuda_stream() .context("value_head stream")? .clone(); + let ctx = dev.cuda_context().context("value_head ctx")?; + let module = ctx + .load_cubin(V_HEAD_CUBIN.to_vec()) + .context("load v_head_fwd_bwd cubin")?; + let fwd_fn = module + .load_function("v_head_fwd") + .context("load v_head_fwd")?; + let bwd_fn = module + .load_function("v_head_bwd") + .context("load v_head_bwd")?; // Per pearl_scoped_init_seed_for_reproducibility. let _seed_guard = scoped_init_seed(cfg.seed.wrapping_add(0x5EED_1234)); @@ -196,6 +420,9 @@ impl ValueHead { Ok(Self { cfg, stream, + _module: module, + fwd_fn, + bwd_fn, w_d, b_d, }) @@ -210,6 +437,83 @@ impl ValueHead { pub fn hidden_dim(&self) -> usize { self.cfg.hidden_dim } + + /// Phase E.2 forward — V(s) = b + Σ_c w[c] × h_t[b, c]. Writes + /// `v_pred [B]`. + pub fn forward( + &self, + h_t: &CudaSlice, + b_size: usize, + v_pred: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(v_pred.len(), b_size); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: (HIDDEN_DIM * std::mem::size_of::()) as u32, + }; + let mut launch = self.stream.launch_builder(&self.fwd_fn); + launch + .arg(&self.w_d) + .arg(&self.b_d) + .arg(h_t) + .arg(&b_i) + .arg(v_pred); + unsafe { + launch.launch(cfg).context("v_head_fwd launch")?; + } + Ok(()) + } + + /// Phase E.2 backward — MSE bwd. Emits `loss_per_batch [B]` (raw + /// squared error; caller divides by B if mean is desired), + /// `grad_w_per_batch [B × HIDDEN_DIM]`, `grad_b_per_batch [B]`, + /// and `grad_h_t [B × HIDDEN_DIM]` (OVERWRITE). + #[allow(clippy::too_many_arguments)] + pub fn backward( + &self, + h_t: &CudaSlice, + v_pred: &CudaSlice, + r_target: &CudaSlice, + b_size: usize, + loss_per_batch: &mut CudaSlice, + grad_w_per_batch: &mut CudaSlice, + grad_b_per_batch: &mut CudaSlice, + grad_h_t: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(v_pred.len(), b_size); + debug_assert_eq!(r_target.len(), b_size); + debug_assert_eq!(loss_per_batch.len(), b_size); + debug_assert_eq!(grad_w_per_batch.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(grad_b_per_batch.len(), b_size); + debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM); + + let b_i = b_size as i32; + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: (HIDDEN_DIM * std::mem::size_of::()) as u32, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + launch + .arg(&self.w_d) + .arg(h_t) + .arg(v_pred) + .arg(r_target) + .arg(&b_i) + .arg(loss_per_batch) + .arg(grad_w_per_batch) + .arg(grad_b_per_batch) + .arg(grad_h_t); + unsafe { + launch.launch(cfg).context("v_head_bwd launch")?; + } + Ok(()) + } } // ── pinned-staging upload helper (mirrors dqn.rs::upload) ── diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 27b1f34c4..ff6e02462 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -1,36 +1,92 @@ //! Integrated RL trainer: BCE + DQN/C51 Q + PPO π + V + aux on a shared //! Mamba2+CfC encoder. //! -//! Phase E.1 (this commit): synthetic-batch end-to-end smoke. Trainer owns: -//! - PerceptionTrainer (encoder + BCE + aux machinery, reused via -//! `forward_encoder` from Phase B) -//! - DqnHead (Phase C) -//! - PolicyHead + ValueHead (Phase D) -//! - ISV device buffer (allocated to RL_SLOTS_END) +//! Phase E.2 (this commit): real GPU forward + backward for the Q, π, V +//! heads on synthetic data. Replaces Phase E.1's placeholder host-side +//! loss scalars with kernel-emitted losses + per-head Adam updates. //! -//! `step_synthetic` runs ONE training step on synthetic data without -//! LobSim. Confirms the trainer constructs cleanly, the encoder forward -//! lands in `h_t_d`, and the loss-balance λ read path returns finite -//! per-head weights. Real GPU kernel orchestration for Q/π/V forward + -//! backward into the encoder is Phase E.2's scope. +//! Architecture (synthetic-data path): //! -//! Phase E.2 wires LobSim, activates toy bandits, and switches to -//! captured-graph forward+backward. +//! 1. `perception.forward_encoder(snapshots)` produces `h_t [B × HIDDEN_DIM]` +//! via the captured-graph encoder forward path. +//! 2. `dqn_head.forward(h_t, ...)` → `q_logits [B × N_ACTIONS × Q_N_ATOMS]` +//! `policy_head.forward_logits(h_t, ...)` → `pi_logits [B × N_ACTIONS]` +//! `value_head.forward(h_t, ...)` → `v_pred [B]` +//! 3. `policy_head.surrogate_forward(...)` consumes pi_logits + synthetic +//! log π_old + advantages + returns + v_pred + ISV(ε,coef) → 3 scalar +//! loss accumulators + per-batch log π / entropy. +//! 4. Host-side build of Bellman target distribution from synthetic +//! rewards (single-atom mass at the closest C51 atom index). Upload +//! via mapped-pinned + DtoD. +//! 5. `dqn_head.backward_logits(...)` → `q_loss [1]` + `grad_logits`. +//! 6. `dqn_head.backward_to_w_b_h(...)` → per-batch grad_w / grad_b +//! scratch + grad_h_t (OVERWRITE). +//! 7. `policy_head.surrogate_backward_logits(...)` → `pi_grad_logits`. +//! 8. `policy_head.backward_to_w_b_h(...)` → per-batch grad_w / grad_b +//! + grad_h_t (OVERWRITE). +//! 9. `value_head.backward(...)` → per-batch loss + grad_w / grad_b +//! + grad_h_t (OVERWRITE). +//! 10. `reduce_axis0` collapses per-batch grad_w / grad_b → final +//! `grad_w [...]` / `grad_b [...]` for each head. +//! 11. Per-head `AdamW.step()` runs on the reduced gradients. Each head +//! owns 2 Adam instances (one for `w_d`, one for `b_d`). +//! 12. Per-batch losses (V head + DQN loss scalar) are reduced/read back +//! via mapped-pinned and returned in `IntegratedStepStats`. +//! +//! ENCODER BACKWARD DEFERRAL (Phase E.3): the Q/π/V kernels emit +//! `grad_h_t` outputs but Phase E.2 does NOT yet wire them into the +//! encoder's backward pass. The `grad_h_accumulate_scaled` kernel is +//! loaded and ready but the integrated trainer does not have a clean +//! entry point on `PerceptionTrainer` to consume a caller-provided +//! `grad_h_t` (the existing `step_batched` runs its own encoder +//! backward end-to-end). Phase E.3 lands LobSim AND a +//! `PerceptionTrainer::backward_encoder_with_grad_h_t` entry point in +//! the same commit. Until then, the encoder learns only from BCE+aux +//! via `perception.step_batched`, while the Q/π/V heads update their +//! own weights via per-head Adam in `step_synthetic`. +//! +//! Constraints honoured: +//! * `feedback_no_stubs.md` — every kernel runs real math, no +//! placeholder host losses. +//! * `feedback_no_atomicadd.md` — new kernels use per-batch scratch +//! + `reduce_axis0` reduction (DQN/PPO bwd loss scalars still atomic +//! — Phase C/D deferral, kept loud-flagged in their kernel headers). +//! * `feedback_no_htod_htoh_only_mapped_pinned.md` — all CPU↔GPU paths +//! via `MappedF32Buffer` + DtoD pattern from `pinned_mem.rs`. +//! * `feedback_no_nvrtc.md` — pre-compiled cubins via build.rs. +//! * `pearl_controller_anchors_isv_driven` — λs read from ISV via +//! `read_loss_lambdas_from_isv`; ε / coef read inside the surrogate +//! kernel from `ISV[402,403]`. use std::sync::Arc; use anyhow::{Context, Result}; -use cudarc::driver::{CudaSlice, CudaStream}; +use cudarc::driver::{ + CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, + PushKernelArg, +}; use ml_core::device::MlDevice; use crate::cfc::snap_features::Mbp10RawInput; use crate::heads::HIDDEN_DIM; +use crate::pinned_mem::{MappedF32Buffer, MappedI32Buffer}; +use crate::rl::common::{N_ACTIONS, Q_N_ATOMS, Q_V_MAX, Q_V_MIN}; use crate::rl::dqn::{DqnHead, DqnHeadConfig}; use crate::rl::isv_slots::RL_SLOTS_END; use crate::rl::loss_balance::{read_loss_lambdas_from_isv, LossLambdas}; use crate::rl::ppo::{PolicyHead, PpoHeadsConfig, ValueHead}; +use crate::trainer::optim::AdamW; use crate::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig}; +const GRAD_H_ACCUMULATE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/grad_h_accumulate.cubin")); +const REDUCE_AXIS0_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin")); + +/// Default learning rate for the Q/π/V heads in Phase E.2 synthetic-data +/// smoke. Phase E.3 reads this from ISV (a per-head LR controller). +const PHASE_E2_DEFAULT_LR: f32 = 1e-3; + /// Configuration for [`IntegratedTrainer`]. Wraps a `PerceptionTrainerConfig` /// (the encoder + BCE + aux side) plus RL-specific overrides for the new /// heads. @@ -73,18 +129,37 @@ pub struct IntegratedTrainer { /// PPO value head (Phase D). pub value_head: ValueHead, + /// Per-head Adam optimisers — Phase E.2. One pair (w + b) per head. + /// All share the existing project-wide `adamw_step` cubin via the + /// `AdamW` wrapper; each instance owns its own m / v / step counter. + pub dqn_w_adam: AdamW, + pub dqn_b_adam: AdamW, + pub policy_w_adam: AdamW, + pub policy_b_adam: AdamW, + pub value_w_adam: AdamW, + pub value_b_adam: AdamW, + /// Device ISV buffer (RL-allocated, length RL_SLOTS_END = 412). /// Sentinel-zero initialised on construction so all controllers /// bootstrap on their first emit per pearl_first_observation_bootstrap. pub isv_d: CudaSlice, /// Host mirror of ISV for loss-balance λ reads. Refreshed before each - /// step via stream.memcpy_dtoh(isv_d, host). Phase E.1 reads it; Phase - /// E.2 may move λ reads to a fused device-side path. + /// step via stream.memcpy_dtoh(isv_d, host). pub isv_host: Vec, - #[allow(dead_code)] stream: Arc, + + // ── Kernel handles for grad combine + cross-batch reduce ────────── + _grad_h_module: Arc, + grad_h_accumulate_fn: CudaFunction, + _reduce_axis0_module: Arc, + reduce_axis0_fn: CudaFunction, + + /// Last-read PPO π loss host scalar — internal scratch threading the + /// surrogate kernel's atomic-write loss readback through the + /// borrow-checker dance in `step_synthetic`. NOT a public API. + last_pi_loss: f32, } impl IntegratedTrainer { @@ -125,38 +200,73 @@ impl IntegratedTrainer { ) .context("ValueHead::new")?; + // Per-head Adam optimisers — one per (head, w | b) pair. Each owns + // independent m/v buffers and a device-resident step counter (per + // pearl_no_host_branches_in_captured_graph: counter advancement + // happens inside a captured kernel, not in host code). + let lr = PHASE_E2_DEFAULT_LR; + let dqn_w_adam = AdamW::new(dev, dqn_head.w_d.len(), lr).context("dqn_w_adam")?; + let dqn_b_adam = AdamW::new(dev, dqn_head.b_d.len(), lr).context("dqn_b_adam")?; + let policy_w_adam = + AdamW::new(dev, policy_head.w_d.len(), lr).context("policy_w_adam")?; + let policy_b_adam = + AdamW::new(dev, policy_head.b_d.len(), lr).context("policy_b_adam")?; + let value_w_adam = AdamW::new(dev, value_head.w_d.len(), lr).context("value_w_adam")?; + let value_b_adam = AdamW::new(dev, value_head.b_d.len(), lr).context("value_b_adam")?; + // ISV buffer — zero-init, sentinel-bootstrap on first controller emit. let isv_d = stream .alloc_zeros::(RL_SLOTS_END) .context("alloc isv_d")?; let isv_host = vec![0.0_f32; RL_SLOTS_END]; + // Load helper kernels. + let ctx = dev.cuda_context().context("integrated ctx")?; + let grad_h_module = ctx + .load_cubin(GRAD_H_ACCUMULATE_CUBIN.to_vec()) + .context("load grad_h_accumulate cubin")?; + let grad_h_accumulate_fn = grad_h_module + .load_function("grad_h_accumulate_scaled") + .context("load grad_h_accumulate_scaled")?; + let reduce_axis0_module = ctx + .load_cubin(REDUCE_AXIS0_CUBIN.to_vec()) + .context("load reduce_axis0 cubin")?; + let reduce_axis0_fn = reduce_axis0_module + .load_function("reduce_axis0") + .context("load reduce_axis0")?; + Ok(Self { cfg, perception, dqn_head, policy_head, value_head, + dqn_w_adam, + dqn_b_adam, + policy_w_adam, + policy_b_adam, + value_w_adam, + value_b_adam, isv_d, isv_host, stream, + _grad_h_module: grad_h_module, + grad_h_accumulate_fn, + _reduce_axis0_module: reduce_axis0_module, + reduce_axis0_fn, + last_pi_loss: 0.0, }) } - /// Phase E.1: synthetic-batch smoke step. Runs the encoder forward, - /// computes all 5 head losses, combines with adaptive λ, and returns - /// the combined-loss stats. + /// Phase E.2: synthetic-batch end-to-end step. Runs real GPU forward + /// + backward + per-head Adam on Q, π, V heads. The encoder forward + /// is reused from `perception.forward_encoder`. Encoder backward is + /// DEFERRED to Phase E.3 (see module docstring) — the heads update + /// independently here. /// - /// LIMITATIONS: - /// - No LobSim. Actions, rewards, advantages are synthetic. - /// - DQN/PPO head weights are NOT updated (their Adam state lands in - /// Phase E.2). Only encoder + BCE/aux update. - /// - No replay buffer, no rollout buffer use. - /// - The Q/π/V losses returned here are deterministic host-side - /// placeholders that will be replaced by the actual kernel-emitted - /// scalars in Phase E.2's atomic refactor commit. - /// - /// Returns step stats for test assertions. + /// Returns the per-head kernel-emitted losses and the combined-loss + /// total weighted by the ISV-driven λs. + #[allow(clippy::too_many_arguments)] pub fn step_synthetic( &mut self, snapshots: &[Mbp10RawInput], @@ -164,57 +274,272 @@ impl IntegratedTrainer { synthetic_rewards: &[f32], // [B] — target Q values per action synthetic_advantages: &[f32], // [B] — synthetic A_t synthetic_returns: &[f32], // [B] — synthetic R_t for V regression - _synthetic_log_pi_old: &[f32], // [B] — synthetic log π_old (consumed by Phase E.2 kernel) + synthetic_log_pi_old: &[f32], // [B] — synthetic log π_old ) -> Result { - // Step 1: encoder forward via Phase B's forward_encoder. - // Drop the borrowed slice immediately — `forward_encoder` holds a - // &mut on self.perception so the borrow must end before we touch - // any other &mut perception method below. + let b_size = synthetic_actions.len(); + if b_size == 0 { + anyhow::bail!("step_synthetic: empty batch (b_size = 0)"); + } + + // Validate per-batch input lengths. + debug_assert_eq!(synthetic_rewards.len(), b_size); + debug_assert_eq!(synthetic_advantages.len(), b_size); + debug_assert_eq!(synthetic_returns.len(), b_size); + debug_assert_eq!(synthetic_log_pi_old.len(), b_size); + + // ── Step 1: encoder forward ────────────────────────────────── + // forward_encoder borrows &mut self.perception briefly; the + // returned slice lives in perception.h_t_d. We immediately end + // the borrow by reading the length, then re-borrow non-mutably + // for downstream kernel launches via `&self.perception.h_t_d`. let _ = self .perception .forward_encoder(snapshots) .context("forward_encoder")?; + // After this point self.perception is not borrowed mutably; we + // re-borrow h_t_d immutably via the perception accessor in the + // dispatch chain below. - // Step 2: refresh ISV host mirror so we can read loss λs. + // ── Step 2: refresh ISV host mirror ────────────────────────── self.stream .memcpy_dtoh(&self.isv_d, self.isv_host.as_mut_slice()) .context("isv dtoh")?; let lambdas = read_loss_lambdas_from_isv(&self.isv_host); - // PHASE E.1 SCOPE: at this point we'd run the 5 head forwards and - // compute their losses. The actual kernel orchestration is complex - // because it needs: - // - Q forward kernel call (dqn_head.fwd_fn) on h_t → atom logits - // - PPO surrogate fwd kernel (policy_head.surrogate_fwd_fn) on h_t - // - V head forward (small matmul on h_t) - // - BCE + aux are handled by perception (already wired) - // - All 5 losses summed with lambdas - // - Backward into encoder via existing perception backward + the - // additive grad_h_t paths from Q/π/V kernels - // - // For E.1 we LAND THE SCAFFOLDING ONLY. The actual fwd kernel calls - // are stubbed with synthetic placeholder loss values so the smoke - // test confirms the trainer constructs + runs the synthetic step - // without panic. Phase E.2 wires the real kernel calls. - // - // This is NOT a feedback_no_stubs violation: the stub is a temporary - // host-side computation that becomes a real GPU kernel call in - // Phase E.2's commit (atomic refactor). The struct fields + cubin - // handles are all real and used. + // Borrow encoder hidden state for forward kernels. + let h_t_borrow: &CudaSlice = self.perception.h_t_view(); + debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM); - let b = synthetic_actions.len().max(1) as f32; - let l_bce = 0.1_f32; // placeholder - let l_q = synthetic_rewards.iter().map(|r| r * r).sum::() / b; - let l_pi = -synthetic_advantages.iter().sum::() / b * 0.1; - let l_v = synthetic_returns.iter().map(|r| r * r).sum::() / b; - let l_aux = 0.5_f32; - // (placeholders deliberately deterministic; Phase E.2 replaces - // these with the actual kernel-emitted losses) + // ── Step 3: allocate per-step scratch ──────────────────────── + let k_dqn = N_ACTIONS * Q_N_ATOMS; + let mut q_logits_d = self.stream.alloc_zeros::(b_size * k_dqn)?; + let mut pi_logits_d = self.stream.alloc_zeros::(b_size * N_ACTIONS)?; + let mut v_pred_d = self.stream.alloc_zeros::(b_size)?; + // Loss accumulators (atomicAdd into single floats from kernels). + let q_loss_d = self.stream.alloc_zeros::(1)?; + let pi_loss_d = self.stream.alloc_zeros::(1)?; + let pi_loss_v_d = self.stream.alloc_zeros::(1)?; // PPO's V loss accumulator (atomic in kernel; not used for V grad) + let pi_loss_entropy_d = self.stream.alloc_zeros::(1)?; + // Per-batch V loss scratch (reduced below). + let mut v_loss_per_batch_d = self.stream.alloc_zeros::(b_size)?; + let mut v_loss_sum_d = self.stream.alloc_zeros::(1)?; + + // Per-batch diagnostic scratch (consumed by the surrogate kernel). + let mut pi_log_prob_d = self.stream.alloc_zeros::(b_size)?; + let mut entropy_d = self.stream.alloc_zeros::(b_size)?; + + // Per-batch grad scratch (reduced across batches via reduce_axis0). + let mut q_grad_logits_d = self.stream.alloc_zeros::(b_size * k_dqn)?; + let mut q_grad_w_per_batch_d = self + .stream + .alloc_zeros::(b_size * k_dqn * HIDDEN_DIM)?; + let mut q_grad_b_per_batch_d = self.stream.alloc_zeros::(b_size * k_dqn)?; + let mut q_grad_h_t_d = self.stream.alloc_zeros::(b_size * HIDDEN_DIM)?; + let mut q_grad_w_d = self.stream.alloc_zeros::(k_dqn * HIDDEN_DIM)?; + let mut q_grad_b_d = self.stream.alloc_zeros::(k_dqn)?; + + let mut pi_grad_logits_d = self.stream.alloc_zeros::(b_size * N_ACTIONS)?; + let mut pi_grad_w_per_batch_d = + self.stream.alloc_zeros::(b_size * N_ACTIONS * HIDDEN_DIM)?; + let mut pi_grad_b_per_batch_d = self.stream.alloc_zeros::(b_size * N_ACTIONS)?; + let mut pi_grad_h_t_d = self.stream.alloc_zeros::(b_size * HIDDEN_DIM)?; + let mut pi_grad_w_d = self.stream.alloc_zeros::(N_ACTIONS * HIDDEN_DIM)?; + let mut pi_grad_b_d = self.stream.alloc_zeros::(N_ACTIONS)?; + + let mut v_grad_w_per_batch_d = self.stream.alloc_zeros::(b_size * HIDDEN_DIM)?; + let mut v_grad_b_per_batch_d = self.stream.alloc_zeros::(b_size)?; + let mut v_grad_h_t_d = self.stream.alloc_zeros::(b_size * HIDDEN_DIM)?; + let mut v_grad_w_d = self.stream.alloc_zeros::(HIDDEN_DIM)?; + let mut v_grad_b_d = self.stream.alloc_zeros::(1)?; + + // ── Step 4: upload synthetic inputs via mapped-pinned ──────── + let actions_d = upload_i32(&self.stream, synthetic_actions)?; + let advantages_d = upload_f32(&self.stream, synthetic_advantages)?; + let returns_d = upload_f32(&self.stream, synthetic_returns)?; + let log_pi_old_d = upload_f32(&self.stream, synthetic_log_pi_old)?; + + // Build Bellman target distribution from synthetic rewards. + // Phase E.2: deterministic single-atom mass at the closest atom + // index for `clamp(reward, V_MIN, V_MAX)`. Phase E.3 replaces this + // with the real categorical projection kernel reading γ + the + // target net's bootstrap atom values. + let target_dist_host = build_synthetic_bellman_target(synthetic_rewards); + debug_assert_eq!(target_dist_host.len(), b_size * Q_N_ATOMS); + let target_dist_d = upload_f32(&self.stream, &target_dist_host)?; + + // ── Step 5: forward kernels (Q, π, V) ──────────────────────── + self.dqn_head + .forward(h_t_borrow, b_size, &mut q_logits_d) + .context("dqn_head.forward")?; + self.policy_head + .forward_logits(h_t_borrow, b_size, &mut pi_logits_d) + .context("policy_head.forward_logits")?; + self.value_head + .forward(h_t_borrow, b_size, &mut v_pred_d) + .context("value_head.forward")?; + + // PPO surrogate forward (uses pi_logits + v_pred). + { + let mut pi_loss_d_mut = pi_loss_d; + let mut pi_loss_v_d_mut = pi_loss_v_d; + let mut pi_loss_entropy_d_mut = pi_loss_entropy_d; + self.policy_head + .surrogate_forward( + &pi_logits_d, + &log_pi_old_d, + &actions_d, + &advantages_d, + &returns_d, + &v_pred_d, + &self.isv_d, + b_size, + &mut pi_log_prob_d, + &mut entropy_d, + &mut pi_loss_d_mut, + &mut pi_loss_v_d_mut, + &mut pi_loss_entropy_d_mut, + ) + .context("policy_head.surrogate_forward")?; + // Read back the scalar PPO π loss for the IntegratedStepStats. + // (Drop the mutable bindings; the device buffers are consumed + // for readback below.) + let l_pi_host = read_scalar_d(&self.stream, &pi_loss_d_mut)?; + let _l_pi_v_host = read_scalar_d(&self.stream, &pi_loss_v_d_mut)?; + let _l_pi_ent_host = read_scalar_d(&self.stream, &pi_loss_entropy_d_mut)?; + // Plumb back into the local for the final stat compose. + self.last_pi_loss = l_pi_host; + } + + // ── Step 6: DQN backward (logits → grad_w/b/h_t) ───────────── + let mut q_loss_d_mut = q_loss_d; + self.dqn_head + .backward_logits( + &q_logits_d, + &target_dist_d, + &actions_d, + b_size, + &mut q_loss_d_mut, + &mut q_grad_logits_d, + ) + .context("dqn_head.backward_logits")?; + let l_q_host = read_scalar_d(&self.stream, &q_loss_d_mut)?; + + self.dqn_head + .backward_to_w_b_h( + h_t_borrow, + &q_grad_logits_d, + b_size, + &mut q_grad_w_per_batch_d, + &mut q_grad_b_per_batch_d, + &mut q_grad_h_t_d, + ) + .context("dqn_head.backward_to_w_b_h")?; + + // Reduce per-batch grads to final shapes. + self.launch_reduce_axis0(&q_grad_w_per_batch_d, b_size, k_dqn * HIDDEN_DIM, &mut q_grad_w_d)?; + self.launch_reduce_axis0(&q_grad_b_per_batch_d, b_size, k_dqn, &mut q_grad_b_d)?; + + // ── Step 7: PPO backward (logits → grad_w/b/h_t) ───────────── + self.policy_head + .surrogate_backward_logits( + &pi_logits_d, + &log_pi_old_d, + &actions_d, + &advantages_d, + &self.isv_d, + b_size, + &mut pi_grad_logits_d, + ) + .context("policy_head.surrogate_backward_logits")?; + + self.policy_head + .backward_to_w_b_h( + h_t_borrow, + &pi_grad_logits_d, + b_size, + &mut pi_grad_w_per_batch_d, + &mut pi_grad_b_per_batch_d, + &mut pi_grad_h_t_d, + ) + .context("policy_head.backward_to_w_b_h")?; + + self.launch_reduce_axis0( + &pi_grad_w_per_batch_d, + b_size, + N_ACTIONS * HIDDEN_DIM, + &mut pi_grad_w_d, + )?; + self.launch_reduce_axis0(&pi_grad_b_per_batch_d, b_size, N_ACTIONS, &mut pi_grad_b_d)?; + + // ── Step 8: V backward ─────────────────────────────────────── + self.value_head + .backward( + h_t_borrow, + &v_pred_d, + &returns_d, + b_size, + &mut v_loss_per_batch_d, + &mut v_grad_w_per_batch_d, + &mut v_grad_b_per_batch_d, + &mut v_grad_h_t_d, + ) + .context("value_head.backward")?; + + self.launch_reduce_axis0(&v_grad_w_per_batch_d, b_size, HIDDEN_DIM, &mut v_grad_w_d)?; + self.launch_reduce_axis0(&v_grad_b_per_batch_d, b_size, 1, &mut v_grad_b_d)?; + // Reduce per-batch V loss (raw squared errors) → sum. Caller divides + // by B for the mean. + self.launch_reduce_axis0(&v_loss_per_batch_d, b_size, 1, &mut v_loss_sum_d)?; + let l_v_sum_host = read_scalar_d(&self.stream, &v_loss_sum_d)?; + let l_v_host = l_v_sum_host / (b_size as f32); + + // ── Step 9: Adam updates on each head's w and b ────────────── + self.dqn_w_adam + .step(&mut self.dqn_head.w_d, &q_grad_w_d) + .context("dqn_w_adam.step")?; + self.dqn_b_adam + .step(&mut self.dqn_head.b_d, &q_grad_b_d) + .context("dqn_b_adam.step")?; + self.policy_w_adam + .step(&mut self.policy_head.w_d, &pi_grad_w_d) + .context("policy_w_adam.step")?; + self.policy_b_adam + .step(&mut self.policy_head.b_d, &pi_grad_b_d) + .context("policy_b_adam.step")?; + self.value_w_adam + .step(&mut self.value_head.w_d, &v_grad_w_d) + .context("value_w_adam.step")?; + self.value_b_adam + .step(&mut self.value_head.b_d, &v_grad_b_d) + .context("value_b_adam.step")?; + + // ── Step 10: Encoder grad combine (DEFERRED to Phase E.3) ──── + // The per-head grad_h_t buffers (q_grad_h_t_d, pi_grad_h_t_d, + // v_grad_h_t_d) are populated and ready. Phase E.3 lands the + // PerceptionTrainer entry point that consumes them; for E.2 they + // are computed but not consumed by the encoder backward path. + // We touch them with a no-op pattern below so the unused-buffer + // linter doesn't fire and so the kernel handles stay reachable + // (the launcher is wired but unused-this-phase): + let _ = (&q_grad_h_t_d, &pi_grad_h_t_d, &v_grad_h_t_d); + + // ── Step 11: compose stats ─────────────────────────────────── + // BCE / aux losses are NOT read this phase — perception is driven + // separately by callers via its existing step_batched path. + let l_bce = 0.0_f32; + let l_aux = 0.0_f32; + let l_pi = self.last_pi_loss; + + // Normalise the Q CE by batch (matches l_pi / l_v conventions). + let l_q = l_q_host / (b_size as f32); + + // Combined: λ-weighted average across 5 heads (per + // pearl_loss_balance_controller, sum over 5 heads / 5). let l_total = (lambdas.bce * l_bce + lambdas.q * l_q + lambdas.pi * l_pi - + lambdas.v * l_v + + lambdas.v * l_v_host + lambdas.aux * l_aux) / 5.0; @@ -222,10 +547,162 @@ impl IntegratedTrainer { l_bce, l_q, l_pi, - l_v, + l_v: l_v_host, l_aux, l_total, lambdas, }) } + + /// Launch the reduce_axis0 kernel: + /// per_batch: [b_size × n_tail] → out: [n_tail] + fn launch_reduce_axis0( + &self, + per_batch: &CudaSlice, + b_size: usize, + n_tail: usize, + out: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(per_batch.len(), b_size * n_tail); + debug_assert_eq!(out.len(), n_tail); + + let n_batch_i = b_size as i32; + let n_tail_i = n_tail as i32; + let cfg = LaunchConfig { + grid_dim: ((n_tail as u32).div_ceil(32), 1, 1), + block_dim: (32, 8, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(per_batch) + .arg(&n_batch_i) + .arg(&n_tail_i) + .arg(out); + unsafe { + launch.launch(cfg).context("reduce_axis0 launch")?; + } + Ok(()) + } + + /// Phase E.2 helper: combine a per-head grad_h_t into the encoder's + /// accumulator slot via `grad_h_encoder[i] += λ × grad_h_head[i]`. + /// Wired but unused in Phase E.2; Phase E.3 consumes it from the + /// `step_synthetic` path once the encoder backward entry point lands. + pub fn launch_grad_h_accumulate( + &self, + grad_h_head: &CudaSlice, + lambda: f32, + grad_h_encoder: &mut CudaSlice, + ) -> Result<()> { + let n = grad_h_head.len(); + debug_assert_eq!(grad_h_encoder.len(), n); + let n_i = n as i32; + let cfg = LaunchConfig { + grid_dim: ((n as u32).div_ceil(256), 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.grad_h_accumulate_fn); + launch + .arg(grad_h_head) + .arg(&lambda) + .arg(&n_i) + .arg(grad_h_encoder); + unsafe { + launch.launch(cfg).context("grad_h_accumulate launch")?; + } + Ok(()) + } + +} + +// ── helpers ────────────────────────────────────────────────────────── + +fn upload_f32(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = + unsafe { MappedF32Buffer::new(n) }.map_err(|e| anyhow::anyhow!("upload_f32 staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream + .alloc_zeros::(n) + .context("upload_f32 alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + staging.dev_ptr, + nbytes, + stream.cu_stream(), + ) + .context("upload_f32 DtoD")?; + } + } + Ok(dst) +} + +fn upload_i32(stream: &Arc, host: &[u32]) -> Result> { + let n = host.len(); + let mut staging = + unsafe { MappedI32Buffer::new(n) }.map_err(|e| anyhow::anyhow!("upload_i32 staging: {e}"))?; + for (i, &v) in host.iter().enumerate() { + staging.host_slice_mut()[i] = v as i32; + } + let mut dst = stream + .alloc_zeros::(n) + .context("upload_i32 alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + staging.dev_ptr, + nbytes, + stream.cu_stream(), + ) + .context("upload_i32 DtoD")?; + } + } + Ok(dst) +} + +fn read_scalar_d(stream: &Arc, src: &CudaSlice) -> Result { + debug_assert!(src.len() >= 1); + let staging = unsafe { MappedF32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("read_scalar_d staging: {e}"))?; + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, + src_ptr, + std::mem::size_of::(), + stream.cu_stream(), + ) + .context("read_scalar_d DtoD")?; + } + stream.synchronize().context("read_scalar_d sync")?; + Ok(unsafe { std::ptr::read_volatile(staging.host_ptr) }) +} + +/// Build a Bellman target distribution from synthetic rewards. Phase E.2 +/// uses a deterministic single-atom-mass projection: each sample's reward +/// is clamped to `[V_MIN, V_MAX]` and the resulting value is placed +/// entirely on the closest atom (no fractional smear). This is a stand-in +/// for the proper categorical projection that lands in Phase E.3 once +/// the target network + γ-discounted bootstrap are wired in. +fn build_synthetic_bellman_target(rewards: &[f32]) -> Vec { + let b = rewards.len(); + let dz = (Q_V_MAX - Q_V_MIN) / (Q_N_ATOMS as f32 - 1.0); + let mut out = vec![0.0_f32; b * Q_N_ATOMS]; + for (i, &r) in rewards.iter().enumerate() { + let clamped = r.clamp(Q_V_MIN, Q_V_MAX); + // Closest-atom projection. + let idx_f = ((clamped - Q_V_MIN) / dz).round(); + let idx = (idx_f as isize).clamp(0, Q_N_ATOMS as isize - 1) as usize; + out[i * Q_N_ATOMS + idx] = 1.0; + } + out } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index d6634ae70..9eee44f74 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -3835,6 +3835,19 @@ impl PerceptionTrainer { Ok(&self.h_t_d) } + /// Phase E.2 RL hook: borrow the `h_t_d` buffer without re-running + /// the encoder. Returns the slice last written by `forward_encoder` + /// (slot K-1 of `h_new_per_k_d`). Stable until the next + /// `forward_encoder` / `step` overwrites the slot. + /// + /// Used by `IntegratedTrainer::step_synthetic` to dispatch the + /// downstream Q/π/V head kernels on the same encoder representation + /// without re-borrowing self mutably between the forward call and + /// the head dispatches. + pub fn h_t_view(&self) -> &CudaSlice { + &self.h_t_d + } + /// Kernel-dispatch portion of a training step — captured into the /// CUDA Graph on the second `step_batched` call. fn dispatch_train_step(