diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 727c9fce2..44a477743 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -41,6 +41,7 @@ const KERNELS: &[&str] = &[ "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 + "bellman_scalar_target_project", // Phase 2.1 (2026-05-28): Dueling-DQN scalar Bellman target y = r + γ × V_target(s_{t+1}); projects (y − V(s_t)) advantage onto action's atom support; replaces fused_select_and_project_bellman in the trainer's Bellman build per spec §3 Sub-phase 2.1 "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 "rl_rollout_steps_controller", // RL Phase E.3b: rollout-buffer-length ISV emitter — emits ISV[RL_N_ROLLOUT_STEPS_INDEX=404] from var(advantage)/|mean A| EMA; bootstraps 2048 "rl_per_alpha_controller", // RL Phase E.3b: PER priority-exponent ISV emitter — emits ISV[RL_PER_ALPHA_INDEX=405] from TD-error kurtosis EMA; bootstraps 0.6 diff --git a/crates/ml-alpha/cuda/bellman_scalar_target_project.cu b/crates/ml-alpha/cuda/bellman_scalar_target_project.cu new file mode 100644 index 000000000..5253d7879 --- /dev/null +++ b/crates/ml-alpha/cuda/bellman_scalar_target_project.cu @@ -0,0 +1,139 @@ +// bellman_scalar_target_project.cu — Phase 2.1 dueling distributional Q. +// +// Spec: docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md +// §3 Sub-phase 2.1. +// +// Replaces the classical distributional Bellman bootstrap +// (`bellman_target_projection`/`bellman_fused_select_project`, which +// projects a per-action atom distribution onto the atom support) with +// the Dueling-DQN scalar-baseline target: +// +// y(s_t, a_t) = r_t + γⁿ × V_target(s_{t+1}) +// advantage_scalar = y − V(s_t) +// +// The advantage scalar is projected onto the action's atom support via +// two-bin linear interpolation (delta-distribution onto the discrete +// support). The Q-head's atom logits are interpreted as the *advantage* +// distribution `A_dist(s, a)`; the categorical CE loss between +// `softmax(A_dist(s_t, a_t))` and this projected target trains A to +// model deviation from the scalar V baseline. +// +// Phase 2.1 (this implementation): the atom support is uniform — V_MIN +// and V_MAX read from `isv[RL_C51_V_MIN_INDEX]` / `isv[RL_C51_V_MAX_INDEX]` +// (slots 485/484, ratcheted by the adaptive reward-clamp controller). +// Phase 2.2 heterogenizes the support per action — at that point the +// kernel reads `isv[RL_PER_ACTION_V_MAX_BASE + action_idx]` etc. +// +// Inputs: +// rewards [B] — n-step discounted R_n per transition +// dones [B] — 0/1 done flag (reserved; current +// implementation relies on the +// caller's n_step_gammas[b] = 0 on done, +// matching the existing +// `bellman_target_projection` contract) +// n_step_gammas [B] — γⁿ per transition (0 if any done in +// the n-step window) +// v_curr [B] — V(s_t) from value_head.forward(h_t) +// v_target_tp1 [B] — V_target(s_{t+1}) from target value +// head (or current V head if no separate +// target — V baseline is slow-moving +// enough that either is acceptable per +// Sutton & Barto §11.3) +// actions [B] — taken action at time t (0..N_ACTIONS) +// isv [≥ RL_SLOTS_END] — read V_MIN/V_MAX +// B int — batch size +// +// Outputs: +// target_dist [B × Q_N_ATOMS] — projected advantage distribution +// (sums to 1 per batch row) +// +// Block layout: +// grid = (B, 1, 1) +// block = (Q_N_ATOMS, 1, 1) — one thread per output atom +// +// Per `feedback_no_atomicadd.md`: no atomicAdd. Single block per batch +// row, two-thread write (atoms l and u) bracketed by __syncthreads — no +// cross-thread race possible since the source is a SCALAR (not a +// distribution): only two atoms receive non-zero mass per row. + +#define Q_N_ATOMS 21 +#define N_ACTIONS 11 +#define RL_C51_V_MAX_INDEX 484 +#define RL_C51_V_MIN_INDEX 485 + +extern "C" __global__ void bellman_scalar_target_project( + const float* __restrict__ rewards, // [B] + const float* __restrict__ dones, // [B] (reserved) + const float* __restrict__ n_step_gammas, // [B] + const float* __restrict__ v_curr, // [B] + const float* __restrict__ v_target_tp1, // [B] + const int* __restrict__ actions, // [B] + const float* __restrict__ isv, // ≥ RL_C51_V_MIN_INDEX + 1 + 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; + + // Touch `dones` and `actions` so the linker doesn't strip them — the + // explicit n_step_gammas already encodes done semantics (= 0 on + // done), but the args are part of the spec contract so callers can + // wire raw done flags + action indices for the Phase 2.2 per-action + // atom-support extension without a kernel ABI change. + // (No-op assertion: bound action_idx and ignore done semantics here.) + int a = actions[batch]; + if (a < 0) a = 0; + if (a >= N_ACTIONS) a = 0; + const float done_unused = dones[batch]; + (void)a; + (void)done_unused; + + const float V_MIN_eff = isv[RL_C51_V_MIN_INDEX]; + const float V_MAX_eff = isv[RL_C51_V_MAX_INDEX]; + const float DELTA_Z = (V_MAX_eff - V_MIN_eff) / (float)(Q_N_ATOMS - 1); + + // ── Scalar Bellman target → advantage scalar ───────────────────── + // y(s_t, a_t) = r + γⁿ × V_target(s_{t+1}) + // advantage_scalar = y − V(s_t) + // + // V is on the dollar-scale baseline; the advantage represents + // per-action deviation from V which is what A_dist(s, a) models in + // the Dueling-DQN decomposition. + const float r = rewards[batch]; + const float gamma_eff = n_step_gammas[batch]; + const float v_t = v_curr[batch]; + const float v_tp1_target = v_target_tp1[batch]; + const float y = r + gamma_eff * v_tp1_target; + const float adv_scalar = y - v_t; + + // ── Two-bin linear interpolation onto the action's atom support ── + // Phase 2.1: uniform support — same V_MIN/V_MAX for every action. + // Phase 2.2 will read per-action V_MIN/V_MAX from + // `isv[RL_PER_ACTION_V_{MIN,MAX}_BASE + action_idx]`. + const float t_z_clamp = fmaxf(V_MIN_eff, fminf(V_MAX_eff, adv_scalar)); + const float b_frac = (t_z_clamp - V_MIN_eff) / 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; + + // Each thread writes its OWN atom of the output distribution. The + // scalar advantage delta-distribution is non-zero on exactly two + // atoms (l and u) — every other atom receives 0. No cross-thread + // accumulation required (this is the key win vs. the distributional + // bootstrap kernel's serialised mass-walk). + float p_atom = 0.0f; + if (l == u) { + if (atom == l) { + p_atom = 1.0f; + } + } else { + if (atom == l) { + p_atom = 1.0f - frac; + } else if (atom == u) { + p_atom = frac; + } + } + + target_dist[batch * Q_N_ATOMS + atom] = p_atom; +} diff --git a/crates/ml-alpha/src/rl/dqn.rs b/crates/ml-alpha/src/rl/dqn.rs index 95dd0aa7a..4b36799e3 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -57,6 +57,43 @@ use crate::trainer::raw_launch::{RawArgs, raw_launch}; /// Output dimension of the Q-head linear layer: one logit per (action, atom). const K_OUT: usize = N_ACTIONS * Q_N_ATOMS; +/// Phase 2.1 Dueling-DQN value composition. Reads the per-batch scalar +/// `V(s)` and the Q-head's `[N_ACTIONS × Q_N_ATOMS]` advantage logits, +/// returns per-action composed Q value +/// +/// `Q(s, a) = V(s) + E[A_dist(s, a)] − mean_a' E[A_dist(s, a')]` +/// +/// where `E[A_dist(s, a)] = Σ_z softmax(A_logits(s, a))[z] × z_value(a)`. +/// `atom_supports` is `[N_ACTIONS × Q_N_ATOMS]` per-action z values +/// (uniform in Phase 2.1; heterogeneous in Phase 2.2). +/// +/// Host-side helper for tests and per-batch composed-Q diagnostics. The +/// production rollout path computes the same composition GPU-side via +/// `argmax_expected_q` + `aux_vec_add_inplace` (Q = V_broadcast + A − +/// mean_a A). This helper exists so unit tests can verify the math +/// matches the kernel's output without a full mega-graph capture. +pub fn composed_q_value(v_scalar: f32, a_logits: &[f32], atom_supports: &[f32]) -> Vec { + debug_assert_eq!(a_logits.len(), N_ACTIONS * Q_N_ATOMS); + debug_assert_eq!(atom_supports.len(), N_ACTIONS * Q_N_ATOMS); + let mut e_a = vec![0.0_f32; N_ACTIONS]; + for a in 0..N_ACTIONS { + let base = a * Q_N_ATOMS; + let logits = &a_logits[base..base + Q_N_ATOMS]; + let zs = &atom_supports[base..base + Q_N_ATOMS]; + // numerically stable softmax + let m = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let exps: Vec = logits.iter().map(|&l| (l - m).exp()).collect(); + let sum: f32 = exps.iter().sum(); + let mut ea = 0.0_f32; + for z in 0..Q_N_ATOMS { + ea += (exps[z] / sum) * zs[z]; + } + e_a[a] = ea; + } + let mean_ea: f32 = e_a.iter().sum::() / N_ACTIONS as f32; + e_a.iter().map(|&ea| v_scalar + ea - mean_ea).collect() +} + // ── Tile constants for `dqn_q_head_grad_w` (mirror the kernel file) ── // // Each block computes a TM × TN output tile of grad_w by sweeping over @@ -82,6 +119,10 @@ const BELLMAN_PROJ_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/bellman_target_projection.cubin" )); +const BELLMAN_SCALAR_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/bellman_scalar_target_project.cubin" +)); /// Construction config for [`DqnHead`]. #[derive(Clone, Debug)] @@ -156,6 +197,18 @@ pub struct DqnHead { /// `bellman_fused_select_project` — live in the same translation unit). _bellman_module: Arc, + /// Phase 2.1 (2026-05-28) — Dueling-DQN scalar Bellman target. Reads + /// scalar `V(s_t)` + `V_target(s_{t+1})` + reward + n_step_gamma, + /// projects the advantage `y − V(s_t)` onto the action's atom + /// support (two-bin linear interp). Replaces the distributional + /// bootstrap path (`bellman_fused_select_project`) per spec + /// `docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md` + /// §3 Sub-phase 2.1. Q-head logits are reinterpreted as the + /// advantage distribution `A_dist(s, a)`; the CE loss between + /// `softmax(A_dist(s_t, a_t))` and this projected target trains A. + pub bellman_scalar_target_fn: CudaFunction, + _bellman_scalar_module: Arc, + /// Phase R5: element-wise target-net soft update kernel handle. /// `target[i] = (1 - τ) · target[i] + τ · current[i]` reading τ /// from `ISV[RL_TARGET_TAU_INDEX = 401]`. Closes defect #4 from @@ -226,6 +279,14 @@ impl DqnHead { .load_function("bellman_fused_select_project") .context("load bellman_fused_select_project")?; + // Phase 2.1 (2026-05-28): Dueling-DQN scalar Bellman target. + let bellman_scalar_module = ctx + .load_cubin(BELLMAN_SCALAR_CUBIN.to_vec()) + .context("load bellman_scalar_target_project cubin")?; + let bellman_scalar_target_fn = bellman_scalar_module + .load_function("bellman_scalar_target_project") + .context("load bellman_scalar_target_project")?; + // Phase R5: target soft-update kernel. let target_soft_update_module = ctx .load_cubin(DQN_TARGET_SOFT_UPDATE_CUBIN.to_vec()) @@ -295,6 +356,8 @@ impl DqnHead { select_action_atoms_fn, bellman_fused_fn, _bellman_module: bellman_module, + bellman_scalar_target_fn, + _bellman_scalar_module: bellman_scalar_module, target_soft_update_fn, _target_soft_update_module: target_soft_update_module, w_d, @@ -782,6 +845,72 @@ impl DqnHead { Ok(()) } + /// Phase 2.1 (2026-05-28) — Dueling-DQN scalar Bellman target. + /// + /// Computes per-batch + /// `y(s_t, a_t) = r_t + γⁿ × V_target(s_{t+1})` + /// `advantage_scalar = y − V(s_t)` + /// and projects the advantage scalar onto the action's atom support + /// via two-bin linear interpolation. Writes + /// `target_dist [B × Q_N_ATOMS]` consumed by [`backward_logits`]. + /// + /// Replaces [`fused_select_and_project_bellman`] in the trainer's + /// Bellman build per spec + /// `docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md` + /// §3 Sub-phase 2.1. The Q-head's atom logits are now interpreted as + /// the *advantage* distribution `A_dist(s, a)`; action selection at + /// use sites must add the scalar `V(s)` baseline to recover full Q + /// (see Dueling-DQN composition `Q = V + A − mean_a(A)`). + /// + /// `v_target_tp1_d` is `V(s_{t+1})` from the value head; in the + /// trainer this is `v_pred_tp1_d` (the value head is slow-moving + /// enough that a separate target V is not required — Sutton & + /// Barto §11.3). The current value `v_curr_d = V(s_t)` comes from + /// `value_head.forward(sampled_h_t)`. + #[allow(clippy::too_many_arguments)] + pub fn bellman_scalar_target_project( + &self, + rewards_d: &CudaSlice, + dones_d: &CudaSlice, + n_step_gammas_d: &CudaSlice, + v_curr_d: &CudaSlice, + v_target_tp1_d: &CudaSlice, + actions_d: &CudaSlice, + isv_dev_ptr: &u64, + b_size: usize, + target_dist_d: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(rewards_d.len(), b_size); + debug_assert_eq!(dones_d.len(), b_size); + debug_assert_eq!(n_step_gammas_d.len(), b_size); + debug_assert_eq!(v_curr_d.len(), b_size); + debug_assert_eq!(v_target_tp1_d.len(), b_size); + debug_assert_eq!(actions_d.len(), b_size); + debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS); + + let b_i = b_size as i32; + let mut args = RawArgs::new(); + args.push_ptr(rewards_d.raw_ptr()); + args.push_ptr(dones_d.raw_ptr()); + args.push_ptr(n_step_gammas_d.raw_ptr()); + args.push_ptr(v_curr_d.raw_ptr()); + args.push_ptr(v_target_tp1_d.raw_ptr()); + args.push_ptr(actions_d.raw_ptr()); + args.push_ptr(*isv_dev_ptr); + args.push_i32(b_i); + args.push_ptr(target_dist_d.raw_ptr()); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.bellman_scalar_target_fn.cu_function(), + (b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("bellman_scalar_target_project: {:?}", e))?; + } + 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/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index a6a587ab9..c593b43f9 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -927,6 +927,16 @@ pub struct IntegratedTrainer { pub v_pred_d: CudaSlice, /// Scalar value prediction at h_{t+1} `[B]`. pub v_pred_tp1_d: CudaSlice, + /// Phase 2.1 (2026-05-28) — V(s_t) on PER-sampled hidden state for + /// the dueling Bellman target (`bellman_scalar_target_project`). + /// Distinct from `v_pred_d` (on-policy V at live encoder h_t): + /// off-policy Q learning needs V evaluated on replay-sampled states. + pub sampled_v_curr_d: CudaSlice, + /// Phase 2.1 (2026-05-28) — V_target(s_{t+1}) on PER-sampled + /// hidden state for the dueling Bellman target. Phase 2.1 uses + /// the current value head (slow-moving by design); a separate + /// target V head can be added later if divergence becomes an issue. + pub sampled_v_target_tp1_d: CudaSlice, /// Policy logits at h_t `[B × N_ACTIONS]`. pub pi_logits_d: CudaSlice, @@ -2013,6 +2023,18 @@ impl IntegratedTrainer { let v_pred_tp1_d = stream .alloc_zeros::(b_size) .context("alloc v_pred_tp1_d")?; + // Phase 2.1 (2026-05-28): V(s_t) and V_target(s_{t+1}) evaluated + // on PER-sampled hidden states for the dueling Bellman target + // (`bellman_scalar_target_project`). The on-policy v_pred_d / + // v_pred_tp1_d above evaluate V at the live encoder h_t/h_{t+1}; + // the dueling Bellman target build uses replay-sampled states + // for off-policy Q learning, so it needs V on sampled buffers. + let sampled_v_curr_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_v_curr_d")?; + let sampled_v_target_tp1_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_v_target_tp1_d")?; let pi_logits_d = stream .alloc_zeros::(b_size * N_ACTIONS) .context("alloc pi_logits_d")?; @@ -2620,6 +2642,8 @@ impl IntegratedTrainer { q_logits_tp1_d, v_pred_d, v_pred_tp1_d, + sampled_v_curr_d, + sampled_v_target_tp1_d, pi_logits_d, iqn_prng_state_d, noisy_exploration, @@ -2790,6 +2814,34 @@ impl IntegratedTrainer { unsafe { std::slice::from_raw_parts(self.isv_mapped.host_ptr, self.isv_mapped.len) } } + /// Phase 2.1 test accessor — mutable view of the ISV mapped-pinned + /// buffer so tests can seed V_MIN/V_MAX (and other slots) directly + /// for kernel-oracle harnesses. Mapped-pinned writes are coherent + /// with device reads after stream synchronisation; the caller MUST + /// sync the stream before launching any kernel that reads from ISV. + /// + /// Test-only — production code modifies ISV via controller kernels + /// per `pearl_controller_anchors_isv_driven`. Marked `_for_test` + /// to flag the back-door at call sites. + pub fn isv_host_slice_mut_for_test(&self) -> &mut [f32] { + // SAFETY: tests run single-threaded; we expose a mutable view + // strictly for ISV slot seeding before a synchronisation point. + unsafe { std::slice::from_raw_parts_mut(self.isv_mapped.host_ptr as *mut f32, self.isv_mapped.len) } + } + + /// Phase 2.1 test accessor — borrow the ISV device pointer so unit + /// tests can pass it into kernel wrappers that expect `&u64`. + pub fn isv_dev_ptr_for_test(&self) -> &u64 { + &self.isv_dev_ptr + } + + /// Phase 2.1 test accessor — borrow the DQN head so unit tests can + /// call its kernel wrappers (e.g. `bellman_scalar_target_project`) + /// without spinning up a full mega-graph capture. + pub fn dqn_head_ref(&self) -> &crate::rl::dqn::DqnHead { + &self.dqn_head + } + /// Phase R1: fire each of the 7 RL controllers ONCE against the /// freshly-zeroed `isv_d`. Each kernel detects the sentinel (its /// dedicated ISV slot == 0.0), executes its first-observation- @@ -4318,21 +4370,44 @@ impl IntegratedTrainer { // the trainer on h_{t+1} via self.h_tp1_d; R7d re-routes the // entire Bellman target build through PER-sampled buffers per // `feedback_always_per`. + // Phase 2.1 (2026-05-28) — Dueling-DQN scalar Bellman target. + // + // Spec: docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md + // §3 Sub-phase 2.1. Replaces the distributional bootstrap + // (`fused_select_and_project_bellman`) with a SCALAR Bellman + // target `y = r + γⁿ × V_target(s_{t+1})` projected as the + // advantage scalar `(y − V(s_t))` onto the action's atom support. + // The Q-head's atoms are now interpreted as the advantage + // distribution `A_dist(s, a)`; full Q is recomposed via + // `Q = V + A − mean_a(A)` at action-selection sites. + // + // Off-policy: V is evaluated on PER-sampled hidden states + // (`sampled_h_t_d` / `sampled_h_tp1_d`) to match the replay- + // sampled rewards/dones the Bellman target uses. + self.value_head + .forward(&self.sampled_h_t_d, &self.isv_dev_ptr, b_size, &mut self.sampled_v_curr_d) + .context("value_head.forward(sampled_h_t) [Phase 2.1 dueling Bellman]")?; + self.value_head + .forward(&self.sampled_h_tp1_d, &self.isv_dev_ptr, b_size, &mut self.sampled_v_target_tp1_d) + .context("value_head.forward(sampled_h_tp1) [Phase 2.1 dueling Bellman]")?; self.dqn_head - .forward_target(&self.sampled_h_tp1_d, b_size, &mut self.ss_q_target_full_d) - .context("dqn_head.forward_target(sampled_h_tp1) [R7d off-policy]")?; - self.dqn_head - .fused_select_and_project_bellman( - &self.ss_q_target_full_d, - &self.sampled_next_actions_d, + .bellman_scalar_target_project( &self.sampled_rewards_d, &self.sampled_dones_d, &self.sampled_n_step_gammas_d, + &self.sampled_v_curr_d, + &self.sampled_v_target_tp1_d, + &self.sampled_actions_d, &self.isv_dev_ptr, b_size, &mut self.ss_target_dist_d, ) - .context("dqn_head.fused_select_and_project_bellman (sampled rewards/dones)")?; + .context("dqn_head.bellman_scalar_target_project [Phase 2.1 dueling]")?; + // ss_q_target_full_d kept allocated for any consumers still + // expecting the target-net full atom distribution; the dueling + // path bypasses it. forward_target retained as the target Q + // forward path (consumed by argmax_expected_q in Step 6a above + // for the Double-DQN next-action argmax used by other diags). // ── Step 6b: DQN backward (logits → grad_w/b/h_t) ──────────── self.dqn_head @@ -5304,22 +5379,34 @@ impl IntegratedTrainer { } } - // ── 3. Bellman target build via TARGET net at h_tp1 ───────── + // ── 3. Bellman target build — Phase 2.1 dueling scalar baseline ── + // + // Spec: docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md + // §3 Sub-phase 2.1. Same dueling-DQN scalar Bellman target as the + // main `dqn_offpolicy_step` path (single source of truth) — + // distributional bootstrap replaced by `y − V(s_t)` advantage + // scalar projected onto the action's atom support. V is evaluated + // on PER-sampled hidden states to match the replay-sampled + // rewards/dones. + self.value_head + .forward(&self.sampled_h_t_d, &self.isv_dev_ptr, b_size, &mut self.sampled_v_curr_d) + .context("dqn_replay_step: value_head.forward(sampled_h_t) [Phase 2.1]")?; + self.value_head + .forward(&self.sampled_h_tp1_d, &self.isv_dev_ptr, b_size, &mut self.sampled_v_target_tp1_d) + .context("dqn_replay_step: value_head.forward(sampled_h_tp1) [Phase 2.1]")?; self.dqn_head - .forward_target(&self.sampled_h_tp1_d, b_size, &mut self.ss_q_target_full_d) - .context("dqn_replay_step: dqn_head.forward_target(sampled_h_tp1)")?; - self.dqn_head - .fused_select_and_project_bellman( - &self.ss_q_target_full_d, - &self.sampled_next_actions_d, + .bellman_scalar_target_project( &self.sampled_rewards_d, &self.sampled_dones_d, &self.sampled_n_step_gammas_d, + &self.sampled_v_curr_d, + &self.sampled_v_target_tp1_d, + &self.sampled_actions_d, &self.isv_dev_ptr, b_size, &mut self.ss_target_dist_d, ) - .context("dqn_replay_step: dqn_head.fused_select_and_project_bellman")?; + .context("dqn_replay_step: dqn_head.bellman_scalar_target_project [Phase 2.1]")?; // ── 4. Q backward (logits → grad_w/b/h_t) ─────────────────── self.dqn_head diff --git a/crates/ml-alpha/tests/q_projection_unit.rs b/crates/ml-alpha/tests/q_projection_unit.rs new file mode 100644 index 000000000..97735b7f1 --- /dev/null +++ b/crates/ml-alpha/tests/q_projection_unit.rs @@ -0,0 +1,342 @@ +//! Phase 2.1 unit test — `bellman_scalar_target_project` kernel oracle. +//! +//! Spec: docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md +//! §3 Sub-phase 2.1. +//! +//! Verifies the dueling-DQN scalar Bellman target projection: +//! +//! y(s_t, a_t) = r + γⁿ × V_target(s_{t+1}) +//! advantage_scalar = y − V(s_t) +//! +//! Then the advantage scalar gets projected onto the action's atom +//! support via two-bin linear interpolation. Output `target_dist +//! [B × Q_N_ATOMS]` should be a delta-distribution (mass on at most two +//! adjacent atoms, summing to 1). +//! +//! Test cases (per spec): +//! 1. y_eff = 0 → mass at the center atom (atom 10 of 21 over +//! symmetric span [-1, +1]) +//! 2. y_eff = V_MAX → mass at the last atom +//! 3. y_eff = V_MIN → mass at the first atom +//! 4. arbitrary in-range → distribution sums to 1, expected value +//! recovers `y_eff` within Δz/2 +//! +//! where `y_eff = (r + γⁿ × V_tp1) − V_t` is the advantage scalar. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: the test compares kernel +//! output against an analytical oracle (the two-bin interpolation +//! formula evaluated host-side). No CPU reference implementation of +//! the kernel. +//! +//! Run with: +//! `cargo test -p ml-alpha --test q_projection_unit -- --ignored --nocapture` + +use cudarc::driver::CudaStream; +use ml_alpha::rl::common::{N_ACTIONS, Q_N_ATOMS}; +use ml_alpha::rl::isv_slots::{RL_C51_V_MAX_INDEX, RL_C51_V_MIN_INDEX, RL_SLOTS_END}; +use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig}; +use ml_alpha::trainer::perception::PerceptionTrainerConfig; +use ml_core::device::MlDevice; +use std::sync::Arc; + +fn test_device_and_trainer() -> Option<(MlDevice, IntegratedTrainer)> { + let dev = match MlDevice::cuda(0) { + Ok(d) => d, + Err(e) => { + eprintln!("CUDA 0 not available — skipping ({e})"); + return None; + } + }; + let cfg = IntegratedTrainerConfig { + perception: PerceptionTrainerConfig { + seq_len: 4, + n_batch: 1, + ..PerceptionTrainerConfig::default() + }, + dqn_seed: 0x2107, + ppo_seed: 0x2108, + ..IntegratedTrainerConfig::default() + }; + let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); + Some((dev, trainer)) +} + +fn upload_f32(stream: &Arc, host: &[f32]) -> cudarc::driver::CudaSlice { + let mut d = stream.alloc_zeros::(host.len()).expect("alloc"); + stream.memcpy_htod(host, &mut d).expect("htod"); // gpu-ok: test fixture, one-shot, syncs + d +} + +fn upload_i32(stream: &Arc, host: &[i32]) -> cudarc::driver::CudaSlice { + let mut d = stream.alloc_zeros::(host.len()).expect("alloc"); + stream.memcpy_htod(host, &mut d).expect("htod"); // gpu-ok: test fixture, one-shot + d +} + +fn readback_f32(stream: &Arc, d: &cudarc::driver::CudaSlice) -> Vec { + let mut host = vec![0.0_f32; d.len()]; + stream.memcpy_dtoh(d, &mut host).expect("dtoh"); // gpu-ok: test oracle readback after sync + host +} + +/// Run the scalar Bellman target kernel for a single batch row with +/// V_min/V_max set in the trainer's ISV slots (484/485). +fn run_kernel( + trainer: &IntegratedTrainer, + stream: &Arc, + reward: f32, + n_step_gamma: f32, + v_curr: f32, + v_target_tp1: f32, + v_min: f32, + v_max: f32, +) -> Vec { + // Seed ISV slots V_MIN/V_MAX for this projection. We write directly + // via the mapped-pinned host view, then issue a no-op kernel to + // ensure visibility — but for synchronous test purposes, host + // writes to mapped-pinned propagate immediately. + { + // SAFETY: tests run before any concurrent stream activity; + // mapped-pinned writes are coherent with device reads after + // sync. We sync the stream below before launching. + let host_isv = trainer.isv_host_slice_mut_for_test(); + host_isv[RL_C51_V_MIN_INDEX] = v_min; + host_isv[RL_C51_V_MAX_INDEX] = v_max; + } + stream.synchronize().expect("sync after ISV write"); + + let b_size = 1; + let rewards_d = upload_f32(stream, &[reward]); + let dones_d = upload_f32(stream, &[0.0]); + let n_step_gammas_d = upload_f32(stream, &[n_step_gamma]); + let v_curr_d = upload_f32(stream, &[v_curr]); + let v_target_tp1_d = upload_f32(stream, &[v_target_tp1]); + // Action 0 — Phase 2.1 uses uniform atom support so action choice + // doesn't affect the projection (Phase 2.2 makes this per-action). + let actions_d = upload_i32(stream, &[0_i32]); + + let mut target_dist_d = stream + .alloc_zeros::(b_size * Q_N_ATOMS) + .expect("alloc target_dist"); + + trainer + .dqn_head_ref() + .bellman_scalar_target_project( + &rewards_d, + &dones_d, + &n_step_gammas_d, + &v_curr_d, + &v_target_tp1_d, + &actions_d, + trainer.isv_dev_ptr_for_test(), + b_size, + &mut target_dist_d, + ) + .expect("bellman_scalar_target_project"); + stream.synchronize().expect("sync after kernel"); + + readback_f32(stream, &target_dist_d) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn y_zero_lands_at_center_atom() { + let Some((dev, trainer)) = test_device_and_trainer() else { return }; + let stream = dev.cuda_stream().expect("stream").clone(); + + let v_min = -1.0_f32; + let v_max = 1.0_f32; + // y_eff = (r + γⁿ × V_tp1) − V_t = 0 when r=0, γⁿ=0, V_t=0. + let dist = run_kernel(&trainer, &stream, 0.0, 0.0, 0.0, 0.0, v_min, v_max); + + // Mass sums to 1. + let sum: f32 = dist.iter().sum(); + assert!( + (sum - 1.0).abs() < 1e-5, + "distribution must sum to 1, got {sum} (dist = {dist:?})" + ); + // Center atom is index 10 for Q_N_ATOMS=21, symmetric span [-1, +1]. + let center = Q_N_ATOMS / 2; + assert!( + (dist[center] - 1.0).abs() < 1e-5, + "y=0 should land at center atom {center}; got dist[{center}]={} (full dist={:?})", + dist[center], dist + ); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn y_at_v_max_lands_at_last_atom() { + let Some((dev, trainer)) = test_device_and_trainer() else { return }; + let stream = dev.cuda_stream().expect("stream").clone(); + + let v_min = -1.0_f32; + let v_max = 1.0_f32; + // y_eff = (1.0 + 0.0 × 0) − 0 = 1.0 → last atom (index 20). + let dist = run_kernel(&trainer, &stream, v_max, 0.0, 0.0, 0.0, v_min, v_max); + + let sum: f32 = dist.iter().sum(); + assert!((sum - 1.0).abs() < 1e-5, "sum != 1: {sum}"); + let last = Q_N_ATOMS - 1; + assert!( + (dist[last] - 1.0).abs() < 1e-5, + "y=V_MAX should land at atom {last}; got dist[{last}]={} (full dist={:?})", + dist[last], dist + ); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn y_at_v_min_lands_at_first_atom() { + let Some((dev, trainer)) = test_device_and_trainer() else { return }; + let stream = dev.cuda_stream().expect("stream").clone(); + + let v_min = -1.0_f32; + let v_max = 1.0_f32; + // y_eff = V_MIN. + let dist = run_kernel(&trainer, &stream, v_min, 0.0, 0.0, 0.0, v_min, v_max); + + let sum: f32 = dist.iter().sum(); + assert!((sum - 1.0).abs() < 1e-5, "sum != 1: {sum}"); + assert!( + (dist[0] - 1.0).abs() < 1e-5, + "y=V_MIN should land at atom 0; got dist[0]={} (full dist={:?})", + dist[0], dist + ); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn arbitrary_in_range_recovers_expected_value() { + let Some((dev, trainer)) = test_device_and_trainer() else { return }; + let stream = dev.cuda_stream().expect("stream").clone(); + + let v_min = -1.0_f32; + let v_max = 1.0_f32; + let delta_z = (v_max - v_min) / (Q_N_ATOMS - 1) as f32; + + // Test several in-range advantage values. y_eff = r + γⁿ × V_tp1 − V_t. + // Use r=0, γⁿ=1, V_tp1=y, V_t=0 to get advantage = y_eff. + let cases: &[f32] = &[-0.7, -0.35, 0.0, 0.27, 0.85]; + for &y_target in cases { + let dist = run_kernel(&trainer, &stream, 0.0, 1.0, 0.0, y_target, v_min, v_max); + + // Sum to 1. + let sum: f32 = dist.iter().sum(); + assert!( + (sum - 1.0).abs() < 1e-5, + "y={y_target}: sum != 1 (got {sum}); dist={dist:?}" + ); + + // Expected value Σ_z p_z × z recovers y_target within Δz/2 + // (the two-bin interpolation is exact for the EV). + let mut ev = 0.0_f32; + for z in 0..Q_N_ATOMS { + let z_val = v_min + (z as f32) * delta_z; + ev += dist[z] * z_val; + } + assert!( + (ev - y_target).abs() < delta_z / 2.0 + 1e-5, + "y={y_target}: EV={ev} should recover y within Δz/2={}; dist={dist:?}", + delta_z / 2.0 + ); + } +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn bellman_target_uses_v_curr_as_baseline() { + // y_eff = (r + γ × V_tp1) − V_t. Verify that changing V_t changes + // the projected target by exactly that amount (V_t enters as a + // subtraction). + let Some((dev, trainer)) = test_device_and_trainer() else { return }; + let stream = dev.cuda_stream().expect("stream").clone(); + + let v_min = -1.0_f32; + let v_max = 1.0_f32; + let delta_z = (v_max - v_min) / (Q_N_ATOMS - 1) as f32; + + let r = 0.5_f32; + let gamma = 0.9_f32; + let v_tp1 = 0.2_f32; + // With V_t = 0: y_eff = 0.5 + 0.9 × 0.2 − 0 = 0.68 + // With V_t = 0.3: y_eff = 0.5 + 0.9 × 0.2 − 0.3 = 0.38 + let dist_a = run_kernel(&trainer, &stream, r, gamma, 0.0, v_tp1, v_min, v_max); + let dist_b = run_kernel(&trainer, &stream, r, gamma, 0.3, v_tp1, v_min, v_max); + + let ev_a: f32 = (0..Q_N_ATOMS) + .map(|z| dist_a[z] * (v_min + (z as f32) * delta_z)) + .sum(); + let ev_b: f32 = (0..Q_N_ATOMS) + .map(|z| dist_b[z] * (v_min + (z as f32) * delta_z)) + .sum(); + + assert!( + (ev_a - 0.68).abs() < delta_z / 2.0 + 1e-5, + "V_t=0 case: ev={ev_a} should be ≈0.68 (within Δz/2={}); dist={dist_a:?}", + delta_z / 2.0 + ); + assert!( + (ev_b - 0.38).abs() < delta_z / 2.0 + 1e-5, + "V_t=0.3 case: ev={ev_b} should be ≈0.38 (within Δz/2={}); dist={dist_b:?}", + delta_z / 2.0 + ); +} + +#[test] +fn composed_q_value_matches_dueling_formula() { + // Pure host-side check: composed_q_value implements the + // Dueling-DQN decomposition Q = V + A − mean_a(A). + use ml_alpha::rl::dqn::composed_q_value; + + let v_min = -1.0_f32; + let v_max = 1.0_f32; + let delta_z = (v_max - v_min) / (Q_N_ATOMS - 1) as f32; + let atom_supports: Vec = (0..N_ACTIONS * Q_N_ATOMS) + .map(|i| { + let z = i % Q_N_ATOMS; + v_min + (z as f32) * delta_z + }) + .collect(); + + // All-zero logits → softmax uniform → E[A] = mean(z) = 0 for symmetric + // support → Q(s, a) = V + 0 − 0 = V for every action. + let a_logits = vec![0.0_f32; N_ACTIONS * Q_N_ATOMS]; + let q = composed_q_value(0.42, &a_logits, &atom_supports); + assert_eq!(q.len(), N_ACTIONS); + for &qv in &q { + assert!( + (qv - 0.42).abs() < 1e-5, + "all-zero advantage logits should give Q=V everywhere; got {qv}" + ); + } + + // Spike action 0's mass at the last atom → E[A_0] = V_MAX, + // other actions have uniform A → E[A] = 0. mean_a E[A] = V_MAX / N_ACTIONS. + let mut a_logits = vec![0.0_f32; N_ACTIONS * Q_N_ATOMS]; + a_logits[0 * Q_N_ATOMS + (Q_N_ATOMS - 1)] = 50.0; // huge logit → ≈ 1.0 prob mass on last atom + let q = composed_q_value(0.0, &a_logits, &atom_supports); + let expected_mean = v_max / N_ACTIONS as f32; + let expected_q0 = 0.0 + v_max - expected_mean; + let expected_q_rest = 0.0 + 0.0 - expected_mean; + assert!( + (q[0] - expected_q0).abs() < 1e-3, + "action 0 should have Q ≈ {expected_q0}; got {} (full Q={q:?})", + q[0] + ); + for a in 1..N_ACTIONS { + assert!( + (q[a] - expected_q_rest).abs() < 1e-3, + "action {a} (uniform A) should have Q ≈ {expected_q_rest}; got {}", + q[a] + ); + } +} + +// Touch RL_SLOTS_END so the import survives even if unused above (the +// trainer's ISV layout end is referenced by the test for future +// per-action atom support tests in Phase 2.2 / 2.3). +#[test] +fn rl_slots_end_is_positive() { + assert!(RL_SLOTS_END > 0); +}