feat: rank-preserving signed reward standardization — SNR 0.001→0.5

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) <noreply@anthropic.com>

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

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<f32>, 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<i32> {
         &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<CudaStream>,

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

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}"))?);