feat(rl): R3 — GPU-resident EMA + advantage/return kernels
Closes defect #5 from the flawed Phase F+G arc (feedback_cpu_is_read_only violation in step_with_lobsim's host advantage + EMA loops) by landing the GPU primitives those loops will become in R6. Three new kernels, each with a GPU-oracle gate test: 1. ema_update_on_done.cu — done-gated EMA producer. - Slot-parameterised (one kernel, 3 callers in R5 covering mean_abs_pnl_ema, q_divergence_ema, td_kurtosis_ema). - Shared-mem tree reduce, no atomicAdd (feedback_no_atomicadd). - Per pearl_first_observation_bootstrap: sentinel-zero ISV → first observation replaces directly. Defers bootstrap if mean_obs == 0 to avoid writing a degenerate sentinel that would be re-bootstrapped next call. - Per pearl_wiener_alpha_floor_for_nonstationary: Wiener-α blend on subsequent calls; caller pre-floors α at 0.4. 2. ema_update_per_step.cu — per-step EMA producer (no done-gate). - Slot-parameterised (kl_pi_ema, entropy_observed_ema, advantage_var_ratio_ema, mean_trade_duration_ema in R5). - Same shared-mem tree reduce + bootstrap discipline as ema_update_on_done. 3. compute_advantage_return.cu — element-wise returns[b] = r + γ(1-done)·V(s_{t+1}); advantages[b] = returns − V(s_t). - Reads γ from ISV[400] (R1 bootstrap = 0.99). - Trivially parallel, one thread per batch entry; no atomics. Rust launchers added to IntegratedTrainer: - launch_ema_update_on_done(slot, alpha, obs_d, dones_d, b_size) - launch_ema_update_per_step(slot, alpha, obs_d, b_size) - launch_compute_advantage_return(rewards_d, dones_d, v_t_d, v_tp1_d, returns_d, advantages_d, b_size) 3 cubin includes, 3 module/function fields, loaders in new() between the rl_reward_scale_controller load and the with_controllers_bootstrapped call so the new fields are populated by struct construction. GPU-oracle tests in tests/r3_ema_advantage.rs (per feedback_no_cpu_test_fallbacks every oracle is either the kernel's documented bootstrap behaviour or an analytical property of the formula, not a CPU reference): R3.1: ema_update_on_done bootstrap path — sentinel-zero ISV + one observation k → ISV[slot] == k exactly. Negative invariant: hold-only step (dones all zero) preserves the EMA. R3.2: ema_update_per_step convergence — feed obs=5.0 for 50 steps with α=0.4 → ISV[slot] → 5.0 within 1e-4 (EMA of constant = constant). R3.3: compute_advantage_return formula — r=0, done=0, v_t=v_tp1=k, γ=0.99 → returns=γk=4.95, advantages=(γ−1)k=−0.05. Negative invariant: done=1 + r=0 zeros the future-value bootstrap (returns=0, advantages=−k). Build cache-bust v26. cargo check + cargo build --test r3_ema_advantage on ml-alpha green. Pre-existing heads_bit_equiv.rs index-out-of-bounds failure persists (unrelated; pre-Phase E). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -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");
|
||||
|
||||
49
crates/ml-alpha/cuda/compute_advantage_return.cu
Normal file
49
crates/ml-alpha/cuda/compute_advantage_return.cu
Normal file
@@ -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;
|
||||
}
|
||||
91
crates/ml-alpha/cuda/ema_update_on_done.cu
Normal file
91
crates/ml-alpha/cuda/ema_update_on_done.cu
Normal file
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
65
crates/ml-alpha/cuda/ema_update_per_step.cu
Normal file
65
crates/ml-alpha/cuda/ema_update_per_step.cu
Normal file
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<CudaModule>,
|
||||
rl_reward_scale_controller_fn: CudaFunction,
|
||||
|
||||
// ── Phase R3: GPU-resident EMA + advantage/return kernels ─────────
|
||||
_ema_update_on_done_module: Arc<CudaModule>,
|
||||
ema_update_on_done_fn: CudaFunction,
|
||||
_ema_update_per_step_module: Arc<CudaModule>,
|
||||
ema_update_per_step_fn: CudaFunction,
|
||||
_compute_advantage_return_module: Arc<CudaModule>,
|
||||
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<f32>,
|
||||
dones_d: &CudaSlice<f32>,
|
||||
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::<f32>()) 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<f32>,
|
||||
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::<f32>()) 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<f32>,
|
||||
dones_d: &CudaSlice<f32>,
|
||||
v_t_d: &CudaSlice<f32>,
|
||||
v_tp1_d: &CudaSlice<f32>,
|
||||
returns_d: &mut CudaSlice<f32>,
|
||||
advantages_d: &mut CudaSlice<f32>,
|
||||
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
|
||||
|
||||
255
crates/ml-alpha/tests/r3_ema_advantage.rs
Normal file
255
crates/ml-alpha/tests/r3_ema_advantage.rs
Normal file
@@ -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<cudarc::driver::CudaStream>,
|
||||
host: &[f32],
|
||||
) -> cudarc::driver::CudaSlice<f32> {
|
||||
let mut d = stream.alloc_zeros::<f32>(host.len()).expect("alloc");
|
||||
stream.memcpy_htod(host, &mut d).expect("htod");
|
||||
d
|
||||
}
|
||||
|
||||
fn readback_isv(
|
||||
dev: &MlDevice,
|
||||
isv_d: &cudarc::driver::CudaSlice<f32>,
|
||||
) -> Vec<f32> {
|
||||
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::<f32>(b_size).expect("alloc returns");
|
||||
let mut advantages_d = stream.alloc_zeros::<f32>(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"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user