foxhunt/crates/ml-ppo/src/entropy_regularization.rs

//! Entropy regularization for PPO to prevent policy collapse
//!
//! This module implements entropy-based reward shaping to maintain action diversity
//! during PPO training. Key features:
//! - Shannon entropy calculation from action probabilities
//! - Bonus/penalty system based on normalized entropy threshold
//! - Numerical stability with epsilon protection
//!
//! # Key Difference from DQN
//! Unlike DQN which computes softmax from Q-values, PPO already has action probabilities
//! from the policy network, making entropy calculation more direct.
//!
//! # Example
//! ```rust,no_run
//! use candle_core::{Tensor, Device};
//! use ml::ppo::entropy_regularization::EntropyRegularizer;
//!
//! let regularizer = EntropyRegularizer::new();
//! // 45 factored actions (5 exposure × 3 order × 3 urgency)
//! let action_probs = Tensor::new(&[1.0_f32 / 45.0; 45], &Device::Cpu).unwrap();
//! let bonus = regularizer.calculate_entropy_bonus(&action_probs).unwrap();
//! ```

use candle_core::{DType, Tensor};

use ml_core::MLError;

/// Entropy regularizer for preventing policy collapse in PPO
///
/// Implements Shannon entropy calculation and entropy-based reward shaping
/// to encourage exploration and maintain action diversity.
#[derive(Debug, Clone)]
pub struct EntropyRegularizer {
    /// Maximum possible entropy: `log(num_actions)`
    max_entropy: f64,
    /// Normalized entropy threshold (0.7) for bonus/penalty
    entropy_threshold: f64,
}

impl EntropyRegularizer {
    /// Number of factored actions (5 exposure × 3 order × 3 urgency)
    const NUM_ACTIONS: usize = 45;

    /// Create a new entropy regularizer for 45 factored actions.
    ///
    /// # Configuration
    /// - `max_entropy`: log(45) ≈ 3.807 for 45 factored actions
    /// - `entropy_threshold`: 0.7 normalized entropy
    ///   - Above 0.7: 2x bonus for high diversity
    ///   - Below 0.7: 3x penalty for low diversity
    pub fn new() -> Self {
        Self {
            max_entropy: (Self::NUM_ACTIONS as f64).ln(),
            entropy_threshold: 0.7,
        }
    }

    /// Calculate Shannon entropy from action probabilities
    ///
    /// # Arguments
    /// * `action_probs` - Action probability tensor, shape [`batch_size`, `num_actions`] or [`num_actions`]
    ///
    /// # Returns
    /// Raw Shannon entropy H(π) = -Σ π(a|s) * log(π(a|s))
    ///
    /// # Numerical Stability
    /// - Adds epsilon (1e-8) to prevent log(0) = -∞
    /// - Handles both single-state and batched inputs
    /// - Averages entropy across batch dimension if present
    pub fn calculate_entropy(&self, action_probs: &Tensor) -> Result<f64, MLError> {
        // Ensure action_probs is F32 to avoid dtype mismatches
        let action_probs_f32 = action_probs.to_dtype(DType::F32)?;

        // Shannon entropy H(π) = -Σ π(a|s) * log(π(a|s))
        // Add epsilon (1e-8) to prevent log(0) = -∞
        let epsilon = Tensor::new(&[1e-8_f32], action_probs.device())?
            .broadcast_as(action_probs_f32.shape())?;
        let action_probs_safe = action_probs_f32.add(&epsilon)?;
        let log_probs = action_probs_safe.log()?;
        let entropy = action_probs_f32
            .mul(&log_probs)?
            .neg()?
            .sum(candle_core::D::Minus1)?;

        // Average across batch dimension (if present)
        let avg_entropy = if entropy.dims().is_empty() {
            entropy.to_scalar::<f32>()? as f64
        } else {
            entropy.mean_all()?.to_scalar::<f32>()? as f64
        };

        Ok(avg_entropy)
    }

    /// Calculate entropy bonus/penalty from action probabilities
    ///
    /// # Arguments
    /// * `action_probs` - Action probability tensor, shape [`batch_size`, `num_actions`] or [`num_actions`]
    ///
    /// # Returns
    /// - Positive value: Bonus for high entropy (> 0.7 normalized)
    /// - Negative value: Penalty for low entropy (< 0.7 normalized)
    ///
    /// # Formula
    /// ```text
    /// Shannon Entropy: H(π) = -Σ π(a|s) * log(π(a|s))
    /// Normalized: H_norm = H(π) / log(num_actions)
    /// Bonus: H_norm * 2.0 if H_norm > 0.7
    /// Penalty: -(0.7 - H_norm) * 3.0 if H_norm <= 0.7
    /// ```
    pub fn calculate_entropy_bonus(&self, action_probs: &Tensor) -> Result<f64, MLError> {
        // Calculate raw entropy
        let avg_entropy = self.calculate_entropy(action_probs)?;

        // Normalize to [0, 1]
        let normalized_entropy = avg_entropy / self.max_entropy;

        // Apply bonus/penalty based on threshold
        if normalized_entropy > self.entropy_threshold {
            Ok(normalized_entropy * 2.0) // 2x bonus for high diversity
        } else {
            Ok(-(self.entropy_threshold - normalized_entropy) * 3.0) // 3x penalty for low diversity
        }
    }
}

impl Default for EntropyRegularizer {
    fn default() -> Self {
        Self::new()
    }
}

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

    /// Helper function to create probability tensor
    fn create_prob_tensor(probs: &[f32]) -> Result<Tensor, MLError> {
        let tensor = Tensor::new(probs, &Device::Cpu)?;
        Ok(tensor.reshape(&[1, probs.len()])?)
    }

    /// Create uniform probability vector over 45 actions
    fn uniform_45() -> Vec<f32> {
        vec![1.0 / 45.0; 45]
    }

    /// Create near-deterministic probability vector over 45 actions
    fn deterministic_45(dominant_action: usize) -> Vec<f32> {
        let mut probs = vec![0.001 / 44.0; 45];
        probs[dominant_action] = 0.999;
        probs
    }

    #[test]
    fn test_entropy_uniform_distribution() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = create_prob_tensor(&uniform_45())?;

        let entropy = regularizer.calculate_entropy(&probs)?;

        // Uniform distribution over 45 actions: entropy = log(45) ≈ 3.807
        let expected = (45.0_f64).ln();
        assert!(
            (entropy - expected).abs() < 0.01,
            "Expected entropy ~{expected:.4}, got {entropy}",
        );
        Ok(())
    }

    #[test]
    fn test_entropy_deterministic() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = create_prob_tensor(&deterministic_45(0))?;

        let entropy = regularizer.calculate_entropy(&probs)?;

        // Near-deterministic policy: entropy ≈ 0
        assert!(entropy < 0.05, "Expected entropy ~0, got {entropy}");
        Ok(())
    }

    #[test]
    fn test_bonus_high_diversity() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = create_prob_tensor(&uniform_45())?;

        let bonus = regularizer.calculate_entropy_bonus(&probs)?;

        // Uniform over 45: normalized entropy ≈ 1.0, bonus = 1.0 × 2.0 = 2.0
        assert!(
            (bonus - 2.0).abs() < 0.05,
            "Expected bonus ~2.0, got {bonus}",
        );
        Ok(())
    }

    #[test]
    fn test_penalty_low_diversity() -> Result<(), MLError> {
        let regularizer = EntropyRegularizer::new();
        let probs = create_prob_tensor(&deterministic_45(38))?;

        let penalty = regularizer.calculate_entropy_bonus(&probs)?;

        // Near-deterministic: normalized entropy ≈ 0, penalty ≈ -(0.7) × 3 = -2.1
        assert!(penalty < 0.0, "Expected penalty < 0.0, got {penalty}");
        assert!(penalty < -1.0, "Expected penalty < -1.0, got {penalty}");
        Ok(())
    }
}