diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 73fcde390..5634c7951 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -57,16 +57,15 @@ const KERNELS: &[&str] = &[ "abs_copy", // RL Phase R7a: element-wise dst[b] = fabsf(src[b]); feeds |reward| into ema_update_on_done for the MEAN_ABS_PNL_EMA slot ]; -// Cache bust v29 (2026-05-23): RL Phase R6 rebuild — three new GPU-pure -// kernels close the last host-roundtrip violations in step_with_lobsim's -// hot path. extract_realized_pnl_delta reads pos.realized_pnl + position_lots -// from the device Pos array and writes per-batch (reward delta, done flag); -// apply_reward_scale multiplies rewards by ISV[406] on device; -// actions_to_market_targets translates the 9-action grid (with conditional -// flat-from-long/short reading current position_lots) into LobSim's -// market_targets format. Together with R3 EMA + R4 Thompson/argmax/log_pi + -// R5 controllers + soft-update, the trainer's step is GPU-pure: only HtoD -// for incoming snapshots (boundary data) + HEALTH_DIAG ISV readback. +// Cache bust v30 (2026-05-23): RL Phase R7d — dqn_distributional_q_bwd +// kernel signature change. Adds a new `loss_per_batch [B]` output so the +// trainer can DtoH per-sample CE loss into ReplayBuffer.update_priorities +// after each off-policy Q backward (feedback_always_per wiring). Single +// writer per batch (atom 0 only), so the new output is non-atomic — the +// existing scalar `loss_out` still accumulates via atomicAdd for the +// diagnostic total. All callers (the integrated trainer's +// dqn_offpolicy_step) migrate atomically in the same R7d commit per +// feedback_no_partial_refactor. fn main() { println!("cargo:rerun-if-changed=build.rs"); diff --git a/crates/ml-alpha/cuda/dqn_distributional_q.cu b/crates/ml-alpha/cuda/dqn_distributional_q.cu index c43b6f40d..a3673e1ec 100644 --- a/crates/ml-alpha/cuda/dqn_distributional_q.cu +++ b/crates/ml-alpha/cuda/dqn_distributional_q.cu @@ -125,12 +125,13 @@ extern "C" __global__ void dqn_distributional_q_fwd( // actions get zero grad. // ───────────────────────────────────────────────────────────────────── extern "C" __global__ void dqn_distributional_q_bwd( - const float* __restrict__ logits, // [B * N_ACTIONS * Q_N_ATOMS] - const float* __restrict__ target_dist, // [B * Q_N_ATOMS] - const int* __restrict__ actions_taken, // [B] + const float* __restrict__ logits, // [B * N_ACTIONS * Q_N_ATOMS] + const float* __restrict__ target_dist, // [B * Q_N_ATOMS] + const int* __restrict__ actions_taken, // [B] int B, - float* __restrict__ loss_out, // [1] - float* __restrict__ grad_logits // [B * N_ACTIONS * Q_N_ATOMS] + float* __restrict__ loss_out, // [1] + float* __restrict__ loss_per_batch, // [B] — R7d: per-sample CE for PER priority update + float* __restrict__ grad_logits // [B * N_ACTIONS * Q_N_ATOMS] ) { const int batch = blockIdx.x; const int atom = threadIdx.x; @@ -148,6 +149,12 @@ extern "C" __global__ void dqn_distributional_q_bwd( for (int a = 0; a < N_ACTIONS; ++a) { grad_logits[base_all + a * Q_N_ATOMS + atom] = 0.0f; } + // R7d: zero per-sample CE for invalid-action samples so the + // PER priority update sees a real (zero) magnitude rather than + // stale buffer contents — bounded-input pearl applies. + if (atom == 0) { + loss_per_batch[batch] = 0.0f; + } return; } @@ -216,6 +223,11 @@ extern "C" __global__ void dqn_distributional_q_bwd( ce -= tz * logf(pz_c); } atomicAdd(loss_out, ce); + // R7d: per-sample CE for PER priority update. Single writer per + // batch (atom 0 only), so non-atomic store. Feeds + // `replay.update_priorities(per_indices, td_per_sample)` via a + // mapped-pinned DtoH after the backward. + loss_per_batch[batch] = ce; } } diff --git a/crates/ml-alpha/src/rl/dqn.rs b/crates/ml-alpha/src/rl/dqn.rs index 9d5255058..6b24a35cd 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -327,12 +327,18 @@ impl DqnHead { actions_taken: &CudaSlice, b_size: usize, loss_out: &mut CudaSlice, + loss_per_batch: &mut CudaSlice, grad_logits: &mut CudaSlice, ) -> Result<()> { debug_assert_eq!(logits.len(), b_size * N_ACTIONS * Q_N_ATOMS); debug_assert_eq!(target_dist.len(), b_size * Q_N_ATOMS); debug_assert_eq!(actions_taken.len(), b_size); debug_assert_eq!(loss_out.len(), 1); + debug_assert_eq!( + loss_per_batch.len(), + b_size, + "R7d: per-sample CE buffer must match batch size for PER priority update" + ); debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS * Q_N_ATOMS); let b_i = b_size as i32; @@ -348,6 +354,7 @@ impl DqnHead { .arg(actions_taken) .arg(&b_i) .arg(loss_out) + .arg(loss_per_batch) .arg(grad_logits); unsafe { launch.launch(cfg).context("dqn_distributional_q_bwd launch")?; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index f24641c23..6f714c612 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -186,6 +186,30 @@ pub struct IntegratedTrainerConfig { pub dqn_seed: u64, /// Seed for PPO heads (policy + value share derived seeds). pub ppo_seed: u64, + /// PER buffer capacity (Phase R7d). Per the canonical replay.rs + /// doc, naive O(N) sampling is acceptable up to ≈ 4096; production + /// scale-up to 100k transitions needs a GPU sum-tree (R-future). + pub per_capacity: usize, + /// PER seed for deterministic sampling per + /// `pearl_scoped_init_seed_for_reproducibility`. Fed to + /// `ReplayBuffer::new(capacity, seed)`. + pub per_seed: u64, +} + +impl Default for IntegratedTrainerConfig { + /// Per `pearl_first_observation_bootstrap` + R1 controller defaults: + /// every adaptive knob sources from ISV; the only knobs here are + /// structural (encoder dims via PerceptionTrainerConfig) or + /// reproducibility (seeds). + fn default() -> Self { + Self { + perception: PerceptionTrainerConfig::default(), + dqn_seed: 0xDA_DA_DA_DA, + ppo_seed: 0xBE_BE_BE_BE, + per_capacity: 4096, + per_seed: 0x9E37_79B9_7F4A_7C15, + } + } } /// Stats returned from `step_synthetic`. Useful for tests + diagnostics. @@ -383,6 +407,45 @@ pub struct IntegratedTrainer { /// h_t-as-proxy comment at the call site). pub h_tp1_d: CudaSlice, + // ── Phase R7d: PER + off-policy Q replay ────────────────────────── + /// Prioritised Experience Replay buffer (off-policy DQN training). + /// Per `feedback_always_per` PER is always-on. Wired in R7d after + /// the existing dead struct sat in `src/rl/replay.rs` since Phase C. + pub replay: crate::rl::replay::ReplayBuffer, + + /// Per-step gather buffer for PER-sampled `h_t`. Populated each step + /// by a per-batch DtoD loop over `replay.transitions[sampled_idx].h_t`. + /// Length `B × HIDDEN_DIM`. + pub sampled_h_t_d: CudaSlice, + + /// Per-step gather buffer for PER-sampled `h_{t+1}`. Same pattern as + /// `sampled_h_t_d`. Length `B × HIDDEN_DIM`. + pub sampled_h_tp1_d: CudaSlice, + + /// Per-step gather buffer for PER-sampled actions (HtoD per step + /// from the host metadata replay stores alongside each transition). + /// Length `B`. + pub sampled_actions_d: CudaSlice, + + /// Per-step gather buffer for PER-sampled rewards. Length `B`. + pub sampled_rewards_d: CudaSlice, + + /// Per-step gather buffer for PER-sampled done flags. Length `B`. + pub sampled_dones_d: CudaSlice, + + /// Per-step gather buffer for PER-sampled next-state argmax actions + /// (Double-DQN online-Q argmax on `sampled_h_tp1`). Computed + /// per-step inside `dqn_offpolicy_step` via `argmax_expected_q` — + /// not stored in the replay buffer (online net weights drift, so + /// the argmax is recomputed each step). Length `B`. + pub sampled_next_actions_d: CudaSlice, + + /// Per-step CE loss output of `dqn_distributional_q_bwd` (R7d + /// kernel signature addition). Length `B`. DtoH'd at the end of + /// `dqn_offpolicy_step` into `td_per_sample_host`; fed to + /// `replay.update_priorities(per_indices, td_per_sample_host)`. + pub td_per_sample_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 @@ -696,6 +759,38 @@ impl IntegratedTrainer { .alloc_zeros::(b_size * HIDDEN_DIM) .context("alloc h_tp1_d")?; + // Phase R7d: PER buffer + sampled-batch device scratch. + // The replay buffer holds per-transition device payloads + // (h_t / next_h_t as small `CudaSlice(HIDDEN_DIM)` slices + // owned by each `Transition`). Sample → gather populates the + // trainer-owned sampled_* buffers via per-batch DtoD copies + // out of the buffer's transitions. + let replay = crate::rl::replay::ReplayBuffer::new( + cfg.per_capacity, + cfg.per_seed, + ); + let sampled_h_t_d = stream + .alloc_zeros::(b_size * HIDDEN_DIM) + .context("alloc sampled_h_t_d")?; + let sampled_h_tp1_d = stream + .alloc_zeros::(b_size * HIDDEN_DIM) + .context("alloc sampled_h_tp1_d")?; + let sampled_actions_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_actions_d")?; + let sampled_rewards_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_rewards_d")?; + let sampled_dones_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_dones_d")?; + let sampled_next_actions_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_next_actions_d")?; + let td_per_sample_d = stream + .alloc_zeros::(b_size) + .context("alloc td_per_sample_d")?; + Ok(Self { cfg, perception, @@ -764,6 +859,14 @@ impl IntegratedTrainer { advantages_d, returns_d, h_tp1_d, + replay, + sampled_h_t_d, + sampled_h_tp1_d, + sampled_actions_d, + sampled_rewards_d, + sampled_dones_d, + sampled_next_actions_d, + td_per_sample_d, grad_h_t_combined_d, last_pi_loss: 0.0, step_counter: 0, @@ -1369,9 +1472,16 @@ impl IntegratedTrainer { // No host uploads here. // ── Step 5: forward kernels (Q, π, V) ──────────────────────── + // R7d: Q forward consumes PER-SAMPLED h_t (off-policy DQN); + // π and V forwards stay on current-step h_t (on-policy PPO + V). + // The shared encoder gets gradient signal ONLY from π + V (and + // BCE/aux via perception's separate step_batched path); Q's + // encoder-direction gradient is stop-grad'd (discarded in + // Step 10 by NOT accumulating q_grad_h_t into the encoder + // grad slot). self.dqn_head - .forward(h_t_borrow, b_size, &mut q_logits_d) - .context("dqn_head.forward")?; + .forward(&self.sampled_h_t_d, b_size, &mut q_logits_d) + .context("dqn_head.forward(sampled_h_t) [R7d off-policy]")?; self.policy_head .forward_logits(h_t_borrow, b_size, &mut pi_logits_d) .context("policy_head.forward_logits")?; @@ -1379,6 +1489,35 @@ impl IntegratedTrainer { .forward(h_t_borrow, b_size, &mut v_pred_d) .context("value_head.forward")?; + // R7d: Online Q at sampled_h_tp1 for the Double-DQN argmax that + // feeds the Bellman target. The argmax is recomputed every step + // because online net weights drift faster than transitions + // recycle through the replay buffer; storing it at push time + // would feed stale-action data into the projection. + let mut q_logits_tp1_sampled_d = self.stream.alloc_zeros::(b_size * k_dqn)?; + self.dqn_head + .forward(&self.sampled_h_tp1_d, b_size, &mut q_logits_tp1_sampled_d) + .context("dqn_head.forward(sampled_h_tp1) [R7d Double-DQN argmax src]")?; + { + let cfg_argmax = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (N_ACTIONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.argmax_expected_q_fn); + launch + .arg(&q_logits_tp1_sampled_d) + .arg(&self.atom_supports_d) + .arg(&mut self.sampled_next_actions_d) + .arg(&b_size_i); + unsafe { + launch + .launch(cfg_argmax) + .context("argmax_expected_q launch (sampled_h_tp1 → sampled_next_actions)")?; + } + } + // PPO surrogate forward — policy loss + entropy bonus only. // V loss flows through the V-head kernels (item 4). { @@ -1404,37 +1543,37 @@ impl IntegratedTrainer { } // ── Step 6a: Bellman target build (item 2) ─────────────────── - // target-net forward at h_{t+1} → full atom logits → per-batch - // selected action's atom row → projection reading γ from ISV. + // R7d off-policy: target Q on PER-SAMPLED h_{t+1}, action + // selector = sampled_next_actions (Double-DQN argmax on + // online_Q(sampled_h_tp1) computed just above), rewards/dones + // from the SAMPLED transitions' metadata. The projection + // itself reads γ from ISV[400] on-device. // - // 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). + // R7c data correctness lift (this commit's predecessor) put + // the trainer on h_{t+1} via self.h_tp1_d; R7d re-routes the + // entire Bellman target build through PER-sampled buffers per + // `feedback_always_per`. self.dqn_head - .forward_target(&self.h_tp1_d, b_size, &mut q_target_full_d) - .context("dqn_head.forward_target(h_tp1)")?; + .forward_target(&self.sampled_h_tp1_d, b_size, &mut q_target_full_d) + .context("dqn_head.forward_target(sampled_h_tp1) [R7d off-policy]")?; self.dqn_head .select_action_atoms( &q_target_full_d, - &self.next_actions_d, + &self.sampled_next_actions_d, b_size, &mut q_target_action_d, ) - .context("dqn_head.select_action_atoms")?; + .context("dqn_head.select_action_atoms (sampled_next_actions)")?; self.dqn_head .project_bellman_target( &q_target_action_d, - &self.rewards_d, - &self.dones_d, + &self.sampled_rewards_d, + &self.sampled_dones_d, &self.isv_d, b_size, &mut target_dist_d, ) - .context("dqn_head.project_bellman_target")?; + .context("dqn_head.project_bellman_target (sampled rewards/dones)")?; // ── Step 6b: DQN backward (logits → grad_w/b/h_t) ──────────── let mut q_loss_d_mut = q_loss_d; @@ -1442,24 +1581,35 @@ impl IntegratedTrainer { .backward_logits( &q_logits_d, &target_dist_d, - &self.actions_d, + &self.sampled_actions_d, b_size, &mut q_loss_d_mut, + &mut self.td_per_sample_d, &mut q_grad_logits_d, ) - .context("dqn_head.backward_logits")?; + .context("dqn_head.backward_logits (sampled_actions) [R7d off-policy]")?; let l_q_host = read_scalar_d(&self.stream, &q_loss_d_mut)?; self.dqn_head .backward_to_w_b_h( - h_t_borrow, + &self.sampled_h_t_d, &q_grad_logits_d, b_size, &mut q_grad_w_per_batch_d, &mut q_grad_b_per_batch_d, &mut q_grad_h_t_d, ) - .context("dqn_head.backward_to_w_b_h")?; + .context("dqn_head.backward_to_w_b_h(sampled_h_t) [R7d off-policy]")?; + // R7d stop-grad: q_grad_h_t_d is the gradient wrt SAMPLED h_t + // (a past-step encoder output). Accumulating it into the + // shared encoder via grad_h_t_combined_d would poison the + // encoder with stale-state gradient signal. Standard pattern + // for off-policy + shared encoder (SAC / R2D2 / IMPALA do the + // same). The buffer is allocated above for backward_to_w_b_h + // to write into; we just don't FOLD it into the encoder grad + // combine below. Encoder learns from π + V (on-policy) + + // BCE/aux (supervised via step_batched) only. + let _ = &q_grad_h_t_d; self.launch_reduce_axis0(&q_grad_w_per_batch_d, b_size, k_dqn * HIDDEN_DIM, &mut q_grad_w_d)?; self.launch_reduce_axis0(&q_grad_b_per_batch_d, b_size, k_dqn, &mut q_grad_b_d)?; @@ -1537,7 +1687,16 @@ impl IntegratedTrainer { .context("value_b_adam.step")?; // ── Step 10: encoder grad combine (Phase E.2-DEFER item 1) ─── - // Zero the combined slot, fold Q+π+V grad_h_t with λ-weights. + // Zero the combined slot, fold π+V grad_h_t with λ-weights. + // + // R7d stop-grad: Q's grad_h_t is OMITTED here per the off-policy + // shared-encoder pattern (Q is trained on PER-sampled past h_t + // values, whose gradient should NOT flow into the encoder + // — accumulating stale-state grad would poison the encoder). + // Encoder learns from π + V (on-policy current-step) + BCE + + // aux (supervised via perception.step_batched). Standard + // off-policy AC pattern (SAC / R2D2 / IMPALA do the same). + // // The combined grad slot feeds `backward_encoder_with_grad_h_t` // below, which runs the shared encoder backward chain that // `dispatch_train_step` also uses (single source of truth per @@ -1545,13 +1704,6 @@ impl IntegratedTrainer { self.stream .memset_zeros(&mut self.grad_h_t_combined_d) .map_err(|e| anyhow::anyhow!("zero grad_h_t_combined_d: {e}"))?; - accumulate_grad_h( - &self.stream, - &self.grad_h_accumulate_fn, - &q_grad_h_t_d, - lambdas.q, - &mut self.grad_h_t_combined_d, - )?; accumulate_grad_h( &self.stream, &self.grad_h_accumulate_fn, @@ -2053,22 +2205,67 @@ impl IntegratedTrainer { // post-R7b (the bootstrap-fallback host read of γ that the // flawed branch did is gone with the host advantage loop). - // ── Step 7: delegate to step_synthetic. ─────────────────────── - // step_synthetic now reads trainer-owned device buffers - // (actions_d, rewards_d, dones_d, next_actions_d, advantages_d, - // returns_d, log_pi_old_d) populated above. No host-slice - // arguments, no DtoH at the boundary. + // ── Step 7a (R7d): PER push + sample + gather. ──────────────── + // Push CURRENT step's b_size transitions (one per batch) to + // ReplayBuffer; sample b_size indices (priority^α from + // ISV[405]) and gather into sampled_*_d buffers that + // step_synthetic's Q forward + Bellman + backward consume. // - // step_synthetic re-runs forward_encoder (a compute redundancy - // the flawed F.4 already noted; lifting requires a deeper - // refactor and is a Phase R-future item). + // ORDER MATTERS: push BEFORE sample so the first + // step_with_lobsim call has ≥ b_size transitions to sample from + // (sample_indices on an empty buffer returns Vec::new(), which + // would force a Q-skip warmup path; pushing first eliminates + // that path and keeps Q always-on per `feedback_always_per`). + // + // Need an ISV host mirror refresh first so sample_and_gather + // reads ISV[405] for per_α. The refresh inside step_synthetic + // is later (Step 2 of step_synthetic) — too late for this + // pre-step_synthetic call. The cost is one DtoH of the full + // ISV slice (424 floats), bounded. + self.stream + .memcpy_dtoh(&self.isv_d, self.isv_host.as_mut_slice()) + .context("step_with_lobsim: isv dtoh (pre-PER)")?; + + self.push_to_replay(b_size) + .context("step_with_lobsim: push_to_replay")?; + let per_indices = self + .sample_and_gather(b_size) + .context("step_with_lobsim: sample_and_gather")?; + + // ── Step 7b: delegate to step_synthetic. ────────────────────── + // step_synthetic now reads: + // * SAMPLED buffers (sampled_h_t_d / sampled_h_tp1_d / + // sampled_actions_d / sampled_rewards_d / sampled_dones_d / + // sampled_next_actions_d) for the off-policy Q forward + + // Bellman target build + Q backward (PER's "always on" per + // `feedback_always_per`). + // * CURRENT-step buffers (h_t via perception.h_t_view(), + // actions_d, log_pi_old_d, advantages_d, returns_d) for + // the on-policy PPO surrogate + V regression. + // + // The encoder backward consumes ONLY π + V grad_h_t + // contributions (R7d stop-grad on Q's encoder grad — see + // step_synthetic Step 10 commentary). let stats = self.step_synthetic(snapshots)?; + // ── Step 7c (R7d): PER priority update. ─────────────────────── + // step_synthetic's Q backward populated self.td_per_sample_d + // (per-sample CE loss = |TD-error| analogue for distributional + // Q). DtoH and feed back to the replay buffer so the next + // step's sample picks the high-TD transitions per the + // canonical Schaul 2015 PER recipe. + if !per_indices.is_empty() { + let mut td_per_sample_host = vec![0.0f32; b_size]; + self.stream + .memcpy_dtoh(&self.td_per_sample_d, td_per_sample_host.as_mut_slice()) + .context("step_with_lobsim: dtoh td_per_sample_d")?; + self.replay + .update_priorities(&per_indices, &td_per_sample_host); + } + // Target-net soft update (Phase R5 + R6) — runs once per step // after the Q-head Adam update inside step_synthetic. Reads τ - // from ISV[401] (R5 reflects this step's q_divergence EMA if - // R7b wires the EMA producer; until then τ stays at R1's - // bootstrap 0.005). + // from ISV[401]. let isv_d_clone = self.isv_d.clone(); self.dqn_head .soft_update_target(&isv_d_clone) @@ -2077,6 +2274,168 @@ impl IntegratedTrainer { Ok(stats) } + /// Phase R7d: push the current step's b_size transitions to the + /// PER replay buffer. Each transition stores per-batch slices of + /// `perception.h_t_view()` and `self.h_tp1_d` as newly-allocated + /// `CudaSlice(HIDDEN_DIM)` device buffers (owned by the + /// Transition; freed when the buffer's random-replacement evicts + /// the transition or when the trainer drops). Per-batch metadata + /// (action, reward, done, log_pi_old) crosses the device→host + /// boundary via mapped-pinned `memcpy_dtoh` calls — accepted cost + /// of host-side PER bookkeeping (the device-side hot path stays + /// pure; only PER's control-plane sees host traffic, per the + /// canonical Phase R7d design call documented in the rebuild + /// plan A9 + `feedback_always_per`). + fn push_to_replay(&mut self, b_size: usize) -> Result<()> { + use crate::rl::common::{N_ACTIONS, Transition}; + + // DtoH per-batch metadata. 4 × b_size floats/ints total. + let mut actions_host = vec![0i32; b_size]; + let mut rewards_host = vec![0.0f32; b_size]; + let mut dones_host = vec![0.0f32; b_size]; + let mut log_pi_old_host = vec![0.0f32; b_size]; + self.stream + .memcpy_dtoh(&self.actions_d, actions_host.as_mut_slice()) + .context("push_to_replay: dtoh actions_d")?; + self.stream + .memcpy_dtoh(&self.rewards_d, rewards_host.as_mut_slice()) + .context("push_to_replay: dtoh rewards_d")?; + self.stream + .memcpy_dtoh(&self.dones_d, dones_host.as_mut_slice()) + .context("push_to_replay: dtoh dones_d")?; + self.stream + .memcpy_dtoh(&self.log_pi_old_d, log_pi_old_host.as_mut_slice()) + .context("push_to_replay: dtoh log_pi_old_d")?; + + // For each batch index, alloc per-transition device buffers + // and DtoD-copy the per-batch slice of h_t / h_tp1 in. + let hidden_nbytes = HIDDEN_DIM * std::mem::size_of::(); + for b in 0..b_size { + let mut h_t_per_b = self + .stream + .alloc_zeros::(HIDDEN_DIM) + .context("push_to_replay: alloc h_t per-batch")?; + let mut h_tp1_per_b = self + .stream + .alloc_zeros::(HIDDEN_DIM) + .context("push_to_replay: alloc h_tp1 per-batch")?; + + // Two separate inner scopes for the per-buffer DtoD so + // each `SyncOnDrop` guard drops before the next + // `device_ptr_mut` borrow / before the move into + // `Transition`. Without this, the borrow checker rejects + // the moves into `Transition { h_t, next_h_t, .. }` as + // overlapping the guard lifetimes (canonical borrow-checker + // dance for the cudarc raw-pointer API). + let off = (b * HIDDEN_DIM * std::mem::size_of::()) as u64; + unsafe { + let s = self.stream.cu_stream(); + let (h_t_base, _g1) = self.perception.h_t_view().device_ptr(&self.stream); + let (dst, _gd) = h_t_per_b.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, h_t_base + off, hidden_nbytes, s, + ) + .context("push_to_replay: DtoD per-batch h_t slice")?; + } + unsafe { + let s = self.stream.cu_stream(); + let (h_tp1_base, _g2) = self.h_tp1_d.device_ptr(&self.stream); + let (dst, _gd) = h_tp1_per_b.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, h_tp1_base + off, hidden_nbytes, s, + ) + .context("push_to_replay: DtoD per-batch h_tp1 slice")?; + } + self.replay.push(Transition { + h_t: h_t_per_b, + action: actions_host[b] as u32, + reward: rewards_host[b], + next_h_t: h_tp1_per_b, + done: dones_host[b] > 0.5, + log_pi_old: log_pi_old_host[b], + q_value_old: [0.0; N_ACTIONS], + }); + } + Ok(()) + } + + /// Phase R7d: sample `b_size` indices from PER (priority^α from + /// `ISV[RL_PER_ALPHA_INDEX=405]`) and gather the corresponding + /// transitions' h_t / h_tp1 device payloads + action / reward / + /// done metadata into trainer-owned `sampled_*_d` buffers via + /// per-batch DtoD + DtoH→HtoD round-trips. Returns the indices + /// so the caller can pass them back to + /// `replay.update_priorities(indices, td_per_sample_host)` after + /// the Q backward writes per-sample CE into `self.td_per_sample_d`. + fn sample_and_gather(&mut self, b_size: usize) -> Result> { + let per_alpha = self.isv_host[crate::rl::isv_slots::RL_PER_ALPHA_INDEX]; + // R1 bootstrap = 0.6; controller can drift it. Bound-check at + // the consumer site so a transient zero doesn't degenerate the + // sampling. ReplayBuffer.sample_indices already falls back to + // uniform when total weights are zero, so this is layered + // defense per `pearl_kernel_input_filter_must_match_validation_surface`. + let per_alpha_safe = if per_alpha > 0.0 { per_alpha } else { 0.6 }; + + let indices = self.replay.sample_indices(b_size, per_alpha_safe); + if indices.len() != b_size { + // Buffer hasn't been pushed yet (cold start, before + // push_to_replay) — return empty to signal "skip Q + // backward this step". Caller handles by zeroing λ_q + // for the off-policy path. + return Ok(indices); + } + + // Gather per-batch metadata into HtoD-stagable Vecs. + let mut actions_host = vec![0i32; b_size]; + let mut rewards_host = vec![0.0f32; b_size]; + let mut dones_host = vec![0.0f32; b_size]; + for (i, &idx) in indices.iter().enumerate() { + let t = &self.replay.transitions[idx]; + actions_host[i] = t.action as i32; + rewards_host[i] = t.reward; + dones_host[i] = if t.done { 1.0 } else { 0.0 }; + } + self.stream + .memcpy_htod(&actions_host, &mut self.sampled_actions_d) + .context("sample_and_gather: htod sampled_actions_d")?; + self.stream + .memcpy_htod(&rewards_host, &mut self.sampled_rewards_d) + .context("sample_and_gather: htod sampled_rewards_d")?; + self.stream + .memcpy_htod(&dones_host, &mut self.sampled_dones_d) + .context("sample_and_gather: htod sampled_dones_d")?; + + // Per-batch DtoD: per-transition h_t / h_tp1 device buffers → + // contiguous `sampled_h_t_d` / `sampled_h_tp1_d` at slot i. + let hidden_nbytes = HIDDEN_DIM * std::mem::size_of::(); + for (i, &idx) in indices.iter().enumerate() { + unsafe { + let s = self.stream.cu_stream(); + let off = (i * HIDDEN_DIM * std::mem::size_of::()) as u64; + let (src_ht, _g1) = self.replay.transitions[idx] + .h_t + .device_ptr(&self.stream); + let (src_htp1, _g2) = self.replay.transitions[idx] + .next_h_t + .device_ptr(&self.stream); + let (dst_ht_base, _gd1) = + self.sampled_h_t_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ht_base + off, src_ht, hidden_nbytes, s, + ) + .context("sample_and_gather: DtoD per-batch h_t")?; + let (dst_htp1_base, _gd2) = + self.sampled_h_tp1_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst_htp1_base + off, src_htp1, hidden_nbytes, s, + ) + .context("sample_and_gather: DtoD per-batch h_tp1")?; + } + } + + Ok(indices) + } + /// Launch `rl_lr_controller` to emit per-head learning rates into /// `ISV[412..417]`. Phase E.2-DEFER item 3. Diagnostic-input args /// are currently zero (the controller emits the bootstrap target diff --git a/crates/ml-alpha/tests/integrated_trainer_smoke.rs b/crates/ml-alpha/tests/integrated_trainer_smoke.rs index 4d514e890..f792aa3cc 100644 --- a/crates/ml-alpha/tests/integrated_trainer_smoke.rs +++ b/crates/ml-alpha/tests/integrated_trainer_smoke.rs @@ -77,6 +77,7 @@ fn integrated_trainer_step_with_lobsim_runs_without_panic() { perception, dqn_seed: 0xC0DE, ppo_seed: 0xC0DF, + ..IntegratedTrainerConfig::default() }; let mut trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); let snapshots = synthetic_window(4); diff --git a/crates/ml-alpha/tests/isv_bootstrap.rs b/crates/ml-alpha/tests/isv_bootstrap.rs index 1d6fdebfe..4c67cff50 100644 --- a/crates/ml-alpha/tests/isv_bootstrap.rs +++ b/crates/ml-alpha/tests/isv_bootstrap.rs @@ -62,6 +62,7 @@ fn g1_isv_bootstrap_writes_canonical_values() { }, dqn_seed: 0xCAFE, ppo_seed: 0xBEEF, + ..IntegratedTrainerConfig::default() }; let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); diff --git a/crates/ml-alpha/tests/r3_ema_advantage.rs b/crates/ml-alpha/tests/r3_ema_advantage.rs index 3f210c93a..32a18f46b 100644 --- a/crates/ml-alpha/tests/r3_ema_advantage.rs +++ b/crates/ml-alpha/tests/r3_ema_advantage.rs @@ -47,6 +47,7 @@ fn build_trainer() -> Option<(MlDevice, IntegratedTrainer)> { }, dqn_seed: 0xA1, ppo_seed: 0xA2, + ..IntegratedTrainerConfig::default() }; let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); Some((dev, trainer)) diff --git a/crates/ml-alpha/tests/r4_action_kernels.rs b/crates/ml-alpha/tests/r4_action_kernels.rs index 163061f53..38b8ae840 100644 --- a/crates/ml-alpha/tests/r4_action_kernels.rs +++ b/crates/ml-alpha/tests/r4_action_kernels.rs @@ -50,6 +50,7 @@ fn build_trainer() -> Option<(MlDevice, IntegratedTrainer)> { }, dqn_seed: 0xB4, ppo_seed: 0xB5, + ..IntegratedTrainerConfig::default() }; let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); Some((dev, trainer)) diff --git a/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs b/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs index f2ffef7cc..5fd254639 100644 --- a/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs +++ b/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs @@ -63,6 +63,7 @@ fn build_trainer() -> Option<(MlDevice, IntegratedTrainer)> { }, dqn_seed: 0xB7, ppo_seed: 0xB8, + ..IntegratedTrainerConfig::default() }; let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); Some((dev, trainer)) diff --git a/crates/ml-alpha/tests/r7d_per_wiring.rs b/crates/ml-alpha/tests/r7d_per_wiring.rs new file mode 100644 index 000000000..ab766609b --- /dev/null +++ b/crates/ml-alpha/tests/r7d_per_wiring.rs @@ -0,0 +1,145 @@ +//! Phase R7d gate G6: PER buffer wired into `step_with_lobsim`. +//! +//! Asserts the load-bearing invariants that distinguish a wired PER +//! buffer from R7c's dead `src/rl/replay.rs` struct: +//! +//! 1. `trainer.replay.len()` grows by exactly `b_size` per +//! `step_with_lobsim` call (the per-batch push order documented +//! in `IntegratedTrainer::push_to_replay`). +//! 2. `replay.sample_indices(b_size, α)` returns a vec of length +//! `b_size` after the first step (buffer non-empty post-push). +//! 3. After N steps with N × b_size ≤ capacity, `replay.len() == +//! N × b_size` (no replacement). After N steps with N × b_size +//! > capacity, `replay.len() == capacity` (ring-with-replacement +//! capped at capacity). +//! +//! Per `pearl_tests_must_prove_not_lock_observations`: asserts +//! buffer-mechanism invariants (length growth, sample size), NOT +//! observed Q values or losses — those vary across runs and lock +//! the test against any future change to Q init or PRNG state. +//! +//! Run with: +//! `cargo test -p ml-alpha --test r7d_per_wiring -- --ignored --nocapture` + +use ml_alpha::cfc::snap_features::Mbp10RawInput; +use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig}; +use ml_alpha::trainer::perception::PerceptionTrainerConfig; +use ml_backtesting::sim::LobSimCuda; +use ml_core::device::MlDevice; + +fn synthetic_window(seq_len: usize, base_mid: f32) -> Vec { + let mut out = Vec::with_capacity(seq_len); + let mut prev_mid = base_mid; + let mut ts_ns = 1_000_000_u64; + for _ in 0..seq_len { + let next_mid = prev_mid + 0.25; + let mut bid_px = [0.0_f32; 10]; + let mut bid_sz = [0.0_f32; 10]; + let mut ask_px = [0.0_f32; 10]; + let mut ask_sz = [0.0_f32; 10]; + for i in 0..10 { + bid_px[i] = next_mid - 0.125 - 0.25 * i as f32; + ask_px[i] = next_mid + 0.125 + 0.25 * i as f32; + bid_sz[i] = 10.0; + ask_sz[i] = 10.0; + } + let prev_ts = ts_ns; + ts_ns += 20_000_000; + out.push(Mbp10RawInput { + bid_px, + bid_sz, + ask_px, + ask_sz, + prev_mid, + trade_signed_vol: 0.0, + trade_count: 0, + ts_ns, + prev_ts_ns: prev_ts, + regime: [0.0; 6], + }); + prev_mid = next_mid; + } + out +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn r7d_per_buffer_grows_one_per_step_at_b_size_1() { + let dev = match MlDevice::cuda(0) { + Ok(d) => d, + Err(_) => { + eprintln!("CUDA 0 not available — skipping r7d_per_wiring"); + return; + } + }; + + // b_size = 1, small capacity so we can also verify the ring-cap + // invariant in the same test (no need for a separate fixture). + let per_capacity = 8usize; + let cfg = IntegratedTrainerConfig { + perception: PerceptionTrainerConfig { + seq_len: 4, + n_batch: 1, + ..PerceptionTrainerConfig::default() + }, + per_capacity, + ..IntegratedTrainerConfig::default() + }; + let mut trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new"); + let mut sim = LobSimCuda::new(1, &dev).expect("LobSimCuda::new"); + + // Invariant 1: buffer starts empty. + assert_eq!( + trainer.replay.len(), + 0, + "PER buffer must start empty before any step_with_lobsim call" + ); + + // Drive N=5 steps; assert linear growth from 0 → 5. + for step in 1..=5usize { + let snapshots = synthetic_window(4, 5500.0 + step as f32 * 0.25); + let next_snapshots = synthetic_window(4, 5500.0 + (step + 1) as f32 * 0.25); + let _stats = trainer + .step_with_lobsim(&snapshots, &next_snapshots, &mut sim) + .unwrap_or_else(|e| panic!("step_with_lobsim step {step}: {e:?}")); + assert_eq!( + trainer.replay.len(), + step, + "PER buffer must grow by exactly 1 per step at b_size=1 (step {step})" + ); + } + + // Invariant 2: sample at α=0.6 returns batch size 1. + let sample = trainer.replay.sample_indices(1, 0.6); + assert_eq!( + sample.len(), + 1, + "sample_indices(1, 0.6) on a non-empty buffer must return 1 index" + ); + assert!( + sample[0] < trainer.replay.len(), + "sampled index must be in [0, replay.len()) — got {} for len {}", + sample[0], + trainer.replay.len() + ); + + // Invariant 3: drive past capacity, buffer caps at `per_capacity`. + // Already at len=5; drive 10 more (total 15 transitions pushed, + // capacity=8 → buffer ends at exactly 8). + for step in 6..=15usize { + let snapshots = synthetic_window(4, 5500.0 + step as f32 * 0.25); + let next_snapshots = synthetic_window(4, 5500.0 + (step + 1) as f32 * 0.25); + trainer + .step_with_lobsim(&snapshots, &next_snapshots, &mut sim) + .unwrap_or_else(|e| panic!("step_with_lobsim step {step}: {e:?}")); + } + assert_eq!( + trainer.replay.len(), + per_capacity, + "PER buffer must cap at per_capacity = {per_capacity} (ring-with-replacement)" + ); + + eprintln!( + "R7d G6 OK — buffer grew 0→5→{per_capacity} across 15 push cycles; sample returns expected size" + ); +}