From 3737feb6649e94f7c49b8cb9fde054ea261f181f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 24 May 2026 10:37:42 +0200 Subject: [PATCH] =?UTF-8?q?audit:=20=CF=80=20drives=20actions=20(proper=20?= =?UTF-8?q?actor-critic)=20+=20bump=20b=5Fsize=201=20=E2=86=92=2016?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/rl_pi_action_kernel.cu | 103 ++++++++++++++++++++ crates/ml-alpha/examples/alpha_rl_train.rs | 10 +- crates/ml-alpha/src/trainer/integrated.rs | 32 +++++- 4 files changed, 139 insertions(+), 7 deletions(-) create mode 100644 crates/ml-alpha/cuda/rl_pi_action_kernel.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 57a8cc50b..477344aa7 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -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 diff --git a/crates/ml-alpha/cuda/rl_pi_action_kernel.cu b/crates/ml-alpha/cuda/rl_pi_action_kernel.cu new file mode 100644 index 000000000..0f4e2af2b --- /dev/null +++ b/crates/ml-alpha/cuda/rl_pi_action_kernel.cu @@ -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 + +#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; +} diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index cc65f19b7..89439c48b 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -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) diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 2fd6c6825..512b69195 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -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, rl_q_pi_agree_b_fn: CudaFunction, + _rl_pi_action_kernel_module: Arc, + 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, @@ -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(π))")?; } }