perf: PPO zero-copy fold loop — train_from_slices + delete train_ppo_fold

Add PpoTrainer::train_from_slices(&[[f64; 42]]) that converts to
Vec<Vec<f32>> 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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-02 22:54:00 +02:00
parent 8d5af977f7
commit 3278cec29b
2 changed files with 83 additions and 164 deletions

View File

@@ -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<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: 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<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
// ---------------------------------------------------------------------------
@@ -968,47 +840,78 @@ fn run_training(args: &Args) -> Result<Vec<RlTrainingResult>> {
}
}
// 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<OHLCVBar> = 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<OHLCVBar> = 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<f64> = 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::<f64>() / 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),
}
}
}

View File

@@ -438,6 +438,22 @@ impl PpoTrainer {
self.train_gpu(market_data, progress_callback).await
}
/// Train from contiguous feature slices (fxcache format).
/// Converts to Vec<Vec<f32>> internally — PPO's train() API requires it.
pub async fn train_from_slices<F>(
&self,
features: &[[f64; 42]],
progress_callback: F,
) -> Result<PpoTrainingMetrics, MLError>
where
F: FnMut(PpoTrainingMetrics) + Send,
{
let market_data: Vec<Vec<f32>> = features.iter()
.map(|f| f.iter().map(|&v| v as f32).collect())
.collect();
self.train(market_data, progress_callback).await
}
async fn train_gpu<F>(
&self,
market_data: Vec<Vec<f32>>,