foxhunt/ml/tests/dqn_gradient_collapse_root_cause_test.rs

// CRITICAL TEST: Isolate DQN gradient collapse root cause
// Purpose: Replicate production failure with dtype conversion hypothesis
// Created: 2025-11-21 - Test-Driven Development Campaign
//
// Based on Zen analysis: Line 1087 dtype conversion likely breaks gradient flow

use candle_core::{DType, Device, Tensor};
use ml::MLError;

/// Test 9: Dtype Conversion Gradient Safety (CRITICAL)
/// Goal: Verify that to_dtype(DType::F32) preserves gradients
/// Expected: Gradients flow through F32 conversion
/// 
/// THIS TEST REPLICATES THE EXACT SUSPECTED BUG FROM PRODUCTION CODE
/// Line 1087 in dqn.rs: state_action_values.to_dtype(DType::F32)?
#[test]
fn test_dtype_conversion_breaks_gradients() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("\n========================================");
    println!("TEST: Dtype Conversion Gradient Safety");
    println!("========================================\n");
    
    println!("[Test] Creating F64 tensor (mimics network output)...");
    
    // Create tensor in F64 (network might output this)
    let input_f64 = Tensor::randn(0f64, 1.0, (32, 51), &device)?;
    
    println!("[Test] Input tensor dtype: {:?}", input_f64.dtype());
    println!("[Test] Input tensor shape: {:?}", input_f64.shape());
    
    // === PATH 1: Direct loss (no conversion) ===
    println!("\n--- PATH 1: No dtype conversion ---");
    let loss1 = input_f64.mean_all()?;
    println!("[Test] Loss value: {:.6}", loss1.to_scalar::<f64>()?);
    
    let grads1 = loss1.backward()?;
    
    if let Some(grad) = grads1.get(&input_f64) {
        let grad_norm = grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f64>()?;
        println!("[Test] ✅ Path 1 (no conversion) - gradient exists: {:.6}", grad_norm);
    } else {
        println!("[Test] ❌ Path 1 (no conversion) - NO GRADIENT (unexpected!)");
    }
    
    // === PATH 2: Convert to F32 then loss (SUSPECTED BUG) ===
    println!("\n--- PATH 2: With to_dtype(DType::F32) conversion ---");
    
    // Re-create tensor for fresh computation graph
    let input_f64_v2 = Tensor::randn(0f64, 1.0, (32, 51), &device)?;
    
    // THIS IS THE SUSPECTED BUG LINE
    // Mimics: state_action_values.to_dtype(DType::F32)?
    let converted_f32 = input_f64_v2.to_dtype(DType::F32)?;
    
    println!("[Test] Converted tensor dtype: {:?}", converted_f32.dtype());
    
    let loss2 = converted_f32.mean_all()?;
    println!("[Test] Loss value: {:.6}", loss2.to_scalar::<f32>()?);
    
    let grads2 = loss2.backward()?;
    
    // CRITICAL TEST: Does gradient exist for ORIGINAL F64 tensor?
    if let Some(grad) = grads2.get(&input_f64_v2) {
        let grad_norm = grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f64>()?;
        println!("[Test] ✅ Path 2 (with conversion) - gradient exists on F64 tensor: {:.6}", grad_norm);
    } else {
        println!("[Test] ❌ Path 2 (with conversion) - NO GRADIENT on F64 tensor");
        println!("[Test] ⚠️  THIS IS THE BUG! to_dtype() breaks gradient link!");
    }
    
    // Also check converted tensor
    if let Some(grad) = grads2.get(&converted_f32) {
        let grad_norm = grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f32>()?;
        println!("[Test] ✅ Path 2 - gradient exists on F32 tensor: {:.6}", grad_norm);
    } else {
        println!("[Test] ❌ Path 2 - NO GRADIENT on F32 tensor either!");
    }
    
    println!("\n========================================");
    println!("ROOT CAUSE ANALYSIS:");
    println!("========================================");
    
    // CRITICAL ASSERTION
    let has_f64_gradient = grads2.get(&input_f64_v2).is_some();
    let has_f32_gradient = grads2.get(&converted_f32).is_some();
    
    if !has_f64_gradient && !has_f32_gradient {
        println!("❌ CATASTROPHIC: No gradients on EITHER tensor!");
        println!("   Candle's to_dtype() completely breaks autograd.");
        println!("   FIX: Ensure network outputs F32 natively.");
        
        panic!(
            "ROOT CAUSE IDENTIFIED: to_dtype(DType::F32) breaks gradient flow.\n\
             This is line 1087 in dqn.rs. Network must output F32 directly."
        );
    } else if !has_f64_gradient {
        println!("⚠️  GRADIENT DETACHMENT: F64 tensor loses gradients after conversion.");
        println!("   F32 tensor has gradients, but link to original computation is severed.");
        println!("   This explains production gradient collapse!");
        
        panic!(
            "ROOT CAUSE IDENTIFIED: to_dtype() detaches gradients from source tensor.\n\
             Original F64 tensor has no gradients after conversion.\n\
             FIX: Remove dtype conversion in gradient path (line 1087 dqn.rs)."
        );
    } else {
        println!("✅ Both tensors have gradients - dtype conversion is safe in Candle.");
        println!("   Root cause must be elsewhere (PER weights, categorical loss, etc.)");
    }
    
    Ok(())
}

/// Test: Network output dtype consistency
/// Goal: Verify DistributionalDuelingQNetwork outputs F32
#[test]
fn test_network_output_dtype_is_f32() -> Result<(), MLError> {
    use ml::dqn::{DistributionalDuelingConfig, DistributionalDuelingQNetwork};
    
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("\n========================================");
    println!("TEST: Network Output Dtype Verification");
    println!("========================================\n");
    
    let config = DistributionalDuelingConfig {
        state_dim: 54,
        num_actions: 45,
        num_atoms: 51,
        shared_hidden_dims: vec![128, 64],
        value_hidden_dim: 64,
        advantage_hidden_dim: 64,
        leaky_relu_alpha: 0.01,
    };
    
    let network = DistributionalDuelingQNetwork::new(config, device.clone())?;
    
    // Create input
    let input = Tensor::randn(0f32, 1.0, (32, 54), &device)?;
    
    println!("[Test] Input dtype: {:?}", input.dtype());
    
    // Forward pass
    let output = network.forward(&input)?;
    
    println!("[Test] Output dtype: {:?}", output.dtype());
    println!("[Test] Output shape: {:?}", output.shape());
    
    // ASSERTION: Network MUST output F32
    assert_eq!(
        output.dtype(),
        DType::F32,
        "Network output dtype mismatch! Expected F32, got {:?}. \
         This causes dtype conversion in training loop (line 1087).",
        output.dtype()
    );
    
    println!("\n✅ Network outputs F32 natively");
    println!("   Dtype conversion on line 1087 should be unnecessary!");
    
    Ok(())
}

/// Test: Gather operation preserves dtype
/// Goal: Verify gather() doesn't change dtype
#[test]
fn test_gather_preserves_dtype() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    println!("\n========================================");
    println!("TEST: Gather Operation Dtype Preservation");
    println!("========================================\n");
    
    // Create F32 tensor
    let tensor_f32 = Tensor::randn(0f32, 1.0, (32, 45, 51), &device)?;
    
    println!("[Test] Input dtype: {:?}", tensor_f32.dtype());
    
    // Create action indices
    let actions = Tensor::zeros((32, 1), DType::U32, &device)?;
    
    // Gather operation (mimics line 1085-1086 in dqn.rs)
    let gathered = tensor_f32.gather(&actions, 1)?.squeeze(1)?;
    
    println!("[Test] After gather+squeeze dtype: {:?}", gathered.dtype());
    
    // ASSERTION: Gather should preserve dtype
    assert_eq!(
        gathered.dtype(),
        DType::F32,
        "Gather changed dtype from F32 to {:?}!",
        gathered.dtype()
    );
    
    println!("\n✅ Gather preserves F32 dtype");
    println!("   Line 1087 conversion is redundant if input is F32!");
    
    Ok(())
}