audit: π drives actions (proper actor-critic) + bump b_size 1 → 16

Two coordinated architectural fixes addressing the deepest blockers
exposed by the audit:

## Option B: π-driven action selection

Per `pearl_q_thompson_actor_makes_pi_dead_weight`: the prior
architecture had Q acting as BOTH actor (via Thompson sample) AND
critic (via Bellman target). π trained by PPO surrogate against
Q's actions but never drove any decision — `q_pi_agree_ema`
decayed to 0 by step 5000 in every smoke because π converged to
Q's Thompson SAMPLING distribution, not Q's argmax. π was
dead-weight: 4 dedicated controllers (ε, ratio_clamp,
entropy_coef, KL EMA), shared encoder gradient interference, and
zero contribution to actor decisions.

### New kernel: rl_pi_action_kernel.cu

Single-thread-per-batch CUDA kernel that:
  1. Computes numerically-stable softmax(pi_logits[b, :])
  2. Draws u ∈ [0, 1) from per-batch xorshift32 PRNG
  3. CDF-walks to pick the multinomial-sampled action

Per-batch xorshift32 PRNG state is the SAME `prng_state_d` buffer
already used by rl_action_kernel — no new state needed. Sampling
deterministic given (seed, b_size, pi_logits).

### Trainer wiring (1 site change in step_with_lobsim)

Replaced `rl_action_kernel(q_logits, atom_supports, ...)`
(Q-Thompson) with `rl_pi_action_kernel(pi_logits, ...)`
(π-multinomial). The argmax_expected_q call on h_{t+1} is
unchanged — Q remains the critic via canonical Double-DQN target.

PPO importance-ratio surrogate now has its canonical actor-critic
semantics: π_new(a|s) / π_old(a|s) where `a` was actually sampled
from π_old. Was nonsensical before (a was sampled from Q-Thompson,
not π, so the ratio measured something incoherent).

The rl_action_kernel (Q-Thompson) cubin + function field are kept
loaded for backward-compat tests and diagnostic comparison; no
longer in the hot path.

## b_size: 1 → 16

Per `pearl_b_size_1_signal_starvation_blocks_q_learning`: at
b_size=1 with 11% done-step rate and 70% loss rate per trade, Q
stayed at uniform baseline ln(21)=3.04 across all 16+ smokes
regardless of controller fixes. The architecture was structurally
signal-starved — 1 gradient sample per Adam step is fundamentally
too noisy.

LobSimCuda already supports b_size>1 (n_backtests parameter at
`crates/ml-backtesting/src/sim/mod.rs:355`). Trainer code is
already b_size-parametric throughout. The blocker was just the
CLI default at `--n-backtests=1`.

Default bumped to 16 (matches the doc note "production sweep at
32-64; L40S 48GB"). 16× more gradient samples per Adam step
gives Q proper batch variance reduction. The K-loop multiplier
(`isv[404]/2048`) will likely settle at K=1 since the
advantage_var_ratio drops with batch size.

## Expected behaviour

  * `q_pi_agree_ema` becomes tautological/dropped (π IS the
    policy now — comparing argmax(Q) to argmax(π) doesn't measure
    a real consistency invariant any more)
  * π gradient flows naturally drive π toward an actor that
    optimises the PPO surrogate — Q's encoder gradient is no
    longer competing with a different policy's gradient
  * l_q should drop meaningfully below 3.04 for the first time
    (was stuck at 2.7-2.9 across all prior smokes)
  * reward/trade should approach 0 (was -$0.5 to -$0.8 across
    every prior run)
  * Wall-clock per env step ~16× slower (b_size=16) but training
    cost per gradient step similar (denser sample = more
    progress per step)

## Verified gates (local sm_86)

  G1 isv_bootstrap   
  G3 controllers     
  G4 target_update   
  integrated_smoke   

## Caveat: integrated_trainer_smoke runs at b_size=1

The default for the CLI is bumped to 16, but the local
`integrated_trainer_smoke` test passes its own b_size=1 to
verify the trainer mechanics. Real-world signal verification
happens via cluster smokes which now use b_size=16 by default.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-24 10:37:42 +02:00
parent 705d6c156b
commit 3737feb664
4 changed files with 139 additions and 7 deletions

View File

@@ -67,6 +67,7 @@ const KERNELS: &[&str] = &[
"rl_streaming_clamp_init", // RL R9: device-side seeder for streaming-kernel output clamp ceilings (ISV[447], ISV[448]) — no HtoD
"rl_isv_write", // RL R9: generic single-slot ISV seeder for tunable design constants — no HtoD
"rl_q_pi_agree_b", // RL R9 audit: per-batch fraction (argmax Q == argmax π) → ISV[407] EMA; wires previously-dead slot
"rl_pi_action_kernel", // audit Option B: π drives action selection via multinomial sampling from softmax(pi_logits); Q becomes pure critic
];
// Cache bust v31 — five new reduce / derive kernels populate the input

View File

@@ -0,0 +1,103 @@
// rl_pi_action_kernel.cu — GPU-resident multinomial sampler over the
// PPO policy head's softmax(pi_logits). Replaces rl_action_kernel
// (which sampled from Q-Thompson) as the canonical actor-side
// decision per `pearl_q_thompson_actor_makes_pi_dead_weight`.
//
// The prior architecture had Q being BOTH the actor (via Thompson)
// AND the critic (via Bellman target). π trained as a third head
// via PPO surrogate against Q's actions but never drove any
// decision — `q_pi_agree_ema` decayed to 0 across every smoke
// because π converged to Q's Thompson SAMPLING distribution mode,
// not Q's argmax mode. The actor-critic design was inverted.
//
// New architecture (this kernel): π drives action selection via
// multinomial sampling from softmax(pi_logits). Q remains the
// critic (Bellman target via `argmax_expected_q` on h_{t+1}). The
// PPO importance-ratio surrogate now has its canonical semantics —
// π_new / π_old at the action that π_old sampled.
//
// PRNG: same xorshift32-per-batch pattern as rl_action_kernel.
// Single thread per block (N_ACTIONS=9 is small enough that the
// per-thread parallelism of the Thompson kernel is unjustified;
// multinomial CDF walk is sequential anyway).
//
// Per `feedback_no_atomicadd`: single-thread kernel, no atomics.
// Per `feedback_cpu_is_read_only`: pure device-side sampling.
#include <stdint.h>
#define N_ACTIONS 9
__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). Single thread per
// block (tid==0). Computes numerically-stable softmax over the
// batch's N_ACTIONS pi_logits, then multinomial-samples via CDF
// walk using the per-batch xorshift32 PRNG state.
//
// Inputs:
// pi_logits [b_size × N_ACTIONS] row-major raw logits
// prng_state [b_size] per-batch xorshift32 state (mutated)
// Outputs:
// actions [b_size] multinomial-sampled action index per batch
extern "C" __global__ void rl_pi_action_kernel(
const float* __restrict__ pi_logits,
uint32_t* __restrict__ prng_state,
int* __restrict__ actions,
int b_size
) {
const int b = blockIdx.x;
if (b >= b_size) return;
if (threadIdx.x != 0) return;
// ── Numerically-stable softmax over pi_logits[b, :]. ──────────
const int base = b * N_ACTIONS;
float logits[N_ACTIONS];
float max_logit = pi_logits[base];
#pragma unroll
for (int a = 0; a < N_ACTIONS; ++a) {
logits[a] = pi_logits[base + a];
if (logits[a] > max_logit) max_logit = logits[a];
}
float sum_exp = 0.0f;
float probs[N_ACTIONS];
#pragma unroll
for (int a = 0; a < N_ACTIONS; ++a) {
probs[a] = expf(logits[a] - max_logit);
sum_exp += probs[a];
}
// Guard against pathological zero-sum (shouldn't happen with
// expf of bounded values, but defensive).
const float inv_sum = (sum_exp > 1e-9f) ? (1.0f / sum_exp) : (1.0f / (float)N_ACTIONS);
#pragma unroll
for (int a = 0; a < N_ACTIONS; ++a) probs[a] *= inv_sum;
// ── Draw uniform u ∈ [0, 1) from per-batch PRNG. ──────────────
uint32_t state = prng_state[b];
// A few warmup advances dispel low-quality seeding correlations.
#pragma unroll
for (int w = 0; w < 4; ++w) xorshift32(&state);
const uint32_t r = xorshift32(&state);
// Convert to [0, 1): divide by 2^32. Use 24-bit mantissa precision
// (good enough for multinomial — Q_N_ATOMS=21 actions max).
const float u = (float)(r >> 8) * (1.0f / 16777216.0f);
prng_state[b] = state;
// ── CDF walk: pick the action whose cumulative prob crosses u. ──
int chosen = N_ACTIONS - 1; // defensive default if numerical drift loses the last bit
float cdf = 0.0f;
#pragma unroll
for (int a = 0; a < N_ACTIONS; ++a) {
cdf += probs[a];
if (u < cdf) { chosen = a; break; }
}
actions[b] = chosen;
}

View File

@@ -111,8 +111,14 @@ struct Cli {
/// Batch dimension — number of parallel LobSim backtests run in
/// lockstep per step. Larger batches amortise the encoder forward
/// cost across more trajectories but push GPU memory pressure
/// linearly. Smoke at 1; production sweep at 32-64 (L40S 48GB).
#[arg(long, default_value_t = 1)]
/// linearly. Default 16 — fixes the b_size=1 signal starvation
/// per `pearl_b_size_1_signal_starvation_blocks_q_learning`
/// (Q stayed at uniform baseline ln(21)=3.04 across all b_size=1
/// smokes regardless of controller tuning; 16 parallel
/// transitions per Adam step give 16× the gradient signal). Was
/// 1 (validation-smoke default); production sweep can push to
/// 32-64 on L40S 48GB.
#[arg(long, default_value_t = 16)]
n_backtests: usize,
/// PER buffer capacity (R7d). Per `replay.rs` doc, naive O(N)

View File

@@ -199,6 +199,8 @@ const RL_ISV_WRITE_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/rl_isv_write.cubin"));
const RL_Q_PI_AGREE_B_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/rl_q_pi_agree_b.cubin"));
const RL_PI_ACTION_KERNEL_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/rl_pi_action_kernel.cubin"));
/// Per-head LR controller Wiener-α target. The controller kernel floors
/// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the
@@ -411,6 +413,8 @@ pub struct IntegratedTrainer {
rl_isv_write_fn: CudaFunction,
_rl_q_pi_agree_b_module: Arc<CudaModule>,
rl_q_pi_agree_b_fn: CudaFunction,
_rl_pi_action_kernel_module: Arc<CudaModule>,
rl_pi_action_kernel_fn: CudaFunction,
/// Per-batch step counter for `trade_duration_ema` — increments
/// every step, resets on done. Zero-init at construction.
pub steps_since_done_d: CudaSlice<i32>,
@@ -851,6 +855,12 @@ impl IntegratedTrainer {
let rl_q_pi_agree_b_fn = rl_q_pi_agree_b_module
.load_function("rl_q_pi_agree_b")
.context("load rl_q_pi_agree_b")?;
let rl_pi_action_kernel_module = ctx
.load_cubin(RL_PI_ACTION_KERNEL_CUBIN.to_vec())
.context("load rl_pi_action_kernel cubin")?;
let rl_pi_action_kernel_fn = rl_pi_action_kernel_module
.load_function("rl_pi_action_kernel")
.context("load rl_pi_action_kernel")?;
// Per-batch PRNG state for the Thompson sampler. Seeded
// deterministically from cfg.dqn_seed via ChaCha8 host RNG so
@@ -1082,6 +1092,8 @@ impl IntegratedTrainer {
rl_isv_write_fn,
_rl_q_pi_agree_b_module: rl_q_pi_agree_b_module,
rl_q_pi_agree_b_fn,
_rl_pi_action_kernel_module: rl_pi_action_kernel_module,
rl_pi_action_kernel_fn,
steps_since_done_d,
trade_duration_emit_d,
prev_realized_pnl_d,
@@ -2497,24 +2509,34 @@ impl IntegratedTrainer {
// increment that drove the old ChaCha8 host RNG is also gone.
self.step_counter = self.step_counter.wrapping_add(1);
// ── π-driven action selection (Option B fix per
// `pearl_q_thompson_actor_makes_pi_dead_weight`). Replaces the
// prior Q-Thompson selector. π is now the actor: actions come
// from multinomial sample of softmax(pi_logits). Q remains the
// critic (Bellman target via argmax_expected_q on h_{t+1}).
// PPO importance-ratio surrogate now has canonical semantics —
// π_new / π_old at the action that π_old sampled.
//
// The rl_action_kernel (Q-Thompson) field is kept loaded for
// backward-compat tests + diagnostic comparison; it's no
// longer in the hot path.
{
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (N_ACTIONS as u32, 1, 1),
block_dim: (1, 1, 1), // single thread per block (multinomial CDF walk is sequential)
shared_mem_bytes: 0,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.rl_action_kernel_fn);
let mut launch = self.stream.launch_builder(&self.rl_pi_action_kernel_fn);
launch
.arg(&q_logits_d)
.arg(&self.atom_supports_d)
.arg(&pi_logits_d)
.arg(&mut self.prng_state_d)
.arg(&mut self.actions_d)
.arg(&b_size_i);
unsafe {
launch
.launch(cfg)
.context("rl_action_kernel launch (Thompson)")?;
.context("rl_pi_action_kernel launch (multinomial from softmax(π))")?;
}
}