335 lines
12 KiB
Rust
335 lines
12 KiB
Rust
#![allow(
|
|
clippy::assertions_on_constants,
|
|
clippy::assertions_on_result_states,
|
|
clippy::clone_on_copy,
|
|
clippy::decimal_literal_representation,
|
|
clippy::doc_markdown,
|
|
clippy::empty_line_after_doc_comments,
|
|
clippy::field_reassign_with_default,
|
|
clippy::get_unwrap,
|
|
clippy::identity_op,
|
|
clippy::inconsistent_digit_grouping,
|
|
clippy::indexing_slicing,
|
|
clippy::integer_division,
|
|
clippy::len_zero,
|
|
clippy::let_underscore_must_use,
|
|
clippy::manual_div_ceil,
|
|
clippy::manual_let_else,
|
|
clippy::manual_range_contains,
|
|
clippy::modulo_arithmetic,
|
|
clippy::needless_range_loop,
|
|
clippy::non_ascii_literal,
|
|
clippy::redundant_clone,
|
|
clippy::shadow_reuse,
|
|
clippy::shadow_same,
|
|
clippy::shadow_unrelated,
|
|
clippy::single_match_else,
|
|
clippy::str_to_string,
|
|
clippy::string_slice,
|
|
clippy::tests_outside_test_module,
|
|
clippy::too_many_lines,
|
|
clippy::unnecessary_wraps,
|
|
clippy::unseparated_literal_suffix,
|
|
clippy::use_debug,
|
|
clippy::useless_vec,
|
|
clippy::wildcard_enum_match_arm,
|
|
clippy::else_if_without_else,
|
|
clippy::expect_used,
|
|
clippy::missing_const_for_fn,
|
|
clippy::similar_names,
|
|
clippy::type_complexity,
|
|
clippy::collapsible_else_if,
|
|
clippy::doc_lazy_continuation,
|
|
clippy::items_after_test_module,
|
|
clippy::map_clone,
|
|
clippy::multiple_unsafe_ops_per_block,
|
|
clippy::unwrap_or_default,
|
|
clippy::assign_op_pattern,
|
|
clippy::needless_borrow,
|
|
clippy::println_empty_string,
|
|
clippy::unnecessary_cast,
|
|
clippy::used_underscore_binding,
|
|
clippy::create_dir,
|
|
clippy::implicit_saturating_sub,
|
|
clippy::exit,
|
|
clippy::expect_fun_call,
|
|
clippy::too_many_arguments,
|
|
clippy::unnecessary_map_or,
|
|
clippy::unwrap_used,
|
|
dead_code,
|
|
unused_imports,
|
|
unused_variables,
|
|
clippy::cloned_ref_to_slice_refs,
|
|
clippy::neg_multiply,
|
|
clippy::while_let_loop,
|
|
clippy::bool_assert_comparison,
|
|
clippy::excessive_precision,
|
|
clippy::trivially_copy_pass_by_ref,
|
|
clippy::op_ref,
|
|
clippy::redundant_closure,
|
|
clippy::unnecessary_lazy_evaluations,
|
|
clippy::if_then_some_else_none,
|
|
clippy::unnecessary_to_owned,
|
|
clippy::single_component_path_imports,
|
|
)]
|
|
//! 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
|
|
|
|
// candle eliminated — test uses native cudarc APIs
|
|
use ml::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
|
|
use std::sync::Arc;
|
|
use cudarc::driver::{CudaContext, CudaSlice, CudaStream};
|
|
|
|
fn require_cuda_stream() -> Arc<CudaStream> {
|
|
let ctx = CudaContext::new(0).expect("CUDA required");
|
|
ctx.new_stream().expect("CUDA stream required")
|
|
}
|
|
|
|
fn test_config(capacity: usize, _state_dim: usize) -> GpuReplayBufferConfig {
|
|
GpuReplayBufferConfig {
|
|
capacity,
|
|
alpha: 0.6,
|
|
beta_start: 0.4,
|
|
beta_max: 1.0,
|
|
beta_annealing_steps: 1000,
|
|
epsilon: 1e-6,
|
|
max_memory_bytes: 4 * 1024 * 1024 * 1024,
|
|
max_batch_size: 1024,
|
|
}
|
|
}
|
|
|
|
/// Insert N random experiences into the buffer via CudaSlice uploads.
|
|
fn fill_buffer(buf: &mut GpuReplayBuffer, n: usize, state_dim: usize) {
|
|
use rand::Rng;
|
|
let stream = buf.stream();
|
|
let mut rng = rand::thread_rng();
|
|
|
|
// Generate random f32 data on host, upload to GPU
|
|
let states_host: Vec<f32> = (0..n * state_dim).map(|_| rng.gen::<f32>()).collect();
|
|
let next_states_host: Vec<f32> = (0..n * state_dim).map(|_| rng.gen::<f32>()).collect();
|
|
let actions_host: Vec<u32> = vec![0_u32; n];
|
|
let rewards_host: Vec<f32> = (0..n).map(|_| rng.gen::<f32>()).collect();
|
|
let dones_host: Vec<f32> = vec![0.0_f32; n];
|
|
|
|
let sf = stream.clone_htod(&states_host).unwrap();
|
|
let nsf = stream.clone_htod(&next_states_host).unwrap();
|
|
let af: CudaSlice<u32> = stream.clone_htod(&actions_host).unwrap();
|
|
let rf = stream.clone_htod(&rewards_host).unwrap();
|
|
let df = stream.clone_htod(&dones_host).unwrap();
|
|
let aux_sign: CudaSlice<i32> = stream.alloc_zeros::<i32>(n).unwrap();
|
|
let aux_conf: CudaSlice<f32> = stream.alloc_zeros::<f32>(n).unwrap();
|
|
|
|
let aux_outcome: CudaSlice<i32> = stream.alloc_zeros::<i32>(n).unwrap();
|
|
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &aux_sign, &aux_conf, &aux_outcome, n).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_per_training_loop_cycle() {
|
|
// End-to-end: insert -> sample -> fake TD error -> update priorities -> repeat
|
|
let stream = require_cuda_stream();
|
|
let state_dim = 16;
|
|
let capacity = 500;
|
|
let batch_size = 32;
|
|
let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).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 via CudaSlice lengths
|
|
assert_eq!(batch.batch_size * batch.state_dim, batch_size * state_dim); // f32 states [bs*sd]
|
|
assert_eq!(batch.batch_size, batch_size);
|
|
assert_eq!(batch.batch_size, batch_size);
|
|
assert_eq!(batch.batch_size, batch_size);
|
|
assert_eq!(batch.batch_size, batch_size);
|
|
|
|
// Simulate TD errors: decreasing magnitude over epochs
|
|
let scale = 1.0_f32 / (epoch as f32 + 1.0);
|
|
let td_host: Vec<f32> = vec![scale; batch_size];
|
|
let td_errors: CudaSlice<f32> = stream.clone_htod(&td_host).unwrap();
|
|
|
|
let idx_ptr = batch.indices_ptr;
|
|
buf.update_priorities_gpu_raw(idx_ptr, &td_errors, batch_size, None).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 stream = require_cuda_stream();
|
|
let state_dim = 4;
|
|
let capacity = 100;
|
|
let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).unwrap();
|
|
|
|
// Fill 100 experiences
|
|
fill_buffer(&mut buf, 100, state_dim);
|
|
|
|
// Set experience 0 to very high priority
|
|
let high_idx: CudaSlice<u32> = stream.clone_htod(&[0_u32]).unwrap();
|
|
let high_td: CudaSlice<f32> = stream.clone_htod(&[100.0]).unwrap();
|
|
buf.update_priorities_gpu(&high_idx, &high_td, 1, None).unwrap();
|
|
|
|
// Sample 1000 times in batches of 10, count how often index 0 appears
|
|
// Download indices to host, count matches
|
|
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();
|
|
// Download u32 indices to host
|
|
let mut indices_host = vec![0_u32; batch_size];
|
|
stream.memcpy_dtoh(&buf.sample_indices_ref(), &mut indices_host).unwrap();
|
|
count_idx_0 += indices_host.iter().filter(|&&idx| idx == 0).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 stream = require_cuda_stream();
|
|
let state_dim = 4;
|
|
let capacity = 50;
|
|
let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).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: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
|
|
let td: CudaSlice<f32> = stream.clone_htod(&[i as f32 + 1.0]).unwrap();
|
|
buf.update_priorities_gpu(&idx, &td, 1, None).unwrap();
|
|
}
|
|
|
|
// Sample 5000 indices via proportional sampling
|
|
// Validate mean index is above uniform mean (24.5),
|
|
// proving the priority gradient skews sampling toward higher-priority (higher-index) items
|
|
let n_samples = 500;
|
|
let batch_size = 10;
|
|
let mut sum_of_means = 0.0_f64;
|
|
for _ in 0..n_samples {
|
|
let batch = buf.sample_proportional(batch_size).unwrap();
|
|
// Download u32 indices to host, compute mean
|
|
let mut indices_host = vec![0_u32; batch_size];
|
|
stream.memcpy_dtoh(&buf.sample_indices_ref(), &mut indices_host).unwrap();
|
|
let mean_idx: f64 = indices_host.iter().map(|&i| i as f64).sum::<f64>() / batch_size as f64;
|
|
sum_of_means += mean_idx;
|
|
}
|
|
let overall_mean = sum_of_means / n_samples as f64;
|
|
|
|
// With priorities i+1 (1..50) and alpha=0.6, higher indices get sampled more.
|
|
// Uniform mean = 24.5. Skewed mean should be significantly above that.
|
|
assert!(
|
|
overall_mean > 27.0,
|
|
"Mean sampled index ({:.2}) not much above uniform 24.5 — PER should skew toward higher-priority items",
|
|
overall_mean
|
|
);
|
|
}
|
|
|
|
#[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 stream = require_cuda_stream();
|
|
let state_dim = 8;
|
|
let capacity = 200;
|
|
let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).unwrap();
|
|
|
|
fill_buffer(&mut buf, 200, state_dim);
|
|
|
|
// Set varied priorities
|
|
for i in 0..200 {
|
|
let idx: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
|
|
let td: CudaSlice<f32> = stream.clone_htod(&[((i % 10) as f32 + 1.0)]).unwrap();
|
|
buf.update_priorities_gpu(&idx, &td, 1, None).unwrap();
|
|
}
|
|
|
|
let batch = buf.sample_proportional(64).unwrap();
|
|
|
|
// Download f32 weights to host for validation
|
|
let mut weights_host = vec![0.0_f32; 64];
|
|
stream.memcpy_dtoh(&buf.sample_weights_ref(), &mut weights_host).unwrap();
|
|
|
|
let w_min = weights_host.iter().copied().fold(f32::INFINITY, f32::min);
|
|
let w_max = weights_host.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
|
|
|
assert!(w_min > 0.0, "Minimum weight is non-positive: {}", w_min);
|
|
assert!(
|
|
w_max <= 1.0 + 1e-6,
|
|
"Maximum weight exceeds 1.0: {} (should be normalized)",
|
|
w_max
|
|
);
|
|
// At least the max weight should be ~1.0 (the max weight normalizes to 1)
|
|
assert!(
|
|
(w_max - 1.0).abs() < 1e-4,
|
|
"Max weight {} not close to 1.0 — normalization may be broken",
|
|
w_max
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_per_ring_buffer_overwrites_correctly() {
|
|
// Fill past capacity, verify size stays at capacity and sampling still works
|
|
let stream = require_cuda_stream();
|
|
let state_dim = 4;
|
|
let capacity = 50;
|
|
let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).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.batch_size * batch.state_dim, 20 * state_dim); // f32 states [bs*sd]
|
|
|
|
// Download u32 indices to host for range validation
|
|
let mut indices_host = vec![0_u32; 20];
|
|
stream.memcpy_dtoh(&buf.sample_indices_ref(), &mut indices_host).unwrap();
|
|
|
|
let idx_min = *indices_host.iter().min().unwrap();
|
|
let idx_max = *indices_host.iter().max().unwrap();
|
|
assert!(
|
|
idx_min < capacity as u32,
|
|
"Minimum index {} is out of range",
|
|
idx_min
|
|
);
|
|
assert!(
|
|
idx_max < capacity as u32,
|
|
"Maximum index {} out of range [0, {})",
|
|
idx_max, capacity
|
|
);
|
|
}
|
|
|