diff --git a/crates/ml/tests/sp17_dueling_oracle_tests.rs b/crates/ml/tests/sp17_dueling_oracle_tests.rs index 3dbff1d00..1c0263027 100644 --- a/crates/ml/tests/sp17_dueling_oracle_tests.rs +++ b/crates/ml/tests/sp17_dueling_oracle_tests.rs @@ -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 = (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 { + 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 = run_q(&v_s1); + let q_s2: Vec = 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], + ); + } }