fix(gpu-per): clamp priorities after index_add to prevent NaN weights

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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-09 23:34:29 +01:00
parent 9502c1e22c
commit 6b3bb26c05

View File

@@ -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(())
}