foxhunt/ml/examples/train_dqn_production.rs

//! Production DQN Training Script
//!
//! Trains a DQN model for 50 epochs using production hyperparameters.

use anyhow::{Context, Result};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use std::path::PathBuf;
use std::time::Instant;

#[tokio::main]
async fn main() -> Result<()> {
    // Setup logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("\n{}", "=".repeat(80));
    println!("🚀 DQN Production Training - 50 Epochs");
    println!("{}", "=".repeat(80));

    let start_time = Instant::now();

    // Get data directory
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .context("Failed to get workspace root")?
        .to_path_buf();

    let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");

    if !data_dir.exists() {
        anyhow::bail!(
            "Data directory not found: {}. Please check the path.",
            data_dir.display()
        );
    }

    // Create checkpoint directory
    let checkpoint_dir = PathBuf::from("/tmp");
    std::fs::create_dir_all(&checkpoint_dir)?;

    println!("\n📋 Configuration:");
    println!("   Data Directory: {}", data_dir.display());
    println!("   Checkpoint Directory: {}", checkpoint_dir.display());

    // Configure production hyperparameters (conservative baseline)
    let mut hyperparams = DQNHyperparameters::conservative();
    hyperparams.epochs = 50;
    hyperparams.batch_size = 64;
    hyperparams.learning_rate = 0.0001;
    hyperparams.gamma = 0.99;
    hyperparams.epsilon_start = 0.3;
    hyperparams.epsilon_end = 0.05;
    hyperparams.epsilon_decay = 0.995;
    hyperparams.checkpoint_frequency = 10;
    hyperparams.early_stopping_enabled = true;
    hyperparams.min_epochs_before_stopping = 50; // Allow all 50 epochs

    println!("\n⚙️  Hyperparameters:");
    println!("   Epochs: {}", hyperparams.epochs);
    println!("   Batch Size: {}", hyperparams.batch_size);
    println!("   Learning Rate: {}", hyperparams.learning_rate);
    println!("   Gamma: {}", hyperparams.gamma);
    println!(
        "   Epsilon: {} → {} (decay: {})",
        hyperparams.epsilon_start, hyperparams.epsilon_end, hyperparams.epsilon_decay
    );

    // Create trainer
    println!("\n🏗️  Initializing DQN trainer...");
    let mut trainer = DQNTrainer::new(hyperparams.clone())?;

    // Train the model
    println!("\n🚀 Starting training...\n");

    let mut best_checkpoint_path = PathBuf::new();

    let metrics = trainer
        .train(
            &data_dir.to_string_lossy().to_string(),
            |epoch, checkpoint_data, is_best| {
                let filename = if is_best {
                    "dqn_prod_best.safetensors".to_string()
                } else {
                    format!("dqn_prod_epoch_{}.safetensors", epoch)
                };
                let path = checkpoint_dir.join(filename);
                std::fs::write(&path, checkpoint_data)?;
                if is_best {
                    best_checkpoint_path = path.clone();
                    println!(
                        "   💾 ⭐ BEST checkpoint saved: epoch {} -> {}",
                        epoch,
                        path.display()
                    );
                } else {
                    println!(
                        "   💾 Checkpoint saved: epoch {} -> {}",
                        epoch,
                        path.display()
                    );
                }
                Ok(path.to_string_lossy().to_string())
            },
        )
        .await?;

    let training_time = start_time.elapsed();

    // Report results
    println!("\n{}", "=".repeat(80));
    println!("✅ TRAINING COMPLETE");
    println!("{}", "=".repeat(80));
    println!("\n📊 Results:");
    println!("   Epochs Completed: {}", metrics.epochs_trained);
    println!("   Final Loss: {:.6}", metrics.loss);
    println!(
        "   Training Time: {:.2}s ({:.1} min)",
        training_time.as_secs_f64(),
        training_time.as_secs_f64() / 60.0
    );
    println!("   Convergence: {}", metrics.convergence_achieved);

    if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
        println!("   Avg Q-value: {:.4}", avg_q_value);
    }

    if let Some(final_epsilon) = metrics.additional_metrics.get("final_epsilon") {
        println!("   Final Epsilon: {:.4}", final_epsilon);
    }

    println!("\n💾 Best Checkpoint: {}", best_checkpoint_path.display());

    let checkpoint_size = std::fs::metadata(&best_checkpoint_path)?.len();
    println!("   Size: {} KB", checkpoint_size / 1024);

    println!("\n{}", "=".repeat(80));

    // Save metrics to JSON
    let metrics_json = serde_json::json!({
        "epochs_trained": metrics.epochs_trained,
        "final_loss": metrics.loss,
        "training_time_seconds": training_time.as_secs_f64(),
        "convergence_achieved": metrics.convergence_achieved,
        "avg_q_value": metrics.additional_metrics.get("avg_q_value"),
        "final_epsilon": metrics.additional_metrics.get("final_epsilon"),
        "checkpoint_path": best_checkpoint_path.to_string_lossy().to_string(),
        "checkpoint_size_kb": checkpoint_size / 1024,
    });

    let metrics_path = PathBuf::from("/tmp/dqn_production_test_training.json");
    std::fs::write(&metrics_path, serde_json::to_string_pretty(&metrics_json)?)?;
    println!("📄 Training metrics saved to: {}", metrics_path.display());

    Ok(())
}