diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 72ebfb8f6..350544294 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -44,9 +44,20 @@ const KERNELS: &[&str] = &[ "rl_rollout_steps_controller", // RL Phase E.3b: rollout-buffer-length ISV emitter — emits ISV[RL_N_ROLLOUT_STEPS_INDEX=404] from var(advantage)/|mean A| EMA; bootstraps 2048 "rl_per_alpha_controller", // RL Phase E.3b: PER priority-exponent ISV emitter — emits ISV[RL_PER_ALPHA_INDEX=405] from TD-error kurtosis EMA; bootstraps 0.6 "rl_reward_scale_controller", // RL Phase R1 (rebuild): reward-standardisation scale ISV emitter — emits ISV[RL_REWARD_SCALE_INDEX=406] from mean |realized_pnl_usd| EMA; bootstraps 1.0 + "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] ]; -// Cache bust v25 (2026-05-23): RL Phase R1 rebuild — port rl_reward_scale_controller.cu (ISV[406] producer, bootstrap 1.0; ported from ml-alpha-phase-f-g-flawed) and wire all 7 RL controllers (gamma/target_tau/ppo_clip/entropy_coef/rollout_steps/per_alpha/reward_scale) into IntegratedTrainer::new via launch-once-with-sentinel-zero-input per pearl_first_observation_bootstrap. Replaces flawed Phase F's host memcpy_htod bootstrap that violated feedback_no_htod_htoh_only_mapped_pinned. +// Cache bust v26 (2026-05-23): RL Phase R3 rebuild — three new GPU-resident +// kernels close defect #5 from the flawed Phase F+G arc +// (feedback_cpu_is_read_only violation in step_with_lobsim's host advantage +// + EMA loops). ema_update_on_done + ema_update_per_step are generic +// slot-parameterised EMA producers feeding the 7 RL controllers' ISV[417..424] +// inputs; compute_advantage_return is element-wise A_t + R_t on device, +// replacing the host loop. All three honor pearl_first_observation_bootstrap +// (bootstrap path defers if mean_obs == 0) + pearl_wiener_alpha_floor_for_nonstationary +// (host pre-floors α at 0.4) + feedback_no_atomicadd (tree reduction). fn main() { println!("cargo:rerun-if-changed=build.rs"); diff --git a/crates/ml-alpha/cuda/compute_advantage_return.cu b/crates/ml-alpha/cuda/compute_advantage_return.cu new file mode 100644 index 000000000..56e7f12f8 --- /dev/null +++ b/crates/ml-alpha/cuda/compute_advantage_return.cu @@ -0,0 +1,49 @@ +// compute_advantage_return.cu — element-wise one-step TD advantage + +// return computation on device (Phase R3 of the integrated RL trainer +// rebuild; see +// docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Replaces the host loop the flawed Phase F shipped in step_with_lobsim +// (which violated `feedback_cpu_is_read_only` by reading v_pred to host +// + computing advantages on CPU + uploading back to device). Per +// `pearl_cold_path_no_exception_to_gpu_drives` the "cold path" (one +// reduction per step) is not a license for CPU compute — even simple +// arithmetic stays on device. +// +// Per-batch formulas (canonical actor-critic): +// +// returns[b] = rewards[b] + γ × (1 − dones[b]) × v_tp1[b] +// advantages[b] = returns[b] − v_t[b] +// +// γ is read from `ISV[RL_GAMMA_INDEX = 400]`. R1 bootstraps the slot +// to 0.99 via the rl_gamma_controller's first-observation-bootstrap +// path; R5 adapts it per step via the trade-duration EMA (ISV[417]). +// +// Element-wise, trivially parallel. One thread per batch entry; no +// reduction, no atomics. Per `feedback_no_atomicadd` not needed. + +#define RL_GAMMA_INDEX 400 + +extern "C" __global__ void compute_advantage_return( + const float* __restrict__ isv, // ISV bus (≥ RL_GAMMA_INDEX + 1) + const float* __restrict__ rewards, // [b_size] + const float* __restrict__ dones, // [b_size] 0.0 / 1.0 + const float* __restrict__ v_t, // [b_size] V(s_t) + const float* __restrict__ v_tp1, // [b_size] V(s_{t+1}) + float* __restrict__ returns, // [b_size] OUT + float* __restrict__ advantages, // [b_size] OUT + int b_size +) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= b_size) return; + + const float gamma = isv[RL_GAMMA_INDEX]; + const float r = rewards[b]; + const float done = dones[b]; + const float vt = v_t[b]; + const float vtp1 = v_tp1[b]; + + const float ret = r + gamma * (1.0f - done) * vtp1; + returns[b] = ret; + advantages[b] = ret - vt; +} diff --git a/crates/ml-alpha/cuda/ema_update_on_done.cu b/crates/ml-alpha/cuda/ema_update_on_done.cu new file mode 100644 index 000000000..44b3f4c52 --- /dev/null +++ b/crates/ml-alpha/cuda/ema_update_on_done.cu @@ -0,0 +1,91 @@ +// ema_update_on_done.cu — done-gated EMA producer for ISV-resident +// adaptive-control EMAs (Phase R3 of the integrated RL trainer rebuild; +// see docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Generic over the ISV slot — one kernel serves all the closed-trade +// EMAs the 7 RL controllers consume (e.g. mean_abs_pnl_ema → ISV[423] +// feeding rl_reward_scale_controller; q_divergence_ema → ISV[418] +// feeding rl_target_tau_controller; etc.). The caller passes the slot +// index and the Wiener-α (already floored at 0.4 per +// `pearl_wiener_alpha_floor_for_nonstationary` host-side). +// +// Done-gated: only the per-batch entries where `dones[b] >= 0.5` +// contribute to the EMA observation. If no batch closed a trade this +// step, the kernel does nothing — the ISV slot retains its prior value +// rather than blending toward zero on hold steps. This is what +// distinguishes "EMA over closed-trade magnitudes" from "EMA over all +// steps" (the latter lives in `ema_update_per_step.cu`). +// +// Bootstrap discipline (per `pearl_first_observation_bootstrap`): if +// the ISV slot reads sentinel `0.0`, the first non-zero observation +// REPLACES the slot directly — no blend with the cold-start zero. If +// `mean_obs` happens to be exactly zero on the first observation +// (vanishingly rare for strictly-positive signals like |reward|), the +// kernel defers the bootstrap to a later call rather than writing zero +// back as a degenerate sentinel that would never be distinguishable +// from the cold-start state. +// +// Per `feedback_no_atomicadd`: shared-memory tree reduction inside a +// single block; no atomics. Caller passes `shared_mem_bytes = +// 2 * b_size * sizeof(float)` (one slot for the gated-sum, one for the +// gate count). + +extern "C" __global__ void ema_update_on_done( + float* __restrict__ isv, // ISV bus (≥ slot_index + 1) + int slot_index, // which ISV slot to update + float alpha, // Wiener-α (host-floored ≥ 0.4) + const float* __restrict__ obs, // [b_size] per-batch raw signal + const float* __restrict__ dones, // [b_size] 0.0 / 1.0 done flag + int b_size +) { + // Single block, b_size threads. The block layout requires + // b_size ≤ blockDim.x; callers (Phase R5) launch with + // block_dim = (max(b_size, 1), 1, 1). For b_size = 1 the tree + // loop degenerates and the if at the end runs unchanged. + extern __shared__ float smem[]; + float* s_obs_gated = smem; // [b_size] + float* s_gate = smem + b_size; // [b_size] + + const int tid = threadIdx.x; + if (tid >= b_size) return; + + const float d = dones[tid]; + const float gate = (d >= 0.5f) ? 1.0f : 0.0f; + s_obs_gated[tid] = obs[tid] * gate; + s_gate[tid] = gate; + __syncthreads(); + + // Tree reduction (block-level, no atomicAdd per + // feedback_no_atomicadd). Power-of-two stride covers b_size ≤ + // blockDim.x; bounds-check `tid + stride < b_size` handles + // non-power-of-two b_size correctly. + for (int stride = 1; stride < b_size; stride *= 2) { + if ((tid % (stride * 2)) == 0 && (tid + stride) < b_size) { + s_obs_gated[tid] += s_obs_gated[tid + stride]; + s_gate[tid] += s_gate[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + const float count = s_gate[0]; + if (count == 0.0f) { + // No closed trade this step — preserve the prior EMA. + return; + } + const float mean_obs = s_obs_gated[0] / count; + const float prev = isv[slot_index]; + if (prev == 0.0f) { + // First-observation bootstrap per pearl_first_observation_bootstrap. + // Defer if mean_obs == 0 so we don't write a sentinel that + // would be re-bootstrapped on the next call. + if (mean_obs != 0.0f) { + isv[slot_index] = mean_obs; + } + } else { + // Wiener-α blend; caller pre-floored α at 0.4 per + // pearl_wiener_alpha_floor_for_nonstationary. + isv[slot_index] = (1.0f - alpha) * prev + alpha * mean_obs; + } + } +} diff --git a/crates/ml-alpha/cuda/ema_update_per_step.cu b/crates/ml-alpha/cuda/ema_update_per_step.cu new file mode 100644 index 000000000..9d553aade --- /dev/null +++ b/crates/ml-alpha/cuda/ema_update_per_step.cu @@ -0,0 +1,65 @@ +// ema_update_per_step.cu — per-step EMA producer for ISV-resident +// adaptive-control EMAs that update every step regardless of episode +// boundary (Phase R3 of the integrated RL trainer rebuild; +// see docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Companion to `ema_update_on_done.cu`. This variant has NO done-gate +// — it consumes the per-batch raw signal on every step. Used for EMAs +// that should update continuously rather than only at trade closes: +// +// * KL(π_new ‖ π_old) EMA (input to rl_ppo_clip_controller, ISV[419]) +// * Observed action entropy H(π) EMA (input to rl_entropy_coef_controller, +// ISV[420]) +// * Advantage variance ratio EMA (input to rl_rollout_steps_controller, +// ISV[421]) +// +// Generic over the ISV slot — caller passes `slot_index` + Wiener-α. +// Bootstrap discipline + reduction strategy are identical to +// `ema_update_on_done.cu`; only the gate differs (here, every batch +// entry contributes, so the per-batch mean is the unweighted average). +// +// Per `feedback_no_atomicadd`: shared-memory tree reduction in a +// single block; no atomics. Caller passes +// `shared_mem_bytes = b_size * sizeof(float)`. + +extern "C" __global__ void ema_update_per_step( + float* __restrict__ isv, // ISV bus (≥ slot_index + 1) + int slot_index, // which ISV slot to update + float alpha, // Wiener-α (host-floored ≥ 0.4) + const float* __restrict__ obs, // [b_size] per-batch raw signal + int b_size +) { + extern __shared__ float smem[]; // [b_size] + + const int tid = threadIdx.x; + if (tid >= b_size) return; + + smem[tid] = obs[tid]; + __syncthreads(); + + // Tree reduction (block-level, no atomicAdd per + // feedback_no_atomicadd). + for (int stride = 1; stride < b_size; stride *= 2) { + if ((tid % (stride * 2)) == 0 && (tid + stride) < b_size) { + smem[tid] += smem[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + const float mean_obs = smem[0] / (float)b_size; + const float prev = isv[slot_index]; + if (prev == 0.0f) { + // First-observation bootstrap per pearl_first_observation_bootstrap. + // Defer if mean_obs == 0 so we don't write a sentinel that + // would be re-bootstrapped on the next call. + if (mean_obs != 0.0f) { + isv[slot_index] = mean_obs; + } + } else { + // Wiener-α blend; caller pre-floored α at 0.4 per + // pearl_wiener_alpha_floor_for_nonstationary. + isv[slot_index] = (1.0f - alpha) * prev + alpha * mean_obs; + } + } +} diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index e5c88fa53..402213e71 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -129,6 +129,17 @@ const RL_PER_ALPHA_CONTROLLER_CUBIN: &[u8] = const RL_REWARD_SCALE_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_reward_scale_controller.cubin")); +// Phase R3: GPU-resident EMA + advantage/return kernels. Generic over +// the ISV slot they update (one kernel serves multiple controller-input +// EMAs). Replace the host-side EMA + advantage loops the flawed Phase F +// shipped in step_with_lobsim (which violated `feedback_cpu_is_read_only`). +const EMA_UPDATE_ON_DONE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/ema_update_on_done.cubin")); +const EMA_UPDATE_PER_STEP_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/ema_update_per_step.cubin")); +const COMPUTE_ADVANTAGE_RETURN_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/compute_advantage_return.cubin")); + /// Per-head LR controller Wiener-α target. The controller kernel floors /// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the /// effective blend rate stays ≥ 0.4 regardless of this seed value. Kept @@ -240,6 +251,14 @@ pub struct IntegratedTrainer { _rl_reward_scale_controller_module: Arc, rl_reward_scale_controller_fn: CudaFunction, + // ── Phase R3: GPU-resident EMA + advantage/return kernels ───────── + _ema_update_on_done_module: Arc, + ema_update_on_done_fn: CudaFunction, + _ema_update_per_step_module: Arc, + ema_update_per_step_fn: CudaFunction, + _compute_advantage_return_module: Arc, + compute_advantage_return_fn: CudaFunction, + /// Combined encoder grad slot — folded from Q + π + V `grad_h_t` /// contributions via `grad_h_accumulate_scaled` (item 1). Consumed by /// `PerceptionTrainer::backward_encoder_with_grad_h_t`. Size @@ -401,6 +420,26 @@ impl IntegratedTrainer { .load_function("rl_reward_scale_controller") .context("load rl_reward_scale_controller")?; + // Phase R3 kernels. + let ema_update_on_done_module = ctx + .load_cubin(EMA_UPDATE_ON_DONE_CUBIN.to_vec()) + .context("load ema_update_on_done cubin")?; + let ema_update_on_done_fn = ema_update_on_done_module + .load_function("ema_update_on_done") + .context("load ema_update_on_done")?; + let ema_update_per_step_module = ctx + .load_cubin(EMA_UPDATE_PER_STEP_CUBIN.to_vec()) + .context("load ema_update_per_step cubin")?; + let ema_update_per_step_fn = ema_update_per_step_module + .load_function("ema_update_per_step") + .context("load ema_update_per_step")?; + let compute_advantage_return_module = ctx + .load_cubin(COMPUTE_ADVANTAGE_RETURN_CUBIN.to_vec()) + .context("load compute_advantage_return cubin")?; + let compute_advantage_return_fn = compute_advantage_return_module + .load_function("compute_advantage_return") + .context("load compute_advantage_return")?; + Ok(Self { cfg, perception, @@ -436,6 +475,12 @@ impl IntegratedTrainer { rl_per_alpha_controller_fn, _rl_reward_scale_controller_module: rl_reward_scale_controller_module, rl_reward_scale_controller_fn, + _ema_update_on_done_module: ema_update_on_done_module, + ema_update_on_done_fn, + _ema_update_per_step_module: ema_update_per_step_module, + ema_update_per_step_fn, + _compute_advantage_return_module: compute_advantage_return_module, + compute_advantage_return_fn, grad_h_t_combined_d, last_pi_loss: 0.0, step_counter: 0, @@ -522,6 +567,132 @@ impl IntegratedTrainer { Ok(()) } + /// Phase R3: launch `ema_update_on_done` — done-gated EMA producer + /// that writes the per-batch closed-trade mean into `isv[slot_index]`. + /// `alpha` is the Wiener-α (caller must pre-floor at 0.4 per + /// `pearl_wiener_alpha_floor_for_nonstationary`). `obs_d` and + /// `dones_d` are `[b_size]` device buffers; only entries where + /// `dones[b] >= 0.5` contribute to the EMA observation. If no + /// batch closed a trade this step, the kernel is a no-op (ISV slot + /// retains its prior value). + /// + /// Shared-memory tree reduce; single block, `b_size` threads. + /// Caller's responsibility: `b_size <= 1024` (one block max). + pub fn launch_ema_update_on_done( + &self, + slot_index: usize, + alpha: f32, + obs_d: &CudaSlice, + dones_d: &CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(obs_d.len(), b_size); + debug_assert_eq!(dones_d.len(), b_size); + debug_assert!(b_size <= 1024, "ema_update_on_done: b_size > 1024 needs multi-block reduce"); + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (b_size as u32, 1, 1), + shared_mem_bytes: (2 * b_size * std::mem::size_of::()) as u32, + }; + let slot_i = slot_index as i32; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.ema_update_on_done_fn); + launch + .arg(&self.isv_d) + .arg(&slot_i) + .arg(&alpha) + .arg(obs_d) + .arg(dones_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("ema_update_on_done launch")?; + } + Ok(()) + } + + /// Phase R3: launch `ema_update_per_step` — per-step EMA producer + /// that writes the per-batch mean into `isv[slot_index]` on every + /// call (no done-gate). Used for continuous EMAs (KL, entropy, + /// advantage variance). Same shared-mem tree-reduce as + /// `ema_update_on_done`. + pub fn launch_ema_update_per_step( + &self, + slot_index: usize, + alpha: f32, + obs_d: &CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(obs_d.len(), b_size); + debug_assert!(b_size <= 1024, "ema_update_per_step: b_size > 1024 needs multi-block reduce"); + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (b_size as u32, 1, 1), + shared_mem_bytes: (b_size * std::mem::size_of::()) as u32, + }; + let slot_i = slot_index as i32; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.ema_update_per_step_fn); + launch + .arg(&self.isv_d) + .arg(&slot_i) + .arg(&alpha) + .arg(obs_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("ema_update_per_step launch")?; + } + Ok(()) + } + + /// Phase R3: launch `compute_advantage_return` — element-wise + /// `returns[b] = r + γ(1-done)·V(s_{t+1})`, + /// `advantages[b] = returns[b] − V(s_t)`. Reads γ from `ISV[400]` + /// (R1 bootstraps to 0.99; R5 adapts via trade-duration EMA). + /// + /// Element-wise; trivially parallel. Block layout: 32-thread blocks + /// covering `b_size`. + #[allow(clippy::too_many_arguments)] + pub fn launch_compute_advantage_return( + &self, + rewards_d: &CudaSlice, + dones_d: &CudaSlice, + v_t_d: &CudaSlice, + v_tp1_d: &CudaSlice, + returns_d: &mut CudaSlice, + advantages_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(rewards_d.len(), b_size); + debug_assert_eq!(dones_d.len(), b_size); + debug_assert_eq!(v_t_d.len(), b_size); + debug_assert_eq!(v_tp1_d.len(), b_size); + debug_assert_eq!(returns_d.len(), b_size); + debug_assert_eq!(advantages_d.len(), b_size); + let b_size_i = b_size as i32; + let grid_x = ((b_size as u32) + 31) / 32; + let cfg = LaunchConfig { + grid_dim: (grid_x.max(1), 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.compute_advantage_return_fn); + launch + .arg(&self.isv_d) + .arg(rewards_d) + .arg(dones_d) + .arg(v_t_d) + .arg(v_tp1_d) + .arg(returns_d) + .arg(advantages_d) + .arg(&b_size_i); + unsafe { + launch + .launch(cfg) + .context("compute_advantage_return launch")?; + } + Ok(()) + } + /// Phase E.2 + E.2-DEFER: synthetic-batch end-to-end step. Runs real /// GPU forward + backward + per-head Adam on Q, π, V heads on a /// SHARED encoder; combined per-head `grad_h_t` flows back into the diff --git a/crates/ml-alpha/tests/r3_ema_advantage.rs b/crates/ml-alpha/tests/r3_ema_advantage.rs new file mode 100644 index 000000000..3f210c93a --- /dev/null +++ b/crates/ml-alpha/tests/r3_ema_advantage.rs @@ -0,0 +1,255 @@ +//! Phase R3 gate — three GPU-oracle tests for the new generic kernels: +//! +//! 1. `ema_update_on_done` — bootstrap path: sentinel-zero ISV slot +//! + first observation k → ISV[slot] == k exactly. +//! 2. `ema_update_per_step` — convergence: constant input k for many +//! steps → ISV[slot] → k within tolerance. +//! 3. `compute_advantage_return` — formula: zero rewards + done=0 + +//! `v_t = v_tp1 = k` + γ=0.99 (from R1 bootstrap) → `returns = 0.99k`, +//! `advantages = (0.99 − 1)·k = −0.01k`. +//! +//! Per `feedback_no_cpu_test_fallbacks` every oracle here is either: +//! - the kernel's documented bootstrap behaviour (kernel writes +//! `mean_obs` directly on first observation), OR +//! - an analytical property of the formula (EMA of constant = +//! constant; advantage = γ·V − V = (γ−1)·V when r=0, done=0, +//! v_t = v_tp1), OR +//! - the kernel's own `*_BOOTSTRAP` value (γ=0.99 from R1). +//! +//! Per `pearl_tests_must_prove_not_lock_observations` each assert is +//! an invariant statement, not a tuned magic number. +//! +//! Run with: +//! `cargo test -p ml-alpha --test r3_ema_advantage -- --ignored --nocapture` + +use ml_alpha::rl::isv_slots::{ + RL_GAMMA_INDEX, RL_KL_PI_EMA_INDEX, RL_MEAN_ABS_PNL_EMA_INDEX, RL_SLOTS_END, +}; +use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig}; +use ml_alpha::trainer::perception::PerceptionTrainerConfig; +use ml_core::device::MlDevice; + +const ALPHA_FLOOR: f32 = 0.4; + +fn build_trainer() -> Option<(MlDevice, IntegratedTrainer)> { + let dev = match MlDevice::cuda(0) { + Ok(d) => d, + Err(e) => { + eprintln!("CUDA 0 not available — skipping ({e})"); + return None; + } + }; + let cfg = IntegratedTrainerConfig { + perception: PerceptionTrainerConfig { + seq_len: 4, + n_batch: 1, + ..PerceptionTrainerConfig::default() + }, + dqn_seed: 0xA1, + ppo_seed: 0xA2, + }; + let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); + Some((dev, trainer)) +} + +fn upload( + stream: &std::sync::Arc, + host: &[f32], +) -> cudarc::driver::CudaSlice { + let mut d = stream.alloc_zeros::(host.len()).expect("alloc"); + stream.memcpy_htod(host, &mut d).expect("htod"); + d +} + +fn readback_isv( + dev: &MlDevice, + isv_d: &cudarc::driver::CudaSlice, +) -> Vec { + let mut isv = vec![0.0_f32; RL_SLOTS_END]; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + stream + .memcpy_dtoh(isv_d, isv.as_mut_slice()) + .expect("isv dtoh"); + isv +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r3_ema_update_on_done_first_observation_bootstrap_replaces_directly() { + let Some((dev, trainer)) = build_trainer() else { return }; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + + // Pre-condition: the EMA-input slot is at sentinel zero (R1 + // bootstraps ISV[400..406] for controllers; ISV[417..423] EMA-input + // slots stay at alloc_zeros per the R1 invariant). + let isv_before = readback_isv(&dev, &trainer.isv_d); + assert_eq!( + isv_before[RL_MEAN_ABS_PNL_EMA_INDEX], 0.0, + "pre-condition: EMA slot must be sentinel zero" + ); + + // One observation: obs = [7.0], dones = [1.0]. Per + // pearl_first_observation_bootstrap the kernel writes the + // observation directly into the slot — no blend with the cold- + // start zero. + let obs_d = upload(&stream, &[7.0]); + let dones_d = upload(&stream, &[1.0]); + trainer + .launch_ema_update_on_done(RL_MEAN_ABS_PNL_EMA_INDEX, ALPHA_FLOOR, &obs_d, &dones_d, 1) + .expect("ema_update_on_done"); + stream.synchronize().expect("sync"); + + let isv_after = readback_isv(&dev, &trainer.isv_d); + let val = isv_after[RL_MEAN_ABS_PNL_EMA_INDEX]; + // Exact equality — bootstrap path is `isv[slot] = mean_obs` with + // no arithmetic. Any drift indicates a wrong code path. + assert_eq!( + val, 7.0, + "first-observation bootstrap should write mean_obs exactly; got {val}" + ); + + // Negative invariant: hold-only step (no dones) MUST NOT touch + // the slot — defends against an off-by-one in the gate count. + let obs2_d = upload(&stream, &[999.0]); + let dones2_d = upload(&stream, &[0.0]); + trainer + .launch_ema_update_on_done(RL_MEAN_ABS_PNL_EMA_INDEX, ALPHA_FLOOR, &obs2_d, &dones2_d, 1) + .expect("ema_update_on_done hold"); + stream.synchronize().expect("sync"); + + let isv_hold = readback_isv(&dev, &trainer.isv_d); + assert_eq!( + isv_hold[RL_MEAN_ABS_PNL_EMA_INDEX], 7.0, + "hold step (no done) must preserve EMA, not blend toward 0" + ); + + eprintln!( + "R3.1 OK — bootstrap wrote 7.0 directly; hold-only step preserved value" + ); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r3_ema_update_per_step_converges_to_constant_input() { + let Some((dev, trainer)) = build_trainer() else { return }; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + + // Property: EMA of constant k converges to k. With α=0.4 and + // bootstrap-then-blend, after N=50 steps with obs=5.0: + // step 0 (bootstrap): isv = 5.0 directly + // step 1+ (blend): isv = (1-α)·prev + α·5.0 = stays at 5.0 + // So the assertion is essentially exact (modulo fp rounding). + let k = 5.0_f32; + let obs_d = upload(&stream, &[k]); + for _ in 0..50 { + trainer + .launch_ema_update_per_step(RL_KL_PI_EMA_INDEX, ALPHA_FLOOR, &obs_d, 1) + .expect("ema_update_per_step"); + } + stream.synchronize().expect("sync"); + + let isv = readback_isv(&dev, &trainer.isv_d); + let val = isv[RL_KL_PI_EMA_INDEX]; + assert!( + (val - k).abs() < 1e-4, + "EMA of constant {k} should converge to {k}; got {val}" + ); + + eprintln!("R3.2 OK — EMA(constant={k}) converged to {val} after 50 steps"); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r3_compute_advantage_return_formula_holds() { + let Some((dev, trainer)) = build_trainer() else { return }; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + + // R1 bootstrapped ISV[γ=400] to 0.99. Verify before the test. + let isv = readback_isv(&dev, &trainer.isv_d); + let gamma = isv[RL_GAMMA_INDEX]; + assert!( + (gamma - 0.99).abs() < 1e-6, + "pre-condition: γ should be R1-bootstrapped to 0.99; got {gamma}" + ); + + // Inputs: r=0, done=0, v_t=v_tp1=k. Expected: + // returns[b] = 0 + γ·1·k = γ·k + // advantages[b] = γ·k − k = (γ − 1)·k + let k = 5.0_f32; + let b_size = 3; + let rewards_d = upload(&stream, &vec![0.0_f32; b_size]); + let dones_d = upload(&stream, &vec![0.0_f32; b_size]); + let v_t_d = upload(&stream, &vec![k; b_size]); + let v_tp1_d = upload(&stream, &vec![k; b_size]); + let mut returns_d = stream.alloc_zeros::(b_size).expect("alloc returns"); + let mut advantages_d = stream.alloc_zeros::(b_size).expect("alloc advantages"); + + trainer + .launch_compute_advantage_return( + &rewards_d, + &dones_d, + &v_t_d, + &v_tp1_d, + &mut returns_d, + &mut advantages_d, + b_size, + ) + .expect("compute_advantage_return"); + + let mut returns = vec![0.0_f32; b_size]; + let mut advantages = vec![0.0_f32; b_size]; + stream.memcpy_dtoh(&returns_d, returns.as_mut_slice()).expect("dtoh returns"); + stream.memcpy_dtoh(&advantages_d, advantages.as_mut_slice()).expect("dtoh advantages"); + stream.synchronize().expect("sync"); + + let expected_ret = gamma * k; // 0.99 · 5 = 4.95 + let expected_adv = (gamma - 1.0) * k; // -0.05 + for b in 0..b_size { + assert!( + (returns[b] - expected_ret).abs() < 1e-5, + "returns[{b}] expected {expected_ret}, got {}", + returns[b] + ); + assert!( + (advantages[b] - expected_adv).abs() < 1e-5, + "advantages[{b}] expected {expected_adv}, got {}", + advantages[b] + ); + } + + // Negative invariant: done=1 zeros the future-value bootstrap. + let dones_done_d = upload(&stream, &vec![1.0_f32; b_size]); + trainer + .launch_compute_advantage_return( + &rewards_d, + &dones_done_d, + &v_t_d, + &v_tp1_d, + &mut returns_d, + &mut advantages_d, + b_size, + ) + .expect("compute_advantage_return done=1"); + stream.memcpy_dtoh(&returns_d, returns.as_mut_slice()).expect("dtoh returns done"); + stream.memcpy_dtoh(&advantages_d, advantages.as_mut_slice()).expect("dtoh advantages done"); + stream.synchronize().expect("sync"); + + for b in 0..b_size { + assert!( + returns[b].abs() < 1e-6, + "returns[{b}] with done=1 + r=0 expected 0, got {}", + returns[b] + ); + assert!( + (advantages[b] - (-k)).abs() < 1e-5, + "advantages[{b}] with done=1 expected −k = {}, got {}", + -k, + advantages[b] + ); + } + + eprintln!( + "R3.3 OK — γ={gamma}: returns=4.95 / advantages=-0.05 when v=5; \ + done=1 zeros bootstrap → returns=0 / advantages=-5" + ); +}