foxhunt/ml/tests/dqn_gradient_collapse_bug41_test.rs

//! BUG #41 Gradient Collapse Fix Validation Test
//!
//! Tests that the `.detach()` fix at `ml/src/dqn/dqn.rs:1207` resolves
//! gradient collapse in Prioritized Experience Replay (PER).
//!
//! **Problem**: TD errors were converted to Vec for PER priority updates without
//! detaching, creating a gradient graph fork that confused Candle's autograd,
//! resulting in zero gradients (0.000000).
//!
//! **Fix**: Added `.detach()` before `.to_vec1()` to make TD errors value-only
//! (no gradients), allowing the main loss gradient path to flow normally.
//!
//! **Expected Results**:
//! - Gradient norms: >0.001 (not 0.000000)
//! - Loss values: Should change across steps (learning is happening)
//! - TD errors: Should be diverse (not all identical)

use ml::dqn::{Experience, WorkingDQN, WorkingDQNConfig};
use ml::MLError;

#[test]
fn test_bug41_gradient_flow_with_per() -> Result<(), MLError> {
    println!("\n🧪 BUG #41 Gradient Collapse Test (PER)");
    println!("========================================");

    // Create DQN config with PER enabled (the critical code path)
    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 32;
    config.num_actions = 3;
    config.use_per = true; // CRITICAL: Test the PER code path
    config.per_alpha = 0.6;
    config.per_beta_start = 0.4;
    config.replay_buffer_capacity = 1000;
    config.batch_size = 32;
    config.min_replay_size = 100;
    config.learning_rate = 1e-4;
    config.warmup_steps = 0; // Disable warmup for fast testing

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

    // Generate synthetic experiences (enough for training)
    println!("\n📝 Generating 200 synthetic experiences...");
    for i in 0..200 {
        let state = vec![0.1 * (i as f32); 32];
        let next_state = vec![0.1 * (i as f32 + 1.0); 32];
        let action = (i % 3) as u8;
        let reward = if i % 10 == 0 { 1.0 } else { -0.1 };
        let done = i % 50 == 49;

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

    println!("✓ Generated 200 experiences (buffer size: {})", dqn.get_replay_buffer_size()?);

    // Run 5 training steps and collect metrics
    println!("\n🔄 Running 5 training steps...");
    let mut gradient_norms = Vec::new();
    let mut losses = Vec::new();

    for step in 1..=5 {
        let (loss, grad_norm) = dqn.train_step(None)?;
        gradient_norms.push(grad_norm);
        losses.push(loss);

        println!(
            "  Step {}: loss={:.6}, grad_norm={:.6}",
            step, loss, grad_norm
        );
    }

    // ========================================
    // VALIDATION 1: Gradient Flow (CRITICAL)
    // ========================================
    println!("\n✅ VALIDATION 1: Gradient Flow");
    let all_gradients_nonzero = gradient_norms.iter().all(|&g| g > 0.001);

    if all_gradients_nonzero {
        println!("  ✓ ALL gradient norms > 0.001 (gradients are flowing)");
        for (i, &g) in gradient_norms.iter().enumerate() {
            println!("    Step {}: {:.6}", i + 1, g);
        }
    } else {
        println!("  ✗ FAILED: Some gradient norms are near zero!");
        for (i, &g) in gradient_norms.iter().enumerate() {
            let status = if g > 0.001 { "✓" } else { "✗" };
            println!("    Step {}: {:.6} {}", i + 1, g, status);
        }
    }

    assert!(
        all_gradients_nonzero,
        "BUG #41 REGRESSION: Gradient collapse detected! All gradients should be > 0.001, got: {:?}",
        gradient_norms
    );

    // ========================================
    // VALIDATION 2: Loss Variation
    // ========================================
    println!("\n✅ VALIDATION 2: Loss Variation");
    let loss_min = losses.iter().cloned().fold(f32::INFINITY, f32::min);
    let loss_max = losses.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let loss_range = loss_max - loss_min;

    println!("  Loss range: [{:.6}, {:.6}] (delta: {:.6})", loss_min, loss_max, loss_range);

    // Loss should vary (not stuck at same value) - use 0.0001 threshold for small synthetic data
    // With real market data, variation will be much larger
    assert!(
        loss_range > 0.0001,
        "BUG #41 REGRESSION: Loss is stuck! Loss should vary across steps (threshold: 0.0001), got range: {:.6}",
        loss_range
    );
    println!("  ✓ Loss is changing (learning is happening)");

    // ========================================
    // VALIDATION 3: Gradient Magnitude
    // ========================================
    println!("\n✅ VALIDATION 3: Gradient Magnitude");
    let grad_min = gradient_norms.iter().cloned().fold(f32::INFINITY, f32::min);
    let grad_max = gradient_norms.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let grad_mean = gradient_norms.iter().sum::<f32>() / gradient_norms.len() as f32;

    println!("  Gradient stats:");
    println!("    Min:  {:.6}", grad_min);
    println!("    Max:  {:.6}", grad_max);
    println!("    Mean: {:.6}", grad_mean);

    // Gradients should be in a reasonable range (not collapsed to zero) - use 0.001 threshold
    // With synthetic data, gradients are small but non-zero (which is correct)
    assert!(
        grad_mean > 0.001,
        "BUG #41 REGRESSION: Average gradient too small ({:.6} < 0.001), indicates gradient collapse",
        grad_mean
    );
    println!("  ✓ Gradients are in healthy range");

    println!("\n========================================");
    println!("🎉 BUG #41 TEST PASSED");
    println!("✓ Gradient flow confirmed with PER");
    println!("✓ .detach() fix is working correctly");
    println!("========================================\n");

    Ok(())
}

#[test]
fn test_bug41_per_vs_uniform_replay() -> Result<(), MLError> {
    println!("\n🧪 BUG #41 Comparison: PER vs Uniform Replay");
    println!("===========================================");

    // Test with PER enabled
    let mut per_config = WorkingDQNConfig::conservative();
    per_config.state_dim = 32;
    per_config.num_actions = 3;
    per_config.use_per = true;
    per_config.batch_size = 32;
    per_config.min_replay_size = 100;
    per_config.warmup_steps = 0;

    let mut dqn_per = WorkingDQN::new(per_config)?;

    // Test with uniform replay (control)
    let mut uniform_config = WorkingDQNConfig::conservative();
    uniform_config.state_dim = 32;
    uniform_config.num_actions = 3;
    uniform_config.use_per = false;
    uniform_config.batch_size = 32;
    uniform_config.min_replay_size = 100;
    uniform_config.warmup_steps = 0;

    let mut dqn_uniform = WorkingDQN::new(uniform_config)?;

    // Generate identical experiences for both
    for i in 0..200 {
        let state = vec![0.1 * (i as f32); 32];
        let next_state = vec![0.1 * (i as f32 + 1.0); 32];
        let action = (i % 3) as u8;
        let reward = if i % 10 == 0 { 1.0 } else { -0.1 };
        let done = i % 50 == 49;

        let experience = Experience::new(state.clone(), action, reward, next_state.clone(), done);
        dqn_per.store_experience(experience.clone())?;
        dqn_uniform.store_experience(experience)?;
    }

    // Train both and compare gradient flow
    println!("\n📊 Training PER...");
    let mut per_gradients = Vec::new();
    for step in 1..=3 {
        let (loss, grad_norm) = dqn_per.train_step(None)?;
        per_gradients.push(grad_norm);
        println!("  Step {}: grad_norm={:.6}", step, grad_norm);
    }

    println!("\n📊 Training Uniform...");
    let mut uniform_gradients = Vec::new();
    for step in 1..=3 {
        let (loss, grad_norm) = dqn_uniform.train_step(None)?;
        uniform_gradients.push(grad_norm);
        println!("  Step {}: grad_norm={:.6}", step, grad_norm);
    }

    // Both should have non-zero gradients
    let per_all_nonzero = per_gradients.iter().all(|&g| g > 0.001);
    let uniform_all_nonzero = uniform_gradients.iter().all(|&g| g > 0.001);

    println!("\n✅ RESULTS:");
    println!("  PER gradients:     {} (all > 0.001)", if per_all_nonzero { "✓" } else { "✗" });
    println!("  Uniform gradients: {} (all > 0.001)", if uniform_all_nonzero { "✓" } else { "✗" });

    assert!(per_all_nonzero, "BUG #41 REGRESSION: PER gradients collapsed!");
    assert!(uniform_all_nonzero, "Uniform replay gradients collapsed (control test failed)!");

    println!("\n🎉 Both PER and Uniform replay show healthy gradients");
    println!("✓ BUG #41 fix does not break uniform replay");
    println!("===========================================\n");

    Ok(())
}