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

test(dqn): Phase 2 Test 2.D — real-batch e2e on synthetic non-degenerate logits

Browse Source 
Plan C Phase 2 T9. Verifies the production kernel's eval-mode argmax
exactly matches a Rust ground-truth E[Q] argmax on a 256-sample batch
with realistic per-direction non-degenerate C51 distributions, and
confirms Thompson explores Long+Short ≥ 40% in training mode.

The plan-prescribed approach (reuse Phase 0 Test 0.F's converged
checkpoint loader) was deferred — Test 0.F itself already exercises
the safetensors load + branching forward path. Replacement strategy
from the dispatch brief: synthetic batch via direct buffer write.

Setup:
- 256 samples, 21 atoms, per-(sample, direction, atom) C51 logits
  drawn from a deterministic hash → uniform [-1, 1] (post-Xavier-init
  scale of fresh C51 head outputs)
- Per-sample adaptive support [v_min ∈ [-1.0, -0.2], v_max ∈ [0.2, 1.0]]
  matching the layout produced by `update_per_sample_support`
- Uniform q_values (mag/ord/urg fall through Boltzmann uniformly)

Rust ground-truth: softmax(b_logits[i, d]) · atom_vals[i, d] computed
with the numerically-stable subtract-max softmax matching the kernel's
softmax_c51_inline.

Assertions:
- Eval mode: per-sample dir_idx EXACTLY matches ground-truth argmax E[Q]
- Train mode: count(d ∈ {Long, Short}) / batch > 0.40

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-29 18:21:45 +02:00
parent cc55c8a25c
commit dfbadf2f3f
2 changed files with 200 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

171
crates/ml/src/trainers/dqn/distributional_q_tests.rs
View File

@@ -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,
);
}

29
docs/dqn-wire-up-audit.md
View File

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