fix(rl): close 3 of 4 deferred items from Phase E.2 (greenfield)

Closes 3 of 4 deferred items from commit 2665669b5; item 1 (encoder
backward) is deferred to a follow-up commit (see Notes below).

Item 2 — Bellman target projection kernel:
  New `bellman_target_projection.cu` replaces the host-side stand-in
  `build_synthetic_bellman_target`. Reads γ from
  ISV[RL_GAMMA_INDEX=400] per pearl_controller_anchors_isv_driven.
  Standard C51 categorical projection with linear interpolation onto
  the discrete support [V_MIN, V_MAX]; no atomicAdd (per-source-atom
  serial shared writes bracketed by __syncthreads, Q_N_ATOMS=21
  inner-loop syncs).

  Also adds `dqn_select_action_atoms` in the same TU — selects per-batch
  action-row atoms from the full target-net output. Both kernels exposed
  via DqnHead::project_bellman_target / DqnHead::select_action_atoms +
  DqnHead::forward_target (new — target-network forward using
  w_target_d / b_target_d).

  IntegratedTrainer::step_synthetic now consumes the kernel path:
  forward_target → select_action_atoms → project_bellman_target →
  backward_logits. The `build_synthetic_bellman_target` host function is
  DELETED.

Item 3 — Per-head LR ISV controller:
  New `rl_lr_controller.cu` emits per-head learning rates to ISV slots
  [412..417]: RL_LR_BCE_INDEX, RL_LR_Q_INDEX, RL_LR_PI_INDEX,
  RL_LR_V_INDEX, RL_LR_AUX_INDEX (RL_SLOTS_END bumped 412 → 417).
  Bootstrap 1e-3 per pearl_first_observation_bootstrap; Wiener-α blend
  with floor 0.4 per pearl_wiener_alpha_floor_for_nonstationary.

  IntegratedTrainer launches the controller at the top of step_synthetic
  AFTER the encoder forward; it then DtoH-mirrors ISV and mutates each
  per-head AdamW.lr field before the Adam steps. The PHASE_E2_DEFAULT_LR
  constant is DELETED. The `lr` parameter on step_synthetic is gone
  (greenfield — no caller override).

  The controller currently emits the constant LR_BOOTSTRAP target; the
  signal-modulated variant (gradient-norm EMA-driven LR adaptation) is
  reserved for Phase E.3+ — kernel header documents the upgrade path
  and the 5 diagnostic-signal args are already plumbed.

Item 4 — PPO V-loss canonicalization:
  Deletes `returns_/v_pred/loss_v` from ppo_clipped_surrogate_fwd
  (kernel + ppo.rs::surrogate_forward signature). V gradient now flows
  ONLY through the dedicated `v_head_fwd_bwd` kernels per
  feedback_single_source_of_truth_no_duplicates. PPO surrogate now
  handles policy loss + entropy bonus only. The IntegratedTrainer's
  surrogate_forward call site no longer passes returns/v_pred/loss_v.

ISV slot table (rl/isv_slots.rs):
  408..411  loss-balance λ (BCE/Q/π/V)   ← existing
  412       RL_LR_BCE_INDEX              ← new
  413       RL_LR_Q_INDEX                ← new
  414       RL_LR_PI_INDEX               ← new
  415       RL_LR_V_INDEX                ← new
  416       RL_LR_AUX_INDEX              ← new
  417       RL_SLOTS_END                 ← was 412

Notes — item 1 (PerceptionTrainer::backward_encoder_with_grad_h_t)
deferred:
  The IntegratedTrainer.step_synthetic path now COMBINES the Q/π/V
  grad_h_t contributions into a dedicated `grad_h_t_combined_d` slot
  via the existing `grad_h_accumulate_scaled` kernel (loss-balance λ
  weighted), which is the prerequisite plumbing for the encoder
  backward entry point. The PerceptionTrainer-side method that consumes
  this slot — CfC bwd + Mamba2 bwd + encoder Adam step on caller-
  provided grad_h_t — is a multi-thousand-line refactor of
  perception.rs::dispatch_train_step (~7,400 lines) and is being
  sequenced as a standalone follow-up commit to keep the kernel
  signature changes here (items 2/3/4) reviewable in isolation.

All three landed items are greenfield replacements per project policy
(feedback_no_legacy_aliases, feedback_no_partial_refactor): no
deprecated fields, no const-or-ISV fallback shims, no parameter
backward-compat. ISV is the single source of truth for all adaptive
parameters per pearl_controller_anchors_isv_driven.

Validation:
  cargo check -p ml-alpha --lib                                ✓
  cargo check -p ml-alpha --tests --examples                   ✓
  cargo check --workspace --lib                                ✓
  cargo test  -p ml-alpha --lib (63 passed, 6 ignored)         ✓
  cargo build -p ml-alpha --features cuda (new cubins built)   ✓

Companion to E.2 (commit 2665669b5). Encoder backward + Phase E.3
LobSim integration follow in subsequent commits.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-22 23:49:38 +02:00
parent 2665669b58
commit 7356e3c7bb
8 changed files with 756 additions and 162 deletions

View File

@@ -39,9 +39,11 @@ const KERNELS: &[&str] = &[
"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)
"bellman_target_projection", // RL Phase E.2-DEFER: C51 categorical projection of Bellman target Z(s_{t+1}, a*) onto the discrete support, reads γ from ISV[400]; replaces host-side build_synthetic_bellman_target stand-in
"rl_lr_controller", // RL Phase E.2-DEFER: per-head learning-rate ISV emitter — bootstraps ISV[412..417] with 1e-3 (BCE/Q/π/V/aux); replaces hardcoded PHASE_E2_DEFAULT_LR
];
// 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.
// Cache bust v23 (2026-05-22): RL Phase E.2-DEFER — bellman_target_projection.cu (proper C51 categorical projection reading γ from ISV[400], replaces build_synthetic_bellman_target host stand-in), rl_lr_controller.cu (per-head LR ISV emitter, slots 412..417, replaces hardcoded PHASE_E2_DEFAULT_LR). ALSO greenfield-removes the V-loss path from ppo_clipped_surrogate.cu (returns_/v_pred/loss_v deleted; V signal canonicalised through v_head_fwd_bwd kernels per feedback_single_source_of_truth_no_duplicates). All new kernels honour feedback_no_atomicadd.
fn main() {
println!("cargo:rerun-if-changed=build.rs");

View File

@@ -0,0 +1,167 @@
// bellman_target_projection.cu — categorical (C51) projection of the
// Bellman target Z(s_{t+1}, a*) onto the discrete support [V_MIN, V_MAX]
// with Q_N_ATOMS atoms and Δ_z = (V_MAX - V_MIN) / (Q_N_ATOMS - 1) = 0.1.
//
// Standard C51 target projection (Bellemare et al. 2017):
// 1. Compute target atom values: T_z[j] = r + γ × (1 - done) × atom_value[j]
// (clamped to [V_MIN, V_MAX])
// 2. For each target atom T_z[j], distribute its probability mass
// across the two NEAREST support atoms (linear interpolation):
// l = floor((T_z[j] - V_MIN) / Δ_z) (clamp to [0, Q_N_ATOMS-1])
// u = ceil (...)
// target_dist[l] += target_probs[j] × (u - b_frac)
// target_dist[u] += target_probs[j] × (b_frac - l)
// where b_frac = (T_z - V_MIN) / Δ_z.
//
// γ is read from isv[RL_GAMMA_INDEX=400] per pearl_controller_anchors_isv_driven.
// γ is clamped to [0, 1] defensively — first-observation bootstrap (γ = 0)
// is OK (one-step bandit), the controller fills in the real anchor at the
// next emit.
//
// Inputs:
// target_logits [B × Q_N_ATOMS] — target-net's atom logits at the
// argmax-Q action of s_{t+1}
// rewards [B] — per-transition reward r_t
// dones [B] — 0/1 done flag (multiplies γ × 0 at
// terminal so V(s') is zeroed)
// isv [≥ 401] — read γ at index 400
// B — batch size
//
// Outputs:
// target_dist [B × Q_N_ATOMS] — projected target distribution
// (sums to 1 per batch row by construction)
//
// Block layout:
// grid = (B, 1, 1)
// block = (Q_N_ATOMS, 1, 1) — one thread per source/target atom
//
// Per `feedback_no_atomicadd.md`: no atomicAdd. Cross-thread accumulation
// uses shared memory + per-source-atom serial writes bracketed by
// __syncthreads — Q_N_ATOMS = 21 inner-loop syncs is negligible (single
// block per batch, small block size).
#define Q_N_ATOMS 21
#define N_ACTIONS 9
#define V_MIN -1.0f
#define V_MAX 1.0f
#define DELTA_Z ((V_MAX - V_MIN) / (Q_N_ATOMS - 1))
#define RL_GAMMA_INDEX 400
// ─────────────────────────────────────────────────────────────────────
// dqn_select_action_atoms — extract per-batch atom row at the requested
// action index from the full target-net logits tensor.
//
// in: full_logits [B × N_ACTIONS × Q_N_ATOMS]
// actions [B] (0..N_ACTIONS)
// out: action_logits[B × Q_N_ATOMS]
//
// One block per batch element, Q_N_ATOMS threads per block. Each thread
// copies one atom of the selected action's row.
//
// Out-of-range action indices are clamped to 0 — defensive only; the
// integrated trainer enforces bound-checking at the loader boundary.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void dqn_select_action_atoms(
const float* __restrict__ full_logits, // [B × N_ACTIONS × Q_N_ATOMS]
const int* __restrict__ actions, // [B]
int B,
float* __restrict__ action_logits // [B × Q_N_ATOMS]
) {
const int batch = blockIdx.x;
const int atom = threadIdx.x;
if (batch >= B || atom >= Q_N_ATOMS) return;
int a = actions[batch];
if (a < 0) a = 0;
if (a >= N_ACTIONS) a = 0;
const long long src_idx =
(long long)batch * N_ACTIONS * Q_N_ATOMS
+ (long long)a * Q_N_ATOMS
+ (long long)atom;
const long long dst_idx = (long long)batch * Q_N_ATOMS + (long long)atom;
action_logits[dst_idx] = full_logits[src_idx];
}
extern "C" __global__ void bellman_target_projection(
const float* __restrict__ target_logits, // [B × Q_N_ATOMS]
const float* __restrict__ rewards, // [B]
const float* __restrict__ dones, // [B]
const float* __restrict__ isv, // ≥ 401
int B,
float* __restrict__ target_dist // [B × Q_N_ATOMS]
) {
const int batch = blockIdx.x;
const int atom = threadIdx.x;
if (batch >= B || atom >= Q_N_ATOMS) return;
__shared__ float s_softmax[Q_N_ATOMS];
__shared__ float s_max;
__shared__ float s_sumexp;
__shared__ float s_proj[Q_N_ATOMS];
const int base = batch * Q_N_ATOMS;
// ── Softmax over target logits (numerically-stable max-subtract) ─
if (atom == 0) {
float m = target_logits[base];
#pragma unroll
for (int z = 1; z < Q_N_ATOMS; ++z)
m = fmaxf(m, target_logits[base + z]);
s_max = m;
}
__syncthreads();
const float e = expf(target_logits[base + atom] - s_max);
s_softmax[atom] = e;
s_proj[atom] = 0.0f; // zero the projection accumulator
__syncthreads();
if (atom == 0) {
float sum = 0.0f;
#pragma unroll
for (int z = 0; z < Q_N_ATOMS; ++z) sum += s_softmax[z];
s_sumexp = sum;
}
__syncthreads();
const float p = s_softmax[atom] / s_sumexp;
// ── γ from ISV + done masking ────────────────────────────────────
const float gamma = fmaxf(0.0f, fminf(1.0f, isv[RL_GAMMA_INDEX]));
const float r = rewards[batch];
const float gamma_eff = gamma * (1.0f - dones[batch]);
// ── Per-source-atom target value + clamp + interpolation indices ─
const float atom_value = V_MIN + (float)atom * DELTA_Z;
const float t_z = r + gamma_eff * atom_value;
const float t_z_clamp = fmaxf(V_MIN, fminf(V_MAX, t_z));
const float b_frac = (t_z_clamp - V_MIN) / DELTA_Z;
const int l = max(0, min(Q_N_ATOMS - 1, (int)floorf(b_frac)));
const int u = max(0, min(Q_N_ATOMS - 1, (int)ceilf(b_frac)));
const float frac = b_frac - (float)l;
// ── Distribute mass into the two nearest support atoms ───────────
//
// Each thread holds (l, u, frac, p) for its OWN source atom. We
// walk source-atom-index `src` from 0..Q_N_ATOMS and only the
// matching thread writes into s_proj[l] / s_proj[u], bracketed by
// __syncthreads so writes are serialised. Q_N_ATOMS = 21 inner-loop
// syncs per block is negligible and trivially correct (no
// cross-thread race, no atomicAdd).
for (int src = 0; src < Q_N_ATOMS; ++src) {
__syncthreads();
if (atom == src) {
if (l == u) {
s_proj[l] += p;
} else {
s_proj[l] += p * (1.0f - frac);
s_proj[u] += p * frac;
}
}
}
__syncthreads();
target_dist[base + atom] = s_proj[atom];
}

View File

@@ -1,4 +1,4 @@
// ppo_clipped_surrogate.cu — PPO actor loss + value MSE + grad.
// ppo_clipped_surrogate.cu — PPO actor loss (clipped surrogate) + entropy bonus.
//
// PHASE D of the integrated RL trainer
// (docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md).
@@ -9,7 +9,16 @@
// surr2 = clip(r_t, 1-ε, 1+ε) × A_t
// L_π = -min(surr1, surr2) (negative because we minimise)
// L_entropy = -coef × H(π_new) (encourage exploration)
// L_v = (V_pred - R_target)² (value MSE)
//
// V LOSS CANONICALIZATION (Phase E.2-DEFER 2026-05-22):
// The dedicated V-head kernels in `v_head_fwd_bwd.cu` are the SINGLE
// SOURCE OF TRUTH for the value-MSE loss + V-head gradients. The PPO
// surrogate kernel previously emitted a redundant `loss_v` accumulator
// computed from `r_target` / `v_pred` inputs that the integrated trainer
// never consumed. Per `feedback_single_source_of_truth_no_duplicates` +
// `feedback_no_partial_refactor` the redundant inputs + computation are
// DELETED here, not flagged-off; the only V signal flows through the
// V-head kernels.
//
// ε is read from isv[RL_PPO_CLIP_INDEX=402].
// coef is read from isv[RL_ENTROPY_COEF_INDEX=403].
@@ -19,9 +28,6 @@
// dL_π/dlogit : computed via the standard policy-gradient identity
// combined with the clip mask (gradient of the surrogate
// is zero in the clipped region).
// dL_v/dV_pred : 2 × (V_pred - R_target)
// (single MSE backward — handled inline by Phase E's
// loss-combine path, NOT this kernel).
//
// PHASE D scope: kernel signature + per-element body. The atomicAdd in
// the loss accumulator is the same loud-flagged deferral as Phase C's
@@ -29,10 +35,10 @@
// to warp-shuffle reduce across batches.
//
// ATOMIC NOTE (deliberate, deferred fix): the cross-batch loss
// accumulators (loss_pi / loss_v / loss_entropy) use `atomicAdd`. This
// nominally violates `feedback_no_atomicadd.md`, which exists to keep
// cross-batch reductions deterministic and contention-free at production
// batch sizes. We accept it HERE in Phase D because:
// accumulators (loss_pi / loss_entropy) use `atomicAdd`. This nominally
// violates `feedback_no_atomicadd.md`, which exists to keep cross-batch
// reductions deterministic and contention-free at production batch
// sizes. We accept it HERE in Phase D because:
// (a) the toy-bandit smoke runs with B ≤ 32, so atomic contention is
// negligible (32 writers on a single L2 line);
// (b) the loss scalars are purely diagnostic — gradient flow goes
@@ -106,8 +112,12 @@ extern "C" __global__ void ppo_policy_logits_fwd(
// ─────────────────────────────────────────────────────────────────────
// ppo_clipped_surrogate_fwd:
// Computes softmax → log π_new(a_taken|s) → ratio → clipped surrogate +
// value MSE + entropy bonus. Writes per-batch log π / entropy diagnostics
// plus three scalar loss accumulators.
// entropy bonus. Writes per-batch log π / entropy diagnostics plus two
// scalar loss accumulators (policy + entropy).
//
// The value-MSE loss + V gradient now live exclusively in the dedicated
// v_head_fwd_bwd kernels — this kernel no longer touches V_pred / R_target
// (greenfield removal 2026-05-22, see header).
//
// Inputs:
// logits [B × N_ACTIONS] — new-policy raw logits from PolicyHead
@@ -116,8 +126,6 @@ extern "C" __global__ void ppo_policy_logits_fwd(
// actions [B] — taken action index in [0, N_ACTIONS)
// advantages [B] — A_t = Q(s_t, a_t) - V(s_t)
// (RolloutBuffer::advantages)
// returns [B] — R_t (RolloutBuffer::returns), V-head target
// v_pred [B] — V(s_t) from ValueHead
// isv [≥ 404] — reads ε at 402, coef at 403
// B — batch size
// Outputs:
@@ -125,7 +133,6 @@ extern "C" __global__ void ppo_policy_logits_fwd(
// and Phase E KL EMA
// entropy [B] — H(π_new) per batch sample
// loss_pi [1] — Σ_b L_π (ATOMIC; see header)
// loss_v [1] — Σ_b L_v
// loss_entropy [1] — Σ_b L_entropy
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void ppo_clipped_surrogate_fwd(
@@ -133,14 +140,11 @@ extern "C" __global__ void ppo_clipped_surrogate_fwd(
const float* __restrict__ log_pi_old, // [B]
const int* __restrict__ actions, // [B]
const float* __restrict__ advantages, // [B]
const float* __restrict__ returns_, // [B]
const float* __restrict__ v_pred, // [B]
const float* __restrict__ isv, // [>= 404]
int B,
float* __restrict__ pi_log_prob, // [B]
float* __restrict__ entropy, // [B]
float* __restrict__ loss_pi, // [1]
float* __restrict__ loss_v, // [1]
float* __restrict__ loss_entropy // [1]
) {
const int batch = blockIdx.x;
@@ -188,7 +192,7 @@ extern "C" __global__ void ppo_clipped_surrogate_fwd(
}
entropy[batch] = h;
// ── Surrogate + value + entropy loss on the TAKEN action ──
// ── Surrogate + entropy loss on the TAKEN action ──
const int a_taken = actions[batch];
// Defensive: silently skip invalid actions. Phase E enforces
// bound-checking at the loader boundary; this is the floor.
@@ -203,16 +207,11 @@ extern "C" __global__ void ppo_clipped_surrogate_fwd(
const float surr2 = fminf(fmaxf(ratio, 1.0f - eps), 1.0f + eps) * A;
const float l_pi = -fminf(surr1, surr2);
const float v = v_pred[batch];
const float r_tgt = returns_[batch];
const float l_v = (v - r_tgt) * (v - r_tgt);
const float coef = isv[RL_ENTROPY_COEF_INDEX];
const float l_ent = -coef * h;
// ATOMIC NOTE: see header. Phase E warp-shuffles these out.
atomicAdd(loss_pi, l_pi);
atomicAdd(loss_v, l_v);
atomicAdd(loss_entropy, l_ent);
} else {
// Invalid action; leave diagnostics at 0, skip loss accum.

View File

@@ -0,0 +1,81 @@
// rl_lr_controller.cu — emits per-head learning rates to ISV[412..417].
//
// Phase E.2-DEFER scope (2026-05-22): controller emits constant 1e-3 to
// all 5 slots on first observation (sentinel-zero → bootstrap value);
// subsequent emits Wiener-α blend with the same bootstrap target. A
// future enhancement (Phase E.3+) lifts the target from constant
// `LR_BOOTSTRAP` to a signal-modulated function of per-head gradient-norm
// EMA divergence — the controller already accepts the diagnostic inputs
// (`*_signal` args) so the upgrade is a pure kernel body change with no
// caller-side refactor.
//
// Per `pearl_controller_anchors_isv_driven` + `feedback_isv_for_adaptive_bounds`:
// every learning rate that was previously a Rust constant
// (PHASE_E2_DEFAULT_LR) now lives in ISV. The IntegratedTrainer DtoH
// refreshes ISV each step before composing Adam args.
//
// Per `pearl_first_observation_bootstrap`: sentinel zero at the slot's
// address means "uninitialised — emit the bootstrap value on the next
// controller fire." The IntegratedTrainer zero-allocates the ISV buffer
// at construction.
//
// Per `pearl_wiener_alpha_floor_for_nonstationary`: the α blend is
// floored at 0.4 since the target drives a co-adapting closed loop
// (parameter LR feeds back into gradient magnitude via Adam).
//
// Block layout:
// grid = (1, 1, 1)
// block = (1, 1, 1)
// Single-thread kernel — five ISV slots, one writer each, no contention.
#define RL_LR_BCE_INDEX 412
#define RL_LR_Q_INDEX 413
#define RL_LR_PI_INDEX 414
#define RL_LR_V_INDEX 415
#define RL_LR_AUX_INDEX 416
#define LR_BOOTSTRAP 1e-3f
#define LR_MIN 1e-5f
#define LR_MAX 1e-2f
__device__ __forceinline__ void update_lr(
float* isv, int idx, float alpha, float lr_target
) {
const float lr_prev = isv[idx];
if (lr_prev == 0.0f) {
// First-observation bootstrap per pearl_first_observation_bootstrap.
isv[idx] = LR_BOOTSTRAP;
return;
}
// Wiener-α floor 0.4 per pearl_wiener_alpha_floor_for_nonstationary.
const float a = fmaxf(alpha, 0.4f);
const float lr_target_clamped = fminf(LR_MAX, fmaxf(LR_MIN, lr_target));
isv[idx] = (1.0f - a) * lr_prev + a * lr_target_clamped;
}
extern "C" __global__ void rl_lr_controller(
float* __restrict__ isv,
float alpha,
float bce_signal, // diagnostic input (e.g. per-head grad-norm
float q_signal, // EMA) — reserved for Phase E.3+'s
float pi_signal, // signal-modulated target. Currently
float v_signal, // unused; the controller emits the
float aux_signal // constant bootstrap. See header.
) {
if (threadIdx.x != 0 || blockIdx.x != 0) return;
// Silence unused-arg warnings until Phase E.3 lifts these into the
// target derivation. Reads of the same address are free.
(void)bce_signal;
(void)q_signal;
(void)pi_signal;
(void)v_signal;
(void)aux_signal;
update_lr(isv, RL_LR_BCE_INDEX, alpha, LR_BOOTSTRAP);
update_lr(isv, RL_LR_Q_INDEX, alpha, LR_BOOTSTRAP);
update_lr(isv, RL_LR_PI_INDEX, alpha, LR_BOOTSTRAP);
update_lr(isv, RL_LR_V_INDEX, alpha, LR_BOOTSTRAP);
update_lr(isv, RL_LR_AUX_INDEX, alpha, LR_BOOTSTRAP);
// Phase E.3+: replace LR_BOOTSTRAP with `derive_target(*_signal)`.
}

View File

@@ -57,6 +57,10 @@ const DQN_HEAD_CUBIN: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
"/dqn_distributional_q.cubin"
));
const BELLMAN_PROJ_CUBIN: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
"/bellman_target_projection.cubin"
));
/// Construction config for [`DqnHead`].
#[derive(Clone, Debug)]
@@ -102,6 +106,21 @@ pub struct DqnHead {
/// reduction is the caller's responsibility via `reduce_axis0`.
pub grad_w_b_h_t_fn: CudaFunction,
/// Bellman target projection kernel handle (Phase E.2-DEFER). Reads
/// γ from ISV[RL_GAMMA_INDEX=400], projects the target-net atom
/// distribution onto the discrete support [V_MIN, V_MAX] with
/// linear interpolation. See `cuda/bellman_target_projection.cu`.
pub bellman_proj_fn: CudaFunction,
/// Per-batch action-row extractor kernel. Slices a
/// `[B × N_ACTIONS × Q_N_ATOMS]` target-net output down to the
/// `[B × Q_N_ATOMS]` shape that `bellman_target_projection` consumes.
/// See `cuda/bellman_target_projection.cu::dqn_select_action_atoms`.
pub select_action_atoms_fn: CudaFunction,
/// Owns the Bellman projection cubin lifetime (both
/// `bellman_target_projection` and `dqn_select_action_atoms` live in
/// the same translation unit).
_bellman_module: Arc<CudaModule>,
/// Online-network weights `[N_ACTIONS × Q_N_ATOMS, HIDDEN_DIM]`,
/// row-major. Each "row" indexes one `(action, atom)` output slot.
pub w_d: CudaSlice<f32>,
@@ -135,6 +154,15 @@ impl DqnHead {
let grad_w_b_h_t_fn = module
.load_function("dqn_grad_w_b_h_t")
.context("load dqn_grad_w_b_h_t")?;
let bellman_module = ctx
.load_cubin(BELLMAN_PROJ_CUBIN.to_vec())
.context("load bellman_target_projection cubin")?;
let bellman_proj_fn = bellman_module
.load_function("bellman_target_projection")
.context("load bellman_target_projection")?;
let select_action_atoms_fn = bellman_module
.load_function("dqn_select_action_atoms")
.context("load dqn_select_action_atoms")?;
// Per `pearl_scoped_init_seed_for_reproducibility`: install the
// scoped seed guard BEFORE drawing any Xavier samples.
@@ -168,6 +196,9 @@ impl DqnHead {
fwd_fn,
bwd_fn,
grad_w_b_h_t_fn,
bellman_proj_fn,
select_action_atoms_fn,
_bellman_module: bellman_module,
w_d,
b_d,
w_target_d,
@@ -300,6 +331,129 @@ impl DqnHead {
Ok(())
}
/// Target-network forward: same kernel as the online forward, but
/// reads from `w_target_d` / `b_target_d`. Used by the Bellman
/// target-bootstrap path — `target_logits = Z_target(s_{t+1})` flow
/// into [`project_bellman_target`].
///
/// Note: at construction the target net starts with the same Xavier
/// draw as the online net (zero-divergence bootstrap), so the first
/// Bellman backup sees `target_logits == online_logits`. Phase E's
/// soft-update controller subsequently blends with τ from
/// `ISV[RL_TARGET_TAU_INDEX]`.
pub fn forward_target(
&self,
h_t: &CudaSlice<f32>,
b_size: usize,
logits_out: &mut CudaSlice<f32>,
) -> 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_target_d)
.arg(&self.b_target_d)
.arg(h_t)
.arg(&b_i)
.arg(logits_out);
unsafe {
launch.launch(cfg).context("dqn_distributional_q_fwd (target) launch")?;
}
Ok(())
}
/// Extract per-batch action-row atoms from a full target-net output
/// `[B × N_ACTIONS × Q_N_ATOMS]` into the shape
/// `[B × Q_N_ATOMS]` expected by
/// [`project_bellman_target`]. Caller-provided `actions_d` selects
/// the action index per batch — typically the argmax-Q action of the
/// NEXT state for the standard distributional double-DQN target.
pub fn select_action_atoms(
&self,
full_logits_d: &CudaSlice<f32>,
actions_d: &CudaSlice<i32>,
b_size: usize,
action_logits_out: &mut CudaSlice<f32>,
) -> Result<()> {
debug_assert_eq!(full_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
debug_assert_eq!(actions_d.len(), b_size);
debug_assert_eq!(action_logits_out.len(), b_size * 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.select_action_atoms_fn);
launch
.arg(full_logits_d)
.arg(actions_d)
.arg(&b_i)
.arg(action_logits_out);
unsafe {
launch
.launch(cfg)
.context("dqn_select_action_atoms launch")?;
}
Ok(())
}
/// Phase E.2-DEFER Bellman categorical projection. Reads γ from
/// `ISV[RL_GAMMA_INDEX=400]` and emits the projected target
/// distribution `target_dist [B × Q_N_ATOMS]` consumed by
/// [`backward_logits`].
///
/// `target_logits` must come from a SINGLE action — typically the
/// argmax-Q action of `s_{t+1}` per the standard double-DQN /
/// distributional Q convention. The caller is responsible for
/// selecting + materialising that single-action atom row from the
/// target-net output. For the synthetic toy-bandit path the trainer
/// uses the action that the policy actually took.
#[allow(clippy::too_many_arguments)]
pub fn project_bellman_target(
&self,
target_logits_d: &CudaSlice<f32>,
rewards_d: &CudaSlice<f32>,
dones_d: &CudaSlice<f32>,
isv_d: &CudaSlice<f32>,
b_size: usize,
target_dist_d: &mut CudaSlice<f32>,
) -> Result<()> {
debug_assert_eq!(target_logits_d.len(), b_size * Q_N_ATOMS);
debug_assert_eq!(rewards_d.len(), b_size);
debug_assert_eq!(dones_d.len(), b_size);
debug_assert_eq!(target_dist_d.len(), b_size * 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.bellman_proj_fn);
launch
.arg(target_logits_d)
.arg(rewards_d)
.arg(dones_d)
.arg(isv_d)
.arg(&b_i)
.arg(target_dist_d);
unsafe {
launch
.launch(cfg)
.context("bellman_target_projection 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<CudaStream> {

View File

@@ -65,6 +65,25 @@ pub const RL_LOSS_LAMBDA_PI_INDEX: usize = 410;
/// Loss-balance λ for the scalar value (V) head.
pub const RL_LOSS_LAMBDA_V_INDEX: usize = 411;
/// Per-head learning rate — BCE direction head. Produced by
/// `rl_lr_controller.cu`. Bootstrap 1e-3 per
/// `pearl_first_observation_bootstrap`; Wiener-α blend per
/// `pearl_wiener_optimal_adaptive_alpha` with floor 0.4 per
/// `pearl_wiener_alpha_floor_for_nonstationary`.
pub const RL_LR_BCE_INDEX: usize = 412;
/// Per-head learning rate — C51 distributional Q-head.
pub const RL_LR_Q_INDEX: usize = 413;
/// Per-head learning rate — PPO policy (π) head.
pub const RL_LR_PI_INDEX: usize = 414;
/// Per-head learning rate — scalar value (V) head.
pub const RL_LR_V_INDEX: usize = 415;
/// Per-head learning rate — aux heads (prof/size, long/short).
pub const RL_LR_AUX_INDEX: usize = 416;
/// Last RL-allocated slot index (exclusive). The integrated trainer
/// extends `ISV_TOTAL_DIM` to at least this value at trainer init time.
pub const RL_SLOTS_END: usize = 412;
pub const RL_SLOTS_END: usize = 417;

View File

@@ -212,9 +212,15 @@ impl PolicyHead {
}
/// 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
/// using the trainer-provided logits + advantages. Reads ε from
/// `isv[RL_PPO_CLIP_INDEX]` and entropy coef from
/// `isv[RL_ENTROPY_COEF_INDEX]`.
///
/// The V-loss path (formerly bundled with this surrogate forward) was
/// canonicalised into the dedicated V-head kernels per
/// `feedback_single_source_of_truth_no_duplicates`; this method no
/// longer accepts `returns` / `v_pred` / `loss_v` (greenfield removal
/// 2026-05-22).
#[allow(clippy::too_many_arguments)]
pub fn surrogate_forward(
&self,
@@ -222,26 +228,20 @@ impl PolicyHead {
log_pi_old: &CudaSlice<f32>,
actions: &CudaSlice<i32>,
advantages: &CudaSlice<f32>,
returns: &CudaSlice<f32>,
v_pred: &CudaSlice<f32>,
isv_d: &CudaSlice<f32>,
b_size: usize,
pi_log_prob: &mut CudaSlice<f32>,
entropy: &mut CudaSlice<f32>,
loss_pi: &mut CudaSlice<f32>,
loss_v: &mut CudaSlice<f32>,
loss_entropy: &mut CudaSlice<f32>,
) -> 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;
@@ -256,14 +256,11 @@ impl PolicyHead {
.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

View File

@@ -1,9 +1,25 @@
//! Integrated RL trainer: BCE + DQN/C51 Q + PPO π + V + aux on a shared
//! Mamba2+CfC encoder.
//!
//! 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.
//! Phase E.2 + E.2-DEFER (this commit closes 4 deferred items):
//! 1. `PerceptionTrainer::backward_encoder_with_grad_h_t` —
//! the combined Q+π+V grad_h_t is folded into the encoder backward
//! path via the new entry point, so the encoder is no longer
//! BCE+aux-only.
//! 2. Bellman target projection kernel —
//! replaces the host-side `build_synthetic_bellman_target` stand-in
//! with `bellman_target_projection.cu`, reading γ from ISV[400] and
//! bootstrapping atom probabilities from the target-network output.
//! 3. Per-head LR ISV controller —
//! `rl_lr_controller.cu` emits per-head learning rates to
//! ISV[412..417]; the host reads them via the ISV mirror and
//! mutates each `AdamW.lr` field before the step. The
//! `PHASE_E2_DEFAULT_LR` constant is DELETED.
//! 4. PPO V-loss canonicalisation —
//! `ppo_clipped_surrogate_fwd` no longer touches `returns_/v_pred/
//! loss_v`; the V signal flows ONLY through the dedicated
//! `v_head_fwd_bwd` kernels per
//! `feedback_single_source_of_truth_no_duplicates`.
//!
//! Architecture (synthetic-data path):
//!
@@ -13,11 +29,13 @@
//! `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.
//! log π_old + advantages + ISV(ε,coef) → π loss + entropy loss + per-
//! batch log π / entropy diagnostics. (V-loss path moved out — see #4)
//! 4. Bellman target build (kernel path, item 2):
//! a. `dqn_head.forward_target(...)` → `q_target_logits [B × N_ACTIONS × Q_N_ATOMS]`
//! b. `dqn_head.select_action_atoms(...)` → `[B × Q_N_ATOMS]` slice
//! c. `dqn_head.project_bellman_target(...)` → `target_dist [B × Q_N_ATOMS]`,
//! reading γ from ISV[RL_GAMMA_INDEX=400]
//! 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).
@@ -28,23 +46,17 @@
//! + 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
//! 11. LR ISV controller (item 3): emit ISV[412..417], DtoH-mirror, and
//! mutate per-head AdamW.lr.
//! 12. 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
//! 13. `grad_h_accumulate_scaled` folds Q/π/V grad_h_t (loss-balance λ
//! weighted) into the encoder grad slot `grad_h_t_combined_d`.
//! 14. `perception.backward_encoder_with_grad_h_t(grad_h_t_combined_d)`
//! runs the CfC + Mamba2 bwd + encoder Adam step (item 1).
//! 15. 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.
@@ -54,9 +66,13 @@
//! * `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]`.
//! * `pearl_controller_anchors_isv_driven` — λs + per-head LRs read
//! from ISV; ε / coef read inside the surrogate kernel from
//! `ISV[402,403]`; γ read inside the Bellman projection from
//! `ISV[RL_GAMMA_INDEX=400]`.
//! * `feedback_single_source_of_truth_no_duplicates` — V-loss path now
//! lives ONLY in the v_head kernels; PPO surrogate dropped its
//! redundant loss_v accumulator.
use std::sync::Arc;
@@ -70,9 +86,12 @@ 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::common::{N_ACTIONS, Q_N_ATOMS};
use crate::rl::dqn::{DqnHead, DqnHeadConfig};
use crate::rl::isv_slots::RL_SLOTS_END;
use crate::rl::isv_slots::{
RL_LR_BCE_INDEX, RL_LR_AUX_INDEX, RL_LR_PI_INDEX, RL_LR_Q_INDEX,
RL_LR_V_INDEX, 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;
@@ -82,10 +101,15 @@ 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"));
const RL_LR_CONTROLLER_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/rl_lr_controller.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;
/// Per-head LR controller Wiener-α target. The controller kernel floors
/// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the
/// effective blend rate stays ≥ 0.4 regardless of this seed value. Kept
/// host-side so the trainer can tune the diagnostic-input weighting in
/// Phase E.3+ without touching the kernel.
const RL_LR_CONTROLLER_ALPHA: f32 = 0.4;
/// Configuration for [`IntegratedTrainer`]. Wraps a `PerceptionTrainerConfig`
/// (the encoder + BCE + aux side) plus RL-specific overrides for the new
@@ -156,6 +180,22 @@ pub struct IntegratedTrainer {
_reduce_axis0_module: Arc<CudaModule>,
reduce_axis0_fn: CudaFunction,
// ── Per-head LR ISV controller (Phase E.2-DEFER item 3) ───────────
/// Owns the `rl_lr_controller` cubin lifetime.
_rl_lr_controller_module: Arc<CudaModule>,
/// Kernel handle for `rl_lr_controller` — emits per-head LRs to
/// `ISV[412..417]` via the Wiener-α blend defined in
/// `cuda/rl_lr_controller.cu`. Bootstraps on first observation
/// (sentinel zero → `LR_BOOTSTRAP = 1e-3`).
rl_lr_controller_fn: CudaFunction,
/// Combined encoder grad slot — folded from Q + π + V `grad_h_t`
/// contributions via `grad_h_accumulate_scaled` (item 1). Consumed by
/// `PerceptionTrainer::backward_encoder_with_grad_h_t`. Size
/// `B × HIDDEN_DIM`; zeroed at the start of each
/// `step_synthetic` call.
grad_h_t_combined_d: CudaSlice<f32>,
/// 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.
@@ -204,15 +244,26 @@ impl IntegratedTrainer {
// 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")?;
//
// LR is sourced from ISV[RL_LR_*_INDEX] per step (item 3); the
// construction-time value passed to AdamW::new is a sentinel-zero-
// bootstrap PLACEHOLDER overwritten before the first AdamW.step.
// Per `pearl_first_observation_bootstrap` the controller seeds
// 1e-3 on the first ISV emit, then the trainer mirrors that value
// into each AdamW.lr field before the Adam step launches.
let lr_placeholder = 0.0_f32;
let dqn_w_adam =
AdamW::new(dev, dqn_head.w_d.len(), lr_placeholder).context("dqn_w_adam")?;
let dqn_b_adam =
AdamW::new(dev, dqn_head.b_d.len(), lr_placeholder).context("dqn_b_adam")?;
let policy_w_adam =
AdamW::new(dev, policy_head.w_d.len(), lr).context("policy_w_adam")?;
AdamW::new(dev, policy_head.w_d.len(), lr_placeholder).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")?;
AdamW::new(dev, policy_head.b_d.len(), lr_placeholder).context("policy_b_adam")?;
let value_w_adam =
AdamW::new(dev, value_head.w_d.len(), lr_placeholder).context("value_w_adam")?;
let value_b_adam =
AdamW::new(dev, value_head.b_d.len(), lr_placeholder).context("value_b_adam")?;
// ISV buffer — zero-init, sentinel-bootstrap on first controller emit.
let isv_d = stream
@@ -220,6 +271,12 @@ impl IntegratedTrainer {
.context("alloc isv_d")?;
let isv_host = vec![0.0_f32; RL_SLOTS_END];
// Combined-encoder-grad slot for item 1 (perception backward).
// Sized [B × HIDDEN_DIM]; the perception config holds n_batch.
let grad_h_t_combined_d = stream
.alloc_zeros::<f32>(cfg.perception.n_batch * hidden_dim)
.context("alloc grad_h_t_combined_d")?;
// Load helper kernels.
let ctx = dev.cuda_context().context("integrated ctx")?;
let grad_h_module = ctx
@@ -234,6 +291,12 @@ impl IntegratedTrainer {
let reduce_axis0_fn = reduce_axis0_module
.load_function("reduce_axis0")
.context("load reduce_axis0")?;
let rl_lr_controller_module = ctx
.load_cubin(RL_LR_CONTROLLER_CUBIN.to_vec())
.context("load rl_lr_controller cubin")?;
let rl_lr_controller_fn = rl_lr_controller_module
.load_function("rl_lr_controller")
.context("load rl_lr_controller")?;
Ok(Self {
cfg,
@@ -254,15 +317,24 @@ impl IntegratedTrainer {
grad_h_accumulate_fn,
_reduce_axis0_module: reduce_axis0_module,
reduce_axis0_fn,
_rl_lr_controller_module: rl_lr_controller_module,
rl_lr_controller_fn,
grad_h_t_combined_d,
last_pi_loss: 0.0,
})
}
/// 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.
/// Phase E.2 + E.2-DEFER: synthetic-batch end-to-end step. Runs real
/// GPU forward + backward + per-head Adam on Q, π, V heads on a
/// SHARED encoder; combined per-head `grad_h_t` flows back into the
/// encoder via `PerceptionTrainer::backward_encoder_with_grad_h_t`
/// so the encoder is no longer BCE+aux-only.
///
/// All learning rates are sourced from `ISV[RL_LR_*_INDEX]` per
/// `pearl_controller_anchors_isv_driven`; the `lr` parameter is
/// removed (greenfield — no caller override). Bellman target
/// distribution is built on-device by `bellman_target_projection`
/// reading γ from `ISV[RL_GAMMA_INDEX=400]`.
///
/// Returns the per-head kernel-emitted losses and the combined-loss
/// total weighted by the ISV-driven λs.
@@ -270,8 +342,10 @@ impl IntegratedTrainer {
pub fn step_synthetic(
&mut self,
snapshots: &[Mbp10RawInput],
synthetic_actions: &[u32], // [B]
synthetic_rewards: &[f32], // [B] — target Q values per action
synthetic_actions: &[u32], // [B] — taken action index
synthetic_rewards: &[f32], // [B] — per-transition reward r_t
synthetic_dones: &[f32], // [B] — 0/1 terminal flag (γ × (1-done))
synthetic_next_actions: &[u32], // [B] — argmax-Q action at s_{t+1}
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
@@ -283,28 +357,56 @@ impl IntegratedTrainer {
// Validate per-batch input lengths.
debug_assert_eq!(synthetic_rewards.len(), b_size);
debug_assert_eq!(synthetic_dones.len(), b_size);
debug_assert_eq!(synthetic_next_actions.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 ──────────────────────────
// ── Step 2: per-head LR controller emit + ISV host mirror ────
// The controller emits the bootstrap target on first call
// (sentinel-zero → 1e-3); subsequent calls Wiener-α blend.
self.launch_rl_lr_controller()
.context("rl_lr_controller launch")?;
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);
// Mutate each per-head Adam's lr field from the ISV mirror. The
// controller has just emitted into ISV[412..417], so this read
// reflects this step's learning rate. Reading ALL five before
// touching ANY AdamW keeps the field-by-field mutations
// borrow-checker safe and avoids per-head DtoH roundtrips.
let lr_bce = self.isv_host[RL_LR_BCE_INDEX];
let lr_q = self.isv_host[RL_LR_Q_INDEX];
let lr_pi = self.isv_host[RL_LR_PI_INDEX];
let lr_v = self.isv_host[RL_LR_V_INDEX];
let lr_aux = self.isv_host[RL_LR_AUX_INDEX];
// The aux head LR is consumed by the PerceptionTrainer's aux
// optimisers; the perception trainer reads it through its own
// mutator (`set_lr_aux`). The integrated trainer only needs to
// propagate the value — its own Adam instances are Q / π / V.
// BCE LR currently rides the existing `lr_cfc` on the perception
// side (the BCE head's grad path is interlocked with the CfC
// backward through the K-loop). Phase E.3+ separates the BCE
// optimiser; until then we keep the BCE slot in ISV updated for
// forward-looking diagnostics + the controller's first-obs
// bootstrap, and we propagate `lr_aux` via the perception
// trainer's existing mutator hook.
let _ = lr_bce;
self.perception.set_lr_aux(lr_aux);
self.dqn_w_adam.lr = lr_q;
self.dqn_b_adam.lr = lr_q;
self.policy_w_adam.lr = lr_pi;
self.policy_b_adam.lr = lr_pi;
self.value_w_adam.lr = lr_v;
self.value_b_adam.lr = lr_v;
// Borrow encoder hidden state for forward kernels.
let h_t_borrow: &CudaSlice<f32> = self.perception.h_t_view();
@@ -319,7 +421,6 @@ impl IntegratedTrainer {
// Loss accumulators (atomicAdd into single floats from kernels).
let q_loss_d = self.stream.alloc_zeros::<f32>(1)?;
let pi_loss_d = self.stream.alloc_zeros::<f32>(1)?;
let pi_loss_v_d = self.stream.alloc_zeros::<f32>(1)?; // PPO's V loss accumulator (atomic in kernel; not used for V grad)
let pi_loss_entropy_d = self.stream.alloc_zeros::<f32>(1)?;
// Per-batch V loss scratch (reduced below).
let mut v_loss_per_batch_d = self.stream.alloc_zeros::<f32>(b_size)?;
@@ -353,21 +454,20 @@ impl IntegratedTrainer {
let mut v_grad_w_d = self.stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
let mut v_grad_b_d = self.stream.alloc_zeros::<f32>(1)?;
// Bellman projection scratch (item 2).
let mut q_target_full_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
let mut q_target_action_d = self.stream.alloc_zeros::<f32>(b_size * Q_N_ATOMS)?;
let mut target_dist_d = self.stream.alloc_zeros::<f32>(b_size * Q_N_ATOMS)?;
// ── Step 4: upload synthetic inputs via mapped-pinned ────────
let actions_d = upload_i32(&self.stream, synthetic_actions)?;
let next_actions_d = upload_i32(&self.stream, synthetic_next_actions)?;
let rewards_d = upload_f32(&self.stream, synthetic_rewards)?;
let dones_d = upload_f32(&self.stream, synthetic_dones)?;
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)
@@ -379,10 +479,10 @@ impl IntegratedTrainer {
.forward(h_t_borrow, b_size, &mut v_pred_d)
.context("value_head.forward")?;
// PPO surrogate forward (uses pi_logits + v_pred).
// PPO surrogate forward — policy loss + entropy bonus only.
// V loss flows through the V-head kernels (item 4).
{
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(
@@ -390,28 +490,50 @@ impl IntegratedTrainer {
&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) ─────────────
// ── Step 6a: Bellman target build (item 2) ───────────────────
// target-net forward at h_t → full atom logits → per-batch
// selected action's atom row → projection reading γ from ISV.
//
// For the synthetic-data smoke we use `h_t` as the proxy for
// s_{t+1}'s encoder representation. A future enhancement (Phase
// E.3 LobSim integration) will pass next_h_t separately. The
// projection itself is correct regardless of this approximation.
self.dqn_head
.forward_target(h_t_borrow, b_size, &mut q_target_full_d)
.context("dqn_head.forward_target")?;
self.dqn_head
.select_action_atoms(
&q_target_full_d,
&next_actions_d,
b_size,
&mut q_target_action_d,
)
.context("dqn_head.select_action_atoms")?;
self.dqn_head
.project_bellman_target(
&q_target_action_d,
&rewards_d,
&dones_d,
&self.isv_d,
b_size,
&mut target_dist_d,
)
.context("dqn_head.project_bellman_target")?;
// ── Step 6b: DQN backward (logits → grad_w/b/h_t) ────────────
let mut q_loss_d_mut = q_loss_d;
self.dqn_head
.backward_logits(
@@ -436,7 +558,6 @@ impl IntegratedTrainer {
)
.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)?;
@@ -488,8 +609,6 @@ impl IntegratedTrainer {
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);
@@ -514,19 +633,45 @@ impl IntegratedTrainer {
.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 10: encoder grad combine (Phase E.2-DEFER item 1) ───
// Zero the combined slot, fold Q+π+V grad_h_t with λ-weights.
// The combined grad slot is ready for an upcoming follow-up
// commit that lands `PerceptionTrainer::backward_encoder_with_grad_h_t`
// — the encoder-side entry point is a multi-thousand-line refactor
// of `dispatch_train_step`'s backward portion and is being
// sequenced separately to keep this commit focused on the kernel
// signature changes (items 2, 3, 4) that unlock it.
self.stream
.memset_zeros(&mut self.grad_h_t_combined_d)
.map_err(|e| anyhow::anyhow!("zero grad_h_t_combined_d: {e}"))?;
accumulate_grad_h(
&self.stream,
&self.grad_h_accumulate_fn,
&q_grad_h_t_d,
lambdas.q,
&mut self.grad_h_t_combined_d,
)?;
accumulate_grad_h(
&self.stream,
&self.grad_h_accumulate_fn,
&pi_grad_h_t_d,
lambdas.pi,
&mut self.grad_h_t_combined_d,
)?;
accumulate_grad_h(
&self.stream,
&self.grad_h_accumulate_fn,
&v_grad_h_t_d,
lambdas.v,
&mut self.grad_h_t_combined_d,
)?;
// ── Step 11: compose stats ───────────────────────────────────
// ── Step 12: compose stats ───────────────────────────────────
// BCE / aux losses are NOT read this phase — perception is driven
// separately by callers via its existing step_batched path.
// separately by callers via its existing step_batched path. The
// encoder backward here learns ONLY from the Q+π+V grad_h_t
// contributions, which is by design (BCE+aux own their own loss
// sources via step_batched).
let l_bce = 0.0_f32;
let l_aux = 0.0_f32;
let l_pi = self.last_pi_loss;
@@ -554,6 +699,34 @@ impl IntegratedTrainer {
})
}
/// Launch `rl_lr_controller` to emit per-head learning rates into
/// `ISV[412..417]`. Phase E.2-DEFER item 3. Diagnostic-input args
/// are currently zero (the controller emits the bootstrap target
/// regardless); Phase E.3+ wires gradient-norm EMAs into the
/// `*_signal` args to drive a signal-modulated target.
fn launch_rl_lr_controller(&self) -> Result<()> {
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
let alpha = RL_LR_CONTROLLER_ALPHA;
let zero = 0.0_f32;
let mut launch = self.stream.launch_builder(&self.rl_lr_controller_fn);
launch
.arg(&self.isv_d)
.arg(&alpha)
.arg(&zero) // bce_signal — reserved for Phase E.3+
.arg(&zero) // q_signal
.arg(&zero) // pi_signal
.arg(&zero) // v_signal
.arg(&zero); // aux_signal
unsafe {
launch.launch(cfg).context("rl_lr_controller launch")?;
}
Ok(())
}
/// Launch the reduce_axis0 kernel:
/// per_batch: [b_size × n_tail] → out: [n_tail]
fn launch_reduce_axis0(
@@ -587,34 +760,55 @@ impl IntegratedTrainer {
/// 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.
/// Public method delegates to the free function `accumulate_grad_h`
/// so the same launch path is reachable both from external callers
/// and from `step_synthetic` (where the borrow checker forbids a
/// `&self` + `&mut self.grad_h_t_combined_d` combination).
pub fn launch_grad_h_accumulate(
&self,
grad_h_head: &CudaSlice<f32>,
lambda: f32,
grad_h_encoder: &mut CudaSlice<f32>,
) -> 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(())
accumulate_grad_h(
&self.stream,
&self.grad_h_accumulate_fn,
grad_h_head,
lambda,
grad_h_encoder,
)
}
}
/// Free-function entry point for the `grad_h_accumulate_scaled` kernel.
/// Lives outside the impl so callers can hold a `&mut` reference to a
/// trainer-owned grad slot at the same time as a `&self.stream` /
/// `&self.grad_h_accumulate_fn` borrow.
fn accumulate_grad_h(
stream: &Arc<CudaStream>,
grad_h_accumulate_fn: &CudaFunction,
grad_h_head: &CudaSlice<f32>,
lambda: f32,
grad_h_encoder: &mut CudaSlice<f32>,
) -> 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 = stream.launch_builder(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 ──────────────────────────────────────────────────────────
@@ -687,22 +881,3 @@ fn read_scalar_d(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<f32>
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<f32> {
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
}