foxhunt/ml/tests/dqn_xavier_init_test.rs

//! Test suite for Xavier/Glorot initialization in DQN networks
//!
//! This test verifies that Xavier initialization is properly applied to DQN network layers,
//! ensuring initial gradient stability and proper weight distribution.

use ml::dqn::network::{QNetwork, QNetworkConfig};

/// Test that network weights follow Xavier distribution after initialization
#[test]
fn test_xavier_weight_distribution() -> Result<(), Box<dyn std::error::Error>> {
    let config = QNetworkConfig {
        state_dim: 64,
        num_actions: 3,
        hidden_dims: vec![128, 64],
        learning_rate: 0.001,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        target_update_freq: 1000,
        dropout_prob: 0.0, // Disable dropout for initialization test
        use_gpu: false,
    };

    let network = QNetwork::new(config)?;
    let vars = network.vars();

    // Check first hidden layer (64 -> 128)
    let vars_data = vars.data().lock().unwrap();
    let layer0_weight = vars_data
        .iter()
        .find(|(name, _)| name.contains("layer_0") && name.contains("weight"))
        .ok_or("Layer 0 weight not found")?
        .1
        .as_tensor();

    let weight_mean = layer0_weight.mean_all()?.to_scalar::<f32>()?;
    // Calculate variance manually: Var(X) = E[X²] - E[X]²
    let squared = layer0_weight.sqr()?;
    let mean_squared = squared.mean_all()?.to_scalar::<f32>()?;
    let weight_var = mean_squared - (weight_mean * weight_mean);

    // For Xavier initialization: Var(W) = 2 / (fan_in + fan_out)
    let fan_in = 64;
    let fan_out = 128;
    let expected_var = 2.0 / (fan_in + fan_out) as f32;

    println!("Layer 0 (64->128) Statistics:");
    println!("  Mean: {:.6}", weight_mean);
    println!("  Variance: {:.6}", weight_var);
    println!("  Expected Xavier variance: {:.6}", expected_var);

    // Mean should be close to zero
    assert!(
        weight_mean.abs() < 0.1,
        "Weight mean {:.6} should be close to 0 (< 0.1)",
        weight_mean
    );

    // Variance should match Xavier formula within 20% tolerance
    // (larger tolerance due to random initialization and small sample size)
    let var_diff = (weight_var - expected_var).abs() / expected_var;
    assert!(
        var_diff < 0.20,
        "Weight variance {:.6} should match Xavier expected {:.6} (diff: {:.2}%)",
        weight_var,
        expected_var,
        var_diff * 100.0
    );

    Ok(())
}

/// Test that Xavier initialization produces weights in expected range
#[test]
fn test_xavier_weight_range() -> Result<(), Box<dyn std::error::Error>> {
    let config = QNetworkConfig {
        state_dim: 64,
        num_actions: 3,
        hidden_dims: vec![128, 64],
        learning_rate: 0.001,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        target_update_freq: 1000,
        dropout_prob: 0.0,
        use_gpu: false,
    };

    let network = QNetwork::new(config)?;
    let vars = network.vars();

    // Check first hidden layer (64 -> 128)
    let vars_data = vars.data().lock().unwrap();
    let layer0_weight = vars_data
        .iter()
        .find(|(name, _)| name.contains("layer_0") && name.contains("weight"))
        .ok_or("Layer 0 weight not found")?
        .1
        .as_tensor();

    // For Xavier uniform: limit = sqrt(6 / (fan_in + fan_out))
    let fan_in = 64.0_f64;
    let fan_out = 128.0_f64;
    let expected_limit = (6.0 / (fan_in + fan_out)).sqrt() as f32;

    let weight_flat = layer0_weight.flatten_all()?;
    let weight_max = weight_flat.max_all()?.to_scalar::<f32>()?;
    let weight_min = weight_flat.min_all()?.to_scalar::<f32>()?;

    println!("Layer 0 (64->128) Weight Range:");
    println!("  Min: {:.6}", weight_min);
    println!("  Max: {:.6}", weight_max);
    println!("  Expected Xavier limit: ±{:.6}", expected_limit);

    // Weights should be approximately within [-limit, +limit]
    // Allow some margin for random variation
    let margin = 1.5; // Allow up to 50% over limit
    assert!(
        weight_max <= expected_limit * margin,
        "Max weight {:.6} should be within Xavier limit {:.6} * {:.1}",
        weight_max,
        expected_limit,
        margin
    );
    assert!(
        weight_min >= -expected_limit * margin,
        "Min weight {:.6} should be within Xavier limit -{:.6} * {:.1}",
        weight_min,
        expected_limit,
        margin
    );

    Ok(())
}

/// Test that Xavier initialization produces forward pass in reasonable range
#[test]
fn test_forward_pass_stability() -> Result<(), Box<dyn std::error::Error>> {
    let config = QNetworkConfig {
        state_dim: 64,
        num_actions: 3,
        hidden_dims: vec![128, 64],
        learning_rate: 0.001,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        target_update_freq: 1000,
        dropout_prob: 0.0,
        use_gpu: false,
    };

    let network = QNetwork::new(config.clone())?;

    // Create test input
    let state = vec![0.5f32; config.state_dim];

    // Forward pass
    let q_values = network.forward(&state)?;

    println!("Initial Q-values: {:?}", q_values);

    // With Xavier initialization, Q-values should be in reasonable range
    // Not exploding (>1000) or vanishing (<0.001)
    for (i, &q_val) in q_values.iter().enumerate() {
        assert!(
            q_val.is_finite(),
            "Q-value {} should be finite, got {}",
            i,
            q_val
        );
        assert!(
            q_val.abs() < 100.0,
            "Q-value {} magnitude should be < 100.0 with Xavier init (Wave 9 gradients were 30,000+), got {}",
            i,
            q_val
        );
    }

    Ok(())
}

/// Test Xavier initialization across all layers
#[test]
fn test_all_layers_xavier_initialized() -> Result<(), Box<dyn std::error::Error>> {
    let config = QNetworkConfig {
        state_dim: 64,
        num_actions: 3,
        hidden_dims: vec![128, 64, 32],
        learning_rate: 0.001,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        target_update_freq: 1000,
        dropout_prob: 0.0,
        use_gpu: false,
    };

    let network = QNetwork::new(config.clone())?;
    let vars = network.vars();

    // Layer dimensions: 64->128, 128->64, 64->32, 32->3
    let layer_dims = vec![
        (64, 128), // layer_0
        (128, 64), // layer_1
        (64, 32),  // layer_2
        (32, 3),   // output
    ];

    let vars_data = vars.data().lock().unwrap();

    for (layer_idx, (fan_in, fan_out)) in layer_dims.iter().enumerate() {
        let layer_name = if layer_idx < config.hidden_dims.len() {
            format!("layer_{}", layer_idx)
        } else {
            "output".to_string()
        };

        let weight = vars_data
            .iter()
            .find(|(name, _)| name.contains(&layer_name) && name.contains("weight"))
            .ok_or_else(|| format!("Layer {} weight not found", layer_name))?
            .1
            .as_tensor();

        // Calculate variance manually: Var(X) = E[X²] - E[X]²
        let weight_mean = weight.mean_all()?.to_scalar::<f32>()?;
        let squared = weight.sqr()?;
        let mean_squared = squared.mean_all()?.to_scalar::<f32>()?;
        let weight_var = mean_squared - (weight_mean * weight_mean);
        let expected_var = 2.0 / (fan_in + fan_out) as f32;

        println!(
            "Layer {} ({}->{}): var={:.6}, expected={:.6}",
            layer_name, fan_in, fan_out, weight_var, expected_var
        );

        // Check variance within 30% tolerance (larger tolerance for output layer)
        let var_diff = (weight_var - expected_var).abs() / expected_var;
        assert!(
            var_diff < 0.30,
            "Layer {} variance {:.6} should match Xavier expected {:.6} (diff: {:.2}%)",
            layer_name,
            weight_var,
            expected_var,
            var_diff * 100.0
        );
    }

    Ok(())
}

/// Test that target network is also Xavier initialized
#[test]
fn test_target_network_xavier_initialized() -> Result<(), Box<dyn std::error::Error>> {
    let config = QNetworkConfig {
        state_dim: 64,
        num_actions: 3,
        hidden_dims: vec![128, 64],
        learning_rate: 0.001,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        target_update_freq: 1000,
        dropout_prob: 0.0,
        use_gpu: false,
    };

    let network = QNetwork::new(config)?;
    let target_vars = network.target_vars();

    // Check target network first layer (64 -> 128)
    let vars_data = target_vars.data().lock().unwrap();
    let target_layer0_weight = vars_data
        .iter()
        .find(|(name, _)| name.contains("layer_0") && name.contains("weight"))
        .ok_or("Target layer 0 weight not found")?
        .1
        .as_tensor();

    // Calculate variance manually: Var(X) = E[X²] - E[X]²
    let weight_mean = target_layer0_weight.mean_all()?.to_scalar::<f32>()?;
    let squared = target_layer0_weight.sqr()?;
    let mean_squared = squared.mean_all()?.to_scalar::<f32>()?;
    let weight_var = mean_squared - (weight_mean * weight_mean);
    let fan_in = 64;
    let fan_out = 128;
    let expected_var = 2.0 / (fan_in + fan_out) as f32;

    println!("Target Network Layer 0 (64->128):");
    println!("  Variance: {:.6}", weight_var);
    println!("  Expected Xavier variance: {:.6}", expected_var);

    let var_diff = (weight_var - expected_var).abs() / expected_var;
    assert!(
        var_diff < 0.20,
        "Target network variance {:.6} should match Xavier expected {:.6} (diff: {:.2}%)",
        weight_var,
        expected_var,
        var_diff * 100.0
    );

    Ok(())
}