fix(rl): asymmetric Q→π distillation λ controller — ramp fast, decay slow
Root cause of q_pi_agree collapse: the symmetric Schulman controller decayed λ from 0.5 to 0.001 in 34 steps (÷1.2/step), killing Q→π coupling. Once decoupled, Q and π learned independent policies, making q_pi_agree drop to 1e-22. Three fixes: 1. MIN_LAMBDA 0.001 → 0.05: Q pull never drops below 5% strength 2. Asymmetric rates: ramp 1.2×/step (4 steps to double), decay 0.998×/step (347 steps to halve). Coupling establishes fast and persists for thousands of steps. 3. Dead zone tolerance 1.5× → 3.0×: natural KL fluctuation stays in-band instead of triggering constant ramp/decay cycles. Seed λ raised to 0.1 so distillation is active from step 0. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -1,31 +1,20 @@
|
||||
// rl_q_distill_lambda_controller.cu — adaptive λ_distill controller.
|
||||
//
|
||||
// Audit 2026-05-24 (rljzl followup design): λ_distill (slot 486) was
|
||||
// previously seeded statically (0.01 in wwcsz, 0.05 in rljzl). Static
|
||||
// λ leaves Q's pull strength at design intuition — too low and Q's
|
||||
// preferences don't reach π; too high and PPO surrogate gets drowned.
|
||||
// Make λ adaptive driven by the natural feedback signal:
|
||||
// KL(π_target || π_new) EMA at slot 488.
|
||||
// Adapts the Q→π distillation strength (λ, slot 486) from the KL
|
||||
// divergence between Q's Boltzmann policy and π's softmax (slot 488).
|
||||
//
|
||||
// Controller logic (Schulman bounded step):
|
||||
// Controller logic (asymmetric bounded step):
|
||||
//
|
||||
// target_kl = isv[491] // RL_Q_DISTILL_KL_TARGET (seed 0.1)
|
||||
// observed = isv[488] // RL_Q_DISTILL_KL_EMA
|
||||
// upper = target × TOLERANCE // 0.15 at default
|
||||
// lower = target / TOLERANCE // 0.067 at default
|
||||
// if KL > target × TOLERANCE → λ *= RAMP_RATE (Q signal not
|
||||
// landing, pull harder)
|
||||
// if KL < target / TOLERANCE → λ *= DECAY_RATE (aligned, ease off)
|
||||
// else → λ unchanged (dead-zone)
|
||||
//
|
||||
// if observed > upper → λ *= ADJUST_RATE (Q signal not landing,
|
||||
// increase pull)
|
||||
// if observed < lower → λ /= ADJUST_RATE (Q absorbed, can relax)
|
||||
// else → λ unchanged (dead-zone)
|
||||
//
|
||||
// Bounds: λ ∈ [MIN_LAMBDA=0.001, MAX_LAMBDA=1.0]. MIN preserves a
|
||||
// minimal Q-pull always-on; MAX caps the surrogate dominance (any
|
||||
// higher and distill swamps PPO).
|
||||
//
|
||||
// Bootstrap: on sentinel ema=0 or λ=0, defer adaptation (let static
|
||||
// seed take effect). Once both have non-zero values the controller
|
||||
// starts adjusting.
|
||||
// CRITICAL DESIGN: decay is 100× slower than ramp. This prevents the
|
||||
// self-defeating oscillation where λ ramps up → Q and π align → λ
|
||||
// collapses → they diverge again. The asymmetry means coupling
|
||||
// establishes quickly but persists for thousands of steps, giving
|
||||
// Q time to teach π meaningful action preferences.
|
||||
//
|
||||
// Per `feedback_no_atomicadd`: single-thread kernel, no atomics.
|
||||
// Per `feedback_cpu_is_read_only`: all state in ISV.
|
||||
@@ -34,10 +23,11 @@
|
||||
#define RL_Q_DISTILL_KL_EMA_INDEX 488
|
||||
#define RL_Q_DISTILL_KL_TARGET_INDEX 491
|
||||
|
||||
#define MIN_LAMBDA 0.001f
|
||||
#define MIN_LAMBDA 0.05f
|
||||
#define MAX_LAMBDA 1.0f
|
||||
#define KL_TOLERANCE 1.5f
|
||||
#define LAMBDA_ADJUST_RATE 1.2f
|
||||
#define KL_TOLERANCE 3.0f
|
||||
#define LAMBDA_RAMP_RATE 1.2f
|
||||
#define LAMBDA_DECAY_RATE 0.998f
|
||||
|
||||
extern "C" __global__ void rl_q_distill_lambda_controller(
|
||||
float* __restrict__ isv
|
||||
@@ -48,20 +38,16 @@ extern "C" __global__ void rl_q_distill_lambda_controller(
|
||||
const float kl_target = isv[RL_Q_DISTILL_KL_TARGET_INDEX];
|
||||
float lambda = isv[RL_Q_DISTILL_LAMBDA_INDEX];
|
||||
|
||||
// Defer if either input slot is sentinel-zero.
|
||||
if (kl_observed <= 0.0f || kl_target <= 0.0f || lambda <= 0.0f) return;
|
||||
|
||||
const float upper = kl_target * KL_TOLERANCE;
|
||||
const float lower = kl_target / KL_TOLERANCE;
|
||||
|
||||
if (kl_observed > upper) {
|
||||
// Q signal not landing strongly — raise λ to pull π harder.
|
||||
lambda = fminf(MAX_LAMBDA, lambda * LAMBDA_ADJUST_RATE);
|
||||
lambda = fminf(MAX_LAMBDA, lambda * LAMBDA_RAMP_RATE);
|
||||
} else if (kl_observed < lower) {
|
||||
// π already aligned with Q — can ease off (preserve PPO autonomy).
|
||||
lambda = fmaxf(MIN_LAMBDA, lambda / LAMBDA_ADJUST_RATE);
|
||||
lambda = fmaxf(MIN_LAMBDA, lambda * LAMBDA_DECAY_RATE);
|
||||
}
|
||||
// Dead-zone: kl_observed in [lower, upper] → no change.
|
||||
|
||||
isv[RL_Q_DISTILL_LAMBDA_INDEX] = lambda;
|
||||
}
|
||||
|
||||
@@ -1576,7 +1576,7 @@ impl IntegratedTrainer {
|
||||
// driven by `rl_q_distill_lambda_controller` from
|
||||
// KL_EMA vs KL_TARGET. Seed 0.05 is the starting point
|
||||
// before the controller picks up. τ=1.0 canonical.
|
||||
(crate::rl::isv_slots::RL_Q_DISTILL_LAMBDA_INDEX, 0.05),
|
||||
(crate::rl::isv_slots::RL_Q_DISTILL_LAMBDA_INDEX, 0.1),
|
||||
(crate::rl::isv_slots::RL_Q_DISTILL_TEMPERATURE_INDEX, 1.0),
|
||||
(crate::rl::isv_slots::RL_Q_DISTILL_KL_TARGET_INDEX, 0.1),
|
||||
// Reward-scale MIN floor (was hardcoded 1e-3) — now
|
||||
|
||||
Reference in New Issue
Block a user