diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index 5fbe49fc9..6eb4d910e 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -141,6 +141,7 @@ causal_intensity = 1.0 [reward] # v8 comprehensive training overhaul popart_enabled = true +popart_robust = true # median/IQR instead of mean/var — robust for bimodal rewards micro_reward_scale = 0.1 td_lambda = 0.9 max_trace_length = 7 diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index d047da9f3..7ed1f6906 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -1513,3 +1513,29 @@ extern "C" __global__ void popart_normalize_kernel( } } } + +/* ══════════════════════════════════════════════════════════════════════ + * ROBUST POPART REWARD NORMALIZATION KERNEL + * + * Normalizes rewards in-place using pre-computed median and IQR + * (interquartile range). More robust than Welford mean/var for + * heavy-tailed bimodal distributions (many ±0.1 micro-rewards + + * few ±5.0 trade exits). + * + * Median/IQR are computed once per epoch from the full sorted + * experience collection rewards (tiny 3-float DtoH), then passed + * as scalars to this kernel for every training batch. + * + * Launch: grid=ceil(N/256), block=256. + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void popart_normalize_robust( + float* __restrict__ rewards, /* [N] in-place normalization */ + float median, /* pre-computed from sorted rewards */ + float iqr, /* interquartile range (Q75 - Q25) */ + int N +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + float safe_iqr = fmaxf(iqr, 1e-6f); + rewards[i] = (rewards[i] - median) / safe_iqr; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 307320f07..edecd3bb8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1336,6 +1336,8 @@ pub struct GpuDqnTrainer { /// v8: PopArt reward normalization kernel (running mean/variance). popart_normalize_kernel: CudaFunction, + /// v8: Robust PopArt kernel (median/IQR normalization). + popart_robust_kernel: CudaFunction, /// v8: PopArt running mean [1]. popart_mean: CudaSlice, /// v8: PopArt running variance [1]. @@ -1787,6 +1789,31 @@ impl GpuDqnTrainer { Ok(()) } + /// Warm-start C51 atom positions from empirical reward quantiles. + /// + /// Called once per epoch after experience collection, before training steps. + /// Writes the same quantile values to all 4 branches — per-branch specialization + /// happens via the existing `step_atom_positions` SGD optimizer. + /// + /// The HtoD transfer is ~1us for 4 * num_atoms floats (typically 4 * 51 = 204). + pub(crate) fn warm_start_atom_positions(&mut self, quantiles: &[f32]) -> Result<(), MLError> { + let na = self.config.num_atoms; + if quantiles.len() != na { + return Ok(()); // skip on mismatch + } + + // Write same quantiles to all 4 branches + let mut host_data = vec![0.0_f32; 4 * na]; + for branch in 0..4 { + host_data[branch * na..(branch + 1) * na].copy_from_slice(quantiles); + } + + self.stream.memcpy_htod(&host_data, &mut self.atom_positions_buf) + .map_err(|e| MLError::ModelError(format!("warm_start_atom_positions HtoD: {e}")))?; + + Ok(()) + } + /// SGD update on adaptive atom spacing_raw parameters. /// Maximizes atom entropy (gradient ascent) — concentrates atoms where mass exists. /// Uses scale_f32_kernel to decay spacing_raw toward zero: softmax(0,...,0) = uniform. @@ -4797,6 +4824,28 @@ impl GpuDqnTrainer { Ok(()) } + /// Normalize the internal rewards_buf in-place using pre-computed median and IQR. + /// + /// More robust than Welford mean/var for heavy-tailed bimodal distributions. + /// Median/IQR are computed once per epoch from sorted experience rewards. + /// Called instead of `normalize_rewards_popart_inplace` when `popart_robust = true`. + pub fn normalize_rewards_robust( + &self, n: usize, median: f32, iqr: f32, + ) -> Result<(), MLError> { + let rewards_ptr = self.rewards_buf.raw_ptr(); + let n_i32 = n as i32; + unsafe { + self.stream.launch_builder(&self.popart_robust_kernel) + .arg(&rewards_ptr) + .arg(&median) + .arg(&iqr) + .arg(&n_i32) + .launch(LaunchConfig::for_num_elems(n as u32)) + .map_err(|e| MLError::ModelError(format!("popart_normalize_robust: {e}")))?; + } + Ok(()) + } + /// Apply spectral normalization to all 10 weight matrices (trunk + 8 heads). /// /// One step of power iteration per call (standard practice — single step @@ -5016,7 +5065,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile utility kernels (grad_norm, adam_update, etc.) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed) = + let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (outside CUDA graph). @@ -7195,6 +7244,7 @@ impl GpuDqnTrainer { curiosity_w2_ptr: u64::MAX, curiosity_b2_ptr: u64::MAX, popart_normalize_kernel, + popart_robust_kernel, popart_mean, popart_var, popart_count, @@ -12205,7 +12255,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -12331,9 +12381,11 @@ fn compile_training_kernels( let popart_normalize = module.load_function("popart_normalize_kernel") .map_err(|e| MLError::ModelError(format!("popart_normalize_kernel load: {e}")))?; + let popart_robust = ungraphed_module.load_function("popart_normalize_robust") + .map_err(|e| MLError::ModelError(format!("popart_normalize_robust load: {e}")))?; - info!("GpuDqnTrainer: 35 utility kernels loaded from precompiled cubin (5 CUmodules)"); - Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed)) + info!("GpuDqnTrainer: 36 utility kernels loaded from precompiled cubin (5 CUmodules)"); + Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 520510f5d..a70fb2b5f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -558,6 +558,12 @@ pub struct GpuExperienceCollector { reward_sort_indices_buf: CudaSlice, /// Allocated capacity of sort buffers (next power-of-2 of max_n). sort_buf_capacity: usize, + /// Kernel: copy raw rewards to sort keys (no |sharpe| transform) for quantile extraction. + reward_copy_raw_kernel: CudaFunction, + /// Kernel: gather evenly-spaced quantile positions from sorted array. + gather_quantiles_kernel: CudaFunction, + /// Scratch buffer for gather_quantiles output: [max_atoms] f32. + quantiles_scratch_buf: CudaSlice, // ── Per-timestep batch buffers (reused each timestep) ─────────── /// #30 Assembled state batch for SGEMM forward: [N, state_dim] f32. @@ -611,6 +617,10 @@ pub struct GpuExperienceCollector { /// True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}. raw_returns_out: CudaSlice, // [alloc_episodes * alloc_timesteps] + /// Last experience count from `collect_experiences_gpu` (episodes * timesteps * cf_mult). + /// Set after each collection for downstream consumers (e.g. quantile extraction). + last_experience_count: usize, + /// Running reward statistics for cross-epoch normalization (Welford's online algorithm) // Persistent epoch state — survives across kernel launches. @@ -1131,6 +1141,10 @@ impl GpuExperienceCollector { .map_err(|e| MLError::ModelError(format!("bitonic_sort_step load: {e}")))?; let reward_scatter_rank_kernel = reward_module.load_function("reward_scatter_rank") .map_err(|e| MLError::ModelError(format!("reward_scatter_rank load: {e}")))?; + let reward_copy_raw_kernel = reward_module.load_function("reward_copy_raw_for_sort") + .map_err(|e| MLError::ModelError(format!("reward_copy_raw_for_sort load: {e}")))?; + let gather_quantiles_kernel = reward_module.load_function("gather_quantiles") + .map_err(|e| MLError::ModelError(format!("gather_quantiles load: {e}")))?; // Sort buffers: sized to next power-of-2 of max possible N (episodes × timesteps × cf_mult). // Bitonic sort requires power-of-2 sized arrays. @@ -1140,6 +1154,9 @@ impl GpuExperienceCollector { .map_err(|e| MLError::ModelError(format!("alloc reward_sort_keys_buf: {e}")))?; let reward_sort_indices_buf = stream.alloc_zeros::(sort_buf_capacity) .map_err(|e| MLError::ModelError(format!("alloc reward_sort_indices_buf: {e}")))?; + // Quantiles scratch: max num_atoms is typically 51-128, allocate 256 as safe upper bound. + let quantiles_scratch_buf = stream.alloc_zeros::(256) + .map_err(|e| MLError::ModelError(format!("alloc quantiles_scratch_buf: {e}")))?; // #31 Bottleneck — always loaded and allocated (one production path) use super::gpu_dqn_trainer::DQN_UTILITY_CUBIN; @@ -1235,6 +1252,7 @@ impl GpuExperienceCollector { rewards_out, done_out, raw_returns_out, + last_experience_count: 0, epoch_state, reset_flags: 0, expert_actions_gpu: None, @@ -1290,6 +1308,9 @@ impl GpuExperienceCollector { reward_sort_keys_buf, reward_sort_indices_buf, sort_buf_capacity, + reward_copy_raw_kernel, + gather_quantiles_kernel, + quantiles_scratch_buf, }) } @@ -1838,6 +1859,8 @@ impl GpuExperienceCollector { "GPU experience collection complete (cuBLAS timestep loop)" ); + self.last_experience_count = total; + Ok(GpuExperienceBatch { states, next_states, @@ -2592,6 +2615,9 @@ impl GpuExperienceCollector { /// Allocated episode buffer capacity. pub fn alloc_episodes(&self) -> usize { self.alloc_episodes } + /// Last experience count from the most recent `collect_experiences_gpu` call. + pub fn last_experience_count(&self) -> usize { self.last_experience_count } + /// Get a reference to the CUDA stream used by this collector. pub fn stream(&self) -> &Arc { &self.stream @@ -2695,6 +2721,140 @@ impl GpuExperienceCollector { Ok(()) } + /// Sort collected rewards and extract `num_atoms` evenly-spaced quantiles. + /// + /// Returns a `Vec` of quantile positions on the HOST (small DtoH: num_atoms floats). + /// Uses the bitonic sort infrastructure to sort raw rewards, then a gather kernel + /// to extract evenly-spaced positions from the sorted array. + /// + /// Used to warm-start C51 adaptive atom positions each epoch. + pub fn compute_reward_quantiles(&mut self, n: usize, num_atoms: usize) -> Result, MLError> { + if n == 0 || num_atoms == 0 { + return Ok(vec![0.0; num_atoms]); + } + if num_atoms > 256 { + return Err(MLError::ModelError(format!( + "compute_reward_quantiles: num_atoms={num_atoms} > 256 scratch limit" + ))); + } + + let n_padded = n.next_power_of_two(); + if n_padded > self.sort_buf_capacity { + return Err(MLError::ModelError(format!( + "reward sort buffer overflow: n_padded={n_padded} > capacity={}", + self.sort_buf_capacity + ))); + } + + let n_i32 = n as i32; + let n_padded_i32 = n_padded as i32; + let num_atoms_i32 = num_atoms as i32; + + // Step 1: Copy raw rewards to sort keys + fill indices (no |sharpe| transform) + let pad_blocks = ((n_padded + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.reward_copy_raw_kernel) + .arg(&self.rewards_out) + .arg(&self.reward_sort_keys_buf) + .arg(&self.reward_sort_indices_buf) + .arg(&n_i32) + .arg(&n_padded_i32) + .launch(LaunchConfig { + grid_dim: (pad_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("reward_copy_raw_for_sort: {e}")))?; + } + + // Step 2: Bitonic sort — O(log^2 N) kernel launches + let half_n = n_padded / 2; + let sort_blocks = ((half_n + 255) / 256) as u32; + + let mut stage: i32 = 1; + while (stage as usize) < n_padded { + let mut step = stage; + while step >= 1 { + unsafe { + self.stream + .launch_builder(&self.bitonic_sort_step_kernel) + .arg(&self.reward_sort_keys_buf) + .arg(&self.reward_sort_indices_buf) + .arg(&n_padded_i32) + .arg(&stage) + .arg(&step) + .launch(LaunchConfig { + grid_dim: (sort_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("bitonic_sort_step (quantiles): {e}")))?; + } + step >>= 1; + } + stage <<= 1; + } + + // Step 3: Gather evenly-spaced quantile positions from sorted array + let gather_blocks = ((num_atoms + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.gather_quantiles_kernel) + .arg(&self.reward_sort_keys_buf) // sorted rewards + .arg(&self.quantiles_scratch_buf) // output [num_atoms] + .arg(&n_i32) + .arg(&num_atoms_i32) + .launch(LaunchConfig { + grid_dim: (gather_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("gather_quantiles: {e}")))?; + } + + // Step 4: DtoH — read num_atoms floats from quantiles scratch + let quantile_slice = self.quantiles_scratch_buf.slice(..num_atoms); + let mut host_quantiles = vec![0.0_f32; num_atoms]; + self.stream.memcpy_dtoh(&quantile_slice, &mut host_quantiles) + .map_err(|e| MLError::ModelError(format!("dtoh quantiles: {e}")))?; + + Ok(host_quantiles) + } + + /// Extract median, Q25, Q75 from the already-sorted reward_sort_keys_buf. + /// + /// MUST be called AFTER `compute_reward_quantiles()` which sorts the rewards. + /// Reads 3 specific positions from the sorted GPU buffer (tiny DtoH: 3 floats). + /// Returns `(median, iqr)` where `iqr = Q75 - Q25`. + pub fn extract_median_iqr(&self, n: usize) -> Result<(f32, f32), MLError> { + if n < 4 { + return Ok((0.0, 1.0)); // degenerate: skip normalization + } + let idx_q25 = n / 4; + let idx_median = n / 2; + let idx_q75 = 3 * n / 4; + + // Read 3 specific floats from sorted GPU buffer + let mut q25_host = [0.0_f32]; + let mut median_host = [0.0_f32]; + let mut q75_host = [0.0_f32]; + + let q25_slice = self.reward_sort_keys_buf.slice(idx_q25..idx_q25 + 1); + let median_slice = self.reward_sort_keys_buf.slice(idx_median..idx_median + 1); + let q75_slice = self.reward_sort_keys_buf.slice(idx_q75..idx_q75 + 1); + + self.stream.memcpy_dtoh(&q25_slice, &mut q25_host) + .map_err(|e| MLError::ModelError(format!("dtoh Q25: {e}")))?; + self.stream.memcpy_dtoh(&median_slice, &mut median_host) + .map_err(|e| MLError::ModelError(format!("dtoh median: {e}")))?; + self.stream.memcpy_dtoh(&q75_slice, &mut q75_host) + .map_err(|e| MLError::ModelError(format!("dtoh Q75: {e}")))?; + + let iqr = q75_host[0] - q25_host[0]; + Ok((median_host[0], iqr)) + } + /// Get a reference to the actions output GPU buffer. pub fn actions_gpu(&self) -> &CudaSlice { &self.actions_out diff --git a/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu b/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu index 9d67778bc..4cb729155 100644 --- a/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu +++ b/crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu @@ -112,6 +112,53 @@ extern "C" __global__ void reward_scatter_rank( rewards_out[orig_idx] = sign * rank; } +/* ── Step 1b: Copy raw rewards to sort keys (for quantile extraction) ─ + * Unlike reward_compute_abs_sharpe which takes |sharpe|, this copies raw + * reward values so that bitonic sort produces a properly-ordered array + * from which empirical quantiles can be read. + * Grid: ceil(N_padded/256), Block: 256. */ +extern "C" __global__ void reward_copy_raw_for_sort( + const float* __restrict__ rewards_in, /* [N] raw rewards */ + float* __restrict__ sort_keys_out, /* [N_padded] sort keys */ + int* __restrict__ indices_out, /* [N_padded] original indices */ + int N, + int N_padded +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N_padded) return; + + if (i < N) { + sort_keys_out[i] = rewards_in[i]; + indices_out[i] = i; + } else { + sort_keys_out[i] = 3.4e38f; /* padding sorts to end */ + indices_out[i] = -1; + } +} + +/* ── Gather quantiles from sorted array ────────────────────────────── + * Extracts num_atoms evenly-spaced quantile positions from a sorted array + * into a contiguous output buffer for a single DtoH transfer. + * Grid: ceil(num_atoms/256), Block: 256. */ +extern "C" __global__ void gather_quantiles( + const float* __restrict__ sorted, /* [N] sorted rewards */ + float* __restrict__ quantiles, /* [num_atoms] output */ + int N, + int num_atoms +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_atoms) return; + int idx; + if (num_atoms <= 1) { + idx = N / 2; /* single atom: median */ + } else { + idx = (int)((long long)i * (long long)(N - 1) / (long long)(num_atoms - 1)); + } + if (idx >= N) idx = N - 1; + if (idx < 0) idx = 0; + quantiles[i] = sorted[idx]; +} + /* ── Legacy O(N²) kernel — kept for reference, no longer called ───── */ #define TILE_SIZE 256 diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 4505d04b4..86dc23700 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -519,6 +519,9 @@ pub struct DQNHyperparameters { pub max_trace_length: usize, /// v8: PopArt reward normalization enabled. pub popart_enabled: bool, + /// v8: Robust PopArt using median/IQR instead of Welford mean/var. + /// More robust for heavy-tailed bimodal reward distributions. + pub popart_robust: bool, /// v8: Curriculum learning enabled. pub curriculum_enabled: bool, /// v8: Fraction of training before full dataset (curriculum). @@ -1423,6 +1426,7 @@ impl DQNHyperparameters { td_lambda: 0.9, max_trace_length: 7, popart_enabled: true, + popart_robust: true, curriculum_enabled: false, // disabled by default until wired curriculum_warmup_fraction: 0.6, hindsight_fraction: 0.1, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 140ed18f9..e53e7c0b1 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -307,6 +307,12 @@ pub(crate) struct FusedTrainingCtx { tau_host: f32, /// v8: PopArt reward normalization enabled (false by default). pub(crate) popart_enabled: bool, + /// v8: Use robust PopArt (median/IQR) instead of Welford mean/var. + pub(crate) popart_robust: bool, + /// Cached median from sorted epoch rewards (for robust PopArt). + pub(crate) cached_median: f32, + /// Cached IQR from sorted epoch rewards (for robust PopArt). + pub(crate) cached_iqr: f32, /// PopArt running variance from previous epoch (for tau reset detection). prev_popart_var: f32, /// v8: Hindsight relabeling fraction (0.0 = disabled). @@ -783,6 +789,9 @@ impl FusedTrainingCtx { eval_forward_exec: None, tau_host: 0.0, popart_enabled: hyperparams.popart_enabled, + popart_robust: hyperparams.popart_robust, + cached_median: 0.0, + cached_iqr: 0.0, prev_popart_var: 0.0, hindsight_fraction: hyperparams.hindsight_fraction, curriculum_enabled: hyperparams.curriculum_enabled, @@ -849,6 +858,8 @@ impl FusedTrainingCtx { // Reset PopArt running statistics for fresh fold if self.popart_enabled { self.prev_popart_var = 0.0; + self.cached_median = 0.0; + self.cached_iqr = 0.0; } Ok(()) } @@ -975,12 +986,19 @@ impl FusedTrainingCtx { batch_size: self.batch_size, }; - // Phase 0a: PopArt normalize rewards in-place (Welford running stats). + // Phase 0a: PopArt normalize rewards in-place. // Rescales raw rewards to unit variance so C51 atom support [-50,50] // covers the actual reward range. Without this, rewards in [-1,1] use // only 2% of the atom range → near-zero C51 gradient. - self.trainer.normalize_rewards_popart_inplace(self.batch_size) - .map_err(|e| anyhow::anyhow!("PopArt normalize: {e}"))?; + if self.popart_robust && self.cached_iqr > 0.0 { + // Robust path: median/IQR from sorted epoch rewards (no running stats). + self.trainer.normalize_rewards_robust(self.batch_size, self.cached_median, self.cached_iqr) + .map_err(|e| anyhow::anyhow!("PopArt robust normalize: {e}"))?; + } else { + // Welford path: running mean/var with warmup. + self.trainer.normalize_rewards_popart_inplace(self.batch_size) + .map_err(|e| anyhow::anyhow!("PopArt normalize: {e}"))?; + } // Phase 0b: Counter increments + stochastic depth mask (ungraphed) self.submit_counters_ops()?; @@ -2142,6 +2160,13 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("Shrink-and-Perturb: {e}")) } + /// Update cached median and IQR for robust PopArt normalization. + /// Called once per epoch after experience collection sorts rewards. + pub(crate) fn set_robust_popart_stats(&mut self, median: f32, iqr: f32) { + self.cached_median = median; + self.cached_iqr = iqr; + } + pub(crate) fn read_popart_variance(&self) -> f32 { if !self.popart_enabled { return 0.0; } let mut var = [0.0_f32]; @@ -2511,6 +2536,10 @@ impl FusedTrainingCtx { /// Per-sample epsilon from IQL expectile gap. pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() } pub(crate) fn num_atoms(&self) -> usize { self.trainer.config().num_atoms } + /// Warm-start C51 atom positions from empirical reward quantiles (once per epoch). + pub(crate) fn warm_start_atom_positions(&mut self, quantiles: &[f32]) -> Result<(), crate::MLError> { + self.trainer.warm_start_atom_positions(quantiles) + } pub(crate) fn update_adaptive_clip(&mut self, grad_norm: f32) { self.trainer.update_adaptive_clip(grad_norm); } pub(crate) fn adaptive_clip_value(&self) -> f32 { self.trainer.adaptive_clip_value() } /// Raw pointer to plan_params_buf [B, 6] for trade plan integration. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 3c948ed43..a284d0115 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -440,6 +440,43 @@ impl DQNTrainer { // v8: TD(λ) is now wired in collect_experiences_gpu — self-bootstrap with // rewards as Q(s') approximation, overwrites n-step when td_lambda > 0.01. + // C51 atom warm-start: extract empirical reward quantiles from the bitonic + // sort pipeline and initialize atom positions for this epoch. The existing + // step_atom_positions SGD optimizer refines from here during training steps. + if let Some(ref mut collector) = self.gpu_experience_collector { + let count = collector.last_experience_count(); + let num_atoms = self.hyperparams.num_atoms; + match collector.compute_reward_quantiles(count, num_atoms) { + Ok(quantiles) => { + if let Some(ref mut fused) = self.fused_ctx { + if let Err(e) = fused.warm_start_atom_positions(&quantiles) { + warn!("C51 atom warm-start failed (non-fatal): {e}"); + } + } + } + Err(e) => { + debug!("C51 reward quantile extraction failed (non-fatal): {e}"); + } + } + + // Robust PopArt: extract median/IQR from the already-sorted reward buffer. + // The bitonic sort from compute_reward_quantiles above leaves reward_sort_keys_buf + // sorted, so extract_median_iqr just reads 3 positions (tiny DtoH: 3 floats). + if self.hyperparams.popart_robust { + match collector.extract_median_iqr(count) { + Ok((median, iqr)) => { + if let Some(ref mut fused) = self.fused_ctx { + fused.set_robust_popart_stats(median, iqr); + } + debug!(epoch, median, iqr, "Robust PopArt: median/IQR cached for epoch"); + } + Err(e) => { + debug!("Robust PopArt median/IQR extraction failed (non-fatal): {e}"); + } + } + } + } + // Periodic shrink-and-perturb let sp_interval = self.hyperparams.shrink_perturb_interval; if sp_interval > 0 && epoch > 0 && epoch % sp_interval == 0 { diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 09e10c51d..88e3b674f 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -269,6 +269,8 @@ pub struct RewardSection { pub reward_scale: Option, /// v8: PopArt reward normalization enabled. pub popart_enabled: Option, + /// v8: Robust PopArt using median/IQR instead of Welford mean/var. + pub popart_robust: Option, /// v8: Dense directional micro-reward scale. pub micro_reward_scale: Option, /// v8: TD(λ) trace parameter. @@ -1076,6 +1078,9 @@ impl DqnTrainingProfile { if let Some(v) = rw.popart_enabled { hp.popart_enabled = v; } + if let Some(v) = rw.popart_robust { + hp.popart_robust = v; + } if let Some(v) = rw.micro_reward_scale { hp.micro_reward_scale = v; }