Plan C as authored assumed:
- c51_probs_dir = post-softmax probabilities (didn't exist on collector;
only raw exp_b_logits is materialised)
- atom_values = single global linear support (production uses per-sample
per-direction adaptive [v_min, v_max, delta_z] + optional atom_positions)
- iqn_quantiles_dir = rollout-side IQN inference (doesn't exist; IQN is
training-only)
Amended kernel signature uses the production architecture:
- b_logits_dir [N, b0_size, n_atoms] — raw direction-branch logits
- per_sample_support [N, b0_size, 3] — adaptive per-direction support
- atom_positions [b0_size, n_atoms] — non-linear positions (NULL = linear)
- n_atoms
Direction-branch Thompson is now single-distribution over C51 (with adaptive
support), not joint C51+IQN. Single-distribution still provides the
principled posterior sample that fixes the UCB selector/target asymmetry —
the goal of Plan C is preserved, just sourced from the existing rollout-time
distribution instead of a non-existent IQN inference path.
New device-inline helpers: softmax_c51_inline, compute_atom_values_inline.
Plan + audit doc updated. Phase 0 standalone test kernel gained two new
entry points (direction_thompson_v2_test, argmax_eq_v2_test) matching the
amended production API; original Plan A entry points retained for tests
0.B-0.F. Tasks 3+4 (buffer wiring) unblocked — collector/evaluator already
have per_sample_support_buf and exp_b_logits in production.
5de5e546a4