95dcc4e312
cvf86 controller_branch diag (commit 708c121f2) revealed:
rollout_steps: 99.99% WIDEN, 0% HOLD, 0% SHRINK
The bounded-step + noise-floor fix was correctly applied, but the
controller WIDENED 99.99% of steps because the hardcoded
ADV_VAR_RATIO_TARGET = 0.1 (`#define` in the kernel) was a b_size>1
design choice. The streaming-EMA regime at b_size=1 has var/|mean|
naturally living in [1, 10] (median 3.9), so input is ALWAYS >> 0.1
and the controller correctly says "noisy advantages → widen". Result:
n_rollout pegs at MAX=8192 within ~5 steps and stays for 50k steps,
PPO update frequency drops 4×, KL stays in numerical noise (median
1.7e-8), Q can't learn (l_q stuck at ~2.7 vs uniform 3.04).
## Fix: ISV-driven target
Per `feedback_isv_for_adaptive_bounds`: ADV_VAR_RATIO_TARGET now
lives in ISV slot 449 (`RL_ADV_VAR_RATIO_TARGET_INDEX`), seeded at
trainer init to 5.0 (matches streaming-regime median 3.9). The
controller reads `isv[RL_ADV_VAR_RATIO_TARGET_INDEX]` each step
instead of a `#define`.
Expected behavior at TARGET=5.0:
* Median input 3.9 lands in-band [3.33, 7.5] → HOLD
* n_rollout stays near BOOTSTRAP=2048 instead of MAX
* 4× more PPO updates per step → policy actually moves
* KL leaves noise floor → ε controller activates
* Q has gradient signal → can learn
Noise floor is now derived multiplicatively from the ISV target
(`target × ADV_VAR_RATIO_NOISE_FLOOR_FRAC = 0.01`) so adjusting
the target proportionally adjusts the floor — no separate slot
needed.
## Wiring
`rl_streaming_clamp_init.cu` extended to seed all three ISV-resident
design constants (adv_var clamp ceiling, td_kurt clamp ceiling, AND
adv_var regression target). Single kernel call at trainer init —
still no HtoD per `feedback_no_htod_htoh_only_mapped_pinned`.
## Diag bake-in
`controller_branch.rollout_steps_target` now reads from
`isv[RL_ADV_VAR_RATIO_TARGET_INDEX]` instead of the prior hardcoded
`0.1f32` literal. The diag shows the current ISV-resident target
so post-hoc branch analysis uses the actual value the controller
saw, and lets us track whether a future adaptive controller (one
that maintains target from observed-input percentile EMA) is
moving the target correctly.
## Slot allocation
RL_SLOTS_END: 449 → 450 (one new design-constant slot).
## Test updates
G1 (isv_bootstrap) + G3 (r5_controllers) skip slot 449 in the
sentinel-zero loop and assert the seeded value (5.0) separately.
G3's `advantage_var_ratio` input bumped from 5.0 → 20.0 so the
WIDEN branch still fires (input > new target × 1.5 = 7.5) and the
test still validates that the controller moves off bootstrap.
## Verified gates (local sm_86)
G1 isv_bootstrap ✅
G3 controllers ✅ (with updated input)
G4 target_update ✅
integrated_smoke ✅
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
3537 lines
158 KiB
Rust
3537 lines
158 KiB
Rust
//! Integrated RL trainer: BCE + DQN/C51 Q + PPO π + V + aux on a shared
|
||
//! Mamba2+CfC encoder.
|
||
//!
|
||
//! Phase E.2 + E.2-DEFER (this commit closes 4 deferred items):
|
||
//! 1. `PerceptionTrainer::backward_encoder_with_grad_h_t` —
|
||
//! the combined Q+π+V grad_h_t is folded into the encoder backward
|
||
//! path via the new entry point, so the encoder is no longer
|
||
//! BCE+aux-only.
|
||
//! 2. Bellman target projection kernel —
|
||
//! replaces the host-side `build_synthetic_bellman_target` stand-in
|
||
//! with `bellman_target_projection.cu`, reading γ from ISV[400] and
|
||
//! bootstrapping atom probabilities from the target-network output.
|
||
//! 3. Per-head LR ISV controller —
|
||
//! `rl_lr_controller.cu` emits per-head learning rates to
|
||
//! ISV[412..417]; the host reads them via the ISV mirror and
|
||
//! mutates each `AdamW.lr` field before the step. The
|
||
//! `PHASE_E2_DEFAULT_LR` constant is DELETED.
|
||
//! 4. PPO V-loss canonicalisation —
|
||
//! `ppo_clipped_surrogate_fwd` no longer touches `returns_/v_pred/
|
||
//! loss_v`; the V signal flows ONLY through the dedicated
|
||
//! `v_head_fwd_bwd` kernels per
|
||
//! `feedback_single_source_of_truth_no_duplicates`.
|
||
//!
|
||
//! Architecture (synthetic-data path):
|
||
//!
|
||
//! 1. `perception.forward_encoder(snapshots)` produces `h_t [B × HIDDEN_DIM]`
|
||
//! via the captured-graph encoder forward path.
|
||
//! 2. `dqn_head.forward(h_t, ...)` → `q_logits [B × N_ACTIONS × Q_N_ATOMS]`
|
||
//! `policy_head.forward_logits(h_t, ...)` → `pi_logits [B × N_ACTIONS]`
|
||
//! `value_head.forward(h_t, ...)` → `v_pred [B]`
|
||
//! 3. `policy_head.surrogate_forward(...)` consumes pi_logits + synthetic
|
||
//! log π_old + advantages + ISV(ε,coef) → π loss + entropy loss + per-
|
||
//! batch log π / entropy diagnostics. (V-loss path moved out — see #4)
|
||
//! 4. Bellman target build (kernel path, item 2):
|
||
//! a. `dqn_head.forward_target(...)` → `q_target_logits [B × N_ACTIONS × Q_N_ATOMS]`
|
||
//! b. `dqn_head.select_action_atoms(...)` → `[B × Q_N_ATOMS]` slice
|
||
//! c. `dqn_head.project_bellman_target(...)` → `target_dist [B × Q_N_ATOMS]`,
|
||
//! reading γ from ISV[RL_GAMMA_INDEX=400]
|
||
//! 5. `dqn_head.backward_logits(...)` → `q_loss [1]` + `grad_logits`.
|
||
//! 6. `dqn_head.backward_to_w_b_h(...)` → per-batch grad_w / grad_b
|
||
//! scratch + grad_h_t (OVERWRITE).
|
||
//! 7. `policy_head.surrogate_backward_logits(...)` → `pi_grad_logits`.
|
||
//! 8. `policy_head.backward_to_w_b_h(...)` → per-batch grad_w / grad_b
|
||
//! + grad_h_t (OVERWRITE).
|
||
//! 9. `value_head.backward(...)` → per-batch loss + grad_w / grad_b
|
||
//! + grad_h_t (OVERWRITE).
|
||
//! 10. `reduce_axis0` collapses per-batch grad_w / grad_b → final
|
||
//! `grad_w [...]` / `grad_b [...]` for each head.
|
||
//! 11. LR ISV controller (item 3): emit ISV[412..417], DtoH-mirror, and
|
||
//! mutate per-head AdamW.lr.
|
||
//! 12. Per-head `AdamW.step()` runs on the reduced gradients. Each head
|
||
//! owns 2 Adam instances (one for `w_d`, one for `b_d`).
|
||
//! 13. `grad_h_accumulate_scaled` folds Q/π/V grad_h_t (loss-balance λ
|
||
//! weighted) into the encoder grad slot `grad_h_t_combined_d`.
|
||
//! 14. `perception.backward_encoder_with_grad_h_t(grad_h_t_combined_d)`
|
||
//! runs the CfC + Mamba2 bwd + encoder Adam step (item 1).
|
||
//! 15. Per-batch losses (V head + DQN loss scalar) are reduced/read back
|
||
//! via mapped-pinned and returned in `IntegratedStepStats`.
|
||
//!
|
||
//! Constraints honoured:
|
||
//! * `feedback_no_stubs.md` — every kernel runs real math, no
|
||
//! placeholder host losses.
|
||
//! * `feedback_no_atomicadd.md` — new kernels use per-batch scratch
|
||
//! + `reduce_axis0` reduction (DQN/PPO bwd loss scalars still atomic
|
||
//! — Phase C/D deferral, kept loud-flagged in their kernel headers).
|
||
//! * `feedback_no_htod_htoh_only_mapped_pinned.md` — all CPU↔GPU paths
|
||
//! via `MappedF32Buffer` + DtoD pattern from `pinned_mem.rs`.
|
||
//! * `feedback_no_nvrtc.md` — pre-compiled cubins via build.rs.
|
||
//! * `pearl_controller_anchors_isv_driven` — λs + per-head LRs read
|
||
//! from ISV; ε / coef read inside the surrogate kernel from
|
||
//! `ISV[402,403]`; γ read inside the Bellman projection from
|
||
//! `ISV[RL_GAMMA_INDEX=400]`.
|
||
//! * `feedback_single_source_of_truth_no_duplicates` — V-loss path now
|
||
//! lives ONLY in the v_head kernels; PPO surrogate dropped its
|
||
//! redundant loss_v accumulator.
|
||
|
||
use std::sync::Arc;
|
||
|
||
use anyhow::{Context, Result};
|
||
use cudarc::driver::{
|
||
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
|
||
PushKernelArg,
|
||
};
|
||
use ml_core::device::MlDevice;
|
||
|
||
use crate::cfc::snap_features::Mbp10RawInput;
|
||
use crate::heads::HIDDEN_DIM;
|
||
use crate::pinned_mem::{MappedF32Buffer, MappedI32Buffer};
|
||
use rand::{Rng, SeedableRng};
|
||
|
||
use crate::rl::common::{N_ACTIONS, Q_N_ATOMS, Q_V_MAX, Q_V_MIN};
|
||
use crate::rl::dqn::{DqnHead, DqnHeadConfig};
|
||
use crate::rl::isv_slots::{
|
||
RL_LR_BCE_INDEX, RL_LR_AUX_INDEX, RL_LR_PI_INDEX, RL_LR_Q_INDEX,
|
||
RL_LR_V_INDEX, RL_SLOTS_END,
|
||
};
|
||
use crate::rl::loss_balance::{read_loss_lambdas_from_isv, LossLambdas};
|
||
use crate::rl::ppo::{PolicyHead, PpoHeadsConfig, ValueHead};
|
||
use crate::rl::reward::RlLobBackend;
|
||
use crate::trainer::optim::AdamW;
|
||
use crate::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
|
||
|
||
const GRAD_H_ACCUMULATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/grad_h_accumulate.cubin"));
|
||
const REDUCE_AXIS0_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin"));
|
||
const RL_LR_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_lr_controller.cubin"));
|
||
|
||
// Phase R1: cubin includes for the 7 RL adaptive controllers
|
||
// (γ / τ / ε / entropy_coef / n_rollout_steps / per_α / reward_scale).
|
||
// All 7 share the `(isv*, alpha, scalar_input)` signature and the
|
||
// `pearl_first_observation_bootstrap` sentinel-zero path. Loaded in
|
||
// `IntegratedTrainer::new` + bootstrap-launched once at construction so
|
||
// every consumer kernel sees the canonical default in its ISV slot
|
||
// before any in-loop access.
|
||
const RL_GAMMA_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_gamma_controller.cubin"));
|
||
const RL_TARGET_TAU_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_target_tau_controller.cubin"));
|
||
const RL_PPO_CLIP_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_ppo_clip_controller.cubin"));
|
||
const RL_ENTROPY_COEF_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_entropy_coef_controller.cubin"));
|
||
const RL_ROLLOUT_STEPS_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_rollout_steps_controller.cubin"));
|
||
const RL_PER_ALPHA_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_per_alpha_controller.cubin"));
|
||
const RL_REWARD_SCALE_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_reward_scale_controller.cubin"));
|
||
|
||
// Phase R3: GPU-resident EMA + advantage/return kernels. Generic over
|
||
// the ISV slot they update (one kernel serves multiple controller-input
|
||
// EMAs). Replace the host-side EMA + advantage loops the flawed Phase F
|
||
// shipped in step_with_lobsim (which violated `feedback_cpu_is_read_only`).
|
||
const EMA_UPDATE_ON_DONE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/ema_update_on_done.cubin"));
|
||
const EMA_UPDATE_PER_STEP_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/ema_update_per_step.cubin"));
|
||
const COMPUTE_ADVANTAGE_RETURN_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/compute_advantage_return.cubin"));
|
||
|
||
// Phase R4: GPU-resident action sampling kernels. Replace the host
|
||
// Thompson + host argmax + host log-softmax loops the flawed Phase F
|
||
// shipped in step_with_lobsim (which violated `feedback_cpu_is_read_only`).
|
||
const RL_ACTION_KERNEL_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_action_kernel.cubin"));
|
||
const ARGMAX_EXPECTED_Q_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/argmax_expected_q.cubin"));
|
||
const LOG_PI_AT_ACTION_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/log_pi_at_action.cubin"));
|
||
|
||
// Phase R6: GPU-pure env interaction kernels. extract_realized_pnl_delta
|
||
// reads pos.realized_pnl from the device Pos array and writes per-batch
|
||
// (reward delta, done flag); apply_reward_scale multiplies rewards by
|
||
// ISV[406]; actions_to_market_targets translates the 9-action grid to
|
||
// LobSim's market_targets[B*2] format on device.
|
||
const EXTRACT_REALIZED_PNL_DELTA_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/extract_realized_pnl_delta.cubin"));
|
||
const APPLY_REWARD_SCALE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/apply_reward_scale.cubin"));
|
||
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"));
|
||
|
||
// Reduction kernels that derive scalar inputs for the controller-input
|
||
// EMAs (variance ratio + kurtosis). entropy_observed reuses
|
||
// ema_update_per_step's built-in mean reduce directly on entropy_d.
|
||
const RL_VAR_OVER_ABS_MEAN_STREAMING_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_var_over_abs_mean_streaming.cubin"));
|
||
const RL_KURTOSIS_STREAMING_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_kurtosis_streaming.cubin"));
|
||
|
||
// Derivation kernels for the EMA inputs that need a new signal
|
||
// computation (not just a reduction over an existing buffer):
|
||
// per-batch KL approximation, Q weight-divergence L2 norm, and the
|
||
// trade-duration step counter.
|
||
const RL_KL_APPROX_B_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_kl_approx_b.cubin"));
|
||
const RL_L2_DIFF_NORM_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_l2_diff_norm.cubin"));
|
||
const RL_STEP_COUNTER_UPDATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_step_counter_update.cubin"));
|
||
const RL_L2_NORM_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_l2_norm.cubin"));
|
||
// R9 — PPO ratio clamp controller (emits ISV[440]) + per-step
|
||
// log-ratio max diagnostic kernel (writes ISV[441]).
|
||
const RL_PPO_RATIO_CLAMP_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_ppo_ratio_clamp_controller.cubin"));
|
||
const PPO_LOG_RATIO_ABS_MAX_B_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/ppo_log_ratio_abs_max_b.cubin"));
|
||
const RL_STREAMING_CLAMP_INIT_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_streaming_clamp_init.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
|
||
/// host-side so the trainer can tune the diagnostic-input weighting in
|
||
/// Phase E.3+ without touching the kernel.
|
||
const RL_LR_CONTROLLER_ALPHA: f32 = 0.4;
|
||
|
||
/// Configuration for [`IntegratedTrainer`]. Wraps a `PerceptionTrainerConfig`
|
||
/// (the encoder + BCE + aux side) plus RL-specific overrides for the new
|
||
/// heads.
|
||
#[derive(Clone, Debug)]
|
||
pub struct IntegratedTrainerConfig {
|
||
/// Inner perception trainer config. Encoder dims (hidden_dim, n_batch,
|
||
/// seq_len) are read from this struct.
|
||
pub perception: PerceptionTrainerConfig,
|
||
/// Seed for DQN head Xavier init.
|
||
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.
|
||
#[derive(Clone, Debug)]
|
||
pub struct IntegratedStepStats {
|
||
pub l_bce: f32,
|
||
pub l_q: f32,
|
||
pub l_pi: f32,
|
||
pub l_v: f32,
|
||
pub l_aux: f32,
|
||
pub l_total: f32,
|
||
pub lambdas: LossLambdas,
|
||
}
|
||
|
||
pub struct IntegratedTrainer {
|
||
#[allow(dead_code)]
|
||
cfg: IntegratedTrainerConfig,
|
||
|
||
/// Inner trainer owns the encoder weights + BCE/aux heads + optimizer.
|
||
pub perception: PerceptionTrainer,
|
||
|
||
/// DQN distributional Q-head (Phase C).
|
||
pub dqn_head: DqnHead,
|
||
|
||
/// PPO policy head (Phase D).
|
||
pub policy_head: PolicyHead,
|
||
|
||
/// PPO value head (Phase D).
|
||
pub value_head: ValueHead,
|
||
|
||
/// Per-head Adam optimisers — Phase E.2. One pair (w + b) per head.
|
||
/// All share the existing project-wide `adamw_step` cubin via the
|
||
/// `AdamW` wrapper; each instance owns its own m / v / step counter.
|
||
pub dqn_w_adam: AdamW,
|
||
pub dqn_b_adam: AdamW,
|
||
pub policy_w_adam: AdamW,
|
||
pub policy_b_adam: AdamW,
|
||
pub value_w_adam: AdamW,
|
||
pub value_b_adam: AdamW,
|
||
|
||
/// Device ISV buffer (RL-allocated, length `RL_SLOTS_END = 424` as
|
||
/// of Phase R1). Allocated zero, then immediately seeded by the 7
|
||
/// RL controllers' bootstrap launches in `new()` so every consumer
|
||
/// kernel sees the canonical default at its slot before any read.
|
||
/// Per `pearl_first_observation_bootstrap` + `feedback_isv_for_adaptive_bounds`.
|
||
pub isv_d: CudaSlice<f32>,
|
||
|
||
/// Host mirror of ISV for loss-balance λ reads. Refreshed before each
|
||
/// step via stream.memcpy_dtoh(isv_d, host).
|
||
pub isv_host: Vec<f32>,
|
||
|
||
stream: Arc<CudaStream>,
|
||
|
||
// ── Kernel handles for grad combine + cross-batch reduce ──────────
|
||
_grad_h_module: Arc<CudaModule>,
|
||
grad_h_accumulate_fn: CudaFunction,
|
||
_reduce_axis0_module: Arc<CudaModule>,
|
||
reduce_axis0_fn: CudaFunction,
|
||
|
||
// ── Per-head LR ISV controller (Phase E.2-DEFER item 3) ───────────
|
||
/// Owns the `rl_lr_controller` cubin lifetime.
|
||
_rl_lr_controller_module: Arc<CudaModule>,
|
||
/// Kernel handle for `rl_lr_controller` — emits per-head LRs to
|
||
/// `ISV[412..417]` via the Wiener-α blend defined in
|
||
/// `cuda/rl_lr_controller.cu`. Bootstraps on first observation
|
||
/// (sentinel zero → `LR_BOOTSTRAP = 1e-3`).
|
||
rl_lr_controller_fn: CudaFunction,
|
||
|
||
// ── 7 RL adaptive controllers (Phase R1) ──────────────────────────
|
||
// Each controller emits ONE float into its dedicated ISV slot
|
||
// (γ→400, τ→401, ε→402, entropy_coef→403, n_rollout_steps→404,
|
||
// per_α→405, reward_scale→406). Bootstrap-launched once at the end
|
||
// of `new()` so every consumer kernel reads a canonical value on
|
||
// first access — eliminates the Phase F defect where sparse-reward
|
||
// production paths read sentinel-zero γ/ε/entropy and trained
|
||
// against degenerate Bellman / surrogate targets. Per-step launches
|
||
// (with real EMA inputs sourced from ISV[417..424]) land in Phase R5.
|
||
_rl_gamma_controller_module: Arc<CudaModule>,
|
||
rl_gamma_controller_fn: CudaFunction,
|
||
_rl_target_tau_controller_module: Arc<CudaModule>,
|
||
rl_target_tau_controller_fn: CudaFunction,
|
||
_rl_ppo_clip_controller_module: Arc<CudaModule>,
|
||
rl_ppo_clip_controller_fn: CudaFunction,
|
||
_rl_entropy_coef_controller_module: Arc<CudaModule>,
|
||
rl_entropy_coef_controller_fn: CudaFunction,
|
||
_rl_rollout_steps_controller_module: Arc<CudaModule>,
|
||
rl_rollout_steps_controller_fn: CudaFunction,
|
||
_rl_per_alpha_controller_module: Arc<CudaModule>,
|
||
rl_per_alpha_controller_fn: CudaFunction,
|
||
_rl_reward_scale_controller_module: Arc<CudaModule>,
|
||
rl_reward_scale_controller_fn: CudaFunction,
|
||
|
||
// ── Phase R3: GPU-resident EMA + advantage/return kernels ─────────
|
||
_ema_update_on_done_module: Arc<CudaModule>,
|
||
ema_update_on_done_fn: CudaFunction,
|
||
_ema_update_per_step_module: Arc<CudaModule>,
|
||
ema_update_per_step_fn: CudaFunction,
|
||
_compute_advantage_return_module: Arc<CudaModule>,
|
||
compute_advantage_return_fn: CudaFunction,
|
||
|
||
// ── Phase R4: GPU action sampling ─────────────────────────────────
|
||
_rl_action_kernel_module: Arc<CudaModule>,
|
||
rl_action_kernel_fn: CudaFunction,
|
||
_argmax_expected_q_module: Arc<CudaModule>,
|
||
argmax_expected_q_fn: CudaFunction,
|
||
_log_pi_at_action_module: Arc<CudaModule>,
|
||
log_pi_at_action_fn: CudaFunction,
|
||
|
||
/// Per-batch xorshift32 PRNG state for `rl_action_kernel`. Allocated
|
||
/// + host-seeded once at trainer init from `cfg.dqn_seed` via
|
||
/// `rand_chacha::ChaCha8Rng` per `pearl_scoped_init_seed_for_reproducibility`.
|
||
/// The kernel advances state in place each call. Length = `n_batch`.
|
||
pub prng_state_d: CudaSlice<u32>,
|
||
|
||
/// Atom support vector `[Q_V_MIN, Q_V_MIN + step, …, Q_V_MAX]` of
|
||
/// length `Q_N_ATOMS`. Allocated + host-uploaded once at init;
|
||
/// read by `rl_action_kernel` (Thompson sample lookup) and
|
||
/// `argmax_expected_q` (expected-value accumulator). Per
|
||
/// `feedback_isv_for_adaptive_bounds`: structural constant, not
|
||
/// adaptive (atom span is fixed [Q_V_MIN, Q_V_MAX] = [-1, +1] per
|
||
/// `crate::rl::common`). Per-branch adaptive atom span (per
|
||
/// `pearl_per_branch_c51_atom_span`) is a Phase R-future extension.
|
||
pub atom_supports_d: CudaSlice<f32>,
|
||
|
||
// ── Phase R6: GPU-pure env-step kernels + per-step buffers ────────
|
||
_extract_realized_pnl_delta_module: Arc<CudaModule>,
|
||
extract_realized_pnl_delta_fn: CudaFunction,
|
||
_apply_reward_scale_module: Arc<CudaModule>,
|
||
apply_reward_scale_fn: CudaFunction,
|
||
_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,
|
||
|
||
// ── Per-batch reducers for controller-input EMAs ──
|
||
_rl_var_over_abs_mean_streaming_module: Arc<CudaModule>,
|
||
rl_var_over_abs_mean_streaming_fn: CudaFunction,
|
||
_rl_kurtosis_streaming_module: Arc<CudaModule>,
|
||
rl_kurtosis_streaming_fn: CudaFunction,
|
||
/// 1-float scratch for the reduce kernels' scalar output. Reused
|
||
/// across per-step launches (reductions serialised by stream order).
|
||
/// Feeds `ema_update_per_step` with `b_size=1` for the controller
|
||
/// EMA slots whose input requires a transform (variance ratio,
|
||
/// kurtosis, KL approximation, weight L2 divergence).
|
||
pub ema_input_scratch_d: CudaSlice<f32>,
|
||
|
||
// ── Derivation kernels + per-batch state for the remaining EMAs ──
|
||
_rl_kl_approx_b_module: Arc<CudaModule>,
|
||
rl_kl_approx_b_fn: CudaFunction,
|
||
_rl_l2_diff_norm_module: Arc<CudaModule>,
|
||
rl_l2_diff_norm_fn: CudaFunction,
|
||
_rl_step_counter_update_module: Arc<CudaModule>,
|
||
rl_step_counter_update_fn: CudaFunction,
|
||
_rl_l2_norm_module: Arc<CudaModule>,
|
||
rl_l2_norm_fn: CudaFunction,
|
||
// R9 PPO ratio clamp + log-ratio diag.
|
||
_rl_ppo_ratio_clamp_controller_module: Arc<CudaModule>,
|
||
rl_ppo_ratio_clamp_controller_fn: CudaFunction,
|
||
_ppo_log_ratio_abs_max_b_module: Arc<CudaModule>,
|
||
ppo_log_ratio_abs_max_b_fn: CudaFunction,
|
||
_rl_streaming_clamp_init_module: Arc<CudaModule>,
|
||
rl_streaming_clamp_init_fn: CudaFunction,
|
||
/// Per-batch step counter for `trade_duration_ema` — increments
|
||
/// every step, resets on done. Zero-init at construction.
|
||
pub steps_since_done_d: CudaSlice<i32>,
|
||
/// Per-batch duration emit buffer — populated by
|
||
/// `rl_step_counter_update` with the counter value AT done
|
||
/// (zero otherwise). Consumed by `ema_update_on_done` gated on
|
||
/// dones_d so only real trade closes contribute to the EMA.
|
||
pub trade_duration_emit_d: CudaSlice<f32>,
|
||
|
||
/// 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
|
||
/// `reward = pos.realized_pnl − prev_realized_pnl_d`, and updates
|
||
/// `prev_realized_pnl_d` for the next step. Separate from
|
||
/// `LobSimCuda::prev_realized_pnl_d` (which serves the production
|
||
/// decision pipeline, not the RL trainer).
|
||
pub prev_realized_pnl_d: CudaSlice<f32>,
|
||
/// 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`,
|
||
/// 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>,
|
||
|
||
/// Phase R7c (data correctness): trainer-owned `h_{t+1}` buffer
|
||
/// (`[B × HIDDEN_DIM]`). Populated by a SECOND
|
||
/// `perception.forward_encoder(next_snapshots)` call in
|
||
/// `step_with_lobsim`, with a stream-ordered DtoD copy out of
|
||
/// `perception.h_t_d` BEFORE the third `forward_encoder(snapshots)`
|
||
/// call overwrites it. Closes the long-standing
|
||
/// `// FUTURE: pass next_h_t separately` approximation noted in
|
||
/// step_synthetic's Bellman target build (which had been using
|
||
/// `h_t` as a stand-in for `s_{t+1}`'s encoder representation since
|
||
/// Phase E.2). Three consumers downstream:
|
||
/// 1. `value_head.forward(&self.h_tp1_d) → v_pred_tp1_d`, which
|
||
/// compute_advantage_return reads as the true `V(s_{t+1})`
|
||
/// (replaces R7b's `v_tp1_d_ref = &v_pred_d` alias).
|
||
/// 2. `dqn_head.forward(&self.h_tp1_d) → q_logits_tp1_d`, fed to
|
||
/// `argmax_expected_q` for `next_actions_d` — the Double-DQN
|
||
/// argmax must be on `h_{t+1}`, not `h_t` (R7b's q_logits_d
|
||
/// was on h_t, so the argmax was wrong since R4).
|
||
/// 3. `dqn_head.forward_target(&self.h_tp1_d)` inside
|
||
/// step_synthetic's Bellman target build (replaces the
|
||
/// h_t-as-proxy comment at the call site).
|
||
pub h_tp1_d: CudaSlice<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
|
||
/// `B × HIDDEN_DIM`; zeroed at the start of each
|
||
/// `step_synthetic` call.
|
||
grad_h_t_combined_d: CudaSlice<f32>,
|
||
|
||
/// Last-read PPO π loss host scalar — internal scratch threading the
|
||
/// surrogate kernel's atomic-write loss readback through the
|
||
/// borrow-checker dance in `step_synthetic`. NOT a public API.
|
||
last_pi_loss: f32,
|
||
/// Last step's Q-head loss (Q CE, host scalar) — fed to the LR
|
||
/// controller at the start of NEXT step_synthetic for plateau
|
||
/// detection. Zero at construction; populated at the end of
|
||
/// step_synthetic after the Q backward's scalar readback.
|
||
last_q_loss: f32,
|
||
/// Last step's V-head loss (MSE, host scalar) — same pattern as
|
||
/// `last_q_loss`. Zero at construction.
|
||
last_v_loss: f32,
|
||
|
||
/// Host-side step counter for Thompson-sampling RNG salt (Phase E.3b).
|
||
/// Increments once per `step_with_lobsim` call so consecutive calls
|
||
/// draw different actions from the same Q-distribution. Per
|
||
/// `pearl_scoped_init_seed_for_reproducibility`: the RNG seed is
|
||
/// `cfg.dqn_seed.wrapping_add(step_counter)`, deterministic given
|
||
/// the same config + step sequence.
|
||
step_counter: u64,
|
||
}
|
||
|
||
impl IntegratedTrainer {
|
||
pub fn new(dev: &MlDevice, cfg: IntegratedTrainerConfig) -> Result<Self> {
|
||
let stream: Arc<CudaStream> = dev.cuda_stream().context("integrated stream")?.clone();
|
||
|
||
let perception =
|
||
PerceptionTrainer::new(dev, &cfg.perception).context("PerceptionTrainer::new")?;
|
||
|
||
// Encoder hidden dim — single source of truth via heads::HIDDEN_DIM
|
||
// (the same const the perception trainer uses).
|
||
let hidden_dim = HIDDEN_DIM;
|
||
|
||
let dqn_head = DqnHead::new(
|
||
dev,
|
||
DqnHeadConfig {
|
||
hidden_dim,
|
||
seed: cfg.dqn_seed,
|
||
},
|
||
)
|
||
.context("DqnHead::new")?;
|
||
|
||
let policy_head = PolicyHead::new(
|
||
dev,
|
||
PpoHeadsConfig {
|
||
hidden_dim,
|
||
seed: cfg.ppo_seed,
|
||
},
|
||
)
|
||
.context("PolicyHead::new")?;
|
||
|
||
let value_head = ValueHead::new(
|
||
dev,
|
||
PpoHeadsConfig {
|
||
hidden_dim,
|
||
seed: cfg.ppo_seed.wrapping_add(0xA1),
|
||
},
|
||
)
|
||
.context("ValueHead::new")?;
|
||
|
||
// Per-head Adam optimisers — one per (head, w | b) pair. Each owns
|
||
// independent m/v buffers and a device-resident step counter (per
|
||
// pearl_no_host_branches_in_captured_graph: counter advancement
|
||
// happens inside a captured kernel, not in host code).
|
||
//
|
||
// LR is sourced from ISV[RL_LR_*_INDEX] per step (item 3); the
|
||
// construction-time value passed to AdamW::new is a sentinel-zero-
|
||
// bootstrap PLACEHOLDER overwritten before the first AdamW.step.
|
||
// Per `pearl_first_observation_bootstrap` the controller seeds
|
||
// 1e-3 on the first ISV emit, then the trainer mirrors that value
|
||
// into each AdamW.lr field before the Adam step launches.
|
||
let lr_placeholder = 0.0_f32;
|
||
let dqn_w_adam =
|
||
AdamW::new(dev, dqn_head.w_d.len(), lr_placeholder).context("dqn_w_adam")?;
|
||
let dqn_b_adam =
|
||
AdamW::new(dev, dqn_head.b_d.len(), lr_placeholder).context("dqn_b_adam")?;
|
||
let policy_w_adam =
|
||
AdamW::new(dev, policy_head.w_d.len(), lr_placeholder).context("policy_w_adam")?;
|
||
let policy_b_adam =
|
||
AdamW::new(dev, policy_head.b_d.len(), lr_placeholder).context("policy_b_adam")?;
|
||
let value_w_adam =
|
||
AdamW::new(dev, value_head.w_d.len(), lr_placeholder).context("value_w_adam")?;
|
||
let value_b_adam =
|
||
AdamW::new(dev, value_head.b_d.len(), lr_placeholder).context("value_b_adam")?;
|
||
|
||
// ISV buffer — zero-init, sentinel-bootstrap on first controller emit.
|
||
let isv_d = stream
|
||
.alloc_zeros::<f32>(RL_SLOTS_END)
|
||
.context("alloc isv_d")?;
|
||
let isv_host = vec![0.0_f32; RL_SLOTS_END];
|
||
// R9 clamp-ceiling seeding happens via the
|
||
// rl_streaming_clamp_init kernel launched alongside the
|
||
// controller bootstraps below — per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned` we never
|
||
// host-seed device-resident state.
|
||
|
||
// Combined-encoder-grad slot for item 1 (perception backward).
|
||
// Sized [B × HIDDEN_DIM]; the perception config holds n_batch.
|
||
let grad_h_t_combined_d = stream
|
||
.alloc_zeros::<f32>(cfg.perception.n_batch * hidden_dim)
|
||
.context("alloc grad_h_t_combined_d")?;
|
||
|
||
// Load helper kernels.
|
||
let ctx = dev.cuda_context().context("integrated ctx")?;
|
||
let grad_h_module = ctx
|
||
.load_cubin(GRAD_H_ACCUMULATE_CUBIN.to_vec())
|
||
.context("load grad_h_accumulate cubin")?;
|
||
let grad_h_accumulate_fn = grad_h_module
|
||
.load_function("grad_h_accumulate_scaled")
|
||
.context("load grad_h_accumulate_scaled")?;
|
||
let reduce_axis0_module = ctx
|
||
.load_cubin(REDUCE_AXIS0_CUBIN.to_vec())
|
||
.context("load reduce_axis0 cubin")?;
|
||
let reduce_axis0_fn = reduce_axis0_module
|
||
.load_function("reduce_axis0")
|
||
.context("load reduce_axis0")?;
|
||
let rl_lr_controller_module = ctx
|
||
.load_cubin(RL_LR_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_lr_controller cubin")?;
|
||
let rl_lr_controller_fn = rl_lr_controller_module
|
||
.load_function("rl_lr_controller")
|
||
.context("load rl_lr_controller")?;
|
||
|
||
// Phase R1: load the 7 RL adaptive controllers.
|
||
let rl_gamma_controller_module = ctx
|
||
.load_cubin(RL_GAMMA_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_gamma_controller cubin")?;
|
||
let rl_gamma_controller_fn = rl_gamma_controller_module
|
||
.load_function("rl_gamma_controller")
|
||
.context("load rl_gamma_controller")?;
|
||
let rl_target_tau_controller_module = ctx
|
||
.load_cubin(RL_TARGET_TAU_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_target_tau_controller cubin")?;
|
||
let rl_target_tau_controller_fn = rl_target_tau_controller_module
|
||
.load_function("rl_target_tau_controller")
|
||
.context("load rl_target_tau_controller")?;
|
||
let rl_ppo_clip_controller_module = ctx
|
||
.load_cubin(RL_PPO_CLIP_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_ppo_clip_controller cubin")?;
|
||
let rl_ppo_clip_controller_fn = rl_ppo_clip_controller_module
|
||
.load_function("rl_ppo_clip_controller")
|
||
.context("load rl_ppo_clip_controller")?;
|
||
let rl_entropy_coef_controller_module = ctx
|
||
.load_cubin(RL_ENTROPY_COEF_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_entropy_coef_controller cubin")?;
|
||
let rl_entropy_coef_controller_fn = rl_entropy_coef_controller_module
|
||
.load_function("rl_entropy_coef_controller")
|
||
.context("load rl_entropy_coef_controller")?;
|
||
let rl_rollout_steps_controller_module = ctx
|
||
.load_cubin(RL_ROLLOUT_STEPS_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_rollout_steps_controller cubin")?;
|
||
let rl_rollout_steps_controller_fn = rl_rollout_steps_controller_module
|
||
.load_function("rl_rollout_steps_controller")
|
||
.context("load rl_rollout_steps_controller")?;
|
||
let rl_per_alpha_controller_module = ctx
|
||
.load_cubin(RL_PER_ALPHA_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_per_alpha_controller cubin")?;
|
||
let rl_per_alpha_controller_fn = rl_per_alpha_controller_module
|
||
.load_function("rl_per_alpha_controller")
|
||
.context("load rl_per_alpha_controller")?;
|
||
let rl_reward_scale_controller_module = ctx
|
||
.load_cubin(RL_REWARD_SCALE_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_reward_scale_controller cubin")?;
|
||
let rl_reward_scale_controller_fn = rl_reward_scale_controller_module
|
||
.load_function("rl_reward_scale_controller")
|
||
.context("load rl_reward_scale_controller")?;
|
||
|
||
// Phase R3 kernels.
|
||
let ema_update_on_done_module = ctx
|
||
.load_cubin(EMA_UPDATE_ON_DONE_CUBIN.to_vec())
|
||
.context("load ema_update_on_done cubin")?;
|
||
let ema_update_on_done_fn = ema_update_on_done_module
|
||
.load_function("ema_update_on_done")
|
||
.context("load ema_update_on_done")?;
|
||
let ema_update_per_step_module = ctx
|
||
.load_cubin(EMA_UPDATE_PER_STEP_CUBIN.to_vec())
|
||
.context("load ema_update_per_step cubin")?;
|
||
let ema_update_per_step_fn = ema_update_per_step_module
|
||
.load_function("ema_update_per_step")
|
||
.context("load ema_update_per_step")?;
|
||
let compute_advantage_return_module = ctx
|
||
.load_cubin(COMPUTE_ADVANTAGE_RETURN_CUBIN.to_vec())
|
||
.context("load compute_advantage_return cubin")?;
|
||
let compute_advantage_return_fn = compute_advantage_return_module
|
||
.load_function("compute_advantage_return")
|
||
.context("load compute_advantage_return")?;
|
||
|
||
// Phase R4 kernels.
|
||
let rl_action_kernel_module = ctx
|
||
.load_cubin(RL_ACTION_KERNEL_CUBIN.to_vec())
|
||
.context("load rl_action_kernel cubin")?;
|
||
let rl_action_kernel_fn = rl_action_kernel_module
|
||
.load_function("rl_action_kernel")
|
||
.context("load rl_action_kernel")?;
|
||
let argmax_expected_q_module = ctx
|
||
.load_cubin(ARGMAX_EXPECTED_Q_CUBIN.to_vec())
|
||
.context("load argmax_expected_q cubin")?;
|
||
let argmax_expected_q_fn = argmax_expected_q_module
|
||
.load_function("argmax_expected_q")
|
||
.context("load argmax_expected_q")?;
|
||
let log_pi_at_action_module = ctx
|
||
.load_cubin(LOG_PI_AT_ACTION_CUBIN.to_vec())
|
||
.context("load log_pi_at_action cubin")?;
|
||
let log_pi_at_action_fn = log_pi_at_action_module
|
||
.load_function("log_pi_at_action")
|
||
.context("load log_pi_at_action")?;
|
||
|
||
// Phase R6 kernels.
|
||
let extract_realized_pnl_delta_module = ctx
|
||
.load_cubin(EXTRACT_REALIZED_PNL_DELTA_CUBIN.to_vec())
|
||
.context("load extract_realized_pnl_delta cubin")?;
|
||
let extract_realized_pnl_delta_fn = extract_realized_pnl_delta_module
|
||
.load_function("extract_realized_pnl_delta")
|
||
.context("load extract_realized_pnl_delta")?;
|
||
let apply_reward_scale_module = ctx
|
||
.load_cubin(APPLY_REWARD_SCALE_CUBIN.to_vec())
|
||
.context("load apply_reward_scale cubin")?;
|
||
let apply_reward_scale_fn = apply_reward_scale_module
|
||
.load_function("apply_reward_scale")
|
||
.context("load apply_reward_scale")?;
|
||
let actions_to_market_targets_module = ctx
|
||
.load_cubin(ACTIONS_TO_MARKET_TARGETS_CUBIN.to_vec())
|
||
.context("load actions_to_market_targets cubin")?;
|
||
let actions_to_market_targets_fn = actions_to_market_targets_module
|
||
.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")?;
|
||
|
||
// Streaming kernels for the variance-ratio + kurtosis EMAs.
|
||
// Per R9 fix: these now fold ACROSS STEPS (Welford-EMA on
|
||
// per-batch means in ISV state slots) instead of across batch
|
||
// — works correctly at b_size=1 where per-batch variance is
|
||
// identically zero. The kernels write the smoothed estimate
|
||
// directly to the controller-input ISV slot; no separate
|
||
// ema_update_per_step is needed downstream.
|
||
let rl_var_over_abs_mean_streaming_module = ctx
|
||
.load_cubin(RL_VAR_OVER_ABS_MEAN_STREAMING_CUBIN.to_vec())
|
||
.context("load rl_var_over_abs_mean_streaming cubin")?;
|
||
let rl_var_over_abs_mean_streaming_fn = rl_var_over_abs_mean_streaming_module
|
||
.load_function("rl_var_over_abs_mean_streaming")
|
||
.context("load rl_var_over_abs_mean_streaming")?;
|
||
let rl_kurtosis_streaming_module = ctx
|
||
.load_cubin(RL_KURTOSIS_STREAMING_CUBIN.to_vec())
|
||
.context("load rl_kurtosis_streaming cubin")?;
|
||
let rl_kurtosis_streaming_fn = rl_kurtosis_streaming_module
|
||
.load_function("rl_kurtosis_streaming")
|
||
.context("load rl_kurtosis_streaming")?;
|
||
let ema_input_scratch_d = stream
|
||
.alloc_zeros::<f32>(1)
|
||
.context("alloc ema_input_scratch_d")?;
|
||
|
||
// Derivation kernels + per-batch state for KL approx, weight
|
||
// L2 divergence, and the trade-duration counter EMAs.
|
||
let rl_kl_approx_b_module = ctx
|
||
.load_cubin(RL_KL_APPROX_B_CUBIN.to_vec())
|
||
.context("load rl_kl_approx_b cubin")?;
|
||
let rl_kl_approx_b_fn = rl_kl_approx_b_module
|
||
.load_function("rl_kl_approx_b")
|
||
.context("load rl_kl_approx_b")?;
|
||
let rl_l2_diff_norm_module = ctx
|
||
.load_cubin(RL_L2_DIFF_NORM_CUBIN.to_vec())
|
||
.context("load rl_l2_diff_norm cubin")?;
|
||
let rl_l2_diff_norm_fn = rl_l2_diff_norm_module
|
||
.load_function("rl_l2_diff_norm")
|
||
.context("load rl_l2_diff_norm")?;
|
||
let rl_step_counter_update_module = ctx
|
||
.load_cubin(RL_STEP_COUNTER_UPDATE_CUBIN.to_vec())
|
||
.context("load rl_step_counter_update cubin")?;
|
||
let rl_step_counter_update_fn = rl_step_counter_update_module
|
||
.load_function("rl_step_counter_update")
|
||
.context("load rl_step_counter_update")?;
|
||
let rl_l2_norm_module = ctx
|
||
.load_cubin(RL_L2_NORM_CUBIN.to_vec())
|
||
.context("load rl_l2_norm cubin")?;
|
||
let rl_l2_norm_fn = rl_l2_norm_module
|
||
.load_function("rl_l2_norm")
|
||
.context("load rl_l2_norm")?;
|
||
|
||
// R9 — PPO ratio clamp controller + per-step log-ratio diag.
|
||
let rl_ppo_ratio_clamp_controller_module = ctx
|
||
.load_cubin(RL_PPO_RATIO_CLAMP_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_ppo_ratio_clamp_controller cubin")?;
|
||
let rl_ppo_ratio_clamp_controller_fn = rl_ppo_ratio_clamp_controller_module
|
||
.load_function("rl_ppo_ratio_clamp_controller")
|
||
.context("load rl_ppo_ratio_clamp_controller")?;
|
||
let ppo_log_ratio_abs_max_b_module = ctx
|
||
.load_cubin(PPO_LOG_RATIO_ABS_MAX_B_CUBIN.to_vec())
|
||
.context("load ppo_log_ratio_abs_max_b cubin")?;
|
||
let ppo_log_ratio_abs_max_b_fn = ppo_log_ratio_abs_max_b_module
|
||
.load_function("ppo_log_ratio_abs_max_b")
|
||
.context("load ppo_log_ratio_abs_max_b")?;
|
||
let rl_streaming_clamp_init_module = ctx
|
||
.load_cubin(RL_STREAMING_CLAMP_INIT_CUBIN.to_vec())
|
||
.context("load rl_streaming_clamp_init cubin")?;
|
||
let rl_streaming_clamp_init_fn = rl_streaming_clamp_init_module
|
||
.load_function("rl_streaming_clamp_init")
|
||
.context("load rl_streaming_clamp_init")?;
|
||
|
||
// 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
|
||
// re-runs of the same trainer init per
|
||
// `pearl_scoped_init_seed_for_reproducibility`. Zero is the
|
||
// only forbidden xorshift32 state (would freeze) so the
|
||
// `.max(1)` guard ensures every batch's seed is non-zero.
|
||
let b_size = cfg.perception.n_batch;
|
||
let mut prng_host_rng =
|
||
<rand_chacha::ChaCha8Rng as SeedableRng>::seed_from_u64(
|
||
cfg.dqn_seed.wrapping_add(0xC0_C0_C0_C0),
|
||
);
|
||
let prng_seeds: Vec<u32> = (0..b_size)
|
||
.map(|_| prng_host_rng.gen::<u32>().max(1))
|
||
.collect();
|
||
let prng_state_d = stream
|
||
.alloc_zeros::<u32>(b_size)
|
||
.context("alloc prng_state_d")?;
|
||
// Mapped-pinned upload via inline DtoD per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned`. One-shot at
|
||
// construction. Cast through i32 for the typed mapped-pinned
|
||
// buffer (u32 and i32 share width); the kernel reads as u32 —
|
||
// wraparound is well-defined.
|
||
{
|
||
let mut staging = unsafe { MappedI32Buffer::new(b_size) }
|
||
.map_err(|e| anyhow::anyhow!("prng_state staging: {e}"))?;
|
||
let dst = staging.host_slice_mut();
|
||
for (i, &v) in prng_seeds.iter().enumerate() {
|
||
dst[i] = v as i32;
|
||
}
|
||
unsafe {
|
||
let s = stream.cu_stream();
|
||
let (dev_dst, _g) = prng_state_d.device_ptr(&stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dev_dst,
|
||
staging.dev_ptr,
|
||
b_size * std::mem::size_of::<u32>(),
|
||
s,
|
||
)
|
||
.context("prng_state DtoD upload")?;
|
||
}
|
||
stream.synchronize().context("prng_state sync")?;
|
||
}
|
||
|
||
// Atom supports — structural constant (atom span [-1, +1]).
|
||
// Single allocation + upload at init; read every step by the
|
||
// Thompson + argmax kernels.
|
||
let atom_step = (Q_V_MAX - Q_V_MIN) / ((Q_N_ATOMS - 1) as f32);
|
||
let atom_supports_host: Vec<f32> = (0..Q_N_ATOMS)
|
||
.map(|i| Q_V_MIN + (i as f32) * atom_step)
|
||
.collect();
|
||
let mut atom_supports_d = stream
|
||
.alloc_zeros::<f32>(Q_N_ATOMS)
|
||
.context("alloc atom_supports_d")?;
|
||
// Mapped-pinned upload per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned`. One-shot.
|
||
write_slice_f32_d(&stream, &atom_supports_host, &mut atom_supports_d)
|
||
.context("atom_supports upload")?;
|
||
|
||
// Phase R6 per-step buffers. prev_* are sentinel-zero at init
|
||
// so the first call's delta is `current - 0 = current`. For a
|
||
// freshly-constructed LobSimCuda, `pos.realized_pnl == 0` and
|
||
// `pos.position_lots == 0` too, so the first delta is 0 and
|
||
// the first done is false (prev_lots == 0 fails the `prev != 0`
|
||
// condition). Subsequent steps see real prev_* values.
|
||
let prev_realized_pnl_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc prev_realized_pnl_d")?;
|
||
let prev_position_lots_d = stream
|
||
.alloc_zeros::<i32>(b_size)
|
||
.context("alloc prev_position_lots_d")?;
|
||
let rewards_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc rewards_d")?;
|
||
let dones_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc dones_d")?;
|
||
|
||
// Per-batch trade-duration counter (mutated by
|
||
// `rl_step_counter_update`: incremented each step, reset on
|
||
// done) and its done-gated emit buffer (consumed by
|
||
// `ema_update_on_done` to feed `mean_trade_duration_ema`).
|
||
let steps_since_done_d = stream
|
||
.alloc_zeros::<i32>(b_size)
|
||
.context("alloc steps_since_done_d")?;
|
||
let trade_duration_emit_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc trade_duration_emit_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")?;
|
||
|
||
// Phase R7c (data correctness): trainer-owned h_{t+1} buffer.
|
||
// Zero at init — first step_with_lobsim call writes it via the
|
||
// forward_encoder(next_snapshots) + DtoD copy out of
|
||
// perception.h_t_d BEFORE any consumer reads it.
|
||
let h_tp1_d = stream
|
||
.alloc_zeros::<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,
|
||
dqn_head,
|
||
policy_head,
|
||
value_head,
|
||
dqn_w_adam,
|
||
dqn_b_adam,
|
||
policy_w_adam,
|
||
policy_b_adam,
|
||
value_w_adam,
|
||
value_b_adam,
|
||
isv_d,
|
||
isv_host,
|
||
stream,
|
||
_grad_h_module: grad_h_module,
|
||
grad_h_accumulate_fn,
|
||
_reduce_axis0_module: reduce_axis0_module,
|
||
reduce_axis0_fn,
|
||
_rl_lr_controller_module: rl_lr_controller_module,
|
||
rl_lr_controller_fn,
|
||
_rl_gamma_controller_module: rl_gamma_controller_module,
|
||
rl_gamma_controller_fn,
|
||
_rl_target_tau_controller_module: rl_target_tau_controller_module,
|
||
rl_target_tau_controller_fn,
|
||
_rl_ppo_clip_controller_module: rl_ppo_clip_controller_module,
|
||
rl_ppo_clip_controller_fn,
|
||
_rl_entropy_coef_controller_module: rl_entropy_coef_controller_module,
|
||
rl_entropy_coef_controller_fn,
|
||
_rl_rollout_steps_controller_module: rl_rollout_steps_controller_module,
|
||
rl_rollout_steps_controller_fn,
|
||
_rl_per_alpha_controller_module: rl_per_alpha_controller_module,
|
||
rl_per_alpha_controller_fn,
|
||
_rl_reward_scale_controller_module: rl_reward_scale_controller_module,
|
||
rl_reward_scale_controller_fn,
|
||
_ema_update_on_done_module: ema_update_on_done_module,
|
||
ema_update_on_done_fn,
|
||
_ema_update_per_step_module: ema_update_per_step_module,
|
||
ema_update_per_step_fn,
|
||
_compute_advantage_return_module: compute_advantage_return_module,
|
||
compute_advantage_return_fn,
|
||
_rl_action_kernel_module: rl_action_kernel_module,
|
||
rl_action_kernel_fn,
|
||
_argmax_expected_q_module: argmax_expected_q_module,
|
||
argmax_expected_q_fn,
|
||
_log_pi_at_action_module: log_pi_at_action_module,
|
||
log_pi_at_action_fn,
|
||
prng_state_d,
|
||
atom_supports_d,
|
||
_extract_realized_pnl_delta_module: extract_realized_pnl_delta_module,
|
||
extract_realized_pnl_delta_fn,
|
||
_apply_reward_scale_module: apply_reward_scale_module,
|
||
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,
|
||
_rl_var_over_abs_mean_streaming_module: rl_var_over_abs_mean_streaming_module,
|
||
rl_var_over_abs_mean_streaming_fn,
|
||
_rl_kurtosis_streaming_module: rl_kurtosis_streaming_module,
|
||
rl_kurtosis_streaming_fn,
|
||
ema_input_scratch_d,
|
||
_rl_kl_approx_b_module: rl_kl_approx_b_module,
|
||
rl_kl_approx_b_fn,
|
||
_rl_l2_diff_norm_module: rl_l2_diff_norm_module,
|
||
rl_l2_diff_norm_fn,
|
||
_rl_step_counter_update_module: rl_step_counter_update_module,
|
||
rl_step_counter_update_fn,
|
||
_rl_l2_norm_module: rl_l2_norm_module,
|
||
rl_l2_norm_fn,
|
||
_rl_ppo_ratio_clamp_controller_module: rl_ppo_ratio_clamp_controller_module,
|
||
rl_ppo_ratio_clamp_controller_fn,
|
||
_ppo_log_ratio_abs_max_b_module: ppo_log_ratio_abs_max_b_module,
|
||
ppo_log_ratio_abs_max_b_fn,
|
||
_rl_streaming_clamp_init_module: rl_streaming_clamp_init_module,
|
||
rl_streaming_clamp_init_fn,
|
||
steps_since_done_d,
|
||
trade_duration_emit_d,
|
||
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,
|
||
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,
|
||
last_q_loss: 0.0,
|
||
last_v_loss: 0.0,
|
||
step_counter: 0,
|
||
}
|
||
.with_controllers_bootstrapped()?)
|
||
}
|
||
|
||
/// Phase R1: fire each of the 7 RL controllers ONCE against the
|
||
/// freshly-zeroed `isv_d`. Each kernel detects the sentinel (its
|
||
/// dedicated ISV slot == 0.0), executes its first-observation-
|
||
/// bootstrap path per `pearl_first_observation_bootstrap`, and
|
||
/// writes its canonical default into the slot:
|
||
///
|
||
/// ISV[400] γ = target(trade_duration_ema) — 0.90 at
|
||
/// sentinel input (after the bootstrap audit; was
|
||
/// hardcoded 0.99 which coincided with
|
||
/// target(d≈69) creating a Wiener-blend
|
||
/// dead-zone at canonical hold time)
|
||
/// ISV[401] τ = 0.005 (TAU_BOOTSTRAP)
|
||
/// ISV[402] ε = 0.2 (EPS_BOOTSTRAP)
|
||
/// ISV[403] entropy_coef = target(entropy_observed_ema) — 0.035
|
||
/// at sentinel input (after the bootstrap audit;
|
||
/// was hardcoded 0.01 which coincided
|
||
/// with target(h≈1.099) creating a
|
||
/// Wiener-blend dead-zone at canonical
|
||
/// mid-range entropy)
|
||
/// ISV[404] n_rollout_steps= 2048 (ROLLOUT_BOOTSTRAP)
|
||
/// ISV[405] per_α = target(td_kurtosis_ema) — 0.4 at
|
||
/// sentinel input (after the bootstrap audit; was
|
||
/// hardcoded 0.6 which coincided with
|
||
/// target(kurt=10) creating a Wiener-
|
||
/// blend dead-zone at canonical market
|
||
/// kurtosis)
|
||
/// ISV[406] reward_scale = 1.0 (REWARD_SCALE_BOOTSTRAP)
|
||
///
|
||
/// The scalar `input` argument is IGNORED on the bootstrap path
|
||
/// (each kernel `return`s immediately after the write). Phase R5
|
||
/// reuses [`Self::launch_isv_controller_3arg`] with real EMA
|
||
/// inputs sourced from ISV[417..424] — at that point the Wiener-α
|
||
/// blend kicks in and the slots adapt from the bootstrap defaults.
|
||
///
|
||
/// Replaces the flawed Phase F approach of host-side `memcpy_htod`
|
||
/// of canonical constants, which violated
|
||
/// `feedback_no_htod_htoh_only_mapped_pinned` (tests not exempt)
|
||
/// and short-circuited `pearl_first_observation_bootstrap`.
|
||
///
|
||
/// Phase R5 update: each controller now reads its EMA input from
|
||
/// the corresponding ISV[417..423] slot via `input_slot`. At
|
||
/// bootstrap time those slots are also at `alloc_zeros` sentinel
|
||
/// zero, but the kernel returns before the input read (the
|
||
/// controller's own output slot is at sentinel zero, triggering
|
||
/// the `pearl_first_observation_bootstrap` return path). The same
|
||
/// helper, `launch_isv_controller_3arg`, serves bootstrap and
|
||
/// per-step launches.
|
||
fn with_controllers_bootstrapped(self) -> Result<Self> {
|
||
let alpha = RL_LR_CONTROLLER_ALPHA;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_gamma_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_MEAN_TRADE_DURATION_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_gamma_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_target_tau_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_Q_DIVERGENCE_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_target_tau_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_ppo_clip_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_KL_PI_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_ppo_clip_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_entropy_coef_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_ENTROPY_OBSERVED_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_entropy_coef_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_rollout_steps_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_rollout_steps_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_per_alpha_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_per_alpha_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_reward_scale_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_MEAN_ABS_PNL_EMA_INDEX as i32,
|
||
)
|
||
.context("R1 bootstrap rl_reward_scale_controller")?;
|
||
// R9 — PPO ratio clamp ceiling. At bootstrap RL_PPO_CLIP_INDEX
|
||
// is still at sentinel zero; the kernel writes
|
||
// PPO_RATIO_CLAMP_BOOTSTRAP=10.0 and returns before reading.
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_ppo_ratio_clamp_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_PPO_CLIP_INDEX as i32,
|
||
)
|
||
.context("R9 bootstrap rl_ppo_ratio_clamp_controller")?;
|
||
// R9 — seed streaming-kernel output clamp ceilings device-side
|
||
// (no HtoD per feedback_no_htod_htoh_only_mapped_pinned).
|
||
// Single-thread kernel writes design constants to ISV[447]
|
||
// and ISV[448] which the streaming variance / kurtosis kernels
|
||
// read each step to clamp their outputs.
|
||
{
|
||
let adv_var_slot: i32 =
|
||
crate::rl::isv_slots::RL_ADV_VAR_RATIO_CLAMP_INDEX as i32;
|
||
let adv_var_val: f32 = 100.0;
|
||
let td_kurt_slot: i32 =
|
||
crate::rl::isv_slots::RL_TD_KURTOSIS_CLAMP_INDEX as i32;
|
||
let td_kurt_val: f32 = 30.0;
|
||
// ISV-driven regression target for rl_rollout_steps_controller.
|
||
// 5.0 matches the b_size=1 streaming regime where
|
||
// var/|mean| naturally lives in [1, 10]. The prior
|
||
// hardcoded #define 0.1 was a b_size>1 design choice and
|
||
// caused 99.99% WIDEN events in alpha-rl-cvf86.
|
||
let adv_var_target_slot: i32 =
|
||
crate::rl::isv_slots::RL_ADV_VAR_RATIO_TARGET_INDEX as i32;
|
||
let adv_var_target_val: f32 = 5.0;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self
|
||
.stream
|
||
.launch_builder(&self.rl_streaming_clamp_init_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(&adv_var_slot)
|
||
.arg(&adv_var_val)
|
||
.arg(&td_kurt_slot)
|
||
.arg(&td_kurt_val)
|
||
.arg(&adv_var_target_slot)
|
||
.arg(&adv_var_target_val);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("R9 launch rl_streaming_clamp_init")?;
|
||
}
|
||
}
|
||
self.stream
|
||
.synchronize()
|
||
.context("R1 controller bootstrap sync")?;
|
||
Ok(self)
|
||
}
|
||
|
||
/// Phase R1 / R5: single-thread launch of any of the 7 RL adaptive
|
||
/// controllers. All 7 share the `(isv*, alpha, input_slot)`
|
||
/// signature after R5; the kernel reads the EMA input from
|
||
/// `isv[input_slot]` directly, eliminating the 7 DtoH-per-step host
|
||
/// roundtrips a scalar-arg signature would have required (per
|
||
/// `feedback_cpu_is_read_only`).
|
||
///
|
||
/// At bootstrap (R1), the controller's output slot is at sentinel
|
||
/// zero; the kernel writes its `*_BOOTSTRAP` value and returns
|
||
/// before the input read, so `input_slot` can point at the (also
|
||
/// sentinel-zero) EMA-input slot safely. At per-step launches
|
||
/// (R5+), the EMA producer (Phase R3 `ema_update_*` kernels) has
|
||
/// populated `isv[input_slot]` with a real observation and the
|
||
/// Wiener-α-with-floor blend per
|
||
/// `pearl_wiener_alpha_floor_for_nonstationary` runs against it.
|
||
fn launch_isv_controller_3arg(
|
||
&self,
|
||
func: &CudaFunction,
|
||
alpha: f32,
|
||
input_slot: i32,
|
||
) -> Result<()> {
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(func);
|
||
launch.arg(&self.isv_d).arg(&alpha).arg(&input_slot);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("isv controller 3-arg launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R5: fire all 7 RL adaptive controllers in sequence,
|
||
/// each reading its EMA input from the corresponding ISV slot
|
||
/// populated by the Phase R3 `ema_update_*` kernels.
|
||
///
|
||
/// Slot pairing (output ← input):
|
||
/// ISV[400] γ ← ISV[417] MEAN_TRADE_DURATION_EMA
|
||
/// ISV[401] τ ← ISV[418] Q_DIVERGENCE_EMA
|
||
/// ISV[402] ε ← ISV[419] KL_PI_EMA
|
||
/// ISV[403] entropy_coef ← ISV[420] ENTROPY_OBSERVED_EMA
|
||
/// ISV[404] n_rollout_steps← ISV[421] ADVANTAGE_VAR_RATIO_EMA
|
||
/// ISV[405] per_α ← ISV[422] TD_KURTOSIS_EMA
|
||
/// ISV[406] reward_scale ← ISV[423] MEAN_ABS_PNL_EMA
|
||
///
|
||
/// R6's `step_with_lobsim` will call this AFTER the EMA producers
|
||
/// have updated the input slots. Idempotent on a single trainer
|
||
/// step but should fire exactly once.
|
||
pub fn launch_rl_controllers_per_step(&self) -> Result<()> {
|
||
let alpha = RL_LR_CONTROLLER_ALPHA;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_gamma_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_MEAN_TRADE_DURATION_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_gamma_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_target_tau_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_Q_DIVERGENCE_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_target_tau_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_ppo_clip_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_KL_PI_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_ppo_clip_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_entropy_coef_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_ENTROPY_OBSERVED_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_entropy_coef_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_rollout_steps_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_rollout_steps_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_per_alpha_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_per_alpha_controller")?;
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_reward_scale_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_MEAN_ABS_PNL_EMA_INDEX as i32,
|
||
)
|
||
.context("R5 launch rl_reward_scale_controller")?;
|
||
// R9 — PPO ratio clamp ceiling, anchored on ε at RL_PPO_CLIP_INDEX
|
||
// (which the rl_ppo_clip_controller just refreshed two calls up).
|
||
self.launch_isv_controller_3arg(
|
||
&self.rl_ppo_ratio_clamp_controller_fn,
|
||
alpha,
|
||
crate::rl::isv_slots::RL_PPO_CLIP_INDEX as i32,
|
||
)
|
||
.context("R9 launch rl_ppo_ratio_clamp_controller")?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R3: launch `ema_update_on_done` — done-gated EMA producer
|
||
/// that writes the per-batch closed-trade mean into `isv[slot_index]`.
|
||
/// `alpha` is the Wiener-α (caller must pre-floor at 0.4 per
|
||
/// `pearl_wiener_alpha_floor_for_nonstationary`). `obs_d` and
|
||
/// `dones_d` are `[b_size]` device buffers; only entries where
|
||
/// `dones[b] >= 0.5` contribute to the EMA observation. If no
|
||
/// batch closed a trade this step, the kernel is a no-op (ISV slot
|
||
/// retains its prior value).
|
||
///
|
||
/// Shared-memory tree reduce; single block, `b_size` threads.
|
||
/// Caller's responsibility: `b_size <= 1024` (one block max).
|
||
pub fn launch_ema_update_on_done(
|
||
&self,
|
||
slot_index: usize,
|
||
alpha: f32,
|
||
obs_d: &CudaSlice<f32>,
|
||
dones_d: &CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(obs_d.len(), b_size);
|
||
debug_assert_eq!(dones_d.len(), b_size);
|
||
debug_assert!(b_size <= 1024, "ema_update_on_done: b_size > 1024 needs multi-block reduce");
|
||
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 = slot_index as i32;
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ema_update_on_done_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(&slot_i)
|
||
.arg(&alpha)
|
||
.arg(obs_d)
|
||
.arg(dones_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("ema_update_on_done launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R3: launch `ema_update_per_step` — per-step EMA producer
|
||
/// that writes the per-batch mean into `isv[slot_index]` on every
|
||
/// call (no done-gate). Used for continuous EMAs (KL, entropy,
|
||
/// advantage variance). Same shared-mem tree-reduce as
|
||
/// `ema_update_on_done`.
|
||
pub fn launch_ema_update_per_step(
|
||
&self,
|
||
slot_index: usize,
|
||
alpha: f32,
|
||
obs_d: &CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(obs_d.len(), b_size);
|
||
debug_assert!(b_size <= 1024, "ema_update_per_step: b_size > 1024 needs multi-block reduce");
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (b_size as u32, 1, 1),
|
||
shared_mem_bytes: (b_size * std::mem::size_of::<f32>()) as u32,
|
||
};
|
||
let slot_i = slot_index as i32;
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ema_update_per_step_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(&slot_i)
|
||
.arg(&alpha)
|
||
.arg(obs_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("ema_update_per_step launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R4: launch `rl_action_kernel` — Thompson sampler over
|
||
/// the C51 Q distribution. Writes per-batch action indices to
|
||
/// `actions_d`. Advances per-batch xorshift32 PRNG state in place.
|
||
/// One block per batch, `N_ACTIONS` threads. `q_logits_d` is
|
||
/// `[b_size, N_ACTIONS, Q_N_ATOMS]` row-major.
|
||
pub fn launch_rl_action_kernel(
|
||
&mut self,
|
||
q_logits_d: &CudaSlice<f32>,
|
||
actions_d: &mut CudaSlice<i32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(q_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
let cfg = 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.rl_action_kernel_fn);
|
||
launch
|
||
.arg(q_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&mut self.prng_state_d)
|
||
.arg(actions_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("rl_action_kernel launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R4: launch `argmax_expected_q` — Bellman-target argmax
|
||
/// over expected Q per action. Deterministic (no PRNG). Pairs with
|
||
/// `launch_rl_action_kernel` per
|
||
/// `pearl_thompson_for_distributional_action_selection` (Thompson
|
||
/// for rollout, argmax for Bellman target).
|
||
pub fn launch_argmax_expected_q(
|
||
&self,
|
||
q_logits_d: &CudaSlice<f32>,
|
||
next_actions_d: &mut CudaSlice<i32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(q_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
|
||
debug_assert_eq!(next_actions_d.len(), b_size);
|
||
let cfg = 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_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(next_actions_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("argmax_expected_q launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R4: launch `log_pi_at_action` — per-batch
|
||
/// `log π(actions[b] | s_b)` via log-softmax + lookup. Feeds the
|
||
/// PPO importance ratio. One thread per batch entry.
|
||
pub fn launch_log_pi_at_action(
|
||
&self,
|
||
pi_logits_d: &CudaSlice<f32>,
|
||
actions_d: &CudaSlice<i32>,
|
||
log_pi_out_d: &mut CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(pi_logits_d.len(), b_size * N_ACTIONS);
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
debug_assert_eq!(log_pi_out_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.log_pi_at_action_fn);
|
||
launch
|
||
.arg(pi_logits_d)
|
||
.arg(actions_d)
|
||
.arg(log_pi_out_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("log_pi_at_action launch")?;
|
||
}
|
||
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
|
||
/// `scaled = rewards[b] * isv[RL_REWARD_SCALE_INDEX = 406]`,
|
||
/// then asymmetric clamp to `[-REWARD_CLAMP_LOSS, +REWARD_CLAMP_WIN]`
|
||
/// = `[-3, +1]` (per
|
||
/// `pearl_audit_unboundedness_for_implicit_asymmetry`).
|
||
/// Also writes per-step pre-clamp `max(|scaled|)` to
|
||
/// `ISV[RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX = 439]` for diag.
|
||
///
|
||
/// Single-block layout (no atomics per `pearl_no_atomicadd`).
|
||
/// Block dim = min(b_size, 256); shared memory = block_dim · 4 B
|
||
/// for the tree reduction.
|
||
pub fn launch_apply_reward_scale(
|
||
&self,
|
||
rewards_d: &mut CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(rewards_d.len(), b_size);
|
||
let block_x: u32 = (b_size as u32).min(256).max(1);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_x, 1, 1),
|
||
shared_mem_bytes: block_x * (std::mem::size_of::<f32>() as u32),
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.apply_reward_scale_fn);
|
||
launch.arg(&self.isv_d).arg(rewards_d).arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("apply_reward_scale launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase R3: launch `compute_advantage_return` — element-wise
|
||
/// `returns[b] = r + γ(1-done)·V(s_{t+1})`,
|
||
/// `advantages[b] = returns[b] − V(s_t)`. Reads γ from `ISV[400]`
|
||
/// (R1 bootstraps to 0.99; R5 adapts via trade-duration EMA).
|
||
///
|
||
/// Element-wise; trivially parallel. Block layout: 32-thread blocks
|
||
/// covering `b_size`.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn launch_compute_advantage_return(
|
||
&self,
|
||
rewards_d: &CudaSlice<f32>,
|
||
dones_d: &CudaSlice<f32>,
|
||
v_t_d: &CudaSlice<f32>,
|
||
v_tp1_d: &CudaSlice<f32>,
|
||
returns_d: &mut CudaSlice<f32>,
|
||
advantages_d: &mut CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(rewards_d.len(), b_size);
|
||
debug_assert_eq!(dones_d.len(), b_size);
|
||
debug_assert_eq!(v_t_d.len(), b_size);
|
||
debug_assert_eq!(v_tp1_d.len(), b_size);
|
||
debug_assert_eq!(returns_d.len(), b_size);
|
||
debug_assert_eq!(advantages_d.len(), b_size);
|
||
let b_size_i = b_size as i32;
|
||
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 mut launch = self.stream.launch_builder(&self.compute_advantage_return_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(rewards_d)
|
||
.arg(dones_d)
|
||
.arg(v_t_d)
|
||
.arg(v_tp1_d)
|
||
.arg(returns_d)
|
||
.arg(advantages_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("compute_advantage_return launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase E.2 + E.2-DEFER: synthetic-batch end-to-end step. Runs real
|
||
/// GPU forward + backward + per-head Adam on Q, π, V heads on a
|
||
/// SHARED encoder; combined per-head `grad_h_t` flows back into the
|
||
/// encoder via `PerceptionTrainer::backward_encoder_with_grad_h_t`
|
||
/// so the encoder is no longer BCE+aux-only.
|
||
///
|
||
/// All learning rates are sourced from `ISV[RL_LR_*_INDEX]` per
|
||
/// `pearl_controller_anchors_isv_driven`; the `lr` parameter is
|
||
/// removed (greenfield — no caller override). Bellman target
|
||
/// distribution is built on-device by `bellman_target_projection`
|
||
/// reading γ from `ISV[RL_GAMMA_INDEX=400]`.
|
||
///
|
||
/// Returns the per-head kernel-emitted losses and the combined-loss
|
||
/// total weighted by the ISV-driven λs.
|
||
/// 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],
|
||
) -> Result<IntegratedStepStats> {
|
||
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 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
|
||
.perception
|
||
.forward_encoder(snapshots)
|
||
.context("forward_encoder")?;
|
||
|
||
// ── Step 2: per-head LR controller emit + ISV host mirror ────
|
||
// The controller emits the bootstrap target on first call
|
||
// (sentinel-zero → 1e-3); subsequent calls Wiener-α blend.
|
||
// Use last step's losses as plateau-decay observations. At
|
||
// construction these are 0.0 (sentinel) and the controller's
|
||
// cold-start gate (early-return on observed_loss == 0) holds
|
||
// LR at LR_BOOTSTRAP for the first step.
|
||
self.launch_rl_lr_controller(
|
||
self.last_q_loss,
|
||
self.last_pi_loss,
|
||
self.last_v_loss,
|
||
)
|
||
.context("rl_lr_controller launch")?;
|
||
// Mapped-pinned ISV mirror refresh per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned` — was a raw
|
||
// `stream.memcpy_dtoh` (forbidden by
|
||
// `feedback_no_htod_htoh_only_mapped_pinned`).
|
||
let isv_host_fresh =
|
||
read_slice_d(&self.stream, &self.isv_d, self.isv_host.len())
|
||
.context("step_synthetic: read isv")?;
|
||
self.isv_host.copy_from_slice(&isv_host_fresh);
|
||
let lambdas = read_loss_lambdas_from_isv(&self.isv_host);
|
||
// Mutate each per-head Adam's lr field from the ISV mirror. The
|
||
// controller has just emitted into ISV[412..417], so this read
|
||
// reflects this step's learning rate. Reading ALL five before
|
||
// touching ANY AdamW keeps the field-by-field mutations
|
||
// borrow-checker safe and avoids per-head DtoH roundtrips.
|
||
let lr_bce = self.isv_host[RL_LR_BCE_INDEX];
|
||
let lr_q = self.isv_host[RL_LR_Q_INDEX];
|
||
let lr_pi = self.isv_host[RL_LR_PI_INDEX];
|
||
let lr_v = self.isv_host[RL_LR_V_INDEX];
|
||
let lr_aux = self.isv_host[RL_LR_AUX_INDEX];
|
||
// The aux head LR is consumed by the PerceptionTrainer's aux
|
||
// optimisers; the perception trainer reads it through its own
|
||
// mutator (`set_lr_aux`). The integrated trainer only needs to
|
||
// propagate the value — its own Adam instances are Q / π / V.
|
||
// BCE LR currently rides the existing `lr_cfc` on the perception
|
||
// side (the BCE head's grad path is interlocked with the CfC
|
||
// backward through the K-loop). Phase E.3+ separates the BCE
|
||
// optimiser; until then we keep the BCE slot in ISV updated for
|
||
// forward-looking diagnostics + the controller's first-obs
|
||
// bootstrap, and we propagate `lr_aux` via the perception
|
||
// trainer's existing mutator hook.
|
||
let _ = lr_bce;
|
||
self.perception.set_lr_aux(lr_aux);
|
||
self.dqn_w_adam.lr = lr_q;
|
||
self.dqn_b_adam.lr = lr_q;
|
||
self.policy_w_adam.lr = lr_pi;
|
||
self.policy_b_adam.lr = lr_pi;
|
||
self.value_w_adam.lr = lr_v;
|
||
self.value_b_adam.lr = lr_v;
|
||
|
||
// Borrow encoder hidden state for forward kernels.
|
||
let h_t_borrow: &CudaSlice<f32> = self.perception.h_t_view();
|
||
debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM);
|
||
|
||
// ── Step 3: allocate per-step scratch ────────────────────────
|
||
let k_dqn = N_ACTIONS * Q_N_ATOMS;
|
||
let mut q_logits_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
|
||
let mut pi_logits_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||
let mut v_pred_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
|
||
// Loss accumulators (atomicAdd into single floats from kernels).
|
||
let q_loss_d = self.stream.alloc_zeros::<f32>(1)?;
|
||
let pi_loss_d = self.stream.alloc_zeros::<f32>(1)?;
|
||
let pi_loss_entropy_d = self.stream.alloc_zeros::<f32>(1)?;
|
||
// Per-batch V loss scratch (reduced below).
|
||
let mut v_loss_per_batch_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
let mut v_loss_sum_d = self.stream.alloc_zeros::<f32>(1)?;
|
||
|
||
// Per-batch diagnostic scratch (consumed by the surrogate kernel).
|
||
let mut pi_log_prob_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
let mut entropy_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
|
||
// Per-batch grad scratch (reduced across batches via reduce_axis0).
|
||
let mut q_grad_logits_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
|
||
let mut q_grad_w_per_batch_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * k_dqn * HIDDEN_DIM)?;
|
||
let mut q_grad_b_per_batch_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
|
||
let mut q_grad_h_t_d = self.stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)?;
|
||
let mut q_grad_w_d = self.stream.alloc_zeros::<f32>(k_dqn * HIDDEN_DIM)?;
|
||
let mut q_grad_b_d = self.stream.alloc_zeros::<f32>(k_dqn)?;
|
||
|
||
let mut pi_grad_logits_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||
let mut pi_grad_w_per_batch_d =
|
||
self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS * HIDDEN_DIM)?;
|
||
let mut pi_grad_b_per_batch_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||
let mut pi_grad_h_t_d = self.stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)?;
|
||
let mut pi_grad_w_d = self.stream.alloc_zeros::<f32>(N_ACTIONS * HIDDEN_DIM)?;
|
||
let mut pi_grad_b_d = self.stream.alloc_zeros::<f32>(N_ACTIONS)?;
|
||
|
||
let mut v_grad_w_per_batch_d = self.stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)?;
|
||
let mut v_grad_b_per_batch_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
let mut v_grad_h_t_d = self.stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)?;
|
||
let mut v_grad_w_d = self.stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
|
||
let mut v_grad_b_d = self.stream.alloc_zeros::<f32>(1)?;
|
||
|
||
// Bellman projection scratch (item 2).
|
||
let mut q_target_full_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
|
||
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 (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) ────────────────────────
|
||
// 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(&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")?;
|
||
self.value_head
|
||
.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).
|
||
{
|
||
let mut pi_loss_d_mut = pi_loss_d;
|
||
let mut pi_loss_entropy_d_mut = pi_loss_entropy_d;
|
||
self.policy_head
|
||
.surrogate_forward(
|
||
&pi_logits_d,
|
||
&self.log_pi_old_d,
|
||
&self.actions_d,
|
||
&self.advantages_d,
|
||
&self.isv_d,
|
||
b_size,
|
||
&mut pi_log_prob_d,
|
||
&mut entropy_d,
|
||
&mut pi_loss_d_mut,
|
||
&mut pi_loss_entropy_d_mut,
|
||
)
|
||
.context("policy_head.surrogate_forward")?;
|
||
let l_pi_host = read_scalar_d(&self.stream, &pi_loss_d_mut)?;
|
||
let _l_pi_ent_host = read_scalar_d(&self.stream, &pi_loss_entropy_d_mut)?;
|
||
self.last_pi_loss = l_pi_host;
|
||
}
|
||
|
||
// ── Step 6a: Bellman target build (item 2) ───────────────────
|
||
// 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 (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.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.sampled_next_actions_d,
|
||
b_size,
|
||
&mut q_target_action_d,
|
||
)
|
||
.context("dqn_head.select_action_atoms (sampled_next_actions)")?;
|
||
self.dqn_head
|
||
.project_bellman_target(
|
||
&q_target_action_d,
|
||
&self.sampled_rewards_d,
|
||
&self.sampled_dones_d,
|
||
&self.isv_d,
|
||
b_size,
|
||
&mut target_dist_d,
|
||
)
|
||
.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;
|
||
self.dqn_head
|
||
.backward_logits(
|
||
&q_logits_d,
|
||
&target_dist_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 (sampled_actions) [R7d off-policy]")?;
|
||
let l_q_host = read_scalar_d(&self.stream, &q_loss_d_mut)?;
|
||
// Mirror for next step's LR controller plateau detection.
|
||
self.last_q_loss = l_q_host / (b_size as f32);
|
||
|
||
self.dqn_head
|
||
.backward_to_w_b_h(
|
||
&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(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)?;
|
||
|
||
// ── Step 7: PPO backward (logits → grad_w/b/h_t) ─────────────
|
||
self.policy_head
|
||
.surrogate_backward_logits(
|
||
&pi_logits_d,
|
||
&self.log_pi_old_d,
|
||
&self.actions_d,
|
||
&self.advantages_d,
|
||
&self.isv_d,
|
||
b_size,
|
||
&mut pi_grad_logits_d,
|
||
)
|
||
.context("policy_head.surrogate_backward_logits")?;
|
||
|
||
self.policy_head
|
||
.backward_to_w_b_h(
|
||
h_t_borrow,
|
||
&pi_grad_logits_d,
|
||
b_size,
|
||
&mut pi_grad_w_per_batch_d,
|
||
&mut pi_grad_b_per_batch_d,
|
||
&mut pi_grad_h_t_d,
|
||
)
|
||
.context("policy_head.backward_to_w_b_h")?;
|
||
|
||
self.launch_reduce_axis0(
|
||
&pi_grad_w_per_batch_d,
|
||
b_size,
|
||
N_ACTIONS * HIDDEN_DIM,
|
||
&mut pi_grad_w_d,
|
||
)?;
|
||
self.launch_reduce_axis0(&pi_grad_b_per_batch_d, b_size, N_ACTIONS, &mut pi_grad_b_d)?;
|
||
|
||
// ── Step 8: V backward ───────────────────────────────────────
|
||
self.value_head
|
||
.backward(
|
||
h_t_borrow,
|
||
&v_pred_d,
|
||
&self.returns_d,
|
||
b_size,
|
||
&mut v_loss_per_batch_d,
|
||
&mut v_grad_w_per_batch_d,
|
||
&mut v_grad_b_per_batch_d,
|
||
&mut v_grad_h_t_d,
|
||
)
|
||
.context("value_head.backward")?;
|
||
|
||
self.launch_reduce_axis0(&v_grad_w_per_batch_d, b_size, HIDDEN_DIM, &mut v_grad_w_d)?;
|
||
self.launch_reduce_axis0(&v_grad_b_per_batch_d, b_size, 1, &mut v_grad_b_d)?;
|
||
self.launch_reduce_axis0(&v_loss_per_batch_d, b_size, 1, &mut v_loss_sum_d)?;
|
||
let l_v_sum_host = read_scalar_d(&self.stream, &v_loss_sum_d)?;
|
||
let l_v_host = l_v_sum_host / (b_size as f32);
|
||
// Mirror for next step's LR controller plateau detection.
|
||
self.last_v_loss = l_v_host;
|
||
|
||
// ── Step 9: Adam updates on each head's w and b ──────────────
|
||
self.dqn_w_adam
|
||
.step(&mut self.dqn_head.w_d, &q_grad_w_d)
|
||
.context("dqn_w_adam.step")?;
|
||
self.dqn_b_adam
|
||
.step(&mut self.dqn_head.b_d, &q_grad_b_d)
|
||
.context("dqn_b_adam.step")?;
|
||
self.policy_w_adam
|
||
.step(&mut self.policy_head.w_d, &pi_grad_w_d)
|
||
.context("policy_w_adam.step")?;
|
||
self.policy_b_adam
|
||
.step(&mut self.policy_head.b_d, &pi_grad_b_d)
|
||
.context("policy_b_adam.step")?;
|
||
self.value_w_adam
|
||
.step(&mut self.value_head.w_d, &v_grad_w_d)
|
||
.context("value_w_adam.step")?;
|
||
self.value_b_adam
|
||
.step(&mut self.value_head.b_d, &v_grad_b_d)
|
||
.context("value_b_adam.step")?;
|
||
|
||
// ── Step 10: encoder grad combine (Phase E.2-DEFER item 1) ───
|
||
// 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
|
||
// feedback_single_source_of_truth_no_duplicates).
|
||
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,
|
||
&pi_grad_h_t_d,
|
||
lambdas.pi,
|
||
&mut self.grad_h_t_combined_d,
|
||
)?;
|
||
accumulate_grad_h(
|
||
&self.stream,
|
||
&self.grad_h_accumulate_fn,
|
||
&v_grad_h_t_d,
|
||
lambdas.v,
|
||
&mut self.grad_h_t_combined_d,
|
||
)?;
|
||
|
||
// ── Step 11: encoder backward (Phase E.3a) ───────────────────
|
||
// Seed the perception trainer's per-K hidden-state grad buffer
|
||
// from `grad_h_t_combined_d` at slot K-1 (the only slot the
|
||
// RL heads consumed `h_t` from) and run the shared encoder
|
||
// backward kernel chain. This updates ALL encoder params:
|
||
// CfC ×4, LN ×2, VSN, attn_q, Mamba2 L1+L2 — closes the last
|
||
// deferred item from Phase E.2.
|
||
self.perception
|
||
.backward_encoder_with_grad_h_t(
|
||
&self.grad_h_t_combined_d,
|
||
self.cfg.perception.n_batch,
|
||
self.cfg.perception.seq_len,
|
||
)
|
||
.context("perception.backward_encoder_with_grad_h_t")?;
|
||
|
||
// ── Step 11b: feed step_synthetic's EMA inputs (entropy,
|
||
// td_kurtosis, KL) into their controller ISV slots. Deferred
|
||
// to AFTER the encoder backward so the `&self.perception`
|
||
// borrow (held by h_t_borrow earlier in this function) is
|
||
// released before the `&mut self` launches here.
|
||
//
|
||
// 1. `entropy_observed_ema` (ISV[420]) ← per-batch entropy
|
||
// from PPO surrogate forward (entropy_d still live here).
|
||
// ema_update_per_step does the mean reduce internally,
|
||
// so entropy_d goes in directly.
|
||
// 2. `td_kurtosis_ema` (ISV[422]) ← kurtosis of per-sample
|
||
// CE loss (td_per_sample_d). rl_kurtosis_b reduces to
|
||
// a scalar, then ema_update_per_step with b_size=1.
|
||
//
|
||
// Both updates take effect on the NEXT step's controller
|
||
// fire (controllers fire in step_with_lobsim BEFORE
|
||
// step_synthetic). One-step lag is acceptable — controllers
|
||
// adapt over many steps.
|
||
self.launch_ema_update_per_step(
|
||
crate::rl::isv_slots::RL_ENTROPY_OBSERVED_EMA_INDEX,
|
||
RL_LR_CONTROLLER_ALPHA,
|
||
&entropy_d,
|
||
b_size,
|
||
)
|
||
.context("ema_update_per_step(entropy_observed)")?;
|
||
// Streaming kurtosis kernel writes directly to
|
||
// RL_TD_KURTOSIS_EMA_INDEX; no separate ema_update_per_step
|
||
// needed (the streaming kernel IS the EMA — folds across STEPS).
|
||
self.launch_kurtosis(&self.td_per_sample_d.clone(), b_size)
|
||
.context("launch_kurtosis(td_per_sample_d)")?;
|
||
|
||
// kl_pi EMA — Schulman-style approximation
|
||
// `mean(log π_old(a) − log π_new(a))` over the sampled action.
|
||
// log_pi_old_d was recorded at action-sample time (per-batch
|
||
// log π under the previous policy at the action taken);
|
||
// pi_log_prob_d (= log π_new at the same action) is the PPO
|
||
// surrogate forward's per-batch output. Mean-reduce to scalar
|
||
// → ema_update_per_step with b_size=1.
|
||
{
|
||
let block_dim = (b_size.next_power_of_two() as u32).max(1);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_dim, 1, 1),
|
||
shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_kl_approx_b_fn);
|
||
launch
|
||
.arg(&self.log_pi_old_d)
|
||
.arg(&pi_log_prob_d)
|
||
.arg(&mut self.ema_input_scratch_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_kl_approx_b launch")?;
|
||
}
|
||
}
|
||
|
||
// R9 diag — per-batch max |log π_new − log π_old| → ISV[441].
|
||
// Same input buffers as rl_kl_approx_b but max-reduce |diff|
|
||
// instead of mean signed-diff. Surfaces the largest single
|
||
// policy excursion in this step's batch for the diag JSONL.
|
||
// Cross-check: log(isv[RL_PPO_RATIO_CLAMP_MAX_INDEX]) is the
|
||
// ceiling — when this slot exceeds that, the ratio clamp fired
|
||
// this step.
|
||
{
|
||
let block_dim = (b_size.next_power_of_two() as u32).max(1);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_dim, 1, 1),
|
||
shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ppo_log_ratio_abs_max_b_fn);
|
||
launch
|
||
.arg(&self.log_pi_old_d)
|
||
.arg(&pi_log_prob_d)
|
||
.arg(&mut self.isv_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("ppo_log_ratio_abs_max_b launch")?;
|
||
}
|
||
}
|
||
self.launch_ema_update_per_step(
|
||
crate::rl::isv_slots::RL_KL_PI_EMA_INDEX,
|
||
RL_LR_CONTROLLER_ALPHA,
|
||
&self.ema_input_scratch_d.clone(),
|
||
1,
|
||
)
|
||
.context("ema_update_per_step(kl_pi)")?;
|
||
|
||
// Per-head gradient-norm EMAs (Q, π, V) — feed the
|
||
// signal-modulated LR controller (rl_lr_controller's target
|
||
// formula: `lr_prev × TARGET_GRAD_NORM / observed_grad_norm`).
|
||
// grad_w_*_d buffers are still in scope here (allocated at
|
||
// top of step_synthetic, consumed by the Adam steps in Step 9
|
||
// but not freed until function return); the l2_norm launches
|
||
// here read them after the Adam step has already used them.
|
||
let k_dqn_hidden = (N_ACTIONS * Q_N_ATOMS * HIDDEN_DIM) as usize;
|
||
self.launch_l2_norm(&q_grad_w_d.clone(), k_dqn_hidden)
|
||
.context("launch_l2_norm(q_grad_w_d)")?;
|
||
self.launch_ema_update_per_step(
|
||
crate::rl::isv_slots::RL_Q_GRAD_NORM_EMA_INDEX,
|
||
RL_LR_CONTROLLER_ALPHA,
|
||
&self.ema_input_scratch_d.clone(),
|
||
1,
|
||
)
|
||
.context("ema_update_per_step(q_grad_norm)")?;
|
||
|
||
let n_actions_hidden = (N_ACTIONS * HIDDEN_DIM) as usize;
|
||
self.launch_l2_norm(&pi_grad_w_d.clone(), n_actions_hidden)
|
||
.context("launch_l2_norm(pi_grad_w_d)")?;
|
||
self.launch_ema_update_per_step(
|
||
crate::rl::isv_slots::RL_PI_GRAD_NORM_EMA_INDEX,
|
||
RL_LR_CONTROLLER_ALPHA,
|
||
&self.ema_input_scratch_d.clone(),
|
||
1,
|
||
)
|
||
.context("ema_update_per_step(pi_grad_norm)")?;
|
||
|
||
self.launch_l2_norm(&v_grad_w_d.clone(), HIDDEN_DIM)
|
||
.context("launch_l2_norm(v_grad_w_d)")?;
|
||
self.launch_ema_update_per_step(
|
||
crate::rl::isv_slots::RL_V_GRAD_NORM_EMA_INDEX,
|
||
RL_LR_CONTROLLER_ALPHA,
|
||
&self.ema_input_scratch_d.clone(),
|
||
1,
|
||
)
|
||
.context("ema_update_per_step(v_grad_norm)")?;
|
||
|
||
// ── Step 12: compose stats ───────────────────────────────────
|
||
// BCE / aux losses are NOT read this phase — perception is driven
|
||
// separately by callers via its existing step_batched path. The
|
||
// encoder backward here learns ONLY from the Q+π+V grad_h_t
|
||
// contributions, which is by design (BCE+aux own their own loss
|
||
// sources via step_batched).
|
||
let l_bce = 0.0_f32;
|
||
let l_aux = 0.0_f32;
|
||
let l_pi = self.last_pi_loss;
|
||
|
||
// Normalise the Q CE by batch (matches l_pi / l_v conventions).
|
||
let l_q = l_q_host / (b_size as f32);
|
||
|
||
// Combined: λ-weighted average across 5 heads (per
|
||
// pearl_loss_balance_controller, sum over 5 heads / 5).
|
||
let l_total = (lambdas.bce * l_bce
|
||
+ lambdas.q * l_q
|
||
+ lambdas.pi * l_pi
|
||
+ lambdas.v * l_v_host
|
||
+ lambdas.aux * l_aux)
|
||
/ 5.0;
|
||
|
||
Ok(IntegratedStepStats {
|
||
l_bce,
|
||
l_q,
|
||
l_pi,
|
||
l_v: l_v_host,
|
||
l_aux,
|
||
l_total,
|
||
lambdas,
|
||
})
|
||
}
|
||
|
||
/// Phase R6: run one integrated training step driven by a real
|
||
/// `LobSimCuda` (via the narrow `RlLobBackend` trait, dep-cycle
|
||
/// break only). GPU-pure for the env interaction: Thompson-
|
||
/// sampled actions land in `lobsim.market_targets_d` via the
|
||
/// `actions_to_market_targets` kernel; the fill runs via
|
||
/// `lobsim.step_fill_from_market_targets`; rewards + done flags
|
||
/// extract from `lobsim.pos_d` via the
|
||
/// `extract_realized_pnl_delta` kernel. No per-batch host loop.
|
||
///
|
||
/// Replaces the synthetic args of [`step_synthetic`] with values
|
||
/// derived from a real action sampling + reward signal:
|
||
///
|
||
/// 1. `forward_encoder(snapshots)` → `h_t [B × HIDDEN_DIM]`.
|
||
/// 2. Forward Q-head + V-head; read `q_logits`, `v_pred` to host.
|
||
/// 3. Thompson-sample one action per batch from the per-action
|
||
/// atom distribution (per
|
||
/// `pearl_thompson_for_distributional_action_selection`).
|
||
/// 4. For each batch:
|
||
/// a. `lobsim.apply_snapshot(snapshot)`,
|
||
/// b. `lobsim.submit_action(batch_idx, action, ts_ns)`,
|
||
/// c. `lobsim.step_event(batch_idx, ts_ns, trade_signed_vol)`
|
||
/// → `(reward, done)`.
|
||
/// 5. Build the synthetic-style arg vectors from the real action /
|
||
/// reward / done values:
|
||
/// * `actions[b] = sampled action`,
|
||
/// * `rewards[b] = lobsim reward`,
|
||
/// * `dones[b] = 1.0 if lobsim done else 0.0`,
|
||
/// * `next_actions[b] = argmax_a Q(s_t, a)` (host-side
|
||
/// expected-value argmax — Phase F may
|
||
/// move this to a kernel),
|
||
/// * `advantages[b] = Q(s_t, a_t) - V(s_t)` (one-step
|
||
/// advantage from the same-state V),
|
||
/// * `returns[b] = reward + γ × (1 - done) × V(s_t)`
|
||
/// (one-step bootstrap target; γ is
|
||
/// read from the host-mirrored
|
||
/// `ISV[RL_GAMMA_INDEX]`),
|
||
/// * `log_pi_old[b] = log π_sampled(a_t | s_t)` recorded at
|
||
/// sample time (note: `step_synthetic`
|
||
/// re-runs the policy forward on the
|
||
/// same `h_t`, so log π_new == log
|
||
/// π_old on the first PPO pass — the
|
||
/// importance ratio is 1 and the
|
||
/// clipped surrogate's gradient flows
|
||
/// from the entropy + advantage path,
|
||
/// not the ratio. Phase F adds a
|
||
/// proper rollout buffer so multiple
|
||
/// PPO epochs see a non-trivial ratio).
|
||
/// 6. Delegate to [`Self::step_synthetic`] for the actual backward
|
||
/// + per-head Adam update + encoder backward — keeping the
|
||
/// training kernel chain a single source of truth per
|
||
/// `feedback_single_source_of_truth_no_duplicates`. The
|
||
/// encoder forward runs TWICE per step (once here for action
|
||
/// sampling, once inside `step_synthetic`); the encoder is
|
||
/// deterministic given the same input, so this is a pure
|
||
/// compute redundancy that Phase F can fuse.
|
||
///
|
||
/// Returns the same stats struct as `step_synthetic`. Tests
|
||
/// (`dqn_toy`, `ppo_toy`, `integrated_trainer_smoke`) drive this in
|
||
/// a loop with `MockLobEnv` to gate that the trainer converges on
|
||
/// the toy bandit optimal action.
|
||
pub fn step_with_lobsim(
|
||
&mut self,
|
||
snapshots: &[Mbp10RawInput],
|
||
next_snapshots: &[Mbp10RawInput],
|
||
lobsim: &mut dyn RlLobBackend,
|
||
) -> Result<IntegratedStepStats> {
|
||
let b_size = self.cfg.perception.n_batch;
|
||
if b_size == 0 {
|
||
anyhow::bail!("step_with_lobsim: empty batch (n_batch = 0)");
|
||
}
|
||
if snapshots.is_empty() {
|
||
anyhow::bail!("step_with_lobsim: snapshots empty (need ≥ 1 for apply_snapshot)");
|
||
}
|
||
if next_snapshots.len() != snapshots.len() {
|
||
anyhow::bail!(
|
||
"step_with_lobsim: next_snapshots.len() ({}) must match snapshots.len() ({}) — \
|
||
R8's CLI binary builds them via MultiHorizonLoader::next_sequence_pair (R2) which \
|
||
guarantees this invariant for adjacent (s_t, s_{{t+1}}) windows",
|
||
next_snapshots.len(),
|
||
snapshots.len()
|
||
);
|
||
}
|
||
|
||
// ── Step 1a (R7c): encoder forward on NEXT snapshots to land
|
||
// h_{t+1} at slot K-1. Done FIRST so the subsequent
|
||
// forward_encoder(snapshots) call leaves perception's internal
|
||
// forward state (h_new_per_k_d / CfC h_state_d) populated for
|
||
// the CURRENT-step encoder backward in step_synthetic — the
|
||
// most recent forward_encoder call wins. We DtoD-copy
|
||
// perception.h_t_d into the trainer-owned self.h_tp1_d
|
||
// immediately so the second forward_encoder doesn't clobber it.
|
||
let _ = self
|
||
.perception
|
||
.forward_encoder(next_snapshots)
|
||
.context("step_with_lobsim: forward_encoder(next_snapshots)")?;
|
||
{
|
||
let nbytes = b_size * HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (src, _gs) = self.perception.h_t_view().device_ptr(&self.stream);
|
||
let (dst, _gd) = self.h_tp1_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes, s)
|
||
.context("step_with_lobsim: DtoD perception.h_t_d → self.h_tp1_d")?;
|
||
}
|
||
}
|
||
|
||
// ── Step 1b: encoder forward on CURRENT snapshots — leaves
|
||
// perception's internal forward state primed for the encoder
|
||
// backward in step_synthetic, AND lands h_t at slot K-1 for
|
||
// the action-sampling forwards below.
|
||
let _ = self
|
||
.perception
|
||
.forward_encoder(snapshots)
|
||
.context("step_with_lobsim: forward_encoder(snapshots)")?;
|
||
|
||
// ── Step 2: Q + V forwards on h_t for action sampling. ────────
|
||
let h_t_borrow: &CudaSlice<f32> = self.perception.h_t_view();
|
||
debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM);
|
||
debug_assert_eq!(self.h_tp1_d.len(), b_size * HIDDEN_DIM);
|
||
|
||
let k_dqn = N_ACTIONS * Q_N_ATOMS;
|
||
let mut q_logits_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
|
||
let mut q_logits_tp1_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
|
||
let mut v_pred_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
let mut v_pred_tp1_d = self.stream.alloc_zeros::<f32>(b_size)?;
|
||
|
||
self.dqn_head
|
||
.forward(h_t_borrow, b_size, &mut q_logits_d)
|
||
.context("step_with_lobsim: dqn_head.forward(h_t)")?;
|
||
self.dqn_head
|
||
.forward(&self.h_tp1_d, b_size, &mut q_logits_tp1_d)
|
||
.context("step_with_lobsim: dqn_head.forward(h_tp1) for Double-DQN argmax")?;
|
||
self.value_head
|
||
.forward(h_t_borrow, b_size, &mut v_pred_d)
|
||
.context("step_with_lobsim: value_head.forward(h_t)")?;
|
||
self.value_head
|
||
.forward(&self.h_tp1_d, b_size, &mut v_pred_tp1_d)
|
||
.context("step_with_lobsim: value_head.forward(h_tp1) for true V(s_{t+1})")?;
|
||
|
||
// ── Step 2b: Forward π logits for log_pi_old. ─────────────────
|
||
let mut pi_logits_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||
self.policy_head
|
||
.forward_logits(h_t_borrow, b_size, &mut pi_logits_d)
|
||
.context("step_with_lobsim: policy_head.forward_logits")?;
|
||
|
||
// ── Step 3 (Phase R7b): GPU-pure action sampling. ─────────────
|
||
// R4's rl_action_kernel performs Thompson sampling on device
|
||
// (per-batch xorshift32 PRNG, per-action softmax + CDF walk,
|
||
// argmax over sampled per-action returns). argmax_expected_q
|
||
// computes the Bellman-target argmax (deterministic, expected
|
||
// value rather than sampled) per
|
||
// `pearl_thompson_for_distributional_action_selection`.
|
||
// log_pi_at_action computes log π_old via log-softmax + lookup.
|
||
//
|
||
// The host Thompson + host argmax + host log-softmax loops the
|
||
// flawed Phase F+G shipped are GONE. The xorshift32 PRNG state
|
||
// is per-batch device-resident; the host-side step_counter
|
||
// increment that drove the old ChaCha8 host RNG is also gone.
|
||
self.step_counter = self.step_counter.wrapping_add(1);
|
||
|
||
{
|
||
let cfg = 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.rl_action_kernel_fn);
|
||
launch
|
||
.arg(&q_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&mut self.prng_state_d)
|
||
.arg(&mut self.actions_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_action_kernel launch (Thompson)")?;
|
||
}
|
||
}
|
||
|
||
// argmax over expected Q for next_actions — Double-DQN: argmax
|
||
// of ONLINE Q at h_{t+1} (per
|
||
// `pearl_thompson_for_distributional_action_selection`'s
|
||
// distinction between Thompson selector (on h_t for rollout)
|
||
// and argmax Bellman selector (on h_{t+1} for target)).
|
||
//
|
||
// R7c fix: was reading q_logits_d (online Q at h_t) which
|
||
// produced an off-by-one-time-index Bellman argmax since R4.
|
||
// Now reads q_logits_tp1_d (online Q at h_{t+1}) which is the
|
||
// canonical Double-DQN target-action selector.
|
||
{
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
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_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&mut self.next_actions_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("argmax_expected_q launch (Bellman target on h_tp1)")?;
|
||
}
|
||
}
|
||
|
||
// log π_old at the sampled action — feeds the PPO importance
|
||
// ratio. Reads actions_d (just written by rl_action_kernel).
|
||
{
|
||
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.log_pi_at_action_fn);
|
||
launch
|
||
.arg(&pi_logits_d)
|
||
.arg(&self.actions_d)
|
||
.arg(&mut self.log_pi_old_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("log_pi_at_action launch")?;
|
||
}
|
||
}
|
||
|
||
// ── Step 5 (Phase R6): GPU-pure env step. ─────────────────────
|
||
// No per-batch host loop — actions land in lobsim's
|
||
// market_targets_d via the `actions_to_market_targets` kernel
|
||
// (which reads current position_lots from lobsim.pos_d to
|
||
// resolve conditional Flat-from-Long/Short sizes), the fill
|
||
// runs via `step_fill_from_market_targets`, and rewards +
|
||
// done flags extract from the post-fill Pos array via
|
||
// `extract_realized_pnl_delta` (reading pos.realized_pnl
|
||
// delta against the trainer's prev_realized_pnl_d snapshot
|
||
// taken at end of the previous step).
|
||
//
|
||
// Phase E.3b broadcasts the LAST snapshot in the window to
|
||
// the book; Phase F may stream every snapshot through.
|
||
let last_snap = snapshots
|
||
.last()
|
||
.expect("snapshots non-empty: guarded above");
|
||
lobsim
|
||
.apply_snapshot(last_snap)
|
||
.context("step_with_lobsim: lobsim.apply_snapshot")?;
|
||
|
||
// GPU translate (action index → market_targets {side, size}).
|
||
// Reads current position_lots from lobsim.pos_d for Flat-from-*
|
||
// actions; reads actions from self.actions_d (just written by
|
||
// rl_action_kernel above).
|
||
let pos_bytes_i = lobsim.pos_bytes() as i32;
|
||
let b_size_i = b_size as i32;
|
||
{
|
||
let (pos_d_ref, market_targets_d) = lobsim.pos_and_market_targets_mut();
|
||
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.actions_to_market_targets_fn);
|
||
launch
|
||
.arg(&self.actions_d)
|
||
.arg(pos_d_ref)
|
||
.arg(market_targets_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("actions_to_market_targets launch")?;
|
||
}
|
||
}
|
||
|
||
// Run the fill + pnl_track on whatever's now in
|
||
// market_targets_d.
|
||
lobsim
|
||
.step_fill_from_market_targets(last_snap.ts_ns)
|
||
.context("step_with_lobsim: lobsim.step_fill_from_market_targets")?;
|
||
|
||
// Extract per-batch (reward, done) deltas vs the trainer's
|
||
// prev_realized_pnl_d snapshot. Updates prev_* in place for
|
||
// the next step's delta.
|
||
{
|
||
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
|
||
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.extract_realized_pnl_delta_fn);
|
||
launch
|
||
.arg(pos_d_ref)
|
||
.arg(&mut self.prev_realized_pnl_d)
|
||
.arg(&mut self.prev_position_lots_d)
|
||
.arg(&mut self.rewards_d)
|
||
.arg(&mut self.dones_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("extract_realized_pnl_delta launch")?;
|
||
}
|
||
}
|
||
|
||
// Trade-duration counter update + done-gated EMA emit. The
|
||
// counter increments every step; on done it emits the count
|
||
// (= trade duration in events) into `trade_duration_emit_d`
|
||
// and resets. `ema_update_on_done` then folds the per-batch
|
||
// duration values into ISV[417] for the rl_gamma controller.
|
||
{
|
||
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.rl_step_counter_update_fn);
|
||
launch
|
||
.arg(&self.dones_d)
|
||
.arg(&mut self.steps_since_done_d)
|
||
.arg(&mut self.trade_duration_emit_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_step_counter_update launch")?;
|
||
}
|
||
}
|
||
{
|
||
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_TRADE_DURATION_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.trade_duration_emit_d)
|
||
.arg(&self.dones_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("ema_update_on_done(mean_trade_duration) launch")?;
|
||
}
|
||
}
|
||
|
||
// ── Step 6 (Phase R7c): GPU-pure post-fill pipeline. ──────────
|
||
// All actions / next_actions / log_pi_old are already in
|
||
// trainer-owned device buffers (R4 kernels wrote them in Step 3).
|
||
// No host uploads here. V(s_t) and V(s_{t+1}) BOTH live on
|
||
// device — v_pred_d from value_head.forward(h_t),
|
||
// v_pred_tp1_d from value_head.forward(h_tp1) (R7c data
|
||
// correctness lift: was an alias of v_pred_d in R7a/R7b).
|
||
|
||
// |reward| → reward_abs_d (input for mean_abs_pnl EMA).
|
||
// Inline launch (vs `launch_abs_copy`) because the per-method
|
||
// 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")?;
|
||
}
|
||
}
|
||
|
||
// Update ISV[423] MEAN_ABS_PNL_EMA from |reward| over ALL
|
||
// non-zero reward events (NOT just trade closes).
|
||
//
|
||
// Originally gated on dones_d so only closed-trade realized
|
||
// PnL contributed to the EMA. Cluster smoke `alpha-rl-9cbpj`
|
||
// diag revealed that 170 of 266 non-zero reward events occur
|
||
// on non-done steps (mid-trade PnL deltas from trail-stop
|
||
// adjustments / mark-to-market / partial fills). These
|
||
// mid-trade swings can be 2-3× larger than realized close
|
||
// PnL — invisible to a done-gated EMA, they end up scaled
|
||
// by `reward_scale = 1/mean_abs_pnl_close` which is calibrated
|
||
// for the smaller magnitude. Result: V regression target
|
||
// spikes ×50-100 on volatile mid-trade steps (l_v=3,456
|
||
// observed at step 590 with raw $2,390 PnL × scale=0.0011).
|
||
//
|
||
// The kernel's "done" parameter is really a generic gate
|
||
// tested as `d >= 0.5`. Passing reward_abs_d as the gate
|
||
// gives "gate on |reward| ≥ 0.5" (= "non-zero reward" for
|
||
// any practical magnitude in dollars) — exactly the
|
||
// semantics we want. Zero-reward steps stay excluded so the
|
||
// EMA isn't biased toward zero on idle steps.
|
||
{
|
||
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.reward_abs_d) // gate on |reward|>0, not just on dones
|
||
.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) +
|
||
// asymmetric clamp [-3, +1] + per-step pre-clamp max diag.
|
||
// Reads ISV[406] which the controller just updated; writes
|
||
// ISV[439] for the diag emitter.
|
||
//
|
||
// Single-block layout: tree reduction needs every thread in
|
||
// the same block to participate in __syncthreads, so we cap
|
||
// block_x at min(b_size, 256). For b_size > 256 the kernel's
|
||
// grid-stride loop walks the rest.
|
||
{
|
||
let block_x: u32 = (b_size as u32).min(256).max(1);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_x, 1, 1),
|
||
shared_mem_bytes: block_x * (std::mem::size_of::<f32>() as u32),
|
||
};
|
||
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_pred_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")?;
|
||
}
|
||
}
|
||
// γ is read on-device by compute_advantage_return from
|
||
// ISV[400] — no host gamma scalar needed in this hot path
|
||
// post-R7b (the bootstrap-fallback host read of γ that the
|
||
// flawed branch did is gone with the host advantage loop).
|
||
|
||
// Feed `advantage_var_ratio_ema` (ISV[421] → rl_rollout_steps
|
||
// controller) from this step's advantages_d. Streaming kernel
|
||
// writes directly to ISV[421] (folds across STEPS — works at
|
||
// b_size=1 where per-batch variance is undefined). No
|
||
// ema_update_per_step needed downstream.
|
||
self.launch_var_over_abs_mean(&self.advantages_d.clone(), b_size)
|
||
.context("launch_var_over_abs_mean(advantages_d)")?;
|
||
|
||
// ── 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.
|
||
//
|
||
// 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.
|
||
// Mapped-pinned staging for ISV mirror refresh per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned`. The full ISV
|
||
// slice is 424 floats; one mapped-pinned alloc + DtoD per
|
||
// step is bounded.
|
||
let isv_host_fresh =
|
||
read_slice_d(&self.stream, &self.isv_d, self.isv_host.len())
|
||
.context("step_with_lobsim: read isv (pre-PER)")?;
|
||
self.isv_host.copy_from_slice(&isv_host_fresh);
|
||
|
||
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 td_per_sample_host =
|
||
read_slice_d(&self.stream, &self.td_per_sample_d, b_size)
|
||
.context("step_with_lobsim: read 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].
|
||
let isv_d_clone = self.isv_d.clone();
|
||
self.dqn_head
|
||
.soft_update_target(&isv_d_clone)
|
||
.context("dqn_head.soft_update_target")?;
|
||
|
||
// q_divergence EMA — `‖W_online − W_target‖₂` after the soft
|
||
// update. Single-block grid-stride reduce over the full
|
||
// weight tensor (24,192 floats), result feeds the
|
||
// rl_target_tau controller via ISV[418].
|
||
{
|
||
let n_total = self.dqn_head.w_d.len() as i32;
|
||
let block_dim: u32 = 256;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_dim, 1, 1),
|
||
shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.rl_l2_diff_norm_fn);
|
||
launch
|
||
.arg(&self.dqn_head.w_d)
|
||
.arg(&self.dqn_head.w_target_d)
|
||
.arg(&mut self.ema_input_scratch_d)
|
||
.arg(&n_total);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_l2_diff_norm(W, W_target) launch")?;
|
||
}
|
||
}
|
||
self.launch_ema_update_per_step(
|
||
crate::rl::isv_slots::RL_Q_DIVERGENCE_EMA_INDEX,
|
||
RL_LR_CONTROLLER_ALPHA,
|
||
&self.ema_input_scratch_d.clone(),
|
||
1,
|
||
)
|
||
.context("ema_update_per_step(q_divergence)")?;
|
||
|
||
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};
|
||
|
||
// Per-batch metadata via mapped-pinned staging (the only
|
||
// permitted CPU↔GPU path per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned`). Each helper
|
||
// allocates a mapped-pinned buffer, issues a DtoD copy from
|
||
// the device source, syncs, and returns a host Vec read from
|
||
// the page-mapped host_ptr.
|
||
let actions_host = read_slice_i32_d(&self.stream, &self.actions_d, b_size)
|
||
.context("push_to_replay: read actions_d")?;
|
||
let rewards_host = read_slice_d(&self.stream, &self.rewards_d, b_size)
|
||
.context("push_to_replay: read rewards_d")?;
|
||
let dones_host = read_slice_d(&self.stream, &self.dones_d, b_size)
|
||
.context("push_to_replay: read dones_d")?;
|
||
let log_pi_old_host = read_slice_d(&self.stream, &self.log_pi_old_d, b_size)
|
||
.context("push_to_replay: read 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 };
|
||
}
|
||
// Mapped-pinned staging — the only permitted CPU→GPU path
|
||
// per `feedback_no_htod_htoh_only_mapped_pinned`.
|
||
write_slice_i32_d(&self.stream, &actions_host, &mut self.sampled_actions_d)
|
||
.context("sample_and_gather: write sampled_actions_d")?;
|
||
write_slice_f32_d(&self.stream, &rewards_host, &mut self.sampled_rewards_d)
|
||
.context("sample_and_gather: write sampled_rewards_d")?;
|
||
write_slice_f32_d(&self.stream, &dones_host, &mut self.sampled_dones_d)
|
||
.context("sample_and_gather: write 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)
|
||
}
|
||
|
||
/// Reduce `var(x) / max(|mean(x)|, 1e-6)` over `input_d[b_size]`
|
||
/// into `self.ema_input_scratch_d[1]`. Single block,
|
||
/// next-pow2(b_size) threads, shared mem reused across the two
|
||
/// passes inside the kernel.
|
||
fn launch_var_over_abs_mean(
|
||
&mut self,
|
||
input_d: &CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert!(input_d.len() >= b_size);
|
||
// Single-thread streaming kernel; folds across STEPS (not batch).
|
||
// Writes var/|mean| directly into the controller-input ISV slot
|
||
// — caller does NOT need a downstream ema_update_per_step.
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let out_slot: i32 =
|
||
crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX as i32;
|
||
let mean_slot: i32 =
|
||
crate::rl::isv_slots::RL_ADV_VAR_STREAM_MEAN_INDEX as i32;
|
||
let m2_slot: i32 =
|
||
crate::rl::isv_slots::RL_ADV_VAR_STREAM_M2_INDEX as i32;
|
||
let clamp_slot: i32 =
|
||
crate::rl::isv_slots::RL_ADV_VAR_RATIO_CLAMP_INDEX as i32;
|
||
let mut launch = self.stream.launch_builder(
|
||
&self.rl_var_over_abs_mean_streaming_fn,
|
||
);
|
||
launch
|
||
.arg(input_d)
|
||
.arg(&mut self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&out_slot)
|
||
.arg(&mean_slot)
|
||
.arg(&m2_slot)
|
||
.arg(&clamp_slot);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_var_over_abs_mean_streaming launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Reduce `‖x‖₂ = sqrt(Σ x²)` over `input_d[n]` into
|
||
/// `self.ema_input_scratch_d[1]`. Single block, 256 threads,
|
||
/// grid-stride loop covers arbitrary n via per-thread partial
|
||
/// accumulation + shared-mem tree-reduce. Used by the LR
|
||
/// controller's grad-norm input path (per-head grad_w_d L2 norm
|
||
/// fed to ema_update_per_step → ISV[424..426]).
|
||
fn launch_l2_norm(
|
||
&mut self,
|
||
input_d: &CudaSlice<f32>,
|
||
n: usize,
|
||
) -> Result<()> {
|
||
debug_assert!(input_d.len() >= n);
|
||
if n == 0 {
|
||
return Ok(());
|
||
}
|
||
let block_dim: u32 = 256;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_dim, 1, 1),
|
||
shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
|
||
};
|
||
let n_i = n as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_l2_norm_fn);
|
||
launch
|
||
.arg(input_d)
|
||
.arg(&mut self.ema_input_scratch_d)
|
||
.arg(&n_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("rl_l2_norm launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Streaming kurtosis `M4/M2²` that folds ACROSS STEPS (Welford-EMA
|
||
/// on per-batch means) rather than across batch. Writes the
|
||
/// smoothed estimate directly into `RL_TD_KURTOSIS_EMA_INDEX` —
|
||
/// caller does NOT need a downstream ema_update_per_step. Replaces
|
||
/// the prior per-batch reduction which returned 0 at b_size=1
|
||
/// (kurtosis undefined) and left the rl_per_alpha controller blind.
|
||
fn launch_kurtosis(
|
||
&mut self,
|
||
input_d: &CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert!(input_d.len() >= b_size);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let out_slot: i32 =
|
||
crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX as i32;
|
||
let mean_slot: i32 =
|
||
crate::rl::isv_slots::RL_TD_KURT_STREAM_MEAN_INDEX as i32;
|
||
let m2_slot: i32 =
|
||
crate::rl::isv_slots::RL_TD_KURT_STREAM_M2_INDEX as i32;
|
||
let m4_slot: i32 =
|
||
crate::rl::isv_slots::RL_TD_KURT_STREAM_M4_INDEX as i32;
|
||
let clamp_slot: i32 =
|
||
crate::rl::isv_slots::RL_TD_KURTOSIS_CLAMP_INDEX as i32;
|
||
let mut launch = self.stream.launch_builder(
|
||
&self.rl_kurtosis_streaming_fn,
|
||
);
|
||
launch
|
||
.arg(input_d)
|
||
.arg(&mut self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&out_slot)
|
||
.arg(&mean_slot)
|
||
.arg(&m2_slot)
|
||
.arg(&m4_slot)
|
||
.arg(&clamp_slot);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_kurtosis_streaming launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Launch `rl_lr_controller` to emit per-head learning rates into
|
||
/// `ISV[412..417]` via ReduceLROnPlateau-style monotone decay.
|
||
/// Per-head observed loss (l_q / l_pi / l_v from step_synthetic's
|
||
/// stats this step) feeds the controller's slow loss EMA; the
|
||
/// controller halves LR when 1000 consecutive steps elapse
|
||
/// without a 1% improvement. BCE / AUX use placeholder zero
|
||
/// losses (the plateau function early-returns for those because
|
||
/// loss_ema_slot = -1).
|
||
///
|
||
/// Loss observations are passed as scalar floats (host-side
|
||
/// values from the trainer's last step). Per-head state lives
|
||
/// in ISV slots 427..439 (loss_ema, best, counter, warmup ×3).
|
||
fn launch_rl_lr_controller(
|
||
&self,
|
||
observed_loss_q: f32,
|
||
observed_loss_pi: f32,
|
||
observed_loss_v: f32,
|
||
) -> Result<()> {
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let observed_loss_bce: f32 = 0.0;
|
||
let observed_loss_aux: f32 = 0.0;
|
||
let q_loss_ema_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_Q_LOSS_EMA_INDEX as i32;
|
||
let q_best_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_Q_BEST_LOSS_INDEX as i32;
|
||
let q_counter_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_Q_STEPS_SINCE_BEST_INDEX as i32;
|
||
let q_warmup_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_Q_WARMUP_COUNTER_INDEX as i32;
|
||
let pi_loss_ema_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_PI_LOSS_EMA_INDEX as i32;
|
||
let pi_best_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_PI_BEST_LOSS_INDEX as i32;
|
||
let pi_counter_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_PI_STEPS_SINCE_BEST_INDEX as i32;
|
||
let pi_warmup_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_PI_WARMUP_COUNTER_INDEX as i32;
|
||
let v_loss_ema_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_V_LOSS_EMA_INDEX as i32;
|
||
let v_best_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_V_BEST_LOSS_INDEX as i32;
|
||
let v_counter_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_V_STEPS_SINCE_BEST_INDEX as i32;
|
||
let v_warmup_slot: i32 =
|
||
crate::rl::isv_slots::RL_LR_V_WARMUP_COUNTER_INDEX as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_lr_controller_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(&observed_loss_bce)
|
||
.arg(&observed_loss_q)
|
||
.arg(&observed_loss_pi)
|
||
.arg(&observed_loss_v)
|
||
.arg(&observed_loss_aux)
|
||
.arg(&q_loss_ema_slot)
|
||
.arg(&q_best_slot)
|
||
.arg(&q_counter_slot)
|
||
.arg(&q_warmup_slot)
|
||
.arg(&pi_loss_ema_slot)
|
||
.arg(&pi_best_slot)
|
||
.arg(&pi_counter_slot)
|
||
.arg(&pi_warmup_slot)
|
||
.arg(&v_loss_ema_slot)
|
||
.arg(&v_best_slot)
|
||
.arg(&v_counter_slot)
|
||
.arg(&v_warmup_slot);
|
||
unsafe {
|
||
launch.launch(cfg).context("rl_lr_controller launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Launch the reduce_axis0 kernel:
|
||
/// per_batch: [b_size × n_tail] → out: [n_tail]
|
||
fn launch_reduce_axis0(
|
||
&self,
|
||
per_batch: &CudaSlice<f32>,
|
||
b_size: usize,
|
||
n_tail: usize,
|
||
out: &mut CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(per_batch.len(), b_size * n_tail);
|
||
debug_assert_eq!(out.len(), n_tail);
|
||
|
||
let n_batch_i = b_size as i32;
|
||
let n_tail_i = n_tail as i32;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: ((n_tail as u32).div_ceil(32), 1, 1),
|
||
block_dim: (32, 8, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||
launch
|
||
.arg(per_batch)
|
||
.arg(&n_batch_i)
|
||
.arg(&n_tail_i)
|
||
.arg(out);
|
||
unsafe {
|
||
launch.launch(cfg).context("reduce_axis0 launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Phase E.2 helper: combine a per-head grad_h_t into the encoder's
|
||
/// accumulator slot via `grad_h_encoder[i] += λ × grad_h_head[i]`.
|
||
/// Public method delegates to the free function `accumulate_grad_h`
|
||
/// so the same launch path is reachable both from external callers
|
||
/// and from `step_synthetic` (where the borrow checker forbids a
|
||
/// `&self` + `&mut self.grad_h_t_combined_d` combination).
|
||
pub fn launch_grad_h_accumulate(
|
||
&self,
|
||
grad_h_head: &CudaSlice<f32>,
|
||
lambda: f32,
|
||
grad_h_encoder: &mut CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
accumulate_grad_h(
|
||
&self.stream,
|
||
&self.grad_h_accumulate_fn,
|
||
grad_h_head,
|
||
lambda,
|
||
grad_h_encoder,
|
||
)
|
||
}
|
||
|
||
/// Phase E.3b evaluation helper: forward the encoder on `snapshots`,
|
||
/// forward the C51 Q-head, and return the expected Q-value E[Q(s, a)]
|
||
/// per action, averaged across the batch. Lets the toy-bandit gate
|
||
/// tests inspect the trained Q-distribution's argmax without
|
||
/// re-implementing the device-to-host readback inside the test
|
||
/// crate.
|
||
///
|
||
/// Returns a `Vec<f32>` of length `N_ACTIONS = 9`. Index = action id
|
||
/// (see `crate::rl::common::Action`). Caller asserts on `argmax`.
|
||
pub fn eval_expected_q_per_action(
|
||
&mut self,
|
||
snapshots: &[Mbp10RawInput],
|
||
) -> Result<Vec<f32>> {
|
||
let b_size = self.cfg.perception.n_batch;
|
||
let _ = self
|
||
.perception
|
||
.forward_encoder(snapshots)
|
||
.context("eval_expected_q: forward_encoder")?;
|
||
let h_t = self.perception.h_t_view();
|
||
let k = N_ACTIONS * Q_N_ATOMS;
|
||
let mut logits_d = self.stream.alloc_zeros::<f32>(b_size * k)?;
|
||
self.dqn_head
|
||
.forward(h_t, b_size, &mut logits_d)
|
||
.context("eval_expected_q: dqn_head.forward")?;
|
||
|
||
let logits_host = read_slice_d(&self.stream, &logits_d, b_size * k)?;
|
||
let atom_step = (Q_V_MAX - Q_V_MIN) / ((Q_N_ATOMS - 1) as f32);
|
||
let atom_supports: Vec<f32> = (0..Q_N_ATOMS)
|
||
.map(|i| Q_V_MIN + (i as f32) * atom_step)
|
||
.collect();
|
||
let mut eq = vec![0.0_f32; N_ACTIONS];
|
||
for b in 0..b_size {
|
||
for a in 0..N_ACTIONS {
|
||
let off = b * k + a * Q_N_ATOMS;
|
||
let row = &logits_host[off..off + Q_N_ATOMS];
|
||
let max_l = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
let mut sum = 0.0_f32;
|
||
let mut probs = [0.0_f32; Q_N_ATOMS];
|
||
for (i, &l) in row.iter().enumerate() {
|
||
probs[i] = (l - max_l).exp();
|
||
sum += probs[i];
|
||
}
|
||
if sum > 0.0 {
|
||
for p in probs.iter_mut() {
|
||
*p /= sum;
|
||
}
|
||
}
|
||
let mut ev = 0.0_f32;
|
||
for i in 0..Q_N_ATOMS {
|
||
ev += probs[i] * atom_supports[i];
|
||
}
|
||
eq[a] += ev;
|
||
}
|
||
}
|
||
for e in eq.iter_mut() {
|
||
*e /= b_size as f32;
|
||
}
|
||
Ok(eq)
|
||
}
|
||
|
||
/// Phase E.3b evaluation helper: forward the encoder on `snapshots`,
|
||
/// forward the PPO policy head, return softmax(π) per action,
|
||
/// averaged across the batch. Used by the `ppo_toy` gate test to
|
||
/// inspect the trained policy's mode without re-implementing the
|
||
/// device-to-host readback in the test crate.
|
||
pub fn eval_policy_probs_per_action(
|
||
&mut self,
|
||
snapshots: &[Mbp10RawInput],
|
||
) -> Result<Vec<f32>> {
|
||
let b_size = self.cfg.perception.n_batch;
|
||
let _ = self
|
||
.perception
|
||
.forward_encoder(snapshots)
|
||
.context("eval_policy_probs: forward_encoder")?;
|
||
let h_t = self.perception.h_t_view();
|
||
let mut logits_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||
self.policy_head
|
||
.forward_logits(h_t, b_size, &mut logits_d)
|
||
.context("eval_policy_probs: policy_head.forward_logits")?;
|
||
|
||
let logits_host = read_slice_d(&self.stream, &logits_d, b_size * N_ACTIONS)?;
|
||
let mut probs = vec![0.0_f32; N_ACTIONS];
|
||
for b in 0..b_size {
|
||
let off = b * N_ACTIONS;
|
||
let row = &logits_host[off..off + N_ACTIONS];
|
||
let max_l = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
let mut sum = 0.0_f32;
|
||
let mut p = [0.0_f32; N_ACTIONS];
|
||
for (i, &l) in row.iter().enumerate() {
|
||
p[i] = (l - max_l).exp();
|
||
sum += p[i];
|
||
}
|
||
if sum > 0.0 {
|
||
for x in p.iter_mut() {
|
||
*x /= sum;
|
||
}
|
||
}
|
||
for i in 0..N_ACTIONS {
|
||
probs[i] += p[i];
|
||
}
|
||
}
|
||
for p in probs.iter_mut() {
|
||
*p /= b_size as f32;
|
||
}
|
||
Ok(probs)
|
||
}
|
||
}
|
||
|
||
/// Free-function entry point for the `grad_h_accumulate_scaled` kernel.
|
||
/// Lives outside the impl so callers can hold a `&mut` reference to a
|
||
/// trainer-owned grad slot at the same time as a `&self.stream` /
|
||
/// `&self.grad_h_accumulate_fn` borrow.
|
||
fn accumulate_grad_h(
|
||
stream: &Arc<CudaStream>,
|
||
grad_h_accumulate_fn: &CudaFunction,
|
||
grad_h_head: &CudaSlice<f32>,
|
||
lambda: f32,
|
||
grad_h_encoder: &mut CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
let n = grad_h_head.len();
|
||
debug_assert_eq!(grad_h_encoder.len(), n);
|
||
let n_i = n as i32;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: ((n as u32).div_ceil(256), 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = stream.launch_builder(grad_h_accumulate_fn);
|
||
launch
|
||
.arg(grad_h_head)
|
||
.arg(&lambda)
|
||
.arg(&n_i)
|
||
.arg(grad_h_encoder);
|
||
unsafe {
|
||
launch.launch(cfg).context("grad_h_accumulate launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
// ── helpers ──────────────────────────────────────────────────────────
|
||
//
|
||
// Phase R7b: `upload_f32`, `upload_i32`, and `argmax_f32` are deleted
|
||
// because all step_with_lobsim host loops they served are GPU-resident
|
||
// now (R3/R4/R5/R6/R7a kernels). Per `feedback_no_hiding` the
|
||
// retired helpers are removed, not `#[allow(dead_code)]`'d. The
|
||
// remaining `read_*_d` helpers are still load-bearing for the per-step
|
||
// loss-scalar readback inside step_synthetic's Q + V backward chain.
|
||
|
||
/// Read `n` `i32`s from a device buffer into a host `Vec<i32>` via
|
||
/// mapped-pinned staging + DtoD + stream sync. Same pattern as
|
||
/// `read_slice_d` (f32 version) — the only permitted CPU↔GPU path
|
||
/// per `feedback_no_htod_htoh_only_mapped_pinned`. Used by R7d's
|
||
/// `push_to_replay` to read the per-step actions device buffer for
|
||
/// PER metadata staging.
|
||
pub fn read_slice_i32_d_pub(
|
||
stream: &Arc<CudaStream>,
|
||
src: &CudaSlice<i32>,
|
||
n: usize,
|
||
) -> Result<Vec<i32>> {
|
||
read_slice_i32_d(stream, src, n)
|
||
}
|
||
|
||
pub fn read_slice_d_pub(
|
||
stream: &Arc<CudaStream>,
|
||
src: &CudaSlice<f32>,
|
||
n: usize,
|
||
) -> Result<Vec<f32>> {
|
||
read_slice_d(stream, src, n)
|
||
}
|
||
|
||
fn read_slice_i32_d(
|
||
stream: &Arc<CudaStream>,
|
||
src: &CudaSlice<i32>,
|
||
n: usize,
|
||
) -> Result<Vec<i32>> {
|
||
debug_assert!(src.len() >= n);
|
||
if n == 0 {
|
||
return Ok(Vec::new());
|
||
}
|
||
let mut staging = unsafe { MappedI32Buffer::new(n) }
|
||
.map_err(|e| anyhow::anyhow!("read_slice_i32_d staging: {e}"))?;
|
||
unsafe {
|
||
let (src_ptr, _g) = src.device_ptr(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
staging.dev_ptr,
|
||
src_ptr,
|
||
n * std::mem::size_of::<i32>(),
|
||
stream.cu_stream(),
|
||
)
|
||
.context("read_slice_i32_d DtoD")?;
|
||
}
|
||
stream.synchronize().context("read_slice_i32_d sync")?;
|
||
Ok(staging.host_slice_mut().to_vec())
|
||
}
|
||
|
||
/// Write a host `&[f32]` slice into a device buffer via mapped-pinned
|
||
/// staging + DtoD. The only permitted CPU→GPU path per
|
||
/// `feedback_no_htod_htoh_only_mapped_pinned` — `stream.memcpy_htod`
|
||
/// on a regular `&[T]` is forbidden (the source slice is not
|
||
/// page-locked, so the driver does an internal HtoD blocking copy).
|
||
/// Mapped-pinned staging keeps the host write zero-copy on the
|
||
/// device side.
|
||
fn write_slice_f32_d(
|
||
stream: &Arc<CudaStream>,
|
||
src: &[f32],
|
||
dst: &mut CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
debug_assert!(dst.len() >= src.len());
|
||
let n = src.len();
|
||
if n == 0 {
|
||
return Ok(());
|
||
}
|
||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||
.map_err(|e| anyhow::anyhow!("write_slice_f32_d staging: {e}"))?;
|
||
staging.write_from_slice(src);
|
||
unsafe {
|
||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr,
|
||
staging.dev_ptr,
|
||
n * std::mem::size_of::<f32>(),
|
||
stream.cu_stream(),
|
||
)
|
||
.context("write_slice_f32_d DtoD")?;
|
||
}
|
||
stream.synchronize().context("write_slice_f32_d sync")?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Write a host `&[i32]` slice into a device buffer via mapped-pinned
|
||
/// staging + DtoD. See `write_slice_f32_d` for rationale.
|
||
fn write_slice_i32_d(
|
||
stream: &Arc<CudaStream>,
|
||
src: &[i32],
|
||
dst: &mut CudaSlice<i32>,
|
||
) -> Result<()> {
|
||
debug_assert!(dst.len() >= src.len());
|
||
let n = src.len();
|
||
if n == 0 {
|
||
return Ok(());
|
||
}
|
||
let mut staging = unsafe { MappedI32Buffer::new(n) }
|
||
.map_err(|e| anyhow::anyhow!("write_slice_i32_d staging: {e}"))?;
|
||
staging.host_slice_mut()[..n].copy_from_slice(src);
|
||
unsafe {
|
||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr,
|
||
staging.dev_ptr,
|
||
n * std::mem::size_of::<i32>(),
|
||
stream.cu_stream(),
|
||
)
|
||
.context("write_slice_i32_d DtoD")?;
|
||
}
|
||
stream.synchronize().context("write_slice_i32_d sync")?;
|
||
Ok(())
|
||
}
|
||
|
||
fn read_scalar_d(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<f32> {
|
||
debug_assert!(src.len() >= 1);
|
||
let staging = unsafe { MappedF32Buffer::new(1) }
|
||
.map_err(|e| anyhow::anyhow!("read_scalar_d staging: {e}"))?;
|
||
unsafe {
|
||
let (src_ptr, _g) = src.device_ptr(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
staging.dev_ptr,
|
||
src_ptr,
|
||
std::mem::size_of::<f32>(),
|
||
stream.cu_stream(),
|
||
)
|
||
.context("read_scalar_d DtoD")?;
|
||
}
|
||
stream.synchronize().context("read_scalar_d sync")?;
|
||
Ok(unsafe { std::ptr::read_volatile(staging.host_ptr) })
|
||
}
|
||
|
||
/// Phase E.3b helper: read `n` floats from a device buffer into a host
|
||
/// `Vec<f32>` via mapped-pinned staging + DtoD copy + stream sync. Used
|
||
/// by `step_with_lobsim` to pull Q-logits / V / π-logits to host for
|
||
/// Thompson sampling. Per `feedback_no_htod_htoh_only_mapped_pinned`:
|
||
/// the staging buffer is mapped-pinned, the device → staging copy uses
|
||
/// `memcpy_dtod_async`, and the host read happens after a stream sync.
|
||
fn read_slice_d(
|
||
stream: &Arc<CudaStream>,
|
||
src: &CudaSlice<f32>,
|
||
n: usize,
|
||
) -> Result<Vec<f32>> {
|
||
debug_assert_eq!(src.len(), n);
|
||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||
.map_err(|e| anyhow::anyhow!("read_slice_d staging: {e}"))?;
|
||
if n > 0 {
|
||
unsafe {
|
||
let (src_ptr, _g) = src.device_ptr(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
staging.dev_ptr,
|
||
src_ptr,
|
||
n * std::mem::size_of::<f32>(),
|
||
stream.cu_stream(),
|
||
)
|
||
.context("read_slice_d DtoD")?;
|
||
}
|
||
}
|
||
stream.synchronize().context("read_slice_d sync")?;
|
||
Ok(staging.read_all())
|
||
}
|
||
|
||
// Phase R7b: `argmax_f32` deleted — R4's `argmax_expected_q` kernel
|
||
// is the canonical Bellman-target argmax now.
|
||
|