diff --git a/crates/ml/tests/gpu_per_integration_test.rs b/crates/ml/tests/gpu_per_integration_test.rs new file mode 100644 index 000000000..3e67d5591 --- /dev/null +++ b/crates/ml/tests/gpu_per_integration_test.rs @@ -0,0 +1,236 @@ +//! Integration test: GPU-resident PER replay buffer +//! +//! Verifies that GpuReplayBuffer produces correct sampling dynamics: +//! 1. Priority-weighted sampling skews towards high-priority experiences +//! 2. Priority updates change sampling distribution +//! 3. Multi-epoch insert + sample + update cycle works end-to-end +//! 4. KS test: GPU PER distribution differs significantly from uniform + +use candle_core::{Device, DType, Tensor}; +use ml::cuda_pipeline::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig}; + +fn test_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfig { + GpuReplayBufferConfig { + capacity, + state_dim, + alpha: 0.6, + beta_start: 0.4, + beta_max: 1.0, + beta_annealing_steps: 1000, + epsilon: 1e-6, + } +} + +/// Insert N random experiences into the buffer. +fn fill_buffer(buf: &mut GpuReplayBuffer, n: usize, state_dim: usize) { + let device = buf.device().clone(); + let states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], &device).unwrap(); + let next_states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], &device).unwrap(); + let actions = Tensor::zeros(&[n], DType::U32, &device).unwrap(); + let rewards = Tensor::randn(0.0_f32, 1.0, &[n], &device).unwrap(); + let dones = Tensor::zeros(&[n], DType::F32, &device).unwrap(); + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones) + .unwrap(); +} + +#[test] +fn test_gpu_per_training_loop_cycle() { + // End-to-end: insert → sample → fake TD error → update priorities → repeat + let device = Device::Cpu; + let state_dim = 16; + let capacity = 500; + let batch_size = 32; + let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap(); + + // Fill buffer + fill_buffer(&mut buf, 200, state_dim); + assert_eq!(buf.len(), 200); + + // Run 5 "training epochs" + for epoch in 0..5 { + let batch = buf.sample_proportional(batch_size).unwrap(); + + // Verify batch shapes + assert_eq!(batch.states.dims(), &[batch_size, state_dim]); + assert_eq!(batch.actions.dims(), &[batch_size]); + assert_eq!(batch.rewards.dims(), &[batch_size]); + assert_eq!(batch.weights.dims(), &[batch_size]); + assert_eq!(batch.indices.dims(), &[batch_size]); + + // Simulate TD errors: decreasing magnitude over epochs + let scale = 1.0_f32 / (epoch as f32 + 1.0); + let td_errors = Tensor::full(scale, &[batch_size], &device).unwrap(); + + buf.update_priorities_gpu(&batch.indices, &td_errors).unwrap(); + buf.step(); + } + + // After updates, priorities should be non-uniform + assert!(buf.len() == 200); +} + +#[test] +fn test_gpu_per_priorities_affect_sampling() { + // Verify that updating priorities actually changes what gets sampled. + // Set one experience to very high priority, sample many times, check it appears often. + let device = Device::Cpu; + let state_dim = 4; + let capacity = 100; + let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap(); + + // Fill 100 experiences + fill_buffer(&mut buf, 100, state_dim); + + // Set experience 0 to very high priority, everything else stays at 1.0 + let high_idx = Tensor::new(&[0_u32], &device).unwrap(); + let high_td = Tensor::new(&[100.0_f32], &device).unwrap(); + buf.update_priorities_gpu(&high_idx, &high_td).unwrap(); + + // Sample 1000 times in batches of 10, count how often index 0 appears + let mut count_idx_0 = 0_usize; + let total_samples = 100; + let batch_size = 10; + for _ in 0..total_samples { + let batch = buf.sample_proportional(batch_size).unwrap(); + let indices: Vec = batch.indices.to_vec1().unwrap(); + count_idx_0 += indices.iter().filter(|&&i| i == 0_u32).count(); + } + + // With uniform sampling (alpha=0), expected = 1000 * 1/100 = 10 + // With priority 100^0.6 ≈ 15.85 vs rest at 1^0.6 = 1, + // expected proportion = 15.85 / (15.85 + 99) ≈ 13.8% + // So out of 1000 samples, index 0 should appear ~138 times + // With alpha=0.6 and high priority, it should appear much more than uniform (10) + assert!( + count_idx_0 > 30, + "Expected index 0 to be sampled frequently (got {}), priority weighting may not work", + count_idx_0 + ); +} + +#[test] +fn test_gpu_per_sampling_distribution_differs_from_uniform() { + // Two-sample KS-style test: compare GPU PER sampling with non-uniform priorities + // against a uniform distribution. They should differ. + let device = Device::Cpu; + let state_dim = 4; + let capacity = 50; + let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap(); + + // Fill buffer + fill_buffer(&mut buf, 50, state_dim); + + // Set skewed priorities: experience i gets priority (i+1) + // This creates a strong left-to-right priority gradient + for i in 0..50 { + let idx = Tensor::new(&[i as u32], &device).unwrap(); + let td = Tensor::new(&[(i as f32 + 1.0)], &device).unwrap(); + buf.update_priorities_gpu(&idx, &td).unwrap(); + } + + // Sample 5000 indices via proportional sampling + let n_samples = 500; + let batch_size = 10; + let mut counts = vec![0_usize; 50]; + for _ in 0..n_samples { + let batch = buf.sample_proportional(batch_size).unwrap(); + let indices: Vec = batch.indices.to_vec1().unwrap(); + for idx in indices { + if let Some(c) = counts.get_mut(idx as usize) { + *c += 1; + } + } + } + + let total = counts.iter().sum::() as f64; + + // KS test: compute max |empirical_CDF - uniform_CDF| + let n = counts.len() as f64; + let mut empirical_cum = 0.0_f64; + let mut max_ks = 0.0_f64; + for (i, &count) in counts.iter().enumerate() { + empirical_cum += count as f64 / total; + let uniform_cum = (i as f64 + 1.0) / n; + let diff = (empirical_cum - uniform_cum).abs(); + if diff > max_ks { + max_ks = diff; + } + } + + // Critical value for KS test at p=0.01 with N=50: ~1.63 / sqrt(N) ≈ 0.23 + // With skewed priorities (1..50)^0.6, the distribution should strongly differ from uniform + assert!( + max_ks > 0.05, + "KS statistic too small ({}), PER should differ significantly from uniform", + max_ks + ); +} + +#[test] +fn test_gpu_per_is_weights_correct_range() { + // IS weights should be in (0, 1] when beta < 1 and normalized by max weight + let device = Device::Cpu; + let state_dim = 8; + let capacity = 200; + let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap(); + + fill_buffer(&mut buf, 200, state_dim); + + // Set varied priorities + for i in 0..200 { + let idx = Tensor::new(&[i as u32], &device).unwrap(); + let td = Tensor::new(&[((i % 10) as f32 + 1.0)], &device).unwrap(); + buf.update_priorities_gpu(&idx, &td).unwrap(); + } + + let batch = buf.sample_proportional(64).unwrap(); + let weights: Vec = batch.weights.to_vec1().unwrap(); + + // All weights must be > 0 and <= 1 (normalized by max weight) + for (i, &w) in weights.iter().enumerate() { + assert!(w > 0.0, "Weight {} is non-positive: {}", i, w); + assert!( + w <= 1.0 + 1e-6, + "Weight {} exceeds 1.0: {} (should be normalized)", + i, w + ); + } + + // At least one weight should be 1.0 (the max weight normalizes to 1) + let has_one = weights.iter().any(|&w| (w - 1.0).abs() < 1e-4); + assert!( + has_one, + "No weight close to 1.0 found — normalization may be broken" + ); +} + +#[test] +fn test_gpu_per_ring_buffer_overwrites_correctly() { + // Fill past capacity, verify size stays at capacity and sampling still works + let device = Device::Cpu; + let state_dim = 4; + let capacity = 50; + let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap(); + + // Insert 120 experiences in 3 batches of 40 + for _ in 0..3 { + fill_buffer(&mut buf, 40, state_dim); + } + + assert_eq!(buf.len(), capacity); // capped at 50 + assert!(buf.can_sample(50)); + + // Sampling should still work + let batch = buf.sample_proportional(20).unwrap(); + assert_eq!(batch.states.dims(), &[20, state_dim]); + + // Indices should all be in [0, capacity) + let indices: Vec = batch.indices.to_vec1().unwrap(); + for &idx in &indices { + assert!( + (idx as usize) < capacity, + "Index {} out of range [0, {})", + idx, capacity + ); + } +}