foxhunt/ml/tests/dqn_full_gradient_flow_integration_test.rs

// Integration tests to replicate DQN gradient collapse in production setting
// Purpose: Verify gradient flow through full training loop
// Created: 2025-11-21 (Test-Driven Development Campaign)

use candle_core::{DType, Device, Tensor};
use ml::dqn::config::DQNConfig;
use ml::dqn::dqn::DistributionalDuelingConfig;
use ml::dqn::replay_buffer_type::ReplayBufferType;
use ml::error::MLError;
use ml::trainers::dqn::{DQNTrainer, DQNTrainerConfig};
use std::sync::Arc;

/// Helper function to create minimal DQN config for testing
fn create_test_dqn_config() -> DQNTrainerConfig {
    DQNTrainerConfig {
        dqn_config: DQNConfig {
            learning_rate: 0.001,
            gamma: 0.99,
            batch_size: 32,
            buffer_size: 1000,
            tau: 0.005,
            epsilon_start: 1.0,
            epsilon_end: 0.01,
            epsilon_decay: 0.995,
            target_update_frequency: 100,
            warmup_steps: 100,
            gradient_clip: 100.0,
            distributional: Some(DistributionalDuelingConfig {
                input_dim: 54,
                hidden_dims: vec![128, 64],
                n_actions: 45,
                n_atoms: 51,
                v_min: -2.0,
                v_max: 2.0,
            }),
            regime_conditional: None,
            replay_buffer_type: ReplayBufferType::PrioritizedWithAdaptive {
                alpha: 0.6,
                beta_start: 0.4,
                beta_frames: 100000,
                adaptive_config: None,
            },
            per_alpha: 0.6,
            per_beta_start: 0.4,
            per_beta_frames: 100000,
            per_epsilon: 1e-6,
            n_step: 3,
            use_soft_update: true,
            use_double_dqn: true,
            use_dueling: true,
        },
        parquet_path: "test_data/ES_FUT_180d.parquet".to_string(),
        checkpoint_dir: None,
        log_interval: 10,
        checkpoint_interval: 100,
        num_epochs: 1,
        reward_config: None,
    }
}

/// Test 6: Full Training Loop (10 Steps)
/// Goal: Quick integration test with real DQN trainer
/// Expected: ALL 10 steps have grad_norm > 0.001
#[test]
#[ignore] // Ignore by default (requires parquet file)
fn test_full_training_loop_gradient_flow() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("[Test 6] Initializing DQN trainer...");
    
    let config = create_test_dqn_config();
    
    // Initialize trainer
    let mut trainer = DQNTrainer::new(config, device)?;
    
    println!("[Test 6] Running 10 training steps...");
    
    let mut grad_norms = Vec::new();
    let mut collapse_step = None;
    
    // Run 10 training steps
    for step in 0..10 {
        let metrics = trainer.train_step()?;
        
        let grad_norm = metrics.grad_norm.unwrap_or(0.0);
        grad_norms.push(grad_norm);
        
        println!(
            "[Test 6] Step {}: loss={:.6}, grad_norm={:.6}, q_value={:.2}",
            step, metrics.loss, grad_norm, metrics.avg_q_value.unwrap_or(0.0)
        );
        
        // Check for gradient collapse
        if grad_norm < 0.001 && collapse_step.is_none() {
            collapse_step = Some(step);
            println!("[Test 6] ⚠️  GRADIENT COLLAPSE detected at step {}", step);
        }
    }
    
    // Summary statistics
    let max_grad = grad_norms.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let min_grad = grad_norms.iter().cloned().fold(f32::INFINITY, f32::min);
    let avg_grad = grad_norms.iter().sum::<f32>() / grad_norms.len() as f32;
    
    println!("[Test 6] Gradient statistics:");
    println!("[Test 6]   Max: {:.6}", max_grad);
    println!("[Test 6]   Min: {:.6}", min_grad);
    println!("[Test 6]   Avg: {:.6}", avg_grad);
    
    // ASSERTION: No gradient collapse should occur
    if let Some(step) = collapse_step {
        panic!(
            "Test 6 FAILED: Gradient collapse at step {} (grad_norm < 0.001). \
             Gradient norms: {:?}",
            step, grad_norms
        );
    }
    
    // ASSERTION: All steps should have grad_norm > 0.001
    let below_threshold: Vec<(usize, f32)> = grad_norms.iter()
        .enumerate()
        .filter(|(_, &gn)| gn < 0.001)
        .map(|(i, &gn)| (i, gn))
        .collect();
    
    assert!(
        below_threshold.is_empty(),
        "Test 6 FAILED: {} steps have grad_norm < 0.001: {:?}",
        below_threshold.len(),
        below_threshold
    );
    
    println!("[Test 6] ✅ PASSED: All 10 steps have healthy gradients");
    Ok(())
}

/// Test 7: Categorical Loss Integration
/// Goal: Test with full C51 categorical loss path in production setting
/// Expected: grad_norm > 0.001 for C51
#[test]
#[ignore] // Ignore by default (requires parquet file)
fn test_categorical_loss_integration() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("[Test 7] Testing C51 categorical loss integration...");
    
    let config = create_test_dqn_config();
    
    // Ensure distributional is enabled
    assert!(
        config.dqn_config.distributional.is_some(),
        "Test 7 requires distributional DQN"
    );
    
    // Initialize trainer
    let mut trainer = DQNTrainer::new(config, device)?;
    
    // Run single training step
    let metrics = trainer.train_step()?;
    
    let grad_norm = metrics.grad_norm.unwrap_or(0.0);
    
    println!("[Test 7] Metrics after C51 training step:");
    println!("[Test 7]   Loss: {:.6}", metrics.loss);
    println!("[Test 7]   Grad norm: {:.6}", grad_norm);
    println!("[Test 7]   Q-value: {:.2}", metrics.avg_q_value.unwrap_or(0.0));
    
    // ASSERTION: Gradient norm should be healthy
    assert!(
        grad_norm > 0.001,
        "Test 7 FAILED: grad_norm {:.6} too small with C51 loss (expected > 0.001)",
        grad_norm
    );
    
    // ASSERTION: Loss should be finite
    assert!(
        metrics.loss.is_finite(),
        "Test 7 FAILED: Loss is not finite: {}",
        metrics.loss
    );
    
    println!("[Test 7] ✅ PASSED: C51 categorical loss preserves gradients");
    Ok(())
}

/// Test 8: 100-Step Endurance Test
/// Goal: Replicate exact production gradient collapse at step 100
/// Expected: NEVER drops below 0.001
#[test]
#[ignore] // Ignore by default (long-running test)
fn test_gradient_endurance_100_steps() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("[Test 8] Running 100-step endurance test...");
    println!("[Test 8] This replicates the exact production failure scenario");
    
    let config = create_test_dqn_config();
    
    // Initialize trainer
    let mut trainer = DQNTrainer::new(config, device)?;
    
    let mut grad_norms = Vec::new();
    let mut collapse_step = None;
    let monitor_interval = 10;
    
    // Run 100 training steps (exact production failure point)
    for step in 0..100 {
        let metrics = trainer.train_step()?;
        
        let grad_norm = metrics.grad_norm.unwrap_or(0.0);
        grad_norms.push(grad_norm);
        
        // Log every 10 steps
        if step % monitor_interval == 0 {
            println!(
                "[Test 8] Step {}: loss={:.6}, grad_norm={:.6}, q_value={:.2}",
                step, metrics.loss, grad_norm, metrics.avg_q_value.unwrap_or(0.0)
            );
        }
        
        // Check for gradient collapse
        if grad_norm < 0.001 && collapse_step.is_none() {
            collapse_step = Some(step);
            println!("[Test 8] ⚠️  GRADIENT COLLAPSE detected at step {}", step);
            
            // Continue running to observe behavior post-collapse
        }
    }
    
    // Calculate gradient statistics over time
    let window_size = 10;
    let mut window_avgs = Vec::new();
    
    for i in 0..(grad_norms.len() / window_size) {
        let start = i * window_size;
        let end = (i + 1) * window_size;
        let window = &grad_norms[start..end];
        let avg = window.iter().sum::<f32>() / window.len() as f32;
        window_avgs.push(avg);
        
        println!(
            "[Test 8] Steps {}-{}: avg_grad_norm={:.6}",
            start, end - 1, avg
        );
    }
    
    // Summary statistics
    let max_grad = grad_norms.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let min_grad = grad_norms.iter().cloned().fold(f32::INFINITY, f32::min);
    let avg_grad = grad_norms.iter().sum::<f32>() / grad_norms.len() as f32;
    
    println!("[Test 8] Overall gradient statistics:");
    println!("[Test 8]   Max: {:.6}", max_grad);
    println!("[Test 8]   Min: {:.6}", min_grad);
    println!("[Test 8]   Avg: {:.6}", avg_grad);
    
    // CRITICAL ASSERTION: Report exact failure step if collapse occurs
    if let Some(step) = collapse_step {
        // Generate detailed failure report
        println!("\n[Test 8] ❌ GRADIENT COLLAPSE DETECTED ❌");
        println!("[Test 8] Collapse occurred at step: {}", step);
        println!("[Test 8] This matches production failure scenario (step 100+)");
        println!("\n[Test 8] Gradient progression:");
        
        // Show gradients around collapse point
        let context_range = 5;
        let start = step.saturating_sub(context_range);
        let end = (step + context_range).min(grad_norms.len());
        
        for i in start..end {
            let marker = if i == step { " <<< COLLAPSE" } else { "" };
            println!("[Test 8]   Step {}: {:.6}{}", i, grad_norms[i], marker);
        }
        
        panic!(
            "Test 8 FAILED: Gradient collapse at step {} (grad_norm={:.6} < 0.001). \
             This replicates the production failure. Root cause investigation needed.",
            step, grad_norms[step]
        );
    }
    
    // ASSERTION: All steps should maintain healthy gradients
    let below_threshold: Vec<(usize, f32)> = grad_norms.iter()
        .enumerate()
        .filter(|(_, &gn)| gn < 0.001)
        .map(|(i, &gn)| (i, gn))
        .collect();
    
    assert!(
        below_threshold.is_empty(),
        "Test 8 FAILED: {} steps have grad_norm < 0.001: {:?}",
        below_threshold.len(),
        below_threshold
    );
    
    println!("[Test 8] ✅ PASSED: All 100 steps maintained healthy gradients");
    println!("[Test 8] Production gradient collapse issue is RESOLVED");
    Ok(())
}

/// Test 9: Dtype Conversion Safety Test
/// Goal: Verify that dtype conversions in gradient path don't break autograd
/// Expected: Gradients should flow through F32 conversion
#[test]
fn test_dtype_conversion_gradient_safety() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("[Test 9] Testing dtype conversion safety in gradient path...");
    
    // Create simple tensor computation with dtype conversion
    let input = Tensor::randn(0f64, 1.0, (32, 64), &device)?;
    
    // Convert to F32 (mimics line 1087 in dqn.rs)
    let converted = input.to_dtype(DType::F32)?;
    
    // Compute loss
    let loss = converted.mean_all()?;
    
    println!("[Test 9] Loss value: {:?}", loss.to_scalar::<f32>()?);
    
    // Backward pass
    let grads = loss.backward()?;
    
    // Check if gradient exists for input tensor
    // NOTE: In Candle, to_dtype may or may not preserve gradient link
    // This test verifies actual behavior
    
    if let Some(grad) = grads.get(&input) {
        let grad_norm = grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f64>()?;
        println!("[Test 9] Input gradient norm: {:.6}", grad_norm);
        
        assert!(
            grad_norm > 1e-6,
            "Test 9: Gradient exists but is too small: {:.6}",
            grad_norm
        );
        
        println!("[Test 9] ✅ PASSED: Dtype conversion preserves gradients");
    } else {
        println!("[Test 9] ⚠️  WARNING: Dtype conversion breaks gradient link");
        println!("[Test 9] This explains production gradient collapse!");
        println!("[Test 9] FIX: Ensure network outputs F32 natively, avoid conversion");
        
        panic!(
            "Test 9 FAILED: to_dtype(DType::F32) breaks autograd graph. \
             This is likely the root cause of gradient collapse in production. \
             RECOMMENDATION: Remove dtype conversion on line 1087 of dqn.rs"
        );
    }
    
    Ok(())
}