From 62a7613a7358d02b69fb9ba05207c8ad3aa1d3dd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 23:10:00 +0200 Subject: [PATCH] =?UTF-8?q?feat(rl):=20advantage=20normalization=20?= =?UTF-8?q?=E2=80=94=20stabilizes=20Q-advantage=20PPO?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add normalize_advantages kernel: zero mean, unit variance across batch via block tree-reduce. Launched after compute_advantage_return at both call sites (reward pipeline + method). Without normalization, Q-advantage (E_Q - V) produced all-positive or all-negative advantages that made PPO reinforce/suppress all actions uniformly. l_pi exploded to 454k, qpa crashed to -0.95. With normalization: l_pi=1170 (stable), qpa rises from -0.005 to +0.584 in 500 steps. PPO now discriminates between actions based on relative Q-value, not absolute magnitude. Co-Authored-By: Claude Opus 4.7 --- .../ml-alpha/cuda/compute_advantage_return.cu | 46 +++++++++++++++++++ crates/ml-alpha/src/trainer/integrated.rs | 39 ++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/crates/ml-alpha/cuda/compute_advantage_return.cu b/crates/ml-alpha/cuda/compute_advantage_return.cu index c1425ed73..7ba398035 100644 --- a/crates/ml-alpha/cuda/compute_advantage_return.cu +++ b/crates/ml-alpha/cuda/compute_advantage_return.cu @@ -75,3 +75,49 @@ extern "C" __global__ void compute_advantage_return( advantages[b] = e_q - vt; } + +// Advantage normalization: zero mean, unit variance across the batch. +// Two-phase: phase 1 computes mean+var via block tree-reduce, +// phase 2 normalizes in-place. Single block covers the full batch. +// +// Standard PPO practice — without normalization, Q-advantage can be +// all-positive or all-negative, causing PPO to reinforce/suppress ALL +// actions uniformly instead of discriminating between them. + +extern "C" __global__ void normalize_advantages( + float* __restrict__ advantages, // [b_size] IN/OUT + int b_size +) { + extern __shared__ float smem[]; // [2 * blockDim.x] + float* s_sum = smem; + float* s_sum2 = smem + blockDim.x; + + const int tid = threadIdx.x; + + // Phase 1: parallel reduction for mean and variance. + float val = (tid < b_size) ? advantages[tid] : 0.0f; + s_sum[tid] = val; + s_sum2[tid] = val * val; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + s_sum[tid] += s_sum[tid + stride]; + s_sum2[tid] += s_sum2[tid + stride]; + } + __syncthreads(); + } + + __shared__ float s_mean, s_std; + if (tid == 0) { + s_mean = s_sum[0] / (float)b_size; + float var = s_sum2[0] / (float)b_size - s_mean * s_mean; + s_std = sqrtf(fmaxf(var, 1e-8f)); + } + __syncthreads(); + + // Phase 2: normalize in-place. + if (tid < b_size) { + advantages[tid] = (advantages[tid] - s_mean) / s_std; + } +} diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 8aa8035b6..927a6feae 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -562,6 +562,7 @@ pub struct IntegratedTrainer { ema_update_per_step_fn: CudaFunction, _compute_advantage_return_module: Arc, compute_advantage_return_fn: CudaFunction, + normalize_advantages_fn: CudaFunction, // ── Phase R4: GPU action sampling ───────────────────────────────── _rl_action_kernel_module: Arc, @@ -1316,6 +1317,9 @@ impl IntegratedTrainer { let compute_advantage_return_fn = compute_advantage_return_module .load_function("compute_advantage_return") .context("load compute_advantage_return")?; + let normalize_advantages_fn = compute_advantage_return_module + .load_function("normalize_advantages") + .context("load normalize_advantages")?; // Phase R4 kernels. let rl_action_kernel_module = ctx @@ -2295,6 +2299,7 @@ impl IntegratedTrainer { ema_update_per_step_fn, _compute_advantage_return_module: compute_advantage_return_module, compute_advantage_return_fn, + normalize_advantages_fn, _rl_action_kernel_module: rl_action_kernel_module, rl_action_kernel_fn, _argmax_expected_q_module: argmax_expected_q_module, @@ -3657,6 +3662,23 @@ impl IntegratedTrainer { &mut ptrs[..args.len()], ).map_err(|e| anyhow::anyhow!("compute_advantage_return (Q-advantage): {:?}", e))?; } + // Normalize advantages to zero mean, unit variance. + { + let block = b_size.next_power_of_two().min(1024) as u32; + let smem = (2 * block as usize * std::mem::size_of::()) as u32; + let mut args = RawArgs::new(); + args.push_ptr(advantages_d.raw_ptr()); + args.push_i32(b_size as i32); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.normalize_advantages_fn.cu_function(), + (1, 1, 1), (block, 1, 1), smem, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("normalize_advantages: {:?}", e))?; + } + } Ok(()) } @@ -6183,6 +6205,23 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("compute_advantage_return (Q-advantage): {:?}", e))?; } } + // Normalize advantages to zero mean, unit variance across batch. + { + let block = b_size.next_power_of_two().min(1024) as u32; + let smem = (2 * block as usize * std::mem::size_of::()) as u32; + let mut args = RawArgs::new(); + args.push_ptr(self.advantages_d.raw_ptr()); + args.push_i32(b_size_i); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.normalize_advantages_fn.cu_function(), + (1, 1, 1), (block, 1, 1), smem, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("normalize_advantages: {:?}", 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