gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(ml-alpha): Phase 1B-A — RolloutBuffer + GAE kernel foundation

Browse Source 
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

Foundation phase of the trainer rollout-buffer + GAE refactor (Phase 1B).
Empirically required after Phase 1A regression confirmed math-agent's
atomicity claim: dropping Phase 5 shaping WITHOUT GAE makes things worse
(eval_pnl regressed -$691k from -$4.46M to -$5.15M, popart sigma CV
0.96 -> 1.58, sign agreement 60% -> 46%). See pearl_reward_signal_
anti_aligned_with_pnl ADDENDUM 2026-06-02d for the mechanism analysis.

This commit establishes the components WITHOUT trainer integration —
subsequent phases (1B-B through 1B-E) wire them into the actual training
loop. Subdivides the multi-week refactor into atomically-verifiable
phases per feedback_investigation_first_falsification_methodology.

New components:

* 3 ISV slots (758-760):
    RL_PPO_ROLLOUT_HORIZON_INDEX = 758 (T_rollout, bootstrap 256)
    RL_PPO_N_EPOCHS_INDEX        = 759 (K_ppo, bootstrap 4)
    RL_PPO_N_MINIBATCHES_INDEX   = 760 (minibatches/epoch, bootstrap 8)
    RL_SLOTS_END 757 -> 761

* crates/ml-alpha/cuda/gae_backward_sweep.cu (58 LOC):
    Single-thread-per-batch sequential backward sweep computing
    A_t = δ_t + γλ·A_{t+1}·(1-done), returns_t = A_t + V_t.
    Deterministic by construction (no parallel reductions, no atomicAdd,
    no nvrtc). Reset on done. v_T_bootstrap parameter for trajectory-end
    V estimate.

* crates/ml-alpha/src/trainer/rollout_buffer.rs (210 LOC):
    RolloutBuffer struct with [B × T] device buffers for rewards, dones,
    v_t, actions, log_pi_old, h_t; plus [B] v_T_bootstrap; plus output
    advantages, returns. compute_gae(γ, λ) invokes the kernel using
    cudarc raw-pointer pattern (`device_ptr(stream).0` resolved into
    local before .arg() — idiomatic for codebase). Loaded module +
    kernel handle stay private; struct fields exposed per spec.

* crates/ml-alpha/tests/rollout_buffer_invariants.rs (466 LOC):
    5 GPU-oracle invariant tests (#[ignore = "requires CUDA"]):
      1. gae_terminal_only_matches_close_event_pnl — single done at T-1
      2. gae_dense_reward_geometric_decay — Σ(γλ)^k geometric series
      3. gae_done_resets_credit — done reset propagation
      4. gae_deterministic_across_runs — bit-equality across contexts
      5. rollout_buffer_alloc_sizes_match_spec — alloc verification
    All 5 PASS in 2.15s.

Validation gates (Phase 1B-A complete when all pass):
* cargo build --release --example alpha_rl_train -p ml-alpha: exit 0
  (1m 00s release build, gae_backward_sweep.cubin compiled with -O3
  --use_fast_math --ftz=true --fmad=true for sm_86)
* cargo test -p ml-alpha --test rollout_buffer_invariants --release: 5/5 PASS
* ./scripts/determinism-check.sh --quick: exit 0 — DETERMINISTIC: all
  checksums.* leaves match across all 200 rows (rel-tol=1e-5, abs-tol=1e-7).
  Pipeline output bit-equal to baseline since new struct/kernel are
  loaded but unused in the training loop.
* Pre-commit hook on staged diff: PASS (0 memcpy_dtoh, 0 atomicAdd).

Next: Phase 1B-B (rollout collection loop behind FOXHUNT_USE_ROLLOUT env
flag) per plan §Phase 1B-B. Dispatch after this commit lands.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-06-02 22:09:47 +02:00
parent cfaa420bd4
commit 5cd2f87039
6 changed files with 768 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

@@ -51,6 +51,7 @@ const KERNELS: &[&str] = &[
"ema_update_on_done", // RL Phase R3: generic done-gated EMA producer (slot-parameterised) for closed-trade-magnitude EMAs (mean_abs_pnl, q_divergence, td_kurtosis)
"ema_update_per_step", // RL Phase R3: generic per-step EMA producer (slot-parameterised) for continuous EMAs (kl_pi, entropy_observed, advantage_var_ratio, trade_duration)
"compute_advantage_return", // RL Phase R3: element-wise A_t = r + γ(1-done)·V(s_{t+1}) − V(s_t), R_t = r + γ(1-done)·V(s_{t+1}); reads γ from ISV[400]
"gae_backward_sweep", // Phase 1B-A (2026-06-02): GAE backward sweep over [B × T] rollout trajectories — A_t = δ_t + γλ·A_{t+1}·non_terminal, returns_t = A_t + V_t; deterministic single-thread-per-batch sequential sweep; replaces compute_advantage_return when wired in Phase 1B-B+
"rl_action_kernel", // RL Phase R4: Thompson sampler over C51 atoms; one block per batch, N_ACTIONS threads; per-batch xorshift32 PRNG state; replaces host Thompson loop per feedback_cpu_is_read_only
"argmax_expected_q", // RL Phase R4: argmax over expected Q per action; Bellman-target argmax (Double-DQN); pairs with rl_action_kernel per pearl_thompson_for_distributional_action_selection
"log_pi_at_action", // RL Phase R4: per-batch log π(action_b) via log-softmax + lookup; PPO importance-ratio path

58
crates/ml-alpha/cuda/gae_backward_sweep.cu Normal file
View File

@@ -0,0 +1,58 @@
// gae_backward_sweep.cu — Generalized Advantage Estimation backward sweep.
//
// Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.3
// Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md
//
// Per batch b, sequentially compute (backward over T):
// A_T = 0
// for t in (T-1)..=0:
// non_terminal = 1.0f - (float)dones[b][t]
// v_tp1 = (t == T-1) ? v_T_bootstrap[b] : v_t[b][t+1]
// δ_t = r[b][t] + γ * v_tp1 * non_terminal - v_t[b][t]
// A_t = δ_t + γ * λ * A_next * non_terminal
// A_next = A_t * non_terminal // reset GAE accumulator on terminal
// returns[b][t] = A_t + v_t[b][t]
//
// Single-thread-per-batch kernel. Sequential backward sweep guarantees
// determinism (no parallel reductions, no inter-thread/block races).
// Memory access is strided over T for each thread — coalescing is
// suboptimal but functional; the rollout T (≤1024) and B (≤1024) keep
// this off the hot path. If profiling later shows it dominating, the
// follow-up is a warp-cooperative tiled implementation that preserves
// determinism via a single-warp reduction.
//
// Per `feedback_no_atomicadd`: no atomic ops anywhere.
// Per `feedback_cpu_is_read_only`: pure device-side, no host roundtrips.
// Per `feedback_no_nvrtc`: pre-compiled cubin (build.rs registers it).
#include <stdint.h>
extern "C" __global__ void gae_backward_sweep(
const float* __restrict__ rewards, // [B × T]
const uint8_t* __restrict__ dones, // [B × T]
const float* __restrict__ v_t, // [B × T]
const float* __restrict__ v_T_bootstrap, // [B] — V(s_T) for trajectory end
float* __restrict__ advantages, // [B × T] (output)
float* __restrict__ returns, // [B × T] (output)
const int B,
const int T,
const float gamma,
const float lambda
) {
const int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
float a_next = 0.0f;
for (int t = T - 1; t >= 0; t--) {
const int idx = b * T + t;
const float r_t = rewards[idx];
const float v_curr = v_t[idx];
const float non_terminal = 1.0f - (float)dones[idx];
const float v_tp1 = (t == T - 1) ? v_T_bootstrap[b] : v_t[idx + 1];
const float delta = r_t + gamma * v_tp1 * non_terminal - v_curr;
const float a_t = delta + gamma * lambda * a_next * non_terminal;
advantages[idx] = a_t;
returns[idx] = a_t + v_curr;
a_next = a_t * non_terminal; // reset on done
}
}

33
crates/ml-alpha/src/rl/isv_slots.rs
View File

@@ -1778,6 +1778,36 @@ pub const RL_SURFER_BREAK_EVEN_WR_INDEX: usize = 754; // default 0.3
pub const RL_SURFER_K_SHARPNESS_INDEX: usize = 755; // default 30.0 (sigmoid slope around break-even)
pub const RL_SURFER_WARMUP_TRADES_INDEX: usize = 756; // default 200.0 (cumulative dones before confidence ≈ 1)
// ─────────────────────────────────────────────────────────────────────
// Phase 1B (2026-06-02) — Trainer rollout-buffer + GAE refactor.
// Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md
//
// These slots are foundation-only in Phase 1B-A (allocated, bootstrap
// values defined). Phase 1B-B+ wires them into the trainer's rollout
// collection + multi-epoch PPO update path. Until then the rollout
// path is unconditional infrastructure that is not yet invoked from
// the step loop.
// ─────────────────────────────────────────────────────────────────────
/// T_rollout — number of env steps the collection loop runs before
/// invoking GAE + training pass. Bootstrap 256 (PPO/GAE standard).
/// Higher → more independent advantage observations per epoch, lower →
/// smaller rollout buffer (less GPU memory). Must divide evenly into
/// `RL_PPO_N_MINIBATCHES_INDEX` so minibatches partition the rollout.
/// Clamp range [32, 1024].
pub const RL_PPO_ROLLOUT_HORIZON_INDEX: usize = 758;
/// K_ppo — number of PPO update epochs per rollout. Bootstrap 4
/// (Schulman 2017 default range 4-10). Higher → more data efficiency
/// per env step, but encoder frozen during all epochs (stale features
/// risk). Clamp range [1, 16].
pub const RL_PPO_N_EPOCHS_INDEX: usize = 759;
/// Minibatches per epoch — divide rollout (B × T) into this many
/// shuffled minibatches per epoch. Bootstrap 8 (standard PPO default).
/// Each minibatch gets one Adam step. Clamp range [1, 64].
pub const RL_PPO_N_MINIBATCHES_INDEX: usize = 760;
/// Last RL-allocated slot index (exclusive).
/// Pre-risk-stack: 662. Post-Fix-B: 685. Post-v9 (warmup boundary): 696.
/// Post-regime-observer (F1.1): 716. Post-warmed-flag addendum: 717.
@@ -1791,4 +1821,5 @@ pub const RL_SURFER_WARMUP_TRADES_INDEX: usize = 756; // default 200
/// Post-edge-decay-detector Phase 1: 753.
/// Post-reward-policy-realignment v4 binary: 754 (slot reused in v5).
/// Post-surfer-scaffold v5 adaptive (3 config slots): 757.
pub const RL_SLOTS_END: usize = 757;
/// Post-rollout-buffer+GAE Phase 1B-A (3 PPO control slots): 760.
pub const RL_SLOTS_END: usize = 761;

1
crates/ml-alpha/src/trainer/mod.rs
View File

@@ -7,3 +7,4 @@ pub mod loss;
pub mod optim;
pub mod perception;
pub mod raw_launch;
pub mod rollout_buffer;

210
crates/ml-alpha/src/trainer/rollout_buffer.rs Normal file
View File

@@ -0,0 +1,210 @@
//! Rollout buffer for on-policy PPO with GAE.
//!
//! Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.1
//! Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md
//!
//! Stores `[B × T_rollout]` trajectories of (rewards, dones, v_t, actions,
//! log_pi_old, h_t encoder snapshot) collected by the rollout-collection
//! loop. The trainer then invokes `compute_gae` (wraps the
//! `gae_backward_sweep` kernel) to compute advantages + returns, and runs
//! multi-epoch minibatch updates over the rollout.
//!
//! Phase 1B-A scope: standalone infrastructure. The struct is allocated,
//! the GAE kernel is loaded, and GPU-oracle invariant tests cover the
//! math. **No trainer integration yet** — wiring into `IntegratedTrainer`
//! lands in Phase 1B-B+.
//!
//! ## Determinism contract
//!
//! Per `pearl_determinism_achieved`: the GAE backward-sweep kernel is
//! single-thread-per-batch with a sequential reverse-T loop, so its
//! outputs are bit-equal across runs with identical inputs. No atomic
//! ops, no parallel reductions inside a trajectory, no inter-block
//! synchronisation needed.
//!
//! ## Constraints honoured
//!
//! * `feedback_no_atomicadd` — kernel has zero atomic operations.
//! * `feedback_cpu_is_read_only` — buffers live on the device; CPU
//! reads only via mapped-pinned `read_slice_d_into<T>` from outside
//! this module.
//! * `feedback_no_htod_htoh_only_mapped_pinned` — no raw `memcpy_*` in
//! this module. All CPU↔GPU paths go through `read_slice_d_into<T>`
//! in the test harness.
//! * `feedback_no_nvrtc` — pre-compiled cubin via `build.rs`.
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream,
DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg,
};
const GAE_BACKWARD_SWEEP_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/gae_backward_sweep.cubin"));
/// PPO rollout buffer + GAE backward-sweep kernel handle.
///
/// All device buffers are sized `[B × T_rollout]` (or `[B]` for the
/// bootstrap V(s_T) vector, or `[B × T × H]` for the encoder snapshot).
/// The collection loop fills the per-step buffers; `compute_gae` reads
/// rewards/dones/v_t/v_T_bootstrap and writes advantages/returns.
pub struct RolloutBuffer {
/// Batch dimension (parallel envs).
pub b_size: usize,
/// Trajectory length T_rollout.
pub t_rollout: usize,
/// Encoder hidden dim (size of `h_t` snapshot per env-step).
pub hidden_dim: usize,
/// Per-step reward `r_t` ([B × T]).
pub rewards_d: CudaSlice<f32>,
/// Per-step done flag `1.0` if terminal at step t else `0.0`, encoded
/// as `u8` to keep the GAE kernel's `non_terminal = 1.0f - (float)d`
/// branch-free and 4× smaller than `f32`. ([B × T])
pub dones_d: CudaSlice<u8>,
/// Per-step value prediction `V(s_t)` snapshot at collection time.
/// Frozen during the multi-epoch PPO update so the advantage targets
/// stay consistent across epochs. ([B × T])
pub v_t_d: CudaSlice<f32>,
/// Per-step action index (output of action selection). ([B × T])
pub actions_d: CudaSlice<i32>,
/// Per-step log π_old(a_t | s_t) snapshot. PPO ratio = exp(log π_new
/// − log π_old). ([B × T])
pub log_pi_old_d: CudaSlice<f32>,
/// Per-step encoder hidden state `h_t`. Multi-epoch PPO re-evaluates
/// the heads on this snapshot rather than re-running the encoder.
/// ([B × T × HIDDEN_DIM])
pub h_t_d: CudaSlice<f32>,
/// Trajectory-end bootstrap V(s_T) — the value prediction at the
/// step immediately after the last collected step. ([B])
pub v_t_bootstrap_d: CudaSlice<f32>,
/// GAE advantages A_t (output of compute_gae). ([B × T])
pub advantages_d: CudaSlice<f32>,
/// GAE returns R_t = A_t + V_t (output of compute_gae). ([B × T])
pub returns_d: CudaSlice<f32>,
/// GAE backward-sweep kernel handle.
gae_kernel: CudaFunction,
/// Module owning the kernel — kept alive for the kernel's lifetime.
_gae_module: Arc<CudaModule>,
/// Write cursor for the collection loop (Phase 1B-B+). Reset at the
/// start of each rollout and advanced by `push_step()` once that
/// path is wired.
pub current_t: usize,
}
impl RolloutBuffer {
/// Allocate all device buffers and load the GAE kernel.
///
/// Buffer sizes are fixed at construction; reuse across rollouts is
/// the expected pattern (no per-rollout cudaMalloc on the hot path).
pub fn new(
ctx: &Arc<CudaContext>,
stream: &Arc<CudaStream>,
b_size: usize,
t_rollout: usize,
hidden_dim: usize,
) -> Result<Self> {
let bt = b_size.checked_mul(t_rollout)
.context("rollout buffer B × T overflow")?;
let bth = bt.checked_mul(hidden_dim)
.context("rollout buffer B × T × HIDDEN_DIM overflow")?;
let rewards_d = stream.alloc_zeros::<f32>(bt).context("rewards_d alloc")?;
let dones_d = stream.alloc_zeros::<u8>(bt).context("dones_d alloc")?;
let v_t_d = stream.alloc_zeros::<f32>(bt).context("v_t_d alloc")?;
let actions_d = stream.alloc_zeros::<i32>(bt).context("actions_d alloc")?;
let log_pi_old_d = stream.alloc_zeros::<f32>(bt).context("log_pi_old_d alloc")?;
let h_t_d = stream.alloc_zeros::<f32>(bth).context("h_t_d alloc")?;
let v_t_bootstrap_d = stream.alloc_zeros::<f32>(b_size).context("v_t_bootstrap_d alloc")?;
let advantages_d = stream.alloc_zeros::<f32>(bt).context("advantages_d alloc")?;
let returns_d = stream.alloc_zeros::<f32>(bt).context("returns_d alloc")?;
let gae_module = ctx
.load_cubin(GAE_BACKWARD_SWEEP_CUBIN.to_vec())
.context("load gae_backward_sweep cubin")?;
let gae_kernel = gae_module
.load_function("gae_backward_sweep")
.context("load gae_backward_sweep function")?;
Ok(Self {
b_size,
t_rollout,
hidden_dim,
rewards_d,
dones_d,
v_t_d,
actions_d,
log_pi_old_d,
h_t_d,
v_t_bootstrap_d,
advantages_d,
returns_d,
gae_kernel,
_gae_module: gae_module,
current_t: 0,
})
}
/// Reset the write cursor. Called at the start of each rollout once
/// the collection path is wired (Phase 1B-B+).
pub fn reset(&mut self) {
self.current_t = 0;
}
/// Invoke the GAE backward-sweep kernel. Reads `rewards_d`,
/// `dones_d`, `v_t_d`, `v_t_bootstrap_d`; writes `advantages_d`,
/// `returns_d`.
///
/// Determinism: single thread per batch with a sequential reverse-T
/// loop, so the output is bit-equal across runs for identical inputs
/// per `pearl_determinism_achieved`.
pub fn compute_gae(
&mut self,
stream: &Arc<CudaStream>,
gamma: f32,
lambda: f32,
) -> Result<()> {
let b = self.b_size as i32;
let t = self.t_rollout as i32;
let threads: u32 = 256;
let blocks: u32 = ((self.b_size as u32) + threads - 1) / threads;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes: 0,
};
// Resolve device pointers up-front so the launch_builder borrows
// are scoped narrowly (the borrow guards drop after `.launch()`).
let rewards_ptr = self.rewards_d.device_ptr(stream).0;
let dones_ptr = self.dones_d.device_ptr(stream).0;
let v_t_ptr = self.v_t_d.device_ptr(stream).0;
let v_t_bootstrap_ptr = self.v_t_bootstrap_d.device_ptr(stream).0;
let advantages_ptr = self.advantages_d.device_ptr_mut(stream).0;
let returns_ptr = self.returns_d.device_ptr_mut(stream).0;
unsafe {
stream
.launch_builder(&self.gae_kernel)
.arg(&rewards_ptr)
.arg(&dones_ptr)
.arg(&v_t_ptr)
.arg(&v_t_bootstrap_ptr)
.arg(&advantages_ptr)
.arg(&returns_ptr)
.arg(&b)
.arg(&t)
.arg(&gamma)
.arg(&lambda)
.launch(cfg)
.context("gae_backward_sweep launch")?;
}
Ok(())
}
}

466
crates/ml-alpha/tests/rollout_buffer_invariants.rs Normal file
View File

@@ -0,0 +1,466 @@
//! 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");
}