foxhunt/ml/tests/dqn_target_network_gradient_isolation_test.rs

//! BUG #41 FIX: Target Network Gradient Isolation Test
//!
//! Tests that gradients flow correctly through the online network while the
//! target network remains frozen (no gradients).
//!
//! **Critical Bug**: Missing `.detach()` calls on target network outputs cause
//! Candle to zero ALL gradients (both online and target), resulting in total
//! gradient collapse.
//!
//! **Fixed Locations** (already implemented in dqn.rs):
//! - Line 1092: dist_dueling_target.forward().detach()
//! - Line 1118: dueling_target.forward().detach()
//! - Line 1122: target_network.forward().detach()
//! - Line 1207: diff.detach() (TD errors for PER)
//! - Line 1310: target_dists.detach() (C51 Bellman target)
//! - Line 1327-1335: Vec→Tensor conversion (strips gradient metadata)
//! - Line 1356: weights_tensor.detach() (PER IS weights)
//! - Line 1392: weights_tensor.detach() (PER IS weights, standard path)
//!
//! **Test Strategy**: Verify non-zero gradient norms across all DQN architectures.

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

#[test]
fn test_standard_dqn_gradient_flow() -> Result<()> {
    println!("\n=== TEST: Standard DQN Gradient Flow ===");

    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 52;
    config.num_actions = 45;
    config.batch_size = 16;
    config.min_replay_size = 16;
    config.use_dueling = false;
    config.use_distributional = false;
    config.use_per = false;
    config.warmup_steps = 0;

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

    // Add experiences
    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            if i % 2 == 0 { 1.0 } else { -0.5 },
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    // Run training step
    let (loss, grad_norm) = dqn.train_step(None)?;

    println!("  Loss: {:.6}, Gradient Norm: {:.6}", loss, grad_norm);

    // ASSERT: Gradient norm should be non-zero
    assert!(
        grad_norm > 1e-6,
        "❌ GRADIENT COLLAPSE: norm={:.10}, expected >1e-6",
        grad_norm
    );

    println!("  ✓ Gradients flowing (norm: {:.6})", grad_norm);
    println!("  ✓ Target network isolation confirmed\n");

    Ok(())
}

#[test]
fn test_dueling_dqn_gradient_flow() -> Result<()> {
    println!("\n=== TEST: Dueling DQN Gradient Flow ===");

    let mut config = WorkingDQNConfig::aggressive();
    config.state_dim = 52;
    config.num_actions = 45;
    config.batch_size = 16;
    config.min_replay_size = 16;
    config.use_dueling = true;
    config.use_distributional = false;
    config.dueling_hidden_dim = 128;
    config.warmup_steps = 0;

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

    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            if i % 2 == 0 { 1.0 } else { -0.5 },
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    let (loss, grad_norm) = dqn.train_step(None)?;

    println!("  Loss: {:.6}, Gradient Norm: {:.6}", loss, grad_norm);

    assert!(
        grad_norm > 1e-6,
        "❌ GRADIENT COLLAPSE (Dueling): norm={:.10}",
        grad_norm
    );

    println!("  ✓ Dueling gradients flowing (norm: {:.6})", grad_norm);
    println!("  ✓ Dueling target network isolation confirmed\n");

    Ok(())
}

#[test]
fn test_hybrid_c51_gradient_flow() -> Result<()> {
    println!("\n=== TEST: Hybrid C51 (Distributional + Dueling) Gradient Flow ===");

    let mut config = WorkingDQNConfig::aggressive();
    config.state_dim = 52;
    config.num_actions = 45;
    config.batch_size = 16;
    config.min_replay_size = 16;
    config.use_dueling = true;
    config.use_distributional = true;
    config.num_atoms = 51;
    config.v_min = -2.0;
    config.v_max = 2.0;
    config.dueling_hidden_dim = 128;
    config.warmup_steps = 0;

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

    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            if i % 2 == 0 { 1.0 } else { -0.5 },
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    let (loss, grad_norm) = dqn.train_step(None)?;

    println!("  Loss: {:.6}, Gradient Norm: {:.6}", loss, grad_norm);

    assert!(
        grad_norm > 1e-6,
        "❌ GRADIENT COLLAPSE (C51): norm={:.10}",
        grad_norm
    );

    println!("  ✓ C51 categorical loss gradients flowing (norm: {:.6})", grad_norm);
    println!("  ✓ C51 target distribution isolation confirmed\n");

    Ok(())
}

#[test]
fn test_per_gradient_flow() -> Result<()> {
    println!("\n=== TEST: PER (Prioritized Replay) Gradient Flow ===");

    let mut config = WorkingDQNConfig::aggressive();
    config.state_dim = 52;
    config.num_actions = 45;
    config.batch_size = 16;
    config.min_replay_size = 16;
    config.use_per = true;
    config.per_alpha = 0.6;
    config.per_beta_start = 0.4;
    config.warmup_steps = 0;

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

    // Vary rewards to create TD error diversity (important for PER)
    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            ((i % 5) as f32 - 2.0) * 0.5, // -1.0, -0.5, 0.0, 0.5, 1.0
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    let (loss, grad_norm) = dqn.train_step(None)?;

    println!("  Loss: {:.6}, Gradient Norm: {:.6}", loss, grad_norm);

    assert!(
        grad_norm > 1e-6,
        "❌ GRADIENT COLLAPSE (PER): norm={:.10}",
        grad_norm
    );

    println!("  ✓ PER gradients flowing with IS weights (norm: {:.6})", grad_norm);
    println!("  ✓ TD errors and IS weights correctly detached\n");

    Ok(())
}

#[test]
fn test_multi_step_gradient_stability() -> Result<()> {
    println!("\n=== TEST: Multi-Step Gradient Stability ===");

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

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

    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            if i % 2 == 0 { 1.0 } else { -0.5 },
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    // Run 5 consecutive training steps
    println!("  Testing gradient flow across 5 consecutive steps:");
    for step in 0..5 {
        let (loss, grad_norm) = dqn.train_step(None)?;

        println!("    Step {}: Loss={:.6}, Norm={:.6}", step + 1, loss, grad_norm);

        assert!(
            grad_norm > 1e-6,
            "❌ GRADIENT COLLAPSE at step {}: norm={:.10}",
            step + 1,
            grad_norm
        );
    }

    println!("  ✓ Gradients stable across 5 consecutive steps");
    println!("  ✓ No gradient accumulation or collapse\n");

    Ok(())
}

#[test]
fn test_c51_target_distribution_vec_conversion() -> Result<()> {
    println!("\n=== TEST: C51 Target Distribution Vec Conversion (Line 1327 Fix) ===");

    let mut config = WorkingDQNConfig::aggressive();
    config.state_dim = 52;
    config.num_actions = 45;
    config.batch_size = 16;
    config.min_replay_size = 16;
    config.use_dueling = true;
    config.use_distributional = true;
    config.num_atoms = 51;
    config.v_min = -2.0;
    config.v_max = 2.0;
    config.warmup_steps = 0;

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

    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            if i % 2 == 0 { 1.0 } else { -0.5 },
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    println!("  Testing Vec→Tensor conversion strips gradient metadata:");
    for step in 0..3 {
        let (loss, grad_norm) = dqn.train_step(None)?;

        println!("    Step {}: Loss={:.6}, Norm={:.6}", step + 1, loss, grad_norm);

        assert!(
            grad_norm > 1e-6,
            "❌ GRADIENT COLLAPSE (Vec conversion failed): step {}, norm={:.10}",
            step + 1,
            grad_norm
        );
    }

    println!("  ✓ Vec→Tensor conversion successfully strips gradient metadata");
    println!("  ✓ No gradient fork detected\n");

    Ok(())
}

#[test]
fn test_double_dqn_with_target_detachment() -> Result<()> {
    println!("\n=== TEST: Double DQN with Target Detachment ===");

    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 52;
    config.num_actions = 45;
    config.batch_size = 16;
    config.min_replay_size = 16;
    config.use_double_dqn = true; // Enable Double DQN (uses both online and target networks)
    config.use_per = false;
    config.warmup_steps = 0;

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

    for i in 0..20 {
        let exp = Experience::new(
            vec![i as f32 * 0.01; 52],
            (i % 45) as u8,
            if i % 2 == 0 { 1.0 } else { -0.5 },
            vec![(i + 1) as f32 * 0.01; 52],
            i == 19,
        );
        dqn.store_experience(exp)?;
    }

    let (loss, grad_norm) = dqn.train_step(None)?;

    println!("  Loss: {:.6}, Gradient Norm: {:.6}", loss, grad_norm);

    assert!(
        grad_norm > 1e-6,
        "❌ GRADIENT COLLAPSE (Double DQN): norm={:.10}",
        grad_norm
    );

    println!("  ✓ Double DQN gradients flowing (online selects, target evaluates)");
    println!("  ✓ Target Q-values correctly detached\n");

    Ok(())
}