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

fix(rl): remove advantage normalization — sparse rewards diluted signal

Browse Source 
With 5-6% of batch elements having real rewards per step, advantage
normalization (zero mean, unit var) turned 94% of the gradient into
noise. The /B gradient normalization already handles batch-size
invariance. Raw V-advantage bounded by reward_scale is sufficient.

Removed: normalize_advantages kernel (#if 0 in .cu), field, load,
both launch sites, q_logits_target_st_d (Q-advantage dead code).
ISV-driven TAU_MAX (slot 573, bootstrap 0.005) retained for future use.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-27 00:20:27 +02:00
parent d743336060
commit 514b04bee3
2 changed files with 54 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
crates/ml-alpha/cuda/compute_advantage_return.cu
View File

@@ -56,40 +56,79 @@ extern "C" __global__ void compute_advantage_return(
// all-positive or all-negative, causing PPO to reinforce/suppress ALL
// actions uniformly instead of discriminating between them.
// DISABLED: advantage normalization causes issues with sparse rewards
// (5% of batch has real rewards, 95% is noise). The /B gradient
// normalization in ppo_clipped_surrogate_bwd already handles batch-
// size invariance. Raw V-advantages are bounded by reward_scale.
//
// Kept for reference but not launched from Rust.
#if 0
// Normalize advantages using only done-steps (where dones > 0.5) for
// mean/variance computation. Non-done steps have reward=0 and advantage
// ≈ γV(s')-V(s) ≈ noise. Normalizing over the full batch dilutes the
// sparse reward signal (5-6% of batch has real rewards) with 94% noise.
//
// Non-done steps still get their advantage normalized (centered/scaled),
// but the centering point and scale come from the DONE steps' advantage
// distribution — so the real reward signal drives the normalization.
// If zero done steps in a batch, fall back to full-batch normalization.
extern "C" __global__ void normalize_advantages(
float* __restrict__ advantages, // [b_size] IN/OUT
int b_size
float* __restrict__ advantages, // [b_size] IN/OUT
const float* __restrict__ dones, // [b_size] 0.0/1.0
int b_size
) {
extern __shared__ float smem[]; // [2 * blockDim.x]
float* s_sum = smem;
float* s_sum2 = smem + blockDim.x;
extern __shared__ float smem[]; // [3 * blockDim.x]
float* s_sum = smem;
float* s_sum2 = smem + blockDim.x;
float* s_count = smem + 2 * 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;
float val = 0.0f;
float is_done = 0.0f;
if (tid < b_size) {
val = advantages[tid];
is_done = (dones[tid] > 0.5f) ? 1.0f : 0.0f;
}
s_sum[tid] = val * is_done;
s_sum2[tid] = val * val * is_done;
s_count[tid] = is_done;
__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];
s_sum[tid] += s_sum[tid + stride];
s_sum2[tid] += s_sum2[tid + stride];
s_count[tid] += s_count[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));
float n = s_count[0];
if (n < 1.0f) {
// No done steps — fall back to full-batch stats.
// Recompute from all elements (rare path).
float sum_all = 0.0f, sum2_all = 0.0f;
for (int i = 0; i < b_size; i++) {
sum_all += advantages[i];
sum2_all += advantages[i] * advantages[i];
}
s_mean = sum_all / (float)b_size;
float var = sum2_all / (float)b_size - s_mean * s_mean;
s_std = sqrtf(fmaxf(var, 1e-8f));
} else {
s_mean = s_sum[0] / n;
float var = s_sum2[0] / n - 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;
}
}
#endif

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

@@ -562,7 +562,6 @@ pub struct IntegratedTrainer {
ema_update_per_step_fn: CudaFunction,
_compute_advantage_return_module: Arc<CudaModule>,
compute_advantage_return_fn: CudaFunction,
normalize_advantages_fn: CudaFunction,
// ── Phase R4: GPU action sampling ─────────────────────────────────
_rl_action_kernel_module: Arc<CudaModule>,
@@ -813,11 +812,6 @@ pub struct IntegratedTrainer {
/// C51 Q logits at h_t `[B × N_ACTIONS × Q_N_ATOMS]`.
pub q_logits_d: CudaSlice<f32>,
/// C51 TARGET Q logits at h_t — for stable PPO advantage signal.
/// Updated via slow-moving target network (τ=0.005 soft update),
/// preventing the phase-lag oscillation where π chases a rapidly
/// shifting online Q ranking.
pub q_logits_target_st_d: CudaSlice<f32>,
/// C51 Q logits at h_{t+1} `[B × N_ACTIONS × Q_N_ATOMS]`.
pub q_logits_tp1_d: CudaSlice<f32>,
/// Scalar value prediction at h_t `[B]`.
@@ -1322,10 +1316,6 @@ 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
.load_cubin(RL_ACTION_KERNEL_CUBIN.to_vec())
@@ -1845,9 +1835,6 @@ impl IntegratedTrainer {
let q_logits_d = stream
.alloc_zeros::<f32>(b_size * k_dqn_alloc)
.context("alloc q_logits_d")?;
let q_logits_target_st_d = stream
.alloc_zeros::<f32>(b_size * k_dqn_alloc)
.context("alloc q_logits_target_st_d")?;
let q_logits_tp1_d = stream
.alloc_zeros::<f32>(b_size * k_dqn_alloc)
.context("alloc q_logits_tp1_d")?;
@@ -2307,7 +2294,6 @@ 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,
@@ -2425,7 +2411,6 @@ impl IntegratedTrainer {
iqn_expected_q_d,
ensemble_q_d,
q_logits_d,
q_logits_target_st_d,
q_logits_tp1_d,
v_pred_d,
v_pred_tp1_d,
@@ -3669,23 +3654,6 @@ 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::<f32>()) 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(())
}
@@ -6202,23 +6170,6 @@ impl IntegratedTrainer {
).map_err(|e| anyhow::anyhow!("compute_advantage_return: {:?}", 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::<f32>()) 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
@@ -6979,9 +6930,6 @@ impl IntegratedTrainer {
self.dqn_head
.forward(h_t_borrow, b_size, &mut self.q_logits_d)
.context("step_with_lobsim_gpu: dqn_head.forward(h_t)")?;
self.dqn_head
.forward_target(h_t_borrow, b_size, &mut self.q_logits_target_st_d)
.context("step_with_lobsim_gpu: dqn_head.forward_target(h_t) for stable advantage")?;
self.dqn_head
.forward(&self.h_tp1_d, b_size, &mut self.q_logits_tp1_d)
.context("step_with_lobsim_gpu: dqn_head.forward(h_tp1)")?;