gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rl): rewrite LR controller as monotone plateau decay (no oscillation)

Browse Source 
Cluster smoke `alpha-rl-tcr5r` confirmed that the grad-norm-driven
multiplicative LR controller — even with the per-step rate cap +
per-head TARGET_GRAD_NORM fixes — could not avoid closed-loop
oscillation when stacked on Adam:

  step 5000 : ALL lrs at MAX (1e-2)
  step 15000: ALL lrs at MIN (1e-5)
  step 25000: lr_pi MAX again
  ...

Result: l_pi max = 1.28e19, l_v max = 701k, l_total mean = 1.15e14.
Q-head benefited (l_q mean 3.24) but π and V destabilised
catastrophically.

## Why the prior design was fundamentally broken

The grad-norm signal that drives the LR controller is itself
*produced* by the LR being applied (via Adam → weights → grads →
norms). When the LR controller reduces lr_pi because grad-norm
spiked, the next-step grad-norm shrinks → controller raises LR →
grad-norm spikes again. Classic two-loop instability when stacked
on Adam (which already does per-parameter LR adaptation via its 2nd
moment). No amount of per-step rate capping breaks the cycle; it
just slows it.

## New design: ReduceLROnPlateau-style monotone decay

The controller now:
  1. Maintains a SLOW loss EMA per head (α = 0.05, half-life ≈ 13
     steps — well below the canonical Wiener 0.4 floor used by the
     per-step EMAs because plateau detection needs smoothness, not
     responsiveness).
  2. Tracks `best_loss_ema` per head — lowest EMA value ever seen.
  3. Per step: if current EMA improves on best by ≥ 1%
     (IMPROVEMENT_THRESHOLD = 0.99), update best + reset counter.
     Otherwise increment counter.
  4. When counter exceeds PLATEAU_PATIENCE (1000 steps ≈ 7 sec at
     145 steps/sec), halve LR (DECAY_FACTOR = 0.5), reset counter,
     keep best.
  5. LR can ONLY decrease — never grows. Bottoms out at LR_MIN = 1e-5.

Closed-loop oscillation is impossible by construction: monotone
decay can't drive 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" — meaningful signal,
not a control failure).

## State storage

9 new ISV slots (3 per head — Q, π, V):
  * RL_LR_Q_LOSS_EMA_INDEX           = 427
  * RL_LR_Q_BEST_LOSS_INDEX          = 428
  * RL_LR_Q_STEPS_SINCE_BEST_INDEX   = 429
  * RL_LR_PI_LOSS_EMA_INDEX          = 430
  * RL_LR_PI_BEST_LOSS_INDEX         = 431
  * RL_LR_PI_STEPS_SINCE_BEST_INDEX  = 432
  * RL_LR_V_LOSS_EMA_INDEX           = 433
  * RL_LR_V_BEST_LOSS_INDEX          = 434
  * RL_LR_V_STEPS_SINCE_BEST_INDEX   = 435
  * RL_SLOTS_END                     = 436 (was 427)

Counters stored as f32 — mantissa precision to 16M is well beyond
any plausible patience threshold.

## Kernel signature change

```cuda
extern "C" __global__ void rl_lr_controller(
    float* isv,
    float observed_loss_bce,   // unused (perception-owned)
    float observed_loss_q,     // host scalar from prior step's Q backward
    float observed_loss_pi,    // host scalar from prior step's PPO surrogate
    float observed_loss_v,     // host scalar from prior step's V backward
    float observed_loss_aux,   // unused
    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
);
```

Grad-norm EMA producers (commit 383b1ad83) remain wired — they're
still useful diagnostics in the JSONL, just not consumed by the
LR controller anymore.

## Trainer wiring

New trainer fields `last_q_loss` + `last_v_loss` mirror per-step
loss scalars (same pattern as the existing `last_pi_loss` from
PPO surrogate forward). Populated at the end of step_synthetic's
backward chain; consumed at the start of NEXT step_synthetic's
launch_rl_lr_controller call. One-step lag is acceptable —
plateau detection operates on 1000-step windows so a 1-step shift
in observations is negligible.

## Verified gates (local sm_86)

  G1, G3, G4, G6, smoke: all ✅

## Expected effect on next 50k smoke

  * lr_q starts at 1e-3, decays monotonically toward 1e-5 if l_q
    plateaus.
  * lr_pi same — but π loss is much noisier, so plateau detection
    may fire more often → faster decay.
  * lr_v starts at 1e-3, decays as l_v approaches its asymptote
    (V regression of ~0 for sparse rewards).
  * NO l_pi explosions (controller can't drive LR up).
  * Final losses should be similar to or better than fixed-LR's
    baseline (mean 2.97 for l_q on `nqd68`; this design's monotone
    decay should produce stable equilibrium at some LR ≤ 1e-3).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-23 18:51:37 +02:00
parent e074c91fb2
commit 042de99e67
3 changed files with 239 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files

261
crates/ml-alpha/cuda/rl_lr_controller.cu
View File

@@ -1,49 +1,48 @@
// rl_lr_controller.cu — emits per-head learning rates to ISV[412..417]
// with signal-modulated targets (per-head gradient-norm EMA divergence
// from a target magnitude).
// 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 learning rate is ISV-resident and driven by an observable
// signal (grad-norm EMA) rather than a hardcoded constant. The
// trainer:
// 1. Computes per-head L2 grad-norm (rl_l2_norm) after each backward.
// 2. Feeds it through ema_update_per_step into the relevant grad-norm
// EMA slot (RL_Q_GRAD_NORM_EMA_INDEX, RL_PI_GRAD_NORM_EMA_INDEX,
// RL_V_GRAD_NORM_EMA_INDEX).
// 3. Launches this kernel passing those slot indices.
// This kernel reads each EMA and derives a per-head LR target:
// every LR is ISV-resident and signal-driven (signal = loss trend).
//
// target_lr = lr_prev × (TARGET_GRAD_NORM / max(observed_grad_norm, ε))
// 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."
//
// Multiplicative pattern (same shape as rl_target_tau / rl_ppo_clip):
// observed > target → ratio < 1 → shrink lr (calm the gradients);
// observed < target → ratio > 1 → grow lr (push more aggressive
// updates). Clamped to [LR_MIN, LR_MAX].
//
// Per `pearl_first_observation_bootstrap`: sentinel zero at the slot's
// address means "uninitialised — emit the bootstrap value on the next
// controller fire." Cold-start gate: when the input EMA is sentinel
// zero (no backward has fired yet), the controller holds prev at
// bootstrap — same anti-pattern fix applied to τ/ε/n_roll/scale.
// Replace-directly on first warm observation: when prev is exactly
// the hardcoded bootstrap value, write target directly instead of
// Wiener-blending (60% bootstrap contamination would otherwise
// distort the first emit).
//
// Per `pearl_wiener_alpha_floor_for_nonstationary`: the α blend is
// floored at 0.4 since the target drives a co-adapting closed loop
// (parameter LR feeds back into gradient magnitude via Adam).
//
// BCE and AUX heads are emitted at their hardcoded bootstrap. Passing
// a `signal_slot < 0` for these heads skips the signal-driven path
// and falls back to `target_lr = LR_BOOTSTRAP` — the perception
// trainer owns those heads, not the RL trainer, so signal wiring is
// out of scope here.
//
// Block layout:
// grid = (1, 1, 1)
// block = (1, 1, 1)
// Single-thread kernel — five ISV slots, one writer each, no contention.
// 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
@@ -53,44 +52,32 @@
#define LR_BOOTSTRAP 1e-3f
#define LR_MIN 1e-5f
#define LR_MAX 1e-2f
#define LR_MAX 1e-2f // unused in monotone-decay design (kept for clamp arithmetic safety)
// Per-head TARGET_GRAD_NORM anchor — the "well-tuned" grad-norm
// magnitude the controller drives toward. Heuristic: scaled to
// √n_params so per-parameter grad magnitude ≈ 1e-2 (Adam-friendly).
//
// Q head w_d: N_ACTIONS × Q_N_ATOMS × HIDDEN_DIM = 9×21×128 = 24,192 params
// √24192 ≈ 156 → target 1.5
// π head w_d: N_ACTIONS × HIDDEN_DIM = 9×128 = 1,152 params
// √1152 ≈ 34 → target 0.3
// V head w_d: HIDDEN_DIM = 128 params
// √128 ≈ 11 → target 0.1
//
// BCE and AUX are owned by the perception trainer; passing
// signal_slot < 0 from the RL trainer's launch bypasses the
// signal-driven path entirely, so the target here is unused.
#define Q_TARGET_GRAD_NORM 1.5f
#define PI_TARGET_GRAD_NORM 0.3f
#define V_TARGET_GRAD_NORM 0.1f
#define BCE_TARGET_GRAD_NORM 1.0f // unused (perception-owned)
#define AUX_TARGET_GRAD_NORM 1.0f // unused
// 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
// Per-step rate-of-change cap on target_lr (cluster smoke
// `alpha-rl-nqd68` showed π controller swinging MIN→MAX 1000× across
// ~10k steps, then getting stuck at MAX after Adam updates wrecked
// the policy weights). With this cap, the per-step target is bounded
// to [lr_prev × 0.5, lr_prev × 2.0] — at most a 2× swing per step,
// so traversing the full [LR_MIN, LR_MAX] range takes ~10 steps each
// direction. That gives downstream gradient signal time to react
// before LR catastrophically overshoots. Same general pattern as
// rl_rollout_steps_controller's `scale ∈ [0.5, 2.0]` cap.
#define LR_RATE_CAP_LO 0.5f
#define LR_RATE_CAP_HI 2.0f
// 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
#define EMA_EPS 1e-6f
__device__ __forceinline__ void update_lr_with_signal(
float* isv, int lr_idx, float alpha, int signal_slot, float head_target_grad_norm
__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];
@@ -100,68 +87,88 @@ __device__ __forceinline__ void update_lr_with_signal(
return;
}
// No signal wiring for this head (BCE / AUX caller-pass -1) —
// hold at bootstrap, Wiener-blend toward LR_BOOTSTRAP (no-op since
// prev is already there in steady state, but the blend gives the
// controller authority to override an externally-edited slot).
if (signal_slot < 0) {
const float a = fmaxf(alpha, 0.4f);
isv[lr_idx] = (1.0f - a) * lr_prev + a * LR_BOOTSTRAP;
// 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;
}
const float observed = isv[signal_slot];
// Cold-start gate: no grad-norm signal observed yet → hold at
// bootstrap. Same fix that closed the cold-start migration in
// τ/ε/n_roll/scale.
if (observed == 0.0f) return;
// Multiplicative target derivation using per-head target anchor.
float target_lr = lr_prev * (head_target_grad_norm / fmaxf(observed, EMA_EPS));
// Per-step rate-of-change cap: target bounded to
// [lr_prev × LR_RATE_CAP_LO, lr_prev × LR_RATE_CAP_HI]. Prevents
// the runaway-swing instability that destabilised the π head in
// alpha-rl-nqd68 (target = lr_prev × 1000 when grad-norm was
// tiny → controller jumped to LR_MAX in a few Wiener blends).
const float lr_lower = lr_prev * LR_RATE_CAP_LO;
const float lr_upper = lr_prev * LR_RATE_CAP_HI;
target_lr = fmaxf(lr_lower, fminf(lr_upper, target_lr));
// Absolute clamp to [LR_MIN, LR_MAX] (runaway protection).
target_lr = fminf(LR_MAX, fmaxf(LR_MIN, target_lr));
// First-observation replace-directly: avoid Wiener-blend's 60%
// bootstrap contamination on the first warm step. Rate cap
// already applied above so the "direct" target respects the
// ≤2× rate constraint from the bootstrap LR.
if (lr_prev == LR_BOOTSTRAP) {
isv[lr_idx] = target_lr;
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;
}
// Wiener-α floor 0.4 per pearl_wiener_alpha_floor_for_nonstationary.
const float a = fmaxf(alpha, 0.4f);
float lr_new = (1.0f - a) * lr_prev + a * target_lr;
lr_new = fminf(LR_MAX, fmaxf(LR_MIN, lr_new));
isv[lr_idx] = lr_new;
// ── 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 alpha,
int bce_signal_slot, // <0 → no signal (BCE owned by perception)
int q_signal_slot, // RL_Q_GRAD_NORM_EMA_INDEX
int pi_signal_slot, // RL_PI_GRAD_NORM_EMA_INDEX
int v_signal_slot, // RL_V_GRAD_NORM_EMA_INDEX
int aux_signal_slot // <0 → no signal (AUX owned by perception)
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;
update_lr_with_signal(isv, RL_LR_BCE_INDEX, alpha, bce_signal_slot, BCE_TARGET_GRAD_NORM);
update_lr_with_signal(isv, RL_LR_Q_INDEX, alpha, q_signal_slot, Q_TARGET_GRAD_NORM);
update_lr_with_signal(isv, RL_LR_PI_INDEX, alpha, pi_signal_slot, PI_TARGET_GRAD_NORM);
update_lr_with_signal(isv, RL_LR_V_INDEX, alpha, v_signal_slot, V_TARGET_GRAD_NORM);
update_lr_with_signal(isv, RL_LR_AUX_INDEX, alpha, aux_signal_slot, AUX_TARGET_GRAD_NORM);
// 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);
}

31
crates/ml-alpha/src/rl/isv_slots.rs
View File

@@ -146,6 +146,35 @@ pub const RL_PI_GRAD_NORM_EMA_INDEX: usize = 425;
/// to `rl_lr_controller`'s V-LR target derivation.
pub const RL_V_GRAD_NORM_EMA_INDEX: usize = 426;
// Plateau-based LR decay state slots (ReduceLROnPlateau-style).
// Each head (Q, π, V) gets 3 slots: internal slow loss EMA, best
// observed loss so far, steps-since-best counter. Monotone decay
// design: lr never INCREASES (closed-loop oscillation impossible by
// construction); decay fires when counter exceeds PLATEAU_PATIENCE
// without 1% improvement. Counter reset on every improvement or
// decay fire; best updates only on improvement.
/// Slow internal loss EMA for the Q head's plateau detector. The
/// controller maintains this with α much lower than the standard
/// 0.4 Wiener floor (LR_LOSS_EMA_ALPHA = 0.05 in the kernel) so the
/// EMA smooths over per-step Q-CE noise.
pub const RL_LR_Q_LOSS_EMA_INDEX: usize = 427;
/// Best (lowest) Q loss EMA observed so far. Updates on improvement;
/// never resets (preserves the best across decay events).
pub const RL_LR_Q_BEST_LOSS_INDEX: usize = 428;
/// Steps-since-best counter for Q. Stored as f32 (mantissa precision
/// to 16M — well beyond any realistic patience). Resets on
/// improvement OR on decay fire.
pub const RL_LR_Q_STEPS_SINCE_BEST_INDEX: usize = 429;
pub const RL_LR_PI_LOSS_EMA_INDEX: usize = 430;
pub const RL_LR_PI_BEST_LOSS_INDEX: usize = 431;
pub const RL_LR_PI_STEPS_SINCE_BEST_INDEX: usize = 432;
pub const RL_LR_V_LOSS_EMA_INDEX: usize = 433;
pub const RL_LR_V_BEST_LOSS_INDEX: usize = 434;
pub const RL_LR_V_STEPS_SINCE_BEST_INDEX: usize = 435;
/// Last RL-allocated slot index (exclusive). The integrated trainer
/// extends `ISV_TOTAL_DIM` to at least this value at trainer init time.
pub const RL_SLOTS_END: usize = 427;
pub const RL_SLOTS_END: usize = 436;

99
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -508,6 +508,14 @@ pub struct IntegratedTrainer {
/// surrogate kernel's atomic-write loss readback through the
/// borrow-checker dance in `step_synthetic`. NOT a public API.
last_pi_loss: f32,
/// Last step's Q-head loss (Q CE, host scalar) — fed to the LR
/// controller at the start of NEXT step_synthetic for plateau
/// detection. Zero at construction; populated at the end of
/// step_synthetic after the Q backward's scalar readback.
last_q_loss: f32,
/// Last step's V-head loss (MSE, host scalar) — same pattern as
/// `last_q_loss`. Zero at construction.
last_v_loss: f32,
/// Host-side step counter for Thompson-sampling RNG salt (Phase E.3b).
/// Increments once per `step_with_lobsim` call so consecutive calls
@@ -1013,6 +1021,8 @@ impl IntegratedTrainer {
td_per_sample_d,
grad_h_t_combined_d,
last_pi_loss: 0.0,
last_q_loss: 0.0,
last_v_loss: 0.0,
step_counter: 0,
}
.with_controllers_bootstrapped()?)
@@ -1538,8 +1548,16 @@ impl IntegratedTrainer {
// ── Step 2: per-head LR controller emit + ISV host mirror ────
// The controller emits the bootstrap target on first call
// (sentinel-zero → 1e-3); subsequent calls Wiener-α blend.
self.launch_rl_lr_controller()
.context("rl_lr_controller launch")?;
// Use last step's losses as plateau-decay observations. At
// construction these are 0.0 (sentinel) and the controller's
// cold-start gate (early-return on observed_loss == 0) holds
// LR at LR_BOOTSTRAP for the first step.
self.launch_rl_lr_controller(
self.last_q_loss,
self.last_pi_loss,
self.last_v_loss,
)
.context("rl_lr_controller launch")?;
// Mapped-pinned ISV mirror refresh per
// `feedback_no_htod_htoh_only_mapped_pinned` — was a raw
// `stream.memcpy_dtoh` (forbidden by
@@ -1752,6 +1770,8 @@ impl IntegratedTrainer {
)
.context("dqn_head.backward_logits (sampled_actions) [R7d off-policy]")?;
let l_q_host = read_scalar_d(&self.stream, &q_loss_d_mut)?;
// Mirror for next step's LR controller plateau detection.
self.last_q_loss = l_q_host / (b_size as f32);
self.dqn_head
.backward_to_w_b_h(
@@ -1828,6 +1848,8 @@ impl IntegratedTrainer {
self.launch_reduce_axis0(&v_loss_per_batch_d, b_size, 1, &mut v_loss_sum_d)?;
let l_v_sum_host = read_scalar_d(&self.stream, &v_loss_sum_d)?;
let l_v_host = l_v_sum_host / (b_size as f32);
// Mirror for next step's LR controller plateau detection.
self.last_v_loss = l_v_host;
// ── Step 9: Adam updates on each head's w and b ──────────────
self.dqn_w_adam
@@ -2914,36 +2936,65 @@ impl IntegratedTrainer {
}
/// Launch `rl_lr_controller` to emit per-head learning rates into
/// `ISV[412..417]`. Reads per-head grad-norm EMAs from
/// `ISV[RL_Q_GRAD_NORM_EMA_INDEX]` etc. and derives multiplicative
/// targets `lr_prev × (TARGET_GRAD_NORM / observed)`. BCE and AUX
/// heads pass sentinel `-1` for their signal slots (those heads
/// are owned by the perception trainer, not the RL trainer; their
/// LRs hold at LR_BOOTSTRAP).
fn launch_rl_lr_controller(&self) -> Result<()> {
/// `ISV[412..417]` via ReduceLROnPlateau-style monotone decay.
/// Per-head observed loss (l_q / l_pi / l_v from step_synthetic's
/// stats this step) feeds the controller's slow loss EMA; the
/// controller halves LR when 1000 consecutive steps elapse
/// without a 1% improvement. BCE / AUX use placeholder zero
/// losses (the plateau function early-returns for those because
/// loss_ema_slot = -1).
///
/// Loss observations are passed as scalar floats (host-side
/// values from the trainer's last step). Per-head state lives
/// in ISV slots 427..436 (loss_ema, best, counter ×3).
fn launch_rl_lr_controller(
&self,
observed_loss_q: f32,
observed_loss_pi: f32,
observed_loss_v: f32,
) -> Result<()> {
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
let alpha = RL_LR_CONTROLLER_ALPHA;
let bce_signal_slot: i32 = -1; // owned by perception
let q_signal_slot: i32 =
crate::rl::isv_slots::RL_Q_GRAD_NORM_EMA_INDEX as i32;
let pi_signal_slot: i32 =
crate::rl::isv_slots::RL_PI_GRAD_NORM_EMA_INDEX as i32;
let v_signal_slot: i32 =
crate::rl::isv_slots::RL_V_GRAD_NORM_EMA_INDEX as i32;
let aux_signal_slot: i32 = -1; // owned by perception
let observed_loss_bce: f32 = 0.0;
let observed_loss_aux: f32 = 0.0;
let q_loss_ema_slot: i32 =
crate::rl::isv_slots::RL_LR_Q_LOSS_EMA_INDEX as i32;
let q_best_slot: i32 =
crate::rl::isv_slots::RL_LR_Q_BEST_LOSS_INDEX as i32;
let q_counter_slot: i32 =
crate::rl::isv_slots::RL_LR_Q_STEPS_SINCE_BEST_INDEX as i32;
let pi_loss_ema_slot: i32 =
crate::rl::isv_slots::RL_LR_PI_LOSS_EMA_INDEX as i32;
let pi_best_slot: i32 =
crate::rl::isv_slots::RL_LR_PI_BEST_LOSS_INDEX as i32;
let pi_counter_slot: i32 =
crate::rl::isv_slots::RL_LR_PI_STEPS_SINCE_BEST_INDEX as i32;
let v_loss_ema_slot: i32 =
crate::rl::isv_slots::RL_LR_V_LOSS_EMA_INDEX as i32;
let v_best_slot: i32 =
crate::rl::isv_slots::RL_LR_V_BEST_LOSS_INDEX as i32;
let v_counter_slot: i32 =
crate::rl::isv_slots::RL_LR_V_STEPS_SINCE_BEST_INDEX as i32;
let mut launch = self.stream.launch_builder(&self.rl_lr_controller_fn);
launch
.arg(&self.isv_d)
.arg(&alpha)
.arg(&bce_signal_slot)
.arg(&q_signal_slot)
.arg(&pi_signal_slot)
.arg(&v_signal_slot)
.arg(&aux_signal_slot);
.arg(&observed_loss_bce)
.arg(&observed_loss_q)
.arg(&observed_loss_pi)
.arg(&observed_loss_v)
.arg(&observed_loss_aux)
.arg(&q_loss_ema_slot)
.arg(&q_best_slot)
.arg(&q_counter_slot)
.arg(&pi_loss_ema_slot)
.arg(&pi_best_slot)
.arg(&pi_counter_slot)
.arg(&v_loss_ema_slot)
.arg(&v_best_slot)
.arg(&v_counter_slot);
unsafe {
launch.launch(cfg).context("rl_lr_controller launch")?;
}