diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 93d2e19ba..4e4327e23 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -29,7 +29,7 @@ use rand::Rng; use serde_json::Value; use tracing::{error, info, warn}; -use ml::dqn::{DQNConfig, Experience, DQN}; +use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer}; #[allow(unreachable_pub)] mod baseline_common; @@ -218,10 +218,31 @@ fn compute_reward(close_current: f64, close_next: f64, action_idx: u8, total_cos // DQN Training // --------------------------------------------------------------------------- -/// Train a DQN model on a single walk-forward fold. +/// Convert pre-aligned features and bars into the format expected by +/// `DQNTrainer::train_with_preloaded_data`: `Vec<([f64; 51], Vec)>`. +/// +/// The `Vec` target contains OHLCV prices for the trainer's internal +/// reward computation and portfolio simulation. +fn features_to_trainer_format( + features: &[[f64; 51]], + bars: &[OHLCVBar], +) -> Vec<([f64; 51], 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)] fn train_dqn_fold( fold: usize, train_features: &[[f64; 51]], @@ -233,7 +254,6 @@ fn train_dqn_fold( hp: &Option, ) -> Result { // Caller must pre-align bars to features (skip warmup period before calling). - // If these differ, bar-index reward computation will silently use wrong prices. debug_assert_eq!( train_bars.len(), train_features.len(), @@ -250,254 +270,64 @@ fn train_dqn_fold( ); info!(" [DQN] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len()); - // Configure DQN -- hyperopt params override compiled defaults, CLI args override both - let config = DQNConfig { - state_dim: args.feature_dim, - num_actions: args.num_actions, - hidden_dims: { - let base = hp_usize(hp, "hidden_dim_base").unwrap_or(128); - if base > 128 { - vec![base, base / 2, base / 4] // GPU-scaled: 3-layer network - } else { - vec![128, 64] // Default small network - } - }, + // 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). + let hyperparams = DQNHyperparameters { learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate), - gamma: hp_f64(hp, "gamma").unwrap_or(0.95) as f32, + batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size), + gamma: hp_f64(hp, "gamma").unwrap_or(0.95), epsilon_start: 1.0, epsilon_end: 0.05, epsilon_decay: 0.995, - replay_buffer_capacity: hp_usize(hp, "buffer_size").unwrap_or(50_000), - batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size), + buffer_size: hp_usize(hp, "buffer_size").unwrap_or(50_000), min_replay_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size), - target_update_freq: 200, + epochs: args.epochs, + checkpoint_frequency: 10, + hidden_dim_base: hp_usize(hp, "hidden_dim_base"), warmup_steps: 0, - use_double_dqn: true, - use_huber_loss: true, - use_per: false, - use_dueling: false, - use_distributional: false, - use_noisy_nets: false, - use_cql: false, - use_iqn: false, - use_cvar_action_selection: false, - ..DQNConfig::default() + // Early stopping is managed by DQNTrainer internally + early_stopping_enabled: true, + ..DQNHyperparameters::default() }; - let mut dqn = DQN::new(config).context("Failed to create DQN model")?; + // Create DQNTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing + let mut trainer = DQNTrainer::new(hyperparams) + .context("Failed to create DQNTrainer")?; - let mut best_val_loss = f64::MAX; - let mut epochs_without_improvement = 0_usize; - let mut rng = rand::thread_rng(); - let fold_str = fold.to_string(); + // 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); - for epoch in 0..args.epochs { - let epoch_start = std::time::Instant::now(); - // --- Training pass --- - let mut epoch_loss = 0.0_f64; - let mut epoch_steps = 0_usize; + // Checkpoint callback: save best model to output directory + let output_dir_owned = output_dir.to_path_buf(); + 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!("dqn_fold{}_{}.safetensors", fold, suffix)); + std::fs::write(&ckpt_path, &data) + .with_context(|| format!("Failed to write checkpoint: {}", ckpt_path.display()))?; + info!(" [DQN] Fold {} saved checkpoint: {}", fold, ckpt_path.display()); + Ok(ckpt_path.to_string_lossy().into_owned()) + }; - // Process training data sequentially - let n_train = train_features.len(); - if n_train < 2 { - warn!(" [DQN] Fold {} -- insufficient training features ({})", fold, n_train); - break; - } + // 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 step_limit = if args.max_steps_per_epoch > 0 { - args.max_steps_per_epoch.min(n_train.saturating_sub(1)) - } else { - n_train.saturating_sub(1) - }; - // Features skip warmup period, so feature[i] corresponds to - // train_bars[i + warmup_offset]. We must use the same offset for - // reward computation to avoid misalignment. - let warmup_offset = train_bars.len().saturating_sub(n_train); - for i in 0..step_limit { - let Some(state_f64) = train_features.get(i) else { - continue; - }; - let Some(next_f64) = train_features.get(i + 1) else { - continue; - }; + let metrics = rt.block_on( + trainer.train_with_preloaded_data(training_data, val_data, checkpoint_callback) + ).context("DQNTrainer training failed")?; - let state: Vec = state_f64.iter().map(|&v| v as f32).collect(); - let next_state: Vec = next_f64.iter().map(|&v| v as f32).collect(); - - // Select action with epsilon-greedy - let action = dqn.select_action(&state) - .map(|fa| fa.to_index().min(args.num_actions.saturating_sub(1)) as u8) - .unwrap_or_else(|_| rng.gen_range(0..args.num_actions as u8)); - - // Compute reward from bar prices (offset by warmup to align with features) - let bar_idx = i + warmup_offset; - let close_cur = train_bars.get(bar_idx).map(|b| b.close).unwrap_or(0.0); - let close_next = train_bars.get(bar_idx + 1).map(|b| b.close).unwrap_or(close_cur); - let total_cost = args.tx_cost_bps + spread_cost_bps(close_cur, args.tick_size, args.spread_ticks); - let reward = compute_reward(close_cur, close_next, action, total_cost); - let done = i + 2 >= n_train; - - // Store experience - let exp = Experience::new(state, action, reward, next_state, done); - if let Err(e) = dqn.store_experience(exp) { - // Non-fatal: buffer may not be ready - if epoch == 0 && i < 5 { - info!(" [DQN] store_experience: {}", e); - } - } - - // Train step (returns (loss, grad_norm)) - if let Ok((loss, _grad_norm)) = dqn.train_step(None) { - if loss.is_nan() || loss.is_infinite() { - metrics::record_nan_detected("dqn", &fold_str); - } - epoch_loss += loss as f64; - epoch_steps += 1; - } - } - - let avg_train_loss = if epoch_steps > 0 { - epoch_loss / epoch_steps as f64 - } else { - f64::MAX - }; - - // --- Validation pass --- - let val_loss = evaluate_dqn_validation(&mut dqn, val_features, val_bars, args); - - // Prometheus metrics - let epoch_elapsed = epoch_start.elapsed(); - metrics::set_epoch("dqn", &fold_str, (epoch + 1) as f64); - metrics::set_epoch_loss("dqn", &fold_str, avg_train_loss); - metrics::set_validation_loss("dqn", &fold_str, val_loss); - metrics::set_batches_processed("dqn", &fold_str, epoch_steps as f64); - metrics::set_iteration_seconds("dqn", &fold_str, epoch_elapsed.as_secs_f64()); - if epoch_elapsed.as_secs_f64() > 0.001 { - metrics::set_batches_per_second( - "dqn", - &fold_str, - epoch_steps as f64 / epoch_elapsed.as_secs_f64(), - ); - } - - // Decay epsilon at epoch level - let current_eps = dqn.get_epsilon(); - let new_eps = (current_eps * 0.995_f32).max(0.05); - dqn.set_epsilon(new_eps as f64); - - info!( - " [DQN] Fold {} Epoch {}/{} -- train_loss={:.6} val_loss={:.6} eps={:.4}", - fold, - epoch + 1, - args.epochs, - avg_train_loss, - val_loss, - dqn.get_epsilon() - ); - - // Early stopping check - if val_loss < best_val_loss { - best_val_loss = val_loss; - epochs_without_improvement = 0; - - // Save best checkpoint - let ckpt_path = output_dir.join(format!("dqn_fold{}_best.safetensors", fold)); - let ckpt_start = std::time::Instant::now(); - if let Err(e) = dqn.get_q_network_vars().save(ckpt_path.to_string_lossy().as_ref()) { - metrics::record_checkpoint_failure("dqn", &fold_str); - warn!(" [DQN] Failed to save checkpoint: {}", e); - } else { - let ckpt_size = ckpt_path.metadata().map(|m| m.len() as f64).unwrap_or(0.0); - metrics::record_checkpoint_save( - "dqn", - &fold_str, - ckpt_start.elapsed().as_secs_f64(), - ckpt_size, - ); - info!(" [DQN] Saved best checkpoint: {}", ckpt_path.display()); - } - } else { - epochs_without_improvement += 1; - if epochs_without_improvement >= args.patience { - info!( - " [DQN] Early stopping at epoch {} (patience {} exhausted)", - epoch + 1, - args.patience - ); - break; - } - } - } - - Ok(best_val_loss) -} - -/// Evaluate DQN on validation features and return average loss proxy. -/// -/// Since `DQN.train_step` uses replay buffer internally, we estimate validation -/// performance via average absolute reward (lower is closer to zero = better). -fn evaluate_dqn_validation( - dqn: &mut DQN, - val_features: &[[f64; 51]], - val_bars: &[OHLCVBar], - args: &Args, -) -> f64 { - // Bars must be pre-aligned to features by the caller. - 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 {} complete -- loss={:.6} epochs_trained={} converged={}", + fold, metrics.loss, metrics.epochs_trained, metrics.convergence_achieved ); - // Run greedy inference on validation set and compute negative average reward. - // Lower values = better model (early stopping minimizes this). - let n = val_features.len(); - if n < 2 { - return f64::MAX; - } - - // Save epsilon and force greedy for validation. - // Note: select_action increments total_steps (side effect), but this is benign - // because warmup_steps=0 in our config, so the warmup branch is never taken. - let saved_epsilon = dqn.get_epsilon(); - dqn.set_epsilon(0.0); - - let warmup_offset = val_bars.len().saturating_sub(n); - let mut total_reward = 0.0_f64; - let mut count = 0_usize; - - for i in 0..n.saturating_sub(1) { - let state: Vec = match val_features.get(i) { - Some(f) => f.iter().map(|&v| v as f32).collect(), - None => break, - }; - - // Greedy action (epsilon=0) - let action = match dqn.select_action(&state) { - Ok(fa) => fa.to_index().min(2) as u8, - Err(_) => 2, // HOLD on error - }; - - let bar_idx = i + warmup_offset; - let close_cur = val_bars.get(bar_idx).map(|b| b.close).unwrap_or(0.0); - let close_next = val_bars.get(bar_idx + 1).map(|b| b.close).unwrap_or(close_cur); - let total_cost = args.tx_cost_bps + spread_cost_bps(close_cur, args.tick_size, args.spread_ticks); - let reward = compute_reward(close_cur, close_next, action, total_cost); - total_reward += reward as f64; - count += 1; - } - - // Restore epsilon - dqn.set_epsilon(saved_epsilon as f64); - - if count > 0 { - // Return negative mean reward: lower = better model - -(total_reward / count as f64) - } else { - f64::MAX - } + Ok(metrics.loss) } // ---------------------------------------------------------------------------