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>
dfbadf2f3f