feat(rl): Phase 2.1 — scalar Bellman target + dueling Q (V baseline)
Spec: docs/superpowers/specs/2026-05-28-state-conditional-q-synthesis.md
§3 Sub-phase 2.1.
Replaces the distributional Bellman bootstrap
(fused_select_and_project_bellman) 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)
projected via two-bin linear interpolation onto the action's atom
support. The Q-head's atom logits are now interpreted as the
advantage distribution A_dist(s, a); the CE loss between
softmax(A_dist(s_t, a_t)) and the projected target trains A to
model deviation from the scalar V baseline.
Changes:
- NEW crates/ml-alpha/cuda/bellman_scalar_target_project.cu —
device-side scalar Bellman target + per-action atom projection.
Grid (B, 1, 1), Block (Q_N_ATOMS=21, 1, 1). No atomicAdd, no
cross-thread accumulation: scalar advantage produces a delta on
exactly two atoms so each thread writes its own atom value.
- crates/ml-alpha/build.rs registers the new kernel.
- crates/ml-alpha/src/rl/dqn.rs:
* Loads bellman_scalar_target_project from its own cubin.
* Adds DqnHead::bellman_scalar_target_project wrapper.
* Adds host-side composed_q_value(v_scalar, a_logits,
atom_supports) -> Vec<f32> implementing the Dueling
composition Q = V + A - mean_a(A) for test oracles.
- crates/ml-alpha/src/trainer/integrated.rs:
* Allocates sampled_v_curr_d + sampled_v_target_tp1_d for
V evaluated on PER-sampled hidden states.
* step_synthetic_body Bellman build: replaces forward_target +
fused_select_and_project_bellman with V forwards on sampled
buffers + the new scalar Bellman projection (single source of
truth — same swap applied to dqn_replay_step for k>0 path).
* Test accessors isv_host_slice_mut_for_test,
isv_dev_ptr_for_test, dqn_head_ref so kernel-oracle unit
tests can drive the new kernel without a full mega-graph.
- NEW crates/ml-alpha/tests/q_projection_unit.rs — pre-impl oracle:
projects y=0 to center atom, y=V_MAX to last, y=V_MIN to first,
arbitrary in-range cases recover EV within Δz/2, verifies V_t
enters as a subtraction. Plus pure host-side check of
composed_q_value.
All Phase 2.1 unit tests pass (5 CUDA --ignored + 2 host-side).
Smoke status: BLOCKED on pre-existing NaN abort at step 4 affecting
HEAD (c52282fb4 + f11bab542 + 7e38e46e6) — l_pi=NaN, l_aux=NaN at
step 4 reproduce IDENTICALLY with and without Phase 2.1 changes.
Per feedback_stop_on_anomaly, the pre-existing NaN must be diagnosed
before the Phase 2.1 smoke gate (l_v<2.0, l_q<3.0,
action_entropy>1.5 for 500 steps) can be validated. The Phase 2.1
changes themselves are structurally sound — unit tests confirm
correct projection and the smoke loss trajectories at steps 0-3 are
within ~1% of HEAD (l_q 2.9484 vs 2.9926, l_v 6.3293 vs 6.3291).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
|
||||
139
crates/ml-alpha/cuda/bellman_scalar_target_project.cu
Normal file
139
crates/ml-alpha/cuda/bellman_scalar_target_project.cu
Normal file
@@ -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;
|
||||
}
|
||||
@@ -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<f32> {
|
||||
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<f32> = 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::<f32>() / 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<CudaModule>,
|
||||
|
||||
/// 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<CudaModule>,
|
||||
|
||||
/// 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<f32>,
|
||||
dones_d: &CudaSlice<f32>,
|
||||
n_step_gammas_d: &CudaSlice<f32>,
|
||||
v_curr_d: &CudaSlice<f32>,
|
||||
v_target_tp1_d: &CudaSlice<f32>,
|
||||
actions_d: &CudaSlice<i32>,
|
||||
isv_dev_ptr: &u64,
|
||||
b_size: usize,
|
||||
target_dist_d: &mut CudaSlice<f32>,
|
||||
) -> 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<CudaStream> {
|
||||
|
||||
@@ -927,6 +927,16 @@ pub struct IntegratedTrainer {
|
||||
pub v_pred_d: CudaSlice<f32>,
|
||||
/// Scalar value prediction at h_{t+1} `[B]`.
|
||||
pub v_pred_tp1_d: CudaSlice<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// Policy logits at h_t `[B × N_ACTIONS]`.
|
||||
pub pi_logits_d: CudaSlice<f32>,
|
||||
|
||||
@@ -2013,6 +2023,18 @@ impl IntegratedTrainer {
|
||||
let v_pred_tp1_d = stream
|
||||
.alloc_zeros::<f32>(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::<f32>(b_size)
|
||||
.context("alloc sampled_v_curr_d")?;
|
||||
let sampled_v_target_tp1_d = stream
|
||||
.alloc_zeros::<f32>(b_size)
|
||||
.context("alloc sampled_v_target_tp1_d")?;
|
||||
let pi_logits_d = stream
|
||||
.alloc_zeros::<f32>(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
|
||||
|
||||
342
crates/ml-alpha/tests/q_projection_unit.rs
Normal file
342
crates/ml-alpha/tests/q_projection_unit.rs
Normal file
@@ -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<CudaStream>, host: &[f32]) -> cudarc::driver::CudaSlice<f32> {
|
||||
let mut d = stream.alloc_zeros::<f32>(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<CudaStream>, host: &[i32]) -> cudarc::driver::CudaSlice<i32> {
|
||||
let mut d = stream.alloc_zeros::<i32>(host.len()).expect("alloc");
|
||||
stream.memcpy_htod(host, &mut d).expect("htod"); // gpu-ok: test fixture, one-shot
|
||||
d
|
||||
}
|
||||
|
||||
fn readback_f32(stream: &Arc<CudaStream>, d: &cudarc::driver::CudaSlice<f32>) -> Vec<f32> {
|
||||
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<CudaStream>,
|
||||
reward: f32,
|
||||
n_step_gamma: f32,
|
||||
v_curr: f32,
|
||||
v_target_tp1: f32,
|
||||
v_min: f32,
|
||||
v_max: f32,
|
||||
) -> Vec<f32> {
|
||||
// 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::<f32>(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<f32> = (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);
|
||||
}
|
||||
Reference in New Issue
Block a user