From 79756a21538bd2b711d806ac0b764b5b5b947eea Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 24 May 2026 13:47:21 +0200 Subject: [PATCH] =?UTF-8?q?fix(rl):=20sparse-aware=20EMA=20+=20Q=E2=86=92?= =?UTF-8?q?=CF=80=20distillation=20breaks=20defensive=20trap?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two coupled fixes addressing vj5f6 findings: (1) WIN_clamp oscillation — sparse-aware EMA vj5f6 showed WIN_clamp oscillating 1.0 ↔ 67.0 across 40k steps. Root cause: the Wiener-α blend in rl_reward_clamp_controller treated pos_max=0 as "no win this step ≡ win magnitude is zero," exponentially decaying the EMA toward 0 during dry-spell windows (no closed winning trades). With α=0.4, ten dry steps decayed EMA by 0.6^10 ≈ 0.006, collapsing WIN back to MIN_WIN=1.0 floor. Fix: only update pos_max_ema AND clip_rate_ema AND MARGIN when pos_max > 0. A dry step is "no signal," not "zero signal." The EMA retains its last winning-period estimate; the controller doesn't ratchet on stale data. (2) Q→π distillation — couples Q's improved calibration to π vj5f6 showed l_q dropping 100× (2.37 → 0.02) but reward economics IDENTICAL to 8xwq8 (no C51 V_MAX lift). Per Option B, π drives action selection but is trained by PPO surrogate using advantage = returns - V. V regression doesn't benefit from C51 calibration, so Q's improved knowledge stays trapped in the critic head. Deep audit revealed a self-reinforcing defensive trap: Q learned "big positions lose money" → π_target favors small actions → π picks a3+a4 (tiny long / Hold) → position lots ≈ 0 → rewards mostly 0 → V learns "everything is 0" → V_pred ≈ 0 → advantage = returns - V_pred ≈ 0 → PPO gradient ≈ 0 → π frozen at defensive attractor → loop. Trade count dropped 3× (rdgzl 25k → 8xwq8/vj5f6 9k closes per 10k steps), win rate inversely correlated with l_q (50% early → 22% late) because only forced closes happen (stops = losses). Fix: new rl_q_pi_distill_grad.cu computes π_target = softmax(E_Q[s,*] / τ) ∂L/∂logits[a] = λ × (π_new(a) - π_target(a)) and ADDS this gradient to pi_grad_logits AFTER the PPO surrogate backward. Couples Q's preferences directly into π's update without going through advantage. λ=0.01 (small, PPO dominant), τ=1.0 (canonical Boltzmann). 3 new ISV slots (λ + τ + KL_ema diag). Diag exposes c51_v_max/v_min, q_distill_lambda/temperature, and q_distill_kl_ema so the adaptation + distillation loop is observable. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/rl_q_pi_distill_grad.cu | 174 ++++++++++++++++++ .../cuda/rl_reward_clamp_controller.cu | 97 +++++----- crates/ml-alpha/examples/alpha_rl_train.rs | 7 + crates/ml-alpha/src/rl/isv_slots.rs | 42 ++++- crates/ml-alpha/src/trainer/integrated.rs | 54 +++++- 6 files changed, 329 insertions(+), 46 deletions(-) create mode 100644 crates/ml-alpha/cuda/rl_q_pi_distill_grad.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index fb1653184..9e042e5a8 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -70,6 +70,7 @@ const KERNELS: &[&str] = &[ "rl_pi_action_kernel", // audit Option B: π drives action selection via multinomial sampling from softmax(pi_logits); Q becomes pure critic "rl_reward_clamp_controller", // audit 2026-05-24: adaptive [-LOSS, +WIN] clamp from positive-tail EMA; replaces static [-3, +1] that crushed winning-trade signal in rmgm5 "rl_atom_support_update", // audit 2026-05-24 followup: refreshes atom_supports_d from ISV V_MIN/V_MAX so C51 atom span adapts with reward clamp (Q learning was capped at V_MAX=1.0) + "rl_q_pi_distill_grad", // audit 2026-05-24 vj5f6 followup: KL(softmax(E_Q/τ) || π_new) gradient ADDED to pi_grad_logits — couples Q's improved C51 calibration to action selection (was decoupled per Option B) ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_q_pi_distill_grad.cu b/crates/ml-alpha/cuda/rl_q_pi_distill_grad.cu new file mode 100644 index 000000000..32cc8e812 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_q_pi_distill_grad.cu @@ -0,0 +1,174 @@ +// rl_q_pi_distill_grad.cu — Q→π KL distillation gradient kernel. +// +// Audit 2026-05-24 (vj5f6 follow-up): the C51 V_MAX lift made Q +// distributional learning calibrated to actual reward magnitudes +// (l_q dropped 100×), but reward economics didn't change because +// per Option B (`pearl_q_thompson_actor_makes_pi_dead_weight`), π +// drives action selection via multinomial sample of softmax(pi_logits) +// and is trained by PPO surrogate using advantage = returns - V. +// V regression doesn't benefit from C51 atom-span calibration, so +// Q's improved knowledge stays trapped in the critic. +// +// Fix: add a KL distillation term to the policy loss that pulls π +// toward a Boltzmann distribution over Q's expected action values: +// +// E_Q[s,a] = sum_z(softmax(q_logits[s,a,*])[z] × atom_supports[z]) +// π_target = softmax(E_Q[s,*] / τ) over actions +// L_distill = KL(π_target || π_new) (forward KL) +// ∂L/∂logits = λ × (π_new(a) - π_target(a)) (xent-like) +// +// The kernel ADDS this gradient to pi_grad_logits (additive — does +// not overwrite the existing PPO surrogate gradient). PPO retains +// the dominant learning signal; distill is a soft bias toward Q's +// argmax. Forward KL chosen over reverse so π_target's high-prob +// actions (Q-preferred) dominate the gradient; reverse KL would +// just keep π broad. +// +// Block layout: one block per batch, N_ACTIONS threads. Each thread +// handles one action's E_Q + softmax row + final gradient write. +// Shared mem holds the E_Q vector + π_target + π_new for the +// per-block softmax reductions. +// +// Per `feedback_no_atomicadd`: shared-mem reductions only, no atomics. +// Per `feedback_cpu_is_read_only`: all state in ISV / device buffers. + +#include + +#define N_ACTIONS 9 +#define Q_N_ATOMS 21 +#define RL_Q_DISTILL_LAMBDA_INDEX 486 +#define RL_Q_DISTILL_TEMPERATURE_INDEX 487 +#define RL_Q_DISTILL_KL_EMA_INDEX 488 + +#define KL_EMA_ALPHA 0.05f + +extern "C" __global__ void rl_q_pi_distill_grad( + const float* __restrict__ q_logits, // [B × N_ACTIONS × Q_N_ATOMS] + const float* __restrict__ pi_logits, // [B × N_ACTIONS] + const float* __restrict__ atom_supports, // [Q_N_ATOMS] + float* __restrict__ isv, // ISV bus (RW — KL_ema) + float* __restrict__ pi_grad_logits, // [B × N_ACTIONS] — ADD to + int B +) { + const int b = blockIdx.x; + const int a = threadIdx.x; + if (b >= B || a >= N_ACTIONS) return; + + const float lambda = isv[RL_Q_DISTILL_LAMBDA_INDEX]; + const float tau = isv[RL_Q_DISTILL_TEMPERATURE_INDEX]; + + // Skip work if distillation is disabled (λ=0). Still need to + // syncthreads to keep the block lockstep — but a single read + + // early return on all threads is safe since they all read the + // same λ. (Actually, returning here means subsequent + // __syncthreads in the kernel hang; safer to keep going with + // grad=0 contribution.) + const bool active = (lambda > 0.0f); + + // ── Step 1: each thread computes E_Q for its OWN action ───── + // + // E_Q[a] = sum_z(softmax(q_logits[b,a,*])[z] × atom_supports[z]) + // + // Numerically-stable softmax: subtract max before exp. + __shared__ float s_eq[N_ACTIONS]; + + const int q_base = (b * N_ACTIONS + a) * Q_N_ATOMS; + float max_q = q_logits[q_base]; + #pragma unroll + for (int z = 1; z < Q_N_ATOMS; ++z) { + max_q = fmaxf(max_q, q_logits[q_base + z]); + } + float sum_exp = 0.0f; + float probs_local[Q_N_ATOMS]; + #pragma unroll + for (int z = 0; z < Q_N_ATOMS; ++z) { + probs_local[z] = expf(q_logits[q_base + z] - max_q); + sum_exp += probs_local[z]; + } + const float inv_sum = (sum_exp > 1e-9f) ? (1.0f / sum_exp) + : (1.0f / (float)Q_N_ATOMS); + float e_q = 0.0f; + #pragma unroll + for (int z = 0; z < Q_N_ATOMS; ++z) { + e_q += probs_local[z] * inv_sum * atom_supports[z]; + } + s_eq[a] = e_q; + __syncthreads(); + + // ── Step 2: thread 0 builds π_target = softmax(E_Q / τ) ───── + __shared__ float s_pi_target[N_ACTIONS]; + if (a == 0) { + float max_eq = s_eq[0]; + #pragma unroll + for (int i = 1; i < N_ACTIONS; ++i) max_eq = fmaxf(max_eq, s_eq[i]); + const float tau_safe = fmaxf(tau, 1e-3f); // guard against τ=0 + float total = 0.0f; + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) { + s_pi_target[i] = expf((s_eq[i] - max_eq) / tau_safe); + total += s_pi_target[i]; + } + const float inv = (total > 1e-9f) ? (1.0f / total) + : (1.0f / (float)N_ACTIONS); + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) s_pi_target[i] *= inv; + } + __syncthreads(); + + // ── Step 3: thread 0 builds π_new = softmax(pi_logits) ────── + __shared__ float s_pi_new[N_ACTIONS]; + if (a == 0) { + const int pi_base = b * N_ACTIONS; + float max_pi = pi_logits[pi_base]; + #pragma unroll + for (int i = 1; i < N_ACTIONS; ++i) { + max_pi = fmaxf(max_pi, pi_logits[pi_base + i]); + } + float total = 0.0f; + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) { + s_pi_new[i] = expf(pi_logits[pi_base + i] - max_pi); + total += s_pi_new[i]; + } + const float inv = (total > 1e-9f) ? (1.0f / total) + : (1.0f / (float)N_ACTIONS); + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) s_pi_new[i] *= inv; + } + __syncthreads(); + + // ── Step 4: ADD distill gradient. ─────────────────────────── + // + // ∂KL/∂logits[a] = λ × (π_new(a) - π_target(a)) + // + // Additive so PPO's surrogate gradient (already written) keeps + // its contribution. Each thread writes one action's grad slot. + if (active) { + const int pi_idx = b * N_ACTIONS + a; + pi_grad_logits[pi_idx] += lambda * (s_pi_new[a] - s_pi_target[a]); + } + + // ── Step 5 (diag): thread 0 of block 0 EMAs the per-batch + // mean KL(π_target || π_new) into ISV. Done by block 0 only so + // we don't race across the grid; we reduce within this block + // first then commit. + __shared__ float s_kl; + if (a == 0) { + float kl = 0.0f; + #pragma unroll + for (int i = 0; i < N_ACTIONS; ++i) { + const float pt = s_pi_target[i]; + const float pn = fmaxf(s_pi_new[i], 1e-12f); + if (pt > 1e-12f) kl += pt * logf(pt / pn); + } + s_kl = kl; + } + __syncthreads(); + + if (b == 0 && a == 0) { + const float kl_prev = isv[RL_Q_DISTILL_KL_EMA_INDEX]; + const float kl_new = (kl_prev == 0.0f) ? s_kl + : (1.0f - KL_EMA_ALPHA) * kl_prev + KL_EMA_ALPHA * s_kl; + isv[RL_Q_DISTILL_KL_EMA_INDEX] = kl_new; + } +} diff --git a/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu b/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu index 6e466f656..db8fb139f 100644 --- a/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu +++ b/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu @@ -101,26 +101,34 @@ extern "C" __global__ void rl_reward_clamp_controller( const float pos_max = isv[RL_POS_SCALED_REWARD_MAX_INDEX]; const float ema_prev = isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX]; - // EMA maintenance — bootstrap on sentinel 0 with first non-zero - // raw observation; otherwise Wiener-α blend with floor. - float ema_new; - if (ema_prev == 0.0f) { - if (pos_max == 0.0f) { - // Both prev and obs sentinels — defer adaptation, leave the - // statically-seeded WIN(1.0) / LOSS(3.0) bounds untouched. - // Next call where the kernel observes a positive reward - // will bootstrap the EMA and start adapting. - return; + // EMA maintenance — sparse-sample-aware. Only update the EMA on + // steps that ACTUALLY observed a positive reward (pos_max > 0). + // + // Audit 2026-05-24 (vj5f6 follow-up): the prior implementation + // blended pos_max=0 into the EMA via the Wiener-α step, which + // exponentially decayed the EMA toward 0 during dry-spell windows + // (no closed winning trades). With α=0.4, ten consecutive + // dry steps decay EMA by 0.6^10 ≈ 0.006, collapsing the WIN + // bound back to the MIN_WIN floor. WIN_clamp was observed + // oscillating 1.0 ↔ 67.0 across 40k steps in vj5f6. + // + // Fix: pos_max=0 means "no signal this step," not "win magnitude + // is zero." Skip EMA updates during dry spells; the EMA retains + // the last winning-period estimate until the next observation. + float ema_new = ema_prev; + if (pos_max > 0.0f) { + if (ema_prev == 0.0f) { + ema_new = pos_max; // first-observation bootstrap + } else { + const float a = fmaxf(alpha, WIENER_ALPHA_FLOOR); + ema_new = (1.0f - a) * ema_prev + a * pos_max; } - ema_new = pos_max; - } else { - const float a = fmaxf(alpha, WIENER_ALPHA_FLOOR); - // If pos_max==0 this step (no positive reward sample), still - // blend toward 0 — captures regimes where wins disappear. The - // EMA decays gracefully toward 0 over ~1/alpha steps. - ema_new = (1.0f - a) * ema_prev + a * pos_max; + isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX] = ema_new; } - isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX] = ema_new; + // If ema_prev was 0 AND pos_max was 0 → ema_new still 0 here → + // WIN_eff below = max(MIN_WIN, MARGIN × 0) = MIN_WIN. That's + // the correct "no data yet" behavior (statically-seeded WIN=1.0 + // remains effective). // ── MARGIN adaptation from clip-rate EMA (audit follow-up). ─── // @@ -136,36 +144,39 @@ extern "C" __global__ void rl_reward_clamp_controller( clip_indicator = (pos_max > win_active) ? 1.0f : 0.0f; } - // Step 2: blend clip indicator into EMA. Fixed α=0.05 here (not - // Wiener-driven — the indicator is bimodal 0/1 so variance-based α - // would saturate; constant α gives ~20-step half-life which - // matches the policy-update cadence). + // Step 2: blend clip indicator into EMA. Sparse-aware — only + // update when pos_max > 0 (matches the pos_max EMA discipline). + // A step with pos_max=0 isn't "0% clipped"; it's "no data." Mixing + // dry steps in would bias clip_rate_ema downward and the MARGIN + // controller would shrink WIN unnecessarily. const float clip_rate_prev = isv[RL_REWARD_CLAMP_CLIP_RATE_EMA_INDEX]; - float clip_rate_new; - if (clip_rate_prev == 0.0f) { - // Bootstrap on sentinel — set to current indicator (avoids - // long warmup where MARGIN stays static). - clip_rate_new = clip_indicator; - } else { - clip_rate_new = (1.0f - CLIP_RATE_EMA_ALPHA) * clip_rate_prev - + CLIP_RATE_EMA_ALPHA * clip_indicator; + float clip_rate_new = clip_rate_prev; + if (pos_max > 0.0f) { + if (clip_rate_prev == 0.0f) { + clip_rate_new = clip_indicator; // bootstrap + } else { + clip_rate_new = (1.0f - CLIP_RATE_EMA_ALPHA) * clip_rate_prev + + CLIP_RATE_EMA_ALPHA * clip_indicator; + } + isv[RL_REWARD_CLAMP_CLIP_RATE_EMA_INDEX] = clip_rate_new; } - isv[RL_REWARD_CLAMP_CLIP_RATE_EMA_INDEX] = clip_rate_new; - // Step 3: Schulman bounded step on MARGIN — if clip rate > target - // × TOLERANCE, raise MARGIN (widen WIN, capture more tail); if - // clip rate < target / TOLERANCE, lower MARGIN (tighten WIN, - // sharper regression). Dead-zone in between prevents oscillation. - const float clip_target = isv[RL_REWARD_CLAMP_CLIP_RATE_TARGET_INDEX]; + // Step 3: Schulman bounded step on MARGIN — only adjust when we + // actually have a fresh clip-rate observation. If pos_max=0 this + // step, MARGIN retains its prior value (consistent with the EMA + // behavior above; the controller doesn't ratchet on stale data). float margin = isv[RL_REWARD_CLAMP_MARGIN_INDEX]; - const float upper = clip_target * MARGIN_TOLERANCE; - const float lower = clip_target / MARGIN_TOLERANCE; - if (clip_rate_new > upper) { - margin = fminf(MAX_MARGIN, margin * MARGIN_ADJUST_RATE); - } else if (clip_rate_new < lower) { - margin = fmaxf(MIN_MARGIN, margin / MARGIN_ADJUST_RATE); + if (pos_max > 0.0f) { + const float clip_target = isv[RL_REWARD_CLAMP_CLIP_RATE_TARGET_INDEX]; + const float upper = clip_target * MARGIN_TOLERANCE; + const float lower = clip_target / MARGIN_TOLERANCE; + if (clip_rate_new > upper) { + margin = fminf(MAX_MARGIN, margin * MARGIN_ADJUST_RATE); + } else if (clip_rate_new < lower) { + margin = fmaxf(MIN_MARGIN, margin / MARGIN_ADJUST_RATE); + } + isv[RL_REWARD_CLAMP_MARGIN_INDEX] = margin; } - isv[RL_REWARD_CLAMP_MARGIN_INDEX] = margin; // Step 4: compute adaptive WIN bound: MARGIN × EMA, floored. // No upper cap per the header note — the MARGIN ∈ [1, 5] × EMA diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 4e3f65b11..22a637f7c 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -68,6 +68,7 @@ use ml_alpha::rl::isv_slots::{ RL_REWARD_CLAMP_MARGIN_INDEX, RL_REWARD_CLAMP_RATIO_INDEX, RL_REWARD_CLAMP_CLIP_RATE_EMA_INDEX, RL_REWARD_CLAMP_CLIP_RATE_TARGET_INDEX, RL_C51_V_MAX_INDEX, RL_C51_V_MIN_INDEX, + RL_Q_DISTILL_LAMBDA_INDEX, RL_Q_DISTILL_TEMPERATURE_INDEX, RL_Q_DISTILL_KL_EMA_INDEX, RL_MEAN_TRADE_DURATION_EMA_INDEX, RL_N_ROLLOUT_STEPS_INDEX, RL_PER_ALPHA_INDEX, RL_PI_GRAD_NORM_EMA_INDEX, RL_PPO_CLIP_INDEX, RL_PPO_LOG_RATIO_ABS_MAX_INDEX, RL_PPO_RATIO_CLAMP_MAX_INDEX, RL_Q_DIVERGENCE_EMA_INDEX, RL_Q_GRAD_NORM_EMA_INDEX, @@ -684,6 +685,10 @@ fn main() -> Result<()> { isv[RL_C51_V_MAX_INDEX], "c51_v_min": isv[RL_C51_V_MIN_INDEX], + // Q→π distillation KL signal — drops toward 0 if π is + // successfully chasing Q's preferences. + "q_distill_kl_ema": + isv[RL_Q_DISTILL_KL_EMA_INDEX], }, // audit — PPO importance-ratio clamp diagnostics. // @@ -795,6 +800,8 @@ fn main() -> Result<()> { "reward_clamp_margin": isv[RL_REWARD_CLAMP_MARGIN_INDEX], "reward_clamp_ratio": isv[RL_REWARD_CLAMP_RATIO_INDEX], "reward_clamp_clip_rate_target": isv[RL_REWARD_CLAMP_CLIP_RATE_TARGET_INDEX], + "q_distill_lambda": isv[RL_Q_DISTILL_LAMBDA_INDEX], + "q_distill_temperature": isv[RL_Q_DISTILL_TEMPERATURE_INDEX], }, // audit — Q vs π action agreement EMA at slot 407 // (previously dead). 1.0 = perfect ranking consistency, diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index a492d60dd..18dab67fa 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -41,8 +41,9 @@ //! | 478-481 | 4 adaptive reward-clamp slots (raw + EMA + margin + ratio) | apply_reward_scale + rl_reward_clamp_controller | //! | 482-483 | 2 MARGIN-adaptation slots (clip-rate EMA + target) | rl_reward_clamp_controller (Schulman step on MARGIN) | //! | 484-485 | 2 C51 atom-span ratchet slots (V_MAX + V_MIN) | rl_reward_clamp_controller (ratchet) + rl_atom_support_update (writes atom_supports_d) | +//! | 486-488 | 3 Q→π distillation slots (λ + τ + KL_ema) | rl_isv_write (seed) + rl_q_pi_distill_grad (ADDS to pi_grad_logits + writes KL_ema) | //! -//! Total: 86 slots; `RL_SLOTS_END = 486`. +//! Total: 89 slots; `RL_SLOTS_END = 489`. /// Discount factor γ. Controller input: mean trade duration / anchor. /// Bootstrap 0.99. @@ -646,6 +647,43 @@ pub const RL_C51_V_MAX_INDEX: usize = 484; /// producers / consumers. pub const RL_C51_V_MIN_INDEX: usize = 485; +// ─── Q→π distillation (vj5f6 follow-up 2026-05-24) ───────────────── +// +// vj5f6 showed that with the C51 atom span lifted, Q learns +// beautifully (best l_q dropped from 2.37 → 0.02 — 100× deeper). +// But reward economics were IDENTICAL to 8xwq8 (no C51 lift): the +// agent's policy is determined by π, which is trained by PPO +// surrogate against advantage = returns - V. V is scalar regression, +// doesn't benefit from C51 calibration. So Q's improved knowledge +// stays trapped in the critic head and doesn't reach action +// selection. +// +// Fix: add a KL distillation term to the policy loss that pulls π +// toward a Boltzmann distribution over Q's expected action values. +// π_target(a|s) = softmax(E_Q[s,a] / τ) where E_Q = sum_z(p_z × atom_z) +// L_distill = KL(π_target || π_new) +// ∂L/∂logits[a] = λ × (π_new(a) - π_target(a)) +// +// The kernel ADDS this gradient to pi_grad_logits AFTER PPO's +// surrogate backward, so PPO retains its role and distill is an +// additive bias toward Q's preferences. Tune τ (sharpness of Q +// target) and λ (distill weight) via ISV. + +/// Q→π distillation weight. Default 0.01 — small so PPO retains +/// dominant learning signal but Q's preferences leak through. +/// Adaptive lift later if needed. +pub const RL_Q_DISTILL_LAMBDA_INDEX: usize = 486; + +/// Q→π distillation temperature τ. Default 1.0 — softmax over +/// E_Q with unit sharpness. Lower τ → sharper π_target (commits +/// to Q's argmax); higher → softer (averaged Q preference). +pub const RL_Q_DISTILL_TEMPERATURE_INDEX: usize = 487; + +/// Diagnostic — EMA of mean KL(π_target || π_new). Surfaces in diag +/// so we can verify the distill term is moving π toward Q +/// (KL should drop over training). +pub const RL_Q_DISTILL_KL_EMA_INDEX: usize = 488; + /// Last RL-allocated slot index (exclusive). The integrated trainer /// extends `ISV_TOTAL_DIM` to at least this value at trainer init time. -pub const RL_SLOTS_END: usize = 486; +pub const RL_SLOTS_END: usize = 489; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 0592592ba..7c460aba6 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -173,6 +173,13 @@ const RL_REWARD_CLAMP_CONTROLLER_CUBIN: &[u8] = // projection in the next step. 21-thread one-block kernel. const RL_ATOM_SUPPORT_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_atom_support_update.cubin")); +// Q→π distillation gradient — audit 2026-05-24 vj5f6 followup. ADDS +// λ × (π_new - π_target) to pi_grad_logits after the PPO surrogate +// backward, where π_target = softmax(E_Q[s,*] / τ). Couples Q's +// improved C51 calibration to π's action selection (was decoupled +// per Option B / `pearl_q_thompson_actor_makes_pi_dead_weight`). +const RL_Q_PI_DISTILL_GRAD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_q_pi_distill_grad.cubin")); const ACTIONS_TO_MARKET_TARGETS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/actions_to_market_targets.cubin")); @@ -394,6 +401,9 @@ pub struct IntegratedTrainer { // C51 atom-support updater (audit 2026-05-24 followup). _rl_atom_support_update_module: Arc, rl_atom_support_update_fn: CudaFunction, + // Q→π distillation gradient (audit 2026-05-24 vj5f6 followup). + _rl_q_pi_distill_grad_module: Arc, + rl_q_pi_distill_grad_fn: CudaFunction, _actions_to_market_targets_module: Arc, actions_to_market_targets_fn: CudaFunction, @@ -791,6 +801,12 @@ impl IntegratedTrainer { let rl_atom_support_update_fn = rl_atom_support_update_module .load_function("rl_atom_support_update") .context("load rl_atom_support_update")?; + let rl_q_pi_distill_grad_module = ctx + .load_cubin(RL_Q_PI_DISTILL_GRAD_CUBIN.to_vec()) + .context("load rl_q_pi_distill_grad cubin")?; + let rl_q_pi_distill_grad_fn = rl_q_pi_distill_grad_module + .load_function("rl_q_pi_distill_grad") + .context("load rl_q_pi_distill_grad")?; let actions_to_market_targets_module = ctx .load_cubin(ACTIONS_TO_MARKET_TARGETS_CUBIN.to_vec()) .context("load actions_to_market_targets cubin")?; @@ -1101,6 +1117,8 @@ impl IntegratedTrainer { rl_reward_clamp_controller_fn, _rl_atom_support_update_module: rl_atom_support_update_module, rl_atom_support_update_fn, + _rl_q_pi_distill_grad_module: rl_q_pi_distill_grad_module, + rl_q_pi_distill_grad_fn, _actions_to_market_targets_module: actions_to_market_targets_module, actions_to_market_targets_fn, _abs_copy_module: abs_copy_module, @@ -1218,7 +1236,7 @@ impl IntegratedTrainer { // (slot, value) pair — pure device write, no HtoD per // `feedback_no_htod_htoh_only_mapped_pinned`. { - let isv_constants: [(usize, f32); 31] = [ + let isv_constants: [(usize, f32); 33] = [ // Static seeds for the adaptive reward-clamp controller — // these are the initial values that // `rl_reward_clamp_controller` will replace once it @@ -1238,6 +1256,10 @@ impl IntegratedTrainer { (crate::rl::isv_slots::RL_REWARD_CLAMP_CLIP_RATE_TARGET_INDEX, 0.05), (crate::rl::isv_slots::RL_C51_V_MAX_INDEX, 1.0), (crate::rl::isv_slots::RL_C51_V_MIN_INDEX, -1.0), + // Q→π distillation (vj5f6 followup) — λ small so PPO + // dominates; τ=1.0 for canonical softmax(E_Q). + (crate::rl::isv_slots::RL_Q_DISTILL_LAMBDA_INDEX, 0.01), + (crate::rl::isv_slots::RL_Q_DISTILL_TEMPERATURE_INDEX, 1.0), (crate::rl::isv_slots::RL_KL_TARGET_INDEX, 0.01), (crate::rl::isv_slots::RL_IMPROVEMENT_THRESHOLD_INDEX, 0.99), (crate::rl::isv_slots::RL_PLATEAU_PATIENCE_INDEX, 1000.0), @@ -2149,6 +2171,36 @@ impl IntegratedTrainer { ) .context("policy_head.surrogate_backward_logits")?; + // Q→π distillation gradient — ADDS λ × (π_new - π_target) to + // pi_grad_logits where π_target = softmax(E_Q[s,*] / τ). + // Couples Q's improved C51 calibration to π's action selection + // (per vj5f6 finding that Q was decoupled from policy under + // Option B). One block per batch, N_ACTIONS=9 threads per block. + // q_logits_d is still in scope from the earlier forward. + { + let cfg_distill = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (N_ACTIONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self + .stream + .launch_builder(&self.rl_q_pi_distill_grad_fn); + launch + .arg(&q_logits_d) + .arg(&pi_logits_d) + .arg(&self.atom_supports_d) + .arg(&self.isv_d) + .arg(&mut pi_grad_logits_d) + .arg(&b_size_i); + unsafe { + launch + .launch(cfg_distill) + .context("rl_q_pi_distill_grad launch")?; + } + } + self.policy_head .backward_to_w_b_h( h_t_borrow,