foxhunt/crates/ml-core/src/xavier_init.rs

//! Xavier/Glorot initialization for DQN networks
//!
//! Implements Xavier uniform initialization to improve initial gradient flow
//! and prevent Q-value collapse during early training.
//!
//! Reference: Glorot & Bengio (2010) - "Understanding the difficulty of training deep feedforward neural networks"

use candle_core::{DType, Device, Result, Tensor};
use candle_nn::{Init, Linear, VarBuilder};

/// Initialize a weight tensor with Xavier uniform distribution
///
/// Xavier initialization samples weights from a uniform distribution U[-limit, limit]
/// where limit = sqrt(6 / (fan_in + fan_out)).
///
/// This produces a variance of Var(W) = 2 / (fan_in + fan_out), which helps maintain
/// stable gradients during backpropagation.
///
/// # Arguments
///
/// * `fan_in` - Number of input units
/// * `fan_out` - Number of output units
/// * `dtype` - Data type for the tensor (typically F32)
/// * `device` - Device to create tensor on (CPU or CUDA)
///
/// # Returns
///
/// A tensor of shape `(fan_out, fan_in)` initialized with Xavier uniform values
///
/// # Example
///
/// ```ignore
/// use candle_core::{DType, Device};
/// use ml::dqn::xavier_init::xavier_uniform;
///
/// let device = Device::Cpu;
/// let weights = xavier_uniform(64, 128, DType::F32, &device)?;
/// // weights shape: (128, 64) with variance ≈ 2/(64+128) = 0.0104
/// ```
pub fn xavier_uniform(
    fan_in: usize,
    fan_out: usize,
    dtype: DType,
    device: &Device,
) -> Result<Tensor> {
    // Calculate Xavier limit: sqrt(6 / (fan_in + fan_out))
    let limit = (6.0 / (fan_in + fan_out) as f64).sqrt();

    // Create tensor with shape (fan_out, fan_in) - Candle's Linear layer convention
    let shape = (fan_out, fan_in);

    // Sample from uniform distribution U(-limit, limit)
    // Note: Tensor::rand samples from U(a, b)
    Tensor::rand(-limit, limit, shape, device)?.to_dtype(dtype)
}

/// Create a candle Init for Xavier uniform initialization
///
/// Returns an Init::Uniform with bounds calculated using Xavier/Glorot formula.
/// This can be used with VarBuilder.get_with_hints() to initialize weights.
pub fn xavier_init(fan_in: usize, fan_out: usize) -> Init {
    let limit = (6.0 / (fan_in + fan_out) as f64).sqrt();
    Init::Uniform {
        lo: -limit,
        up: limit,
    }
}

/// Create a Linear layer with Xavier initialization and register in VarMap
///
/// This is a replacement for `candle_nn::linear()` that uses Xavier initialization
/// instead of the default Kaiming initialization. The weights and biases are properly
/// registered in the VarBuilder's VarMap for checkpointing and gradient tracking.
///
/// # Arguments
///
/// * `fan_in` - Number of input features
/// * `fan_out` - Number of output features
/// * `vb` - VarBuilder for registering variables
///
/// # Returns
///
/// A Linear layer with Xavier-initialized weights
pub fn linear_xavier(fan_in: usize, fan_out: usize, vb: VarBuilder<'_>) -> Result<Linear> {
    // Create Xavier initialization
    let init_ws = xavier_init(fan_in, fan_out);
    let ws = vb.get_with_hints((fan_out, fan_in), "weight", init_ws)?;

    // Bias initialization: uniform(-bound, bound) where bound = 1/sqrt(fan_in)
    let bound = 1.0 / (fan_in as f64).sqrt();
    let init_bs = Init::Uniform {
        lo: -bound,
        up: bound,
    };
    let bs = vb.get_with_hints(fan_out, "bias", init_bs)?;

    Ok(Linear::new(ws, Some(bs)))
}

/// Verify Xavier initialization statistics for testing
///
/// Returns (mean, variance) of the provided weight tensor.
/// For Xavier initialization:
/// - Mean should be ≈ 0
/// - Variance should be ≈ 2 / (fan_in + fan_out)
pub fn verify_xavier_stats(
    weights: &Tensor,
    fan_in: usize,
    fan_out: usize,
) -> Result<(f32, f32, f32)> {
    let weight_mean = weights.mean_all()?.to_scalar::<f32>()?;

    // Calculate variance manually: Var(X) = E[X²] - E[X]²
    let squared = weights.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;

    Ok((weight_mean, weight_var, expected_var))
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::Device;

    #[test]
    fn test_xavier_uniform_shape() -> Result<()> {
        let device = Device::Cpu;
        let weights = xavier_uniform(64, 128, DType::F32, &device)?;

        // Verify shape is (fan_out, fan_in)
        assert_eq!(weights.shape().dims(), &[128, 64]);
        Ok(())
    }

    #[test]
    fn test_xavier_uniform_statistics() -> Result<()> {
        let device = Device::Cpu;
        let fan_in = 64;
        let fan_out = 128;
        let weights = xavier_uniform(fan_in, fan_out, DType::F32, &device)?;

        let (mean, variance, expected_var) = verify_xavier_stats(&weights, fan_in, fan_out)?;

        // Mean should be close to zero
        assert!(
            mean.abs() < 0.05,
            "Mean {} should be close to 0 (< 0.05)",
            mean
        );

        // Variance should match Xavier formula within 20% tolerance
        let var_diff = (variance - expected_var).abs() / expected_var;
        assert!(
            var_diff < 0.20,
            "Variance {} should match expected {} (diff: {:.2}%)",
            variance,
            expected_var,
            var_diff * 100.0
        );

        Ok(())
    }

    #[test]
    fn test_xavier_uniform_range() -> Result<()> {
        let device = Device::Cpu;
        let fan_in = 64;
        let fan_out = 128;
        let weights = xavier_uniform(fan_in, fan_out, DType::F32, &device)?;

        // Calculate expected limit
        let expected_limit = (6.0 / (fan_in + fan_out) as f64).sqrt() as f32;

        // Get min/max values
        let weight_max = weights.flatten_all()?.max(0)?.to_scalar::<f32>()?;
        let weight_min = weights.flatten_all()?.min(0)?.to_scalar::<f32>()?;

        // Allow 50% margin for random variation
        let margin = 1.5;
        assert!(
            weight_max <= expected_limit * margin,
            "Max weight {} should be within Xavier limit {} * {}",
            weight_max,
            expected_limit,
            margin
        );
        assert!(
            weight_min >= -expected_limit * margin,
            "Min weight {} should be within Xavier limit -{} * {}",
            weight_min,
            expected_limit,
            margin
        );

        Ok(())
    }
}