foxhunt/ml/tests/dqn_c51_gradient_flow_bug36_test.rs

//! BUG #36: C51 project_distribution() Gradient Flow Test
//!
//! Tests that the distributional projection maintains gradient flow.
//! This is a regression test for the CPU scatter loop bug.

use candle_core::{Device, Tensor, Var};
use candle_nn::{VarBuilder, VarMap};
use ml::dqn::distributional::{CategoricalDistribution, DistributionalConfig};
use ml::MLError;

#[test]
fn test_project_distribution_gradient_flow() -> Result<(), MLError> {
    println!("\n=== BUG #36: Testing project_distribution Gradient Flow ===\n");

    let device = Device::cuda_if_available(0)?;

    // Create categorical distribution
    let config = DistributionalConfig {
        num_atoms: 11,
        v_min: -5.0,
        v_max: 5.0,
    };
    let cat_dist = CategoricalDistribution::new(&config, &device)?;

    // Create target support and probabilities
    let batch_size = 4;
    let num_atoms = 11;

    // Target support: shifted support values
    let target_support = Tensor::randn(0f32, 1.0, (batch_size, num_atoms), &device)?;

    // Create a VarMap and use it to create learnable probabilities
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    // Create learnable probabilities
    let probs_raw = vb.get((batch_size, num_atoms), "probs")?;
    let probs_sum = probs_raw.sum_keepdim(1)?;
    let probabilities = probs_raw.broadcast_div(&probs_sum)?;

    // Project distribution
    let projected = cat_dist.project_distribution(&target_support, &probabilities)?;

    // Compute loss (sum of projected distribution)
    let loss = projected.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // CRITICAL TEST: Check if gradients were computed
    let all_vars = varmap.all_vars();
    let mut total_grad_norm = 0.0f32;
    let mut non_zero_count = 0;
    let mut zero_count = 0;

    for var in all_vars.iter() {
        if let Some(grad) = grads.get(var) {
            let grad_vec: Vec<f32> = grad.flatten_all()?.to_vec1()?;
            for &g in grad_vec.iter() {
                if g.abs() > 1e-9 {
                    non_zero_count += 1;
                    total_grad_norm += g * g;
                } else {
                    zero_count += 1;
                }
            }
        }
    }

    let grad_norm_sqrt = total_grad_norm.sqrt();
    let total_elements = non_zero_count + zero_count;

    println!("Gradient norm: {:.6}", grad_norm_sqrt);
    println!("Non-zero gradients: {}/{}", non_zero_count, total_elements);
    println!("Zero gradients: {}/{}", zero_count, total_elements);

    // ASSERTION: At least 50% of gradients should be non-zero
    assert!(
        grad_norm_sqrt > 1e-6,
        "BUG #36: Gradient flow broken! Expected non-zero gradients, got norm={:.9}",
        grad_norm_sqrt
    );

    assert!(
        non_zero_count > total_elements / 2,
        "BUG #36: Most gradients are zero! Expected >50% non-zero, got {}/{}",
        non_zero_count,
        total_elements
    );

    println!("\n✅ BUG #36 TEST PASSED: Gradient flow maintained through project_distribution\n");

    Ok(())
}

#[test]
fn test_project_distribution_vs_original() -> Result<(), MLError> {
    println!("\n=== BUG #36: Comparing New vs Original Implementation ===\n");

    let device = Device::cuda_if_available(0)?;

    let config = DistributionalConfig {
        num_atoms: 11,
        v_min: -5.0,
        v_max: 5.0,
    };
    let cat_dist = CategoricalDistribution::new(&config, &device)?;

    // Create deterministic test case
    let batch_size = 2;
    let num_atoms = 11;

    let target_support = Tensor::new(
        vec![
            // Batch 0: support values -5.0 to 5.0
            -5.0f32, -4.0, -3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0,
            // Batch 1: shifted support
            -3.0f32, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
        ],
        &device,
    )?.reshape((batch_size, num_atoms))?;

    // Uniform probabilities
    let probs = Tensor::full(1.0 / num_atoms as f32, (batch_size, num_atoms), &device)?;

    // Project
    let projected = cat_dist.project_distribution(&target_support, &probs)?;

    // Check properties
    let sums: Vec<f32> = projected.sum(1)?.to_vec1()?;
    println!("Projected distribution sums: {:?}", sums);

    for (i, &sum) in sums.iter().enumerate() {
        assert!(
            (sum - 1.0).abs() < 1e-5,
            "BUG #36: Projected distribution batch {} doesn't sum to 1.0: {}",
            i,
            sum
        );
    }

    println!("\n✅ BUG #36 CORRECTNESS TEST PASSED: Projected distributions are valid\n");

    Ok(())
}