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

//! Generalized Advantage Estimation (GAE) implementation
//!
//! This module implements GAE for computing advantages in PPO training.
//! GAE provides a good trade-off between bias and variance in advantage estimation.

use serde::{Deserialize, Serialize};

use super::trajectories::Trajectory;
use ml_core::MLError;

/// Configuration for GAE computation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GAEConfig {
    /// Discount factor (gamma)
    pub gamma: f32,
    /// GAE parameter (lambda) for bias-variance trade-off
    pub lambda: f32,
    /// Whether to normalize advantages
    pub normalize_advantages: bool,
}

impl Default for GAEConfig {
    fn default() -> Self {
        Self {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        }
    }
}

/// Compute GAE advantages for a single trajectory
///
/// GAE(γ, λ) computes advantages as:
/// `A_t` = `δ_t` + (γλ)δ_{t+1} + (γλ)^2δ_{t+2} + ...
/// where `δ_t` = `r_t` + γV(s_{t+1}) - `V(s_t)`
pub fn compute_gae_single_trajectory(
    rewards: &[f32],
    values: &[f32],
    dones: &[bool],
    next_value: f32,
    config: &GAEConfig,
) -> Result<(Vec<f32>, Vec<f32>), MLError> {
    if rewards.len() != values.len() || rewards.len() != dones.len() {
        return Err(MLError::ValidationError {
            message: "Mismatched lengths in GAE computation".to_owned(),
        });
    }

    let length = rewards.len();
    let mut advantages = vec![0.0; length];
    let mut returns = vec![0.0; length];

    let mut next_advantage = 0.0;
    let mut next_return = next_value;

    // Compute GAE backwards through the trajectory
    for i in (0..length).rev() {
        // Safe indexing with bounds checking
        let reward = *rewards.get(i).ok_or_else(|| MLError::ValidationError {
            message: format!("Index {} out of bounds for rewards", i),
        })?;
        let value = *values.get(i).ok_or_else(|| MLError::ValidationError {
            message: format!("Index {} out of bounds for values", i),
        })?;
        let done = *dones.get(i).ok_or_else(|| MLError::ValidationError {
            message: format!("Index {} out of bounds for dones", i),
        })?;

        // Compute TD error (temporal difference)
        let next_non_terminal = if done { 0.0 } else { 1.0 };
        let next_value_estimate = if i == length - 1 {
            next_value
        } else {
            *values.get(i + 1).ok_or_else(|| MLError::ValidationError {
                message: format!("Index {} out of bounds for next value", i + 1),
            })?
        };

        let delta = reward + config.gamma * next_value_estimate * next_non_terminal - value;

        // Compute advantage using GAE
        *advantages
            .get_mut(i)
            .ok_or_else(|| MLError::ValidationError {
                message: format!("Index {} out of bounds for advantages", i),
            })? = delta + config.gamma * config.lambda * next_non_terminal * next_advantage;

        // Compute return
        *returns.get_mut(i).ok_or_else(|| MLError::ValidationError {
            message: format!("Index {} out of bounds for returns", i),
        })? = reward + config.gamma * next_non_terminal * next_return;

        // Update for next iteration (going backwards)
        let advantage_i = *advantages.get(i).ok_or_else(|| MLError::ValidationError {
            message: format!("Index {} out of bounds for advantages", i),
        })?;
        let return_i = *returns.get(i).ok_or_else(|| MLError::ValidationError {
            message: format!("Index {} out of bounds for returns", i),
        })?;

        next_advantage = advantage_i;
        next_return = return_i;
    }

    Ok((advantages, returns))
}

/// Compute GAE for multiple trajectories
pub fn compute_gae(
    trajectories: &[Trajectory],
    config: &GAEConfig,
) -> Result<(Vec<f32>, Vec<f32>), MLError> {
    let mut all_advantages = Vec::new();
    let mut all_returns = Vec::new();

    for trajectory in trajectories {
        let rewards = trajectory.get_rewards();
        let values = trajectory.get_values();
        let dones = trajectory.get_dones();

        // For terminal trajectories, next value is 0
        // For non-terminal trajectories, we use the last value estimate
        let next_value = if trajectory.is_complete() {
            0.0
        } else {
            values.last().copied().unwrap_or(0.0)
        };

        let (advantages, returns) =
            compute_gae_single_trajectory(&rewards, &values, &dones, next_value, config)?;

        all_advantages.extend(advantages);
        all_returns.extend(returns);
    }

    // NOTE: Advantage normalization is NOT performed here.
    // It is done once in PPO::update() via TrajectoryBatch::normalize_advantages()
    // to avoid double-normalization (correctness bug + wasted CPU).
    // The config.normalize_advantages flag is intentionally ignored at this level.

    Ok((all_advantages, all_returns))
}

/// Normalize advantages to have zero mean and unit variance
pub fn normalize_advantages(advantages: &mut [f32]) -> Result<(), MLError> {
    if advantages.is_empty() {
        return Ok(());
    }

    // Compute mean
    let mean = advantages.iter().sum::<f32>() / advantages.len() as f32;

    // Compute variance
    let variance =
        advantages.iter().map(|a| (a - mean).powi(2)).sum::<f32>() / advantages.len() as f32;

    let std_dev = (variance + 1e-8).sqrt(); // Add small epsilon for numerical stability

    // Normalize
    for advantage in advantages {
        *advantage = (*advantage - mean) / std_dev;
    }

    Ok(())
}

/// Compute discounted returns without GAE (simple Monte Carlo)
pub fn compute_discounted_returns(
    trajectories: &[Trajectory],
    gamma: f32,
) -> Result<Vec<f32>, MLError> {
    let mut all_returns = Vec::new();

    for trajectory in trajectories {
        let returns = trajectory.compute_returns(gamma);
        all_returns.extend(returns);
    }

    Ok(all_returns)
}

/// Compute temporal difference (TD) advantages
/// `A_t` = `r_t` + γV(s_{t+1}) - `V(s_t)`
pub fn compute_td_advantages(
    trajectories: &[Trajectory],
    gamma: f32,
    normalize: bool,
) -> Result<Vec<f32>, MLError> {
    let mut all_advantages = Vec::new();

    for trajectory in trajectories {
        let rewards = trajectory.get_rewards();
        let values = trajectory.get_values();
        let dones = trajectory.get_dones();

        let mut advantages = Vec::with_capacity(rewards.len());

        for i in 0..rewards.len() {
            let reward = *rewards.get(i).ok_or_else(|| MLError::ValidationError {
                message: format!("Index {} out of bounds for rewards", i),
            })?;
            let value = *values.get(i).ok_or_else(|| MLError::ValidationError {
                message: format!("Index {} out of bounds for values", i),
            })?;
            let done = *dones.get(i).ok_or_else(|| MLError::ValidationError {
                message: format!("Index {} out of bounds for dones", i),
            })?;

            let next_value = if i == rewards.len() - 1 {
                if done {
                    0.0
                } else {
                    value
                }
            } else {
                *values.get(i + 1).ok_or_else(|| MLError::ValidationError {
                    message: format!("Index {} out of bounds for next value", i + 1),
                })?
            };

            let next_non_terminal = if done { 0.0 } else { 1.0 };
            let advantage = reward + gamma * next_value * next_non_terminal - value;
            advantages.push(advantage);
        }

        all_advantages.extend(advantages);
    }

    // Normalize if requested
    if normalize {
        normalize_advantages(&mut all_advantages)?;
    }

    Ok(all_advantages)
}

/// Configuration for different advantage estimation methods
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AdvantageMethod {
    /// Generalized Advantage Estimation
    GAE(GAEConfig),
    /// Simple temporal difference
    TemporalDifference { gamma: f32, normalize: bool },
    /// Monte Carlo returns minus baseline
    MonteCarlo { gamma: f32, normalize: bool },
}

impl Default for AdvantageMethod {
    fn default() -> Self {
        Self::GAE(GAEConfig::default())
    }
}

/// Compute advantages using the specified method
pub fn compute_advantages(
    trajectories: &[Trajectory],
    method: &AdvantageMethod,
) -> Result<(Vec<f32>, Vec<f32>), MLError> {
    match method {
        AdvantageMethod::GAE(config) => compute_gae(trajectories, config),
        AdvantageMethod::TemporalDifference { gamma, normalize } => {
            let advantages = compute_td_advantages(trajectories, *gamma, *normalize)?;
            let returns = compute_discounted_returns(trajectories, *gamma)?;
            Ok((advantages, returns))
        },
        AdvantageMethod::MonteCarlo { gamma, normalize } => {
            let returns = compute_discounted_returns(trajectories, *gamma)?;

            // Compute advantages as returns minus values
            let mut advantages = Vec::new();
            let mut return_idx = 0;

            for trajectory in trajectories {
                let values = trajectory.get_values();
                for value in values {
                    advantages.push(returns[return_idx] - value);
                    return_idx += 1;
                }
            }

            if *normalize {
                normalize_advantages(&mut advantages)?;
            }

            Ok((advantages, returns))
        },
    }
}

#[cfg(test)]
#[allow(
    clippy::assertions_on_result_states,
    clippy::doc_markdown
)]
mod tests {
    use super::*;
    use ml_core::action_space::FactoredAction;
    use ml_core::trading_action::TradingAction;
    use crate::trajectories::{Trajectory, TrajectoryStep};

    /// Helper: build a FactoredAction from a TradingAction for test brevity.
    fn fa(ta: TradingAction) -> FactoredAction {
        FactoredAction::from_trading_action(ta)
    }

    fn create_test_trajectory() -> Trajectory {
        let mut trajectory = Trajectory::new();

        // Add some test steps
        trajectory.add_step(TrajectoryStep::new(
            vec![1.0, 2.0],
            fa(TradingAction::Buy),
            -0.5,
            10.0, // value
            1.0,  // reward
            false,
        ));

        trajectory.add_step(TrajectoryStep::new(
            vec![2.0, 3.0],
            fa(TradingAction::Sell),
            -0.3,
            8.0, // value
            2.0, // reward
            false,
        ));

        trajectory.add_step(TrajectoryStep::new(
            vec![3.0, 4.0],
            fa(TradingAction::Hold),
            -0.7,
            5.0,  // value
            0.0,  // reward
            true, // done
        ));

        trajectory
    }

    #[test]
    fn test_gae_single_trajectory() -> Result<(), Box<dyn std::error::Error>> {
        let rewards = vec![1.0, 2.0, 0.0];
        let values = vec![10.0, 8.0, 5.0];
        let dones = vec![false, false, true];
        let next_value = 0.0; // Terminal state

        let config = GAEConfig {
            gamma: 0.9,
            lambda: 0.95,
            normalize_advantages: false,
        };

        let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
        assert!(result.is_ok());

        let (advantages, returns) = result?;
        assert_eq!(advantages.len(), 3);
        assert_eq!(returns.len(), 3);

        // Basic sanity checks
        assert!(advantages.iter().any(|&a| a != 0.0)); // Should have non-zero advantages
        assert!(returns.iter().any(|&r| r != 0.0)); // Should have non-zero returns
        Ok(())
    }

    #[test]
    fn test_gae_multiple_trajectories() -> Result<(), Box<dyn std::error::Error>> {
        let trajectory = create_test_trajectory();
        let trajectories = vec![trajectory];

        let config = GAEConfig::default();
        let result = compute_gae(&trajectories, &config);
        assert!(result.is_ok());

        let (advantages, returns) = result?;
        assert_eq!(advantages.len(), 3); // Should match trajectory length
        assert_eq!(returns.len(), 3);
        Ok(())
    }

    #[test]
    fn test_advantage_normalization() {
        let mut advantages = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let result = normalize_advantages(&mut advantages);
        assert!(result.is_ok());

        // Check zero mean (approximately)
        let mean = advantages.iter().sum::<f32>() / advantages.len() as f32;
        assert!(mean.abs() < 1e-6);

        // Check unit variance (approximately)
        let variance = advantages.iter().map(|a| a.powi(2)).sum::<f32>() / advantages.len() as f32;
        assert!((variance - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_discounted_returns() -> Result<(), Box<dyn std::error::Error>> {
        let trajectory = create_test_trajectory();
        let trajectories = vec![trajectory];

        let result = compute_discounted_returns(&trajectories, 0.9);
        assert!(result.is_ok());

        let returns = result?;
        assert_eq!(returns.len(), 3);

        // Returns should be discounted properly
        // Last return should be just the reward (0.0 since it's terminal)
        assert!((returns[2] - 0.0).abs() < 1e-6);
        // Second return should be reward + gamma * next_return
        assert!((returns[1] - (2.0 + 0.9 * 0.0)).abs() < 1e-6);
        // First return should include both future rewards
        assert!((returns[0] - (1.0 + 0.9 * 2.0)).abs() < 1e-6);
        Ok(())
    }

    #[test]
    fn test_td_advantages() -> Result<(), Box<dyn std::error::Error>> {
        let trajectory = create_test_trajectory();
        let trajectories = vec![trajectory];

        let result = compute_td_advantages(&trajectories, 0.9, false);
        assert!(result.is_ok());

        let advantages = result?;
        assert_eq!(advantages.len(), 3);

        // TD advantages should be r + γV(s') - V(s)
        // For terminal state: advantage = reward + 0 - value = 0.0 + 0 - 5.0 = -5.0
        assert!((advantages[2] - (-5.0)).abs() < 1e-6);
        Ok(())
    }

    #[test]
    fn test_advantage_methods() {
        let trajectory = create_test_trajectory();
        let trajectories = vec![trajectory];

        // Test GAE method
        let gae_method = AdvantageMethod::GAE(GAEConfig::default());
        let result = compute_advantages(&trajectories, &gae_method);
        assert!(result.is_ok());

        // Test TD method
        let td_method = AdvantageMethod::TemporalDifference {
            gamma: 0.9,
            normalize: true,
        };
        let result = compute_advantages(&trajectories, &td_method);
        assert!(result.is_ok());

        // Test Monte Carlo method
        let mc_method = AdvantageMethod::MonteCarlo {
            gamma: 0.9,
            normalize: false,
        };
        let result = compute_advantages(&trajectories, &mc_method);
        assert!(result.is_ok());
    }

    #[test]
    fn test_empty_trajectory_handling() -> Result<(), Box<dyn std::error::Error>> {
        let trajectories: Vec<Trajectory> = vec![];
        let config = GAEConfig::default();

        let result = compute_gae(&trajectories, &config);
        assert!(result.is_ok());

        let (advantages, returns) = result?;
        assert!(advantages.is_empty());
        assert!(returns.is_empty());
        Ok(())
    }

    #[test]
    fn test_mismatched_lengths_error() {
        let rewards = vec![1.0, 2.0];
        let values = vec![10.0]; // Mismatched length
        let dones = vec![false, false];
        let config = GAEConfig::default();

        let result = compute_gae_single_trajectory(&rewards, &values, &dones, 0.0, &config);
        assert!(result.is_err());
    }
}