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>
27feb94a49