diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 350544294..456c0f6eb 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -47,17 +47,20 @@ 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] + "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 ]; -// 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). +// Cache bust v27 (2026-05-23): RL Phase R4 rebuild — three new GPU-resident +// action-sampling kernels close defect #5 prerequisites for the GPU-pure +// step_with_lobsim that lands in R6. rl_action_kernel uses a per-batch +// xorshift32 PRNG state (allocated + host-seeded at trainer init per +// pearl_scoped_init_seed_for_reproducibility, no cuRAND dep); +// argmax_expected_q is the Bellman-target deterministic argmax per +// pearl_thompson_for_distributional_action_selection (Thompson for rollout, +// argmax for Bellman); log_pi_at_action computes log π for the PPO +// importance ratio. All three element-wise or single-block; no atomicAdd. fn main() { println!("cargo:rerun-if-changed=build.rs"); diff --git a/crates/ml-alpha/cuda/argmax_expected_q.cu b/crates/ml-alpha/cuda/argmax_expected_q.cu new file mode 100644 index 000000000..2011a7893 --- /dev/null +++ b/crates/ml-alpha/cuda/argmax_expected_q.cu @@ -0,0 +1,86 @@ +// argmax_expected_q.cu — GPU-resident argmax over expected Q-values +// per action (Phase R4 of the integrated RL trainer rebuild; +// see docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Companion to `rl_action_kernel.cu`. Per +// `pearl_thompson_for_distributional_action_selection`: Thompson is +// used for ROLLOUT action selection (rl_action_kernel) and argmax over +// EXPECTED Q is used for the Bellman-target argmax (this kernel). The +// distinction matters for the canonical Double-DQN-style backup: +// +// target Q = Q_target(s_{t+1}, argmax_a E[Q_online(s_{t+1}, a)]) +// +// Replaces the host-side `argmax_f32(&action_expected)` loop the +// flawed Phase F shipped in step_with_lobsim — which violated +// `feedback_cpu_is_read_only` by DtoH-copying q_logits and reducing +// on CPU. +// +// Per-batch: softmax over each action's atoms, accumulate +// `expected[a] = Σ_i p_i · support[i]`, then argmax over the 9 +// per-action expected values. Deterministic (no PRNG). +// +// Per `feedback_no_atomicadd`: thread 0 writes to `next_actions[b]` +// after __syncthreads. + +#define N_ACTIONS 9 +#define Q_N_ATOMS 21 + +// One block per batch (grid_dim.x = b_size). N_ACTIONS threads per +// block; each thread computes its action's expected Q and writes to +// shared mem; thread 0 argmaxes + writes next_actions[b]. +// +// Inputs: +// q_logits [b_size, N_ACTIONS, Q_N_ATOMS] row-major +// atom_supports [Q_N_ATOMS] +// Outputs: +// next_actions [b_size] argmax over expected Q per batch +extern "C" __global__ void argmax_expected_q( + const float* __restrict__ q_logits, + const float* __restrict__ atom_supports, + int* __restrict__ next_actions, + int b_size +) { + const int b = blockIdx.x; + const int a = threadIdx.x; + if (b >= b_size) return; + if (a >= N_ACTIONS) return; + + __shared__ float expected_q[N_ACTIONS]; + + // Softmax + expected value over this action's atoms. + const int row_off = (b * N_ACTIONS + a) * Q_N_ATOMS; + float max_l = -INFINITY; + #pragma unroll + for (int i = 0; i < Q_N_ATOMS; ++i) { + const float l = q_logits[row_off + i]; + if (l > max_l) max_l = l; + } + float sum_exp = 0.0f; + #pragma unroll + for (int i = 0; i < Q_N_ATOMS; ++i) { + sum_exp += expf(q_logits[row_off + i] - max_l); + } + if (!isfinite(sum_exp) || sum_exp <= 0.0f) sum_exp = 1.0f; + + float ev = 0.0f; + #pragma unroll + for (int i = 0; i < Q_N_ATOMS; ++i) { + const float p = expf(q_logits[row_off + i] - max_l) / sum_exp; + ev += p * atom_supports[i]; + } + expected_q[a] = ev; + __syncthreads(); + + if (a == 0) { + int best_a = 0; + float best_v = expected_q[0]; + #pragma unroll + for (int i = 1; i < N_ACTIONS; ++i) { + if (expected_q[i] > best_v) { + best_v = expected_q[i]; + best_a = i; + } + } + next_actions[b] = best_a; + } +} diff --git a/crates/ml-alpha/cuda/log_pi_at_action.cu b/crates/ml-alpha/cuda/log_pi_at_action.cu new file mode 100644 index 000000000..566a17070 --- /dev/null +++ b/crates/ml-alpha/cuda/log_pi_at_action.cu @@ -0,0 +1,55 @@ +// log_pi_at_action.cu — GPU-resident per-batch log π(action | s) +// computation for the PPO importance-ratio path (Phase R4 of the +// integrated RL trainer rebuild; +// see docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Replaces the host-side `pi_logits[action] − log_sum_exp(pi_logits)` +// loop the flawed Phase F shipped in step_with_lobsim — which +// violated `feedback_cpu_is_read_only` by DtoH-copying π logits and +// computing log-softmax on CPU. +// +// Per-batch: log_pi[b] = pi_logits[b, actions[b]] − log Σ_a exp(pi_logits[b, a]). +// Numerically stabilised via max-shift before exp. +// +// One thread per batch entry; element-wise. No reduction across +// batches, no atomics. N_ACTIONS = 9 fits in a per-thread sequential +// loop comfortably. + +#define N_ACTIONS 9 + +// Inputs: +// pi_logits [b_size, N_ACTIONS] row-major +// actions [b_size] index into [0, N_ACTIONS) +// Outputs: +// log_pi_out [b_size] +extern "C" __global__ void log_pi_at_action( + const float* __restrict__ pi_logits, + const int* __restrict__ actions, + float* __restrict__ log_pi_out, + int b_size +) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= b_size) return; + + const int row_off = b * N_ACTIONS; + float max_l = -INFINITY; + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) { + const float l = pi_logits[row_off + i]; + if (l > max_l) max_l = l; + } + float sum_exp = 0.0f; + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) { + sum_exp += expf(pi_logits[row_off + i] - max_l); + } + // log-sum-exp identity: log Σ exp(x) = max + log Σ exp(x − max) + const float log_norm = max_l + logf(sum_exp); + + const int chosen = actions[b]; + // No bounds check on `chosen` per `feedback_no_quickfixes`: the + // upstream Thompson kernel (rl_action_kernel) writes into + // [0, N_ACTIONS) by construction. A bounds check here would mask + // an upstream contract violation; let it crash visibly. + log_pi_out[b] = pi_logits[row_off + chosen] - log_norm; +} diff --git a/crates/ml-alpha/cuda/rl_action_kernel.cu b/crates/ml-alpha/cuda/rl_action_kernel.cu new file mode 100644 index 000000000..19bfed0ee --- /dev/null +++ b/crates/ml-alpha/cuda/rl_action_kernel.cu @@ -0,0 +1,132 @@ +// rl_action_kernel.cu — GPU-resident Thompson sampler over the C51 +// distributional Q-head output (Phase R4 of the integrated RL trainer +// rebuild; see +// docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Replaces the host-side Thompson loop the flawed Phase F shipped in +// step_with_lobsim (which violated `feedback_cpu_is_read_only` by +// DtoH-copying q_logits + sampling on CPU + HtoD-copying actions back +// to device). +// +// Per `pearl_thompson_for_distributional_action_selection`: Thompson +// is the canonical rollout selector under C51. For each batch + each +// action, sample ONE atom from that action's categorical (softmax +// over atom logits) and use the sampled atom's support value as the +// "sampled return"; pick the action with the highest sampled return. +// The companion `argmax_expected_q.cu` kernel does the Bellman-target +// argmax (over expected Q, not sampled). +// +// PRNG: per-batch xorshift32 state in `prng_state_d` (allocated + +// host-seeded once at trainer init per +// `pearl_scoped_init_seed_for_reproducibility`). To avoid the +// inter-thread race that a shared per-batch state would create, each +// per-action thread XORs the action index (via golden-ratio constant) +// into a thread-local copy of the per-batch state, advances locally +// for the atom sample, and only thread 0 writes the advanced per-batch +// state back. Reproducibility holds: identical (per-batch seed, +// blockDim, action grid, q_logits) → identical actions. +// +// Per `feedback_no_atomicadd`: thread 0 writes to `actions[b]` after +// __syncthreads. Per `pearl_no_host_branches_in_captured_graph`: no +// host branches; the only branches inside the kernel are device-side +// data-dependent (CDF walk, argmax). + +#include + +#define N_ACTIONS 9 +#define Q_N_ATOMS 21 + +__device__ static uint32_t xorshift32(uint32_t* state) { + uint32_t x = *state; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *state = x; + return x; +} + +// One block per batch (grid_dim.x = b_size). N_ACTIONS threads per +// block; each thread handles one action's atom sample, writes its +// sampled return to shared mem, then thread 0 argmaxes + writes +// actions[b] + advances the per-batch PRNG state. +// +// Inputs: +// q_logits [b_size, N_ACTIONS, Q_N_ATOMS] row-major +// atom_supports [Q_N_ATOMS] e.g. linspace(Q_V_MIN, Q_V_MAX, Q_N_ATOMS) +// prng_state [b_size] per-batch xorshift32 state (mutated in place) +// Outputs: +// actions [b_size] Thompson-sampled action index per batch +extern "C" __global__ void rl_action_kernel( + const float* __restrict__ q_logits, + const float* __restrict__ atom_supports, + uint32_t* __restrict__ prng_state, + int* __restrict__ actions, + int b_size +) { + const int b = blockIdx.x; + const int a = threadIdx.x; + if (b >= b_size) return; + if (a >= N_ACTIONS) return; + + __shared__ float sampled_returns[N_ACTIONS]; + + // Derive thread-local PRNG state from the per-batch seed. + // Golden-ratio constant decorrelates per-thread streams. + const uint32_t base = prng_state[b]; + uint32_t local_state = base ^ ((uint32_t)a * 0x9E3779B1u); + // A few warmup advances dispel correlations from the cheap mixing. + #pragma unroll + for (int w = 0; w < 4; ++w) xorshift32(&local_state); + + // Softmax over this action's atoms (numerically stabilised). + const int row_off = (b * N_ACTIONS + a) * Q_N_ATOMS; + float max_l = -INFINITY; + #pragma unroll + for (int i = 0; i < Q_N_ATOMS; ++i) { + const float l = q_logits[row_off + i]; + if (l > max_l) max_l = l; + } + float sum_exp = 0.0f; + #pragma unroll + for (int i = 0; i < Q_N_ATOMS; ++i) { + sum_exp += expf(q_logits[row_off + i] - max_l); + } + // Guard against numerical degeneracy (all-NaN or all-inf logits). + if (!isfinite(sum_exp) || sum_exp <= 0.0f) sum_exp = 1.0f; + + // Categorical CDF walk against u ∈ [0, 1). Walk-to-end fallback + // catches the case where cumulative-sum rounding edges below 1.0. + const uint32_t r = xorshift32(&local_state); + const float u = (float)r / 4294967296.0f; + float cum = 0.0f; + int chosen = Q_N_ATOMS - 1; + for (int i = 0; i < Q_N_ATOMS; ++i) { + cum += expf(q_logits[row_off + i] - max_l) / sum_exp; + if (cum >= u) { + chosen = i; + break; + } + } + + sampled_returns[a] = atom_supports[chosen]; + __syncthreads(); + + // Thread 0: argmax over per-action sampled returns; write action; + // advance per-batch PRNG state for next call. + if (a == 0) { + int best_a = 0; + float best_v = sampled_returns[0]; + #pragma unroll + for (int i = 1; i < N_ACTIONS; ++i) { + if (sampled_returns[i] > best_v) { + best_v = sampled_returns[i]; + best_a = i; + } + } + actions[b] = best_a; + + uint32_t adv = base; + xorshift32(&adv); + prng_state[b] = adv; + } +} diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 402213e71..c14057272 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -140,6 +140,16 @@ const EMA_UPDATE_PER_STEP_CUBIN: &[u8] = const COMPUTE_ADVANTAGE_RETURN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/compute_advantage_return.cubin")); +// Phase R4: GPU-resident action sampling kernels. Replace the host +// Thompson + host argmax + host log-softmax loops the flawed Phase F +// shipped in step_with_lobsim (which violated `feedback_cpu_is_read_only`). +const RL_ACTION_KERNEL_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_action_kernel.cubin")); +const ARGMAX_EXPECTED_Q_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/argmax_expected_q.cubin")); +const LOG_PI_AT_ACTION_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/log_pi_at_action.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 @@ -259,6 +269,30 @@ pub struct IntegratedTrainer { _compute_advantage_return_module: Arc, compute_advantage_return_fn: CudaFunction, + // ── Phase R4: GPU action sampling ───────────────────────────────── + _rl_action_kernel_module: Arc, + rl_action_kernel_fn: CudaFunction, + _argmax_expected_q_module: Arc, + argmax_expected_q_fn: CudaFunction, + _log_pi_at_action_module: Arc, + log_pi_at_action_fn: CudaFunction, + + /// Per-batch xorshift32 PRNG state for `rl_action_kernel`. Allocated + /// + host-seeded once at trainer init from `cfg.dqn_seed` via + /// `rand_chacha::ChaCha8Rng` per `pearl_scoped_init_seed_for_reproducibility`. + /// The kernel advances state in place each call. Length = `n_batch`. + pub prng_state_d: CudaSlice, + + /// Atom support vector `[Q_V_MIN, Q_V_MIN + step, …, Q_V_MAX]` of + /// length `Q_N_ATOMS`. Allocated + host-uploaded once at init; + /// read by `rl_action_kernel` (Thompson sample lookup) and + /// `argmax_expected_q` (expected-value accumulator). Per + /// `feedback_isv_for_adaptive_bounds`: structural constant, not + /// adaptive (atom span is fixed [Q_V_MIN, Q_V_MAX] = [-1, +1] per + /// `crate::rl::common`). Per-branch adaptive atom span (per + /// `pearl_per_branch_c51_atom_span`) is a Phase R-future extension. + pub atom_supports_d: CudaSlice, + /// 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 @@ -440,6 +474,62 @@ impl IntegratedTrainer { .load_function("compute_advantage_return") .context("load compute_advantage_return")?; + // Phase R4 kernels. + let rl_action_kernel_module = ctx + .load_cubin(RL_ACTION_KERNEL_CUBIN.to_vec()) + .context("load rl_action_kernel cubin")?; + let rl_action_kernel_fn = rl_action_kernel_module + .load_function("rl_action_kernel") + .context("load rl_action_kernel")?; + let argmax_expected_q_module = ctx + .load_cubin(ARGMAX_EXPECTED_Q_CUBIN.to_vec()) + .context("load argmax_expected_q cubin")?; + let argmax_expected_q_fn = argmax_expected_q_module + .load_function("argmax_expected_q") + .context("load argmax_expected_q")?; + let log_pi_at_action_module = ctx + .load_cubin(LOG_PI_AT_ACTION_CUBIN.to_vec()) + .context("load log_pi_at_action cubin")?; + let log_pi_at_action_fn = log_pi_at_action_module + .load_function("log_pi_at_action") + .context("load log_pi_at_action")?; + + // Per-batch PRNG state for the Thompson sampler. Seeded + // deterministically from cfg.dqn_seed via ChaCha8 host RNG so + // (cfg.dqn_seed, b_size) → identical Thompson draws across + // re-runs of the same trainer init per + // `pearl_scoped_init_seed_for_reproducibility`. Zero is the + // only forbidden xorshift32 state (would freeze) so the + // `.max(1)` guard ensures every batch's seed is non-zero. + let b_size = cfg.perception.n_batch; + let mut prng_host_rng = + ::seed_from_u64( + cfg.dqn_seed.wrapping_add(0xC0_C0_C0_C0), + ); + let prng_seeds: Vec = (0..b_size) + .map(|_| prng_host_rng.gen::().max(1)) + .collect(); + let mut prng_state_d = stream + .alloc_zeros::(b_size) + .context("alloc prng_state_d")?; + stream + .memcpy_htod(&prng_seeds, &mut prng_state_d) + .context("memcpy_htod prng_state_d")?; + + // Atom supports — structural constant (atom span [-1, +1]). + // Single allocation + upload at init; read every step by the + // Thompson + argmax kernels. + let atom_step = (Q_V_MAX - Q_V_MIN) / ((Q_N_ATOMS - 1) as f32); + let atom_supports_host: Vec = (0..Q_N_ATOMS) + .map(|i| Q_V_MIN + (i as f32) * atom_step) + .collect(); + let mut atom_supports_d = stream + .alloc_zeros::(Q_N_ATOMS) + .context("alloc atom_supports_d")?; + stream + .memcpy_htod(&atom_supports_host, &mut atom_supports_d) + .context("memcpy_htod atom_supports_d")?; + Ok(Self { cfg, perception, @@ -481,6 +571,14 @@ impl IntegratedTrainer { ema_update_per_step_fn, _compute_advantage_return_module: compute_advantage_return_module, compute_advantage_return_fn, + _rl_action_kernel_module: rl_action_kernel_module, + rl_action_kernel_fn, + _argmax_expected_q_module: argmax_expected_q_module, + argmax_expected_q_fn, + _log_pi_at_action_module: log_pi_at_action_module, + log_pi_at_action_fn, + prng_state_d, + atom_supports_d, grad_h_t_combined_d, last_pi_loss: 0.0, step_counter: 0, @@ -644,6 +742,101 @@ impl IntegratedTrainer { Ok(()) } + /// Phase R4: launch `rl_action_kernel` — Thompson sampler over + /// the C51 Q distribution. Writes per-batch action indices to + /// `actions_d`. Advances per-batch xorshift32 PRNG state in place. + /// One block per batch, `N_ACTIONS` threads. `q_logits_d` is + /// `[b_size, N_ACTIONS, Q_N_ATOMS]` row-major. + pub fn launch_rl_action_kernel( + &mut self, + q_logits_d: &CudaSlice, + actions_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(q_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS); + debug_assert_eq!(actions_d.len(), b_size); + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (N_ACTIONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.rl_action_kernel_fn); + launch + .arg(q_logits_d) + .arg(&self.atom_supports_d) + .arg(&mut self.prng_state_d) + .arg(actions_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("rl_action_kernel launch")?; + } + Ok(()) + } + + /// Phase R4: launch `argmax_expected_q` — Bellman-target argmax + /// over expected Q per action. Deterministic (no PRNG). Pairs with + /// `launch_rl_action_kernel` per + /// `pearl_thompson_for_distributional_action_selection` (Thompson + /// for rollout, argmax for Bellman target). + pub fn launch_argmax_expected_q( + &self, + q_logits_d: &CudaSlice, + next_actions_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(q_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS); + debug_assert_eq!(next_actions_d.len(), b_size); + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (N_ACTIONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.argmax_expected_q_fn); + launch + .arg(q_logits_d) + .arg(&self.atom_supports_d) + .arg(next_actions_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("argmax_expected_q launch")?; + } + Ok(()) + } + + /// Phase R4: launch `log_pi_at_action` — per-batch + /// `log π(actions[b] | s_b)` via log-softmax + lookup. Feeds the + /// PPO importance ratio. One thread per batch entry. + pub fn launch_log_pi_at_action( + &self, + pi_logits_d: &CudaSlice, + actions_d: &CudaSlice, + log_pi_out_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(pi_logits_d.len(), b_size * N_ACTIONS); + debug_assert_eq!(actions_d.len(), b_size); + debug_assert_eq!(log_pi_out_d.len(), b_size); + 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 b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.log_pi_at_action_fn); + launch + .arg(pi_logits_d) + .arg(actions_d) + .arg(log_pi_out_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("log_pi_at_action 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]` diff --git a/crates/ml-alpha/tests/r4_action_kernels.rs b/crates/ml-alpha/tests/r4_action_kernels.rs new file mode 100644 index 000000000..163061f53 --- /dev/null +++ b/crates/ml-alpha/tests/r4_action_kernels.rs @@ -0,0 +1,246 @@ +//! Phase R4 gate — three GPU-oracle tests for the new action-sampling +//! kernels: +//! +//! 1. `rl_action_kernel` (Thompson sampler) — sharp Q distribution +//! (action 5 dominates with `logit=20` on the highest support +//! atom; other actions dominate with `logit=20` on the lowest +//! atom) collapses Thompson sampling to deterministic argmax, +//! so action 5 wins every trial. +//! 2. `argmax_expected_q` — same sharp distribution → next_action = 5. +//! Cross-kernel oracle: argmax matches Thompson when distribution +//! is sharp per `pearl_thompson_for_distributional_action_selection`. +//! 3. `log_pi_at_action` — π logits with `logit=20` at action 3 and +//! `0` elsewhere → `log π(3) ≈ 0` (probability mass concentrates +//! on action 3). +//! +//! Per `feedback_no_cpu_test_fallbacks` the oracles are analytical: +//! - Test 1: deterministic-collapse property of Thompson under sharp +//! distribution (no CPU sampling). +//! - Test 2: action-with-highest-EV is the index with the highest- +//! support dominant atom (no CPU softmax + argmax oracle). +//! - Test 3: log-softmax algebra (log(exp(20) / (exp(20) + 8·1))) ≈ 0 +//! to within `8·exp(-20) ≈ 4e-8`. +//! +//! Per `pearl_tests_must_prove_not_lock_observations` each assert is +//! an invariant. +//! +//! Run with: +//! `cargo test -p ml-alpha --test r4_action_kernels -- --ignored --nocapture` + +use cudarc::driver::CudaStream; +use ml_alpha::rl::common::{N_ACTIONS, Q_N_ATOMS}; +use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig}; +use ml_alpha::trainer::perception::PerceptionTrainerConfig; +use ml_core::device::MlDevice; +use std::sync::Arc; + +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: 0xB4, + ppo_seed: 0xB5, + }; + let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); + Some((dev, trainer)) +} + +fn upload_f32(stream: &Arc, host: &[f32]) -> cudarc::driver::CudaSlice { + let mut d = stream.alloc_zeros::(host.len()).expect("alloc f32"); + stream.memcpy_htod(host, &mut d).expect("htod f32"); + d +} + +fn upload_i32(stream: &Arc, host: &[i32]) -> cudarc::driver::CudaSlice { + let mut d = stream.alloc_zeros::(host.len()).expect("alloc i32"); + stream.memcpy_htod(host, &mut d).expect("htod i32"); + d +} + +/// Build a `[b_size, N_ACTIONS, Q_N_ATOMS]` Q-logits buffer with +/// `dominant_logit` at the highest-support atom (index `Q_N_ATOMS−1`) +/// for `rewarded_action` and at the lowest-support atom (index `0`) +/// for every other action. +fn build_sharp_q_logits( + b_size: usize, + rewarded_action: usize, + dominant_logit: f32, +) -> Vec { + let mut out = vec![0.0_f32; b_size * N_ACTIONS * Q_N_ATOMS]; + for b in 0..b_size { + for a in 0..N_ACTIONS { + let dominant_atom = if a == rewarded_action { + Q_N_ATOMS - 1 + } else { + 0 + }; + let row_off = (b * N_ACTIONS + a) * Q_N_ATOMS; + out[row_off + dominant_atom] = dominant_logit; + } + } + out +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r4_thompson_sharp_distribution_collapses_to_argmax() { + let Some((dev, mut trainer)) = build_trainer() else { return }; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + + // Sharp distribution: action 5's mass concentrates on atom 20 + // (support +1.0); other actions' mass concentrates on atom 0 + // (support −1.0). exp(20) ≈ 4.85e8 vs other atoms' exp(0) = 1 → + // per-action dominant-atom probability ≈ 1 − 20·exp(−20) ≈ + // 1 − 4e-8. So Thompson collapses to argmax: action 5 wins every + // trial. + let rewarded_action: i32 = 5; + let b_size = 1; + let q_logits = build_sharp_q_logits(b_size, rewarded_action as usize, 20.0); + let q_logits_d = upload_f32(&stream, &q_logits); + let mut actions_d = stream.alloc_zeros::(b_size).expect("alloc actions"); + + let n_trials = 100; + let mut wins = 0; + for _ in 0..n_trials { + trainer + .launch_rl_action_kernel(&q_logits_d, &mut actions_d, b_size) + .expect("rl_action_kernel"); + stream.synchronize().expect("sync"); + let mut actions_host = vec![0_i32; b_size]; + stream + .memcpy_dtoh(&actions_d, actions_host.as_mut_slice()) + .expect("dtoh actions"); + if actions_host[0] == rewarded_action { + wins += 1; + } + } + + // Theoretical loss rate < 1e-6 per trial. 100/100 expected; allow + // 99/100 for fp-rounding edge in CDF walk near `u ≈ 1.0`. + assert!( + wins >= 99, + "Thompson under sharp distribution should pick rewarded action ~100% of the time; got {wins}/{n_trials}" + ); + + eprintln!( + "R4.1 OK — Thompson collapsed to argmax under sharp distribution: {wins}/{n_trials} picked action {rewarded_action}" + ); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r4_argmax_expected_q_picks_highest_ev() { + let Some((dev, trainer)) = build_trainer() else { return }; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + + // Same sharp distribution as R4.1. + let rewarded_action: i32 = 5; + let b_size = 1; + let q_logits = build_sharp_q_logits(b_size, rewarded_action as usize, 20.0); + let q_logits_d = upload_f32(&stream, &q_logits); + let mut next_actions_d = stream.alloc_zeros::(b_size).expect("alloc next"); + + trainer + .launch_argmax_expected_q(&q_logits_d, &mut next_actions_d, b_size) + .expect("argmax_expected_q"); + stream.synchronize().expect("sync"); + + let mut next_actions_host = vec![0_i32; b_size]; + stream + .memcpy_dtoh(&next_actions_d, next_actions_host.as_mut_slice()) + .expect("dtoh next_actions"); + + // Action 5's expected Q ≈ +1.0 (mass on +1 atom); others ≈ −1.0. + // Argmax is deterministic — the assertion is exact. + assert_eq!( + next_actions_host[0], rewarded_action, + "argmax_expected_q should pick the action with the highest-support dominant atom" + ); + + // Negative invariant: swap the dominant atoms — now action 2 wins. + let swapped_action: i32 = 2; + let q_swapped = build_sharp_q_logits(b_size, swapped_action as usize, 20.0); + let q_swapped_d = upload_f32(&stream, &q_swapped); + trainer + .launch_argmax_expected_q(&q_swapped_d, &mut next_actions_d, b_size) + .expect("argmax_expected_q swapped"); + stream.synchronize().expect("sync"); + stream + .memcpy_dtoh(&next_actions_d, next_actions_host.as_mut_slice()) + .expect("dtoh next_actions swapped"); + assert_eq!( + next_actions_host[0], swapped_action, + "argmax_expected_q should follow the rewarded-action swap" + ); + + eprintln!( + "R4.2 OK — argmax_expected_q picked action {rewarded_action} then action {swapped_action} after swap" + ); +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r4_log_pi_at_action_dominant_returns_near_zero() { + let Some((dev, trainer)) = build_trainer() else { return }; + let stream = dev.cuda_stream().expect("cuda_stream").clone(); + + // π logits: action 3 has logit 20, others 0. softmax: p[3] ≈ + // 1 − 8·exp(−20) ≈ 1 − 1.6e-8. log p[3] = 20 − log(exp(20) + + // 8·exp(0)) ≈ 20 − 20 = 0 within 1e-7. + let b_size = 1; + let dominant_action: usize = 3; + let mut pi_logits = vec![0.0_f32; b_size * N_ACTIONS]; + pi_logits[dominant_action] = 20.0; + + let pi_logits_d = upload_f32(&stream, &pi_logits); + let actions_d = upload_i32(&stream, &[dominant_action as i32]); + let mut log_pi_out_d = stream.alloc_zeros::(b_size).expect("alloc log_pi"); + + trainer + .launch_log_pi_at_action(&pi_logits_d, &actions_d, &mut log_pi_out_d, b_size) + .expect("log_pi_at_action"); + stream.synchronize().expect("sync"); + + let mut log_pi_host = vec![0.0_f32; b_size]; + stream + .memcpy_dtoh(&log_pi_out_d, log_pi_host.as_mut_slice()) + .expect("dtoh log_pi"); + + let val = log_pi_host[0]; + assert!( + val.abs() < 1e-4, + "log π(dominant action with logit=20) should be ≈ 0; got {val}" + ); + + // Negative invariant: query a NON-dominant action — log π should + // be ≈ −20 (massive penalty). + let other_action: usize = 7; + let actions_other_d = upload_i32(&stream, &[other_action as i32]); + trainer + .launch_log_pi_at_action(&pi_logits_d, &actions_other_d, &mut log_pi_out_d, b_size) + .expect("log_pi_at_action other"); + stream.synchronize().expect("sync"); + stream + .memcpy_dtoh(&log_pi_out_d, log_pi_host.as_mut_slice()) + .expect("dtoh log_pi other"); + + let val_other = log_pi_host[0]; + assert!( + (val_other + 20.0).abs() < 1e-3, + "log π(non-dominant action) should be ≈ −20 (logit 0 minus log-norm ≈ 20); got {val_other}" + ); + + eprintln!( + "R4.3 OK — log π(dominant)={val:.6} ≈ 0; log π(other)={val_other:.4} ≈ −20" + ); +}