foxhunt/ml/examples/train_dqn.rs

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

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

use ml::checkpoint::{CheckpointConfig, CheckpointManager};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

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

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

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

    /// Discount factor (gamma)
    #[structopt(long, default_value = "0.99")]
    gamma: f64,

    /// Checkpoint save frequency (epochs)
    #[structopt(long, default_value = "10")]
    checkpoint_frequency: usize,

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

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

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

    /// Enable early stopping (recommended, use --no-early-stopping to disable)
    #[structopt(long)]
    early_stopping: bool,

    /// Disable early stopping
    #[structopt(long)]
    no_early_stopping: bool,

    /// Q-value floor threshold for early stopping
    #[structopt(long, default_value = "0.5")]
    q_value_floor: f64,

    /// Minimum loss improvement percentage for plateau detection
    #[structopt(long, default_value = "2.0")]
    min_loss_improvement: f64,

    /// Plateau detection window size (epochs)
    #[structopt(long, default_value = "30")]
    plateau_window: usize,
}

#[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 DQN Training");
    info!("Configuration:");
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • Gamma: {}", opts.gamma);
    info!("  • Checkpoint frequency: {} epochs", opts.checkpoint_frequency);
    info!("  • Output directory: {}", opts.output_dir);
    info!("  • Data directory: {}", opts.data_dir);

    // Determine early stopping (enabled by default, unless --no-early-stopping is specified)
    let early_stopping_enabled = !opts.no_early_stopping;
    info!("  • Early stopping: {}", if early_stopping_enabled { "enabled" } else { "disabled" });
    if early_stopping_enabled {
        info!("    - Q-value floor: {}", opts.q_value_floor);
        info!("    - Min loss improvement: {}%", opts.min_loss_improvement);
        info!("    - Plateau window: {} epochs", opts.plateau_window);
    }

    // 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 DQN hyperparameters
    let hyperparams = DQNHyperparameters {
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        gamma: opts.gamma,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        buffer_size: 100_000,
        epochs: opts.epochs,
        checkpoint_frequency: opts.checkpoint_frequency,
        early_stopping_enabled,
        q_value_floor: opts.q_value_floor,
        min_loss_improvement_pct: opts.min_loss_improvement,
        plateau_window: opts.plateau_window,
        min_epochs_before_stopping: 50,
    };

    // Create DQN trainer
    let mut trainer = DQNTrainer::new(hyperparams)
        .context("Failed to create DQN trainer")?;

    info!("✅ DQN trainer initialized");

    // Setup checkpoint manager
    let checkpoint_config = CheckpointConfig {
        base_dir: output_path.clone(),
        max_checkpoints_per_model: 10,
        auto_cleanup: true,
        validate_checksums: true,
        ..Default::default()
    };

    let checkpoint_manager = CheckpointManager::new(checkpoint_config)
        .context("Failed to create checkpoint manager")?;

    // Track checkpoint count
    let mut checkpoint_count = 0;

    // Create checkpoint callback
    let output_dir_for_callback = opts.output_dir.clone();
    let checkpoint_callback = move |epoch: usize, model_data: Vec<u8>| -> Result<String> {
        let checkpoint_path = PathBuf::from(&output_dir_for_callback)
            .join(format!("dqn_epoch_{}.safetensors", epoch));

        // Save checkpoint to disk
        std::fs::write(&checkpoint_path, &model_data)
            .context(format!("Failed to save checkpoint: {:?}", checkpoint_path))?;

        info!(
            "💾 Checkpoint saved: {} ({} bytes)",
            checkpoint_path.display(),
            model_data.len()
        );

        Ok(checkpoint_path.to_string_lossy().to_string())
    };

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

    let metrics = trainer
        .train(&opts.data_dir, checkpoint_callback)
        .await
        .context("Training failed")?;

    let training_duration = start_time.elapsed();

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    info!("  • Final loss: {:.6}", metrics.loss);
    info!("  • Epochs trained: {}", metrics.epochs_trained);
    info!("  • Training time: {:.1}s ({:.1} min)",
          metrics.training_time_seconds,
          metrics.training_time_seconds / 60.0);
    info!("  • Convergence: {}", if metrics.convergence_achieved { "✅ Yes" } else { "❌ No" });

    // Additional metrics from training
    if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
        info!("  • Average Q-value: {:.4}", avg_q_value);
    }
    if let Some(final_epsilon) = metrics.additional_metrics.get("final_epsilon") {
        info!("  • Final epsilon: {:.4}", final_epsilon);
    }
    if let Some(grad_norm) = metrics.additional_metrics.get("avg_gradient_norm") {
        info!("  • Average gradient norm: {:.6}", grad_norm);
    }

    // Save final model
    let final_model_path = output_path.join(format!("dqn_final_epoch{}.safetensors", opts.epochs));
    info!("\n💾 Saving final model to: {}", final_model_path.display());

    // Get final model state
    let final_checkpoint_data = trainer.serialize_model().await
        .context("Failed to serialize final model")?;

    std::fs::write(&final_model_path, &final_checkpoint_data)
        .context("Failed to save final model")?;

    info!("✅ Final model saved: {} ({} bytes)",
          final_model_path.display(),
          final_checkpoint_data.len());

    info!("\n🎉 DQN training complete!");
    info!("📁 Model files saved to: {}", opts.output_dir);

    Ok(())
}