gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(rl): advantage normalization — scale-invariant PPO learning

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-28 10:39:12 +02:00
parent 6f21513319
commit 6d33f18b75
5 changed files with 122 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

@@ -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

50
crates/ml-alpha/cuda/compute_advantage_rms.cu Normal file
View File

@@ -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);
}
}

34
crates/ml-alpha/cuda/ppo_clipped_surrogate.cu
View File

@@ -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)

10
crates/ml-alpha/src/rl/isv_slots.rs
View File

@@ -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;

31
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -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<CudaModule>,
compute_advantage_return_fn: CudaFunction,
_compute_advantage_rms_module: Arc<CudaModule>,
compute_advantage_rms_fn: CudaFunction,
_action_entropy_per_step_module: Arc<CudaModule>,
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