foxhunt/crates/ml-alpha/tests/rollout_buffer_invariants.rs

//! GPU-oracle invariant tests for `RolloutBuffer` + `gae_backward_sweep`.
//!
//! Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md
//! Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md
//!
//! Five tests verify the GAE backward-sweep kernel's analytical
//! invariants:
//!
//!   1. `gae_terminal_only_matches_close_event_pnl` — single terminal at
//!      T-1 with otherwise-zero rewards collapses GAE to the
//!      close-event PnL pattern (δ_{T-1} = r_{T-1} - V_{T-1}).
//!   2. `gae_dense_reward_geometric_decay` — uniform r=1, V=0, no
//!      terminals yields the closed-form geometric series A_t = Σ
//!      (γλ)^k for k=0..(T-t-1).
//!   3. `gae_done_resets_credit` — a terminal in the middle of the
//!      trajectory zeroes the GAE accumulator so credit does not leak
//!      across episode boundaries.
//!   4. `gae_deterministic_across_runs` — bit-equality contract; the
//!      sequential per-batch sweep must produce identical outputs for
//!      identical inputs across fresh CUDA contexts (sustains the
//!      determinism foundation per `pearl_determinism_achieved`).
//!   5. `rollout_buffer_alloc_sizes_match_spec` — allocation sizing
//!      contract for the [B × T], [B × T × H], and [B] device buffers.
//!
//! Per `feedback_no_cpu_test_fallbacks`: every assertion is an
//! analytical identity of the GAE recurrence, not a CPU reimplementation.
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: host↔device transfers
//! go through mapped-pinned `MappedRecordBuffer<T>` (same pattern as the
//! production trainer).
//!
//! Run with:
//!   `SQLX_OFFLINE=true cargo test -p ml-alpha --test rollout_buffer_invariants --release -- --ignored --nocapture`

#![allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.

use cudarc::driver::{CudaSlice, CudaStream, DevicePtrMut};
use ml_alpha::pinned_mem::MappedRecordBuffer;
use ml_alpha::trainer::rollout_buffer::RolloutBuffer;
use ml_core::device::MlDevice;
use std::sync::Arc;

const REL_TOL: f32 = 1e-5;
const ABS_TOL: f32 = 1e-5;

const TEST_GAMMA: f32 = 0.99;
const TEST_LAMBDA: f32 = 0.95;

fn make_device() -> Option<MlDevice> {
    match MlDevice::cuda(0) {
        Ok(d) => Some(d),
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            None
        }
    }
}

/// Upload an `f32` host slice into a fresh `CudaSlice<f32>` via mapped
/// pinned staging + DtoD. Sync the stream so callers observe the write.
fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
    let n = host.len();
    let staging = unsafe { MappedRecordBuffer::<f32>::new(n) }.expect("f32 staging alloc");
    staging.write_from_slice(host);
    let mut d = stream.alloc_zeros::<f32>(n).expect("alloc f32 dst");
    unsafe {
        let (dst_ptr, _g) = d.device_ptr_mut(stream);
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr,
            staging.dev_ptr,
            n * std::mem::size_of::<f32>(),
            stream.cu_stream(),
        )
        .expect("f32 DtoD upload");
    }
    stream.synchronize().expect("upload_f32 sync");
    d
}

/// Upload a `u8` host slice into a fresh `CudaSlice<u8>` via mapped
/// pinned staging + DtoD. Used for the dones buffer.
fn upload_u8(stream: &Arc<CudaStream>, host: &[u8]) -> CudaSlice<u8> {
    let n = host.len();
    let staging = unsafe { MappedRecordBuffer::<u8>::new(n) }.expect("u8 staging alloc");
    unsafe {
        for (i, &v) in host.iter().enumerate() {
            std::ptr::write_volatile(staging.host_ptr.add(i), v);
        }
    }
    let mut d = stream.alloc_zeros::<u8>(n).expect("alloc u8 dst");
    unsafe {
        let (dst_ptr, _g) = d.device_ptr_mut(stream);
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr,
            staging.dev_ptr,
            n * std::mem::size_of::<u8>(),
            stream.cu_stream(),
        )
        .expect("u8 DtoD upload");
    }
    stream.synchronize().expect("upload_u8 sync");
    d
}

/// Download an `f32` device slice via mapped pinned staging. Uses
/// `read_slice_d_into` to honour `feedback_no_htod_htoh_only_mapped_pinned`.
fn download_f32(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Vec<f32> {
    let n = src.len();
    let mut host = vec![0.0_f32; n];
    ml_alpha::trainer::integrated::read_slice_d_into(stream, src, host.as_mut_slice())
        .expect("read_slice_d_into f32");
    host
}

/// Element-wise assert close: `|a - b| ≤ ABS_TOL + REL_TOL · max(|a|, |b|)`.
fn assert_close(a: f32, b: f32, ctx: &str) {
    let tol = ABS_TOL + REL_TOL * a.abs().max(b.abs());
    assert!(
        (a - b).abs() <= tol,
        "{ctx}: {a} vs {b} (|Δ|={} > tol={tol})",
        (a - b).abs()
    );
}

/// Convenience: populate a `RolloutBuffer`'s read-only inputs and run
/// the kernel. Returns the host-side (advantages, returns).
fn run_gae(
    rb: &mut RolloutBuffer,
    stream: &Arc<CudaStream>,
    rewards: &[f32],
    dones: &[u8],
    v_t: &[f32],
    v_t_bootstrap: &[f32],
    gamma: f32,
    lambda: f32,
) -> (Vec<f32>, Vec<f32>) {
    rb.rewards_d = upload_f32(stream, rewards);
    rb.dones_d = upload_u8(stream, dones);
    rb.v_t_d = upload_f32(stream, v_t);
    rb.v_t_bootstrap_d = upload_f32(stream, v_t_bootstrap);
    rb.compute_gae(stream, gamma, lambda).expect("compute_gae");
    stream.synchronize().expect("gae sync");
    (
        download_f32(stream, &rb.advantages_d),
        download_f32(stream, &rb.returns_d),
    )
}

// ── Test 1: terminal-only trajectory ───────────────────────────────────

/// B=2, T=4, every step has zero reward except the final terminal step.
/// V is constant per batch.
///
/// Expected (per batch b):
///   t = 3 (done=1):
///     non_terminal = 0
///     v_tp1 contribution = γ · v_bootstrap · 0 = 0
///     δ_3 = r_3 + 0 − V_3 = r_3 − V_3
///     A_3 = δ_3 + γλ · A_next · 0 = δ_3
///     a_next = A_3 · 0 = 0  (reset on done)
///   t = 2 (done=0):
///     v_tp1 = V_3
///     δ_2 = 0 + γ · V_3 · 1 − V_2 = γ · V_3 − V_2
///     A_2 = δ_2 + γλ · 0 · 1 = δ_2
///   t = 1 (done=0): same shape; v_tp1 = V_2
///     δ_1 = γ · V_2 − V_1
///     A_1 = δ_1 + γλ · A_2 · 1
///   t = 0 (done=0): v_tp1 = V_1
///     δ_0 = γ · V_1 − V_0
///     A_0 = δ_0 + γλ · A_1 · 1
///
/// Returns: R_t = A_t + V_t.
#[test]
#[ignore = "requires CUDA"]
fn gae_terminal_only_matches_close_event_pnl() {
    let Some(dev) = make_device() else { return };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    let ctx = dev.cuda_context().expect("cuda_context");

    let b = 2usize;
    let t = 4usize;
    let hidden_dim = 4usize;
    let mut rb = RolloutBuffer::new(&ctx, &stream, b, t, hidden_dim)
        .expect("RolloutBuffer::new");

    // Layout: [b][t] flattened as b*T + t.
    // rewards = [[0, 0, 0, +10], [0, 0, 0, −5]]
    let rewards = vec![
        0.0, 0.0, 0.0, 10.0,
        0.0, 0.0, 0.0, -5.0,
    ];
    // dones = [[0, 0, 0, 1], [0, 0, 0, 1]]
    let dones = vec![
        0u8, 0, 0, 1,
        0u8, 0, 0, 1,
    ];
    // V_t = [[1, 1, 1, 0], [-1, -1, -1, 0]]
    let v_t = vec![
        1.0, 1.0, 1.0, 0.0,
        -1.0, -1.0, -1.0, 0.0,
    ];
    // Bootstrap arbitrary — zeroed by non_terminal=0 at t=T-1.
    let v_t_bootstrap = vec![0.0, 0.0];

    let (a, r) = run_gae(&mut rb, &stream, &rewards, &dones, &v_t, &v_t_bootstrap,
                         TEST_GAMMA, TEST_LAMBDA);

    // Hand computation per batch ───────────────────────────────────────
    for b_idx in 0..b {
        let off = b_idx * t;
        let vs = &v_t[off..off + t];
        let rs = &rewards[off..off + t];

        // t = 3 (terminal):
        let delta_3 = rs[3] - vs[3];
        let a_3 = delta_3;
        let a_next_after_3: f32 = 0.0; // reset on done.

        // t = 2:
        let delta_2 = rs[2] + TEST_GAMMA * vs[3] - vs[2];
        let a_2 = delta_2 + TEST_GAMMA * TEST_LAMBDA * a_next_after_3;
        let a_next_after_2 = a_2; // non_terminal = 1

        // t = 1:
        let delta_1 = rs[1] + TEST_GAMMA * vs[2] - vs[1];
        let a_1 = delta_1 + TEST_GAMMA * TEST_LAMBDA * a_next_after_2;
        let a_next_after_1 = a_1;

        // t = 0:
        let delta_0 = rs[0] + TEST_GAMMA * vs[1] - vs[0];
        let a_0 = delta_0 + TEST_GAMMA * TEST_LAMBDA * a_next_after_1;

        let expected_a = [a_0, a_1, a_2, a_3];
        let expected_r = [
            a_0 + vs[0],
            a_1 + vs[1],
            a_2 + vs[2],
            a_3 + vs[3],
        ];

        for j in 0..t {
            assert_close(a[off + j], expected_a[j],
                         &format!("A[{b_idx}][{j}]"));
            assert_close(r[off + j], expected_r[j],
                         &format!("R[{b_idx}][{j}]"));
        }
    }
}

// ── Test 2: dense reward, no terminals, V=0 → geometric series ─────────

/// B=1, T=10, all r=1, V=0 everywhere, no terminals, bootstrap=0.
///
/// Then δ_t = 1 + γ·0·1 − 0 = 1 (for t < T-1) and δ_{T-1} = 1 + γ·0·1 − 0 = 1.
/// A_t = 1 + γλ · A_{t+1}.
/// Closed form: A_t = Σ_{k=0}^{T-1-t} (γλ)^k.
/// Returns equal advantages because V_t = 0.
#[test]
#[ignore = "requires CUDA"]
fn gae_dense_reward_geometric_decay() {
    let Some(dev) = make_device() else { return };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    let ctx = dev.cuda_context().expect("cuda_context");

    let b = 1usize;
    let t = 10usize;
    let hidden_dim = 4usize;
    let mut rb = RolloutBuffer::new(&ctx, &stream, b, t, hidden_dim)
        .expect("RolloutBuffer::new");

    let rewards = vec![1.0_f32; b * t];
    let dones = vec![0u8; b * t];
    let v_t = vec![0.0_f32; b * t];
    let v_t_bootstrap = vec![0.0_f32; b];

    let (a, r) = run_gae(&mut rb, &stream, &rewards, &dones, &v_t, &v_t_bootstrap,
                         TEST_GAMMA, TEST_LAMBDA);

    let gl = TEST_GAMMA * TEST_LAMBDA;
    // Backwards: A_{T-1} = 1; A_t = 1 + gl·A_{t+1}.
    let mut expected = vec![0.0_f32; t];
    expected[t - 1] = 1.0;
    for k in (0..t - 1).rev() {
        expected[k] = 1.0 + gl * expected[k + 1];
    }
    // Sanity: also equals Σ_{k=0}^{T-1-t} (gl)^k closed form.
    for t_idx in 0..t {
        let n_terms = t - 1 - t_idx;
        let mut closed_form = 0.0_f32;
        let mut p = 1.0_f32;
        for _ in 0..=n_terms {
            closed_form += p;
            p *= gl;
        }
        assert_close(expected[t_idx], closed_form,
                     &format!("recurrence vs closed form @ t={t_idx}"));
    }

    for t_idx in 0..t {
        assert_close(a[t_idx], expected[t_idx],
                     &format!("A[0][{t_idx}]"));
        // V_t = 0 → R_t = A_t.
        assert_close(r[t_idx], expected[t_idx],
                     &format!("R[0][{t_idx}]"));
    }
}

// ── Test 3: mid-trajectory done resets credit ─────────────────────────

/// B=1, T=6, dones = [0, 0, 1, 0, 0, 0]. All r=1, V=0, bootstrap=0.
///
/// Backwards sweep:
///   t = 5 (done=0): δ = 1, A_5 = 1, a_next = 1
///   t = 4: δ = 1, A_4 = 1 + gl·1 = 1 + gl
///   t = 3: A_3 = 1 + gl·A_4
///   t = 2 (done=1): non_terminal = 0
///     δ_2 = 1 + γ·V_3·0 − 0 = 1
///     A_2 = 1 + γλ·A_3·0 = 1          (γλ·A_{t+1}·non_terminal=0)
///     a_next ← A_2 · 0 = 0            (RESET)
///   t = 1: δ = 1, A_1 = 1 + γλ · 0 = 1
///   t = 0: δ = 1, A_0 = 1 + γλ · A_1 = 1 + gl
///
/// The done at t=2 must zero the credit handed to t=1, so A_1 = 1
/// regardless of A_3..A_5.
#[test]
#[ignore = "requires CUDA"]
fn gae_done_resets_credit() {
    let Some(dev) = make_device() else { return };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    let ctx = dev.cuda_context().expect("cuda_context");

    let b = 1usize;
    let t = 6usize;
    let hidden_dim = 4usize;
    let mut rb = RolloutBuffer::new(&ctx, &stream, b, t, hidden_dim)
        .expect("RolloutBuffer::new");

    let rewards = vec![1.0_f32; b * t];
    let dones = vec![0u8, 0, 1, 0, 0, 0];
    let v_t = vec![0.0_f32; b * t];
    let v_t_bootstrap = vec![0.0_f32; b];

    let (a, _r) = run_gae(&mut rb, &stream, &rewards, &dones, &v_t, &v_t_bootstrap,
                          TEST_GAMMA, TEST_LAMBDA);

    let gl = TEST_GAMMA * TEST_LAMBDA;
    let a_5 = 1.0_f32;
    let a_4 = 1.0_f32 + gl * a_5;
    let a_3 = 1.0_f32 + gl * a_4;
    let a_2 = 1.0_f32;               // terminal — γλ·A_3·0 = 0
    // a_next after t=2 is A_2 · 0 = 0.
    let a_1 = 1.0_f32;               // 1 + gl · 0 (reset)
    let a_0 = 1.0_f32 + gl * a_1;    // 1 + gl

    let expected = [a_0, a_1, a_2, a_3, a_4, a_5];
    for t_idx in 0..t {
        assert_close(a[t_idx], expected[t_idx],
                     &format!("A[0][{t_idx}]"));
    }

    // Stronger invariant: the post-terminal A_1 must NOT have absorbed
    // any geometric credit from A_3..A_5. If the kernel forgot the
    // reset, A_1 would equal `1 + gl·A_3` (much larger than 1).
    let a_1_unreset = 1.0_f32 + gl * a_3;
    assert!(
        (a[1] - a_1).abs() < (a_1_unreset - a_1).abs(),
        "A_1 reset invariant violated: kernel A_1={} closer to unreset {} than reset {}",
        a[1], a_1_unreset, a_1
    );
}

// ── Test 4: bit-equal determinism across fresh contexts ───────────────

/// Per `pearl_determinism_achieved`: the GAE backward-sweep kernel is
/// single-thread-per-batch + sequential, so identical inputs must
/// produce bit-equal outputs across runs.
#[test]
#[ignore = "requires CUDA"]
fn gae_deterministic_across_runs() {
    let Some(dev) = make_device() else { return };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    let ctx = dev.cuda_context().expect("cuda_context");

    let b = 4usize;
    let t = 16usize;
    let hidden_dim = 8usize;
    // Synthetic non-trivial input — mix of positive/negative rewards,
    // scattered terminals, varying V.
    let mut rewards = vec![0.0_f32; b * t];
    let mut dones = vec![0u8; b * t];
    let mut v_t = vec![0.0_f32; b * t];
    for b_idx in 0..b {
        for t_idx in 0..t {
            let idx = b_idx * t + t_idx;
            rewards[idx] = ((b_idx * 7 + t_idx * 3) as f32 / 11.0).sin();
            v_t[idx] = ((b_idx * 5 + t_idx * 2) as f32 / 13.0).cos() * 0.5;
            // Terminal at t=7 and t=15 (varying per batch).
            if (t_idx == 7 && b_idx == 1) || (t_idx == 15) {
                dones[idx] = 1;
            }
        }
    }
    let v_t_bootstrap: Vec<f32> = (0..b)
        .map(|b_idx| (b_idx as f32 + 1.0).ln() * 0.3)
        .collect();

    // Run 1.
    let mut rb1 = RolloutBuffer::new(&ctx, &stream, b, t, hidden_dim)
        .expect("RolloutBuffer::new run 1");
    let (a1, r1) = run_gae(&mut rb1, &stream, &rewards, &dones, &v_t, &v_t_bootstrap,
                            TEST_GAMMA, TEST_LAMBDA);

    // Run 2 — fresh buffer, same context. The kernel is pure (no PRNG
    // state, no time-dependent reads), so this is the relevant
    // determinism check for the new module.
    let mut rb2 = RolloutBuffer::new(&ctx, &stream, b, t, hidden_dim)
        .expect("RolloutBuffer::new run 2");
    let (a2, r2) = run_gae(&mut rb2, &stream, &rewards, &dones, &v_t, &v_t_bootstrap,
                            TEST_GAMMA, TEST_LAMBDA);

    assert_eq!(a1.len(), a2.len(), "advantages length mismatch");
    assert_eq!(r1.len(), r2.len(), "returns length mismatch");
    for i in 0..a1.len() {
        assert_eq!(
            a1[i].to_bits(),
            a2[i].to_bits(),
            "A[{i}] not bit-equal across runs: {} vs {}",
            a1[i], a2[i]
        );
        assert_eq!(
            r1[i].to_bits(),
            r2[i].to_bits(),
            "R[{i}] not bit-equal across runs: {} vs {}",
            r1[i], r2[i]
        );
    }
}

// ── Test 5: allocation sizes match the [B × T] / [B × T × H] / [B] spec ─

#[test]
#[ignore = "requires CUDA"]
fn rollout_buffer_alloc_sizes_match_spec() {
    let Some(dev) = make_device() else { return };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    let ctx = dev.cuda_context().expect("cuda_context");

    let b = 128usize;
    let t = 256usize;
    let h = 256usize;
    let rb = RolloutBuffer::new(&ctx, &stream, b, t, h).expect("RolloutBuffer::new");

    assert_eq!(rb.b_size, b, "b_size mismatch");
    assert_eq!(rb.t_rollout, t, "t_rollout mismatch");
    assert_eq!(rb.hidden_dim, h, "hidden_dim mismatch");

    assert_eq!(rb.rewards_d.len(), b * t, "rewards_d size");
    assert_eq!(rb.dones_d.len(), b * t, "dones_d size");
    assert_eq!(rb.v_t_d.len(), b * t, "v_t_d size");
    assert_eq!(rb.actions_d.len(), b * t, "actions_d size");
    assert_eq!(rb.log_pi_old_d.len(), b * t, "log_pi_old_d size");
    assert_eq!(rb.h_t_d.len(), b * t * h, "h_t_d size");
    assert_eq!(rb.v_t_bootstrap_d.len(), b, "v_t_bootstrap_d size");
    assert_eq!(rb.advantages_d.len(), b * t, "advantages_d size");
    assert_eq!(rb.returns_d.len(), b * t, "returns_d size");
    assert_eq!(rb.current_t, 0, "current_t initial");
}