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

feat(rl): R7a — lift remaining host work in step_with_lobsim to GPU

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-23 11:30:14 +02:00
parent 7e7c8b5d90
commit aba8ec61b2
3 changed files with 319 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

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

23
crates/ml-alpha/cuda/abs_copy.cu Normal file
View File

@@ -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]);
}

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

@@ -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<CudaModule>,
actions_to_market_targets_fn: CudaFunction,
// ── Phase R7a: abs-copy + persistent per-step device buffers ──────
_abs_copy_module: Arc<CudaModule>,
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<i32>,
/// 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<f32>,
/// Per-step dones buffer (output of `extract_realized_pnl_delta`).
pub dones_d: CudaSlice<f32>,
/// 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<f32>,
/// 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<i32>,
/// 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<i32>,
/// 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<f32>,
/// Per-step advantages A_t. Output of R3's
/// `compute_advantage_return`.
pub advantages_d: CudaSlice<f32>,
/// Per-step returns R_t (V regression target). Output of R3's
/// `compute_advantage_return`.
pub returns_d: CudaSlice<f32>,
/// 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::<f32>(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::<f32>(b_size)
.context("alloc reward_abs_d")?;
let actions_d = stream
.alloc_zeros::<i32>(b_size)
.context("alloc actions_d")?;
let next_actions_d = stream
.alloc_zeros::<i32>(b_size)
.context("alloc next_actions_d")?;
let log_pi_old_d = stream
.alloc_zeros::<f32>(b_size)
.context("alloc log_pi_old_d")?;
let advantages_d = stream
.alloc_zeros::<f32>(b_size)
.context("alloc advantages_d")?;
let returns_d = stream
.alloc_zeros::<f32>(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<f32>,
dst_d: &mut CudaSlice<f32>,
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<IntegratedStepStats> {
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::<f32>(b_size * Q_N_ATOMS)?;
let mut target_dist_d = self.stream.alloc_zeros::<f32>(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<i32> = actions.iter().map(|&a| a as i32).collect();
let next_actions_i32: Vec<i32> = 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::<f32>()) 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