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

// ISV-driven LR numerics per `feedback_isv_for_adaptive_bounds`.
// Seeded by rl_isv_write at trainer init; tunable at runtime.
#define RL_LR_BOOTSTRAP_INDEX        462
#define RL_LR_MIN_INDEX              463
#define RL_LR_MAX_INDEX              464  // unused in monotone-decay; defensive
#define RL_LR_LOSS_EMA_ALPHA_INDEX   465
#define RL_LR_DECAY_FACTOR_INDEX     466

// (LR_LOSS_EMA_ALPHA — was 0.05f #define — now read from
// isv[RL_LR_LOSS_EMA_ALPHA_INDEX] per kernel call.)

// 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).
//   LR_WARMUP_STEPS: number of EMA updates before `best` is locked in
//     and plateau detection turns on. Without this, a head whose
//     first non-zero loss observation is a cold-start artifact (e.g.
//     V loss ≈ machine-epsilon at step 0 before any reward has
//     flowed) pins `best` to an unreachable value, and the controller
//     decays on the patience clock alone. At α=0.05 the EMA
//     half-life is ≈14 steps; 500 updates leaves ≈35 half-lives,
//     which is well past convergence. Canonical incident:
//     alpha-rl-rzltn fold0 (commit 13d81dc5e) — V `best` stuck at
//     7.12e-10 across all 50k steps, only 2 unique best values
//     across the entire run.
// ISV-driven plateau-detection thresholds per
// `feedback_isv_for_adaptive_bounds`. Defaults match prior `#define`s
// but tunable at runtime via re-seeding (rl_isv_write). Adjusting
// these knobs lets us calibrate the plateau-decay aggressiveness to
// the trainer's noise floor without recompiling the kernel.
#define RL_IMPROVEMENT_THRESHOLD_INDEX 455
#define RL_PLATEAU_PATIENCE_INDEX      456
#define RL_LR_WARMUP_STEPS_INDEX       461
// (DECAY_FACTOR — was 0.5f #define — now read from
// isv[RL_LR_DECAY_FACTOR_INDEX] per kernel call.)

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

    // 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];
    const float loss_ema_alpha = isv[RL_LR_LOSS_EMA_ALPHA_INDEX];
    float loss_ema_new;
    if (loss_ema_prev == 0.0f) {
        loss_ema_new = observed_loss;
    } else {
        loss_ema_new = (1.0f - loss_ema_alpha) * loss_ema_prev
                     + loss_ema_alpha * observed_loss;
    }
    isv[loss_ema_slot] = loss_ema_new;

    // ── 2. Warmup: unconditionally overwrite best while EMA settles.
    //
    // During the first LR_WARMUP_STEPS observations, `best` tracks
    // the current loss_ema. Without this, sparse-signal heads (V
    // before any trade closes) lock `best` to a near-zero first
    // observation that subsequent EMA values can never beat.
    const float warmup_prev   = isv[warmup_slot];
    const float warmup_target = isv[RL_LR_WARMUP_STEPS_INDEX];
    if (warmup_prev < warmup_target) {
        isv[best_slot]    = loss_ema_new;
        isv[counter_slot] = 0.0f;
        isv[warmup_slot]  = warmup_prev + 1.0f;
        return;
    }

    // ── 3. Compare to best — improvement or plateau? ────────────────
    // ISV-driven thresholds (was hardcoded). Read once per call.
    const float improvement_threshold = isv[RL_IMPROVEMENT_THRESHOLD_INDEX];
    const float plateau_patience      = isv[RL_PLATEAU_PATIENCE_INDEX];
    const float best_prev = isv[best_slot];

    if (loss_ema_new < best_prev * improvement_threshold) {
        isv[best_slot]    = loss_ema_new;
        isv[counter_slot] = 0.0f;
        return;
    }

    // ── 4. Plateau — increment counter; maybe decay LR. ─────────────
    const float counter_next = isv[counter_slot] + 1.0f;
    if (counter_next >= plateau_patience) {
        const float lr_min       = isv[RL_LR_MIN_INDEX];
        const float decay_factor = isv[RL_LR_DECAY_FACTOR_INDEX];
        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          q_warmup_slot,
    int          pi_loss_ema_slot,
    int          pi_best_slot,
    int          pi_counter_slot,
    int          pi_warmup_slot,
    int          v_loss_ema_slot,
    int          v_best_slot,
    int          v_counter_slot,
    int          v_warmup_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, -1);
    plateau_decay_head(isv, RL_LR_Q_INDEX, observed_loss_q,
                       q_loss_ema_slot, q_best_slot, q_counter_slot, q_warmup_slot);
    plateau_decay_head(isv, RL_LR_PI_INDEX, observed_loss_pi,
                       pi_loss_ema_slot, pi_best_slot, pi_counter_slot, pi_warmup_slot);
    plateau_decay_head(isv, RL_LR_V_INDEX, observed_loss_v,
                       v_loss_ema_slot, v_best_slot, v_counter_slot, v_warmup_slot);
    plateau_decay_head(isv, RL_LR_AUX_INDEX, observed_loss_aux, -1, -1, -1, -1);
}