Files
foxhunt/crates/ml/examples/alpha_compose_backtest.rs
jgrusewski d090685ca9 feat(phase-e-4-a-T16): vol-regime detection + cost-aware training
Adds two coupled interventions on the regime fragility exposed by
walk-forward CV (mean Sharpe +27 ± 56 at half-tick across 3 folds —
std-dev ≈ mean means the strategy is regime-dependent).

(1) Vol-EMA regime detector (new ISV slots 549/550)
- New alpha_regime_vol_update.cu kernel: per inference step, reads B
  parallel-env mid prices, computes cross-env mean squared log-return,
  and maintains ISV[549]=REGIME_VOL_EMA (Wiener-α with 0.4 floor +
  Pearl A bootstrap) and ISV[550]=REGIME_VOL_REF (slow tracker β=0.005,
  ≈200-step horizon).
- Block-tree-reduce (no atomicAdd), guards against zero/non-finite mids.

(2) Pre-emptive Kelly attenuation (modified stacker controller)
- stacker_threshold_controller.cu takes 3 new args: regime_vol_ema_idx,
  regime_vol_ref_idx, regime_scale_floor.
- Multiplies its reactive Sharpe-error Kelly output by
    regime_scale = clamp(vol_ref / vol_ema, 0.25, 1.0)
- Disabled when indices = -1 (backward-compatible smoke + kernel test).

(3) Cost-aware training (--train-cost-hi)
- alpha_compose_backtest --train-cost-hi: when > --train-cost, each
  training epoch samples cost ~ U[lo, hi] so the Q-network learns
  cost-conservative behaviour across the realistic ES range.

(4) Wiring
- alpha_compose_backtest --regime-scale enables both per-step regime
  kernel firing during eval AND the regime hookup in the per-episode
  controller call. Mapped-pinned mids buffers, all compute device-side.
- ExecutionEnv exposes current_mid() so the host gather reads the
  active snapshot mid per env without leaking the private cursor field.

Smoke + test sites pass -1/-1 for regime indices (backward compat).
Doc: docs/isv-slots.md ledger for slots 549/550.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-16 00:28:01 +02:00

1313 lines
59 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Phase E.3 Task 23 — Composition backtest with cost sweep.
//!
//! Trains the Phase E execution-policy DQN (linear Q + Phase 1d.3
//! alpha-cache + stabilizers) on the first `--train-frac` of the
//! fxcache, then evaluates the FROZEN policy (no SGD, ε=0 greedy) over
//! the held-out remainder at multiple transaction costs. Compares
//! per-cost annualized Sharpe vs the Phase 1d.4 "always-market-when-
//! confident" baseline (+4.4 frictionless, -4.0 at half-tick).
//!
//! Goal per the plan: lift the half-tick Sharpe above 0 — i.e., let
//! the execution-policy intelligence offset the cost the
//! threshold-only baseline can't.
//!
//! ## What's swept and what's frozen
//!
//! Frozen across costs: trained Q-network weights, fill-model
//! coefficients (from `alpha_fill_coeffs.json`), alpha-logit cache
//! (from `alpha_logits_cache.bin`). One policy evaluated at multiple
//! costs.
//!
//! Swept: only `ExecutionEnvConfig.cost_per_contract`. The env's
//! mutable config is updated between cost levels (snapshots stay in
//! place, cursor re-seeded each episode).
//!
//! ## Run
//!
//! ```bash
//! cargo run -p ml --release --example alpha_compose_backtest -- \
//! --fxcache-path /home/jgrusewski/Work/foxhunt/test_data/feature-cache/9297....fxcache \
//! --alpha-cache config/ml/alpha_logits_cache.bin \
//! --fill-coeffs config/ml/alpha_fill_coeffs.json
//! ```
use std::fs::File;
use std::io::Write;
use std::mem::MaybeUninit;
use std::path::PathBuf;
use anyhow::{Context, Result};
use clap::Parser;
use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut};
use tracing::info;
// Phase E.4.A.8/T12: Mamba2 temporal encoder (ml-alpha Phase 1d.1).
// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights.
use ml_alpha::mamba2_block::{
Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig,
};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
// ── Mapped-pinned helpers (mirror gpu_training_guard.rs MappedBuffer) ──
struct MappedI32 {
host_ptr: *mut i32,
dev_ptr: cudarc::driver::sys::CUdeviceptr,
len: usize,
}
impl MappedI32 {
unsafe fn new(len: usize) -> Result<Self> {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
| cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
let bytes = len * std::mem::size_of::<i32>();
let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
.map_err(|e| anyhow::anyhow!("mapped i32 alloc: {e}"))?
as *mut i32;
std::ptr::write_bytes(host_ptr, 0, len);
let mut dev_raw = MaybeUninit::uninit();
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
dev_raw.as_mut_ptr(),
host_ptr as *mut std::ffi::c_void,
0,
)
.result()
.map_err(|e| anyhow::anyhow!("cuMemHostGetDevicePointer: {e}"))?;
Ok(Self { host_ptr, dev_ptr: dev_raw.assume_init(), len })
}
fn dev_u64(&self) -> u64 { self.dev_ptr as u64 }
fn read(&self) -> i32 {
debug_assert!(self.len >= 1);
unsafe { std::ptr::read_volatile(self.host_ptr) }
}
/// Read all `len` entries as a slice. Caller must ensure the
/// associated GPU kernel has issued `__threadfence_system()` before
/// the host accesses these values (volatile reads still happen
/// per-element under the hood via `*host_ptr.add(i)`).
fn read_all(&self) -> Vec<i32> {
(0..self.len)
.map(|i| unsafe { std::ptr::read_volatile(self.host_ptr.add(i)) })
.collect()
}
}
impl Drop for MappedI32 {
fn drop(&mut self) {
unsafe {
let _ = cudarc::driver::result::free_host(self.host_ptr as *mut std::ffi::c_void);
}
}
}
struct MappedF32 {
host_ptr: *mut f32,
dev_ptr: cudarc::driver::sys::CUdeviceptr,
len: usize,
}
impl MappedF32 {
unsafe fn new(len: usize) -> Result<Self> {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
| cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
let bytes = len * std::mem::size_of::<f32>();
let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
.map_err(|e| anyhow::anyhow!("mapped f32 alloc({len}): {e}"))?
as *mut f32;
std::ptr::write_bytes(host_ptr, 0, len);
let mut dev_raw = MaybeUninit::uninit();
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
dev_raw.as_mut_ptr(),
host_ptr as *mut std::ffi::c_void,
0,
)
.result()
.map_err(|e| anyhow::anyhow!("cuMemHostGetDevicePointer: {e}"))?;
Ok(Self { host_ptr, dev_ptr: dev_raw.assume_init(), len })
}
fn dev_u64(&self) -> u64 { self.dev_ptr as u64 }
fn write(&self, data: &[f32]) {
debug_assert_eq!(data.len(), self.len);
unsafe { std::ptr::copy_nonoverlapping(data.as_ptr(), self.host_ptr, self.len); }
}
}
impl Drop for MappedF32 {
fn drop(&mut self) {
unsafe {
let _ = cudarc::driver::result::free_host(self.host_ptr as *mut std::ffi::c_void);
}
}
}
use ml::cuda_pipeline::alpha_isv_slots::{
RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX,
};
use ml::env::action_space::N_ACTIONS;
use ml::env::execution_env::{
EpisodeState, ExecutionEnv, ExecutionEnvConfig, ReplayRng, SnapshotRow,
};
const STATE_DIM: usize = 10;
const N_WEIGHTS: usize = N_ACTIONS * STATE_DIM;
const N_BIASES: usize = N_ACTIONS;
/// Action allow-list for the pruning experiment. Same shape as the smoke
/// binary: Q-net is unchanged (9 outputs), only the selector restricts.
const PRUNED_ACTIONS: [u8; 4] = [0, 1, 4, 7]; // Wait, BuyMarket, SellMarket, FlatMarket
const FULL_ACTIONS: [u8; 9] = [0, 1, 2, 3, 4, 5, 6, 7, 8];
#[derive(Debug, Parser)]
#[command(
name = "alpha_compose_backtest",
about = "Phase E.3 Task 23 — composition backtest with cost sweep"
)]
struct Cli {
#[arg(long)]
fxcache_path: PathBuf,
#[arg(long, default_value = "config/ml/alpha_fill_coeffs.json")]
fill_coeffs: PathBuf,
#[arg(long)]
alpha_cache: PathBuf,
/// Train segment fraction. First N% of snapshots used for DQN training,
/// the rest for evaluation.
#[arg(long, default_value_t = 0.8)]
train_frac: f32,
/// Snapshots to load from the fxcache.
#[arg(long, default_value_t = 1_500_000)]
max_snapshots: usize,
/// Bars to skip at the front of the fxcache before windowing. Used by
/// walk-forward CV scripts to slide a fixed-size train+eval window
/// across the fxcache so each fold sees a distinct chronological
/// region.
#[arg(long, default_value_t = 0)]
data_start_offset: usize,
/// Episode horizon in snapshots.
#[arg(long, default_value_t = 600)]
horizon: usize,
/// DQN training episodes (on train segment).
#[arg(long, default_value_t = 1_000)]
n_train_episodes: usize,
/// Frozen-policy evaluation episodes per cost level.
#[arg(long, default_value_t = 500)]
n_eval_episodes: usize,
/// Comma-separated cost grid (price units per contract round-turn).
/// Phase 1d.4 used [0.0, 0.0625, 0.125, 0.25, 0.50].
#[arg(long, value_delimiter = ',', default_value = "0.0,0.0625,0.125,0.25,0.5")]
cost_grid: Vec<f32>,
/// Comma-separated alpha-confidence threshold grid for the gate at eval.
/// Direct analogue of Phase 1d.4's `--threshold` sweep — at each
/// threshold, the policy is forced to Wait when |sigmoid(alpha)0.5|
/// < threshold. Threshold 0 = no gate (original Task 23 behaviour).
#[arg(long, value_delimiter = ',', default_value = "0.0,0.05,0.10,0.15,0.20,0.25")]
threshold_grid: Vec<f32>,
#[arg(long, default_value_t = 1)]
trade_size: i32,
#[arg(long, default_value_t = 0xCAFEBABE_u64)]
seed: u64,
/// Training-time cost (the policy LEARNED against this cost). Used as
/// a fixed cost when `--train-cost-hi == train_cost`; otherwise serves
/// as the lower bound of a uniform[lo, hi] random cost sampled per
/// training episode (Phase E.4.A T16 cost-aware training).
#[arg(long, default_value_t = 0.0625)]
train_cost: f32,
/// Upper bound on the per-episode random training cost. When > `train_cost`,
/// each training episode samples its `cost_per_contract` uniformly from
/// [train_cost, train_cost_hi] so the Q-network learns cost-conservative
/// behaviour across the realistic ES futures range (quarter-tick to
/// quad-tick). Defaults to the same value as `train_cost` (no randomization,
/// backward-compatible).
#[arg(long, default_value_t = 0.0625)]
train_cost_hi: f32,
/// Training-time alpha-confidence threshold for the gate. Forces
/// Wait during training when `|sigmoid(alpha)0.5| < this`, so the
/// Q-network learns weights for the gated policy class. Default
/// 0.39 — the equilibrium the controller stabilized to in the
/// Phase E.2 smoke (alpha_dqn_h600_smoke at ep 200+). Set to 0.0
/// for the original ungated-training behavior.
#[arg(long, default_value_t = 0.39)]
train_threshold: f32,
/// SGD learning rate during training.
#[arg(long, default_value_t = 1.0e-4)]
lr: f32,
#[arg(long, default_value_t = 0.50)]
eps_start: f32,
#[arg(long, default_value_t = 0.05)]
eps_end: f32,
#[arg(long, default_value_t = 0.99)]
gamma: f32,
#[arg(long, default_value_t = 0.9)]
alpha_m: f32,
#[arg(long, default_value_t = 0.03)]
tau: f32,
#[arg(long, default_value_t = -1.0)]
log_clip_min: f32,
#[arg(long, default_value_t = 1000.0)]
reward_scale: f32,
#[arg(long, default_value_t = 10)]
target_update_every: usize,
#[arg(long, default_value_t = 1.0)]
grad_clip: f32,
/// Restrict action selection to {Wait, BuyMarket, SellMarket, FlatMarket}.
/// Applies during BOTH training and eval. **FALSIFIED 2026-05-15** —
/// kept for reproducibility (see `pearl_action_pruning_falsified`).
#[arg(long, default_value_t = false)]
pruned_actions: bool,
/// Use C51 distributional Q-network (Phase E.3 follow-up). Same
/// semantics as the alpha_dqn_h600_smoke `--c51` flag.
#[arg(long, default_value_t = false)]
c51: bool,
/// C51 atom-support lower bound (normalized reward units).
#[arg(long, default_value_t = -10.0)]
c51_vmin: f32,
/// C51 atom-support upper bound.
#[arg(long, default_value_t = 10.0)]
c51_vmax: f32,
/// C51 atom count (canonical 51; kernel caps at 64).
#[arg(long, default_value_t = 51)]
c51_n_atoms: usize,
/// Derive bid/ask from real spread_bps in fxcache instead of fixed
/// ±0.125-tick. Phase E.3 Path 3 follow-up.
#[arg(long, default_value_t = false)]
real_spread: bool,
/// Phase E.4.A: enable the temporal-encoder stack (sliding window +
/// Mamba2 + C51).
#[arg(long, default_value_t = false)]
temporal: bool,
#[arg(long, default_value_t = 16)]
window_k: usize,
#[arg(long, default_value_t = 32)]
mamba2_hidden_dim: usize,
#[arg(long, default_value_t = 16)]
mamba2_state_dim: usize,
/// Phase E.4.A Pillar B: fire stacker-threshold controller at the
/// end of EVERY eval episode (not just training). Slot 543
/// (threshold) adapts during the 2D sweep based on observed trade
/// rate / Sharpe.
#[arg(long, default_value_t = false)]
isv_continual: bool,
/// Phase E.4.A T16: enable vol-EMA regime detection + pre-emptive
/// Kelly attenuation. When set:
/// - per inference step, `alpha_regime_vol_update_kernel` reads
/// the parallel envs' current/previous mids and maintains
/// ISV[549]=vol_ema (Wiener-α with 0.4 floor + Pearl A
/// bootstrap) and ISV[550]=vol_ref (slow tracker, β=0.005);
/// - per eval episode (requires `--isv-continual`), the existing
/// stacker controller multiplies its Sharpe-reactive Kelly
/// output by `clamp(vol_ref/vol_ema, 0.25, 1.0)` so vol spikes
/// cut Kelly *before* losses materialise.
/// Off by default to keep the T10 baseline comparable.
#[arg(long, default_value_t = false)]
regime_scale: bool,
/// Phase E.4.A.T15.batched (2026-05-15): parallel envs per training
/// "epoch". n_train_episodes / n_train_par epochs total. Each epoch
/// runs N_par envs in lockstep H steps + ONE batched C51 update at
/// B = N_par * H. Default 50 keeps SGD semantics (W is updated 20
/// times over 1000 training eps); raise to 1000 for one-epoch full-
/// batch training.
#[arg(long, default_value_t = 50)]
n_train_par: usize,
#[arg(long, default_value = "config/ml/alpha_compose_backtest.json")]
out_path: PathBuf,
}
struct SmokeRng {
state: u64,
}
impl SmokeRng {
fn new(seed: u64) -> Self {
Self { state: seed }
}
fn next_u64(&mut self) -> u64 {
self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
let mut z = self.state;
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
fn next_f32(&mut self) -> f32 {
((self.next_u64() >> 40) as f32) / ((1u64 << 24) as f32)
}
}
fn epsilon_greedy(q: &[f32], eps: f32, allowed: &[u8], rng: &mut SmokeRng) -> u8 {
if rng.next_f32() < eps {
allowed[(rng.next_u64() as usize) % allowed.len()]
} else {
let mut best_a: u8 = allowed[0];
let mut best_v: f32 = q[allowed[0] as usize];
for &a in &allowed[1..] {
let v = q[a as usize];
if v > best_v {
best_v = v;
best_a = a;
}
}
best_a
}
}
/// Phase E.3 confidence-gated ε-greedy. If `alpha_confidence < threshold`,
/// force action=0 (Wait). Otherwise standard ε-greedy over `allowed`.
fn epsilon_greedy_gated(
q: &[f32],
alpha_confidence: f32,
threshold: f32,
eps: f32,
allowed: &[u8],
rng: &mut SmokeRng,
) -> u8 {
if alpha_confidence < threshold {
return 0;
}
epsilon_greedy(q, eps, allowed, rng)
}
#[derive(Debug, Clone, serde::Serialize)]
struct CostBin {
cost: f32,
threshold: f32,
n_episodes: usize,
mean_reward: f32,
std_reward: f32,
sharpe_per_episode: f32,
/// Sharpe scaled by sqrt(episodes_per_year). Time span derived from
/// the eval window's mid_price index span × bar_seconds.
sharpe_annualised: f32,
win_rate: f32,
avg_n_trades: f32,
p05: f32,
p50: f32,
p95: f32,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();
let cli = Cli::parse();
info!("Phase E.3 Task 23 — composition backtest starting");
let allowed_actions: &[u8] = if cli.pruned_actions {
info!(" ACTION SET: pruned ({} actions: Wait, BuyMarket, SellMarket, FlatMarket)", PRUNED_ACTIONS.len());
&PRUNED_ACTIONS
} else {
info!(" ACTION SET: full ({} actions)", FULL_ACTIONS.len());
&FULL_ACTIONS
};
let ctx = CudaContext::new(0).context("CUDA init")?;
let stream = ctx.default_stream();
// --- Load cubins ---
let lq_module = ctx
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_LINEAR_Q_CUBIN.to_vec())
.context("alpha_linear_q cubin")?;
let lq_fwd = lq_module.load_function("alpha_linear_q_forward_kernel")?;
let lq_grad = lq_module.load_function("alpha_linear_q_grad_kernel")?;
let lq_sgd = lq_module.load_function("alpha_linear_q_sgd_step_kernel")?;
let lq_clip = lq_module.load_function("alpha_clip_inplace_kernel")?;
let munch_cubin: Vec<u8> = std::fs::read(concat!(
env!("OUT_DIR"),
"/alpha_munchausen_target.cubin"
))?;
let munch_module = ctx.load_cubin(munch_cubin)?;
let munch_kernel =
munch_module.load_function("alpha_munchausen_target_kernel")?;
// Phase E.3 follow-up: C51 kernels (always loaded; only used when --c51).
let c51_module = ctx
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_C51_CUBIN.to_vec())
.context("alpha_c51 cubin")?;
let c51_fwd_kernel = c51_module.load_function("alpha_c51_forward_kernel")?;
let c51_project_kernel = c51_module.load_function("alpha_c51_project_kernel")?;
let c51_grad_kernel = c51_module.load_function("alpha_c51_grad_kernel")?;
let c51_thompson_kernel = c51_module.load_function("alpha_c51_thompson_select_kernel")?;
// Phase E.4.A.6: window push kernel (used in --temporal path).
let push_module = ctx
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_WINDOW_PUSH_CUBIN.to_vec())
.context("alpha_window_push cubin")?;
let push_kernel = push_module.load_function("alpha_window_push_kernel")?;
let h_store_kernel = push_module.load_function("alpha_h_enriched_store_kernel")?;
let push_batched_kernel = push_module.load_function("alpha_window_push_batched_kernel")?;
let h_store_batched_kernel = push_module.load_function("alpha_h_enriched_store_batched_kernel")?;
// Phase E.4.A.7 controller (used in --isv-continual eval path).
let ctl_module = ctx
.load_cubin(ml::cuda_pipeline::alpha_kernels::STACKER_THRESHOLD_CONTROLLER_CUBIN.to_vec())
.context("stacker controller cubin")?;
let ctl_kernel = ctl_module
.load_function("stacker_threshold_controller_update")?;
let c51_n_atoms = cli.c51_n_atoms;
let c51_v_min = cli.c51_vmin;
let c51_v_max = cli.c51_vmax;
let c51_delta_z = (c51_v_max - c51_v_min) / (c51_n_atoms.saturating_sub(1).max(1) as f32);
if cli.c51 {
info!(
" Q-network: C51 ({} atoms, [{:.2}, {:.2}], Δz={:.4})",
c51_n_atoms, c51_v_min, c51_v_max, c51_delta_z
);
if c51_n_atoms == 0 || c51_n_atoms > 64 {
anyhow::bail!("--c51-n-atoms must be in (0, 64]");
}
if c51_v_max <= c51_v_min {
anyhow::bail!("--c51-vmax must exceed --c51-vmin");
}
} else {
info!(" Q-network: linear scalar (9 outputs)");
}
// --- Load env data ---
let fill_model = ml::env::loaders::load_fill_model_from_json(&cli.fill_coeffs)?;
info!("Loaded fill model");
let alpha_cache = ml::env::loaders::load_alpha_cache(&cli.alpha_cache)?;
info!("Loaded alpha cache: {} entries", alpha_cache.len());
let rows = ml::env::loaders::load_snapshots_from_fxcache_at(
&cli.fxcache_path,
cli.data_start_offset,
cli.max_snapshots,
Some(&alpha_cache),
cli.real_spread,
None, // E.4.A T4: mbp10_dir wiring + --use-real-depth flag lands in T5
)?;
let n_total = rows.len();
info!("Loaded {} snapshots", n_total);
let n_train = ((n_total as f32) * cli.train_frac) as usize;
let n_eval = n_total - n_train;
if n_train <= cli.horizon || n_eval <= cli.horizon {
anyhow::bail!(
"insufficient snapshots: train={}, eval={}, horizon={}",
n_train,
n_eval,
cli.horizon
);
}
info!(
"Split: train={} bars (cursor 0..{}), eval={} bars (cursor {}..{})",
n_train, n_train, n_eval, n_train, n_total
);
let mut env = ExecutionEnv::new(
ExecutionEnvConfig {
horizon_snapshots: cli.horizon,
trade_size_contracts: cli.trade_size,
cost_per_contract: cli.train_cost,
},
fill_model,
rows,
cli.seed,
);
// --- Initialize Q-network ---
let c51_input_dim: usize = if cli.c51 && cli.temporal {
cli.mamba2_hidden_dim
} else {
STATE_DIM
};
let n_weights_eff: usize = if cli.c51 {
N_ACTIONS * c51_n_atoms * c51_input_dim
} else {
N_WEIGHTS
};
let n_biases_eff: usize = if cli.c51 { N_ACTIONS * c51_n_atoms } else { N_BIASES };
let mut rng = SmokeRng::new(cli.seed.wrapping_add(0xDEAD_BEEF));
let xavier_scale = (2.0_f32 / c51_input_dim as f32).sqrt();
let w_init: Vec<f32> = (0..n_weights_eff)
.map(|_| xavier_scale * 2.0 * (rng.next_f32() - 0.5))
.collect();
let b_init: Vec<f32> = vec![0.0; n_biases_eff];
let mut w_dev = stream.clone_htod(&w_init)?;
let mut b_dev = stream.clone_htod(&b_init)?;
let mut w_target_dev = stream.clone_htod(&w_init)?;
let mut b_target_dev = stream.clone_htod(&b_init)?;
let mut dw_dev = stream.alloc_zeros::<f32>(n_weights_eff)?;
let mut db_dev = stream.alloc_zeros::<f32>(n_biases_eff)?;
let state_dim_i = STATE_DIM as i32;
let n_act_i = N_ACTIONS as i32;
let n_atoms_i = c51_n_atoms as i32;
let mut states_dev = stream.alloc_zeros::<f32>(cli.horizon * STATE_DIM)?;
let mut next_states_dev = stream.alloc_zeros::<f32>(cli.horizon * STATE_DIM)?;
let mut actions_dev = stream.alloc_zeros::<i32>(cli.horizon)?;
let mut rewards_dev = stream.alloc_zeros::<f32>(cli.horizon)?;
let mut dones_dev = stream.alloc_zeros::<f32>(cli.horizon)?;
let mut q_current_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS)?;
let mut q_next_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS)?;
let mut target_dev = stream.alloc_zeros::<f32>(cli.horizon)?;
let mut single_state_dev = stream.alloc_zeros::<f32>(STATE_DIM)?;
let mut single_q_dev = stream.alloc_zeros::<f32>(N_ACTIONS)?;
// C51 device buffers (always allocated; cheap).
let mut probs_current_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS * c51_n_atoms)?;
let mut probs_next_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS * c51_n_atoms)?;
let mut m_dev = stream.alloc_zeros::<f32>(cli.horizon * c51_n_atoms)?;
let mut single_probs_dev = stream.alloc_zeros::<f32>(N_ACTIONS * c51_n_atoms)?;
// Mapped-pinned per-step inference buffers (C51 path).
let state_pinned = unsafe { MappedF32::new(STATE_DIM)? };
let action_pinned = unsafe { MappedI32::new(1)? };
// Phase E.4.A.8 (T12): temporal buffers, allocated unconditionally.
let mut window_tensor = GpuTensor::zeros(
&[1, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc window: {e}"))?;
let h_enriched_buf_capacity = (cli.horizon + 1) * cli.mamba2_hidden_dim;
let mut h_enriched_buf_dev = stream
.alloc_zeros::<f32>(h_enriched_buf_capacity)
.context("alloc h_enriched buffer")?;
let mut mamba2_block: Option<Mamba2Block> = if cli.temporal {
let cfg = Mamba2BlockConfig {
in_dim: STATE_DIM,
hidden_dim: cli.mamba2_hidden_dim,
state_dim: cli.mamba2_state_dim,
seq_len: cli.window_k,
};
info!("TEMPORAL Mamba2: in={} hidden={} state={} K={}",
cfg.in_dim, cfg.hidden_dim, cfg.state_dim, cfg.seq_len);
Some(Mamba2Block::new(cfg, stream.clone())
.map_err(|e| anyhow::anyhow!("Mamba2Block init: {e}"))?)
} else {
None
};
// Phase E.4.A.T10: AdamW for Mamba2 weights (only when --temporal).
let mut mamba2_adamw: Option<Mamba2AdamW> = if let Some(block) = mamba2_block.as_ref() {
Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default())
.map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?)
} else {
None
};
// Load the C51 grad-input kernel (for backward chain into Mamba2).
let c51_grad_input_kernel = c51_module
.load_function("alpha_c51_grad_input_kernel")
.context("c51 grad_input load")?;
// T10: train-time window store (captures windows during inference for
// the end-of-epoch batched Mamba2 forward).
let train_window_store_kernel = push_module
.load_function("alpha_train_window_store_batched_kernel")
.context("train_window_store kernel load")?;
let n_atoms_i = c51_n_atoms as i32;
// ISV buffer + Wiener state for the eval-time controller path.
let mut isv_host: Vec<f32> = vec![0.0; 552];
isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_MEAN_INDEX] = -5191.53;
isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX] = 4963.62;
isv_host[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX] = 0.08;
let mut isv_dev = stream.clone_htod(&isv_host).context("upload isv")?;
let mut ctl_wiener_dev = stream.alloc_zeros::<f32>(3).context("alloc ctl wiener")?;
// ============================================================
// Phase 1: TRAIN DQN — BATCHED PARALLEL ENVS (Phase E.4.A.T15)
// ============================================================
let train_n_par = cli.n_train_par;
let train_n_epochs = (cli.n_train_episodes + train_n_par - 1) / train_n_par;
info!("=== Training phase: {} epochs × N_par={} parallel envs × H={} steps (BATCHED) ===",
train_n_epochs, train_n_par, cli.horizon);
let train_max_start = (n_train.saturating_sub(cli.horizon + 1)).max(1);
let train_h = cli.horizon;
let train_batch = train_n_par * train_h;
let mut states_train_dev = stream.alloc_zeros::<f32>(train_batch * STATE_DIM)?;
let mut next_states_train_dev = stream.alloc_zeros::<f32>(train_batch * STATE_DIM)?;
let mut actions_train_dev = stream.alloc_zeros::<i32>(train_batch)?;
let mut rewards_train_dev = stream.alloc_zeros::<f32>(train_batch)?;
let mut dones_train_dev = stream.alloc_zeros::<f32>(train_batch)?;
let mut probs_curr_train_dev = stream.alloc_zeros::<f32>(train_batch * N_ACTIONS * c51_n_atoms)?;
let mut probs_next_train_dev = stream.alloc_zeros::<f32>(train_batch * N_ACTIONS * c51_n_atoms)?;
let mut m_train_dev = stream.alloc_zeros::<f32>(train_batch * c51_n_atoms)?;
let h_enriched_train_capacity = (train_h + 1) * train_n_par * cli.mamba2_hidden_dim;
let mut h_enriched_train_dev = stream.alloc_zeros::<f32>(h_enriched_train_capacity)?;
let batched_state_pinned_train = unsafe { MappedF32::new(train_n_par * STATE_DIM)? };
let batched_action_pinned_train = unsafe { MappedI32::new(train_n_par)? };
let mut batched_window_tensor_train = GpuTensor::zeros(
&[train_n_par, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc batched window train: {e}"))?;
// T10 buffers: all-step train windows and d_h_enriched for the backward chain.
let mut train_windows_tensor_train = GpuTensor::zeros(
&[train_batch, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc train windows train: {e}"))?;
let mut d_h_enriched_tensor_train = GpuTensor::zeros(
&[train_batch, cli.mamba2_hidden_dim], &stream
).map_err(|e| anyhow::anyhow!("alloc d_h_enriched train: {e}"))?;
let mut batched_probs_inference_train = stream
.alloc_zeros::<f32>(train_n_par * N_ACTIONS * c51_n_atoms)?;
let snapshots_arc_train = env.snapshots_arc();
let fill_model_for_train = env.fill_model.clone();
let env_config_template_train = env.config.clone();
let mut episode_rng = SmokeRng::new(cli.seed.wrapping_add(0xFEED));
for epoch in 0..train_n_epochs {
let _eps = cli.eps_start + (cli.eps_end - cli.eps_start)
* (epoch as f32 / train_n_epochs.max(1) as f32);
// T16 cost-aware training: when --train-cost-hi > --train-cost,
// sample the per-episode cost uniformly from that range, so the
// policy learns to behave robustly across the realistic ES cost
// spectrum (quarter-tick to quad-tick) instead of memorizing a
// single low-cost world.
let cost_for_epoch: f32 = if cli.train_cost_hi > cli.train_cost {
let u = (episode_rng.next_u64() as f32) / (u64::MAX as f32);
cli.train_cost + u * (cli.train_cost_hi - cli.train_cost)
} else {
cli.train_cost
};
let mut par_envs: Vec<ExecutionEnv> = (0..train_n_par).map(|_| {
let mut cfg = env_config_template_train.clone();
cfg.cost_per_contract = cost_for_epoch;
ExecutionEnv::new_arc(
cfg, fill_model_for_train.clone(),
std::sync::Arc::clone(&snapshots_arc_train),
0,
)
}).collect();
for env_i in par_envs.iter_mut() {
let start = (episode_rng.next_u64() as usize) % train_max_start;
let seed = episode_rng.next_u64();
env_i.reset_at(seed, start);
}
let mut par_states: Vec<EpisodeState> = vec![EpisodeState::new(); train_n_par];
let mut done_flags = vec![false; train_n_par];
if cli.temporal {
stream.memset_zeros(batched_window_tensor_train.data_mut())?;
stream.memset_zeros(train_windows_tensor_train.data_mut())?;
stream.memset_zeros(d_h_enriched_tensor_train.data_mut())?;
}
stream.memset_zeros(&mut h_enriched_train_dev)?;
let mut states_host_train: Vec<f32> = vec![0.0; train_batch * STATE_DIM];
let mut next_states_host_train: Vec<f32> = vec![0.0; train_batch * STATE_DIM];
let mut actions_host_train: Vec<i32> = vec![0; train_batch];
let mut rewards_host_train: Vec<f32> = vec![0.0; train_batch];
let mut dones_host_train: Vec<f32> = vec![0.0; train_batch];
for step in 0..train_h {
let mut batched_states_host = vec![0.0_f32; train_n_par * STATE_DIM];
for i in 0..train_n_par {
if !done_flags[i] {
let s = par_envs[i].state(&par_states[i]);
batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM].copy_from_slice(&s);
}
}
batched_state_pinned_train.write(&batched_states_host);
if cli.temporal {
{
let (w_ptr, _g) = batched_window_tensor_train.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push_batched(
&stream, &push_batched_kernel,
batched_state_pinned_train.dev_u64(), w_ptr,
train_n_par as i32, cli.window_k as i32, state_dim_i,
)?;
}
}
// T10: capture post-push windows into [H*N_par, K, in_dim] buffer
// at step row offset `step * N_par`.
{
let (src_ptr, _g_src) = batched_window_tensor_train.data().device_ptr(&stream);
let (dst_ptr, _g_dst) = train_windows_tensor_train.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched(
&stream, &train_window_store_kernel,
src_ptr, dst_ptr, (step * train_n_par) as i32,
train_n_par as i32, cli.window_k as i32, state_dim_i,
)?;
}
}
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&batched_window_tensor_train)
.map_err(|e| anyhow::anyhow!("train mamba2: {e}"))?;
{
let (src_p, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (buf_p, _g_buf) = h_enriched_train_dev.device_ptr_mut(&stream);
let step_row_offset = (step * train_n_par) as i32;
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store_batched(
&stream, &h_store_batched_kernel,
src_p, buf_p, step_row_offset,
train_n_par as i32, cli.mamba2_hidden_dim as i32,
)?;
}
}
let (h_ptr, _g_h) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = batched_probs_inference_train.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, h_ptr, p_ptr,
train_n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
} else {
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = batched_probs_inference_train.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, batched_state_pinned_train.dev_u64(), p_ptr,
train_n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
}
let step_seed = episode_rng.next_u64() as u32;
{
let (p_ptr, _g0) = batched_probs_inference_train.device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
&stream, &c51_thompson_kernel,
p_ptr,
batched_state_pinned_train.dev_u64(),
cli.train_threshold,
1, state_dim_i, c51_v_min, c51_delta_z, step_seed,
batched_action_pinned_train.dev_u64(),
train_n_par as i32, n_act_i, n_atoms_i,
)?;
}
}
stream.synchronize()?;
let actions = batched_action_pinned_train.read_all();
for i in 0..train_n_par {
if done_flags[i] { continue; }
let s = &batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM];
let row = step * train_n_par + i;
states_host_train[row * STATE_DIM..(row + 1) * STATE_DIM].copy_from_slice(s);
actions_host_train[row] = actions[i];
let action = actions[i] as u8;
let (next_state_arr, reward, done) = par_envs[i].step(action, &mut par_states[i])
.ok_or_else(|| anyhow::anyhow!("env step None"))?;
next_states_host_train[row * STATE_DIM..(row + 1) * STATE_DIM]
.copy_from_slice(&next_state_arr);
rewards_host_train[row] = reward / cli.reward_scale;
dones_host_train[row] = if done { 1.0 } else { 0.0 };
if done {
done_flags[i] = true;
}
}
}
stream.memcpy_htod(&states_host_train, &mut states_train_dev)?;
stream.memcpy_htod(&next_states_host_train, &mut next_states_train_dev)?;
stream.memcpy_htod(&actions_host_train, &mut actions_train_dev)?;
stream.memcpy_htod(&rewards_host_train, &mut rewards_train_dev)?;
stream.memcpy_htod(&dones_host_train, &mut dones_train_dev)?;
let bt = train_batch as i32;
// T10: re-forward Mamba2 on collected train windows to get a cache
// that backward_from_h_enriched can consume. cache.h_enriched
// replaces h_enriched_train_dev curr offset (bit-identical since
// Mamba2 weights haven't been updated yet this epoch).
let cache_train: Option<ml_alpha::mamba2_block::Mamba2ForwardCache> = if cli.temporal {
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&train_windows_tensor_train)
.map_err(|e| anyhow::anyhow!("train mamba2 (epoch forward): {e}"))?;
Some(cache)
} else {
None
};
let (curr_input_ptr_guard, next_input_ptr_guard);
let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal {
let cache = cache_train.as_ref().expect("cache_train missing");
let (h_curr, g) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (h_base_next, g_next) = h_enriched_train_dev.device_ptr(&stream);
curr_input_ptr_guard = g;
next_input_ptr_guard = g_next;
let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
(
h_curr,
h_base_next + (train_n_par as u64) * (cli.mamba2_hidden_dim as u64) * 4u64,
)
} else {
let (s_ptr, g_c) = states_train_dev.device_ptr(&stream);
let (n_ptr, g_n) = next_states_train_dev.device_ptr(&stream);
curr_input_ptr_guard = g_c;
next_input_ptr_guard = g_n;
let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
(s_ptr, n_ptr)
};
{
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = probs_curr_train_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, curr_input_ptr, p_ptr,
bt, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
}
{
let (w_ptr, _g0) = w_target_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_target_dev.device_ptr(&stream);
let (p_ptr, _g3) = probs_next_train_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, next_input_ptr, p_ptr,
bt, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
}
{
let (pn_ptr, _g0) = probs_next_train_dev.device_ptr(&stream);
let (r_ptr, _g1) = rewards_train_dev.device_ptr(&stream);
let (d_ptr, _g2) = dones_train_dev.device_ptr(&stream);
let (m_ptr, _g3) = m_train_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_project(
&stream, &c51_project_kernel,
pn_ptr, r_ptr, d_ptr,
c51_v_min, c51_v_max, cli.gamma, c51_delta_z,
m_ptr, bt, n_act_i, n_atoms_i,
)?;
}
}
{
let (p_ptr, _g0) = probs_curr_train_dev.device_ptr(&stream);
let (m_ptr, _g1) = m_train_dev.device_ptr(&stream);
let (a_ptr, _g2) = actions_train_dev.device_ptr(&stream);
let (s_ptr, _g3) = states_train_dev.device_ptr(&stream);
let (dw_ptr, _g4) = dw_dev.device_ptr_mut(&stream);
let (db_ptr, _g5) = db_dev.device_ptr_mut(&stream);
let grad_input_ptr = if cli.temporal { curr_input_ptr } else { s_ptr };
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad(
&stream, &c51_grad_kernel,
p_ptr, m_ptr, a_ptr, grad_input_ptr, dw_ptr, db_ptr,
bt, c51_input_dim as i32, n_act_i, n_atoms_i,
1.0 / bt as f32,
)?;
}
}
// T10 backward chain: C51 grad_input → Mamba2 backward → AdamW step.
if cli.temporal {
{
let (p_ptr, _g0) = probs_curr_train_dev.device_ptr(&stream);
let (m_ptr, _g1) = m_train_dev.device_ptr(&stream);
let (a_ptr, _g2) = actions_train_dev.device_ptr(&stream);
let (w_ptr, _g3) = w_dev.device_ptr(&stream);
let (dh_ptr, _g4) = d_h_enriched_tensor_train.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input(
&stream, &c51_grad_input_kernel,
p_ptr, m_ptr, a_ptr, w_ptr, dh_ptr,
bt, c51_input_dim as i32, n_act_i, n_atoms_i,
1.0 / bt as f32,
)?;
}
}
let cache = cache_train.as_ref().expect("cache_train missing");
let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing");
let grads = block_ref
.backward_from_h_enriched(cache, &d_h_enriched_tensor_train)
.map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?;
let _ = cache_train;
let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing");
let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing");
adamw
.step(block_mut, &grads)
.map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?;
}
{
let (dw_ptr, _g0) = dw_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
&stream, &lq_clip, dw_ptr, cli.grad_clip, n_weights_eff as i32,
)?;
}
}
{
let (db_ptr, _g0) = db_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
&stream, &lq_clip, db_ptr, cli.grad_clip, n_biases_eff as i32,
)?;
}
}
{
let (w_ptr, _g0) = w_dev.device_ptr_mut(&stream);
let (dw_ptr, _g1) = dw_dev.device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_sgd_step(
&stream, &lq_sgd, w_ptr, dw_ptr, cli.lr, n_weights_eff as i32,
)?;
}
}
{
let (b_ptr, _g0) = b_dev.device_ptr_mut(&stream);
let (db_ptr, _g1) = db_dev.device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_sgd_step(
&stream, &lq_sgd, b_ptr, db_ptr, cli.lr, n_biases_eff as i32,
)?;
}
}
if (epoch + 1) % cli.target_update_every == 0 {
stream.synchronize()?;
let w_now = stream.clone_dtoh(&w_dev)?;
let b_now = stream.clone_dtoh(&b_dev)?;
stream.memcpy_htod(&w_now, &mut w_target_dev)?;
stream.memcpy_htod(&b_now, &mut b_target_dev)?;
}
if (epoch + 1) % 5 == 0 || epoch + 1 == train_n_epochs {
info!(" train epoch {}/{}", epoch + 1, train_n_epochs);
}
}
info!("Training complete. Frozen policy ready for eval.");
// ---------------------------------------------------------------
// Phase 2: 2D SWEEP — frozen-policy greedy eval per (threshold, cost).
// Direct analogue of Phase 1d.4's threshold × cost table.
//
// Phase E.4.A.T14.batched (2026-05-15): batched eval — N=n_eval_episodes
// environments run in LOCKSTEP per cell, with ONE GPU forward + ONE
// sync per step (instead of N forwards × N syncs). Snapshots shared
// across envs via Arc to avoid 50MB × N duplication.
// ---------------------------------------------------------------
info!("=== Eval phase: {} episodes × {} thresholds × {} costs (batched) ===",
cli.n_eval_episodes, cli.threshold_grid.len(), cli.cost_grid.len());
let eval_max_start = (n_total - n_train).saturating_sub(cli.horizon + 1).max(1);
let n_par = cli.n_eval_episodes;
// Batched-eval GPU buffers (allocated once, reused across cells).
let batched_state_pinned = unsafe { MappedF32::new(n_par * STATE_DIM)? };
let batched_action_pinned = unsafe { MappedI32::new(n_par)? };
let mut batched_window_tensor = GpuTensor::zeros(
&[n_par, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc batched window: {e}"))?;
let mut batched_probs_dev = stream
.alloc_zeros::<f32>(n_par * N_ACTIONS * c51_n_atoms)
.context("alloc batched probs")?;
// T16 regime kernel state. mids buffers are mapped-pinned so the host
// gather (one f32 per env per step) writes directly into device-mapped
// memory with no htod copy.
let mids_curr_pinned = unsafe { MappedF32::new(n_par)? };
let mids_prev_pinned = unsafe { MappedF32::new(n_par)? };
let regime_module = if cli.regime_scale {
Some(ctx.load_cubin(
ml::cuda_pipeline::alpha_kernels::ALPHA_REGIME_VOL_UPDATE_CUBIN.to_vec()
).context("alpha_regime_vol_update cubin load")?)
} else { None };
let regime_kernel = regime_module.as_ref().map(|m|
m.load_function("alpha_regime_vol_update_kernel")
.expect("alpha_regime_vol_update_kernel load")
);
let snapshots_arc = env.snapshots_arc();
let fill_model_for_par = env.fill_model.clone();
let env_config_template = env.config.clone();
let mut bins: Vec<CostBin> =
Vec::with_capacity(cli.cost_grid.len() * cli.threshold_grid.len());
for &threshold in &cli.threshold_grid {
for &cost in &cli.cost_grid {
env.config.cost_per_contract = cost;
// Build N parallel envs sharing snapshots; reset each at a
// random eval-start cursor + seed.
let mut par_envs: Vec<ExecutionEnv> = (0..n_par).map(|_| {
let mut cfg = env_config_template.clone();
cfg.cost_per_contract = cost;
ExecutionEnv::new_arc(
cfg, fill_model_for_par.clone(),
std::sync::Arc::clone(&snapshots_arc),
0, // seed re-set below
)
}).collect();
for env_i in par_envs.iter_mut() {
let start = n_train + (episode_rng.next_u64() as usize) % eval_max_start;
let seed = episode_rng.next_u64();
env_i.reset_at(seed, start);
}
let mut par_states: Vec<EpisodeState> = vec![EpisodeState::new(); n_par];
let mut terminal_rs = vec![0.0_f32; n_par];
let mut ep_n_trades_par = vec![0_u32; n_par];
let mut done_flags = vec![false; n_par];
// Zero batched windows.
if cli.temporal {
stream.memset_zeros(batched_window_tensor.data_mut())
.context("zero batched windows")?;
}
// T16: reset prev_mids so the first-step Pearl-A bootstrap fires
// cleanly (kernel guards against curr<=0 || prev<=0 internally).
if cli.regime_scale {
let zeros = vec![0.0_f32; n_par];
mids_curr_pinned.write(&zeros);
mids_prev_pinned.write(&zeros);
}
// Lockstep step loop.
for _step in 0..cli.horizon {
// Gather current states (CPU; <100μs for N=500).
let mut batched_states_host = vec![0.0_f32; n_par * STATE_DIM];
for i in 0..n_par {
if !done_flags[i] {
let s = par_envs[i].state(&par_states[i]);
batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM].copy_from_slice(&s);
}
}
batched_state_pinned.write(&batched_states_host);
// T16: capture current mids + fire regime kernel BEFORE the
// env step advances. Regime signal therefore reflects vol
// observed up to and including this step, which is what the
// controller's per-episode Kelly update will read.
if let Some(rk) = regime_kernel.as_ref() {
let mut mids_host = vec![0.0_f32; n_par];
for i in 0..n_par {
if !done_flags[i] {
mids_host[i] = par_envs[i].current_mid();
}
}
mids_curr_pinned.write(&mids_host);
unsafe {
let (isv_ptr, _g) = isv_dev.device_ptr_mut(&stream);
ml::cuda_pipeline::alpha_kernels::launch_alpha_regime_vol_update(
&stream, rk,
mids_curr_pinned.dev_u64(),
mids_prev_pinned.dev_u64(),
isv_ptr,
n_par as i32,
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_EMA_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_INDEX as i32,
0.005, // β: 200-step horizon for vol_ref tracking
)?;
}
// Swap prev ← curr for the next step. Mapped-pinned write
// is a CPU-side fill, no device-to-device copy needed.
mids_prev_pinned.write(&mids_host);
}
if cli.temporal {
// Batched window push.
{
let (w_ptr, _g) = batched_window_tensor.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push_batched(
&stream, &push_batched_kernel,
batched_state_pinned.dev_u64(), w_ptr,
n_par as i32, cli.window_k as i32, state_dim_i,
)?;
}
}
// Batched Mamba2 forward → cache.h_enriched [N, hidden].
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&batched_window_tensor)
.map_err(|e| anyhow::anyhow!("batched mamba2 forward (eval): {e}"))?;
// Batched C51 forward.
let (h_ptr, _g_h) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = batched_probs_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, h_ptr, p_ptr,
n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
} else {
// C51 reads directly from batched_state_pinned (B=N).
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = batched_probs_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, batched_state_pinned.dev_u64(), p_ptr,
n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
}
// Batched Thompson select.
let step_seed = episode_rng.next_u64() as u32;
{
let (p_ptr, _g0) = batched_probs_dev.device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
&stream, &c51_thompson_kernel,
p_ptr,
batched_state_pinned.dev_u64(),
threshold,
1,
state_dim_i,
c51_v_min, c51_delta_z,
step_seed,
batched_action_pinned.dev_u64(),
n_par as i32, n_act_i, n_atoms_i,
)?;
}
}
// ONE sync per step (instead of N).
stream.synchronize()?;
let actions = batched_action_pinned.read_all();
// Step all envs on CPU.
for i in 0..n_par {
if done_flags[i] { continue; }
let action = actions[i] as u8;
if action != 0 { ep_n_trades_par[i] += 1; }
if let Some((_, r, done)) = par_envs[i].step(action, &mut par_states[i]) {
if done {
terminal_rs[i] = r;
done_flags[i] = true;
}
}
}
}
// Aggregate stats. Declared mut to keep the dead-code legacy
// block (under `if false`) borrow-check clean.
#[allow(unused_mut)]
let mut rewards: Vec<f32> = terminal_rs;
#[allow(unused_mut)]
let mut trade_count_total: u64 = ep_n_trades_par.iter().map(|&t| t as u64).sum();
#[allow(unused_mut)]
let mut win_count: usize = rewards.iter().filter(|&&r| r > 0.0).count();
// ISV-continual at cell-level (single fire with aggregate stats).
if cli.isv_continual {
let mean_terminal_r = rewards.iter().sum::<f32>() / (n_par as f32);
let trade_count_avg = trade_count_total as f32 / (n_par as f32);
let decisions_avg = cli.horizon as f32;
let baseline_std = isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX].max(1e-6);
let rollout_sharpe = mean_terminal_r / baseline_std;
unsafe {
let (isv_ptr, _gi) = isv_dev.device_ptr_mut(&stream);
let (wv_ptr, _gw) = ctl_wiener_dev.device_ptr_mut(&stream);
ml::cuda_pipeline::alpha_kernels::launch_stacker_threshold_controller(
&stream, &ctl_kernel,
trade_count_avg, decisions_avg, rollout_sharpe,
0.5, 0.01, 0.005, 0.4, 0.1,
ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX as i32,
// T16 regime hookup. Active when --isv-continual is on;
// the regime kernel writes slots 549/550 per-step during
// inference. Floor 0.25 = max defensive 4× Kelly cut.
if cli.regime_scale {
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_EMA_INDEX as i32
} else { -1 },
if cli.regime_scale {
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_INDEX as i32
} else { -1 },
0.25,
isv_ptr, wv_ptr,
)?;
}
}
let n = rewards.len() as f64;
let mean = rewards.iter().map(|r| *r as f64).sum::<f64>() / n;
let var = rewards
.iter()
.map(|r| (*r as f64 - mean).powi(2))
.sum::<f64>()
/ n;
let std = var.sqrt().max(1e-6);
let sharpe_per = (mean / std) as f32;
let episodes_per_year = 252.0 * 6.5 * 3600.0 / (cli.horizon as f64 * 12.0);
let sharpe_ann = (sharpe_per as f64 * episodes_per_year.sqrt()) as f32;
let win_rate = win_count as f32 / cli.n_eval_episodes as f32;
let avg_n_trades = trade_count_total as f32 / cli.n_eval_episodes as f32;
let mut sorted = rewards.clone();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let pick = |q: f64| -> f32 {
let idx = (q * (n - 1.0)).round() as usize;
sorted[idx.min(sorted.len() - 1)]
};
let bin = CostBin {
cost,
threshold,
n_episodes: cli.n_eval_episodes,
mean_reward: mean as f32,
std_reward: std as f32,
sharpe_per_episode: sharpe_per,
sharpe_annualised: sharpe_ann,
win_rate,
avg_n_trades,
p05: pick(0.05),
p50: pick(0.50),
p95: pick(0.95),
};
info!(
" τ={:.2} cost={:>6.4} mean={:>+9.2} Sharpe/ep={:+.3} Sharpe_ann={:+.3} win={:.3} trades/ep={:.1}",
threshold, cost, mean, sharpe_per, sharpe_ann, win_rate, avg_n_trades
);
bins.push(bin);
}
}
// --- Print table ---
info!("");
info!("=== Phase E.3 2D sweep (threshold × cost, vs Phase 1d.4 baseline) ===");
info!(
" {:>6} {:>6} {:>9} {:>9} {:>10} {:>11} {:>8} {:>10}",
"τ", "cost", "mean_R", "std_R", "Sharpe/ep", "Sharpe_ann", "win_rate", "trades/ep"
);
for b in &bins {
info!(
" {:>6.3} {:>6.4} {:>9.2} {:>9.2} {:>10.4} {:>11.4} {:>8.3} {:>10.2}",
b.threshold, b.cost, b.mean_reward, b.std_reward,
b.sharpe_per_episode, b.sharpe_annualised, b.win_rate, b.avg_n_trades
);
}
info!("");
info!("BEST per-cost (max Sharpe_ann across all τ at each cost):");
for &cost in &cli.cost_grid {
let best = bins
.iter()
.filter(|b| (b.cost - cost).abs() < 1e-6)
.max_by(|a, b| {
a.sharpe_annualised
.partial_cmp(&b.sharpe_annualised)
.unwrap_or(std::cmp::Ordering::Equal)
});
if let Some(b) = best {
info!(
" cost={:>6.4} best τ={:.3} Sharpe_ann={:+.3} win={:.3} trades/ep={:.1}",
b.cost, b.threshold, b.sharpe_annualised, b.win_rate, b.avg_n_trades
);
}
}
info!("");
info!("Phase 1d.4 baseline for comparison: +4.4 annualised at cost=0,");
info!(" -4.0 annualised at cost=0.125.");
// --- Save JSON ---
let json = serde_json::json!({
"phase": "E.3 Task 23 (2D sweep)",
"pruned_actions": cli.pruned_actions,
"n_allowed_actions": allowed_actions.len(),
"c51": cli.c51,
"c51_n_atoms": if cli.c51 { c51_n_atoms } else { 0 },
"c51_vmin": if cli.c51 { c51_v_min } else { 0.0 },
"c51_vmax": if cli.c51 { c51_v_max } else { 0.0 },
"temporal": cli.temporal,
"window_k": if cli.temporal { cli.window_k } else { 0 },
"mamba2_hidden_dim": if cli.temporal { cli.mamba2_hidden_dim } else { 0 },
"mamba2_state_dim": if cli.temporal { cli.mamba2_state_dim } else { 0 },
"isv_continual": cli.isv_continual,
"horizon": cli.horizon,
"train_frac": cli.train_frac,
"n_train_episodes": cli.n_train_episodes,
"n_eval_episodes": cli.n_eval_episodes,
"train_cost": cli.train_cost,
"cost_grid": cli.cost_grid,
"threshold_grid": cli.threshold_grid,
"bins": bins,
});
let mut f = File::create(&cli.out_path)?;
write!(f, "{}", serde_json::to_string_pretty(&json)?)?;
info!("Wrote table to {}", cli.out_path.display());
// Touch unused ISV slot constants so they're imported for future expansion.
let _ = (RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX);
let _: Option<ReplayRng> = None;
let _ = SnapshotRow {
mid_price: 0.0,
bid_l: [0.0; 10],
ask_l: [0.0; 10],
alpha_logit: 0.0,
alpha_confidence: 0.0,
spread_bps: 0.0,
l1_imbalance: 0.0,
ofi_sum_5: 0.0,
mid_drift_5: 0.0,
time_since_trade_s: 0.0,
book_event_rate: 0.0,
};
Ok(())
}