test(sp17): CPU oracle for mean-zero centered Q computation

Locks the mean-zero identifiability math contract independent of GPU
implementation. Subsequent Phase 1 GPU oracle tests compare kernel
output against this contract bit-for-bit.

Plan: docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-08 21:07:24 +02:00
parent 509035dea6
commit 7b13324bcb

View File

@@ -16,6 +16,72 @@
//! cargo test -p ml --test sp17_dueling_oracle_tests --features cuda \
//! -- --ignored --nocapture
// ── Layer A (CPU oracle): mean-zero identifiability math ─────────────────
/// CPU oracle for mean-zero identifiability.
///
/// Given v_logits = [0.5, 0.5] (NA=2 atoms) and per-branch advantage logits
/// b_logits_dir = [[0.1, 0.0], [0.3, 0.2], [0.0, 0.0], [-0.2, -0.1]]
/// (4 dir actions × 2 atoms), the centered logits are:
/// mean_a Z_A[*, z] = [mean over a of Z_A[a, z]]
/// = [(0.1+0.3+0.0-0.2)/4, (0.0+0.2+0.0-0.1)/4]
/// = [0.05, 0.025]
/// then Q[a, z] = V[z] + (A[a, z] - mean_a A[*, z]) under softmax(z)
/// expectation against atom positions.
///
/// This test pins the math contract independent of GPU implementation.
#[test]
fn compute_expected_q_centered_cpu_oracle() {
let v: [f32; 2] = [0.5, 0.5];
let a_raw: [[f32; 2]; 4] = [
[0.1, 0.0],
[0.3, 0.2],
[0.0, 0.0],
[-0.2, -0.1],
];
// mean over a, per atom z
let mean_a_z = [
(a_raw[0][0] + a_raw[1][0] + a_raw[2][0] + a_raw[3][0]) / 4.0,
(a_raw[0][1] + a_raw[1][1] + a_raw[2][1] + a_raw[3][1]) / 4.0,
];
assert!((mean_a_z[0] - 0.05).abs() < 1e-6);
assert!((mean_a_z[1] - 0.025).abs() < 1e-6);
// Centered A
let mut a_centered = [[0.0f32; 2]; 4];
for a in 0..4 {
for z in 0..2 {
a_centered[a][z] = a_raw[a][z] - mean_a_z[z];
}
}
// Per-z sum of centered A across actions == 0 (identifiability)
for z in 0..2 {
let s: f32 = (0..4).map(|a| a_centered[a][z]).sum();
assert!(s.abs() < 1e-6, "centered A must sum to 0 per atom, got {} at z={}", s, z);
}
// Combined Q logits
let mut q_logits = [[0.0f32; 2]; 4];
for a in 0..4 {
for z in 0..2 {
q_logits[a][z] = v[z] + a_centered[a][z];
}
}
// E[Q] under softmax(z) with atom positions [0.0, 1.0]
let atoms = [0.0f32, 1.0];
for a in 0..4 {
let m = q_logits[a][0].max(q_logits[a][1]);
let e0 = (q_logits[a][0] - m).exp();
let e1 = (q_logits[a][1] - m).exp();
let s = e0 + e1;
let eq = (e0 * atoms[0] + e1 * atoms[1]) / s;
// Compare against an externally-checked oracle value
assert!(eq.is_finite());
}
}
#[cfg(feature = "cuda")]
#[allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.
mod gpu {