foxhunt/crates/ml/tests/ppo_lstm_training_loop_tests.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! 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;
use tracing::info;

/// 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::new_cuda(0).expect("CUDA required");
    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();
    info!(policy_loss, value_loss, "MLP mode");

    // 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::new_cuda(0).expect("CUDA required");
    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();
    info!(policy_loss, value_loss, "LSTM mode");

    // 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::new_cuda(0).expect("CUDA required");

    // 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"
    );
}