foxhunt/ml/tests/rainbow_capacity_tests.rs

#[cfg(test)]
mod rainbow_capacity_tests {
    use anyhow::Result;
    use candle_core::{DType, Device, Tensor};
    use candle_nn::{Module, VarBuilder, VarMap};
    use ml::dqn::rainbow_network::{RainbowNetwork, RainbowNetworkConfig};

    #[test]
    fn test_default_hidden_sizes_reduced() -> Result<()> {
        let config = RainbowNetworkConfig::default();
        assert_eq!(
            config.hidden_sizes,
            vec![256, 128],
            "Default hidden sizes should be [256, 128] for reduced capacity (was [512, 512])"
        );
        Ok(())
    }

    #[test]
    fn test_default_dropout_increased() -> Result<()> {
        let config = RainbowNetworkConfig::default();
        assert!(
            (config.dropout_rate - 0.3).abs() < 0.001,
            "Default dropout rate should be 0.3 for better regularization (was 0.1)"
        );
        Ok(())
    }

    #[test]
    fn test_layer_norm_enabled_by_default() -> Result<()> {
        let config = RainbowNetworkConfig::default();
        assert!(
            config.use_layer_norm,
            "Layer normalization should be enabled by default for anti-overfitting"
        );
        assert!(
            (config.layer_norm_eps - 1e-5).abs() < 1e-10,
            "Layer norm epsilon should be 1e-5"
        );
        Ok(())
    }

    #[test]
    fn test_network_forward_pass_with_reduced_capacity() -> Result<()> {
        // Ensure network still functions correctly with reduced capacity
        let device = Device::Cpu;
        let varmap = VarMap::new();
        let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

        let config = RainbowNetworkConfig {
            input_size: 64,
            num_actions: 3,
            hidden_sizes: vec![256, 128],
            dropout_rate: 0.3,
            ..Default::default()
        };

        let network = RainbowNetwork::new(&vs, config)
            .map_err(|e| anyhow::anyhow!("Failed to create network: {}", e))?;

        let state = Tensor::randn(0.0f32, 1.0, (1, 64), &device)?;
        let output = network.forward(&state)?;

        // Output should be [batch_size, num_actions, num_atoms]
        let output_shape = output.shape().dims();
        assert_eq!(
            output_shape,
            &[1, 3, 51],
            "Output shape should be [1, 3, 51] for batch=1, actions=3, atoms=51"
        );
        Ok(())
    }

    #[test]
    fn test_capacity_comparison_with_legacy() -> Result<()> {
        // Compare new reduced capacity with old larger network
        let device = Device::Cpu;

        // Reduced capacity network
        let reduced_varmap = VarMap::new();
        let reduced_vs = VarBuilder::from_varmap(&reduced_varmap, DType::F32, &device);
        let reduced_config = RainbowNetworkConfig {
            input_size: 64,
            num_actions: 3,
            hidden_sizes: vec![256, 128], // New reduced
            dropout_rate: 0.3,
            ..Default::default()
        };
        let _reduced_net = RainbowNetwork::new(&reduced_vs, reduced_config.clone())
            .map_err(|e| anyhow::anyhow!("Failed to create reduced network: {}", e))?;

        // Legacy larger network
        let legacy_varmap = VarMap::new();
        let legacy_vs = VarBuilder::from_varmap(&legacy_varmap, DType::F32, &device);
        let legacy_config = RainbowNetworkConfig {
            input_size: 64,
            num_actions: 3,
            hidden_sizes: vec![512, 512], // Old larger
            dropout_rate: 0.1,
            ..Default::default()
        };
        let _legacy_net = RainbowNetwork::new(&legacy_vs, legacy_config.clone())
            .map_err(|e| anyhow::anyhow!("Failed to create legacy network: {}", e))?;

        // Count parameters (approximate calculation)
        let reduced_params = estimate_parameter_count(&reduced_config);
        let legacy_params = estimate_parameter_count(&legacy_config);

        assert!(
            reduced_params < legacy_params,
            "Reduced network ({} params) should have fewer parameters than legacy ({} params)",
            reduced_params,
            legacy_params
        );

        // Should be roughly 3-4x reduction
        let ratio = legacy_params as f64 / reduced_params as f64;
        assert!(
            ratio > 2.0 && ratio < 5.0,
            "Parameter reduction ratio should be 2-5x, got {:.2}x",
            ratio
        );

        println!("✓ Reduced network: {} params", reduced_params);
        println!("✓ Legacy network: {} params", legacy_params);
        println!("✓ Reduction ratio: {:.2}x", ratio);

        Ok(())
    }

    #[test]
    fn test_parameter_count_reasonable() -> Result<()> {
        // Verify total params < 500K for 64-dim input, 3 actions, 51 atoms
        let config = RainbowNetworkConfig {
            input_size: 64,
            num_actions: 3,
            hidden_sizes: vec![256, 128],
            dropout_rate: 0.3,
            ..Default::default()
        };

        let param_count = estimate_parameter_count(&config);

        assert!(
            param_count < 500_000,
            "Parameter count {} should be < 500K to prevent overfitting",
            param_count
        );

        println!("✓ Rainbow network parameter count: {}", param_count);

        Ok(())
    }

    #[test]
    fn test_dueling_architecture_enabled() -> Result<()> {
        let config = RainbowNetworkConfig::default();
        assert!(
            config.dueling,
            "Dueling architecture should be enabled by default"
        );
        Ok(())
    }

    #[test]
    fn test_noisy_layers_enabled() -> Result<()> {
        let config = RainbowNetworkConfig::default();
        assert!(
            config.use_noisy_layers,
            "Noisy layers should be enabled by default for exploration"
        );
        Ok(())
    }

    /// Estimate parameter count for Rainbow network
    /// This is an approximation of the actual parameter count
    fn estimate_parameter_count(config: &RainbowNetworkConfig) -> usize {
        let mut total_params = 0;

        // Feature extraction layers
        let mut current_size = config.input_size;
        for &hidden_size in &config.hidden_sizes {
            // Linear layer: (input * output) + bias
            total_params += current_size * hidden_size + hidden_size;
            // LayerNorm: 2 * hidden_size (weight + bias)
            if config.use_layer_norm {
                total_params += 2 * hidden_size;
            }
            current_size = hidden_size;
        }

        let final_feature_size = current_size;
        let num_atoms = config.distributional.num_atoms;

        if config.dueling {
            // Value stream
            let value_hidden = final_feature_size / 2;
            total_params += final_feature_size * value_hidden + value_hidden;
            if config.use_layer_norm {
                total_params += 2 * value_hidden;
            }
            // Value distribution
            total_params += value_hidden * num_atoms + num_atoms;

            // Advantage stream
            let advantage_hidden = final_feature_size / 2;
            total_params += final_feature_size * advantage_hidden + advantage_hidden;
            if config.use_layer_norm {
                total_params += 2 * advantage_hidden;
            }
            // Advantage distribution
            let advantage_output = config.num_actions * num_atoms;
            total_params += advantage_hidden * advantage_output + advantage_output;
        } else {
            // Single action distribution output
            let output_size = config.num_actions * num_atoms;
            total_params += final_feature_size * output_size + output_size;
        }

        total_params
    }
}