foxhunt/ml/tests/dqn_gradient_clipping_integration.rs

//! Integration test for DQN gradient clipping (WAVE 26 - Agent 16)
//!
//! Verifies that gradient clipping is enabled and prevents gradient explosion

use ml::dqn::{DQNAgent, DQNConfig, Experience, TradingState};
use ml::MLError;

#[tokio::test]
async fn test_gradient_clipping_enabled() -> Result<(), MLError> {
    // Create agent with default config
    let config = DQNConfig::default();
    let mut agent = DQNAgent::new(config)?;

    // Add enough experiences to enable training
    for i in 0..500 {
        let state = vec![i as f32 / 100.0; 52];
        let next_state = vec![(i + 1) as f32 / 100.0; 52];

        let experience = Experience::new(
            state,
            0, // Buy action
            100.0, // reward
            next_state,
            i % 100 == 0, // terminal every 100 steps
        );

        agent.store_experience(experience)?;
    }

    // Verify agent is ready for training
    assert!(agent.is_ready_for_training());

    // Perform training steps and verify no gradient explosion
    let mut loss_values = Vec::new();

    for _ in 0..10 {
        let loss = agent.train()?;
        loss_values.push(loss);

        // Verify loss is finite (not NaN or Inf from gradient explosion)
        assert!(loss.is_finite(), "Loss should be finite with gradient clipping");

        // Verify loss doesn't explode beyond reasonable bounds
        assert!(loss < 1e6, "Loss should not explode with gradient clipping: {}", loss);
    }

    // Verify gradient clipping prevents divergence
    // With clipping, losses should not increase exponentially
    let first_half_avg = loss_values.iter().take(5).sum::<f64>() / 5.0;
    let second_half_avg = loss_values.iter().skip(5).sum::<f64>() / 5.0;

    // Loss shouldn't explode by 100x (gradient clipping should prevent this)
    assert!(
        second_half_avg < first_half_avg * 100.0,
        "Gradient clipping should prevent loss explosion: first={}, second={}",
        first_half_avg,
        second_half_avg
    );

    Ok(())
}

#[tokio::test]
async fn test_gradient_clipping_with_extreme_rewards() -> Result<(), MLError> {
    // Test that gradient clipping handles extreme reward scenarios
    let config = DQNConfig::default();
    let mut agent = DQNAgent::new(config)?;

    // Create experiences with extreme rewards (which could cause gradient explosion)
    for i in 0..500 {
        let state = vec![i as f32 / 100.0; 52];
        let next_state = vec![(i + 1) as f32 / 100.0; 52];

        // Alternate between extreme positive and negative rewards
        let reward = if i % 2 == 0 { 10000.0 } else { -10000.0 };

        let experience = Experience::new(
            state,
            (i % 3) as u8, // Cycle through actions
            reward,
            next_state,
            i % 100 == 0,
        );

        agent.store_experience(experience)?;
    }

    // Train with extreme rewards - gradient clipping should prevent explosion
    for iteration in 0..10 {
        let loss = agent.train()?;

        // Verify stability even with extreme rewards
        assert!(
            loss.is_finite() && loss < 1e8,
            "Iteration {}: Loss should remain stable with gradient clipping (got {})",
            iteration,
            loss
        );
    }

    Ok(())
}

#[tokio::test]
async fn test_gradient_clipping_max_norm_10() -> Result<(), MLError> {
    // Verify that max_norm=10.0 is being used (standard for DQN)
    let config = DQNConfig::default();
    let mut agent = DQNAgent::new(config)?;

    // Add training data
    for i in 0..500 {
        let experience = Experience::new(
            vec![i as f32; 52],
            (i % 3) as u8,
            (i as f32).sin() * 100.0, // Oscillating rewards
            vec![(i + 1) as f32; 52],
            i % 100 == 0,
        );
        agent.store_experience(experience)?;
    }

    // Multiple training steps should converge with max_norm=10.0
    let initial_loss = agent.train()?;

    for _ in 0..20 {
        agent.train()?;
    }

    let final_loss = agent.train()?;

    // With proper gradient clipping, training should improve or stabilize
    // (not explode to infinity)
    assert!(
        final_loss.is_finite(),
        "Final loss should be finite with gradient clipping"
    );

    // Loss shouldn't increase by more than 10x (clipping prevents explosion)
    assert!(
        final_loss < initial_loss * 10.0,
        "Gradient clipping should prevent excessive loss increase: initial={}, final={}",
        initial_loss,
        final_loss
    );

    Ok(())
}

#[tokio::test]
async fn test_gradient_norm_monitoring() -> Result<(), MLError> {
    // Test that gradient norms are being monitored (via backward_step_with_monitoring)
    let config = DQNConfig::default();
    let mut agent = DQNAgent::new(config)?;

    // Add diverse training experiences
    for i in 0..500 {
        let state = TradingState::default();
        let next_state = TradingState::default();

        let experience = Experience::new(
            state.to_vector(),
            (i % 3) as u8,
            (i as f32 / 10.0).sin() * 50.0,
            next_state.to_vector(),
            i % 50 == 0,
        );

        agent.store_experience(experience)?;
    }

    // Train multiple times - monitoring should log gradient norms
    // (We can't directly verify logs, but we verify training completes successfully)
    for iteration in 0..5 {
        let loss = agent.train()?;

        // Verify training succeeds with monitoring enabled
        assert!(
            loss.is_finite(),
            "Iteration {}: Training should succeed with gradient monitoring",
            iteration
        );
    }

    Ok(())
}

#[tokio::test]
async fn test_no_gradient_explosion_during_convergence() -> Result<(), MLError> {
    // Test that gradient clipping prevents explosion during convergence phase
    let mut config = DQNConfig::default();
    config.learning_rate = 0.01; // Higher learning rate to stress-test clipping

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

    // Generate consistent training pattern (should converge with clipping)
    for i in 0..1000 {
        let state_value = (i % 100) as f32 / 100.0;
        let state = vec![state_value; 52];
        let next_state = vec![state_value + 0.01; 52];

        let experience = Experience::new(
            state,
            0, // Always buy
            10.0, // Consistent reward
            next_state,
            i % 100 == 99,
        );

        agent.store_experience(experience)?;
    }

    // Track loss progression
    let mut losses = Vec::new();
    for _ in 0..30 {
        let loss = agent.train()?;
        losses.push(loss);
    }

    // Verify no gradient explosion occurred
    for (i, &loss) in losses.iter().enumerate() {
        assert!(
            loss.is_finite() && loss < 1e6,
            "Step {}: Loss exploded despite gradient clipping: {}",
            i,
            loss
        );
    }

    // Verify losses are bounded (no exponential growth)
    let max_loss = losses.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    let min_loss = losses.iter().copied().fold(f64::INFINITY, f64::min);

    assert!(
        max_loss < min_loss * 1000.0,
        "Loss range too large (possible gradient explosion): min={}, max={}",
        min_loss,
        max_loss
    );

    Ok(())
}