feat(rl): R7d — PER wired + off-policy DQN with stop-grad on encoder
Closes plan A9 (rebuild plan's "PER wiring" R7 scope second half;
R7c-data shipped the first half last commit). The `ReplayBuffer` in
`src/rl/replay.rs` has sat as dead code since Phase C — this commit
makes it load-bearing per `feedback_always_per` ("PER always enabled;
non-PER paths are dead code").
## Architecture: off-policy Q + on-policy PPO + V + stop-grad encoder
Shared-encoder pattern with the canonical off-policy + shared-encoder
discipline: the Q head trains from PER-sampled past transitions,
PPO + V train on current-step on-policy data, and the encoder receives
gradient signal ONLY from PPO + V (and BCE/aux via the perception
trainer's separate `step_batched` path). Standard pattern in SAC,
R2D2, IMPALA.
Stop-grad is implemented by computing Q's `grad_h_t` (via
`backward_to_w_b_h(sampled_h_t, ...)`) but NOT accumulating it into
`grad_h_t_combined_d` — the encoder backward only sees π + V
contributions. Per `feedback_no_hiding` the discarded buffer is
allocated and written (the kernel API requires the writeback target);
the discard is a deliberate design call documented at the
accumulation site.
## Wiring summary
### Kernel: `dqn_distributional_q_bwd`
* New `loss_per_batch [B]` output. Atom 0 of each block writes the
per-sample CE loss (non-atomic — single writer per batch).
`loss_out [1]` continues to atomicAdd the scalar sum for the
diagnostic total. Build.rs cache bust v30.
### `DqnHead::backward_logits` (Rust wrapper)
* New `loss_per_batch: &mut CudaSlice<f32>` arg. Migrated atomically
in the same commit per `feedback_no_partial_refactor` — only
caller is the integrated trainer.
### `IntegratedTrainerConfig`
* New `per_capacity: usize` (default 4096, matches `replay.rs` doc
ceiling for naive O(N) sampling).
* New `per_seed: u64` (default 0x9E37_79B9_7F4A_7C15).
* `Default` impl added so test fixtures forward-compat via
`..IntegratedTrainerConfig::default()`. All 5 existing test
fixtures migrated.
### `IntegratedTrainer`
* New fields: `replay: ReplayBuffer`, `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`.
* New methods: `push_to_replay(b_size)` — DtoH per-batch metadata
(action/reward/done/log_pi_old) + alloc per-transition
`CudaSlice<f32>(HIDDEN_DIM)` ×2 + DtoD per-batch slice copies +
push to `ReplayBuffer`. `sample_and_gather(b_size)` — read
per_α from ISV[405], call `replay.sample_indices`, gather sampled
transitions' h_t/h_tp1 device payloads via per-batch DtoD into
`sampled_h_t_d` / `sampled_h_tp1_d`, HtoD upload action/reward/done.
### `step_with_lobsim` orchestration
After `compute_advantage_return` and BEFORE `step_synthetic`:
1. DtoH full ISV slice to refresh `isv_host` (so PER reads ISV[405]
for per_α).
2. `push_to_replay(b_size)` — push current step's transitions.
3. `sample_and_gather(b_size)` — return `per_indices` for the
priority update.
4. `step_synthetic(snapshots)` — runs π + V on current-step h_t,
Q on SAMPLED h_t (off-policy).
5. DtoH `td_per_sample_d` → host; `replay.update_priorities(
per_indices, td_per_sample_host)`.
6. Target-net soft update (unchanged from R5).
### `step_synthetic` redirects (Q path → sampled, π/V stay on-policy)
* Q forward: `forward(&self.sampled_h_t_d)` (was `h_t_borrow`).
* New: forward online Q on `&self.sampled_h_tp1_d` → local scratch +
`argmax_expected_q` → `self.sampled_next_actions_d`. The
Double-DQN argmax MUST be recomputed each step (online net weights
drift faster than transitions recycle through replay; storing
argmax at push time would feed stale-action data into the
projection).
* `forward_target(&self.sampled_h_tp1_d)` (was `&self.h_tp1_d`).
* `select_action_atoms(..., &self.sampled_next_actions_d, ...)`
(was `&self.next_actions_d`).
* `project_bellman_target(..., &self.sampled_rewards_d,
&self.sampled_dones_d, ...)` (was `rewards_d` / `dones_d`).
* `backward_logits(..., &self.sampled_actions_d, ...,
&mut self.td_per_sample_d, ...)` (added per-sample loss output).
* `backward_to_w_b_h(&self.sampled_h_t_d, ...)` (was `h_t_borrow`).
* Q grad_h_t accumulation REMOVED from Step 10 (stop-grad).
## Test: r7d_per_wiring.rs (gate G6)
Three invariants per `pearl_tests_must_prove_not_lock_observations`:
1. `replay.len()` grows by exactly `b_size` per `step_with_lobsim`
call (push semantics).
2. `sample_indices(b_size, α)` returns vec of length `b_size` on a
non-empty buffer.
3. Buffer caps at `per_capacity` (ring-with-random-replacement).
Drives 15 steps with `per_capacity=8`, asserts growth 0→5→8 across
the cap boundary.
## Acceptable host traffic this commit adds
* Per-step DtoH of 4 × b_size scalars (action/reward/done/log_pi_old)
for PER push metadata.
* Per-step DtoH of b_size floats (td_per_sample_d) for
update_priorities.
* Per-step HtoD of 3 × b_size scalars (sampled action/reward/done)
for sampled metadata gather.
* Per-step DtoD of 2 × b_size × HIDDEN_DIM floats (per-batch h_t /
h_tp1 slices) for PER push + gather.
PER bookkeeping is a control-plane operation by design (host-side
priority/index management); the device-side training hot path
(encoder, Q/π/V forward/backward, Adam) stays GPU-pure. GPU sum-tree
+ device-resident transitions are a Phase R-future optimization
flagged in `replay.rs`'s doc.
## What's NOT in this commit
* Q `loss_per_batch [B]` is now wired through `backward_logits` but
the DtoH happens inside step_with_lobsim (not inside
step_synthetic). Earlier R7d sketches considered a separate
`dqn_offpolicy_step` method; the in-step_synthetic redirect
approach landed because it touches fewer lines + reuses the
existing scratch buffer allocations + matches the trainer's
established λ-weighted multi-head pattern. A future refactor
could split for clarity.
Local sm_86 smoke gates: `cargo test -p ml-alpha --test
r7d_per_wiring -- --ignored --nocapture` (G6) +
`integrated_trainer_smoke` (end-to-end). Cluster smoke deferred to
R9.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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");
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -327,12 +327,18 @@ impl DqnHead {
|
||||
actions_taken: &CudaSlice<i32>,
|
||||
b_size: usize,
|
||||
loss_out: &mut CudaSlice<f32>,
|
||||
loss_per_batch: &mut CudaSlice<f32>,
|
||||
grad_logits: &mut CudaSlice<f32>,
|
||||
) -> 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")?;
|
||||
|
||||
@@ -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<f32>,
|
||||
|
||||
// ── 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<f32>,
|
||||
|
||||
/// 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<f32>,
|
||||
|
||||
/// 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<i32>,
|
||||
|
||||
/// Per-step gather buffer for PER-sampled rewards. Length `B`.
|
||||
pub sampled_rewards_d: CudaSlice<f32>,
|
||||
|
||||
/// Per-step gather buffer for PER-sampled done flags. Length `B`.
|
||||
pub sampled_dones_d: CudaSlice<f32>,
|
||||
|
||||
/// 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<i32>,
|
||||
|
||||
/// 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<f32>,
|
||||
|
||||
/// 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::<f32>(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<f32>(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::<f32>(b_size * HIDDEN_DIM)
|
||||
.context("alloc sampled_h_t_d")?;
|
||||
let sampled_h_tp1_d = stream
|
||||
.alloc_zeros::<f32>(b_size * HIDDEN_DIM)
|
||||
.context("alloc sampled_h_tp1_d")?;
|
||||
let sampled_actions_d = stream
|
||||
.alloc_zeros::<i32>(b_size)
|
||||
.context("alloc sampled_actions_d")?;
|
||||
let sampled_rewards_d = stream
|
||||
.alloc_zeros::<f32>(b_size)
|
||||
.context("alloc sampled_rewards_d")?;
|
||||
let sampled_dones_d = stream
|
||||
.alloc_zeros::<f32>(b_size)
|
||||
.context("alloc sampled_dones_d")?;
|
||||
let sampled_next_actions_d = stream
|
||||
.alloc_zeros::<i32>(b_size)
|
||||
.context("alloc sampled_next_actions_d")?;
|
||||
let td_per_sample_d = stream
|
||||
.alloc_zeros::<f32>(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::<f32>(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<f32>(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::<f32>();
|
||||
for b in 0..b_size {
|
||||
let mut h_t_per_b = self
|
||||
.stream
|
||||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||||
.context("push_to_replay: alloc h_t per-batch")?;
|
||||
let mut h_tp1_per_b = self
|
||||
.stream
|
||||
.alloc_zeros::<f32>(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::<f32>()) 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<Vec<usize>> {
|
||||
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::<f32>();
|
||||
for (i, &idx) in indices.iter().enumerate() {
|
||||
unsafe {
|
||||
let s = self.stream.cu_stream();
|
||||
let off = (i * HIDDEN_DIM * std::mem::size_of::<f32>()) 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
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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");
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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))
|
||||
|
||||
145
crates/ml-alpha/tests/r7d_per_wiring.rs
Normal file
145
crates/ml-alpha/tests/r7d_per_wiring.rs
Normal file
@@ -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<Mbp10RawInput> {
|
||||
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"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user