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:
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user