foxhunt/ml/tests/dqn_gradient_clipping_bug32.rs

// Bug #32: DQN Gradient Clipping - TDD Test Suite
//
// CRITICAL BUG: Gradient clipping is completely disabled (NO-OP stub)
// - Current: `fn clip_gradients(&self, max_norm: f32) -> Result<(), MLError> { Ok(()) }`
// - Reality: Gradients exploding to 20K-25K instead of being clipped to 10.0
// - Impact: Weight corruption → Q-value collapse → training failure
//
// These tests will FAIL with current code, PASS after implementing:
// `candle_nn::optim::clip_grad_norm(&varmap.all_vars(), max_norm as f64)?`

use candle_core::{Device, Tensor};
use candle_nn::{AdamW, Optimizer, VarBuilder, VarMap, optim, Init};
use ml::gradient_utils::clip_grad_norm;
use ml::MLError;

/// Test 1: Verify gradient norm calculation is working
///
/// This test ensures we can correctly compute the global L2 norm of all gradients.
/// With NO-OP clipping: This test will PASS (norm calculation is separate)
/// After fix: This test will PASS (norm calculation unchanged)
#[test]
fn test_gradient_norm_calculation() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create simple VarMap with known parameters
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Create a 2x2 weight tensor initialized to 1.0 to ensure non-zero gradients
    let weight = vb.get_with_hints((2, 2), "weight", Init::Const(1.0))?;

    // Create a simple loss: sum of squares
    let loss = weight.sqr()?.sum_all()?;

    // Compute gradients
    let mut grads = loss.backward()?;

    // Calculate gradient norm manually
    // gradient of sum(w^2) w.r.t. w = 2w
    // If w is initialized to ~0.01, gradient ~0.02, norm should be small
    let (grad_norm, _) = clip_grad_norm(&varmap.all_vars(), &mut grads, 999999.0)?;

    // Gradient norm should be positive and finite
    assert!(grad_norm > 0.0, "Gradient norm should be positive, got {}", grad_norm);
    assert!(grad_norm.is_finite(), "Gradient norm should be finite, got {}", grad_norm);
    assert!(grad_norm < 100.0, "Gradient norm should be reasonable for this simple case, got {}", grad_norm);

    println!("✓ Test 1 PASS: Gradient norm = {:.6}", grad_norm);
    Ok(())
}

/// Test 2: Verify gradient clipping enforcement (CRITICAL - will FAIL with NO-OP)
///
/// This test creates artificially large gradients and verifies they get clipped.
/// With NO-OP clipping: This test will FAIL (gradients unchanged)
/// After fix: This test will PASS (gradients scaled down)
#[test]
fn test_gradient_clipping_enforcement() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Create weight initialized to large values to produce large gradients
    let weight = vb.get((10, 10), "large_weight")?;

    // Create loss that will produce gradients ~100x larger than clip threshold
    // Loss = 1000 * sum(w^2), gradient = 2000 * w
    let loss = (weight.sqr()? * 1000.0)?.sum_all()?;

    // Compute gradients
    let mut grads = loss.backward()?;

    // Measure gradient norm BEFORE clipping (use very large max_norm to just measure, not clip)
    let (norm_before, _) = clip_grad_norm(&varmap.all_vars(), &mut grads, 999999.0)?;
    println!("Gradient norm before clipping: {:.2}", norm_before);

    // Re-compute gradients since we need to clip them fresh
    // (backward() consumes the computational graph, so we need a fresh loss)
    let loss2 = (weight.sqr()? * 1000.0)?.sum_all()?;
    let mut grads2 = loss2.backward()?;

    // Now clip to max_norm = 10.0
    let max_norm = 10.0;
    let (_, norm_after) = clip_grad_norm(&varmap.all_vars(), &mut grads2, max_norm)?;
    println!("Gradient norm after clipping to {}: {:.2}", max_norm, norm_after);

    // CRITICAL ASSERTION: After clipping, norm should be <= max_norm
    assert!(
        norm_after <= max_norm + 1e-5,
        "FAIL: Gradient norm ({}) exceeds max_norm ({}) - clipping not working!",
        norm_after,
        max_norm
    );

    // If norm_before was > max_norm, then norm_after should be reduced
    if norm_before > max_norm {
        assert!(
            norm_after < norm_before,
            "FAIL: Gradient norm should decrease after clipping (before: {}, after: {})",
            norm_before,
            norm_after
        );
    }

    println!("✓ Test 2 PASS: Gradients clipped from {:.2} to {:.2}", norm_before, norm_after);
    Ok(())
}

/// Test 3: Verify no clipping when below threshold
///
/// This test ensures gradients below max_norm are NOT modified.
/// With NO-OP clipping: This test will PASS (no clipping occurs)
/// After fix: This test will PASS (correct behavior)
#[test]
fn test_no_clipping_when_below_threshold() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Create small weight to produce small gradients
    let weight = vb.get((2, 2), "small_weight")?;

    // Small loss → small gradients
    let loss = (weight.sqr()? * 0.01)?.sum_all()?;
    let mut grads = loss.backward()?;

    // Measure gradient norm (use very large max_norm to just measure, not clip)
    let (norm_before, _) = clip_grad_norm(&varmap.all_vars(), &mut grads, 999999.0)?;
    println!("Gradient norm (small): {:.6}", norm_before);

    // Re-compute gradients for second measurement
    let loss2 = (weight.sqr()? * 0.01)?.sum_all()?;
    let mut grads2 = loss2.backward()?;

    // Clip with large max_norm (should have no effect)
    let max_norm = 100.0;
    let (norm_after, _) = clip_grad_norm(&varmap.all_vars(), &mut grads2, max_norm)?;

    // Norm should be unchanged (within floating point tolerance)
    let tolerance = 1e-6;
    assert!(
        (norm_after - norm_before).abs() < tolerance,
        "FAIL: Gradient norm changed when it shouldn't have (before: {}, after: {}, diff: {})",
        norm_before,
        norm_after,
        (norm_after - norm_before).abs()
    );

    println!("✓ Test 3 PASS: Small gradients unchanged ({:.6})", norm_after);
    Ok(())
}

/// Test 4: Edge case - zero gradients
///
/// This test ensures clipping doesn't break when gradients are zero.
/// With NO-OP clipping: This test will PASS (no operation = no crash)
/// After fix: This test will PASS (clip_grad_norm handles zero gracefully)
#[test]
fn test_zero_gradients_edge_case() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Create weight
    let weight = vb.get((3, 3), "weight")?;

    // Zero loss → zero gradients
    let loss = (weight * 0.0)?.sum_all()?;
    let mut grads = loss.backward()?;

    // Clip should handle zero gradients gracefully
    let (norm, _) = clip_grad_norm(&varmap.all_vars(), &mut grads, 10.0)?;

    // Norm should be zero or very close
    assert!(
        norm < 1e-10,
        "FAIL: Expected near-zero gradient norm, got {}",
        norm
    );
    assert!(norm.is_finite(), "FAIL: Gradient norm should be finite, got {}", norm);

    println!("✓ Test 4 PASS: Zero gradients handled correctly (norm = {:.10})", norm);
    Ok(())
}

/// Test 5: Integration - verify clipping in DQN training loop order
///
/// This test ensures gradient clipping happens in the correct order:
/// 1. loss.backward() - compute gradients
/// 2. clip_gradients() - clip gradients IN-PLACE
/// 3. optimizer.step() - apply clipped gradients
///
/// With NO-OP clipping: This test will FAIL (large update will occur)
/// After fix: This test will PASS (update size constrained)
#[test]
fn test_clipping_integration_order() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Create weight
    let weight = vb.get((5, 5), "integration_weight")?;

    // Get initial weight values
    let weight_before = weight.to_vec2::<f32>()?;

    // Create large loss → large gradients
    let loss = (weight.sqr()? * 5000.0)?.sum_all()?;

    // Step 1: Backward pass (compute gradients)
    let mut grads = loss.backward()?;

    // Step 2: Clip gradients (max_norm = 1.0, very strict)
    let max_norm = 1.0;
    let (_, grad_norm) = clip_grad_norm(&varmap.all_vars(), &mut grads, max_norm)?;
    println!("Gradient norm after clipping: {:.6}", grad_norm);

    // Verify clipping occurred
    assert!(
        grad_norm <= max_norm + 1e-5,
        "FAIL: Gradient norm ({}) exceeds max_norm ({})",
        grad_norm,
        max_norm
    );

    // Step 3: Optimizer step (apply clipped gradients)
    let learning_rate = 0.1;
    let mut optimizer = AdamW::new(varmap.all_vars(), optim::ParamsAdamW {
        lr: learning_rate,
        ..Default::default()
    })?;
    optimizer.step(&grads)?;

    // Get updated weight values
    let weight_after = weight.to_vec2::<f32>()?;

    // Calculate max absolute change
    let mut max_change = 0.0f32;
    for i in 0..weight_before.len() {
        for j in 0..weight_before[i].len() {
            let change = (weight_after[i][j] - weight_before[i][j]).abs();
            max_change = max_change.max(change);
        }
    }

    println!("Max weight change: {:.6}", max_change);

    // With gradient clipping, weight changes should be bounded
    // max_change ≈ learning_rate * max_norm / sqrt(num_params)
    // For our case: ~0.1 * 1.0 / sqrt(25) = ~0.02
    let expected_max_change = (learning_rate * max_norm * 2.0) as f32; // 2x safety factor
    assert!(
        max_change < expected_max_change,
        "FAIL: Weight change ({}) too large, suggests clipping not working (expected < {})",
        max_change,
        expected_max_change
    );

    println!("✓ Test 5 PASS: Integration test - clipping constrained weight updates");
    Ok(())
}

/// Test 6: Regression test - verify fix for Bug #32
///
/// This test simulates the exact bug scenario: gradients exploding to 20K+
/// With NO-OP clipping: This test will FAIL (gradients > 10K)
/// After fix: This test will PASS (gradients clipped to ≤10.0)
#[test]
fn test_bug32_regression_gradient_explosion() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Simulate DQN network size (3 layers, ~1000 params total)
    // Initialize to 1.0 to ensure non-zero gradients for this test
    let layer1 = vb.get_with_hints((20, 20), "layer1", Init::Const(1.0))?;
    let layer2 = vb.get_with_hints((20, 10), "layer2", Init::Const(1.0))?;
    let layer3 = vb.get_with_hints((10, 5), "layer3", Init::Const(1.0))?;

    // Create loss that produces exploding gradients (simulates large TD-error)
    // In real DQN, this happens when Q-values are unstable
    let huge_multiplier = 10000.0; // Simulate TD-error of 10K
    let loss = ((layer1.sqr()?.sum_all()? + layer2.sqr()?.sum_all()? + layer3.sqr()?.sum_all()?)? * huge_multiplier)?;

    // Compute gradients
    let mut grads = loss.backward()?;

    // Measure gradient norm before clipping (use very large max_norm to just measure)
    let (norm_before, _) = clip_grad_norm(&varmap.all_vars(), &mut grads, 999999.0)?;
    println!("Bug #32 scenario - Gradient norm before clipping: {:.2}", norm_before);

    // This should be HUGE (>1000) simulating the bug
    assert!(
        norm_before > 100.0,
        "Test setup error: Expected large gradients to simulate bug, got {}",
        norm_before
    );

    // Re-compute gradients for clipping test
    let loss2 = ((layer1.sqr()?.sum_all()? + layer2.sqr()?.sum_all()? + layer3.sqr()?.sum_all()?)? * huge_multiplier)?;
    let mut grads2 = loss2.backward()?;

    // Apply DQN's gradient clipping threshold
    let dqn_max_norm = 10.0;
    let (_, norm_after) = clip_grad_norm(&varmap.all_vars(), &mut grads2, dqn_max_norm)?;
    println!("Bug #32 scenario - Gradient norm after clipping: {:.2}", norm_after);

    // CRITICAL: This is the bug fix verification
    assert!(
        norm_after <= dqn_max_norm + 1e-5,
        "BUG #32 NOT FIXED: Gradient norm ({}) exceeds DQN max_norm ({})",
        norm_after,
        dqn_max_norm
    );

    // Verify significant reduction occurred
    let reduction_ratio = norm_after / norm_before;
    assert!(
        reduction_ratio < 0.1,
        "BUG #32 NOT FIXED: Gradient norm reduction insufficient ({}x → {}x, ratio: {:.4})",
        norm_before,
        norm_after,
        reduction_ratio
    );

    println!(
        "✓ Test 6 PASS: BUG #32 FIXED - Gradients clipped from {:.2} to {:.2} ({:.2}% reduction)",
        norm_before,
        norm_after,
        (1.0 - reduction_ratio) * 100.0
    );

    Ok(())
}