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); }