foxhunt/ml/tests/dqn_per_gradient_isolation_test.rs

//! BUG #41 FIX VALIDATION: PER Gradient Isolation Test
//!
//! **Purpose**: Validate that PER (Prioritized Experience Replay) metadata tensors
//! are properly detached and do not pollute the autograd graph.
//!
//! **Critical Boundaries**:
//! 1. IS weights (buffer → loss): Must be detached before multiplying loss
//! 2. TD errors (training → buffer): Must be detached before priority updates
//! 3. Target network outputs: Must be detached immediately after forward pass
//!
//! **Expected Behavior**:
//! - Gradient norms > 0.001 for ALL steps (no collapse)
//! - Priorities update correctly (non-zero, changing over time)
//! - PER vs uniform replay should have similar gradient magnitudes (within 2x)
//!
//! **What We're Testing**:
//! - Line 1210: `td_errors_vec` detached before buffer update ✅
//! - Line 1362-1364: IS weights detached before loss multiplication (C51 path) ✅
//! - Line 1374-1376: IS weights detached before loss multiplication (scalar path) ✅
//! - Lines 1088-1146: Target network outputs detached immediately ✅

use ml::dqn::{WorkingDQN, WorkingDQNConfig, Experience};
use anyhow::Result;

#[test]
fn test_per_gradient_isolation_no_collapse() -> Result<()> {
    // Create PER-enabled DQN config
    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 52; // Match production state dim
    config.num_actions = 45; // 45-action space
    config.use_per = true; // Enable PER
    config.per_alpha = 0.6;
    config.per_beta_start = 0.4;
    config.batch_size = 32;
    config.min_replay_size = 50;
    config.warmup_steps = 0; // Disable warmup for testing
    config.use_distributional = false; // Test scalar path first
    config.use_dueling = false;

    let mut dqn = WorkingDQN::new(config)?;

    // Fill replay buffer with experiences
    for i in 0..100 {
        let state = vec![i as f32 * 0.01; 52];
        let next_state = vec![(i + 1) as f32 * 0.01; 52];
        let action = (i % 45) as u8;
        let reward = (i as f32 * 0.1).sin(); // Varying rewards
        let done = i % 20 == 19; // Episode every 20 steps

        let experience = Experience::new(state, action, reward as f32, next_state, done);
        dqn.store_experience(experience)?;
    }

    // Train for 100 steps and validate gradients at EVERY step
    let mut gradient_norms = Vec::new();
    let mut losses = Vec::new();

    for step in 0..100 {
        let result = dqn.train_step(None);
        
        // Validate training succeeds
        assert!(
            result.is_ok(),
            "Training failed at step {}: {:?}",
            step,
            result.err()
        );

        let (loss, grad_norm) = result?;
        gradient_norms.push(grad_norm);
        losses.push(loss);

        // CRITICAL: Gradient norm must be > 0.001 at EVERY step
        // BUG #41 would cause gradients to collapse to zero at random steps (e.g., step 1900)
        assert!(
            grad_norm > 0.001,
            "❌ GRADIENT COLLAPSE at step {}: grad_norm={:.6} (threshold: 0.001)",
            step,
            grad_norm
        );

        // Loss must be finite
        assert!(
            loss.is_finite(),
            "❌ Loss not finite at step {}: {}",
            step,
            loss
        );
    }

    // Statistical validation: gradient norms should be stable
    let mean_grad = gradient_norms.iter().sum::<f32>() / gradient_norms.len() as f32;
    let min_grad = gradient_norms.iter().cloned().fold(f32::INFINITY, f32::min);
    let max_grad = gradient_norms.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

    println!("✅ PER Gradient Isolation Test PASSED");
    println!("   Gradient norms: mean={:.4}, range=[{:.4}, {:.4}]", mean_grad, min_grad, max_grad);
    println!("   All 100 steps had non-zero gradients (>0.001)");

    // Validate gradient stability (max shouldn't be >10x min)
    assert!(
        max_grad < min_grad * 10.0,
        "Gradient instability detected: max/min ratio = {:.2}x (should be <10x)",
        max_grad / min_grad
    );

    Ok(())
}

#[test]
fn test_per_priorities_update_correctly() -> Result<()> {
    // Create PER-enabled DQN
    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 52;
    config.num_actions = 45;
    config.use_per = true;
    config.per_alpha = 0.6;
    config.batch_size = 16;
    config.min_replay_size = 20;
    config.warmup_steps = 0;

    let mut dqn = WorkingDQN::new(config)?;

    // Add experiences with varying TD errors (rewards)
    for i in 0..30 {
        let state = vec![i as f32 * 0.01; 52];
        let next_state = vec![(i + 1) as f32 * 0.01; 52];
        let action = (i % 45) as u8;
        // Create varying rewards (some high TD error, some low)
        let reward = if i % 5 == 0 { 10.0 } else { 0.1 };
        let done = false;

        let experience = Experience::new(state, action, reward, next_state, done);
        dqn.store_experience(experience)?;
    }

    // Train several steps to update priorities
    for _ in 0..10 {
        dqn.train_step(None)?;
    }

    println!("✅ PER priorities update correctly (no errors during training)");
    Ok(())
}

#[test]
fn test_per_vs_uniform_replay_gradient_parity() -> Result<()> {
    // Test that PER and uniform replay produce similar gradient magnitudes
    // (PER IS weights should scale loss but not kill gradients)

    let mut per_config = WorkingDQNConfig::conservative();
    per_config.state_dim = 52;
    per_config.num_actions = 45;
    per_config.use_per = true;
    per_config.batch_size = 32;
    per_config.min_replay_size = 50;
    per_config.warmup_steps = 0;

    let mut uniform_config = per_config.clone();
    uniform_config.use_per = false;

    let mut dqn_per = WorkingDQN::new(per_config)?;
    let mut dqn_uniform = WorkingDQN::new(uniform_config)?;

    // Add identical experiences to both
    for i in 0..100 {
        let state = vec![i as f32 * 0.01; 52];
        let next_state = vec![(i + 1) as f32 * 0.01; 52];
        let action = (i % 45) as u8;
        let reward = (i as f32 * 0.1).sin();
        let done = i % 20 == 19;

        let experience = Experience::new(state, action, reward as f32, next_state, done);
        dqn_per.store_experience(experience.clone())?;
        dqn_uniform.store_experience(experience)?;
    }

    // Train both and compare gradient norms
    let mut per_grads = Vec::new();
    let mut uniform_grads = Vec::new();

    for _ in 0..50 {
        let (_, grad_per) = dqn_per.train_step(None)?;
        let (_, grad_uniform) = dqn_uniform.train_step(None)?;

        per_grads.push(grad_per);
        uniform_grads.push(grad_uniform);
    }

    let mean_per = per_grads.iter().sum::<f32>() / per_grads.len() as f32;
    let mean_uniform = uniform_grads.iter().sum::<f32>() / uniform_grads.len() as f32;

    // PER and uniform should have similar gradient magnitudes (within 2x)
    let ratio = mean_per / mean_uniform;
    println!(
        "✅ PER vs Uniform gradient parity: PER={:.4}, Uniform={:.4}, ratio={:.2}x",
        mean_per, mean_uniform, ratio
    );

    assert!(
        ratio > 0.5 && ratio < 2.0,
        "Gradient magnitude mismatch: PER/Uniform={:.2}x (should be 0.5-2.0x)",
        ratio
    );

    Ok(())
}

#[test]
fn test_per_c51_distributional_gradient_flow() -> Result<()> {
    // Test PER with C51 distributional RL (categorical loss path)
    let mut config = WorkingDQNConfig::aggressive();
    config.state_dim = 52;
    config.num_actions = 45;
    config.use_per = true;
    config.use_distributional = true;
    config.use_dueling = true; // Hybrid mode
    config.batch_size = 32;
    config.min_replay_size = 50;
    config.warmup_steps = 0;

    let mut dqn = WorkingDQN::new(config)?;

    // Fill buffer
    for i in 0..100 {
        let state = vec![i as f32 * 0.01; 52];
        let next_state = vec![(i + 1) as f32 * 0.01; 52];
        let action = (i % 45) as u8;
        let reward = (i as f32 * 0.1).sin();
        let done = i % 20 == 19;

        let experience = Experience::new(state, action, reward as f32, next_state, done);
        dqn.store_experience(experience)?;
    }

    // Train and validate gradients (C51 categorical loss path)
    for step in 0..50 {
        let result = dqn.train_step(None);
        
        assert!(
            result.is_ok(),
            "C51+PER training failed at step {}: {:?}",
            step,
            result.err()
        );

        let (loss, grad_norm) = result?;

        // Critical: C51 categorical loss with PER must have non-zero gradients
        assert!(
            grad_norm > 0.001,
            "❌ C51+PER gradient collapse at step {}: grad_norm={:.6}",
            step,
            grad_norm
        );

        assert!(
            loss.is_finite(),
            "❌ C51+PER loss not finite at step {}: {}",
            step,
            loss
        );
    }

    println!("✅ PER + C51 distributional gradient flow validated (50 steps, all non-zero)");
    Ok(())
}