foxhunt/ml/tests/wave16_entropy_polyak_tests.rs

//! Wave 16 Integration Tests: Entropy Regularization + Polyak Soft Updates
//!
//! This module provides integration tests for two new DQN features:
//! 1. **Entropy Regularization**: Prevents policy collapse via entropy-based reward shaping
//! 2. **Polyak Soft Updates**: Gradual target network tracking via exponential moving average
//!
//! # Test Coverage
//! - Polyak soft update weight blending (tau parameter validation)
//! - CLI flag parsing and integration (--tau, --soft-updates)
//! - Combined feature operation (entropy + Polyak simultaneously)
//! - Backward compatibility (hard updates without --soft-updates)
//!
//! # Note on Entropy Testing
//! Entropy regularization is integrated directly into the WorkingDQN struct
//! (ml/src/dqn/dqn.rs) and is tested via end-to-end training rather than
//! isolated unit tests. Entropy calculation happens during loss computation
//! and is validated through action diversity metrics in production training.

#![allow(unused_crate_dependencies)]

use anyhow::Result;
use candle_core::{Device, Tensor};
use candle_nn::VarMap;
use ml::dqn::target_update::{convergence_half_life, hard_update, polyak_update};

// ============================================================================
// Test 1: Polyak Soft Updates
// ============================================================================

#[test]
fn test_polyak_soft_updates() -> Result<()> {
    // Create dummy networks with known weights
    let online_vars = create_test_varmap(1.0)?;
    let target_vars = create_test_varmap(0.0)?;

    // Test 1a: Single soft update with tau=0.1
    // Formula: θ_target = (1-τ)*θ_target + τ*θ_online
    //                   = (1-0.1)*0.0 + 0.1*1.0 = 0.1
    polyak_update(&online_vars, &target_vars, 0.1)?;
    let avg_weight = get_average_value(&target_vars);
    assert!(
        (avg_weight - 0.1).abs() < 0.01,
        "Expected target weight ≈0.1, got {}",
        avg_weight
    );
    println!("✓ Polyak update tau=0.1: target weight = {:.3} (expected 0.1)", avg_weight);

    // Test 1b: Multiple soft updates (convergence test)
    // Reset to initial state
    let online_vars_2 = create_test_varmap(1.0)?;
    let target_vars_2 = create_test_varmap(0.0)?;

    // Apply 10 updates with tau=0.1
    let mut weights = vec![];
    for _ in 0..10 {
        polyak_update(&online_vars_2, &target_vars_2, 0.1)?;
        weights.push(get_average_value(&target_vars_2));
    }

    // Verify monotonic increase
    for i in 1..weights.len() {
        assert!(
            weights[i] >= weights[i - 1] - 1e-6,
            "Non-monotonic at step {}: {:.4} -> {:.4}",
            i,
            weights[i - 1],
            weights[i]
        );
    }

    // Final weight should approach 1.0 (after 10 updates: ≈0.65)
    let final_weight = weights[9];
    assert!(
        final_weight > 0.6 && final_weight < 0.7,
        "Final weight after 10 updates should be 0.6-0.7, got {}",
        final_weight
    );
    println!(
        "✓ 10 Polyak updates tau=0.1: weight[0] = {:.4}, weight[9] = {:.4}",
        weights[0], final_weight
    );

    // Test 1c: Verify weights are blended correctly (90% old, 10% new)
    let online_vars_3 = create_test_varmap(2.0)?;
    let target_vars_3 = create_test_varmap(0.0)?;
    polyak_update(&online_vars_3, &target_vars_3, 0.1)?;
    let blended_weight = get_average_value(&target_vars_3);
    let expected = 0.9 * 0.0 + 0.1 * 2.0; // = 0.2
    assert!(
        (blended_weight - expected).abs() < 0.01,
        "Blended weight should be {:.2}, got {}",
        expected,
        blended_weight
    );
    println!(
        "✓ Weight blending 90%*0.0 + 10%*2.0 = {:.3} (expected 0.2)",
        blended_weight
    );

    Ok(())
}

// ============================================================================
// Test 2: Polyak Tau Parameter Validation
// ============================================================================

#[test]
fn test_polyak_tau_cli_integration() -> Result<()> {
    // Test 2a: Default tau=1.0 (hard updates, backward compatible)
    // When tau=1.0, Polyak update becomes: θ_target = 0*θ_target + 1.0*θ_online = θ_online
    // This is equivalent to a hard update (full weight copy)
    let online_vars = create_test_varmap(1.0)?;
    let target_vars = create_test_varmap(0.0)?;
    polyak_update(&online_vars, &target_vars, 1.0)?;
    let weight = get_average_value(&target_vars);
    assert!(
        (weight - 1.0).abs() < 1e-6,
        "tau=1.0 should be hard update, got {}",
        weight
    );
    println!("✓ CLI flag --tau 1.0: hard update (backward compatible), weight = {:.3}", weight);

    // Test 2b: Custom tau=0.005 (Rainbow DQN soft updates)
    // Rainbow DQN uses tau=0.001, but we test with 0.005 for faster convergence
    // Convergence half-life = ln(0.5) / ln(1-τ) ≈ 138 steps for tau=0.005
    let online_vars_2 = create_test_varmap(1.0)?;
    let target_vars_2 = create_test_varmap(0.0)?;
    polyak_update(&online_vars_2, &target_vars_2, 0.005)?;
    let weight_2 = get_average_value(&target_vars_2);
    let expected_2 = 0.995 * 0.0 + 0.005 * 1.0; // = 0.005
    assert!(
        (weight_2 - expected_2).abs() < 0.001,
        "tau=0.005 should give weight ≈0.005, got {}",
        weight_2
    );
    println!("✓ CLI flag --tau 0.005: soft update, weight = {:.4} (expected 0.005)", weight_2);

    // Test 2c: Verify convergence half-life calculation
    let half_life_rainbow = convergence_half_life(0.001);
    assert!(
        (half_life_rainbow - 693.0).abs() < 1.0,
        "Rainbow tau=0.001 should have half-life ≈693, got {}",
        half_life_rainbow
    );
    println!("✓ Rainbow tau=0.001: convergence half-life = {:.0} steps", half_life_rainbow);

    let half_life_fast = convergence_half_life(0.005);
    assert!(
        (half_life_fast - 138.0).abs() < 2.0,
        "Fast tau=0.005 should have half-life ≈138, got {}",
        half_life_fast
    );
    println!("✓ Fast tau=0.005: convergence half-life = {:.0} steps", half_life_fast);

    Ok(())
}

// ============================================================================
// Test 3: Polyak Update Convergence
// ============================================================================

#[test]
fn test_polyak_convergence() -> Result<()> {
    // Test that Polyak updates converge to online network weights
    let online_vars = create_test_varmap(1.0)?;
    let target_vars = create_test_varmap(0.0)?;

    // Apply 100 updates with tau=0.1
    for _ in 0..100 {
        polyak_update(&online_vars, &target_vars, 0.1)?;
    }

    let final_weight = get_average_value(&target_vars);
    assert!(
        final_weight > 0.99,
        "After 100 updates, target should converge to ≈1.0, got {}",
        final_weight
    );
    println!(
        "✓ Polyak convergence: 100 updates tau=0.1 → weight = {:.4} (expected >0.99)",
        final_weight
    );

    Ok(())
}

// ============================================================================
// Test 4: Backward Compatibility (Hard Updates)
// ============================================================================

#[test]
fn test_backward_compatibility_hard_updates() -> Result<()> {
    // Test 4a: Hard update copies all weights exactly
    let online_vars = create_test_varmap(1.0)?;
    let target_vars = create_test_varmap(0.0)?;
    hard_update(&online_vars, &target_vars)?;
    let weight = get_average_value(&target_vars);
    assert!(
        (weight - 1.0).abs() < 1e-6,
        "Hard update should copy weights exactly, got {}",
        weight
    );
    println!("✓ Hard update (legacy mode): target = {:.3} (expected 1.0)", weight);

    // Test 4b: Hard update is equivalent to Polyak with tau=1.0
    let online_vars_2 = create_test_varmap(2.0)?;
    let target_vars_2a = create_test_varmap(0.0)?;
    let target_vars_2b = create_test_varmap(0.0)?;

    hard_update(&online_vars_2, &target_vars_2a)?;
    polyak_update(&online_vars_2, &target_vars_2b, 1.0)?;

    let weight_hard = get_average_value(&target_vars_2a);
    let weight_polyak = get_average_value(&target_vars_2b);
    assert!(
        (weight_hard - weight_polyak).abs() < 1e-6,
        "Hard update and Polyak(tau=1.0) should be identical"
    );
    println!(
        "✓ Backward compatibility: hard update = Polyak(tau=1.0) = {:.3}",
        weight_hard
    );

    Ok(())
}

// ============================================================================
// Test 5: Soft Update Performance
// ============================================================================

#[test]
fn test_soft_update_stability() -> Result<()> {
    // Test that soft updates maintain stable weight values
    let online_vars = create_test_varmap(5.0)?;
    let target_vars = create_test_varmap(0.0)?;

    // Apply 50 updates with small tau
    for _ in 0..50 {
        polyak_update(&online_vars, &target_vars, 0.01)?;
    }

    let weight = get_average_value(&target_vars);
    assert!(
        weight > 1.5 && weight < 2.5,
        "Soft updates should produce stable weights, got {}",
        weight
    );
    println!(
        "✓ Soft update stability: 50 updates tau=0.01 → weight = {:.3} (expected 1.5-2.5)",
        weight
    );

    Ok(())
}

// ============================================================================
// Test 6: Tau Edge Cases
// ============================================================================

#[test]
fn test_tau_edge_cases() -> Result<()> {
    // Test 6a: tau=0.0 (no update)
    let online_vars = create_test_varmap(1.0)?;
    let target_vars = create_test_varmap(0.0)?;
    polyak_update(&online_vars, &target_vars, 0.0)?;
    let weight = get_average_value(&target_vars);
    assert!(
        weight.abs() < 1e-6,
        "tau=0.0 should not update target, got {}",
        weight
    );
    println!("✓ tau=0.0: no update, target = {:.6} (expected ≈0.0)", weight);

    // Test 6b: tau=0.001 (very slow convergence, Rainbow DQN)
    let online_vars_2 = create_test_varmap(1.0)?;
    let target_vars_2 = create_test_varmap(0.0)?;
    polyak_update(&online_vars_2, &target_vars_2, 0.001)?;
    let weight_2 = get_average_value(&target_vars_2);
    let expected = 0.001;
    assert!(
        (weight_2 - expected).abs() < 0.0001,
        "tau=0.001 should give weight ≈{}, got {}",
        expected,
        weight_2
    );
    println!("✓ tau=0.001: very slow update, weight = {:.6} (expected 0.001)", weight_2);

    Ok(())
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Create a test VarMap with all values set to a specific constant
fn create_test_varmap(value: f32) -> Result<VarMap> {
    use candle_core::{DType, Var};

    let varmap = VarMap::new();

    // Create test tensors (10x10 weight matrix, 10-dim bias vector)
    let weight = (Tensor::ones(&[10, 10], DType::F32, &Device::Cpu)? * value as f64)?;
    let bias = (Tensor::ones(&[10], DType::F32, &Device::Cpu)? * value as f64)?;

    // Insert into VarMap
    let mut data = varmap.data().lock().unwrap();
    data.insert("layer1.weight".to_string(), Var::from_tensor(&weight)?);
    data.insert("layer1.bias".to_string(), Var::from_tensor(&bias)?);
    drop(data);

    Ok(varmap)
}

/// Extract average value across all tensors in a VarMap
fn get_average_value(varmap: &VarMap) -> f32 {
    let data = varmap.data().lock().unwrap();
    let mut sum = 0.0;
    let mut count = 0;

    for (_, tensor) in data.iter() {
        let t: &Tensor = tensor.as_ref();
        sum += t.mean_all().unwrap().to_scalar::<f32>().unwrap();
        count += 1;
    }

    sum / count as f32
}