foxhunt/ml/tests/dqn_q_value_stability_test.rs

//! Q-Value Stability Tests for DQN
//!
//! Tests for gradient clipping and Huber loss implementation to prevent
//! extreme Q-value variance and training instability.

use candle_core::{Device, Tensor};
use ml::dqn::{Experience, WorkingDQN, WorkingDQNConfig};

/// Test 1: Gradient clipping ensures gradient norm <= max_norm
#[test]
fn test_gradient_clipping_norm() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 4;
    config.batch_size = 4;
    config.state_dim = 52;
    config.learning_rate = 0.001; // Higher LR to trigger gradient clipping

    let mut dqn = WorkingDQN::new(config)?;
    let device = dqn.device().clone();

    // Add experiences with extreme rewards to trigger large gradients
    for i in 0..10 {
        let extreme_reward = if i % 2 == 0 { 10000.0 } else { -10000.0 };
        let experience = Experience::new(
            vec![i as f32 * 0.1; 52],
            (i % 3) as u8,
            extreme_reward,
            vec![(i + 1) as f32 * 0.1; 52],
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Train step should apply gradient clipping
    let (loss, _grad_norm) = dqn.train_step(None)?;

    // After clipping, loss should be finite and bounded
    assert!(
        loss.is_finite(),
        "Loss should be finite after gradient clipping"
    );
    assert!(loss >= 0.0, "Loss should be non-negative");

    // Verify gradients are clipped by checking loss doesn't explode
    // With extreme rewards, unclipped gradients would cause NaN/Inf
    assert!(loss < 1e6, "Loss should not explode with gradient clipping");

    Ok(())
}

/// Test 2: Huber loss vs MSE - Huber is more robust to outliers
#[test]
fn test_huber_loss_vs_mse() -> anyhow::Result<()> {
    let device = Device::cuda_if_available(0)?;

    // Create prediction and target with one outlier
    let prediction = Tensor::from_vec(
        vec![1.0_f32, 2.0, 3.0, 100.0], // 100.0 is outlier
        4,
        &device,
    )?;
    let target = Tensor::from_vec(
        vec![1.1_f32, 2.1, 3.1, 3.5], // Target for outlier is 3.5
        4,
        &device,
    )?;

    // MSE loss
    let diff = prediction.sub(&target)?;
    let mse_loss = (&diff * &diff)?.mean_all()?;
    let mse_value = mse_loss.to_scalar::<f32>()?;

    // Huber loss (delta=1.0)
    let huber_loss = huber_loss_fn(&prediction, &target, 1.0)?;
    let huber_value = huber_loss.to_scalar::<f32>()?;

    // Huber loss should be smaller than MSE for outliers
    // MSE squares the error (96.5^2 = 9312), Huber clips it
    assert!(
        huber_value < mse_value,
        "Huber loss ({:.2}) should be less than MSE ({:.2}) with outliers",
        huber_value,
        mse_value
    );

    println!("MSE: {:.4}, Huber: {:.4}", mse_value, huber_value);

    Ok(())
}

/// Helper: Huber loss implementation
fn huber_loss_fn(
    prediction: &Tensor,
    target: &Tensor,
    delta: f64,
) -> Result<Tensor, candle_core::Error> {
    let diff = (prediction - target)?;
    let abs_diff = diff.abs()?;

    // MSE region: |diff| <= delta
    let mse_mask = abs_diff.le(delta)?;
    let mse_loss = (diff.sqr()? * 0.5)?;

    // MAE region: |diff| > delta
    let mae_mask = abs_diff.gt(delta)?;
    let mae_loss = (abs_diff * delta - delta * delta * 0.5)?;

    // Combine losses
    let loss = ((mse_loss * mse_mask.to_dtype(prediction.dtype())?)?
        + (mae_loss * mae_mask.to_dtype(prediction.dtype())?)?)?;

    loss.mean_all()
}

/// Test 3: Q-values remain bounded after training
#[test]
fn test_q_values_bounded() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 4;
    config.batch_size = 4;
    config.state_dim = 52;

    let mut dqn = WorkingDQN::new(config)?;
    let device = dqn.device().clone();

    // Add normal experiences
    for i in 0..50 {
        let experience = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 3) as u8,
            (i as f32) * 0.1, // Normal rewards
            vec![(i + 1) as f32 * 0.01; 52],
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Train for several steps
    for _ in 0..10 {
        let _ = dqn.train_step(None)?;
    }

    // Check Q-values for a sample state
    let test_state = Tensor::from_vec(vec![0.5_f32; 52], (1, 52), &device)?;

    let q_values = dqn.forward(&test_state)?;
    let q_vec = q_values.to_vec2::<f32>()?;

    // Q-values should be bounded (not extreme)
    for &q in &q_vec[0] {
        assert!(
            q.abs() < 1000.0,
            "Q-value {:.2} exceeds reasonable bounds",
            q
        );
        assert!(q.is_finite(), "Q-value should be finite");
    }

    println!("Q-values after training: {:?}", q_vec[0]);

    Ok(())
}

/// Test 4: No NaN/Inf values in Q-values
#[test]
fn test_no_nan_inf_q_values() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 4;
    config.batch_size = 4;
    config.state_dim = 52;
    config.learning_rate = 0.01; // Aggressive LR to stress test

    let mut dqn = WorkingDQN::new(config)?;
    let device = dqn.device().clone();

    // Add experiences with varied rewards
    for i in 0..20 {
        let experience = Experience::new(
            vec![i as f32 * 0.1; 52],
            (i % 3) as u8,
            (i as f32 - 10.0) * 10.0, // Rewards from -100 to +90
            vec![(i + 1) as f32 * 0.1; 52],
            i == 19,
        );
        dqn.store_experience(experience)?;
    }

    // Train for multiple steps
    for step in 0..20 {
        let (loss, _grad_norm) = dqn.train_step(None)?;

        // Loss should always be finite
        assert!(loss.is_finite(), "Loss is NaN/Inf at step {}", step);

        // Check Q-values periodically
        if step % 5 == 0 {
            let test_state = Tensor::from_vec(vec![0.0_f32; 52], (1, 52), &device)?;

            let q_values = dqn.forward(&test_state)?;
            let q_vec = q_values.to_vec2::<f32>()?;

            for (i, &q) in q_vec[0].iter().enumerate() {
                assert!(
                    q.is_finite(),
                    "Q-value[{}] is NaN/Inf at step {}: {}",
                    i,
                    step,
                    q
                );
            }
        }
    }

    Ok(())
}

/// Test 5: Gradients don't explode with large Q-value updates
#[test]
fn test_gradients_no_explosion() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 4;
    config.batch_size = 4;
    config.state_dim = 52;
    config.learning_rate = 0.1; // Very high LR to test gradient clipping

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

    // Add experiences designed to cause large TD errors
    for i in 0..10 {
        let experience = Experience::new(
            vec![0.0_f32; 52], // Same state
            0,                 // Same action
            1000.0,            // Large reward
            vec![1.0_f32; 52], // Different next state
            false,
        );
        dqn.store_experience(experience)?;
    }

    // First training step (large initial error)
    let (loss1, _grad_norm) = dqn.train_step(None)?;
    assert!(loss1.is_finite(), "Initial loss should be finite");

    // Second training step (should be stable, not explode)
    let (loss2, _grad_norm) = dqn.train_step(None)?;
    assert!(loss2.is_finite(), "Second loss should be finite");

    // Loss shouldn't explode exponentially
    assert!(
        loss2 < loss1 * 10.0,
        "Loss exploded: {:.2} -> {:.2}",
        loss1,
        loss2
    );

    // Third step should remain stable
    let (loss3, _grad_norm) = dqn.train_step(None)?;
    assert!(loss3.is_finite(), "Third loss should be finite");
    assert!(
        loss3 < loss1 * 20.0,
        "Loss continues to explode: {:.2} -> {:.2} -> {:.2}",
        loss1,
        loss2,
        loss3
    );

    println!(
        "Loss trajectory: {:.4} -> {:.4} -> {:.4}",
        loss1, loss2, loss3
    );

    Ok(())
}

/// Test 6: Huber loss implementation correctness
#[test]
fn test_huber_loss_correctness() -> anyhow::Result<()> {
    let device = Device::cuda_if_available(0)?;
    let delta = 1.0;

    // Case 1: Small error (use MSE)
    let pred_small = Tensor::from_vec(vec![1.0_f32], 1, &device)?;
    let target_small = Tensor::from_vec(vec![1.5_f32], 1, &device)?;
    let huber_small = huber_loss_fn(&pred_small, &target_small, delta)?;

    // Expected: 0.5 * (0.5)^2 = 0.125
    let expected_small = 0.125_f32;
    let actual_small = huber_small.to_scalar::<f32>()?;
    assert!(
        (actual_small - expected_small).abs() < 0.01,
        "Huber loss for small error: expected {:.3}, got {:.3}",
        expected_small,
        actual_small
    );

    // Case 2: Large error (use MAE)
    let pred_large = Tensor::from_vec(vec![1.0_f32], 1, &device)?;
    let target_large = Tensor::from_vec(vec![5.0_f32], 1, &device)?;
    let huber_large = huber_loss_fn(&pred_large, &target_large, delta)?;

    // Expected: |4.0| * 1.0 - 0.5 * 1.0^2 = 4.0 - 0.5 = 3.5
    let expected_large = 3.5_f32;
    let actual_large = huber_large.to_scalar::<f32>()?;
    assert!(
        (actual_large - expected_large).abs() < 0.01,
        "Huber loss for large error: expected {:.3}, got {:.3}",
        expected_large,
        actual_large
    );

    Ok(())
}

/// Test 7: Gradient clipping preserves learning direction
#[test]
fn test_gradient_clipping_preserves_direction() -> anyhow::Result<()> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.min_replay_size = 4;
    config.batch_size = 4;
    config.state_dim = 52;
    config.learning_rate = 0.001;

    let mut dqn = WorkingDQN::new(config)?;
    let device = dqn.device().clone();

    // Add experiences with consistent positive rewards
    for i in 0..20 {
        let experience = Experience::new(
            vec![i as f32 * 0.1; 52],
            0,    // Always Buy action
            10.0, // Consistent positive reward
            vec![(i + 1) as f32 * 0.1; 52],
            false,
        );
        dqn.store_experience(experience)?;
    }

    // Get initial Q-value for Buy action
    let test_state = Tensor::from_vec(vec![0.5_f32; 52], (1, 52), &device)?;

    let q_before = dqn.forward(&test_state)?;
    let q_buy_before = q_before.to_vec2::<f32>()?[0][0];

    // Train for several steps
    for _ in 0..10 {
        let _ = dqn.train_step(None)?;
    }

    // Q-value for Buy should increase (positive rewards)
    let q_after = dqn.forward(&test_state)?;
    let q_buy_after = q_after.to_vec2::<f32>()?[0][0];

    println!(
        "Q(Buy) before: {:.4}, after: {:.4}",
        q_buy_before, q_buy_after
    );

    // With gradient clipping, learning should still progress
    // (may be slower but direction preserved)
    // Allow for some variance due to exploration
    assert!(
        q_buy_after > q_buy_before - 1.0,
        "Q-value should not decrease significantly with positive rewards"
    );

    Ok(())
}