foxhunt/crates/ml/tests/ppo_lstm_architecture_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,
)]
//! LSTM Architecture Tests for Recurrent PPO
//!
//! This module provides comprehensive tests for LSTM-augmented PPO networks:
//! - Policy network with LSTM layers
//! - Value network with LSTM layers
//! - Hidden state propagation across timesteps
//! - Sequence batching for temporal modeling

#![allow(unused_crate_dependencies)]

// candle eliminated — LSTM networks use native cuda_nn with flat f32 slices
use ml::ppo::lstm_networks::{LSTMPolicyNetwork, LSTMValueNetwork};
use ml::ppo::PPOConfig;
use rand::Rng;
use tracing::info;

/// Generate a random f32 vector from a normal distribution (host-side).
fn randn(len: usize) -> Vec<f32> {
    let mut rng = rand::thread_rng();
    (0..len).map(|_| rng.gen::<f32>() * 2.0 - 1.0).collect()
}

// ============================================================================
// LSTM Policy Network Tests
// ============================================================================

#[test]
fn test_lstm_policy_network_output() {
    // Test that LSTM policy network outputs correct shapes for logits and hidden states
    let input_dim = 64;
    let hidden_dim = 128;
    let num_actions = 45;
    let num_layers = 1;
    let batch_size = 32;

    let network = LSTMPolicyNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        num_actions,
    )
    .expect("Failed to create LSTM policy network");

    // Create flat input state: [batch_size * input_dim]
    let state = randn(batch_size * input_dim);

    // Initial hidden and cell states (zeros): [num_layers * batch_size * hidden_dim]
    let hc_size = num_layers * batch_size * hidden_dim;
    let h0 = vec![0.0f32; hc_size];
    let c0 = vec![0.0f32; hc_size];

    // Forward pass — returns (logits_host, new_h_flat, new_c_flat)
    let (logits, h_t, c_t) = network
        .forward(&state, &h0, &c0, batch_size)
        .expect("Forward pass failed");

    // Verify output sizes (flat vectors)
    assert_eq!(logits.len(), batch_size * num_actions, "Logits size mismatch");
    assert_eq!(h_t.len(), num_layers * batch_size * hidden_dim, "Hidden state size mismatch");
    assert_eq!(c_t.len(), num_layers * batch_size * hidden_dim, "Cell state size mismatch");

    info!("LSTM policy network output test passed");
}

#[test]
fn test_lstm_value_network_output() {
    // Test that LSTM value network outputs correct shapes for values and hidden states
    let input_dim = 64;
    let hidden_dim = 256;
    let num_layers = 2;
    let batch_size = 16;

    let network = LSTMValueNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
    )
    .expect("Failed to create LSTM value network");

    // Create flat input state: [batch_size * input_dim]
    let state = randn(batch_size * input_dim);

    // Initial hidden and cell states (zeros): [num_layers * batch_size * hidden_dim]
    let hc_size = num_layers * batch_size * hidden_dim;
    let h0 = vec![0.0f32; hc_size];
    let c0 = vec![0.0f32; hc_size];

    // Forward pass — returns (values_host, new_h_flat, new_c_flat)
    let (value, h_t, c_t) = network
        .forward(&state, &h0, &c0, batch_size)
        .expect("Forward pass failed");

    // Verify output sizes (flat vectors)
    assert_eq!(value.len(), batch_size, "Value size mismatch");
    assert_eq!(h_t.len(), num_layers * batch_size * hidden_dim, "Hidden state size mismatch");
    assert_eq!(c_t.len(), num_layers * batch_size * hidden_dim, "Cell state size mismatch");

    info!("LSTM value network output test passed");
}

// ============================================================================
// Hidden State Propagation Tests
// ============================================================================

#[test]
fn test_lstm_hidden_state_propagation() {
    // Test that hidden states correctly propagate across multiple timesteps
    let input_dim = 32;
    let hidden_dim = 64;
    let num_layers = 1;
    let num_actions = 3;
    let batch_size = 8;
    let seq_len = 5;

    let network = LSTMPolicyNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        num_actions,
    )
    .expect("Failed to create LSTM policy network");

    // Initialize hidden states (zeros)
    let hc_size = num_layers * batch_size * hidden_dim;
    let mut h_t = vec![0.0f32; hc_size];
    let mut c_t = vec![0.0f32; hc_size];

    // Process sequence of states
    for t in 0..seq_len {
        let state = randn(batch_size * input_dim);

        let (logits, new_h, new_c) = network
            .forward(&state, &h_t, &c_t, batch_size)
            .expect(&format!("Forward pass failed at timestep {}", t));

        // Verify sizes remain consistent
        assert_eq!(logits.len(), batch_size * num_actions);
        assert_eq!(new_h.len(), num_layers * batch_size * hidden_dim);
        assert_eq!(new_c.len(), num_layers * batch_size * hidden_dim);

        // Hidden states should NOT be all zeros after first timestep (indicating information flow)
        if t > 0 {
            let h_sum: f32 = new_h.iter().sum();
            assert!(h_sum.abs() > 1e-6, "Hidden state is all zeros at timestep {}", t);
        }

        // Update hidden states for next timestep
        h_t = new_h;
        c_t = new_c;
    }

    info!("LSTM hidden state propagation test passed");
}

// ============================================================================
// Sequence Batching Tests
// ============================================================================

#[test]
fn test_lstm_sequence_batching() {
    // Test that LSTM networks can process sequences in batch format
    let input_dim = 16;
    let hidden_dim = 32;
    let num_layers = 1;
    let num_actions = 5;
    let seq_len = 10;
    let batch_size = 4;

    let network = LSTMPolicyNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        num_actions,
    )
    .expect("Failed to create LSTM policy network");

    // Initialize hidden states (zeros)
    let hc_size = num_layers * batch_size * hidden_dim;
    let mut h_t = vec![0.0f32; hc_size];
    let mut c_t = vec![0.0f32; hc_size];

    let mut all_logits = Vec::new();

    // Process sequence one timestep at a time
    for t in 0..seq_len {
        // Each timestep: flat [batch_size * input_dim]
        let state = randn(batch_size * input_dim);

        let (logits, new_h, new_c) = network
            .forward(&state, &h_t, &c_t, batch_size)
            .expect(&format!("Forward pass failed at timestep {}", t));

        all_logits.push(logits);
        h_t = new_h;
        c_t = new_c;
    }

    // Verify we processed all timesteps
    assert_eq!(all_logits.len(), seq_len);

    // Verify each timestep has correct batch dimension
    for (t, logits) in all_logits.iter().enumerate() {
        assert_eq!(logits.len(), batch_size * num_actions, "Logits size mismatch at timestep {}", t);
    }

    info!("LSTM sequence batching test passed");
}

// ============================================================================
// Configuration Tests
// ============================================================================

#[test]
fn test_ppo_config_lstm_fields() {
    // Test that PPOConfig includes LSTM configuration fields
    let config = PPOConfig {
        use_lstm: true,
        lstm_hidden_dim: 256,
        lstm_num_layers: 2,
        ..PPOConfig::default()
    };

    assert_eq!(config.use_lstm, true, "use_lstm field not set correctly");
    assert_eq!(config.lstm_hidden_dim, 256, "lstm_hidden_dim field not set correctly");
    assert_eq!(config.lstm_num_layers, 2, "lstm_num_layers field not set correctly");

    info!("PPOConfig LSTM fields test passed");
}

#[test]
fn test_ppo_config_backward_compatible() {
    // Test that default PPOConfig has LSTM disabled for backward compatibility
    let config = PPOConfig::default();

    assert_eq!(config.use_lstm, false, "LSTM should be disabled by default");
    assert!(config.lstm_hidden_dim > 0, "LSTM hidden dim should have valid default");
    assert!(config.lstm_num_layers > 0, "LSTM num layers should have valid default");

    info!("PPOConfig backward compatibility test passed");
}