From 3278cec29b2c5feec5ea411bf54b1f29ba1fa141 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 22:54:00 +0200 Subject: [PATCH] =?UTF-8?q?perf:=20PPO=20zero-copy=20fold=20loop=20?= =?UTF-8?q?=E2=80=94=20train=5Ffrom=5Fslices=20+=20delete=20train=5Fppo=5F?= =?UTF-8?q?fold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add PpoTrainer::train_from_slices(&[[f64; 42]]) that converts to Vec> internally (PPO train() API requires ownership). Replace the PPO fold loop in train_baseline_rl.rs: eliminates the OHLCVBar construction shim and calls train_from_slices directly on fxcache feature slices. Delete the standalone train_ppo_fold function entirely. Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/examples/train_baseline_rl.rs | 231 +++++++----------------- crates/ml/src/trainers/ppo.rs | 16 ++ 2 files changed, 83 insertions(+), 164 deletions(-) diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 779c74681..862cb288f 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -103,7 +103,7 @@ use serde_json::Value; use tracing::{error, info, warn}; use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer}; -use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer}; +use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer, PpoTrainingMetrics}; #[allow(unreachable_pub)] mod baseline_common; @@ -495,134 +495,6 @@ fn train_dqn_fold( 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: ml_core::gpu::profile::GpuProfile::load().training.batch_size, - 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(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 // --------------------------------------------------------------------------- @@ -968,47 +840,78 @@ fn run_training(args: &Args) -> Result> { } } - // Train PPO (uses old path — constructs minimal bars from fxcache targets) + // Train PPO (zero-copy: pass fxcache feature slices directly) if train_ppo { - let hp = load_hyperopt_params(&args.hyperopt_params, "ppo"); + let hp_ppo = load_hyperopt_params(&args.hyperopt_params, "ppo"); let fold_str = fold_idx.to_string(); let fold_start = std::time::Instant::now(); - // PPO compatibility: construct minimal OHLCVBars from fxcache targets - let train_bars: Vec = fxcache.timestamps[range.train_start..range.train_end] - .iter().zip(fxcache.targets[range.train_start..range.train_end].iter()) - .map(|(&ts, t)| OHLCVBar { - timestamp: chrono::DateTime::from_timestamp_nanos(ts), - open: t[0], high: t[0], low: t[0], close: t[2], volume: 0.0, - }).collect(); - let val_bars: Vec = fxcache.timestamps[range.val_start..range.val_end] - .iter().zip(fxcache.targets[range.val_start..range.val_end].iter()) - .map(|(&ts, t)| OHLCVBar { - timestamp: chrono::DateTime::from_timestamp_nanos(ts), - open: t[0], high: t[0], low: t[0], close: t[2], volume: 0.0, - }).collect(); + // Compute per-fold tx cost from fxcache targets (raw close at index 2) + let train_closes: Vec = fxcache.targets[range.train_start..range.train_end] + .iter().map(|t| t[2]).collect(); + let avg_price = if train_closes.is_empty() { + 0.0 + } else { + train_closes.iter().sum::() / train_closes.len() as f64 + }; + let avg_spread_bps_ppo = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks); + let total_cost_bps_ppo = args.tx_cost_bps + avg_spread_bps_ppo; - match train_ppo_fold( - range.fold, - &train_norm, - &val_norm, - &train_bars, - &val_bars, - 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((range.fold, best_loss)); - } - Err(e) => { - error!(" [PPO] Fold {} failed: {:#}", range.fold, e); + // Build PpoHyperparameters inline (same as old train_ppo_fold) + let hp_ppo_hidden_base = hp_usize(&hp_ppo, "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 ppo_hp = PpoHyperparameters { + learning_rate: hp_f64(&hp_ppo, "learning_rate").unwrap_or(args.learning_rate), + actor_learning_rate: Some(hp_f64(&hp_ppo, "policy_learning_rate").unwrap_or(args.learning_rate)), + critic_learning_rate: Some(hp_f64(&hp_ppo, "value_learning_rate").unwrap_or(args.learning_rate * 3.0)), + batch_size: ml_core::gpu::profile::GpuProfile::load().training.batch_size, + gamma: hp_f64(&hp_ppo, "gamma").unwrap_or(0.99), + clip_epsilon: hp_f64(&hp_ppo, "clip_epsilon").unwrap_or(0.2) as f32, + vf_coef: hp_f64(&hp_ppo, "value_loss_coeff").unwrap_or(0.5) as f32, + ent_coef: hp_f64(&hp_ppo, "entropy_coeff").unwrap_or(0.01) as f32, + gae_lambda: hp_f64(&hp_ppo, "gae_lambda").unwrap_or(0.95) as f32, + rollout_steps: hp_usize(&hp_ppo, "rollout_steps").unwrap_or(2048), + minibatch_size: hp_usize(&hp_ppo, "minibatch_size").unwrap_or(64), + epochs: args.epochs, + early_stopping_enabled: true, + transaction_cost_bps: total_cost_bps_ppo / 100.0, + hidden_dim_base: ppo_hidden_base, + ..PpoHyperparameters::conservative() + }; + + let fold_ckpt_dir = args.output_dir.join(format!("ppo_fold{}", range.fold)); + let ppo_trainer = PpoTrainer::new( + ppo_hp, + args.feature_dim, + &fold_ckpt_dir, + true, + None, + ); + match ppo_trainer { + Ok(trainer) => { + let fold_idx_cap = range.fold; + let epochs_cap = args.epochs; + let progress_cb = move |m: PpoTrainingMetrics| { + info!("[PPO] Fold {} Epoch {}/{} -- value_loss={:.6}", fold_idx_cap, m.epoch, epochs_cap, m.value_loss); + }; + match rt.block_on(trainer.train_from_slices(&train_norm, progress_cb)) { + Ok(metrics) => { + let best_loss = metrics.value_loss as f64; + 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((range.fold, best_loss)); + } + Err(e) => error!("PPO fold {} failed: {:#}", range.fold, e), + } } + Err(e) => error!(" [PPO] Fold {} trainer init failed: {:#}", range.fold, e), } } } diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index 694652389..6932a2433 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -438,6 +438,22 @@ impl PpoTrainer { self.train_gpu(market_data, progress_callback).await } + /// Train from contiguous feature slices (fxcache format). + /// Converts to Vec> internally — PPO's train() API requires it. + pub async fn train_from_slices( + &self, + features: &[[f64; 42]], + progress_callback: F, + ) -> Result + where + F: FnMut(PpoTrainingMetrics) + Send, + { + let market_data: Vec> = features.iter() + .map(|f| f.iter().map(|&v| v as f32).collect()) + .collect(); + self.train(market_data, progress_callback).await + } + async fn train_gpu( &self, market_data: Vec>,