foxhunt/ml/tests/polyak_integration_test.rs

//! Polyak Averaging Integration Tests
//!
//! Validates that Polyak averaging (soft target updates) is properly integrated
//! into the DQN training pipeline and reduces Q-value oscillations compared to
//! hard updates.
//!
//! Test Coverage:
//! 1. Soft updates reduce Q-oscillations vs hard updates (50-70% variance reduction)
//! 2. Rainbow τ=0.001 produces expected convergence half-life (~693 steps)
//! 3. Hard update fallback works when use_soft_updates=false
//! 4. Convergence half-life calculation is accurate

use anyhow::Result;
use candle_core::Device;
use ml::dqn::{convergence_half_life, hard_update, polyak_update, WorkingDQN, WorkingDQNConfig};
use std::sync::Arc;
use tokio::sync::RwLock;

/// Test 1: Soft updates reduce Q-value oscillations compared to hard updates
///
/// Expectation: Q-value variance should be 50-70% lower with Polyak averaging
/// compared to periodic hard updates.
///
/// Method:
/// 1. Train 2 identical DQN agents for 100 steps
/// 2. Agent A: Soft updates every step (τ=0.001)
/// 3. Agent B: Hard updates every 10 steps
/// 4. Measure Q-value variance for both
/// 5. Assert: variance_soft < 0.7 * variance_hard (30% reduction)
#[tokio::test]
async fn test_soft_updates_reduce_q_oscillations() -> Result<()> {
    // Create two identical DQN configurations
    let config_soft = WorkingDQNConfig {
        state_dim: 225,
        hidden_dims: vec![128, 64, 32],
        num_actions: 3,
        learning_rate: 0.0001,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 10000,
        batch_size: 32,
        min_replay_size: 100,
        target_update_freq: 1, // Update every step (soft updates)
        use_double_dqn: true,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
    };

    let config_hard = WorkingDQNConfig {
        target_update_freq: 10, // Update every 10 steps (hard updates)
        ..config_soft.clone()
    };

    // Create agents
    let mut agent_soft = WorkingDQN::new(config_soft)?;
    let mut agent_hard = WorkingDQN::new(config_hard)?;

    // Generate random training data (225 features → 3 actions)
    let mut q_values_soft = Vec::new();
    let mut q_values_hard = Vec::new();

    for step in 0..100 {
        // Generate random state
        let state: Vec<f64> = (0..225).map(|_| rand::random::<f64>()).collect();

        // Get Q-values before training (to measure variance)
        let q_soft = agent_soft.get_q_values(&state)?;
        let q_hard = agent_hard.get_q_values(&state)?;

        q_values_soft.push(q_soft.iter().sum::<f64>() / q_soft.len() as f64);
        q_values_hard.push(q_hard.iter().sum::<f64>() / q_hard.len() as f64);

        // Simulate training step (add experience, train if buffer ready)
        let action = rand::random::<usize>() % 3;
        let reward = rand::random::<f64>() - 0.5; // -0.5 to 0.5
        let next_state: Vec<f64> = (0..225).map(|_| rand::random::<f64>()).collect();
        let done = false;

        agent_soft.add_experience(state.clone(), action, reward, next_state.clone(), done)?;
        agent_hard.add_experience(state.clone(), action, reward, next_state.clone(), done)?;

        // Train if buffer is ready
        if step >= 100 {
            let _ = agent_soft.train_step();
            let _ = agent_hard.train_step();
        }

        // Apply updates (soft vs hard)
        if step >= 100 {
            // Soft update (every step)
            let tau = 0.001;
            let online_vars = agent_soft.get_q_network_vars();
            let target_vars = agent_soft.get_target_network_vars();
            polyak_update(&online_vars, &target_vars, tau)?;

            // Hard update (every 10 steps)
            if step % 10 == 0 {
                let online_vars = agent_hard.get_q_network_vars();
                let target_vars = agent_hard.get_target_network_vars();
                hard_update(&online_vars, &target_vars)?;
            }
        }
    }

    // Calculate Q-value variance
    let mean_soft = q_values_soft.iter().sum::<f64>() / q_values_soft.len() as f64;
    let mean_hard = q_values_hard.iter().sum::<f64>() / q_values_hard.len() as f64;

    let variance_soft = q_values_soft
        .iter()
        .map(|q| (q - mean_soft).powi(2))
        .sum::<f64>()
        / q_values_soft.len() as f64;

    let variance_hard = q_values_hard
        .iter()
        .map(|q| (q - mean_hard).powi(2))
        .sum::<f64>()
        / q_values_hard.len() as f64;

    println!("Soft update variance: {:.6}", variance_soft);
    println!("Hard update variance: {:.6}", variance_hard);
    println!(
        "Variance reduction: {:.1}%",
        (1.0 - variance_soft / variance_hard) * 100.0
    );

    // Assert: Soft updates reduce variance by at least 40%
    assert!(
        variance_soft < 0.6 * variance_hard,
        "Soft updates should reduce Q-value variance by ≥40%: {:.6} vs {:.6}",
        variance_soft,
        variance_hard
    );

    Ok(())
}

/// Test 2: Rainbow τ=0.001 produces expected convergence half-life (~693 steps)
///
/// Expectation: With τ=0.001, target network should reach 50% of online network
/// distance after ~693 training steps.
#[test]
fn test_rainbow_tau_convergence_half_life() {
    let tau = 0.001;
    let expected_half_life = 693.0;

    let actual_half_life = convergence_half_life(tau);

    println!(
        "Rainbow τ={}: half-life = {:.0} steps (expected: {:.0})",
        tau, actual_half_life, expected_half_life
    );

    assert!(
        (actual_half_life - expected_half_life).abs() < 1.0,
        "Half-life should be ~693 steps for τ=0.001: {:.0}",
        actual_half_life
    );
}

/// Test 3: Hard update fallback works when use_soft_updates=false
///
/// Expectation: When soft updates are disabled, periodic hard updates should
/// still synchronize the target network with the online network.
#[tokio::test]
async fn test_hard_update_fallback() -> Result<()> {
    let config = WorkingDQNConfig {
        state_dim: 225,
        hidden_dims: vec![128, 64, 32],
        num_actions: 3,
        learning_rate: 0.0001,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 10000,
        batch_size: 32,
        min_replay_size: 100,
        target_update_freq: 10, // Hard update every 10 steps
        use_double_dqn: true,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
    };

    let mut agent = WorkingDQN::new(config)?;

    // Train for 20 steps (2 hard updates expected at steps 10, 20)
    for step in 0..20 {
        let state: Vec<f64> = (0..225).map(|_| rand::random::<f64>()).collect();
        let action = rand::random::<usize>() % 3;
        let reward = rand::random::<f64>() - 0.5;
        let next_state: Vec<f64> = (0..225).map(|_| rand::random::<f64>()).collect();
        let done = false;

        agent.add_experience(state, action, reward, next_state, done)?;

        if step >= 100 {
            let _ = agent.train_step();

            // Apply hard update every 10 steps
            if step % 10 == 0 {
                let online_vars = agent.get_q_network_vars();
                let target_vars = agent.get_target_network_vars();
                hard_update(&online_vars, &target_vars)?;
                println!("✓ Hard update applied at step {}", step);
            }
        }
    }

    println!("✓ Hard update fallback works correctly");
    Ok(())
}

/// Test 4: Convergence half-life calculation is accurate for various τ values
///
/// Expectation: Half-life formula should produce correct values for:
/// - τ=0.001 → ~693 steps (Rainbow)
/// - τ=0.01  → ~69 steps (faster convergence)
/// - τ=0.1   → ~7 steps (very fast convergence)
#[test]
fn test_convergence_half_life_accuracy() {
    let test_cases = vec![(0.001, 693.0), (0.01, 69.0), (0.1, 7.0)];

    for (tau, expected) in test_cases {
        let actual = convergence_half_life(tau);
        let error = (actual - expected).abs();

        println!(
            "τ={}: half-life = {:.1} steps (expected: {:.0}, error: {:.1})",
            tau, actual, expected, error
        );

        assert!(
            error < 1.0,
            "Half-life calculation error too large for τ={}: {:.1} steps",
            tau,
            error
        );
    }
}

/// Test 5: Polyak averaging parameters can be configured via DQN trainer
///
/// Expectation: DQN trainer should accept τ and use_soft_updates parameters
/// and apply them correctly during training.
#[tokio::test]
async fn test_dqn_trainer_polyak_configuration() -> Result<()> {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

    // Create hyperparameters with Polyak averaging enabled
    let hyperparams = DQNHyperparameters {
        learning_rate: 0.0001,
        batch_size: 32,
        gamma: 0.99,
        epsilon_start: 0.3,
        epsilon_end: 0.05,
        epsilon_decay: 0.995,
        buffer_size: 10000,
        min_replay_size: 100,
        epochs: 1, // Just test initialization
        checkpoint_frequency: 10,
        early_stopping_enabled: false,
        q_value_floor: 0.5,
        min_loss_improvement_pct: 2.0,
        plateau_window: 5,
        min_epochs_before_stopping: 10,
        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,
        tau: 0.001,             // Rainbow's τ
        use_soft_updates: true, // Enable Polyak averaging
    };

    // Create trainer (should not panic)
    let trainer = DQNTrainer::new(hyperparams)?;

    println!("✓ DQN trainer accepts Polyak averaging parameters");
    println!("  • τ = 0.001 (Rainbow)");
    println!("  • use_soft_updates = true");

    Ok(())
}