From aba8ec61b2eda0bbe1a1b221b47df632a4f1475f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 23 May 2026 11:30:14 +0200 Subject: [PATCH] =?UTF-8?q?feat(rl):=20R7a=20=E2=80=94=20lift=20remaining?= =?UTF-8?q?=20host=20work=20in=20step=5Fwith=5Flobsim=20to=20GPU?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Honors R6's commit-message promise to close the host-work boundary the partial GPU-purity left behind. After R7a, step_with_lobsim's post-fill pipeline is GPU-resident through the entire training step except for 3 small host-slice uploads (actions, next_actions, log_pi_old) that R7b lifts via R4's Thompson/argmax/log_pi kernels. CHANGES: 1. New cuda/abs_copy.cu — element-wise dst[b] = fabsf(src[b]). Feeds the |reward| signal into ema_update_on_done for the MEAN_ABS_PNL_EMA slot without mutating signed rewards_d. 2. New trainer-owned per-step device buffers (allocated once in new(), reused every step — no per-call churn): reward_abs_d, actions_d, next_actions_d, log_pi_old_d, advantages_d, returns_d 3. step_synthetic signature change: drops the 7 synthetic_* host slice args (synthetic_actions, synthetic_rewards, synthetic_dones, synthetic_next_actions, synthetic_advantages, synthetic_returns, synthetic_log_pi_old). The body now reads from trainer-owned device buffers (self.actions_d, self.rewards_d, etc.) via the disjoint-field borrow rule. The 7 upload_i32/upload_f32 calls are deleted — caller (step_with_lobsim) populates the buffers via GPU kernels before invoking step_synthetic. 4. step_with_lobsim Step 6 rewritten end-to-end: - Upload host Thompson outputs (actions, next_actions, log_pi_old) to trainer buffers — R7b removes these via R4's GPU kernels. - GPU abs_copy(rewards_d) → reward_abs_d. - GPU ema_update_on_done(MEAN_ABS_PNL_EMA, reward_abs_d, dones_d). - GPU launch_rl_controllers_per_step (R5) — all 7 controllers adapt to this step's EMA inputs. - GPU apply_reward_scale(rewards_d) in-place — reads the freshly updated ISV[406] from the controller. - GPU compute_advantage_return(rewards_d, dones_d, v_t, v_tp1) → returns_d, advantages_d. - step_synthetic(snapshots) consumes all trainer device buffers, runs the training kernel chain, returns stats. - dqn_head.soft_update_target(isv_d) — R5 target-net Polyak update with τ from ISV[401]. R7a PARTIAL — work R7b lifts: - Host Thompson sampler still produces actions/next_actions/log_pi (R7b replaces with R4 rl_action_kernel + argmax_expected_q + log_pi_at_action — eliminating the 3 remaining HtoD uploads). - v_pred_host → v_pred_d HtoD round-trip (the host already has v_pred_host from the action-sampling Thompson read; R7b keeps V on device throughout). - v_tp1_d uses v_pred_host (V(s_t) approximated as V(s_{t+1}) until R7b wires next_snapshots + forward_encoder(next_snapshots)). The 4 final DtoH copies the R6 commit message warned about are GONE. Hot path: snapshot upload (boundary HtoD), ISV diagnostic readback (HEALTH_DIAG), and 3 host-Thompson uploads (R7b removes). cargo check + cargo build --tests on ml-alpha green. Tests still compile against the new step_synthetic signature (no test calls it directly — only step_with_lobsim does). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/abs_copy.cu | 23 ++ crates/ml-alpha/src/trainer/integrated.rs | 369 +++++++++++++++++----- 3 files changed, 319 insertions(+), 74 deletions(-) create mode 100644 crates/ml-alpha/cuda/abs_copy.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 6114fc203..73fcde390 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -54,6 +54,7 @@ const KERNELS: &[&str] = &[ "extract_realized_pnl_delta", // RL Phase R6: GPU-pure reward + done extraction from device Pos array; replaces host read_pos loop per feedback_cpu_is_read_only "apply_reward_scale", // RL Phase R6: element-wise rewards *= ISV[RL_REWARD_SCALE_INDEX=406]; closes the F.3b host roundtrip "actions_to_market_targets", // RL Phase R6: 9-action grid → LobSim market_targets[B*2] on device; replaces host submit_market loop per feedback_cpu_is_read_only + "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 diff --git a/crates/ml-alpha/cuda/abs_copy.cu b/crates/ml-alpha/cuda/abs_copy.cu new file mode 100644 index 000000000..0b9a73fe9 --- /dev/null +++ b/crates/ml-alpha/cuda/abs_copy.cu @@ -0,0 +1,23 @@ +// abs_copy.cu — element-wise absolute-value copy: dst[b] = fabsf(src[b]) +// (Phase R7a of the integrated RL trainer rebuild; +// see docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md). +// +// Feeds the |reward| signal into `ema_update_on_done` for the +// `MEAN_ABS_PNL_EMA` slot (ISV[423], input to +// rl_reward_scale_controller). The trainer keeps signed `rewards_d` +// for the advantage / return / reward-scale-apply pipeline; this +// kernel writes the magnitude into a separate `reward_abs_d` scratch +// so the EMA producer reads the magnitude without mutating the signed +// reward stream. +// +// Element-wise, trivially parallel. No reduction, no atomics. + +extern "C" __global__ void abs_copy( + const float* __restrict__ src, + float* __restrict__ dst, + int b_size +) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= b_size) return; + dst[b] = fabsf(src[b]); +} diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 1c6758a67..abfcd5206 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -162,6 +162,11 @@ const APPLY_REWARD_SCALE_CUBIN: &[u8] = const ACTIONS_TO_MARKET_TARGETS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/actions_to_market_targets.cubin")); +// Phase R7a: element-wise abs-copy kernel for the |reward| signal +// feeding ema_update_on_done's MEAN_ABS_PNL_EMA slot. +const ABS_COPY_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/abs_copy.cubin")); + /// Per-head LR controller Wiener-α target. The controller kernel floors /// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the /// effective blend rate stays ≥ 0.4 regardless of this seed value. Kept @@ -313,6 +318,10 @@ pub struct IntegratedTrainer { _actions_to_market_targets_module: Arc, actions_to_market_targets_fn: CudaFunction, + // ── Phase R7a: abs-copy + persistent per-step device buffers ────── + _abs_copy_module: Arc, + abs_copy_fn: CudaFunction, + /// Trainer-owned snapshot of `lobsim.pos.realized_pnl` taken at end /// of previous step. `extract_realized_pnl_delta` reads /// `lobsim.pos_d` post-fill, computes @@ -324,10 +333,33 @@ pub struct IntegratedTrainer { /// Trainer-owned snapshot of `lobsim.pos.position_lots` for the /// done-flag detection: `done = (prev != 0 && current == 0)`. pub prev_position_lots_d: CudaSlice, - /// Per-step rewards buffer (output of `extract_realized_pnl_delta`). + /// Per-step rewards buffer (output of `extract_realized_pnl_delta`, + /// in-place modified by `apply_reward_scale`). pub rewards_d: CudaSlice, /// Per-step dones buffer (output of `extract_realized_pnl_delta`). pub dones_d: CudaSlice, + /// Per-step |reward| scratch (output of `abs_copy(rewards_d)`). + /// Feeds `ema_update_on_done` for the `MEAN_ABS_PNL_EMA` slot — + /// keeping signed `rewards_d` intact for the + /// `compute_advantage_return` + `apply_reward_scale` pipeline. + pub reward_abs_d: CudaSlice, + /// Per-step actions (Thompson sample). R7a keeps the host Thompson + /// loop and uploads here; R7b lifts to R4's `rl_action_kernel`. + pub actions_d: CudaSlice, + /// Per-step next-state argmax actions for the Bellman target. + /// R7a keeps the host argmax loop and uploads here; R7b lifts to + /// R4's `argmax_expected_q` on h_{t+1}. + pub next_actions_d: CudaSlice, + /// Per-step log π_old at the sampled action. R7a keeps the host + /// log-softmax loop and uploads here; R7b lifts to R4's + /// `log_pi_at_action`. + pub log_pi_old_d: CudaSlice, + /// Per-step advantages A_t. Output of R3's + /// `compute_advantage_return`. + pub advantages_d: CudaSlice, + /// Per-step returns R_t (V regression target). Output of R3's + /// `compute_advantage_return`. + pub returns_d: CudaSlice, /// Combined encoder grad slot — folded from Q + π + V `grad_h_t` /// contributions via `grad_h_accumulate_scaled` (item 1). Consumed by @@ -550,6 +582,14 @@ impl IntegratedTrainer { .load_function("actions_to_market_targets") .context("load actions_to_market_targets")?; + // Phase R7a kernel. + let abs_copy_module = ctx + .load_cubin(ABS_COPY_CUBIN.to_vec()) + .context("load abs_copy cubin")?; + let abs_copy_fn = abs_copy_module + .load_function("abs_copy") + .context("load abs_copy")?; + // Per-batch PRNG state for the Thompson sampler. Seeded // deterministically from cfg.dqn_seed via ChaCha8 host RNG so // (cfg.dqn_seed, b_size) → identical Thompson draws across @@ -605,6 +645,27 @@ impl IntegratedTrainer { .alloc_zeros::(b_size) .context("alloc dones_d")?; + // Phase R7a per-step buffers (persistent so step_with_lobsim + // doesn't churn allocations every call). + let reward_abs_d = stream + .alloc_zeros::(b_size) + .context("alloc reward_abs_d")?; + let actions_d = stream + .alloc_zeros::(b_size) + .context("alloc actions_d")?; + let next_actions_d = stream + .alloc_zeros::(b_size) + .context("alloc next_actions_d")?; + let log_pi_old_d = stream + .alloc_zeros::(b_size) + .context("alloc log_pi_old_d")?; + let advantages_d = stream + .alloc_zeros::(b_size) + .context("alloc advantages_d")?; + let returns_d = stream + .alloc_zeros::(b_size) + .context("alloc returns_d")?; + Ok(Self { cfg, perception, @@ -660,10 +721,18 @@ impl IntegratedTrainer { apply_reward_scale_fn, _actions_to_market_targets_module: actions_to_market_targets_module, actions_to_market_targets_fn, + _abs_copy_module: abs_copy_module, + abs_copy_fn, prev_realized_pnl_d, prev_position_lots_d, rewards_d, dones_d, + reward_abs_d, + actions_d, + next_actions_d, + log_pi_old_d, + advantages_d, + returns_d, grad_h_t_combined_d, last_pi_loss: 0.0, step_counter: 0, @@ -1025,6 +1094,32 @@ impl IntegratedTrainer { Ok(()) } + /// Phase R7a: launch `abs_copy` — element-wise `dst[b] = fabsf(src[b])`. + /// Feeds the `|reward|` signal into `ema_update_on_done` for the + /// `MEAN_ABS_PNL_EMA` slot without mutating the signed `rewards_d`. + pub fn launch_abs_copy( + &self, + src_d: &CudaSlice, + dst_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(src_d.len(), b_size); + debug_assert_eq!(dst_d.len(), b_size); + let grid_x = ((b_size as u32) + 31) / 32; + let cfg = LaunchConfig { + grid_dim: (grid_x.max(1), 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.abs_copy_fn); + launch.arg(src_d).arg(dst_d).arg(&b_size_i); + unsafe { + launch.launch(cfg).context("abs_copy launch")?; + } + Ok(()) + } + /// Phase R6: launch `apply_reward_scale` — element-wise /// `rewards[b] *= isv[RL_REWARD_SCALE_INDEX = 406]`. R7 wires /// this into `step_with_lobsim` between @@ -1115,30 +1210,32 @@ impl IntegratedTrainer { /// /// Returns the per-head kernel-emitted losses and the combined-loss /// total weighted by the ISV-driven λs. - #[allow(clippy::too_many_arguments)] + /// Phase R7a refactor: step_synthetic reads the per-step inputs from + /// trainer-owned device buffers (`self.actions_d`, `self.rewards_d`, + /// `self.dones_d`, `self.next_actions_d`, `self.advantages_d`, + /// `self.returns_d`, `self.log_pi_old_d`) instead of taking host + /// slices. Caller's responsibility (step_with_lobsim) is to + /// populate those buffers via GPU kernels before calling. This + /// eliminates the 7 host-slice uploads (+4 DtoH copies in + /// step_with_lobsim's old handoff) that violated + /// `feedback_cpu_is_read_only` in the flawed Phase F+G branch. pub fn step_synthetic( &mut self, snapshots: &[Mbp10RawInput], - synthetic_actions: &[u32], // [B] — taken action index - synthetic_rewards: &[f32], // [B] — per-transition reward r_t - synthetic_dones: &[f32], // [B] — 0/1 terminal flag (γ × (1-done)) - synthetic_next_actions: &[u32], // [B] — argmax-Q action at s_{t+1} - synthetic_advantages: &[f32], // [B] — synthetic A_t - synthetic_returns: &[f32], // [B] — synthetic R_t for V regression - synthetic_log_pi_old: &[f32], // [B] — synthetic log π_old ) -> Result { - let b_size = synthetic_actions.len(); + let b_size = self.cfg.perception.n_batch; if b_size == 0 { anyhow::bail!("step_synthetic: empty batch (b_size = 0)"); } - // Validate per-batch input lengths. - debug_assert_eq!(synthetic_rewards.len(), b_size); - debug_assert_eq!(synthetic_dones.len(), b_size); - debug_assert_eq!(synthetic_next_actions.len(), b_size); - debug_assert_eq!(synthetic_advantages.len(), b_size); - debug_assert_eq!(synthetic_returns.len(), b_size); - debug_assert_eq!(synthetic_log_pi_old.len(), b_size); + // Validate per-batch input buffer lengths. + debug_assert_eq!(self.actions_d.len(), b_size); + debug_assert_eq!(self.rewards_d.len(), b_size); + debug_assert_eq!(self.dones_d.len(), b_size); + debug_assert_eq!(self.next_actions_d.len(), b_size); + debug_assert_eq!(self.advantages_d.len(), b_size); + debug_assert_eq!(self.returns_d.len(), b_size); + debug_assert_eq!(self.log_pi_old_d.len(), b_size); // ── Step 1: encoder forward ────────────────────────────────── let _ = self @@ -1236,14 +1333,9 @@ impl IntegratedTrainer { let mut q_target_action_d = self.stream.alloc_zeros::(b_size * Q_N_ATOMS)?; let mut target_dist_d = self.stream.alloc_zeros::(b_size * Q_N_ATOMS)?; - // ── Step 4: upload synthetic inputs via mapped-pinned ──────── - let actions_d = upload_i32(&self.stream, synthetic_actions)?; - let next_actions_d = upload_i32(&self.stream, synthetic_next_actions)?; - let rewards_d = upload_f32(&self.stream, synthetic_rewards)?; - let dones_d = upload_f32(&self.stream, synthetic_dones)?; - let advantages_d = upload_f32(&self.stream, synthetic_advantages)?; - let returns_d = upload_f32(&self.stream, synthetic_returns)?; - let log_pi_old_d = upload_f32(&self.stream, synthetic_log_pi_old)?; + // ── Step 4 (Phase R7a): inputs are trainer-owned device buffers, + // populated by step_with_lobsim's GPU pipeline before this call. + // No host uploads here. // ── Step 5: forward kernels (Q, π, V) ──────────────────────── self.dqn_head @@ -1264,9 +1356,9 @@ impl IntegratedTrainer { self.policy_head .surrogate_forward( &pi_logits_d, - &log_pi_old_d, - &actions_d, - &advantages_d, + &self.log_pi_old_d, + &self.actions_d, + &self.advantages_d, &self.isv_d, b_size, &mut pi_log_prob_d, @@ -1294,7 +1386,7 @@ impl IntegratedTrainer { self.dqn_head .select_action_atoms( &q_target_full_d, - &next_actions_d, + &self.next_actions_d, b_size, &mut q_target_action_d, ) @@ -1302,8 +1394,8 @@ impl IntegratedTrainer { self.dqn_head .project_bellman_target( &q_target_action_d, - &rewards_d, - &dones_d, + &self.rewards_d, + &self.dones_d, &self.isv_d, b_size, &mut target_dist_d, @@ -1316,7 +1408,7 @@ impl IntegratedTrainer { .backward_logits( &q_logits_d, &target_dist_d, - &actions_d, + &self.actions_d, b_size, &mut q_loss_d_mut, &mut q_grad_logits_d, @@ -1342,9 +1434,9 @@ impl IntegratedTrainer { self.policy_head .surrogate_backward_logits( &pi_logits_d, - &log_pi_old_d, - &actions_d, - &advantages_d, + &self.log_pi_old_d, + &self.actions_d, + &self.advantages_d, &self.isv_d, b_size, &mut pi_grad_logits_d, @@ -1375,7 +1467,7 @@ impl IntegratedTrainer { .backward( h_t_borrow, &v_pred_d, - &returns_d, + &self.returns_d, b_size, &mut v_loss_per_batch_d, &mut v_grad_w_per_batch_d, @@ -1795,47 +1887,176 @@ impl IntegratedTrainer { } } - // Phase R6 hand-off boundary: DtoH rewards + dones for the - // host-side advantage/return loop below + the eventual - // step_after_encoder_forward call (which currently takes - // host slices). The host-side handoff is bounded — R7 lifts - // step_after_encoder_forward to take device buffers, - // eliminating this DtoH per - // `feedback_cpu_is_read_only`. - let rewards = read_slice_d(&self.stream, &self.rewards_d, b_size)?; - let dones = read_slice_d(&self.stream, &self.dones_d, b_size)?; + // ── Step 6 (Phase R7a): GPU-pure post-fill pipeline. ────────── + // All advantage/return + reward-scaling + EMA + per-step + // controller launches stay on device. Replaces the host loops + // the flawed Phase F+G shipped per `feedback_cpu_is_read_only`. - // ── Step 6: derive advantages + returns from real reward + V. ─ - // One-step TD advantage: A_t = r_t + γ(1-done)·V(s_t) - V(s_t) - // = r_t - γ·done·V(s_t) - V(s_t)·(1-γ(1-done)) - // (we treat s_{t+1} ≈ s_t for Phase E.3b — Phase F threads the - // next-state encoder forward). - // R_t = r_t + γ(1-done)·V(s_t) - let mut advantages = vec![0.0_f32; b_size]; - let mut returns = vec![0.0_f32; b_size]; - for b in 0..b_size { - let v_b = v_pred_host[b]; - let r_b = rewards[b]; - let done_b = dones[b]; - let r_t = r_b + gamma * (1.0 - done_b) * v_b; - returns[b] = r_t; - advantages[b] = r_t - v_b; + // Upload Thompson-sampled actions + log_pi + next_actions to + // trainer-owned device buffers. R7b lifts Thompson + log_pi + + // next_actions to R4's GPU kernels; the host upload here is + // bounded (3 small buffers of size b_size) and goes away + // entirely after R7b. The u32 → i32 cast matches the kernel + // signature (CUDA `int`). + let actions_i32: Vec = actions.iter().map(|&a| a as i32).collect(); + let next_actions_i32: Vec = next_actions.iter().map(|&a| a as i32).collect(); + self.stream + .memcpy_htod(actions_i32.as_slice(), &mut self.actions_d) + .context("htod actions_d")?; + self.stream + .memcpy_htod(next_actions_i32.as_slice(), &mut self.next_actions_d) + .context("htod next_actions_d")?; + self.stream + .memcpy_htod(log_pi_old.as_slice(), &mut self.log_pi_old_d) + .context("htod log_pi_old_d")?; + + // Upload V(s_t) into a temporary device buffer so + // compute_advantage_return can read it. R-future: keep V(s_t) + // on device throughout (value_head.forward already writes to + // a device buffer; this round-trip exists because v_pred_host + // was already DtoH'd above for the Thompson loop's + // diagnostics). + let v_pred_d = upload_f32(&self.stream, &v_pred_host)?; + // R7a partial: V(s_{t+1}) approximated as V(s_t) until R7b + // wires next_snapshots + forward_encoder(next_snapshots). + // The compute_advantage_return kernel will then produce the + // canonical A_t = r + γV(s_{t+1}) − V(s_t) with the real + // h_{t+1}-derived value. + let v_tp1_d = upload_f32(&self.stream, &v_pred_host)?; + + // |reward| → reward_abs_d (input for mean_abs_pnl EMA). + // Inline launch (vs `launch_abs_copy`) because the per-method + // helper would take `&self` + `&mut CudaSlice` separately, + // conflicting with the disjoint-field borrow split here. + { + let cfg = LaunchConfig { + grid_dim: (((b_size as u32) + 31) / 32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.abs_copy_fn); + launch + .arg(&self.rewards_d) + .arg(&mut self.reward_abs_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("abs_copy launch")?; + } } - // ── Step 7: delegate to step_synthetic for the train step. ─── - // step_synthetic re-runs forward_encoder (compute redundancy, - // see method docs); we accept that for Phase E.3b in service of - // a single backward + Adam codepath. - self.step_synthetic( - snapshots, - &actions, - &rewards, - &dones, - &next_actions, - &advantages, - &returns, - &log_pi_old, - ) + // Update ISV[423] MEAN_ABS_PNL_EMA from |reward| over closed + // trades (dones_d gating). R5 controllers consume this on + // their next per-step launch below. + { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (b_size as u32, 1, 1), + shared_mem_bytes: (2 * b_size * std::mem::size_of::()) as u32, + }; + let slot_i = crate::rl::isv_slots::RL_MEAN_ABS_PNL_EMA_INDEX as i32; + let alpha = RL_LR_CONTROLLER_ALPHA; + let mut launch = self.stream.launch_builder(&self.ema_update_on_done_fn); + launch + .arg(&self.isv_d) + .arg(&slot_i) + .arg(&alpha) + .arg(&self.reward_abs_d) + .arg(&self.dones_d) + .arg(&b_size_i); + unsafe { + launch + .launch(cfg) + .context("ema_update_on_done(MEAN_ABS_PNL) launch")?; + } + } + + // Fire all 7 RL controllers — each reads its ISV EMA-input + // slot and Wiener-blends its output. Critical: this happens + // BEFORE apply_reward_scale so the scale ISV[406] reflects + // this step's mean_abs_pnl EMA update. + self.launch_rl_controllers_per_step() + .context("launch_rl_controllers_per_step")?; + + // Apply the reward scale on device (in place on rewards_d). + // Reads ISV[406] which the controller just updated. + { + let rewards_clone = self.rewards_d.clone(); + let _ = rewards_clone; // keep clone alive for the duration + } + // Borrow-checker note: launch_apply_reward_scale takes &self + // + &mut CudaSlice; calling self.launch_apply_reward_scale + // on self.rewards_d directly works under disjoint-field rule. + // Use a scoped block so the &mut on self.rewards_d ends + // before the next self.* call. + { + // Have to break the borrow: take a local mutable handle + // via std::mem::replace temporarily. Cleaner: just inline + // the launch here. + let cfg = LaunchConfig { + grid_dim: (((b_size as u32) + 31) / 32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.apply_reward_scale_fn); + launch + .arg(&self.isv_d) + .arg(&mut self.rewards_d) + .arg(&b_size_i); + unsafe { + launch.launch(cfg).context("apply_reward_scale launch")?; + } + } + + // GPU compute_advantage_return: returns_d, advantages_d + // populated for step_synthetic to consume. + { + let cfg = LaunchConfig { + grid_dim: (((b_size as u32) + 31) / 32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self + .stream + .launch_builder(&self.compute_advantage_return_fn); + launch + .arg(&self.isv_d) + .arg(&self.rewards_d) + .arg(&self.dones_d) + .arg(&v_pred_d) + .arg(&v_tp1_d) + .arg(&mut self.returns_d) + .arg(&mut self.advantages_d) + .arg(&b_size_i); + unsafe { + launch + .launch(cfg) + .context("compute_advantage_return launch")?; + } + } + let _ = gamma; // gamma is now ISV-resident; host fallback no longer needed here + + // ── 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_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). + let stats = self.step_synthetic(snapshots)?; + + // 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). + let isv_d_clone = self.isv_d.clone(); + self.dqn_head + .soft_update_target(&isv_d_clone) + .context("dqn_head.soft_update_target")?; + + Ok(stats) } /// Launch `rl_lr_controller` to emit per-head learning rates into