foxhunt/crates/ml-alpha/cuda/rl_lr_controller.cu

// rl_lr_controller.cu — emits per-head learning rates to ISV[412..417]
// via ReduceLROnPlateau-style monotone decay.
//
// PRIOR DESIGN (grad-norm-driven multiplicative): observed
// closed-loop oscillation in cluster smoke `alpha-rl-tcr5r`. The
// grad-norm EMA tracked recent gradients which were themselves shaped
// by the LR being applied — classic two-loop instability when stacked
// on Adam (which already does per-parameter adaptive LR via its 2nd
// moment). Even with per-step rate caps the controller cycled
// MIN→MAX→MIN across the run, blowing up l_pi to 1.3e19.
//
// NEW DESIGN (plateau-based monotone decay):
//   1. Maintain a slow loss EMA per head (α = LR_LOSS_EMA_ALPHA, well
//      below the canonical Wiener 0.4 floor so the EMA smooths over
//      per-step loss noise).
//   2. Track `best_loss_ema` per head — lowest EMA value ever seen.
//   3. Each step: if current EMA improves on best by ≥
//      IMPROVEMENT_THRESHOLD ratio, update best + reset counter.
//      Else increment counter.
//   4. When counter exceeds PLATEAU_PATIENCE without improvement,
//      halve the LR (DECAY_FACTOR), reset counter, keep best
//      (preserves the absolute reference across decay events).
//   5. LR can ONLY decrease — never grows. Bottoms out at LR_MIN.
//
// Closed-loop oscillation is impossible by construction: monotone
// decay can't drive the LR up in response to its own induced
// gradient changes. Worst case: LR decays to MIN and stays there
// (interpretable as "model has stopped learning at any LR scale").
//
// State storage: 3 ISV slots per head (loss_ema, best, counter) —
// see `isv_slots.rs::RL_LR_Q_LOSS_EMA_INDEX` etc. Stored as f32
// (counter mantissa precision to 16M — well beyond any realistic
// patience setting).
//
// Per `pearl_controller_anchors_isv_driven` + `feedback_isv_for_adaptive_bounds`:
// every LR is ISV-resident and signal-driven (signal = loss trend).
//
// Per `pearl_first_observation_bootstrap`: sentinel-zero slot values
// trigger bootstrap paths — best=0 means "first loss observation;
// initialise best from it", loss_ema=0 means "first loss arrived;
// EMA seeds from it directly without blending against zero."
//
// Per `pearl_blend_formulas_must_have_permanent_floor`: LR has hard
// floor LR_MIN; counter has implicit floor of 0; best is
// monotonically improving (no floor needed).

#define RL_LR_BCE_INDEX 412
#define RL_LR_Q_INDEX   413
#define RL_LR_PI_INDEX  414
#define RL_LR_V_INDEX   415
#define RL_LR_AUX_INDEX 416

#define LR_BOOTSTRAP    1e-3f
#define LR_MIN          1e-5f
#define LR_MAX          1e-2f  // unused in monotone-decay design (kept for clamp arithmetic safety)

// Slow EMA α for the controller's internal loss tracking. Well below
// the Wiener α=0.4 floor used by ema_update_per_step — plateau
// detection needs smoothness, not responsiveness. At α=0.05, the EMA
// half-life is log(0.5)/log(0.95) ≈ 13.5 steps.
#define LR_LOSS_EMA_ALPHA  0.05f

// Plateau detection thresholds.
//   IMPROVEMENT_THRESHOLD: current EMA must be < best × this to count
//     as improvement. 0.99 = "must improve by at least 1%."
//   PLATEAU_PATIENCE: steps without improvement before LR decays.
//     1000 steps ≈ 7 seconds at 145 steps/sec on L40S.
//   DECAY_FACTOR: shrink LR by this on each plateau fire.
//     0.5 = "halve LR" (canonical ReduceLROnPlateau default).
#define IMPROVEMENT_THRESHOLD 0.99f
#define PLATEAU_PATIENCE      1000.0f
#define DECAY_FACTOR          0.5f

__device__ __forceinline__ void plateau_decay_head(
    float* isv,
    int    lr_idx,
    float  observed_loss,
    int    loss_ema_slot,
    int    best_slot,
    int    counter_slot
) {
    const float lr_prev = isv[lr_idx];

    // Bootstrap on sentinel 0.0 per pearl_first_observation_bootstrap.
    if (lr_prev == 0.0f) {
        isv[lr_idx] = LR_BOOTSTRAP;
        return;
    }

    // No observed loss this step (sentinel-zero observation, e.g. on
    // a step where the head's backward didn't fire). Hold LR.
    if (observed_loss == 0.0f) return;

    // Cold-start gate when loss_ema_slot is not yet wired (caller
    // passes -1 for BCE/AUX which are perception-owned).
    if (loss_ema_slot < 0) return;

    // ── 1. Update the slow loss EMA. ────────────────────────────────
    const float loss_ema_prev = isv[loss_ema_slot];
    float loss_ema_new;
    if (loss_ema_prev == 0.0f) {
        // First observation: replace directly per
        // pearl_first_observation_bootstrap (don't blend against the
        // sentinel zero).
        loss_ema_new = observed_loss;
    } else {
        loss_ema_new = (1.0f - LR_LOSS_EMA_ALPHA) * loss_ema_prev
                     + LR_LOSS_EMA_ALPHA * observed_loss;
    }
    isv[loss_ema_slot] = loss_ema_new;

    // ── 2. Compare to best — improvement or plateau? ────────────────
    const float best_prev = isv[best_slot];

    if (best_prev == 0.0f) {
        // First-observation bootstrap for best (and counter).
        isv[best_slot]    = loss_ema_new;
        isv[counter_slot] = 0.0f;
        return;
    }

    if (loss_ema_new < best_prev * IMPROVEMENT_THRESHOLD) {
        // Improvement — update best, reset counter.
        isv[best_slot]    = loss_ema_new;
        isv[counter_slot] = 0.0f;
        return;
    }

    // ── 3. Plateau — increment counter; maybe decay LR. ─────────────
    const float counter_next = isv[counter_slot] + 1.0f;
    if (counter_next >= PLATEAU_PATIENCE) {
        // Decay fires. Halve LR (clamped to MIN), reset counter,
        // KEEP best (decay events don't invalidate the historical
        // optimum — only further improvement does).
        const float lr_new = fmaxf(LR_MIN, lr_prev * DECAY_FACTOR);
        isv[lr_idx]       = lr_new;
        isv[counter_slot] = 0.0f;
    } else {
        isv[counter_slot] = counter_next;
    }
}

extern "C" __global__ void rl_lr_controller(
    float*       __restrict__ isv,
    float        observed_loss_bce,  // unused: perception-owned head
    float        observed_loss_q,
    float        observed_loss_pi,
    float        observed_loss_v,
    float        observed_loss_aux,  // unused: perception-owned head
    int          q_loss_ema_slot,
    int          q_best_slot,
    int          q_counter_slot,
    int          pi_loss_ema_slot,
    int          pi_best_slot,
    int          pi_counter_slot,
    int          v_loss_ema_slot,
    int          v_best_slot,
    int          v_counter_slot
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    // BCE and AUX heads are owned by the perception trainer; the RL
    // trainer doesn't compute per-step loss observations for them.
    // Pass loss_ema_slot=-1 to the plateau function which then
    // early-returns, holding lr_bce / lr_aux at LR_BOOTSTRAP.
    plateau_decay_head(isv, RL_LR_BCE_INDEX, observed_loss_bce, -1, -1, -1);
    plateau_decay_head(isv, RL_LR_Q_INDEX, observed_loss_q,
                       q_loss_ema_slot, q_best_slot, q_counter_slot);
    plateau_decay_head(isv, RL_LR_PI_INDEX, observed_loss_pi,
                       pi_loss_ema_slot, pi_best_slot, pi_counter_slot);
    plateau_decay_head(isv, RL_LR_V_INDEX, observed_loss_v,
                       v_loss_ema_slot, v_best_slot, v_counter_slot);
    plateau_decay_head(isv, RL_LR_AUX_INDEX, observed_loss_aux, -1, -1, -1);
}