foxhunt/ml/tests/ppo_checkpoint_roundtrip_test.rs

//! PPO Checkpoint Round-Trip Validation Test
//!
//! Proves that saving a PPO model checkpoint and loading it back
//! produces identical action probabilities for the same input state.
//!
//! Run manually:
//! ```sh
//! SQLX_OFFLINE=true cargo test -p ml --test ppo_checkpoint_roundtrip_test -- --ignored --nocapture
//! ```

#![allow(unused_crate_dependencies)]

use anyhow::Result;
use candle_core::Device;
use ml::ppo::ppo::{PPOConfig, PPO};

#[tokio::test]
#[ignore]
async fn test_ppo_checkpoint_roundtrip() -> Result<()> {
    let checkpoint_dir = tempfile::tempdir()?;
    let state_dim = 54;
    let num_actions = 45;

    // Create PPO model with known config
    let config = PPOConfig {
        state_dim,
        num_actions,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        use_lstm: false,
        ..PPOConfig::default()
    };

    let ppo = PPO::new(config.clone())?;

    // Create test state
    let test_state: Vec<f32> = (0..state_dim).map(|i| (i as f32 * 0.1).sin()).collect();

    // Get predictions before save
    let predictions_before = ppo.predict(&test_state)?;

    // Save checkpoint (API requires three &str paths: actor, critic, metadata)
    let actor_path = checkpoint_dir.path().join("actor.safetensors");
    let critic_path = checkpoint_dir.path().join("critic.safetensors");
    let metadata_path = checkpoint_dir.path().join("metadata.json");

    let actor_str = actor_path.to_str().ok_or_else(|| anyhow::anyhow!("Invalid actor path"))?;
    let critic_str = critic_path.to_str().ok_or_else(|| anyhow::anyhow!("Invalid critic path"))?;
    let metadata_str =
        metadata_path.to_str().ok_or_else(|| anyhow::anyhow!("Invalid metadata path"))?;

    ppo.save_checkpoint(actor_str, critic_str, metadata_str)?;

    // Verify files exist and are non-empty
    assert!(actor_path.exists(), "Actor checkpoint not saved");
    assert!(critic_path.exists(), "Critic checkpoint not saved");
    assert!(metadata_path.exists(), "Metadata file not saved");
    assert!(std::fs::metadata(&actor_path)?.len() > 0);
    assert!(std::fs::metadata(&critic_path)?.len() > 0);

    // Load into a fresh model (API takes &str, &str, PPOConfig, Device)
    let loaded_ppo = PPO::load_checkpoint(actor_str, critic_str, config, Device::Cpu)?;

    // Get predictions after load
    let predictions_after = loaded_ppo.predict(&test_state)?;

    // Compare predictions - should be identical (not just close, IDENTICAL)
    assert_eq!(
        predictions_before.len(),
        predictions_after.len(),
        "Prediction vector length mismatch"
    );

    for (i, (before, after)) in predictions_before
        .iter()
        .zip(predictions_after.iter())
        .enumerate()
    {
        let diff = (before - after).abs();
        assert!(
            diff < 1e-6,
            "Prediction mismatch at index {}: before={}, after={}, diff={}",
            i,
            before,
            after,
            diff
        );
    }

    println!("Checkpoint round-trip validation passed!");
    println!(
        "  Actor checkpoint: {} bytes",
        std::fs::metadata(&actor_path)?.len()
    );
    println!(
        "  Critic checkpoint: {} bytes",
        std::fs::metadata(&critic_path)?.len()
    );
    println!(
        "  Predictions compared: {} values",
        predictions_before.len()
    );
    println!(
        "  Max difference: {:.2e}",
        predictions_before
            .iter()
            .zip(predictions_after.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max)
    );

    Ok(())
}