diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 4451adcd3..c29f0e379 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -45,7 +45,8 @@ use ml_alpha::cfc::snap_features::Mbp10RawInput; // emerging meta-pattern: any constant that mirrors a kernel-side // structural dimension MUST reference the Rust const, not duplicate // the literal value. -use ml_alpha::rl::common::N_ACTIONS; +use ml_alpha::rl::common::{N_ACTIONS, FRD_N_ATOMS, FRD_N_HORIZONS}; +use ml_alpha::rl::frd::FRD_OUT_DIM; use ml_alpha::data::loader::{ discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig, DEFAULT_OUTCOME_LABEL_COST_ES, @@ -587,6 +588,50 @@ fn main() -> Result<()> { let done_count: u32 = dones_host.iter().map(|&d| if d > 0.5 { 1 } else { 0 }).sum(); let isv = &trainer.isv_host; + + // SP20 P3 FRD diag — per-horizon softmax entropy + argmax index, + // averaged across the batch. Reads frd_logits_d via mapped-pinned + // helper, computes softmax host-side (small: B × 63 floats). + let frd_diag = { + let frd_logits = read_slice_d_pub( + dev_stream, + &trainer.frd_logits_d, + cli.n_backtests * FRD_OUT_DIM, + ) + .context("diag: read frd_logits_d")?; + let mut per_h_entropy = [0.0_f32; FRD_N_HORIZONS]; + let mut per_h_argmax_sum = [0.0_f32; FRD_N_HORIZONS]; + for b in 0..cli.n_backtests { + for h in 0..FRD_N_HORIZONS { + let off = b * FRD_OUT_DIM + h * FRD_N_ATOMS; + let row = &frd_logits[off..off + FRD_N_ATOMS]; + let max_l = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let denom: f32 = row.iter().map(|x| (x - max_l).exp()).sum(); + let mut entropy = 0.0_f32; + let mut argmax_idx = 0; + let mut argmax_v = f32::NEG_INFINITY; + for (a, x) in row.iter().enumerate() { + let p = (x - max_l).exp() / denom; + if p > 1e-9 { + entropy -= p * p.ln(); + } + if *x > argmax_v { + argmax_v = *x; + argmax_idx = a; + } + } + per_h_entropy[h] += entropy; + per_h_argmax_sum[h] += argmax_idx as f32; + } + } + let b_f = cli.n_backtests.max(1) as f32; + json!({ + "h1": { "entropy_mean": per_h_entropy[0] / b_f, "argmax_mean": per_h_argmax_sum[0] / b_f }, + "h2": { "entropy_mean": per_h_entropy[1] / b_f, "argmax_mean": per_h_argmax_sum[1] / b_f }, + "h3": { "entropy_mean": per_h_entropy[2] / b_f, "argmax_mean": per_h_argmax_sum[2] / b_f }, + }) + }; + let record = json!({ "step": step, "elapsed_s": t_start.elapsed().as_secs_f32(), @@ -856,6 +901,13 @@ fn main() -> Result<()> { "position": { "lots": position_lots_host, }, + // SP20 P3 FRD head diag — per-horizon softmax entropy + argmax-mode + // bucket index (averaged across the batch). At init, Xavier × 0.1 + // weights → logits ≈ 0 → entropy ≈ ln(21) = 3.044 and the argmax + // is dominated by tiny initial-weight noise. Watching these drift + // away from the uniform baseline is the supervised-learning signal + // for the FRD head (CE loss + bwd land in F.3). + "frd": frd_diag, }); writeln!(diag, "{}", record).context("diag: writeln jsonl")?; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index c37faba9a..de3b10590 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -631,6 +631,23 @@ pub struct IntegratedTrainer { /// `cfg.dqn_seed.wrapping_add(step_counter)`, deterministic given /// the same config + step sequence. step_counter: u64, + + /// SP20 P3 Forward-Return-Distribution head (forecast over 3 + /// horizons × 21 return-bucket atoms). Forward kernel runs every + /// `step_with_lobsim`; backward + label-supervised loss arrive in + /// F.3 / loader label generation. + pub frd_head: crate::rl::frd::FrdHead, + + /// SP20 P3 — cached post-ReLU hidden activation buffer for the FRD + /// head's backward pass (`B × FRD_HIDDEN_DIM`). Owned per-trainer + /// so we don't allocate per step. + pub frd_hidden_d: CudaSlice, + + /// SP20 P3 — FRD output logits (`B × FRD_OUT_DIM` where + /// FRD_OUT_DIM = FRD_N_HORIZONS × FRD_N_ATOMS = 63). Public so + /// callers can read for diag/inference; backward kernel writes + /// gradients into this buffer in F.3. + pub frd_logits_d: CudaSlice, } impl IntegratedTrainer { @@ -1154,6 +1171,22 @@ impl IntegratedTrainer { .alloc_zeros::(b_size) .context("alloc td_per_sample_d")?; + // SP20 P3 FRD head — construct + own per-step buffers. Seed + // distinct from dqn/ppo to keep Xavier draws independent. + let frd_head = crate::rl::frd::FrdHead::new( + dev, + crate::rl::frd::FrdHeadConfig { + seed: cfg.dqn_seed.wrapping_add(0xF8D), + }, + ) + .context("FrdHead::new")?; + let frd_hidden_d = stream + .alloc_zeros::(b_size * crate::rl::common::FRD_HIDDEN_DIM) + .context("alloc frd_hidden_d")?; + let frd_logits_d = stream + .alloc_zeros::(b_size * crate::rl::frd::FRD_OUT_DIM) + .context("alloc frd_logits_d")?; + Ok(Self { cfg, perception, @@ -1285,6 +1318,9 @@ impl IntegratedTrainer { last_v_loss: 0.0, last_k_updates: 0, step_counter: 0, + frd_head, + frd_hidden_d, + frd_logits_d, } .with_controllers_bootstrapped()?) } @@ -3089,6 +3125,20 @@ impl IntegratedTrainer { debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM); debug_assert_eq!(self.h_tp1_d.len(), b_size * HIDDEN_DIM); + // ── SP20 P3: FRD head forward — reads h_t, writes 3 horizons × + // 21 atom logits per batch into self.frd_logits_d (consumed by + // the FRD gate in P9 and the supervised CE loss in F.3). The + // post-ReLU hidden activations land in self.frd_hidden_d for + // the bwd kernel to consume without recomputing the W1 product. + self.frd_head + .forward( + h_t_borrow, + &mut self.frd_hidden_d, + &mut self.frd_logits_d, + b_size, + ) + .context("frd_head.forward in step_with_lobsim")?; + let k_dqn = N_ACTIONS * Q_N_ATOMS; let mut q_logits_d = self.stream.alloc_zeros::(b_size * k_dqn)?; let mut q_logits_tp1_d = self.stream.alloc_zeros::(b_size * k_dqn)?;