96abc364b5
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>
54 lines
2.1 KiB
Plaintext
54 lines
2.1 KiB
Plaintext
// rl_q_distill_lambda_controller.cu — adaptive λ_distill controller.
|
||
//
|
||
// Adapts the Q→π distillation strength (λ, slot 486) from the KL
|
||
// divergence between Q's Boltzmann policy and π's softmax (slot 488).
|
||
//
|
||
// Controller logic (asymmetric bounded step):
|
||
//
|
||
// 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)
|
||
//
|
||
// 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.
|
||
|
||
#define RL_Q_DISTILL_LAMBDA_INDEX 486
|
||
#define RL_Q_DISTILL_KL_EMA_INDEX 488
|
||
#define RL_Q_DISTILL_KL_TARGET_INDEX 491
|
||
|
||
#define MIN_LAMBDA 0.05f
|
||
#define MAX_LAMBDA 1.0f
|
||
#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
|
||
) {
|
||
if (threadIdx.x != 0 || blockIdx.x != 0) return;
|
||
|
||
const float kl_observed = isv[RL_Q_DISTILL_KL_EMA_INDEX];
|
||
const float kl_target = isv[RL_Q_DISTILL_KL_TARGET_INDEX];
|
||
float lambda = isv[RL_Q_DISTILL_LAMBDA_INDEX];
|
||
|
||
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) {
|
||
lambda = fminf(MAX_LAMBDA, lambda * LAMBDA_RAMP_RATE);
|
||
} else if (kl_observed < lower) {
|
||
lambda = fmaxf(MIN_LAMBDA, lambda * LAMBDA_DECAY_RATE);
|
||
}
|
||
|
||
isv[RL_Q_DISTILL_LAMBDA_INDEX] = lambda;
|
||
}
|