diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index 5ffaebeb1..5bdf746ec 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -486,6 +486,12 @@ pub struct GradientResult { pub td_errors: Vec, /// Replay buffer indices for PER priority updates. pub indices: Vec, + /// GPU-resident TD errors (GpuPrioritized path — avoids to_vec1). + #[cfg(feature = "cuda")] + pub td_errors_gpu: Option, + /// GPU-resident buffer indices (GpuPrioritized path). + #[cfg(feature = "cuda")] + pub indices_gpu: Option, } /// Internal result from forward pass + loss computation (no backward pass). @@ -500,6 +506,12 @@ struct ComputeLossResult { indices: Vec, /// States tensor needed for Q-value diagnostic logging. states_tensor: Tensor, + /// GPU-resident TD errors (GpuPrioritized path). + #[cfg(feature = "cuda")] + td_errors_gpu: Option, + /// GPU-resident buffer indices (GpuPrioritized path). + #[cfg(feature = "cuda")] + indices_gpu: Option, } /// Experience replay buffer for `DQN` @@ -1969,6 +1981,10 @@ impl DQN { td_errors: td_errors_vec, indices, states_tensor, + #[cfg(feature = "cuda")] + td_errors_gpu: None, + #[cfg(feature = "cuda")] + indices_gpu: None, }) } @@ -2064,6 +2080,10 @@ impl DQN { grads, td_errors: result.td_errors, indices: result.indices, + #[cfg(feature = "cuda")] + td_errors_gpu: result.td_errors_gpu, + #[cfg(feature = "cuda")] + indices_gpu: result.indices_gpu, }) } diff --git a/crates/ml/src/dqn/regime_conditional.rs b/crates/ml/src/dqn/regime_conditional.rs index 4267a5067..bda45318d 100644 --- a/crates/ml/src/dqn/regime_conditional.rs +++ b/crates/ml/src/dqn/regime_conditional.rs @@ -560,6 +560,10 @@ impl RegimeConditionalDQN { })?, td_errors: all_td_errors, indices: all_indices, + #[cfg(feature = "cuda")] + td_errors_gpu: None, + #[cfg(feature = "cuda")] + indices_gpu: None, }) } diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 0880d1785..c95b01bc7 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -332,6 +332,16 @@ impl DQNAgentType { self.memory().update_priorities(indices, td_errors) } + /// Update replay buffer priorities from GPU-resident tensors (GpuPrioritized only). + #[cfg(feature = "cuda")] + pub fn update_priorities_gpu( + &self, + indices: &candle_core::Tensor, + td_errors: &candle_core::Tensor, + ) -> Result<(), MLError> { + self.memory().update_priorities_gpu(indices, td_errors) + } + /// Step the replay buffer training counter for beta annealing (PER only). pub fn step_replay_buffer(&self) { self.memory().step(); diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 8318963bf..f540343b3 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -3433,6 +3433,10 @@ impl DQNTrainer { let mut all_td_errors = Vec::new(); let mut all_indices = Vec::new(); let mut final_grad_norm = 0.0_f32; + #[cfg(feature = "cuda")] + let mut gpu_td_errors: Vec = Vec::new(); + #[cfg(feature = "cuda")] + let mut gpu_indices: Vec = Vec::new(); for (step, batch) in pre_sampled.into_iter().enumerate() { // Compute forward pass + backward WITHOUT optimizer step @@ -3456,6 +3460,15 @@ impl DQNTrainer { final_grad_norm = result.grad_norm; all_td_errors.extend(result.td_errors); all_indices.extend(result.indices); + #[cfg(feature = "cuda")] + { + if let Some(td_gpu) = result.td_errors_gpu { + gpu_td_errors.push(td_gpu); + } + if let Some(idx_gpu) = result.indices_gpu { + gpu_indices.push(idx_gpu); + } + } // Emergency brake for extreme losses if result.loss > 1e6 { @@ -3490,6 +3503,24 @@ impl DQNTrainer { } // === Phase 3: Bookkeeping === + #[cfg(feature = "cuda")] + { + // GPU PER path: concatenate GPU tensors and update in one shot + if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() { + let td_cat = candle_core::Tensor::cat(&gpu_td_errors, 0) + .map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?; + let idx_cat = candle_core::Tensor::cat(&gpu_indices, 0) + .map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?; + agent + .update_priorities_gpu(&idx_cat, &td_cat) + .map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?; + } else if !all_indices.is_empty() { + agent + .update_priorities(&all_indices, &all_td_errors) + .map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?; + } + } + #[cfg(not(feature = "cuda"))] if !all_indices.is_empty() { agent .update_priorities(&all_indices, &all_td_errors)