foxhunt/crates/ml/examples/benchmark_ppo_optimization.rs

//! PPO Training Optimization Benchmark
//!
//! Compares baseline PPO trainer vs. optimized trainer with vectorized environments.
//!
//! # Expected Improvements
//!
//! - **Vectorized rollouts**: 2-4x faster rollout collection
//! - **Batch GAE**: 3-5x faster advantage computation
//! - **Parallel updates**: 1.5-2x faster network updates
//! - **Overall**: 2-3x end-to-end training speedup
//!
//! # Usage
//!
//! ```bash
//! # Benchmark both trainers (20 epochs each)
//! cargo run -p ml --example benchmark_ppo_optimization --release --features cuda
//!
//! # Quick test (5 epochs)
//! cargo run -p ml --example benchmark_ppo_optimization --release --features cuda -- --epochs 5
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use std::time::Instant;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::data_loaders::BarSamplingMethod;
use ml::features::extraction::{extract_ml_features, OHLCVBar};
use ml::real_data_loader::RealDataLoader;
use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer, PpoTrainingMetrics};

#[derive(Debug, Parser)]
#[command(name = "benchmark_ppo", about = "Benchmark PPO optimization")]
struct Opts {
    /// Number of training epochs (default: 20)
    #[arg(long, default_value = "20")]
    epochs: usize,

    /// Learning rate
    #[arg(long, default_value = "0.0003")]
    learning_rate: f64,

    /// Batch size (max 230 for RTX 3050 Ti 4GB)
    #[arg(long, default_value = "64")]
    batch_size: usize,

    /// Output directory for checkpoints
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// Data directory containing DBN files
    #[arg(long, default_value = "test_data/real/databento")]
    data_dir: String,

    /// Symbol to train on
    #[arg(long, default_value = "ZN.FUT")]
    symbol: String,

    /// Number of parallel environments for optimized trainer
    #[arg(long, default_value = "4")]
    num_envs: usize,

    /// Disable early stopping for fair comparison
    #[arg(long)]
    no_early_stopping: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    let opts = Opts::parse();

    // Setup logging
    let subscriber = FmtSubscriber::builder()
        .with_max_level(tracing::Level::INFO)
        .finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    info!("🚀 PPO Training Optimization Benchmark");
    info!("Configuration:");
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • Symbol: {}", opts.symbol);
    info!("  • Num envs (optimized): {}", opts.num_envs);
    info!("  • Early stopping: {}", !opts.no_early_stopping);

    // Load market data (once for both trainers)
    info!("\n📊 Loading market data...");
    let mut loader = RealDataLoader::new(&opts.data_dir);
    let bars = loader
        .load_symbol_data(&opts.symbol)
        .await
        .context(format!("Failed to load data for symbol: {}", opts.symbol))?;

    info!("✅ Loaded {} OHLCV bars for {}", bars.len(), opts.symbol);

    // Extract features
    info!("\n🔧 Extracting 54-dimensional features...");
    let ohlcv_bars: Vec<OHLCVBar> = bars
        .iter()
        .map(|bar| OHLCVBar {
            timestamp: bar.timestamp,
            open: bar.open,
            high: bar.high,
            low: bar.low,
            close: bar.close,
            volume: bar.volume,
        })
        .collect();

    let feature_vectors =
        extract_ml_features(&ohlcv_bars).context("Failed to extract 54-dimensional features")?;

    info!(
        "✅ Extracted {} feature vectors (dim=54)",
        feature_vectors.len()
    );

    let market_data: Vec<Vec<f32>> = feature_vectors
        .iter()
        .map(|fv| fv.iter().map(|&v| v as f32).collect())
        .collect();

    let state_dim = 54;

    // Shared hyperparameters
    let hyperparams = PpoHyperparameters {
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        gamma: 0.99,
        clip_epsilon: 0.2,
        vf_coef: 0.5,
        ent_coef: 0.01,
        gae_lambda: 0.95,
        rollout_steps: 2048,
        minibatch_size: opts.batch_size,
        epochs: opts.epochs,
        early_stopping_enabled: !opts.no_early_stopping,
        min_value_loss_improvement_pct: 2.0,
        min_explained_variance: 0.4,
        plateau_window: 30,
        min_epochs_before_stopping: 50,
    };

    // ========================================================================
    // BENCHMARK 1: Baseline PPO Trainer
    // ========================================================================
    info!("\n═══════════════════════════════════════════════════════════════");
    info!("📊 BENCHMARK 1: Baseline PPO Trainer");
    info!("═══════════════════════════════════════════════════════════════\n");

    let baseline_trainer = PpoTrainer::new(
        hyperparams.clone(),
        state_dim,
        format!("{}/baseline", opts.output_dir),
        true, // CUDA
        None, // Standard mode (no vectorization)
    )
    .context("Failed to create baseline trainer")?;

    let mut baseline_epochs_completed = 0;
    let baseline_callback = |metrics: PpoTrainingMetrics| {
        baseline_epochs_completed = metrics.epoch;
        if metrics.epoch % 5 == 0 {
            info!(
                "Baseline Epoch {}/{}: policy_loss={:.4}, value_loss={:.4}, expl_var={:.4}",
                metrics.epoch,
                hyperparams.epochs,
                metrics.policy_loss,
                metrics.value_loss,
                metrics.explained_variance
            );
        }
    };

    let baseline_start = Instant::now();
    let baseline_metrics = baseline_trainer
        .train(market_data.clone(), baseline_callback)
        .await
        .context("Baseline training failed")?;
    let baseline_duration = baseline_start.elapsed();

    info!("\n✅ Baseline Training Complete!");
    info!(
        "  • Duration: {:.2}s ({:.2} min)",
        baseline_duration.as_secs_f64(),
        baseline_duration.as_secs_f64() / 60.0
    );
    info!("  • Epochs completed: {}", baseline_epochs_completed);
    info!("  • Final policy loss: {:.6}", baseline_metrics.policy_loss);
    info!("  • Final value loss: {:.6}", baseline_metrics.value_loss);
    info!(
        "  • Explained variance: {:.4}",
        baseline_metrics.explained_variance
    );

    // ========================================================================
    // BENCHMARK 2: Optimized PPO Trainer
    // ========================================================================
    info!("\n═══════════════════════════════════════════════════════════════");
    info!("📊 BENCHMARK 2: Optimized PPO Trainer (Vectorized Environments)");
    info!("═══════════════════════════════════════════════════════════════\n");

    let optimized_trainer = PpoTrainer::new(
        hyperparams.clone(),
        state_dim,
        format!("{}/optimized", opts.output_dir),
        true,                // CUDA
        Some(opts.num_envs), // Vectorized mode with parallel environments
    )
    .context("Failed to create optimized trainer")?;

    let mut optimized_epochs_completed = 0;
    let optimized_callback = |metrics: PpoTrainingMetrics| {
        optimized_epochs_completed = metrics.epoch;
        if metrics.epoch % 5 == 0 {
            info!(
                "Optimized Epoch {}/{}: policy_loss={:.4}, value_loss={:.4}, expl_var={:.4}",
                metrics.epoch,
                hyperparams.epochs,
                metrics.policy_loss,
                metrics.value_loss,
                metrics.explained_variance
            );
        }
    };

    let optimized_start = Instant::now();
    let optimized_metrics = optimized_trainer
        .train(market_data.clone(), optimized_callback)
        .await
        .context("Optimized training failed")?;
    let optimized_duration = optimized_start.elapsed();

    info!("\n✅ Optimized Training Complete!");
    info!(
        "  • Duration: {:.2}s ({:.2} min)",
        optimized_duration.as_secs_f64(),
        optimized_duration.as_secs_f64() / 60.0
    );
    info!("  • Epochs completed: {}", optimized_epochs_completed);
    info!(
        "  • Final policy loss: {:.6}",
        optimized_metrics.policy_loss
    );
    info!("  • Final value loss: {:.6}", optimized_metrics.value_loss);
    info!(
        "  • Explained variance: {:.4}",
        optimized_metrics.explained_variance
    );

    // ========================================================================
    // PERFORMANCE COMPARISON
    // ========================================================================
    info!("\n═══════════════════════════════════════════════════════════════");
    info!("📊 PERFORMANCE COMPARISON");
    info!("═══════════════════════════════════════════════════════════════\n");

    let speedup = baseline_duration.as_secs_f64() / optimized_duration.as_secs_f64();
    let time_saved = baseline_duration.as_secs_f64() - optimized_duration.as_secs_f64();
    let time_saved_pct = (time_saved / baseline_duration.as_secs_f64()) * 100.0;

    info!("Timing Comparison:");
    info!(
        "  • Baseline:  {:.2}s ({:.2} min)",
        baseline_duration.as_secs_f64(),
        baseline_duration.as_secs_f64() / 60.0
    );
    info!(
        "  • Optimized: {:.2}s ({:.2} min)",
        optimized_duration.as_secs_f64(),
        optimized_duration.as_secs_f64() / 60.0
    );
    info!("  • Speedup:   {:.2}x", speedup);
    info!(
        "  • Time saved: {:.2}s ({:.1}%)",
        time_saved, time_saved_pct
    );

    info!("\nQuality Comparison:");
    info!(
        "  • Baseline policy loss:   {:.6}",
        baseline_metrics.policy_loss
    );
    info!(
        "  • Optimized policy loss:  {:.6}",
        optimized_metrics.policy_loss
    );
    info!(
        "  • Baseline value loss:    {:.6}",
        baseline_metrics.value_loss
    );
    info!(
        "  • Optimized value loss:   {:.6}",
        optimized_metrics.value_loss
    );
    info!(
        "  • Baseline expl_var:      {:.4}",
        baseline_metrics.explained_variance
    );
    info!(
        "  • Optimized expl_var:     {:.4}",
        optimized_metrics.explained_variance
    );

    info!("\nOptimization Breakdown (Estimated):");
    info!("  • Vectorized rollouts:     2-4x speedup");
    info!("  • Batch GAE computation:   3-5x speedup");
    info!("  • Parallel updates:        1.5-2x speedup");
    info!("  • Total (measured):        {:.2}x speedup", speedup);

    // Validate target achieved
    if speedup >= 2.0 {
        info!("\n✅ SUCCESS: Achieved target 2x speedup!");
        info!(
            "   Actual speedup: {:.2}x ({}% faster)",
            speedup, time_saved_pct
        );
    } else {
        warn!("\n⚠️  WARNING: Did not achieve target 2x speedup");
        warn!(
            "   Actual speedup: {:.2}x ({}% faster)",
            speedup, time_saved_pct
        );
        warn!("   Target: 2.0x or higher");
    }

    info!("\n🎉 Benchmark complete!");
    info!("📁 Checkpoints saved to: {}", opts.output_dir);

    Ok(())
}