diff --git a/crates/ml/src/trainers/dqn/distributional_q_tests.rs b/crates/ml/src/trainers/dqn/distributional_q_tests.rs index c640a652c..b16edd96f 100644 --- a/crates/ml/src/trainers/dqn/distributional_q_tests.rs +++ b/crates/ml/src/trainers/dqn/distributional_q_tests.rs @@ -1626,3 +1626,174 @@ fn test_2c_mag_ord_urg_branches_unaffected() { ); } +/// Phase 2 Test 2.D — Real-batch e2e: synthetic batch with eval-mode argmax +/// matches Rust ground-truth E[Q] argmax per sample; train-mode active_frac +/// > 40%. +/// +/// The plan-prescribed approach reuses the Phase 0 Test 0.F converged +/// checkpoint loader. Wiring up that loader (DQN::new + safetensors load +/// + branching forward eval) is heavyweight for a 1-2-hour dispatch and +/// is already exercised by Test 0.F itself. Replacement strategy from the +/// dispatch brief: synthetic batch via direct buffer write, with random +/// per-(sample, direction) C51 logits and per-sample adaptive support +/// covering a realistic v_min/v_max range. +/// +/// Setup: +/// - 256 distinct samples; per-sample `b_logits` filled with deterministic +/// pseudorandom values via a hash of `(sample_idx, direction, atom)`. +/// - Per-sample adaptive support: each sample gets its own (v_min, v_max, +/// delta_z) drawn from realistic ranges seen in production +/// (v_min ∈ [-1.0, -0.2], v_max ∈ [0.2, 1.0]). +/// +/// Assertions: +/// - Eval mode: per-sample dir_idx EXACTLY matches the argmax over Rust +/// ground-truth E[Q]_d = sum_a softmax(b_logits[i, d])[a] * atom_vals[i, d, a]. +/// - Train mode: count(d ∈ {Long=2, Short=0}) / batch > 0.40 — Thompson +/// over realistic non-degenerate distributions explores actively. +#[test] +#[ignore] +fn test_2d_real_batch_e2e() { + let stream = make_test_stream(); + let n_atoms: i32 = 21; + let batch: usize = 256; + + // Deterministic pseudorandom: simple integer hash for reproducibility. + fn hash_to_unit(a: u32, b: u32, c: u32) -> f32 { + let mut h = a.wrapping_mul(0x9E3779B1); + h ^= b.wrapping_add(0x85EBCA6B).wrapping_mul(0xC2B2AE35); + h ^= c.wrapping_add(0x27D4EB2F).wrapping_mul(0x165667B1); + h ^= h >> 16; + // Map to [-1, 1]. + let u = (h as f32) / (u32::MAX as f32); + 2.0_f32 * u - 1.0 + } + + // Per-sample adaptive support — realistic [v_min, v_max] ranges. + let mut support_host = vec![0.0_f32; batch * (PROD_B0 as usize) * 3]; + for i in 0..batch { + let v_min_offset = hash_to_unit(i as u32, 0xAA, 0).abs() * 0.8; // [0, 0.8] + let v_max_offset = hash_to_unit(i as u32, 0xBB, 0).abs() * 0.8; // [0, 0.8] + let v_min = -0.2_f32 - v_min_offset; + let v_max = 0.2_f32 + v_max_offset; + let delta_z = (v_max - v_min) / (n_atoms as f32 - 1.0); + for d in 0..(PROD_B0 as usize) { + let base = i * (PROD_B0 as usize) * 3 + d * 3; + support_host[base] = v_min; + support_host[base + 1] = v_max; + support_host[base + 2] = delta_z; + } + } + + // Per-(sample, direction, atom) random logits in [-1, 1] — matches the + // post-Xavier-init scale of fresh C51 head outputs. + let mut b_logits_host = + vec![0.0_f32; batch * (PROD_B0 as usize) * (n_atoms as usize)]; + for i in 0..batch { + for d in 0..(PROD_B0 as usize) { + for a in 0..(n_atoms as usize) { + let base = i * (PROD_B0 as usize) * (n_atoms as usize) + + d * (n_atoms as usize) + a; + b_logits_host[base] = hash_to_unit(i as u32, d as u32, a as u32); + } + } + } + + // q_values uniform; mag/ord/urg fall through Boltzmann. + let q_values_host = vec![0.0_f32; batch * PROD_Q_STRIDE]; + + // ── Compute Rust ground-truth E[Q] argmax per sample ───────────────── + let mut ground_truth_dir = vec![0_i32; batch]; + let mut ground_truth_eq = vec![[0.0_f32; 4]; batch]; + for i in 0..batch { + let support_base = i * (PROD_B0 as usize) * 3; + let v_min = support_host[support_base]; + let delta_z = support_host[support_base + 2]; + let mut best_q = f32::NEG_INFINITY; + let mut best_d = 0_i32; + for d in 0..(PROD_B0 as i32) { + let logits_base = i * (PROD_B0 as usize) * (n_atoms as usize) + + (d as usize) * (n_atoms as usize); + let logits = + &b_logits_host[logits_base..logits_base + n_atoms as usize]; + // Numerically-stable softmax (matches softmax_c51_inline). + let max_logit = logits + .iter() + .cloned() + .fold(f32::NEG_INFINITY, f32::max); + let mut sum = 0.0_f32; + let mut probs = vec![0.0_f32; n_atoms as usize]; + for a in 0..(n_atoms as usize) { + probs[a] = (logits[a] - max_logit).exp(); + sum += probs[a]; + } + let inv_sum = if sum > 0.0 { 1.0 / sum } else { 1.0 / n_atoms as f32 }; + for p in probs.iter_mut() { + *p *= inv_sum; + } + // Linear support (atom_positions = NULL in our launch). + let mut e = 0.0_f32; + for a in 0..(n_atoms as usize) { + let atom_v = v_min + (a as f32) * delta_z; + e += probs[a] * atom_v; + } + ground_truth_eq[i][d as usize] = e; + if e > best_q { + best_q = e; + best_d = d; + } + } + ground_truth_dir[i] = best_d; + } + + // ── Eval mode: dir_idx must match ground truth per sample ──────────── + let mut fixture = ProdActionSelectFixture::new( + Arc::clone(&stream), + batch, + n_atoms, + &q_values_host, + &b_logits_host, + &support_host, + ); + let eval_actions = fixture.launch(/*eval_mode=*/ true, /*timestep=*/ 0); + let mut mismatches = 0_usize; + for i in 0..batch { + let kernel_dir = decode_dir(eval_actions[i]); + if kernel_dir != ground_truth_dir[i] { + if mismatches < 5 { + eprintln!( + "Test 2.D mismatch sample {}: kernel={} gt={} E[Q]={:?}", + i, kernel_dir, ground_truth_dir[i], ground_truth_eq[i], + ); + } + mismatches += 1; + } + } + assert_eq!( + mismatches, 0, + "Eval mode: {} of {} samples have dir_idx ≠ ground-truth argmax E[Q]", + mismatches, batch, + ); + + // ── Train mode: P(Long+Short) > 0.40 across the batch ──────────────── + // + // Realistic random Xavier-init logits produce non-degenerate per-direction + // distributions; Thompson should sample the active directions ≥ 40% of the + // time. (Per spec Task 9: "count(d ∈ {Long, Short}) / batch_size > 0.40".) + let train_actions = fixture.launch(/*eval_mode=*/ false, /*timestep=*/ 0); + let mut train_counts = [0_u32; PROD_B0 as usize]; + for action in &train_actions { + train_counts[decode_dir(*action) as usize] += 1; + } + let active = train_counts[0] + train_counts[2]; + let p_active = active as f32 / batch as f32; + println!( + "Test 2.D train counts (S/H/L/F): {:?} active={:.4}", + train_counts, p_active, + ); + assert!( + p_active > 0.40, + "Train mode active_frac {:.4} expected > 0.40 (Thompson on realistic \ + random logits should explore Long+Short ≥ 40%). Counts: {:?}", + p_active, train_counts, + ); +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 83f29cd58..d96b67616 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2143,5 +2143,34 @@ Algorithmic divergence in mag/ord/urg branches (e.g. an inadvertent strict-argmax substitution) would shift P(best) to 1.0, which the ±5% tolerance trivially detects. +### Test 2.D — Real-batch e2e on synthetic non-degenerate logits (Task T9) + +`crates/ml/src/trainers/dqn/distributional_q_tests.rs::test_2d_real_batch_e2e` +verifies the production kernel matches a Rust ground-truth E[Q] argmax +on a synthetic batch with realistic non-degenerate per-direction C51 +distributions. The plan-prescribed approach (reuse Phase 0 Test 0.F's +converged-checkpoint loader) was deferred — wiring the safetensors +load + branching forward path is heavyweight and Test 0.F itself +already exercises that code path. Replacement: synthetic batch via +direct buffer write. + +Setup: +- 256 distinct samples; per-(sample, direction, atom) C51 logits drawn + from a deterministic hash → uniform [-1, 1] (matches post-Xavier-init + scale of fresh C51 head outputs). +- Per-sample adaptive support drawn from realistic ranges: v_min + ∈ [-1.0, -0.2], v_max ∈ [0.2, 1.0]. Each sample's (v_min, v_max, + delta_z) varies; this is the layout `update_per_sample_support` + produces in the live trainer. + +Assertions: +- Eval mode: per-sample dir_idx EXACTLY matches the argmax over a + Rust ground-truth E[Q]_d = sum_a softmax(b_logits[i, d])[a] · + atom_vals[i, d, a]. The ground-truth softmax mirrors the kernel's + numerically-stable subtract-max path. +- Train mode: count(d ∈ {Long=2, Short=0}) / batch > 0.40 — Thompson + on realistic non-degenerate distributions explores actively. Per + spec Task 9. + Plan C T11 follow-up A.1 (2026-04-29): reset `prev_epoch_q_mean` and `adaptive_tau` in `DQNTrainer::reset_for_fold` (`crates/ml/src/trainers/dqn/trainer/mod.rs`). Q-drift kill criterion's prev-baseline carried across fold boundaries because `q_value_history` was cleared but its derived `prev_epoch_q_mean` was not — the ratio at fold-N epoch 0 was computed against fold-(N-1) epoch 5's q_mean. Companion `adaptive_tau` modulation state also reset to `hyperparams.tau` baseline. Independent bug fix; applies regardless of Plan C Thompson outcome. Identified by the Goal-1 q-drift research (researcher report 2026-04-29) when reconciling a52d99613's q_mean=5.003 fold-1 spike (no kill — criterion was added later in 1c917e3cb) against Plan C's fold-0 ep2 trigger.