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

diag(rl): emit grad_norm_ema + lr_plateau state in alpha_rl_train JSONL

Browse Source 
Adds two new top-level keys to each diag.jsonl row:

  "grad_norm_ema": {q, pi, v}   — slots 424-426
  "lr_plateau": {q,pi,v} × {loss_ema, best, stale}  — slots 427-435

With these in place we can independently verify each plateau-decay
event in `mjgsj`'s diag (and all future runs):
  * `loss_ema` traces the controller's slow EMA of head loss
    (α=0.05); confirms the EMA actually moves and isn't stuck on the
    bootstrap zero
  * `best` shows the rolling minimum the controller compares against;
    confirms it improves early then plateaus
  * `stale` is the steps-since-best counter; should hit
    PLATEAU_PATIENCE = 1000 exactly when an LR halving fires; reset to
    0 after every decay event or every improvement

The `grad_norm_ema` block is kept because the grad-norm producers are
still wired (commit 383b1ad83) even though the LR controller no
longer consumes them — useful for correlating LR-decay events with
gradient-magnitude trajectory.

All R-phase gates green on local sm_86:
  G1 isv_bootstrap   ✅
  G3 controllers     ✅
  G4 target_update   ✅
  G6 r7d_per_wiring  ✅
  integrated_smoke   ✅

No new imports beyond the 9 new plateau-state slot constants + 3
grad-norm slot constants from `ml_alpha::rl::isv_slots`.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-23 19:19:53 +02:00
parent 042de99e67
commit 13d81dc5e6
1 changed files with 34 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
crates/ml-alpha/examples/alpha_rl_train.rs
View File

@@ -39,10 +39,15 @@ use ml_alpha::data::loader::{
use ml_alpha::heads::HORIZONS;
use ml_alpha::rl::isv_slots::{
RL_ADVANTAGE_VAR_RATIO_EMA_INDEX, RL_ENTROPY_COEF_INDEX, RL_ENTROPY_OBSERVED_EMA_INDEX,
RL_GAMMA_INDEX, RL_KL_PI_EMA_INDEX, RL_LR_AUX_INDEX, RL_LR_BCE_INDEX, RL_LR_PI_INDEX,
RL_LR_Q_INDEX, RL_LR_V_INDEX, RL_MEAN_ABS_PNL_EMA_INDEX, RL_MEAN_TRADE_DURATION_EMA_INDEX,
RL_N_ROLLOUT_STEPS_INDEX, RL_PER_ALPHA_INDEX, RL_PPO_CLIP_INDEX, RL_Q_DIVERGENCE_EMA_INDEX,
RL_REWARD_SCALE_INDEX, RL_TARGET_TAU_INDEX, RL_TD_KURTOSIS_EMA_INDEX,
RL_GAMMA_INDEX, RL_KL_PI_EMA_INDEX, RL_LR_AUX_INDEX, RL_LR_BCE_INDEX, RL_LR_PI_BEST_LOSS_INDEX,
RL_LR_PI_INDEX, RL_LR_PI_LOSS_EMA_INDEX, RL_LR_PI_STEPS_SINCE_BEST_INDEX,
RL_LR_Q_BEST_LOSS_INDEX, RL_LR_Q_INDEX, RL_LR_Q_LOSS_EMA_INDEX,
RL_LR_Q_STEPS_SINCE_BEST_INDEX, RL_LR_V_BEST_LOSS_INDEX, RL_LR_V_INDEX,
RL_LR_V_LOSS_EMA_INDEX, RL_LR_V_STEPS_SINCE_BEST_INDEX, RL_MEAN_ABS_PNL_EMA_INDEX,
RL_MEAN_TRADE_DURATION_EMA_INDEX, RL_N_ROLLOUT_STEPS_INDEX, RL_PER_ALPHA_INDEX,
RL_PI_GRAD_NORM_EMA_INDEX, RL_PPO_CLIP_INDEX, RL_Q_DIVERGENCE_EMA_INDEX,
RL_Q_GRAD_NORM_EMA_INDEX, RL_REWARD_SCALE_INDEX, RL_TARGET_TAU_INDEX, RL_TD_KURTOSIS_EMA_INDEX,
RL_V_GRAD_NORM_EMA_INDEX,
};
use ml_alpha::trainer::integrated::{
read_slice_d_pub, read_slice_i32_d_pub, IntegratedStepStats, IntegratedTrainer,
@@ -506,6 +511,31 @@ fn main() -> Result<()> {
"td_kurtosis": isv[RL_TD_KURTOSIS_EMA_INDEX],
"mean_abs_pnl": isv[RL_MEAN_ABS_PNL_EMA_INDEX],
},
// Grad-norm EMAs (diagnostic — no longer drive the LR
// controller as of commit 042de99e6 but still useful to
// correlate with plateau-decay events).
"grad_norm_ema": {
"q": isv[RL_Q_GRAD_NORM_EMA_INDEX],
"pi": isv[RL_PI_GRAD_NORM_EMA_INDEX],
"v": isv[RL_V_GRAD_NORM_EMA_INDEX],
},
// Plateau-decay LR controller state (per head). loss_ema is
// the controller's internal slow EMA at α=0.05; best is the
// lowest loss_ema observed so far; steps_since_best is the
// staleness counter that fires LR halving at PLATEAU_PATIENCE
// (1000). Confirms whether the controller is seeing plateaus
// and decaying for the right reason.
"lr_plateau": {
"q": { "loss_ema": isv[RL_LR_Q_LOSS_EMA_INDEX],
"best": isv[RL_LR_Q_BEST_LOSS_INDEX],
"stale": isv[RL_LR_Q_STEPS_SINCE_BEST_INDEX] },
"pi": { "loss_ema": isv[RL_LR_PI_LOSS_EMA_INDEX],
"best": isv[RL_LR_PI_BEST_LOSS_INDEX],
"stale": isv[RL_LR_PI_STEPS_SINCE_BEST_INDEX] },
"v": { "loss_ema": isv[RL_LR_V_LOSS_EMA_INDEX],
"best": isv[RL_LR_V_BEST_LOSS_INDEX],
"stale": isv[RL_LR_V_STEPS_SINCE_BEST_INDEX] },
},
"replay_len": trainer.replay.len(),
"rewards": {
"sum": reward_sum,