diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 43633b29d..4d073d14c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1021,6 +1021,11 @@ impl GpuDqnTrainer { &self.rewards_buf } + /// Mutable reference to the rewards buffer on GPU (for in-place normalization). + pub fn rewards_buf_mut(&mut self) -> &mut CudaSlice { + &mut self.rewards_buf + } + /// Reference to the dones buffer on GPU. /// /// Shape: `[B]` f32 — done flags (0.0/1.0). Valid after `train_step()` or @@ -1992,6 +1997,34 @@ impl GpuDqnTrainer { Ok(()) } + /// v8: Normalize the internal rewards_buf in-place using PopArt. + /// + /// Convenience wrapper that avoids double-borrow by using `self.rewards_buf` + /// directly. Called from `FusedTrainingCtx::run_full_step()`. + pub fn normalize_rewards_popart_inplace( + &mut self, n: usize, + ) -> Result<(), MLError> { + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + let mean_ptr = self.popart_mean.raw_ptr(); + let var_ptr = self.popart_var.raw_ptr(); + let count_ptr = self.popart_count.raw_ptr(); + let rewards_ptr = self.rewards_buf.raw_ptr(); + let warmup = 100_i32; + let n_i32 = n as i32; + unsafe { + self.stream.launch_builder(&self.popart_normalize_kernel) + .arg(&rewards_ptr) + .arg(&mean_ptr) + .arg(&var_ptr) + .arg(&count_ptr) + .arg(&n_i32) + .arg(&warmup) + .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("popart_normalize_inplace: {e}")))?; + } + Ok(()) + } + /// v8: Normalize rewards in-place using running mean/variance (PopArt). /// /// The kernel updates running statistics with Welford's algorithm and diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 1e25a3eff..05f52fcd7 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -209,6 +209,12 @@ pub(crate) struct FusedTrainingCtx { pub(crate) exposure_aux_weight: f64, /// v7.1: Exposure targets buffer [batch_size] -- from experience collector. pub(crate) exposure_targets: cudarc::driver::CudaSlice, + /// v8: PopArt reward normalization enabled (false by default). + pub(crate) popart_enabled: bool, + /// v8: Hindsight relabeling fraction (0.0 = disabled). + pub(crate) hindsight_fraction: f64, + /// v8: Curriculum learning enabled (false by default). + pub(crate) curriculum_enabled: bool, } impl Drop for FusedTrainingCtx { @@ -602,6 +608,9 @@ impl FusedTrainingCtx { tau_host: 0.0, exposure_aux_weight: hyperparams.exposure_aux_weight, exposure_targets, + popart_enabled: hyperparams.popart_enabled, + hindsight_fraction: hyperparams.hindsight_fraction, + curriculum_enabled: hyperparams.curriculum_enabled, }) } @@ -616,6 +625,26 @@ impl FusedTrainingCtx { self.batch_size } + /// v8: Run one supervised pre-training step for the exposure branch. + /// Delegates to `GpuDqnTrainer::run_pretrain_step`. + pub(crate) fn run_pretrain_step( + &mut self, + targets: &cudarc::driver::CudaSlice, + bar_indices: &cudarc::driver::CudaSlice, + batch_size: usize, + max_position: f32, + ) -> Result<()> { + self.trainer.run_pretrain_step(targets, bar_indices, batch_size, max_position) + .map_err(|e| anyhow::anyhow!("pretrain_step: {e}")) + } + + /// v8: Run Adam optimizer and readback scalars (for pre-training gradient apply). + pub(crate) fn replay_adam_and_readback_pretrain(&mut self) -> Result<()> { + self.trainer.replay_adam_and_readback() + .map_err(|e| anyhow::anyhow!("adam_readback: {e}"))?; + Ok(()) + } + /// Reset per-fold state. Keeps GPU infrastructure + compiled kernels. /// Invalidates graph_aux (fold data range changed). pub(crate) fn reset_for_fold(&mut self) -> anyhow::Result<()> { @@ -720,7 +749,29 @@ impl FusedTrainingCtx { self.trainer.run_nan_checks_post_forward(self.batch_size)?; - // ── Step 2b: HER donor computation (outside graph_aux) ─────────── + // ── Step 2a: PopArt reward normalization (in-place, before C51 loss) ── + // Normalizes the rewards buffer using running mean/variance. The forward + // graph has already captured the rewards pointer — PopArt updates the same + // buffer in-place so C51 loss on the NEXT step uses normalized rewards. + // Disabled by default (popart_enabled: false). + if self.popart_enabled { + let bs = self.batch_size; + if let Err(e) = self.trainer.normalize_rewards_popart_inplace(bs) { + tracing::warn!("PopArt normalize failed (non-fatal): {e}"); + } + } + + // ── Step 2b: v8 stubs (hindsight + curriculum — disabled by default) ── + // v8: Hindsight relabeling (wired but requires bar_indices in PER batch — future work) + if self.hindsight_fraction > 0.0 { + tracing::debug!("Hindsight relabeling: kernel loaded, awaiting bar_indices in PER batch"); + } + // v8: Curriculum learning (wired but requires sorted episode starts — future work) + if self.curriculum_enabled { + tracing::debug!("Curriculum learning: kernel loaded, awaiting sorted episode starts"); + } + + // ── Step 2c: HER donor computation (outside graph_aux) ─────────── // Donor indices vary per step (random/future/final). The computation // fills her.donor_indices GPU buffer. graph_aux captures the relabel // kernel that READS from this stable buffer address. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index c74c63d7f..e470983cf 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -190,6 +190,59 @@ impl DQNTrainer { } } + // v8: Ensure fused context exists for pre-training (normally lazy-init in Phase 3) + if self.hyperparams.exposure_aux_weight > 0.0 && self.fused_ctx.is_none() && self.device.is_cuda() { + if let Some(ref stream) = self.cuda_stream { + let agent = self.agent.read().await; + match super::super::fused_training::FusedTrainingCtx::new( + &self.device, &*agent, &self.hyperparams, self.current_batch_size, Arc::clone(stream), + ) { + Ok(ctx) => { self.fused_ctx = Some(ctx); } + Err(e) => { tracing::warn!("Fused CUDA context init for pre-training failed: {e}"); } + } + } + } + + // v8: Supervised pre-training (epoch 0) — initialize exposure branch with directional knowledge + if self.hyperparams.exposure_aux_weight > 0.0 { + let pretrain_total_bars = self.gpu_data.as_ref().map(|d| d.num_bars); + let pretrain_batch_size = pretrain_total_bars.map(|t| self.current_batch_size.min(t)); + let max_pos = self.max_position as f32; + + if let (Some(total_bars), Some(batch_size)) = (pretrain_total_bars, pretrain_batch_size) { + if let (Some(ref mut fused), Some(ref stream), Some(ref targets_cuda)) = + (&mut self.fused_ctx, &self.cuda_stream, &self.targets_raw_cuda) + { + info!("v8: Supervised pre-training — 50 batches for exposure direction initialization"); + let mut bar_indices = vec![0_i32; batch_size]; + for batch_idx in 0..50_usize { + // Sequential bar indices starting at different offsets + let offset = (batch_idx * batch_size) % total_bars.saturating_sub(batch_size); + for i in 0..batch_size { + bar_indices[i] = (offset + i).min(total_bars - 2) as i32; + } + let mut bar_indices_gpu = stream.alloc_zeros::(batch_size) + .map_err(|e| anyhow::anyhow!("alloc bar_indices: {e}"))?; + stream.memcpy_htod(&bar_indices, &mut bar_indices_gpu) + .map_err(|e| anyhow::anyhow!("upload bar_indices: {e}"))?; + + if let Err(e) = fused.run_pretrain_step( + targets_cuda, &bar_indices_gpu, batch_size, max_pos, + ) { + tracing::warn!("Pre-training batch {batch_idx} failed (non-fatal): {e}"); + break; + } + // Run Adam to apply the pre-training gradients + if let Err(e) = fused.replay_adam_and_readback_pretrain() { + tracing::warn!("Pre-training Adam failed (non-fatal): {e}"); + break; + } + } + info!("v8: Pre-training complete — exposure branch initialized"); + } + } + } + // Training loop for epoch in 0..self.hyperparams.epochs { self.current_epoch = epoch; @@ -282,6 +335,36 @@ impl DQNTrainer { )); } + // v8: Copy best_exposure targets from collector to fused context for aux GEMM + { + let exposure_copy_info = self.gpu_experience_collector.as_ref().map(|c| { + let src = c.best_exposure_buf(); + (src.raw_ptr(), src.len()) + }); + if let (Some((src_ptr, src_len)), Some(ref mut fused), Some(ref stream)) = + (exposure_copy_info, &mut self.fused_ctx, &self.cuda_stream) + { + let batch_size = self.current_batch_size; + if src_len >= batch_size && fused.exposure_targets.len() >= batch_size { + let dst_ptr = fused.exposure_targets.raw_ptr(); + let bytes = batch_size * std::mem::size_of::(); + unsafe { + cudarc::driver::sys::cuMemcpyDtoDAsync_v2( + dst_ptr, src_ptr, bytes, stream.cu_stream(), + ); + } + } + } + } + + // v8: TD(lambda) returns — kernel loaded, awaiting value estimates in collector + if self.hyperparams.td_lambda > 0.01 { + tracing::debug!( + td_lambda = self.hyperparams.td_lambda, + "TD(lambda): kernel loaded, awaiting V(s) estimates for full wiring" + ); + } + // Periodic shrink-and-perturb let sp_interval = self.hyperparams.shrink_perturb_interval; if sp_interval > 0 && epoch > 0 && epoch % sp_interval == 0 {