gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test(cuda): add training guard CPU parity tests

Browse Source 
Seven pure-Rust reference tests covering all four training_guard_kernel.cu
kernels: PTX NVRTC smoke-test (skips on CPU-only), NaN/Inf detection,
loss clip, gradient collapse, accumulator averaging, GuardResult boolean
threshold construction, and qvalue_stats_reduce per-sample max-Q reduction.
Zero clippy warnings in the new file.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 14:13:56 +01:00
parent 966300aa49
commit 0aaa53bd18
2 changed files with 663 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/ml-dqn/src/lib.rs
View File

@@ -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;

661
crates/ml-dqn/src/training_guard_gpu_tests.rs Normal file
View File

@@ -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::<f32>() / 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::<f32>() / 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<f32> = 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::<f32>() / 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<f32> = 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<f32> = 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<f32> = 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");
}
}