From a776fab31f6119adba3c527960fb82521a97d7f7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 24 May 2026 11:38:37 +0200 Subject: [PATCH] =?UTF-8?q?fix(rl-cli):=20build=20B=C3=97K=20snapshot=20te?= =?UTF-8?q?nsor=20per=20step=20at=20b=5Fsize>1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Crash in alpha-rl-ljn8k (commit 9c6c280bd) at step 0: forward_only: expected 512 snapshots (B=16 * K=32); got 32 Root cause: CLI binary called next_sequence_pair() once per step and passed the resulting K-snapshot window to step_with_lobsim. At b_size=1 the encoder's B×K=32 contract matched; at b_size=16 it silently expected B×K=512 and bailed. Fix: sample n_backtests INDEPENDENT pairs per step and concat into B×K row-major layout. This is the "proper" per-batch market diversity that delivers the gradient-variance reduction promised by `pearl_b_size_1_signal_starvation_blocks_q_learning` — tiling one window B times would be bit-identical encoder input across batch slots and contribute zero encoder gradient diversity. Bumps loader cap from n_steps×2 to n_steps×n_backtests×2 so the extra pair-sampling doesn't EOF mid-run. Eval loop fixed identically. Staging buffers preallocated outside the step loop to avoid per-step Vec alloc thrash at the hot path. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/examples/alpha_rl_train.rs | 103 ++++++++++++++------- 1 file changed, 72 insertions(+), 31 deletions(-) diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 89439c48b..96d1a859f 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -32,6 +32,7 @@ use anyhow::{Context, Result}; use clap::Parser; use data::providers::databento::dbn_parser::InstrumentFilter; +use ml_alpha::cfc::snap_features::Mbp10RawInput; use ml_alpha::data::loader::{ discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig, DEFAULT_OUTCOME_LABEL_COST_ES, @@ -353,10 +354,16 @@ fn main() -> Result<()> { predecoded_dir: cli.predecoded_dir.clone(), multi_resolution: multi_resolution.clone(), horizons: HORIZONS, - // Soft cap above n_steps × 2 so next_sequence_pair never - // returns Ok(None) mid-run. Loader's per-file random anchor - // sampling reuses files indefinitely up to this cap. - n_max_sequences: cli.n_steps.saturating_mul(2).max(64), + // Soft cap above n_steps × n_backtests × 2 so next_sequence_pair + // never returns Ok(None) mid-run. Each step samples n_backtests + // pairs (one per batch slot) to build the B×K snapshot tensor + // that forward_encoder consumes — per-batch market diversity is + // what gives the gradient-variance reduction promised by + // `pearl_b_size_1_signal_starvation_blocks_q_learning`. + n_max_sequences: cli.n_steps + .saturating_mul(cli.n_backtests.max(1)) + .saturating_mul(2) + .max(64), seed: cli.seed, inference_only: false, outcome_label_cost: cli.outcome_label_cost, @@ -423,23 +430,47 @@ fn main() -> Result<()> { let mut windowed_act_hist: [f32; 9] = [0.0; 9]; const WINDOWED_ACT_ALPHA: f32 = 1.0 / 1000.0; + // Per-step staging buffers for the B×K snapshot tensor that + // forward_encoder consumes. Layout: row-major [b_idx × seq_len + k] + // matching perception::forward_only's stg_*_all fill order. + // Pre-allocate once and clear+extend per step. + let n_batch = cli.n_backtests.max(1); + let bk = n_batch.saturating_mul(cli.seq_len); + let mut s_t_bk: Vec = Vec::with_capacity(bk); + let mut s_tp1_bk: Vec = Vec::with_capacity(bk); + let t_start = std::time::Instant::now(); - for step in 0..cli.n_steps { - let pair = loader - .next_sequence_pair() - .with_context(|| format!("next_sequence_pair at step {step}"))?; - let (s_t, s_tp1) = match pair { - Some(p) => p, - None => { - eprintln!( - "loader EOF at step {step} (n_max_sequences exhausted) — exiting early" - ); - break; - } - }; + 'outer: for step in 0..cli.n_steps { + // Sample n_backtests INDEPENDENT pairs per step. Each pair + // becomes one batch slot's K-window — distinct market segments + // per batch slot give true gradient-variance reduction (vs + // tiling one window B times, which is bit-identical encoder + // input and contributes no encoder gradient diversity). + s_t_bk.clear(); + s_tp1_bk.clear(); + for b_idx in 0..n_batch { + let pair = loader + .next_sequence_pair() + .with_context(|| format!("next_sequence_pair at step {step} batch {b_idx}"))?; + let (s_t, s_tp1) = match pair { + Some(p) => p, + None => { + eprintln!( + "loader EOF at step {step} batch {b_idx} (n_max_sequences exhausted) — exiting early" + ); + break 'outer; + } + }; + debug_assert_eq!(s_t.snapshots.len(), cli.seq_len); + debug_assert_eq!(s_tp1.snapshots.len(), cli.seq_len); + s_t_bk.extend_from_slice(&s_t.snapshots); + s_tp1_bk.extend_from_slice(&s_tp1.snapshots); + } + debug_assert_eq!(s_t_bk.len(), bk); + debug_assert_eq!(s_tp1_bk.len(), bk); let stats = trainer - .step_with_lobsim(&s_t.snapshots, &s_tp1.snapshots, &mut sim) + .step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim) .with_context(|| format!("step_with_lobsim at step {step}"))?; // Gate G8: NaN abort. Per `feedback_stop_on_anomaly` + @@ -824,7 +855,10 @@ fn main() -> Result<()> { predecoded_dir: cli.predecoded_dir.clone(), multi_resolution, horizons: HORIZONS, - n_max_sequences: cli.n_eval_steps.saturating_mul(2).max(64), + n_max_sequences: cli.n_eval_steps + .saturating_mul(cli.n_backtests.max(1)) + .saturating_mul(2) + .max(64), // Different seed for eval so sequence sampling is independent // of the train seed (deterministic given the eval seed). seed: cli.seed.wrapping_add(0xE7AE), @@ -835,19 +869,26 @@ fn main() -> Result<()> { let mut eval_loader = MultiHorizonLoader::new(&eval_loader_cfg) .context("MultiHorizonLoader::new (eval)")?; - for eval_step in 0..cli.n_eval_steps { - let pair = eval_loader - .next_sequence_pair() - .with_context(|| format!("eval next_sequence_pair at step {eval_step}"))?; - let (s_t, s_tp1) = match pair { - Some(p) => p, - None => { - eprintln!("eval loader EOF at step {eval_step}"); - break; - } - }; + // Reuse staging buffers from train phase (B×K snapshot layout). + 'eval_outer: for eval_step in 0..cli.n_eval_steps { + s_t_bk.clear(); + s_tp1_bk.clear(); + for b_idx in 0..n_batch { + let pair = eval_loader + .next_sequence_pair() + .with_context(|| format!("eval next_sequence_pair at step {eval_step} batch {b_idx}"))?; + let (s_t, s_tp1) = match pair { + Some(p) => p, + None => { + eprintln!("eval loader EOF at step {eval_step} batch {b_idx}"); + break 'eval_outer; + } + }; + s_t_bk.extend_from_slice(&s_t.snapshots); + s_tp1_bk.extend_from_slice(&s_tp1.snapshots); + } let stats = trainer - .step_with_lobsim(&s_t.snapshots, &s_tp1.snapshots, &mut sim) + .step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim) .with_context(|| format!("step_with_lobsim eval step {eval_step}"))?; if !stats.l_total.is_finite() { eprintln!("G8 NaN ABORT during eval at step {eval_step}");