feat(rl): R4 — GPU-resident action sampling kernels

Closes defect #5 prerequisites for the GPU-pure step_with_lobsim that
lands in R6. Replaces 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 with 5 DtoH copies + host
per-action loops + HtoD action upload per training step).

Three new kernels:

1. rl_action_kernel.cu — Thompson sampler over the C51 atom
   distribution. One block per batch, N_ACTIONS=9 threads. Each thread
   softmaxes its action's Q_N_ATOMS=21 atoms, samples one atom via CDF
   walk, writes sampled return to shared mem. Thread 0 argmaxes over
   per-action sampled returns + writes actions[b] + advances per-batch
   PRNG state.

   PRNG: per-batch xorshift32 state in prng_state_d (allocated +
   host-seeded from cfg.dqn_seed via ChaCha8 at trainer init per
   pearl_scoped_init_seed_for_reproducibility, with .max(1) guard
   since xorshift32 freezes at 0). Each per-action thread XORs its
   action index (golden-ratio mixed) into a thread-local copy of the
   per-batch state — no inter-thread race, reproducible by
   (cfg.dqn_seed, b_size, step_count). No cuRAND dep.

2. argmax_expected_q.cu — Bellman-target argmax over expected Q per
   action. Same layout as rl_action_kernel but deterministic (no
   PRNG). Per pearl_thompson_for_distributional_action_selection:
   Thompson for rollout (rl_action_kernel), argmax for Bellman target
   (this kernel) — distinct kernels, distinct ISV consumers.

3. log_pi_at_action.cu — per-batch log π(actions[b] | s_b) via
   log-softmax + lookup. One thread per batch entry (N_ACTIONS=9 is
   small enough for a per-thread sequential loop). Feeds the PPO
   importance ratio in R6.

IntegratedTrainer gains:
- 3 cubin includes (rl_action_kernel, argmax_expected_q, log_pi_at_action)
- 3 module/function field pairs
- 2 new device buffers populated at init:
    prng_state_d: CudaSlice<u32> of length n_batch
    atom_supports_d: CudaSlice<f32> of length Q_N_ATOMS=21,
      values [Q_V_MIN, Q_V_MIN + step, …, Q_V_MAX] = linspace(-1, +1, 21)
- 3 launcher methods:
    launch_rl_action_kernel(q_logits_d, actions_d, b_size)
    launch_argmax_expected_q(q_logits_d, next_actions_d, b_size)
    launch_log_pi_at_action(pi_logits_d, actions_d, log_pi_out_d, b_size)

GPU-oracle tests in tests/r4_action_kernels.rs (per
feedback_no_cpu_test_fallbacks every oracle is analytical, not a CPU
reference):

  R4.1: Thompson under sharp distribution (action 5 has logit=20 on
        atom 20 / support +1.0; others have logit=20 on atom 0 /
        support −1.0) collapses to argmax — per-action dominant-atom
        probability ≈ 1 − 4e-8, so 100/100 trials should pick action
        5. Assert ≥99/100 (tolerates one fp-rounding edge near
        u ≈ 1.0 in CDF walk).
  R4.2: argmax_expected_q picks the rewarded action under the same
        sharp distribution. Negative invariant: swap dominant atom to
        action 2 → next_action follows.
  R4.3: log_pi_at_action with π logits dominant at action 3 (logit=20,
        others=0) → log π(3) ≈ 0 within 1e-4. Negative invariant: log
        π(other action) ≈ −20 within 1e-3.

Build cache-bust v27.

cargo check + cargo build --tests on ml-alpha green (heads_bit_equiv
pre-existing failure persists).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

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");

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;
+    }
+}

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;
+}

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 <stdint.h>
+
+#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;
+    }
+}

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<CudaModule>,
     compute_advantage_return_fn: CudaFunction,
 
+    // ── Phase R4: GPU action sampling ─────────────────────────────────
+    _rl_action_kernel_module: Arc<CudaModule>,
+    rl_action_kernel_fn: CudaFunction,
+    _argmax_expected_q_module: Arc<CudaModule>,
+    argmax_expected_q_fn: CudaFunction,
+    _log_pi_at_action_module: Arc<CudaModule>,
+    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<u32>,
+
+    /// 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<f32>,
+
     /// 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 =
+            <rand_chacha::ChaCha8Rng as SeedableRng>::seed_from_u64(
+                cfg.dqn_seed.wrapping_add(0xC0_C0_C0_C0),
+            );
+        let prng_seeds: Vec<u32> = (0..b_size)
+            .map(|_| prng_host_rng.gen::<u32>().max(1))
+            .collect();
+        let mut prng_state_d = stream
+            .alloc_zeros::<u32>(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<f32> = (0..Q_N_ATOMS)
+            .map(|i| Q_V_MIN + (i as f32) * atom_step)
+            .collect();
+        let mut atom_supports_d = stream
+            .alloc_zeros::<f32>(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<f32>,
+        actions_d: &mut CudaSlice<i32>,
+        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<f32>,
+        next_actions_d: &mut CudaSlice<i32>,
+        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<f32>,
+        actions_d: &CudaSlice<i32>,
+        log_pi_out_d: &mut CudaSlice<f32>,
+        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]`

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<CudaStream>, host: &[f32]) -> cudarc::driver::CudaSlice<f32> {
+    let mut d = stream.alloc_zeros::<f32>(host.len()).expect("alloc f32");
+    stream.memcpy_htod(host, &mut d).expect("htod f32");
+    d
+}
+
+fn upload_i32(stream: &Arc<CudaStream>, host: &[i32]) -> cudarc::driver::CudaSlice<i32> {
+    let mut d = stream.alloc_zeros::<i32>(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<f32> {
+    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::<i32>(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::<i32>(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::<f32>(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"
+    );
+}