diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 61cf7dae7..4414969e2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -9145,6 +9145,24 @@ impl GpuDqnTrainer { Ok(dst) } + /// Hard-copy online params into target params. At fold boundaries the online + /// weights are shrink-and-perturb'd; the EMA target still holds the end-of- + /// prior-fold values and would otherwise produce a large TD error gap in the + /// first steps of the new fold, driving runaway gradients before Polyak + /// averaging can close the divergence. DtoD copy — no CPU staging. + pub fn sync_target_from_online(&mut self) -> Result<(), MLError> { + let n_bytes = self.total_params * std::mem::size_of::(); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + self.target_params_buf.raw_ptr(), + self.params_buf.raw_ptr(), + n_bytes, + self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("sync_target_from_online DtoD: {e}")))?; + } + Ok(()) + } + /// Clone target_params_buf to a new CudaSlice via DtoD copy — checkpoint stays on GPU, zero CPU. pub fn clone_target_params_gpu(&self) -> Result, MLError> { let len = self.target_params_buf.len(); diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 3d8000d49..2495fff9a 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -830,6 +830,14 @@ impl FusedTrainingCtx { } else { tracing::info!(alpha = sp_alpha, sigma = sp_sigma, "Fold-boundary shrink-and-perturb applied"); } + // Hard-copy the shrink-and-perturb'd online weights into target params. + // Without this, target_params_buf retains end-of-previous-fold weights + // while online was just modified — the Bellman target would use stale + // weights against perturbed online predictions, producing an outsized + // TD error gap in the first fold-N+1 steps. Polyak averaging alone is + // too slow (tau=0.005) to close that gap before the oversized gradients + // compound through Adam. This pairs with reset_adam_state below. + self.trainer.sync_target_from_online()?; // Reset Adam optimizer state — each fold starts with fresh momentum. // Stale momentum from a previous fold causes weight explosion → Q-value -1e30. self.trainer.reset_adam_state()?;