0b7895a4fb
Move q_logits_d, q_logits_tp1_d, v_pred_d, v_pred_tp1_d, pi_logits_d from per-step alloc_zeros to persistent trainer fields allocated at init. Stable device pointers are a prerequisite for CUDA Graph capture. Updated step_with_lobsim, step_synthetic, and dqn_replay_step to use self.field instead of local allocations. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
6177 lines
275 KiB
Rust
6177 lines
275 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::iqn::EMBED_DIM;
|
||
use crate::rl::noisy::{NoisyLinear, NoisyLinearConfig};
|
||
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"));
|
||
// Adaptive reward-clamp controller — audit fix 2026-05-24. Consumes the
|
||
// per-step positive-tail max written by apply_reward_scale into slot
|
||
// 478 and emits adaptive WIN/LOSS bounds into slots 452/453. See the
|
||
// kernel header for the design rationale (rmgm5 clamp-firing audit).
|
||
const RL_REWARD_CLAMP_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_reward_clamp_controller.cubin"));
|
||
// C51 atom-support updater — audit 2026-05-24 followup. Refreshes
|
||
// `atom_supports_d` from the ISV V_MIN/V_MAX slots that the reward
|
||
// clamp controller writes (with ratchet semantics). Launched right
|
||
// after the reward clamp controller so atom_supports_d reflects the
|
||
// just-updated V_MIN/V_MAX before any C51 forward / Bellman target
|
||
// projection in the next step. 21-thread one-block kernel.
|
||
const RL_ATOM_SUPPORT_UPDATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_atom_support_update.cubin"));
|
||
// Q→π distillation gradient — audit 2026-05-24 vj5f6 followup. ADDS
|
||
// λ × (π_new - π_target) to pi_grad_logits after the PPO surrogate
|
||
// backward, where π_target = softmax(E_Q[s,*] / τ). Couples Q's
|
||
// improved C51 calibration to π's action selection (was decoupled
|
||
// per Option B / `pearl_q_thompson_actor_makes_pi_dead_weight`).
|
||
// C51+IQN ensemble action-value kernel (audit 2026-05-25). Combines
|
||
// E_C51 and E_IQN into a single ensemble Q vector via ISV α blend.
|
||
// Feeds Q→π agreement diag and future ensemble-level selection.
|
||
const RL_ENSEMBLE_ACTION_VALUE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_ensemble_action_value.cubin"));
|
||
const RL_Q_PI_DISTILL_GRAD_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_q_pi_distill_grad.cubin"));
|
||
// λ_distill adaptive controller (rljzl followup 2026-05-24).
|
||
// Drives λ via Schulman bounded step on KL_EMA toward target.
|
||
const RL_Q_DISTILL_LAMBDA_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_q_distill_lambda_controller.cubin"));
|
||
// SP20 P1+P5 — audit-wiring catch on TrailTighten/Loosen + foundational
|
||
// per-unit trade state machine. Three kernels work together: state
|
||
// machine on position transitions (open/close/reverse), trail mutation
|
||
// from a7/a8 actions, trail breach check that overrides action to flat.
|
||
const RL_UNIT_STATE_UPDATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_unit_state_update.cubin"));
|
||
const RL_TRAIL_MUTATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_trail_mutate.cubin"));
|
||
const RL_POSITION_HEAT_CHECK_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_position_heat_check.cubin"));
|
||
const RL_CONFIDENCE_GATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_confidence_gate.cubin"));
|
||
const RL_FRD_GATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_gate.cubin"));
|
||
const RL_RECENT_OUTCOME_UPDATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_recent_outcome_update.cubin"));
|
||
const RL_GATE_THRESHOLD_CONTROLLER_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_gate_threshold_controller.cubin"));
|
||
// Fused kernel: 10 RL ISV controllers in one launch — saves 9 kernel
|
||
// launch overheads (~40-80μs/step). Replaces individual launches of
|
||
// gamma, tau, ppo_clip, entropy_coef, rollout_steps, per_alpha,
|
||
// reward_scale, ppo_ratio_clamp, gate_threshold, q_distill_lambda.
|
||
const RL_FUSED_CONTROLLERS_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_fused_controllers.cubin"));
|
||
const RL_REWARD_SHAPING_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_reward_shaping.cubin"));
|
||
const RL_ASYMMETRIC_TRAIL_DECAY_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_asymmetric_trail_decay.cubin"));
|
||
const RL_SESSION_RISK_CHECK_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_session_risk_check.cubin"));
|
||
const RL_MIN_HOLD_CHECK_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_min_hold_check.cubin"));
|
||
const RL_TRADE_CONTEXT_UPDATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_trade_context_update.cubin"));
|
||
const RL_MULTIRES_FEATURES_UPDATE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_multires_features_update.cubin"));
|
||
const RL_TRAIL_STOP_CHECK_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_trail_stop_check.cubin"));
|
||
const ACTIONS_TO_MARKET_TARGETS_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/actions_to_market_targets.cubin"));
|
||
// Device-resident step counter bump — ISV[548] += 1.0 each step.
|
||
// Prerequisite for CUDA Graph capture: removes scalar current_step
|
||
// from all downstream kernel argument lists.
|
||
const RL_INCREMENT_STEP_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_increment_step.cubin"));
|
||
const RL_SAMPLE_TAU_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_sample_tau.cubin"));
|
||
|
||
// Element-wise in-place add: dst[i] += src[i]. Used by noisy
|
||
// exploration to fold NoisyLinear output into ensemble_q_d before
|
||
// action selection. Loaded from the same cubin that perception.rs
|
||
// uses for aux→encoder gradient accumulation.
|
||
const AUX_VEC_ADD_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/aux_vec_add.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"));
|
||
// audit — 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"));
|
||
const RL_ISV_WRITE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_isv_write.cubin"));
|
||
const RL_Q_PI_AGREE_B_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_q_pi_agree_b.cubin"));
|
||
const RL_PI_ACTION_KERNEL_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_pi_action_kernel.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. At b_size=16, capacity/16 = unique steps of
|
||
/// history retained. 32768/16 = 2048 steps (~34 minutes of ES data
|
||
/// at event rate). Naive O(N) sampling is fast enough at this scale.
|
||
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: 32768,
|
||
per_seed: 0x9E37_79B9_7F4A_7C15,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Stats returned from `step_synthetic`. Useful for tests + diagnostics.
|
||
#[derive(Clone, Debug, Default)]
|
||
pub struct IntegratedStepStats {
|
||
pub l_bce: f32,
|
||
pub l_q: f32,
|
||
pub l_pi: f32,
|
||
pub l_v: f32,
|
||
pub l_aux: f32,
|
||
/// SP20 P3 FRD head loss — sum of CE across (batch, horizon),
|
||
/// normalized by `B × FRD_N_HORIZONS`. Zero whenever all labels
|
||
/// are sentinel (-1); becomes the head's training signal once
|
||
/// the loader feeds real forward-return labels.
|
||
pub l_frd: f32,
|
||
pub l_total: f32,
|
||
pub lambdas: LossLambdas,
|
||
}
|
||
|
||
/// Entry in the n-step return ring buffer. Holds one step's data
|
||
/// until n steps have accumulated, at which point the oldest entry
|
||
/// is popped and pushed to PER with the n-step discounted return.
|
||
struct NStepEntry {
|
||
h_t: CudaSlice<f32>,
|
||
action: u32,
|
||
raw_reward: f32,
|
||
scaled_reward: f32,
|
||
done: bool,
|
||
log_pi_old: f32,
|
||
}
|
||
|
||
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,
|
||
|
||
/// IQN distributional Q-head — runs alongside C51 for ensemble
|
||
/// action selection. Forward produces `Q(s, τ, a)` for N_TAU
|
||
/// quantile samples; expected Q is combined with C51's E[Q] via
|
||
/// ISV[RL_IQN_ENSEMBLE_ALPHA_INDEX].
|
||
pub iqn_head: crate::rl::iqn::IqnHead,
|
||
|
||
/// 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,
|
||
/// IQN head Adam optimisers — w_embed/b_embed/w_out/b_out.
|
||
pub iqn_w_embed_adam: AdamW,
|
||
pub iqn_b_embed_adam: AdamW,
|
||
pub iqn_w_out_adam: AdamW,
|
||
pub iqn_b_out_adam: AdamW,
|
||
pub policy_w_adam: AdamW,
|
||
pub policy_b_adam: AdamW,
|
||
pub value_w_adam: AdamW,
|
||
pub value_b_adam: AdamW,
|
||
/// SP20 P3 FRD head Adam optimisers (W1/b1/W2/b2). LR from
|
||
/// `RL_FRD_LR_INDEX` (slot 499), default 1e-3.
|
||
pub frd_w1_adam: AdamW,
|
||
pub frd_b1_adam: AdamW,
|
||
pub frd_w2_adam: AdamW,
|
||
pub frd_b2_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,
|
||
// Adaptive reward-clamp controller (audit 2026-05-24).
|
||
_rl_reward_clamp_controller_module: Arc<CudaModule>,
|
||
rl_reward_clamp_controller_fn: CudaFunction,
|
||
// C51 atom-support updater (audit 2026-05-24 followup).
|
||
_rl_atom_support_update_module: Arc<CudaModule>,
|
||
rl_atom_support_update_fn: CudaFunction,
|
||
// Q→π distillation gradient (audit 2026-05-24 vj5f6 followup).
|
||
_rl_q_pi_distill_grad_module: Arc<CudaModule>,
|
||
rl_q_pi_distill_grad_fn: CudaFunction,
|
||
// C51+IQN ensemble action-value (audit 2026-05-25).
|
||
_rl_ensemble_action_value_module: Arc<CudaModule>,
|
||
rl_ensemble_action_value_fn: CudaFunction,
|
||
// λ_distill adaptive controller (rljzl followup 2026-05-24).
|
||
// Retained for bootstrap / testing; per-step launch fused into
|
||
// rl_fused_controllers.
|
||
_rl_q_distill_lambda_controller_module: Arc<CudaModule>,
|
||
#[allow(dead_code)]
|
||
rl_q_distill_lambda_controller_fn: CudaFunction,
|
||
// SP20 P1+P5 (audit-wiring fix for dead a7/a8).
|
||
_rl_unit_state_update_module: Arc<CudaModule>,
|
||
rl_unit_state_update_fn: CudaFunction,
|
||
_rl_trail_mutate_module: Arc<CudaModule>,
|
||
rl_trail_mutate_fn: CudaFunction,
|
||
_rl_trail_stop_check_module: Arc<CudaModule>,
|
||
rl_trail_stop_check_fn: CudaFunction,
|
||
_rl_position_heat_check_module: Arc<CudaModule>,
|
||
rl_position_heat_check_fn: CudaFunction,
|
||
_rl_confidence_gate_module: Arc<CudaModule>,
|
||
rl_confidence_gate_fn: CudaFunction,
|
||
_rl_frd_gate_module: Arc<CudaModule>,
|
||
rl_frd_gate_fn: CudaFunction,
|
||
_rl_recent_outcome_update_module: Arc<CudaModule>,
|
||
rl_recent_outcome_update_fn: CudaFunction,
|
||
_rl_gate_threshold_controller_module: Arc<CudaModule>,
|
||
// Retained for testing; per-step launch fused into
|
||
// rl_fused_controllers.
|
||
#[allow(dead_code)]
|
||
rl_gate_threshold_controller_fn: CudaFunction,
|
||
// Fused kernel: 10 RL ISV controllers in one launch.
|
||
_rl_fused_controllers_module: Arc<CudaModule>,
|
||
rl_fused_controllers_fn: CudaFunction,
|
||
/// Device buffer holding the 7 ISV input-slot indices for the fused
|
||
/// controller kernel. Allocated once at init, never mutated.
|
||
fused_ctrl_input_slots_d: CudaSlice<i32>,
|
||
_rl_reward_shaping_module: Arc<CudaModule>,
|
||
rl_reward_shaping_fn: CudaFunction,
|
||
_rl_asymmetric_trail_decay_module: Arc<CudaModule>,
|
||
rl_asymmetric_trail_decay_fn: CudaFunction,
|
||
_rl_session_risk_check_module: Arc<CudaModule>,
|
||
rl_session_risk_check_fn: CudaFunction,
|
||
_rl_min_hold_check_module: Arc<CudaModule>,
|
||
rl_min_hold_check_fn: CudaFunction,
|
||
_rl_trade_context_update_module: Arc<CudaModule>,
|
||
rl_trade_context_update_fn: CudaFunction,
|
||
_rl_multires_features_update_module: Arc<CudaModule>,
|
||
rl_multires_features_update_fn: CudaFunction,
|
||
// Device-resident step counter (ISV[548] += 1.0 per step).
|
||
_rl_increment_step_module: Arc<CudaModule>,
|
||
rl_increment_step_fn: CudaFunction,
|
||
// Device-side tau ~ U(0,1) for IQN (graph-safe, no host RNG).
|
||
_rl_sample_tau_module: Arc<CudaModule>,
|
||
rl_sample_tau_fn: CudaFunction,
|
||
pub outcome_ema_d: CudaSlice<f32>,
|
||
pub trade_context_d: CudaSlice<f32>,
|
||
pub multires_output_d: CudaSlice<f32>,
|
||
multires_state_d: CudaSlice<f32>,
|
||
multires_prev_mid_d: CudaSlice<f32>,
|
||
multires_prev_ts_ns_d: CudaSlice<u64>,
|
||
|
||
// Per-batch per-unit trade state (SP20 P1).
|
||
// MAX_UNITS=4 reserved; only slot 0 used until SP20 P7 ships
|
||
// pyramiding semantics. Diag exposes all 4 slots — slots 1-3
|
||
// hold zeros pre-P7 (audit-diag tolerates per-batch-buffer
|
||
// arrays of any shape so long as values are non-NaN).
|
||
pub unit_entry_price_d: CudaSlice<f32>, // [B * 4]
|
||
pub unit_entry_step_d: CudaSlice<i32>, // [B * 4]
|
||
pub unit_lots_d: CudaSlice<i32>, // [B * 4]
|
||
pub unit_initial_r_d: CudaSlice<f32>, // [B * 4]
|
||
pub unit_trail_distance_d: CudaSlice<f32>, // [B * 4]
|
||
pub unit_active_d: CudaSlice<u8>, // [B * 4]
|
||
pub pyramid_units_count_d: CudaSlice<i32>, // [B]
|
||
pub close_unit_index_d: CudaSlice<i32>, // [B] — trail-stop override
|
||
// (-1 = auto oldest). Reset to
|
||
// -1 each step before trail check.
|
||
pub unit_prev_pos_lots_d: CudaSlice<i32>, // [B] — separate from
|
||
// extract_realized_pnl_delta's
|
||
// prev_position_lots_d so unit
|
||
// state machine isn't coupled
|
||
// to reward extraction order.
|
||
_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,
|
||
// 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,
|
||
_rl_isv_write_module: Arc<CudaModule>,
|
||
rl_isv_write_fn: CudaFunction,
|
||
_rl_q_pi_agree_b_module: Arc<CudaModule>,
|
||
rl_q_pi_agree_b_fn: CudaFunction,
|
||
_rl_pi_action_kernel_module: Arc<CudaModule>,
|
||
rl_pi_action_kernel_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>,
|
||
raw_rewards_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>,
|
||
|
||
// ── IQN per-step device buffers ──────────────────────────────────
|
||
/// Per-step quantile fractions `[B × N_TAU]` from U(0,1).
|
||
/// Populated each step by `rl_sample_tau` kernel on device.
|
||
pub iqn_tau_d: CudaSlice<f32>,
|
||
/// IQN forward output `[B × N_TAU × N_ACTIONS]`.
|
||
pub iqn_q_values_d: CudaSlice<f32>,
|
||
/// IQN expected Q `[B × N_ACTIONS]` = mean over tau dimension.
|
||
pub iqn_expected_q_d: CudaSlice<f32>,
|
||
/// Ensemble Q `[B × N_ACTIONS]` = α × E_C51 + (1-α) × E_IQN.
|
||
/// Output of `rl_ensemble_action_value` kernel. Feeds the Q→π
|
||
/// agreement diagnostic (`rl_q_pi_agree_b`).
|
||
pub ensemble_q_d: CudaSlice<f32>,
|
||
|
||
// ── Persistent per-step head output buffers (CUDA Graph stable) ──
|
||
// Pre-allocated at init so device pointers are stable across steps,
|
||
// enabling CUDA Graph capture of the RL step pipeline.
|
||
|
||
/// C51 Q logits at h_t `[B × N_ACTIONS × Q_N_ATOMS]`.
|
||
pub q_logits_d: CudaSlice<f32>,
|
||
/// C51 Q logits at h_{t+1} `[B × N_ACTIONS × Q_N_ATOMS]`.
|
||
pub q_logits_tp1_d: CudaSlice<f32>,
|
||
/// Scalar value prediction at h_t `[B]`.
|
||
pub v_pred_d: CudaSlice<f32>,
|
||
/// Scalar value prediction at h_{t+1} `[B]`.
|
||
pub v_pred_tp1_d: CudaSlice<f32>,
|
||
/// Policy logits at h_t `[B × N_ACTIONS]`.
|
||
pub pi_logits_d: CudaSlice<f32>,
|
||
|
||
/// Device-resident xorshift32 PRNG `[B]` for tau sampling.
|
||
iqn_prng_state_d: CudaSlice<u32>,
|
||
|
||
// ── NoisyNet state-dependent exploration ─────────────────────────
|
||
// Factored noisy linear layer (Fortunato et al. 2017) that maps
|
||
// h_t → [B, N_ACTIONS] state-dependent exploration noise. Added
|
||
// to ensemble_q_d before action selection so the policy sees
|
||
// noisy Q values during rollout. Replaces the P_MIN probability
|
||
// floor in rl_pi_action_kernel. Target network stays deterministic
|
||
// (noise buffers zeroed = no resample_noise calls on target path).
|
||
|
||
/// NoisyLinear layer: h_t [B, HIDDEN_DIM] → noise [B, N_ACTIONS].
|
||
/// Resampled once per step_with_lobsim via resample_noise().
|
||
pub noisy_exploration: NoisyLinear,
|
||
/// Output buffer for the NoisyLinear forward [B × N_ACTIONS].
|
||
pub noisy_output_d: CudaSlice<f32>,
|
||
/// Adam optimiser for NoisyLinear mu_w [N_ACTIONS × HIDDEN_DIM].
|
||
pub noisy_mu_w_adam: AdamW,
|
||
/// Adam optimiser for NoisyLinear sigma_w [N_ACTIONS × HIDDEN_DIM].
|
||
pub noisy_sigma_w_adam: AdamW,
|
||
/// Adam optimiser for NoisyLinear mu_b [N_ACTIONS].
|
||
pub noisy_mu_b_adam: AdamW,
|
||
/// Adam optimiser for NoisyLinear sigma_b [N_ACTIONS].
|
||
pub noisy_sigma_b_adam: AdamW,
|
||
/// Kernel handle for `aux_vec_add_inplace` — folds noisy_output_d
|
||
/// into ensemble_q_d before the π action sampler reads it.
|
||
_aux_vec_add_module: Arc<CudaModule>,
|
||
aux_vec_add_fn: CudaFunction,
|
||
|
||
// ── 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,
|
||
n_step_buffer: Vec<Vec<NStepEntry>>,
|
||
|
||
/// 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>,
|
||
pub sampled_n_step_gammas_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,
|
||
|
||
/// audit — number of training iterations performed in the most-recent
|
||
/// `step_with_lobsim` call. Derived from `isv[RL_N_ROLLOUT_STEPS_INDEX]`
|
||
/// via `K = clamp(isv[404] / 1024, 1, 8)`. Surfaces in the diag
|
||
/// JSONL via `pub` accessor so post-hoc analysis can see how the
|
||
/// `n_rollout_steps` controller is driving training intensity per
|
||
/// env step. Zero at construction (first step overwrites).
|
||
pub last_k_updates: usize,
|
||
|
||
/// 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,
|
||
|
||
/// SP20 P3 Forward-Return-Distribution head (forecast over 3
|
||
/// horizons × 21 return-bucket atoms). Forward kernel runs every
|
||
/// `step_with_lobsim`; backward + label-supervised loss arrive in
|
||
/// F.3 / loader label generation.
|
||
pub frd_head: crate::rl::frd::FrdHead,
|
||
|
||
/// SP20 P3 — cached post-ReLU hidden activation buffer for the FRD
|
||
/// head's backward pass (`B × FRD_HIDDEN_DIM`). Owned per-trainer
|
||
/// so we don't allocate per step.
|
||
pub frd_hidden_d: CudaSlice<f32>,
|
||
|
||
/// SP20 P3 — FRD output logits (`B × FRD_OUT_DIM` where
|
||
/// FRD_OUT_DIM = FRD_N_HORIZONS × FRD_N_ATOMS = 63). Public so
|
||
/// callers can read for diag/inference; backward kernel writes
|
||
/// gradients into this buffer in F.3.
|
||
pub frd_logits_d: CudaSlice<f32>,
|
||
|
||
/// SP20 P3 — FRD per-(batch, horizon) labels for the supervised
|
||
/// CE loss. Atom index `[0, FRD_N_ATOMS)` selects the correct
|
||
/// return bucket; sentinel `-1` marks "missing horizon" (the
|
||
/// forward return at h-ticks-ahead isn't realized yet at the
|
||
/// rightmost edge of the snapshot stream). All entries initialise
|
||
/// to -1; F.5 will populate before each `step_with_lobsim` from
|
||
/// the loader's forward-snapshot lookahead.
|
||
pub frd_labels_d: CudaSlice<i32>,
|
||
}
|
||
|
||
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 iqn_head = crate::rl::iqn::IqnHead::new(
|
||
dev,
|
||
crate::rl::iqn::IqnHeadConfig {
|
||
hidden_dim: HIDDEN_DIM,
|
||
seed: cfg.dqn_seed.wrapping_add(100),
|
||
n_tau: 32,
|
||
},
|
||
)
|
||
.context("IqnHead::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 iqn_w_embed_adam =
|
||
AdamW::new(dev, iqn_head.w_embed_d.len(), lr_placeholder).context("iqn_w_embed_adam")?;
|
||
let iqn_b_embed_adam =
|
||
AdamW::new(dev, iqn_head.b_embed_d.len(), lr_placeholder).context("iqn_b_embed_adam")?;
|
||
let iqn_w_out_adam =
|
||
AdamW::new(dev, iqn_head.w_out_d.len(), lr_placeholder).context("iqn_w_out_adam")?;
|
||
let iqn_b_out_adam =
|
||
AdamW::new(dev, iqn_head.b_out_d.len(), lr_placeholder).context("iqn_b_out_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];
|
||
// 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 rl_reward_clamp_controller_module = ctx
|
||
.load_cubin(RL_REWARD_CLAMP_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_reward_clamp_controller cubin")?;
|
||
let rl_reward_clamp_controller_fn = rl_reward_clamp_controller_module
|
||
.load_function("rl_reward_clamp_controller")
|
||
.context("load rl_reward_clamp_controller")?;
|
||
let rl_atom_support_update_module = ctx
|
||
.load_cubin(RL_ATOM_SUPPORT_UPDATE_CUBIN.to_vec())
|
||
.context("load rl_atom_support_update cubin")?;
|
||
let rl_atom_support_update_fn = rl_atom_support_update_module
|
||
.load_function("rl_atom_support_update")
|
||
.context("load rl_atom_support_update")?;
|
||
let rl_q_pi_distill_grad_module = ctx
|
||
.load_cubin(RL_Q_PI_DISTILL_GRAD_CUBIN.to_vec())
|
||
.context("load rl_q_pi_distill_grad cubin")?;
|
||
let rl_q_pi_distill_grad_fn = rl_q_pi_distill_grad_module
|
||
.load_function("rl_q_pi_distill_grad")
|
||
.context("load rl_q_pi_distill_grad")?;
|
||
let rl_ensemble_action_value_module = ctx
|
||
.load_cubin(RL_ENSEMBLE_ACTION_VALUE_CUBIN.to_vec())
|
||
.context("load rl_ensemble_action_value cubin")?;
|
||
let rl_ensemble_action_value_fn = rl_ensemble_action_value_module
|
||
.load_function("rl_ensemble_action_value")
|
||
.context("load rl_ensemble_action_value")?;
|
||
let rl_q_distill_lambda_controller_module = ctx
|
||
.load_cubin(RL_Q_DISTILL_LAMBDA_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_q_distill_lambda_controller cubin")?;
|
||
let rl_q_distill_lambda_controller_fn = rl_q_distill_lambda_controller_module
|
||
.load_function("rl_q_distill_lambda_controller")
|
||
.context("load rl_q_distill_lambda_controller")?;
|
||
|
||
// SP20 P1+P5 cubin loads.
|
||
let rl_unit_state_update_module = ctx
|
||
.load_cubin(RL_UNIT_STATE_UPDATE_CUBIN.to_vec())
|
||
.context("load rl_unit_state_update cubin")?;
|
||
let rl_unit_state_update_fn = rl_unit_state_update_module
|
||
.load_function("rl_unit_state_update")
|
||
.context("load rl_unit_state_update")?;
|
||
let rl_trail_mutate_module = ctx
|
||
.load_cubin(RL_TRAIL_MUTATE_CUBIN.to_vec())
|
||
.context("load rl_trail_mutate cubin")?;
|
||
let rl_trail_mutate_fn = rl_trail_mutate_module
|
||
.load_function("rl_trail_mutate")
|
||
.context("load rl_trail_mutate")?;
|
||
let rl_trail_stop_check_module = ctx
|
||
.load_cubin(RL_TRAIL_STOP_CHECK_CUBIN.to_vec())
|
||
.context("load rl_trail_stop_check cubin")?;
|
||
let rl_trail_stop_check_fn = rl_trail_stop_check_module
|
||
.load_function("rl_trail_stop_check")
|
||
.context("load rl_trail_stop_check")?;
|
||
let rl_position_heat_check_module = ctx
|
||
.load_cubin(RL_POSITION_HEAT_CHECK_CUBIN.to_vec())
|
||
.context("load rl_position_heat_check cubin")?;
|
||
let rl_position_heat_check_fn = rl_position_heat_check_module
|
||
.load_function("rl_position_heat_check")
|
||
.context("load rl_position_heat_check")?;
|
||
let rl_confidence_gate_module = ctx
|
||
.load_cubin(RL_CONFIDENCE_GATE_CUBIN.to_vec())
|
||
.context("load rl_confidence_gate cubin")?;
|
||
let rl_confidence_gate_fn = rl_confidence_gate_module
|
||
.load_function("rl_confidence_gate")
|
||
.context("load rl_confidence_gate")?;
|
||
let rl_frd_gate_module = ctx
|
||
.load_cubin(RL_FRD_GATE_CUBIN.to_vec())
|
||
.context("load rl_frd_gate cubin")?;
|
||
let rl_frd_gate_fn = rl_frd_gate_module
|
||
.load_function("rl_frd_gate")
|
||
.context("load rl_frd_gate")?;
|
||
let rl_recent_outcome_update_module = ctx
|
||
.load_cubin(RL_RECENT_OUTCOME_UPDATE_CUBIN.to_vec())
|
||
.context("load rl_recent_outcome_update cubin")?;
|
||
let rl_recent_outcome_update_fn = rl_recent_outcome_update_module
|
||
.load_function("rl_recent_outcome_update")
|
||
.context("load rl_recent_outcome_update")?;
|
||
let rl_gate_threshold_controller_module = ctx
|
||
.load_cubin(RL_GATE_THRESHOLD_CONTROLLER_CUBIN.to_vec())
|
||
.context("load rl_gate_threshold_controller cubin")?;
|
||
let rl_gate_threshold_controller_fn = rl_gate_threshold_controller_module
|
||
.load_function("rl_gate_threshold_controller")
|
||
.context("load rl_gate_threshold_controller")?;
|
||
let rl_fused_controllers_module = ctx
|
||
.load_cubin(RL_FUSED_CONTROLLERS_CUBIN.to_vec())
|
||
.context("load rl_fused_controllers cubin")?;
|
||
let rl_fused_controllers_fn = rl_fused_controllers_module
|
||
.load_function("rl_fused_controllers")
|
||
.context("load rl_fused_controllers")?;
|
||
// Device buffer for the 7 ISV input-slot indices consumed by
|
||
// the fused controller kernel. Allocated once, never mutated.
|
||
let fused_ctrl_input_slots_h: [i32; 7] = [
|
||
crate::rl::isv_slots::RL_MEAN_TRADE_DURATION_EMA_INDEX as i32,
|
||
crate::rl::isv_slots::RL_Q_DIVERGENCE_EMA_INDEX as i32,
|
||
crate::rl::isv_slots::RL_KL_PI_EMA_INDEX as i32,
|
||
crate::rl::isv_slots::RL_ENTROPY_OBSERVED_EMA_INDEX as i32,
|
||
crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX as i32,
|
||
crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX as i32,
|
||
crate::rl::isv_slots::RL_MEAN_ABS_PNL_EMA_INDEX as i32,
|
||
];
|
||
let mut fused_ctrl_input_slots_d = stream
|
||
.alloc_zeros::<i32>(7)
|
||
.context("alloc fused_ctrl_input_slots_d")?;
|
||
write_slice_i32_d_pub(&stream, &fused_ctrl_input_slots_h, &mut fused_ctrl_input_slots_d)?;
|
||
let rl_reward_shaping_module = ctx
|
||
.load_cubin(RL_REWARD_SHAPING_CUBIN.to_vec())
|
||
.context("load rl_reward_shaping cubin")?;
|
||
let rl_reward_shaping_fn = rl_reward_shaping_module
|
||
.load_function("rl_reward_shaping")
|
||
.context("load rl_reward_shaping")?;
|
||
let rl_asymmetric_trail_decay_module = ctx
|
||
.load_cubin(RL_ASYMMETRIC_TRAIL_DECAY_CUBIN.to_vec())
|
||
.context("load rl_asymmetric_trail_decay cubin")?;
|
||
let rl_asymmetric_trail_decay_fn = rl_asymmetric_trail_decay_module
|
||
.load_function("rl_asymmetric_trail_decay")
|
||
.context("load rl_asymmetric_trail_decay")?;
|
||
let rl_session_risk_check_module = ctx
|
||
.load_cubin(RL_SESSION_RISK_CHECK_CUBIN.to_vec())
|
||
.context("load rl_session_risk_check cubin")?;
|
||
let rl_session_risk_check_fn = rl_session_risk_check_module
|
||
.load_function("rl_session_risk_check")
|
||
.context("load rl_session_risk_check")?;
|
||
let rl_min_hold_check_module = ctx
|
||
.load_cubin(RL_MIN_HOLD_CHECK_CUBIN.to_vec())
|
||
.context("load rl_min_hold_check cubin")?;
|
||
let rl_min_hold_check_fn = rl_min_hold_check_module
|
||
.load_function("rl_min_hold_check")
|
||
.context("load rl_min_hold_check")?;
|
||
let rl_trade_context_update_module = ctx
|
||
.load_cubin(RL_TRADE_CONTEXT_UPDATE_CUBIN.to_vec())
|
||
.context("load rl_trade_context_update cubin")?;
|
||
let rl_trade_context_update_fn = rl_trade_context_update_module
|
||
.load_function("rl_trade_context_update")
|
||
.context("load rl_trade_context_update")?;
|
||
let rl_multires_features_update_module = ctx
|
||
.load_cubin(RL_MULTIRES_FEATURES_UPDATE_CUBIN.to_vec())
|
||
.context("load rl_multires_features_update cubin")?;
|
||
let rl_multires_features_update_fn = rl_multires_features_update_module
|
||
.load_function("rl_multires_features_update")
|
||
.context("load rl_multires_features_update")?;
|
||
let rl_increment_step_module = ctx
|
||
.load_cubin(RL_INCREMENT_STEP_CUBIN.to_vec())
|
||
.context("load rl_increment_step cubin")?;
|
||
let rl_increment_step_fn = rl_increment_step_module
|
||
.load_function("rl_increment_step")
|
||
.context("load rl_increment_step")?;
|
||
let rl_sample_tau_module = ctx
|
||
.load_cubin(RL_SAMPLE_TAU_CUBIN.to_vec())
|
||
.context("load rl_sample_tau cubin")?;
|
||
let rl_sample_tau_fn = rl_sample_tau_module
|
||
.load_function("rl_sample_tau")
|
||
.context("load rl_sample_tau")?;
|
||
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 the streaming refactor: 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")?;
|
||
|
||
// audit — 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")?;
|
||
let rl_isv_write_module = ctx
|
||
.load_cubin(RL_ISV_WRITE_CUBIN.to_vec())
|
||
.context("load rl_isv_write cubin")?;
|
||
let rl_isv_write_fn = rl_isv_write_module
|
||
.load_function("rl_isv_write")
|
||
.context("load rl_isv_write")?;
|
||
let rl_q_pi_agree_b_module = ctx
|
||
.load_cubin(RL_Q_PI_AGREE_B_CUBIN.to_vec())
|
||
.context("load rl_q_pi_agree_b cubin")?;
|
||
let rl_q_pi_agree_b_fn = rl_q_pi_agree_b_module
|
||
.load_function("rl_q_pi_agree_b")
|
||
.context("load rl_q_pi_agree_b")?;
|
||
let rl_pi_action_kernel_module = ctx
|
||
.load_cubin(RL_PI_ACTION_KERNEL_CUBIN.to_vec())
|
||
.context("load rl_pi_action_kernel cubin")?;
|
||
let rl_pi_action_kernel_fn = rl_pi_action_kernel_module
|
||
.load_function("rl_pi_action_kernel")
|
||
.context("load rl_pi_action_kernel")?;
|
||
|
||
// 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")?;
|
||
// SP20 P1+P5 per-unit trade state (slot 0 used; slots 1-3
|
||
// reserved for P7 pyramid).
|
||
let unit_entry_price_d = stream
|
||
.alloc_zeros::<f32>(b_size * 4)
|
||
.context("alloc unit_entry_price_d")?;
|
||
let unit_entry_step_d = stream
|
||
.alloc_zeros::<i32>(b_size * 4)
|
||
.context("alloc unit_entry_step_d")?;
|
||
let unit_lots_d = stream
|
||
.alloc_zeros::<i32>(b_size * 4)
|
||
.context("alloc unit_lots_d")?;
|
||
let unit_initial_r_d = stream
|
||
.alloc_zeros::<f32>(b_size * 4)
|
||
.context("alloc unit_initial_r_d")?;
|
||
let unit_trail_distance_d = stream
|
||
.alloc_zeros::<f32>(b_size * 4)
|
||
.context("alloc unit_trail_distance_d")?;
|
||
let unit_active_d = stream
|
||
.alloc_zeros::<u8>(b_size * 4)
|
||
.context("alloc unit_active_d")?;
|
||
let pyramid_units_count_d = stream
|
||
.alloc_zeros::<i32>(b_size)
|
||
.context("alloc pyramid_units_count_d")?;
|
||
let close_unit_index_d = stream
|
||
.alloc_zeros::<i32>(b_size)
|
||
.context("alloc close_unit_index_d")?;
|
||
let outcome_ema_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc outcome_ema_d")?;
|
||
let trade_context_d = stream
|
||
.alloc_zeros::<f32>(b_size * 4)
|
||
.context("alloc trade_context_d")?;
|
||
let multires_output_d = stream
|
||
.alloc_zeros::<f32>(b_size * 12)
|
||
.context("alloc multires_output_d")?;
|
||
let multires_state_d = stream
|
||
.alloc_zeros::<f32>(b_size * 12)
|
||
.context("alloc multires_state_d")?;
|
||
let multires_prev_mid_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc multires_prev_mid_d")?;
|
||
let multires_prev_ts_ns_d = stream
|
||
.alloc_zeros::<u64>(b_size)
|
||
.context("alloc multires_prev_ts_ns_d")?;
|
||
let unit_prev_pos_lots_d = stream
|
||
.alloc_zeros::<i32>(b_size)
|
||
.context("alloc unit_prev_pos_lots_d")?;
|
||
let rewards_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc rewards_d")?;
|
||
let raw_rewards_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc raw_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")?;
|
||
|
||
// IQN per-step buffers: tau samples, Q(s,τ,a), E[Q(s,a)].
|
||
let iqn_n_tau = iqn_head.n_tau();
|
||
let iqn_tau_d = stream
|
||
.alloc_zeros::<f32>(b_size * iqn_n_tau)
|
||
.context("alloc iqn_tau_d")?;
|
||
let iqn_q_values_d = stream
|
||
.alloc_zeros::<f32>(b_size * iqn_n_tau * N_ACTIONS)
|
||
.context("alloc iqn_q_values_d")?;
|
||
let iqn_expected_q_d = stream
|
||
.alloc_zeros::<f32>(b_size * N_ACTIONS)
|
||
.context("alloc iqn_expected_q_d")?;
|
||
let ensemble_q_d = stream
|
||
.alloc_zeros::<f32>(b_size * N_ACTIONS)
|
||
.context("alloc ensemble_q_d")?;
|
||
|
||
// Persistent per-step head output buffers (CUDA Graph stable).
|
||
let k_dqn_alloc = N_ACTIONS * Q_N_ATOMS;
|
||
let q_logits_d = stream
|
||
.alloc_zeros::<f32>(b_size * k_dqn_alloc)
|
||
.context("alloc q_logits_d")?;
|
||
let q_logits_tp1_d = stream
|
||
.alloc_zeros::<f32>(b_size * k_dqn_alloc)
|
||
.context("alloc q_logits_tp1_d")?;
|
||
let v_pred_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc v_pred_d")?;
|
||
let v_pred_tp1_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc v_pred_tp1_d")?;
|
||
let pi_logits_d = stream
|
||
.alloc_zeros::<f32>(b_size * N_ACTIONS)
|
||
.context("alloc pi_logits_d")?;
|
||
|
||
// Device-resident PRNG seeds for IQN tau sampling. Starts at
|
||
// zero; the rl_sample_tau kernel self-seeds on first call using
|
||
// batch index (no memcpy needed).
|
||
let iqn_prng_state_d = stream
|
||
.alloc_zeros::<u32>(b_size)
|
||
.context("alloc iqn_prng_state_d")?;
|
||
|
||
// NoisyNet exploration layer: h_t → [B, N_ACTIONS] noise added
|
||
// to ensemble_q_d before action selection. Replaces the P_MIN
|
||
// probability floor with learned, state-dependent exploration.
|
||
// Per pearl_scoped_init_seed_for_reproducibility: seed offset
|
||
// +200 from dqn_seed keeps Xavier draws independent of C51/IQN.
|
||
let noisy_exploration = NoisyLinear::new(
|
||
dev,
|
||
NoisyLinearConfig {
|
||
in_dim: HIDDEN_DIM,
|
||
out_dim: N_ACTIONS,
|
||
sigma_init: 0.5, // canonical NoisyNet; ISV slot 547 overrides at step time
|
||
seed: cfg.dqn_seed.wrapping_add(200),
|
||
},
|
||
)
|
||
.context("NoisyLinear::new (exploration)")?;
|
||
let noisy_output_d = stream
|
||
.alloc_zeros::<f32>(b_size * N_ACTIONS)
|
||
.context("alloc noisy_output_d")?;
|
||
// Adam state for the 4 NoisyLinear weight tensors. LR mirrored
|
||
// from ISV[RL_LR_Q_INDEX] (same as the Q head — the noisy layer
|
||
// is part of the Q exploration pipeline).
|
||
let noisy_mu_w_adam =
|
||
AdamW::new(dev, noisy_exploration.mu_w.len(), lr_placeholder)
|
||
.context("noisy_mu_w_adam")?;
|
||
let noisy_sigma_w_adam =
|
||
AdamW::new(dev, noisy_exploration.sigma_w.len(), lr_placeholder)
|
||
.context("noisy_sigma_w_adam")?;
|
||
let noisy_mu_b_adam =
|
||
AdamW::new(dev, noisy_exploration.mu_b.len(), lr_placeholder)
|
||
.context("noisy_mu_b_adam")?;
|
||
let noisy_sigma_b_adam =
|
||
AdamW::new(dev, noisy_exploration.sigma_b.len(), lr_placeholder)
|
||
.context("noisy_sigma_b_adam")?;
|
||
// aux_vec_add_inplace for folding noisy output into ensemble_q_d.
|
||
let aux_vec_add_module = ctx
|
||
.load_cubin(AUX_VEC_ADD_CUBIN.to_vec())
|
||
.context("load aux_vec_add cubin")?;
|
||
let aux_vec_add_fn = aux_vec_add_module
|
||
.load_function("aux_vec_add_inplace")
|
||
.context("load aux_vec_add_inplace fn")?;
|
||
|
||
// 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_n_step_gammas_d = stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("alloc sampled_n_step_gammas_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")?;
|
||
|
||
// SP20 P3 FRD head — construct + own per-step buffers. Seed
|
||
// distinct from dqn/ppo to keep Xavier draws independent.
|
||
let frd_head = crate::rl::frd::FrdHead::new(
|
||
dev,
|
||
crate::rl::frd::FrdHeadConfig {
|
||
seed: cfg.dqn_seed.wrapping_add(0xF8D),
|
||
},
|
||
)
|
||
.context("FrdHead::new")?;
|
||
let frd_hidden_d = stream
|
||
.alloc_zeros::<f32>(b_size * crate::rl::common::FRD_HIDDEN_DIM)
|
||
.context("alloc frd_hidden_d")?;
|
||
let frd_logits_d = stream
|
||
.alloc_zeros::<f32>(b_size * crate::rl::frd::FRD_OUT_DIM)
|
||
.context("alloc frd_logits_d")?;
|
||
// SP20 P3 Adam state for the 4 FRD weight tensors. LR mirrored
|
||
// from ISV[RL_FRD_LR_INDEX] (slot 499) on every step before
|
||
// the Adam launches — same pattern as Q/π/V (placeholder=0
|
||
// here, updated per-step).
|
||
let frd_w1_adam = AdamW::new(dev, frd_head.w1_d.len(), lr_placeholder)
|
||
.context("frd_w1_adam")?;
|
||
let frd_b1_adam = AdamW::new(dev, frd_head.b1_d.len(), lr_placeholder)
|
||
.context("frd_b1_adam")?;
|
||
let frd_w2_adam = AdamW::new(dev, frd_head.w2_d.len(), lr_placeholder)
|
||
.context("frd_w2_adam")?;
|
||
let frd_b2_adam = AdamW::new(dev, frd_head.b2_d.len(), lr_placeholder)
|
||
.context("frd_b2_adam")?;
|
||
// FRD labels buffer — sentinel-initialized to -1 (missing
|
||
// horizon). F.5 loader integration overwrites per step.
|
||
let mut frd_labels_d = stream
|
||
.alloc_zeros::<i32>(b_size * crate::rl::common::FRD_N_HORIZONS)
|
||
.context("alloc frd_labels_d")?;
|
||
let neg_ones = vec![-1_i32; b_size * crate::rl::common::FRD_N_HORIZONS];
|
||
write_slice_i32_d_pub(&stream, &neg_ones, &mut frd_labels_d)?;
|
||
|
||
Ok(Self {
|
||
cfg,
|
||
perception,
|
||
dqn_head,
|
||
iqn_head,
|
||
policy_head,
|
||
value_head,
|
||
dqn_w_adam,
|
||
dqn_b_adam,
|
||
iqn_w_embed_adam,
|
||
iqn_b_embed_adam,
|
||
iqn_w_out_adam,
|
||
iqn_b_out_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,
|
||
_rl_reward_clamp_controller_module: rl_reward_clamp_controller_module,
|
||
rl_reward_clamp_controller_fn,
|
||
_rl_atom_support_update_module: rl_atom_support_update_module,
|
||
rl_atom_support_update_fn,
|
||
_rl_q_pi_distill_grad_module: rl_q_pi_distill_grad_module,
|
||
rl_q_pi_distill_grad_fn,
|
||
_rl_ensemble_action_value_module: rl_ensemble_action_value_module,
|
||
rl_ensemble_action_value_fn,
|
||
_rl_q_distill_lambda_controller_module: rl_q_distill_lambda_controller_module,
|
||
rl_q_distill_lambda_controller_fn,
|
||
_rl_unit_state_update_module: rl_unit_state_update_module,
|
||
rl_unit_state_update_fn,
|
||
_rl_trail_mutate_module: rl_trail_mutate_module,
|
||
rl_trail_mutate_fn,
|
||
_rl_trail_stop_check_module: rl_trail_stop_check_module,
|
||
rl_trail_stop_check_fn,
|
||
_rl_position_heat_check_module: rl_position_heat_check_module,
|
||
rl_position_heat_check_fn,
|
||
_rl_confidence_gate_module: rl_confidence_gate_module,
|
||
rl_confidence_gate_fn,
|
||
_rl_frd_gate_module: rl_frd_gate_module,
|
||
rl_frd_gate_fn,
|
||
_rl_recent_outcome_update_module: rl_recent_outcome_update_module,
|
||
rl_recent_outcome_update_fn,
|
||
_rl_gate_threshold_controller_module: rl_gate_threshold_controller_module,
|
||
rl_gate_threshold_controller_fn,
|
||
_rl_fused_controllers_module: rl_fused_controllers_module,
|
||
rl_fused_controllers_fn,
|
||
fused_ctrl_input_slots_d,
|
||
_rl_reward_shaping_module: rl_reward_shaping_module,
|
||
rl_reward_shaping_fn,
|
||
_rl_asymmetric_trail_decay_module: rl_asymmetric_trail_decay_module,
|
||
rl_asymmetric_trail_decay_fn,
|
||
_rl_session_risk_check_module: rl_session_risk_check_module,
|
||
rl_session_risk_check_fn,
|
||
_rl_min_hold_check_module: rl_min_hold_check_module,
|
||
rl_min_hold_check_fn,
|
||
_rl_trade_context_update_module: rl_trade_context_update_module,
|
||
rl_trade_context_update_fn,
|
||
_rl_multires_features_update_module: rl_multires_features_update_module,
|
||
rl_multires_features_update_fn,
|
||
_rl_increment_step_module: rl_increment_step_module,
|
||
rl_increment_step_fn,
|
||
_rl_sample_tau_module: rl_sample_tau_module,
|
||
rl_sample_tau_fn,
|
||
outcome_ema_d,
|
||
trade_context_d,
|
||
multires_output_d,
|
||
multires_state_d,
|
||
multires_prev_mid_d,
|
||
multires_prev_ts_ns_d,
|
||
unit_entry_price_d,
|
||
unit_entry_step_d,
|
||
unit_lots_d,
|
||
unit_initial_r_d,
|
||
unit_trail_distance_d,
|
||
unit_active_d,
|
||
pyramid_units_count_d,
|
||
close_unit_index_d,
|
||
unit_prev_pos_lots_d,
|
||
_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,
|
||
_rl_isv_write_module: rl_isv_write_module,
|
||
rl_isv_write_fn,
|
||
_rl_q_pi_agree_b_module: rl_q_pi_agree_b_module,
|
||
rl_q_pi_agree_b_fn,
|
||
_rl_pi_action_kernel_module: rl_pi_action_kernel_module,
|
||
rl_pi_action_kernel_fn,
|
||
steps_since_done_d,
|
||
trade_duration_emit_d,
|
||
prev_realized_pnl_d,
|
||
prev_position_lots_d,
|
||
rewards_d,
|
||
raw_rewards_d,
|
||
dones_d,
|
||
reward_abs_d,
|
||
actions_d,
|
||
next_actions_d,
|
||
log_pi_old_d,
|
||
advantages_d,
|
||
returns_d,
|
||
h_tp1_d,
|
||
iqn_tau_d,
|
||
iqn_q_values_d,
|
||
iqn_expected_q_d,
|
||
ensemble_q_d,
|
||
q_logits_d,
|
||
q_logits_tp1_d,
|
||
v_pred_d,
|
||
v_pred_tp1_d,
|
||
pi_logits_d,
|
||
iqn_prng_state_d,
|
||
noisy_exploration,
|
||
noisy_output_d,
|
||
noisy_mu_w_adam,
|
||
noisy_sigma_w_adam,
|
||
noisy_mu_b_adam,
|
||
noisy_sigma_b_adam,
|
||
_aux_vec_add_module: aux_vec_add_module,
|
||
aux_vec_add_fn,
|
||
replay,
|
||
n_step_buffer: (0..b_size).map(|_| Vec::with_capacity(16)).collect(),
|
||
sampled_h_t_d,
|
||
sampled_h_tp1_d,
|
||
sampled_actions_d,
|
||
sampled_rewards_d,
|
||
sampled_dones_d,
|
||
sampled_n_step_gammas_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,
|
||
last_k_updates: 0,
|
||
step_counter: 0,
|
||
frd_head,
|
||
frd_hidden_d,
|
||
frd_logits_d,
|
||
frd_labels_d,
|
||
frd_w1_adam,
|
||
frd_b1_adam,
|
||
frd_w2_adam,
|
||
frd_b2_adam,
|
||
}
|
||
.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> {
|
||
// audit — seed ISV-resident design constants FIRST, before any
|
||
// controller bootstrap fires. The controllers' bootstrap paths
|
||
// read these slots (e.g. rl_entropy_coef_controller reads
|
||
// RL_ENTROPY_TARGET_FRAC_INDEX at bootstrap to derive its
|
||
// target); if seeded after, they'd see sentinel 0.0 and
|
||
// bootstrap to wrong values. Generic rl_isv_write per
|
||
// (slot, value) pair — pure device write, no HtoD per
|
||
// `feedback_no_htod_htoh_only_mapped_pinned`.
|
||
{
|
||
let isv_constants: [(usize, f32); 83] = [
|
||
// Static seeds for the adaptive reward-clamp controller —
|
||
// these are the initial values that
|
||
// `rl_reward_clamp_controller` will replace once it
|
||
// observes the first positive reward. Until then, kernels
|
||
// reading slots 452/453 see the documented [-3, +1]
|
||
// defaults. MARGIN (480) starts at 1.5 then adapts via
|
||
// the controller's Schulman bounded-step from clip-rate
|
||
// EMA (482) vs target (483). CLIP_RATE_TARGET seeded to
|
||
// 0.05 = "let 5% of wins ride the clamp as tail
|
||
// outliers." V_MAX/V_MIN seeded to the original C51
|
||
// [-1, +1] span as the canonical floor — the ratchet
|
||
// in `rl_reward_clamp_controller` only grows magnitude.
|
||
(crate::rl::isv_slots::RL_REWARD_CLAMP_WIN_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_REWARD_CLAMP_LOSS_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_REWARD_CLAMP_MARGIN_INDEX, 1.5),
|
||
(crate::rl::isv_slots::RL_REWARD_CLAMP_RATIO_INDEX, 3.0),
|
||
(crate::rl::isv_slots::RL_REWARD_CLAMP_CLIP_RATE_TARGET_INDEX, 0.05),
|
||
(crate::rl::isv_slots::RL_C51_V_MAX_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_C51_V_MIN_INDEX, -0.5),
|
||
// Q→π distillation (vj5f6 followup) — λ now controller-
|
||
// driven by `rl_q_distill_lambda_controller` from
|
||
// KL_EMA vs KL_TARGET. Seed 0.05 is the starting point
|
||
// before the controller picks up. τ=1.0 canonical.
|
||
(crate::rl::isv_slots::RL_Q_DISTILL_LAMBDA_INDEX, 0.1),
|
||
(crate::rl::isv_slots::RL_Q_DISTILL_TEMPERATURE_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_Q_DISTILL_KL_TARGET_INDEX, 0.1),
|
||
// Reward-scale MIN floor (was hardcoded 1e-3) — now
|
||
// ISV-driven, lowered to 1e-4 per wwcsz audit to give
|
||
// one more order of magnitude before pegging when
|
||
// mean_abs_pnl spikes.
|
||
(crate::rl::isv_slots::RL_REWARD_SCALE_MIN_INDEX, 1e-4),
|
||
// Q→π distill KL_EMA blend α (was hardcoded 0.05f in
|
||
// rl_q_pi_distill_grad.cu, caught by audit-isv.sh dogfood).
|
||
(crate::rl::isv_slots::RL_Q_DISTILL_KL_EMA_ALPHA_INDEX, 0.05),
|
||
// SP20 P1+P5 trail-stop bounds (audit-wiring fix for a7/a8 dead).
|
||
(crate::rl::isv_slots::RL_TRAIL_MIN_INDEX, 0.001),
|
||
(crate::rl::isv_slots::RL_TRAIL_MAX_INDEX, 100.0),
|
||
(crate::rl::isv_slots::RL_TRAIL_K_INIT_INDEX, 2.0),
|
||
(crate::rl::isv_slots::RL_TRAIL_ADJUST_RATE_INDEX, 0.9),
|
||
// SP20 P3 FRD head seeds — λ, LR, h1-3 tick offsets, ±σ atom range.
|
||
// Per feedback_isv_for_adaptive_bounds + single-source-of-truth:
|
||
// horizon ticks + bucket-range σ derive from the canonical
|
||
// `rl::common` consts (same source the loader's
|
||
// `compute_frd_labels` reads via `MultiHorizonLoaderConfig`),
|
||
// so the two sides can't drift across literal-vs-const edits.
|
||
(crate::rl::isv_slots::RL_FRD_LAMBDA_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_FRD_LR_INDEX, 1e-3),
|
||
(crate::rl::isv_slots::RL_FRD_HORIZON_1_TICKS_INDEX,
|
||
crate::rl::common::FRD_HORIZON_TICKS[0] as f32),
|
||
(crate::rl::isv_slots::RL_FRD_HORIZON_2_TICKS_INDEX,
|
||
crate::rl::common::FRD_HORIZON_TICKS[1] as f32),
|
||
(crate::rl::isv_slots::RL_FRD_HORIZON_3_TICKS_INDEX,
|
||
crate::rl::common::FRD_HORIZON_TICKS[2] as f32),
|
||
(crate::rl::isv_slots::RL_FRD_BUCKET_RANGE_SIGMA_INDEX,
|
||
crate::rl::common::FRD_BUCKET_RANGE_SIGMA),
|
||
(crate::rl::isv_slots::RL_HEAT_CAP_MAX_LOTS_INDEX, 8.0),
|
||
(crate::rl::isv_slots::RL_PYRAMID_THRESHOLD_INDEX, 0.50),
|
||
(crate::rl::isv_slots::RL_ANTIMARTINGALE_KAPPA_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_ANTIMARTINGALE_MIN_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_ANTIMARTINGALE_MAX_INDEX, 2.0),
|
||
(crate::rl::isv_slots::RL_CONF_GATE_THRESHOLD_INDEX, 0.01),
|
||
(crate::rl::isv_slots::RL_CONF_GATE_LAMBDA_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_CONF_GATE_SIGMA_NORM_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_FRD_GATE_THR_LONG_INDEX, 0.05),
|
||
(crate::rl::isv_slots::RL_FRD_GATE_THR_SHORT_INDEX, 0.05),
|
||
// P_MIN set to 0.0: NoisyNet state-dependent exploration
|
||
// replaces the fixed probability floor. Per
|
||
// `pearl_pi_actor_collapses_without_entropy_floor`:
|
||
// NoisyLinear noise on ensemble_q_d provides learned
|
||
// exploration, so the P_MIN anti-collapse floor in
|
||
// rl_pi_action_kernel is no longer needed.
|
||
(crate::rl::isv_slots::RL_PI_SAMPLE_P_MIN_INDEX, 0.0),
|
||
(crate::rl::isv_slots::RL_OUTCOME_ALPHA_INDEX, 0.1),
|
||
(crate::rl::isv_slots::RL_GATE_WARMUP_STEPS_INDEX, 10000.0),
|
||
(crate::rl::isv_slots::RL_GATE_DONES_TARGET_INDEX, 0.10),
|
||
(crate::rl::isv_slots::RL_GATE_CONF_MIN_INDEX, 0.0),
|
||
(crate::rl::isv_slots::RL_GATE_CONF_MAX_INDEX, 0.50),
|
||
(crate::rl::isv_slots::RL_GATE_FRD_MIN_INDEX, 0.0),
|
||
(crate::rl::isv_slots::RL_GATE_FRD_MAX_INDEX, 0.50),
|
||
(crate::rl::isv_slots::RL_GATE_ADJUST_RATE_INDEX, 0.95),
|
||
(crate::rl::isv_slots::RL_ENTRY_COST_INDEX, 15.0),
|
||
(crate::rl::isv_slots::RL_SHORT_HOLD_MIN_STEPS_INDEX, 100.0),
|
||
(crate::rl::isv_slots::RL_SHORT_HOLD_PENALTY_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_HOLD_BONUS_INDEX, 2.0),
|
||
(crate::rl::isv_slots::RL_MIN_HOLD_STEPS_INDEX, 100.0),
|
||
(crate::rl::isv_slots::RL_ASYM_LOSS_DECAY_RATE_INDEX, 0.995),
|
||
(crate::rl::isv_slots::RL_ASYM_WIN_TRAIL_FACTOR_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_SESSION_LOSS_LIMIT_INDEX, -50.0),
|
||
(crate::rl::isv_slots::RL_SESSION_EMA_ALPHA_INDEX, 0.02),
|
||
(crate::rl::isv_slots::RL_ASYM_WIN_THRESHOLD_INDEX, 0.25),
|
||
(crate::rl::isv_slots::RL_N_STEP_INDEX, 10.0),
|
||
(crate::rl::isv_slots::RL_MULTIRES_HORIZON_1_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_MULTIRES_HORIZON_2_INDEX, 10.0),
|
||
(crate::rl::isv_slots::RL_MULTIRES_HORIZON_3_INDEX, 600.0),
|
||
(crate::rl::isv_slots::RL_KL_TARGET_INDEX, 0.01),
|
||
(crate::rl::isv_slots::RL_IMPROVEMENT_THRESHOLD_INDEX, 0.99),
|
||
(crate::rl::isv_slots::RL_PLATEAU_PATIENCE_INDEX, 1000.0),
|
||
(crate::rl::isv_slots::RL_DIV_TARGET_INDEX, 0.01),
|
||
(crate::rl::isv_slots::RL_ENTROPY_TARGET_FRAC_INDEX, 0.7),
|
||
(crate::rl::isv_slots::RL_KURT_LIFT_SCALE_INDEX, 7.0),
|
||
(crate::rl::isv_slots::RL_PPO_CLAMP_MARGIN_INDEX, 10.0),
|
||
(crate::rl::isv_slots::RL_LR_WARMUP_STEPS_INDEX, 2000.0),
|
||
(crate::rl::isv_slots::RL_LR_BOOTSTRAP_INDEX, 1e-3),
|
||
(crate::rl::isv_slots::RL_LR_MIN_INDEX, 1e-4),
|
||
(crate::rl::isv_slots::RL_LR_MAX_INDEX, 1e-2),
|
||
(crate::rl::isv_slots::RL_LR_LOSS_EMA_ALPHA_INDEX, 0.05),
|
||
(crate::rl::isv_slots::RL_LR_DECAY_FACTOR_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_LOSS_LAMBDA_AUX_INDEX, 1.0),
|
||
// Batch 3 — Schulman pattern + streaming α + kurt refs + bootstraps.
|
||
(crate::rl::isv_slots::RL_SCHULMAN_TOLERANCE_INDEX, 1.5),
|
||
(crate::rl::isv_slots::RL_SCHULMAN_ADJUST_RATE_INDEX, 1.5),
|
||
(crate::rl::isv_slots::RL_STREAM_ALPHA_INDEX, 0.05),
|
||
(crate::rl::isv_slots::RL_KURT_GAUSSIAN_INDEX, 3.0),
|
||
(crate::rl::isv_slots::RL_KURT_NOISE_FLOOR_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_TAU_BOOTSTRAP_INDEX, 0.005),
|
||
(crate::rl::isv_slots::RL_EPS_BOOTSTRAP_INDEX, 0.2),
|
||
(crate::rl::isv_slots::RL_ROLLOUT_BOOTSTRAP_INDEX, 2048.0),
|
||
(crate::rl::isv_slots::RL_REWARD_SCALE_BOOTSTRAP_INDEX, 1.0),
|
||
(crate::rl::isv_slots::RL_PPO_RATIO_CLAMP_BOOTSTRAP_INDEX, 10.0),
|
||
// IQN ensemble seeds — alpha=0.5 (equal C51/IQN blend),
|
||
// N_TAU=32, LR=1e-3 (canonical).
|
||
(crate::rl::isv_slots::RL_IQN_N_TAU_INDEX, 32.0),
|
||
(crate::rl::isv_slots::RL_IQN_ENSEMBLE_ALPHA_INDEX, 0.5),
|
||
(crate::rl::isv_slots::RL_IQN_LR_INDEX, 1e-3),
|
||
// NoisyNet σ_init — canonical Fortunato et al. 2017 value.
|
||
// Read by NoisyLinear at construction (fallback if ISV is
|
||
// sentinel-zero) and available for runtime re-tuning via
|
||
// diagnostic JSONL visibility.
|
||
(crate::rl::isv_slots::RL_NOISY_SIGMA_INIT_INDEX, 0.5),
|
||
];
|
||
let cfg_isv = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
for (slot, value) in isv_constants.iter() {
|
||
let slot_i32 = *slot as i32;
|
||
let value_f32 = *value;
|
||
let mut launch = self.stream.launch_builder(&self.rl_isv_write_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(&slot_i32)
|
||
.arg(&value_f32);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_isv)
|
||
.context("launch rl_isv_write (design constant)")?;
|
||
}
|
||
}
|
||
}
|
||
|
||
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")?;
|
||
// audit — 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("bootstrap rl_ppo_ratio_clamp_controller")?;
|
||
// audit — 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;
|
||
// K-loop config — also ISV-driven per
|
||
// `feedback_isv_for_adaptive_bounds`. Divisor 2048 matches
|
||
// ROLLOUT_BOOTSTRAP (so K=1 at controller bootstrap, before
|
||
// any adaptation). K_MAX=4 prevents gradient overtraining
|
||
// at b_size=1; raised from default 8 after alpha-rl-f2ggr
|
||
// confirmed K=8 firing on 22% of steps over-trained
|
||
// (KL excursions to 12.44, reward/trade -$0.585 → -$0.723).
|
||
let k_loop_divisor_slot: i32 =
|
||
crate::rl::isv_slots::RL_K_LOOP_DIVISOR_INDEX as i32;
|
||
let k_loop_divisor_val: f32 = 2048.0;
|
||
let k_loop_max_slot: i32 =
|
||
crate::rl::isv_slots::RL_K_LOOP_MAX_INDEX as i32;
|
||
let k_loop_max_val: f32 = 4.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)
|
||
.arg(&k_loop_divisor_slot)
|
||
.arg(&k_loop_divisor_val)
|
||
.arg(&k_loop_max_slot)
|
||
.arg(&k_loop_max_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(())
|
||
}
|
||
|
||
/// Fire all 10 RL ISV controllers in a single fused kernel launch.
|
||
///
|
||
/// Replaces 10 individual kernel launches (gamma, tau, ppo_clip,
|
||
/// entropy_coef, rollout_steps, per_alpha, reward_scale,
|
||
/// ppo_ratio_clamp, gate_threshold, q_distill_lambda) with one
|
||
/// launch of `rl_fused_controllers`, saving ~40-80us/step of
|
||
/// kernel-launch overhead.
|
||
///
|
||
/// Slot pairing (output <- input):
|
||
/// ISV[400] gamma <- ISV[417] MEAN_TRADE_DURATION_EMA
|
||
/// ISV[401] tau <- ISV[418] Q_DIVERGENCE_EMA
|
||
/// ISV[402] eps <- 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_alpha <- ISV[422] TD_KURTOSIS_EMA
|
||
/// ISV[406] reward_scale <- ISV[423] MEAN_ABS_PNL_EMA
|
||
/// ISV[440] ratio_clamp <- ISV[402] (reads eps just written)
|
||
/// ISV[512,516,517] gates <- dones[B]
|
||
/// ISV[486] lambda_distill <- ISV[488] KL_EMA
|
||
pub fn launch_rl_fused_controllers(
|
||
&self,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let alpha = RL_LR_CONTROLLER_ALPHA;
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_fused_controllers_fn);
|
||
launch
|
||
.arg(&self.isv_d)
|
||
.arg(&self.dones_d)
|
||
.arg(&alpha)
|
||
.arg(&b_size_i)
|
||
.arg(&self.fused_ctrl_input_slots_d);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_fused_controllers launch")?;
|
||
}
|
||
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(())
|
||
}
|
||
|
||
/// SP20 P4: launch `actions_to_market_targets` — convert per-batch
|
||
/// action index → (side, size) pair written to `market_targets_d`
|
||
/// as `[b_size * 2]` row-major. Reads `pos_state_d` for actions
|
||
/// that depend on current position (a3/a4 flat, a9/a10 half-flat).
|
||
/// Public for `sp20_trade_management_kernels` GPU-oracle tests.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn launch_actions_to_market_targets(
|
||
&self,
|
||
actions_d: &CudaSlice<i32>,
|
||
pos_state_d: &CudaSlice<u8>,
|
||
market_targets_d: &mut CudaSlice<i32>,
|
||
bid_px_d: &CudaSlice<f32>,
|
||
ask_px_d: &CudaSlice<f32>,
|
||
unit_entry_price_d: &CudaSlice<f32>,
|
||
unit_active_d: &CudaSlice<u8>,
|
||
unit_lots_d: &CudaSlice<i32>,
|
||
pyramid_units_count_d: &CudaSlice<i32>,
|
||
close_unit_index_d: &CudaSlice<i32>,
|
||
outcome_ema_d: &CudaSlice<f32>,
|
||
b_size: usize,
|
||
pos_bytes: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
debug_assert_eq!(market_targets_d.len(), b_size * 2);
|
||
debug_assert_eq!(pos_state_d.len(), b_size * pos_bytes);
|
||
debug_assert_eq!(unit_entry_price_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_active_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_lots_d.len(), b_size * 4);
|
||
debug_assert_eq!(pyramid_units_count_d.len(), b_size);
|
||
debug_assert_eq!(close_unit_index_d.len(), b_size);
|
||
debug_assert_eq!(outcome_ema_d.len(), b_size);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((b_size as u32) + 31) / 32, 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = pos_bytes as i32;
|
||
let mut launch = self.stream.launch_builder(&self.actions_to_market_targets_fn);
|
||
launch
|
||
.arg(actions_d)
|
||
.arg(pos_state_d)
|
||
.arg(market_targets_d)
|
||
.arg(&self.isv_d)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(unit_entry_price_d)
|
||
.arg(unit_active_d)
|
||
.arg(unit_lots_d)
|
||
.arg(pyramid_units_count_d)
|
||
.arg(close_unit_index_d)
|
||
.arg(outcome_ema_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("actions_to_market_targets launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// SP20 P5: launch `rl_trail_mutate` — mutates per-(batch, unit)
|
||
/// `unit_trail_distance_d` for active units when the chosen
|
||
/// action is a7 (TrailTighten) or a8 (TrailLoosen). Uses ISV
|
||
/// slots 494/495/497 (MIN/MAX/ADJUST_RATE) from `self.isv_d`.
|
||
/// Public for SP20 P14 GPU-oracle tests.
|
||
pub fn launch_rl_trail_mutate(
|
||
&self,
|
||
actions_d: &CudaSlice<i32>,
|
||
unit_active_d: &CudaSlice<u8>,
|
||
unit_trail_distance_d: &mut CudaSlice<f32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
debug_assert_eq!(unit_active_d.len(), b_size * 4); // MAX_UNITS
|
||
debug_assert_eq!(unit_trail_distance_d.len(), b_size * 4);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (4, 1, 1), // MAX_UNITS threads per block
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_trail_mutate_fn);
|
||
launch
|
||
.arg(actions_d)
|
||
.arg(&self.isv_d)
|
||
.arg(unit_active_d)
|
||
.arg(unit_trail_distance_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch.launch(cfg).context("rl_trail_mutate launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Launch `rl_trail_stop_check` — for each active unit, check if
|
||
/// mid breached the unit's stop (entry ± trail_distance). On breach
|
||
/// with multiple units: overrides to HalfFlat + writes
|
||
/// close_unit_index. With single unit: overrides to full flat.
|
||
/// Resets close_unit_index_d to -1 device-side each invocation.
|
||
/// Public for GPU-oracle tests.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn launch_rl_trail_stop_check(
|
||
&self,
|
||
actions_d: &mut CudaSlice<i32>,
|
||
bid_px_d: &CudaSlice<f32>,
|
||
ask_px_d: &CudaSlice<f32>,
|
||
unit_active_d: &CudaSlice<u8>,
|
||
unit_entry_price_d: &CudaSlice<f32>,
|
||
unit_lots_d: &CudaSlice<i32>,
|
||
unit_trail_distance_d: &CudaSlice<f32>,
|
||
pyramid_units_count_d: &CudaSlice<i32>,
|
||
close_unit_index_d: &mut CudaSlice<i32>,
|
||
b_size: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
debug_assert_eq!(unit_active_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_entry_price_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_lots_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_trail_distance_d.len(), b_size * 4);
|
||
debug_assert_eq!(pyramid_units_count_d.len(), b_size);
|
||
debug_assert_eq!(close_unit_index_d.len(), b_size);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (4, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_trail_stop_check_fn);
|
||
launch
|
||
.arg(actions_d)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(unit_active_d)
|
||
.arg(unit_entry_price_d)
|
||
.arg(unit_lots_d)
|
||
.arg(unit_trail_distance_d)
|
||
.arg(pyramid_units_count_d)
|
||
.arg(close_unit_index_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_trail_stop_check launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// SP20 P6: launch `rl_position_heat_check` — if `|position_lots|`
|
||
/// exceeds the ISV-driven cap (`RL_HEAT_CAP_MAX_LOTS_INDEX`, slot
|
||
/// 504, default 8), OVERRIDE `actions[b]` to FlatFromLong (a3) or
|
||
/// FlatFromShort (a4). Last-defense guard — full flat, no partial.
|
||
/// Also writes `RL_HEAT_CAP_FIRED_COUNT_INDEX` (slot 505) with the
|
||
/// per-step fired count for diag.
|
||
pub fn launch_rl_position_heat_check(
|
||
&self,
|
||
actions_d: &mut CudaSlice<i32>,
|
||
pos_state_d: &CudaSlice<u8>,
|
||
b_size: usize,
|
||
pos_bytes: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
debug_assert_eq!(pos_state_d.len(), b_size * pos_bytes);
|
||
let block = (b_size as u32).min(256);
|
||
let grid = ((b_size as u32) + block - 1) / block;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (grid.max(1), 1, 1),
|
||
block_dim: (block.max(1), 1, 1),
|
||
shared_mem_bytes: 256 * std::mem::size_of::<i32>() as u32,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = pos_bytes as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_position_heat_check_fn);
|
||
launch
|
||
.arg(actions_d)
|
||
.arg(pos_state_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_position_heat_check launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Launch `rl_confidence_gate` — for opening actions on flat
|
||
/// positions, compute C51 distributional LCB and override to Hold
|
||
/// if below threshold. Public for GPU-oracle tests.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn launch_rl_confidence_gate(
|
||
&self,
|
||
actions_d: &mut CudaSlice<i32>,
|
||
q_logits_d: &CudaSlice<f32>,
|
||
pos_state_d: &CudaSlice<u8>,
|
||
b_size: usize,
|
||
pos_bytes: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
debug_assert_eq!(q_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
|
||
debug_assert_eq!(pos_state_d.len(), b_size * pos_bytes);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = pos_bytes as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_confidence_gate_fn);
|
||
launch
|
||
.arg(actions_d)
|
||
.arg(q_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(pos_state_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_confidence_gate launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Launch `rl_frd_gate` — for opening actions on flat positions,
|
||
/// check FRD head's horizon-2 favorable probability mass. Override
|
||
/// to Hold if below threshold. Public for GPU-oracle tests.
|
||
pub fn launch_rl_frd_gate(
|
||
&self,
|
||
actions_d: &mut CudaSlice<i32>,
|
||
frd_logits_d: &CudaSlice<f32>,
|
||
pos_state_d: &CudaSlice<u8>,
|
||
b_size: usize,
|
||
pos_bytes: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(actions_d.len(), b_size);
|
||
let frd_out_dim = crate::rl::frd::FRD_OUT_DIM;
|
||
debug_assert_eq!(frd_logits_d.len(), b_size * frd_out_dim);
|
||
debug_assert_eq!(pos_state_d.len(), b_size * pos_bytes);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = pos_bytes as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_frd_gate_fn);
|
||
launch
|
||
.arg(actions_d)
|
||
.arg(frd_logits_d)
|
||
.arg(pos_state_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_frd_gate launch")?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Launch `rl_unit_state_update` — detects per-batch position
|
||
/// transitions vs `prev_pos_lots_d` and updates per-unit state
|
||
/// (entry_price, entry_step, lots, initial_r, trail_distance,
|
||
/// active, pyramid_units_count). Handles pyramid add (allocate next
|
||
/// slot) and partial flat (deactivate oldest slot). Uses ISV slot
|
||
/// 496 (K_INIT) from `self.isv_d` to scale the bootstrap trail.
|
||
/// Public for GPU-oracle tests.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn launch_rl_unit_state_update(
|
||
&self,
|
||
pos_state_d: &CudaSlice<u8>,
|
||
prev_pos_lots_d: &mut CudaSlice<i32>,
|
||
unit_entry_price_d: &mut CudaSlice<f32>,
|
||
unit_entry_step_d: &mut CudaSlice<i32>,
|
||
unit_lots_d: &mut CudaSlice<i32>,
|
||
unit_initial_r_d: &mut CudaSlice<f32>,
|
||
unit_trail_distance_d: &mut CudaSlice<f32>,
|
||
unit_active_d: &mut CudaSlice<u8>,
|
||
pyramid_units_count_d: &mut CudaSlice<i32>,
|
||
b_size: usize,
|
||
pos_bytes: usize,
|
||
) -> Result<()> {
|
||
debug_assert_eq!(pos_state_d.len(), b_size * pos_bytes);
|
||
debug_assert_eq!(prev_pos_lots_d.len(), b_size);
|
||
debug_assert_eq!(unit_entry_price_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_entry_step_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_lots_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_initial_r_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_trail_distance_d.len(), b_size * 4);
|
||
debug_assert_eq!(unit_active_d.len(), b_size * 4);
|
||
debug_assert_eq!(pyramid_units_count_d.len(), b_size);
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((b_size as u32) + 31) / 32, 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = pos_bytes as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_unit_state_update_fn);
|
||
launch
|
||
.arg(pos_state_d)
|
||
.arg(&self.isv_d)
|
||
.arg(prev_pos_lots_d)
|
||
.arg(unit_entry_price_d)
|
||
.arg(unit_entry_step_d)
|
||
.arg(unit_lots_d)
|
||
.arg(unit_initial_r_d)
|
||
.arg(unit_trail_distance_d)
|
||
.arg(unit_active_d)
|
||
.arg(pyramid_units_count_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_unit_state_update 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);
|
||
// Shared mem: 2 × block × f32 — kernel runs two parallel
|
||
// reductions (max|scaled| + max(positive scaled)) per audit
|
||
// 2026-05-24, both stored in `extern __shared__ float smem[]`.
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_x, 1, 1),
|
||
shared_mem_bytes: 3 * 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")?;
|
||
}
|
||
|
||
// Adaptive reward-clamp controller — reads the per-step
|
||
// positive-tail max just published into slot 478, maintains
|
||
// an EMA in slot 479, writes adaptive WIN/LOSS bounds into
|
||
// slots 452/453 for the NEXT step's apply_reward_scale to
|
||
// consume. Single-thread kernel; α driven by the same
|
||
// controller-cohort floor as the other R5 EMAs.
|
||
{
|
||
let cfg_ctrl = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let alpha = RL_LR_CONTROLLER_ALPHA;
|
||
let mut launch = self
|
||
.stream
|
||
.launch_builder(&self.rl_reward_clamp_controller_fn);
|
||
launch.arg(&self.isv_d).arg(&alpha);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_ctrl)
|
||
.context("rl_reward_clamp_controller launch")?;
|
||
}
|
||
}
|
||
|
||
// C51 atom support refresh — the reward clamp controller just
|
||
// ratcheted V_MIN/V_MAX; rewrite atom_supports_d so downstream
|
||
// C51 kernels (argmax_expected_q, rl_action_kernel,
|
||
// dqn_distributional_q) see the updated span. 21-thread
|
||
// single-block kernel; trivially cheap.
|
||
{
|
||
let cfg_atom = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (21, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self
|
||
.stream
|
||
.launch_builder(&self.rl_atom_support_update_fn);
|
||
launch.arg(&self.isv_d).arg(&self.atom_supports_d);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_atom)
|
||
.context("rl_atom_support_update 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;
|
||
// SP20 P3 FRD head — LR from dedicated ISV slot (498-503 block).
|
||
let lr_frd = self.isv_host[crate::rl::isv_slots::RL_FRD_LR_INDEX];
|
||
self.frd_w1_adam.lr = lr_frd;
|
||
self.frd_b1_adam.lr = lr_frd;
|
||
self.frd_w2_adam.lr = lr_frd;
|
||
self.frd_b2_adam.lr = lr_frd;
|
||
|
||
// 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;
|
||
|
||
// 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)?;
|
||
// SP20 P3 FRD layer-1 backward produces grad_h_t into this
|
||
// buffer (the encoder-upstream gradient). Folded into
|
||
// grad_h_t_combined_d at Step 10 with the lambdas.frd weight.
|
||
let mut frd_grad_h_t_d = self.stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)?;
|
||
|
||
// 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 self.q_logits_d)
|
||
.context("dqn_head.forward(sampled_h_t) [R7d off-policy]")?;
|
||
self.policy_head
|
||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||
.context("policy_head.forward_logits")?;
|
||
self.value_head
|
||
.forward(h_t_borrow, &self.isv_d, b_size, &mut self.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.
|
||
self.dqn_head
|
||
.forward(&self.sampled_h_tp1_d, b_size, &mut self.q_logits_tp1_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(&self.q_logits_tp1_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(
|
||
&self.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.sampled_n_step_gammas_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(
|
||
&self.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(
|
||
&self.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")?;
|
||
|
||
// Q→π distillation gradient — ADDS λ × (π_new - π_target) to
|
||
// pi_grad_logits where π_target = softmax(E_Q[s,*] / τ).
|
||
// Couples Q's improved C51 calibration to π's action selection
|
||
// (per vj5f6 finding that Q was decoupled from policy under
|
||
// Option B). One block per batch, N_ACTIONS=9 threads per block.
|
||
// q_logits_d is a persistent field from the earlier forward.
|
||
{
|
||
let cfg_distill = 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_q_pi_distill_grad_fn);
|
||
launch
|
||
.arg(&self.q_logits_d)
|
||
.arg(&self.pi_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&mut pi_grad_logits_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_distill)
|
||
.context("rl_q_pi_distill_grad launch")?;
|
||
}
|
||
}
|
||
|
||
// NOTE: q_distill_lambda controller is now fused into
|
||
// rl_fused_controllers (fired in the per-step block). Reads
|
||
// KL_EMA from the previous step's backward pass — one-step
|
||
// delay is negligible for a 0.998/1.2 step-size controller.
|
||
|
||
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,
|
||
&self.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")?;
|
||
|
||
// ── SP20 P3 FRD backward chain (F.4) ─────────────────────────
|
||
// Reads frd_hidden_d + frd_logits_d (populated by
|
||
// step_with_lobsim's FRD fwd) + self.frd_labels_d (sentinel -1
|
||
// until F.5 loader integration populates from forward-snapshot
|
||
// lookahead). On all-sentinel labels, softmax_ce_grad emits
|
||
// zero gradient + zero loss → entire chain produces no signal,
|
||
// Adam steps are no-ops modulo β decay. Once F.5 lands real
|
||
// labels, this path activates the supervised learning signal
|
||
// without any further trainer-side changes.
|
||
let l_frd_host: f32 = {
|
||
let frd_n_h = crate::rl::common::FRD_N_HORIZONS;
|
||
let frd_h_dim = crate::rl::common::FRD_HIDDEN_DIM;
|
||
let frd_out_dim = crate::rl::frd::FRD_OUT_DIM;
|
||
let mut frd_grad_logits_d = self.stream.alloc_zeros::<f32>(b_size * frd_out_dim)?;
|
||
let mut frd_loss_per_b_h_d = self.stream.alloc_zeros::<f32>(b_size * frd_n_h)?;
|
||
self.frd_head.softmax_ce_grad(
|
||
&self.frd_logits_d,
|
||
&self.frd_labels_d,
|
||
&mut frd_grad_logits_d,
|
||
&mut frd_loss_per_b_h_d,
|
||
b_size,
|
||
)?;
|
||
let mut frd_grad_w2_pb_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * frd_h_dim * frd_out_dim)?;
|
||
let mut frd_grad_b2_pb_d = self.stream.alloc_zeros::<f32>(b_size * frd_out_dim)?;
|
||
let mut frd_grad_hidden_d = self.stream.alloc_zeros::<f32>(b_size * frd_h_dim)?;
|
||
self.frd_head.layer2_bwd(
|
||
&self.frd_hidden_d,
|
||
&frd_grad_logits_d,
|
||
&mut frd_grad_w2_pb_d,
|
||
&mut frd_grad_b2_pb_d,
|
||
&mut frd_grad_hidden_d,
|
||
b_size,
|
||
)?;
|
||
let mut frd_grad_w1_pb_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * HIDDEN_DIM * frd_h_dim)?;
|
||
let mut frd_grad_b1_pb_d = self.stream.alloc_zeros::<f32>(b_size * frd_h_dim)?;
|
||
self.frd_head.layer1_bwd(
|
||
h_t_borrow,
|
||
&self.frd_hidden_d,
|
||
&frd_grad_hidden_d,
|
||
&mut frd_grad_w1_pb_d,
|
||
&mut frd_grad_b1_pb_d,
|
||
&mut frd_grad_h_t_d,
|
||
b_size,
|
||
)?;
|
||
// Reduce-axis-0: batch-summed weight grads.
|
||
let mut frd_grad_w1_d = self.stream.alloc_zeros::<f32>(HIDDEN_DIM * frd_h_dim)?;
|
||
let mut frd_grad_b1_d = self.stream.alloc_zeros::<f32>(frd_h_dim)?;
|
||
let mut frd_grad_w2_d = self.stream.alloc_zeros::<f32>(frd_h_dim * frd_out_dim)?;
|
||
let mut frd_grad_b2_d = self.stream.alloc_zeros::<f32>(frd_out_dim)?;
|
||
self.launch_reduce_axis0(
|
||
&frd_grad_w1_pb_d,
|
||
b_size,
|
||
HIDDEN_DIM * frd_h_dim,
|
||
&mut frd_grad_w1_d,
|
||
)?;
|
||
self.launch_reduce_axis0(
|
||
&frd_grad_b1_pb_d,
|
||
b_size,
|
||
frd_h_dim,
|
||
&mut frd_grad_b1_d,
|
||
)?;
|
||
self.launch_reduce_axis0(
|
||
&frd_grad_w2_pb_d,
|
||
b_size,
|
||
frd_h_dim * frd_out_dim,
|
||
&mut frd_grad_w2_d,
|
||
)?;
|
||
self.launch_reduce_axis0(
|
||
&frd_grad_b2_pb_d,
|
||
b_size,
|
||
frd_out_dim,
|
||
&mut frd_grad_b2_d,
|
||
)?;
|
||
// Adam steps — sentinel labels yield zero grads → no
|
||
// weight movement (only momentum decay).
|
||
self.frd_w1_adam
|
||
.step(&mut self.frd_head.w1_d, &frd_grad_w1_d)
|
||
.context("frd_w1_adam.step")?;
|
||
self.frd_b1_adam
|
||
.step(&mut self.frd_head.b1_d, &frd_grad_b1_d)
|
||
.context("frd_b1_adam.step")?;
|
||
self.frd_w2_adam
|
||
.step(&mut self.frd_head.w2_d, &frd_grad_w2_d)
|
||
.context("frd_w2_adam.step")?;
|
||
self.frd_b2_adam
|
||
.step(&mut self.frd_head.b2_d, &frd_grad_b2_d)
|
||
.context("frd_b2_adam.step")?;
|
||
// Per-row CE loss → mean over (b, h). Mapped-pinned read.
|
||
let loss_pb_h =
|
||
read_slice_d(&self.stream, &frd_loss_per_b_h_d, b_size * frd_n_h)?;
|
||
loss_pb_h.iter().sum::<f32>() / ((b_size * frd_n_h) as f32)
|
||
};
|
||
|
||
// ── 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,
|
||
)?;
|
||
// SP20 P3 FRD encoder-upstream gradient — same accumulator
|
||
// pattern as π / V. With sentinel labels (F.4 state) the
|
||
// buffer is all-zeros so this is a no-op; once F.5 lands real
|
||
// labels it injects supervised-learning signal into the
|
||
// encoder via the standard λ-weighted accumulator.
|
||
accumulate_grad_h(
|
||
&self.stream,
|
||
&self.grad_h_accumulate_fn,
|
||
&frd_grad_h_t_d,
|
||
lambdas.frd,
|
||
&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 6 heads per
|
||
// pearl_loss_balance_controller (SP20 P3 added FRD as 6th).
|
||
let l_total = (lambdas.bce * l_bce
|
||
+ lambdas.q * l_q
|
||
+ lambdas.pi * l_pi
|
||
+ lambdas.v * l_v_host
|
||
+ lambdas.aux * l_aux
|
||
+ lambdas.frd * l_frd_host)
|
||
/ 6.0;
|
||
|
||
Ok(IntegratedStepStats {
|
||
l_bce,
|
||
l_q,
|
||
l_pi,
|
||
l_v: l_v_host,
|
||
l_aux,
|
||
l_frd: l_frd_host,
|
||
l_total,
|
||
lambdas,
|
||
})
|
||
}
|
||
|
||
/// DQN-only replay step — used by step_with_lobsim's K-loop for
|
||
/// iterations after the first. Skips PPO surrogate, V regression,
|
||
/// encoder backward, LR controller emit, ISV mirror refresh — all
|
||
/// of those run once per env step in the first iter via
|
||
/// step_synthetic. This helper just does:
|
||
///
|
||
/// 1. Forward Q on `sampled_h_t` and `sampled_h_tp1`
|
||
/// 2. Double-DQN argmax on online Q at h_tp1 →
|
||
/// `sampled_next_actions_d`
|
||
/// 3. Bellman target via `forward_target` (target net) +
|
||
/// `select_action_atoms` + `project_bellman_target`
|
||
/// 4. Q backward: logits → grad_w/b/h_t. Q's grad_h_t is
|
||
/// discarded (R7d stop-grad: SAMPLED h_t is a past-step
|
||
/// encoder output; folding its gradient into the encoder
|
||
/// would poison live training).
|
||
/// 5. Per-batch grad reduce → grad_w/grad_b
|
||
/// 6. Q Adam (uses LR already set in step_synthetic via LR
|
||
/// controller emit)
|
||
/// 7. Writes `td_per_sample_d` for PER priority update by the
|
||
/// caller.
|
||
///
|
||
/// Assumes `sample_and_gather` has been called externally —
|
||
/// reads `sampled_h_t_d / sampled_h_tp1_d / sampled_actions_d /
|
||
/// sampled_rewards_d / sampled_dones_d` and writes
|
||
/// `sampled_next_actions_d`.
|
||
///
|
||
/// Returns mean Q loss (scalar). Caller uses this for the
|
||
/// LR-controller's plateau detection (but DOES NOT mirror it into
|
||
/// `last_q_loss` — that field is the env-step-paired iter's Q
|
||
/// loss, the canonical reference for plateau detection).
|
||
pub fn dqn_replay_step(&mut self, b_size: usize) -> Result<f32> {
|
||
if b_size == 0 {
|
||
anyhow::bail!("dqn_replay_step: empty batch (b_size = 0)");
|
||
}
|
||
debug_assert_eq!(self.sampled_h_t_d.len(), b_size * HIDDEN_DIM);
|
||
debug_assert_eq!(self.sampled_h_tp1_d.len(), b_size * HIDDEN_DIM);
|
||
|
||
let k_dqn = N_ACTIONS * Q_N_ATOMS;
|
||
|
||
// Per-iter scratch — small relative to the GPU allocator
|
||
// amortisation, and freed when the function returns.
|
||
let q_loss_d = self.stream.alloc_zeros::<f32>(1)?;
|
||
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)?;
|
||
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 b_size_i = b_size as i32;
|
||
|
||
// ── 1. Forward Q on sampled_h_t and sampled_h_tp1 ───────────
|
||
self.dqn_head
|
||
.forward(&self.sampled_h_t_d, b_size, &mut self.q_logits_d)
|
||
.context("dqn_replay_step: dqn_head.forward(sampled_h_t)")?;
|
||
self.dqn_head
|
||
.forward(&self.sampled_h_tp1_d, b_size, &mut self.q_logits_tp1_d)
|
||
.context("dqn_replay_step: dqn_head.forward(sampled_h_tp1)")?;
|
||
|
||
// ── 1b. IQN forward + loss + backward on replay transitions.
|
||
// Online forward on sampled_h_t, target forward on sampled_h_tp1,
|
||
// quantile Huber loss, backward through forward to weight grads,
|
||
// reduce_axis0 + Adam step.
|
||
{
|
||
let n_tau = self.iqn_head.n_tau();
|
||
let tau_len = b_size * n_tau;
|
||
|
||
// Sample online tau on device.
|
||
{
|
||
let b_i = b_size as i32;
|
||
let n_tau_i = n_tau as i32;
|
||
let tau_cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (n_tau as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut tau_launch = self.stream.launch_builder(&self.rl_sample_tau_fn);
|
||
tau_launch
|
||
.arg(&mut self.iqn_prng_state_d)
|
||
.arg(&mut self.iqn_tau_d)
|
||
.arg(&b_i)
|
||
.arg(&n_tau_i);
|
||
unsafe {
|
||
tau_launch
|
||
.launch(tau_cfg)
|
||
.context("dqn_replay_step: rl_sample_tau online")?;
|
||
}
|
||
}
|
||
|
||
// Online forward on sampled_h_t.
|
||
self.iqn_head
|
||
.forward(
|
||
&self.sampled_h_t_d,
|
||
&self.iqn_tau_d,
|
||
b_size,
|
||
n_tau,
|
||
&mut self.iqn_q_values_d,
|
||
)
|
||
.context("dqn_replay_step: iqn_head.forward(sampled_h_t)")?;
|
||
self.iqn_head
|
||
.expected_q(
|
||
&self.iqn_q_values_d,
|
||
b_size,
|
||
n_tau,
|
||
&mut self.iqn_expected_q_d,
|
||
)
|
||
.context("dqn_replay_step: iqn_head.expected_q")?;
|
||
|
||
// Sample target tau on device and forward target network on sampled_h_tp1.
|
||
let mut iqn_tau_target_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(tau_len)
|
||
.context("dqn_replay_step: alloc iqn_tau_target")?;
|
||
{
|
||
let b_i = b_size as i32;
|
||
let n_tau_i = n_tau as i32;
|
||
let tau_cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (n_tau as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut tau_launch = self.stream.launch_builder(&self.rl_sample_tau_fn);
|
||
tau_launch
|
||
.arg(&mut self.iqn_prng_state_d)
|
||
.arg(&mut iqn_tau_target_d)
|
||
.arg(&b_i)
|
||
.arg(&n_tau_i);
|
||
unsafe {
|
||
tau_launch
|
||
.launch(tau_cfg)
|
||
.context("dqn_replay_step: rl_sample_tau target")?;
|
||
}
|
||
}
|
||
|
||
let mut iqn_target_q_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * n_tau * N_ACTIONS)
|
||
.context("dqn_replay_step: alloc iqn_target_q")?;
|
||
self.iqn_head
|
||
.forward_target(
|
||
&self.sampled_h_tp1_d,
|
||
&iqn_tau_target_d,
|
||
b_size,
|
||
n_tau,
|
||
&mut iqn_target_q_d,
|
||
)
|
||
.context("dqn_replay_step: iqn_head.forward_target(sampled_h_tp1)")?;
|
||
|
||
// Quantile Huber loss + gradient w.r.t. online Q values.
|
||
let mut iqn_loss_pb_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size)
|
||
.context("dqn_replay_step: alloc iqn_loss_pb")?;
|
||
let mut iqn_grad_q_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * n_tau * N_ACTIONS)
|
||
.context("dqn_replay_step: alloc iqn_grad_q")?;
|
||
self.iqn_head
|
||
.compute_loss(
|
||
&self.iqn_q_values_d,
|
||
&iqn_target_q_d,
|
||
&self.iqn_tau_d,
|
||
&self.sampled_actions_d,
|
||
b_size,
|
||
n_tau,
|
||
n_tau,
|
||
&mut iqn_loss_pb_d,
|
||
&mut iqn_grad_q_d,
|
||
)
|
||
.context("dqn_replay_step: iqn_head.compute_loss")?;
|
||
|
||
// Backward through IQN forward: grad_q → grad_w_out/b_out/w_embed/b_embed.
|
||
let k_w_out = HIDDEN_DIM * N_ACTIONS;
|
||
let k_w_embed = EMBED_DIM * HIDDEN_DIM;
|
||
let mut iqn_grad_w_out_pb = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * k_w_out)
|
||
.context("dqn_replay_step: alloc iqn_grad_w_out_pb")?;
|
||
let mut iqn_grad_b_out_pb = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * N_ACTIONS)
|
||
.context("dqn_replay_step: alloc iqn_grad_b_out_pb")?;
|
||
let mut iqn_grad_w_embed_pb = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * k_w_embed)
|
||
.context("dqn_replay_step: alloc iqn_grad_w_embed_pb")?;
|
||
let mut iqn_grad_b_embed_pb = self
|
||
.stream
|
||
.alloc_zeros::<f32>(b_size * HIDDEN_DIM)
|
||
.context("dqn_replay_step: alloc iqn_grad_b_embed_pb")?;
|
||
|
||
self.iqn_head
|
||
.backward(
|
||
&self.sampled_h_t_d,
|
||
&self.iqn_tau_d,
|
||
&iqn_grad_q_d,
|
||
b_size,
|
||
n_tau,
|
||
&mut iqn_grad_w_out_pb,
|
||
&mut iqn_grad_b_out_pb,
|
||
&mut iqn_grad_w_embed_pb,
|
||
&mut iqn_grad_b_embed_pb,
|
||
)
|
||
.context("dqn_replay_step: iqn_head.backward")?;
|
||
|
||
// Reduce across batches.
|
||
let mut iqn_grad_w_out_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(k_w_out)
|
||
.context("dqn_replay_step: alloc iqn_grad_w_out")?;
|
||
let mut iqn_grad_b_out_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(N_ACTIONS)
|
||
.context("dqn_replay_step: alloc iqn_grad_b_out")?;
|
||
let mut iqn_grad_w_embed_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(k_w_embed)
|
||
.context("dqn_replay_step: alloc iqn_grad_w_embed")?;
|
||
let mut iqn_grad_b_embed_d = self
|
||
.stream
|
||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||
.context("dqn_replay_step: alloc iqn_grad_b_embed")?;
|
||
|
||
self.launch_reduce_axis0(
|
||
&iqn_grad_w_out_pb, b_size, k_w_out, &mut iqn_grad_w_out_d,
|
||
)?;
|
||
self.launch_reduce_axis0(
|
||
&iqn_grad_b_out_pb, b_size, N_ACTIONS, &mut iqn_grad_b_out_d,
|
||
)?;
|
||
self.launch_reduce_axis0(
|
||
&iqn_grad_w_embed_pb, b_size, k_w_embed, &mut iqn_grad_w_embed_d,
|
||
)?;
|
||
self.launch_reduce_axis0(
|
||
&iqn_grad_b_embed_pb, b_size, HIDDEN_DIM, &mut iqn_grad_b_embed_d,
|
||
)?;
|
||
|
||
// IQN Adam steps (LR from same ISV[RL_IQN_LR_INDEX] as the Q head).
|
||
self.iqn_w_out_adam
|
||
.step(&mut self.iqn_head.w_out_d, &iqn_grad_w_out_d)
|
||
.context("dqn_replay_step: iqn_w_out_adam.step")?;
|
||
self.iqn_b_out_adam
|
||
.step(&mut self.iqn_head.b_out_d, &iqn_grad_b_out_d)
|
||
.context("dqn_replay_step: iqn_b_out_adam.step")?;
|
||
self.iqn_w_embed_adam
|
||
.step(&mut self.iqn_head.w_embed_d, &iqn_grad_w_embed_d)
|
||
.context("dqn_replay_step: iqn_w_embed_adam.step")?;
|
||
self.iqn_b_embed_adam
|
||
.step(&mut self.iqn_head.b_embed_d, &iqn_grad_b_embed_d)
|
||
.context("dqn_replay_step: iqn_b_embed_adam.step")?;
|
||
}
|
||
|
||
// ── 2. Double-DQN argmax on online Q at h_tp1 ───────────────
|
||
{
|
||
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 mut launch = self.stream.launch_builder(&self.argmax_expected_q_fn);
|
||
launch
|
||
.arg(&self.q_logits_tp1_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&mut self.sampled_next_actions_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_argmax)
|
||
.context("dqn_replay_step: argmax_expected_q launch")?;
|
||
}
|
||
}
|
||
|
||
// ── 3. Bellman target build via TARGET net at h_tp1 ─────────
|
||
self.dqn_head
|
||
.forward_target(&self.sampled_h_tp1_d, b_size, &mut q_target_full_d)
|
||
.context("dqn_replay_step: dqn_head.forward_target(sampled_h_tp1)")?;
|
||
self.dqn_head
|
||
.select_action_atoms(
|
||
&q_target_full_d,
|
||
&self.sampled_next_actions_d,
|
||
b_size,
|
||
&mut q_target_action_d,
|
||
)
|
||
.context("dqn_replay_step: dqn_head.select_action_atoms")?;
|
||
self.dqn_head
|
||
.project_bellman_target(
|
||
&q_target_action_d,
|
||
&self.sampled_rewards_d,
|
||
&self.sampled_dones_d,
|
||
&self.sampled_n_step_gammas_d,
|
||
&self.isv_d,
|
||
b_size,
|
||
&mut target_dist_d,
|
||
)
|
||
.context("dqn_replay_step: dqn_head.project_bellman_target")?;
|
||
|
||
// ── 4. Q backward (logits → grad_w/b/h_t) ───────────────────
|
||
let mut q_loss_d_mut = q_loss_d;
|
||
self.dqn_head
|
||
.backward_logits(
|
||
&self.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_replay_step: dqn_head.backward_logits")?;
|
||
let l_q_host = read_scalar_d(&self.stream, &q_loss_d_mut)?;
|
||
let l_q = 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_replay_step: dqn_head.backward_to_w_b_h")?;
|
||
// R7d stop-grad: discard q_grad_h_t (sampled h_t is past-step
|
||
// encoder output; folding its gradient into the encoder would
|
||
// poison live training). Same semantics as step_synthetic.
|
||
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)?;
|
||
|
||
// ── 5. Q Adam (uses LR set by step_synthetic; we don't re-fire
|
||
// the LR controller here — that runs once per env step).
|
||
self.dqn_w_adam
|
||
.step(&mut self.dqn_head.w_d, &q_grad_w_d)
|
||
.context("dqn_replay_step: dqn_w_adam.step")?;
|
||
self.dqn_b_adam
|
||
.step(&mut self.dqn_head.b_d, &q_grad_b_d)
|
||
.context("dqn_replay_step: dqn_b_adam.step")?;
|
||
|
||
Ok(l_q)
|
||
}
|
||
|
||
/// 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 0: bump device-resident step counter (ISV[548]).
|
||
// Must run BEFORE any kernel that reads current_step from ISV.
|
||
// Single thread, single block — graph-safe (no scalar args change).
|
||
{
|
||
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_increment_step_fn);
|
||
launch.arg(&self.isv_d);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_increment_step launch")?;
|
||
}
|
||
}
|
||
|
||
// ── 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 (tc_ptr, _g1) = self.trade_context_d.device_ptr(&self.stream);
|
||
let (mr_ptr, _g2) = self.multires_output_d.device_ptr(&self.stream);
|
||
self.perception.encoder_ctx_trade_ptr = tc_ptr;
|
||
self.perception.encoder_ctx_multires_ptr = mr_ptr;
|
||
}
|
||
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);
|
||
|
||
// ── SP20 P3: FRD head forward — reads h_t, writes 3 horizons ×
|
||
// 21 atom logits per batch into self.frd_logits_d (consumed by
|
||
// the FRD gate in P9 and the supervised CE loss in F.3). The
|
||
// post-ReLU hidden activations land in self.frd_hidden_d for
|
||
// the bwd kernel to consume without recomputing the W1 product.
|
||
self.frd_head
|
||
.forward(
|
||
h_t_borrow,
|
||
&mut self.frd_hidden_d,
|
||
&mut self.frd_logits_d,
|
||
b_size,
|
||
)
|
||
.context("frd_head.forward in step_with_lobsim")?;
|
||
|
||
self.dqn_head
|
||
.forward(h_t_borrow, b_size, &mut self.q_logits_d)
|
||
.context("step_with_lobsim: dqn_head.forward(h_t)")?;
|
||
self.dqn_head
|
||
.forward(&self.h_tp1_d, b_size, &mut self.q_logits_tp1_d)
|
||
.context("step_with_lobsim: dqn_head.forward(h_tp1) for Double-DQN argmax")?;
|
||
self.value_head
|
||
.forward(h_t_borrow, &self.isv_d, b_size, &mut self.v_pred_d)
|
||
.context("step_with_lobsim: value_head.forward(h_t)")?;
|
||
self.value_head
|
||
.forward(&self.h_tp1_d, &self.isv_d, b_size, &mut self.v_pred_tp1_d)
|
||
.context("step_with_lobsim: value_head.forward(h_tp1) for true V(s_{t+1})")?;
|
||
|
||
// ── IQN forward alongside C51 ────────────────────────────────
|
||
// Sample tau ~ U(0,1) on device via rl_sample_tau kernel
|
||
// (graph-safe, no host RNG), then forward + expected_q.
|
||
{
|
||
let n_tau = self.iqn_head.n_tau();
|
||
let b_i = b_size as i32;
|
||
let n_tau_i = n_tau as i32;
|
||
let tau_cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (n_tau as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut tau_launch = self.stream.launch_builder(&self.rl_sample_tau_fn);
|
||
tau_launch
|
||
.arg(&mut self.iqn_prng_state_d)
|
||
.arg(&mut self.iqn_tau_d)
|
||
.arg(&b_i)
|
||
.arg(&n_tau_i);
|
||
unsafe {
|
||
tau_launch
|
||
.launch(tau_cfg)
|
||
.context("step_with_lobsim: rl_sample_tau launch")?;
|
||
}
|
||
self.iqn_head
|
||
.forward(
|
||
h_t_borrow,
|
||
&self.iqn_tau_d,
|
||
b_size,
|
||
n_tau,
|
||
&mut self.iqn_q_values_d,
|
||
)
|
||
.context("step_with_lobsim: iqn_head.forward(h_t)")?;
|
||
self.iqn_head
|
||
.expected_q(
|
||
&self.iqn_q_values_d,
|
||
b_size,
|
||
n_tau,
|
||
&mut self.iqn_expected_q_d,
|
||
)
|
||
.context("step_with_lobsim: iqn_head.expected_q")?;
|
||
}
|
||
|
||
// ── Ensemble action-value: α × E_C51 + (1-α) × E_IQN ────────
|
||
// Launched after both C51 and IQN forwards complete. The ensemble
|
||
// Q feeds the Q→π agreement diagnostic below.
|
||
{
|
||
let cfg_ensemble = 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_ensemble_action_value_fn);
|
||
launch
|
||
.arg(&self.q_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&self.iqn_expected_q_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&mut self.ensemble_q_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_ensemble)
|
||
.context("step_with_lobsim: rl_ensemble_action_value launch")?;
|
||
}
|
||
}
|
||
|
||
// ── NoisyNet state-dependent exploration on ensemble Q ─────────
|
||
// Resample factored noise, forward h_t through the NoisyLinear
|
||
// layer → noisy_output_d [B, N_ACTIONS], then fold into
|
||
// ensemble_q_d in-place via aux_vec_add_inplace. The perturbed
|
||
// ensemble Q feeds the Q→π agreement diagnostic and distillation
|
||
// gradient, so the policy indirectly sees noisy Q preferences
|
||
// and learns to explore without a fixed P_MIN floor.
|
||
//
|
||
// Target network path (Bellman argmax / bellman_target_projection)
|
||
// does NOT see noise — it reads the un-perturbed C51/IQN target
|
||
// weights. This matches the NoisyNet paper's prescription:
|
||
// noise on the online network only during rollout.
|
||
self.noisy_exploration
|
||
.resample_noise()
|
||
.context("step_with_lobsim: noisy_exploration.resample_noise")?;
|
||
self.noisy_exploration
|
||
.forward(h_t_borrow, b_size, &mut self.noisy_output_d)
|
||
.context("step_with_lobsim: noisy_exploration.forward")?;
|
||
{
|
||
let n_elems = (b_size * N_ACTIONS) as i32;
|
||
let cfg_add = LaunchConfig {
|
||
grid_dim: (((n_elems as u32) + 255) / 256, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.aux_vec_add_fn);
|
||
launch
|
||
.arg(&mut self.ensemble_q_d)
|
||
.arg(&self.noisy_output_d)
|
||
.arg(&n_elems);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_add)
|
||
.context("step_with_lobsim: aux_vec_add_inplace (noisy → ensemble_q)")?;
|
||
}
|
||
}
|
||
|
||
// ── Step 2b: Forward π logits for log_pi_old. ─────────────────
|
||
self.policy_head
|
||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||
.context("step_with_lobsim: policy_head.forward_logits")?;
|
||
|
||
// R9 audit — Q-vs-π agreement diag fires into ISV[407] EMA.
|
||
// Needs both q_logits_d (from step 2) and pi_logits_d (just
|
||
// computed). Single-block tree-reduce; block_dim must be
|
||
// power of 2 ≥ b_size for the in-block sum.
|
||
{
|
||
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_q_pi_agree_b_fn);
|
||
launch
|
||
.arg(&self.q_logits_d)
|
||
.arg(&self.pi_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("step_with_lobsim: rl_q_pi_agree_b launch")?;
|
||
}
|
||
}
|
||
|
||
// ── 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);
|
||
|
||
// ── π-driven action selection (Option B fix per
|
||
// `pearl_q_thompson_actor_makes_pi_dead_weight`). Replaces the
|
||
// prior Q-Thompson selector. π is now the actor: actions come
|
||
// from multinomial sample of softmax(pi_logits). Q remains the
|
||
// critic (Bellman target via argmax_expected_q on h_{t+1}).
|
||
// PPO importance-ratio surrogate now has canonical semantics —
|
||
// π_new / π_old at the action that π_old sampled.
|
||
//
|
||
// The rl_action_kernel (Q-Thompson) field is kept loaded for
|
||
// backward-compat tests + diagnostic comparison; it's no
|
||
// longer in the hot path.
|
||
{
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (1, 1, 1), // single thread per block (multinomial CDF walk is sequential)
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_pi_action_kernel_fn);
|
||
launch
|
||
.arg(&self.pi_logits_d)
|
||
.arg(&mut self.prng_state_d)
|
||
.arg(&mut self.actions_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_pi_action_kernel launch (multinomial from softmax(π))")?;
|
||
}
|
||
}
|
||
|
||
// 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(&self.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(&self.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")?;
|
||
}
|
||
}
|
||
|
||
// ── Confidence gate — override opening actions to Hold when C51
|
||
// distributional Q has high uncertainty for the chosen action.
|
||
// Fires AFTER π samples an action but BEFORE trail/heat/market
|
||
// pipeline. Reads pos_state from lobsim (pre-snapshot = current
|
||
// position) to only gate on flat entries.
|
||
{
|
||
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = lobsim.pos_bytes() as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_confidence_gate_fn);
|
||
launch
|
||
.arg(&mut self.actions_d)
|
||
.arg(&self.q_logits_d)
|
||
.arg(&self.atom_supports_d)
|
||
.arg(pos_d_ref)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_confidence_gate launch")?;
|
||
}
|
||
}
|
||
|
||
// FRD gate — same warmup as confidence gate.
|
||
{
|
||
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = lobsim.pos_bytes() as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_frd_gate_fn);
|
||
launch
|
||
.arg(&mut self.actions_d)
|
||
.arg(&self.frd_logits_d)
|
||
.arg(pos_d_ref)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_frd_gate 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;
|
||
|
||
// Session risk circuit breaker — block opening actions when
|
||
// cumulative PnL EMA drops below threshold.
|
||
{
|
||
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 mut launch = self.stream.launch_builder(&self.rl_session_risk_check_fn);
|
||
launch
|
||
.arg(&mut self.actions_d)
|
||
.arg(&self.rewards_d)
|
||
.arg(&self.dones_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_session_risk_check launch")?;
|
||
}
|
||
}
|
||
|
||
// Min hold check — override closing actions to Hold when
|
||
// hold time < ISV-driven minimum. Trail stops still fire
|
||
// (they run AFTER this, safety overrides patience).
|
||
{
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((b_size as u32) + 31) / 32, 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
|
||
let mut launch = self.stream.launch_builder(&self.rl_min_hold_check_fn);
|
||
launch
|
||
.arg(&mut self.actions_d)
|
||
.arg(&self.steps_since_done_d)
|
||
.arg(pos_d_ref)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_min_hold_check launch")?;
|
||
}
|
||
}
|
||
|
||
// Asymmetric trail decay — auto-tighten losers, widen winners.
|
||
// Runs BEFORE trail_mutate so the structural decay is applied
|
||
// first, then agent's a7/a8 fine-tuning on top.
|
||
{
|
||
let bid_px_d = lobsim.bid_px_d();
|
||
let ask_px_d = lobsim.ask_px_d();
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (4, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch =
|
||
self.stream.launch_builder(&self.rl_asymmetric_trail_decay_fn);
|
||
launch
|
||
.arg(&mut self.unit_trail_distance_d)
|
||
.arg(&self.unit_active_d)
|
||
.arg(&self.unit_entry_price_d)
|
||
.arg(&self.unit_initial_r_d)
|
||
.arg(&self.unit_lots_d)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_asymmetric_trail_decay launch")?;
|
||
}
|
||
}
|
||
|
||
// Trail mutate (a7/a8) — BEFORE actions_to_market_targets.
|
||
// Mutates unit_trail_distance for active units if the chosen action
|
||
// is TrailTighten/Loosen. Non-trail actions pass through.
|
||
{
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (4, 1, 1), // MAX_UNITS threads per block
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.rl_trail_mutate_fn);
|
||
launch
|
||
.arg(&self.actions_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&self.unit_active_d)
|
||
.arg(&mut self.unit_trail_distance_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_trail_mutate launch")?;
|
||
}
|
||
}
|
||
|
||
// ── SP20 P1+P5 trail stop check — AFTER trail_mutate, BEFORE
|
||
// actions_to_market_targets. May OVERRIDE actions_d[b] to
|
||
// FlatFromLong/Short on per-unit breach. Reads shared lobsim
|
||
// best book for current mid.
|
||
{
|
||
let bid_px_d = lobsim.bid_px_d();
|
||
let ask_px_d = lobsim.ask_px_d();
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (b_size as u32, 1, 1),
|
||
block_dim: (4, 1, 1), // MAX_UNITS threads per block
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.rl_trail_stop_check_fn);
|
||
launch
|
||
.arg(&mut self.actions_d)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(&self.unit_active_d)
|
||
.arg(&self.unit_entry_price_d)
|
||
.arg(&self.unit_lots_d)
|
||
.arg(&self.unit_trail_distance_d)
|
||
.arg(&self.pyramid_units_count_d)
|
||
.arg(&mut self.close_unit_index_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_trail_stop_check launch")?;
|
||
}
|
||
}
|
||
|
||
// ── SP20 P6 position heat cap — AFTER trail_stop_check, BEFORE
|
||
// actions_to_market_targets. If |position_lots| exceeds the ISV-
|
||
// driven max (slot 504, default 8), OVERRIDE to FlatFromLong/Short.
|
||
// Reads pos_state for aggregate position; writes diag fired-count
|
||
// to ISV slot 505.
|
||
{
|
||
let pos_d_ref_heat: &CudaSlice<u8> = lobsim.pos_d();
|
||
let block = (b_size as u32).min(256);
|
||
let grid = ((b_size as u32) + block - 1) / block;
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (grid.max(1), 1, 1),
|
||
block_dim: (block.max(1), 1, 1),
|
||
shared_mem_bytes: 256 * std::mem::size_of::<i32>() as u32,
|
||
};
|
||
let mut launch =
|
||
self.stream.launch_builder(&self.rl_position_heat_check_fn);
|
||
launch
|
||
.arg(&mut self.actions_d)
|
||
.arg(pos_d_ref_heat)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_position_heat_check launch")?;
|
||
}
|
||
}
|
||
|
||
{
|
||
let (pos_d_ref, bid_px_d, ask_px_d, market_targets_d) =
|
||
lobsim.pos_book_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(&self.isv_d)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(&self.unit_entry_price_d)
|
||
.arg(&self.unit_active_d)
|
||
.arg(&self.unit_lots_d)
|
||
.arg(&self.pyramid_units_count_d)
|
||
.arg(&self.close_unit_index_d)
|
||
.arg(&self.outcome_ema_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")?;
|
||
}
|
||
}
|
||
|
||
// ── SP20 P1 per-unit state machine update — runs AFTER fill
|
||
// + extract_realized_pnl_delta so position transitions are
|
||
// visible. Detects open/close/reverse and populates unit slot
|
||
// 0 with entry_price, entry_step, lots, initial_r, trail.
|
||
// Maintains its own prev_pos tracker (unit_prev_pos_lots_d)
|
||
// separate from extract_realized_pnl_delta's prev_position_lots_d
|
||
// so kernels are independently composable.
|
||
{
|
||
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.rl_unit_state_update_fn);
|
||
launch
|
||
.arg(pos_d_ref)
|
||
.arg(&self.isv_d)
|
||
.arg(&mut self.unit_prev_pos_lots_d)
|
||
.arg(&mut self.unit_entry_price_d)
|
||
.arg(&mut self.unit_entry_step_d)
|
||
.arg(&mut self.unit_lots_d)
|
||
.arg(&mut self.unit_initial_r_d)
|
||
.arg(&mut self.unit_trail_distance_d)
|
||
.arg(&mut self.unit_active_d)
|
||
.arg(&mut self.pyramid_units_count_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_unit_state_update launch")?;
|
||
}
|
||
}
|
||
|
||
// Trade context features — derived from unit state + current mid.
|
||
{
|
||
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
|
||
let bid_px_d = lobsim.bid_px_d();
|
||
let ask_px_d = lobsim.ask_px_d();
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((b_size as u32) + 31) / 32, 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let pos_bytes_i = lobsim.pos_bytes() as i32;
|
||
let mut launch =
|
||
self.stream.launch_builder(&self.rl_trade_context_update_fn);
|
||
launch
|
||
.arg(&mut self.trade_context_d)
|
||
.arg(&self.unit_active_d)
|
||
.arg(&self.unit_entry_price_d)
|
||
.arg(&self.unit_entry_step_d)
|
||
.arg(&self.unit_initial_r_d)
|
||
.arg(&self.unit_lots_d)
|
||
.arg(pos_d_ref)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_trade_context_update launch")?;
|
||
}
|
||
}
|
||
|
||
// Multi-resolution streaming features — time-weighted EMA at 3 horizons.
|
||
{
|
||
let bid_px_d = lobsim.bid_px_d();
|
||
let ask_px_d = lobsim.ask_px_d();
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((b_size as u32) + 31) / 32, 1, 1),
|
||
block_dim: (32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let b_size_i = b_size as i32;
|
||
let ts_ns: u64 = last_snap.ts_ns;
|
||
let mut launch =
|
||
self.stream.launch_builder(&self.rl_multires_features_update_fn);
|
||
launch
|
||
.arg(&mut self.multires_state_d)
|
||
.arg(&mut self.multires_output_d)
|
||
.arg(&mut self.multires_prev_mid_d)
|
||
.arg(&mut self.multires_prev_ts_ns_d)
|
||
.arg(bid_px_d)
|
||
.arg(ask_px_d)
|
||
.arg(bid_px_d) // bid_sz uses bid_px (sizes at same depth levels)
|
||
.arg(ask_px_d) // ask_sz uses ask_px
|
||
.arg(&self.isv_d)
|
||
.arg(&ts_ns)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_multires_features_update 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")?;
|
||
}
|
||
}
|
||
|
||
// Signed outcome EMA (ISV[508]) — updated on trade close only.
|
||
// Feeds anti-martingale sizing in actions_to_market_targets.
|
||
{
|
||
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_OUTCOME_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.rewards_d)
|
||
.arg(&self.dones_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("ema_update_on_done(OUTCOME_EMA) launch")?;
|
||
}
|
||
}
|
||
|
||
// Per-batch outcome EMA — feeds anti-martingale sizing.
|
||
{
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((b_size as u32) + 31) / 32, 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.rl_recent_outcome_update_fn);
|
||
launch
|
||
.arg(&self.rewards_d)
|
||
.arg(&self.dones_d)
|
||
.arg(&mut self.outcome_ema_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_recent_outcome_update launch")?;
|
||
}
|
||
}
|
||
|
||
// Fire all 10 RL ISV controllers in a single fused kernel
|
||
// launch (gamma, tau, ppo_clip, entropy_coef, rollout_steps,
|
||
// per_alpha, reward_scale, ppo_ratio_clamp, gate_threshold,
|
||
// q_distill_lambda). Critical: this happens BEFORE
|
||
// apply_reward_scale so the scale ISV[406] reflects this
|
||
// step's mean_abs_pnl EMA update.
|
||
self.launch_rl_fused_controllers(b_size)
|
||
.context("launch_rl_fused_controllers")?;
|
||
|
||
// Apply the reward scale on device (in place on rewards_d) +
|
||
// adaptive clamp + per-step pre-clamp diag + per-step
|
||
// positive-tail max for the adaptive clamp controller.
|
||
// Reads ISV[406] (scale) + ISV[452]/ISV[453] (current clamp
|
||
// bounds, refreshed by rl_reward_clamp_controller via the
|
||
// helper below); writes ISV[439] (diag) + ISV[478] (positive
|
||
// tail input for the controller).
|
||
//
|
||
// 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. Shared mem now 2× block ×
|
||
// f32 to fit the parallel max(|scaled|) + max(positive)
|
||
// reductions.
|
||
// Reward shaping: entry cost + short-hold penalty + hold bonus.
|
||
// Runs on raw rewards BEFORE scaling so costs are in USD terms.
|
||
{
|
||
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 b_size_i = b_size as i32;
|
||
let pos_bytes_i = lobsim.pos_bytes() as i32;
|
||
let mut launch = self.stream.launch_builder(&self.rl_reward_shaping_fn);
|
||
launch
|
||
.arg(&mut self.rewards_d)
|
||
.arg(&self.prev_position_lots_d)
|
||
.arg(pos_d_ref)
|
||
.arg(&self.dones_d)
|
||
.arg(&self.steps_since_done_d)
|
||
.arg(&self.isv_d)
|
||
.arg(&b_size_i)
|
||
.arg(&pos_bytes_i);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg)
|
||
.context("rl_reward_shaping launch")?;
|
||
}
|
||
}
|
||
|
||
// Snapshot raw rewards before scaling — replay buffer stores
|
||
// raw_reward for re-normalization at sample time.
|
||
{
|
||
let nbytes = (b_size * std::mem::size_of::<f32>()) as u64;
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (src, _g1) = self.rewards_d.device_ptr(&self.stream);
|
||
let (dst, _g2) = self.raw_rewards_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes as usize, s)
|
||
.context("snapshot raw_rewards_d before scale")?;
|
||
}
|
||
}
|
||
{
|
||
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: 3 * 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")?;
|
||
}
|
||
}
|
||
|
||
// Adaptive reward-clamp controller — refresh slots 452/453 from
|
||
// the just-published positive-tail max in slot 478. The new
|
||
// bounds take effect on the NEXT step's apply_reward_scale.
|
||
{
|
||
let cfg_ctrl = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let alpha = RL_LR_CONTROLLER_ALPHA;
|
||
let mut launch = self
|
||
.stream
|
||
.launch_builder(&self.rl_reward_clamp_controller_fn);
|
||
launch.arg(&self.isv_d).arg(&alpha);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_ctrl)
|
||
.context("rl_reward_clamp_controller launch (step_with_lobsim path)")?;
|
||
}
|
||
|
||
// C51 atom support refresh from updated V_MIN/V_MAX.
|
||
let cfg_atom = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (21, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self
|
||
.stream
|
||
.launch_builder(&self.rl_atom_support_update_fn);
|
||
launch.arg(&self.isv_d).arg(&self.atom_supports_d);
|
||
unsafe {
|
||
launch
|
||
.launch(cfg_atom)
|
||
.context("rl_atom_support_update launch (step_with_lobsim path)")?;
|
||
}
|
||
}
|
||
|
||
// 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(&self.v_pred_d)
|
||
.arg(&self.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")?;
|
||
|
||
// ── Step 7b: K-loop training intensification driven by
|
||
// `n_rollout_steps` (ISV[404]). Wires the previously-write-only
|
||
// controller to actual training behavior.
|
||
//
|
||
// The controller widens its emitted value when advantages are
|
||
// noisy (var/|mean| high → "need more samples to learn"). We
|
||
// map that to the number of training iterations performed per
|
||
// env step: K = clamp(isv[404] / K_DIVISOR, 1, K_MAX). Each
|
||
// iteration is a fresh PER sample + full step_synthetic call
|
||
// (Q + π + V forward, backward, Adam) + PER priority update.
|
||
//
|
||
// * isv[404] = 256 (MIN) → K = 1 (single update per env step, current behavior)
|
||
// * isv[404] = 2048 (BOOTSTRAP) → K = 2
|
||
// * isv[404] = 8192 (MAX) → K = 8
|
||
//
|
||
// Adapts the training:env ratio so noisy regimes get more
|
||
// gradient samples without slowing down env stepping. b_size=1
|
||
// means each iteration is cheap; the K-multiplier directly
|
||
// addresses the per-step gradient signal starvation that left
|
||
// l_q stuck at 2.82 (uniform=3.04, only 7% information gain)
|
||
// in alpha-rl-frt7s.
|
||
// K-loop config — both divisor and max are ISV-driven per
|
||
// `feedback_isv_for_adaptive_bounds` so they're discoverable
|
||
// in diag and tunable at runtime by re-launching
|
||
// rl_streaming_clamp_init. Defaults seeded at trainer init:
|
||
// divisor=2048 (matches ROLLOUT_BOOTSTRAP so K=1 at controller
|
||
// bootstrap), max=4 (prevents gradient overtraining at
|
||
// b_size=1).
|
||
let n_rollout_steps =
|
||
self.isv_host[crate::rl::isv_slots::RL_N_ROLLOUT_STEPS_INDEX];
|
||
let k_divisor = self
|
||
.isv_host[crate::rl::isv_slots::RL_K_LOOP_DIVISOR_INDEX]
|
||
.max(1.0);
|
||
let k_max = self
|
||
.isv_host[crate::rl::isv_slots::RL_K_LOOP_MAX_INDEX]
|
||
.max(1.0) as usize;
|
||
let k_updates = ((n_rollout_steps / k_divisor).round() as usize)
|
||
.clamp(1, k_max);
|
||
self.last_k_updates = k_updates;
|
||
|
||
// K-loop split (Path B): first iter does the full env-step-
|
||
// paired Q+π+V+encoder update via step_synthetic. Subsequent
|
||
// iters do DQN-ONLY replay via dqn_replay_step — fresh PER
|
||
// sample + Q forward + Bellman + Q backward + Q Adam. π+V+
|
||
// encoder updates stay at 1× per env step regardless of K, so
|
||
// the K-loop intensifies Q learning without overshooting PPO
|
||
// (avoids the f2ggr KL=12.44 pathology where K=8 ran 8× PPO
|
||
// updates per env-step). At b_size=16 the K-loop typically
|
||
// settles at K=1 (advantage_var_ratio drops with batch size),
|
||
// but the split protects against pathological regimes and
|
||
// makes the actor/critic separation explicit.
|
||
let mut stats = IntegratedStepStats::default();
|
||
for k_iter in 0..k_updates {
|
||
// Fresh PER sample per iteration — different transitions
|
||
// each gradient step.
|
||
let per_indices = self
|
||
.sample_and_gather(b_size)
|
||
.context("step_with_lobsim: sample_and_gather (k-loop)")?;
|
||
if k_iter == 0 {
|
||
// First iter: full step_synthetic (env-step-paired).
|
||
stats = self.step_synthetic(snapshots)?;
|
||
} else {
|
||
// Subsequent iters: Q-only replay (no π, no V, no
|
||
// encoder backward, no LR controller refresh).
|
||
let _l_q_extra = self
|
||
.dqn_replay_step(b_size)
|
||
.context("step_with_lobsim: dqn_replay_step (k-loop iter > 0)")?;
|
||
}
|
||
// PER priority update per iteration so the next sample
|
||
// picks transitions with updated TD errors.
|
||
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 (k-loop)")?;
|
||
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};
|
||
|
||
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 raw_rewards_host = read_slice_d(&self.stream, &self.raw_rewards_d, b_size)
|
||
.context("push_to_replay: read raw_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")?;
|
||
|
||
let n_step = self.isv_host[crate::rl::isv_slots::RL_N_STEP_INDEX] as usize;
|
||
let n_step = n_step.max(1);
|
||
let gamma = self.isv_host[crate::rl::isv_slots::RL_GAMMA_INDEX];
|
||
|
||
let hidden_nbytes = HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
for b in 0..b_size {
|
||
// Alloc h_t for this step's n-step entry.
|
||
let mut h_t_per_b = self
|
||
.stream
|
||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||
.context("push_to_replay: alloc h_t per-batch")?;
|
||
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")?;
|
||
}
|
||
|
||
let is_done = dones_host[b] > 0.5;
|
||
self.n_step_buffer[b].push(NStepEntry {
|
||
h_t: h_t_per_b,
|
||
action: actions_host[b] as u32,
|
||
raw_reward: raw_rewards_host[b],
|
||
scaled_reward: rewards_host[b],
|
||
done: is_done,
|
||
log_pi_old: log_pi_old_host[b],
|
||
});
|
||
|
||
// Flush when buffer reaches n OR a done event truncates.
|
||
let should_flush = self.n_step_buffer[b].len() >= n_step || is_done;
|
||
if !should_flush {
|
||
continue;
|
||
}
|
||
|
||
let buf = &self.n_step_buffer[b];
|
||
let buf_len = buf.len();
|
||
|
||
// Compute n-step discounted return: R_n = Σ γᵏ rₖ
|
||
let mut r_n_scaled = 0.0_f32;
|
||
let mut r_n_raw = 0.0_f32;
|
||
let mut gamma_acc = 1.0_f32;
|
||
let mut any_done = false;
|
||
for entry in buf.iter() {
|
||
r_n_scaled += gamma_acc * entry.scaled_reward;
|
||
r_n_raw += gamma_acc * entry.raw_reward;
|
||
if entry.done {
|
||
any_done = true;
|
||
}
|
||
gamma_acc *= gamma;
|
||
}
|
||
|
||
// γⁿ for the bootstrap term. Zero if any done in window
|
||
// (terminal state zeroes the bootstrap).
|
||
let n_step_gamma_val = if any_done { 0.0 } else { gamma_acc };
|
||
|
||
// h_t from oldest entry, h_tp1 from current step.
|
||
let oldest = &buf[0];
|
||
let mut h_tp1_per_b = self
|
||
.stream
|
||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||
.context("push_to_replay: alloc h_tp1 per-batch")?;
|
||
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")?;
|
||
}
|
||
|
||
// Clone oldest h_t (it's about to be consumed by the drain).
|
||
let mut h_t_oldest = self
|
||
.stream
|
||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||
.context("push_to_replay: alloc h_t_oldest")?;
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (src, _g1) = oldest.h_t.device_ptr(&self.stream);
|
||
let (dst, _gd) = h_t_oldest.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst, src, hidden_nbytes, s,
|
||
)
|
||
.context("push_to_replay: DtoD clone h_t_oldest")?;
|
||
}
|
||
let oldest_action = oldest.action;
|
||
let oldest_log_pi = oldest.log_pi_old;
|
||
|
||
// Drain the buffer.
|
||
self.n_step_buffer[b].clear();
|
||
|
||
self.replay.push(Transition {
|
||
h_t: h_t_oldest,
|
||
action: oldest_action,
|
||
reward: r_n_scaled,
|
||
raw_reward: r_n_raw,
|
||
next_h_t: h_tp1_per_b,
|
||
done: any_done,
|
||
n_step_gamma: n_step_gamma_val,
|
||
log_pi_old: oldest_log_pi,
|
||
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];
|
||
let mut n_step_gammas_host = vec![0.0f32; b_size];
|
||
let current_scale = self.isv_host[crate::rl::isv_slots::RL_REWARD_SCALE_INDEX];
|
||
for (i, &idx) in indices.iter().enumerate() {
|
||
let t = &self.replay.transitions[idx];
|
||
actions_host[i] = t.action as i32;
|
||
rewards_host[i] = t.raw_reward * current_scale;
|
||
dones_host[i] = if t.done { 1.0 } else { 0.0 };
|
||
n_step_gammas_host[i] = t.n_step_gamma;
|
||
}
|
||
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")?;
|
||
write_slice_f32_d(&self.stream, &n_step_gammas_host, &mut self.sampled_n_step_gammas_d)
|
||
.context("sample_and_gather: write sampled_n_step_gammas_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)
|
||
}
|
||
|
||
/// Public mapped-pinned write helper (`&[f32]` → device).
|
||
/// Per `feedback_no_htod_htoh_only_mapped_pinned`: regular un-pinned
|
||
/// host→device transfers are forbidden because the host slice isn't
|
||
/// page-locked. Use this for tests and any other caller outside this
|
||
/// module that needs a host→device write.
|
||
pub fn write_slice_f32_d_pub(
|
||
stream: &Arc<CudaStream>,
|
||
src: &[f32],
|
||
dst: &mut CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
write_slice_f32_d(stream, src, dst)
|
||
}
|
||
|
||
/// Public mapped-pinned write helper (`&[i32]` → device).
|
||
/// See `write_slice_f32_d_pub`.
|
||
pub fn write_slice_i32_d_pub(
|
||
stream: &Arc<CudaStream>,
|
||
src: &[i32],
|
||
dst: &mut CudaSlice<i32>,
|
||
) -> Result<()> {
|
||
write_slice_i32_d(stream, src, dst)
|
||
}
|
||
|
||
/// Public mapped-pinned write helper for raw byte buffers
|
||
/// (`&[u8]` → device). Used for `PosFlat`-encoded test fixtures.
|
||
/// Stages through a `MappedI32Buffer` (4-byte alignment), so `src.len()`
|
||
/// MUST be a multiple of 4. The byte→i32 reinterpretation is
|
||
/// little-endian on x86_64 + L40S/H100 — matches the device's int
|
||
/// load semantics for the same byte sequence.
|
||
pub fn write_slice_u8_d_pub(
|
||
stream: &Arc<CudaStream>,
|
||
src: &[u8],
|
||
dst: &mut CudaSlice<u8>,
|
||
) -> Result<()> {
|
||
debug_assert!(dst.len() >= src.len());
|
||
let n_bytes = src.len();
|
||
if n_bytes == 0 {
|
||
return Ok(());
|
||
}
|
||
assert_eq!(
|
||
n_bytes % 4,
|
||
0,
|
||
"write_slice_u8_d_pub requires 4-byte-aligned length (got {n_bytes})"
|
||
);
|
||
let n_i32 = n_bytes / 4;
|
||
let mut staging = unsafe { crate::pinned_mem::MappedI32Buffer::new(n_i32) }
|
||
.map_err(|e| anyhow::anyhow!("write_slice_u8_d_pub staging: {e}"))?;
|
||
// Reinterpret &[u8] as &[i32] for the staging copy.
|
||
let host_i32 = staging.host_slice_mut();
|
||
for (i, chunk) in src.chunks_exact(4).enumerate() {
|
||
host_i32[i] = i32::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
|
||
}
|
||
unsafe {
|
||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr,
|
||
staging.dev_ptr,
|
||
n_bytes,
|
||
stream.cu_stream(),
|
||
)
|
||
.context("write_slice_u8_d_pub DtoD")?;
|
||
}
|
||
stream.synchronize().context("write_slice_u8_d_pub sync")?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Public mapped-pinned read helper for raw byte buffers
|
||
/// (device → `Vec<u8>`). Staging via `MappedI32Buffer`; `n_bytes`
|
||
/// MUST be a multiple of 4. Counterpart to `write_slice_u8_d_pub`.
|
||
pub fn read_slice_u8_d_pub(
|
||
stream: &Arc<CudaStream>,
|
||
src: &CudaSlice<u8>,
|
||
n_bytes: usize,
|
||
) -> Result<Vec<u8>> {
|
||
debug_assert!(src.len() >= n_bytes);
|
||
if n_bytes == 0 {
|
||
return Ok(Vec::new());
|
||
}
|
||
assert_eq!(
|
||
n_bytes % 4,
|
||
0,
|
||
"read_slice_u8_d_pub requires 4-byte-aligned length (got {n_bytes})"
|
||
);
|
||
let n_i32 = n_bytes / 4;
|
||
let mut staging = unsafe { crate::pinned_mem::MappedI32Buffer::new(n_i32) }
|
||
.map_err(|e| anyhow::anyhow!("read_slice_u8_d_pub staging: {e}"))?;
|
||
unsafe {
|
||
let (src_ptr, _g) = src.device_ptr(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
staging.dev_ptr,
|
||
src_ptr,
|
||
n_bytes,
|
||
stream.cu_stream(),
|
||
)
|
||
.context("read_slice_u8_d_pub DtoD")?;
|
||
}
|
||
stream.synchronize().context("read_slice_u8_d_pub sync")?;
|
||
let mut out = Vec::with_capacity(n_bytes);
|
||
for word in staging.host_slice_mut().iter() {
|
||
out.extend_from_slice(&word.to_ne_bytes());
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
|
||
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.
|
||
|