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

feat(rl): FRD head trainer integration — Adam + bwd chain + loss (F.4)

Browse Source 
Wires the F.3a/b/c backward kernels into IntegratedTrainer's per-step
flow so the FRD head trains end-to-end as a 6th loss-balanced head
alongside BCE/Q/π/V/aux. With labels currently sentinel-initialized to
-1 (F.5 loader will populate from forward-snapshot lookahead), the
chain produces zero gradients + zero loss — Adam steps are no-ops
modulo β decay, and the encoder receives no FRD-derived signal yet.
The wiring is complete and the path is exercised end-to-end; F.5 just
needs to swap the labels in for the head to start training.

IntegratedTrainer state additions:
  * frd_w1_adam / frd_b1_adam / frd_w2_adam / frd_b2_adam — AdamW
    instances for the 4 FRD weight tensors (LR mirrored per-step from
    ISV[RL_FRD_LR_INDEX=499], seed 1e-3 per F.1).
  * frd_labels_d — owned [B × FRD_N_HORIZONS] i32 buffer, sentinel-
    initialized to -1 (every entry "missing horizon" → softmax_ce_grad
    zeros loss + grad for every row). F.5 loader integration overwrites
    pre-step from forward-return-bucketed labels.

LossLambdas extension:
  * Added `frd: f32` field, default 1.0
  * read_loss_lambdas_from_isv reads slot 498 (RL_FRD_LAMBDA_INDEX)
    with the standard zero-sentinel bootstrap path
  * Doc-comment updated: "5 heads / 5.0" → "6 heads / 6.0"

IntegratedStepStats extension:
  * Added `l_frd: f32` — mean CE across (B × FRD_N_HORIZONS) rows
  * step_synthetic returns the real l_frd from the bwd chain; the
    new combined l_total formula includes `lambdas.frd × l_frd / 6`

step_synthetic bwd chain — inserted between Step 9 (Q/π/V Adam) and
Step 10 (grad_h_t_combined zero+accumulate):
  1. softmax_ce_grad → frd_grad_logits_d + frd_loss_per_b_h_d
  2. layer2_bwd → frd_grad_w2_pb_d, frd_grad_b2_pb_d, frd_grad_hidden_d
  3. layer1_bwd → frd_grad_w1_pb_d, frd_grad_b1_pb_d, frd_grad_h_t_d
  4. 4× reduce_axis0 to collapse per-batch scratch → final grads
  5. 4× AdamW.step on w1/b1/w2/b2
  6. read loss_per_b_h via mapped-pinned, average → l_frd_host

Step 10 grad_h_t_combined accumulation adds a third λ-weighted call:
  accumulate_grad_h(frd_grad_h_t_d, lambdas.frd, &mut combined)

With sentinel labels (F.4 state) this contributes zero gradient to the
encoder backward — the wiring is exercised but silent. F.5 makes it
active by providing real labels.

alpha_rl_train diag JSON gains:
  * "loss": { ..., "frd": stats.l_frd, ... }
  * "lambdas": { ..., "frd": stats.lambdas.frd, ... }

Verification (RTX 3050 Ti):
  * cargo check -p ml-alpha + --examples → clean
  * integrated_trainer_step_with_lobsim_runs_without_panic → ok
    (l_total 0.5073 vs prior 0.6087 — ÷6 instead of ÷5 expected;
    l_frd=0 confirms sentinel labels are passing through cleanly)
  * frd_head 10/10 tests still pass (no regression)
  * trade_management_kernels 5/5 → no regression
  * audit-rust-consts → 0 flags

F.5 (next, separate scope):
  * Loader-side forward-return label generation (mid[i+h] - mid[i])/σ
    bucketed into FRD_N_ATOMS=21 atoms over the ISV-driven ±range_σ
  * Populate trainer.frd_labels_d before each step_with_lobsim call
  * That unlocks the supervised learning signal; FRD entropy_mean
    should start dropping below ln(21) in diag as the head trains.
This commit is contained in:
jgrusewski
2026-05-24 19:03:20 +02:00
parent 0f75d6bb7b
commit 935433850c
3 changed files with 186 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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": {

17
crates/ml-alpha/src/rl/loss_balance.rs
View File

@@ -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]

173
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -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<f32>,
/// 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<i32>,
}
impl IntegratedTrainer {
@@ -1186,6 +1206,25 @@ impl IntegratedTrainer {
let frd_logits_d = stream
.alloc_zeros::<f32>(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::<i32>(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<f32> = self.perception.h_t_view();
@@ -2350,6 +2400,10 @@ impl IntegratedTrainer {
let mut v_grad_h_t_d = self.stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)?;
let mut v_grad_w_d = self.stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
let mut v_grad_b_d = self.stream.alloc_zeros::<f32>(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::<f32>(b_size * HIDDEN_DIM)?;
// Bellman projection scratch (item 2).
let mut q_target_full_d = self.stream.alloc_zeros::<f32>(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::<f32>(b_size * frd_out_dim)?;
let mut frd_loss_per_b_h_d = self.stream.alloc_zeros::<f32>(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::<f32>(b_size * frd_h_dim * frd_out_dim)?;
let mut frd_grad_b2_pb_d = self.stream.alloc_zeros::<f32>(b_size * frd_out_dim)?;
let mut frd_grad_hidden_d = self.stream.alloc_zeros::<f32>(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::<f32>(b_size * HIDDEN_DIM * frd_h_dim)?;
let mut frd_grad_b1_pb_d = self.stream.alloc_zeros::<f32>(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::<f32>(HIDDEN_DIM * frd_h_dim)?;
let mut frd_grad_b1_d = self.stream.alloc_zeros::<f32>(frd_h_dim)?;
let mut frd_grad_w2_d = self.stream.alloc_zeros::<f32>(frd_h_dim * frd_out_dim)?;
let mut frd_grad_b2_d = self.stream.alloc_zeros::<f32>(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::<f32>() / ((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,
})