diff --git a/crates/ml-alpha/cuda/compute_advantage_return.cu b/crates/ml-alpha/cuda/compute_advantage_return.cu index a4d26b714..3775c2dd6 100644 --- a/crates/ml-alpha/cuda/compute_advantage_return.cu +++ b/crates/ml-alpha/cuda/compute_advantage_return.cu @@ -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 diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index efa7844a2..7f7c2eaa4 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -562,7 +562,6 @@ 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, @@ -813,11 +812,6 @@ pub struct IntegratedTrainer { /// C51 Q logits at h_t `[B × N_ACTIONS × Q_N_ATOMS]`. pub q_logits_d: CudaSlice, - /// 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, /// C51 Q logits at h_{t+1} `[B × N_ACTIONS × Q_N_ATOMS]`. pub q_logits_tp1_d: CudaSlice, /// 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::(b_size * k_dqn_alloc) .context("alloc q_logits_d")?; - let q_logits_target_st_d = stream - .alloc_zeros::(b_size * k_dqn_alloc) - .context("alloc q_logits_target_st_d")?; let q_logits_tp1_d = stream .alloc_zeros::(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::()) 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::()) 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)")?;