diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index b3dc766ac..010d4e8b1 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -595,13 +595,15 @@ impl GpuReplayBuffer { // then index_add the deltas back. Single batched GPU kernel, no CPU loop. // // For rare duplicate indices (batch_size << capacity), the deltas may - // stack incorrectly, but this is benign: priorities will be slightly off - // and self-correct on the next update. This matches the old CPU loop's - // "last-write-wins" semantics in expectation. + // stack — index_add adds ALL deltas, so duplicates double-apply. + // Post-clamp prevents priorities from going to zero/negative (→ NaN weights). let indices_i64 = indices.to_dtype(DType::I64)?; let old_prios = self.priorities.index_select(&indices_i64, 0)?; let delta = (&clamped - &old_prios)?; - self.priorities = self.priorities.index_add(&indices_i64, &delta, 0)?; + self.priorities = self + .priorities + .index_add(&indices_i64, &delta, 0)? + .clamp(self.config.epsilon, 1e6)?; Ok(()) }