diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 010d4e8b1..ab0739096 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -594,15 +594,48 @@ impl GpuReplayBuffer { // Delta trick: gather old priorities at indices, compute delta = new - old, // then index_add the deltas back. Single batched GPU kernel, no CPU loop. // - // For rare duplicate indices (batch_size << capacity), the deltas may - // stack — index_add adds ALL deltas, so duplicates double-apply. - // Post-clamp prevents priorities from going to zero/negative (→ NaN weights). + // Dedup: index_add accumulates ALL deltas at each position, so duplicate + // indices would double-apply. Pull the small indices tensor (batch_size + // = 256 → 1 KB) to CPU, deduplicate keeping the last td_error per index, + // then rebuild GPU tensors with unique entries only. let indices_i64 = indices.to_dtype(DType::I64)?; - let old_prios = self.priorities.index_select(&indices_i64, 0)?; - let delta = (&clamped - &old_prios)?; + let (unique_idx, unique_prios) = { + let idx_vec: Vec = indices_i64.to_vec1().map_err(|e| { + MLError::ModelError(format!("indices to_vec1: {e}")) + })?; + let prio_vec: Vec = clamped.to_vec1().map_err(|e| { + MLError::ModelError(format!("clamped to_vec1: {e}")) + })?; + // Last-write-wins: later occurrences overwrite earlier ones + let mut seen = std::collections::HashMap::with_capacity(idx_vec.len()); + for (i, &idx) in idx_vec.iter().enumerate() { + seen.insert(idx, i); + } + if seen.len() < idx_vec.len() { + // Duplicates found — rebuild with unique entries + let mut u_idx = Vec::with_capacity(seen.len()); + let mut u_prio = Vec::with_capacity(seen.len()); + for (&idx, &pos) in &seen { + u_idx.push(idx); + if let Some(&p) = prio_vec.get(pos) { + u_prio.push(p); + } + } + let len = u_idx.len(); + ( + Tensor::from_vec(u_idx, &[len], &self.device)?, + Tensor::from_vec(u_prio, &[len], &self.device)?, + ) + } else { + // No duplicates — fast path, reuse existing tensors + (indices_i64, clamped) + } + }; + let old_prios = self.priorities.index_select(&unique_idx, 0)?; + let delta = (&unique_prios - &old_prios)?; self.priorities = self .priorities - .index_add(&indices_i64, &delta, 0)? + .index_add(&unique_idx, &delta, 0)? .clamp(self.config.epsilon, 1e6)?; Ok(()) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index c178c2d5c..c31e1b2eb 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -812,13 +812,17 @@ impl GpuExperienceCollector { } else { // Out of bounds — copy current state as fallback let s_start = idx * sd; - next_states.extend_from_slice(&states[s_start..s_start + sd]); + if let Some(s) = states.get(s_start..s_start + sd) { + next_states.extend_from_slice(s); + } } } else { // Terminal or last timestep — next_state = current state // (masked by done flag in Bellman equation) let s_start = idx * sd; - next_states.extend_from_slice(&states[s_start..s_start + sd]); + if let Some(s) = states.get(s_start..s_start + sd) { + next_states.extend_from_slice(s); + } } } } diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index f5d78a802..dbbb98bd1 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -281,6 +281,50 @@ impl DQNAgentType { } } + /// Convert GPU tensors to `Vec` for CPU replay buffer insertion. + /// + /// Used as fallback when a replay buffer is CPU-based (Prioritized or Uniform) + /// but experiences were collected on GPU. + #[cfg(feature = "cuda")] + fn tensors_to_experiences( + states: &candle_core::Tensor, + next_states: &candle_core::Tensor, + actions: &candle_core::Tensor, + rewards: &candle_core::Tensor, + dones: &candle_core::Tensor, + ) -> Result, MLError> { + let batch_size = states.dim(0).map_err(|e| { + MLError::TrainingError(format!("states dim: {e}")) + })?; + let state_dim = states.dim(1).map_err(|e| { + MLError::TrainingError(format!("states dim1: {e}")) + })?; + let s_flat: Vec = states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("states to_vec1: {e}")))?; + let ns_flat: Vec = next_states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("next_states to_vec1: {e}")))?; + let a_vec: Vec = actions.to_dtype(candle_core::DType::U32)?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("actions to_vec1: {e}")))?; + let r_vec: Vec = rewards.to_dtype(candle_core::DType::F32)?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("rewards to_vec1: {e}")))?; + let d_vec: Vec = dones.to_dtype(candle_core::DType::F32)?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("dones to_vec1: {e}")))?; + let mut experiences = Vec::with_capacity(batch_size); + for i in 0..batch_size { + let start = i * state_dim; + let end = start + state_dim; + experiences.push(Experience { + state: s_flat.get(start..end).unwrap_or(&[]).to_vec(), + action: a_vec.get(i).copied().unwrap_or(0) as u8, + reward: (r_vec.get(i).copied().unwrap_or(0.0) * 1_000_000.0) as i32, + next_state: ns_flat.get(start..end).unwrap_or(&[]).to_vec(), + done: d_vec.get(i).map_or(false, |&v| v > 0.5), + timestamp: 0, // fallback path — timestamp not meaningful for replay + }); + } + Ok(experiences) + } + /// Insert a batch of experience tensors directly into the GPU replay buffer. /// /// When the replay buffer is `GpuPrioritized`, this feeds tensors straight to @@ -300,63 +344,40 @@ impl DQNAgentType { if let Some(mut gpu_buf) = agent.memory.as_gpu_buffer() { gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones) } else { - // CPU PER fallback: convert tensors back to experiences - let batch_size = states.dim(0).map_err(|e| { - MLError::TrainingError(format!("states dim: {e}")) - })?; - let state_dim = states.dim(1).map_err(|e| { - MLError::TrainingError(format!("states dim1: {e}")) - })?; - let s_flat: Vec = states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1() - .map_err(|e| MLError::TrainingError(format!("states to_vec1: {e}")))?; - let ns_flat: Vec = next_states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1() - .map_err(|e| MLError::TrainingError(format!("next_states to_vec1: {e}")))?; - let a_vec: Vec = actions.to_dtype(candle_core::DType::U32)?.to_vec1() - .map_err(|e| MLError::TrainingError(format!("actions to_vec1: {e}")))?; - let r_vec: Vec = rewards.to_dtype(candle_core::DType::F32)?.to_vec1() - .map_err(|e| MLError::TrainingError(format!("rewards to_vec1: {e}")))?; - let d_vec: Vec = dones.to_dtype(candle_core::DType::F32)?.to_vec1() - .map_err(|e| MLError::TrainingError(format!("dones to_vec1: {e}")))?; - let mut experiences = Vec::with_capacity(batch_size); - for i in 0..batch_size { - let start = i * state_dim; - let end = start + state_dim; - experiences.push(crate::dqn::Experience { - state: s_flat.get(start..end).unwrap_or(&[]).to_vec(), - action: a_vec.get(i).copied().unwrap_or(0) as u8, - reward: (r_vec.get(i).copied().unwrap_or(0.0) * 1_000_000.0) as i32, - next_state: ns_flat.get(start..end).unwrap_or(&[]).to_vec(), - done: d_vec.get(i).map_or(false, |&v| v > 0.5), - timestamp: 0, // fallback path — timestamp not meaningful for replay - }); - } + let experiences = Self::tensors_to_experiences( + states, next_states, actions, rewards, dones, + )?; agent.memory.add_batch(experiences) } } Self::RegimeConditional(agent) => { // Insert into all head buffers — regime routing happens at train time // via GPU tensor mask operations (zero CPU roundtrip). - agent.get_trending_head().map(|h| &h.memory).map_or(Ok(()), |m| { - if let Some(mut gpu_buf) = m.as_gpu_buffer() { - gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones) - } else { - Ok(()) - } - })?; - agent.get_ranging_head().map(|h| &h.memory).map_or(Ok(()), |m| { - if let Some(mut gpu_buf) = m.as_gpu_buffer() { - gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones) - } else { - Ok(()) - } - })?; - agent.get_volatile_head().map(|h| &h.memory).map_or(Ok(()), |m| { - if let Some(mut gpu_buf) = m.as_gpu_buffer() { - gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones) - } else { - Ok(()) - } - })?; + // CPU fallback converts tensors→experiences once, shared across heads. + let mut cpu_experiences: Option> = None; + + macro_rules! insert_head { + ($getter:ident) => { + if let Some(head) = agent.$getter() { + if let Some(mut gpu_buf) = head.memory.as_gpu_buffer() { + gpu_buf.gpu.insert_batch(states, next_states, actions, rewards, dones)?; + } else { + if cpu_experiences.is_none() { + cpu_experiences = Some(Self::tensors_to_experiences( + states, next_states, actions, rewards, dones, + )?); + } + if let Some(ref exps) = cpu_experiences { + head.memory.add_batch(exps.clone())?; + } + } + } + }; + } + + insert_head!(get_trending_head); + insert_head!(get_ranging_head); + insert_head!(get_volatile_head); Ok(()) } }