foxhunt/ml/tests/dqn_per_integration_test.rs

//! Integration tests for Prioritized Experience Replay (PER) in DQN
//!
//! Tests the following PER functionality:
//! 1. PER sampling - high TD-error transitions are sampled more frequently
//! 2. Priority updates - TD errors correctly update sample priorities
//! 3. Importance sampling weights - correct bias correction
//! 4. Beta annealing - beta increases from start to 1.0 over training
//!
//! Expected behavior:
//! - High-error experiences sampled 2-5× more than low-error experiences
//! - Priorities updated after each training step
//! - IS weights compensate for biased sampling
//! - Beta anneals linearly over training steps

use candle_core::{DType, Device, Tensor};
use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::dqn::experience::Experience;
use ml::dqn::prioritized_replay::{PrioritizedReplayBuffer, PrioritizedReplayConfig, PrioritizationStrategy};
use ml::MLError;

/// Helper to create test DQN config with PER enabled
fn create_per_config(use_per: bool) -> WorkingDQNConfig {
    WorkingDQNConfig {
        state_dim: 128,
        num_actions: 45,
        hidden_dims: vec![64, 32],
        learning_rate: 0.001,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.05,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 10000,
        batch_size: 32,
        min_replay_size: 100,
        target_update_freq: 1000,
        use_double_dqn: true,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        tau: 0.001,
        use_soft_updates: true,
        warmup_steps: 0,
        initial_capital: 100_000.0,
        use_per,
        per_alpha: 0.6,
        per_beta_start: 0.4,
        per_beta_max: 1.0,
        per_beta_annealing_steps: 100_000,
    }
}

/// Test 1: PER samples high-error experiences more frequently
#[test]
fn test_per_sampling_prioritization() -> Result<(), MLError> {
    let config = PrioritizedReplayConfig {
        capacity: 1000,
        alpha: 0.6,
        beta: 0.4,
        beta_max: 1.0,
        beta_annealing_steps: 100_000,
        initial_priority: 1.0,
        min_priority: 1e-6,
        strategy: PrioritizationStrategy::Proportional,
    };

    let mut buffer = PrioritizedReplayBuffer::new(config)?;

    // Add 100 experiences with varying priorities (simulating different TD errors)
    for i in 0..100 {
        let state = vec![i as f32; 128];
        let action = (i % 45) as u8;
        let reward = i as f32;
        let next_state = vec![(i + 1) as f32; 128];
        let done = false;

        let exp = Experience::new(state, action, reward, next_state, done);
        buffer.push(exp)?;

        // Set priority based on index (simulate high TD errors for indices 90-99)
        let priority = if i >= 90 {
            10.0 // High priority
        } else {
            1.0 // Low priority
        };
        buffer.update_priorities(&[i], &[priority])?;
    }

    // Sample 1000 times and count how many times high-priority experiences are sampled
    let mut high_priority_count = 0;
    let mut low_priority_count = 0;

    for _ in 0..1000 {
        let (samples, _weights, indices) = buffer.sample(10)?;

        for &idx in &indices {
            if idx >= 90 {
                high_priority_count += 1;
            } else {
                low_priority_count += 1;
            }
        }
    }

    // High-priority experiences (10% of buffer) should be sampled significantly more
    // Expected ratio: ~50% of samples (5× more than uniform 10%)
    let high_ratio = high_priority_count as f64 / (high_priority_count + low_priority_count) as f64;

    println!("High-priority sample ratio: {:.2}% (expected: >30%)", high_ratio * 100.0);
    println!("High-priority count: {}, Low-priority count: {}", high_priority_count, low_priority_count);

    // Assert high-priority experiences are sampled at least 3× more than uniform
    assert!(
        high_ratio > 0.3,
        "High-priority experiences should be sampled >30% of the time, got {:.2}%",
        high_ratio * 100.0
    );

    Ok(())
}

/// Test 2: Priority updates work correctly
#[test]
fn test_per_priority_updates() -> Result<(), MLError> {
    let config = PrioritizedReplayConfig {
        capacity: 100,
        alpha: 0.6,
        beta: 0.4,
        beta_max: 1.0,
        beta_annealing_steps: 100_000,
        initial_priority: 1.0,
        min_priority: 1e-6,
        strategy: PrioritizationStrategy::Proportional,
    };

    let mut buffer = PrioritizedReplayBuffer::new(config)?;

    // Add 10 experiences
    for i in 0..10 {
        let state = vec![i as f32; 128];
        let action = (i % 45) as u8;
        let reward = i as f32;
        let next_state = vec![(i + 1) as f32; 128];
        let done = false;

        let exp = Experience::new(state, action, reward, next_state, done);
        buffer.push(exp)?;
    }

    // Initially all priorities are equal (1.0)
    let (_, weights_before, indices_before) = buffer.sample(10)?;

    // Update priorities for experiences 0-4 to high values (10.0)
    let high_priority_indices: Vec<usize> = (0..5).collect();
    let high_priorities = vec![10.0; 5];
    buffer.update_priorities(&high_priority_indices, &high_priorities)?;

    // Sample again and verify high-priority experiences are sampled more
    let mut high_priority_sampled = 0;
    for _ in 0..100 {
        let (_, _, indices) = buffer.sample(5)?;
        high_priority_sampled += indices.iter().filter(|&&idx| idx < 5).count();
    }

    // Expect high-priority experiences (50% of buffer) to be sampled >60% of the time
    let high_ratio = high_priority_sampled as f64 / 500.0;

    println!("High-priority experiences sampled: {:.2}% (expected: >60%)", high_ratio * 100.0);

    assert!(
        high_ratio > 0.6,
        "After priority update, high-priority experiences should be sampled >60%, got {:.2}%",
        high_ratio * 100.0
    );

    Ok(())
}

/// Test 3: Importance sampling weights compensate for biased sampling
#[test]
fn test_per_importance_sampling_weights() -> Result<(), MLError> {
    let config = PrioritizedReplayConfig {
        capacity: 100,
        alpha: 0.6,
        beta: 0.8, // High beta for strong bias correction
        beta_max: 1.0,
        beta_annealing_steps: 100_000,
        initial_priority: 1.0,
        min_priority: 1e-6,
        strategy: PrioritizationStrategy::Proportional,
    };

    let mut buffer = PrioritizedReplayBuffer::new(config)?;

    // Add experiences with varying priorities
    for i in 0..100 {
        let state = vec![i as f32; 128];
        let action = (i % 45) as u8;
        let reward = i as f32;
        let next_state = vec![(i + 1) as f32; 128];
        let done = false;

        let exp = Experience::new(state, action, reward, next_state, done);
        buffer.push(exp)?;

        // Set varying priorities (1.0 to 10.0)
        let priority = 1.0 + (i as f32 / 10.0).min(9.0);
        buffer.update_priorities(&[i], &[priority])?;
    }

    // Sample and verify IS weights
    let (samples, is_weights, indices) = buffer.sample(32)?;

    // IS weights should:
    // 1. All be between 0 and 1 (since beta < 1.0, max weight is normalized to 1.0)
    // 2. Low-priority samples get higher weights (compensate for under-sampling)
    // 3. High-priority samples get lower weights (compensate for over-sampling)

    for (i, &weight) in is_weights.iter().enumerate() {
        assert!(
            weight > 0.0 && weight <= 1.0,
            "IS weight {} should be in (0, 1], got {}",
            i,
            weight
        );
    }

    // Check that low-priority experiences have higher IS weights
    let low_priority_idx = indices.iter().position(|&idx| idx < 10).unwrap_or(0);
    let high_priority_idx = indices.iter().position(|&idx| idx > 90).unwrap_or(1);

    let low_priority_weight = is_weights[low_priority_idx];
    let high_priority_weight = is_weights[high_priority_idx];

    println!("Low-priority IS weight: {:.4}", low_priority_weight);
    println!("High-priority IS weight: {:.4}", high_priority_weight);

    // Low-priority samples should have higher weights (less likely to be sampled → higher correction)
    assert!(
        low_priority_weight >= high_priority_weight * 0.8,
        "Low-priority weight ({:.4}) should be >= 0.8× high-priority weight ({:.4})",
        low_priority_weight,
        high_priority_weight
    );

    Ok(())
}

/// Test 4: Beta annealing works correctly
#[test]
fn test_per_beta_annealing() -> Result<(), MLError> {
    let config = PrioritizedReplayConfig {
        capacity: 100,
        alpha: 0.6,
        beta: 0.4, // Start at 0.4
        beta_max: 1.0, // Anneal to 1.0
        beta_annealing_steps: 1000,
        initial_priority: 1.0,
        min_priority: 1e-6,
        strategy: PrioritizationStrategy::Proportional,
    };

    let mut buffer = PrioritizedReplayBuffer::new(config)?;

    // Add experiences
    for i in 0..100 {
        let state = vec![i as f32; 128];
        let action = (i % 45) as u8;
        let reward = i as f32;
        let next_state = vec![(i + 1) as f32; 128];
        let done = false;

        let exp = Experience::new(state, action, reward, next_state, done);
        buffer.push(exp)?;
    }

    // Sample and step through training to test beta annealing
    let (_, weights_start, _) = buffer.sample(10)?;

    // Step 500 times (half-way through annealing)
    for _ in 0..500 {
        buffer.step();
    }

    let (_, weights_mid, _) = buffer.sample(10)?;

    // Step another 500 times (complete annealing)
    for _ in 0..500 {
        buffer.step();
    }

    let (_, weights_end, _) = buffer.sample(10)?;

    // Calculate average weight variance at each stage
    let var_start = weights_start.iter().map(|&w| (w - 1.0).powi(2)).sum::<f32>() / weights_start.len() as f32;
    let var_mid = weights_mid.iter().map(|&w| (w - 1.0).powi(2)).sum::<f32>() / weights_mid.len() as f32;
    let var_end = weights_end.iter().map(|&w| (w - 1.0).powi(2)).sum::<f32>() / weights_end.len() as f32;

    println!("Weight variance at start (beta=0.4): {:.6}", var_start);
    println!("Weight variance at mid (beta≈0.7): {:.6}", var_mid);
    println!("Weight variance at end (beta=1.0): {:.6}", var_end);

    // Weight variance should increase as beta increases (stronger bias correction)
    assert!(
        var_end >= var_start * 0.9,
        "Weight variance should increase with beta annealing, start: {:.6}, end: {:.6}",
        var_start,
        var_end
    );

    Ok(())
}

/// Test 5: DQN agent with PER enabled trains without errors
#[test]
fn test_dqn_with_per_trains() -> Result<(), MLError> {
    let config = create_per_config(true);
    let mut agent = WorkingDQN::new(config)?;

    // Add enough experiences for training
    for i in 0..150 {
        let state = vec![i as f32 * 0.01; 128];
        let action = (i % 45) as u8;
        let reward = (i as f32) * 0.1;
        let next_state = vec![(i + 1) as f32 * 0.01; 128];
        let done = i % 50 == 49;

        let exp = Experience::new(state, action, reward, next_state, done);
        agent.store_experience(exp)?;
    }

    // Train for 10 steps
    for _ in 0..10 {
        match agent.train_step(None) {
            Ok((loss, grad_norm)) => {
                assert!(loss >= 0.0, "Loss should be non-negative, got: {}", loss);
                assert!(grad_norm >= 0.0, "Gradient norm should be non-negative, got: {}", grad_norm);
            }
            Err(e) => {
                panic!("Training step failed: {:?}", e);
            }
        }
    }

    println!("✓ DQN with PER trained successfully for 10 steps");

    Ok(())
}

/// Test 6: Verify PER vs standard replay performance (convergence speed)
#[test]
fn test_per_vs_standard_convergence() -> Result<(), MLError> {
    // Train two agents: one with PER, one without
    let config_per = create_per_config(true);
    let config_standard = create_per_config(false);

    let mut agent_per = WorkingDQN::new(config_per)?;
    let mut agent_standard = WorkingDQN::new(config_standard)?;

    // Add same experiences to both agents
    for i in 0..200 {
        let state = vec![i as f32 * 0.01; 128];
        let action = (i % 45) as u8;
        let reward = (i as f32) * 0.1;
        let next_state = vec![(i + 1) as f32 * 0.01; 128];
        let done = i % 50 == 49;

        let exp = Experience::new(state.clone(), action, reward, next_state.clone(), done);
        agent_per.store_experience(exp.clone())?;
        agent_standard.store_experience(exp)?;
    }

    // Train both agents for 20 steps and track loss
    let mut loss_per_total = 0.0;
    let mut loss_standard_total = 0.0;

    for _ in 0..20 {
        if let Ok((loss, _)) = agent_per.train_step(None) {
            loss_per_total += loss as f64;
        }
        if let Ok((loss, _)) = agent_standard.train_step(None) {
            loss_standard_total += loss as f64;
        }
    }

    let avg_loss_per = loss_per_total / 20.0;
    let avg_loss_standard = loss_standard_total / 20.0;

    println!("Average loss with PER: {:.6}", avg_loss_per);
    println!("Average loss with standard replay: {:.6}", avg_loss_standard);

    // Note: We can't guarantee PER always has lower loss in this short test,
    // but we can verify both agents train successfully
    assert!(avg_loss_per >= 0.0 && avg_loss_per < 1000.0);
    assert!(avg_loss_standard >= 0.0 && avg_loss_standard < 1000.0);

    println!("✓ Both PER and standard replay agents trained successfully");

    Ok(())
}