5546a45bc3
gpu_n_episodes was manually overridden in GPU profiles, training configs, test files, and hyperopt — all set to 0 or small fixed values that bypassed the auto-scaling logic, causing a div-by-zero crash in train_baseline_rl. Now: single auto-scaling path via optimal_n_episodes() from VRAM/SM count. No manual override field. Cap at 16384 (consistent with AutoBatchSizer's 8192 cap pattern). Floor at 32 for small GPUs. Removed gpu_n_episodes from: - DQNHyperparameters, PpoHyperparameters structs - All 4 GPU profiles (rtx3050, h100, a100, default) - Training profiles (smoketest, localdev) - ExperienceProfile struct + serde - Hyperopt adapter - All test overrides Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1283 lines
52 KiB
Rust
1283 lines
52 KiB
Rust
#![allow(
|
|
clippy::assertions_on_constants,
|
|
clippy::assertions_on_result_states,
|
|
clippy::clone_on_copy,
|
|
clippy::decimal_literal_representation,
|
|
clippy::doc_markdown,
|
|
clippy::empty_line_after_doc_comments,
|
|
clippy::field_reassign_with_default,
|
|
clippy::get_unwrap,
|
|
clippy::identity_op,
|
|
clippy::inconsistent_digit_grouping,
|
|
clippy::indexing_slicing,
|
|
clippy::integer_division,
|
|
clippy::len_zero,
|
|
clippy::let_underscore_must_use,
|
|
clippy::manual_div_ceil,
|
|
clippy::manual_let_else,
|
|
clippy::manual_range_contains,
|
|
clippy::modulo_arithmetic,
|
|
clippy::needless_range_loop,
|
|
clippy::non_ascii_literal,
|
|
clippy::redundant_clone,
|
|
clippy::shadow_reuse,
|
|
clippy::shadow_same,
|
|
clippy::shadow_unrelated,
|
|
clippy::single_match_else,
|
|
clippy::str_to_string,
|
|
clippy::string_slice,
|
|
clippy::tests_outside_test_module,
|
|
clippy::too_many_lines,
|
|
clippy::unnecessary_wraps,
|
|
clippy::unseparated_literal_suffix,
|
|
clippy::use_debug,
|
|
clippy::useless_vec,
|
|
clippy::wildcard_enum_match_arm,
|
|
clippy::else_if_without_else,
|
|
clippy::expect_used,
|
|
clippy::missing_const_for_fn,
|
|
clippy::similar_names,
|
|
clippy::type_complexity,
|
|
clippy::collapsible_else_if,
|
|
clippy::doc_lazy_continuation,
|
|
clippy::items_after_test_module,
|
|
clippy::map_clone,
|
|
clippy::multiple_unsafe_ops_per_block,
|
|
clippy::unwrap_or_default,
|
|
clippy::assign_op_pattern,
|
|
clippy::needless_borrow,
|
|
clippy::println_empty_string,
|
|
clippy::unnecessary_cast,
|
|
clippy::used_underscore_binding,
|
|
clippy::create_dir,
|
|
clippy::implicit_saturating_sub,
|
|
clippy::exit,
|
|
clippy::expect_fun_call,
|
|
clippy::too_many_arguments,
|
|
clippy::unnecessary_map_or,
|
|
clippy::unwrap_used,
|
|
dead_code,
|
|
unused_imports,
|
|
unused_variables,
|
|
clippy::cloned_ref_to_slice_refs,
|
|
clippy::neg_multiply,
|
|
clippy::while_let_loop,
|
|
clippy::bool_assert_comparison,
|
|
clippy::excessive_precision,
|
|
clippy::trivially_copy_pass_by_ref,
|
|
clippy::op_ref,
|
|
clippy::redundant_closure,
|
|
clippy::unnecessary_lazy_evaluations,
|
|
clippy::if_then_some_else_none,
|
|
clippy::unnecessary_to_owned,
|
|
clippy::single_component_path_imports,
|
|
)]
|
|
//! Walk-forward RL training binary for DQN and PPO models.
|
|
//!
|
|
//! Trains models using expanding walk-forward windows on real OHLCV data loaded
|
|
//! from Databento DBN files. Supports early stopping, checkpoint saving, and
|
|
//! normalization statistics export for reproducible evaluation.
|
|
//!
|
|
//! # Usage
|
|
//!
|
|
//! ```bash
|
|
//! SQLX_OFFLINE=true cargo run -p ml --example train_baseline -- \
|
|
//! --model both --epochs 50 --batch-size 128 \
|
|
//! --data-dir test_data/futures-baseline \
|
|
//! --output-dir ml/trained_models
|
|
//! ```
|
|
|
|
#![allow(unused_crate_dependencies)]
|
|
#![deny(
|
|
clippy::unwrap_used,
|
|
clippy::expect_used,
|
|
clippy::panic,
|
|
clippy::indexing_slicing
|
|
)]
|
|
|
|
use std::path::{Path, PathBuf};
|
|
|
|
use anyhow::{Context, Result};
|
|
use clap::Parser;
|
|
use serde_json::Value;
|
|
use tracing::{error, info, warn};
|
|
|
|
use ml::cuda_pipeline::DqnGpuData;
|
|
use std::sync::Arc;
|
|
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
|
use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer};
|
|
|
|
#[allow(unreachable_pub)]
|
|
mod baseline_common;
|
|
use baseline_common::completion::{write_failure_marker, write_success_marker, CompletionMetrics};
|
|
use baseline_common::{load_all_bars, spread_cost_bps};
|
|
use common::metrics::{server as metrics_server, training_metrics as metrics};
|
|
use ml::features::extraction::{extract_ml_features, FeatureVector};
|
|
use ml_core::gpu::profile::GpuProfile;
|
|
use ml::types::OHLCVBar;
|
|
use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CLI Arguments
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Walk-forward training binary for DQN and PPO baseline models.
|
|
#[derive(Parser, Debug)]
|
|
#[command(name = "train_baseline_rl", about = "Train DQN/PPO with walk-forward RL windows")]
|
|
struct Args {
|
|
/// Which model(s) to train: "dqn", "ppo", or "both"
|
|
#[arg(long, default_value = "both")]
|
|
model: String,
|
|
|
|
/// Maximum training epochs per fold
|
|
#[arg(long, default_value_t = 50)]
|
|
epochs: usize,
|
|
|
|
/// Training batch size
|
|
#[arg(long, default_value_t = 0)]
|
|
batch_size: usize,
|
|
|
|
/// Path to directory containing .dbn.zst files (env: FOXHUNT_DATA_DIR)
|
|
#[arg(long, env = "FOXHUNT_DATA_DIR")]
|
|
data_dir: PathBuf,
|
|
|
|
/// Output directory for trained model checkpoints
|
|
#[arg(long, default_value = "ml/trained_models")]
|
|
output_dir: PathBuf,
|
|
|
|
/// Optional path to hyperopt results JSON -- overrides matching config fields
|
|
#[arg(long)]
|
|
hyperopt_params: Option<PathBuf>,
|
|
|
|
/// Feature dimension (42 market + 3 portfolio = 45, or 53 with OFI; must match trainer `state_dim`)
|
|
#[arg(long, default_value_t = 43)]
|
|
feature_dim: usize,
|
|
|
|
/// Early stopping patience (epochs without improvement)
|
|
#[arg(long, default_value_t = 10)]
|
|
patience: usize,
|
|
|
|
/// Number of actions (5 exposure levels for DQN, pass --num-actions 45 for PPO)
|
|
#[arg(long, default_value_t = 5)]
|
|
num_actions: usize,
|
|
|
|
/// Walk-forward: initial training window in months
|
|
#[arg(long, default_value_t = 12)]
|
|
train_months: u32,
|
|
|
|
/// Walk-forward: validation window in months
|
|
#[arg(long, default_value_t = 3)]
|
|
val_months: u32,
|
|
|
|
/// Walk-forward: test window in months
|
|
#[arg(long, default_value_t = 3)]
|
|
test_months: u32,
|
|
|
|
/// Walk-forward: step size in months between folds
|
|
#[arg(long, default_value_t = 3)]
|
|
step_months: u32,
|
|
|
|
/// Learning rate for optimizer
|
|
#[arg(long, default_value_t = 1e-4)]
|
|
learning_rate: f64,
|
|
|
|
/// Max environment steps per epoch (caps trajectory length for OOM safety;
|
|
/// 0 = use all bars, but this can use >1GB RAM per fold on large datasets)
|
|
#[arg(long, default_value_t = 2000)]
|
|
max_steps_per_epoch: usize,
|
|
|
|
/// Symbol subdirectory to load (e.g. "ES.FUT", "NQ.FUT")
|
|
#[arg(long, default_value = "ES.FUT")]
|
|
symbol: String,
|
|
|
|
/// Maximum absolute per-bar return; larger moves are clamped (contract roll filter)
|
|
#[arg(long, default_value_t = 0.01)]
|
|
max_bar_return: f64,
|
|
|
|
/// Round-trip commission cost in basis points (1 bps = 0.01%)
|
|
/// Applied to BUY/SELL rewards; HOLD is free.
|
|
/// Default 1.0 bps covers ~$4 exchange+broker for ES e-mini.
|
|
#[arg(long, default_value_t = 1.0)]
|
|
tx_cost_bps: f64,
|
|
|
|
/// Instrument tick size in price units (ES=0.25, NQ=0.25, ZN=1/64)
|
|
#[arg(long, default_value_t = 0.25)]
|
|
tick_size: f64,
|
|
|
|
/// Typical bid-ask spread in ticks (ES=1.0, ZN=1.0, 6E=2.0)
|
|
/// Half-spread slippage is added to `tx_cost_bps` per trade.
|
|
#[arg(long, default_value_t = 1.0)]
|
|
spread_ticks: f64,
|
|
|
|
/// Number of top hyperopt configs to train as ensemble (default 1 = best only).
|
|
/// When > 1, loads `top_k_params` from the hyperopt JSON and trains a separate
|
|
/// model for each param set, saving as `dqn_ensemble_{k}_fold_{fold}.safetensors`.
|
|
#[arg(long, default_value_t = 1)]
|
|
ensemble_top_k: usize,
|
|
|
|
/// Optional path to MBP-10 order book data directory for OFI features.
|
|
/// When set, enables 8 OFI features (OFI L1/L5, depth imbalance, VPIN, etc.)
|
|
/// expanding state dimension from 43 to 51.
|
|
#[arg(long)]
|
|
mbp10_data_dir: Option<PathBuf>,
|
|
|
|
/// Optional path to trade data directory (.dbn.zst files with Schema::Trades).
|
|
/// When set, real trade buy/sell classification feeds VPIN and Kyle's Lambda
|
|
/// instead of the tick-rule proxy. Requires --mbp10-data-dir to take effect.
|
|
#[arg(long)]
|
|
trades_data_dir: Option<PathBuf>,
|
|
|
|
/// Enable offline RL mode: train exclusively from a pre-collected dataset,
|
|
/// skipping online experience collection. Requires --dataset-path.
|
|
#[arg(long, default_value_t = false)]
|
|
offline: bool,
|
|
|
|
/// Path to a pre-collected experience dataset (bincode format).
|
|
/// Used with --offline to load a fixed dataset into the replay buffer.
|
|
#[arg(long)]
|
|
dataset_path: Option<PathBuf>,
|
|
|
|
/// Collect and save a dataset from the current policy, then exit.
|
|
/// Use this to generate datasets for offline training.
|
|
#[arg(long)]
|
|
collect_dataset: Option<PathBuf>,
|
|
|
|
/// Disable Branching DQN (3-head: exposure, order, urgency). Enabled by default.
|
|
#[arg(long)]
|
|
no_branching: bool,
|
|
|
|
/// Initial trading capital in dollars. Lower capital teaches conservative
|
|
/// position sizing. Must match hyperopt --initial-capital for consistency.
|
|
#[arg(long, default_value_t = 35_000.0)]
|
|
initial_capital: f64,
|
|
|
|
/// Minimum bars to hold a position before allowing exit (churn prevention).
|
|
/// Lower = more trades, higher = fewer. Default from TOML (typically 5).
|
|
#[arg(long)]
|
|
min_hold_bars: Option<usize>,
|
|
|
|
/// Named training profile to load from config/training/<profile>.toml.
|
|
/// Profile values are applied after hyperopt JSON but before explicit CLI args.
|
|
/// Known profiles: dqn-production, dqn-smoketest, dqn-hyperopt.
|
|
#[arg(long, default_value = "dqn-production")]
|
|
training_profile: String,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Hyperopt parameter loading
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Load `best_params` from a hyperopt results JSON file.
|
|
///
|
|
/// Expected format: `{ "model_key": { "best_params": { ... }, ... } }`
|
|
/// Returns `None` if the file doesn't exist or can't be parsed.
|
|
#[allow(clippy::cognitive_complexity)]
|
|
fn load_hyperopt_params(hp_path: &Option<PathBuf>, model_key: &str) -> Option<Value> {
|
|
let file_path = hp_path.as_ref()?;
|
|
if !file_path.exists() {
|
|
info!("Hyperopt params file not found: {}, using defaults", file_path.display());
|
|
return None;
|
|
}
|
|
let contents = match std::fs::read_to_string(file_path) {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
warn!("Failed to read hyperopt params {}: {}", file_path.display(), e);
|
|
return None;
|
|
}
|
|
};
|
|
let json: Value = match serde_json::from_str(&contents) {
|
|
Ok(v) => v,
|
|
Err(e) => {
|
|
warn!("Failed to parse hyperopt params JSON: {}", e);
|
|
return None;
|
|
}
|
|
};
|
|
let params = json.get(model_key)
|
|
.and_then(|m| m.get("best_params"))
|
|
.cloned();
|
|
if params.is_some() {
|
|
info!("Loaded hyperopt params for '{}' from {}", model_key, file_path.display());
|
|
} else {
|
|
warn!("No best_params found for '{}' in {}", model_key, file_path.display());
|
|
}
|
|
params
|
|
}
|
|
|
|
/// Load top-K param sets from a hyperopt results JSON file.
|
|
///
|
|
/// Expected format: `{ "model_key": { "top_k_params": [{ "params": {...}, ... }, ...], "best_params": {...} } }`
|
|
/// Falls back to `best_params` as a single entry when `top_k_params` is absent.
|
|
/// Returns a vec of length `k` (or fewer if not enough entries), each entry `Some(params)` or `None`.
|
|
fn load_top_k_params(hp_path: &Option<PathBuf>, model_key: &str, k: usize) -> Vec<Option<Value>> {
|
|
let file_path = match hp_path.as_ref() {
|
|
Some(p) if p.exists() => p,
|
|
_ => return vec![None; k],
|
|
};
|
|
let Ok(contents) = std::fs::read_to_string(file_path) else {
|
|
return vec![None; k];
|
|
};
|
|
let Ok(json): Result<Value, _> = serde_json::from_str(&contents) else {
|
|
return vec![None; k];
|
|
};
|
|
let top_k = json
|
|
.get(model_key)
|
|
.and_then(|m| m.get("top_k_params"))
|
|
.and_then(|v| v.as_array());
|
|
if let Some(arr) = top_k {
|
|
arr.iter()
|
|
.take(k)
|
|
.map(|entry| entry.get("params").cloned())
|
|
.collect()
|
|
} else {
|
|
// Fallback: just use best_params as the single entry
|
|
let best = json
|
|
.get(model_key)
|
|
.and_then(|m| m.get("best_params"))
|
|
.cloned();
|
|
vec![best]
|
|
}
|
|
}
|
|
|
|
fn hp_f64(params: &Option<Value>, key: &str) -> Option<f64> {
|
|
params.as_ref()?.get(key)?.as_f64()
|
|
}
|
|
|
|
fn hp_usize(params: &Option<Value>, key: &str) -> Option<usize> {
|
|
params.as_ref()?.get(key)?.as_u64().map(|v| v as usize)
|
|
}
|
|
|
|
fn hp_bool(params: &Option<Value>, key: &str) -> Option<bool> {
|
|
params.as_ref()?.get(key)?.as_bool()
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Fold data preparation (extracted for prefetching)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Prepared fold data: (`train_features`, `val_features`, `train_bars`, `val_bars`)
|
|
type FoldData = (Vec<FeatureVector>, Vec<FeatureVector>, Vec<OHLCVBar>, Vec<OHLCVBar>);
|
|
|
|
#[allow(clippy::cognitive_complexity)]
|
|
/// Prepare a fold's data for training: extract features, normalize, align bars.
|
|
///
|
|
/// Returns `None` if feature extraction fails or produces empty features.
|
|
fn prepare_fold_data(
|
|
window: &ml::walk_forward::WalkForwardWindow,
|
|
output_dir: &Path,
|
|
) -> Option<FoldData> {
|
|
let train_feat = match extract_ml_features(&window.train) {
|
|
Ok(f) => f,
|
|
Err(e) => {
|
|
warn!(" Fold {} -- train feature extraction failed: {}", window.fold, e);
|
|
return None;
|
|
}
|
|
};
|
|
let val_feat = match extract_ml_features(&window.val) {
|
|
Ok(f) => f,
|
|
Err(e) => {
|
|
warn!(" Fold {} -- val feature extraction failed: {}", window.fold, e);
|
|
return None;
|
|
}
|
|
};
|
|
if train_feat.is_empty() || val_feat.is_empty() {
|
|
warn!(" Fold {} -- empty features, skipping", window.fold);
|
|
return None;
|
|
}
|
|
|
|
let norm_stats = NormStats::from_features(&train_feat);
|
|
let train_norm = norm_stats.normalize_batch(&train_feat);
|
|
let val_norm = norm_stats.normalize_batch(&val_feat);
|
|
|
|
// Save NormStats for this fold
|
|
let norm_path = output_dir.join(format!("norm_stats_fold{}.json", window.fold));
|
|
match serde_json::to_string_pretty(&norm_stats) {
|
|
Ok(json) => {
|
|
let norm_tmp = norm_path.with_extension("json.tmp");
|
|
if let Err(e) = std::fs::write(&norm_tmp, &json) {
|
|
error!(" Failed to write NormStats tmp {}: {}", norm_tmp.display(), e);
|
|
return None;
|
|
}
|
|
if let Err(e) = std::fs::rename(&norm_tmp, &norm_path) {
|
|
error!(" Failed to rename NormStats {} -> {}: {}", norm_tmp.display(), norm_path.display(), e);
|
|
drop(std::fs::remove_file(&norm_tmp));
|
|
return None;
|
|
}
|
|
info!(" Saved NormStats to {}", norm_path.display());
|
|
}
|
|
Err(e) => {
|
|
error!(" Failed to serialize NormStats: {}", e);
|
|
return None;
|
|
}
|
|
}
|
|
|
|
let train_warmup = window.train.len().saturating_sub(train_norm.len());
|
|
let train_bars_aligned = window.train.get(train_warmup..).unwrap_or(&window.train).to_vec();
|
|
let val_warmup = window.val.len().saturating_sub(val_norm.len());
|
|
let val_bars_aligned = window.val.get(val_warmup..).unwrap_or(&window.val).to_vec();
|
|
|
|
Some((train_norm, val_norm, train_bars_aligned, val_bars_aligned))
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// DQN Training
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Convert pre-aligned features and bars into the format expected by
|
|
/// `DQNTrainer::train_with_preloaded_data`: `Vec<(FeatureVector, Vec<f64>)>`.
|
|
///
|
|
/// The `Vec<f64>` target contains OHLCV prices for the trainer's internal
|
|
/// reward computation and portfolio simulation.
|
|
fn features_to_trainer_format(
|
|
features: &[FeatureVector],
|
|
bars: &[OHLCVBar],
|
|
) -> Vec<(FeatureVector, Vec<f64>)> {
|
|
features
|
|
.iter()
|
|
.zip(bars.iter())
|
|
.map(|(feat, bar)| {
|
|
(*feat, vec![bar.open, bar.high, bar.low, bar.close])
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Train a DQN model on a single walk-forward fold using `DQNTrainer`.
|
|
///
|
|
/// Delegates all GPU optimizations (mixed precision, dynamic batching, gradient
|
|
/// accumulation, Rainbow DQN components) to the production trainer. The walk-forward
|
|
/// fold structure stays in this binary; only per-fold training delegates to `DQNTrainer`.
|
|
///
|
|
/// Returns the best validation loss achieved.
|
|
#[allow(clippy::cognitive_complexity, clippy::too_many_arguments)]
|
|
fn train_dqn_fold(
|
|
fold: usize,
|
|
train_features: &[FeatureVector],
|
|
val_features: &[FeatureVector],
|
|
train_bars: &[OHLCVBar],
|
|
val_bars: &[OHLCVBar],
|
|
args: &Args,
|
|
output_dir: &Path,
|
|
hp: &Option<Value>,
|
|
pre_uploaded_gpu_data: Option<DqnGpuData>,
|
|
checkpoint_prefix: &str,
|
|
) -> Result<f64> {
|
|
// Caller must pre-align bars to features (skip warmup period before calling).
|
|
debug_assert_eq!(
|
|
train_bars.len(),
|
|
train_features.len(),
|
|
"train_bars ({}) must be pre-aligned to train_features ({})",
|
|
train_bars.len(),
|
|
train_features.len(),
|
|
);
|
|
debug_assert_eq!(
|
|
val_bars.len(),
|
|
val_features.len(),
|
|
"val_bars ({}) must be pre-aligned to val_features ({})",
|
|
val_bars.len(),
|
|
val_features.len(),
|
|
);
|
|
info!(" [DQN] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len());
|
|
|
|
// Compute average spread slippage in bps from training bar prices.
|
|
// Same pattern as evaluate_baseline: total_cost = commission + spread.
|
|
let avg_price = if train_bars.is_empty() {
|
|
0.0
|
|
} else {
|
|
train_bars.iter().map(|b| b.close).sum::<f64>() / train_bars.len() as f64
|
|
};
|
|
let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks);
|
|
let total_cost_bps = args.tx_cost_bps + avg_spread_bps;
|
|
info!(" [DQN] Fold {} total tx cost: {:.2} bps (commission {:.1} + spread {:.2})",
|
|
fold, total_cost_bps, args.tx_cost_bps, avg_spread_bps);
|
|
|
|
// Build DQNHyperparameters -- hyperopt JSON overrides defaults, CLI args override both.
|
|
// `..DQNHyperparameters::default()` enables all Rainbow components (PER, dueling,
|
|
// distributional C51, noisy nets, gradient accumulation, mixed precision auto-detect).
|
|
// transaction_cost_multiplier scales internal reward cost by total trading friction.
|
|
//
|
|
// hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware
|
|
// scaling so that large GPUs (L40S 48GB, H100 80GB) get proportionally wider
|
|
// networks instead of the tiny CPU defaults.
|
|
// Load GPU profile for VRAM-aware defaults (replaces scattered if/else chains)
|
|
let gpu_profile = GpuProfile::load();
|
|
|
|
let hp_hidden_base = hp_usize(hp, "hidden_dim_base");
|
|
let dqn_hidden_base = hp_hidden_base.or_else(|| {
|
|
info!(" [DQN] hidden_dim_base: {} (from GPU profile)", gpu_profile.training.hidden_dim_base);
|
|
Some(gpu_profile.training.hidden_dim_base)
|
|
});
|
|
|
|
// Noisy nets are always on (mandatory feature). Epsilon-greedy must be
|
|
// low — otherwise epsilon=1.0 forces pure random actions for the entire
|
|
// run, preventing the model from ever learning.
|
|
let epsilon_start = hp_f64(hp, "epsilon_start").unwrap_or(0.05);
|
|
|
|
let mut hyperparams = DQNHyperparameters {
|
|
learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
|
|
batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
|
|
gamma: hp_f64(hp, "gamma").unwrap_or(0.95),
|
|
epsilon_start,
|
|
epsilon_end: hp_f64(hp, "epsilon_end").unwrap_or(0.01),
|
|
epsilon_decay: hp_f64(hp, "epsilon_decay").unwrap_or(0.995),
|
|
buffer_size: hp_usize(hp, "buffer_size").unwrap_or(gpu_profile.training.buffer_size),
|
|
min_replay_size: hp_usize(hp, "min_replay_size").unwrap_or(1000),
|
|
epochs: args.epochs,
|
|
checkpoint_frequency: 10,
|
|
hidden_dim_base: dqn_hidden_base,
|
|
warmup_steps: 0,
|
|
early_stopping_enabled: true,
|
|
transaction_cost_multiplier: total_cost_bps,
|
|
// Pass through hyperopt architectural choices
|
|
per_alpha: hp_f64(hp, "per_alpha").unwrap_or(0.6),
|
|
per_beta_start: hp_f64(hp, "per_beta_start").unwrap_or(0.4),
|
|
dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128),
|
|
n_steps: hp_usize(hp, "n_steps").unwrap_or(3),
|
|
tau: hp_f64(hp, "tau").unwrap_or(0.005),
|
|
num_atoms: hp_usize(hp, "num_atoms")
|
|
.unwrap_or(gpu_profile.training.num_atoms),
|
|
v_min: hp_f64(hp, "v_min").unwrap_or_else(|| {
|
|
let gamma = hp_f64(hp, "gamma").unwrap_or(0.95);
|
|
-(10.0_f64 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0)
|
|
}),
|
|
v_max: hp_f64(hp, "v_max").unwrap_or_else(|| {
|
|
let gamma = hp_f64(hp, "gamma").unwrap_or(0.95);
|
|
(10.0_f64 / (1.0 - gamma) * 1.2).clamp(20.0, 300.0)
|
|
}),
|
|
noisy_sigma_init: hp_f64(hp, "noisy_sigma_init").unwrap_or(0.5),
|
|
num_quantiles: hp_usize(hp, "num_quantiles").unwrap_or(64),
|
|
noisy_epsilon_floor: hp_f64(hp, "noisy_epsilon_floor").unwrap_or(0.05).into(),
|
|
// Reward-shaping & environment params — critical for trade generation
|
|
hold_penalty_weight: hp_f64(hp, "hold_penalty_weight").unwrap_or(0.01),
|
|
max_position_absolute: hp_f64(hp, "max_position_absolute").unwrap_or(2.0),
|
|
huber_delta: hp_f64(hp, "huber_delta").unwrap_or(10.0),
|
|
entropy_coefficient: hp_f64(hp, "entropy_coefficient").unwrap_or(0.01),
|
|
curiosity_weight: hp_f64(hp, "curiosity_weight").unwrap_or(0.1),
|
|
weight_decay: hp_f64(hp, "weight_decay").unwrap_or(1e-4),
|
|
kelly_fractional: hp_f64(hp, "kelly_fractional").unwrap_or(0.5),
|
|
kelly_max_fraction: hp_f64(hp, "kelly_max_fraction").unwrap_or(0.25),
|
|
mbp10_data_dir: args.mbp10_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()).unwrap_or_else(|| "test_data/futures-baseline-mbp10".to_string()),
|
|
trades_data_dir: args.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()).unwrap_or_else(|| "test_data/futures-baseline-trades".to_string()),
|
|
offline_mode: args.offline,
|
|
dataset_path: args.dataset_path.as_ref().map(|p| p.to_string_lossy().into_owned()),
|
|
// GPU PER is mandatory on CUDA. VRAM fraction controls AutoReplaySizer.
|
|
// Small GPUs (RTX 3050 profile: buffer_size=5000 < 100K threshold) bypass
|
|
// AutoReplaySizer entirely, so VRAM fraction is irrelevant for them.
|
|
replay_buffer_vram_fraction: gpu_profile.training.replay_buffer_vram_fraction,
|
|
// GPU experience collection: n_episodes auto-scales from VRAM
|
|
gpu_timesteps_per_episode: gpu_profile.experience.gpu_timesteps_per_episode,
|
|
max_training_steps_per_epoch: args.max_steps_per_epoch,
|
|
..DQNHyperparameters::default()
|
|
};
|
|
|
|
// Load training profile (TOML) and apply to hyperparams.
|
|
// Profile values override the struct defaults above; CLI args re-applied below win.
|
|
let profile = ml::training_profile::DqnTrainingProfile::load(&args.training_profile);
|
|
profile.apply_to(&mut hyperparams);
|
|
// CLI args override profile: re-apply any arg that the user can set explicitly.
|
|
hyperparams.epochs = args.epochs;
|
|
// Only override batch_size from CLI/hyperopt if explicitly non-zero.
|
|
// batch_size=0 is the auto-compute sentinel — let the constructor handle it.
|
|
let hp_batch = hp_usize(hp, "batch_size").unwrap_or(args.batch_size);
|
|
if hp_batch > 0 {
|
|
hyperparams.batch_size = hp_batch;
|
|
}
|
|
hyperparams.learning_rate = hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate);
|
|
hyperparams.max_training_steps_per_epoch = args.max_steps_per_epoch;
|
|
hyperparams.initial_capital = args.initial_capital as f32;
|
|
if let Some(mhb) = args.min_hold_bars {
|
|
hyperparams.min_hold_bars = mhb;
|
|
}
|
|
|
|
// Create DQNTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing
|
|
let mut trainer = DQNTrainer::new(hyperparams)
|
|
.context("Failed to create DQNTrainer")?;
|
|
|
|
// Inject pre-uploaded GPU data from DoubleBufferedLoader (overlapped with previous fold)
|
|
if let Some(gpu_data) = pre_uploaded_gpu_data {
|
|
info!(" [DQN] Fold {} -- injecting pre-uploaded GPU data ({} bars), skipping lazy upload",
|
|
fold, gpu_data.num_bars);
|
|
trainer.set_gpu_data(gpu_data);
|
|
}
|
|
|
|
// Convert features + bars to trainer format
|
|
let training_data = features_to_trainer_format(train_features, train_bars);
|
|
let val_data = features_to_trainer_format(val_features, val_bars);
|
|
|
|
// Checkpoint callback: save best model to output directory
|
|
let output_dir_owned = output_dir.to_path_buf();
|
|
let prefix_owned = checkpoint_prefix.to_owned();
|
|
let checkpoint_callback = move |epoch: usize, data: Vec<u8>, is_best: bool| -> Result<String> {
|
|
let suffix = if is_best { "best" } else { &format!("epoch{}", epoch) };
|
|
let ckpt_path = output_dir_owned.join(format!("{}_fold{}_{}.safetensors", prefix_owned, fold, suffix));
|
|
// Atomic write: write to .tmp then rename (POSIX rename is atomic)
|
|
let tmp_path = ckpt_path.with_extension("safetensors.tmp");
|
|
std::fs::write(&tmp_path, &data)
|
|
.with_context(|| format!("Failed to write checkpoint tmp: {}", tmp_path.display()))?;
|
|
if let Err(e) = std::fs::rename(&tmp_path, &ckpt_path) {
|
|
drop(std::fs::remove_file(&tmp_path));
|
|
return Err(e).with_context(|| format!("Failed to rename checkpoint: {} -> {}", tmp_path.display(), ckpt_path.display()));
|
|
}
|
|
info!(" [DQN] Fold {} saved checkpoint: {} (prefix: {})", fold, ckpt_path.display(), prefix_owned);
|
|
Ok(ckpt_path.to_string_lossy().into_owned())
|
|
};
|
|
|
|
// Run the async training loop via a tokio runtime.
|
|
// The trainer handles epochs, early stopping, epsilon decay, validation, and
|
|
// all Rainbow DQN components internally.
|
|
let rt = tokio::runtime::Builder::new_current_thread()
|
|
.enable_all()
|
|
.build()
|
|
.context("Failed to create tokio runtime for DQN training")?;
|
|
|
|
let metrics = rt.block_on(
|
|
trainer.train_with_preloaded_data(training_data, val_data, checkpoint_callback)
|
|
).map_err(|e| {
|
|
error!(" [DQN] Fold {} training error chain: {:#}", fold, e);
|
|
e
|
|
}).context("DQNTrainer training failed")?;
|
|
|
|
info!(
|
|
" [DQN] Fold {} complete -- loss={:.6} epochs_trained={} converged={}",
|
|
fold, metrics.loss, metrics.epochs_trained, metrics.convergence_achieved
|
|
);
|
|
|
|
Ok(metrics.loss)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// PPO Training
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Train a PPO model on a single walk-forward fold using `PpoTrainer`.
|
|
///
|
|
/// Delegates all GPU optimizations (dynamic batch sizing, mixed precision when
|
|
/// available, GAE computation, trajectory collection, early stopping) to the
|
|
/// production trainer. The walk-forward fold structure stays in this binary;
|
|
/// only per-fold training delegates to `PpoTrainer`.
|
|
///
|
|
/// Returns the best validation loss achieved (approximated by final `value_loss`).
|
|
#[allow(clippy::cognitive_complexity)]
|
|
fn train_ppo_fold(
|
|
fold: usize,
|
|
train_features: &[FeatureVector],
|
|
val_features: &[FeatureVector],
|
|
train_bars: &[OHLCVBar],
|
|
_val_bars: &[OHLCVBar],
|
|
args: &Args,
|
|
output_dir: &Path,
|
|
hp: &Option<Value>,
|
|
) -> Result<f64> {
|
|
info!(" [PPO] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len());
|
|
|
|
let n_train = train_features.len();
|
|
if n_train < 2 {
|
|
warn!(" [PPO] Fold {} -- insufficient training features ({})", fold, n_train);
|
|
return Ok(f64::MAX);
|
|
}
|
|
|
|
// Compute average spread slippage in bps from training bar prices.
|
|
// Same pattern as evaluate_baseline: total_cost = commission + spread.
|
|
let avg_price = if train_bars.is_empty() {
|
|
0.0
|
|
} else {
|
|
train_bars.iter().map(|b| b.close).sum::<f64>() / train_bars.len() as f64
|
|
};
|
|
let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks);
|
|
let total_cost_bps = args.tx_cost_bps + avg_spread_bps;
|
|
info!(" [PPO] Fold {} total tx cost: {:.2} bps (commission {:.1} + spread {:.2})",
|
|
fold, total_cost_bps, args.tx_cost_bps, avg_spread_bps);
|
|
|
|
// Build PpoHyperparameters -- hyperopt JSON overrides defaults, CLI args override both.
|
|
// Start from conservative() baseline so all fields have sane values.
|
|
//
|
|
// hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware
|
|
// scaling so that large GPUs get proportionally wider networks. PPO uses the
|
|
// "ppo_policy" model type for base dim since the trainer internally scales
|
|
// the value network from the same base (value: [4*base, 3*base, 2*base, base, base/2]).
|
|
let hp_ppo_hidden_base = hp_usize(hp, "hidden_dim_base");
|
|
let ppo_hidden_base = hp_ppo_hidden_base.or_else(|| {
|
|
let profile = ml_core::gpu::profile::GpuProfile::load();
|
|
info!(" [PPO] hidden_dim_base: {} (from GPU profile)", profile.training.hidden_dim_base);
|
|
Some(profile.training.hidden_dim_base)
|
|
});
|
|
|
|
let hyperparams = PpoHyperparameters {
|
|
learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
|
|
actor_learning_rate: Some(hp_f64(hp, "policy_learning_rate").unwrap_or(args.learning_rate)),
|
|
critic_learning_rate: Some(hp_f64(hp, "value_learning_rate").unwrap_or(args.learning_rate * 3.0)),
|
|
batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size.max(64)),
|
|
gamma: hp_f64(hp, "gamma").unwrap_or(0.99),
|
|
clip_epsilon: hp_f64(hp, "clip_epsilon").unwrap_or(0.2) as f32,
|
|
vf_coef: hp_f64(hp, "value_loss_coeff").unwrap_or(0.5) as f32,
|
|
ent_coef: hp_f64(hp, "entropy_coeff").unwrap_or(0.01) as f32,
|
|
gae_lambda: hp_f64(hp, "gae_lambda").unwrap_or(0.95) as f32,
|
|
rollout_steps: hp_usize(hp, "rollout_steps").unwrap_or(
|
|
if args.max_steps_per_epoch > 0 { args.max_steps_per_epoch } else { 2048 }
|
|
),
|
|
minibatch_size: hp_usize(hp, "minibatch_size").unwrap_or(64),
|
|
epochs: args.epochs,
|
|
early_stopping_enabled: true,
|
|
transaction_cost_bps: total_cost_bps / 100.0, // total (commission + spread) in bps → pct for trainer
|
|
hidden_dim_base: ppo_hidden_base,
|
|
..PpoHyperparameters::conservative()
|
|
};
|
|
|
|
// Checkpoint directory for this fold
|
|
let fold_ckpt_dir = output_dir.join(format!("ppo_fold{}", fold));
|
|
|
|
// Create PpoTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing
|
|
let trainer = PpoTrainer::new(
|
|
hyperparams,
|
|
args.feature_dim,
|
|
&fold_ckpt_dir,
|
|
true, // use_gpu = true, trainer auto-detects availability
|
|
None, // num_envs: standard single-env collection
|
|
).context("Failed to create PpoTrainer")?;
|
|
|
|
// Convert features from FeatureVector to Vec<Vec<f32>> for the trainer
|
|
let market_data: Vec<Vec<f32>> = train_features
|
|
.iter()
|
|
.map(|feat| feat.iter().map(|&v| v as f32).collect())
|
|
.collect();
|
|
|
|
// Progress callback: log per-epoch metrics
|
|
let progress_callback = move |metrics: ml::trainers::ppo::PpoTrainingMetrics| {
|
|
info!(
|
|
" [PPO] Fold {} Epoch {}/{} -- policy_loss={:.6} value_loss={:.6} expl_var={:.4} mean_reward={:.4}",
|
|
fold,
|
|
metrics.epoch,
|
|
args.epochs,
|
|
metrics.policy_loss,
|
|
metrics.value_loss,
|
|
metrics.explained_variance,
|
|
metrics.mean_reward,
|
|
);
|
|
};
|
|
|
|
// Run the async training loop via a tokio runtime.
|
|
// PpoTrainer handles epochs, early stopping, trajectory collection, GAE,
|
|
// checkpointing, and all GPU optimizations internally.
|
|
let rt = tokio::runtime::Builder::new_current_thread()
|
|
.enable_all()
|
|
.build()
|
|
.context("Failed to create tokio runtime for PPO training")?;
|
|
|
|
let metrics = rt.block_on(
|
|
trainer.train(market_data, progress_callback)
|
|
).context("PpoTrainer training failed")?;
|
|
|
|
info!(
|
|
" [PPO] Fold {} complete -- value_loss={:.6} policy_loss={:.6} expl_var={:.4} epochs_trained={}",
|
|
fold, metrics.value_loss, metrics.policy_loss, metrics.explained_variance, metrics.epoch
|
|
);
|
|
|
|
Ok(metrics.value_loss as f64)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Training orchestration
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Container for per-model RL results collected during training.
|
|
struct RlTrainingResult {
|
|
model_name: String,
|
|
fold_results: Vec<(usize, f64)>,
|
|
total_epochs: usize,
|
|
}
|
|
|
|
/// Run the full walk-forward RL training pipeline.
|
|
///
|
|
/// Returns per-model results so `main()` can write completion markers.
|
|
#[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
|
|
fn run_training(args: &Args) -> Result<Vec<RlTrainingResult>> {
|
|
let train_dqn = args.model == "dqn" || args.model == "both";
|
|
let train_ppo = args.model == "ppo" || args.model == "both";
|
|
|
|
info!("=== Walk-Forward Baseline Training ===");
|
|
info!(" Model(s): {}", args.model);
|
|
info!(" Symbol: {}", args.symbol);
|
|
info!(" Epochs: {}", args.epochs);
|
|
info!(" Batch size: {}", args.batch_size);
|
|
info!(" Data dir: {}", args.data_dir.display());
|
|
info!(" Output dir: {}", args.output_dir.display());
|
|
info!(" Feature dim: {}", args.feature_dim);
|
|
info!(" Num actions: {}", args.num_actions);
|
|
info!(" Learning rate: {:.1e}", args.learning_rate);
|
|
info!(" Tx cost: {:.1} bps commission + {:.1} tick spread (tick_size={:.4})",
|
|
args.tx_cost_bps, args.spread_ticks, args.tick_size);
|
|
info!(" Patience: {}", args.patience);
|
|
if let Some(ref hp_path) = args.hyperopt_params {
|
|
info!(" Hyperopt params: {}", hp_path.display());
|
|
}
|
|
if args.ensemble_top_k > 1 {
|
|
info!(" Ensemble top-K: {} (training multiple models per fold)", args.ensemble_top_k);
|
|
}
|
|
|
|
// 1. Try fxcache first, fall back to DBN loading + feature extraction
|
|
info!("Step 1/5: Loading data...");
|
|
let data_load_start = std::time::Instant::now();
|
|
|
|
// Auto-discover fxcache: env var > sibling feature-cache/ dir
|
|
let fxcache_dir = std::env::var("FOXHUNT_FEATURE_CACHE_DIR").ok().map(PathBuf::from)
|
|
.or_else(|| {
|
|
let symbol_dir = args.data_dir.join(&args.symbol);
|
|
let mut dir = symbol_dir.as_path();
|
|
loop {
|
|
if let Some(parent) = dir.parent() {
|
|
let candidate = parent.join("feature-cache");
|
|
if candidate.exists() { return Some(candidate); }
|
|
if parent == dir { break; }
|
|
dir = parent;
|
|
} else { break; }
|
|
}
|
|
None
|
|
});
|
|
|
|
// Try loading from fxcache
|
|
let fxcache_data = fxcache_dir.and_then(|cache_dir| {
|
|
let symbol_dir = args.data_dir.join(&args.symbol);
|
|
let default_mbp10 = PathBuf::from("test_data/futures-baseline-mbp10");
|
|
let default_trades = PathBuf::from("test_data/futures-baseline-trades");
|
|
let mbp10 = args.mbp10_data_dir.as_deref().unwrap_or(&default_mbp10);
|
|
let trades = args.trades_data_dir.as_deref().unwrap_or(&default_trades);
|
|
let mbp10: Option<&Path> = if mbp10.exists() { Some(mbp10) } else { None };
|
|
let trades: Option<&Path> = if trades.exists() { Some(trades) } else { None };
|
|
|
|
let key_hex = ml::feature_cache::calculate_dbn_cache_key_full(
|
|
&symbol_dir, mbp10, trades,
|
|
).ok()?;
|
|
let key: [u8; 32] = hex::decode(&key_hex).ok()?.try_into().ok()?;
|
|
let path = ml::fxcache::find_fxcache(&cache_dir, &key)?;
|
|
match ml::fxcache::load_fxcache(&path) {
|
|
Ok(data) => {
|
|
info!("fxcache hit: {} bars from {:?}", data.bar_count, path);
|
|
Some(data)
|
|
}
|
|
Err(e) => {
|
|
warn!("fxcache load failed: {e}");
|
|
None
|
|
}
|
|
}
|
|
});
|
|
|
|
let (bars, all_features) = if let Some(cached) = fxcache_data {
|
|
// Reconstruct bars from cached targets (raw_close at index 2) + real timestamps
|
|
let n = cached.bar_count;
|
|
let mut bars = Vec::with_capacity(n);
|
|
for i in 0..n {
|
|
let close = cached.targets[i][2]; // raw_close
|
|
bars.push(ml::features::extraction::OHLCVBar {
|
|
timestamp: chrono::DateTime::from_timestamp_nanos(cached.timestamps[i]),
|
|
open: close,
|
|
high: close,
|
|
low: close,
|
|
close,
|
|
volume: 0.0,
|
|
});
|
|
}
|
|
info!(" Loaded {} bars + features from fxcache in {:.1}s",
|
|
n, data_load_start.elapsed().as_secs_f64());
|
|
(bars, cached.features)
|
|
} else {
|
|
// Fall back to DBN loading
|
|
info!(" Loading OHLCV bars from DBN files...");
|
|
let bars = load_all_bars(&args.data_dir, &args.symbol)?;
|
|
if bars.is_empty() {
|
|
anyhow::bail!("No bars loaded from {}", args.data_dir.display());
|
|
}
|
|
info!(" Loaded {} bars ({} to {})",
|
|
bars.len(),
|
|
bars.first().map(|b| b.timestamp.to_string()).unwrap_or_default(),
|
|
bars.last().map(|b| b.timestamp.to_string()).unwrap_or_default(),
|
|
);
|
|
|
|
// 2. Extract features
|
|
info!(" Extracting {}-dimensional features...", args.feature_dim);
|
|
let all_features = extract_ml_features(&bars)
|
|
.context("Feature extraction failed")?;
|
|
info!(" Extracted {} feature vectors (warmup period consumed {} bars)",
|
|
all_features.len(),
|
|
bars.len().saturating_sub(all_features.len()),
|
|
);
|
|
// Trim bars to align with features (skip warmup)
|
|
let warmup_offset = bars.len().saturating_sub(all_features.len());
|
|
let bars = bars[warmup_offset..].to_vec();
|
|
(bars, all_features)
|
|
};
|
|
|
|
// Since features skip the warmup period, we need bars aligned to features.
|
|
// Features start at bar index warmup_offset (typically 50).
|
|
let warmup_offset = bars.len().saturating_sub(all_features.len());
|
|
let aligned_bars = bars.get(warmup_offset..).unwrap_or(&bars);
|
|
|
|
// 3. Generate walk-forward windows
|
|
info!("Step 3/5: Generating walk-forward windows...");
|
|
let wf_config = WalkForwardConfig {
|
|
initial_train_months: args.train_months,
|
|
val_months: args.val_months,
|
|
test_months: args.test_months,
|
|
step_months: args.step_months,
|
|
};
|
|
let windows = generate_walk_forward_windows(aligned_bars, &wf_config);
|
|
if windows.is_empty() {
|
|
anyhow::bail!(
|
|
"No walk-forward windows generated. Need at least {} months of data.",
|
|
wf_config.initial_train_months + wf_config.val_months + wf_config.test_months
|
|
);
|
|
}
|
|
info!(" Generated {} walk-forward folds", windows.len());
|
|
|
|
// Record data loading + feature extraction time
|
|
if train_dqn {
|
|
metrics::record_data_load("dqn", data_load_start.elapsed().as_secs_f64());
|
|
}
|
|
if train_ppo {
|
|
metrics::record_data_load("ppo", data_load_start.elapsed().as_secs_f64());
|
|
}
|
|
|
|
// Create output directory
|
|
std::fs::create_dir_all(&args.output_dir)
|
|
.with_context(|| format!("Failed to create output dir: {}", args.output_dir.display()))?;
|
|
|
|
// 4. Train each fold
|
|
info!("Step 4/5: Training models on each fold...");
|
|
let mut dqn_results: Vec<(usize, f64)> = Vec::new();
|
|
let mut ppo_results: Vec<(usize, f64)> = Vec::new();
|
|
|
|
let mut prefetched_data: Option<FoldData> = None;
|
|
// GPU double-buffer: pre-uploaded GPU data for the next DQN fold.
|
|
// While fold N trains on GPU, fold N+1's data is uploaded in background.
|
|
let mut dqn_gpu_staged: Option<DqnGpuData> = None;
|
|
|
|
for (fold_idx, window) in windows.iter().enumerate() {
|
|
info!("--- Fold {} ---", window.fold);
|
|
info!(
|
|
" Train: {} bars (up to {}), Val: {} bars (up to {}), Test: {} bars (up to {})",
|
|
window.train.len(),
|
|
window.train_end,
|
|
window.val.len(),
|
|
window.val_end,
|
|
window.test.len(),
|
|
window.test_end,
|
|
);
|
|
|
|
// Get fold data: from prefetcher (fold 1+) or prepare fresh (fold 0)
|
|
let fold_data = if let Some(data) = prefetched_data.take() {
|
|
info!(" Using prefetched data for fold {}", window.fold);
|
|
Some(data)
|
|
} else {
|
|
prepare_fold_data(window, &args.output_dir)
|
|
};
|
|
|
|
// Kick off prefetch for NEXT fold on background thread
|
|
let prefetch_rx = if fold_idx + 1 < windows.len() {
|
|
let next_window = windows.get(fold_idx + 1).cloned();
|
|
if let Some(nw) = next_window {
|
|
let out_dir = args.output_dir.clone();
|
|
let (tx, rx) = std::sync::mpsc::channel::<Option<FoldData>>();
|
|
let _handle = std::thread::Builder::new()
|
|
.name("fold-prefetch".into())
|
|
.spawn(move || {
|
|
info!(" [Prefetch] Loading fold {} data on background thread", nw.fold);
|
|
let result = prepare_fold_data(&nw, &out_dir);
|
|
drop(tx.send(result));
|
|
});
|
|
Some(rx)
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let Some((train_norm, val_norm, train_bars_aligned, val_bars_aligned)) = fold_data else {
|
|
// Collect prefetch result even if this fold failed
|
|
if let Some(rx) = prefetch_rx {
|
|
if let Ok(data) = rx.recv() {
|
|
prefetched_data = data;
|
|
}
|
|
}
|
|
continue;
|
|
};
|
|
|
|
// Train DQN
|
|
if train_dqn {
|
|
let fold_str = fold_idx.to_string();
|
|
let fold_start = std::time::Instant::now();
|
|
|
|
// Take pre-uploaded GPU data from previous fold's background upload
|
|
let mut gpu_data_for_fold = dqn_gpu_staged.take();
|
|
|
|
if args.ensemble_top_k > 1 && args.hyperopt_params.is_some() {
|
|
// Ensemble mode: train one model per top-K hyperopt param set
|
|
let param_sets = load_top_k_params(
|
|
&args.hyperopt_params,
|
|
"dqn",
|
|
args.ensemble_top_k,
|
|
);
|
|
for (k, hp) in param_sets.iter().enumerate() {
|
|
info!(
|
|
" [DQN] Training ensemble member {}/{} on fold {}",
|
|
k + 1,
|
|
param_sets.len(),
|
|
window.fold
|
|
);
|
|
let prefix = format!("dqn_ensemble_{}", k);
|
|
// Only the first ensemble member uses the pre-uploaded GPU data
|
|
let gpu_data = if k == 0 { gpu_data_for_fold.take() } else { None };
|
|
match train_dqn_fold(
|
|
window.fold,
|
|
&train_norm,
|
|
&val_norm,
|
|
&train_bars_aligned,
|
|
&val_bars_aligned,
|
|
args,
|
|
&args.output_dir,
|
|
hp,
|
|
gpu_data,
|
|
&prefix,
|
|
) {
|
|
Ok(best_loss) => {
|
|
info!(
|
|
" [DQN] Ensemble member {} fold {} best_loss={:.6}",
|
|
k, window.fold, best_loss
|
|
);
|
|
// Record the best ensemble member's loss as the fold result
|
|
if k == 0 {
|
|
let elapsed = fold_start.elapsed().as_secs_f64();
|
|
metrics::set_epoch("dqn", &fold_str, fold_idx as f64);
|
|
metrics::set_epoch_loss("dqn", &fold_str, best_loss);
|
|
metrics::set_validation_loss("dqn", &fold_str, best_loss);
|
|
metrics::set_iteration_seconds("dqn", &fold_str, elapsed);
|
|
dqn_results.push((window.fold, best_loss));
|
|
}
|
|
}
|
|
Err(e) => {
|
|
error!(
|
|
" [DQN] Ensemble member {} fold {} failed: {}",
|
|
k, window.fold, e
|
|
);
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
// Single-model mode (default)
|
|
let hp = load_hyperopt_params(&args.hyperopt_params, "dqn");
|
|
match train_dqn_fold(
|
|
window.fold,
|
|
&train_norm,
|
|
&val_norm,
|
|
&train_bars_aligned,
|
|
&val_bars_aligned,
|
|
args,
|
|
&args.output_dir,
|
|
&hp,
|
|
gpu_data_for_fold,
|
|
"dqn",
|
|
) {
|
|
Ok(best_loss) => {
|
|
let elapsed = fold_start.elapsed().as_secs_f64();
|
|
metrics::set_epoch("dqn", &fold_str, fold_idx as f64);
|
|
metrics::set_epoch_loss("dqn", &fold_str, best_loss);
|
|
metrics::set_validation_loss("dqn", &fold_str, best_loss);
|
|
metrics::set_iteration_seconds("dqn", &fold_str, elapsed);
|
|
dqn_results.push((window.fold, best_loss));
|
|
}
|
|
Err(e) => {
|
|
error!(" [DQN] Fold {} failed: {:#}", window.fold, e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Train PPO
|
|
if train_ppo {
|
|
let hp = load_hyperopt_params(&args.hyperopt_params, "ppo");
|
|
let fold_str = fold_idx.to_string();
|
|
let fold_start = std::time::Instant::now();
|
|
match train_ppo_fold(
|
|
window.fold,
|
|
&train_norm,
|
|
&val_norm,
|
|
&train_bars_aligned,
|
|
&val_bars_aligned,
|
|
args,
|
|
&args.output_dir,
|
|
&hp,
|
|
) {
|
|
Ok(best_loss) => {
|
|
let elapsed = fold_start.elapsed().as_secs_f64();
|
|
metrics::set_epoch("ppo", &fold_str, fold_idx as f64);
|
|
metrics::set_epoch_loss("ppo", &fold_str, best_loss);
|
|
metrics::set_validation_loss("ppo", &fold_str, best_loss);
|
|
metrics::set_iteration_seconds("ppo", &fold_str, elapsed);
|
|
ppo_results.push((window.fold, best_loss));
|
|
}
|
|
Err(e) => {
|
|
error!(" [PPO] Fold {} failed: {:#}", window.fold, e);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Collect prefetch result for next iteration
|
|
if let Some(rx) = prefetch_rx {
|
|
match rx.recv() {
|
|
Ok(Some(data)) => {
|
|
info!(" Prefetched fold {} data ready ({} train features)",
|
|
windows.get(fold_idx + 1).map_or(0, |w| w.fold), data.0.len());
|
|
|
|
// GPU double-buffer: pre-upload next fold's DQN data to GPU now
|
|
// while the result is fresh and GPU is idle between folds.
|
|
if train_dqn {
|
|
let next_train_data = features_to_trainer_format(&data.0, &data.2);
|
|
let ctx = cudarc::driver::CudaContext::new(0);
|
|
if let Ok(ctx) = ctx {
|
|
let stream = ctx.new_stream();
|
|
if let Ok(stream) = stream {
|
|
let stream = Arc::new(stream);
|
|
match DqnGpuData::upload(&next_train_data, &stream) {
|
|
Ok(gpu_data) => {
|
|
info!(" [DoubleBuffer] Pre-uploaded {} bars to GPU for next fold ({:.1} MB)",
|
|
gpu_data.num_bars,
|
|
gpu_data.vram_bytes() as f64 / 1_048_576.0);
|
|
dqn_gpu_staged = Some(gpu_data);
|
|
}
|
|
Err(e) => {
|
|
warn!(" [DoubleBuffer] GPU pre-upload failed, will upload lazily: {}", e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
prefetched_data = Some(data);
|
|
}
|
|
Ok(None) => {
|
|
warn!(" Prefetch for next fold returned None (will load synchronously)");
|
|
prefetched_data = None;
|
|
}
|
|
Err(_) => {
|
|
warn!(" Prefetch thread disconnected (will load synchronously)");
|
|
prefetched_data = None;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 5. Summary
|
|
info!("Step 5/5: Training Summary");
|
|
info!(" ===================================");
|
|
|
|
let mut all_results = Vec::new();
|
|
|
|
if train_dqn {
|
|
info!(" DQN Results ({} folds):", dqn_results.len());
|
|
for (fold, loss) in &dqn_results {
|
|
info!(" Fold {}: best_val_metric = {:.6}", fold, loss);
|
|
}
|
|
if !dqn_results.is_empty() {
|
|
let avg: f64 = dqn_results.iter().map(|(_, l)| l).sum::<f64>()
|
|
/ dqn_results.len() as f64;
|
|
info!(" Average: {:.6}", avg);
|
|
}
|
|
all_results.push(RlTrainingResult {
|
|
model_name: "dqn".to_owned(),
|
|
fold_results: dqn_results,
|
|
total_epochs: windows.len() * args.epochs,
|
|
});
|
|
}
|
|
if train_ppo {
|
|
info!(" PPO Results ({} folds):", ppo_results.len());
|
|
for (fold, loss) in &ppo_results {
|
|
info!(" Fold {}: best_val_metric = {:.6}", fold, loss);
|
|
}
|
|
if !ppo_results.is_empty() {
|
|
let avg: f64 = ppo_results.iter().map(|(_, l)| l).sum::<f64>()
|
|
/ ppo_results.len() as f64;
|
|
info!(" Average: {:.6}", avg);
|
|
}
|
|
all_results.push(RlTrainingResult {
|
|
model_name: "ppo".to_owned(),
|
|
fold_results: ppo_results,
|
|
total_epochs: windows.len() * args.epochs,
|
|
});
|
|
}
|
|
info!(" Checkpoints saved to: {}", args.output_dir.display());
|
|
info!(" ===================================");
|
|
|
|
Ok(all_results)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Main
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn main() -> Result<()> {
|
|
// Initialize tracing with optional OTLP export to Tempo
|
|
let otlp_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok();
|
|
if let Err(e) = common::observability::init_observability(
|
|
"train_baseline_rl",
|
|
otlp_endpoint.as_deref(),
|
|
) {
|
|
eprintln!("Observability init failed (non-fatal): {e}");
|
|
}
|
|
|
|
// Pre-allocate CUBLAS workspace for deterministic + faster tensor core ops.
|
|
// Enable TF32 for all FP32 matmuls — ~8x throughput on H100 tensor cores.
|
|
// SAFETY: called once at startup before any multi-threading or CUDA work begins.
|
|
#[allow(unsafe_code)]
|
|
unsafe {
|
|
std::env::set_var("CUBLAS_WORKSPACE_CONFIG", ":4096:8");
|
|
std::env::set_var("NVIDIA_TF32_OVERRIDE", "1");
|
|
}
|
|
|
|
metrics::init();
|
|
metrics_server::start_metrics_server(9094);
|
|
common::metrics::questdb_sink::init(None);
|
|
metrics::set_active_workers(1.0);
|
|
|
|
let args = Args::parse();
|
|
|
|
// Ensure output directory exists before training so markers can always be written.
|
|
if let Err(e) = std::fs::create_dir_all(&args.output_dir) {
|
|
error!("Failed to create output dir {}: {}", args.output_dir.display(), e);
|
|
}
|
|
|
|
let result = run_training(&args);
|
|
metrics::set_active_workers(0.0);
|
|
|
|
// Push final metrics to pushgateway so they persist after pod termination
|
|
if let Err(e) = metrics_server::push_to_gateway(None, "train_baseline_rl") {
|
|
tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}");
|
|
}
|
|
common::metrics::questdb_sink::flush();
|
|
|
|
match result {
|
|
Ok(results) => {
|
|
for training_result in &results {
|
|
let best_val = training_result
|
|
.fold_results
|
|
.iter()
|
|
.map(|(_, loss)| *loss)
|
|
.fold(f64::MAX, f64::min);
|
|
|
|
let metrics = CompletionMetrics {
|
|
model: training_result.model_name.clone(),
|
|
symbol: args.symbol.clone(),
|
|
best_val_loss: (best_val < f64::MAX).then_some(best_val),
|
|
sharpe_ratio: None,
|
|
epochs_completed: training_result.total_epochs,
|
|
folds_completed: training_result.fold_results.len(),
|
|
};
|
|
write_success_marker(&args.output_dir, &metrics);
|
|
}
|
|
Ok(())
|
|
}
|
|
Err(e) => {
|
|
let msg = format!("{:#}", e);
|
|
error!("Training failed: {}", msg);
|
|
write_failure_marker(&args.output_dir, &msg);
|
|
Err(e)
|
|
}
|
|
}
|
|
}
|