diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index c29f0e379..d04c99847 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -641,6 +641,7 @@ fn main() -> Result<()> { "pi": stats.l_pi, "v": stats.l_v, "aux": stats.l_aux, + "frd": stats.l_frd, "total": stats.l_total, }, "lambdas": { @@ -649,6 +650,7 @@ fn main() -> Result<()> { "pi": stats.lambdas.pi, "v": stats.lambdas.v, "aux": stats.lambdas.aux, + "frd": stats.lambdas.frd, }, // 7 R5 controller outputs (the adaptive knobs). "isv_out": { diff --git a/crates/ml-alpha/src/rl/loss_balance.rs b/crates/ml-alpha/src/rl/loss_balance.rs index b283184a8..4cb1b7d52 100644 --- a/crates/ml-alpha/src/rl/loss_balance.rs +++ b/crates/ml-alpha/src/rl/loss_balance.rs @@ -29,12 +29,16 @@ pub struct LossLambdas { pub pi: f32, pub v: f32, pub aux: f32, + /// SP20 P3 FRD head loss weight. Reads from `RL_FRD_LAMBDA_INDEX` + /// (slot 498); also feeds `grad_h_accumulate` for the FRD branch + /// of the combined encoder-upstream gradient. + pub frd: f32, } impl Default for LossLambdas { - /// Initial equal-weight allocation. Sums to 5.0; per-head divide of 1/5 + /// Initial equal-weight allocation. Sums to 6.0; per-head divide of 1/6 /// is performed at the trainer's loss-combine site so each head's loss - /// contribution is weighted as `lambda / 5.0`. + /// contribution is weighted as `lambda / 6.0`. fn default() -> Self { Self { bce: 1.0, @@ -42,6 +46,7 @@ impl Default for LossLambdas { pi: 1.0, v: 1.0, aux: 1.0, + frd: 1.0, } } } @@ -76,6 +81,12 @@ pub fn read_loss_lambdas_from_isv(isv_host_slice: &[f32]) -> LossLambdas { out.v = v; } } + if isv_host_slice.len() > crate::rl::isv_slots::RL_FRD_LAMBDA_INDEX { + let v = isv_host_slice[crate::rl::isv_slots::RL_FRD_LAMBDA_INDEX]; + if v != 0.0 { + out.frd = v; + } + } out } @@ -92,6 +103,7 @@ mod tests { assert_eq!(d.pi, 1.0); assert_eq!(d.v, 1.0); assert_eq!(d.aux, 1.0); + assert_eq!(d.frd, 1.0); } #[test] @@ -103,6 +115,7 @@ mod tests { assert_eq!(l.q, 1.0); assert_eq!(l.pi, 1.0); assert_eq!(l.v, 1.0); + assert_eq!(l.frd, 1.0); } #[test] diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index de3b10590..39ca0c529 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -290,6 +290,11 @@ pub struct IntegratedStepStats { pub l_pi: f32, pub l_v: f32, pub l_aux: f32, + /// SP20 P3 FRD head loss — sum of CE across (batch, horizon), + /// normalized by `B × FRD_N_HORIZONS`. Zero whenever all labels + /// are sentinel (-1); becomes the head's training signal once + /// the loader feeds real forward-return labels. + pub l_frd: f32, pub l_total: f32, pub lambdas: LossLambdas, } @@ -319,6 +324,12 @@ pub struct IntegratedTrainer { pub policy_b_adam: AdamW, pub value_w_adam: AdamW, pub value_b_adam: AdamW, + /// SP20 P3 FRD head Adam optimisers (W1/b1/W2/b2). LR from + /// `RL_FRD_LR_INDEX` (slot 499), default 1e-3. + pub frd_w1_adam: AdamW, + pub frd_b1_adam: AdamW, + pub frd_w2_adam: AdamW, + pub frd_b2_adam: AdamW, /// Device ISV buffer (RL-allocated, length `RL_SLOTS_END = 424` as /// of Phase R1). Allocated zero, then immediately seeded by the 7 @@ -648,6 +659,15 @@ pub struct IntegratedTrainer { /// callers can read for diag/inference; backward kernel writes /// gradients into this buffer in F.3. pub frd_logits_d: CudaSlice, + + /// SP20 P3 — FRD per-(batch, horizon) labels for the supervised + /// CE loss. Atom index `[0, FRD_N_ATOMS)` selects the correct + /// return bucket; sentinel `-1` marks "missing horizon" (the + /// forward return at h-ticks-ahead isn't realized yet at the + /// rightmost edge of the snapshot stream). All entries initialise + /// to -1; F.5 will populate before each `step_with_lobsim` from + /// the loader's forward-snapshot lookahead. + pub frd_labels_d: CudaSlice, } impl IntegratedTrainer { @@ -1186,6 +1206,25 @@ impl IntegratedTrainer { let frd_logits_d = stream .alloc_zeros::(b_size * crate::rl::frd::FRD_OUT_DIM) .context("alloc frd_logits_d")?; + // SP20 P3 Adam state for the 4 FRD weight tensors. LR mirrored + // from ISV[RL_FRD_LR_INDEX] (slot 499) on every step before + // the Adam launches — same pattern as Q/π/V (placeholder=0 + // here, updated per-step). + let frd_w1_adam = AdamW::new(dev, frd_head.w1_d.len(), lr_placeholder) + .context("frd_w1_adam")?; + let frd_b1_adam = AdamW::new(dev, frd_head.b1_d.len(), lr_placeholder) + .context("frd_b1_adam")?; + let frd_w2_adam = AdamW::new(dev, frd_head.w2_d.len(), lr_placeholder) + .context("frd_w2_adam")?; + let frd_b2_adam = AdamW::new(dev, frd_head.b2_d.len(), lr_placeholder) + .context("frd_b2_adam")?; + // FRD labels buffer — sentinel-initialized to -1 (missing + // horizon). F.5 loader integration overwrites per step. + let mut frd_labels_d = stream + .alloc_zeros::(b_size * crate::rl::common::FRD_N_HORIZONS) + .context("alloc frd_labels_d")?; + let neg_ones = vec![-1_i32; b_size * crate::rl::common::FRD_N_HORIZONS]; + write_slice_i32_d_pub(&stream, &neg_ones, &mut frd_labels_d)?; Ok(Self { cfg, @@ -1321,6 +1360,11 @@ impl IntegratedTrainer { frd_head, frd_hidden_d, frd_logits_d, + frd_labels_d, + frd_w1_adam, + frd_b1_adam, + frd_w2_adam, + frd_b2_adam, } .with_controllers_bootstrapped()?) } @@ -2304,6 +2348,12 @@ impl IntegratedTrainer { self.policy_b_adam.lr = lr_pi; self.value_w_adam.lr = lr_v; self.value_b_adam.lr = lr_v; + // SP20 P3 FRD head — LR from dedicated ISV slot (498-503 block). + let lr_frd = self.isv_host[crate::rl::isv_slots::RL_FRD_LR_INDEX]; + self.frd_w1_adam.lr = lr_frd; + self.frd_b1_adam.lr = lr_frd; + self.frd_w2_adam.lr = lr_frd; + self.frd_b2_adam.lr = lr_frd; // Borrow encoder hidden state for forward kernels. let h_t_borrow: &CudaSlice = self.perception.h_t_view(); @@ -2350,6 +2400,10 @@ impl IntegratedTrainer { let mut v_grad_h_t_d = self.stream.alloc_zeros::(b_size * HIDDEN_DIM)?; let mut v_grad_w_d = self.stream.alloc_zeros::(HIDDEN_DIM)?; let mut v_grad_b_d = self.stream.alloc_zeros::(1)?; + // SP20 P3 FRD layer-1 backward produces grad_h_t into this + // buffer (the encoder-upstream gradient). Folded into + // grad_h_t_combined_d at Step 10 with the lambdas.frd weight. + let mut frd_grad_h_t_d = self.stream.alloc_zeros::(b_size * HIDDEN_DIM)?; // Bellman projection scratch (item 2). let mut q_target_full_d = self.stream.alloc_zeros::(b_size * k_dqn)?; @@ -2631,6 +2685,103 @@ impl IntegratedTrainer { .step(&mut self.value_head.b_d, &v_grad_b_d) .context("value_b_adam.step")?; + // ── SP20 P3 FRD backward chain (F.4) ───────────────────────── + // Reads frd_hidden_d + frd_logits_d (populated by + // step_with_lobsim's FRD fwd) + self.frd_labels_d (sentinel -1 + // until F.5 loader integration populates from forward-snapshot + // lookahead). On all-sentinel labels, softmax_ce_grad emits + // zero gradient + zero loss → entire chain produces no signal, + // Adam steps are no-ops modulo β decay. Once F.5 lands real + // labels, this path activates the supervised learning signal + // without any further trainer-side changes. + let l_frd_host: f32 = { + let frd_n_h = crate::rl::common::FRD_N_HORIZONS; + let frd_h_dim = crate::rl::common::FRD_HIDDEN_DIM; + let frd_out_dim = crate::rl::frd::FRD_OUT_DIM; + let mut frd_grad_logits_d = self.stream.alloc_zeros::(b_size * frd_out_dim)?; + let mut frd_loss_per_b_h_d = self.stream.alloc_zeros::(b_size * frd_n_h)?; + self.frd_head.softmax_ce_grad( + &self.frd_logits_d, + &self.frd_labels_d, + &mut frd_grad_logits_d, + &mut frd_loss_per_b_h_d, + b_size, + )?; + let mut frd_grad_w2_pb_d = self + .stream + .alloc_zeros::(b_size * frd_h_dim * frd_out_dim)?; + let mut frd_grad_b2_pb_d = self.stream.alloc_zeros::(b_size * frd_out_dim)?; + let mut frd_grad_hidden_d = self.stream.alloc_zeros::(b_size * frd_h_dim)?; + self.frd_head.layer2_bwd( + &self.frd_hidden_d, + &frd_grad_logits_d, + &mut frd_grad_w2_pb_d, + &mut frd_grad_b2_pb_d, + &mut frd_grad_hidden_d, + b_size, + )?; + let mut frd_grad_w1_pb_d = self + .stream + .alloc_zeros::(b_size * HIDDEN_DIM * frd_h_dim)?; + let mut frd_grad_b1_pb_d = self.stream.alloc_zeros::(b_size * frd_h_dim)?; + self.frd_head.layer1_bwd( + h_t_borrow, + &self.frd_hidden_d, + &frd_grad_hidden_d, + &mut frd_grad_w1_pb_d, + &mut frd_grad_b1_pb_d, + &mut frd_grad_h_t_d, + b_size, + )?; + // Reduce-axis-0: batch-summed weight grads. + let mut frd_grad_w1_d = self.stream.alloc_zeros::(HIDDEN_DIM * frd_h_dim)?; + let mut frd_grad_b1_d = self.stream.alloc_zeros::(frd_h_dim)?; + let mut frd_grad_w2_d = self.stream.alloc_zeros::(frd_h_dim * frd_out_dim)?; + let mut frd_grad_b2_d = self.stream.alloc_zeros::(frd_out_dim)?; + self.launch_reduce_axis0( + &frd_grad_w1_pb_d, + b_size, + HIDDEN_DIM * frd_h_dim, + &mut frd_grad_w1_d, + )?; + self.launch_reduce_axis0( + &frd_grad_b1_pb_d, + b_size, + frd_h_dim, + &mut frd_grad_b1_d, + )?; + self.launch_reduce_axis0( + &frd_grad_w2_pb_d, + b_size, + frd_h_dim * frd_out_dim, + &mut frd_grad_w2_d, + )?; + self.launch_reduce_axis0( + &frd_grad_b2_pb_d, + b_size, + frd_out_dim, + &mut frd_grad_b2_d, + )?; + // Adam steps — sentinel labels yield zero grads → no + // weight movement (only momentum decay). + self.frd_w1_adam + .step(&mut self.frd_head.w1_d, &frd_grad_w1_d) + .context("frd_w1_adam.step")?; + self.frd_b1_adam + .step(&mut self.frd_head.b1_d, &frd_grad_b1_d) + .context("frd_b1_adam.step")?; + self.frd_w2_adam + .step(&mut self.frd_head.w2_d, &frd_grad_w2_d) + .context("frd_w2_adam.step")?; + self.frd_b2_adam + .step(&mut self.frd_head.b2_d, &frd_grad_b2_d) + .context("frd_b2_adam.step")?; + // Per-row CE loss → mean over (b, h). Mapped-pinned read. + let loss_pb_h = + read_slice_d(&self.stream, &frd_loss_per_b_h_d, b_size * frd_n_h)?; + loss_pb_h.iter().sum::() / ((b_size * frd_n_h) as f32) + }; + // ── Step 10: encoder grad combine (Phase E.2-DEFER item 1) ─── // Zero the combined slot, fold π+V grad_h_t with λ-weights. // @@ -2663,6 +2814,18 @@ impl IntegratedTrainer { lambdas.v, &mut self.grad_h_t_combined_d, )?; + // SP20 P3 FRD encoder-upstream gradient — same accumulator + // pattern as π / V. With sentinel labels (F.4 state) the + // buffer is all-zeros so this is a no-op; once F.5 lands real + // labels it injects supervised-learning signal into the + // encoder via the standard λ-weighted accumulator. + accumulate_grad_h( + &self.stream, + &self.grad_h_accumulate_fn, + &frd_grad_h_t_d, + lambdas.frd, + &mut self.grad_h_t_combined_d, + )?; // ── Step 11: encoder backward (Phase E.3a) ─────────────────── // Seed the perception trainer's per-K hidden-state grad buffer @@ -2825,14 +2988,15 @@ impl IntegratedTrainer { // Normalise the Q CE by batch (matches l_pi / l_v conventions). let l_q = l_q_host / (b_size as f32); - // Combined: λ-weighted average across 5 heads (per - // pearl_loss_balance_controller, sum over 5 heads / 5). + // Combined: λ-weighted average across 6 heads per + // pearl_loss_balance_controller (SP20 P3 added FRD as 6th). let l_total = (lambdas.bce * l_bce + lambdas.q * l_q + lambdas.pi * l_pi + lambdas.v * l_v_host - + lambdas.aux * l_aux) - / 5.0; + + lambdas.aux * l_aux + + lambdas.frd * l_frd_host) + / 6.0; Ok(IntegratedStepStats { l_bce, @@ -2840,6 +3004,7 @@ impl IntegratedTrainer { l_pi, l_v: l_v_host, l_aux, + l_frd: l_frd_host, l_total, lambdas, })