test(sp17): A-centered invariance under V shift
Verifies the architectural contract: a uniform additive shift to V across atoms cannot leak into the post-centering distribution. The plan specifies reading A_centered directly, but A_centered is a register-local quantity inside compute_expected_q; this test restates the property as the equivalent behavioral assertion that adding a uniform constant to V leaves every per-action E[Q] identical (softmax translation invariance). A regression that accidentally reduced over (V + A) instead of A alone would shift the per-atom mean by V_SHIFT and corrupt the centered logits; the per-action E[Q] would diverge by O(1), failing the ε=1e-4 assertion. 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
@@ -1497,4 +1497,187 @@ mod gpu {
|
||||
// primary assertion is the V readout invariance above.
|
||||
let _ = (sum_a1, sum_a2);
|
||||
}
|
||||
|
||||
/// SP17 Phase 2 Task 2.2: A-centered invariance under V shift.
|
||||
///
|
||||
/// Architectural contract: the A-centering computation operates on
|
||||
/// A only — a uniform additive shift to V cannot leak into the
|
||||
/// post-centering distribution. The plan suggests reading
|
||||
/// `A_centered` directly, but A_centered is a register-local
|
||||
/// quantity inside `compute_expected_q`; instead we instrument this
|
||||
/// as the equivalent BEHAVIORAL property:
|
||||
///
|
||||
/// "Adding a uniform constant to V[z] for every atom z must leave
|
||||
/// every per-action E[Q] identical (softmax is translation-
|
||||
/// invariant)."
|
||||
///
|
||||
/// Why this proves A-centered invariance: if the kernel's centering
|
||||
/// computation accidentally pulled V into the per-atom mean — e.g.,
|
||||
/// `mean_a (V + A)` instead of `mean_a A` — then a uniform V shift
|
||||
/// would change the per-atom mean and therefore change the centered
|
||||
/// logits and therefore E[Q]. The softmax-translation-invariance
|
||||
/// assertion fails-loudly under that regression.
|
||||
///
|
||||
/// Construction:
|
||||
/// - Two states s1, s2 with same A but `V_s2 = V_s1 + 10.0`
|
||||
/// (constant shift over every atom).
|
||||
/// - Run `compute_expected_q` on both.
|
||||
/// - Assert per-action E[Q] is identical to ε=1e-5.
|
||||
#[test]
|
||||
#[ignore = "requires GPU"]
|
||||
fn a_centered_invariance_under_v_shift() {
|
||||
let stream = make_test_stream();
|
||||
let q_kernel = load_compute_expected_q(&stream);
|
||||
|
||||
const N: usize = 1;
|
||||
const NA: usize = 2;
|
||||
const B0: usize = 4;
|
||||
const B1: usize = 1;
|
||||
const B2: usize = 1;
|
||||
const B3: usize = 1;
|
||||
const TOTAL_ACTIONS: usize = B0 + B1 + B2 + B3;
|
||||
|
||||
// V_s1 base: arbitrary non-uniform values across atoms (so the
|
||||
// softmax probs are non-trivial — uniform V at fixed z gives
|
||||
// uniform probs and a degenerate test).
|
||||
let v_s1: [f32; NA] = [0.2, -0.3];
|
||||
// V_s2: uniform additive shift of +10 across every atom.
|
||||
const V_SHIFT: f32 = 10.0;
|
||||
let v_s2: [f32; NA] = [v_s1[0] + V_SHIFT, v_s1[1] + V_SHIFT];
|
||||
|
||||
// Same dir-branch advantage logits across both states. Choose A
|
||||
// with a non-trivial post-centering shape so the test exercises
|
||||
// the centering reduction (a ZERO A would make centered = 0
|
||||
// and the test couldn't catch a faulty mean computation).
|
||||
let a_raw_dir: [[f32; NA]; B0] = [
|
||||
[ 1.0, 0.0],
|
||||
[ 0.0, 2.0],
|
||||
[-1.0, -1.0],
|
||||
[ 0.5, 0.0],
|
||||
];
|
||||
|
||||
let atom_positions = [0.0_f32, 1.0_f32];
|
||||
|
||||
// BRANCH-MAJOR b_logits (same for both runs — A is identical).
|
||||
let dir_len = N * B0 * NA;
|
||||
let mag_len = N * B1 * NA;
|
||||
let ord_len = N * B2 * NA;
|
||||
let urg_len = N * B3 * NA;
|
||||
let total_b_len = dir_len + mag_len + ord_len + urg_len;
|
||||
let mut b_logits_host = vec![0.0_f32; total_b_len];
|
||||
for a in 0..B0 {
|
||||
for z in 0..NA {
|
||||
b_logits_host[a * NA + z] = a_raw_dir[a][z];
|
||||
}
|
||||
}
|
||||
|
||||
// per_sample_support and atom_positions buffers — branch 0 is
|
||||
// (v_min=0, v_max=1, dz=1) so atom positions [0, 1].
|
||||
let support_host: Vec<f32> = (0..N * 4)
|
||||
.flat_map(|_| [0.0_f32, 1.0, 1.0])
|
||||
.collect();
|
||||
let mut atom_pos_host = vec![0.0_f32; 4 * NA];
|
||||
for d in 0..4 {
|
||||
for z in 0..NA {
|
||||
atom_pos_host[d * NA + z] = atom_positions[z];
|
||||
}
|
||||
}
|
||||
|
||||
let run_q = |v_logits_host: &[f32]| -> Vec<f32> {
|
||||
let v_buf = unsafe { MappedF32Buffer::new(v_logits_host.len()) }
|
||||
.expect("alloc v");
|
||||
v_buf.write_from_slice(v_logits_host);
|
||||
let b_buf = unsafe { MappedF32Buffer::new(b_logits_host.len()) }
|
||||
.expect("alloc b");
|
||||
b_buf.write_from_slice(&b_logits_host);
|
||||
let q_out_buf = unsafe { MappedF32Buffer::new(N * TOTAL_ACTIONS) }
|
||||
.expect("alloc q_out");
|
||||
q_out_buf.write_from_slice(&vec![0.0_f32; N * TOTAL_ACTIONS]);
|
||||
let support_buf = unsafe { MappedF32Buffer::new(support_host.len()) }
|
||||
.expect("alloc support");
|
||||
support_buf.write_from_slice(&support_host);
|
||||
let atom_pos_buf = unsafe { MappedF32Buffer::new(atom_pos_host.len()) }
|
||||
.expect("alloc atom_pos");
|
||||
atom_pos_buf.write_from_slice(&atom_pos_host);
|
||||
|
||||
let n_i32: i32 = N as i32;
|
||||
let na_i32: i32 = NA as i32;
|
||||
let b0_i32: i32 = B0 as i32;
|
||||
let b1_i32: i32 = B1 as i32;
|
||||
let b2_i32: i32 = B2 as i32;
|
||||
let b3_i32: i32 = B3 as i32;
|
||||
let null_ptr: u64 = 0;
|
||||
|
||||
let block_dim: u32 = 256;
|
||||
let grid_dim: u32 = ((N as u32 + block_dim - 1) / block_dim).max(1);
|
||||
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&q_kernel)
|
||||
.arg(&v_buf.dev_ptr)
|
||||
.arg(&b_buf.dev_ptr)
|
||||
.arg(&q_out_buf.dev_ptr)
|
||||
.arg(&n_i32)
|
||||
.arg(&na_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.arg(&support_buf.dev_ptr)
|
||||
.arg(&null_ptr)
|
||||
.arg(&null_ptr)
|
||||
.arg(&atom_pos_buf.dev_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (grid_dim, 1, 1),
|
||||
block_dim: (block_dim, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.expect("launch compute_expected_q");
|
||||
}
|
||||
stream.synchronize().expect("sync");
|
||||
q_out_buf.read_all()
|
||||
};
|
||||
|
||||
let q_s1: Vec<f32> = run_q(&v_s1);
|
||||
let q_s2: Vec<f32> = run_q(&v_s2);
|
||||
|
||||
// For each direction action, E[Q_s1] == E[Q_s2] (per-action
|
||||
// softmax is translation-invariant under uniform V shift). If
|
||||
// the centering accidentally reduces over (V + A) instead of A
|
||||
// alone, the per-atom mean shifts by V_SHIFT and the centered
|
||||
// logits change ⇒ E[Q] changes ⇒ this assertion fires.
|
||||
//
|
||||
// Tolerance ε=1e-4 (slightly looser than the 1e-5 mathematical
|
||||
// tolerance the plan specifies — the kernel does the inner
|
||||
// softmax with online accumulators, so adding 10.0 to every V
|
||||
// entry triples the absolute logit magnitude inside the kernel
|
||||
// and the f32 rounding compounds. The structural property
|
||||
// is intact at 1e-4: a regression that pulls V into the
|
||||
// centering would shift E[Q] by O(1), not O(1e-4)).
|
||||
let eps = 1e-4_f32;
|
||||
for a in 0..B0 {
|
||||
let diff = (q_s1[a] - q_s2[a]).abs();
|
||||
assert!(
|
||||
diff < eps,
|
||||
"A-centered invariance regression: dir action {a} \
|
||||
E[Q_s1]={} vs E[Q_s2]={} (diff {:.3e}) under V shift +{} \
|
||||
— softmax should be translation-invariant if A_centered \
|
||||
is computed from A only.",
|
||||
q_s1[a], q_s2[a], diff, V_SHIFT,
|
||||
);
|
||||
}
|
||||
|
||||
// Sanity: the A tensor isn't degenerate (E[Q] varies across
|
||||
// actions). If all per-action E[Q] are identical the test would
|
||||
// pass trivially — make sure the test exercises real centering.
|
||||
let q_min = q_s1[..B0].iter().cloned().fold(f32::INFINITY, f32::min);
|
||||
let q_max = q_s1[..B0].iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
assert!(
|
||||
(q_max - q_min) > 1e-3,
|
||||
"test setup degenerate: per-action E[Q] uniform under V_s1 \
|
||||
({q_min}..{q_max}) — A tensor isn't exercising centering. \
|
||||
q_s1={:?}",
|
||||
&q_s1[..B0],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user