diff --git a/crates/ml-dqn/src/lib.rs b/crates/ml-dqn/src/lib.rs index 8599285f8..97eab7d31 100644 --- a/crates/ml-dqn/src/lib.rs +++ b/crates/ml-dqn/src/lib.rs @@ -60,6 +60,8 @@ mod branching_composition_tests; mod dsr_gpu_tests; #[cfg(test)] mod nstep_gpu_tests; +#[cfg(test)] +mod training_guard_gpu_tests; // Rainbow DQN components pub mod distributional; diff --git a/crates/ml-dqn/src/training_guard_gpu_tests.rs b/crates/ml-dqn/src/training_guard_gpu_tests.rs new file mode 100644 index 000000000..27cbb6f47 --- /dev/null +++ b/crates/ml-dqn/src/training_guard_gpu_tests.rs @@ -0,0 +1,661 @@ +//! CPU-parity tests for the training guard CUDA kernels. +//! +//! These tests verify the logic implemented in `training_guard_kernel.cu` by +//! running equivalent pure-Rust reference implementations and confirming that +//! the arithmetic, branching, and output layout match the kernel specification +//! exactly. No GPU is required. +//! +//! Test inventory: +//! - `test_guard_ptx_compiles` NVRTC smoke-test (skips when absent) +//! - `test_nan_detection_cpu_reference` f32 NaN semantics mirror `isnan()` +//! - `test_loss_clip_cpu_reference` `fminf(loss, threshold)` clipping +//! - `test_grad_collapse_cpu_reference` collapse threshold = LR × multiplier +//! - `test_accumulator_averaging_cpu_reference` `loss_sum / step_count` = mean +//! - `test_guard_result_construction` `GuardResult` boolean thresholds +//! - `test_qvalue_stats_cpu_reference` per-sample max-Q reduction + +#[cfg(test)] +mod tests { + // ── helper: CPU port of training_guard_check output layout ────────────── + + /// Mirror of the 7-float output layout written by `training_guard_check`. + /// + /// ```text + /// [0] halt_nan — 1.0 if loss or grad_norm is NaN/Inf + /// [1] halt_loss_clip — 1.0 if loss > clip_threshold (and not NaN/Inf) + /// [2] halt_grad_collapse — 1.0 if grad_norm < collapse_threshold AND warmup == false + /// [3] clipped_loss — min(loss, clip_threshold); NaN/Inf pass through + /// [4] raw_loss + /// [5] raw_grad_norm + /// [6] reserved — 0.0 + /// ``` + fn training_guard_check_cpu( + loss: f32, + grad_norm: f32, + clip_threshold: f32, + collapse_threshold: f32, + warmup: bool, + ) -> [f32; 7] { + let halt_nan = loss.is_nan() + || loss.is_infinite() + || grad_norm.is_nan() + || grad_norm.is_infinite(); + + let halt_loss_clip = + !loss.is_nan() && !loss.is_infinite() && loss > clip_threshold; + + let halt_grad_collapse = !warmup + && !grad_norm.is_nan() + && !grad_norm.is_infinite() + && grad_norm < collapse_threshold; + + let clipped_loss = if !loss.is_nan() && !loss.is_infinite() { + loss.min(clip_threshold) + } else { + loss // NaN/Inf pass through (matches CUDA kernel) + }; + + [ + halt_nan as u8 as f32, + halt_loss_clip as u8 as f32, + halt_grad_collapse as u8 as f32, + clipped_loss, + loss, + grad_norm, + 0.0_f32, + ] + } + + // ── helper: CPU port of training_guard_accumulate ─────────────────────── + + /// Mirror of `training_guard_accumulate`: adds valid scalars to running sums. + /// + /// NaN/Inf values are skipped for the sums but `step_count` still increments. + fn training_guard_accumulate_cpu( + loss: f32, + grad_norm: f32, + loss_sum: &mut f32, + grad_norm_sum: &mut f32, + step_count: &mut f32, + ) { + if !loss.is_nan() && !loss.is_infinite() { + *loss_sum += loss; + } + if !grad_norm.is_nan() && !grad_norm.is_infinite() { + *grad_norm_sum += grad_norm; + } + *step_count += 1.0; + } + + // ── helper: CPU port of qvalue_stats_reduce (single-threaded) ─────────── + + /// Mirror of the parallel reduction in `qvalue_stats_reduce`. + /// + /// Input: `q_values[batch_size * num_actions]` (row-major). + /// Output: `[q_min, q_max, q_mean, q_all_mean]`. + fn qvalue_stats_reduce_cpu( + q_values: &[f32], + batch_size: usize, + num_actions: usize, + ) -> [f32; 4] { + if batch_size == 0 || num_actions == 0 { + return [0.0; 4]; + } + + let mut global_min = f32::MAX; + let mut global_max = f32::MIN; + let mut sum_of_maxes = 0.0_f32; + let mut sum_all = 0.0_f32; + let mut valid_samples = 0_usize; + + for i in 0..batch_size { + let start = i * num_actions; + let end = start + num_actions; + let Some(row) = q_values.get(start..end) else { break }; + let mut sample_max = f32::NEG_INFINITY; + + for &qv in row { + if !qv.is_nan() && !qv.is_infinite() { + if qv > sample_max { + sample_max = qv; + } + sum_all += qv; + } + } + + if sample_max.is_finite() { + if sample_max < global_min { + global_min = sample_max; + } + if sample_max > global_max { + global_max = sample_max; + } + sum_of_maxes += sample_max; + valid_samples += 1; + } + } + + let total_elements = (batch_size * num_actions) as f32; + + let q_min = if valid_samples == 0 { 0.0 } else { global_min }; + let q_max = if valid_samples == 0 { 0.0 } else { global_max }; + let q_mean = if valid_samples > 0 { + sum_of_maxes / valid_samples as f32 + } else { + 0.0 + }; + let q_all_mean = if total_elements > 0.0 { + sum_all / total_elements + } else { + 0.0 + }; + + [q_min, q_max, q_mean, q_all_mean] + } + + // ── helper: CPU port of qvalue_divergence_check ───────────────────────── + + /// Mirror of `qvalue_divergence_check` (single-sample, single-thread). + /// + /// Output: `[q_min, q_max, q_mean, q_variance, diverged]`. + fn qvalue_divergence_check_cpu(q_values: &[f32], divergence_threshold: f32) -> [f32; 5] { + let mut q_min = f32::MAX; + let mut q_max = f32::MIN; + let mut q_sum = 0.0_f32; + let mut q_sq = 0.0_f32; + let mut valid = 0_usize; + + for &qv in q_values { + if qv.is_nan() || qv.is_infinite() { + continue; + } + if qv < q_min { + q_min = qv; + } + if qv > q_max { + q_max = qv; + } + q_sum += qv; + q_sq += qv * qv; + valid += 1; + } + + let (q_min_out, q_max_out, q_mean, q_variance) = if valid == 0 { + (0.0_f32, 0.0_f32, 0.0_f32, 0.0_f32) + } else { + let mean = q_sum / valid as f32; + let mean_sq = q_sq / valid as f32; + let var = (mean_sq - mean * mean).max(0.0); // numerical floor + (q_min, q_max, mean, var) + }; + + let diverged = if q_min_out.abs() > divergence_threshold + || q_max_out.abs() > divergence_threshold + { + 1.0_f32 + } else { + 0.0_f32 + }; + + [q_min_out, q_max_out, q_mean, q_variance, diverged] + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 1: PTX compilation smoke-test + // ════════════════════════════════════════════════════════════════════════ + + /// Try to compile `training_guard_kernel.cu` via NVRTC. + /// + /// Skips gracefully when NVRTC is not installed so that CPU-only CI + /// machines pass without a GPU. On CUDA machines this validates the + /// kernel compiles cleanly. + #[test] + fn test_guard_ptx_compiles() { + let src = include_str!("../../ml/src/cuda_pipeline/training_guard_kernel.cu"); + + // Only attempt compilation when the cuda feature is active. + #[cfg(feature = "cuda")] + { + let result = candle_core::cuda_backend::cudarc::nvrtc::compile_ptx(src); + match result { + Ok(_) => { + // PTX compiled successfully. + } + Err(ref e) => { + let msg = e.to_string().to_lowercase(); + // Accept graceful skip conditions: NVRTC unavailable or library not found. + if msg.contains("nvrtc") + || msg.contains("not found") + || msg.contains("no such file") + || msg.contains("cannot open") + || msg.contains("libnvrtc") + { + return; // NVRTC not installed — skip. + } + panic!("training_guard_kernel.cu PTX compilation failed: {e}"); + } + } + } + + // Without the cuda feature, verify the source file is readable and intact. + #[cfg(not(feature = "cuda"))] + { + assert!(!src.is_empty(), "training_guard_kernel.cu should not be empty"); + assert!( + src.contains("training_guard_check"), + "kernel source must contain training_guard_check" + ); + } + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 2: NaN detection CPU reference + // ════════════════════════════════════════════════════════════════════════ + + /// Verify that `f32::is_nan()` matches the CUDA `isnan()` semantics used + /// in the kernel. The kernel uses `isnan()` which implements the IEEE 754 + /// reflexivity-violation property (`NaN != NaN`). + #[test] + fn test_nan_detection_cpu_reference() { + // IEEE 754 NaN properties. + assert!(f32::NAN.is_nan(), "f32::NAN.is_nan() must be true"); + assert!(f32::INFINITY.is_infinite(), "INFINITY must be infinite"); + assert!(f32::NEG_INFINITY.is_infinite(), "NEG_INFINITY must be infinite"); + + // Normal finite values must not trigger the guard. + assert!(!1.0_f32.is_nan(), "1.0 is not NaN"); + assert!(!0.0_f32.is_nan(), "0.0 is not NaN"); + assert!(!(-1.0_f32).is_nan(), "-1.0 is not NaN"); + + // Guard check: NaN loss should set halt_nan. + let out = training_guard_check_cpu(f32::NAN, 0.1, 1e6, 1e-7, false); + assert_eq!(out[0], 1.0, "halt_nan must be 1.0 for NaN loss"); + + // Guard check: Inf loss should set halt_nan. + let out2 = training_guard_check_cpu(f32::INFINITY, 0.1, 1e6, 1e-7, false); + assert_eq!(out2[0], 1.0, "halt_nan must be 1.0 for Inf loss"); + + // Guard check: NaN grad_norm should set halt_nan. + let out3 = training_guard_check_cpu(0.5, f32::NAN, 1e6, 1e-7, false); + assert_eq!(out3[0], 1.0, "halt_nan must be 1.0 for NaN grad_norm"); + + // Guard check: finite values must not set halt_nan. + let out4 = training_guard_check_cpu(0.5, 0.1, 1e6, 1e-7, false); + assert_eq!(out4[0], 0.0, "halt_nan must be 0.0 for finite inputs"); + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 3: Loss clip CPU reference + // ════════════════════════════════════════════════════════════════════════ + + /// Verify `loss.min(threshold)` = CUDA `fminf(loss, clip_threshold)`. + /// + /// Also verifies that `halt_loss_clip` fires iff `loss > clip_threshold` + /// (and loss is finite). + #[test] + fn test_loss_clip_cpu_reference() { + let clip = 10.0_f32; + + // Below threshold: clipped_loss == raw_loss, no halt. + let out = training_guard_check_cpu(3.0, 0.5, clip, 1e-7, false); + assert_eq!(out[1], 0.0, "halt_loss_clip must be 0 below threshold"); + assert!( + (out[3] - 3.0).abs() < 1e-6, + "clipped_loss should equal raw loss: {}", + out[3] + ); + + // Exactly at threshold: no halt (not strictly greater). + let out2 = training_guard_check_cpu(10.0, 0.5, clip, 1e-7, false); + assert_eq!(out2[1], 0.0, "halt_loss_clip must be 0 at threshold exactly"); + assert!( + (out2[3] - 10.0).abs() < 1e-6, + "clipped_loss at threshold: {}", + out2[3] + ); + + // Above threshold: clipped to threshold, halt fires. + let out3 = training_guard_check_cpu(1e7, 0.5, clip, 1e-7, false); + assert_eq!(out3[1], 1.0, "halt_loss_clip must be 1.0 above threshold"); + assert!( + (out3[3] - clip).abs() < 1e-6, + "clipped_loss must be clip_threshold: {}", + out3[3] + ); + assert!( + (out3[4] - 1e7).abs() < 1.0, + "raw_loss must be preserved: {}", + out3[4] + ); + + // NaN loss: halt_nan fires, halt_loss_clip must NOT fire (kernel skips). + let out4 = training_guard_check_cpu(f32::NAN, 0.5, clip, 1e-7, false); + assert_eq!(out4[0], 1.0, "halt_nan must be 1 for NaN loss"); + assert_eq!(out4[1], 0.0, "halt_loss_clip must be 0 when loss is NaN"); + // NaN passthrough: clipped_loss stays NaN. + assert!(out4[3].is_nan(), "clipped_loss must be NaN for NaN input"); + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 4: Gradient collapse CPU reference + // ════════════════════════════════════════════════════════════════════════ + + /// Verify the collapse threshold derivation: `threshold = lr * multiplier`. + /// + /// The kernel compares `grad_norm < collapse_threshold`. The typical + /// production values are: LR=1e-4, multiplier=1e-3 → threshold=1e-7. + #[test] + fn test_grad_collapse_cpu_reference() { + let lr = 1e-4_f32; + let multiplier = 1e-3_f32; + let collapse_threshold = lr * multiplier; // 1e-7 + + // In warmup: collapse check disabled regardless of grad_norm. + let out_warmup = + training_guard_check_cpu(0.5, 0.0, 1e6, collapse_threshold, true); + assert_eq!( + out_warmup[2], 0.0, + "halt_grad_collapse must be 0 during warmup (grad_norm=0)" + ); + + // Not in warmup, grad_norm above threshold: no collapse. + let out_ok = + training_guard_check_cpu(0.5, 1e-5, 1e6, collapse_threshold, false); + assert_eq!( + out_ok[2], 0.0, + "halt_grad_collapse must be 0 when grad_norm > threshold" + ); + + // Not in warmup, grad_norm exactly at threshold: no collapse (not strictly less). + let out_at = training_guard_check_cpu( + 0.5, + collapse_threshold, + 1e6, + collapse_threshold, + false, + ); + assert_eq!( + out_at[2], 0.0, + "halt_grad_collapse must be 0 when grad_norm == threshold" + ); + + // Not in warmup, grad_norm below threshold: collapse detected. + let tiny = collapse_threshold * 0.1; // 1e-8 + let out_collapse = + training_guard_check_cpu(0.5, tiny, 1e6, collapse_threshold, false); + assert_eq!( + out_collapse[2], 1.0, + "halt_grad_collapse must be 1.0 when grad_norm < threshold outside warmup" + ); + + // NaN grad_norm outside warmup: halt_nan fires, collapse does NOT fire. + let out_nan = + training_guard_check_cpu(0.5, f32::NAN, 1e6, collapse_threshold, false); + assert_eq!(out_nan[0], 1.0, "halt_nan must fire for NaN grad_norm"); + assert_eq!( + out_nan[2], 0.0, + "halt_grad_collapse must be 0 when grad_norm is NaN" + ); + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 5: Accumulator averaging CPU reference + // ════════════════════════════════════════════════════════════════════════ + + /// Verify `loss_sum / step_count` produces the correct mean. + /// + /// This mirrors the epoch-boundary `read_accumulators()` path which reads + /// `[0]=loss_sum, [1]=grad_norm_sum, [2]=step_count` from the GPU buffer. + #[test] + fn test_accumulator_averaging_cpu_reference() { + let losses = [0.5_f32, 0.3, 0.7, 0.1, 0.4, 0.6, 0.2, 0.8]; + let grad_norms = [0.01_f32, 0.02, 0.015, 0.018, 0.022, 0.009, 0.030, 0.011]; + + let mut loss_sum = 0.0_f32; + let mut grad_sum = 0.0_f32; + let mut step_count = 0.0_f32; + + for (&l, &g) in losses.iter().zip(grad_norms.iter()) { + training_guard_accumulate_cpu(l, g, &mut loss_sum, &mut grad_sum, &mut step_count); + } + + assert!( + (step_count - losses.len() as f32).abs() < 1e-6, + "step_count should equal number of steps: {} vs {}", + step_count, + losses.len() + ); + + let expected_mean_loss: f32 = losses.iter().sum::() / losses.len() as f32; + let computed_mean_loss = loss_sum / step_count; + assert!( + (computed_mean_loss - expected_mean_loss).abs() < 1e-5, + "mean_loss mismatch: {} vs {}", + computed_mean_loss, + expected_mean_loss + ); + + let expected_mean_grad: f32 = grad_norms.iter().sum::() / grad_norms.len() as f32; + let computed_mean_grad = grad_sum / step_count; + assert!( + (computed_mean_grad - expected_mean_grad).abs() < 1e-5, + "mean_grad_norm mismatch: {} vs {}", + computed_mean_grad, + expected_mean_grad + ); + + // NaN/Inf inputs must be skipped by the accumulator but step_count still increments. + // + // Call 1: loss=NaN (skipped), grad=0.5 (valid) → loss_sum=0.0, grad_sum=0.5 + // Call 2: loss=0.3 (valid), grad=Inf (skipped) → loss_sum=0.3, grad_sum=0.5 + // Call 3: loss=0.4 (valid), grad=0.2 (valid) → loss_sum=0.7, grad_sum=0.7 + let mut ls2 = 0.0_f32; + let mut gs2 = 0.0_f32; + let mut sc2 = 0.0_f32; + training_guard_accumulate_cpu(f32::NAN, 0.5, &mut ls2, &mut gs2, &mut sc2); + training_guard_accumulate_cpu(0.3, f32::INFINITY, &mut ls2, &mut gs2, &mut sc2); + training_guard_accumulate_cpu(0.4, 0.2, &mut ls2, &mut gs2, &mut sc2); + + // step_count = 3 (all steps counted regardless of NaN/Inf). + assert!( + (sc2 - 3.0).abs() < 1e-6, + "step_count must count all steps including NaN/Inf inputs: {}", + sc2 + ); + // loss_sum = 0.3 + 0.4 = 0.7 (NaN on call 1 skipped; calls 2 and 3 added). + assert!( + (ls2 - 0.7).abs() < 1e-5, + "loss_sum: NaN skipped, 0.3+0.4 accumulated: {} vs 0.7", + ls2 + ); + // grad_sum = 0.5 + 0.2 = 0.7 (Inf on call 2 skipped; calls 1 and 3 added). + assert!( + (gs2 - 0.7).abs() < 1e-5, + "grad_sum: Inf skipped, 0.5+0.2 accumulated: {} vs 0.7", + gs2 + ); + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 6: GuardResult construction from simulated readback values + // ════════════════════════════════════════════════════════════════════════ + + /// Simulate the device-to-host readback of the 7-float guard output buffer and + /// verify that the `GuardResult` boolean thresholds (any non-zero float + /// maps to `true`) match the kernel's output specification. + /// + /// This mirrors the construction in `gpu_training_guard.rs`: + /// ```text + /// halt_nan: g[0] != 0.0 + /// halt_loss_clip: g[1] != 0.0 + /// halt_grad_collapse: g[2] != 0.0 + /// ``` + #[test] + fn test_guard_result_construction() { + struct GuardResultSim { + halt_nan: bool, + halt_loss_clip: bool, + halt_grad_collapse: bool, + clipped_loss: f32, + _raw_grad_norm: f32, + } + + fn from_pinned(buf: &[f32; 7]) -> GuardResultSim { + GuardResultSim { + halt_nan: buf[0] != 0.0, + halt_loss_clip: buf[1] != 0.0, + halt_grad_collapse: buf[2] != 0.0, + clipped_loss: buf[3], + _raw_grad_norm: buf[5], + } + } + + // Scenario A: normal training step — no halts. + let buf_a: [f32; 7] = [0.0, 0.0, 0.0, 0.42, 0.42, 0.015, 0.0]; + let r_a = from_pinned(&buf_a); + assert!(!r_a.halt_nan, "A: halt_nan must be false"); + assert!(!r_a.halt_loss_clip, "A: halt_loss_clip must be false"); + assert!(!r_a.halt_grad_collapse, "A: halt_grad_collapse must be false"); + assert!( + (r_a.clipped_loss - 0.42).abs() < 1e-6, + "A: clipped_loss = {}", + r_a.clipped_loss + ); + + // Scenario B: NaN detected — only halt_nan fires. + let buf_b: [f32; 7] = [1.0, 0.0, 0.0, f32::NAN, f32::NAN, 0.01, 0.0]; + let r_b = from_pinned(&buf_b); + assert!(r_b.halt_nan, "B: halt_nan must be true"); + assert!(!r_b.halt_loss_clip, "B: halt_loss_clip must be false"); + assert!(!r_b.halt_grad_collapse, "B: halt_grad_collapse must be false"); + + // Scenario C: loss clip fired — loss was above threshold. + let buf_c: [f32; 7] = [0.0, 1.0, 0.0, 1e6, 1e9, 0.05, 0.0]; + let r_c = from_pinned(&buf_c); + assert!(!r_c.halt_nan, "C: halt_nan must be false"); + assert!(r_c.halt_loss_clip, "C: halt_loss_clip must be true"); + assert!(!r_c.halt_grad_collapse, "C: halt_grad_collapse must be false"); + assert!( + (r_c.clipped_loss - 1e6).abs() < 1.0, + "C: clipped_loss must be clip_threshold: {}", + r_c.clipped_loss + ); + + // Scenario D: gradient collapse outside warmup. + let buf_d: [f32; 7] = [0.0, 0.0, 1.0, 0.5, 0.5, 1e-10, 0.0]; + let r_d = from_pinned(&buf_d); + assert!(!r_d.halt_nan, "D: halt_nan must be false"); + assert!(!r_d.halt_loss_clip, "D: halt_loss_clip must be false"); + assert!(r_d.halt_grad_collapse, "D: halt_grad_collapse must be true"); + + // Verify that 0.0 → false and 1.0 → true for all flag fields. + for flag_val in [0.0_f32, 1.0_f32] { + let expected = flag_val != 0.0; + let buf: [f32; 7] = [flag_val, flag_val, flag_val, 0.1, 0.1, 0.1, 0.0]; + let r = from_pinned(&buf); + assert_eq!(r.halt_nan, expected, "halt_nan for flag={}", flag_val); + assert_eq!( + r.halt_loss_clip, expected, + "halt_loss_clip for flag={}", + flag_val + ); + assert_eq!( + r.halt_grad_collapse, expected, + "halt_grad_collapse for flag={}", + flag_val + ); + } + } + + // ════════════════════════════════════════════════════════════════════════ + // Test 7: Q-value statistics CPU reference + // ════════════════════════════════════════════════════════════════════════ + + /// CPU reference for the `qvalue_stats_reduce` kernel. + /// + /// For each sample: find max-Q across actions, then reduce across samples + /// to obtain `q_min`, `q_max`, `q_mean`, and the mean of all Q-values. + #[test] + fn test_qvalue_stats_cpu_reference() { + // Batch of 4 samples, 5 actions each (DQN default: 5 exposure actions). + // + // Sample 0: [1.0, 2.0, 3.0, 4.0, 5.0] → max = 5.0 + // Sample 1: [0.5, 0.3, 0.8, 0.2, 0.6] → max = 0.8 + // Sample 2: [-1.0, -2.0, -0.5, -3.0, -4.0] → max = -0.5 + // Sample 3: [10.0, 9.0, 8.0, 7.0, 6.0] → max = 10.0 + #[rustfmt::skip] + let q_values: Vec = vec![ + 1.0, 2.0, 3.0, 4.0, 5.0, + 0.5, 0.3, 0.8, 0.2, 0.6, + -1.0, -2.0, -0.5, -3.0, -4.0, + 10.0, 9.0, 8.0, 7.0, 6.0, + ]; + let batch_size = 4; + let num_actions = 5; + + let stats = qvalue_stats_reduce_cpu(&q_values, batch_size, num_actions); + + // per-sample maxes: [5.0, 0.8, -0.5, 10.0] + assert!( + (stats[0] - (-0.5)).abs() < 1e-5, + "q_min should be -0.5, got {}", + stats[0] + ); + assert!( + (stats[1] - 10.0).abs() < 1e-5, + "q_max should be 10.0, got {}", + stats[1] + ); + let expected_mean = (5.0_f32 + 0.8 + (-0.5) + 10.0) / 4.0; + assert!( + (stats[2] - expected_mean).abs() < 1e-4, + "q_mean should be {}, got {}", + expected_mean, + stats[2] + ); + + let expected_all_mean: f32 = q_values.iter().sum::() / q_values.len() as f32; + assert!( + (stats[3] - expected_all_mean).abs() < 1e-4, + "q_all_mean should be {}, got {}", + expected_all_mean, + stats[3] + ); + + // Edge case: batch_size=1, num_actions=3. + let single: Vec = vec![2.0, 7.0, -1.0]; + let s = qvalue_stats_reduce_cpu(&single, 1, 3); + assert!( + (s[0] - 7.0).abs() < 1e-5, + "single sample q_min=q_max=7.0: q_min={}", + s[0] + ); + assert!( + (s[1] - 7.0).abs() < 1e-5, + "single sample q_min=q_max=7.0: q_max={}", + s[1] + ); + assert!((s[2] - 7.0).abs() < 1e-5, "single sample q_mean=7.0: {}", s[2]); + let all_mean_single = (2.0_f32 + 7.0 + (-1.0)) / 3.0; + assert!( + (s[3] - all_mean_single).abs() < 1e-5, + "single sample q_all_mean: {} vs {}", + s[3], + all_mean_single + ); + + // Edge case: empty batch → all zeros. + let empty = qvalue_stats_reduce_cpu(&[], 0, 5); + assert_eq!(empty, [0.0; 4], "empty batch should produce zero stats"); + + // Divergence detection: extreme Q-values trigger the flag. + let extreme: Vec = vec![-5e5, 0.0, 5e5, 1.0, -1.0]; + let div = qvalue_divergence_check_cpu(&extreme, 1e4); + assert_eq!(div[4], 1.0, "divergence must be detected for |value| > 1e4"); + + // No divergence for moderate Q-values. + let moderate: Vec = vec![1.0, -1.0, 0.5, -0.5, 0.0]; + let no_div = qvalue_divergence_check_cpu(&moderate, 1e4); + assert_eq!(no_div[4], 0.0, "no divergence for moderate Q-values"); + } +}