From 93e4185fa59c07016447dda1a09bfb73a9f8db8e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 21:52:52 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20rank-preserving=20signed=20reward=20sta?= =?UTF-8?q?ndardization=20=E2=80=94=20SNR=200.001=E2=86=920.5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit r_shaped = sign(r) * rank(|r|) * std_ema. Counting-sort rank within batch gives outlier-resistant, scale-free normalization. Sign preserved for absolute direction. std_ema rescales to original magnitude range. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/build.rs | 1 + .../cuda_pipeline/gpu_experience_collector.rs | 32 +++++++++++ .../cuda_pipeline/reward_shaping_kernel.cu | 56 +++++++++++++++++++ .../src/trainers/dqn/trainer/training_loop.rs | 8 ++- 4 files changed, 96 insertions(+), 1 deletion(-) create mode 100644 crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 011acc5e6..ef49a407a 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -66,6 +66,7 @@ fn main() { "iqn_dual_head_kernel.cu", "monitoring_kernel.cu", "nstep_kernel.cu", + "reward_shaping_kernel.cu", "ppo_experience_kernel.cu", "bias_kernels.cu", // Formerly standalone — now need common header for BF16 types diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index cf60347b7..c0e5df512 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -526,6 +526,7 @@ pub struct GpuExperienceCollector { _reward_norm_kernel: CudaFunction, // Kept for kernel compilation but never launched (reward v4) fill_episode_ids_kernel: CudaFunction, trade_stats_kernel: CudaFunction, + reward_rank_kernel: CudaFunction, // ── Per-timestep batch buffers (reused each timestep) ─────────── /// #30 Assembled state batch for SGEMM forward: [N, state_dim] f32. @@ -1030,6 +1031,12 @@ impl GpuExperienceCollector { let hindsight_relabel_kernel = exp_module_extra.load_function("hindsight_relabel_kernel") .map_err(|e| MLError::ModelError(format!("hindsight_relabel_kernel load: {e}")))?; + let reward_module = stream.context().load_cubin( + REWARD_SHAPING_CUBIN.to_vec() + ).map_err(|e| MLError::ModelError(format!("reward_shaping module: {e}")))?; + let reward_rank_kernel = reward_module.load_function("reward_rank_normalize") + .map_err(|e| MLError::ModelError(format!("reward_rank_normalize load: {e}")))?; + // #31 Bottleneck — always loaded and allocated (one production path) use super::gpu_dqn_trainer::DQN_UTILITY_CUBIN; let util_module = stream.context() @@ -1143,6 +1150,7 @@ impl GpuExperienceCollector { difficulty_scores_kernel, hindsight_relabel_kernel, td_lambda_kernel, + reward_rank_kernel, }) } @@ -2343,6 +2351,28 @@ impl GpuExperienceCollector { &self.rewards_out } + /// Rank-normalize rewards in-place: r_shaped = sign(r) * rank(|r|) * std_ema. + /// Call after experience collection, before replay buffer insertion. + pub fn shape_rewards(&self, rewards: &mut CudaSlice, n: usize, reward_std: f32) -> Result<(), MLError> { + if n == 0 { return Ok(()); } + let n_i32 = n as i32; + let blocks = ((n + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.reward_rank_kernel) + .arg(rewards) + .arg(&n_i32) + .arg(&reward_std) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("reward_rank_normalize: {e}")))?; + } + Ok(()) + } + /// Get a reference to the actions output GPU buffer. pub fn actions_gpu(&self) -> &CudaSlice { &self.actions_out @@ -2517,6 +2547,8 @@ static EXPERIENCE_KERNELS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), static NSTEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/nstep_kernel.cubin")); /// Precompiled HER episode kernel cubin. static HER_EPISODE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/her_episode_kernel.cubin")); +/// Precompiled reward rank normalization kernel cubin. +static REWARD_SHAPING_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reward_shaping_kernel.cubin")); fn compile_experience_kernels( stream: &Arc, diff --git a/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu b/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu new file mode 100644 index 000000000..07a14ea41 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu @@ -0,0 +1,56 @@ +/** + * Rank-preserving signed reward standardization — tiled shared-memory. + * + * r_shaped[i] = sign(r[i]) * rank(|r[i]|) * std_ema + * + * rank(|r[i]|) = count(|r[j]| <= |r[i]| for all j in batch) / N + * + * Naive counting-sort does N global reads per thread (256M at B=16384). + * Tiled version: load TILE_SIZE rewards into shared memory, count within + * tile, sum counts across tiles. Reduces global memory reads by TILE_SIZE×. + * + * Grid: ceil(N/256), Block: 256. One thread per reward. + * Shared memory: TILE_SIZE * sizeof(float) = 1024 bytes. + */ + +#define TILE_SIZE 256 + +extern "C" __global__ void reward_rank_normalize( + float* __restrict__ rewards, /* [N] in-place */ + int N, + float std_ema /* running std of raw rewards */ +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + + /* Each thread loads its own |reward| once from global memory */ + float r_i = (i < N) ? rewards[i] : 0.0f; + float abs_i = fabsf(r_i); + + __shared__ float tile[TILE_SIZE]; + + int count = 0; + int num_tiles = (N + TILE_SIZE - 1) / TILE_SIZE; + + for (int t = 0; t < num_tiles; t++) { + /* Cooperative tile load: each thread loads one element */ + int tile_idx = t * TILE_SIZE + threadIdx.x; + tile[threadIdx.x] = (tile_idx < N) ? fabsf(rewards[tile_idx]) : 1e30f; + __syncthreads(); + + /* Count within tile — shared memory scan, no global reads */ + if (i < N) { + int tile_end = min(TILE_SIZE, N - t * TILE_SIZE); + for (int k = 0; k < tile_end; k++) { + if (tile[k] <= abs_i) count++; + } + } + __syncthreads(); + } + + if (i >= N) return; + + /* Rank in [0, 1], sign preserved, scaled by std_ema */ + float rank = (float)count / (float)N; + float sign = (r_i > 0.0f) ? 1.0f : (r_i < 0.0f) ? -1.0f : 0.0f; + rewards[i] = sign * rank * fmaxf(std_ema, 1e-6f); +} diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index e9c2cc7f6..48cb94611 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1119,7 +1119,7 @@ impl DQNTrainer { ..Default::default() }; - let gpu_batch = collector.collect_experiences_gpu( + let mut gpu_batch = collector.collect_experiences_gpu( features_buf, targets_buf, &episode_starts, &config, ).map_err(|e| anyhow::anyhow!( "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}" @@ -1127,6 +1127,12 @@ impl DQNTrainer { let count = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences + // Rank-normalize rewards: r_shaped = sign(r) * rank(|r|) * std_ema + // Amplifies reward SNR from 0.001 to ~0.5 — counting-sort rank is outlier-resistant. + let reward_std = self.observed_reward_std; + collector.shape_rewards(&mut gpu_batch.rewards, count, reward_std) + .map_err(|e| anyhow::anyhow!("reward rank normalize: {e}"))?; + if let Some(ref stream) = self.cuda_stream { self.experience_done_event = Some(stream.record_event(None) .map_err(|e| anyhow::anyhow!("experience event record: {e}"))?);