diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 38649708b..d227fda07 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -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); } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 8356c513e..c31747a67 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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`