foxhunt/ml/tests/polyak_averaging_test.rs

#[cfg(test)]
mod polyak_tests {
    use ml::dqn::dqn::DQNConfig;
    use tch::{nn, Device, Tensor};

    /// Helper: Build a simple 2-layer network for testing
    fn build_test_network(vs: &nn::Path, input_dim: i64, output_dim: i64) -> nn::Sequential {
        nn::seq()
            .add(nn::linear(vs / "fc1", input_dim, 64, Default::default()))
            .add_fn(|x| x.relu())
            .add(nn::linear(vs / "fc2", 64, output_dim, Default::default()))
    }

    /// Helper: Get average weight value across all parameters
    fn get_average_weight(vs: &nn::VarStore) -> f64 {
        let mut sum = 0.0;
        let mut count = 0;

        for (_, param) in vs.variables() {
            sum += f64::try_from(param.mean(tch::Kind::Float)).unwrap();
            count += 1;
        }

        sum / count as f64
    }

    /// Helper: Polyak averaging function (to be implemented in main code)
    fn polyak_update(online_vs: &nn::VarStore, target_vs: &nn::VarStore, tau: f64) {
        for ((_, online_param), (_, target_param)) in
            online_vs.variables().zip(target_vs.variables())
        {
            // θ_target = (1-τ)*θ_target + τ*θ_online
            let new_target = (1.0 - tau) * &*target_param + tau * &*online_param;
            target_param.copy_(&new_target);
        }
    }

    #[test]
    fn test_polyak_single_update() {
        // GIVEN: Q-network and target network with different weights
        let vs_q = nn::VarStore::new(Device::Cpu);
        let vs_target = nn::VarStore::new(Device::Cpu);

        let _q_net = build_test_network(&vs_q.root(), 10, 3);
        let _target_net = build_test_network(&vs_target.root(), 10, 3);

        // Set Q-network weights to 1.0
        for (_, param) in vs_q.variables() {
            let _ = param.fill_(1.0);
        }

        // Set target weights to 0.0
        for (_, param) in vs_target.variables() {
            let _ = param.fill_(0.0);
        }

        // WHEN: Polyak update with τ=0.1
        polyak_update(&vs_q, &vs_target, 0.1);

        // THEN: Target should be 0.1 * 1.0 + 0.9 * 0.0 = 0.1
        for (_, param) in vs_target.variables() {
            let value = f64::try_from(param.mean(tch::Kind::Float)).unwrap();
            assert!(
                (value - 0.1).abs() < 0.01,
                "Expected target weight ≈0.1, got {}",
                value
            );
        }

        println!("✓ Single Polyak update: target weights = 0.1 (expected)");
    }

    #[test]
    fn test_gradual_convergence() {
        // GIVEN: Q-net at 1.0, target at 0.0
        let vs_q = nn::VarStore::new(Device::Cpu);
        let vs_target = nn::VarStore::new(Device::Cpu);

        let _q_net = build_test_network(&vs_q.root(), 10, 3);
        let _target_net = build_test_network(&vs_target.root(), 10, 3);

        // Initialize weights
        for (_, param) in vs_q.variables() {
            let _ = param.fill_(1.0);
        }
        for (_, param) in vs_target.variables() {
            let _ = param.fill_(0.0);
        }

        // WHEN: Apply Polyak updates for 100 steps (τ=0.01)
        let mut target_weights = vec![];
        for _step in 0..100 {
            polyak_update(&vs_q, &vs_target, 0.01);
            let w = get_average_weight(&vs_target);
            target_weights.push(w);
        }

        // THEN: Check monotonic increase
        for i in 1..target_weights.len() {
            assert!(
                target_weights[i] >= target_weights[i - 1] - 1e-6,
                "Target weights should increase monotonically at step {}: {} -> {}",
                i,
                target_weights[i - 1],
                target_weights[i]
            );
        }

        // Final weight should be close to 1.0 (but not exactly)
        let final_weight = target_weights[99];
        assert!(
            final_weight > 0.6 && final_weight < 1.0,
            "Final weight should be 0.6-1.0, got {}",
            final_weight
        );

        println!(
            "✓ Gradual convergence: weight[0] = {:.4}, weight[99] = {:.4}",
            target_weights[0], final_weight
        );
    }

    #[test]
    fn test_rainbow_tau_value() {
        // GIVEN: Rainbow's τ=0.001
        let tau = 0.001;

        // WHEN: Calculate convergence half-life
        // Formula: t_half = ln(0.5) / ln(1 - τ)
        let half_life = (-0.5_f64.ln()) / (-(1.0 - tau).ln());

        // THEN: Should converge slowly (half-life ≈ 693 steps)
        assert!(
            half_life > 600.0 && half_life < 800.0,
            "Expected half-life ≈693, got {:.0}",
            half_life
        );

        println!(
            "✓ Rainbow τ=0.001 gives half-life = {:.0} steps (expected ≈693)",
            half_life
        );
    }

    #[test]
    fn test_polyak_vs_hard_update_stability() {
        // GIVEN: Networks with noisy weight updates
        let vs_q = nn::VarStore::new(Device::Cpu);
        let vs_target_soft = nn::VarStore::new(Device::Cpu);
        let vs_target_hard = nn::VarStore::new(Device::Cpu);

        let _q_net = build_test_network(&vs_q.root(), 10, 3);
        let _target_soft = build_test_network(&vs_target_soft.root(), 10, 3);
        let _target_hard = build_test_network(&vs_target_hard.root(), 10, 3);

        // Initialize all to 0.0
        for vs in [&vs_q, &vs_target_soft, &vs_target_hard] {
            for (_, param) in vs.variables() {
                let _ = param.fill_(0.0);
            }
        }

        // WHEN: Simulate 100 training steps with noisy Q-network updates
        let mut soft_variance = 0.0;
        let mut hard_variance = 0.0;
        let mut prev_soft = 0.0;
        let mut prev_hard = 0.0;

        for step in 0..100 {
            // Add noise to Q-network
            for (_, param) in vs_q.variables() {
                let noise = Tensor::randn(&param.size(), (tch::Kind::Float, Device::Cpu)) * 0.1;
                let _ = param.add_(&noise);
            }

            // Soft update (every step)
            polyak_update(&vs_q, &vs_target_soft, 0.001);
            let soft_weight = get_average_weight(&vs_target_soft);
            if step > 0 {
                soft_variance += (soft_weight - prev_soft).powi(2);
            }
            prev_soft = soft_weight;

            // Hard update (every 10 steps)
            if step % 10 == 0 {
                for ((_, q_param), (_, target_param)) in
                    vs_q.variables().zip(vs_target_hard.variables())
                {
                    target_param.copy_(&q_param);
                }
            }
            let hard_weight = get_average_weight(&vs_target_hard);
            if step > 0 {
                hard_variance += (hard_weight - prev_hard).powi(2);
            }
            prev_hard = hard_weight;
        }

        // THEN: Soft updates should have lower variance
        soft_variance /= 99.0;
        hard_variance /= 99.0;

        assert!(
            soft_variance < hard_variance,
            "Soft updates should have lower variance: soft={:.6} vs hard={:.6}",
            soft_variance,
            hard_variance
        );

        let reduction = ((hard_variance - soft_variance) / hard_variance) * 100.0;
        println!(
            "✓ Polyak reduces variance by {:.1}% (soft={:.6}, hard={:.6})",
            reduction, soft_variance, hard_variance
        );
    }

    #[test]
    fn test_extreme_tau_values() {
        // Test boundary conditions
        let vs_q = nn::VarStore::new(Device::Cpu);
        let vs_target = nn::VarStore::new(Device::Cpu);

        let _q_net = build_test_network(&vs_q.root(), 10, 3);
        let _target_net = build_test_network(&vs_target.root(), 10, 3);

        // Initialize
        for (_, param) in vs_q.variables() {
            let _ = param.fill_(1.0);
        }
        for (_, param) in vs_target.variables() {
            let _ = param.fill_(0.0);
        }

        // Test τ=0.0 (no update)
        polyak_update(&vs_q, &vs_target, 0.0);
        let weight_tau_0 = get_average_weight(&vs_target);
        assert!(
            (weight_tau_0 - 0.0).abs() < 1e-6,
            "τ=0.0 should not update target"
        );

        // Test τ=1.0 (full copy)
        polyak_update(&vs_q, &vs_target, 1.0);
        let weight_tau_1 = get_average_weight(&vs_target);
        assert!(
            (weight_tau_1 - 1.0).abs() < 1e-6,
            "τ=1.0 should copy Q-network"
        );

        println!(
            "✓ Extreme τ values: τ=0.0 → {:.6}, τ=1.0 → {:.6}",
            weight_tau_0, weight_tau_1
        );
    }
}