foxhunt/ml/tests/ppo_trainer_continuous_tests.rs

//! Continuous PPO Trainer Tests
//!
//! Integration tests for continuous action space PPO training.

use anyhow::Result;
use ml::ppo::continuous_ppo::{
    ContinuousAction, ContinuousPPO, ContinuousPPOConfig, ContinuousTrajectory,
    ContinuousTrajectoryBatch, ContinuousTrajectoryStep,
};
use ml::ppo::continuous_policy::ContinuousPolicyConfig;
use ml::ppo::gae::GAEConfig;

#[test]
fn test_trainer_continuous_ppo_creation() -> Result<()> {
    // Create continuous PPO config
    let policy_config = ContinuousPolicyConfig {
        state_dim: 64,
        hidden_dims: vec![128, 64],
        min_log_std: -5.0,
        max_log_std: 2.0,
        init_log_std: -1.0,
        learnable_std: true,
        action_bounds: (-1.0, 1.0),
    };

    let config = ContinuousPPOConfig {
        state_dim: 64,
        policy_config,
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 0.000001,
        value_learning_rate: 0.001,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        batch_size: 64,
        mini_batch_size: 16,
        num_epochs: 5,
        max_grad_norm: 0.5,
    };

    // Create agent
    let agent = ContinuousPPO::new(config)?;

    // Verify initialization
    assert_eq!(agent.get_training_steps(), 0);
    assert_eq!(agent.get_config().state_dim, 64);

    Ok(())
}

#[test]
fn test_trainer_continuous_training_step() -> Result<()> {
    // Create continuous PPO config
    let policy_config = ContinuousPolicyConfig {
        state_dim: 10,
        hidden_dims: vec![16, 8],
        min_log_std: -5.0,
        max_log_std: 2.0,
        init_log_std: -1.0,
        learnable_std: true,
        action_bounds: (-1.0, 1.0),
    };

    let config = ContinuousPPOConfig {
        state_dim: 10,
        policy_config,
        value_hidden_dims: vec![16, 8],
        policy_learning_rate: 0.000001,
        value_learning_rate: 0.001,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        batch_size: 32,
        mini_batch_size: 8,
        num_epochs: 2,
        max_grad_norm: 0.5,
    };

    // Create agent
    let mut agent = ContinuousPPO::new(config)?;

    // Create simple trajectory
    let mut trajectory = ContinuousTrajectory::new();

    for i in 0..10 {
        let state = vec![0.1 * i as f32; 10];
        let action = ContinuousAction::new(0.5);
        let log_prob = -1.0;
        let reward = 1.0;
        let value = 0.5;
        let done = i == 9;

        let step =
            ContinuousTrajectoryStep::new(state, action, log_prob, reward, value, done);
        trajectory.add_step(step);
    }

    // Compute GAE advantages
    let steps = trajectory.steps();
    let rewards: Vec<f32> = steps.iter().map(|s| s.reward).collect();
    let values: Vec<f32> = steps.iter().map(|s| s.value).collect();
    let dones: Vec<bool> = steps.iter().map(|s| s.done).collect();

    let advantages = compute_gae_advantages(&rewards, &values, &dones, 0.99, 0.95);
    let returns = compute_returns(&rewards, 0.99);

    // Create batch
    let mut batch = ContinuousTrajectoryBatch::from_trajectories(
        vec![trajectory],
        advantages,
        returns,
    );

    // Training step
    let initial_steps = agent.get_training_steps();
    let (policy_loss, value_loss) = agent.update(&mut batch)?;

    // Verify training occurred
    assert_eq!(agent.get_training_steps(), initial_steps + 1);
    assert!(policy_loss.is_finite());
    assert!(value_loss.is_finite());

    Ok(())
}

#[test]
fn test_trainer_continuous_epoch_completion() -> Result<()> {
    // Create continuous PPO config
    let policy_config = ContinuousPolicyConfig {
        state_dim: 8,
        hidden_dims: vec![16],
        min_log_std: -5.0,
        max_log_std: 2.0,
        init_log_std: -1.0,
        learnable_std: true,
        action_bounds: (-1.0, 1.0),
    };

    let config = ContinuousPPOConfig {
        state_dim: 8,
        policy_config,
        value_hidden_dims: vec![16],
        policy_learning_rate: 0.000001,
        value_learning_rate: 0.001,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        batch_size: 32,
        mini_batch_size: 8,
        num_epochs: 3,
        max_grad_norm: 0.5,
    };

    // Create agent
    let mut agent = ContinuousPPO::new(config)?;

    // Create multiple trajectories
    let mut trajectories = Vec::new();

    for traj_idx in 0..3 {
        let mut trajectory = ContinuousTrajectory::new();

        for i in 0..10 {
            let state = vec![0.1 * (traj_idx * 10 + i) as f32; 8];
            let action = ContinuousAction::new(0.3 + traj_idx as f32 * 0.2);
            let log_prob = -1.5;
            let reward = 0.5 + i as f32 * 0.1;
            let value = 0.3;
            let done = i == 9;

            let step =
                ContinuousTrajectoryStep::new(state, action, log_prob, reward, value, done);
            trajectory.add_step(step);
        }

        trajectories.push(trajectory);
    }

    // Compute batch advantages
    let mut all_advantages = Vec::new();
    let mut all_returns = Vec::new();

    for trajectory in &trajectories {
        let steps = trajectory.steps();
        let rewards: Vec<f32> = steps.iter().map(|s| s.reward).collect();
        let values: Vec<f32> = steps.iter().map(|s| s.value).collect();
        let dones: Vec<bool> = steps.iter().map(|s| s.done).collect();

        let advantages = compute_gae_advantages(&rewards, &values, &dones, 0.99, 0.95);
        let returns = compute_returns(&rewards, 0.99);

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

    // Create batch
    let mut batch = ContinuousTrajectoryBatch::from_trajectories(
        trajectories,
        all_advantages,
        all_returns,
    );

    // Run multiple training steps
    for epoch in 0..3 {
        let (policy_loss, value_loss) = agent.update(&mut batch)?;

        assert!(policy_loss.is_finite(), "Policy loss NaN at epoch {}", epoch);
        assert!(value_loss.is_finite(), "Value loss NaN at epoch {}", epoch);
    }

    assert_eq!(agent.get_training_steps(), 3);

    Ok(())
}

#[test]
fn test_trainer_continuous_checkpoint_save() -> Result<()> {
    use tempfile::TempDir;

    // Create temporary directory for checkpoints
    let temp_dir = TempDir::new()?;
    let checkpoint_dir = temp_dir.path();

    // Create continuous PPO config
    let policy_config = ContinuousPolicyConfig {
        state_dim: 10,
        hidden_dims: vec![16],
        min_log_std: -5.0,
        max_log_std: 2.0,
        init_log_std: -1.0,
        learnable_std: true,
        action_bounds: (-1.0, 1.0),
    };

    let config = ContinuousPPOConfig {
        state_dim: 10,
        policy_config,
        value_hidden_dims: vec![16],
        policy_learning_rate: 0.000001,
        value_learning_rate: 0.001,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        batch_size: 32,
        mini_batch_size: 8,
        num_epochs: 2,
        max_grad_norm: 0.5,
    };

    // Create agent
    let agent = ContinuousPPO::new(config)?;

    // Save checkpoints
    let actor_path = checkpoint_dir.join("actor_test.safetensors");
    let critic_path = checkpoint_dir.join("critic_test.safetensors");

    agent
        .actor
        .vars()
        .save(&actor_path)
        .expect("Failed to save actor");
    agent
        .critic
        .vars()
        .save(&critic_path)
        .expect("Failed to save critic");

    // Verify checkpoint files exist
    assert!(actor_path.exists(), "Actor checkpoint not created");
    assert!(critic_path.exists(), "Critic checkpoint not created");

    // Verify files have non-zero size
    let actor_metadata = std::fs::metadata(&actor_path)?;
    let critic_metadata = std::fs::metadata(&critic_path)?;

    assert!(actor_metadata.len() > 0, "Actor checkpoint is empty");
    assert!(critic_metadata.len() > 0, "Critic checkpoint is empty");

    Ok(())
}

// Helper functions

fn compute_gae_advantages(
    rewards: &[f32],
    values: &[f32],
    dones: &[bool],
    gamma: f32,
    lambda: f32,
) -> Vec<f32> {
    let n = rewards.len();
    let mut advantages = vec![0.0; n];
    let mut gae = 0.0;

    for t in (0..n).rev() {
        let reward = rewards[t];
        let value = values[t];
        let next_value = if t + 1 < n { values[t + 1] } else { 0.0 };
        let done = dones[t];

        let mask = if done { 0.0 } else { 1.0 };
        let delta = reward + gamma * next_value * mask - value;
        gae = delta + gamma * lambda * mask * gae;

        advantages[t] = gae;
    }

    advantages
}

fn compute_returns(rewards: &[f32], gamma: f32) -> Vec<f32> {
    let n = rewards.len();
    let mut returns = vec![0.0; n];
    let mut cumulative = 0.0;

    for t in (0..n).rev() {
        cumulative = rewards[t] + gamma * cumulative;
        returns[t] = cumulative;
    }

    returns
}