foxhunt/ml/examples/train_ppo.rs

//! PPO Training Example
//!
//! Trains a PPO model on market data and saves checkpoints to disk.
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (100 epochs)
//! cargo run -p ml --example train_ppo --release --features cuda
//!
//! # Custom epochs and output path
//! cargo run -p ml --example train_ppo --release --features cuda -- \
//!   --epochs 500 \
//!   --output-dir ml/trained_models
//! ```

use anyhow::{Context, Result};
use std::path::PathBuf;
use structopt::StructOpt;
use tracing::{info};
use tracing_subscriber::FmtSubscriber;

use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer, PpoTrainingMetrics};

#[derive(Debug, StructOpt)]
#[structopt(name = "train_ppo", about = "Train PPO model on market data")]
struct Opts {
    /// Number of training epochs
    #[structopt(long, default_value = "100")]
    epochs: usize,

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

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

    /// Output directory for trained model
    #[structopt(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// Use GPU
    #[structopt(long)]
    use_gpu: bool,

    /// Verbose logging
    #[structopt(short, long)]
    verbose: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Parse CLI options
    let opts = Opts::from_args();

    // Setup logging
    let level = if opts.verbose {
        tracing::Level::DEBUG
    } else {
        tracing::Level::INFO
    };

    let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    info!("🚀 Starting PPO Training");
    info!("Configuration:");
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • GPU enabled: {}", opts.use_gpu);
    info!("  • Output directory: {}", opts.output_dir);

    // Create output directory
    let output_path = PathBuf::from(&opts.output_dir);
    if !output_path.exists() {
        std::fs::create_dir_all(&output_path)
            .context("Failed to create output directory")?;
        info!("✅ Created output directory: {}", opts.output_dir);
    }

    // Configure PPO 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,
    };

    // Create PPO trainer
    let trainer = PpoTrainer::new(
        hyperparams.clone(),
        64,  // state_dim - inferred from data in production
        &opts.output_dir,
        opts.use_gpu,
    ).context("Failed to create PPO trainer")?;

    info!("✅ PPO trainer initialized");

    // Generate synthetic market data for training
    info!("\n📊 Generating training data...");
    let num_samples = 10000;
    let mut market_data = Vec::with_capacity(num_samples);

    for i in 0..num_samples {
        // Generate synthetic 64-dimensional state
        let price_base = 4000.0 + (i as f32 * 0.1);
        let mut state = vec![price_base; 64];

        // Add some variation
        for j in 0..64 {
            state[j] += (i as f32 * 0.01 * (j as f32).sin());
        }

        market_data.push(state);
    }

    info!("✅ Generated {} samples", num_samples);

    // Create progress callback
    let progress_callback = |metrics: PpoTrainingMetrics| {
        if metrics.epoch % 10 == 0 {
            info!(
                "📊 Epoch {}/{}: policy_loss={:.4}, value_loss={:.4}, kl_div={:.4}",
                metrics.epoch,
                hyperparams.epochs,
                metrics.policy_loss,
                metrics.value_loss,
                metrics.kl_divergence
            );
        }
    };

    // Train the model
    info!("\n🏋️  Starting training...\n");
    let start_time = std::time::Instant::now();

    let final_metrics = trainer
        .train(market_data, progress_callback)
        .await
        .context("Training failed")?;

    let training_duration = start_time.elapsed();

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    info!("  • Policy loss: {:.6}", final_metrics.policy_loss);
    info!("  • Value loss: {:.6}", final_metrics.value_loss);
    info!("  • KL divergence: {:.6}", final_metrics.kl_divergence);
    info!("  • Explained variance: {:.4}", final_metrics.explained_variance);
    info!("  • Mean reward: {:.4}", final_metrics.mean_reward);
    info!("  • Training time: {:.1}s ({:.1} min)",
          training_duration.as_secs_f64(),
          training_duration.as_secs_f64() / 60.0);

    // Checkpoint is already saved by trainer (every 10 epochs)
    let final_checkpoint = output_path.join(format!("ppo_checkpoint_epoch_{}.safetensors", hyperparams.epochs));

    info!("\n💾 Final checkpoint saved to: {}", final_checkpoint.display());
    info!("\n🎉 PPO training complete!");
    info!("📁 Model files saved to: {}", opts.output_dir);

    Ok(())
}