From e074c91fb2bc5fb560daf7972b00e025d35e1653 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 23 May 2026 18:31:06 +0200 Subject: [PATCH] fix(rl): LR controller per-step rate cap + per-head TARGET_GRAD_NORM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cluster smoke `alpha-rl-nqd68` showed the signal-driven LR controller working mechanically but destabilising the π head: lr_pi swung MIN→MAX (1000×) over ~10k steps, then got stuck at MAX after the catastrophic Adam updates wrecked the policy weights. Aggregate l_pi max = 2.4e17, l_v max = 2,083,330 (no NaN abort, but useless for learning). Q head was fine (lr_q correctly stuck at MIN throughout, l_q dropped 34% vs fixed-LR). ## Two fixes ### 1. Per-step rate-of-change cap The original target formula `target = lr_prev × (TARGET/observed)` allows arbitrary swing magnitude. When `observed` is tiny (e.g. quiescent π grad-norm between trade closes), target = lr_prev × 1000, which the Wiener α=0.4 blend drags toward LR_MAX in a few steps. Once at MAX, the next real reward signal applies catastrophic Adam updates → policy explodes → grad-norm spikes 10⁵× → controller sees this and tries to shrink, but the damage is done. Adds `target_lr ∈ [lr_prev × 0.5, lr_prev × 2.0]` constraint post-formula, pre-clamp. The controller can now at most halve or double LR per step, taking ~10 steps to traverse the full [LR_MIN, LR_MAX] range. Downstream gradient signal has time to react before LR overshoots. Same pattern as `rl_rollout_steps_controller`'s `scale ∈ [0.5, 2.0]` cap (which was added for the same class of multiplicative-controller instability). ### 2. Per-head TARGET_GRAD_NORM The single `TARGET_GRAD_NORM = 1.0` anchor was wrong for π and V: those heads have far fewer parameters than Q (1,152 and 128 vs 24,192). A "well-tuned" grad-norm magnitude scales with √n_params (so per-parameter grad magnitude stays Adam-friendly ≈ 1e-2). Q head w_d: 9 × 21 × 128 = 24,192 params → √ ≈ 156 → target 1.5 π head w_d: 9 × 128 = 1,152 params → √ ≈ 34 → target 0.3 V head w_d: 128 = 128 params → √ ≈ 11 → target 0.1 Without this scaling, the controller was pushing π LR up because its grad-norm (typically 0.1-0.3) was always "below the 1.0 target" — interpreted as "model coasting, grow LR" when really the smaller grad-norm just reflected the smaller parameter count. `update_lr_with_signal` now takes `head_target_grad_norm` as a parameter. BCE and AUX heads (owned by perception, signal_slot=-1) get target=1.0 but it's unused because the early-return at `signal_slot < 0` short-circuits past the target derivation. ## Verified gates (local sm_86) G1 isv_bootstrap ✅ G3 controllers_emit ✅ G4 target_soft_update ✅ G6 r7d_per_wiring ✅ smoke ✅ all losses finite ## Expected effect on next 50k smoke * lr_q stays near MIN (already worked — Q grad-norm > target_Q typically) — unchanged. * lr_pi should NOT runaway to MAX — rate cap limits 1000× swing to at most 2× per step; per-head π target 0.3 puts the multiplicative ratio closer to 1.0 (no extreme target). * lr_v should also stabilise via the V-specific target 0.1. * Aggregate l_pi / l_v max values should drop from the 1e17 / 2e6 range to O(1). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/rl_lr_controller.cu | 70 +++++++++++++++++++----- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/crates/ml-alpha/cuda/rl_lr_controller.cu b/crates/ml-alpha/cuda/rl_lr_controller.cu index b3b326c4c..37c5ec525 100644 --- a/crates/ml-alpha/cuda/rl_lr_controller.cu +++ b/crates/ml-alpha/cuda/rl_lr_controller.cu @@ -54,16 +54,43 @@ #define LR_BOOTSTRAP 1e-3f #define LR_MIN 1e-5f #define LR_MAX 1e-2f -// Target grad-norm magnitude. Chosen so a typical post-backward -// grad-norm at LR=LR_BOOTSTRAP lands near this value; observed > -// target → controller shrinks LR; observed < target → grows LR. -// Magnitude derived empirically from the 50k smoke (Q grad_w -// L2 norms typically O(1) at LR=1e-3). -#define TARGET_GRAD_NORM 1.0f + +// Per-head TARGET_GRAD_NORM anchor — the "well-tuned" grad-norm +// magnitude the controller drives toward. Heuristic: scaled to +// √n_params so per-parameter grad magnitude ≈ 1e-2 (Adam-friendly). +// +// Q head w_d: N_ACTIONS × Q_N_ATOMS × HIDDEN_DIM = 9×21×128 = 24,192 params +// √24192 ≈ 156 → target 1.5 +// π head w_d: N_ACTIONS × HIDDEN_DIM = 9×128 = 1,152 params +// √1152 ≈ 34 → target 0.3 +// V head w_d: HIDDEN_DIM = 128 params +// √128 ≈ 11 → target 0.1 +// +// BCE and AUX are owned by the perception trainer; passing +// signal_slot < 0 from the RL trainer's launch bypasses the +// signal-driven path entirely, so the target here is unused. +#define Q_TARGET_GRAD_NORM 1.5f +#define PI_TARGET_GRAD_NORM 0.3f +#define V_TARGET_GRAD_NORM 0.1f +#define BCE_TARGET_GRAD_NORM 1.0f // unused (perception-owned) +#define AUX_TARGET_GRAD_NORM 1.0f // unused + +// Per-step rate-of-change cap on target_lr (cluster smoke +// `alpha-rl-nqd68` showed π controller swinging MIN→MAX 1000× across +// ~10k steps, then getting stuck at MAX after Adam updates wrecked +// the policy weights). With this cap, the per-step target is bounded +// to [lr_prev × 0.5, lr_prev × 2.0] — at most a 2× swing per step, +// so traversing the full [LR_MIN, LR_MAX] range takes ~10 steps each +// direction. That gives downstream gradient signal time to react +// before LR catastrophically overshoots. Same general pattern as +// rl_rollout_steps_controller's `scale ∈ [0.5, 2.0]` cap. +#define LR_RATE_CAP_LO 0.5f +#define LR_RATE_CAP_HI 2.0f + #define EMA_EPS 1e-6f __device__ __forceinline__ void update_lr_with_signal( - float* isv, int lr_idx, float alpha, int signal_slot + float* isv, int lr_idx, float alpha, int signal_slot, float head_target_grad_norm ) { const float lr_prev = isv[lr_idx]; @@ -90,12 +117,25 @@ __device__ __forceinline__ void update_lr_with_signal( // τ/ε/n_roll/scale. if (observed == 0.0f) return; - // Multiplicative target derivation. - float target_lr = lr_prev * (TARGET_GRAD_NORM / fmaxf(observed, EMA_EPS)); + // Multiplicative target derivation using per-head target anchor. + float target_lr = lr_prev * (head_target_grad_norm / fmaxf(observed, EMA_EPS)); + + // Per-step rate-of-change cap: target bounded to + // [lr_prev × LR_RATE_CAP_LO, lr_prev × LR_RATE_CAP_HI]. Prevents + // the runaway-swing instability that destabilised the π head in + // alpha-rl-nqd68 (target = lr_prev × 1000 when grad-norm was + // tiny → controller jumped to LR_MAX in a few Wiener blends). + const float lr_lower = lr_prev * LR_RATE_CAP_LO; + const float lr_upper = lr_prev * LR_RATE_CAP_HI; + target_lr = fmaxf(lr_lower, fminf(lr_upper, target_lr)); + + // Absolute clamp to [LR_MIN, LR_MAX] (runaway protection). target_lr = fminf(LR_MAX, fmaxf(LR_MIN, target_lr)); // First-observation replace-directly: avoid Wiener-blend's 60% - // bootstrap contamination on the first warm step. + // bootstrap contamination on the first warm step. Rate cap + // already applied above so the "direct" target respects the + // ≤2× rate constraint from the bootstrap LR. if (lr_prev == LR_BOOTSTRAP) { isv[lr_idx] = target_lr; return; @@ -119,9 +159,9 @@ extern "C" __global__ void rl_lr_controller( ) { if (threadIdx.x != 0 || blockIdx.x != 0) return; - update_lr_with_signal(isv, RL_LR_BCE_INDEX, alpha, bce_signal_slot); - update_lr_with_signal(isv, RL_LR_Q_INDEX, alpha, q_signal_slot); - update_lr_with_signal(isv, RL_LR_PI_INDEX, alpha, pi_signal_slot); - update_lr_with_signal(isv, RL_LR_V_INDEX, alpha, v_signal_slot); - update_lr_with_signal(isv, RL_LR_AUX_INDEX, alpha, aux_signal_slot); + update_lr_with_signal(isv, RL_LR_BCE_INDEX, alpha, bce_signal_slot, BCE_TARGET_GRAD_NORM); + update_lr_with_signal(isv, RL_LR_Q_INDEX, alpha, q_signal_slot, Q_TARGET_GRAD_NORM); + update_lr_with_signal(isv, RL_LR_PI_INDEX, alpha, pi_signal_slot, PI_TARGET_GRAD_NORM); + update_lr_with_signal(isv, RL_LR_V_INDEX, alpha, v_signal_slot, V_TARGET_GRAD_NORM); + update_lr_with_signal(isv, RL_LR_AUX_INDEX, alpha, aux_signal_slot, AUX_TARGET_GRAD_NORM); }