From 6d33f18b751e96939bcd7917bc24f34e941c051b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 28 May 2026 10:39:12 +0200 Subject: [PATCH] =?UTF-8?q?fix(rl):=20advantage=20normalization=20?= =?UTF-8?q?=E2=80=94=20scale-invariant=20PPO=20learning?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Standard PPO practice missing in this codebase. Without it, lifting the WIN clamp from 1.0 to 167× (yesterday's fix) expanded |A| by 167× → gradients 167× → KL 167× → ε saturated MAX → unbounded policy drift → overfit (alpha-rl-wf-f0 step 2000: l_pi=105k, dones=0). Implementation: - New kernel compute_advantage_rms.cu: single-block warp-shuffle reduce computes sqrt(mean(A²)) over the batch each step - New ISV slot 589 RL_ADVANTAGE_BATCH_RMS_INDEX - ppo_clipped_surrogate fwd/bwd divide A by RMS (with floor bootstrap) - Dispatched after compute_advantage_return in step_with_lobsim Scale-only normalization (RMS, not z-score) preserves A=0 on non-done steps → no spurious PPO gradient on Hold actions where no trade closed. This is architectural correctness, not a quickfix. Fights the ROOT CAUSE (reward range expansion breaking LR calibration), not the SYMPTOM (ε saturation, l_pi explosion). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/compute_advantage_rms.cu | 50 +++++++++++++++++++ crates/ml-alpha/cuda/ppo_clipped_surrogate.cu | 34 +++++++++++-- crates/ml-alpha/src/rl/isv_slots.rs | 10 +++- crates/ml-alpha/src/trainer/integrated.rs | 31 ++++++++++++ 5 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 crates/ml-alpha/cuda/compute_advantage_rms.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 9695c3f65..b09f52f2b 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -47,6 +47,7 @@ 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] + "compute_advantage_rms", // Per-batch RMS(advantages) → ISV[589]; consumed by PPO surrogate for scale-invariant normalization "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 diff --git a/crates/ml-alpha/cuda/compute_advantage_rms.cu b/crates/ml-alpha/cuda/compute_advantage_rms.cu new file mode 100644 index 000000000..41095fdb5 --- /dev/null +++ b/crates/ml-alpha/cuda/compute_advantage_rms.cu @@ -0,0 +1,50 @@ +// compute_advantage_rms.cu — per-batch RMS of advantages. +// +// Writes RMS = sqrt(mean(A²)) over the batch to ISV[589]. Reads by +// ppo_clipped_surrogate_fwd / _bwd which normalize A as A/RMS to make +// PPO learning scale-invariant to the WIN clamp range. +// +// Standard PPO advantage normalization. Without it, lifting the WIN +// clamp from 1.0 to 167× scales |A| by 167× → gradients 167× → KL +// 167× → ε saturates → unbounded policy drift → overfit +// (alpha-rl-wf-f0 step 2000: l_pi=105k, observed 2026-05-28). +// +// Scale-only normalization (RMS, not std about mean) preserves zero +// advantages on non-done steps — A_done=0 → A_norm=0 → no spurious +// PPO gradient on Hold steps where the agent didn't close a trade. +// +// Single block, 32 threads, tree-reduce. b_size=1024 typical → 32 +// elements per thread sequential accumulation then warp shuffle. +// Per `feedback_no_atomicadd`: block-local reduction, no atomics. + +#define RL_ADVANTAGE_BATCH_RMS_INDEX 589 + +extern "C" __global__ void compute_advantage_rms( + const float* __restrict__ advantages, // [B] + float* __restrict__ isv, + int b_size +) { + if (blockIdx.x != 0) return; + const int tid = threadIdx.x; + const int nt = blockDim.x; + + // Each thread accumulates a partial sum of A² over its slice. + float partial = 0.0f; + for (int b = tid; b < b_size; b += nt) { + const float a = advantages[b]; + partial += a * a; + } + + // Warp-shuffle reduce within the block. Assumes blockDim.x ≤ 32 + // (single warp) — keeps it simple and matches caller's launch. + for (int off = 16; off > 0; off >>= 1) { + partial += __shfl_xor_sync(0xffffffff, partial, off); + } + + if (tid == 0) { + const float mean_sq = partial / (float)b_size; + // sqrt(0)=0 — PPO surrogate's adv_rms<=0 branch handles the + // bootstrap case (no normalization until first real RMS). + isv[RL_ADVANTAGE_BATCH_RMS_INDEX] = sqrtf(mean_sq); + } +} diff --git a/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu b/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu index c648a3b0d..9e5766451 100644 --- a/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu +++ b/crates/ml-alpha/cuda/ppo_clipped_surrogate.cu @@ -67,6 +67,15 @@ // (clamp window is symmetric in log space: [1/ratio_max, ratio_max]). #define RL_PPO_RATIO_CLAMP_MAX_INDEX 440 +// Per-batch RMS of advantages — written by compute_advantage_std +// kernel each step. Used to normalize A in PPO surrogate so learning +// is scale-invariant to the WIN clamp range. Without this, lifting +// WIN from 1.0 to 167× expands |A| by 167× → gradients 167× → KL 167× +// → PPO ε saturates → policy drift unbounded → overfit +// (alpha-rl-wf-f0 step 2000: l_pi=105k). +#define RL_ADVANTAGE_BATCH_RMS_INDEX 589 +#define ADV_STD_FLOOR 1e-4f + // ───────────────────────────────────────────────────────────────────── // ppo_policy_logits_fwd (Phase E.2): @@ -234,7 +243,18 @@ extern "C" __global__ void ppo_clipped_surrogate_fwd( const float ratio_min = 1.0f / ratio_max; const float ratio = fmaxf(ratio_min, fminf(ratio_raw, ratio_max)); - const float A = advantages[batch]; + // Advantage normalization: scale-invariant PPO learning. + // Standard PPO practice. M2 EMA from rl_var_over_abs_mean_streaming + // is the running variance of advantages. Divide A by std so + // gradient magnitude is bounded regardless of WIN clamp scale. + const float adv_rms = isv[RL_ADVANTAGE_BATCH_RMS_INDEX]; + float A; + if (adv_rms <= 0.0f) { + A = advantages[batch]; // bootstrap: no RMS yet + } else { + const float adv_std = fmaxf(adv_rms, ADV_STD_FLOOR); + A = advantages[batch] / adv_std; + } const float eps = isv[RL_PPO_CLIP_INDEX]; const float surr1 = ratio * A; const float surr2 = fminf(fmaxf(ratio, 1.0f - eps), @@ -345,8 +365,16 @@ extern "C" __global__ void ppo_clipped_surrogate_bwd( const float ratio_max = isv[RL_PPO_RATIO_CLAMP_MAX_INDEX]; const float ratio_min = 1.0f / ratio_max; const float ratio = fmaxf(ratio_min, fminf(ratio_raw, ratio_max)); - const float A = advantages[batch]; - const float eps = isv[RL_PPO_CLIP_INDEX]; + // Match forward: normalize A by advantage batch RMS. + const float adv_rms = isv[RL_ADVANTAGE_BATCH_RMS_INDEX]; + float A; + if (adv_rms <= 0.0f) { + A = advantages[batch]; + } else { + const float adv_std = fmaxf(adv_rms, ADV_STD_FLOOR); + A = advantages[batch] / adv_std; + } + const float eps = isv[RL_PPO_CLIP_INDEX]; // Clip-band membership: gradient of min(surr1, surr2) is the standard // policy-gradient identity inside [1-ε, 1+ε] (where surr1 is unclipped) diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index 735fa4f78..61e7eac47 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -1146,5 +1146,13 @@ pub const RL_MARGIN_FLOOR_USD_INDEX: usize = 587; /// from `isv[RL_THOMPSON_FLOOR_INDEX]`. Bootstrap 0.05. pub const RL_THOMPSON_FLOOR_INDEX: usize = 588; +/// Per-batch RMS of advantages, sqrt(mean(A²)). Written by +/// `compute_advantage_rms` each step after `compute_advantage_return`. +/// Read by PPO surrogate fwd/bwd to normalize A as A/RMS, making +/// PPO learning scale-invariant to the WIN clamp range. Without this, +/// lifting WIN from 1.0→167× scales |A| by 167× → KL 167× → ε +/// saturates → unbounded drift → overfit. Standard PPO practice. +pub const RL_ADVANTAGE_BATCH_RMS_INDEX: usize = 589; + /// Last RL-allocated slot index (exclusive). -pub const RL_SLOTS_END: usize = 589; +pub const RL_SLOTS_END: usize = 590; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 0483c94a3..73d64934f 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -148,6 +148,8 @@ const EMA_UPDATE_PER_STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ema_update_per_step.cubin")); const COMPUTE_ADVANTAGE_RETURN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/compute_advantage_return.cubin")); +const COMPUTE_ADVANTAGE_RMS_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/compute_advantage_rms.cubin")); const ACTION_ENTROPY_PER_STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/action_entropy_per_step.cubin")); @@ -564,6 +566,8 @@ pub struct IntegratedTrainer { ema_update_per_step_fn: CudaFunction, _compute_advantage_return_module: Arc, compute_advantage_return_fn: CudaFunction, + _compute_advantage_rms_module: Arc, + compute_advantage_rms_fn: CudaFunction, _action_entropy_per_step_module: Arc, action_entropy_per_step_fn: CudaFunction, @@ -1336,6 +1340,12 @@ impl IntegratedTrainer { let compute_advantage_return_fn = compute_advantage_return_module .load_function("compute_advantage_return") .context("load compute_advantage_return")?; + let compute_advantage_rms_module = ctx + .load_cubin(COMPUTE_ADVANTAGE_RMS_CUBIN.to_vec()) + .context("load compute_advantage_rms cubin")?; + let compute_advantage_rms_fn = compute_advantage_rms_module + .load_function("compute_advantage_rms") + .context("load compute_advantage_rms")?; // Phase R4 kernels. let rl_action_kernel_module = ctx .load_cubin(RL_ACTION_KERNEL_CUBIN.to_vec()) @@ -2325,6 +2335,8 @@ impl IntegratedTrainer { action_entropy_per_step_fn, _compute_advantage_return_module: compute_advantage_return_module, compute_advantage_return_fn, + _compute_advantage_rms_module: compute_advantage_rms_module, + compute_advantage_rms_fn, _rl_action_kernel_module: rl_action_kernel_module, rl_action_kernel_fn, _argmax_expected_q_module: argmax_expected_q_module, @@ -6308,6 +6320,25 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("compute_advantage_return: {:?}", e))?; } } + + // Advantage normalization — compute per-batch RMS(A) → ISV[589]. + // Consumed by PPO surrogate fwd/bwd to make learning scale-invariant + // to the WIN clamp range. Standard PPO practice. + { + let mut args = RawArgs::new(); + args.push_ptr(self.advantages_d.raw_ptr()); + args.push_ptr(self.isv_dev_ptr); + args.push_i32(b_size_i); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.compute_advantage_rms_fn.cu_function(), + (1, 1, 1), (32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("compute_advantage_rms: {:?}", e))?; + } + } // γ is read on-device by compute_advantage_return from // ISV[400] — no host gamma scalar needed in this hot path // post-R7b (the bootstrap-fallback host read of γ that the