From 52ea2ac36bcfbd4580ff63f2426c73096bf06332 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 3 Mar 2026 00:09:27 +0100 Subject: [PATCH] feat(ml): ReplayBufferType::GpuPrioritized variant with full dispatch New variant delegates to GpuReplayBuffer for all operations. Sample returns BatchSample with gpu_batch populated (empty CPU vecs). Added update_priorities_gpu() for tensor-based priority updates, as_gpu_buffer() for direct access from trainer. Existing tests unchanged. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/dqn/replay_buffer_type.rs | 120 +++++++++++++++++++++++- 1 file changed, 117 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/dqn/replay_buffer_type.rs b/crates/ml/src/dqn/replay_buffer_type.rs index f9d8359fb..e95e6db93 100644 --- a/crates/ml/src/dqn/replay_buffer_type.rs +++ b/crates/ml/src/dqn/replay_buffer_type.rs @@ -60,6 +60,9 @@ pub enum ReplayBufferType { Uniform(Arc>), /// Prioritized sampling based on TD errors (Rainbow DQN) Prioritized(Arc), + /// GPU-resident prioritized replay — all sampling on GPU (zero CPU round-trips) + #[cfg(feature = "cuda")] + GpuPrioritized(Arc>), } impl std::fmt::Debug for ReplayBufferType { @@ -67,6 +70,8 @@ impl std::fmt::Debug for ReplayBufferType { match self { Self::Uniform(_) => write!(f, "ReplayBufferType::Uniform"), Self::Prioritized(_) => write!(f, "ReplayBufferType::Prioritized"), + #[cfg(feature = "cuda")] + Self::GpuPrioritized(_) => write!(f, "ReplayBufferType::GpuPrioritized"), } } } @@ -102,6 +107,36 @@ impl ReplayBufferType { Ok(Self::Prioritized(Arc::new(buffer))) } + /// Create GPU-resident prioritized replay buffer. + /// + /// All experience data lives as contiguous GPU tensors. Sampling and + /// priority updates happen entirely on GPU — zero CPU round-trips. + #[cfg(feature = "cuda")] + pub fn new_gpu_prioritized( + capacity: usize, + state_dim: usize, + alpha: f64, + beta: f64, + beta_max: f64, + beta_annealing_steps: usize, + device: &candle_core::Device, + ) -> Result { + use crate::cuda_pipeline::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig}; + + let config = GpuReplayBufferConfig { + capacity, + state_dim, + alpha: alpha as f32, + beta_start: beta as f32, + beta_max: beta_max as f32, + beta_annealing_steps, + epsilon: 1e-6, + }; + + let buffer = GpuReplayBuffer::new(config, device)?; + Ok(Self::GpuPrioritized(Arc::new(Mutex::new(buffer)))) + } + /// Sample a batch from the buffer pub fn sample(&self, batch_size: usize) -> Result { match self { @@ -126,6 +161,17 @@ impl ReplayBufferType { gpu_batch: None, }) } + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => { + let buf = buffer.lock(); + let gpu_batch = buf.sample_proportional(batch_size)?; + Ok(BatchSample { + experiences: vec![], // Empty — data is on GPU + weights: vec![], // Empty — weights on GPU + indices: vec![], // Empty — indices on GPU + gpu_batch: Some(gpu_batch), + }) + } } } @@ -140,6 +186,11 @@ impl ReplayBufferType { Self::Prioritized(buffer) => { buffer.push(experience) } + #[cfg(feature = "cuda")] + Self::GpuPrioritized(_) => { + // No-op: GPU buffer uses insert_batch() from GPU experience collector + Ok(()) + } } } @@ -164,6 +215,11 @@ impl ReplayBufferType { } Ok(()) } + #[cfg(feature = "cuda")] + Self::GpuPrioritized(_) => { + // No-op: GPU buffer uses insert_batch() with tensors directly + Ok(()) + } } } @@ -176,21 +232,49 @@ impl ReplayBufferType { let priorities: Vec = td_errors.iter().map(|&td| td.abs()).collect(); buffer.update_priorities(indices, &priorities) } + #[cfg(feature = "cuda")] + Self::GpuPrioritized(_) => { + // Use update_priorities_gpu() with tensor inputs instead + Ok(()) + } + } + } + + /// Update priorities from GPU-resident TD error tensors (GpuPrioritized only). + /// + /// Takes tensor indices and TD errors directly — no `to_vec1()` needed. + #[cfg(feature = "cuda")] + pub fn update_priorities_gpu( + &self, + indices: &candle_core::Tensor, + td_errors: &candle_core::Tensor, + ) -> Result<(), MLError> { + match self { + Self::GpuPrioritized(buffer) => { + let mut buf = buffer.lock(); + buf.update_priorities_gpu(indices, td_errors) + } + _ => Ok(()), // No-op for non-GPU buffers } } /// Step training counter for beta annealing (PER only) pub fn step(&self) { - if let Self::Prioritized(buffer) = self { - buffer.step(); + match self { + Self::Prioritized(buffer) => buffer.step(), + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => buffer.lock().step(), + _ => {} } } /// Get current beta value (PER only) pub fn current_beta(&self) -> f32 { match self { - Self::Uniform(_) => 1.0, // No importance sampling for uniform + Self::Uniform(_) => 1.0, Self::Prioritized(buffer) => buffer.current_beta(), + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => buffer.lock().current_beta(), } } @@ -199,6 +283,8 @@ impl ReplayBufferType { match self { Self::Uniform(buffer) => buffer.lock().len(), Self::Prioritized(buffer) => buffer.len(), + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => buffer.lock().len(), } } @@ -217,6 +303,10 @@ impl ReplayBufferType { Self::Prioritized(buffer) => { buffer.clear(); } + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => { + buffer.lock().clear(); + } } } @@ -225,6 +315,8 @@ impl ReplayBufferType { match self { Self::Uniform(buffer) => buffer.lock().len() >= batch_size, Self::Prioritized(buffer) => buffer.can_sample(batch_size), + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => buffer.lock().can_sample(batch_size), } } @@ -233,6 +325,12 @@ impl ReplayBufferType { matches!(self, Self::Prioritized(_)) } + /// Check if using GPU-resident prioritized replay + #[cfg(feature = "cuda")] + pub fn is_gpu_prioritized(&self) -> bool { + matches!(self, Self::GpuPrioritized(_)) + } + /// Adaptive buffer sizing based on epsilon decay /// /// Resizes buffer capacity using threshold-based strategy to reduce memory usage @@ -294,6 +392,11 @@ impl ReplayBufferType { // PER resize not yet implemented (requires rebuilding sum tree) return Ok(false); } + #[cfg(feature = "cuda")] + Self::GpuPrioritized(_) => { + // GPU buffer is fixed-capacity (pre-allocated tensors) + return Ok(false); + } } tracing::info!( @@ -312,6 +415,17 @@ impl ReplayBufferType { match self { Self::Uniform(buffer) => buffer.lock().capacity(), Self::Prioritized(buffer) => buffer.capacity(), + #[cfg(feature = "cuda")] + Self::GpuPrioritized(buffer) => buffer.lock().capacity(), + } + } + + /// Get a locked reference to the inner GPU replay buffer (GpuPrioritized only). + #[cfg(feature = "cuda")] + pub fn as_gpu_buffer(&self) -> Option> { + match self { + Self::GpuPrioritized(buffer) => Some(buffer.lock()), + _ => None, } }