diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 0f686f256..0b783e65d 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -41,7 +41,6 @@ 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 "mean_b_reduce", // 2026-05-29: deterministic [B]→[1] mean reducer replacing atomicAdd in dqn_distributional_q (l_q) and ppo_clipped_surrogate (l_pi, l_ent); single block + warp-shuffle per feedback_no_atomicadd "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 diff --git a/crates/ml-alpha/cuda/bellman_scalar_target_project.cu b/crates/ml-alpha/cuda/bellman_scalar_target_project.cu deleted file mode 100644 index 5253d7879..000000000 --- a/crates/ml-alpha/cuda/bellman_scalar_target_project.cu +++ /dev/null @@ -1,139 +0,0 @@ -// 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 fef49d51c..b3cd392d3 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -57,43 +57,6 @@ 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 @@ -126,10 +89,6 @@ const MEAN_B_REDUCE_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/mean_b_reduce.cubin" )); -const BELLMAN_SCALAR_CUBIN: &[u8] = include_bytes!(concat!( - env!("OUT_DIR"), - "/bellman_scalar_target_project.cubin" -)); /// Construction config for [`DqnHead`]. #[derive(Clone, Debug)] @@ -204,17 +163,6 @@ 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, /// 2026-05-29: deterministic mean reducer for cross-batch loss /// accumulation in `backward_logits`. Replaces atomicAdd into /// `loss_out` (which was non-deterministic across warp execution @@ -292,13 +240,6 @@ 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")?; // 2026-05-29: deterministic mean reducer (no-atomicAdd fix). let mean_b_reduce_module = ctx .load_cubin(MEAN_B_REDUCE_CUBIN.to_vec()) @@ -376,8 +317,6 @@ impl DqnHead { select_action_atoms_fn, bellman_fused_fn, _bellman_module: bellman_module, - bellman_scalar_target_fn, - _bellman_scalar_module: bellman_scalar_module, mean_b_reduce_fn, _mean_b_reduce_module: mean_b_reduce_module, target_soft_update_fn, @@ -890,72 +829,6 @@ 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 1a57201c3..82e4803c3 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -935,16 +935,6 @@ 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, @@ -2071,18 +2061,6 @@ 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")?; @@ -2695,8 +2673,6 @@ 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, @@ -2907,34 +2883,6 @@ 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- @@ -4527,44 +4475,21 @@ 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 - .bellman_scalar_target_project( + .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, &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.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). + .context("dqn_head.fused_select_and_project_bellman (sampled rewards/dones)")?; // ── Step 6b: DQN backward (logits → grad_w/b/h_t) ──────────── self.dqn_head @@ -5536,34 +5461,22 @@ impl IntegratedTrainer { } } - // ── 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]")?; + // ── 3. Bellman target build via TARGET net at h_tp1 ───────── self.dqn_head - .bellman_scalar_target_project( + .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, &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.bellman_scalar_target_project [Phase 2.1]")?; + .context("dqn_replay_step: dqn_head.fused_select_and_project_bellman")?; // ── 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 deleted file mode 100644 index 97735b7f1..000000000 --- a/crates/ml-alpha/tests/q_projection_unit.rs +++ /dev/null @@ -1,342 +0,0 @@ -//! 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); -}