feat(ml): wire GpuReplayBuffer into DQN trainer hot path

GradientResult and ComputeLossResult gain td_errors_gpu/indices_gpu
Option<Tensor> fields for GPU-resident TD errors. Trainer accumulates
GPU tensors and calls update_priorities_gpu() when available, falling
back to CPU path otherwise. DqnAgent delegates via memory().

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/ml/src/dqn/dqn.rs

@@ -486,6 +486,12 @@ pub struct GradientResult {
     pub td_errors: Vec<f32>,
     /// Replay buffer indices for PER priority updates.
     pub indices: Vec<usize>,
+    /// GPU-resident TD errors (GpuPrioritized path — avoids to_vec1).
+    #[cfg(feature = "cuda")]
+    pub td_errors_gpu: Option<Tensor>,
+    /// GPU-resident buffer indices (GpuPrioritized path).
+    #[cfg(feature = "cuda")]
+    pub indices_gpu: Option<Tensor>,
 }
 
 /// Internal result from forward pass + loss computation (no backward pass).
@@ -500,6 +506,12 @@ struct ComputeLossResult {
     indices: Vec<usize>,
     /// States tensor needed for Q-value diagnostic logging.
     states_tensor: Tensor,
+    /// GPU-resident TD errors (GpuPrioritized path).
+    #[cfg(feature = "cuda")]
+    td_errors_gpu: Option<Tensor>,
+    /// GPU-resident buffer indices (GpuPrioritized path).
+    #[cfg(feature = "cuda")]
+    indices_gpu: Option<Tensor>,
 }
 
 /// 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,
         })
     }
 

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,
         })
     }
 

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();

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<candle_core::Tensor> = Vec::new();
+        #[cfg(feature = "cuda")]
+        let mut gpu_indices: Vec<candle_core::Tensor> = 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)