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:
jgrusewski
2026-05-25 11:54:17 +02:00
parent 0251656be4
commit 96abc364b5
2 changed files with 19 additions and 33 deletions

View File

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

View File

@@ -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