d49003de6d
Closes the TrainingPersist loop for the regime defense per the KELLY_F_SMOOTH precedent (pearl_kelly_cap_signal_driven_floors). Three layers: (1) Per-step: alpha_regime_vol_update_kernel tracks min(vol_obs > 1e-12) in new slot 553 (REGIME_VOL_OBS_MIN_INDEX). Filtered against artifact- zero observations from stationary snapshots where curr == prev. (2) Per-cell: stacker_threshold_controller_update gains a fifth branch that reads vol_ref (slot 550) at cell-end, computes `target = floor_target_ratio × vol_ref`, slow-EMAs the floor anchor (slot 552) toward the target with rate `floor_update_rate` (0.1 per cell). Subfloor 1e-12 inside the kernel guards against the anchor collapsing to zero. (3) Per-invocation: alpha_baseline reads `config/ml/alpha_baseline_state.json` at startup and seeds slot 552 from the `regime_vol_ref_floor` field. At end of main(), the learned floor is written back via tmp+rename atomic write so concurrent walk-forward invocations see a consistent file. Matches the cross-fold-persistent shape of KELLY_F_SMOOTH. Block extended to 15 slots (539..=553). Smoke + kernel unit test pass -1/-1 for the new floor-controller indices (backward compat). Walk-forward CV verdict (Q1 fxcache, 3 sequential folds): iteration fold-A fold-B fold-C mean ± SD pre-defense (no regime) +91.52 -21.44 +46.74 +38.94 ± 56.88 hardcoded 1e-9 floor -19.77 +65.04 +6.45 +17.24 ± 43.42 learned floor (0.5 × cell_min) +74.78 -12.72 +8.80 +23.62 ± 45.59 learned floor (0.1 × vol_ref) -19.53 -26.51 +15.46 -10.20 ± 22.49 Controller infrastructure is structurally correct (loop closes, floor persists across invocations, kernel + disk + ISV all roundtrip). The TUNING is data-dependent — single-quarter CV doesn't have enough regime diversity to anchor the floor against. Multi-quarter fxcache validation is the next step (built cluster-side on the 9-quarter 2024-Q1..2026-Q1 ES futures dataset, downloaded as a single artifact for local CV). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1203 lines
54 KiB
Rust
1203 lines
54 KiB
Rust
//! 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 FULL_ACTIONS: [u8; 9] = [0, 1, 2, 3, 4, 5, 6, 7, 8];
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(
|
||
name = "alpha_baseline",
|
||
about = "Alpha baseline — trained Mamba2 + C51 + ISV-continual + vol-regime defense"
|
||
)]
|
||
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,
|
||
// pruned_actions — FALSIFIED 2026-05-15, hardcoded false.
|
||
// c51, temporal, isv_continual, regime_scale — always on now,
|
||
// removed from CLI surface.
|
||
/// 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,
|
||
/// Temporal-encoder window length (number of historical snapshots
|
||
/// fed to Mamba2 per step).
|
||
#[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.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)
|
||
}
|
||
}
|
||
|
||
// The linear-Q ε-greedy + confidence-gated ε-greedy helpers used by the
|
||
// legacy scalar Q path were deleted with the rest of that path. All
|
||
// action selection now happens on-GPU via `alpha_c51_thompson_select`.
|
||
|
||
#[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!("Alpha baseline — trained Mamba2 + C51 + ISV-continual + vol-regime defense");
|
||
info!(" ACTION SET: full ({} actions)", FULL_ACTIONS.len());
|
||
let allowed_actions: &[u8] = &FULL_ACTIONS;
|
||
|
||
let ctx = CudaContext::new(0).context("CUDA init")?;
|
||
let stream = ctx.default_stream();
|
||
|
||
// --- Load cubins ---
|
||
// alpha_linear_q.cubin provides two kernels we still need: the
|
||
// in-place clip and the SGD step (both used on the C51 head's
|
||
// dW / dB after gradients). The forward/grad linear-Q kernels are
|
||
// unused (C51 is the only Q-network in this binary).
|
||
let lq_module = ctx
|
||
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_LINEAR_Q_CUBIN.to_vec())
|
||
.context("alpha_linear_q cubin")?;
|
||
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 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")?;
|
||
let push_module = ctx
|
||
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_WINDOW_PUSH_CUBIN.to_vec())
|
||
.context("alpha_window_push cubin")?;
|
||
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);
|
||
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");
|
||
}
|
||
|
||
// --- 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 (C51 over Mamba2 h_enriched) ---
|
||
let c51_input_dim: usize = cli.mamba2_hidden_dim;
|
||
let n_weights_eff: usize = N_ACTIONS * c51_n_atoms * c51_input_dim;
|
||
let n_biases_eff: usize = N_ACTIONS * c51_n_atoms;
|
||
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;
|
||
// All per-step inference + per-episode transition buffers live in
|
||
// the batched parallel-env section below (states_train_dev etc.,
|
||
// batched_window_tensor_train, batched_probs_dev, ...). The
|
||
// single-env legacy buffers and the scalar linear-Q transition
|
||
// buffers were removed when C51 + temporal + batched became the
|
||
// only execution path.
|
||
|
||
let mut mamba2_block: Mamba2Block = {
|
||
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!("Mamba2: in={} hidden={} state={} K={}",
|
||
cfg.in_dim, cfg.hidden_dim, cfg.state_dim, cfg.seq_len);
|
||
Mamba2Block::new(cfg, stream.clone())
|
||
.map_err(|e| anyhow::anyhow!("Mamba2Block init: {e}"))?
|
||
};
|
||
let mut mamba2_adamw: Mamba2AdamW = Mamba2AdamW::new(&mamba2_block, Mamba2AdamWConfig::default())
|
||
.map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?;
|
||
// 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; 554];
|
||
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;
|
||
// T16 vol_ref floor seed (TrainingPersist anchor). Read from disk if
|
||
// a previous training run wrote a learned value; otherwise seed with
|
||
// a sensible default based on ES MBP-10 squared-log-return range
|
||
// (typical vol_obs ~1e-10..1e-6). Hardcoded sub-floor 1e-12 inside
|
||
// the regime kernel still catches the slot collapsing to zero.
|
||
let floor_state_path = std::path::PathBuf::from("config/ml/alpha_baseline_state.json");
|
||
let seeded_floor: f32 = std::fs::read_to_string(&floor_state_path)
|
||
.ok()
|
||
.and_then(|s| serde_json::from_str::<serde_json::Value>(&s).ok())
|
||
.and_then(|v| v.get("regime_vol_ref_floor").and_then(|x| x.as_f64()))
|
||
.map(|x| x as f32)
|
||
.unwrap_or(1.0e-9);
|
||
isv_host[ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_FLOOR_INDEX] = seeded_floor;
|
||
info!(
|
||
"Regime floor seed: {:.3e} ({})",
|
||
seeded_floor,
|
||
if floor_state_path.exists() { "from disk" } else { "default" },
|
||
);
|
||
// T16 vol_obs running-min slot: reset to sentinel so the first cell
|
||
// accumulates min from any positive vol_obs observation.
|
||
isv_host[ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_INDEX] =
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_SENTINEL;
|
||
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];
|
||
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);
|
||
|
||
{
|
||
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 (_logit, cache) = mamba2_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,
|
||
)?;
|
||
}
|
||
}
|
||
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 (_logit, cache_train) = mamba2_block.forward_train(&train_windows_tensor_train)
|
||
.map_err(|e| anyhow::anyhow!("train mamba2 (epoch forward): {e}"))?;
|
||
let (curr_input_ptr, next_input_ptr): (u64, u64) = {
|
||
let (h_curr, _g) = cache_train.h_enriched.cuda_data().device_ptr(&stream);
|
||
let (h_base_next, _g_next) = h_enriched_train_dev.device_ptr(&stream);
|
||
(
|
||
h_curr,
|
||
h_base_next + (train_n_par as u64) * (cli.mamba2_hidden_dim as u64) * 4u64,
|
||
)
|
||
};
|
||
|
||
{
|
||
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 (dw_ptr, _g4) = dw_dev.device_ptr_mut(&stream);
|
||
let (db_ptr, _g5) = db_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad(
|
||
&stream, &c51_grad_kernel,
|
||
p_ptr, m_ptr, a_ptr, curr_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.
|
||
{
|
||
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 grads = mamba2_block
|
||
.backward_from_h_enriched(&cache_train, &d_h_enriched_tensor_train)
|
||
.map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?;
|
||
drop(cache_train);
|
||
mamba2_adamw
|
||
.step(&mut mamba2_block, &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 = ctx.load_cubin(
|
||
ml::cuda_pipeline::alpha_kernels::ALPHA_REGIME_VOL_UPDATE_CUBIN.to_vec()
|
||
).context("alpha_regime_vol_update cubin load")?;
|
||
let regime_kernel = regime_module
|
||
.load_function("alpha_regime_vol_update_kernel")
|
||
.context("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.
|
||
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).
|
||
{
|
||
let zeros = vec![0.0_f32; n_par];
|
||
mids_curr_pinned.write(&zeros);
|
||
mids_prev_pinned.write(&zeros);
|
||
}
|
||
// T16: reset the vol_obs min slot to its sentinel at cell start
|
||
// so each cell's min accumulates independently. The floor
|
||
// controller (fired at cell-end inside the stacker controller)
|
||
// will consume this slot's final value to update the floor
|
||
// anchor.
|
||
{
|
||
let sentinel = vec![ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_SENTINEL];
|
||
let start = ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_INDEX;
|
||
let mut slot_view = isv_dev.slice_mut(start..start + 1);
|
||
stream.memcpy_htod(&sentinel, &mut slot_view)
|
||
.context("reset vol_obs min slot")?;
|
||
}
|
||
// 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.
|
||
{
|
||
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, ®ime_kernel,
|
||
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,
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_SAMPLES_INDEX as i32,
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_FLOOR_INDEX as i32,
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_INDEX as i32,
|
||
100, // bootstrap samples: running mean over first 100 obs
|
||
0.005, // β: 200-step horizon for vol_ref tracking after bootstrap
|
||
)?;
|
||
}
|
||
// 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);
|
||
}
|
||
// 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 (_logit, cache) = mamba2_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,
|
||
)?;
|
||
}
|
||
}
|
||
// 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).
|
||
// Threshold + Kelly controllers adapt during the 2D sweep; the
|
||
// T16 regime kernel feeds the pre-emptive Kelly attenuation
|
||
// multiplier via slots 549/550.
|
||
{
|
||
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,
|
||
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.25,
|
||
// T16 floor controller: refine REGIME_VOL_REF_FLOOR
|
||
// toward 0.1 × vol_ref (slot 550). Floor sits at
|
||
// 10% of the trained-regime baseline — small
|
||
// enough that the spike detector still fires on
|
||
// realistic vol_ema excursions, large enough to
|
||
// prevent the deadband deadlock.
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_FLOOR_INDEX as i32,
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_INDEX as i32,
|
||
0.1, // floor sits at 10% of vol_ref
|
||
0.1, // slow EMA toward target (10% per cell)
|
||
ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_OBS_MIN_SENTINEL,
|
||
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!({
|
||
"binary": "alpha_baseline",
|
||
"n_allowed_actions": allowed_actions.len(),
|
||
"c51_n_atoms": c51_n_atoms,
|
||
"c51_vmin": c51_v_min,
|
||
"c51_vmax": c51_v_max,
|
||
"window_k": cli.window_k,
|
||
"mamba2_hidden_dim": cli.mamba2_hidden_dim,
|
||
"mamba2_state_dim": cli.mamba2_state_dim,
|
||
"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());
|
||
|
||
// T16: persist the learned vol_ref floor across invocations so the
|
||
// KELLY_F_SMOOTH-equivalent cross-fold-persistent pattern works for
|
||
// walk-forward CV. Read the final floor value off the device, write
|
||
// to disk; next invocation reloads via the seed-from-disk path at
|
||
// startup. Atomic write: tmp file + rename so concurrent reads see
|
||
// either the old or the new file, never a partial one.
|
||
{
|
||
stream.synchronize().context("sync before floor persist")?;
|
||
let isv_final = stream.clone_dtoh(&isv_dev).context("dtoh isv for floor persist")?;
|
||
let learned_floor = isv_final
|
||
[ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_FLOOR_INDEX];
|
||
info!(
|
||
"Regime floor learned: {:.3e} (was {:.3e})",
|
||
learned_floor, seeded_floor,
|
||
);
|
||
let json = serde_json::json!({ "regime_vol_ref_floor": learned_floor as f64 });
|
||
let tmp_path = floor_state_path.with_extension("json.tmp");
|
||
std::fs::write(&tmp_path, serde_json::to_string_pretty(&json)?)
|
||
.context("write floor state tmp")?;
|
||
std::fs::rename(&tmp_path, &floor_state_path)
|
||
.context("rename floor state into place")?;
|
||
}
|
||
|
||
// 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(())
|
||
}
|