From 779c03b9d76149aa49c484e2a6dc7d12f680c0ba Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 2 Jun 2026 22:42:39 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20Phase=201B-B=20=E2=80=94=20ro?= =?UTF-8?q?llout=20collection=20loop=20+=20GAE=20behind=20FOXHUNT=5FUSE=5F?= =?UTF-8?q?ROLLOUT=20flag?= 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 §Phase 1B-B Builds on commit 5cd2f8703 (1B-A foundation). Wires the RolloutBuffer into the per-step training loop via a new RolloutCollection helper, gated behind FOXHUNT_USE_ROLLOUT=1 env flag. Mode 0 (default unset) is bit-equal to HEAD; Mode 1 (flag=1) snapshots per-step data into the rollout buffer and invokes GAE every T_rollout steps. This phase is a STRUCTURAL validation, not an empirical pnl test. Phase 1B-C will replace the per-step training in Mode 1 with multi- epoch PPO over the rollout buffer. New components: * crates/ml-alpha/cuda/rollout_pack.cu (42 LOC): Single-thread-per-batch deterministic f32 -> u8 packer for dones. Used because RolloutBuffer.dones_d is u8 (per spec for memory) but the trainer's dones_d is f32 (per existing pipeline). * crates/ml-alpha/src/trainer/rollout_collection.rs (375 LOC): RolloutCollection helper with three methods: - new(ctx): loads rollout_pack cubin - snapshot_step(trainer, buf, t, b_size, t_max): DtoD scatter via cuMemcpy2DAsync_v2 for rewards/v_t/actions/log_pi/h_t at offset b*t_max+t in the [B × T] buffer; rollout_pack kernel for dones - copy_bootstrap_v(trainer, buf, b_size): DtoD V(s_T) from trainer.v_pred_tp1_d into buf.v_t_bootstrap_d at rollout end All transfers use mapped-pinned / DtoD only — no raw memcpy_dtoh on regular slices (`feedback_no_htod_htoh_only_mapped_pinned`). * crates/ml-alpha/examples/alpha_rl_train.rs (+137 LOC at lines 59-66, 410-466, 591-666): - FOXHUNT_USE_ROLLOUT env-flag parse at startup - Conditional RolloutBuffer + RolloutCollection allocation - Per-step snapshot_step() after step_with_lobsim_gpu - Every T_rollout steps: copy_bootstrap_v + compute_gae + stderr log `[rollout] step=N GAE complete (T=…, γ=…, λ=0.95): adv_mean=… adv_abs_mean=… returns_mean=…` Adaptations from the dispatch (documented inline): 1. T_rollout fallback to 32 when slot 758 is unbootstrapped. Slot 758 (RL_PPO_ROLLOUT_HORIZON_INDEX) was allocated in 1B-A but no kernel bootstraps it — `read_isv_host(758)` returns 0.0. Binary uses `.clamp(32, 1024)`, flooring to T=32 for the smoke. Phase 1B-C will add the trainer bootstrap entry (or a controller kernel) to make T_rollout=256 the production default. 2. Mode 1 still runs per-step training as a side effect; the rollout buffer is populated as a parallel snapshot. The "stop training in Mode 1" separation would require a ~1500 LOC refactor of step_with_lobsim_gpu (single-step/training-coupled per pearl_foxhunt_trainer_is_genuinely_single_step). The dispatch explicitly authorized this "skeleton with correct semantics" adaptation — the load-bearing 1B-B gate is "collection + GAE pipeline works structurally", which is verified. Phase 1B-C will replace the per-step training with multi-epoch PPO over the rollout buffer. Validation (all gates PASS): * cargo build --release --example alpha_rl_train -p ml-alpha: exit 0 (~60s) * cargo test rollout_buffer_invariants --release: 5/5 PASS (no regression) * Mode 0 (flag unset, default): ./scripts/determinism-check.sh --quick exit 0 - DETERMINISTIC: all checksums.* leaves match across all 200 rows - pipeline output bit-equal to HEAD 5cd2f8703 baseline * Mode 1 (FOXHUNT_USE_ROLLOUT=1): smoke completes in 2m 53s - 500 train + 100 eval steps, exit 0, completed_clean=true - 16 GAE windows logged (T=32 each) - adv_abs_mean range [4.4e-3, 1.92e-1] — always > 1e-3 sanity gate - eval_summary written: total_pnl_usd=-$101,487.70, n_trades=301, wr=0.346 * Cross-source consistency BOTH modes within $50 (Mode 0: $0.00 diff; Mode 1: $0.01 diff). Phase 0 contract preserved. * Pre-commit hook PASS (0 memcpy_dtoh raw violations, 0 atomicAdd). Next: Phase 1B-C will (a) bootstrap slot 758 to production T_rollout=256, (b) replace Mode 1's per-step training with multi-epoch PPO over the rollout buffer using the existing PPO surrogate + V regression kernels operating on minibatches of [B × T] flattened to [B*T]. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/rollout_pack.cu | 42 ++ crates/ml-alpha/examples/alpha_rl_train.rs | 137 ++++++- crates/ml-alpha/src/trainer/mod.rs | 1 + .../src/trainer/rollout_collection.rs | 375 ++++++++++++++++++ 5 files changed, 554 insertions(+), 2 deletions(-) create mode 100644 crates/ml-alpha/cuda/rollout_pack.cu create mode 100644 crates/ml-alpha/src/trainer/rollout_collection.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 22b74e726..871264a09 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -52,6 +52,7 @@ const KERNELS: &[&str] = &[ "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+ + "rollout_pack", // Phase 1B-B (2026-06-02): per-step f32→u8 dones packer; writes `dones_f32 [B]` into rollout buffer's `dones_u8_bt [B × T]` at offset `b * T + t_cursor`; single-thread-per-batch, deterministic "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/rollout_pack.cu b/crates/ml-alpha/cuda/rollout_pack.cu new file mode 100644 index 000000000..49a5ab32f --- /dev/null +++ b/crates/ml-alpha/cuda/rollout_pack.cu @@ -0,0 +1,42 @@ +// rollout_pack.cu — Per-step pack helpers for Phase 1B-B rollout collection. +// +// Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.2 +// Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B +// +// The trainer's per-step `dones_d` is `f32` in {0.0, 1.0} (legacy contract +// driven by `compute_advantage_return` and the K-loop sampled buffers). +// The rollout buffer's `dones_d` is `u8` (compact storage — 4× smaller, +// matches the `gae_backward_sweep` kernel's `uint8_t*` input). This kernel +// packs a per-step `[B]` f32 done vector into the rollout buffer's +// `[B × T]` u8 slice at offset `bt = b * T + t` for the current write +// cursor `t`. +// +// Single-thread-per-batch kernel (one thread per env). Deterministic by +// construction — no atomic ops, no reductions, no shared memory. The +// fmaxf clamp + ≥0.5 thresholding makes the conversion robust to fp32 +// noise (the trainer never produces non-{0, 1} dones, but the pack is +// defensive). +// +// Determinism contract: each thread writes one distinct element, so the +// kernel is bit-equal across runs for identical inputs. Compatible with +// `pearl_determinism_achieved`. +// +// Per `feedback_no_atomicadd`: no atomic ops. +// Per `feedback_cpu_is_read_only`: pure device-side, no host roundtrip. +// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs. + +#include + +extern "C" __global__ void rollout_pack_dones_f32_to_u8( + const float* __restrict__ dones_f32, // [B] — per-step trainer dones (f32 in {0, 1}) + uint8_t* __restrict__ dones_u8_bt, // [B × T] — rollout buffer dones (u8) + const int B, + const int T, + const int t_cursor // current write cursor in [0, T) +) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + const int idx = b * T + t_cursor; + const float d = dones_f32[b]; + dones_u8_bt[idx] = (d >= 0.5f) ? (uint8_t)1 : (uint8_t)0; +} diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 830dde07e..cf8831165 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -58,14 +58,17 @@ use ml_alpha::heads::HORIZONS; // the per-N-step stderr ticker. use ml_alpha::rl::isv_slots::{ RL_CONF_GATE_FIRED_COUNT_INDEX, RL_FRD_GATE_FIRED_COUNT_INDEX, RL_GAMMA_INDEX, - RL_HEAT_CAP_FIRED_COUNT_INDEX, RL_PER_ALPHA_INDEX, RL_PPO_CLIP_INDEX, RL_PYRAMID_ADD_COUNT_INDEX, - RL_REWARD_SCALE_INDEX, + RL_HEAT_CAP_FIRED_COUNT_INDEX, RL_PER_ALPHA_INDEX, RL_PPO_CLIP_INDEX, + RL_PPO_ROLLOUT_HORIZON_INDEX, RL_PYRAMID_ADD_COUNT_INDEX, RL_REWARD_SCALE_INDEX, }; +use ml_alpha::heads::HIDDEN_DIM as ROLLOUT_HIDDEN_DIM; use ml_alpha::trainer::diag_staging::DiagStaging; use ml_alpha::trainer::integrated::{ read_slice_d_into, DiagInputs, IntegratedStepStats, IntegratedTrainer, IntegratedTrainerConfig, }; use ml_alpha::trainer::perception::PerceptionTrainerConfig; +use ml_alpha::trainer::rollout_buffer::RolloutBuffer; +use ml_alpha::trainer::rollout_collection::RolloutCollection; use ml_backtesting::sim::LobSimCuda; use ml_core::device::MlDevice; use serde::Serialize; @@ -420,6 +423,70 @@ fn main() -> Result<()> { cli.seq_len, cli.n_backtests, cli.per_capacity ); + // ── Phase 1B-B rollout collection gate (FOXHUNT_USE_ROLLOUT). ──── + // + // Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.2 + // Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B + // + // TEMPORARY VALIDATION FLAG — per `feedback_no_feature_flags`, + // permanent feature flags are forbidden. This env toggle is the + // same pattern as `FOXHUNT_DETERMINISTIC` (Phase 2.3 determinism + // foundation). REMOVED in Phase 1B-E once rollout-buffer + GAE is + // the default path. + // + // Gate semantics: + // * `FOXHUNT_USE_ROLLOUT=0` (default, or unset): legacy path + // unchanged. RolloutBuffer NOT allocated. Bit-equal to HEAD. + // * `FOXHUNT_USE_ROLLOUT=1`: per-step snapshot into RolloutBuffer + // after each `step_with_lobsim_gpu`. Every T_rollout steps run + // GAE backward sweep + log advantage / returns stats to stderr. + // Training pipeline runs UNCHANGED (1B-B is structural validation + // of collection + GAE wiring, not training replacement — that's + // 1B-C). + let use_rollout: bool = std::env::var("FOXHUNT_USE_ROLLOUT") + .ok() + .and_then(|v| v.parse::().ok()) + .map(|n| n > 0) + .unwrap_or(false); + + let t_rollout_default: usize = { + let v = trainer.read_isv_host(RL_PPO_ROLLOUT_HORIZON_INDEX) as usize; + v.clamp(32, 1024) + }; + + let (mut rollout_buf, rollout_collection): (Option, Option) = + if use_rollout { + let ctx = dev + .cuda_context() + .context("FOXHUNT_USE_ROLLOUT=1: dev.cuda_context()")?; + let stream = trainer.stream.clone(); + let buf = RolloutBuffer::new( + ctx, + &stream, + cli.n_backtests, + t_rollout_default, + ROLLOUT_HIDDEN_DIM, + ) + .context("FOXHUNT_USE_ROLLOUT=1: RolloutBuffer::new")?; + let coll = RolloutCollection::new(ctx) + .context("FOXHUNT_USE_ROLLOUT=1: RolloutCollection::new")?; + eprintln!( + "[rollout] FOXHUNT_USE_ROLLOUT=1 → RolloutBuffer + GAE active \ + (B={} T={} HIDDEN_DIM={}, γ from ISV[{}], λ=0.95)", + cli.n_backtests, + t_rollout_default, + ROLLOUT_HIDDEN_DIM, + RL_GAMMA_INDEX, + ); + (Some(buf), Some(coll)) + } else { + eprintln!( + "[rollout] FOXHUNT_USE_ROLLOUT unset/=0 → legacy per-step path \ + (RolloutBuffer not allocated; bit-equal to HEAD)" + ); + (None, None) + }; + let mut summary = AlphaRlTrainSummary { n_steps_planned: cli.n_steps, seq_len: cli.seq_len, @@ -505,6 +572,10 @@ fn main() -> Result<()> { let mut heat_cap_total: u64 = 0; let t_start = std::time::Instant::now(); + // Phase 1B-B: rollout write-cursor (counts steps within the CURRENT + // T_rollout window; resets to 0 after each GAE invocation). Only + // advances when `use_rollout` is true. + let mut rollout_t: usize = 0; for step in 0..cli.n_steps { // GPU-resident data loading: sample_and_gather + gather_next + // gather_current + gather_frd_labels all run as GPU kernels @@ -513,6 +584,68 @@ fn main() -> Result<()> { .step_with_lobsim_gpu(&mut gpu_data_loader, &gpu_dataset, &mut sim) .with_context(|| format!("step_with_lobsim_gpu at step {step}"))?; + // ── Phase 1B-B rollout snapshot + GAE (gated). ─────────────── + // + // Stream-ordered after step_with_lobsim_gpu's writes (same + // trainer.stream). The snapshot is DtoD-only (cuMemcpy2DAsync + + // pack kernel) — zero host roundtrip, zero atomic ops. + // + // Every T_rollout steps: invoke GAE backward sweep, log the + // advantage / returns summary stats to stderr, then reset the + // write cursor and `current_t`. This is STRUCTURAL VALIDATION + // per the 1B-B dispatch — the resulting advantages_d / returns_d + // are not yet consumed by training (1B-C wires that in). + if let (Some(buf), Some(coll)) = (rollout_buf.as_mut(), rollout_collection.as_ref()) { + coll.snapshot_step( + &mut trainer, + buf, + rollout_t, + cli.n_backtests, + t_rollout_default, + ) + .with_context(|| { + format!("rollout snapshot_step at step={step} rollout_t={rollout_t}") + })?; + buf.current_t = rollout_t + 1; + rollout_t += 1; + + if rollout_t >= t_rollout_default { + // Bootstrap V(s_T) from the last step's v_pred_tp1_d + // (V on h_{t+1}), then run GAE backward sweep. + coll.copy_bootstrap_v(&mut trainer, buf, cli.n_backtests) + .context("rollout copy_bootstrap_v")?; + let gamma_isv = trainer.read_isv_host(RL_GAMMA_INDEX); + // λ = 0.95 hardcoded for 1B-B (1B-C allocates a slot). + buf.compute_gae(&trainer.stream, gamma_isv, 0.95_f32) + .context("rollout compute_gae")?; + + // Diagnostic: read advantages / returns means to verify + // the buffer was populated non-trivially. Uses + // `read_slice_d_into` (mapped-pinned, stream-ordered + // sync per `feedback_no_htod_htoh_only_mapped_pinned`). + let bt = cli.n_backtests * t_rollout_default; + let mut adv_host: Vec = vec![0.0; bt]; + let mut ret_host: Vec = vec![0.0; bt]; + read_slice_d_into(&trainer.stream, &buf.advantages_d, &mut adv_host) + .context("rollout read advantages_d")?; + read_slice_d_into(&trainer.stream, &buf.returns_d, &mut ret_host) + .context("rollout read returns_d")?; + let adv_mean: f64 = + adv_host.iter().map(|&x| x as f64).sum::() / bt as f64; + let adv_abs_mean: f64 = + adv_host.iter().map(|&x| (x as f64).abs()).sum::() / bt as f64; + let ret_mean: f64 = + ret_host.iter().map(|&x| x as f64).sum::() / bt as f64; + eprintln!( + "[rollout] step={step} GAE complete (T={t_rollout_default}, γ={gamma_isv:.4}, λ=0.95): \ + adv_mean={adv_mean:.4e} adv_abs_mean={adv_abs_mean:.4e} returns_mean={ret_mean:.4e}" + ); + + rollout_t = 0; + buf.reset(); + } + } + // Gate G8: NaN abort. Per `feedback_stop_on_anomaly` + // `feedback_kill_runs_on_anomaly_quickly` the cluster smoke // tail-watcher kills the workflow on non-zero exit; here we diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index f22c8de42..d4f269b27 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -8,3 +8,4 @@ pub mod optim; pub mod perception; pub mod raw_launch; pub mod rollout_buffer; +pub mod rollout_collection; diff --git a/crates/ml-alpha/src/trainer/rollout_collection.rs b/crates/ml-alpha/src/trainer/rollout_collection.rs new file mode 100644 index 000000000..de40a8948 --- /dev/null +++ b/crates/ml-alpha/src/trainer/rollout_collection.rs @@ -0,0 +1,375 @@ +//! Phase 1B-B: T-step env collection populating the RolloutBuffer. +//! +//! Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.2 +//! Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B +//! +//! Runs `T_rollout` env steps via the existing `step_with_lobsim_gpu` +//! pipeline and snapshots per-step `(h_t, a_t, log_pi_old_t, V(s_t), +//! r_t, done_t)` into the `RolloutBuffer` at offset `b * T + t`. After +//! the loop completes, copies `V(s_{t+1})` from the last step's +//! `v_pred_tp1_d` into `v_t_bootstrap_d` (the GAE boundary condition). +//! Caller then invokes `RolloutBuffer::compute_gae(γ, λ)`. +//! +//! ## STRUCTURAL note — Phase 1B-B vs 1B-C split +//! +//! `step_with_lobsim_gpu` is single-step / training-coupled per +//! `pearl_foxhunt_trainer_is_genuinely_single_step` — it does +//! forward + env step + PER replay training + Adam in a single call. +//! There is no clean "forward-only" path without a ~1500 LOC refactor. +//! Per `feedback_investigation_first_falsification_methodology` STOP-on- +//! unexpected-finding rule, we do not refactor in 1B-B. +//! +//! Consequently: +//! * 1B-B (this module): training STILL runs per env-step (as it does +//! today). The rollout buffer is populated as a SIDE EFFECT after +//! each `step_with_lobsim_gpu` call. GAE runs at the end. +//! This validates the COLLECTION + GAE pipeline structurally. +//! * 1B-C (future): adds multi-epoch PPO update over the rollout +//! buffer. THAT is where the per-step training is split — multi- +//! epoch PPO replaces the per-step PPO update path. +//! +//! Phase 1B-B's falsification therefore validates: +//! 1. Rollout buffer fills correctly (advantages/returns non-trivial) +//! 2. GAE backward sweep produces sensible distributions +//! 3. Default `FOXHUNT_USE_ROLLOUT=0` is BIT-EQUAL to HEAD +//! 4. `FOXHUNT_USE_ROLLOUT=1` smoke completes without NaN +//! +//! Per `feedback_no_atomicadd`: no atomic ops. +//! Per `feedback_cpu_is_read_only`: pure device-side population. +//! Per `feedback_no_htod_htoh_only_mapped_pinned`: all CPU↔GPU paths via +//! mapped-pinned helpers; this module uses DtoD-only via the existing +//! `raw_memcpy_dtod_async` (same pattern as the trainer hot path). + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaContext, CudaFunction, CudaModule, DevicePtr, DevicePtrMut, LaunchConfig, + PushKernelArg, +}; + +use crate::heads::HIDDEN_DIM; +use crate::rl::reward::RlLobBackend; +use crate::trainer::integrated::{IntegratedStepStats, IntegratedTrainer}; +use crate::trainer::rollout_buffer::RolloutBuffer; + +const ROLLOUT_PACK_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rollout_pack.cubin")); + +/// Helpers for the rollout-collection path. Holds the per-step +/// `rollout_pack_dones_f32_to_u8` kernel handle (loaded once at +/// construction) and exposes `collect()` as the public entry point. +pub struct RolloutCollection { + /// Per-step `dones_f32 [B] → dones_u8_bt [B × T]` packer. + pack_dones_kernel: CudaFunction, + /// Module owning the pack kernel — kept alive for the kernel's lifetime. + _pack_module: Arc, +} + +impl RolloutCollection { + /// Load the rollout-pack kernel cubin. Allocates no device memory of + /// its own; the `RolloutBuffer` is the data owner. + pub fn new(ctx: &Arc) -> Result { + let pack_module = ctx + .load_cubin(ROLLOUT_PACK_CUBIN.to_vec()) + .context("load rollout_pack cubin")?; + let pack_dones_kernel = pack_module + .load_function("rollout_pack_dones_f32_to_u8") + .context("load rollout_pack_dones_f32_to_u8")?; + Ok(Self { + pack_dones_kernel, + _pack_module: pack_module, + }) + } + + /// Run `t_rollout` env steps via `step_with_lobsim_gpu`, snapshotting + /// per-step state into `buf`. After this returns, `buf.v_t_bootstrap_d` + /// holds `V(s_T)` and the caller should invoke + /// `buf.compute_gae(stream, γ, λ)`. + /// + /// Returns the final step's `IntegratedStepStats` so the caller can + /// drive its existing diag / NaN-abort logic without restructuring. + /// + /// The per-step training pipeline runs unchanged (single-step PPO + + /// PER replay) — see module docs for the 1B-B/1B-C split rationale. + pub fn collect( + &self, + trainer: &mut IntegratedTrainer, + gpu_loader: &mut crate::data::gpu_dataset::GpuDataLoader, + gpu_dataset: &crate::data::gpu_dataset::GpuDataset, + sim: &mut dyn RlLobBackend, + buf: &mut RolloutBuffer, + t_rollout: usize, + ) -> Result { + anyhow::ensure!( + t_rollout > 0 && t_rollout <= buf.t_rollout, + "RolloutCollection::collect: t_rollout={t_rollout} out of range \ + (buf.t_rollout={})", + buf.t_rollout + ); + anyhow::ensure!( + buf.hidden_dim == HIDDEN_DIM, + "RolloutCollection::collect: buf.hidden_dim={} != HIDDEN_DIM={}", + buf.hidden_dim, + HIDDEN_DIM + ); + anyhow::ensure!( + buf.b_size == trainer.perception.h_t_view().len() / HIDDEN_DIM, + "RolloutCollection::collect: buf.b_size={} != trainer h_t batch \ + (h_t.len={} / HIDDEN_DIM={})", + buf.b_size, + trainer.perception.h_t_view().len(), + HIDDEN_DIM + ); + + buf.reset(); + let b_size = buf.b_size; + let t_max = buf.t_rollout; + + let mut last_stats: Option = None; + for t in 0..t_rollout { + // Drive ONE env step via the existing single-step pipeline. + // This populates the trainer's per-step `[B]` buffers + // (rewards_d / dones_d / actions_d / log_pi_old_d / + // v_pred_d / v_pred_tp1_d) and `perception.h_t_view()` with + // values from this step's forward + env transition. + let stats = trainer + .step_with_lobsim_gpu(gpu_loader, gpu_dataset, sim) + .with_context(|| { + format!("RolloutCollection::collect: step_with_lobsim_gpu at t={t}") + })?; + + // Snapshot the per-step `[B]` buffers into the rollout + // buffer at offset `b * T + t` for each batch element. + // All copies are DtoD-async on the trainer's main stream, + // so they're stream-ordered after the step's writes — no + // cross-stream sync needed. + self.snapshot_step(trainer, buf, t, b_size, t_max)?; + + buf.current_t = t + 1; + last_stats = Some(stats); + } + + // Bootstrap V(s_T): the last step's `v_pred_tp1_d` IS V(s_{t+1}) + // for the final action — exactly the GAE bootstrap target. Copy + // `[B]` f32 DtoD into `buf.v_t_bootstrap_d`. + self.copy_bootstrap_v(trainer, buf, b_size)?; + + Ok(last_stats.expect("t_rollout > 0 guarantees ≥1 iteration")) + } + + /// DtoD-snapshot the trainer's per-step `[B]` buffers into the + /// rollout buffer's `[B × T]` slots at write cursor `t`. + /// + /// Each copy is `B × sizeof(T)` bytes on the trainer's main stream. + /// Stream-ordered after the producing `step_with_lobsim_gpu` writes. + /// + /// PUBLIC for integration with the existing main training loop — + /// the loop drives `step_with_lobsim_gpu` itself (so the diag / + /// NaN-abort / counter machinery stays intact) and calls + /// `snapshot_step` after each step. + pub fn snapshot_step( + &self, + trainer: &mut IntegratedTrainer, + buf: &mut RolloutBuffer, + t: usize, + b_size: usize, + t_max: usize, + ) -> Result<()> { + let stream = &trainer.stream; + let raw_stream = stream.cu_stream(); + + // `gae_backward_sweep.cu` indexes `rewards[b * T + t]` — i.e. + // `[B, T]` row-major (batch-major). Per-step writes therefore + // scatter a contiguous `[B]` buffer into column `t` of `[B, T]`: + // one element per batch row, with destination stride + // `T * sizeof(elem)`. cuMemcpy2D is the canonical DtoD-only idiom + // for this layout: + // * src pitch = elem_sz (contiguous [B] source) + // * dst pitch = T * elem_sz (per-batch row pitch of [B, T]) + // * width = elem_sz (one element per row) + // * height = B (one row per batch element) + // No host roundtrip, no atomic ops; deterministic per + // `pearl_determinism_achieved` (single-direction DtoD). + // The u8 dones use a dedicated pack kernel (`rollout_pack.cu`) + // that does the f32→u8 conversion + `b * T + t` scatter in one + // launch. + + let sz_f32 = std::mem::size_of::(); + let sz_i32 = std::mem::size_of::(); + + // Helper: scatter `[B]` of `elem_sz` bytes into `[B, T]` column t. + // Implements cuMemcpy2D with src stride = elem_sz (contiguous), + // dst stride = T * elem_sz (row pitch of the [B, T] buffer). + let scatter_to_column = |dst_base: u64, src: u64, elem_sz: usize| -> Result<()> { + // src offset 0, dst offset = t * elem_sz (column t in row b=0) + let dst_off = (t as u64) * (elem_sz as u64); + unsafe { + let mut copy = std::mem::MaybeUninit::< + cudarc::driver::sys::CUDA_MEMCPY2D_v2, + >::zeroed(); + let cref = &mut *copy.as_mut_ptr(); + cref.srcMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE; + cref.srcDevice = src; + cref.srcPitch = elem_sz; // contiguous source [B] → row pitch = elem_sz + cref.dstMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE; + cref.dstDevice = dst_base + dst_off; + cref.dstPitch = (t_max as usize) * elem_sz; // row pitch of [B, T] = T * elem_sz + cref.WidthInBytes = elem_sz; // one element per row + cref.Height = b_size; + let r = cudarc::driver::sys::cuMemcpy2DAsync_v2( + copy.as_ptr(), + raw_stream, + ); + if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + anyhow::bail!( + "RolloutCollection::scatter_to_column cuMemcpy2DAsync \ + (t={t}, elem_sz={elem_sz}): {:?}", + r + ); + } + } + Ok(()) + }; + + // rewards_d: f32 [B] → rollout.rewards_d[b, t] + let rewards_dst = buf.rewards_d.device_ptr_mut(stream).0; + scatter_to_column(rewards_dst, trainer.rewards_d.device_ptr(stream).0, sz_f32)?; + + // v_t_d: f32 [B] (= trainer.v_pred_d) → rollout.v_t_d[b, t] + let v_t_dst = buf.v_t_d.device_ptr_mut(stream).0; + scatter_to_column(v_t_dst, trainer.v_pred_d.device_ptr(stream).0, sz_f32)?; + + // actions_d: i32 [B] → rollout.actions_d[b, t] + let actions_dst = buf.actions_d.device_ptr_mut(stream).0; + scatter_to_column(actions_dst, trainer.actions_d.device_ptr(stream).0, sz_i32)?; + + // log_pi_old_d: f32 [B] → rollout.log_pi_old_d[b, t] + let log_pi_dst = buf.log_pi_old_d.device_ptr_mut(stream).0; + scatter_to_column(log_pi_dst, trainer.log_pi_old_d.device_ptr(stream).0, sz_f32)?; + + // h_t_d: f32 [B × HIDDEN_DIM] → rollout.h_t_d[b, t, :] + // + // Layout: rollout.h_t_d is `[B × T × HIDDEN_DIM]` row-major in + // (b, t, h). One batch row at offset (b, t, *) starts at byte + // `(b * T + t) * HIDDEN_DIM * sizeof(f32)`. Per-batch rows of + // length HIDDEN_DIM are not contiguous across `b` (stride = + // T * HIDDEN_DIM * 4 bytes). Source `perception.h_t_view()` is + // contiguous `[B × HIDDEN_DIM]` (stride = HIDDEN_DIM * 4 bytes). + // + // cuMemcpy2D with: + // * src stride = HIDDEN_DIM * sz_f32 (per-batch row pitch) + // * dst stride = T * HIDDEN_DIM * sz_f32 (per-batch row pitch in [B, T, H]) + // * WidthInBytes = HIDDEN_DIM * sz_f32 (full hidden row) + // * Height = B + // copies B rows of HIDDEN_DIM floats from contiguous src to + // strided dst at column t. + { + let h_src = trainer.perception.h_t_view().device_ptr(stream).0; + let h_dst_base = buf.h_t_d.device_ptr_mut(stream).0; + let row_bytes = HIDDEN_DIM * sz_f32; + let dst_off = (t as u64) * (row_bytes as u64); + unsafe { + let mut copy = std::mem::MaybeUninit::< + cudarc::driver::sys::CUDA_MEMCPY2D_v2, + >::zeroed(); + let cref = &mut *copy.as_mut_ptr(); + cref.srcMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE; + cref.srcDevice = h_src; + cref.srcPitch = row_bytes; + cref.dstMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE; + cref.dstDevice = h_dst_base + dst_off; + cref.dstPitch = (t_max as usize) * row_bytes; + cref.WidthInBytes = row_bytes; + cref.Height = b_size; + let r = cudarc::driver::sys::cuMemcpy2DAsync_v2( + copy.as_ptr(), + raw_stream, + ); + if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + anyhow::bail!( + "RolloutCollection::snapshot_step h_t cuMemcpy2DAsync \ + (t={t}): {:?}", + r + ); + } + } + } + + // dones_d: f32 [B] {0.0, 1.0} → rollout.dones_d[b, t] (u8 {0, 1}) + // Pack kernel writes `b * T + t` directly per the GAE indexing. + self.launch_pack_dones(trainer, buf, t, b_size, t_max)?; + + Ok(()) + } + + /// Launch `rollout_pack_dones_f32_to_u8` to convert + scatter the + /// per-step `[B]` f32 dones into the `[B × T]` u8 rollout buffer at + /// column `t`. Single-thread-per-batch kernel; one launch per step. + fn launch_pack_dones( + &self, + trainer: &mut IntegratedTrainer, + buf: &mut RolloutBuffer, + t: usize, + b_size: usize, + t_max: usize, + ) -> Result<()> { + let stream = &trainer.stream; + let b = b_size as i32; + let t_i = t_max as i32; + let t_cursor = t as i32; + + let threads: u32 = 256; + let blocks: u32 = ((b_size as u32) + threads - 1) / threads; + let cfg = LaunchConfig { + grid_dim: (blocks.max(1), 1, 1), + block_dim: (threads, 1, 1), + shared_mem_bytes: 0, + }; + + let src = trainer.dones_d.device_ptr(stream).0; + let dst = buf.dones_d.device_ptr_mut(stream).0; + + unsafe { + stream + .launch_builder(&self.pack_dones_kernel) + .arg(&src) + .arg(&dst) + .arg(&b) + .arg(&t_i) + .arg(&t_cursor) + .launch(cfg) + .with_context(|| { + format!("rollout_pack_dones_f32_to_u8 launch (t={t})") + })?; + } + Ok(()) + } + + /// Copy the last step's `v_pred_tp1_d` ([B] f32 — `V(s_{t+1})` for + /// the final action) into `buf.v_t_bootstrap_d`. This is the + /// boundary condition for the GAE backward sweep at `t = T - 1`. + /// + /// Contiguous DtoD; `B × sizeof(f32)` bytes. PUBLIC for the same + /// reason as `snapshot_step`. + pub fn copy_bootstrap_v( + &self, + trainer: &mut IntegratedTrainer, + buf: &mut RolloutBuffer, + b_size: usize, + ) -> Result<()> { + use crate::trainer::raw_launch::raw_memcpy_dtod_async; + + let stream = &trainer.stream; + let raw_stream = stream.cu_stream(); + let src = trainer.v_pred_tp1_d.device_ptr(stream).0; + let dst = buf.v_t_bootstrap_d.device_ptr_mut(stream).0; + let nbytes = b_size * std::mem::size_of::(); + unsafe { + raw_memcpy_dtod_async(dst, src, nbytes, raw_stream).map_err(|e| { + anyhow::anyhow!("copy_bootstrap_v DtoD: {:?}", e) + })?; + } + Ok(()) + } +}