cc40780b80
CUDA kernel launches via cudarc::launch_builder and MappedF32::new (cuMemHostAlloc DEVICEMAP FFI) inherently require unsafe blocks. The workspace-wide '-W unsafe-code' lint produced ~80 warnings across these files, all structurally identical. Match the established pattern from ml-backtesting/src/harness.rs and ml-backtesting/src/sim/mod.rs: single file-level allow with a comment explaining the rationale. Files: alpha_dqn_h600_smoke, alpha_baseline, cublaslt_debug, gpu_walk_forward, cuda_pipeline/mod, hyperopt adapters (mamba2, ppo), DQN smoke_tests (helpers, td_propagation), trainers/ppo, and two ml/tests files. Documented in dqn-wire-up-audit.md.
1724 lines
78 KiB
Rust
1724 lines
78 KiB
Rust
//! Phase E.1 Task 12 — H=600 DQN smoke for the kill-criteria gate.
|
||
//!
|
||
//! Trains a linear Q-network (W [9×10] + b [9], no hidden layer) on the
|
||
//! Phase E ExecutionEnv at horizon H=600. Uses ε-greedy action selection
|
||
//! with Munchausen target augmentation (alpha_munchausen_target_kernel)
|
||
//! and reports the four kill criteria at end of training:
|
||
//!
|
||
//! ISV[539] Q_SPREAD_EMA must be ≥ 0.05
|
||
//! ISV[540] ACTION_ENTROPY_EMA must be ≥ 0.5 · ln(9) ≈ 1.0986
|
||
//! ISV[541] RETURN_VS_RANDOM_EMA must be ≥ 0.0
|
||
//! ISV[542] EARLY_Q_MOVEMENT_EMA must be ≥ 0.01
|
||
//!
|
||
//! If ALL FOUR pass: H=6000 scale-up (Task 13) is viable. If ANY fail,
|
||
//! the plan calls for pivot to NoisyNet (Task 19).
|
||
//!
|
||
//! ## Architecture
|
||
//!
|
||
//! Single linear layer — no hidden layer. The 10-dim state vector
|
||
//! `[alpha_logit, alpha_confidence, spread_bps, l1_imbalance, ofi_sum_5,
|
||
//! mid_drift_5, position, step_normalized, log_tau, log_event_rate]`
|
||
//! has meaningful direct features, so linear Q captures real relations
|
||
//! like `Q[Buy] ∝ alpha_logit`. If linear can't pass the gate, no
|
||
//! architecture upgrade will save it — pivot to NoisyNet.
|
||
//!
|
||
//! All compute on GPU: forward, grad, weight update via alpha_linear_q
|
||
//! kernels; target augmentation via alpha_munchausen_target_kernel;
|
||
//! kill criteria via alpha_kill_criteria + apply_pearls_ad chain.
|
||
//! Action selection is read-only on CPU (per feedback_cpu_is_read_only,
|
||
//! reading 9 Q-values to pick argmax is not compute).
|
||
//!
|
||
//! ## Run
|
||
//!
|
||
//! ```bash
|
||
//! SQLX_OFFLINE=true cargo run -p ml --release --example alpha_dqn_h600_smoke -- \
|
||
//! --mbp10-dir /home/jgrusewski/Work/foxhunt/test_data/futures-baseline-mbp10/ES.FUT \
|
||
//! --fill-coeffs config/ml/alpha_fill_coeffs.json \
|
||
//! --horizon 600 \
|
||
//! --n-episodes 1000
|
||
//! ```
|
||
|
||
// CUDA kernel launches via `cudarc::launch_builder.launch(...)` and
|
||
// `MappedF32::new` (`cuMemHostAlloc(DEVICEMAP)` FFI) require `unsafe`. The
|
||
// workspace-wide `-W unsafe-code` lint is therefore noise here. Same rationale
|
||
// as `ml-backtesting/src/harness.rs` and `sim/mod.rs`'s file-level allows.
|
||
#![allow(unsafe_code)]
|
||
|
||
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, warn};
|
||
|
||
// Phase E.4.A.8: Mamba2 temporal encoder (ml-alpha Phase 1d.1).
|
||
// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights.
|
||
use ml_alpha::mamba2_block::{
|
||
Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig,
|
||
};
|
||
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
|
||
|
||
// ── Mapped-pinned helpers ──────────────────────────────────────────
|
||
//
|
||
// `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,
|
||
len: usize,
|
||
}
|
||
|
||
impl MappedI32 {
|
||
/// # 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::<i32>();
|
||
let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
|
||
.map_err(|e| anyhow::anyhow!("mapped i32 alloc: {e}"))?
|
||
as *mut i32;
|
||
std::ptr::write_bytes(host_ptr, 0, len);
|
||
let mut dev_raw = MaybeUninit::uninit();
|
||
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
|
||
dev_raw.as_mut_ptr(),
|
||
host_ptr as *mut std::ffi::c_void,
|
||
0,
|
||
)
|
||
.result()
|
||
.map_err(|e| anyhow::anyhow!("cuMemHostGetDevicePointer: {e}"))?;
|
||
Ok(Self {
|
||
host_ptr,
|
||
dev_ptr: dev_raw.assume_init(),
|
||
len,
|
||
})
|
||
}
|
||
fn dev_u64(&self) -> u64 {
|
||
self.dev_ptr as u64
|
||
}
|
||
fn read(&self) -> i32 {
|
||
debug_assert!(self.len >= 1);
|
||
unsafe { std::ptr::read_volatile(self.host_ptr) }
|
||
}
|
||
}
|
||
|
||
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, InstrumentFilter},
|
||
mbp10::Mbp10Snapshot,
|
||
};
|
||
use ml::cuda_pipeline::alpha_isv_slots::{
|
||
ACTION_ENTROPY_EMA_INDEX, ALPHA_ISV_BLOCK_LO, EARLY_Q_MOVEMENT_EMA_INDEX,
|
||
Q_SPREAD_EMA_INDEX, RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX,
|
||
RETURN_VS_RANDOM_EMA_INDEX,
|
||
};
|
||
// `COL_RAW_CLOSE` / `FEAT_DIM` are no longer used directly in this
|
||
// binary; the snapshot loader handles fxcache column layout internally.
|
||
use ml::env::action_space::N_ACTIONS;
|
||
use ml::env::execution_env::{
|
||
EpisodeState, ExecutionEnv, ExecutionEnvConfig, SnapshotRow,
|
||
};
|
||
// `FillCoeffs` / `FillModel` are loaded via the env::loaders module
|
||
// and consumed inside ExecutionEnv — no direct references here.
|
||
use ml::trainers::dqn::collect_dbn_files_recursive;
|
||
|
||
const STATE_DIM: usize = 10;
|
||
const TICK: f32 = 0.25;
|
||
/// Number of weight floats: n_actions × state_dim = 9 × 10.
|
||
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
|
||
}
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(
|
||
name = "alpha_dqn_h600_smoke",
|
||
about = "Phase E.1 Task 12 — H=600 DQN smoke for kill-criteria gate"
|
||
)]
|
||
struct Cli {
|
||
/// Snapshot source 1: raw MBP-10 directory (bid/ask from real LOB,
|
||
/// alpha_logit=0 placeholder unless --alpha-cache also set).
|
||
#[arg(long)]
|
||
mbp10_dir: Option<PathBuf>,
|
||
/// Snapshot source 2: precomputed fxcache (bid/ask synthesized from
|
||
/// mid at fixed half-tick, 81-dim feature row available). Required when
|
||
/// --alpha-cache is set since the cache indexing is fxcache-aligned.
|
||
#[arg(long)]
|
||
fxcache_path: Option<PathBuf>,
|
||
/// Optional alpha-logits cache produced by `alpha_train_stacker
|
||
/// --alpha-cache-out`. Binary file: [u32 n] [f32 logits[n]]. When set,
|
||
/// SnapshotRow.alpha_logit is populated from the cache instead of 0.0.
|
||
/// Requires --fxcache-path (cache indices align to fxcache bars).
|
||
#[arg(long)]
|
||
alpha_cache: Option<PathBuf>,
|
||
#[arg(long, default_value = "config/ml/alpha_fill_coeffs.json")]
|
||
fill_coeffs: PathBuf,
|
||
#[arg(long, default_value_t = 600)]
|
||
horizon: usize,
|
||
#[arg(long, default_value_t = 1_000)]
|
||
n_episodes: usize,
|
||
#[arg(long, default_value_t = 1)]
|
||
trade_size: i32,
|
||
#[arg(long, default_value_t = 0.0625)]
|
||
cost_per_contract: f32,
|
||
#[arg(long, default_value_t = 0xCAFEBABE_u64)]
|
||
seed: u64,
|
||
#[arg(long, default_value_t = 50)]
|
||
snapshot_interval: usize,
|
||
#[arg(long, default_value_t = 500_000)]
|
||
max_snapshots: usize,
|
||
/// SGD learning rate. With the stabilizer trio (--reward-scale,
|
||
/// --target-update-every, --grad-clip), 1e-4 is the right operating
|
||
/// point — gradients are O(1) in normalized units, target net
|
||
/// breaks the V_soft chase-tail loop, grad clip catches bursts.
|
||
/// Without the stabilizers, lr=1e-4 diverges to NaN within 20
|
||
/// episodes; lr=1e-6 stays finite but undertrains.
|
||
#[arg(long, default_value_t = 1.0e-4)]
|
||
lr: f32,
|
||
/// ε-greedy: ε at start of training.
|
||
#[arg(long, default_value_t = 0.50)]
|
||
eps_start: f32,
|
||
/// ε-greedy: ε at end of training.
|
||
#[arg(long, default_value_t = 0.05)]
|
||
eps_end: f32,
|
||
/// DQN discount factor.
|
||
#[arg(long, default_value_t = 0.99)]
|
||
gamma: f32,
|
||
/// Munchausen scale (Vieillard 2020 default 0.9).
|
||
#[arg(long, default_value_t = 0.9)]
|
||
alpha_m: f32,
|
||
/// Boltzmann temperature for Munchausen (Vieillard default 0.03).
|
||
#[arg(long, default_value_t = 0.03)]
|
||
tau: f32,
|
||
/// Lower clip for τ·log π(a|s) (Vieillard default -1.0).
|
||
#[arg(long, default_value_t = -1.0)]
|
||
log_clip_min: f32,
|
||
/// Episodes between kill-criteria pipeline launches.
|
||
#[arg(long, default_value_t = 50)]
|
||
kill_criteria_every: usize,
|
||
/// Phase E.2 stacker-threshold controller — target trade rate (ISV[544]).
|
||
/// Set once at training start, never reset (TrainingPersist anchor).
|
||
#[arg(long, default_value_t = 0.08)]
|
||
trade_rate_target: f32,
|
||
/// Stacker-threshold controller P-gain on rate error.
|
||
#[arg(long, default_value_t = 0.01)]
|
||
k_threshold: f32,
|
||
/// Kelly-attenuation controller P-gain on Sharpe error.
|
||
#[arg(long, default_value_t = 0.005)]
|
||
k_atten: f32,
|
||
/// Target rollout Sharpe for the Kelly-attenuation controller.
|
||
#[arg(long, default_value_t = 0.5)]
|
||
target_sharpe: f32,
|
||
/// Pearl D Wiener-α floor on the observed-rate EMA (slot 545). 0.4 per
|
||
/// `pearl_wiener_alpha_floor_for_nonstationary` — controller co-adapts
|
||
/// with the policy, so the EMA needs to stay responsive.
|
||
#[arg(long, default_value_t = 0.4)]
|
||
wiener_alpha_floor: f32,
|
||
/// Meta-α for the Wiener variance updates inside the controller.
|
||
#[arg(long, default_value_t = 0.1)]
|
||
ctl_alpha_meta: f32,
|
||
/// Reward normalization scale. Rewards are divided by this before
|
||
/// being passed to the Munchausen target — the network learns
|
||
/// normalized Q-values. Default 1000 ≈ |median reward| at H=600
|
||
/// (≈-2800 median random-baseline reward / sane order). Action
|
||
/// selection and rollout-R reporting use ORIGINAL rewards.
|
||
#[arg(long, default_value_t = 1000.0)]
|
||
reward_scale: f32,
|
||
/// Episodes between hard updates of the target network (W_target,
|
||
/// b_target ← W_online, b_online). Decouples the V_soft(s')
|
||
/// bootstrap target from the policy's per-step drift — essential
|
||
/// for Munchausen DQN stability.
|
||
#[arg(long, default_value_t = 10)]
|
||
target_update_every: usize,
|
||
/// Per-element gradient clip bound. After alpha_linear_q_grad and
|
||
/// before the SGD step, |dW[i]|, |db[i]| are clamped to ±this value.
|
||
/// Default 1.0 — combined with lr=1e-4 caps per-step weight delta
|
||
/// 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,
|
||
/// Phase E.4.A: enable the temporal-encoder stack (sliding window
|
||
/// buffer + Mamba2 SSM + C51 head). When false, runs the existing
|
||
/// C51-flat path on the current state vector only. The buffer is
|
||
/// maintained but unused until Task 8 wires Mamba2 over it.
|
||
#[arg(long, default_value_t = false)]
|
||
temporal: bool,
|
||
/// Sliding-window length K (in snapshots). 16 starting point; the
|
||
/// kernel max is `MAMBA2_KERNEL_SEQ_MAX = 32`. Sweep 16 / 32 in
|
||
/// later iterations.
|
||
#[arg(long, default_value_t = 16)]
|
||
window_k: usize,
|
||
/// Mamba2 hidden dimension (= sh2 in kernel). 32 starting point.
|
||
/// Also the input dim to the C51 head when --temporal is set
|
||
/// (C51 W shape: [N_ACTIONS * n_atoms, mamba2_hidden_dim]).
|
||
#[arg(long, default_value_t = 32)]
|
||
mamba2_hidden_dim: usize,
|
||
/// Mamba2 SSM state dimension. Kernel max is
|
||
/// `MAMBA2_KERNEL_STATE_MAX = 16`.
|
||
#[arg(long, default_value_t = 16)]
|
||
mamba2_state_dim: usize,
|
||
/// Output JSON path for final verdict + ISV readings.
|
||
#[arg(long, default_value = "config/ml/alpha_dqn_h600_smoke.json")]
|
||
out_path: PathBuf,
|
||
}
|
||
|
||
/// Build the env::SnapshotRow stream from MBP-10 files. Same loader pattern
|
||
/// as `alpha_random_baseline.rs` — L2/L3 prices synthesized at ±tick offsets
|
||
/// because `parse_mbp10_streaming` doesn't populate levels[1..10].
|
||
fn load_snapshots(
|
||
parser: &DbnParser,
|
||
mbp10_dir: &std::path::Path,
|
||
snapshot_interval: usize,
|
||
max_snapshots: usize,
|
||
) -> Result<Vec<SnapshotRow>> {
|
||
let files = collect_dbn_files_recursive(mbp10_dir);
|
||
if files.is_empty() {
|
||
anyhow::bail!("no MBP-10 .dbn[.zst] files in {:?}", mbp10_dir);
|
||
}
|
||
info!("Found {} MBP-10 file(s)", files.len());
|
||
let mut rows: Vec<SnapshotRow> = Vec::with_capacity(max_snapshots);
|
||
'files: for file in &files {
|
||
let mut hit_limit = false;
|
||
let result = parser.parse_mbp10_streaming(file, snapshot_interval, InstrumentFilter::All, |snap: &Mbp10Snapshot| {
|
||
if rows.len() >= max_snapshots {
|
||
hit_limit = true;
|
||
return;
|
||
}
|
||
if snap.levels.is_empty() {
|
||
return;
|
||
}
|
||
let l1 = &snap.levels[0];
|
||
let bid_l1 = raw_price_to_f32(l1.bid_px);
|
||
let ask_l1 = raw_price_to_f32(l1.ask_px);
|
||
if bid_l1 <= 0.0 || ask_l1 <= 0.0 || bid_l1 >= ask_l1 {
|
||
return;
|
||
}
|
||
let mid = 0.5 * (bid_l1 + ask_l1);
|
||
let spread_bps = 10_000.0 * (ask_l1 - bid_l1) / mid;
|
||
let bid_sz = l1.bid_sz as f32;
|
||
let ask_sz = l1.ask_sz as f32;
|
||
let l1_imbalance = if bid_sz + ask_sz > 0.0 {
|
||
bid_sz / (bid_sz + ask_sz)
|
||
} else {
|
||
0.5
|
||
};
|
||
let ofi_sum_5 = bid_sz - ask_sz;
|
||
// 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,
|
||
ask_l,
|
||
alpha_logit: 0.0,
|
||
alpha_confidence: 0.5,
|
||
spread_bps,
|
||
l1_imbalance,
|
||
ofi_sum_5,
|
||
mid_drift_5: 0.0,
|
||
time_since_trade_s: 0.0,
|
||
book_event_rate: 5.0,
|
||
});
|
||
});
|
||
match result {
|
||
Ok(c) => info!(
|
||
" {} → {} streamed snapshots",
|
||
file.file_name().unwrap_or_default().to_string_lossy(),
|
||
c
|
||
),
|
||
Err(e) => warn!("parser error on {}: {}", file.display(), e),
|
||
}
|
||
if hit_limit {
|
||
break 'files;
|
||
}
|
||
}
|
||
Ok(rows)
|
||
}
|
||
|
||
/// Compute env-state observation as Vec<f32> for direct upload.
|
||
fn state_as_vec(env: &ExecutionEnv, ep: &EpisodeState) -> Vec<f32> {
|
||
env.state(ep).to_vec()
|
||
}
|
||
|
||
/// L2 (Frobenius) norm of W concatenated with b — diagnostic anchor for
|
||
/// early-Q-movement. Pure read-side reduction; computed CPU-side after
|
||
/// dtoh (not on the hot loop, only at episode boundaries).
|
||
fn weight_norm(w: &[f32], b: &[f32]) -> f32 {
|
||
let s: f32 = w.iter().map(|x| x * x).sum::<f32>() + b.iter().map(|x| x * x).sum::<f32>();
|
||
s.sqrt()
|
||
}
|
||
|
||
/// L2 distance from init: `||W_now − W_init||_F + ||b_now − b_init||_F` (combined).
|
||
/// Captures *direction change* in weight space, not just growth — the kill-
|
||
/// criteria kernel computes `|q_early − q_init| / |q_init|`, so the call
|
||
/// site sets `q_early = q_init + weight_distance_from_init(...)`, making
|
||
/// the kernel's ratio = `distance_from_init / ||W_init||_F`. A pure
|
||
/// rotation that keeps `||W||` constant still produces a non-zero distance.
|
||
fn weight_distance_from_init(w_now: &[f32], b_now: &[f32], w_init: &[f32], b_init: &[f32]) -> f32 {
|
||
let mut s: f32 = 0.0;
|
||
for (a, b) in w_now.iter().zip(w_init.iter()) {
|
||
let d = a - b;
|
||
s += d * d;
|
||
}
|
||
for (a, b) in b_now.iter().zip(b_init.iter()) {
|
||
let d = a - b;
|
||
s += d * d;
|
||
}
|
||
s.sqrt()
|
||
}
|
||
|
||
/// Snapshot of all trainable Mamba2 parameters at init time. Concatenated
|
||
/// flat Vec<f32> in deterministic order so a later snapshot subtracts
|
||
/// element-wise. Used by the EARLY_Q_MOVEMENT_EMA diagnostic so the
|
||
/// kill criterion captures encoder movement (not just C51 head), which
|
||
/// would otherwise under-report system learning when --temporal is on.
|
||
fn mamba2_snapshot(
|
||
stream: &std::sync::Arc<cudarc::driver::CudaStream>,
|
||
block: &Mamba2Block,
|
||
) -> anyhow::Result<Vec<f32>> {
|
||
let mut out = Vec::new();
|
||
out.extend(stream.clone_dtoh(&block.w_in.weight).context("dtoh w_in")?);
|
||
out.extend(stream.clone_dtoh(&block.w_in.bias).context("dtoh w_in.b")?);
|
||
out.extend(stream.clone_dtoh(&block.w_a.weight).context("dtoh w_a")?);
|
||
out.extend(stream.clone_dtoh(&block.w_a.bias).context("dtoh w_a.b")?);
|
||
out.extend(stream.clone_dtoh(&block.w_b.weight).context("dtoh w_b")?);
|
||
out.extend(stream.clone_dtoh(&block.w_b.bias).context("dtoh w_b.b")?);
|
||
out.extend(stream.clone_dtoh(&block.w_c).context("dtoh w_c")?);
|
||
out.extend(stream.clone_dtoh(&block.w_out.weight).context("dtoh w_out")?);
|
||
out.extend(stream.clone_dtoh(&block.w_out.bias).context("dtoh w_out.b")?);
|
||
Ok(out)
|
||
}
|
||
|
||
/// L2 distance ||M_now − M_init||_F across the same concatenated layout
|
||
/// produced by `mamba2_snapshot`. Returns 0.0 if init has different
|
||
/// length (shouldn't happen unless model topology changed).
|
||
fn mamba2_weight_distance(now: &[f32], init: &[f32]) -> f32 {
|
||
if now.len() != init.len() {
|
||
return 0.0;
|
||
}
|
||
let mut s: f32 = 0.0;
|
||
for (a, b) in now.iter().zip(init.iter()) {
|
||
let d = a - b;
|
||
s += d * d;
|
||
}
|
||
s.sqrt()
|
||
}
|
||
|
||
/// SplitMix64 RNG — same as ExecutionEnv's ReplayRng so seeds compose
|
||
/// cleanly when the smoke is reproduced.
|
||
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)
|
||
}
|
||
}
|
||
|
||
/// Phase E.3 confidence-gated ε-greedy. If `alpha_confidence < threshold`,
|
||
/// 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, allowed, rng)
|
||
}
|
||
|
||
/// 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 {
|
||
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
|
||
}
|
||
}
|
||
|
||
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.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")?;
|
||
let stream = ctx.default_stream();
|
||
|
||
// Linear Q kernels
|
||
let lq_module = ctx
|
||
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_LINEAR_Q_CUBIN.to_vec())
|
||
.context("alpha_linear_q cubin load")?;
|
||
let lq_fwd = lq_module.load_function("alpha_linear_q_forward_kernel")
|
||
.context("forward load")?;
|
||
let lq_grad = lq_module.load_function("alpha_linear_q_grad_kernel")
|
||
.context("grad load")?;
|
||
let lq_sgd = lq_module.load_function("alpha_linear_q_sgd_step_kernel")
|
||
.context("sgd load")?;
|
||
let lq_clip = lq_module.load_function("alpha_clip_inplace_kernel")
|
||
.context("clip load")?;
|
||
|
||
// Munchausen target
|
||
let munch_cubin: Vec<u8> = std::fs::read(
|
||
concat!(env!("OUT_DIR"), "/alpha_munchausen_target.cubin")
|
||
).map_err(|e| anyhow::anyhow!("read munch cubin: {e}"))?;
|
||
let munch_module = ctx.load_cubin(munch_cubin).context("munch cubin load")?;
|
||
let munch_kernel = munch_module.load_function("alpha_munchausen_target_kernel")
|
||
.context("munch load")?;
|
||
|
||
// Kill criteria + apply_pearls
|
||
let kc_cubin: Vec<u8> = std::fs::read(
|
||
concat!(env!("OUT_DIR"), "/alpha_kill_criteria.cubin")
|
||
).map_err(|e| anyhow::anyhow!("read kc cubin: {e}"))?;
|
||
let kc_module = ctx.load_cubin(kc_cubin).context("kc cubin load")?;
|
||
let kc_kernel = kc_module.load_function("alpha_kill_criteria_compute_kernel")
|
||
.context("kc kernel load")?;
|
||
|
||
let pearls_cubin: Vec<u8> = std::fs::read(
|
||
concat!(env!("OUT_DIR"), "/apply_pearls_kernel.cubin")
|
||
).map_err(|e| anyhow::anyhow!("read pearls cubin: {e}"))?;
|
||
let pearls_module = ctx.load_cubin(pearls_cubin).context("pearls cubin load")?;
|
||
let pearls_kernel = pearls_module.load_function("apply_pearls_ad_kernel")
|
||
.context("pearls load")?;
|
||
|
||
// Phase E.2 Task 16: stacker-threshold + Kelly-attenuation controller
|
||
let ctl_module = ctx
|
||
.load_cubin(ml::cuda_pipeline::alpha_kernels::STACKER_THRESHOLD_CONTROLLER_CUBIN.to_vec())
|
||
.context("controller cubin load")?;
|
||
let ctl_kernel = ctl_module
|
||
.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")?;
|
||
// Phase E.4.A.6: circular-buffer push kernel for sliding window.
|
||
let push_module = ctx
|
||
.load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_WINDOW_PUSH_CUBIN.to_vec())
|
||
.context("alpha_window_push cubin load")?;
|
||
let push_kernel = push_module
|
||
.load_function("alpha_window_push_kernel")
|
||
.context("window_push kernel load")?;
|
||
// Phase E.4.A.8.fix: GPU-side h_enriched slot copy (replaces
|
||
// dtoh/htod that was violating feedback_cpu_is_read_only).
|
||
let h_store_kernel = push_module
|
||
.load_function("alpha_h_enriched_store_kernel")
|
||
.context("h_enriched_store kernel load")?;
|
||
if cli.temporal {
|
||
info!(
|
||
" TEMPORAL: sliding window K={} × state_dim={} = {} floats",
|
||
cli.window_k, STATE_DIM, cli.window_k * STATE_DIM,
|
||
);
|
||
}
|
||
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());
|
||
|
||
// Load alpha cache first (if any) so we can pass it to the fxcache loader.
|
||
let alpha_cache_vec: Option<Vec<f32>> = if let Some(p) = cli.alpha_cache.as_ref() {
|
||
let cache = ml::env::loaders::load_alpha_cache(p)?;
|
||
info!("Loaded alpha cache: {} entries from {}", cache.len(), p.display());
|
||
Some(cache)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
// Choose snapshot source. Exactly one of --fxcache-path / --mbp10-dir must
|
||
// be set (or both, in which case fxcache wins because alpha_cache needs
|
||
// fxcache-aligned indices).
|
||
let rows: Vec<SnapshotRow> = match (cli.fxcache_path.as_ref(), cli.mbp10_dir.as_ref()) {
|
||
(Some(fxc), _) => {
|
||
info!("Loading snapshots from fxcache: {}", fxc.display());
|
||
ml::env::loaders::load_snapshots_from_fxcache(
|
||
fxc,
|
||
cli.max_snapshots,
|
||
alpha_cache_vec.as_deref(),
|
||
cli.real_spread,
|
||
None, // E.4.A T4: mbp10_dir wiring + --use-real-depth flag lands in T5
|
||
)?
|
||
}
|
||
(None, Some(mbp10)) => {
|
||
if cli.alpha_cache.is_some() {
|
||
anyhow::bail!(
|
||
"--alpha-cache requires --fxcache-path (cache indices align to fxcache bars)"
|
||
);
|
||
}
|
||
info!("Loading snapshots from MBP-10 dir: {}", mbp10.display());
|
||
let parser = DbnParser::new().context("DbnParser::new")?;
|
||
load_snapshots(&parser, mbp10, cli.snapshot_interval, cli.max_snapshots)?
|
||
}
|
||
(None, None) => {
|
||
anyhow::bail!("must set either --fxcache-path or --mbp10-dir");
|
||
}
|
||
};
|
||
info!("Loaded {} snapshots into env", rows.len());
|
||
if rows.len() <= cli.horizon {
|
||
anyhow::bail!("not enough snapshots ({}) for horizon ({})", rows.len(), cli.horizon);
|
||
}
|
||
let n_rows = rows.len();
|
||
|
||
let mut env = ExecutionEnv::new(
|
||
ExecutionEnvConfig {
|
||
horizon_snapshots: cli.horizon,
|
||
trade_size_contracts: cli.trade_size,
|
||
cost_per_contract: cli.cost_per_contract,
|
||
},
|
||
fill_model,
|
||
rows,
|
||
cli.seed,
|
||
);
|
||
|
||
// --- 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.
|
||
// --temporal additionally swaps the C51 input feature dim from
|
||
// STATE_DIM=10 to mamba2_hidden_dim (typically 32) — Q-net consumes
|
||
// Mamba2's `h_enriched` instead of raw state.
|
||
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 xavier_scale = (2.0_f32 / c51_input_dim as f32).sqrt();
|
||
let mut rng = SmokeRng::new(cli.seed.wrapping_add(0xDEAD_BEEF));
|
||
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 q_init_norm = weight_norm(&w_init, &b_init);
|
||
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.
|
||
// 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_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];
|
||
// Updated 2026-05-15 with the new FillModel (L2-regularized fit + MBP-10
|
||
// parser full-levels-copy fix). Negligible drift from the pre-fix run
|
||
// (-5185.13 / 4952.85) — confirms the smoke verdict is robust to these
|
||
// bug fixes. See alpha_random_baseline.json for the source numbers.
|
||
isv_host[RANDOM_BASELINE_MEAN_INDEX] = -5191.53;
|
||
isv_host[RANDOM_BASELINE_STD_INDEX] = 4963.62;
|
||
// Phase E.2 Task 17: trade-rate target anchor for the stacker-threshold
|
||
// controller. Set at training start, never reset across folds
|
||
// (TrainingPersist). Default 0.08 = ~8% per-step trade rate.
|
||
isv_host[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX] =
|
||
cli.trade_rate_target;
|
||
let mut isv_dev = stream.clone_htod(&isv_host).context("upload isv")?;
|
||
|
||
// Wiener state for the 4 kill-criteria slots: 4 × [sample_var, diff_var, x_lag] = 12 floats.
|
||
let mut wiener_dev = stream.alloc_zeros::<f32>(12).context("alloc wiener")?;
|
||
// Scratch buffer for kill-criteria producer output: 4 floats.
|
||
let mut kc_scratch_dev = stream.alloc_zeros::<f32>(4).context("alloc kc_scratch")?;
|
||
// Phase E.2: Wiener state for slot 545 (observed-rate EMA): 3 floats.
|
||
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; // 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(1)? };
|
||
|
||
// Phase E.4.A.7: GPU-resident sliding-window state buffer for the
|
||
// temporal encoder. Allocated unconditionally; only populated when
|
||
// cli.temporal is set. Capacity: window_k × STATE_DIM floats.
|
||
// Phase E.4.A.6+A.8: shift+insert chronological window, allocated
|
||
// as a GpuTensor of shape [1, K, in_dim] so Mamba2Block can consume
|
||
// it directly via `forward_train(&window_tensor)`. The push kernel
|
||
// writes via the underlying CudaSlice (no reshape needed — same
|
||
// 160 floats either way).
|
||
let mut window_tensor = GpuTensor::zeros(
|
||
&[1, cli.window_k, STATE_DIM], &stream
|
||
).map_err(|e| anyhow::anyhow!("alloc window tensor: {e}"))?;
|
||
// Phase E.4.A.8 batched training: store each step's h_enriched
|
||
// (Mamba2 output) so the batched C51 forward at episode end reads
|
||
// pre-computed temporal representations instead of recomputing
|
||
// Mamba2 over each step's window. Size = (horizon + 1) so the
|
||
// terminal next-state h_enriched fits at the end.
|
||
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")?;
|
||
|
||
// Phase E.4.A.8: construct the Mamba2 temporal encoder when
|
||
// --temporal is set. Hidden_dim = c51_input_dim (so the encoder's
|
||
// output flows directly into the C51 head without an additional
|
||
// projection). State_dim ≤ 16 (kernel max).
|
||
let mut mamba2_block: Option<Mamba2Block> = if cli.temporal {
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: STATE_DIM,
|
||
hidden_dim: cli.mamba2_hidden_dim,
|
||
state_dim: cli.mamba2_state_dim,
|
||
seq_len: cli.window_k,
|
||
};
|
||
info!(
|
||
" TEMPORAL Mamba2: in={} hidden={} state={} K={}",
|
||
cfg.in_dim, cfg.hidden_dim, cfg.state_dim, cfg.seq_len
|
||
);
|
||
Some(Mamba2Block::new(cfg, stream.clone())
|
||
.map_err(|e| anyhow::anyhow!("Mamba2Block init: {e}"))?)
|
||
} else {
|
||
None
|
||
};
|
||
// Phase E.4.A.T10: AdamW for Mamba2 weights (active when --temporal).
|
||
let mut mamba2_adamw: Option<Mamba2AdamW> = if let Some(block) = mamba2_block.as_ref() {
|
||
Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default())
|
||
.map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?)
|
||
} else {
|
||
None
|
||
};
|
||
// Snapshot Mamba2 weights at init so EARLY_Q_MOVEMENT can include
|
||
// encoder movement (not just C51 head). Without this, the diagnostic
|
||
// under-reports system learning when --temporal: the encoder absorbs
|
||
// the gradient signal, the head moves little, and the KC reads ~0.
|
||
let mamba2_init_snap: Option<Vec<f32>> = if let Some(block) = mamba2_block.as_ref() {
|
||
Some(mamba2_snapshot(&stream, block)?)
|
||
} else {
|
||
None
|
||
};
|
||
// d_h_enriched as GpuTensor so it can feed `backward_from_h_enriched`
|
||
// after the C51 grad-input kernel writes per-row gradients into it.
|
||
let mut d_h_enriched_tensor_smoke = GpuTensor::zeros(
|
||
&[cli.horizon, cli.mamba2_hidden_dim], &stream
|
||
).map_err(|e| anyhow::anyhow!("alloc d_h_enriched smoke: {e}"))?;
|
||
// Per-step train-time windows, written by the new
|
||
// `alpha_train_window_store_batched_kernel` during inference and
|
||
// re-forwarded as one big batched Mamba2 call at end-of-episode.
|
||
let mut train_windows_tensor_smoke = GpuTensor::zeros(
|
||
&[cli.horizon, cli.window_k, STATE_DIM], &stream
|
||
).map_err(|e| anyhow::anyhow!("alloc train windows smoke: {e}"))?;
|
||
// Load the C51 grad-input kernel and the train-window store kernel.
|
||
let c51_grad_input_kernel = c51_module
|
||
.load_function("alpha_c51_grad_input_kernel")
|
||
.context("c51 grad_input load")?;
|
||
let train_window_store_kernel = push_module
|
||
.load_function("alpha_train_window_store_batched_kernel")
|
||
.context("train_window_store kernel load")?;
|
||
|
||
// 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")?;
|
||
let mut kc_scalar_dev = stream.alloc_zeros::<f32>(3).context("alloc kc scalar")?;
|
||
|
||
// --- Training loop ---
|
||
let mut episode_rng = SmokeRng::new(cli.seed.wrapping_add(0xFEED));
|
||
let mut action_counts_host = vec![0_i32; N_ACTIONS];
|
||
// Phase E.3: cached threshold from controller (slot 543). Updated after
|
||
// each rollout-end controller invocation; used in action selection for
|
||
// the NEXT episode. Initialised to 0 (no gate) so the first episode's
|
||
// exploration is unrestricted, letting the controller accumulate enough
|
||
// observations to drive a meaningful threshold.
|
||
let mut current_threshold: f32 = 0.0;
|
||
let mut recent_returns: Vec<f32> = Vec::new();
|
||
let max_start = n_rows.saturating_sub(cli.horizon + 1).max(1);
|
||
|
||
let mut kc_logs: Vec<(usize, [f32; 4])> = Vec::new();
|
||
|
||
for ep in 0..cli.n_episodes {
|
||
let eps = cli.eps_start
|
||
+ (cli.eps_end - cli.eps_start) * (ep as f32 / cli.n_episodes.max(1) as f32);
|
||
let start_cursor = (episode_rng.next_u64() as usize) % max_start;
|
||
let env_seed = episode_rng.next_u64();
|
||
env.reset_at(env_seed, start_cursor);
|
||
let mut state = EpisodeState::new();
|
||
|
||
// Phase E.4.A.7: clear the temporal window on episode start so
|
||
// Mamba2 sees zero-history for the first window_k-1 steps.
|
||
// Phase E.4.A.T10: also zero the train-time window buffer and the
|
||
// d_h_enriched buffer so terminal slots (past ep_len) contribute
|
||
// zero gradient through the Mamba2 backward.
|
||
if cli.temporal {
|
||
stream
|
||
.memset_zeros(window_tensor.data_mut())
|
||
.context("zero window on episode reset")?;
|
||
stream
|
||
.memset_zeros(train_windows_tensor_smoke.data_mut())
|
||
.context("zero train windows on episode reset")?;
|
||
stream
|
||
.memset_zeros(d_h_enriched_tensor_smoke.data_mut())
|
||
.context("zero d_h_enriched on episode reset")?;
|
||
}
|
||
|
||
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);
|
||
|
||
loop {
|
||
// 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);
|
||
let action: u8 = if cli.c51 {
|
||
state_pinned.write(&s_vec);
|
||
// Phase E.4.A.7: shift-and-insert push of current state.
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
// T10: capture post-push window for end-of-episode re-forward.
|
||
{
|
||
let (src_ptr, _g_src) = window_tensor.data().device_ptr(&stream);
|
||
let (dst_ptr, _g_dst) = train_windows_tensor_smoke.data_mut().device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched(
|
||
&stream, &train_window_store_kernel,
|
||
src_ptr, dst_ptr, state.step as i32,
|
||
1, cli.window_k as i32, state_dim_i,
|
||
)?;
|
||
}
|
||
}
|
||
}
|
||
// Phase E.4.A.8: Mamba2 forward over the window → h_enriched
|
||
// (`cache.h_enriched [1, hidden_dim]`). The C51 head reads
|
||
// this device buffer in place of `state_pinned`.
|
||
//
|
||
// The cache is dropped at the end of this scope —
|
||
// inference doesn't need backward (we only train via
|
||
// the batched update at end-of-episode). Backward
|
||
// wiring with cache retention lands in T10.
|
||
// Phase E.4.A.8 + .8.fix: per-step Mamba2 forward; store
|
||
// h_enriched into per-step buffer slot via GPU kernel
|
||
// (no CPU roundtrip); feed h_enriched to C51 forward.
|
||
if cli.temporal {
|
||
let block = mamba2_block.as_ref()
|
||
.expect("temporal flag set but Mamba2Block missing");
|
||
let (_logit, cache) = block.forward_train(&window_tensor)
|
||
.map_err(|e| anyhow::anyhow!("mamba2 forward: {e}"))?;
|
||
// GPU-side slot copy: cache.h_enriched → h_enriched_buf_dev[slot..]
|
||
{
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
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(),
|
||
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"))?;
|
||
|
||
// Buffer transition for batched update
|
||
states_host.extend_from_slice(&s_vec);
|
||
// next-state observation: get the env state AT THE NEW CURSOR.
|
||
// env.state(&state) gives current state; after env.step, cursor
|
||
// and state are updated to s'.
|
||
let s_next_vec = state_as_vec(&env, &state);
|
||
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 });
|
||
action_counts_host[action as usize] += 1;
|
||
|
||
if done {
|
||
recent_returns.push(reward);
|
||
break;
|
||
}
|
||
}
|
||
|
||
let ep_len = actions_host.len() as i32;
|
||
if ep_len < 2 {
|
||
continue;
|
||
}
|
||
|
||
// --- Batched update on this episode's transitions ---
|
||
// Upload
|
||
stream.memcpy_htod(&states_host, &mut states_dev)
|
||
.context("htod states")?;
|
||
stream.memcpy_htod(&next_states_host, &mut next_states_dev)
|
||
.context("htod next_states")?;
|
||
stream.memcpy_htod(&actions_host, &mut actions_dev)
|
||
.context("htod actions")?;
|
||
// Reward normalization: divide by cli.reward_scale so the
|
||
// Munchausen target produces normalized targets. Original
|
||
// rewards stay in rewards_host for reporting / kill-criteria.
|
||
let rewards_normalized: Vec<f32> = rewards_host
|
||
.iter()
|
||
.map(|r| r / cli.reward_scale)
|
||
.collect();
|
||
stream.memcpy_htod(&rewards_normalized, &mut rewards_dev)
|
||
.context("htod rewards")?;
|
||
stream.memcpy_htod(&dones_host, &mut dones_dev)
|
||
.context("htod dones")?;
|
||
|
||
if cli.c51 {
|
||
// ── C51 batched compute ─────────────────────────────────
|
||
// Phase E.4.A.8 + .8.fix: in temporal mode, read pre-stored
|
||
// h_enriched from h_enriched_buf_dev (populated GPU-side
|
||
// during per-step inference via alpha_h_enriched_store).
|
||
// Run ONE extra Mamba2 forward on the terminal window to
|
||
// fill slot ep_len (the "next-state" representation for
|
||
// the last training transition).
|
||
//
|
||
// Phase E.4.A.T10: also run a second Mamba2 forward on the
|
||
// collected train_windows_tensor_smoke to get a cache that
|
||
// covers all ep_len visited windows. We use the cache's
|
||
// h_enriched as the curr_input for C51 (so the C51 grad can
|
||
// flow back into Mamba2 via backward_from_h_enriched), while
|
||
// keeping h_enriched_buf_dev offset for next_input (frozen
|
||
// target side, no backward).
|
||
let cache_train: Option<ml_alpha::mamba2_block::Mamba2ForwardCache> = if cli.temporal {
|
||
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
|
||
// Terminal forward fills h_enriched_buf_dev[ep_len].
|
||
{
|
||
let (_logit, cache) = block.forward_train(&window_tensor)
|
||
.map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?;
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
// Training-time forward on the full [horizon, K, in_dim]
|
||
// batch (slots ep_len..horizon are zero, contribute zero
|
||
// gradient since d_h_enriched stays zero there).
|
||
let (_logit, cache) = block.forward_train(&train_windows_tensor_smoke)
|
||
.map_err(|e| anyhow::anyhow!("mamba2 train forward: {e}"))?;
|
||
Some(cache)
|
||
} else {
|
||
None
|
||
};
|
||
let (curr_input_ptr_guard, next_input_ptr_guard);
|
||
let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal {
|
||
let cache = cache_train.as_ref().expect("cache_train missing");
|
||
let (h_ptr_curr, g_curr) = cache.h_enriched.cuda_data().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)
|
||
};
|
||
// 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,
|
||
)?;
|
||
}
|
||
}
|
||
// 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 (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,
|
||
)?;
|
||
}
|
||
}
|
||
// 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). With --temporal
|
||
// the input to the C51 layer is h_enriched (hidden_dim);
|
||
// gradient w.r.t. W computes dW = (p - m) ⊗ input^T, so
|
||
// input pointer must match the forward's input.
|
||
{
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
// T10 backward chain: C51 grad_input → Mamba2 backward → AdamW.
|
||
// The C51 grad_input kernel writes dL/d_h_enriched into the
|
||
// first ep_len rows of d_h_enriched_tensor_smoke; remaining
|
||
// rows stay zero (zeroed on episode reset).
|
||
if cli.temporal {
|
||
{
|
||
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 (w_ptr, _g3) = w_dev.device_ptr(&stream);
|
||
let (dh_ptr, _g4) = d_h_enriched_tensor_smoke.data_mut().device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input(
|
||
&stream, &c51_grad_input_kernel,
|
||
p_ptr, m_ptr, a_ptr, w_ptr,
|
||
dh_ptr,
|
||
ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
|
||
1.0 / ep_len as f32,
|
||
)?;
|
||
}
|
||
}
|
||
let cache = cache_train.as_ref().expect("cache_train missing");
|
||
let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing");
|
||
let grads = block_ref
|
||
.backward_from_h_enriched(cache, &d_h_enriched_tensor_smoke)
|
||
.map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?;
|
||
let _ = cache_train;
|
||
let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing");
|
||
let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing");
|
||
adamw
|
||
.step(block_mut, &grads)
|
||
.map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?;
|
||
}
|
||
} 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,
|
||
)?;
|
||
}
|
||
}
|
||
}
|
||
// ── 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_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,
|
||
)?;
|
||
}
|
||
}
|
||
// SGD step on W
|
||
{
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
// SGD step on b
|
||
{
|
||
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,
|
||
)?;
|
||
}
|
||
}
|
||
|
||
// Phase E.2 Task 17: Stacker-threshold controller at rollout end.
|
||
// Trade rate = (count of non-Wait actions) / total decisions in the
|
||
// episode. Realized Sharpe ≈ mean(returns) / std(returns) over the
|
||
// window — for a single rollout this collapses to the rollout's
|
||
// terminal reward divided by the random baseline std (a per-rollout
|
||
// analog of the rvr metric, lets the controller decide whether to
|
||
// tighten or loosen Kelly).
|
||
let trade_count_ep: f32 = actions_host
|
||
.iter()
|
||
.filter(|&&a| a != 0) // 0 == Wait
|
||
.count() as f32;
|
||
let decisions_ep = actions_host.len() as f32;
|
||
// Single-episode Sharpe proxy: terminal reward / baseline std.
|
||
// Note: terminal reward is the LAST element of rewards_host
|
||
// (sum of step rewards from env.step into ep.cumulative_pnl).
|
||
let ep_terminal_r = *rewards_host.last().unwrap_or(&0.0);
|
||
let baseline_std = isv_host[RANDOM_BASELINE_STD_INDEX].max(1e-6);
|
||
let rollout_sharpe = ep_terminal_r / baseline_std;
|
||
unsafe {
|
||
let (isv_ptr, _g0) = isv_dev.device_ptr_mut(&stream);
|
||
let (w_ptr, _g1) = 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,
|
||
cli.target_sharpe,
|
||
cli.k_threshold,
|
||
cli.k_atten,
|
||
cli.wiener_alpha_floor,
|
||
cli.ctl_alpha_meta,
|
||
ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX as i32,
|
||
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX as i32,
|
||
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX as i32,
|
||
ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX as i32,
|
||
// T16 regime hookup. Smoke doesn't currently run the vol
|
||
// update kernel so we pass -1/-1 to disable the regime
|
||
// scale here (backward compat).
|
||
-1,
|
||
-1,
|
||
0.25,
|
||
// T16 floor controller — also disabled in smoke.
|
||
-1,
|
||
-1,
|
||
0.5,
|
||
0.1,
|
||
1.0e10,
|
||
isv_ptr,
|
||
w_ptr,
|
||
)?;
|
||
}
|
||
// Phase E.3: refresh CPU-side threshold cache from device after the
|
||
// controller update, for the NEXT episode's gating decisions.
|
||
stream.synchronize().context("sync after controller")?;
|
||
let isv_now = stream.clone_dtoh(&isv_dev).context("dtoh isv post-controller")?;
|
||
current_threshold =
|
||
isv_now[ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX];
|
||
|
||
// Target network hard update every K episodes: w_target ← w_online.
|
||
// Standard DQN stabilizer — V_soft(s') bootstrap reads w_target,
|
||
// so updating it slowly breaks the chase-its-own-tail divergence
|
||
// of online-only Munchausen.
|
||
if (ep + 1) % cli.target_update_every == 0 {
|
||
stream.synchronize().context("sync before target update")?;
|
||
let w_now = stream.clone_dtoh(&w_dev).context("dtoh W for target update")?;
|
||
let b_now = stream.clone_dtoh(&b_dev).context("dtoh b for target update")?;
|
||
stream.memcpy_htod(&w_now, &mut w_target_dev)
|
||
.context("htod W_target hard update")?;
|
||
stream.memcpy_htod(&b_now, &mut b_target_dev)
|
||
.context("htod b_target hard update")?;
|
||
}
|
||
|
||
// Periodic kill-criteria pipeline launch
|
||
if (ep + 1) % cli.kill_criteria_every == 0 {
|
||
stream.synchronize().context("sync before kc")?;
|
||
// Snapshot weights for q_early_norm
|
||
let w_now = stream.clone_dtoh(&w_dev).context("dtoh W")?;
|
||
let b_now = stream.clone_dtoh(&b_dev).context("dtoh b")?;
|
||
// q_early = q_init + ||W_now − W_init||_F so the kernel's
|
||
// |q_early − q_init| / |q_init| ratio yields the relative
|
||
// weight-space distance from init (direction-sensitive).
|
||
//
|
||
// When --temporal is on, include Mamba2 weight movement too:
|
||
// the encoder absorbs most of the gradient signal so the C51
|
||
// head moves little on its own. Without this addition, the
|
||
// diagnostic systematically under-reports system learning.
|
||
let mut dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init);
|
||
if let (Some(block), Some(init_snap)) =
|
||
(mamba2_block.as_ref(), mamba2_init_snap.as_ref())
|
||
{
|
||
let now_snap = mamba2_snapshot(&stream, block)?;
|
||
dist += mamba2_weight_distance(&now_snap, init_snap);
|
||
}
|
||
let q_early_norm = q_init_norm + dist;
|
||
let rollout_r_mean: f32 = if recent_returns.is_empty() {
|
||
0.0
|
||
} else {
|
||
let n = recent_returns.len();
|
||
let take = n.min(cli.kill_criteria_every);
|
||
let slice = &recent_returns[n - take..];
|
||
slice.iter().sum::<f32>() / take as f32
|
||
};
|
||
let scalar_inputs = vec![rollout_r_mean, q_init_norm, q_early_norm];
|
||
stream.memcpy_htod(&scalar_inputs, &mut kc_scalar_dev)
|
||
.context("htod kc scalars")?;
|
||
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);
|
||
let (ac_ptr, _g1) = kc_action_counts_dev.device_ptr(&stream);
|
||
let (sc_ptr, _g2) = kc_scalar_dev.device_ptr(&stream);
|
||
let (isv_ptr, _g3) = isv_dev.device_ptr(&stream);
|
||
let (out_ptr, _g4) = kc_scratch_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::alpha_kernels::launch_alpha_kill_criteria(
|
||
&stream, &kc_kernel,
|
||
q_ptr, ac_ptr, sc_ptr, isv_ptr,
|
||
ep_len, n_act_i, out_ptr,
|
||
)?;
|
||
}
|
||
}
|
||
// Chained Pearls A+D applicator → ISV[539..543]
|
||
{
|
||
let (sc_ptr, _g0) = kc_scratch_dev.device_ptr(&stream);
|
||
let (isv_ptr, _g1) = isv_dev.device_ptr_mut(&stream);
|
||
let (w_ptr, _g2) = wiener_dev.device_ptr_mut(&stream);
|
||
unsafe {
|
||
ml::cuda_pipeline::sp4_wiener_ema::launch_apply_pearls(
|
||
&stream, &pearls_kernel,
|
||
sc_ptr, 0,
|
||
isv_ptr, ALPHA_ISV_BLOCK_LO as i32,
|
||
w_ptr, 0,
|
||
4,
|
||
ml::cuda_pipeline::sp4_wiener_ema::ALPHA_META,
|
||
)?;
|
||
}
|
||
}
|
||
stream.synchronize().context("sync after kc")?;
|
||
|
||
let isv = stream.clone_dtoh(&isv_dev).context("dtoh isv")?;
|
||
let kc = [
|
||
isv[Q_SPREAD_EMA_INDEX],
|
||
isv[ACTION_ENTROPY_EMA_INDEX],
|
||
isv[RETURN_VS_RANDOM_EMA_INDEX],
|
||
isv[EARLY_Q_MOVEMENT_EMA_INDEX],
|
||
];
|
||
// Phase E.2: controller-driven slots for visibility.
|
||
let ctl_threshold = isv[ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX];
|
||
let ctl_obs_rate = isv[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX];
|
||
let ctl_atten = isv[ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX];
|
||
info!(
|
||
"ep {:>4} | ε={:.3} | R_mean={:>+9.1} | KC q={:.3} H={:.3} rvr={:+.3} ΔQ={:.3} | CTL thresh={:.4} obs={:.3} atten={:.3}",
|
||
ep + 1, eps, rollout_r_mean, kc[0], kc[1], kc[2], kc[3],
|
||
ctl_threshold, ctl_obs_rate, ctl_atten,
|
||
);
|
||
kc_logs.push((ep + 1, kc));
|
||
}
|
||
}
|
||
|
||
stream.synchronize().context("final sync")?;
|
||
let isv_final = stream.clone_dtoh(&isv_dev).context("dtoh isv final")?;
|
||
let kc_final = [
|
||
isv_final[Q_SPREAD_EMA_INDEX],
|
||
isv_final[ACTION_ENTROPY_EMA_INDEX],
|
||
isv_final[RETURN_VS_RANDOM_EMA_INDEX],
|
||
isv_final[EARLY_Q_MOVEMENT_EMA_INDEX],
|
||
];
|
||
|
||
// --- Verdict ---
|
||
let entropy_threshold = 0.5 * (N_ACTIONS as f32).ln();
|
||
let pass_q_spread = kc_final[0] >= 0.05;
|
||
let pass_entropy = kc_final[1] >= entropy_threshold;
|
||
let pass_rvr = kc_final[2] >= 0.0;
|
||
let pass_early = kc_final[3] >= 0.01;
|
||
let all_pass = pass_q_spread && pass_entropy && pass_rvr && pass_early;
|
||
|
||
info!("=== Final kill-criteria verdict ===");
|
||
info!(
|
||
" Q_SPREAD_EMA = {:.4} threshold ≥ 0.05 [{}]",
|
||
kc_final[0], if pass_q_spread { "PASS" } else { "FAIL" }
|
||
);
|
||
info!(
|
||
" ACTION_ENTROPY_EMA = {:.4} threshold ≥ {:.4} [{}]",
|
||
kc_final[1], entropy_threshold,
|
||
if pass_entropy { "PASS" } else { "FAIL" }
|
||
);
|
||
info!(
|
||
" RETURN_VS_RANDOM_EMA = {:+.4} threshold ≥ 0.0 [{}]",
|
||
kc_final[2], if pass_rvr { "PASS" } else { "FAIL" }
|
||
);
|
||
info!(
|
||
" EARLY_Q_MOVEMENT_EMA = {:.4} threshold ≥ 0.01 [{}]",
|
||
kc_final[3], if pass_early { "PASS" } else { "FAIL" }
|
||
);
|
||
info!(
|
||
" Overall: {} (H=6000 scale-up {})",
|
||
if all_pass { "PASS" } else { "FAIL" },
|
||
if all_pass { "VIABLE" } else { "BLOCKED → pivot to NoisyNet (Task 19)" }
|
||
);
|
||
|
||
// --- 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,
|
||
"eps_start": cli.eps_start,
|
||
"eps_end": cli.eps_end,
|
||
"gamma": cli.gamma,
|
||
"alpha_m": cli.alpha_m,
|
||
"tau": cli.tau,
|
||
"reward_scale": cli.reward_scale,
|
||
"target_update_every": cli.target_update_every,
|
||
"grad_clip": cli.grad_clip,
|
||
"q_init_norm": q_init_norm,
|
||
"final_stacker_threshold": isv_final[ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX],
|
||
"final_trade_rate_observed_ema": isv_final[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX],
|
||
"final_stacker_kelly_attenuation": isv_final[ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX],
|
||
"trade_rate_target": cli.trade_rate_target,
|
||
"q_spread_ema": kc_final[0],
|
||
"action_entropy_ema": kc_final[1],
|
||
"return_vs_random_ema": kc_final[2],
|
||
"early_q_movement_ema": kc_final[3],
|
||
"pass_q_spread": pass_q_spread,
|
||
"pass_entropy": pass_entropy,
|
||
"pass_rvr": pass_rvr,
|
||
"pass_early": pass_early,
|
||
"all_pass": all_pass,
|
||
"kc_log": kc_logs.iter().map(|(ep, kc)| {
|
||
serde_json::json!({
|
||
"episode": ep,
|
||
"q_spread": kc[0],
|
||
"entropy": kc[1],
|
||
"rvr": kc[2],
|
||
"early_mvmt": kc[3],
|
||
})
|
||
}).collect::<Vec<_>>(),
|
||
});
|
||
let mut f = File::create(&cli.out_path).context("create out")?;
|
||
write!(f, "{}", serde_json::to_string_pretty(&json)?)?;
|
||
info!("Wrote verdict + KC trajectory to {}", cli.out_path.display());
|
||
|
||
Ok(())
|
||
}
|