foxhunt/ml/tests/dqn_gradient_flow_isolation_test.rs

// Test suite to isolate DQN gradient collapse root cause
// Purpose: Systematically verify gradient flow through individual components
// Created: 2025-11-21 (Test-Driven Development Campaign)

use candle_core::{DType, Device, Tensor, Var};
use candle_nn::{linear, AdamW, Linear, Optimizer, VarBuilder, VarMap};
use ml::dqn::dqn::{DistributionalDuelingConfig, DistributionalDuelingQNetwork};
use ml::error::MLError;

/// Test 1: Network Forward Pass Gradient Flow
/// Goal: Verify gradients flow through DistributionalDuelingQNetwork.forward()
/// Expected: grad_norm ∈ [0.01, 1.0]
#[test]
fn test_network_forward_has_gradients() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Create network with minimal config
    let config = DistributionalDuelingConfig {
        input_dim: 54,  // Standard DQN feature count
        hidden_dims: vec![128, 64],  // Small network for fast test
        n_actions: 45,
        n_atoms: 51,
        v_min: -2.0,
        v_max: 2.0,
    };
    
    let network = DistributionalDuelingQNetwork::new(config, device.clone())?;
    
    // Create dummy input [batch=8, features=54]
    let batch_size = 8;
    let input = Tensor::randn(0f32, 1.0, (batch_size, 54), &device)?;
    
    // Forward pass
    let output = network.forward(&input)?;
    
    // Compute simple loss: mean of all outputs
    let loss = output.mean_all()?;
    
    println!("[Test 1] Loss value: {:?}", loss.to_scalar::<f32>()?);
    
    // Backward pass
    let grads = loss.backward()?;
    
    // Check all network parameters have non-zero gradients
    let vars = network.vars();
    let all_vars = vars.all_vars();
    
    let mut total_grad_norm = 0.0f32;
    let mut zero_grad_count = 0;
    let mut total_param_count = 0;
    
    for (name, var) in all_vars.iter() {
        total_param_count += 1;
        
        if let Some(grad) = grads.get(var) {
            let grad_norm = grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f32>()?;
            total_grad_norm += grad_norm;
            
            if grad_norm < 1e-9 {
                zero_grad_count += 1;
                println!("[Test 1] WARNING: Zero gradient for parameter: {}", name);
            } else {
                println!("[Test 1] Parameter: {}, grad_norm: {:.6}", name, grad_norm);
            }
        } else {
            zero_grad_count += 1;
            println!("[Test 1] ERROR: No gradient for parameter: {}", name);
        }
    }
    
    let avg_grad_norm = total_grad_norm / total_param_count as f32;
    
    println!("[Test 1] Total parameters: {}", total_param_count);
    println!("[Test 1] Zero gradient count: {}", zero_grad_count);
    println!("[Test 1] Average grad_norm: {:.6}", avg_grad_norm);
    
    // ASSERTION: ALL parameters should have non-zero gradients
    assert_eq!(
        zero_grad_count, 0,
        "Test 1 FAILED: {} parameters have zero gradients (expected 0)",
        zero_grad_count
    );
    
    // ASSERTION: Average gradient norm should be in expected range
    assert!(
        avg_grad_norm >= 0.01 && avg_grad_norm <= 1.0,
        "Test 1 FAILED: avg_grad_norm {:.6} outside expected range [0.01, 1.0]",
        avg_grad_norm
    );
    
    println!("[Test 1] ✅ PASSED: All parameters have non-zero gradients");
    Ok(())
}

/// Test 2: Categorical Loss Gradient Flow
/// Goal: Verify categorical cross-entropy preserves gradients
/// Expected: grad_norm ∈ [0.1, 1.0]
#[test]
fn test_categorical_loss_gradients() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Create simple linear layer
    let vars = VarMap::new();
    let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);
    
    let input_dim = 64;
    let output_dim = 51;  // 51 atoms (C51 distribution)
    let batch_size = 32;
    
    let linear_layer = linear(input_dim, output_dim, vb.pp("test_layer"))?;
    
    // Create dummy input
    let input = Tensor::randn(0f32, 1.0, (batch_size, input_dim), &device)?;
    
    // Forward pass
    let logits = linear_layer.forward(&input)?;
    
    // Apply softmax to get distribution
    let probs = candle_nn::ops::softmax(&logits, 1)?;
    
    // Create target distribution (uniform for simplicity)
    let target = Tensor::ones((batch_size, output_dim), DType::F32, &device)?
        .broadcast_div(&Tensor::new(output_dim as f32, &device)?)?;
    
    // Categorical cross-entropy loss: -sum(target * log(probs))
    let log_probs = probs.log()?;
    let loss = target.mul(&log_probs)?.sum_all()?.neg()?;
    
    println!("[Test 2] Loss value: {:?}", loss.to_scalar::<f32>()?);
    
    // Backward pass
    let grads = loss.backward()?;
    
    // Check linear layer parameters have non-zero gradients
    let all_vars = vars.all_vars();
    
    let mut total_grad_norm = 0.0f32;
    let mut zero_grad_count = 0;
    let mut total_param_count = 0;
    
    for (name, var) in all_vars.iter() {
        total_param_count += 1;
        
        if let Some(grad) = grads.get(var) {
            let grad_norm = grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f32>()?;
            total_grad_norm += grad_norm;
            
            if grad_norm < 1e-9 {
                zero_grad_count += 1;
                println!("[Test 2] WARNING: Zero gradient for parameter: {}", name);
            } else {
                println!("[Test 2] Parameter: {}, grad_norm: {:.6}", name, grad_norm);
            }
        } else {
            zero_grad_count += 1;
            println!("[Test 2] ERROR: No gradient for parameter: {}", name);
        }
    }
    
    let avg_grad_norm = total_grad_norm / total_param_count as f32;
    
    println!("[Test 2] Total parameters: {}", total_param_count);
    println!("[Test 2] Zero gradient count: {}", zero_grad_count);
    println!("[Test 2] Average grad_norm: {:.6}", avg_grad_norm);
    
    // ASSERTION: ALL parameters should have non-zero gradients
    assert_eq!(
        zero_grad_count, 0,
        "Test 2 FAILED: {} parameters have zero gradients (expected 0)",
        zero_grad_count
    );
    
    // ASSERTION: Average gradient norm should be in expected range
    assert!(
        avg_grad_norm >= 0.001 && avg_grad_norm <= 10.0,  // Wider range for categorical loss
        "Test 2 FAILED: avg_grad_norm {:.6} outside expected range [0.001, 10.0]",
        avg_grad_norm
    );
    
    println!("[Test 2] ✅ PASSED: Categorical loss preserves gradients");
    Ok(())
}

/// Test 3: PER Weight Application
/// Goal: Verify importance sampling weights don't kill gradients
/// Expected: grad_norm identical ±10% WITH/WITHOUT IS weights
#[test]
fn test_per_weights_preserve_gradients() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Create simple network
    let vars = VarMap::new();
    let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);
    
    let linear_layer = linear(64, 1, vb.pp("test_layer"))?;
    
    let batch_size = 32;
    let input = Tensor::randn(0f32, 1.0, (batch_size, 64), &device)?;
    
    // === WITHOUT IS WEIGHTS ===
    let output1 = linear_layer.forward(&input)?;
    let loss1 = output1.sqr()?.mean_all()?;
    let grads1 = loss1.backward()?;
    
    let all_vars = vars.all_vars();
    let var1 = &all_vars[0].1;
    let grad1_norm = grads1.get(var1)
        .unwrap()
        .sqr()?
        .sum_all()?
        .sqrt()?
        .to_scalar::<f32>()?;
    
    println!("[Test 3] WITHOUT IS weights - grad_norm: {:.6}", grad1_norm);
    
    // === WITH IS WEIGHTS (DETACHED) ===
    // Re-create network for fresh computation graph
    let vars2 = VarMap::new();
    let vb2 = VarBuilder::from_varmap(&vars2, DType::F32, &device);
    let linear_layer2 = linear(64, 1, vb2.pp("test_layer"))?;
    
    let output2 = linear_layer2.forward(&input)?;
    
    // Create IS weights (should be detached)
    let is_weights = Tensor::ones((batch_size, 1), DType::F32, &device)?
        .mul(&Tensor::new(0.5f32, &device)?)?
        .detach();  // CRITICAL: Must detach weights
    
    // Apply IS weights
    let weighted_output = output2.broadcast_mul(&is_weights)?;
    let loss2 = weighted_output.sqr()?.mean_all()?;
    let grads2 = loss2.backward()?;
    
    let all_vars2 = vars2.all_vars();
    let var2 = &all_vars2[0].1;
    let grad2_norm = grads2.get(var2)
        .unwrap()
        .sqr()?
        .sum_all()?
        .sqrt()?
        .to_scalar::<f32>()?;
    
    println!("[Test 3] WITH IS weights (detached) - grad_norm: {:.6}", grad2_norm);
    
    // Calculate percentage difference
    let pct_diff = ((grad2_norm - grad1_norm) / grad1_norm).abs() * 100.0;
    println!("[Test 3] Percentage difference: {:.2}%", pct_diff);
    
    // ASSERTION: Gradients should exist for both cases
    assert!(
        grad1_norm > 1e-6 && grad2_norm > 1e-6,
        "Test 3 FAILED: Gradients too small (grad1: {:.6}, grad2: {:.6})",
        grad1_norm, grad2_norm
    );
    
    // ASSERTION: Gradient norms should be similar (accounting for weight scaling)
    // Since weights are 0.5, grad2 should be ~0.5 * grad1
    let expected_ratio = 0.5;
    let actual_ratio = grad2_norm / grad1_norm;
    let ratio_diff = ((actual_ratio - expected_ratio) / expected_ratio).abs() * 100.0;
    
    println!("[Test 3] Expected ratio: {:.2}, Actual ratio: {:.2}, Diff: {:.2}%", 
             expected_ratio, actual_ratio, ratio_diff);
    
    assert!(
        ratio_diff < 50.0,  // Allow 50% deviation due to implementation details
        "Test 3 FAILED: Gradient ratio {:.2} differs from expected {:.2} by {:.2}%",
        actual_ratio, expected_ratio, ratio_diff
    );
    
    println!("[Test 3] ✅ PASSED: IS weights preserve gradient flow");
    Ok(())
}

/// Test 4: Target Network Isolation
/// Goal: Verify target network doesn't receive gradients
/// Expected: main grad_norm > 0.001, target grad_norm = 0.0
#[test]
fn test_target_network_isolation() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Create main network
    let main_vars = VarMap::new();
    let main_vb = VarBuilder::from_varmap(&main_vars, DType::F32, &device);
    let main_layer = linear(64, 1, main_vb.pp("main_layer"))?;
    
    // Create target network
    let target_vars = VarMap::new();
    let target_vb = VarBuilder::from_varmap(&target_vars, DType::F32, &device);
    let target_layer = linear(64, 1, target_vb.pp("target_layer"))?;
    
    let batch_size = 32;
    let input = Tensor::randn(0f32, 1.0, (batch_size, 64), &device)?;
    
    // Forward through both networks
    let main_output = main_layer.forward(&input)?;
    let target_output = target_layer.forward(&input)?.detach();  // CRITICAL: Detach target
    
    // Compute TD error
    let td_error = (&main_output - &target_output)?;
    let loss = td_error.sqr()?.mean_all()?;
    
    println!("[Test 4] Loss value: {:?}", loss.to_scalar::<f32>()?);
    
    // Backward pass
    let grads = loss.backward()?;
    
    // Check main network gradients (should be non-zero)
    let main_all_vars = main_vars.all_vars();
    let main_var = &main_all_vars[0].1;
    
    let main_grad_norm = if let Some(grad) = grads.get(main_var) {
        grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f32>()?
    } else {
        0.0
    };
    
    println!("[Test 4] Main network grad_norm: {:.6}", main_grad_norm);
    
    // Check target network gradients (should be zero)
    let target_all_vars = target_vars.all_vars();
    let target_var = &target_all_vars[0].1;
    
    let target_grad_norm = if let Some(grad) = grads.get(target_var) {
        grad.sqr()?.sum_all()?.sqrt()?.to_scalar::<f32>()?
    } else {
        0.0
    };
    
    println!("[Test 4] Target network grad_norm: {:.6}", target_grad_norm);
    
    // ASSERTION: Main network should have gradients
    assert!(
        main_grad_norm > 0.001,
        "Test 4 FAILED: Main network grad_norm {:.6} too small (expected > 0.001)",
        main_grad_norm
    );
    
    // ASSERTION: Target network should have ZERO gradients
    assert!(
        target_grad_norm < 1e-9,
        "Test 4 FAILED: Target network grad_norm {:.6} not zero (expected < 1e-9)",
        target_grad_norm
    );
    
    println!("[Test 4] ✅ PASSED: Target network properly isolated");
    Ok(())
}

/// Test 5: Optimizer Parameter Consistency
/// Goal: Verify optimizer updates same params as forward pass
/// Expected: max_param_delta > 1e-6
#[test]
fn test_optimizer_parameter_consistency() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Create network
    let vars = VarMap::new();
    let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);
    let layer = linear(64, 1, vb.pp("test_layer"))?;
    
    // Create optimizer with network's VarMap
    let learning_rate = 0.001;
    let mut optimizer = AdamW::new(vars.all_vars(), learning_rate)?;
    
    // Clone parameters BEFORE training
    let all_vars = vars.all_vars();
    let before_params: Vec<Tensor> = all_vars.iter()
        .map(|(_, var)| var.as_tensor().clone())
        .collect();
    
    // Training step
    let batch_size = 32;
    let input = Tensor::randn(0f32, 1.0, (batch_size, 64), &device)?;
    let output = layer.forward(&input)?;
    let loss = output.sqr()?.mean_all()?;
    
    println!("[Test 5] Loss before step: {:?}", loss.to_scalar::<f32>()?);
    
    // Backward pass
    let grads = loss.backward()?;
    
    // Optimizer step
    optimizer.step(&grads)?;
    
    // Get parameters AFTER training
    let after_params: Vec<Tensor> = all_vars.iter()
        .map(|(_, var)| var.as_tensor().clone())
        .collect();
    
    // Calculate parameter changes
    let mut max_delta = 0.0f32;
    let mut total_delta = 0.0f32;
    let mut unchanged_count = 0;
    
    for (i, (before, after)) in before_params.iter().zip(after_params.iter()).enumerate() {
        let delta = (after - before)?;
        let delta_norm = delta.sqr()?.sum_all()?.sqrt()?.to_scalar::<f32>()?;
        
        total_delta += delta_norm;
        max_delta = max_delta.max(delta_norm);
        
        println!("[Test 5] Parameter {}: delta_norm = {:.9}", i, delta_norm);
        
        if delta_norm < 1e-9 {
            unchanged_count += 1;
        }
    }
    
    let avg_delta = total_delta / before_params.len() as f32;
    
    println!("[Test 5] Max parameter delta: {:.9}", max_delta);
    println!("[Test 5] Avg parameter delta: {:.9}", avg_delta);
    println!("[Test 5] Unchanged parameters: {}/{}", unchanged_count, before_params.len());
    
    // ASSERTION: ALL parameters should have changed
    assert_eq!(
        unchanged_count, 0,
        "Test 5 FAILED: {} parameters unchanged after optimizer step",
        unchanged_count
    );
    
    // ASSERTION: Max delta should be significant
    assert!(
        max_delta > 1e-6,
        "Test 5 FAILED: max_param_delta {:.9} too small (expected > 1e-6)",
        max_delta
    );
    
    println!("[Test 5] ✅ PASSED: Optimizer updates correct parameters");
    Ok(())
}