From 5cd2f87039bfb153545f65f9ae8d1d5a534a0ce0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 2 Jun 2026 22:09:47 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20Phase=201B-A=20=E2=80=94=20Ro?= =?UTF-8?q?lloutBuffer=20+=20GAE=20kernel=20foundation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/gae_backward_sweep.cu | 58 +++ crates/ml-alpha/src/rl/isv_slots.rs | 33 +- crates/ml-alpha/src/trainer/mod.rs | 1 + crates/ml-alpha/src/trainer/rollout_buffer.rs | 210 ++++++++ .../tests/rollout_buffer_invariants.rs | 466 ++++++++++++++++++ 6 files changed, 768 insertions(+), 1 deletion(-) create mode 100644 crates/ml-alpha/cuda/gae_backward_sweep.cu create mode 100644 crates/ml-alpha/src/trainer/rollout_buffer.rs create mode 100644 crates/ml-alpha/tests/rollout_buffer_invariants.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 210ab2c57..22b74e726 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -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 diff --git a/crates/ml-alpha/cuda/gae_backward_sweep.cu b/crates/ml-alpha/cuda/gae_backward_sweep.cu new file mode 100644 index 000000000..ea224aa72 --- /dev/null +++ b/crates/ml-alpha/cuda/gae_backward_sweep.cu @@ -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 + +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 + } +} diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index b2fd82a62..e08cff8cf 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -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; diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index d8b1ee485..f22c8de42 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -7,3 +7,4 @@ pub mod loss; pub mod optim; pub mod perception; pub mod raw_launch; +pub mod rollout_buffer; diff --git a/crates/ml-alpha/src/trainer/rollout_buffer.rs b/crates/ml-alpha/src/trainer/rollout_buffer.rs new file mode 100644 index 000000000..c8c839b19 --- /dev/null +++ b/crates/ml-alpha/src/trainer/rollout_buffer.rs @@ -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` 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` +//! 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, + /// 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, + /// 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, + /// Per-step action index (output of action selection). ([B × T]) + pub actions_d: CudaSlice, + /// Per-step log π_old(a_t | s_t) snapshot. PPO ratio = exp(log π_new + /// − log π_old). ([B × T]) + pub log_pi_old_d: CudaSlice, + /// 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, + + /// 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, + + /// GAE advantages A_t (output of compute_gae). ([B × T]) + pub advantages_d: CudaSlice, + /// GAE returns R_t = A_t + V_t (output of compute_gae). ([B × T]) + pub returns_d: CudaSlice, + + /// GAE backward-sweep kernel handle. + gae_kernel: CudaFunction, + /// Module owning the kernel — kept alive for the kernel's lifetime. + _gae_module: Arc, + + /// 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, + stream: &Arc, + b_size: usize, + t_rollout: usize, + hidden_dim: usize, + ) -> Result { + 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::(bt).context("rewards_d alloc")?; + let dones_d = stream.alloc_zeros::(bt).context("dones_d alloc")?; + let v_t_d = stream.alloc_zeros::(bt).context("v_t_d alloc")?; + let actions_d = stream.alloc_zeros::(bt).context("actions_d alloc")?; + let log_pi_old_d = stream.alloc_zeros::(bt).context("log_pi_old_d alloc")?; + let h_t_d = stream.alloc_zeros::(bth).context("h_t_d alloc")?; + let v_t_bootstrap_d = stream.alloc_zeros::(b_size).context("v_t_bootstrap_d alloc")?; + let advantages_d = stream.alloc_zeros::(bt).context("advantages_d alloc")?; + let returns_d = stream.alloc_zeros::(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, + 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(()) + } +} diff --git a/crates/ml-alpha/tests/rollout_buffer_invariants.rs b/crates/ml-alpha/tests/rollout_buffer_invariants.rs new file mode 100644 index 000000000..f0d7bc7b9 --- /dev/null +++ b/crates/ml-alpha/tests/rollout_buffer_invariants.rs @@ -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` (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 { + 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` via mapped +/// pinned staging + DtoD. Sync the stream so callers observe the write. +fn upload_f32(stream: &Arc, host: &[f32]) -> CudaSlice { + let n = host.len(); + let staging = unsafe { MappedRecordBuffer::::new(n) }.expect("f32 staging alloc"); + staging.write_from_slice(host); + let mut d = stream.alloc_zeros::(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::(), + stream.cu_stream(), + ) + .expect("f32 DtoD upload"); + } + stream.synchronize().expect("upload_f32 sync"); + d +} + +/// Upload a `u8` host slice into a fresh `CudaSlice` via mapped +/// pinned staging + DtoD. Used for the dones buffer. +fn upload_u8(stream: &Arc, host: &[u8]) -> CudaSlice { + let n = host.len(); + let staging = unsafe { MappedRecordBuffer::::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::(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::(), + 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, src: &CudaSlice) -> Vec { + 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, + rewards: &[f32], + dones: &[u8], + v_t: &[f32], + v_t_bootstrap: &[f32], + gamma: f32, + lambda: f32, +) -> (Vec, Vec) { + 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 = (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"); +}