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

feat(dqn): direction-branch Thompson sampling in experience_action_select

Browse Source 
Replaces eps-greedy + Boltzmann on direction with:
  - Training: argmax of (sample_C51 + sample_IQN) per direction (Thompson)
  - Eval: argmax of (E_C51 + E_IQN) per direction (no exploration)

Conviction calculation now reads e_dir[] (joint E[Q] computed inline)
preserving the spec's requirement that conviction uses raw E[Q] not
samples (avoids Kelly cap jitter).

Magnitude/order/urgency branches: unchanged.
This commit is contained in:
jgrusewski
2026-04-29 16:22:43 +02:00
parent 2cca2edbc6
commit 52a2663a2e
2 changed files with 148 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -893,7 +893,19 @@ extern "C" __global__ void experience_action_select(
* what the magnitude head WOULD pick if direction gating and
* post-policy physics didn't mask it. */
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. */
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
* (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 */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
@@ -989,85 +1001,90 @@ extern "C" __global__ void experience_action_select(
int dir_idx, mag_idx, a2, a3;
/* Per-direction joint E[Q] (E_C51 + E_IQN), populated by the
* Thompson direction branch below and reused for the conviction
* block at the bottom of this kernel. Allocated at outer scope so
* the conviction read site (Plan C Phase 2 Task 2 Sub-change C)
* can consume it without recomputation. b0_size = 4 in the current
* 4-direction action space (Short/Hold/Long/Flat). */
float e_dir[4];
/* Hold enforcement removed — replaced by cost-driven hold timing.
* The model can always take any action; costs make bad actions expensive.
* min_hold_bars and portfolio_states params kept for ABI compatibility. */
(void)min_hold_bars;
(void)portfolio_states;
/* Branch 0: direction — Boltzmann softmax with ISV-adaptive tau floor.
/* Branch 0: direction — Thompson sampling at training, argmax E[Q] at
* eval. Plan C Phase 2 Task 2 replacement for the prior eps-greedy +
* ISV-adaptive Boltzmann path (which itself replaced a UCB-bonus
* direction selector that destabilised training; see
* project_ff00af68a_regression_root_cause and the Plan C spec). The
* eps_dir floor / passive-pressure adaptive boost / contrarian sign
* 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.
*
* tau = max(q_range, q_dir_abs_ref) where q_dir_abs_ref is ISV[21], the
* EMA of mean |Q| across direction bins. This makes the tau floor scale
* with the network's actual Q magnitude rather than relying on a tuned
* 0.01 constant.
* 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.
*
* Why the old `max(q_range, 0.01)` was dangerous: C51's expected Q
* structurally biases Flat above directional actions when the policy has
* no measurable edge — Flat returns are exactly 0 (no position change),
* collapsing C51's distribution to a delta at 0; directional returns
* concentrate slightly below 0 from tx_cost without compensating edge,
* giving E[Q_directional] < E[Q_flat] = 0. With Q-spread on the order
* of 0.01-0.1 (typical at fresh init or post-collapse), the static 0.01
* floor produced exp(Q/tau) ratios approaching deterministic argmax,
* locking the policy into Hold/Flat within ~1 epoch and preventing the
* exploration needed to discover real edge. Validated empirically: the
* Kelly val-Flat-collapse fix (commit 0c9d1ee39) exposed val_dir_dist
* shifting from S=.23/H=.17/L=.42/F=.18 (epoch 1) to S=.14/H=.35/L=.15/
* F=.37 (epoch 2) — a 35→72% Hold+Flat shift in one epoch.
* 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).
*
* With the ISV-driven floor: when Q-magnitudes grow during training,
* the floor scales with them, preserving the SAME relative spread for
* sampling. The policy never collapses to deterministic argmax until
* q_range exceeds the network's typical Q magnitude — i.e., until the
* spread represents real, scale-significant edge over the baseline.
* Both conviction (computed below) and tau use the SAME ISV reference,
* keeping the two adaptive mechanisms coherent.
* Eval: argmax_d (E_C51(d) + E_IQN(d)) — pure exploitation, no RNG.
*
* Cold-start fallback: when ISV[21] is < 1e-6 (pre-first-update), use
* the legacy 0.01 floor so the kernel doesn't divide by zero. Once the
* ISV producer fires (every epoch), the adaptive floor takes over.
* 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.
*
* Philox stream is seeded by (i, timestep) so eval is bit-reproducible
* across runs at the same checkpoint. */
if (!eval_mode && philox_uniform(i, timestep, rng_ctr++) < eps_dir) {
int r = (int)(philox_uniform(i, timestep, rng_ctr++) * (float)b0_size);
dir_idx = (r >= b0_size) ? b0_size - 1 : r;
* 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`. */
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 */
}
if (eval_mode) {
/* EVAL: argmax of joint E[Q]. Pure exploitation, no exploration. */
float best_q = -1e30f;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
if (e_dir[d] > best_q) { best_q = e_dir[d]; best_d = d; }
}
dir_idx = best_d;
} else {
/* Boltzmann softmax over direction Q-values (train + eval). */
float q_max_d = q_sign * (q_b0[0]);
float q_min_d = q_max_d;
for (int a = 1; a < b0_size; a++) {
float qv = q_sign * (q_b0[a]);
q_max_d = fmaxf(q_max_d, qv);
q_min_d = fminf(q_min_d, qv);
}
{
/* ISV-adaptive tau floor: same q_dir_abs_ref signal that drives
* the conviction-based Kelly warmup floor (kelly_position_cap).
* Scales the exploration temperature with the network's actual
* Q magnitude so softmax never becomes deterministic for tiny
* absolute Q-spreads — preventing the C51 Hold/Flat attractor. */
float q_dir_abs_ref = (isv_signals_ptr != NULL)
? isv_signals_ptr[21]
: 0.0f;
float tau_floor = fmaxf(q_dir_abs_ref, 0.01f);
float tau_d = fmaxf(q_max_d - q_min_d, tau_floor);
float sum_e = 0.0f;
float exps_d[4]; /* b0_size=4: Short/Hold/Long/Flat */
for (int a = 0; a < b0_size; a++) {
float qv = q_sign * (q_b0[a]);
exps_d[a] = expf((qv - q_max_d) / tau_d);
sum_e += exps_d[a];
}
float r = philox_uniform(i, timestep, rng_ctr++) * sum_e;
float cum = 0.0f;
dir_idx = b0_size - 1;
for (int a = 0; a < b0_size; a++) {
cum += exps_d[a];
if (r < cum) { dir_idx = a; break; }
}
/* TRAINING: Thompson sample. One C51 inverse-CDF draw and one
* IQN tau-interpolation draw per direction; argmax of the joint
* sum picks the direction. */
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) */
if (q_sample > best_sample) { best_sample = q_sample; best_d = d; }
}
dir_idx = best_d;
}
/* Hold/Flat detection shortcut — when direction=Hold(1) or Flat(3),
@@ -1273,10 +1290,30 @@ extern "C" __global__ void experience_action_select(
* confidence in the decision, regardless of which side of the distribution
* argmax lands on. */
if (out_conviction != NULL) {
float q_max_raw = q_b0[0];
float q_min_raw = q_b0[0];
/* Plan C Phase 2 Task 2 — conviction reads `e_dir` (joint
* E[C51] + E[IQN] computed in the Thompson direction branch
* above) instead of the legacy `q_b0` cuBLAS Q-head output.
*
* Why: the spec requires conviction to remain E[Q]-based — using
* Thompson per-step samples here would feed the Kelly cap a
* jittery confidence signal that re-shapes target position with
* every draw, defeating the warmup floor's purpose. The joint
* E[Q] across the C51 and IQN posteriors is the same quantity
* eval-mode argmax uses for direction selection, so dir-pick and
* conviction stay coherent.
*
* Note: `e_dir` is the joint sum (E_C51 + E_IQN), not the legacy
* `q_b0` value-head + advantage cuBLAS sum. Magnitude (and the
* ISV[21] q_dir_abs_ref EMA the denom is normalised by) was
* previously fit to `q_b0` scale. Both Q signals are intended to
* track the same expected return so the empirical scale should
* track; if Phase 3 reveals a systematic mismatch, this
* normaliser becomes the lever to retune. The `q_range > 0`
* cold-start fallback still holds (denom = max(q_range, 1e-6)). */
float q_max_raw = e_dir[0];
float q_min_raw = e_dir[0];
for (int a = 1; a < b0_size; a++) {
float qv = q_b0[a];
float qv = e_dir[a];
q_max_raw = fmaxf(q_max_raw, qv);
q_min_raw = fminf(q_min_raw, qv);
}

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

@@ -971,6 +971,48 @@ selector inside `experience_action_select` that consumes them. Plan:
`docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-C-phase-2.md`.
Comments expanded post-review to preserve reference parameter-contract docs.
Plan C Phase 2 Task 2 direction-branch Thompson swap (2026-04-29):
`cuda_pipeline/experience_kernels.cu::experience_action_select` —
direction selection (Branch 0) replaced with Thompson sampling at
training and argmax E[Q] at eval, consuming the four device-inline
math primitives Task 1 installed. Kernel signature gains four trailing
arguments after `out_magnitude_conviction`: `c51_probs_dir
[N, b0_size, n_atoms]`, `atom_values [n_atoms]`, `iqn_quantiles_dir
[N, b0_size, N_IQN_QUANTILES]`, `n_atoms`. Existing direction-related
parameters (`eps_dir` derivation, `eps_exp_mult`, `contrarian_active`'s
`q_sign` flip on direction, ISV[21] tau-floor, ISV[71]/[72] adaptive
passive-pressure boost) become unreachable on this branch — Thompson
draws supply exploration directly from the C51 + IQN posteriors and
the adaptive-floor / contrarian-on-direction mechanisms are subsumed
by sample noise. They remain coded for the magnitude branch (`eps_mag`
still consumes `eps_exp_mult`) and unchanged for order/urgency
(`eps_ord`/`eps_urg`). Phase 3 / Plan D removes the direction-only
dead branches.
Per-direction joint E[Q] (`e_dir[4]`) is computed once at the top of
the direction branch and reused for: (a) eval-mode argmax;
(b) the conviction calculation (E[Q]-based per spec — Thompson per-
step samples here would feed the Kelly cap a jittery confidence
signal that re-shapes target position with every draw, defeating the
warmup floor's purpose). The conviction normaliser still reads
ISV[21] (`q_dir_abs_ref`) — it was previously fit to the cuBLAS
`q_b0` scale rather than the C51+IQN joint scale, so Phase 3 may
need to retune if a systematic mismatch surfaces. The cold-start
`fmaxf(q_range, 1e-6)` denominator fallback is preserved.
Magnitude (Branch 1), order (Branch 2), and urgency (Branch 3)
selection paths are untouched — they still read `q_b0`/`q_b1`/`q_b2`/
`q_b3` from the cuBLAS Q-head and apply the existing Boltzmann-with-
adaptive-tau path. The `q_b0`-based `out_q_gaps` writer below is also
unchanged (still computes gap from cuBLAS direction Q, not `e_dir`).
Until Tasks 3 (`gpu_dqn_trainer.rs`) and 4 (`gpu_backtest_evaluator
.rs`) land, the kernel signature is intentionally non-callable end-
to-end — the new trailing arguments have no Rust-side passers yet.
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.
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`