diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d349a2f61..b447d49a5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -7957,9 +7957,12 @@ impl GpuDqnTrainer { if ring.snapshots.len() >= crate::cuda_pipeline::q_snapshot::MAX_SNAPSHOTS { let worst_idx = ring.snapshots.iter() .enumerate() - .min_by(|(_, a), (_, b)| a.q_gap.partial_cmp(&b.q_gap).unwrap()) + .min_by(|(_, a), (_, b)| { + // NaN-safe: NaN q_gaps compare as Equal instead of panicking. + a.q_gap.partial_cmp(&b.q_gap).unwrap_or(std::cmp::Ordering::Equal) + }) .map(|(i, _)| i) - .unwrap(); + .expect("MAX_SNAPSHOTS >= 1 and len >= MAX_SNAPSHOTS, so at least one snapshot exists"); ring.snapshots.swap_remove(worst_idx); } diff --git a/crates/ml/src/cuda_pipeline/q_snapshot.rs b/crates/ml/src/cuda_pipeline/q_snapshot.rs index 7cae71a92..6b0ab6fd9 100644 --- a/crates/ml/src/cuda_pipeline/q_snapshot.rs +++ b/crates/ml/src/cuda_pipeline/q_snapshot.rs @@ -107,9 +107,13 @@ impl SnapshotRing { if self.snapshots.len() >= MAX_SNAPSHOTS { let worst_idx = self.snapshots.iter() .enumerate() - .min_by(|(_, a), (_, b)| a.q_gap.partial_cmp(&b.q_gap).unwrap()) + .min_by(|(_, a), (_, b)| { + // NaN-safe: treat NaN q_gaps as equal so a snapshot with NaN + // does not panic the swap-out path. + a.q_gap.partial_cmp(&b.q_gap).unwrap_or(std::cmp::Ordering::Equal) + }) .map(|(i, _)| i) - .unwrap(); + .expect("MAX_SNAPSHOTS >= 1 and len >= MAX_SNAPSHOTS, so at least one snapshot exists"); self.snapshots.swap_remove(worst_idx); } @@ -120,7 +124,7 @@ impl SnapshotRing { /// Best (highest q_gap) snapshot raw device pointer, if any. pub fn best_raw_ptr(&self) -> Option { self.snapshots.iter() - .max_by(|a, b| a.q_gap.partial_cmp(&b.q_gap).unwrap()) + .max_by(|a, b| a.q_gap.partial_cmp(&b.q_gap).unwrap_or(std::cmp::Ordering::Equal)) .map(|s| s.weights.raw_ptr()) } }