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

feat(rl): LobEnv trait + step_with_lobsim + toy bandit activation (E.3b)

Browse Source 
Closes Phase E of the integrated RL trainer plan
(docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md). The
integrated trainer can now be exercised end-to-end on a real (or mock)
LobSim environment.

What this commit lands:
- LobEnv trait in src/rl/reward.rs — narrow contract over apply_snapshot
  + submit_action + step_event. Chosen over a direct
  ml-backtesting → ml-alpha dep because ml-backtesting already depends
  on ml-alpha (loader + trunk reuse); a reverse direct dep would cycle.
  The simulator-side `impl LobEnv for LobSimCuda` completes the wire
  from ml-backtesting in a follow-up (the trait surface is intentionally
  small and lobsim-agnostic so other env implementations can plug in).
- MockLobEnv in the same module — deterministic toy bandit
  (action 5 → +1, else → -1, done = true every step). Powers the
  dqn_toy / ppo_toy / integrated_trainer_smoke gate tests.
- IntegratedTrainer::step_with_lobsim — one training step driven by a
  real LobEnv. Forwards encoder, forwards Q + V + π, reads logits to
  host, Thompson-samples action per batch
  (pearl_thompson_for_distributional_action_selection), drives the env,
  derives real reward / advantages / returns / log_pi_old, and delegates
  to step_synthetic for the per-head backward + Adam + encoder backward
  (single source of truth per feedback_single_source_of_truth_no_duplicates).
- IntegratedTrainer::eval_expected_q_per_action +
  IntegratedTrainer::eval_policy_probs_per_action — public eval helpers
  the gate tests use to inspect the trained Q-distribution / policy
  without re-implementing device readback in test crates.
- dqn_toy / ppo_toy / integrated_trainer_smoke — bodies filled with the
  real training loop driven by MockLobEnv::toy_bandit. The convergence
  gates assert argmax_a E[Q(s, a)] == LongSmall and mode π(s) == LongSmall
  after 300 step_with_lobsim calls. Tests are #[ignore]-gated for CUDA
  availability per the project's test discipline; activate via
  `cargo test -p ml-alpha --test dqn_toy -- --ignored` on a CUDA host.
- rl_rollout_steps_controller.cu — ISV[404] producer. Rollout length
  adapts to var(advantage)/|mean A|: noisy advantages → grow rollout
  (more samples per PPO update), stable → shrink (fresher data).
  Bootstrap 2048 (PPO default), bounds [256, 8192], Wiener-α blend with
  floor 0.4 per pearl_wiener_alpha_floor_for_nonstationary.
- rl_per_alpha_controller.cu — ISV[405] producer. PER priority exponent
  α adapts to TD-error kurtosis EMA: heavy tails → raise α to concentrate
  on the informative tails, light tails → keep α near the canonical 0.6.
  Bootstrap 0.6 (Schaul 2016), bounds [0.3, 1.0].

Both controllers honour pearl_first_observation_bootstrap (sentinel
0.0 → first-emit writes bootstrap, subsequent emits Wiener-α blend) and
pearl_controller_anchors_isv_driven (no hardcoded constants — every
adaptive hyperparameter sourced from ISV).

build.rs registers both kernels and bumps cache-bust v23 → v24.

Phase F follows with the reward-shaping ISV controller (RL_REWARD_SCALE_INDEX=406)
+ per-trade PnL extraction calibration. Phase G adds the Argo workflow
+ dispatcher. Phase H runs the actual training + backtest smoke and tests
the G8 gate (profit_factor > 1.0).

Verification:
- cargo check --workspace --lib clean (no warnings)
- cargo test -p ml-alpha --lib: 66 passed (was 63; +3 mock_bandit_*
  tests in rl::reward), 0 failed, 6 ignored
- cargo build -p ml-alpha --test dqn_toy --test ppo_toy
  --test integrated_trainer_smoke clean
- integrated_trainer_loss_lambdas_default_equal_weight (non-ignored)
  still passes
This commit is contained in:
jgrusewski
2026-05-23 00:40:09 +02:00
parent 4b5ef093de
commit 9114374d25
9 changed files with 1184 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -41,9 +41,11 @@ const KERNELS: &[&str] = &[
"grad_h_accumulate", // RL Phase E.2: element-wise grad_h_encoder += λ × grad_h_head accumulator (one head at a time, serialised by stream)
"bellman_target_projection", // RL Phase E.2-DEFER: C51 categorical projection of Bellman target Z(s_{t+1}, a*) onto the discrete support, reads γ from ISV[400]; replaces host-side build_synthetic_bellman_target stand-in
"rl_lr_controller", // RL Phase E.2-DEFER: per-head learning-rate ISV emitter — bootstraps ISV[412..417] with 1e-3 (BCE/Q/π/V/aux); replaces hardcoded PHASE_E2_DEFAULT_LR
"rl_rollout_steps_controller", // RL Phase E.3b: rollout-buffer-length ISV emitter — emits ISV[RL_N_ROLLOUT_STEPS_INDEX=404] from var(advantage)/|mean A| EMA; bootstraps 2048
"rl_per_alpha_controller", // RL Phase E.3b: PER priority-exponent ISV emitter — emits ISV[RL_PER_ALPHA_INDEX=405] from TD-error kurtosis EMA; bootstraps 0.6
];
// Cache bust v23 (2026-05-22): RL Phase E.2-DEFER — bellman_target_projection.cu (proper C51 categorical projection reading γ from ISV[400], replaces build_synthetic_bellman_target host stand-in), rl_lr_controller.cu (per-head LR ISV emitter, slots 412..417, replaces hardcoded PHASE_E2_DEFAULT_LR). ALSO greenfield-removes the V-loss path from ppo_clipped_surrogate.cu (returns_/v_pred/loss_v deleted; V signal canonicalised through v_head_fwd_bwd kernels per feedback_single_source_of_truth_no_duplicates). All new kernels honour feedback_no_atomicadd.
// Cache bust v24 (2026-05-23): RL Phase E.3b — rl_rollout_steps_controller.cu (ISV[404] producer; advantage-variance-driven rollout length, bootstrap 2048) + rl_per_alpha_controller.cu (ISV[405] producer; TD-kurtosis-driven PER priority exponent α, bootstrap 0.6). Both follow pearl_first_observation_bootstrap + pearl_wiener_alpha_floor_for_nonstationary. Companion to LobEnv trait + step_with_lobsim in src/rl/reward.rs and the integrated trainer.
fn main() {
println!("cargo:rerun-if-changed=build.rs");

97
crates/ml-alpha/cuda/rl_per_alpha_controller.cu Normal file
View File

@@ -0,0 +1,97 @@
// rl_per_alpha_controller.cu — emits PER (Prioritized Experience Replay)
// priority exponent to ISV[RL_PER_ALPHA_INDEX=405].
//
// Phase E.3b of the integrated RL trainer
// (docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md).
//
// PER's priority exponent α controls how aggressively the replay buffer
// concentrates sampling on high-TD-error transitions. Per
// `pearl_controller_anchors_isv_driven` and
// `feedback_isv_for_adaptive_bounds`, α is NOT a hardcoded constant:
// it adapts to the tail thickness of the TD-error distribution.
//
// Controller logic:
// * High TD-error kurtosis (heavy tails) → a few transitions
// dominate the loss → raise α to sharpen the priority distribution
// and concentrate sampling on the informative tails.
// * Low kurtosis (light tails, ≈ Gaussian) → TD-errors are uniform →
// keep α near the standard PER default (0.6) — over-sharpening on a
// light-tailed distribution wastes samples on noise.
//
// Bootstrap discipline (per `pearl_first_observation_bootstrap`): the
// ISV slot starts at 0.0 sentinel. First emit writes
// `PER_ALPHA_BOOTSTRAP = 0.6` directly (the canonical PER default per
// Schaul et al. 2016). Subsequent emits Wiener-α blend with floor 0.4
// (per `pearl_wiener_alpha_floor_for_nonstationary` — the TD-error
// distribution drifts as Q co-adapts, breaking stationarity).
//
// Bounds: α ∈ [0.3, 1.0]. Below 0.3 the priority distribution
// approaches uniform (PER degenerates into vanilla replay); at 1.0
// sampling is perfectly proportional to priority.
#define RL_PER_ALPHA_INDEX 405
#define PER_ALPHA_MIN 0.3f
#define PER_ALPHA_MAX 1.0f
#define PER_ALPHA_BOOTSTRAP 0.6f
// Kurtosis of a standard normal = 3.0 ("excess kurtosis 0" with the
// alternative convention). Used as the breakpoint above which we start
// raising α.
#define KURT_GAUSSIAN 3.0f
// Width of the kurtosis → α-lift logistic map. At kurt=10 (heavy-tailed
// market data), the lift fully spans 0.2 of the [PER_ALPHA_MIN,
// PER_ALPHA_MAX] band. Beyond ~30 the mapping saturates at α_max.
#define KURT_LIFT_SCALE 7.0f
#define WIENER_ALPHA_FLOOR 0.4f
// ─────────────────────────────────────────────────────────────────────
// rl_per_alpha_controller:
// Single-thread controller — writes ONE float to
// isv[RL_PER_ALPHA_INDEX].
//
// Inputs:
// isv [≥ RL_PER_ALPHA_INDEX+1] — ISV bus
// alpha_step — Wiener-α from the controller's signal stats;
// floored at WIENER_ALPHA_FLOOR before the blend.
// Named `alpha_step` to disambiguate from the
// output (PER's α).
// td_kurtosis_ema — caller-computed kurtosis EMA of the per-
// transition |TD-error| series. Phase F produces
// this via a 4-moment reduce over the replay
// priority column.
//
// Outputs:
// isv[RL_PER_ALPHA_INDEX] — PER priority exponent α ∈
// [PER_ALPHA_MIN, PER_ALPHA_MAX]
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_per_alpha_controller(
float* __restrict__ isv,
float alpha_step,
float td_kurtosis_ema
) {
if (threadIdx.x != 0 || blockIdx.x != 0) return;
const float prev = isv[RL_PER_ALPHA_INDEX];
// Bootstrap on sentinel 0.0 per pearl_first_observation_bootstrap.
if (prev == 0.0f) {
isv[RL_PER_ALPHA_INDEX] = PER_ALPHA_BOOTSTRAP;
return;
}
// Map kurtosis → target α via a piecewise linear lift.
// kurt ≤ KURT_GAUSSIAN → target = 0.4 (PER_ALPHA_MIN + 0.1)
// kurt = KURT_GAUSSIAN + KURT_LIFT_SCALE (≈10) → target = 0.6 (default)
// kurt → ∞ → target → PER_ALPHA_MAX
//
// The 0.4-0.6 baseline keeps the default close to PER's canonical
// 0.6 while leaving headroom to lift toward 1.0 under heavy tails.
const float kurt_excess = fmaxf(0.0f, td_kurtosis_ema - KURT_GAUSSIAN);
float target = 0.4f + 0.2f * (kurt_excess / KURT_LIFT_SCALE);
target = fmaxf(PER_ALPHA_MIN, fminf(target, PER_ALPHA_MAX));
// Wiener-α blend with floor per pearl_wiener_alpha_floor_for_nonstationary.
const float a = fmaxf(alpha_step, WIENER_ALPHA_FLOOR);
float out = (1.0f - a) * prev + a * target;
out = fmaxf(PER_ALPHA_MIN, fminf(out, PER_ALPHA_MAX));
isv[RL_PER_ALPHA_INDEX] = out;
}

88
crates/ml-alpha/cuda/rl_rollout_steps_controller.cu Normal file
View File

@@ -0,0 +1,88 @@
// rl_rollout_steps_controller.cu — emits PPO rollout buffer flush size
// to ISV[RL_N_ROLLOUT_STEPS_INDEX=404].
//
// Phase E.3b of the integrated RL trainer
// (docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md).
//
// The rollout length sets how many on-policy transitions PPO accumulates
// before flushing into an update batch. Per
// `pearl_controller_anchors_isv_driven` and
// `feedback_isv_for_adaptive_bounds`, this is NOT a hardcoded constant:
// it adapts to the noise level of the advantage estimator.
//
// Controller logic:
// * High var(A) / |mean A| → advantage estimates are noisy → grow
// rollout length so PPO has more samples per update and the
// surrogate gradient averages over the noise.
// * Low var(A) / |mean A| → advantages are stable → shrink rollout
// so the data is fresher and policy update lag stays small.
//
// Bootstrap discipline (per `pearl_first_observation_bootstrap`): the
// ISV slot starts at 0.0 sentinel. First emit writes
// `ROLLOUT_BOOTSTRAP = 2048` directly (the canonical PPO default).
// Subsequent emits Wiener-α blend with floor 0.4 (per
// `pearl_wiener_alpha_floor_for_nonstationary` — the advantage
// distribution drifts as π_new co-adapts, breaking stationarity).
//
// Bounds: rollout ∈ [256, 8192]. Below 256 PPO has too little data per
// update and the surrogate is dominated by sampling noise; above 8192
// updates lag the data so far behind that the importance ratios blow
// past the clip band.
#define RL_N_ROLLOUT_STEPS_INDEX 404
#define ROLLOUT_MIN 256.0f
#define ROLLOUT_MAX 8192.0f
#define ROLLOUT_BOOTSTRAP 2048.0f
#define ADV_VAR_RATIO_TARGET 0.1f
#define WIENER_ALPHA_FLOOR 0.4f
// ─────────────────────────────────────────────────────────────────────
// rl_rollout_steps_controller:
// Single-thread controller — writes ONE float to
// isv[RL_N_ROLLOUT_STEPS_INDEX].
//
// Inputs:
// isv [≥ RL_N_ROLLOUT_STEPS_INDEX+1] — ISV bus
// alpha — Wiener-α from controller stats;
// floored at WIENER_ALPHA_FLOOR before
// the blend.
// advantage_var_over_abs_mean — caller-computed
// `var(A) / max(|mean A|, ε)` EMA;
// Phase F produces this via a small
// reduce kernel over the rollout
// buffer's advantage column.
//
// Outputs:
// isv[RL_N_ROLLOUT_STEPS_INDEX] — rollout length ∈ [ROLLOUT_MIN, ROLLOUT_MAX]
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_rollout_steps_controller(
float* __restrict__ isv,
float alpha,
float advantage_var_over_abs_mean
) {
if (threadIdx.x != 0 || blockIdx.x != 0) return;
const float prev = isv[RL_N_ROLLOUT_STEPS_INDEX];
// Bootstrap on sentinel 0.0 per pearl_first_observation_bootstrap.
if (prev == 0.0f) {
isv[RL_N_ROLLOUT_STEPS_INDEX] = ROLLOUT_BOOTSTRAP;
return;
}
// Multiplicative adaptation: if the var-ratio exceeds the target
// (noisy advantages), scale up the rollout; if it falls below, scale
// down. Clamped to [0.5, 2.0] per step so we never double-halve in a
// single emit (per pearl_blend_formulas_must_have_permanent_floor:
// step-size bounds even before the Wiener blend).
const float ratio_raw = advantage_var_over_abs_mean / ADV_VAR_RATIO_TARGET;
const float scale = fmaxf(0.5f, fminf(2.0f, ratio_raw));
float target = prev * scale;
target = fmaxf(ROLLOUT_MIN, fminf(target, ROLLOUT_MAX));
// Wiener-α blend with floor per pearl_wiener_alpha_floor_for_nonstationary.
const float a = fmaxf(alpha, WIENER_ALPHA_FLOOR);
float out = (1.0f - a) * prev + a * target;
out = fmaxf(ROLLOUT_MIN, fminf(out, ROLLOUT_MAX));
isv[RL_N_ROLLOUT_STEPS_INDEX] = out;
}

1
crates/ml-alpha/src/rl/mod.rs
View File

@@ -21,4 +21,5 @@ pub mod isv_slots;
pub mod loss_balance;
pub mod ppo;
pub mod replay;
pub mod reward;
pub mod rollout;

220
crates/ml-alpha/src/rl/reward.rs Normal file
View File

@@ -0,0 +1,220 @@
//! LobSim adapter trait for the integrated RL trainer (Phase E.3b).
//!
//! Defines the minimal interface the trainer needs from any LOB
//! simulator: apply a snapshot, submit an action, and step the
//! environment one event forward to receive a reward + done flag.
//!
//! The actual `LobSimCuda` implementation lives in `ml-backtesting`,
//! which already depends on `ml-alpha` (loader + trunk reuse). A direct
//! `ml-alpha → ml-backtesting` dep would cycle, so we expose the
//! contract here as a trait and let the simulator-side `impl LobEnv for
//! LobSimCuda` complete the wire from the other direction. If/when
//! the dep direction is one-way, the trait can be retired in favour of
//! a concrete `LobSimCuda` field on `IntegratedTrainer`.
//!
//! The trait is intentionally narrow: `IntegratedTrainer::step_with_lobsim`
//! must NOT reach into LobSim's bytecode-program or per-trade audit
//! plumbing — it only needs:
//!
//! 1. push a new MBP-10 snapshot into the book,
//! 2. submit the chosen discrete action,
//! 3. step pnl_track + position accounting one event,
//! 4. read back per-step realized PnL + done flag.
//!
//! Phase F+ may extend this trait with reward-shaping hooks (per-step
//! cost integration, max-hold force-close) once the LobSim side exposes
//! those signals on a per-call basis.
use anyhow::Result;
use crate::cfc::snap_features::Mbp10RawInput;
/// Trait the integrated trainer requires of any LOB-style RL environment.
///
/// `LobSimCuda` (in `ml-backtesting`) implements this for the production
/// path: snapshots flow from the MBP-10 loader into `apply_snapshot`,
/// the trainer's Thompson-sampled action lands via `submit_action`,
/// and `step_event` advances the event-time clock with resting-order
/// matching + position accounting.
///
/// All methods are `&mut self`: the env owns the device-resident book +
/// position state and mutates it in place per call.
pub trait LobEnv {
/// Apply the most recent MBP-10 snapshot to the order book of every
/// parallel backtest batch slot. Phase E.3b broadcasts a single
/// snapshot across all batches (matches `LobSimCuda::apply_snapshot`
/// semantics — see ml-backtesting/sim/mod.rs §`apply_snapshot`).
///
/// `snapshot.bid_px/bid_sz/ask_px/ask_sz` are forwarded to the
/// underlying book-update kernel; other fields (regime, ts_ns, …)
/// are consumed by the trainer's encoder, not the env.
fn apply_snapshot(&mut self, snapshot: &Mbp10RawInput) -> Result<()>;
/// Submit a market order for `batch_idx` matching the discrete
/// 9-action grid (see `crate::rl::common::Action`):
/// * `0 = ShortLarge`, `1 = ShortSmall`,
/// * `2 = Hold` (no-op),
/// * `3 = FlatFromLong`, `4 = FlatFromShort`,
/// * `5 = LongSmall`, `6 = LongLarge`,
/// * `7 = TrailTighten`, `8 = TrailLoosen`.
///
/// The implementation translates the action index into the
/// simulator's `(side, size)` market-order vocabulary. Action `2`
/// (Hold) MUST be a no-op (no order submitted, no fill). Trail-
/// tighten / loosen DO NOT submit a market order — they nudge the
/// trailing-stop distance via the ISV-driven stop controller; Phase F
/// wires the actual mutation, Phase E.3b accepts them as no-ops
/// while the surface is shaken out.
fn submit_action(&mut self, batch_idx: usize, action: u32, ts_ns: u64) -> Result<()>;
/// Advance one event for `batch_idx`: run resting-order matching +
/// step_pnl_track + position accounting. Returns
/// `(reward, done)`:
///
/// * `reward` is the realized PnL delta in USD since the previous
/// `step_event` call. Zero between trade closes (the reward is
/// sparse and fires only when a position returns to flat).
/// * `done` is `true` exactly when a trade closed during this step
/// (position transitioned through zero); `false` otherwise.
///
/// `trade_signed_vol` is the signed traded volume from the next
/// snapshot (positive = buyer-initiated, negative = seller-initiated),
/// used by the resting-order matching kernel.
fn step_event(
&mut self,
batch_idx: usize,
ts_ns: u64,
trade_signed_vol: f32,
) -> Result<(f32, bool)>;
}
// ─────────────────────────────────────────────────────────────────────
// MockLobEnv — toy-bandit test fixture
// ─────────────────────────────────────────────────────────────────────
/// Deterministic toy-bandit environment for activating the
/// `dqn_toy` / `ppo_toy` / `integrated_trainer_smoke` tests.
///
/// Reward function: action 5 (`LongSmall`) always returns `+1.0`; every
/// other action returns `-1.0`. State is uninformative (state-
/// independent bandit). `done = true` after every step (each step is a
/// trade close). `apply_snapshot` is a no-op — the bandit ignores the
/// book entirely.
///
/// This lives in production code (not `#[cfg(test)]`) because the toy
/// bandit tests live in `tests/` (integration-test crates), which can
/// only depend on `ml_alpha::…` public items. Phase F may move it
/// behind a `cfg(test)` gate once the smoke tests share a fixtures
/// module.
pub struct MockLobEnv {
/// Reward returned for the "good" action. Pinned at +1.0 by default
/// (toy bandit gate) but exposed so future tests can drive the
/// trainer with arbitrary deterministic reward functions.
pub good_reward: f32,
/// Reward returned for every other action. Pinned at -1.0 by
/// default.
pub bad_reward: f32,
/// The action that pays `good_reward`. Default = 5 (`LongSmall`)
/// to match the canonical toy-bandit gate.
pub good_action: u32,
/// Counter of submitted actions, indexed by action value. Lets the
/// test inspect post-hoc which actions the trainer actually picked.
pub action_counts: Vec<u64>,
/// Last reward written by `step_event`. Inspected by tests to
/// assert end-of-training optimal-action selection.
pub last_reward: f32,
/// Last submitted action. `step_event` reads this to compute the
/// reward (the action submitted right before this `step_event`
/// call drives the reward).
last_action: u32,
}
impl MockLobEnv {
/// Construct a bandit with the canonical toy-gate reward function:
/// `r = +1` if `action == 5`, else `-1`. `done = true` every step.
pub fn toy_bandit() -> Self {
Self {
good_reward: 1.0,
bad_reward: -1.0,
good_action: 5,
action_counts: vec![0; crate::rl::common::N_ACTIONS],
last_reward: 0.0,
last_action: u32::MAX,
}
}
}
impl LobEnv for MockLobEnv {
fn apply_snapshot(&mut self, _snapshot: &Mbp10RawInput) -> Result<()> {
// Toy bandit ignores market state — reward is state-independent.
Ok(())
}
fn submit_action(
&mut self,
_batch_idx: usize,
action: u32,
_ts_ns: u64,
) -> Result<()> {
if (action as usize) < self.action_counts.len() {
self.action_counts[action as usize] += 1;
}
self.last_action = action;
Ok(())
}
fn step_event(
&mut self,
_batch_idx: usize,
_ts_ns: u64,
_trade_signed_vol: f32,
) -> Result<(f32, bool)> {
let r = if self.last_action == self.good_action {
self.good_reward
} else {
self.bad_reward
};
self.last_reward = r;
Ok((r, true))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mock_bandit_rewards_action_5() {
let mut env = MockLobEnv::toy_bandit();
env.submit_action(0, 5, 0).unwrap();
let (r, done) = env.step_event(0, 0, 0.0).unwrap();
assert_eq!(r, 1.0);
assert!(done);
assert_eq!(env.action_counts[5], 1);
}
#[test]
fn mock_bandit_penalises_other_actions() {
let mut env = MockLobEnv::toy_bandit();
for a in 0..9u32 {
if a == 5 {
continue;
}
env.submit_action(0, a, 0).unwrap();
let (r, done) = env.step_event(0, 0, 0.0).unwrap();
assert_eq!(r, -1.0, "action {a} should pay -1");
assert!(done);
}
}
#[test]
fn mock_bandit_apply_snapshot_is_noop() {
let mut env = MockLobEnv::toy_bandit();
let snap = Mbp10RawInput::default();
env.apply_snapshot(&snap).unwrap();
// Submitting action 5 still pays +1 (snapshot is ignored).
env.submit_action(0, 5, 0).unwrap();
let (r, _done) = env.step_event(0, 0, 0.0).unwrap();
assert_eq!(r, 1.0);
}
}

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

@@ -86,7 +86,9 @@ use ml_core::device::MlDevice;
use crate::cfc::snap_features::Mbp10RawInput;
use crate::heads::HIDDEN_DIM;
use crate::pinned_mem::{MappedF32Buffer, MappedI32Buffer};
use crate::rl::common::{N_ACTIONS, Q_N_ATOMS};
use rand::{Rng, SeedableRng};
use crate::rl::common::{N_ACTIONS, Q_N_ATOMS, Q_V_MAX, Q_V_MIN};
use crate::rl::dqn::{DqnHead, DqnHeadConfig};
use crate::rl::isv_slots::{
RL_LR_BCE_INDEX, RL_LR_AUX_INDEX, RL_LR_PI_INDEX, RL_LR_Q_INDEX,
@@ -94,6 +96,7 @@ use crate::rl::isv_slots::{
};
use crate::rl::loss_balance::{read_loss_lambdas_from_isv, LossLambdas};
use crate::rl::ppo::{PolicyHead, PpoHeadsConfig, ValueHead};
use crate::rl::reward::LobEnv;
use crate::trainer::optim::AdamW;
use crate::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
@@ -200,6 +203,14 @@ pub struct IntegratedTrainer {
/// surrogate kernel's atomic-write loss readback through the
/// borrow-checker dance in `step_synthetic`. NOT a public API.
last_pi_loss: f32,
/// Host-side step counter for Thompson-sampling RNG salt (Phase E.3b).
/// Increments once per `step_with_lobsim` call so consecutive calls
/// draw different actions from the same Q-distribution. Per
/// `pearl_scoped_init_seed_for_reproducibility`: the RNG seed is
/// `cfg.dqn_seed.wrapping_add(step_counter)`, deterministic given
/// the same config + step sequence.
step_counter: u64,
}
impl IntegratedTrainer {
@@ -321,6 +332,7 @@ impl IntegratedTrainer {
rl_lr_controller_fn,
grad_h_t_combined_d,
last_pi_loss: 0.0,
step_counter: 0,
})
}
@@ -712,6 +724,270 @@ impl IntegratedTrainer {
})
}
/// Phase E.3b: run one integrated training step driven by a real
/// LobSim environment (`impl LobEnv`). 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],
lobsim: &mut dyn LobEnv,
) -> 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)");
}
// ── Step 1: encoder forward to land h_t at slot K-1. ──────────
let _ = self
.perception
.forward_encoder(snapshots)
.context("step_with_lobsim: forward_encoder")?;
// ── 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);
let k_dqn = N_ACTIONS * Q_N_ATOMS;
let mut q_logits_d = self.stream.alloc_zeros::<f32>(b_size * k_dqn)?;
let mut v_pred_d = self.stream.alloc_zeros::<f32>(b_size)?;
self.dqn_head
.forward(h_t_borrow, b_size, &mut q_logits_d)
.context("step_with_lobsim: dqn_head.forward")?;
self.value_head
.forward(h_t_borrow, b_size, &mut v_pred_d)
.context("step_with_lobsim: value_head.forward")?;
// ── Step 2b: Forward π logits for log_pi_old. ─────────────────
let mut pi_logits_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
self.policy_head
.forward_logits(h_t_borrow, b_size, &mut pi_logits_d)
.context("step_with_lobsim: policy_head.forward_logits")?;
// ── Step 3: read Q logits / V / π logits back to host. ────────
// Mapped-pinned staging buffers per
// `feedback_no_htod_htoh_only_mapped_pinned`.
let q_logits_host = read_slice_d(&self.stream, &q_logits_d, b_size * k_dqn)?;
let v_pred_host = read_slice_d(&self.stream, &v_pred_d, b_size)?;
let pi_logits_host = read_slice_d(&self.stream, &pi_logits_d, b_size * N_ACTIONS)?;
// Refresh ISV host mirror so we can read γ for one-step return.
// (`launch_rl_lr_controller` fires inside `step_synthetic` later;
// for the return-bootstrap we only need γ, which the trainer's
// bootstrap path seeds at 0.99 on the first γ-controller emit.
// Until that fire, the mirror reads 0.0 — we fall back to the
// canonical 0.99 default rather than zeroing the bootstrap.)
self.stream
.memcpy_dtoh(&self.isv_d, self.isv_host.as_mut_slice())
.context("step_with_lobsim: isv dtoh")?;
let gamma = {
let v = self.isv_host[crate::rl::isv_slots::RL_GAMMA_INDEX];
if v > 0.0 { v } else { 0.99 }
};
// ── Step 4: Thompson sample action per batch. ─────────────────
// For each batch, sample one return draw from each action's atom
// distribution and pick the argmax. Bypass `pearl_no_host_branches_
// in_captured_graph`: this is a host-side sampling step OUTSIDE
// any captured graph; Phase F may move it to a GPU kernel.
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();
// Deterministic per-step RNG seeded from the dqn config seed
// so the toy-bandit test is reproducible per
// `pearl_scoped_init_seed_for_reproducibility`. The salt is the
// host-side step counter (incremented at the end of this method)
// so consecutive calls draw different actions from the same
// Q-distribution without losing reproducibility.
let salt = self.step_counter;
self.step_counter = self.step_counter.wrapping_add(1);
let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(
self.cfg.dqn_seed.wrapping_add(salt),
);
let mut actions = vec![0u32; b_size];
let mut next_actions = vec![0u32; b_size];
let mut log_pi_old = vec![0.0_f32; b_size];
for b in 0..b_size {
// Softmax per-action atom probabilities (apply numerical
// stabilisation per-action row).
let mut action_expected = [0.0_f32; N_ACTIONS];
let mut sampled_returns = [0.0_f32; N_ACTIONS];
for a in 0..N_ACTIONS {
let off = b * k_dqn + a * Q_N_ATOMS;
let row = &q_logits_host[off..off + Q_N_ATOMS];
// Softmax to atom probabilities.
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;
}
}
// Expected value of action a — used both for next_actions
// (host-side argmax) and Thompson-sampled-vs-expected
// diagnostics (Phase F).
let mut ev = 0.0_f32;
for i in 0..Q_N_ATOMS {
ev += probs[i] * atom_supports[i];
}
action_expected[a] = ev;
// Thompson: sample ONE atom from the categorical, that
// atom's support value is the sampled return.
let u: f32 = rng.gen_range(0.0..1.0);
let mut cum = 0.0_f32;
let mut chosen = Q_N_ATOMS - 1;
for (i, &p) in probs.iter().enumerate() {
cum += p;
if cum >= u {
chosen = i;
break;
}
}
sampled_returns[a] = atom_supports[chosen];
}
// Thompson selector: argmax over sampled per-action returns.
let action_sampled = argmax_f32(&sampled_returns) as u32;
// Bellman target uses argmax of EXPECTED Q (not Thompson sample)
// per the canonical distributional-Q Bellman backup.
let next_action = argmax_f32(&action_expected) as u32;
// log π_old at the sampled action — softmax of pi_logits row.
let pi_off = b * N_ACTIONS;
let pi_row = &pi_logits_host[pi_off..pi_off + N_ACTIONS];
let pi_max = pi_row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mut pi_sum = 0.0_f32;
for &l in pi_row.iter() {
pi_sum += (l - pi_max).exp();
}
let log_norm = pi_max + pi_sum.ln();
log_pi_old[b] = pi_row[action_sampled as usize] - log_norm;
actions[b] = action_sampled;
next_actions[b] = next_action;
}
// ── Step 5: drive the env, collect reward + done per batch. ───
// Phase E.3b broadcasts the LAST snapshot in the window to the
// env (matches `LobSimCuda::apply_snapshot` which operates on
// current top-of-book state, not the history window the encoder
// consumed). Phase F may stream every snapshot through the env.
let last_snap = snapshots
.last()
.expect("snapshots non-empty: guarded above");
lobsim
.apply_snapshot(last_snap)
.context("step_with_lobsim: lobsim.apply_snapshot")?;
let mut rewards = vec![0.0_f32; b_size];
let mut dones = vec![0.0_f32; b_size];
for b in 0..b_size {
lobsim
.submit_action(b, actions[b], last_snap.ts_ns)
.with_context(|| format!("step_with_lobsim: lobsim.submit_action b={b}"))?;
let (r, done) = lobsim
.step_event(b, last_snap.ts_ns, last_snap.trade_signed_vol)
.with_context(|| format!("step_with_lobsim: lobsim.step_event b={b}"))?;
rewards[b] = r;
dones[b] = if done { 1.0 } else { 0.0 };
}
// ── Step 6: derive advantages + returns from real reward + V. ─
// One-step TD advantage: A_t = r_t + γ(1-done)·V(s_t) - V(s_t)
// = r_t - γ·done·V(s_t) - V(s_t)·(1-γ(1-done))
// (we treat s_{t+1} ≈ s_t for Phase E.3b — Phase F threads the
// next-state encoder forward).
// R_t = r_t + γ(1-done)·V(s_t)
let mut advantages = vec![0.0_f32; b_size];
let mut returns = vec![0.0_f32; b_size];
for b in 0..b_size {
let v_b = v_pred_host[b];
let r_b = rewards[b];
let done_b = dones[b];
let r_t = r_b + gamma * (1.0 - done_b) * v_b;
returns[b] = r_t;
advantages[b] = r_t - v_b;
}
// ── Step 7: delegate to step_synthetic for the train step. ───
// step_synthetic re-runs forward_encoder (compute redundancy,
// see method docs); we accept that for Phase E.3b in service of
// a single backward + Adam codepath.
self.step_synthetic(
snapshots,
&actions,
&rewards,
&dones,
&next_actions,
&advantages,
&returns,
&log_pi_old,
)
}
/// Launch `rl_lr_controller` to emit per-head learning rates into
/// `ISV[412..417]`. Phase E.2-DEFER item 3. Diagnostic-input args
/// are currently zero (the controller emits the bootstrap target
@@ -791,6 +1067,113 @@ impl IntegratedTrainer {
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.
@@ -894,3 +1277,50 @@ fn read_scalar_d(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<f32>
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 E.3b helper: argmax over a small slice (used for Thompson-
/// selector + Bellman-target next-action). Returns the first index of
/// the maximum value (deterministic tie-break for reproducibility per
/// `pearl_scoped_init_seed_for_reproducibility`).
fn argmax_f32(xs: &[f32]) -> usize {
debug_assert!(!xs.is_empty());
let mut best_i = 0;
let mut best_v = xs[0];
for (i, &v) in xs.iter().enumerate().skip(1) {
if v > best_v {
best_v = v;
best_i = i;
}
}
best_i
}

167
crates/ml-alpha/tests/dqn_toy.rs
View File

@@ -1,45 +1,136 @@
//! Phase C toy bandit gate — falsifiable test that the C51 DQN head can
//! learn a simple state-conditional optimal action.
//! learn a simple state-conditional optimal action when driven by the
//! integrated trainer.
//!
//! Synthetic setup (when Phase E activates this test):
//! * `h_t ~ uniform[-1, 1]^HIDDEN_DIM` — synthetic encoder hidden state.
//! * Reward function: action 5 (LongSmall) always returns +1; all
//! other actions return -1. State is uninformative (the bandit is
//! state-independent) — this gates that the head's bias path is
//! sufficient to learn the optimal arm.
//! * Run 1000 training steps, sampling from the replay buffer with
//! priority-weighted PER (α from `ISV[RL_PER_ALPHA_INDEX]`).
//! * Bellman target uses γ from `ISV[RL_GAMMA_INDEX]` and target-net
//! τ from `ISV[RL_TARGET_TAU_INDEX]` (NOT hardcoded — that's the
//! whole point of the ISV-driven design).
//! * Gate: `argmax_a E[Q(s, a)] == Action::LongSmall` for ≥ 90% of
//! fresh held-out states.
//! Phase E.3b activation:
//! * `h_t` derived from a fixed synthetic snapshot window (uninformative
//! to the bandit reward — the bandit is state-independent).
//! * Reward function: action 5 (LongSmall) → +1; all others → -1
//! (`MockLobEnv::toy_bandit()`).
//! * Run N=300 training steps via `IntegratedTrainer::step_with_lobsim`
//! (drives the full per-head fwd+bwd+Adam + encoder backward via the
//! shared training kernel chain).
//! * Bellman target uses γ from `ISV[RL_GAMMA_INDEX]` (bootstrapped
//! 0.99 on first emit), τ from `ISV[RL_TARGET_TAU_INDEX]` (Phase F),
//! PER α from `ISV[RL_PER_ALPHA_INDEX]` (bootstrapped 0.6).
//! * Gate: argmax_a E[Q(s, a)] equals `Action::LongSmall` (= 5) by end
//! of training.
//!
//! This file currently holds the test SKELETON only — Phase C lands
//! the type contract (DqnHead, ReplayBuffer, kernels), but the training
//! loop that drives end-to-end learning lives in Phase E (which wires
//! `LobSimCuda` for the reward signal + the categorical projection
//! kernel + Adam stepping). The test is `#[ignore]` until Phase E
//! activates it so `cargo test --workspace` stays green in the interim.
//! `#[ignore]`-gated for CUDA availability per the project's CUDA test
//! discipline (`MlDevice::cuda(0)` may not be available in pure-CPU CI
//! runs; the Argo `ci-compile-cpu` pool skips ignored tests by design).
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::rl::common::Action;
use ml_alpha::rl::reward::MockLobEnv;
use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig};
use ml_alpha::trainer::perception::PerceptionTrainerConfig;
use ml_core::device::MlDevice;
/// Build a deterministic single-direction price ramp matching the
/// shape consumed by the encoder. The reward is state-independent in
/// `MockLobEnv::toy_bandit` so the exact snapshot contents don't
/// matter; we just need shapes the snap_feature_assemble kernel
/// accepts.
fn synthetic_window(seq_len: usize) -> Vec<Mbp10RawInput> {
let mut out = Vec::with_capacity(seq_len);
let mut prev_mid = 5500.0_f32;
let mut ts_ns = 1_000_000_u64;
for _ in 0..seq_len {
let next_mid = prev_mid + 0.25;
let mut bid_px = [0.0_f32; 10];
let mut bid_sz = [0.0_f32; 10];
let mut ask_px = [0.0_f32; 10];
let mut ask_sz = [0.0_f32; 10];
for i in 0..10 {
bid_px[i] = next_mid - 0.125 - 0.25 * i as f32;
ask_px[i] = next_mid + 0.125 + 0.25 * i as f32;
bid_sz[i] = 10.0;
ask_sz[i] = 10.0;
}
let prev_ts = ts_ns;
ts_ns += 20_000_000;
out.push(Mbp10RawInput {
bid_px,
bid_sz,
ask_px,
ask_sz,
prev_mid,
trade_signed_vol: 0.0,
trade_count: 0,
ts_ns,
prev_ts_ns: prev_ts,
regime: [0.0; 6],
});
prev_mid = next_mid;
}
out
}
#[test]
#[ignore = "Phase C wires types; Phase E activates the training loop. \
Re-enable after Phase E lands the categorical projection \
kernel + soft-update + lobsim reward path."]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn toy_bandit_q_argmax_converges_to_optimal_action() {
// Phase E will populate this test with:
// 1. Initialise MlDevice + DqnHead + ReplayBuffer.
// 2. Seed ISV[RL_GAMMA_INDEX] / ISV[RL_TARGET_TAU_INDEX] /
// ISV[RL_PER_ALPHA_INDEX] by firing each controller once with
// the bootstrap sentinel (0.0) so the slot picks up its
// first-observation value.
// 3. Run a synthetic rollout: for each of 1000 steps, sample
// a random h_t, pick an action via Thompson over atoms,
// record (r = +1 if a == 5 else -1), push the Transition.
// 4. Every 4 rollout steps: sample a batch from the replay,
// run the categorical projection (Phase E kernel) + Bellman
// backward, Adam-step the Q weights, soft-update the target.
// 5. Assert: across 256 fresh h_t draws, argmax_a E[Q(s, a)]
// equals Action::LongSmall (= 5) in ≥ 90% of samples.
panic!("Phase E wires the toy bandit training loop");
let dev = match MlDevice::cuda(0) {
Ok(d) => d,
Err(_) => {
eprintln!("CUDA 0 not available — skipping toy_bandit_q");
return;
}
};
// Minimal config — B=1, seq_len=4 — keeps the test light.
let perception = PerceptionTrainerConfig {
seq_len: 4,
n_batch: 1,
..PerceptionTrainerConfig::default()
};
let cfg = IntegratedTrainerConfig {
perception,
dqn_seed: 0xC51D,
ppo_seed: 0xC51E,
};
let mut trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new");
let snapshots = synthetic_window(4);
let mut env = MockLobEnv::toy_bandit();
let n_steps = 300;
for _ in 0..n_steps {
trainer
.step_with_lobsim(&snapshots, &mut env)
.expect("step_with_lobsim");
}
// After training, argmax_a E[Q(s, a)] should be Action::LongSmall (5).
let eq = trainer
.eval_expected_q_per_action(&snapshots)
.expect("eval_expected_q_per_action");
let argmax_a = eq
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(i, _)| i)
.unwrap();
println!(
"toy_bandit_q: E[Q] = {:?}, argmax = {} (expected {})",
eq, argmax_a, Action::LongSmall as u32
);
assert_eq!(
argmax_a,
Action::LongSmall as usize,
"C51 head should learn argmax-Q = LongSmall (5); got action {argmax_a} with E[Q] = {eq:?}"
);
// Sanity: the trainer selected the good action more than uniform-random
// (1/9 ≈ 0.111) by end of training. Thompson sampling concentrates on
// the high-EV arm as the distribution sharpens.
let total: u64 = env.action_counts.iter().sum();
let good_share = env.action_counts[Action::LongSmall as usize] as f32 / total as f32;
println!(
"toy_bandit_q: action_counts = {:?}, good_share = {:.3}",
env.action_counts, good_share
);
assert!(
good_share > 0.111,
"Thompson selector should beat uniform 1/9; got {good_share:.3}"
);
}

115
crates/ml-alpha/tests/integrated_trainer_smoke.rs
View File

@@ -1,19 +1,106 @@
//! Phase E.1 smoke: confirm IntegratedTrainer constructs and runs a
//! synthetic step without panic. Real kernel wiring + LobSim integration
//! is Phase E.2.
//! Phase E.3b smoke: confirm `IntegratedTrainer` constructs and runs a
//! single `step_with_lobsim` step (real action sampling + reward path)
//! without panic. Distinct from the toy-bandit gate tests (`dqn_toy`,
//! `ppo_toy`) which assert convergence — this test only proves the
//! wiring compiles + runs end-to-end on a single training step.
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::rl::reward::MockLobEnv;
use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig};
use ml_alpha::trainer::perception::PerceptionTrainerConfig;
use ml_core::device::MlDevice;
fn synthetic_window(seq_len: usize) -> Vec<Mbp10RawInput> {
let mut out = Vec::with_capacity(seq_len);
let mut prev_mid = 5500.0_f32;
let mut ts_ns = 1_000_000_u64;
for _ in 0..seq_len {
let next_mid = prev_mid + 0.25;
let mut bid_px = [0.0_f32; 10];
let mut bid_sz = [0.0_f32; 10];
let mut ask_px = [0.0_f32; 10];
let mut ask_sz = [0.0_f32; 10];
for i in 0..10 {
bid_px[i] = next_mid - 0.125 - 0.25 * i as f32;
ask_px[i] = next_mid + 0.125 + 0.25 * i as f32;
bid_sz[i] = 10.0;
ask_sz[i] = 10.0;
}
let prev_ts = ts_ns;
ts_ns += 20_000_000;
out.push(Mbp10RawInput {
bid_px,
bid_sz,
ask_px,
ask_sz,
prev_mid,
trade_signed_vol: 0.0,
trade_count: 0,
ts_ns,
prev_ts_ns: prev_ts,
regime: [0.0; 6],
});
prev_mid = next_mid;
}
out
}
#[test]
#[ignore = "Phase E.1 smoke — requires CUDA; activated in Phase E.2 with the real kernel wiring"]
fn integrated_trainer_constructs_and_runs_synthetic_step() {
// Phase E.2 fills this in alongside the real GPU kernel calls:
// 1. Construct MlDevice::cuda(0). Skip gracefully if unavailable.
// 2. Build IntegratedTrainerConfig with minimal perception cfg (B=1, K=4).
// 3. Build IntegratedTrainer.
// 4. Run step_synthetic with all-zero synthetic inputs.
// 5. Assert: stats.l_total is finite, lambdas are equal-weight default.
//
// The #[ignore] attribute keeps this test from running until the body
// is filled in — the empty body is intentional, not a deferred marker.
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn integrated_trainer_step_with_lobsim_runs_without_panic() {
let dev = match MlDevice::cuda(0) {
Ok(d) => d,
Err(_) => {
eprintln!("CUDA 0 not available — skipping integrated_trainer smoke");
return;
}
};
let perception = PerceptionTrainerConfig {
seq_len: 4,
n_batch: 1,
..PerceptionTrainerConfig::default()
};
let cfg = IntegratedTrainerConfig {
perception,
dqn_seed: 0xC0DE,
ppo_seed: 0xC0DF,
};
let mut trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new");
let snapshots = synthetic_window(4);
let mut env = MockLobEnv::toy_bandit();
// One step end-to-end (encoder fwd, Q/π/V fwd, Thompson sample,
// lobsim step, per-head bwd, Adam, encoder bwd). The internal
// assertions inside the trainer + each kernel guard the shapes.
let stats = trainer
.step_with_lobsim(&snapshots, &mut env)
.expect("step_with_lobsim");
// Loss components are finite.
assert!(stats.l_q.is_finite(), "l_q is not finite: {}", stats.l_q);
assert!(stats.l_pi.is_finite(), "l_pi is not finite: {}", stats.l_pi);
assert!(stats.l_v.is_finite(), "l_v is not finite: {}", stats.l_v);
assert!(stats.l_total.is_finite(), "l_total is not finite: {}", stats.l_total);
// λs sum to ~1.0 (loss-balance default before any controller fires).
let lambda_sum = stats.lambdas.bce
+ stats.lambdas.q
+ stats.lambdas.pi
+ stats.lambdas.v
+ stats.lambdas.aux;
assert!(
(lambda_sum - 1.0).abs() < 1e-3,
"loss-balance λs should sum to ~1; got {lambda_sum} (lambdas={:?})",
stats.lambdas
);
// Env should have observed exactly one submitted action.
let total_submitted: u64 = env.action_counts.iter().sum();
assert_eq!(
total_submitted, 1,
"env should have observed exactly one submit_action call; got {total_submitted}"
);
}
#[test]

160
crates/ml-alpha/tests/ppo_toy.rs
View File

@@ -1,53 +1,121 @@
//! Phase D toy bandit gate — falsifiable test that the PPO policy head
//! can learn a simple state-conditional optimal action.
//! can learn a simple state-conditional optimal action when driven by
//! the integrated trainer.
//!
//! Synthetic setup (when Phase E activates this test):
//! * `h_t ~ uniform[-1, 1]^HIDDEN_DIM` — synthetic encoder hidden state.
//! * Reward function: action 5 (LongSmall) always returns +1; all
//! other actions return -1. State is uninformative (the bandit is
//! state-independent) — this gates that the policy head can learn
//! the optimal arm via the clipped surrogate.
//! * Roll out 256 steps per update, action sampled from π_new (categorical
//! over the 9-action grid). Bootstrap advantage from the Q-head's
//! expected value: A_t = Q(s_t, a_t) - V(s_t).
//! * Run 1000 PPO updates, sampling minibatches from the rollout buffer
//! for `PPO_N_EPOCHS = 4` epochs each.
//! * ε from `ISV[RL_PPO_CLIP_INDEX]` (NOT hardcoded); entropy coef
//! from `ISV[RL_ENTROPY_COEF_INDEX]` (NOT hardcoded) — that's the
//! whole point of the ISV-driven design.
//! * Gate: `mode π(s) == Action::LongSmall` for ≥ 90% of fresh held-out
//! states.
//! Phase E.3b activation:
//! * `h_t` derived from a fixed synthetic snapshot window (uninformative
//! to the bandit reward — the bandit is state-independent).
//! * Reward function: action 5 (LongSmall) → +1; all others → -1
//! (`MockLobEnv::toy_bandit()`).
//! * Run N=300 training steps via `IntegratedTrainer::step_with_lobsim`.
//! * ε from `ISV[RL_PPO_CLIP_INDEX]` (bootstrap 0.2); entropy coef
//! from `ISV[RL_ENTROPY_COEF_INDEX]` (bootstrap 0.01).
//! * Gate: mode π(s) — argmax over the post-softmax action probability
//! vector — equals `Action::LongSmall` (= 5) by end of training.
//!
//! This file currently holds the test SKELETON only — Phase D lands the
//! type contract (PolicyHead, ValueHead, RolloutBuffer, surrogate kernel,
//! 2 controllers), but the training loop that drives end-to-end learning
//! lives in Phase E (which wires `LobSimCuda` for the reward signal +
//! the integrated trainer + Adam stepping). The test is `#[ignore]`
//! until Phase E activates it so `cargo test --workspace` stays green
//! in the interim.
//! `#[ignore]`-gated for CUDA availability.
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::rl::common::Action;
use ml_alpha::rl::reward::MockLobEnv;
use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig};
use ml_alpha::trainer::perception::PerceptionTrainerConfig;
use ml_core::device::MlDevice;
/// Deterministic synthetic snapshot window. State is uninformative to
/// the toy bandit reward.
fn synthetic_window(seq_len: usize) -> Vec<Mbp10RawInput> {
let mut out = Vec::with_capacity(seq_len);
let mut prev_mid = 5500.0_f32;
let mut ts_ns = 1_000_000_u64;
for _ in 0..seq_len {
let next_mid = prev_mid + 0.25;
let mut bid_px = [0.0_f32; 10];
let mut bid_sz = [0.0_f32; 10];
let mut ask_px = [0.0_f32; 10];
let mut ask_sz = [0.0_f32; 10];
for i in 0..10 {
bid_px[i] = next_mid - 0.125 - 0.25 * i as f32;
ask_px[i] = next_mid + 0.125 + 0.25 * i as f32;
bid_sz[i] = 10.0;
ask_sz[i] = 10.0;
}
let prev_ts = ts_ns;
ts_ns += 20_000_000;
out.push(Mbp10RawInput {
bid_px,
bid_sz,
ask_px,
ask_sz,
prev_mid,
trade_signed_vol: 0.0,
trade_count: 0,
ts_ns,
prev_ts_ns: prev_ts,
regime: [0.0; 6],
});
prev_mid = next_mid;
}
out
}
#[test]
#[ignore = "Phase D wires types; Phase E activates the training loop. \
Re-enable after Phase E lands the integrated trainer + \
lobsim reward path + KL/entropy EMA + Adam stepping."]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn toy_bandit_policy_mode_converges_to_optimal_action() {
// Phase E will populate this test with:
// 1. Initialise MlDevice + PolicyHead + ValueHead + RolloutBuffer.
// 2. Seed ISV[RL_PPO_CLIP_INDEX] / ISV[RL_ENTROPY_COEF_INDEX] /
// ISV[RL_GAMMA_INDEX] by firing each controller once with the
// bootstrap sentinel (0.0) so the slot picks up its
// first-observation value.
// 3. Run a synthetic rollout: for each of 256 steps, sample a
// random h_t, sample an action from π_new, record (r = +1 if
// a == 5 else -1), push the Transition.
// 4. Pre-compute Q(s, ·) + V(s) over the rollout via DqnHead /
// ValueHead forward; call
// `rollout.compute_advantages_and_returns(q, v, gamma)`.
// 5. For `PPO_N_EPOCHS` epochs: sample a minibatch, run the
// surrogate fwd (writes loss_pi/loss_v/loss_entropy), run the
// surrogate bwd, Adam-step the policy and value heads.
// 6. After 1000 updates: assert that across 256 fresh h_t draws,
// argmax_a π_new(a | s) equals Action::LongSmall (= 5) in
// ≥ 90% of samples.
panic!("Phase E wires the toy bandit training loop");
let dev = match MlDevice::cuda(0) {
Ok(d) => d,
Err(_) => {
eprintln!("CUDA 0 not available — skipping toy_bandit_policy");
return;
}
};
let perception = PerceptionTrainerConfig {
seq_len: 4,
n_batch: 1,
..PerceptionTrainerConfig::default()
};
let cfg = IntegratedTrainerConfig {
perception,
dqn_seed: 0xC51F,
ppo_seed: 0xCAFE,
};
let mut trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new");
let snapshots = synthetic_window(4);
let mut env = MockLobEnv::toy_bandit();
let n_steps = 300;
for _ in 0..n_steps {
trainer
.step_with_lobsim(&snapshots, &mut env)
.expect("step_with_lobsim");
}
// After training, mode(π(s)) should be Action::LongSmall (5).
let probs = trainer
.eval_policy_probs_per_action(&snapshots)
.expect("eval_policy_probs_per_action");
let mode_a = probs
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(i, _)| i)
.unwrap();
println!(
"toy_bandit_policy: π = {:?}, mode = {} (expected {})",
probs, mode_a, Action::LongSmall as u32
);
assert_eq!(
mode_a,
Action::LongSmall as usize,
"PPO policy should learn mode(π) = LongSmall (5); got action {mode_a} with π = {probs:?}"
);
// Sanity: π(good_action) > 1/9 (uniform).
let good_prob = probs[Action::LongSmall as usize];
println!("toy_bandit_policy: π(LongSmall) = {:.3}", good_prob);
assert!(
good_prob > 1.0 / 9.0,
"policy should put more than uniform mass on the good action; got {good_prob:.3}"
);
}