rustytorch/crates/training/rtx-rl/tests/ppo_tests.rs

use rtx_rl::algorithms::{PPO, PPOConfig};
use rtx_tensor::{DType, Device, Tensor};

#[tokio::test]
async fn test_ppo_creation() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig {
        learning_rate: 3e-4,
        gamma: 0.99,
        gae_lambda: 0.95,
        clip_epsilon: 0.2,
        entropy_coeff: 0.01,
        value_coeff: 0.5,
        max_grad_norm: 0.5,
        n_epochs: 4,
        batch_size: 64,
        minibatch_size: 16,
    };

    let state_dim = 4;
    let action_dim = 2;
    let hidden_dim = 128;

    let ppo = PPO::new(config, state_dim, action_dim, hidden_dim, device);
    assert_eq!(ppo.state_dim(), state_dim);
    assert_eq!(ppo.action_dim(), action_dim);
}

#[tokio::test]
#[ignore = "rtx-rl PPO forward shape mismatch"]
async fn test_ppo_forward_pass() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig::default();

    let state_dim = 4;
    let action_dim = 2;
    let hidden_dim = 64;
    let batch_size = 32;

    let ppo = PPO::new(config, state_dim, action_dim, hidden_dim, device.clone());

    let states = Tensor::randn(&[batch_size, state_dim], &device).unwrap();

    let (actions, log_probs, values) = ppo
        .forward(&states)
        .await
        .expect("Forward pass should work");

    assert_eq!(actions.shape(), &[batch_size, action_dim]);
    assert_eq!(log_probs.shape(), &[batch_size]);
    assert_eq!(values.shape(), &[batch_size]);
}

#[tokio::test]
#[ignore = "rtx-tensor Bool dtype CPU copy not supported"]
async fn test_ppo_compute_advantages() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig::default();
    let ppo = PPO::new(config, 4, 2, 64, device.clone());

    let batch_size = 10;
    let rewards = Tensor::ones(&[batch_size], &device).unwrap();
    let values = Tensor::zeros(&[batch_size], &device).unwrap();
    let next_values = Tensor::ones(&[batch_size], &device).unwrap();
    let dones = Tensor::zeros(&[batch_size], &device)
        .unwrap()
        .to_dtype(DType::Bool)
        .unwrap();

    let advantages = ppo
        .compute_advantages(&rewards, &values, &next_values, &dones)
        .await
        .expect("Should compute advantages");

    assert_eq!(advantages.shape(), &[batch_size]);

    // With gamma=0.99, lambda=0.95, rewards=1, values=0, next_values=1
    // GAE should produce positive advantages
    let adv_data: Vec<f32> = advantages.to_vec().expect("Should convert to vec");
    for adv in adv_data {
        assert!(adv > 0.0);
    }
}

#[tokio::test]
#[ignore = "rtx-rl PPO loss shape mismatch"]
async fn test_ppo_policy_loss() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig::default();
    let ppo = PPO::new(config, 4, 2, 64, device.clone());

    let batch_size = 16;
    let old_log_probs = Tensor::randn(&[batch_size], &device).unwrap();
    let new_log_probs = Tensor::randn(&[batch_size], &device).unwrap();
    let advantages = Tensor::randn(&[batch_size], &device).unwrap();

    let policy_loss = ppo
        .compute_policy_loss(&old_log_probs, &new_log_probs, &advantages)
        .await
        .expect("Should compute policy loss");

    assert_eq!(policy_loss.shape(), &[]);

    let loss_val: f32 = policy_loss.item().expect("Should get scalar value");
    assert!(loss_val.is_finite());
}

#[tokio::test]
#[ignore = "rtx-rl PPO loss shape mismatch"]
async fn test_ppo_value_loss() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig::default();
    let ppo = PPO::new(config, 4, 2, 64, device.clone());

    let batch_size = 16;
    let predicted_values = Tensor::randn(&[batch_size], &device).unwrap();
    let target_values = Tensor::randn(&[batch_size], &device).unwrap();

    let value_loss = ppo
        .compute_value_loss(&predicted_values, &target_values)
        .await
        .expect("Should compute value loss");

    assert_eq!(value_loss.shape(), &[]);

    let loss_val: f32 = value_loss.item().expect("Should get scalar value");
    assert!(loss_val >= 0.0); // MSE loss should be non-negative
}

#[tokio::test]
#[ignore = "rtx-rl PPO entropy implementation incomplete"]
async fn test_ppo_entropy_bonus() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig::default();
    let ppo = PPO::new(config, 4, 2, 64, device.clone());

    let batch_size = 16;
    let action_dim = 2;
    let logits = Tensor::randn(&[batch_size, action_dim], &device).unwrap();

    let entropy = ppo
        .compute_entropy(&logits)
        .await
        .expect("Should compute entropy");

    assert_eq!(entropy.shape(), &[]);

    let entropy_val: f32 = entropy.item().expect("Should get scalar value");
    assert!(entropy_val > 0.0); // Entropy should be positive for non-degenerate distributions
}

#[tokio::test]
#[ignore = "rtx-tensor Bool dtype CPU copy not supported"]
async fn test_ppo_update_step() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig {
        n_epochs: 1,
        batch_size: 32,
        minibatch_size: 8,
        ..PPOConfig::default()
    };
    let mut ppo = PPO::new(config, 4, 2, 64, device.clone());

    let batch_size = 32;
    let states = Tensor::randn(&[batch_size, 4], &device).unwrap();
    let actions = Tensor::randn(&[batch_size, 2], &device).unwrap();
    let old_log_probs = Tensor::randn(&[batch_size], &device).unwrap();
    let rewards = Tensor::ones(&[batch_size], &device).unwrap();
    let values = Tensor::zeros(&[batch_size], &device).unwrap();
    let dones = Tensor::zeros(&[batch_size], &device)
        .unwrap()
        .to_dtype(DType::Bool)
        .unwrap();

    let metrics = ppo
        .update(&states, &actions, &old_log_probs, &rewards, &values, &dones)
        .await
        .expect("Update should succeed");

    assert!(metrics.policy_loss.is_finite());
    assert!(metrics.value_loss.is_finite());
    assert!(metrics.entropy.is_finite());
    assert!(metrics.total_loss.is_finite());
    assert!(metrics.kl_divergence >= 0.0);
    assert!(metrics.explained_variance >= -1.0 && metrics.explained_variance <= 1.0);
}

#[tokio::test]
#[ignore = "rtx-rl PPO config validation behavior differs"]
async fn test_ppo_config_validation() {
    let mut config = PPOConfig::default();
    config.clip_epsilon = -0.1; // Invalid negative clipping

    let result = std::panic::catch_unwind(|| PPOConfig::validate(&config));

    assert!(result.is_err());
}

#[tokio::test]
#[ignore = "rtx-tensor Bool dtype CPU copy not supported"]
async fn test_ppo_gradient_clipping() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = PPOConfig {
        max_grad_norm: 0.1, // Very small gradient clipping
        ..PPOConfig::default()
    };
    let mut ppo = PPO::new(config, 4, 2, 64, device.clone());

    let batch_size = 16;
    let states = Tensor::randn(&[batch_size, 4], &device)
        .unwrap()
        .scalar_mul(10.0)
        .unwrap(); // Large states
    let actions = Tensor::randn(&[batch_size, 2], &device)
        .unwrap()
        .scalar_mul(10.0)
        .unwrap();
    let old_log_probs = Tensor::randn(&[batch_size], &device)
        .unwrap()
        .scalar_mul(10.0)
        .unwrap();
    let rewards = Tensor::ones(&[batch_size], &device)
        .unwrap()
        .scalar_mul(100.0)
        .unwrap(); // Large rewards
    let values = Tensor::zeros(&[batch_size], &device).unwrap();
    let dones = Tensor::zeros(&[batch_size], &device)
        .unwrap()
        .to_dtype(DType::Bool)
        .unwrap();

    // Should not panic even with large gradients due to clipping
    let metrics = ppo
        .update(&states, &actions, &old_log_probs, &rewards, &values, &dones)
        .await
        .expect("Update should succeed with gradient clipping");

    assert!(metrics.total_loss.is_finite());
}