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 std::path::PathBuf;
use ml::hyperopt::ArgminOptimizer;
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 Parquet file with OHLCV data
    #[arg(long)]
    parquet_file: 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,

    /// Early stopping plateau window (epochs to check for improvement)
    #[arg(long, default_value = "5")]
    early_stopping_plateau_window: usize,

    /// Early stopping minimum epochs (minimum epochs before early stopping can trigger)
    /// Default: 1000 (effectively disabled - Wave 7 validation proved early stopping kills 8-10 profitable trials)
    #[arg(long, default_value = "1000")]
    early_stopping_min_epochs: usize,
    
    /// Optional path to feature cache directory for faster hyperopt
    /// 
    /// Pre-compute cache with: cargo run -p ml --example cache_dqn_features
    /// 
    /// Expected speedup: 2m 25s → <1s per trial (99% reduction)
    #[arg(long)]
    feature_cache_dir: Option<PathBuf>,
}

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!("  Parquet file: {}", args.parquet_file);
    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!("");

    // Validate Parquet file exists
    let parquet_path = std::path::Path::new(&args.parquet_file);
    if !parquet_path.exists() {
        anyhow::bail!("Parquet file not found: {}", args.parquet_file);
    }

    // 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 with parquet file path directly
    info!(
        "Creating DQN trainer with parquet file: {}",
        args.parquet_file
    );
    let mut trainer = DQNTrainer::new(&args.parquet_file, args.epochs)?
        .with_early_stopping(
            args.early_stopping_plateau_window,
            args.early_stopping_min_epochs,
        );
    
    // Enable feature cache if provided
    if let Some(cache_dir) = args.feature_cache_dir {
        info!("🚀 Enabling feature cache: {:?}", cache_dir);
        trainer = trainer.with_feature_cache(cache_dir);
    }
    
    let trainer = trainer.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!("  Buffer size: {}", result.best_params.buffer_size);
    info!("");
    info!("Performance:");
    info!("  Best episode reward: {:.6}", -result.best_objective); // Negate to get actual reward
    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 (sorted by reward descending = objective ascending)
    if result.all_trials.len() >= 5 {
        info!("Top 5 Trials (by episode reward):");
        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!(
                "  {}. Reward: {:.6} (LR: {:.6}, BS: {}, Gamma: {:.3})",
                i + 1,
                -trial.objective, // Negate to show actual reward
                trial.params.learning_rate,
                trial.params.batch_size,
                trial.params.gamma
            );
        }
        info!("");
    }

    // Validation check: Verify reward variance (objectives are negated rewards)
    let objectives: Vec<f64> = result.all_trials.iter().map(|t| t.objective).collect();
    let rewards: Vec<f64> = objectives.iter().map(|o| -o).collect();
    let mean_reward = rewards.iter().sum::<f64>() / rewards.len() as f64;
    let variance = rewards
        .iter()
        .map(|r| (r - mean_reward).powi(2))
        .sum::<f64>()
        / rewards.len() as f64;
    let std_dev = variance.sqrt();

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

    if std_dev < 1e-6 {
        info!("⚠️  WARNING: Reward 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: Reward values vary across trials (real training confirmed)");
        info!(
            "   Coefficient of variation: {:.2}%",
            (std_dev / mean_reward.abs()) * 100.0
        );
    }
    info!("");

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

    info!("Improvement:");
    info!("  Initial (near-default): {:.6}", default_reward);
    info!("  Best (optimized): {:.6}", best_reward);
    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(())
}