diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 4d43a4b2e..f24641c23 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -78,7 +78,8 @@ use std::sync::Arc; use anyhow::{Context, Result}; use cudarc::driver::{ - CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg, + CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, + PushKernelArg, }; use ml_core::device::MlDevice; @@ -360,6 +361,28 @@ pub struct IntegratedTrainer { /// `compute_advantage_return`. pub returns_d: CudaSlice, + /// Phase R7c (data correctness): trainer-owned `h_{t+1}` buffer + /// (`[B × HIDDEN_DIM]`). Populated by a SECOND + /// `perception.forward_encoder(next_snapshots)` call in + /// `step_with_lobsim`, with a stream-ordered DtoD copy out of + /// `perception.h_t_d` BEFORE the third `forward_encoder(snapshots)` + /// call overwrites it. Closes the long-standing + /// `// FUTURE: pass next_h_t separately` approximation noted in + /// step_synthetic's Bellman target build (which had been using + /// `h_t` as a stand-in for `s_{t+1}`'s encoder representation since + /// Phase E.2). Three consumers downstream: + /// 1. `value_head.forward(&self.h_tp1_d) → v_pred_tp1_d`, which + /// compute_advantage_return reads as the true `V(s_{t+1})` + /// (replaces R7b's `v_tp1_d_ref = &v_pred_d` alias). + /// 2. `dqn_head.forward(&self.h_tp1_d) → q_logits_tp1_d`, fed to + /// `argmax_expected_q` for `next_actions_d` — the Double-DQN + /// argmax must be on `h_{t+1}`, not `h_t` (R7b's q_logits_d + /// was on h_t, so the argmax was wrong since R4). + /// 3. `dqn_head.forward_target(&self.h_tp1_d)` inside + /// step_synthetic's Bellman target build (replaces the + /// h_t-as-proxy comment at the call site). + pub h_tp1_d: CudaSlice, + /// Combined encoder grad slot — folded from Q + π + V `grad_h_t` /// contributions via `grad_h_accumulate_scaled` (item 1). Consumed by /// `PerceptionTrainer::backward_encoder_with_grad_h_t`. Size @@ -665,6 +688,14 @@ impl IntegratedTrainer { .alloc_zeros::(b_size) .context("alloc returns_d")?; + // Phase R7c (data correctness): trainer-owned h_{t+1} buffer. + // Zero at init — first step_with_lobsim call writes it via the + // forward_encoder(next_snapshots) + DtoD copy out of + // perception.h_t_d BEFORE any consumer reads it. + let h_tp1_d = stream + .alloc_zeros::(b_size * HIDDEN_DIM) + .context("alloc h_tp1_d")?; + Ok(Self { cfg, perception, @@ -732,6 +763,7 @@ impl IntegratedTrainer { log_pi_old_d, advantages_d, returns_d, + h_tp1_d, grad_h_t_combined_d, last_pi_loss: 0.0, step_counter: 0, @@ -1372,16 +1404,19 @@ impl IntegratedTrainer { } // ── Step 6a: Bellman target build (item 2) ─────────────────── - // target-net forward at h_t → full atom logits → per-batch + // target-net forward at h_{t+1} → full atom logits → per-batch // selected action's atom row → projection reading γ from ISV. // - // For the synthetic-data smoke we use `h_t` as the proxy for - // s_{t+1}'s encoder representation. A future enhancement (Phase - // E.3 LobSim integration) will pass next_h_t separately. The - // projection itself is correct regardless of this approximation. + // R7c data correctness lift: was using `h_t` as proxy for + // s_{t+1}'s encoder representation since Phase E.2. Now reads + // the trainer-owned self.h_tp1_d populated by + // step_with_lobsim's `forward_encoder(next_snapshots)` + + // DtoD copy out of perception.h_t_d (R7c.A). Pairs with the + // Double-DQN argmax on h_{t+1} that step_with_lobsim writes + // into self.next_actions_d (also R7c — was on h_t since R4). self.dqn_head - .forward_target(h_t_borrow, b_size, &mut q_target_full_d) - .context("dqn_head.forward_target")?; + .forward_target(&self.h_tp1_d, b_size, &mut q_target_full_d) + .context("dqn_head.forward_target(h_tp1)")?; self.dqn_head .select_action_atoms( &q_target_full_d, @@ -1643,6 +1678,7 @@ impl IntegratedTrainer { pub fn step_with_lobsim( &mut self, snapshots: &[Mbp10RawInput], + next_snapshots: &[Mbp10RawInput], lobsim: &mut dyn RlLobBackend, ) -> Result { let b_size = self.cfg.perception.n_batch; @@ -1652,27 +1688,71 @@ impl IntegratedTrainer { if snapshots.is_empty() { anyhow::bail!("step_with_lobsim: snapshots empty (need ≥ 1 for apply_snapshot)"); } + if next_snapshots.len() != snapshots.len() { + anyhow::bail!( + "step_with_lobsim: next_snapshots.len() ({}) must match snapshots.len() ({}) — \ + R8's CLI binary builds them via MultiHorizonLoader::next_sequence_pair (R2) which \ + guarantees this invariant for adjacent (s_t, s_{{t+1}}) windows", + next_snapshots.len(), + snapshots.len() + ); + } - // ── Step 1: encoder forward to land h_t at slot K-1. ────────── + // ── Step 1a (R7c): encoder forward on NEXT snapshots to land + // h_{t+1} at slot K-1. Done FIRST so the subsequent + // forward_encoder(snapshots) call leaves perception's internal + // forward state (h_new_per_k_d / CfC h_state_d) populated for + // the CURRENT-step encoder backward in step_synthetic — the + // most recent forward_encoder call wins. We DtoD-copy + // perception.h_t_d into the trainer-owned self.h_tp1_d + // immediately so the second forward_encoder doesn't clobber it. + let _ = self + .perception + .forward_encoder(next_snapshots) + .context("step_with_lobsim: forward_encoder(next_snapshots)")?; + { + let nbytes = b_size * HIDDEN_DIM * std::mem::size_of::(); + unsafe { + let s = self.stream.cu_stream(); + let (src, _gs) = self.perception.h_t_view().device_ptr(&self.stream); + let (dst, _gd) = self.h_tp1_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes, s) + .context("step_with_lobsim: DtoD perception.h_t_d → self.h_tp1_d")?; + } + } + + // ── Step 1b: encoder forward on CURRENT snapshots — leaves + // perception's internal forward state primed for the encoder + // backward in step_synthetic, AND lands h_t at slot K-1 for + // the action-sampling forwards below. let _ = self .perception .forward_encoder(snapshots) - .context("step_with_lobsim: forward_encoder")?; + .context("step_with_lobsim: forward_encoder(snapshots)")?; // ── Step 2: Q + V forwards on h_t for action sampling. ──────── let h_t_borrow: &CudaSlice = self.perception.h_t_view(); debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(self.h_tp1_d.len(), b_size * HIDDEN_DIM); 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)?; let mut v_pred_d = self.stream.alloc_zeros::(b_size)?; + let mut v_pred_tp1_d = self.stream.alloc_zeros::(b_size)?; self.dqn_head .forward(h_t_borrow, b_size, &mut q_logits_d) - .context("step_with_lobsim: dqn_head.forward")?; + .context("step_with_lobsim: dqn_head.forward(h_t)")?; + self.dqn_head + .forward(&self.h_tp1_d, b_size, &mut q_logits_tp1_d) + .context("step_with_lobsim: dqn_head.forward(h_tp1) for Double-DQN argmax")?; self.value_head .forward(h_t_borrow, b_size, &mut v_pred_d) - .context("step_with_lobsim: value_head.forward")?; + .context("step_with_lobsim: value_head.forward(h_t)")?; + self.value_head + .forward(&self.h_tp1_d, b_size, &mut v_pred_tp1_d) + .context("step_with_lobsim: value_head.forward(h_tp1) for true V(s_{t+1})")?; // ── Step 2b: Forward π logits for log_pi_old. ───────────────── let mut pi_logits_d = self.stream.alloc_zeros::(b_size * N_ACTIONS)?; @@ -1716,9 +1796,16 @@ impl IntegratedTrainer { } } - // argmax over expected Q for next_actions (Bellman-target - // action selection — distinct from Thompson per - // `pearl_thompson_for_distributional_action_selection`). + // argmax over expected Q for next_actions — Double-DQN: argmax + // of ONLINE Q at h_{t+1} (per + // `pearl_thompson_for_distributional_action_selection`'s + // distinction between Thompson selector (on h_t for rollout) + // and argmax Bellman selector (on h_{t+1} for target)). + // + // R7c fix: was reading q_logits_d (online Q at h_t) which + // produced an off-by-one-time-index Bellman argmax since R4. + // Now reads q_logits_tp1_d (online Q at h_{t+1}) which is the + // canonical Double-DQN target-action selector. { let cfg = LaunchConfig { grid_dim: (b_size as u32, 1, 1), @@ -1728,14 +1815,14 @@ impl IntegratedTrainer { let b_size_i = b_size as i32; let mut launch = self.stream.launch_builder(&self.argmax_expected_q_fn); launch - .arg(&q_logits_d) + .arg(&q_logits_tp1_d) .arg(&self.atom_supports_d) .arg(&mut self.next_actions_d) .arg(&b_size_i); unsafe { launch .launch(cfg) - .context("argmax_expected_q launch (Bellman target)")?; + .context("argmax_expected_q launch (Bellman target on h_tp1)")?; } } @@ -1844,15 +1931,13 @@ impl IntegratedTrainer { } } - // ── Step 6 (Phase R7b): GPU-pure post-fill pipeline. ────────── + // ── Step 6 (Phase R7c): GPU-pure post-fill pipeline. ────────── // All actions / next_actions / log_pi_old are already in // trainer-owned device buffers (R4 kernels wrote them in Step 3). - // No host uploads here. V(s_t) stays in v_pred_d throughout; - // R7c adds next_snapshots + forward_encoder(next_snapshots) - // for true V(s_{t+1}). Until then v_tp1 reuses v_pred_d (the - // h_t approximation — compute_advantage_return still produces - // a valid TD target, just with the bootstrap approximation). - let v_tp1_d_ref = &v_pred_d; + // No host uploads here. V(s_t) and V(s_{t+1}) BOTH live on + // device — v_pred_d from value_head.forward(h_t), + // v_pred_tp1_d from value_head.forward(h_tp1) (R7c data + // correctness lift: was an alias of v_pred_d in R7a/R7b). // |reward| → reward_abs_d (input for mean_abs_pnl EMA). // Inline launch (vs `launch_abs_copy`) because the per-method @@ -1953,7 +2038,7 @@ impl IntegratedTrainer { .arg(&self.rewards_d) .arg(&self.dones_d) .arg(&v_pred_d) - .arg(v_tp1_d_ref) + .arg(&v_pred_tp1_d) .arg(&mut self.returns_d) .arg(&mut self.advantages_d) .arg(&b_size_i); diff --git a/crates/ml-alpha/tests/integrated_trainer_smoke.rs b/crates/ml-alpha/tests/integrated_trainer_smoke.rs index 6f96f8514..4d514e890 100644 --- a/crates/ml-alpha/tests/integrated_trainer_smoke.rs +++ b/crates/ml-alpha/tests/integrated_trainer_smoke.rs @@ -82,12 +82,20 @@ fn integrated_trainer_step_with_lobsim_runs_without_panic() { let snapshots = synthetic_window(4); let mut sim = LobSimCuda::new(1, &dev).expect("LobSimCuda::new"); - // One step end-to-end (encoder fwd, Q/π/V fwd, host Thompson + - // actions upload, actions_to_market_targets kernel, - // step_fill_from_market_targets, extract_realized_pnl_delta, - // host advantage/return, step_synthetic delegation). + // One step end-to-end (encoder fwd ×2 [snapshots + next_snapshots + // for R7c data correctness], Q/π/V fwd on h_t + Q/V fwd on h_tp1, + // GPU Thompson + Double-DQN argmax kernels, actions_to_market_targets + // kernel, step_fill_from_market_targets, extract_realized_pnl_delta, + // GPU compute_advantage_return with true V(s_{t+1}), step_synthetic + // delegation). + // + // next_snapshots is the same window shifted by one step. R8's CLI + // binary uses MultiHorizonLoader::next_sequence_pair (R2) to load + // adjacent (s_t, s_{t+1}) windows from real MBP-10 data; for the + // smoke we synthesise a +1-tick window directly. + let next_snapshots = synthetic_window(4); let stats = trainer - .step_with_lobsim(&snapshots, &mut sim) + .step_with_lobsim(&snapshots, &next_snapshots, &mut sim) .expect("step_with_lobsim"); // Loss components are finite.