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

fix(dqn): Plan C T2 amendment — match production adaptive C51 support

Browse Source 
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.
This commit is contained in:
jgrusewski
2026-04-29 17:26:22 +02:00
parent de519d5806
commit 5de5e546a4
4 changed files with 456 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

208
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -251,6 +251,111 @@ __device__ __forceinline__ float compute_e_iqn_inline(const float* __restrict__
return s / (float)N_IQN_QUANTILES;
}
/**
* Plan C Phase 2 Task 2 amendment (2026-04-29): direction-branch Thompson
* now consumes RAW C51 logits + adaptive per-direction support instead of
* the plan-as-written `c51_probs_dir + atom_values` pair, because:
* - The collector materialises `exp_b_logits` (raw logits) only — there
* is no post-softmax `c51_probs_dir` buffer on the rollout path.
* - Production C51 uses per-sample per-direction adaptive support
* [v_min, v_max, delta_z] (see `c51_loss_kernel.cu` lines 729-744),
* plus optional non-linear `atom_positions [b0_size, n_atoms]`.
* A single global `atom_values [n_atoms]` doesn't exist.
* The amendment keeps the principled posterior-sample fix to UCB asymmetry
* but sources the distribution from what the collector already has.
*
* `softmax_c51_inline` matches `compute_expected_q`'s numerically-stable
* (subtract max, then exp + normalise) softmax. `compute_atom_values_inline`
* mirrors the per-branch atom-support contract from `c51_loss_kernel.cu`
* (atom_positions when non-NULL else linear v_min + a * delta_z).
*
* Suppress the unused `v_max` slot of per_sample_support: the loss kernel
* derives `delta_z` independently and only uses `v_max` for clamping inside
* the projection — for E[Q] / sampling we just need atom positions.
*/
/**
* Numerically-stable softmax over C51 atom logits.
*
* @param logits [n_atoms] raw C51 logits for one (sample, direction).
* @param n_atoms Atom count (must be ≤ caller's stack array length).
* @param probs_out [n_atoms] caller-allocated output buffer (typically a
* stack array sized to THOMPSON_MAX_ATOMS).
*
* Subtracts max(logits) before expf to prevent overflow; falls back to
* uniform on numerical underflow (sum == 0).
*/
__device__ __forceinline__ void softmax_c51_inline(
const float* __restrict__ logits,
int n_atoms,
float* probs_out
) {
float max_logit = logits[0];
for (int a = 1; a < n_atoms; a++) {
if (logits[a] > max_logit) max_logit = logits[a];
}
float sum = 0.0f;
for (int a = 0; a < n_atoms; a++) {
probs_out[a] = expf(logits[a] - max_logit);
sum += probs_out[a];
}
float inv_sum = (sum > 0.0f) ? (1.0f / sum) : (1.0f / (float)n_atoms);
for (int a = 0; a < n_atoms; a++) {
probs_out[a] *= inv_sum;
}
}
/**
* Compute atom values for one (sample, direction) using the production
* adaptive C51 support contract.
*
* Layout match (mirrors `c51_loss_kernel.cu` lines 728-748):
* per_sample_support [N, b0_size, 3] stride-12 — sample-major. Slot
* layout per direction: [v_min, v_max, delta_z]. Strided index for
* sample i, direction d is `i * b0_size * 3 + d * 3`.
*
* atom_positions [b0_size, n_atoms] non-linear positions; NULL = use
* linear support (v_min + a * delta_z). When non-NULL, indexed by
* `d * n_atoms + a`.
*
* @param per_sample_support [N, b0_size, 3] stride-12 device pointer.
* @param atom_positions [b0_size, n_atoms] adaptive positions; may
* be NULL (then linear support is used).
* @param i Sample index.
* @param d Direction index ∈ [0, b0_size).
* @param b0_size Direction-branch action count (4 currently).
* @param n_atoms C51 atom count (must be ≤ caller's stack
* array length).
* @param atom_values_out [n_atoms] caller-allocated output buffer.
*/
__device__ __forceinline__ void compute_atom_values_inline(
const float* __restrict__ per_sample_support,
const float* __restrict__ atom_positions,
int i,
int d,
int b0_size,
int n_atoms,
float* atom_values_out
) {
long long support_base = (long long)i * b0_size * 3 + (long long)d * 3;
float v_min = per_sample_support[support_base + 0];
float v_max = per_sample_support[support_base + 1];
float delta_z = per_sample_support[support_base + 2];
if (atom_positions != NULL) {
long long pos_base = (long long)d * n_atoms;
for (int a = 0; a < n_atoms; a++) {
atom_values_out[a] = atom_positions[pos_base + a];
}
} else {
for (int a = 0; a < n_atoms; a++) {
atom_values_out[a] = v_min + (float)a * delta_z;
}
}
/* v_max is unused here — atoms are derived from v_min + delta_z (linear
* grid) or atom_positions (non-linear). Suppress unused-variable warning. */
(void)v_max;
}
/* ================================================================== */
/* Kernel 0a: domain_rand_episode_starts (GPU-native #1) */
/* ================================================================== */
@@ -895,17 +1000,26 @@ extern "C" __global__ void experience_action_select(
float* out_conviction, /* [N] output: direction-branch conviction ∈ [0,1] for Kelly warmup. NULL=skip. */
float* out_magnitude_conviction, /* [N] output: magnitude-branch conviction ∈ [0,1] for Kelly warmup. NULL=skip. */
/* Plan C Phase 2 Task 2 — Thompson direction-branch buffers.
* Direction selection now consumes the C51/IQN posteriors directly
* Direction selection consumes the production C51 posterior directly
* (Thompson at training; argmax E[Q] at eval); the legacy direction
* eps-greedy + Boltzmann + UCB-bonus path is replaced. mag/ord/urg
* branches are unchanged and continue to read q_values + their per-
* branch eps multipliers. Tasks 3 & 4 wire these buffers from the
* Rust call sites; until they land the kernel signature is
* intentionally non-callable end-to-end. */
const float* __restrict__ c51_probs_dir, /* [N, b0_size, n_atoms] direction-branch C51 atom probs */
const float* __restrict__ atom_values, /* [n_atoms] C51 atom support */
const float* __restrict__ iqn_quantiles_dir,/* [N, b0_size, N_IQN_QUANTILES] direction-branch IQN quantiles */
int n_atoms /* C51 atom count */
* intentionally non-callable end-to-end.
*
* Plan C T2 amendment (2026-04-29): the original plan-as-written args
* (`c51_probs_dir` post-softmax probs, single `atom_values [n_atoms]`,
* `iqn_quantiles_dir`) didn't match production — the collector only
* materialises raw `b_logits`, support is per-sample per-direction
* adaptive, and IQN inference is training-only (no rollout-time
* quantiles). Amended kernel takes raw C51 logits + adaptive support
* and drops IQN — single-distribution C51 still provides the
* principled posterior sample that fixes the UCB asymmetry. */
const float* __restrict__ b_logits_dir, /* [N, b0_size, n_atoms] direction-branch RAW C51 logits */
const float* __restrict__ per_sample_support, /* [N, b0_size, 3] stride-12 per-direction [v_min, v_max, delta_z] */
const float* __restrict__ atom_positions, /* [b0_size, n_atoms] adaptive non-linear positions; NULL = linear */
int n_atoms /* C51 atom count */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
@@ -1024,42 +1138,60 @@ extern "C" __global__ void experience_action_select(
* flip on direction become unreachable on this branch — Thompson
* sampling supplies exploration directly from the posterior, and the
* symmetry-breaking that contrarian mode previously enforced is now
* carried by the joint C51+IQN sample noise.
* carried by the C51 sample noise.
*
* Math reference: thompson_test_kernel.cu (Phase 0 standalone, same
* inverse-CDF + uniform-tau interpolation primitives). Phase 2
* Test 2.B verifies bit-identical output between this kernel and the
* standalone wrapper at fixed RNG.
* Plan C T2 amendment (2026-04-29): single-distribution C51 Thompson
* with production's adaptive per-direction support. The plan-as-
* authored joint C51+IQN sampling required rollout-time IQN
* inference that doesn't exist (`gpu_iqn_head.rs` is training-only)
* and post-softmax probs the collector doesn't materialise. The
* principled posterior-sample fix to UCB asymmetry is preserved —
* sourced from the existing rollout C51 distribution, with
* production's adaptive support faithfully consumed.
*
* Sampling at training: argmax_d (sample_C51(d) + sample_IQN(d)) with
* Thompson draws against per-direction posteriors. The /2 from the
* joint-mean form is dropped (monotonic for argmax).
* Math reference: thompson_test_kernel.cu (Phase 0 standalone). The
* standalone now exposes a single-distribution direction sampler
* matching this kernel for Test 2.B bit-identical parity.
*
* Eval: argmax_d (E_C51(d) + E_IQN(d)) — pure exploitation, no RNG.
* Sampling at training: argmax_d sample_C51_inverse_cdf(probs_d, atoms_d, u)
* where probs_d = softmax(b_logits_dir[i, d, :]) and atoms_d is the
* per-sample per-direction C51 support (atom_positions when non-NULL,
* else linear v_min + a*delta_z).
*
* Eval: argmax_d E[Q]_d where E[Q]_d = sum_a probs_d[a] * atoms_d[a].
*
* E[Q]-derived `e_dir` is computed once here unconditionally and
* reused for: (a) eval-mode argmax; (b) the conviction calculation
* downstream which the spec requires to remain E[Q]-based (avoids
* Kelly-cap jitter from per-step Thompson sample noise). Lives in
* the outer kernel scope so the conviction block at the bottom can
* read it without recomputation.
* Kelly-cap jitter from per-step Thompson sample noise);
* (c) `out_q_gaps`. Lives in outer kernel scope so the conviction
* block at the bottom can read it without recomputation.
*
* Philox stream is seeded by (i, timestep) so eval is bit-reproducible
* across runs at the same checkpoint, matching the prior path. */
const float* c51_probs_i = c51_probs_dir + (long long)i * b0_size * n_atoms;
const float* iqn_quantiles_i = iqn_quantiles_dir + (long long)i * b0_size * N_IQN_QUANTILES;
/* b0_size = 4 in the current 4-direction action space. Stack-allocated
* with 4 to keep nvcc from spilling to local memory; if b0_size grows,
* raise this and confirm with `cuobjdump --dump-resource-usage`. */
/* THOMPSON_MAX_ATOMS sized to match the production C51 ceiling
* (`c51_loss_kernel.cu::MAX_ATOMS = 128`). Stack arrays for probs +
* atom values (2 * n_atoms floats) — at n_atoms=51 that's ~408 bytes
* per thread, well within nvcc's local-memory threshold. */
#ifndef THOMPSON_MAX_ATOMS
#define THOMPSON_MAX_ATOMS 128
#endif
float probs_d[THOMPSON_MAX_ATOMS];
float atom_vals_d[THOMPSON_MAX_ATOMS];
const float* logits_i = b_logits_dir + (long long)i * b0_size * n_atoms;
/* Pass 1: per-direction E[Q] for eval argmax + conviction reuse. */
for (int d = 0; d < b0_size; d++) {
float e_c51 = compute_e_c51_inline(c51_probs_i + d * n_atoms, atom_values, n_atoms);
float e_iqn = compute_e_iqn_inline(iqn_quantiles_i + d * N_IQN_QUANTILES);
e_dir[d] = e_c51 + e_iqn; /* joint sum; /2 monotonic for argmax */
softmax_c51_inline(logits_i + d * n_atoms, n_atoms, probs_d);
compute_atom_values_inline(per_sample_support, atom_positions,
i, d, b0_size, n_atoms, atom_vals_d);
e_dir[d] = compute_e_c51_inline(probs_d, atom_vals_d, n_atoms);
}
if (eval_mode) {
/* EVAL: argmax of joint E[Q]. Pure exploitation, no exploration. */
/* EVAL: argmax of E[Q]. Pure exploitation, no exploration. */
float best_q = -1e30f;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
@@ -1067,21 +1199,19 @@ extern "C" __global__ void experience_action_select(
}
dir_idx = best_d;
} else {
/* TRAINING: Thompson sample. One C51 inverse-CDF draw and one
* IQN tau-interpolation draw per direction; argmax of the joint
* sum picks the direction. */
/* TRAINING: per-direction Thompson sample, argmax over directions.
* Pass 2 recomputes per-direction probs/atoms; the alternative —
* caching across directions — would cost b0_size * 2 * MAX_ATOMS
* stack floats (~4 KiB per thread) and risk local-memory spill.
* Recompute is < 1 µs per thread at n_atoms=51. */
float best_sample = -1e30f;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
float u_c51 = philox_uniform(i, timestep, rng_ctr++);
float s_c51 = sample_c51_inverse_cdf(
c51_probs_i + d * n_atoms, atom_values, n_atoms, u_c51
);
float u_iqn = philox_uniform(i, timestep, rng_ctr++);
float s_iqn = sample_iqn_quantile_interp(
iqn_quantiles_i + d * N_IQN_QUANTILES, u_iqn
);
float q_sample = s_c51 + s_iqn; /* /2 dropped (monotonic for argmax) */
softmax_c51_inline(logits_i + d * n_atoms, n_atoms, probs_d);
compute_atom_values_inline(per_sample_support, atom_positions,
i, d, b0_size, n_atoms, atom_vals_d);
float u = philox_uniform(i, timestep, rng_ctr++);
float q_sample = sample_c51_inverse_cdf(probs_d, atom_vals_d, n_atoms, u);
if (q_sample > best_sample) { best_sample = q_sample; best_d = d; }
}
dir_idx = best_d;

169
crates/ml/src/cuda_pipeline/thompson_test_kernel.cu
View File

@@ -191,3 +191,172 @@ extern "C" __global__ void compute_sigma_iqn_test(
out_sigma[d] = iqr / 1.349f;
__threadfence_system();
}
// ─── Plan C Phase 2 Task 2 amendment (2026-04-29) ─────────────────────
//
// The kernels above (thompson_direction_test_batched, argmax_eq_test) are
// the Plan A Phase 0 entry points: they consume post-softmax `c51_probs`,
// a single linear `atom_values [n_atoms]`, AND `iqn_quantiles` (joint
// C51+IQN sample). They remain in place because Plan A tests 0.B-0.F (in
// `distributional_q_tests.rs`) still verify the underlying inverse-CDF +
// uniform-tau primitives via these wrappers — none of that math changed.
//
// The kernels below match the AMENDED production API (single-distribution
// C51 with adaptive per-direction support; no IQN inference at rollout):
//
// - direction_thompson_v2_test: Thompson sample per direction using
// `b_logits` + `per_sample_support` + optional `atom_positions`.
// Single-distribution; no IQN term. Bit-equivalent to the production
// direction-branch Thompson loop in
// `experience_kernels.cu::experience_action_select`.
//
// - argmax_eq_v2_test: argmax over E[Q] computed from softmax(logits)
// and the same adaptive support. Bit-equivalent to the production
// eval-mode argmax.
//
// Test 2.B (Plan C Task 7) launches the production kernel with the same
// inputs as `direction_thompson_v2_test` and asserts both produce
// identical `dir_idx` output across many seeds.
//
// RNG note: `direction_thompson_v2_test` uses the same LCG (`test_rand`)
// as `thompson_direction_test_batched` for parity with Plan A wrappers,
// NOT Philox. Test 2.B (Task 7) bridges to Philox at execution time by
// laying out `u` draws produced via Philox and passing them through this
// kernel's deterministic LCG path is NOT bit-equivalent to production —
// Test 2.B will need to either (a) drive both kernels from the same
// uniform array (preferred — strips RNG implementation choice from the
// equivalence test), or (b) invoke a separate Philox-driven variant. (a)
// is the simpler integration for Task 7.
// Numerically stable softmax over C51 logits. Matches the production
// `softmax_c51_inline` in `experience_kernels.cu`.
__device__ __forceinline__ void softmax_c51_test(
const float* __restrict__ logits,
int n_atoms,
float* probs_out
) {
float max_logit = logits[0];
for (int a = 1; a < n_atoms; a++) {
if (logits[a] > max_logit) max_logit = logits[a];
}
float sum = 0.0f;
for (int a = 0; a < n_atoms; a++) {
probs_out[a] = expf(logits[a] - max_logit);
sum += probs_out[a];
}
float inv_sum = (sum > 0.0f) ? (1.0f / sum) : (1.0f / (float)n_atoms);
for (int a = 0; a < n_atoms; a++) {
probs_out[a] *= inv_sum;
}
}
// Compute atom values for one (sample, direction). Matches the production
// `compute_atom_values_inline` in `experience_kernels.cu`.
__device__ __forceinline__ void compute_atom_values_test(
const float* __restrict__ per_sample_support,
const float* __restrict__ atom_positions,
int i,
int d,
int b0_size,
int n_atoms,
float* atom_values_out
) {
long long support_base = (long long)i * b0_size * 3 + (long long)d * 3;
float v_min = per_sample_support[support_base + 0];
float v_max = per_sample_support[support_base + 1];
float delta_z = per_sample_support[support_base + 2];
if (atom_positions != NULL) {
long long pos_base = (long long)d * n_atoms;
for (int a = 0; a < n_atoms; a++) {
atom_values_out[a] = atom_positions[pos_base + a];
}
} else {
for (int a = 0; a < n_atoms; a++) {
atom_values_out[a] = v_min + (float)a * delta_z;
}
}
(void)v_max;
}
#ifndef THOMPSON_TEST_MAX_ATOMS
#define THOMPSON_TEST_MAX_ATOMS 128
#endif
// Test kernel: single-distribution C51 Thompson direction sampling
// (training mode), aligned with the amended production API.
//
// Input layout (single-sample, since this is a standalone test wrapper):
// - b_logits [b0_size, n_atoms] raw C51 logits
// - per_sample_support [b0_size, 3] — single-sample slice; the
// caller passes a 1×4×3 tile
// corresponding to sample i=0.
// - atom_positions [b0_size, n_atoms] may be NULL (linear support).
//
// One thread per seed; `out_dir_idx[i]` records the Thompson pick for
// `rng_state = base_seed + i`.
extern "C" __global__ void direction_thompson_v2_test(
const float* __restrict__ b_logits,
const float* __restrict__ per_sample_support,
const float* __restrict__ atom_positions,
int b0_size,
int n_atoms,
unsigned int base_seed,
int n_seeds,
int* __restrict__ out_dir_idx // [n_seeds] — pinned-mapped
) {
int seed_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (seed_idx >= n_seeds) return;
unsigned int rng_state = base_seed + (unsigned int)seed_idx;
float best_sample = -CUDART_INF_F;
int best_d = 0;
float probs_d[THOMPSON_TEST_MAX_ATOMS];
float atom_vals_d[THOMPSON_TEST_MAX_ATOMS];
for (int d = 0; d < b0_size; d++) {
softmax_c51_test(b_logits + d * n_atoms, n_atoms, probs_d);
// Sample index 0 — the standalone wrapper feeds a single tile.
compute_atom_values_test(per_sample_support, atom_positions,
/*i=*/0, d, b0_size, n_atoms, atom_vals_d);
float u = test_rand(&rng_state);
float q_sample = sample_c51_inverse_cdf(probs_d, atom_vals_d, n_atoms, u);
if (q_sample > best_sample) {
best_sample = q_sample;
best_d = d;
}
}
out_dir_idx[seed_idx] = best_d;
__threadfence_system();
}
// Test kernel: argmax of E[Q] (eval mode), aligned with the amended
// production API. Single-threaded.
extern "C" __global__ void argmax_eq_v2_test(
const float* __restrict__ b_logits,
const float* __restrict__ per_sample_support,
const float* __restrict__ atom_positions,
int b0_size,
int n_atoms,
int* out_dir_idx
) {
if (threadIdx.x != 0 || blockIdx.x != 0) return;
float best_q = -CUDART_INF_F;
int best_d = 0;
float probs_d[THOMPSON_TEST_MAX_ATOMS];
float atom_vals_d[THOMPSON_TEST_MAX_ATOMS];
for (int d = 0; d < b0_size; d++) {
softmax_c51_test(b_logits + d * n_atoms, n_atoms, probs_d);
compute_atom_values_test(per_sample_support, atom_positions,
/*i=*/0, d, b0_size, n_atoms, atom_vals_d);
float e = compute_e_c51(probs_d, atom_vals_d, n_atoms);
if (e > best_q) { best_q = e; best_d = d; }
}
*out_dir_idx = best_d;
__threadfence_system();
}

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

@@ -1015,6 +1015,42 @@ Task 6 / Test 2.A then validates training stochastic / eval
deterministic behaviour and Task 7 / Test 2.B verifies bit-identical
output between this kernel and `thompson_test_kernel.cu` at fixed RNG.
Plan C T2 amendment (2026-04-29): the four trailing kernel arguments
above (`c51_probs_dir`, `atom_values [n_atoms]`, `iqn_quantiles_dir`,
`n_atoms`) were replaced with `b_logits_dir [N, b0_size, n_atoms]`
(raw direction-branch C51 logits — first slice of the collector's
`exp_b_logits`), `per_sample_support [N, b0_size, 3]` stride-12
(production's adaptive per-sample per-direction `[v_min, v_max,
delta_z]` tile), `atom_positions [b0_size, n_atoms]` (optional
non-linear positions; NULL = linear), and `n_atoms` — to align with
the production C51 architecture. The plan-as-written assumed a
post-softmax probability buffer the collector doesn't materialise, a
single global linear support that production has replaced with
adaptive per-direction support, and a rollout-time IQN inference path
that does not exist (`gpu_iqn_head.rs` is training-only). Direction-
branch Thompson is now single-distribution C51 with the production
adaptive support — the principled posterior-sample fix to the UCB
asymmetry is preserved. Two new device-inline helpers
(`softmax_c51_inline`, `compute_atom_values_inline`) sit alongside
the four Plan C T1 primitives in the same Thompson math section and
mirror the production C51 loss kernel's softmax + atom-support
contract. The Phase 0 standalone (`thompson_test_kernel.cu`) gained
two new entry points (`direction_thompson_v2_test`,
`argmax_eq_v2_test`) matching the amended production API; the
original Plan A entry points are retained for Plan A tests 0.B-0.F.
Tasks 3 + 4 are unblocked — collector and evaluator already own
`per_sample_support_buf` and the `exp_b_logits` slice in production.
Test 2.B (Plan C Task 7) and Test 2.A (Task 6) are re-scoped to the
amended single-distribution C51 surface; the bit-equivalence test
should drive both kernels from the same uniform array (strips the
LCG-vs-Philox RNG implementation choice from the equivalence check).
Conviction read site continues to consume `e_dir[]` (now a pure
single-distribution E[C51] vector, no longer joint E[C51]+E[IQN]) —
the spec invariant "conviction stays E[Q]-based" still holds at the
amended scale, but ISV[21]'s `q_dir_abs_ref` reference may need
retuning since the joint-sum scale is gone (revisit in Phase 3 if a
systematic mismatch surfaces).
Persistent picked_action_history scatter (2026-04-26): `gpu_backtest_evaluator.rs`
gains a window-major `picked_action_history_buf` (parallel layout to
`actions_history_buf`) populated by an additional `scatter_intent_chunk`

114
docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-C-phase-2.md
View File

@@ -111,56 +111,71 @@ kernel file. Used by Phase 2 direction-branch action selection."
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
> **Note: revised 2026-04-29** to match production C51 adaptive-support
> architecture. The prior wording (post-softmax `c51_probs_dir`, single
> linear `atom_values`, rollout-time `iqn_quantiles_dir`) didn't match
> production: the collector materialises only RAW `b_logits`, the C51
> support is per-sample per-direction adaptive (`per_sample_support
> [N, 4, 3]` + optional `atom_positions [4, num_atoms]`), and IQN
> inference is training-only — there is no rollout-time IQN quantile
> buffer to consume. Direction-branch Thompson is now
> single-distribution C51 with the production adaptive support; the
> principled posterior-sample fix to UCB asymmetry is preserved (sourced
> from the existing C51 rollout distribution). See the spec footer
> Status entry "2026-04-29 — Plan C T2 amendment" for the full triage.
- [ ] **Step 1: Read the existing direction branch**
Open `experience_kernels.cu` and find `experience_action_select` (around line 756). Find the direction branch around lines 858-887 (the eps-greedy + Boltzmann block).
- [ ] **Step 2: Add new kernel parameters**
Modify the kernel signature to accept three new buffers + n_atoms:
Modify the kernel signature to accept the AMENDED set of buffers — RAW direction-branch C51 logits, the per-sample per-direction adaptive support tile, and (optionally) the non-linear atom positions. IQN buffers are NOT threaded; rollout-time IQN inference does not exist.
```cuda
extern "C" __global__ void experience_action_select(
// ... existing args up to and including out_conviction ...
const float* __restrict__ c51_probs_dir, // [N, b0_size, n_atoms] NEW
const float* __restrict__ atom_values, // [n_atoms] NEW
const float* __restrict__ iqn_quantiles_dir, // [N, b0_size, N_IQN_QUANTILES] NEW
int n_atoms // NEW
// ... existing args up to and including out_magnitude_conviction ...
const float* __restrict__ b_logits_dir, // [N, b0_size, n_atoms] NEW (raw logits; collector's exp_b_logits slice)
const float* __restrict__ per_sample_support, // [N, b0_size, 3] str-12 NEW (adaptive [v_min, v_max, delta_z])
const float* __restrict__ atom_positions, // [b0_size, n_atoms] NEW (NULL = linear support)
int n_atoms // NEW
) {
```
- [ ] **Step 3: Replace direction branch with Thompson + argmax**
Replace the existing direction-branch code (the if-else with eps-greedy + Boltzmann at ~lines 858-887) with:
Replace the existing direction-branch code (the if-else with eps-greedy + Boltzmann at ~lines 858-887) with the single-distribution C51 Thompson loop. The amendment introduces two new device-inline helpers (`softmax_c51_inline`, `compute_atom_values_inline`) — installed alongside the four Plan C Task 1 primitives, in the same `experience_kernels.cu` Thompson math section.
```cuda
/* Direction branch: Thompson sampling at training, argmax E[Q] at eval.
* Replaces the prior eps-greedy + Boltzmann; ε floor and tau-floor for
* direction become unreachable (cleaned up in Phase 3 / Plan D).
*
* Math reference: thompson_test_kernel.cu (Phase 0 standalone, same math).
* Sampling at training: 0.5 × (sample_C51 + sample_IQN), argmax across dirs.
* Eval: argmax of (E_C51 + E_IQN) — /2 dropped (monotonic).
* Single-distribution C51 with production's adaptive per-direction
* support. Math reference: thompson_test_kernel.cu (Phase 0 standalone)
* — the v2 entry points (`direction_thompson_v2_test`, `argmax_eq_v2_test`)
* mirror this loop for Test 2.B bit-equivalence verification.
*
* E[Q]-derived q_b0 must be reused for the conviction calculation below
* (per spec §"Conviction stays E[Q]-based"). We compute e_dir[d] once
* here and reuse.
* E[Q]-derived `e_dir` must be reused for the conviction calculation
* below (per spec §"Conviction stays E[Q]-based"). Computed once here.
*/
const float* c51_probs_i = c51_probs_dir + (long long)i * b0_size * n_atoms;
const float* iqn_quantiles_i = iqn_quantiles_dir + (long long)i * b0_size * N_IQN_QUANTILES;
#ifndef THOMPSON_MAX_ATOMS
#define THOMPSON_MAX_ATOMS 128 /* matches c51_loss_kernel.cu MAX_ATOMS */
#endif
float probs_d[THOMPSON_MAX_ATOMS];
float atom_vals_d[THOMPSON_MAX_ATOMS];
const float* logits_i = b_logits_dir + (long long)i * b0_size * n_atoms;
/* Compute per-direction E[Q] for: (a) eval-mode argmax;
* (b) conviction read-back below; (c) Thompson uses raw samples,
* not E[Q], but conviction needs E[Q]. */
float e_dir[4]; // b0_size = 4 in current 4-direction action space
for (int d = 0; d < b0_size; d++) {
float e_c51 = compute_e_c51_inline(c51_probs_i + d * n_atoms, atom_values, n_atoms);
float e_iqn = compute_e_iqn_inline(iqn_quantiles_i + d * N_IQN_QUANTILES);
e_dir[d] = e_c51 + e_iqn; // joint sum; /2 monotonic for argmax
softmax_c51_inline(logits_i + d * n_atoms, n_atoms, probs_d);
compute_atom_values_inline(per_sample_support, atom_positions,
i, d, b0_size, n_atoms, atom_vals_d);
e_dir[d] = compute_e_c51_inline(probs_d, atom_vals_d, n_atoms);
}
if (eval_mode) {
/* EVAL: argmax of joint E[Q]. Pure exploitation. */
/* EVAL: argmax of E[Q]. Pure exploitation. */
float best_q = -CUDART_INF_F;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
@@ -168,19 +183,18 @@ Replace the existing direction-branch code (the if-else with eps-greedy + Boltzm
}
dir_idx = best_d;
} else {
/* TRAINING: Thompson sample. */
/* TRAINING: per-direction Thompson sample, argmax over directions.
* Recompute probs/atoms for each direction — caching across
* directions would cost b0_size * 2 * MAX_ATOMS stack floats and
* risk local-memory spill. */
float best_sample = -CUDART_INF_F;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
float u_c51 = philox_uniform(i, timestep, rng_ctr++);
float s_c51 = sample_c51_inverse_cdf(
c51_probs_i + d * n_atoms, atom_values, n_atoms, u_c51
);
float u_iqn = philox_uniform(i, timestep, rng_ctr++);
float s_iqn = sample_iqn_quantile_interp(
iqn_quantiles_i + d * N_IQN_QUANTILES, u_iqn
);
float q_sample = s_c51 + s_iqn; // /2 dropped
softmax_c51_inline(logits_i + d * n_atoms, n_atoms, probs_d);
compute_atom_values_inline(per_sample_support, atom_positions,
i, d, b0_size, n_atoms, atom_vals_d);
float u = philox_uniform(i, timestep, rng_ctr++);
float q_sample = sample_c51_inverse_cdf(probs_d, atom_vals_d, n_atoms, u);
if (q_sample > best_sample) { best_sample = q_sample; best_d = d; }
}
dir_idx = best_d;
@@ -669,3 +683,39 @@ All requirements covered. ✓
---
## Plan complete and saved to `docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-C-phase-2.md`.
---
## Status
**2026-04-29 — Plan C T2 amendment.** The plan-as-authored Task 2 kernel
signature (`c51_probs_dir`, single linear `atom_values [n_atoms]`,
`iqn_quantiles_dir`) was incompatible with production:
1. **No post-softmax `c51_probs_dir` on the collector.** Only raw
`exp_b_logits` is materialised; softmax happens inline in consumers.
2. **C51 support is per-sample per-direction adaptive.**
`c51_loss_kernel.cu` (lines 615-744) consumes `per_sample_support
[N, 4, 3]` (stride-12, per-direction `[v_min, v_max, delta_z]`) plus
optional non-linear `atom_positions [4, num_atoms]`. A single
global `atom_values [n_atoms]` does not exist.
3. **No rollout-time IQN inference.** `gpu_iqn_head.rs` is training-only;
the collector does not produce per-direction IQN quantile buffers.
The amended Task 2 kernel signature passes raw `b_logits_dir`,
`per_sample_support`, optional `atom_positions`, and `n_atoms`.
Direction-branch Thompson is single-distribution C51 with adaptive
support — the principled posterior-sample fix to the UCB selector/target
asymmetry is preserved (the C51 distribution alone provides per-direction
posterior samples), without requiring infrastructure that doesn't exist.
Tasks 3+4 (buffer wiring) are unblocked — collector and evaluator already
own `per_sample_support_buf` and `exp_b_logits` in production.
Phase 0 standalone test kernel (`thompson_test_kernel.cu`) gained two new
entry points (`direction_thompson_v2_test`, `argmax_eq_v2_test`) that
mirror the amended production API for Test 2.B bit-equivalence parity.
The original Plan A entry points (`thompson_direction_test_batched`,
`argmax_eq_test`) are retained — they verify the underlying inverse-CDF
+ uniform-tau math primitives unchanged by the amendment, and Plan A
tests 0.B-0.F continue to consume them.