diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index d04c99847..f10f2ac24 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -88,8 +88,8 @@ use ml_alpha::rl::isv_slots::{ RL_TD_KURT_STREAM_MEAN_INDEX, RL_V_GRAD_NORM_EMA_INDEX, }; use ml_alpha::trainer::integrated::{ - read_slice_d_pub, read_slice_i32_d_pub, IntegratedStepStats, IntegratedTrainer, - IntegratedTrainerConfig, + read_slice_d_pub, read_slice_i32_d_pub, write_slice_i32_d_pub, IntegratedStepStats, + IntegratedTrainer, IntegratedTrainerConfig, }; use ml_alpha::trainer::perception::PerceptionTrainerConfig; use ml_backtesting::sim::LobSimCuda; @@ -462,6 +462,10 @@ fn main() -> Result<()> { 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); + // SP20 P3 F.5 — per-step FRD label staging buffer. + // Layout: row-major [b_idx × FRD_N_HORIZONS + h] matching the + // trainer's frd_labels_d. Sentinel -1 for missing horizons. + let mut frd_labels_bh: Vec = vec![-1_i32; n_batch * FRD_N_HORIZONS]; let t_start = std::time::Instant::now(); 'outer: for step in 0..cli.n_steps { @@ -489,10 +493,31 @@ fn main() -> Result<()> { 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); + // SP20 P3 F.5 — FRD labels from s_t (the current-snapshot + // sequence's anchor). LabeledSequence::frd_labels has 1 + // entry per horizon (the newest snapshot's bucket index; + // see loader's slicing comment). Sentinel -1 if absent. + for h in 0..FRD_N_HORIZONS { + let row = s_t.frd_labels.get(h); + let v = row.and_then(|r| r.first().copied()).unwrap_or(-1); + frd_labels_bh[b_idx * FRD_N_HORIZONS + h] = v; + } } debug_assert_eq!(s_t_bk.len(), bk); debug_assert_eq!(s_tp1_bk.len(), bk); + // Push the per-step labels into the trainer's owned device + // buffer. The bwd chain in step_synthetic reads from + // self.frd_labels_d — sentinel -1 entries no-op the row, + // bucket-index entries activate supervised CE for that row. + let labels_stream = dev.cuda_stream().context("cuda_stream for frd labels")?; + write_slice_i32_d_pub( + labels_stream, + &frd_labels_bh, + &mut trainer.frd_labels_d, + ) + .context("write FRD labels to trainer.frd_labels_d")?; + let stats = trainer .step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim) .with_context(|| format!("step_with_lobsim at step {step}"))?; diff --git a/crates/ml-alpha/src/data/loader.rs b/crates/ml-alpha/src/data/loader.rs index 71a8bca7d..40be2f217 100644 --- a/crates/ml-alpha/src/data/loader.rs +++ b/crates/ml-alpha/src/data/loader.rs @@ -308,6 +308,15 @@ pub struct LabeledSequence { /// log-returns at each snapshot. NaN during the warmup window. Used /// downstream for σ-normalized size targets and any in-trainer scaling. pub sigma_k: [Vec; N_HORIZONS], + /// SP20 P3 F.5 — per-(horizon, snapshot) FRD return-bucket label + /// indices `[0, FRD_N_ATOMS)`. Sentinel `-1` at the rightmost edge + /// where `snapshot_idx + h_ticks >= snapshots.len()`. Length per + /// horizon equals `snapshots.len()`. Caller writes the row at + /// the anchor index into `trainer.frd_labels_d` before each + /// `step_with_lobsim`. See `compute_frd_labels` for the bucket + /// math; the trainer-side softmax+CE interpretation MUST match + /// (same `FRD_BUCKET_RANGE_SIGMA` const). + pub frd_labels: [Vec; crate::rl::common::FRD_N_HORIZONS], /// Per-file class imbalance for `pos_weight` BCE balancing. /// Length = `2 * N_HORIZONS`, laid out as /// `[long_h0, long_h1, ..., long_hN, short_h0, ..., short_hN]`. @@ -341,6 +350,10 @@ struct LoadedFile { /// `sigma_k_full[h]` rolling 1000-step Welford std of K-step log-returns. /// NaN during the per-horizon warmup window. sigma_k_full: [Vec; N_HORIZONS], + /// SP20 P3 F.5 — full-file FRD per-horizon bucket-index labels. + /// See `LabeledSequence::frd_labels` doc for shape + semantics. + /// Sliced per anchor on each `next_sequence()` call. + frd_labels_full: [Vec; crate::rl::common::FRD_N_HORIZONS], /// File-level positive-class fraction per (direction, horizon) for BCE /// pos_weight balancing. Length = `2 * N_HORIZONS`, laid out as /// `[long_h0..hN, short_h0..hN]`. Copied verbatim into every yielded @@ -562,6 +575,7 @@ impl MultiHorizonLoader { regime_full, ) }; + let frd_labels_full = compute_frd_labels(&snapshots); Ok(Some(LoadedFile { snapshots, labels_full, @@ -570,6 +584,7 @@ impl MultiHorizonLoader { outcome_size_long_full, outcome_size_short_full, sigma_k_full, + frd_labels_full, pos_fraction, regime_full, })) @@ -844,6 +859,26 @@ impl MultiHorizonLoader { let sequence: Vec = sequence_rev.into_iter().rev().collect(); debug_assert_eq!(sequence.len(), total); + // SP20 P3 F.5 — FRD labels for THIS sequence. The encoder + // forward emits h_t aligned to the newest snapshot (= the + // last position in `sequence` = file index `edge - 1`). + // That single per-horizon bucket index becomes the + // supervised target the trainer writes into frd_labels_d + // for the current batch slot. + let newest_idx = edge - 1; + let mut frd_labels: [Vec; crate::rl::common::FRD_N_HORIZONS] = Default::default(); + for h in 0..crate::rl::common::FRD_N_HORIZONS { + // Yield a length-1 vec per horizon (the trainer + // consumes labels[K-1] only — FRD supervises h_t at + // the rightmost K position, not the full K window). + let v = lf + .frd_labels_full + .get(h) + .and_then(|col| col.get(newest_idx).copied()) + .unwrap_or(-1); + frd_labels[h] = vec![v]; + } + Ok(LabeledSequence { snapshots: sequence, labels, @@ -852,6 +887,7 @@ impl MultiHorizonLoader { outcome_size_long, outcome_size_short, sigma_k, + frd_labels, pos_fraction: lf.pos_fraction.clone(), }) } @@ -864,6 +900,99 @@ fn mid_price_f32(s: &Mbp10Snapshot) -> f32 { (0.5 * (bid + ask)) as f32 } +/// SP20 P3 F.5 — per-file FRD label generation. +/// +/// For each snapshot index `i` and each FRD horizon `h_ticks ∈ +/// FRD_HORIZON_TICKS = [60, 300, 1800]`, compute the σ-normalized +/// forward return: +/// +/// ```text +/// r = (mid[i + h_ticks] - mid[i]) / sigma_per_step +/// ``` +/// +/// where `sigma_per_step` is the file-level sample std of single- +/// tick mid increments. Bucket `r` into `FRD_N_ATOMS = 21` atoms +/// uniform over `[-FRD_BUCKET_RANGE_SIGMA, +FRD_BUCKET_RANGE_SIGMA]` +/// (default ±3σ) — clip to `[0, FRD_N_ATOMS-1]` for tail returns. +/// +/// Sentinel `-1` marks `i + h_ticks >= snapshots.len()` (the +/// forward return isn't realized at the rightmost edge of the +/// file) — the `rl_frd_softmax_ce_grad` kernel zeros loss + grad +/// for those rows. +/// +/// Output shape: `[FRD_N_HORIZONS][snapshots.len()]` row-major +/// per horizon. The caller slices into this at the anchor index +/// when yielding a sequence (same pattern as the BCE labels). +/// +/// Per `pearl_glm_fitter_link_must_match_inference`: the bucket +/// edges here MUST match the trainer-side FRD head's softmax + CE +/// interpretation (atom `a` = bucket center `(a / (N-1) × 2 - 1) × +/// range_σ`). Both reference `FRD_BUCKET_RANGE_SIGMA` const so the +/// label and forward interpretations stay aligned. +fn compute_frd_labels(snapshots: &[Mbp10Snapshot]) -> [Vec; crate::rl::common::FRD_N_HORIZONS] { + use crate::rl::common::{FRD_BUCKET_RANGE_SIGMA, FRD_HORIZON_TICKS, FRD_N_ATOMS, FRD_N_HORIZONS}; + + let n = snapshots.len(); + let mids: Vec = snapshots.iter().map(mid_price_f32).collect(); + + // File-level sample std of single-tick mid increments. Bootstraps + // off the loaded file alone (no warmup blackout — for the very + // first file this is the only signal). At ~3M ticks per file the + // sample is statistically reliable. + let sigma_per_step: f32 = { + if mids.len() < 2 { + // Pathological 1-snapshot file — sentinel everything. + 1.0 + } else { + let mut sum_sq = 0.0_f64; + let mut cnt = 0_usize; + for i in 1..mids.len() { + let d = (mids[i] - mids[i - 1]) as f64; + if d.is_finite() { + sum_sq += d * d; + cnt += 1; + } + } + if cnt == 0 { + 1.0 + } else { + ((sum_sq / cnt as f64).sqrt() as f32).max(1e-6) + } + } + }; + + let mut out: [Vec; FRD_N_HORIZONS] = Default::default(); + let n_atoms_f = FRD_N_ATOMS as f32; + let half = (FRD_N_ATOMS - 1) as f32 / 2.0; + let scale = half / FRD_BUCKET_RANGE_SIGMA; + + for (h_idx, &h_ticks) in FRD_HORIZON_TICKS.iter().enumerate() { + // sqrt(h_ticks) Brownian scaling — variance of an h-step + // sum of iid increments is h × per-step variance, so std + // scales with sqrt(h). + let sigma_h = sigma_per_step * (h_ticks as f32).sqrt(); + let mut labels = vec![-1_i32; n]; + if h_ticks < n { + for i in 0..(n - h_ticks) { + let r = (mids[i + h_ticks] - mids[i]) / sigma_h.max(1e-6); + // Bucket center mapping: atom `a` → center `(a - half) / scale`. + // Inverse: `a = round(r × scale + half)`. + let a_raw = (r * scale + half).round() as i32; + let a_clamped = a_raw.clamp(0, (FRD_N_ATOMS - 1) as i32); + labels[i] = a_clamped; + } + } + // Final safety — guard against pathological NaN propagation. + for v in labels.iter_mut() { + if *v < -1 || *v >= n_atoms_f as i32 { + *v = -1; + } + } + out[h_idx] = labels; + } + out +} + fn convert( cur: &Mbp10Snapshot, prev: &Mbp10Snapshot, @@ -1015,6 +1144,95 @@ mod trade_flow_tests { } } +#[cfg(test)] +mod frd_label_tests { + use super::*; + use crate::rl::common::{FRD_HORIZON_TICKS, FRD_N_ATOMS, FRD_N_HORIZONS}; + + /// Construct a synthetic snapshot stream with a deterministic + /// mid-price function. Only L1 bid/ask matter — `compute_frd_labels` + /// reads `mid_price_f32` which uses only the first level. + fn snap_with_mid(mid: f64) -> Mbp10Snapshot { + let half_spread = 0.005; + let bid = mid - half_spread; + let ask = mid + half_spread; + let mut levels = vec![BidAskPair::empty(); BOOK_LEVELS]; + levels[0] = BidAskPair { + bid_px: BidAskPair::price_from_f64(bid), + bid_sz: 1, + bid_ct: 1, + ask_px: BidAskPair::price_from_f64(ask), + ask_sz: 1, + ask_ct: 1, + }; + Mbp10Snapshot::new("ES.FUT".into(), 0, levels, 0, 0) + } + + #[test] + fn frd_labels_flat_price_maps_to_mid_bucket() { + // Constant mid → all forward returns 0 → bucket = round(half) = 10. + let snaps: Vec = (0..2500).map(|_| snap_with_mid(100.0)).collect(); + let labels = compute_frd_labels(&snaps); + let mid_bucket = (FRD_N_ATOMS as i32 - 1) / 2; + for h in 0..FRD_N_HORIZONS { + let h_ticks = FRD_HORIZON_TICKS[h]; + // Non-sentinel range: [0, n - h_ticks) + for i in 0..(snaps.len() - h_ticks) { + assert_eq!( + labels[h][i], mid_bucket, + "flat-price horizon {} idx {} should be mid bucket {}; got {}", + h, i, mid_bucket, labels[h][i] + ); + } + // Sentinel range: [n - h_ticks, n) + for i in (snaps.len() - h_ticks)..snaps.len() { + assert_eq!( + labels[h][i], -1, + "right-edge sentinel at h={} idx={} should be -1; got {}", + h, i, labels[h][i] + ); + } + } + } + + #[test] + fn frd_labels_monotonic_ramp_lands_in_upper_buckets() { + // Monotone ramp: mid[i] = 100 + i × 0.01. Forward return at + // each horizon is positive → bucket > mid (= 10). Exact bucket + // depends on σ_per_step (= 0.01 here, std of constant + // increments), so a 60-tick forward return is +0.6 ticks. + // σ_h = 0.01 × sqrt(60) ≈ 0.0775 → r = 0.6 / 0.0775 ≈ 7.75. + // 7.75 is way > range_σ=3, so bucket clamps to FRD_N_ATOMS-1 = 20. + let snaps: Vec = + (0..2500).map(|i| snap_with_mid(100.0 + 0.01 * i as f64)).collect(); + let labels = compute_frd_labels(&snaps); + let max_bucket = (FRD_N_ATOMS - 1) as i32; + for h in 0..FRD_N_HORIZONS { + let h_ticks = FRD_HORIZON_TICKS[h]; + for i in 0..(snaps.len() - h_ticks) { + assert_eq!( + labels[h][i], max_bucket, + "monotonic-up horizon {} idx {} should clamp to max bucket {}; got {}", + h, i, max_bucket, labels[h][i] + ); + } + } + } + + #[test] + fn frd_labels_short_input_below_h_ticks_all_sentinel() { + // n=10 < h_ticks[0]=60 → every label sentinel for h0; same for + // h1 (=300) and h2 (=1800). + let snaps: Vec = (0..10).map(|_| snap_with_mid(100.0)).collect(); + let labels = compute_frd_labels(&snaps); + for h in 0..FRD_N_HORIZONS { + for (i, v) in labels[h].iter().enumerate() { + assert_eq!(*v, -1, "sub-h_ticks input must be all -1; got {} at h={} i={}", v, h, i); + } + } + } +} + #[cfg(test)] mod inference_mode_tests { use super::*; diff --git a/crates/ml-alpha/src/rl/common.rs b/crates/ml-alpha/src/rl/common.rs index 47a2b1458..6fbe85476 100644 --- a/crates/ml-alpha/src/rl/common.rs +++ b/crates/ml-alpha/src/rl/common.rs @@ -28,6 +28,18 @@ pub const FRD_N_HORIZONS: usize = 3; /// only structural compile-time dim per §0.1 audit. pub const FRD_N_ATOMS: usize = 21; +/// FRD head — forward-return horizons in tick units, matching the ISV +/// seeds at `RL_FRD_HORIZON_{1,2,3}_TICKS_INDEX` (slots 500/501/502). +/// Used by the loader's label-generation pass at file-load time, where +/// ISV isn't available. F.5+ caller is responsible for keeping these +/// in sync with the runtime ISV slots if the latter are overridden. +pub const FRD_HORIZON_TICKS: [usize; FRD_N_HORIZONS] = [60, 300, 1800]; + +/// FRD head — bucket-range default, matching the ISV seed at +/// `RL_FRD_BUCKET_RANGE_SIGMA_INDEX` (slot 503). See `FRD_HORIZON_TICKS` +/// comment on loader-side vs trainer-side ISV alignment. +pub const FRD_BUCKET_RANGE_SIGMA: f32 = 3.0; + /// C51 atom support `v_min`. After standardisation by the reward-scale /// controller (`RL_REWARD_SCALE_INDEX`, Phase F), per-step reward + /// γ-discounted returns are expected to fit inside `[V_MIN, V_MAX]`.