20c7852b66
The xv66n smoke (commit d5c29fb4f) confirmed plateau-decay LR works
across all 3 heads — but exposed a residual V instability: 10 trade-
close steps with l_v > 1e4, max 9.46e4. Root cause: the
reward_scale controller's Wiener-α blend cannot adapt fast enough
to a sudden fat-tail trade outcome, so a single closed trade with
realised PnL well outside `1 / mean_abs_pnl_ema`'s current estimate
produces a scaled reward 100s of times the C51 atom span.
Since `returns = scaled_reward + γ(1-done) v_tp1` and the spike
happens on done=1 steps, returns equals the unbounded scaled
reward, and V regression `(v_pred - returns)²` blows up.
## Fix: asymmetric clamp at apply_reward_scale boundary
`apply_reward_scale.cu` is rewritten to:
1. Scale `rewards[b] *= isv[RL_REWARD_SCALE_INDEX]` as before.
2. Asymmetric-clamp scaled to `[-REWARD_CLAMP_LOSS, +REWARD_CLAMP_WIN]`
= `[-3.0, +1.0]` per `pearl_audit_unboundedness_for_implicit_asymmetry`:
* `WIN = +1.0` matches the C51 atom span on the win side.
* `LOSS = -3.0` preserves loss-aversion asymmetry — fat-tail
losses remain visible up to 3 atom-units before flattening,
matching typical HFT P&L distributions where losses run
2-3× larger than wins per close.
3. Write back the clamped value to `rewards[b]`.
Single-block layout (block_x = min(b_size, 256), grid_x = 1,
shared = block_x × 4 B) per `pearl_no_atomicadd` — tree reduction
inside the block, no inter-block atomic.
## Diagnostic: pre-clamp max ISV slot
New ISV slot `RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX = 439`
holds `max(|scaled|)` over the current batch BEFORE the clamp
fires (each step overwrites — point measurement, not EMA).
Surfaced in diag.jsonl as `rewards.scaled_pre_clamp_max`.
Interpretation:
* pre_clamp_max ≤ 1.0 most steps → reward_scale controller is
tracking typical magnitudes correctly; clamp is a no-op.
* pre_clamp_max > 1.0 frequently → controller is failing to
track magnitudes; clamp is doing load-bearing work shaping V
target.
* pre_clamp_max > 100 ever → controller is grossly mis-scaled
(likely cold-start before mean_abs_pnl_ema converged).
RL_SLOTS_END: 439 → 440 (one new diagnostic slot).
## Why a clamp instead of fixing the controller
The reward_scale controller IS doing its job — it Wiener-blends
toward `1 / mean_abs_pnl_ema` with α floor 0.4. The problem is
that a single closed trade represents one observation in the EMA
denominator, so a sudden 10× excursion in trade magnitude takes
~3-5 closes to fully reflect in the scale. During those 3-5
steps, scaled rewards can be 5-10× the atom span.
A faster controller (smaller EMA floor, lookahead, etc.) would
oscillate. A clamp is the principled bound:
* source signal (raw PnL) remains unbounded — controller
continues to track magnitudes
* downstream signal (V/Q target) is bounded — no catastrophic
backward gradient
* pre-clamp diagnostic surfaces clamp activity so we know when
the controller is failing vs handling the regime fine
## Verified gates (local sm_86)
G1 isv_bootstrap ✅
G3 controllers ✅
G4 target_update ✅
G6 r7d_per_wiring ✅
integrated_smoke ✅
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
740 lines
33 KiB
Rust
740 lines
33 KiB
Rust
//! Integrated RL trainer CLI (DQN + PPO on shared Mamba2 -> CfC encoder).
|
||
//!
|
||
//! Drives `IntegratedTrainer::step_with_lobsim` against a real `LobSimCuda`
|
||
//! backed by MBP-10 windows loaded via `MultiHorizonLoader::next_sequence_pair`
|
||
//! (R2) for true `(s_t, s_{t+1})` adjacency. The trainer combines:
|
||
//! * R3: GPU-resident EMA + advantage/return kernels.
|
||
//! * R4: GPU Thompson + Double-DQN argmax + log π_old kernels.
|
||
//! * R5: 7 RL controllers + target-net soft update.
|
||
//! * R6: GPU-pure env step (extract_realized_pnl_delta, apply_reward_scale,
|
||
//! actions_to_market_targets).
|
||
//! * R7a/R7b: host-lift of every per-step orchestration loop.
|
||
//! * R7c-data: true `h_{t+1}` / `V(s_{t+1})` / Bellman target on `h_{t+1}`.
|
||
//! * R7d: PER push/sample + off-policy DQN with stop-grad on encoder.
|
||
//!
|
||
//! Emits `alpha_rl_train_summary.json` on exit. Gate G8 (NaN abort): if any
|
||
//! per-head loss is non-finite at any step, the binary exits with code 2 so
|
||
//! the cluster smoke harness can terminate the workflow on first divergence.
|
||
//!
|
||
//! Usage:
|
||
//! alpha_rl_train \
|
||
//! --mbp10-data-dir /data/futures-baseline-mbp10/ES.FUT \
|
||
//! --predecoded-dir /feature-cache/predecoded \
|
||
//! --out /feature-cache/alpha-rl-runs/<sha> \
|
||
//! --n-steps 1000 \
|
||
//! --seed 16962 \
|
||
//! --instrument-mode front-month \
|
||
//! --gpu-idx 0
|
||
//!
|
||
//! Plan-doc reference: docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md
|
||
//! Section R8 (CLI binary + Argo template + dispatcher).
|
||
|
||
use anyhow::{Context, Result};
|
||
use clap::Parser;
|
||
use data::providers::databento::dbn_parser::InstrumentFilter;
|
||
use ml_alpha::data::loader::{
|
||
discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig,
|
||
DEFAULT_OUTCOME_LABEL_COST_ES,
|
||
};
|
||
use ml_alpha::heads::HORIZONS;
|
||
use ml_alpha::rl::isv_slots::{
|
||
RL_ADVANTAGE_VAR_RATIO_EMA_INDEX, RL_ENTROPY_COEF_INDEX, RL_ENTROPY_OBSERVED_EMA_INDEX,
|
||
RL_GAMMA_INDEX, RL_KL_PI_EMA_INDEX, RL_LR_AUX_INDEX, RL_LR_BCE_INDEX, RL_LR_PI_BEST_LOSS_INDEX,
|
||
RL_LR_PI_INDEX, RL_LR_PI_LOSS_EMA_INDEX, RL_LR_PI_STEPS_SINCE_BEST_INDEX,
|
||
RL_LR_PI_WARMUP_COUNTER_INDEX, RL_LR_Q_BEST_LOSS_INDEX, RL_LR_Q_INDEX,
|
||
RL_LR_Q_LOSS_EMA_INDEX, RL_LR_Q_STEPS_SINCE_BEST_INDEX, RL_LR_Q_WARMUP_COUNTER_INDEX,
|
||
RL_LR_V_BEST_LOSS_INDEX, RL_LR_V_INDEX, RL_LR_V_LOSS_EMA_INDEX,
|
||
RL_LR_V_STEPS_SINCE_BEST_INDEX, RL_LR_V_WARMUP_COUNTER_INDEX,
|
||
RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX, RL_MEAN_ABS_PNL_EMA_INDEX,
|
||
RL_MEAN_TRADE_DURATION_EMA_INDEX, RL_N_ROLLOUT_STEPS_INDEX, RL_PER_ALPHA_INDEX,
|
||
RL_PI_GRAD_NORM_EMA_INDEX, RL_PPO_CLIP_INDEX, RL_Q_DIVERGENCE_EMA_INDEX,
|
||
RL_Q_GRAD_NORM_EMA_INDEX, RL_REWARD_SCALE_INDEX, RL_TARGET_TAU_INDEX, RL_TD_KURTOSIS_EMA_INDEX,
|
||
RL_V_GRAD_NORM_EMA_INDEX,
|
||
};
|
||
use ml_alpha::trainer::integrated::{
|
||
read_slice_d_pub, read_slice_i32_d_pub, IntegratedStepStats, IntegratedTrainer,
|
||
IntegratedTrainerConfig,
|
||
};
|
||
use ml_alpha::trainer::perception::PerceptionTrainerConfig;
|
||
use ml_backtesting::sim::LobSimCuda;
|
||
use ml_core::device::MlDevice;
|
||
use serde::Serialize;
|
||
use serde_json::json;
|
||
use std::io::{BufWriter, Write};
|
||
use std::path::PathBuf;
|
||
use std::process;
|
||
|
||
#[derive(Parser)]
|
||
#[command(name = "alpha_rl_train")]
|
||
struct Cli {
|
||
/// Directory containing MBP-10 .dbn.zst files. Per
|
||
/// `feedback_mbp10_mandatory`, this argument is required — no
|
||
/// synthetic-data fallback in the production training path.
|
||
#[arg(long)]
|
||
mbp10_data_dir: PathBuf,
|
||
|
||
/// Directory holding predecoded sidecar caches (created on miss).
|
||
#[arg(long)]
|
||
predecoded_dir: PathBuf,
|
||
|
||
/// Output directory for summary JSON + per-step log artefacts.
|
||
#[arg(long)]
|
||
out: PathBuf,
|
||
|
||
/// Total `step_with_lobsim` calls to run. The validation smoke
|
||
/// (R9 prerequisite) ships with 1000; production scale-up runs
|
||
/// 50k+. Each call processes one `(s_t, s_{t+1})` sequence pair
|
||
/// across `n_backtests` parallel LobSim instances.
|
||
#[arg(long, default_value_t = 1000)]
|
||
n_steps: usize,
|
||
|
||
/// Encoder sequence length (snapshots per step). Must match
|
||
/// `PerceptionTrainerConfig::seq_len`. Default 32 matches the
|
||
/// validated single-scale `1:32` window from alpha_train.
|
||
#[arg(long, default_value_t = 32)]
|
||
seq_len: usize,
|
||
|
||
/// Batch dimension — number of parallel LobSim backtests run in
|
||
/// lockstep per step. Larger batches amortise the encoder forward
|
||
/// cost across more trajectories but push GPU memory pressure
|
||
/// linearly. Smoke at 1; production sweep at 32-64 (L40S 48GB).
|
||
#[arg(long, default_value_t = 1)]
|
||
n_backtests: usize,
|
||
|
||
/// PER buffer capacity (R7d). Per `replay.rs` doc, naive O(N)
|
||
/// sampling is acceptable up to ~4096. Production scale-up to
|
||
/// 100k+ needs a GPU sum-tree (R-future).
|
||
#[arg(long, default_value_t = 4096)]
|
||
per_capacity: usize,
|
||
|
||
/// Random seed pinning RNG + Xavier init + ISV controller bootstrap
|
||
/// state. Per `pearl_scoped_init_seed_for_reproducibility`.
|
||
#[arg(long, default_value_t = 16962)]
|
||
seed: u64,
|
||
|
||
/// MBP-10 instrument filter — `all`, `front-month`, or `id=<N>`.
|
||
/// `front-month` is the canonical setting for ES (auto-detects the
|
||
/// dominant contract per file, handles quarterly roll).
|
||
#[arg(long, default_value = "front-month", value_parser = parse_instrument_mode)]
|
||
instrument_mode: InstrumentFilter,
|
||
|
||
/// CUDA device index. Single-GPU dev (RTX 3050 Ti) uses 0; cluster
|
||
/// L40S/H100 also 0 (one GPU per pod per the alpha-rl template).
|
||
#[arg(long, default_value_t = 0)]
|
||
gpu_idx: usize,
|
||
|
||
/// Round-trip trade cost in price units for the D-style outcome
|
||
/// labels (consumed by the perception aux head, not by the RL
|
||
/// path directly — but the loader needs it). Default 0.5 = 2 ticks
|
||
/// × $0.25/tick for ES.
|
||
#[arg(long, default_value_t = DEFAULT_OUTCOME_LABEL_COST_ES)]
|
||
outcome_label_cost: f32,
|
||
|
||
/// How often (in steps) to flush a progress line to stderr. The
|
||
/// summary JSON is written ONCE at the end; this controls in-flight
|
||
/// visibility only.
|
||
#[arg(long, default_value_t = 100)]
|
||
log_every: usize,
|
||
|
||
/// Walk-forward fold index for the multi-fold G8 gate (per
|
||
/// `pearl_single_window_oos_is_not_oos` — a single window is NOT
|
||
/// out-of-sample). Slices the MBP-10 file list into K equal-sized
|
||
/// blocks; fold k trains on blocks [0..k+1) and evaluates on block
|
||
/// [k+1] (purged walk-forward). At `fold_idx=0` the first block
|
||
/// is the train window, second is eval. At `fold_idx=K-2` the
|
||
/// trainer sees almost all the data and evaluates on the final
|
||
/// block. The dispatcher (`scripts/argo-alpha-rl.sh --n-folds K`)
|
||
/// fans out K-1 workflows, one per fold; the aggregator computes
|
||
/// mean ± SD of the per-fold profit_factor.
|
||
///
|
||
/// Default 0 + n-folds=1 = single-window mode (no eval split).
|
||
#[arg(long, default_value_t = 0)]
|
||
fold_idx: usize,
|
||
|
||
/// Total walk-forward fold count. Set with `--fold-idx` for the
|
||
/// multi-fold G8 gate. n_folds=1 disables the train/eval split
|
||
/// (single-window mode — only meaningful for smoke validation).
|
||
/// n_folds≥3 enables a real evaluation: train on first
|
||
/// `fold_idx + 1` of `n_folds` data blocks, evaluate on the next.
|
||
#[arg(long, default_value_t = 1)]
|
||
n_folds: usize,
|
||
|
||
/// Eval-phase step count. After the train phase completes its
|
||
/// `--n-steps`, the eval phase runs `--n-eval-steps` additional
|
||
/// steps on the held-out fold using the same trainer
|
||
/// (acknowledged minor in-eval learning contamination at b_size=1
|
||
/// — pure eval mode is a follow-up). LobSim trade records are
|
||
/// drained at the END so the summary reflects eval-phase trades
|
||
/// only (train-phase trades are flushed when the eval phase
|
||
/// starts).
|
||
///
|
||
/// 0 disables the eval phase entirely (single-window training
|
||
/// only — the summary's profit_factor will reflect the full run
|
||
/// including train).
|
||
#[arg(long, default_value_t = 0)]
|
||
n_eval_steps: usize,
|
||
|
||
/// Per-step diagnostic JSONL path. One record per step capturing:
|
||
/// * step number, wall time
|
||
/// * loss components (l_bce/q/pi/v/aux/total) + λs (5)
|
||
/// * all 7 RL controller outputs (γ τ ε coef n_roll per_α scale)
|
||
/// * all 5 per-head learning rates (lr_bce/q/pi/v/aux)
|
||
/// * all 7 EMA inputs the controllers consume
|
||
/// * replay buffer length
|
||
/// * per-step reward + done aggregates (DtoH b_size floats per stat)
|
||
/// * per-step action histogram (counts per action class)
|
||
///
|
||
/// Default: `<out>/diag.jsonl`. Pass `--diag-jsonl /dev/null` (or
|
||
/// any unwritable path that errors at open) to disable. The dump
|
||
/// happens AFTER each `step_with_lobsim` so it sees the freshly-
|
||
/// adapted ISV values. Writes are line-buffered + flushed every
|
||
/// `log_every` steps so a `tail -f` shows in-flight progress.
|
||
///
|
||
/// Critically: this provides the per-step diagnostic visibility
|
||
/// the cluster smoke needs to detect anomalies early. Without it,
|
||
/// the smoke is "blind" — only the final summary is observable,
|
||
/// and a controller drifting off-anchor mid-run is invisible until
|
||
/// it produces a non-finite loss (G8 NaN abort).
|
||
#[arg(long)]
|
||
diag_jsonl: Option<PathBuf>,
|
||
}
|
||
|
||
fn parse_instrument_mode(s: &str) -> Result<InstrumentFilter, String> {
|
||
let trimmed = s.trim();
|
||
if trimmed.eq_ignore_ascii_case("all") {
|
||
return Ok(InstrumentFilter::All);
|
||
}
|
||
if trimmed.eq_ignore_ascii_case("front-month")
|
||
|| trimmed.eq_ignore_ascii_case("front_month")
|
||
{
|
||
return Ok(InstrumentFilter::FrontMonth);
|
||
}
|
||
if let Some(num) = trimmed.strip_prefix("id=") {
|
||
let id: u32 = num.parse().map_err(|e| {
|
||
format!("instrument-mode: expected `id=<u32>` but failed to parse '{num}': {e}")
|
||
})?;
|
||
return Ok(InstrumentFilter::Id(id));
|
||
}
|
||
Err(format!(
|
||
"instrument-mode: expected `all`, `front-month`, or `id=<N>`; got '{s}'"
|
||
))
|
||
}
|
||
|
||
#[derive(Serialize, Default)]
|
||
struct AlphaRlTrainSummary {
|
||
n_steps_planned: usize,
|
||
n_steps_completed: usize,
|
||
seq_len: usize,
|
||
n_backtests: usize,
|
||
per_capacity: usize,
|
||
seed: u64,
|
||
|
||
/// True iff `n_steps_completed == n_steps_planned` AND no NaN
|
||
/// abort fired. False if early-terminated for any reason (NaN,
|
||
/// loader EOF, panic — though panics also exit nonzero).
|
||
completed_clean: bool,
|
||
|
||
/// Final-step loss components (raw, not λ-weighted). All NaN if
|
||
/// trainer never returned a finite step.
|
||
final_l_bce: f32,
|
||
final_l_q: f32,
|
||
final_l_pi: f32,
|
||
final_l_v: f32,
|
||
final_l_aux: f32,
|
||
final_l_total: f32,
|
||
|
||
/// Replay buffer state at exit.
|
||
final_replay_len: usize,
|
||
|
||
/// Step at which the first NaN was observed, if any. -1 means
|
||
/// no NaN observed.
|
||
nan_abort_step: i64,
|
||
}
|
||
|
||
fn main() -> Result<()> {
|
||
tracing_subscriber::fmt()
|
||
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
|
||
.init();
|
||
|
||
let cli = Cli::parse();
|
||
|
||
std::fs::create_dir_all(&cli.out)
|
||
.with_context(|| format!("mkdir -p {}", cli.out.display()))?;
|
||
std::fs::create_dir_all(&cli.predecoded_dir)
|
||
.with_context(|| format!("mkdir -p {}", cli.predecoded_dir.display()))?;
|
||
|
||
let dev = MlDevice::cuda(cli.gpu_idx)
|
||
.with_context(|| format!("MlDevice::cuda({})", cli.gpu_idx))?;
|
||
eprintln!("CUDA device {} initialised", cli.gpu_idx);
|
||
|
||
// Loader: paired (s_t, s_{t+1}) sequences. inference_only=false so
|
||
// next_sequence_pair is available (the labels it preloads are
|
||
// unused by step_with_lobsim, but the API requires the label
|
||
// precomputation pass per `feedback_no_partial_refactor`).
|
||
let all_files = discover_mbp10_files_sorted(&cli.mbp10_data_dir)
|
||
.with_context(|| format!("discover_mbp10_files_sorted({})", cli.mbp10_data_dir.display()))?;
|
||
anyhow::ensure!(
|
||
!all_files.is_empty(),
|
||
"no MBP-10 .dbn.zst files under {}",
|
||
cli.mbp10_data_dir.display()
|
||
);
|
||
eprintln!("loader: {} MBP-10 files total", all_files.len());
|
||
|
||
// Walk-forward fold split per `pearl_single_window_oos_is_not_oos`.
|
||
// Slice the file list into `n_folds` equal blocks. fold k uses
|
||
// blocks [0..=k] for training and block [k+1] for evaluation. With
|
||
// n_folds=1 (default) train_files = all_files and eval_files is
|
||
// empty (single-window mode, no eval phase).
|
||
let (train_files, eval_files) = if cli.n_folds <= 1 || cli.n_eval_steps == 0 {
|
||
eprintln!("walk-forward: disabled (n_folds={}, n_eval_steps={})", cli.n_folds, cli.n_eval_steps);
|
||
(all_files, Vec::new())
|
||
} else {
|
||
anyhow::ensure!(
|
||
cli.fold_idx + 1 < cli.n_folds,
|
||
"fold_idx {} requires fold_idx + 1 < n_folds ({}); the last block must be the eval window",
|
||
cli.fold_idx,
|
||
cli.n_folds
|
||
);
|
||
let n_files = all_files.len();
|
||
let block_size = n_files / cli.n_folds;
|
||
anyhow::ensure!(
|
||
block_size >= 1,
|
||
"not enough MBP-10 files ({}) to split into n_folds={} (need ≥ n_folds files)",
|
||
n_files,
|
||
cli.n_folds
|
||
);
|
||
// Train: blocks 0..=fold_idx, Eval: block fold_idx+1.
|
||
let train_end = (cli.fold_idx + 1) * block_size;
|
||
let eval_end = (cli.fold_idx + 2) * block_size;
|
||
let train: Vec<_> = all_files[..train_end].to_vec();
|
||
let eval: Vec<_> = all_files[train_end..eval_end.min(n_files)].to_vec();
|
||
eprintln!(
|
||
"walk-forward fold {}/{}: train={} files ({}..{}), eval={} files ({}..{})",
|
||
cli.fold_idx, cli.n_folds, train.len(), 0, train_end,
|
||
eval.len(), train_end, eval_end.min(n_files)
|
||
);
|
||
(train, eval)
|
||
};
|
||
|
||
let multi_resolution: ml_alpha::data::aggregation::MultiResolutionConfig =
|
||
format!("1:{}", cli.seq_len)
|
||
.parse()
|
||
.with_context(|| format!("multi-resolution `1:{}`", cli.seq_len))?;
|
||
anyhow::ensure!(
|
||
multi_resolution.total_positions() == cli.seq_len,
|
||
"internal: single-scale 1:{} should produce total_positions={} but got {}",
|
||
cli.seq_len,
|
||
cli.seq_len,
|
||
multi_resolution.total_positions()
|
||
);
|
||
|
||
let loader_cfg = MultiHorizonLoaderConfig {
|
||
files: train_files,
|
||
predecoded_dir: cli.predecoded_dir.clone(),
|
||
multi_resolution: multi_resolution.clone(),
|
||
horizons: HORIZONS,
|
||
// Soft cap above n_steps × 2 so next_sequence_pair never
|
||
// returns Ok(None) mid-run. Loader's per-file random anchor
|
||
// sampling reuses files indefinitely up to this cap.
|
||
n_max_sequences: cli.n_steps.saturating_mul(2).max(64),
|
||
seed: cli.seed,
|
||
inference_only: false,
|
||
outcome_label_cost: cli.outcome_label_cost,
|
||
instrument_filter: cli.instrument_mode,
|
||
};
|
||
let mut loader = MultiHorizonLoader::new(&loader_cfg).context("MultiHorizonLoader::new")?;
|
||
|
||
// Trainer + LobSim.
|
||
let perception_cfg = PerceptionTrainerConfig {
|
||
seq_len: cli.seq_len,
|
||
n_batch: cli.n_backtests,
|
||
..PerceptionTrainerConfig::default()
|
||
};
|
||
let trainer_cfg = IntegratedTrainerConfig {
|
||
perception: perception_cfg,
|
||
dqn_seed: cli.seed.wrapping_add(0xDA),
|
||
ppo_seed: cli.seed.wrapping_add(0xBE),
|
||
per_capacity: cli.per_capacity,
|
||
per_seed: cli.seed.wrapping_add(0x9E37_79B9),
|
||
};
|
||
let mut trainer =
|
||
IntegratedTrainer::new(&dev, trainer_cfg).context("IntegratedTrainer::new")?;
|
||
let mut sim = LobSimCuda::new(cli.n_backtests, &dev).context("LobSimCuda::new")?;
|
||
eprintln!(
|
||
"trainer + LobSimCuda initialised: seq_len={} n_backtests={} per_capacity={}",
|
||
cli.seq_len, cli.n_backtests, cli.per_capacity
|
||
);
|
||
|
||
let mut summary = AlphaRlTrainSummary {
|
||
n_steps_planned: cli.n_steps,
|
||
seq_len: cli.seq_len,
|
||
n_backtests: cli.n_backtests,
|
||
per_capacity: cli.per_capacity,
|
||
seed: cli.seed,
|
||
nan_abort_step: -1,
|
||
..Default::default()
|
||
};
|
||
let mut last_stats: Option<IntegratedStepStats> = None;
|
||
|
||
// ── Per-step JSONL diag writer. ──────────────────────────────────
|
||
// Default path: `<out>/diag.jsonl`. The smoke run's failure-mode
|
||
// catalogue (NaN abort, controller drift, replay stagnation, reward
|
||
// explosion) is detected via post-hoc inspection of this file —
|
||
// without it, the run is blind beyond the eprintln every-N-step
|
||
// ticks.
|
||
let diag_path = cli
|
||
.diag_jsonl
|
||
.clone()
|
||
.unwrap_or_else(|| cli.out.join("diag.jsonl"));
|
||
let diag_file = std::fs::File::create(&diag_path)
|
||
.with_context(|| format!("create diag jsonl {}", diag_path.display()))?;
|
||
let mut diag = BufWriter::new(diag_file);
|
||
eprintln!("per-step diag JSONL: {}", diag_path.display());
|
||
|
||
// Per-step scratch (replaced each step via mapped-pinned helpers).
|
||
|
||
// ── Training loop. ───────────────────────────────────────────────
|
||
let t_start = std::time::Instant::now();
|
||
for step in 0..cli.n_steps {
|
||
let pair = loader
|
||
.next_sequence_pair()
|
||
.with_context(|| format!("next_sequence_pair at step {step}"))?;
|
||
let (s_t, s_tp1) = match pair {
|
||
Some(p) => p,
|
||
None => {
|
||
eprintln!(
|
||
"loader EOF at step {step} (n_max_sequences exhausted) — exiting early"
|
||
);
|
||
break;
|
||
}
|
||
};
|
||
|
||
let stats = trainer
|
||
.step_with_lobsim(&s_t.snapshots, &s_tp1.snapshots, &mut sim)
|
||
.with_context(|| format!("step_with_lobsim at step {step}"))?;
|
||
|
||
// Gate G8: NaN abort. Per `feedback_stop_on_anomaly` +
|
||
// `feedback_kill_runs_on_anomaly_quickly` the cluster smoke
|
||
// tail-watcher kills the workflow on non-zero exit; here we
|
||
// signal that with exit code 2 (distinct from clean exit 0
|
||
// and panic-derived exit codes).
|
||
let non_finite = !stats.l_total.is_finite()
|
||
|| !stats.l_bce.is_finite()
|
||
|| !stats.l_q.is_finite()
|
||
|| !stats.l_pi.is_finite()
|
||
|| !stats.l_v.is_finite()
|
||
|| !stats.l_aux.is_finite();
|
||
if non_finite {
|
||
eprintln!(
|
||
"G8 NaN ABORT at step {step}: l_bce={} l_q={} l_pi={} l_v={} l_aux={} l_total={}",
|
||
stats.l_bce, stats.l_q, stats.l_pi, stats.l_v, stats.l_aux, stats.l_total
|
||
);
|
||
summary.n_steps_completed = step;
|
||
summary.nan_abort_step = step as i64;
|
||
summary.completed_clean = false;
|
||
write_summary(&cli.out, &summary)?;
|
||
process::exit(2);
|
||
}
|
||
|
||
// ── Per-step diag dump (mapped-pinned reads). ────────────────
|
||
// Per `feedback_no_htod_htoh_only_mapped_pinned`: raw
|
||
// `stream.memcpy_dtoh` on a regular `&mut [T]` is forbidden;
|
||
// the only permitted CPU↔GPU path is mapped-pinned staging
|
||
// (cuMemHostAlloc DEVICEMAP → host writes via host_ptr +
|
||
// kernel reads dev_ptr + DtoD between them). The
|
||
// `read_slice_*_d_pub` helpers in `trainer::integrated`
|
||
// encapsulate the pattern.
|
||
let dev_stream = dev.cuda_stream().context("cuda_stream for diag")?;
|
||
let actions_host = read_slice_i32_d_pub(dev_stream, &trainer.actions_d, cli.n_backtests)
|
||
.context("diag: read actions_d")?;
|
||
let rewards_host = read_slice_d_pub(dev_stream, &trainer.rewards_d, cli.n_backtests)
|
||
.context("diag: read rewards_d")?;
|
||
let dones_host = read_slice_d_pub(dev_stream, &trainer.dones_d, cli.n_backtests)
|
||
.context("diag: read dones_d")?;
|
||
|
||
let mut act_hist = [0u32; 9];
|
||
for &a in &actions_host {
|
||
if (0..9).contains(&a) {
|
||
act_hist[a as usize] += 1;
|
||
}
|
||
}
|
||
let reward_sum: f32 = rewards_host.iter().sum();
|
||
let reward_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
let reward_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
|
||
let reward_abs_max = rewards_host.iter().map(|r| r.abs()).fold(0.0f32, f32::max);
|
||
let done_count: u32 = dones_host.iter().map(|&d| if d > 0.5 { 1 } else { 0 }).sum();
|
||
|
||
let isv = &trainer.isv_host;
|
||
let record = json!({
|
||
"step": step,
|
||
"elapsed_s": t_start.elapsed().as_secs_f32(),
|
||
"loss": {
|
||
"bce": stats.l_bce,
|
||
"q": stats.l_q,
|
||
"pi": stats.l_pi,
|
||
"v": stats.l_v,
|
||
"aux": stats.l_aux,
|
||
"total": stats.l_total,
|
||
},
|
||
"lambdas": {
|
||
"bce": stats.lambdas.bce,
|
||
"q": stats.lambdas.q,
|
||
"pi": stats.lambdas.pi,
|
||
"v": stats.lambdas.v,
|
||
"aux": stats.lambdas.aux,
|
||
},
|
||
// 7 R5 controller outputs (the adaptive knobs).
|
||
"isv_out": {
|
||
"gamma": isv[RL_GAMMA_INDEX],
|
||
"target_tau": isv[RL_TARGET_TAU_INDEX],
|
||
"ppo_clip_eps": isv[RL_PPO_CLIP_INDEX],
|
||
"entropy_coef": isv[RL_ENTROPY_COEF_INDEX],
|
||
"n_rollout_steps": isv[RL_N_ROLLOUT_STEPS_INDEX],
|
||
"per_alpha": isv[RL_PER_ALPHA_INDEX],
|
||
"reward_scale": isv[RL_REWARD_SCALE_INDEX],
|
||
},
|
||
// 5 per-head learning rates (rl_lr_controller emits).
|
||
"isv_lr": {
|
||
"bce": isv[RL_LR_BCE_INDEX],
|
||
"q": isv[RL_LR_Q_INDEX],
|
||
"pi": isv[RL_LR_PI_INDEX],
|
||
"v": isv[RL_LR_V_INDEX],
|
||
"aux": isv[RL_LR_AUX_INDEX],
|
||
},
|
||
// 7 EMA inputs the controllers consume — knowing these
|
||
// makes controller behaviour debuggable (was target wrong?
|
||
// was input wrong? was alpha wrong?).
|
||
"isv_ema_in": {
|
||
"mean_trade_duration": isv[RL_MEAN_TRADE_DURATION_EMA_INDEX],
|
||
"q_divergence": isv[RL_Q_DIVERGENCE_EMA_INDEX],
|
||
"kl_pi": isv[RL_KL_PI_EMA_INDEX],
|
||
"entropy_observed": isv[RL_ENTROPY_OBSERVED_EMA_INDEX],
|
||
"advantage_var_ratio": isv[RL_ADVANTAGE_VAR_RATIO_EMA_INDEX],
|
||
"td_kurtosis": isv[RL_TD_KURTOSIS_EMA_INDEX],
|
||
"mean_abs_pnl": isv[RL_MEAN_ABS_PNL_EMA_INDEX],
|
||
},
|
||
// Grad-norm EMAs (diagnostic — no longer drive the LR
|
||
// controller as of commit 042de99e6 but still useful to
|
||
// correlate with plateau-decay events).
|
||
"grad_norm_ema": {
|
||
"q": isv[RL_Q_GRAD_NORM_EMA_INDEX],
|
||
"pi": isv[RL_PI_GRAD_NORM_EMA_INDEX],
|
||
"v": isv[RL_V_GRAD_NORM_EMA_INDEX],
|
||
},
|
||
// Plateau-decay LR controller state (per head). loss_ema is
|
||
// the controller's internal slow EMA at α=0.05; best is the
|
||
// lowest loss_ema observed so far; steps_since_best is the
|
||
// staleness counter that fires LR halving at PLATEAU_PATIENCE
|
||
// (1000). Confirms whether the controller is seeing plateaus
|
||
// and decaying for the right reason.
|
||
"lr_plateau": {
|
||
"q": { "loss_ema": isv[RL_LR_Q_LOSS_EMA_INDEX],
|
||
"best": isv[RL_LR_Q_BEST_LOSS_INDEX],
|
||
"stale": isv[RL_LR_Q_STEPS_SINCE_BEST_INDEX],
|
||
"warmup": isv[RL_LR_Q_WARMUP_COUNTER_INDEX] },
|
||
"pi": { "loss_ema": isv[RL_LR_PI_LOSS_EMA_INDEX],
|
||
"best": isv[RL_LR_PI_BEST_LOSS_INDEX],
|
||
"stale": isv[RL_LR_PI_STEPS_SINCE_BEST_INDEX],
|
||
"warmup": isv[RL_LR_PI_WARMUP_COUNTER_INDEX] },
|
||
"v": { "loss_ema": isv[RL_LR_V_LOSS_EMA_INDEX],
|
||
"best": isv[RL_LR_V_BEST_LOSS_INDEX],
|
||
"stale": isv[RL_LR_V_STEPS_SINCE_BEST_INDEX],
|
||
"warmup": isv[RL_LR_V_WARMUP_COUNTER_INDEX] },
|
||
},
|
||
"replay_len": trainer.replay.len(),
|
||
// Post-scale, POST-clamp reward stats — what V regression
|
||
// and Q distributional projection actually saw this step.
|
||
// The `scaled_pre_clamp_max` field (below) is the SAME
|
||
// batch's max BEFORE the [-3, +1] clamp fired. If
|
||
// scaled_pre_clamp_max > 1.0 (clamp threshold on the win
|
||
// side) repeatedly, the clamp is doing load-bearing work
|
||
// — i.e. the reward_scale controller is failing to keep
|
||
// typical trade magnitudes within the C51 atom support.
|
||
"rewards": {
|
||
"sum": reward_sum,
|
||
"max": reward_max,
|
||
"min": reward_min,
|
||
"abs_max": reward_abs_max,
|
||
"scaled_pre_clamp_max":
|
||
isv[RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX],
|
||
},
|
||
"done_count": done_count,
|
||
// Action histogram: indices match `Action` enum
|
||
// (0=ShortLarge, 1=ShortSmall, 2=Hold, 3=FlatFromLong,
|
||
// 4=FlatFromShort, 5=LongSmall, 6=LongLarge,
|
||
// 7=TrailTighten, 8=TrailLoosen).
|
||
"action_hist": act_hist.to_vec(),
|
||
});
|
||
writeln!(diag, "{}", record).context("diag: writeln jsonl")?;
|
||
|
||
if step % cli.log_every == 0 || step == cli.n_steps - 1 {
|
||
diag.flush().context("diag: flush")?;
|
||
eprintln!(
|
||
"step {:>6}/{}: l_q={:.4} l_pi={:.4} l_v={:.4} l_total={:.4} \
|
||
γ={:.4} ε={:.4} per_α={:.4} scale={:.4} \
|
||
replay={} dones={} rew_sum={:.3} elapsed={:.1}s",
|
||
step,
|
||
cli.n_steps,
|
||
stats.l_q,
|
||
stats.l_pi,
|
||
stats.l_v,
|
||
stats.l_total,
|
||
isv[RL_GAMMA_INDEX],
|
||
isv[RL_PPO_CLIP_INDEX],
|
||
isv[RL_PER_ALPHA_INDEX],
|
||
isv[RL_REWARD_SCALE_INDEX],
|
||
trainer.replay.len(),
|
||
done_count,
|
||
reward_sum,
|
||
t_start.elapsed().as_secs_f32()
|
||
);
|
||
}
|
||
last_stats = Some(stats);
|
||
summary.n_steps_completed = step + 1;
|
||
}
|
||
diag.flush().context("diag: final flush")?;
|
||
|
||
// ── Eval phase (walk-forward G8). ────────────────────────────────
|
||
// After the train phase, run n_eval_steps additional step_with_lobsim
|
||
// calls on the held-out eval files. Trade records accumulated by
|
||
// LobSim during eval are isolated from train-phase records via the
|
||
// `head_before_eval` checkpoint — eval records are sliced as
|
||
// `all_records[head_before_eval..]` post-eval.
|
||
//
|
||
// Note: this is NOT pure-eval mode — step_with_lobsim still runs
|
||
// its learning machinery (Adam steps, PER updates). At b_size=1
|
||
// the per-step learning effect is small relative to the policy
|
||
// already accumulated in train; the eval profit_factor approximates
|
||
// the train-end policy's OOS performance. True pure-eval (forward
|
||
// only, no backward) is a follow-up architectural change.
|
||
if cli.n_eval_steps > 0 && !eval_files.is_empty() {
|
||
let head_before_eval = sim
|
||
.read_total_trade_count()
|
||
.context("read trade count pre-eval")?;
|
||
eprintln!(
|
||
"── eval phase: {} steps on {} held-out files (trade-record checkpoint head={}) ──",
|
||
cli.n_eval_steps, eval_files.len(), head_before_eval
|
||
);
|
||
|
||
let eval_loader_cfg = MultiHorizonLoaderConfig {
|
||
files: eval_files.clone(),
|
||
predecoded_dir: cli.predecoded_dir.clone(),
|
||
multi_resolution,
|
||
horizons: HORIZONS,
|
||
n_max_sequences: cli.n_eval_steps.saturating_mul(2).max(64),
|
||
// Different seed for eval so sequence sampling is independent
|
||
// of the train seed (deterministic given the eval seed).
|
||
seed: cli.seed.wrapping_add(0xE7AE),
|
||
inference_only: false,
|
||
outcome_label_cost: cli.outcome_label_cost,
|
||
instrument_filter: cli.instrument_mode,
|
||
};
|
||
let mut eval_loader = MultiHorizonLoader::new(&eval_loader_cfg)
|
||
.context("MultiHorizonLoader::new (eval)")?;
|
||
|
||
for eval_step in 0..cli.n_eval_steps {
|
||
let pair = eval_loader
|
||
.next_sequence_pair()
|
||
.with_context(|| format!("eval next_sequence_pair at step {eval_step}"))?;
|
||
let (s_t, s_tp1) = match pair {
|
||
Some(p) => p,
|
||
None => {
|
||
eprintln!("eval loader EOF at step {eval_step}");
|
||
break;
|
||
}
|
||
};
|
||
let stats = trainer
|
||
.step_with_lobsim(&s_t.snapshots, &s_tp1.snapshots, &mut sim)
|
||
.with_context(|| format!("step_with_lobsim eval step {eval_step}"))?;
|
||
if !stats.l_total.is_finite() {
|
||
eprintln!("G8 NaN ABORT during eval at step {eval_step}");
|
||
process::exit(2);
|
||
}
|
||
if eval_step % cli.log_every == 0 || eval_step == cli.n_eval_steps - 1 {
|
||
eprintln!(
|
||
"eval {:>5}/{}: l_total={:.4} elapsed={:.1}s",
|
||
eval_step, cli.n_eval_steps, stats.l_total,
|
||
t_start.elapsed().as_secs_f32()
|
||
);
|
||
}
|
||
}
|
||
|
||
// Drain trade records, slice to eval-only, compute summary.
|
||
let all_records = sim
|
||
.read_trade_records(0)
|
||
.context("read trade records post-eval")?;
|
||
let head_before_usize = head_before_eval as usize;
|
||
let eval_records: Vec<_> = if all_records.len() > head_before_usize {
|
||
all_records[head_before_usize..].to_vec()
|
||
} else {
|
||
// Trade log wrapped past TRADE_LOG_CAP — use what we have.
|
||
// For smoke (b_size=1, ≤200 trades) this branch never fires.
|
||
eprintln!(
|
||
"warning: trade log wrapped — head_before={} but only {} records readable",
|
||
head_before_usize, all_records.len()
|
||
);
|
||
all_records
|
||
};
|
||
|
||
// Synthesise a pnl_curve from per-trade cumulative PnL for
|
||
// compute_summary's max_drawdown calc.
|
||
let mut pnl_curve = Vec::with_capacity(eval_records.len());
|
||
let mut cum = 0.0f32;
|
||
for r in &eval_records {
|
||
cum += (r.realised_pnl_usd_fp as f32) / 100.0;
|
||
pnl_curve.push(cum);
|
||
}
|
||
let eval_summary =
|
||
ml_backtesting::artifacts::compute_summary(&eval_records, &pnl_curve);
|
||
|
||
eprintln!(
|
||
"eval summary: n_trades={} pnl_usd={:.2} pf={:.3} sharpe_ann={:.3} max_dd_usd={:.2} win_rate={:.3}",
|
||
eval_summary.n_trades,
|
||
eval_summary.total_pnl_usd,
|
||
eval_summary.profit_factor,
|
||
eval_summary.sharpe_ann,
|
||
eval_summary.max_drawdown_usd,
|
||
eval_summary.win_rate,
|
||
);
|
||
|
||
let eval_summary_path = cli.out.join("eval_summary.json");
|
||
let f = std::fs::File::create(&eval_summary_path)
|
||
.with_context(|| format!("create {}", eval_summary_path.display()))?;
|
||
serde_json::to_writer_pretty(f, &eval_summary)
|
||
.with_context(|| format!("write {}", eval_summary_path.display()))?;
|
||
eprintln!("eval summary written: {}", eval_summary_path.display());
|
||
}
|
||
|
||
// ── Compose + emit summary. ──────────────────────────────────────
|
||
if let Some(s) = last_stats {
|
||
summary.final_l_bce = s.l_bce;
|
||
summary.final_l_q = s.l_q;
|
||
summary.final_l_pi = s.l_pi;
|
||
summary.final_l_v = s.l_v;
|
||
summary.final_l_aux = s.l_aux;
|
||
summary.final_l_total = s.l_total;
|
||
}
|
||
summary.final_replay_len = trainer.replay.len();
|
||
summary.completed_clean = summary.n_steps_completed == summary.n_steps_planned;
|
||
write_summary(&cli.out, &summary)?;
|
||
|
||
eprintln!(
|
||
"alpha_rl_train complete: {} / {} steps in {:.1}s (replay len = {})",
|
||
summary.n_steps_completed,
|
||
summary.n_steps_planned,
|
||
t_start.elapsed().as_secs_f32(),
|
||
summary.final_replay_len
|
||
);
|
||
Ok(())
|
||
}
|
||
|
||
fn write_summary(out_dir: &std::path::Path, summary: &AlphaRlTrainSummary) -> Result<()> {
|
||
let path = out_dir.join("alpha_rl_train_summary.json");
|
||
let f = std::fs::File::create(&path)
|
||
.with_context(|| format!("create {}", path.display()))?;
|
||
serde_json::to_writer_pretty(f, summary)
|
||
.with_context(|| format!("write {}", path.display()))?;
|
||
eprintln!("summary written: {}", path.display());
|
||
Ok(())
|
||
}
|