#![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 = 128)] 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, /// 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, /// 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, /// 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, /// Collect and save a dataset from the current policy, then exit. /// Use this to generate datasets for offline training. #[arg(long)] collect_dataset: Option, /// Disable Branching DQN (3-head: exposure, order, urgency). Enabled by default. #[arg(long)] no_branching: bool, /// Named training profile to load from config/training/.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, model_key: &str) -> Option { 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, model_key: &str, k: usize) -> Vec> { 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 = 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, key: &str) -> Option { params.as_ref()?.get(key)?.as_f64() } fn hp_usize(params: &Option, key: &str) -> Option { params.as_ref()?.get(key)?.as_u64().map(|v| v as usize) } fn hp_bool(params: &Option, key: &str) -> Option { 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, Vec, Vec, Vec); #[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 { 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)>`. /// /// The `Vec` 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)> { 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, pre_uploaded_gpu_data: Option, checkpoint_prefix: &str, ) -> Result { // 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::() / 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, "batch_size").unwrap_or(args.batch_size), 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()), trades_data_dir: args.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()), 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: profile-driven defaults replace VRAM if/else chains gpu_n_episodes: gpu_profile.experience.gpu_n_episodes, 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; hyperparams.batch_size = hp_usize(hp, "batch_size").unwrap_or(args.batch_size); hyperparams.learning_rate = hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate); hyperparams.max_training_steps_per_epoch = args.max_steps_per_epoch; // 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, is_best: bool| -> Result { 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, ) -> Result { 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::() / 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> for the trainer let market_data: Vec> = 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> { 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. Load all OHLCV bars from DBN files info!("Step 1/5: Loading OHLCV bars from DBN files..."); let data_load_start = std::time::Instant::now(); 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!("Step 2/5: 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()), ); // 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 = 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 = 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::>(); 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::() / 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::() / 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) } } }