From 6b3bb26c053407cf85628a20bca9e778eda9f8b4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 9 Mar 2026 23:34:29 +0100 Subject: [PATCH] fix(gpu-per): clamp priorities after index_add to prevent NaN weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit index_add with duplicate indices double-applies deltas — e.g. if index i is sampled twice with delta -0.5, priority goes 1.0 + (-0.5) + (-0.5) = 0.0, causing NaN in importance sampling weights. Post-clamp to [epsilon, 1e6] ensures priorities stay positive. Co-Authored-By: Claude Opus 4.6 --- crates/ml-dqn/src/gpu_replay_buffer.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) 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(()) }