foxhunt/ml/examples/hyperopt_dqn_demo.rs

//! DQN Hyperparameter Optimization Demo
//!
//! This example demonstrates how to use the argmin-based hyperparameter
//! optimization framework with DQN. It runs a small-scale optimization
//! to show the complete workflow with REAL training (not mock metrics).
//!
//! ## Usage
//!
//! ```bash
//! # Quick test with small DBN directory (5-10 minutes)
//! cargo run -p ml --example hyperopt_dqn_demo --release --features cuda -- \
//!   --dbn-data-dir test_data/real/databento/ml_training_small \
//!   --trials 3 \
//!   --epochs 5
//!
//! # Production run with full optimization (1-2 hours)
//! cargo run -p ml --example hyperopt_dqn_demo --release --features cuda -- \
//!   --dbn-data-dir test_data/real/databento/ml_training \
//!   --trials 30 \
//!   --epochs 50
//! ```
//!
//! ## Output
//!
//! The example will:
//! 1. Initialize DQN trainer with specified Parquet file
//! 2. Run argmin optimization with Nelder-Mead simplex
//! 3. Display trial results including loss and parameter values
//! 4. Report best hyperparameters found
//! 5. Show convergence and top trials
//!
//! ## Verification
//!
//! This example uses REAL training via `InternalDQNTrainer`, not mock metrics.
//! You should see:
//! - Loss values VARY across trials (not identical)
//! - Training takes time (not instant)
//! - GPU utilization visible (if CUDA available)
//! - Convergence over trials (best loss improves)

use anyhow::Result;
use clap::Parser;
use ml::hyperopt::adapters::dqn::DQNTrainer;
use ml::hyperopt::paths::{generate_run_id, TrainingPaths};
use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
use tracing::{info, Level};
use tracing_subscriber;

#[derive(Parser, Debug)]
#[command(name = "DQN Hyperparameter Optimization Demo")]
#[command(about = "Demonstrates argmin-based hyperparameter optimization for DQN")]
struct Args {
    /// Path to directory containing DBN data files
    #[arg(long)]
    dbn_data_dir: String,

    /// Number of optimization trials (default: 10)
    #[arg(long, default_value = "10")]
    trials: usize,

    /// Number of training epochs per trial (default: 20)
    #[arg(long, default_value = "20")]
    epochs: usize,

    /// Number of initial random samples (default: 2)
    #[arg(long, default_value = "2")]
    n_initial: usize,

    /// Random seed for reproducibility (default: 42)
    #[arg(long, default_value = "42")]
    seed: u64,

    /// Base directory for training outputs (default: /tmp/ml_training)
    #[arg(long, default_value = "/tmp/ml_training")]
    base_dir: String,

    /// Run ID for organizing outputs (default: auto-generated)
    #[arg(long)]
    run_id: Option<String>,

    /// Run type for run ID generation (default: hyperopt)
    #[arg(long, default_value = "hyperopt")]
    run_type: String,
}

fn estimate_runtime(trials: usize, epochs: usize) -> usize {
    // DQN training is faster than MAMBA-2
    // Rough estimate: ~0.5 min per trial per 10 epochs
    let mins_per_trial = (epochs as f64 / 10.0) * 0.5;
    (trials as f64 * mins_per_trial).ceil() as usize
}

fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .with_target(false)
        .init();

    // Parse arguments
    let args = Args::parse();

    info!("========================================");
    info!("DQN Hyperparameter Optimization Demo");
    info!("========================================");
    info!("Configuration:");
    info!("  DBN data directory: {}", args.dbn_data_dir);
    info!("  Trials: {}", args.trials);
    info!("  Epochs per trial: {}", args.epochs);
    info!("  Initial samples: {}", args.n_initial);
    info!("  Random seed: {}", args.seed);
    info!("  Base directory: {}", args.base_dir);
    info!("");

    // Verify DBN data directory exists
    if !std::path::Path::new(&args.dbn_data_dir).exists() {
        anyhow::bail!("DBN data directory not found: {}", args.dbn_data_dir);
    }

    // Generate run ID
    let run_id = args.run_id.unwrap_or_else(|| generate_run_id(&args.run_type));
    info!("Run ID: {}", run_id);

    // Create training paths
    let training_paths = TrainingPaths::new(&args.base_dir, "dqn", &run_id);
    info!("Training paths:");
    info!("  Run directory: {:?}", training_paths.run_dir());
    info!("  Checkpoints: {:?}", training_paths.checkpoints_dir());
    info!("  Hyperopt: {:?}", training_paths.hyperopt_dir());
    info!("");

    // Create trainer
    info!("Creating DQN trainer with REAL training (not mock metrics)...");
    let trainer = DQNTrainer::new(&args.dbn_data_dir, args.epochs)?
        .with_training_paths(training_paths);

    // Create optimizer
    info!("Initializing argmin optimizer...");
    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    // Run optimization
    info!("");
    info!("Starting optimization (this may take a while)...");
    info!("Expected runtime: ~{} minutes", estimate_runtime(args.trials, args.epochs));
    info!("");
    info!("VERIFICATION CHECKS:");
    info!("  ✓ Each trial should take >1 second (real training)");
    info!("  ✓ Loss values should VARY across trials");
    info!("  ✓ Best loss should improve over trials");
    info!("  ✓ GPU utilization should be visible (if CUDA available)");
    info!("");

    let result = optimizer.optimize(trainer)?;

    // Display results
    info!("");
    info!("========================================");
    info!("Optimization Complete!");
    info!("========================================");
    info!("");
    info!("Best Hyperparameters:");
    info!("  Learning rate: {:.6}", result.best_params.learning_rate);
    info!("  Batch size: {}", result.best_params.batch_size);
    info!("  Gamma: {:.3}", result.best_params.gamma);
    info!("  Epsilon decay: {:.5}", result.best_params.epsilon_decay);
    info!("  Buffer size: {}", result.best_params.buffer_size);
    info!("");
    info!("Performance:");
    info!("  Best training loss: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());

    // Find convergence trial (where best was found)
    let convergence_trial = result
        .all_trials
        .iter()
        .position(|t| (t.objective - result.best_objective).abs() < 1e-10)
        .unwrap_or(0);
    info!("  Convergence: {} trials to best", convergence_trial + 1);
    info!("");

    // Show top 5 trials
    if result.all_trials.len() >= 5 {
        info!("Top 5 Trials:");
        let mut sorted_trials = result.all_trials.clone();
        sorted_trials.sort_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap());

        for (i, trial) in sorted_trials.iter().take(5).enumerate() {
            info!(
                "  {}. Loss: {:.6} (LR: {:.6}, BS: {}, Gamma: {:.3}, Eps: {:.5})",
                i + 1,
                trial.objective,
                trial.params.learning_rate,
                trial.params.batch_size,
                trial.params.gamma,
                trial.params.epsilon_decay
            );
        }
        info!("");
    }

    // Validation check: Verify loss variance
    let losses: Vec<f64> = result.all_trials.iter().map(|t| t.objective).collect();
    let mean_loss = losses.iter().sum::<f64>() / losses.len() as f64;
    let variance = losses
        .iter()
        .map(|l| (l - mean_loss).powi(2))
        .sum::<f64>()
        / losses.len() as f64;
    let std_dev = variance.sqrt();

    info!("VERIFICATION RESULTS:");
    info!("  Mean loss: {:.6}", mean_loss);
    info!("  Std deviation: {:.6}", std_dev);
    info!("  Min loss: {:.6}", losses.iter().cloned().fold(f64::INFINITY, f64::min));
    info!("  Max loss: {:.6}", losses.iter().cloned().fold(f64::NEG_INFINITY, f64::max));
    info!("");

    if std_dev < 1e-6 {
        info!("⚠️  WARNING: Loss values are identical across trials!");
        info!("   This suggests mock metrics are being used instead of real training.");
        info!("   Expected: std_dev > 0.001 for real training");
    } else {
        info!("✅ VERIFIED: Loss values vary across trials (real training confirmed)");
        info!("   Coefficient of variation: {:.2}%", (std_dev / mean_loss) * 100.0);
    }
    info!("");

    // Calculate improvement over default
    let default_loss = losses[0]; // First trial uses near-default params
    let improvement_pct = ((default_loss - result.best_objective) / default_loss) * 100.0;

    info!("Improvement:");
    info!("  Initial (near-default): {:.6}", default_loss);
    info!("  Best (optimized): {:.6}", result.best_objective);
    info!("  Improvement: {:.2}%", improvement_pct);
    info!("");

    info!("========================================");
    info!("Next Steps:");
    info!("  1. Review hyperparameters above");
    info!("  2. Run full optimization with --trials 30 --epochs 50");
    info!("  3. Deploy best params to production DQN config");
    info!("========================================");

    Ok(())
}