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

feat(alpha): Phase E.3 follow-up — C51 distributional Q + Thompson + L1-L10 depth + falsifications

Browse Source 
C51 distributional Q-network with GPU Thompson selection borrowed
minimally from production (alpha_c51.cu: forward, project, grad,
expected_q, thompson_select kernels; ~260 lines). Uses Huber
negative-tail compression in projection per production
block_bellman_project_f. Action selection 100% GPU via mapped-pinned
i32 output + __threadfence_system + host volatile read (matches
gpu_training_guard MappedBuffer pattern).

Backtest result (2D sweep, 500 episodes per cell, 30 cells):
  cost=0    C51 +10.41 vs linear-Q -15.72  (+26pt, BEATS Phase 1d.4
                                            no-RL baseline +4.4 by 6pt)
  cost=0.125 C51 -13.81 vs -29.17  (+15pt closes half-tick gap)
Win rate at cost=0 best τ: linear-Q 0.008 → C51 0.552.

Calibration hypothesis vindicated; documented in
memory/pearl_c51_thompson_closed_phase_e3_gap.md.

Also in this commit (Phase E.3 follow-up cleanup):
- --pruned-actions falsified (2.4× worse Sharpe). Documented in
  memory/pearl_action_pruning_falsified.md.
- --real-spread falsified for ES futures (76% of bars at 1-tick floor).
- SnapshotRow bid_l/ask_l extended from [f32; 3] to [f32; 10].
  L4-L10 synthesized in this commit; real MBP-10 peek lands in E.4.A T5.
- docs/isv-slots.md updated per kernel-audit-doc hook requirement.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-15 20:43:57 +02:00
parent eb9047fc30
commit eb49e2a0f7
15 changed files with 3602 additions and 257 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml/build.rs
View File

@@ -1615,6 +1615,12 @@ fn main() {
// at 0.4 per pearl_wiener_alpha_floor_for_nonstationary) on the
// observed-rate EMA at slot 545.
"stacker_threshold_controller.cu",
// Phase E.3 follow-up (2026-05-15): vanilla C51 distributional
// Q-network kernels. Forward (logits → softmax over atoms),
// categorical Bellman projection, cross-entropy gradient. Fixed
// atom support, single network, no per-branch / Adam machinery —
// testing the calibration hypothesis from `pearl_action_pruning_falsified`.
"alpha_c51.cu",
];
// ALL kernels get common header (BF16 types + wrappers)

544
crates/ml/examples/alpha_compose_backtest.rs
View File

@@ -33,6 +33,7 @@
use std::fs::File;
use std::io::Write;
use std::mem::MaybeUninit;
use std::path::PathBuf;
use anyhow::{Context, Result};
@@ -40,6 +41,86 @@ use clap::Parser;
use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut};
use tracing::info;
// ── 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,
};
@@ -52,6 +133,11 @@ 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",
@@ -126,6 +212,28 @@ struct Cli {
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,
#[arg(long, default_value = "config/ml/alpha_compose_backtest.json")]
out_path: PathBuf,
}
@@ -150,35 +258,37 @@ impl SmokeRng {
}
}
fn epsilon_greedy(q: &[f32], eps: f32, rng: &mut SmokeRng) -> u8 {
fn epsilon_greedy(q: &[f32], eps: f32, allowed: &[u8], rng: &mut SmokeRng) -> u8 {
if rng.next_f32() < eps {
(rng.next_u64() % N_ACTIONS as u64) as u8
allowed[(rng.next_u64() as usize) % allowed.len()]
} else {
let mut best_i: usize = 0;
let mut best_v: f32 = q[0];
for i in 1..N_ACTIONS {
if q[i] > best_v {
best_v = q[i];
best_i = i;
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_i as u8
best_a
}
}
/// Phase E.3 confidence-gated ε-greedy. If `alpha_confidence < threshold`,
/// force action=0 (Wait). Otherwise standard ε-greedy.
/// 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, rng)
epsilon_greedy(q, eps, allowed, rng)
}
#[derive(Debug, Clone, serde::Serialize)]
@@ -208,6 +318,13 @@ fn main() -> Result<()> {
.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();
@@ -228,6 +345,33 @@ fn main() -> Result<()> {
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")?;
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");
@@ -237,6 +381,7 @@ fn main() -> Result<()> {
&cli.fxcache_path,
cli.max_snapshots,
Some(&alpha_cache),
cli.real_spread,
)?;
let n_total = rows.len();
info!("Loaded {} snapshots", n_total);
@@ -268,21 +413,28 @@ fn main() -> Result<()> {
);
// --- Initialize Q-network ---
let n_weights_eff: usize = if cli.c51 {
N_ACTIONS * c51_n_atoms * STATE_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 / STATE_DIM as f32).sqrt();
let w_init: Vec<f32> = (0..N_WEIGHTS)
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];
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)?;
let mut db_dev = stream.alloc_zeros::<f32>(N_BIASES)?;
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)?;
@@ -293,6 +445,14 @@ fn main() -> Result<()> {
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 1: TRAIN DQN on train segment.
@@ -316,33 +476,65 @@ fn main() -> Result<()> {
let mut dones_host: Vec<f32> = Vec::with_capacity(cli.horizon);
loop {
let s_vec = env.state(&state).to_vec();
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,
)?;
let action: u8 = if cli.c51 {
state_pinned.write(&s_vec);
{
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, state_dim_i, n_act_i, n_atoms_i,
)?;
}
}
}
stream.synchronize()?;
let q_host = stream.clone_dtoh(&single_q_dev)?;
// Phase E.3: gate during training so the Q-network learns the
// value function for the gated policy class. Using a FIXED
// threshold (--train-threshold) keeps the backtest self-contained
// — no controller invocation needed. The default 0.39 is the
// equilibrium the smoke's controller stabilized on.
let action = epsilon_greedy_gated(
&q_host,
s_vec[1],
cli.train_threshold,
eps,
&mut episode_rng,
);
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"))?;
@@ -370,68 +562,129 @@ fn main() -> Result<()> {
stream.memcpy_htod(&actions_host, &mut actions_dev)?;
stream.memcpy_htod(&rewards_norm, &mut rewards_dev)?;
stream.memcpy_htod(&dones_host, &mut dones_dev)?;
{
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,
)?;
if cli.c51 {
// C51 forward/project/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 (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, s_ptr, p_ptr,
ep_len, state_dim_i, 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 (s_ptr, _g2) = next_states_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, s_ptr, p_ptr,
ep_len, state_dim_i, 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);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad(
&stream, &c51_grad_kernel,
p_ptr, m_ptr, a_ptr, s_ptr, dw_ptr, db_ptr,
ep_len, state_dim_i, 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,
)?;
}
}
}
{
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
// 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 as i32,
&stream, &lq_clip, dw_ptr, cli.grad_clip, n_weights_eff as i32,
)?;
}
}
@@ -439,7 +692,7 @@ fn main() -> Result<()> {
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 as i32,
&stream, &lq_clip, db_ptr, cli.grad_clip, n_biases_eff as i32,
)?;
}
}
@@ -448,7 +701,7 @@ fn main() -> Result<()> {
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 as i32,
&stream, &lq_sgd, w_ptr, dw_ptr, cli.lr, n_weights_eff as i32,
)?;
}
}
@@ -457,7 +710,7 @@ fn main() -> Result<()> {
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 as i32,
&stream, &lq_sgd, b_ptr, db_ptr, cli.lr, n_biases_eff as i32,
)?;
}
}
@@ -500,29 +753,66 @@ fn main() -> Result<()> {
let mut ep_n_trades = 0_u32;
loop {
let s_vec = env.state(&state).to_vec();
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,
)?;
let action: u8 = if cli.c51 {
state_pinned.write(&s_vec);
{
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, state_dim_i, n_act_i, n_atoms_i,
)?;
}
}
}
stream.synchronize()?;
let q_host = stream.clone_dtoh(&single_q_dev)?;
let mut greedy_rng = SmokeRng::new(0); // unused — eps=0
let action = epsilon_greedy_gated(
&q_host,
s_vec[1], // alpha_confidence
threshold,
0.0,
&mut greedy_rng,
);
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;
}
@@ -622,6 +912,12 @@ fn main() -> Result<()> {
// --- 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 },
"horizon": cli.horizon,
"train_frac": cli.train_frac,
"n_train_episodes": cli.n_train_episodes,
@@ -640,8 +936,8 @@ fn main() -> Result<()> {
let _: Option<ReplayRng> = None;
let _ = SnapshotRow {
mid_price: 0.0,
bid_l: [0.0; 3],
ask_l: [0.0; 3],
bid_l: [0.0; 10],
ask_l: [0.0; 10],
alpha_logit: 0.0,
alpha_confidence: 0.0,
spread_bps: 0.0,

621
crates/ml/examples/alpha_dqn_h600_smoke.rs
View File

@@ -40,6 +40,7 @@
use std::fs::File;
use std::io::Write;
use std::mem::MaybeUninit;
use std::path::PathBuf;
use anyhow::{Context, Result};
@@ -47,6 +48,120 @@ use clap::Parser;
use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut};
use tracing::{info, warn};
// ── Mapped-pinned helpers ──────────────────────────────────────────
//
// `cuMemHostAlloc(DEVICEMAP|PORTABLE)` allocations: host and device see
// the SAME physical memory; the device pointer obtained via
// `cuMemHostGetDevicePointer_v2` lets a kernel read/write directly,
// while the host pointer lets us write inputs (per-step state) and
// read outputs (per-step action) without `memcpy_htod` / `memcpy_dtoh`.
//
// Kernels writing to mapped-pinned outputs MUST call
// `__threadfence_system()` after the write (alpha_c51.cu does this for
// the Thompson selector) so the value is PCIe-visible to the host.
//
// Mirrors the production `MappedBuffer` pattern in `gpu_training_guard.rs`
// per `feedback_no_htod_htoh_only_mapped_pinned`.
struct MappedI32 {
host_ptr: *mut i32,
dev_ptr: cudarc::driver::sys::CUdeviceptr,
}
impl MappedI32 {
/// # Safety
/// Caller must ensure a CUDA context is active on the current thread.
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 {
/// # Safety
/// Caller must ensure a CUDA context is active on the current thread.
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, "MappedF32 write length mismatch");
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 data::providers::databento::{dbn_parser::DbnParser, mbp10::Mbp10Snapshot};
use ml::cuda_pipeline::alpha_isv_slots::{
ACTION_ENTROPY_EMA_INDEX, ALPHA_ISV_BLOCK_LO, EARLY_Q_MOVEMENT_EMA_INDEX,
@@ -68,6 +183,15 @@ const N_WEIGHTS: usize = N_ACTIONS * STATE_DIM;
/// Number of bias floats: n_actions.
const N_BIASES: usize = N_ACTIONS;
/// Action allow-list for the pruning experiment (Phase E.3 follow-up,
/// Hypothesis A: action-variance is the alpha-extraction bottleneck).
/// Indices into the full 9-action space, kept as u8 so they drop straight
/// into the env's `step(action_id: u8, …)`. The Q-network is unchanged
/// (still 9 outputs); the selector below restricts argmax + ε-random to
/// this subset. Unused outputs get stale gradients but are never executed.
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];
#[inline]
fn raw_price_to_f32(fixed: i64) -> f32 {
(fixed as f64 * 1e-9) as f32
@@ -179,6 +303,40 @@ struct Cli {
/// at 1e-4 per element, which is safe.
#[arg(long, default_value_t = 1.0)]
grad_clip: f32,
/// Restrict action selection to {Wait, BuyMarket, SellMarket, FlatMarket}.
/// Q-network still outputs 9 actions; selector + ε-greedy filter to the
/// 4-subset. Tests Hypothesis A from the Phase E.3 close-out memo
/// (action variance is the alpha-extraction bottleneck, not capacity
/// or fill economics). **FALSIFIED 2026-05-15** — kept for reproducibility.
#[arg(long, default_value_t = false)]
pruned_actions: bool,
/// Use C51 distributional Q-network instead of scalar linear Q.
/// Output is 9 × n_atoms probability distributions per state; action
/// selection uses Thompson sampling (inverse-CDF) on GPU; target is
/// the categorical Bellman projection with Huber negative-tail
/// compression. Tests the calibration hypothesis (Phase E.3 close-out).
/// Replaces Munchausen target — C51 has its own implicit entropy bonus.
#[arg(long, default_value_t = false)]
c51: bool,
/// C51 atom-support lower bound (in normalized reward units, i.e.
/// reward / `reward_scale`). Default -10 covers ~2σ of the random
/// baseline's normalized terminal-reward distribution.
#[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; the kernel caps at 64. delta_z is
/// derived: `(vmax - vmin) / (n_atoms - 1)`.
#[arg(long, default_value_t = 51)]
c51_n_atoms: usize,
/// Path 3 (Phase E.3 follow-up, 2026-05-15): derive bid/ask from
/// the fxcache's real `spread_bps` feature instead of fixed
/// ±0.125-tick synthesis. L2/L3 still synthesized at ±tick from
/// real L1. Tests whether the variable-spread env improves the
/// half-tick gap to the Phase 1d.4 baseline.
#[arg(long, default_value_t = false)]
real_spread: bool,
/// Output JSON path for final verdict + ISV readings.
#[arg(long, default_value = "config/ml/alpha_dqn_h600_smoke.json")]
out_path: PathBuf,
@@ -225,8 +383,16 @@ fn load_snapshots(
0.5
};
let ofi_sum_5 = bid_sz - ask_sz;
let bid_l = [bid_l1, bid_l1 - TICK, bid_l1 - 2.0 * TICK];
let ask_l = [ask_l1, ask_l1 + TICK, ask_l1 + 2.0 * TICK];
// Phase E.4.A.3: L1 from real MBP-10; L2-L10 synthesized.
// Real L2-L10 lookup (per snapshot) is also available from
// `snap.levels[1..10]` after the parser fix 5c0bcb1fd —
// wire that in Task 5 follow-on.
let mut bid_l = [0.0_f32; 10];
let mut ask_l = [0.0_f32; 10];
for k in 0..10 {
bid_l[k] = bid_l1 - (k as f32) * TICK;
ask_l[k] = ask_l1 + (k as f32) * TICK;
}
rows.push(SnapshotRow {
mid_price: mid,
bid_l,
@@ -311,38 +477,40 @@ impl SmokeRng {
}
/// Phase E.3 confidence-gated ε-greedy. If `alpha_confidence < threshold`,
/// force action=0 (Wait) regardless of Q. Otherwise standard ε-greedy.
/// Both args are expected in [0, 0.5] (the state's
/// `alpha_confidence = |sigmoid(alpha_logit) − 0.5|`, and the controller
/// slot 543 which is clamped to [0, 0.5]). Threshold=0 disables the gate.
/// force action=0 (Wait) regardless of Q. Otherwise standard ε-greedy
/// restricted to `allowed` (the action allow-list — full 9 or pruned 4).
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; // Wait
}
epsilon_greedy(q, eps, rng)
epsilon_greedy(q, eps, allowed, rng)
}
/// Picks an action via ε-greedy: with probability `eps` random, otherwise
/// argmax over `q`. Returns the action index in `[0, N_ACTIONS)`.
fn epsilon_greedy(q: &[f32], eps: f32, rng: &mut SmokeRng) -> u8 {
/// Picks an action via ε-greedy over the `allowed` subset. With probability
/// `eps` uniform-random from `allowed`; otherwise argmax of `q[a]` for
/// `a ∈ allowed`. Returns the raw u8 action index (in the full 9-action
/// numbering, so it drops straight into `env.step`).
fn epsilon_greedy(q: &[f32], eps: f32, allowed: &[u8], rng: &mut SmokeRng) -> u8 {
if rng.next_f32() < eps {
(rng.next_u64() % N_ACTIONS as u64) as u8
allowed[(rng.next_u64() as usize) % allowed.len()]
} else {
let mut best_i: usize = 0;
let mut best_v: f32 = q[0];
for i in 1..N_ACTIONS {
if q[i] > best_v {
best_v = q[i];
best_i = i;
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_i as u8
best_a
}
}
@@ -357,6 +525,13 @@ fn main() -> Result<()> {
info!("Phase E.1 Task 12 — H=600 DQN smoke starting");
info!(" horizon={}, n_episodes={}, lr={}, eps={:.2}→{:.2}",
cli.horizon, cli.n_episodes, cli.lr, cli.eps_start, cli.eps_end);
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
};
// --- CUDA + cubins ---
let ctx = CudaContext::new(0).context("CUDA context init")?;
@@ -406,6 +581,51 @@ fn main() -> Result<()> {
.load_function("stacker_threshold_controller_update")
.context("controller kernel load")?;
// Phase E.3 follow-up: C51 distributional Q kernels. Loaded
// unconditionally — cubin is small and the linear-Q path simply
// ignores it. When --c51 is set, the training loop swaps in the
// C51 forward / project / grad / Thompson chain.
let c51_module = ctx
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_C51_CUBIN.to_vec())
.context("alpha_c51 cubin load")?;
let c51_fwd_kernel = c51_module
.load_function("alpha_c51_forward_kernel")
.context("c51 forward load")?;
let c51_project_kernel = c51_module
.load_function("alpha_c51_project_kernel")
.context("c51 project load")?;
let c51_grad_kernel = c51_module
.load_function("alpha_c51_grad_kernel")
.context("c51 grad load")?;
let c51_expq_kernel = c51_module
.load_function("alpha_c51_expected_q_kernel")
.context("c51 expected_q load")?;
let c51_thompson_kernel = c51_module
.load_function("alpha_c51_thompson_select_kernel")
.context("c51 thompson load")?;
// C51 atom-support derived (constant across episodes; ISV-adaptive
// would lift this to a controller — out of scope for the minimal
// borrow per the Phase E.3 wisdom rationale).
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 distributional ({} atoms, support [{:.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], got {}", c51_n_atoms);
}
if c51_v_max <= c51_v_min {
anyhow::bail!("--c51-vmax ({}) must exceed --c51-vmin ({})", c51_v_max, c51_v_min);
}
} 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 FillModel from {}", cli.fill_coeffs.display());
@@ -429,6 +649,7 @@ fn main() -> Result<()> {
fxc,
cli.max_snapshots,
alpha_cache_vec.as_deref(),
cli.real_spread,
)?
}
(None, Some(mbp10)) => {
@@ -463,23 +684,38 @@ fn main() -> Result<()> {
);
// --- Initialize Q-network: Xavier ---
// C51 lifts the output dimension to N_ACTIONS × n_atoms (logits per
// atom, softmax-normalized across atoms). Linear-Q stays at N_ACTIONS.
let n_weights_eff: usize = if cli.c51 {
N_ACTIONS * c51_n_atoms * STATE_DIM
} else {
N_WEIGHTS
};
let n_biases_eff: usize = if cli.c51 {
N_ACTIONS * c51_n_atoms
} else {
N_BIASES
};
let xavier_scale = (2.0_f32 / STATE_DIM as f32).sqrt();
let mut rng = SmokeRng::new(cli.seed.wrapping_add(0xDEAD_BEEF));
let w_init: Vec<f32> = (0..N_WEIGHTS)
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];
let b_init: Vec<f32> = vec![0.0; n_biases_eff];
let q_init_norm = weight_norm(&w_init, &b_init);
info!("Q-net init: ||W||₂ = {:.4}", q_init_norm);
info!(
"Q-net init: ||W||₂ = {:.4} (n_weights={}, n_biases={})",
q_init_norm, n_weights_eff, n_biases_eff
);
let mut w_dev = stream.clone_htod(&w_init).context("upload W")?;
let mut b_dev = stream.clone_htod(&b_init).context("upload b")?;
// Target network: identical to online at init; periodic hard-update.
// Munchausen V_soft(s') bootstrap reads from this, not w_dev/b_dev.
// V_soft / projection bootstrap reads from this.
let mut w_target_dev = stream.clone_htod(&w_init).context("upload W_target")?;
let mut b_target_dev = stream.clone_htod(&b_init).context("upload b_target")?;
let mut dw_dev = stream.alloc_zeros::<f32>(N_WEIGHTS).context("alloc dW")?;
let mut db_dev = stream.alloc_zeros::<f32>(N_BIASES).context("alloc db")?;
let mut dw_dev = stream.alloc_zeros::<f32>(n_weights_eff).context("alloc dW")?;
let mut db_dev = stream.alloc_zeros::<f32>(n_biases_eff).context("alloc db")?;
// --- ISV buffer (552 floats) with TrainingPersist anchors set ---
let mut isv_host: Vec<f32> = vec![0.0; 552];
@@ -504,20 +740,42 @@ fn main() -> Result<()> {
let mut ctl_wiener_dev = stream.alloc_zeros::<f32>(3).context("alloc ctl wiener")?;
// --- Allocate per-episode batch buffers (sized to horizon, reused) ---
let h = cli.horizon as i32;
let _h = cli.horizon as i32; // unused but kept for log diagnostics
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).context("alloc states")?;
let mut next_states_dev = stream.alloc_zeros::<f32>(cli.horizon * STATE_DIM).context("alloc next_states")?;
let mut actions_dev = stream.alloc_zeros::<i32>(cli.horizon).context("alloc actions")?;
let mut rewards_dev = stream.alloc_zeros::<f32>(cli.horizon).context("alloc rewards")?;
let mut dones_dev = stream.alloc_zeros::<f32>(cli.horizon).context("alloc dones")?;
// Linear-Q scalar Q outputs (used by linear-Q path and as the
// input to kill-criteria in BOTH paths — in C51 the expected_q
// kernel writes E[Z] into the same buffer shape).
let mut q_current_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS).context("alloc q_current")?;
let mut q_next_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS).context("alloc q_next")?;
// Linear-Q TD target (scalar per sample). Unused in C51 path.
let mut target_dev = stream.alloc_zeros::<f32>(cli.horizon).context("alloc target")?;
let mut single_state_dev = stream.alloc_zeros::<f32>(STATE_DIM).context("alloc single_state")?;
let mut single_q_dev = stream.alloc_zeros::<f32>(N_ACTIONS).context("alloc single_q")?;
// C51 distributional buffers — allocated unconditionally; small
// overhead when --c51 is off.
let probs_size_full = cli.horizon * N_ACTIONS * c51_n_atoms;
let mut probs_current_dev = stream.alloc_zeros::<f32>(probs_size_full)
.context("alloc probs_current (C51)")?;
let mut probs_next_dev = stream.alloc_zeros::<f32>(probs_size_full)
.context("alloc probs_next (C51)")?;
let mut m_dev = stream.alloc_zeros::<f32>(cli.horizon * c51_n_atoms)
.context("alloc m (C51 projection target)")?;
let mut single_probs_dev = stream.alloc_zeros::<f32>(N_ACTIONS * c51_n_atoms)
.context("alloc single_probs (C51)")?;
// Mapped-pinned per-step inference buffers (C51 path). Host writes
// state, kernel reads via dev_ptr; kernel writes action with
// __threadfence_system(), host reads via volatile.
let state_pinned = unsafe { MappedF32::new(STATE_DIM)? };
let action_pinned = unsafe { MappedI32::new()? };
// Kill-criteria inputs: action_counts (i32 × n_actions),
// scalar_inputs (f32 × 3: rollout_R_mean, q_init_norm, q_early_norm).
let mut kc_action_counts_dev = stream.alloc_zeros::<i32>(N_ACTIONS).context("alloc kc actions")?;
@@ -552,32 +810,72 @@ fn main() -> Result<()> {
let mut dones_host: Vec<f32> = Vec::with_capacity(cli.horizon);
loop {
// Per-step forward pass (batch=1)
// Per-step forward pass (batch=1). C51 path uses mapped-pinned
// state input + GPU Thompson selector + mapped-pinned action
// output — zero memcpy_htod / memcpy_dtoh per step. Linear-Q
// path retains the legacy CPU-readback + ε-greedy selector.
let s_vec = state_as_vec(&env, &state);
stream.memcpy_htod(&s_vec, &mut single_state_dev)
.context("htod single state")?;
{
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,
)?;
let action: u8 = if cli.c51 {
state_pinned.write(&s_vec);
{
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, state_dim_i, n_act_i, n_atoms_i,
)?;
}
}
}
stream.synchronize().context("sync after per-step fwd")?;
let q_host = stream.clone_dtoh(&single_q_dev).context("dtoh Q")?;
let action = epsilon_greedy_gated(
&q_host,
s_vec[1], // alpha_confidence (state index 1)
current_threshold,
eps,
&mut episode_rng,
);
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(),
current_threshold,
1, // alpha_confidence is state index 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().context("sync after C51 inference")?;
action_pinned.read() as u8
} else {
stream.memcpy_htod(&s_vec, &mut single_state_dev)
.context("htod single state")?;
{
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().context("sync after per-step fwd")?;
let q_host = stream.clone_dtoh(&single_q_dev).context("dtoh Q")?;
epsilon_greedy_gated(
&q_host,
s_vec[1],
current_threshold,
eps,
allowed_actions,
&mut episode_rng,
)
};
let (_next_state_arr, reward, done) = env.step(action, &mut state)
.ok_or_else(|| anyhow::anyhow!("step returned None"))?;
@@ -625,81 +923,142 @@ fn main() -> Result<()> {
stream.memcpy_htod(&dones_host, &mut dones_dev)
.context("htod dones")?;
// Forward Q_current on states
{
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,
)?;
if cli.c51 {
// ── C51 batched compute ─────────────────────────────────
// Forward online → probs_current
{
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 (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, s_ptr, p_ptr,
ep_len, state_dim_i, n_act_i, n_atoms_i,
)?;
}
}
// Forward target net → probs_next
{
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 (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, s_ptr, p_ptr,
ep_len, state_dim_i, n_act_i, n_atoms_i,
)?;
}
}
// Bellman categorical projection → m
{
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,
)?;
}
}
// CE gradient on (probs_current, m, actions)
{
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);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad(
&stream, &c51_grad_kernel,
p_ptr, m_ptr, a_ptr, s_ptr,
dw_ptr, db_ptr,
ep_len, state_dim_i, n_act_i, n_atoms_i,
1.0 / ep_len as f32,
)?;
}
}
} else {
// ── Linear-Q batched compute (Munchausen target) ────────
// Forward Q_current on states
{
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,
)?;
}
}
// Forward Q_next on next_states USING TARGET WEIGHTS.
{
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,
)?;
}
}
// Munchausen target → target_dev[0..ep_len]
{
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,
)?;
}
}
// Linear-Q gradient
{
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,
)?;
}
}
}
// Forward Q_next on next_states USING TARGET WEIGHTS.
// Decouples the V_soft(s') bootstrap target from per-step policy
// drift — the Munchausen target then has a stationary reference.
{
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,
)?;
}
}
// Munchausen target → target_dev[0..ep_len]
{
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,
)?;
}
}
// Gradient
{
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,
)?;
}
}
// Gradient clip — element-wise |g| ≤ cli.grad_clip on dW and db.
// Safety net against gradient bursts that survive reward
// normalization + target net (e.g., rare large-PnL terminal
// rewards). Identity transform when |g| < bound.
// ── Gradient clip + SGD (shared kernels; size depends on mode) ──
{
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 as i32,
dw_ptr, cli.grad_clip, n_weights_eff as i32,
)?;
}
}
@@ -708,7 +1067,7 @@ fn main() -> Result<()> {
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
&stream, &lq_clip,
db_ptr, cli.grad_clip, N_BIASES as i32,
db_ptr, cli.grad_clip, n_biases_eff as i32,
)?;
}
}
@@ -719,7 +1078,7 @@ fn main() -> Result<()> {
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_sgd_step(
&stream, &lq_sgd,
w_ptr, dw_ptr, cli.lr, N_WEIGHTS as i32,
w_ptr, dw_ptr, cli.lr, n_weights_eff as i32,
)?;
}
}
@@ -730,7 +1089,7 @@ fn main() -> Result<()> {
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_sgd_step(
&stream, &lq_sgd,
b_ptr, db_ptr, cli.lr, N_BIASES as i32,
b_ptr, db_ptr, cli.lr, n_biases_eff as i32,
)?;
}
}
@@ -821,6 +1180,22 @@ fn main() -> Result<()> {
stream.memcpy_htod(&action_counts_host, &mut kc_action_counts_dev)
.context("htod kc action_counts")?;
// C51 path: convert distributions to scalar E[Z(s,a)] so the
// existing kill-criteria kernel (which expects scalar Q[B, A])
// can consume them. Reads probs_current_dev (filled by the
// most recent batched forward), writes q_current_dev.
if cli.c51 {
let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream);
let (q_ptr, _g1) = q_current_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_expected_q(
&stream, &c51_expq_kernel,
p_ptr, c51_v_min, c51_delta_z, q_ptr,
ep_len, n_act_i, n_atoms_i,
)?;
}
}
// Launch kill-criteria producer → kc_scratch[0..4]
{
let (q_ptr, _g0) = q_current_dev.device_ptr(&stream);
@@ -918,6 +1293,12 @@ fn main() -> Result<()> {
// --- Save JSON ---
let json = serde_json::json!({
"phase": "E.1 Task 12",
"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 },
"horizon": cli.horizon,
"n_episodes": cli.n_episodes,
"lr": cli.lr,

14
crates/ml/examples/alpha_random_baseline.rs
View File

@@ -192,10 +192,16 @@ fn main() -> Result<()> {
0.5
};
let ofi_sum_5 = bid_sz - ask_sz;
// L2/L3 posting prices synthesised at ±tick offsets (parser
// doesn't populate levels[1..10]; document limitation).
let bid_l = [bid_l1, bid_l1 - TICK, bid_l1 - 2.0 * TICK];
let ask_l = [ask_l1, ask_l1 + TICK, ask_l1 + 2.0 * TICK];
// Phase E.4.A.3: L1-L10 depth — L1 from real MBP-10,
// L2-L10 synthesized at ±tick offsets (parser fix
// 5c0bcb1fd makes real L2-L10 available; wire in
// Task 5 follow-on).
let mut bid_l = [0.0_f32; 10];
let mut ask_l = [0.0_f32; 10];
for k in 0..10 {
bid_l[k] = bid_l1 - (k as f32) * TICK;
ask_l[k] = ask_l1 + (k as f32) * TICK;
}
rows.push(SnapshotRow {
mid_price: mid,

371
crates/ml/src/cuda_pipeline/alpha_c51.cu Normal file
View File

@@ -0,0 +1,371 @@
// crates/ml/src/cuda_pipeline/alpha_c51.cu
//
// Phase E.3 follow-up (2026-05-15): minimal C51 distributional Q for the
// alpha execution policy. Borrows the technique from the production DQN's
// per-branch C51 (per `pearl_per_branch_c51_atom_span` /
// `pearl_thompson_for_distributional_action_selection`) but without the
// branched-action / Adam-optimizer / ISV-driven atom-span machinery. Fixed
// atom support, single network, vanilla SGD — testing the calibration
// hypothesis cheap before committing to MLP or real-LOB lifts.
//
// Layout convention:
// W: [n_actions * n_atoms, state_dim] row-major
// b: [n_actions * n_atoms]
// X: [batch, state_dim] row-major
// probs: [batch, n_actions, n_atoms] row-major; softmax over atoms
// m: [batch, n_atoms] projection target
// dW: [n_actions * n_atoms, state_dim] row-major (gradient)
// db: [n_actions * n_atoms]
//
// Three kernels:
// 1. alpha_c51_forward_kernel — logits = W·x + b, softmax across atoms
// 2. alpha_c51_project_kernel — Bellman categorical projection target
// 3. alpha_c51_grad_kernel — CE gradient w.r.t. W and b
//
// **GPU contract:** no atomicAdd (block-tree-reduce or per-element loops);
// no host branches inside kernels; fixed compile-time MAX_N_ATOMS lets us
// use register-resident scratch instead of dynamic shared memory.
#include <cuda_runtime.h>
#include <math.h>
// Fixed compile-time ceiling on the atom count. Caller's runtime
// `n_atoms` MUST be ≤ MAX_N_ATOMS. 64 covers the canonical C51=51 +
// headroom for an "extended" experiment (e.g., 64-atom variant).
#define MAX_N_ATOMS 64
// ----------------------------------------------------------------------
// (1) Forward: logits = W·x + b, softmax over atoms axis
// ----------------------------------------------------------------------
//
// One thread per (sample, action) cell. Each thread:
// - computes n_atoms logits (sequential inner loop over state_dim per atom)
// - finds max for numerical stability
// - computes exp(logit - max), sums, normalizes
// - stores probs[sample, action, *]
//
// Per-thread scratch lives in registers/local memory (MAX_N_ATOMS floats).
// Total threads = batch · n_actions. For batch=600, n_actions=9 → 5400
// threads, fits comfortably in a single launch.
extern "C" __global__ void alpha_c51_forward_kernel(
const float* __restrict__ W, // [n_actions * n_atoms, state_dim]
const float* __restrict__ b, // [n_actions * n_atoms]
const float* __restrict__ X, // [batch, state_dim]
float* __restrict__ probs, // [batch, n_actions, n_atoms]
int batch,
int state_dim,
int n_actions,
int n_atoms
) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total = batch * n_actions;
if (idx >= total) return;
const int sample = idx / n_actions;
const int action = idx % n_actions;
float logits[MAX_N_ATOMS];
float maxv = -INFINITY;
// 1) Compute all logits, track max.
const int x_base = sample * state_dim;
const int wb_base = action * n_atoms;
for (int k = 0; k < n_atoms; ++k) {
const int w_row = wb_base + k;
float acc = b[w_row];
const int w_off = w_row * state_dim;
#pragma unroll 4
for (int j = 0; j < state_dim; ++j) {
acc += W[w_off + j] * X[x_base + j];
}
logits[k] = acc;
if (acc > maxv) maxv = acc;
}
// 2) exp(logit - max), accumulate sum.
float sum = 0.0f;
for (int k = 0; k < n_atoms; ++k) {
const float e = __expf(logits[k] - maxv);
logits[k] = e;
sum += e;
}
// 3) Normalize and store.
const float inv = 1.0f / sum;
const int p_base = (sample * n_actions + action) * n_atoms;
for (int k = 0; k < n_atoms; ++k) {
probs[p_base + k] = logits[k] * inv;
}
}
// ----------------------------------------------------------------------
// (2) Bellman categorical projection
// ----------------------------------------------------------------------
//
// For each batch sample b:
// - find a* = argmax_a E_p(s'_b, a)[Z] (greedy action under target net)
// - for each atom k: T̂z_k = r_b + (1 - done_b) · γ · z_k
// z_k = v_min + k · Δz; clip T̂z_k to [v_min, v_max]
// distribute p_k(s'_b, a*) onto floor/ceil bins of (T̂z_k - v_min) / Δz
//
// One thread per batch sample. No atomicAdd: each thread writes to its
// own m[b, :] row sequentially. Grid sized for B threads at block=256.
extern "C" __global__ void alpha_c51_project_kernel(
const float* __restrict__ probs_next, // [batch, n_actions, n_atoms] (from TARGET net)
const float* __restrict__ rewards, // [batch]
const float* __restrict__ dones, // [batch] (1.0 = terminal, else 0.0)
float v_min,
float v_max,
float gamma,
float delta_z,
float* __restrict__ m, // [batch, n_atoms] (overwritten)
int batch,
int n_actions,
int n_atoms
) {
const int b_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (b_idx >= batch) return;
const float r = rewards[b_idx];
const float done = dones[b_idx];
const int probs_base_per_action = b_idx * n_actions * n_atoms;
const int m_base = b_idx * n_atoms;
// (a) Greedy action under target distribution.
int best_a = 0;
float best_q = -INFINITY;
for (int a = 0; a < n_actions; ++a) {
const int pb = probs_base_per_action + a * n_atoms;
float ez = 0.0f;
for (int k = 0; k < n_atoms; ++k) {
const float z_k = v_min + (float)k * delta_z;
ez += z_k * probs_next[pb + k];
}
if (ez > best_q) {
best_q = ez;
best_a = a;
}
}
const int probs_base = probs_base_per_action + best_a * n_atoms;
// (b) Zero target row.
for (int k = 0; k < n_atoms; ++k) {
m[m_base + k] = 0.0f;
}
// (c) Project each support atom forward and distribute.
// Borrowed from production `block_bellman_project_f` (c51_loss_kernel.cu):
// Huber-style negative-tail compression applied BEFORE the v_min clamp.
// Smooth exponential squeeze on negatives so a single catastrophic
// terminal reward can't dominate the gradient via long-tail accumulation
// at v_min. Positives untouched; the v_max clamp afterwards still caps
// the upper end. f(-5) ≈ -3.93, f(-15) ≈ -7.77, f(0) = 0.
const float gamma_eff = (1.0f - done) * gamma;
for (int k = 0; k < n_atoms; ++k) {
const float z_k = v_min + (float)k * delta_z;
float tz = r + gamma_eff * z_k;
if (tz < 0.0f) {
tz = -10.0f * (1.0f - __expf(tz / 10.0f));
}
if (tz < v_min) tz = v_min;
if (tz > v_max) tz = v_max;
const float bin = (tz - v_min) / delta_z;
int lo = (int)floorf(bin);
int hi = (int)ceilf(bin);
if (lo < 0) lo = 0;
if (hi > n_atoms - 1) hi = n_atoms - 1;
const float p = probs_next[probs_base + k];
if (lo == hi) {
m[m_base + lo] += p;
} else {
const float frac = bin - (float)lo;
m[m_base + lo] += p * (1.0f - frac);
m[m_base + hi] += p * frac;
}
}
}
// ----------------------------------------------------------------------
// (3) C51 cross-entropy gradient
// ----------------------------------------------------------------------
//
// Loss per sample: L_b = -Σ_k m[b,k] · log p[b, a_taken, k]
// Gradient on logit[b, a_taken, k]: dL/dlogit = p[b, a_taken, k] - m[b, k]
// Then dL/dW[a*K+k, j] = Σ_b 𝟙{a_taken==a} · (p[b,a,k] - m[b,k]) · X[b,j]
// dL/db[a*K+k] = Σ_b 𝟙{a_taken==a} · (p[b,a,k] - m[b,k])
//
// Threads index a flat (n_actions·n_atoms·state_dim + n_actions·n_atoms)
// range. For batch=600 the inner loop is light enough to be sequential.
// `scale` (typ. 1/batch) is applied at write time so the consumer SGD step
// uses raw `lr` without renormalization.
extern "C" __global__ void alpha_c51_grad_kernel(
const float* __restrict__ probs, // [batch, n_actions, n_atoms]
const float* __restrict__ m, // [batch, n_atoms]
const int* __restrict__ actions, // [batch]
const float* __restrict__ X, // [batch, state_dim]
float* __restrict__ dW, // [n_actions * n_atoms, state_dim] (overwritten)
float* __restrict__ db, // [n_actions * n_atoms] (overwritten)
int batch,
int state_dim,
int n_actions,
int n_atoms,
float scale
) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total_w = n_actions * n_atoms * state_dim;
const int total_b = n_actions * n_atoms;
if (idx >= total_w + total_b) return;
if (idx < total_w) {
// dW[ak, j]
const int ak = idx / state_dim;
const int j = idx % state_dim;
const int a = ak / n_atoms;
const int k = ak % n_atoms;
float grad = 0.0f;
for (int b_idx = 0; b_idx < batch; ++b_idx) {
if (actions[b_idx] == a) {
const float p = probs[(b_idx * n_actions + a) * n_atoms + k];
const float mv = m[b_idx * n_atoms + k];
grad += (p - mv) * X[b_idx * state_dim + j];
}
}
dW[idx] = grad * scale;
} else {
const int ak = idx - total_w;
const int a = ak / n_atoms;
const int k = ak % n_atoms;
float grad = 0.0f;
for (int b_idx = 0; b_idx < batch; ++b_idx) {
if (actions[b_idx] == a) {
const float p = probs[(b_idx * n_actions + a) * n_atoms + k];
const float mv = m[b_idx * n_atoms + k];
grad += (p - mv);
}
}
db[ak] = grad * scale;
}
}
// ----------------------------------------------------------------------
// (4) Thompson-select with confidence gate
// ----------------------------------------------------------------------
//
// GPU action selector: replaces the CPU epsilon_greedy_gated path entirely
// (per `feedback_cpu_is_read_only` — action selection is not "read-only";
// it's a sampling reduction).
//
// For each batch sample b:
// 1. If `states[b, conf_idx] < threshold` → out_action[b] = 0 (Wait)
// 2. Otherwise: per-action Thompson sample via inverse-CDF on probs[b,a,:],
// then pick the action whose sample has the largest z-value.
//
// Borrowed from production `thompson_direction_test_batched`
// (thompson_test_kernel.cu): LCG RNG seeded per-sample, inverse-CDF over
// the categorical distribution. One thread per batch sample.
//
// **Pinned-mapped output contract**: `out_action` MUST point to a
// `cuMemHostAlloc(DEVICEMAP|PORTABLE)` buffer's device-visible pointer
// (`cuMemHostGetDevicePointer_v2`). The kernel issues
// `__threadfence_system()` after the final write so the host can read
// the value through the corresponding host pointer using
// `std::ptr::read_volatile` without a `memcpy_dtoh`. Same pattern as
// `MappedBuffer` in `gpu_training_guard.rs`.
//
// The threshold scalar lives in ISV[543] on GPU but is also cached host-
// side and passed as a kernel arg (cold-path readback at rollout boundary,
// not per-step CPU compute). Same pattern as `current_threshold` in the
// existing linear-Q smoke.
extern "C" __global__ void alpha_c51_thompson_select_kernel(
const float* __restrict__ probs, // [batch, n_actions, n_atoms]
const float* __restrict__ states, // [batch, state_dim]
float threshold, // confidence gate
int conf_idx, // index into state for alpha_confidence
int state_dim,
float v_min,
float delta_z,
unsigned int base_seed, // per-call deterministic seed
int* __restrict__ out_action, // [batch] — overwritten
int batch,
int n_actions,
int n_atoms
) {
const int b_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (b_idx >= batch) return;
// Gate: alpha_confidence < threshold → Wait.
const float alpha_conf = states[b_idx * state_dim + conf_idx];
if (alpha_conf < threshold) {
out_action[b_idx] = 0;
return;
}
// LCG RNG seeded per (call_seed, batch_idx) — same primitive as
// production thompson_test_kernel.cu::test_rand. Distinct streams
// per (call, sample) by mixing the index into the seed.
unsigned int rng_state = base_seed + (unsigned int)b_idx;
float best_z = -INFINITY;
int best_a = 0;
for (int a = 0; a < n_actions; ++a) {
// One uniform draw per action.
rng_state = rng_state * 1664525u + 1013904223u;
const float u = (float)(rng_state >> 8) / (float)(1u << 24);
// Inverse-CDF over probs[b_idx, a, :].
const int p_base = (b_idx * n_actions + a) * n_atoms;
float cum = 0.0f;
// Numerical safety: if u >= sum (shouldn't happen post-softmax but
// floating-point can produce u slightly > cumulative sum) fall
// through to last atom.
float sampled_z = v_min + (float)(n_atoms - 1) * delta_z;
for (int k = 0; k < n_atoms; ++k) {
cum += probs[p_base + k];
if (u < cum) {
sampled_z = v_min + (float)k * delta_z;
break;
}
}
if (sampled_z > best_z) {
best_z = sampled_z;
best_a = a;
}
}
out_action[b_idx] = best_a;
// Mapped-pinned write: fence makes the value PCIe-visible to the host
// without an explicit memcpy_dtoh. Host reads via read_volatile.
__threadfence_system();
}
// ----------------------------------------------------------------------
// (5) Expected Q = Σ_k z_k · p_k(s, a)
// ----------------------------------------------------------------------
//
// Converts the C51 categorical distribution into a scalar Q-value per
// (sample, action). Needed to feed the existing `alpha_kill_criteria`
// kernel (which expects scalar Q[B, A]) when running in C51 mode, and
// for diagnostics. One thread per (sample, action).
extern "C" __global__ void alpha_c51_expected_q_kernel(
const float* __restrict__ probs, // [batch, n_actions, n_atoms]
float v_min,
float delta_z,
float* __restrict__ q_out, // [batch, n_actions] (overwritten)
int batch,
int n_actions,
int n_atoms
) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total = batch * n_actions;
if (idx >= total) return;
const int p_base = idx * n_atoms;
float ez = 0.0f;
for (int k = 0; k < n_atoms; ++k) {
const float z_k = v_min + (float)k * delta_z;
ez += z_k * probs[p_base + k];
}
q_out[idx] = ez;
}

277
crates/ml/src/cuda_pipeline/alpha_kernels.rs
View File

@@ -437,6 +437,283 @@ pub static ALPHA_LINEAR_Q_CUBIN: &[u8] =
pub static STACKER_THRESHOLD_CONTROLLER_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/stacker_threshold_controller.cubin"));
/// Precompiled vanilla C51 distributional Q cubin. Phase E.3 follow-up
/// (2026-05-15). Three kernels: forward (logits → softmax across atoms),
/// Bellman categorical projection, CE gradient.
pub static ALPHA_C51_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/alpha_c51.cubin"));
/// Launch `alpha_c51_forward_kernel`. One thread per (sample, action);
/// each thread sequentially computes n_atoms logits + softmax. Per-thread
/// scratch lives in local memory (MAX_N_ATOMS = 64 floats hardcoded in
/// the kernel — caller's `n_atoms` MUST be ≤ 64).
///
/// # Safety
/// All buffer pointers MUST be valid device pointers. `probs_dev` must
/// point at a writable `[batch * n_actions * n_atoms]`-float region.
pub unsafe fn launch_alpha_c51_forward(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
w_dev: u64,
b_dev: u64,
x_dev: u64,
probs_dev: u64,
batch: i32,
state_dim: i32,
n_actions: i32,
n_atoms: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(batch > 0, "batch must be positive");
debug_assert!(state_dim > 0, "state_dim must be positive");
debug_assert!(n_actions > 0, "n_actions must be positive");
debug_assert!(n_atoms > 0 && n_atoms <= 64, "n_atoms must be in (0, 64]");
const BLOCK: u32 = 256;
let total = (batch as u32) * (n_actions as u32);
let grid_x = (total + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), 1, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&w_dev)
.arg(&b_dev)
.arg(&x_dev)
.arg(&probs_dev)
.arg(&batch)
.arg(&state_dim)
.arg(&n_actions)
.arg(&n_atoms)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_c51_forward launch: {e}")))?;
Ok(())
}
/// Launch `alpha_c51_project_kernel`. Categorical Bellman projection: for
/// each batch sample, computes greedy action under the target distribution
/// then projects the shifted/scaled support `r + γ·z` onto the fixed atom
/// grid. One thread per sample.
///
/// # Safety
/// All buffer pointers MUST be valid device pointers. `m_dev` must point
/// at a writable `[batch * n_atoms]`-float region — overwritten.
pub unsafe fn launch_alpha_c51_project(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
probs_next_dev: u64,
rewards_dev: u64,
dones_dev: u64,
v_min: f32,
v_max: f32,
gamma: f32,
delta_z: f32,
m_dev: u64,
batch: i32,
n_actions: i32,
n_atoms: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(batch > 0, "batch must be positive");
debug_assert!(n_actions > 0, "n_actions must be positive");
debug_assert!(n_atoms > 0 && n_atoms <= 64, "n_atoms must be in (0, 64]");
debug_assert!(v_max > v_min, "v_max must exceed v_min");
debug_assert!(delta_z > 0.0, "delta_z must be positive");
debug_assert!((0.0..=1.0).contains(&gamma), "gamma must be in [0, 1]");
const BLOCK: u32 = 256;
let grid_x = ((batch as u32) + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), 1, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&probs_next_dev)
.arg(&rewards_dev)
.arg(&dones_dev)
.arg(&v_min)
.arg(&v_max)
.arg(&gamma)
.arg(&delta_z)
.arg(&m_dev)
.arg(&batch)
.arg(&n_actions)
.arg(&n_atoms)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_c51_project launch: {e}")))?;
Ok(())
}
/// Launch `alpha_c51_expected_q_kernel`. Computes scalar E[Z(s,a)] from
/// the C51 categorical distribution; one thread per (sample, action).
/// Used to bridge C51 outputs to the legacy scalar-Q kill-criteria kernel.
///
/// # Safety
/// All buffer pointers MUST be valid device pointers. `q_out_dev` must
/// point at a writable `[batch * n_actions]`-float region.
pub unsafe fn launch_alpha_c51_expected_q(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
probs_dev: u64,
v_min: f32,
delta_z: f32,
q_out_dev: u64,
batch: i32,
n_actions: i32,
n_atoms: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(batch > 0, "batch must be positive");
debug_assert!(n_actions > 0, "n_actions must be positive");
debug_assert!(n_atoms > 0 && n_atoms <= 64, "n_atoms must be in (0, 64]");
debug_assert!(delta_z > 0.0, "delta_z must be positive");
const BLOCK: u32 = 256;
let total = (batch as u32) * (n_actions as u32);
let grid_x = (total + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), 1, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&probs_dev)
.arg(&v_min)
.arg(&delta_z)
.arg(&q_out_dev)
.arg(&batch)
.arg(&n_actions)
.arg(&n_atoms)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_c51_expected_q launch: {e}")))?;
Ok(())
}
/// Launch `alpha_c51_thompson_select_kernel`. GPU action selector with
/// confidence gate. Replaces CPU `epsilon_greedy_gated` per
/// `feedback_cpu_is_read_only`. One thread per batch sample; writes the
/// chosen action index into `out_action_dev`.
///
/// # Safety
/// All buffer pointers MUST be valid device pointers. `out_action_dev`
/// must point at a writable `[batch]`-int region — overwritten.
pub unsafe fn launch_alpha_c51_thompson_select(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
probs_dev: u64,
states_dev: u64,
threshold: f32,
conf_idx: i32,
state_dim: i32,
v_min: f32,
delta_z: f32,
base_seed: u32,
out_action_dev: u64,
batch: i32,
n_actions: i32,
n_atoms: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(batch > 0, "batch must be positive");
debug_assert!(state_dim > 0, "state_dim must be positive");
debug_assert!(n_actions > 0, "n_actions must be positive");
debug_assert!(n_atoms > 0 && n_atoms <= 64, "n_atoms must be in (0, 64]");
debug_assert!(conf_idx >= 0 && conf_idx < state_dim, "conf_idx in range");
debug_assert!(delta_z > 0.0, "delta_z must be positive");
const BLOCK: u32 = 256;
let grid_x = ((batch as u32) + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), 1, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&probs_dev)
.arg(&states_dev)
.arg(&threshold)
.arg(&conf_idx)
.arg(&state_dim)
.arg(&v_min)
.arg(&delta_z)
.arg(&base_seed)
.arg(&out_action_dev)
.arg(&batch)
.arg(&n_actions)
.arg(&n_atoms)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_c51_thompson_select launch: {e}")))?;
Ok(())
}
/// Launch `alpha_c51_grad_kernel`. CE gradient: `dW[ak, j] += scale ·
/// 𝟙{a_taken[b]==a} · (p[b,a,k] − m[b,k]) · X[b,j]`. Flat-thread layout
/// over (`n_actions·n_atoms·state_dim` + `n_actions·n_atoms`) range.
///
/// # Safety
/// All buffer pointers MUST be valid device pointers. `dW_dev` must point
/// at a writable `[n_actions * n_atoms * state_dim]`-float region; `db_dev`
/// at `[n_actions * n_atoms]`. Both overwritten.
pub unsafe fn launch_alpha_c51_grad(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
probs_dev: u64,
m_dev: u64,
actions_dev: u64,
x_dev: u64,
dw_dev: u64,
db_dev: u64,
batch: i32,
state_dim: i32,
n_actions: i32,
n_atoms: i32,
scale: f32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(batch > 0, "batch must be positive");
debug_assert!(state_dim > 0, "state_dim must be positive");
debug_assert!(n_actions > 0, "n_actions must be positive");
debug_assert!(n_atoms > 0 && n_atoms <= 64, "n_atoms must be in (0, 64]");
const BLOCK: u32 = 256;
let total_w = (n_actions as u32) * (n_atoms as u32) * (state_dim as u32);
let total_b = (n_actions as u32) * (n_atoms as u32);
let total = total_w + total_b;
let grid_x = (total + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), 1, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&probs_dev)
.arg(&m_dev)
.arg(&actions_dev)
.arg(&x_dev)
.arg(&dw_dev)
.arg(&db_dev)
.arg(&batch)
.arg(&state_dim)
.arg(&n_actions)
.arg(&n_atoms)
.arg(&scale)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_c51_grad launch: {e}")))?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;

47
crates/ml/src/env/execution_env.rs vendored
View File

@@ -73,8 +73,13 @@ impl Default for EpisodeState {
#[derive(Debug, Clone)]
pub struct SnapshotRow {
pub mid_price: f32,
pub bid_l: [f32; 3],
pub ask_l: [f32; 3],
/// Bid prices at levels L1..L10 (index 0 = L1 = top of book; deeper
/// indices are further from mid). Phase E.4.A.3 extension from
/// `[f32; 3]` to enable real depth from MBP-10 raw data while
/// keeping env.step's L1/L2 access (indices 0/1) unchanged.
pub bid_l: [f32; 10],
/// Ask prices at levels L1..L10 (index 0 = L1; deeper = further).
pub ask_l: [f32; 10],
pub alpha_logit: f32,
pub alpha_confidence: f32,
pub spread_bps: f32,
@@ -365,10 +370,16 @@ mod tests {
(0..n)
.map(|i| {
let drift = i as f32 * slope;
let mut bid_l = [0.0_f32; 10];
let mut ask_l = [0.0_f32; 10];
for k in 0..10 {
bid_l[k] = 4499.875 + drift - (k as f32) * 0.25;
ask_l[k] = 4500.125 + drift + (k as f32) * 0.25;
}
SnapshotRow {
mid_price: 4500.0 + drift,
bid_l: [4499.875 + drift, 4499.625 + drift, 4499.375 + drift],
ask_l: [4500.125 + drift, 4500.375 + drift, 4500.625 + drift],
bid_l,
ask_l,
alpha_logit: 1.0,
alpha_confidence: 0.3,
spread_bps: 0.25,
@@ -396,6 +407,34 @@ mod tests {
)
}
#[test]
fn snapshot_row_carries_ten_levels_of_depth() {
// Phase E.4.A.3: SnapshotRow.bid_l / ask_l extended from
// [f32; 3] to [f32; 10] to carry real MBP-10 depth.
let mut bid_l = [0.0_f32; 10];
let mut ask_l = [0.0_f32; 10];
for k in 0..10 {
bid_l[k] = 4500.0 - 0.125 - (k as f32) * 0.25;
ask_l[k] = 4500.0 + 0.125 + (k as f32) * 0.25;
}
let row = SnapshotRow {
mid_price: 4500.0,
bid_l, ask_l,
alpha_logit: 0.0, alpha_confidence: 0.0,
spread_bps: 5.0, l1_imbalance: 0.5, ofi_sum_5: 0.0,
mid_drift_5: 0.0, time_since_trade_s: 0.0, book_event_rate: 5.0,
};
assert_eq!(row.bid_l.len(), 10);
assert_eq!(row.ask_l.len(), 10);
// L1 < L10 for bid (deeper = lower price)
assert!(row.bid_l[0] > row.bid_l[9]);
// L1 < L10 for ask (deeper = higher price)
assert!(row.ask_l[0] < row.ask_l[9]);
// L1/L2 indices (env.step's actual reads) remain accessible
assert!((row.bid_l[0] - 4499.875).abs() < 1e-5);
assert!((row.ask_l[1] - 4500.375).abs() < 1e-5);
}
#[test]
fn state_has_expected_dim_and_is_finite() {
let env = make_env(100, 0.1, 100);

86
crates/ml/src/env/loaders.rs vendored
View File

@@ -97,10 +97,36 @@ pub fn load_fill_model_from_json(path: &Path) -> Result<FillModel> {
/// Block-S features feed the runtime feature fields. If `alpha_cache`
/// is `Some`, each `SnapshotRow.alpha_logit` is populated from the
/// cache (and `alpha_confidence = |sigmoid(z) − 0.5|`); otherwise 0.0.
/// Maximum spread we'll accept from `spread_bps` (in price units) when
/// `use_real_spread` is true. Above this we cap to the max — protects
/// against fxcache feature outliers / NaN / unrealistic wide ticks.
/// 10 ticks = 2.50 in ES futures price = ~5.6 bps at mid=4500.
const MAX_REAL_SPREAD_PRICE: f32 = 10.0 * TICK;
/// Build `Vec<SnapshotRow>` from a precomputed fxcache. Mid from
/// `raw_close`; 81-dim Block-S features feed the runtime feature fields.
///
/// **Bid/ask synthesis:**
/// - `use_real_spread = false` (legacy): bid_l1 = mid − 0.125, ask_l1
/// = mid + 0.125. Fixed half-tick spread, equivalent to the original
/// loader behaviour and matched the Phase E.1/2/3 smoke/backtest runs.
/// - `use_real_spread = true` (Path 3): spread_price derived from
/// fxcache `features[78]` (`spread_bps`); bid_l1 = mid − spread/2,
/// ask_l1 = mid + spread/2. Variable per-bar spread reflecting
/// actual market state. Floored at 1 tick (=0.25), capped at 10 ticks
/// to protect against feature outliers.
///
/// L2/L3 are always synthesized at ±TICK offsets from L1 — the fxcache
/// doesn't store depth beyond L1 spread+imbalance, and L2/L3 fills are
/// rare (most policy decisions hinge on L1 + market crosses).
///
/// If `alpha_cache` is `Some`, each row's `alpha_logit` is populated
/// from the cache (and `alpha_confidence = |sigmoid(z) − 0.5|`).
pub fn load_snapshots_from_fxcache(
fxcache_path: &Path,
max_snapshots: usize,
alpha_cache: Option<&[f32]>,
use_real_spread: bool,
) -> Result<Vec<SnapshotRow>> {
let reader = FxCacheReader::open(fxcache_path)
.with_context(|| format!("open fxcache {}", fxcache_path.display()))?;
@@ -116,6 +142,11 @@ pub fn load_snapshots_from_fxcache(
let n_bars_total = reader.bar_count();
let n = n_bars_total.min(max_snapshots);
info!("fxcache: {} total bars, taking {} for the env", n_bars_total, n);
info!(
"fxcache loader: spread mode = {}",
if use_real_spread { "REAL (derived from features[78] spread_bps)" }
else { "FIXED ±0.125-tick" }
);
if let Some(cache) = alpha_cache {
if cache.len() < n {
@@ -129,6 +160,12 @@ pub fn load_snapshots_from_fxcache(
let mut rows: Vec<SnapshotRow> = Vec::with_capacity(n);
let mut n_degenerate = 0_usize;
// Diagnostic: spread distribution (real mode only).
let mut sum_spread = 0.0_f64;
let mut min_spread = f32::INFINITY;
let mut max_spread = f32::NEG_INFINITY;
let mut n_floor_hits = 0_usize;
let mut n_cap_hits = 0_usize;
for i in 0..n {
let rec = reader.record(i);
let mid = rec.targets[COL_RAW_CLOSE - FEAT_DIM];
@@ -140,11 +177,6 @@ pub fn load_snapshots_from_fxcache(
.alpha_features(i)
.ok_or_else(|| anyhow::anyhow!("missing alpha row at bar {}", i))?;
let bid_l1 = mid - 0.125;
let ask_l1 = mid + 0.125;
let bid_l = [bid_l1, bid_l1 - TICK, bid_l1 - 2.0 * TICK];
let ask_l = [ask_l1, ask_l1 + TICK, ask_l1 + 2.0 * TICK];
let spread_bps = features[78];
let l1_imbalance = features[79];
let ofi_sum_5 = features[0..5].iter().sum::<f32>();
@@ -152,6 +184,40 @@ pub fn load_snapshots_from_fxcache(
let time_since_trade_s = features[75];
let book_event_rate = features[77];
let (bid_l1, ask_l1) = if use_real_spread {
// spread_bps = 10000 × (ask − bid) / mid → spread_price = bps × mid / 10000
let raw = if spread_bps.is_finite() && spread_bps > 0.0 {
spread_bps / 10_000.0 * mid
} else {
TICK // sentinel: fall back to 1-tick spread if bps is non-finite
};
let mut clamped = raw;
if clamped < TICK {
clamped = TICK;
n_floor_hits += 1;
}
if clamped > MAX_REAL_SPREAD_PRICE {
clamped = MAX_REAL_SPREAD_PRICE;
n_cap_hits += 1;
}
sum_spread += clamped as f64;
if clamped < min_spread { min_spread = clamped; }
if clamped > max_spread { max_spread = clamped; }
let half = clamped * 0.5;
(mid - half, mid + half)
} else {
(mid - 0.125, mid + 0.125)
};
// Phase E.4.A.3: extend to L1-L10 depth. L2-L10 synthesized at
// ±TICK offsets from L1; real L4-L10 from MBP-10 lands in
// Task 5 follow-on (the loader doesn't have MBP-10 access yet).
let mut bid_l = [0.0_f32; 10];
let mut ask_l = [0.0_f32; 10];
for k in 0..10 {
bid_l[k] = bid_l1 - (k as f32) * TICK;
ask_l[k] = ask_l1 + (k as f32) * TICK;
}
let alpha_logit = alpha_cache.map(|c| c[i]).unwrap_or(0.0);
let alpha_confidence = {
let p = 1.0_f32 / (1.0 + (-alpha_logit.clamp(-50.0, 50.0)).exp());
@@ -178,5 +244,15 @@ pub fn load_snapshots_from_fxcache(
n_degenerate
);
}
if use_real_spread && !rows.is_empty() {
let mean_spread = (sum_spread / rows.len() as f64) as f32;
info!(
"fxcache loader: real-spread stats — mean={:.4} ({:.2} ticks), min={:.4}, max={:.4}, floor_hits={}/{}, cap_hits={}/{}",
mean_spread, mean_spread / TICK,
min_spread, max_spread,
n_floor_hits, rows.len(),
n_cap_hits, rows.len(),
);
}
Ok(rows)
}