foxhunt/crates/ml/tests/ppo_lstm_training_loop_tests.rs

//! Integration tests for PPO LSTM training loop
//!
//! Tests verify that:
//! 1. Training works with LSTM enabled (use_lstm=true)
//! 2. Training works with standard MLP (use_lstm=false) - backward compatibility
//! 3. Hidden state management is properly integrated
//! 4. Networks are correctly initialized based on config

use ml::ppo::{PPOConfig, PPO};
use ml::ppo::trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep};
use ml_core::common::action::{ExposureLevel, FactoredAction, OrderType, Urgency};
use candle_core::Device;

/// Create a small dummy trajectory batch for testing
fn create_dummy_trajectory_batch(num_steps: usize, state_dim: usize) -> TrajectoryBatch {
    let mut trajectory = Trajectory::new();

    for i in 0..num_steps {
        let state = vec![0.1 * i as f32; state_dim];
        let reward = if i % 2 == 0 { 1.0 } else { -0.5 };

        trajectory.add_step(TrajectoryStep {
            state,
            action: FactoredAction::new(ExposureLevel::Flat, OrderType::Market, Urgency::Normal),
            log_prob: -1.5,
            value: 0.5,
            reward,
            done: i == num_steps - 1,
        });
    }

    // Create dummy advantages and returns (same length as num_steps)
    let advantages = vec![0.1; num_steps];
    let returns = vec![0.5; num_steps];

    TrajectoryBatch::from_trajectories(vec![trajectory], advantages, returns)
}

#[test]
fn test_ppo_training_with_lstm_disabled() {
    // Test backward compatibility: standard MLP networks should work
    let config = PPOConfig {
        state_dim: 32,
        num_actions: 45,
        policy_hidden_dims: vec![64, 32],
        value_hidden_dims: vec![64, 32],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        batch_size: 64,
        mini_batch_size: 32,
        num_epochs: 2,
        use_lstm: false, // Standard MLP mode
        lstm_hidden_dim: 128,
        lstm_num_layers: 1,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut ppo = PPO::with_device(config.clone(), device).expect("Failed to create PPO");

    // Verify LSTM is disabled
    assert!(
        ppo.hidden_state_manager.is_none(),
        "Hidden state manager should be None when use_lstm=false"
    );

    // Create dummy trajectory batch
    let mut batch = create_dummy_trajectory_batch(10, config.state_dim);

    // Run single training update
    let result = ppo.update(&mut batch);
    assert!(
        result.is_ok(),
        "Training update failed with LSTM disabled: {:?}",
        result.err()
    );

    let (policy_loss, value_loss) = result.unwrap();
    println!("MLP mode - Policy loss: {}, Value loss: {}", policy_loss, value_loss);

    // Verify losses are reasonable (not NaN or Inf)
    assert!(
        policy_loss.is_finite(),
        "Policy loss should be finite, got: {}",
        policy_loss
    );
    assert!(
        value_loss.is_finite(),
        "Value loss should be finite, got: {}",
        value_loss
    );
}

#[test]
fn test_ppo_training_with_lstm_enabled() {
    // Test LSTM mode: LSTM networks should be used when enabled
    // NOTE: LSTM integration now complete via enum-based architecture
    let config = PPOConfig {
        state_dim: 32,
        num_actions: 45,
        policy_hidden_dims: vec![64, 32],
        value_hidden_dims: vec![64, 32],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        batch_size: 64,
        mini_batch_size: 32,
        num_epochs: 2,
        use_lstm: true, // Enable LSTM
        lstm_hidden_dim: 64,
        lstm_num_layers: 2,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut ppo = PPO::with_device(config.clone(), device).expect("Failed to create PPO");

    // Verify LSTM is enabled
    assert!(
        ppo.hidden_state_manager.is_some(),
        "Hidden state manager should be initialized when use_lstm=true"
    );

    // TODO: Add verification that LSTM networks are actually being used
    // This requires checking network types or tracking LSTM state updates

    // Create dummy trajectory batch
    let mut batch = create_dummy_trajectory_batch(10, config.state_dim);

    // Run single training update
    let result = ppo.update(&mut batch);
    assert!(
        result.is_ok(),
        "Training update failed with LSTM enabled: {:?}",
        result.err()
    );

    let (policy_loss, value_loss) = result.unwrap();
    println!("LSTM mode - Policy loss: {}, Value loss: {}", policy_loss, value_loss);

    // Verify losses are reasonable (not NaN or Inf)
    assert!(
        policy_loss.is_finite(),
        "Policy loss should be finite, got: {}",
        policy_loss
    );
    assert!(
        value_loss.is_finite(),
        "Value loss should be finite, got: {}",
        value_loss
    );
}

#[test]
fn test_lstm_network_initialization() {
    // Test that LSTM networks are correctly initialized based on config
    let lstm_config = PPOConfig {
        state_dim: 32,
        num_actions: 45,
        use_lstm: true,
        lstm_hidden_dim: 128,
        lstm_num_layers: 2,
        ..PPOConfig::default()
    };

    let mlp_config = PPOConfig {
        state_dim: 32,
        num_actions: 45,
        use_lstm: false,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    // Create LSTM-based PPO
    let lstm_ppo = PPO::with_device(lstm_config, device.clone())
        .expect("Failed to create LSTM PPO");
    assert!(
        lstm_ppo.hidden_state_manager.is_some(),
        "LSTM PPO should have hidden state manager"
    );

    // Create MLP-based PPO
    let mlp_ppo = PPO::with_device(mlp_config, device)
        .expect("Failed to create MLP PPO");
    assert!(
        mlp_ppo.hidden_state_manager.is_none(),
        "MLP PPO should NOT have hidden state manager"
    );
}