2feeeda8bb
The Phase E.4.A T8 wiring stored Mamba2's per-step cache.h_enriched
into h_enriched_buf_dev via a dtoh+htod sequence:
let h_host = stream.clone_dtoh(cache.h_enriched.cuda_data())?;
let mut buf_host = stream.clone_dtoh(&h_enriched_buf_dev)?; // <- whole buffer
for j in 0..hidden_dim { buf_host[slot_offset + j] = h_host[j]; }
stream.memcpy_htod(&buf_host, &mut h_enriched_buf_dev)?; // <- whole buffer
This violates feedback_cpu_is_read_only AND
feedback_no_htod_htoh_only_mapped_pinned. Worse, the buffer-wide
dtoh+htod every step is ~20K floats × 600 steps × 500 eps × 30 cells
= ~9M roundtrips totaling significant PCIe latency in the backtest.
Fix: new tiny CUDA kernel alpha_h_enriched_store_kernel in
alpha_window_push.cu (one thread per hidden-dim feature, writes
src[j] → buf[slot_offset + j]). Replaces the dtoh/htod sequence
in both smoke and backtest binaries.
Estimated speed-up at backtest scale: 3-6× on the temporal eval
path. Pure-GPU per-step inference restored — no synchronisation
points on the hot path.
docs/isv-slots.md updated per kernel-audit-doc hook.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1174 lines
50 KiB
Rust
1174 lines
50 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).
|
||
use ml_alpha::mamba2_block::{Mamba2Block, Mamba2BlockConfig};
|
||
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,
|
||
}
|
||
|
||
impl MappedI32 {
|
||
unsafe fn new() -> Result<Self> {
|
||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
|
||
| cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
|
||
let bytes = 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(host_ptr, 0);
|
||
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() })
|
||
}
|
||
fn dev_u64(&self) -> u64 { self.dev_ptr as u64 }
|
||
fn read(&self) -> i32 {
|
||
unsafe { std::ptr::read_volatile(self.host_ptr) }
|
||
}
|
||
}
|
||
|
||
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,
|
||
/// 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).
|
||
#[arg(long, default_value_t = 0.0625)]
|
||
train_cost: 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,
|
||
#[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")?;
|
||
// 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(
|
||
&cli.fxcache_path,
|
||
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()? };
|
||
|
||
// 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 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
|
||
};
|
||
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 on train segment.
|
||
// ---------------------------------------------------------------
|
||
info!("=== Training phase: {} episodes on train segment ===", cli.n_train_episodes);
|
||
let mut episode_rng = SmokeRng::new(cli.seed.wrapping_add(0xFEED));
|
||
let train_max_start = (n_train.saturating_sub(cli.horizon + 1)).max(1);
|
||
for ep in 0..cli.n_train_episodes {
|
||
let eps = cli.eps_start
|
||
+ (cli.eps_end - cli.eps_start)
|
||
* (ep as f32 / cli.n_train_episodes.max(1) as f32);
|
||
let start_cursor = (episode_rng.next_u64() as usize) % train_max_start;
|
||
let env_seed = episode_rng.next_u64();
|
||
env.reset_at(env_seed, start_cursor);
|
||
let mut state = EpisodeState::new();
|
||
|
||
let mut states_host: Vec<f32> = Vec::with_capacity(cli.horizon * STATE_DIM);
|
||
let mut next_states_host: Vec<f32> = Vec::with_capacity(cli.horizon * STATE_DIM);
|
||
let mut actions_host: Vec<i32> = Vec::with_capacity(cli.horizon);
|
||
let mut rewards_host: Vec<f32> = Vec::with_capacity(cli.horizon);
|
||
let mut dones_host: Vec<f32> = Vec::with_capacity(cli.horizon);
|
||
// Phase E.4.A.7 (T12): zero window at episode start.
|
||
if cli.temporal {
|
||
stream.memset_zeros(window_tensor.data_mut())
|
||
.context("zero window on episode reset (train)")?;
|
||
}
|
||
loop {
|
||
let s_vec = env.state(&state).to_vec();
|
||
let action: u8 = if cli.c51 {
|
||
state_pinned.write(&s_vec);
|
||
// Phase E.4.A.7 (T12): push and run Mamba2.
|
||
if cli.temporal {
|
||
let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push(
|
||
&stream, &push_kernel,
|
||
state_pinned.dev_u64(), w_ptr,
|
||
cli.window_k as i32, state_dim_i,
|
||
)?;
|
||
}
|
||
}
|
||
if cli.temporal {
|
||
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
|
||
let (_logit, cache) = block.forward_train(&window_tensor)
|
||
.map_err(|e| anyhow::anyhow!("mamba2 forward: {e}"))?;
|
||
// Phase E.4.A.8.fix: GPU-side slot copy.
|
||
{
|
||
let slot_offset = (state.step * cli.mamba2_hidden_dim) as i32;
|
||
let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
|
||
let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store(
|
||
&stream, &h_store_kernel,
|
||
src_ptr, buf_ptr, slot_offset,
|
||
cli.mamba2_hidden_dim as i32,
|
||
)?;
|
||
}
|
||
}
|
||
// C51 forward on h_enriched.
|
||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||
let (p_ptr, _g2) = single_probs_dev.device_ptr_mut(&stream);
|
||
let (h_ptr, _g3) = cache.h_enriched.cuda_data().device_ptr(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
|
||
&stream, &c51_fwd_kernel,
|
||
w_ptr, b_ptr, h_ptr, p_ptr,
|
||
1, 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, _g2) = single_probs_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
|
||
&stream, &c51_fwd_kernel,
|
||
w_ptr, b_ptr, state_pinned.dev_u64(), p_ptr,
|
||
1, c51_input_dim as i32, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
let step_seed = episode_rng.next_u64() as u32;
|
||
{
|
||
let (p_ptr, _g0) = single_probs_dev.device_ptr(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
|
||
&stream, &c51_thompson_kernel,
|
||
p_ptr,
|
||
state_pinned.dev_u64(),
|
||
cli.train_threshold,
|
||
1,
|
||
state_dim_i,
|
||
c51_v_min, c51_delta_z,
|
||
step_seed,
|
||
action_pinned.dev_u64(),
|
||
1, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
stream.synchronize()?;
|
||
action_pinned.read() as u8
|
||
} else {
|
||
stream.memcpy_htod(&s_vec, &mut single_state_dev)?;
|
||
{
|
||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||
let (s_ptr, _g2) = single_state_dev.device_ptr(&stream);
|
||
let (q_ptr, _g3) = single_q_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
|
||
&stream, &lq_fwd, w_ptr, b_ptr, s_ptr, q_ptr,
|
||
1, state_dim_i, n_act_i,
|
||
)?;
|
||
}
|
||
}
|
||
stream.synchronize()?;
|
||
let q_host = stream.clone_dtoh(&single_q_dev)?;
|
||
epsilon_greedy_gated(
|
||
&q_host,
|
||
s_vec[1],
|
||
cli.train_threshold,
|
||
eps,
|
||
allowed_actions,
|
||
&mut episode_rng,
|
||
)
|
||
};
|
||
let (_s_next, reward, done) = env
|
||
.step(action, &mut state)
|
||
.ok_or_else(|| anyhow::anyhow!("step returned None"))?;
|
||
let s_next_vec = env.state(&state).to_vec();
|
||
states_host.extend_from_slice(&s_vec);
|
||
next_states_host.extend_from_slice(&s_next_vec);
|
||
actions_host.push(action as i32);
|
||
rewards_host.push(reward);
|
||
dones_host.push(if done { 1.0 } else { 0.0 });
|
||
if done {
|
||
break;
|
||
}
|
||
}
|
||
let ep_len = actions_host.len() as i32;
|
||
if ep_len < 2 {
|
||
continue;
|
||
}
|
||
// Batched train update (same logic as alpha_dqn_h600_smoke).
|
||
let rewards_norm: Vec<f32> = rewards_host
|
||
.iter()
|
||
.map(|r| r / cli.reward_scale)
|
||
.collect();
|
||
stream.memcpy_htod(&states_host, &mut states_dev)?;
|
||
stream.memcpy_htod(&next_states_host, &mut next_states_dev)?;
|
||
stream.memcpy_htod(&actions_host, &mut actions_dev)?;
|
||
stream.memcpy_htod(&rewards_norm, &mut rewards_dev)?;
|
||
stream.memcpy_htod(&dones_host, &mut dones_dev)?;
|
||
if cli.c51 {
|
||
// Phase E.4.A.8 (T12): in temporal mode, the batched C51
|
||
// forwards read h_enriched_buf_dev populated during the
|
||
// per-step inference. Run one extra Mamba2 forward on the
|
||
// terminal window to fill slot ep_len.
|
||
if cli.temporal {
|
||
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
|
||
let (_logit, cache) = block.forward_train(&window_tensor)
|
||
.map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?;
|
||
// Phase E.4.A.8.fix: GPU-side slot copy for terminal h_enriched.
|
||
let term_offset = ((ep_len as usize) * cli.mamba2_hidden_dim) as i32;
|
||
let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
|
||
let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store(
|
||
&stream, &h_store_kernel,
|
||
src_ptr, buf_ptr, term_offset,
|
||
cli.mamba2_hidden_dim as i32,
|
||
)?;
|
||
}
|
||
}
|
||
let (curr_input_ptr_guard, next_input_ptr_guard);
|
||
let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal {
|
||
let (h_ptr_curr, g_curr) = h_enriched_buf_dev.device_ptr(&stream);
|
||
let (h_ptr_next_base, g_next) = h_enriched_buf_dev.device_ptr(&stream);
|
||
curr_input_ptr_guard = g_curr;
|
||
next_input_ptr_guard = g_next;
|
||
let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
|
||
(
|
||
h_ptr_curr,
|
||
h_ptr_next_base + (cli.mamba2_hidden_dim as u64) * 4u64,
|
||
)
|
||
} else {
|
||
let (s_ptr, g_curr) = states_dev.device_ptr(&stream);
|
||
let (n_ptr, g_next) = next_states_dev.device_ptr(&stream);
|
||
curr_input_ptr_guard = g_curr;
|
||
next_input_ptr_guard = g_next;
|
||
let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
|
||
(s_ptr, n_ptr)
|
||
};
|
||
// C51 forward online → probs_current
|
||
{
|
||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||
let (p_ptr, _g3) = probs_current_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,
|
||
ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
// C51 forward target → probs_next
|
||
{
|
||
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_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,
|
||
ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
{
|
||
let (pn_ptr, _g0) = probs_next_dev.device_ptr(&stream);
|
||
let (r_ptr, _g1) = rewards_dev.device_ptr(&stream);
|
||
let (d_ptr, _g2) = dones_dev.device_ptr(&stream);
|
||
let (m_ptr, _g3) = m_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, ep_len, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
{
|
||
let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream);
|
||
let (m_ptr, _g1) = m_dev.device_ptr(&stream);
|
||
let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
|
||
let (s_ptr, _g3) = states_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,
|
||
ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
|
||
1.0 / ep_len as f32,
|
||
)?;
|
||
}
|
||
}
|
||
} else {
|
||
// Linear-Q forward/Munchausen/grad
|
||
{
|
||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||
let (s_ptr, _g2) = states_dev.device_ptr(&stream);
|
||
let (q_ptr, _g3) = q_current_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
|
||
&stream, &lq_fwd, w_ptr, b_ptr, s_ptr, q_ptr,
|
||
ep_len, state_dim_i, n_act_i,
|
||
)?;
|
||
}
|
||
}
|
||
{
|
||
let (w_ptr, _g0) = w_target_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_target_dev.device_ptr(&stream);
|
||
let (s_ptr, _g2) = next_states_dev.device_ptr(&stream);
|
||
let (q_ptr, _g3) = q_next_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
|
||
&stream, &lq_fwd, w_ptr, b_ptr, s_ptr, q_ptr,
|
||
ep_len, state_dim_i, n_act_i,
|
||
)?;
|
||
}
|
||
}
|
||
{
|
||
let (qn_ptr, _g0) = q_next_dev.device_ptr(&stream);
|
||
let (qc_ptr, _g1) = q_current_dev.device_ptr(&stream);
|
||
let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
|
||
let (r_ptr, _g3) = rewards_dev.device_ptr(&stream);
|
||
let (d_ptr, _g4) = dones_dev.device_ptr(&stream);
|
||
let (t_ptr, _g5) = target_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_munchausen_target(
|
||
&stream, &munch_kernel,
|
||
qn_ptr, qc_ptr, a_ptr, r_ptr, d_ptr,
|
||
cli.gamma, cli.alpha_m, cli.tau, cli.log_clip_min,
|
||
t_ptr, ep_len, n_act_i,
|
||
)?;
|
||
}
|
||
}
|
||
{
|
||
let (qc_ptr, _g0) = q_current_dev.device_ptr(&stream);
|
||
let (t_ptr, _g1) = target_dev.device_ptr(&stream);
|
||
let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
|
||
let (s_ptr, _g3) = states_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_linear_q_grad(
|
||
&stream, &lq_grad,
|
||
qc_ptr, t_ptr, a_ptr, s_ptr, dw_ptr, db_ptr,
|
||
ep_len, state_dim_i, n_act_i,
|
||
1.0 / ep_len as f32,
|
||
)?;
|
||
}
|
||
}
|
||
}
|
||
// Clip + SGD (shared kernels; sizes from n_weights_eff / n_biases_eff)
|
||
{
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
// Target net hard update
|
||
if (ep + 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 (ep + 1) % 200 == 0 {
|
||
info!(" train ep {}/{}", ep + 1, cli.n_train_episodes);
|
||
}
|
||
}
|
||
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.
|
||
// ---------------------------------------------------------------
|
||
info!("=== Eval phase: {} episodes × {} thresholds × {} costs ===",
|
||
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 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;
|
||
let mut rewards: Vec<f32> = Vec::with_capacity(cli.n_eval_episodes);
|
||
let mut win_count = 0_usize;
|
||
let mut trade_count_total: u64 = 0;
|
||
for _ in 0..cli.n_eval_episodes {
|
||
let start_cursor =
|
||
n_train + (episode_rng.next_u64() as usize) % eval_max_start;
|
||
let env_seed = episode_rng.next_u64();
|
||
env.reset_at(env_seed, start_cursor);
|
||
let mut state = EpisodeState::new();
|
||
let mut terminal_r = 0.0_f32;
|
||
let mut ep_n_trades = 0_u32;
|
||
// Phase E.4.A.7 (T12): zero window at eval episode start.
|
||
if cli.temporal {
|
||
stream.memset_zeros(window_tensor.data_mut())
|
||
.context("zero window on eval episode reset")?;
|
||
}
|
||
loop {
|
||
let s_vec = env.state(&state).to_vec();
|
||
let action: u8 = if cli.c51 {
|
||
state_pinned.write(&s_vec);
|
||
// Phase E.4.A (T12): push + Mamba2 + C51. The push
|
||
// is scoped in its own block so the mutable borrow
|
||
// guard drops before the &window_tensor read.
|
||
if cli.temporal {
|
||
{
|
||
let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push(
|
||
&stream, &push_kernel,
|
||
state_pinned.dev_u64(), w_ptr,
|
||
cli.window_k as i32, state_dim_i,
|
||
)?;
|
||
}
|
||
}
|
||
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
|
||
let (_logit, cache) = block.forward_train(&window_tensor)
|
||
.map_err(|e| anyhow::anyhow!("mamba2 forward (eval): {e}"))?;
|
||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||
let (p_ptr, _g2) = single_probs_dev.device_ptr_mut(&stream);
|
||
let (h_ptr, _g3) = cache.h_enriched.cuda_data().device_ptr(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
|
||
&stream, &c51_fwd_kernel,
|
||
w_ptr, b_ptr, h_ptr, p_ptr,
|
||
1, 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, _g2) = single_probs_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
|
||
&stream, &c51_fwd_kernel,
|
||
w_ptr, b_ptr, state_pinned.dev_u64(), p_ptr,
|
||
1, c51_input_dim as i32, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
let step_seed = episode_rng.next_u64() as u32;
|
||
{
|
||
let (p_ptr, _g0) = single_probs_dev.device_ptr(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
|
||
&stream, &c51_thompson_kernel,
|
||
p_ptr,
|
||
state_pinned.dev_u64(),
|
||
threshold,
|
||
1,
|
||
state_dim_i,
|
||
c51_v_min, c51_delta_z,
|
||
step_seed,
|
||
action_pinned.dev_u64(),
|
||
1, n_act_i, n_atoms_i,
|
||
)?;
|
||
}
|
||
}
|
||
stream.synchronize()?;
|
||
action_pinned.read() as u8
|
||
} else {
|
||
stream.memcpy_htod(&s_vec, &mut single_state_dev)?;
|
||
{
|
||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||
let (s_ptr, _g2) = single_state_dev.device_ptr(&stream);
|
||
let (q_ptr, _g3) = single_q_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
|
||
&stream, &lq_fwd, w_ptr, b_ptr, s_ptr, q_ptr,
|
||
1, state_dim_i, n_act_i,
|
||
)?;
|
||
}
|
||
}
|
||
stream.synchronize()?;
|
||
let q_host = stream.clone_dtoh(&single_q_dev)?;
|
||
let mut greedy_rng = SmokeRng::new(0); // unused — eps=0
|
||
epsilon_greedy_gated(
|
||
&q_host,
|
||
s_vec[1],
|
||
threshold,
|
||
0.0,
|
||
allowed_actions,
|
||
&mut greedy_rng,
|
||
)
|
||
};
|
||
if action != 0 {
|
||
ep_n_trades += 1;
|
||
}
|
||
match env.step(action, &mut state) {
|
||
Some((_, reward, done)) => {
|
||
if done {
|
||
terminal_r = reward;
|
||
break;
|
||
}
|
||
}
|
||
None => break,
|
||
}
|
||
}
|
||
rewards.push(terminal_r);
|
||
trade_count_total += ep_n_trades as u64;
|
||
if terminal_r > 0.0 {
|
||
win_count += 1;
|
||
}
|
||
// Phase E.4.A Pillar B (T11): fire the stacker-threshold
|
||
// controller at end of each eval episode when
|
||
// --isv-continual. Adapts ISV slot 543/545/546 based on
|
||
// observed trade-rate and Sharpe; the τ-grid still
|
||
// overrides via `threshold` for the actual gate, but the
|
||
// ISV evolution is the "live deployment" simulation.
|
||
if cli.isv_continual {
|
||
let trade_count_ep = ep_n_trades as f32;
|
||
let decisions_ep = (state.step as f32).max(1.0);
|
||
let baseline_std = isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX].max(1e-6);
|
||
let rollout_sharpe = 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_ep, decisions_ep, rollout_sharpe,
|
||
/* target_sharpe */ 0.5,
|
||
/* k_threshold */ 0.01,
|
||
/* k_atten */ 0.005,
|
||
/* wiener_alpha_floor */ 0.4,
|
||
/* ctl_alpha_meta */ 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,
|
||
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(())
|
||
}
|