foxhunt/ml/examples/train_dqn_ensemble_demo.rs

//! DQN Ensemble Training Demo
//!
//! Demonstrates multi-agent ensemble training with 5 DQN agents.
//! Shows both shared and independent replay buffer modes.

use anyhow::Result;
use ml::dqn::Experience;
use ml::trainers::dqn::DQNHyperparameters;
use ml::trainers::dqn_ensemble::{BufferMode, DQNEnsembleTrainer, EnsembleConfig};
use tracing::{info, Level};
use tracing_subscriber::FmtSubscriber;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .finish();
    tracing::subscriber::set_global_default(subscriber)?;

    info!("🚀 DQN Ensemble Training Demo");

    // Configure hyperparameters
    let hyperparams = DQNHyperparameters {
        learning_rate: 0.0001,
        batch_size: 64,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        buffer_size: 10000,
        min_replay_size: 500,
        epochs: 10,
        checkpoint_frequency: 5,
        early_stopping_enabled: false,
        q_value_floor: 0.5,
        min_loss_improvement_pct: 2.0,
        plateau_window: 30,
        min_epochs_before_stopping: 50,
        hold_penalty: -0.001,
        use_huber_loss: true,
        huber_delta: 1.0,
        use_double_dqn: true,
        gradient_clip_norm: Some(10.0),
        hold_penalty_weight: 0.01,
        movement_threshold: 0.02,
        diversity_penalty_weight: 0.05,
        enable_preprocessing: true,
        preprocessing_window: 50,
        preprocessing_clip_sigma: 5.0,
        td_error_clip: 10.0,
        tau: 0.001,
        target_update_mode: ml::trainers::TargetUpdateMode::Hard,
        target_update_frequency: 1000,
        warmup_steps: 0,
        use_regime_adaptation: false,
        regime_temperature_multipliers: std::collections::HashMap::new(),
        temperature_start: 1.0,
        temperature_min: 0.1,
        temperature_decay: 0.995,
        target_temperature_fraction: 0.75,
        reward_scale: 1000.0,
    };

    // Demo 1: Shared Buffer Mode (5 agents, shared replay)
    info!("\n📊 Demo 1: Shared Buffer Mode (5 agents)");
    demo_shared_buffer(hyperparams.clone()).await?;

    // Demo 2: Independent Buffer Mode (3 agents, independent replays)
    info!("\n📊 Demo 2: Independent Buffer Mode (3 agents)");
    demo_independent_buffer(hyperparams.clone()).await?;

    // Demo 3: Ensemble Prediction (majority vote)
    info!("\n📊 Demo 3: Ensemble Prediction (majority vote)");
    demo_ensemble_prediction(hyperparams).await?;

    info!("\n✅ All demos completed successfully!");
    Ok(())
}

/// Demo 1: Shared buffer mode - all agents sample from the same replay buffer
async fn demo_shared_buffer(hyperparams: DQNHyperparameters) -> Result<()> {
    let config = EnsembleConfig {
        num_agents: 5,
        buffer_mode: BufferMode::Shared,
        sync_target_updates: true,
        target_update_frequency: 1000,
        ..Default::default()
    };

    let mut trainer = DQNEnsembleTrainer::new(config, hyperparams)?;
    info!(
        "✓ Ensemble initialized: {} agents, {:?} buffer mode",
        trainer.num_agents(),
        trainer.buffer_mode()
    );

    // Generate synthetic experiences
    info!("Collecting experiences...");
    for i in 0..1000 {
        let state = vec![i as f32 * 0.001; 128];
        let action = (i % 3) as u8; // Cycle through BUY, SELL, HOLD
        let reward = (i as f32 * 0.1).sin(); // Synthetic reward
        let next_state = vec![(i + 1) as f32 * 0.001; 128];
        let done = false;

        let experience = Experience::new(state, action, reward, next_state, done);
        trainer.store_experience(experience, None).await?;
    }

    let buffer_size = trainer.get_replay_buffer_size().await?;
    info!("✓ Buffer size: {} experiences", buffer_size);

    // Train for 10 steps
    info!("Training for 10 steps...");
    for step in 1..=10 {
        let (avg_loss, avg_grad) = trainer.train_step(None).await?;
        info!(
            "Step {}: avg_loss={:.6}, avg_grad={:.6}",
            step, avg_loss, avg_grad
        );

        // Show per-agent metrics every 5 steps
        if step % 5 == 0 {
            for agent_id in 0..trainer.num_agents() {
                let agent_loss = trainer.get_agent_avg_loss(agent_id, 5).unwrap_or(0.0);
                let agent_grad = trainer.get_agent_avg_grad(agent_id, 5).unwrap_or(0.0);
                info!(
                    "  Agent {}: loss={:.6}, grad={:.6}",
                    agent_id, agent_loss, agent_grad
                );
            }
        }
    }

    // Update exploration parameters
    trainer.update_epsilon().await;
    info!(
        "✓ Updated epsilon: {:.4}",
        trainer.get_agent_epsilon(0).await.unwrap()
    );

    Ok(())
}

/// Demo 2: Independent buffer mode - each agent has its own replay buffer
async fn demo_independent_buffer(hyperparams: DQNHyperparameters) -> Result<()> {
    let config = EnsembleConfig {
        num_agents: 3,
        buffer_mode: BufferMode::Independent,
        sync_target_updates: true,
        target_update_frequency: 500,
        ..Default::default()
    };

    let mut trainer = DQNEnsembleTrainer::new(config, hyperparams)?;
    info!(
        "✓ Ensemble initialized: {} agents, {:?} buffer mode",
        trainer.num_agents(),
        trainer.buffer_mode()
    );

    // Store experiences in each agent's buffer
    info!("Collecting experiences for each agent...");
    for agent_id in 0..trainer.num_agents() {
        for i in 0..600 {
            let state = vec![(agent_id as f32 + i as f32 * 0.001); 128];
            let action = ((agent_id + i) % 3) as u8;
            let reward = ((agent_id + i) as f32 * 0.1).sin();
            let next_state = vec![(agent_id as f32 + (i + 1) as f32 * 0.001); 128];
            let done = false;

            let experience = Experience::new(state, action, reward, next_state, done);
            trainer.store_experience(experience, Some(agent_id)).await?;
        }
        info!("  Agent {}: {} experiences collected", agent_id, 600);
    }

    // Train for 5 steps
    info!("Training for 5 steps...");
    for step in 1..=5 {
        let (avg_loss, avg_grad) = trainer.train_step(None).await?;
        info!(
            "Step {}: avg_loss={:.6}, avg_grad={:.6}",
            step, avg_loss, avg_grad
        );
    }

    Ok(())
}

/// Demo 3: Ensemble prediction using majority vote
async fn demo_ensemble_prediction(hyperparams: DQNHyperparameters) -> Result<()> {
    let config = EnsembleConfig {
        num_agents: 7,
        buffer_mode: BufferMode::Shared,
        ..Default::default()
    };

    let trainer = DQNEnsembleTrainer::new(config, hyperparams)?;
    info!(
        "✓ Ensemble initialized: {} agents for prediction",
        trainer.num_agents()
    );

    // Test ensemble prediction on 5 sample states
    info!("Testing ensemble predictions (majority vote):");
    for i in 0..5 {
        let state = vec![i as f32 * 0.1; 128];
        let action = trainer.predict_ensemble(&state).await?;
        info!("  State {}: ensemble action = {:?}", i, action);
    }

    Ok(())
}