foxhunt/crates/ml/tests/ppo_hidden_state_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,
)]
//! Tests for PPO Hidden State Management
//!
//! Validates LSTM hidden state initialization, propagation, and reset logic.
//! All state storage is host-side Vec<f32> (no Candle tensors).

// candle eliminated -- HiddenStateManager uses host-side Vec<f32>
use ml::ppo::hidden_state_manager::HiddenStateManager;

#[test]
fn test_hidden_state_initialization() -> Result<(), Box<dyn std::error::Error>> {
    // Test that hidden states are initialized to zeros
    let num_layers = 2;
    let batch_size = 4;
    let hidden_dim = 128;

    let manager = HiddenStateManager::with_defaults(num_layers, batch_size, hidden_dim)?;

    let (policy_h, policy_c) = manager.get_policy_state();
    let (value_h, value_c) = manager.get_value_state();

    // Verify lengths (flat: num_layers * batch_size * hidden_dim)
    let expected_len = num_layers * batch_size * hidden_dim;
    assert_eq!(policy_h.len(), expected_len);
    assert_eq!(policy_c.len(), expected_len);
    assert_eq!(value_h.len(), expected_len);
    assert_eq!(value_c.len(), expected_len);

    // Verify all zeros
    let policy_h_sum: f32 = policy_h.iter().sum();
    let policy_c_sum: f32 = policy_c.iter().sum();
    let value_h_sum: f32 = value_h.iter().sum();
    let value_c_sum: f32 = value_c.iter().sum();

    assert_eq!(policy_h_sum, 0.0, "Policy hidden state should be zeros");
    assert_eq!(policy_c_sum, 0.0, "Policy cell state should be zeros");
    assert_eq!(value_h_sum, 0.0, "Value hidden state should be zeros");
    assert_eq!(value_c_sum, 0.0, "Value cell state should be zeros");

    Ok(())
}

#[test]
fn test_hidden_state_propagation() -> Result<(), Box<dyn std::error::Error>> {
    // Test that hidden states carry across timesteps
    let num_layers = 2;
    let batch_size = 4;
    let hidden_dim = 128;

    let mut manager = HiddenStateManager::with_defaults(num_layers, batch_size, hidden_dim)?;

    let total = num_layers * batch_size * hidden_dim;

    // Create new states with non-zero values
    let new_h = vec![1.0_f32; total];
    let new_c = vec![2.0_f32; total]; // c = 2.0 (was affine(2.0, 0.0) on ones)

    // Update policy states
    manager.update_policy_state(new_h.clone(), new_c.clone())?;

    // Verify states were updated
    let (policy_h, policy_c) = manager.get_policy_state();
    let policy_h_sum: f32 = policy_h.iter().sum();
    let policy_c_sum: f32 = policy_c.iter().sum();

    let expected_h_sum = total as f32;
    let expected_c_sum = total as f32 * 2.0;

    assert!(
        (policy_h_sum - expected_h_sum).abs() < 0.01,
        "Policy hidden state should be updated. Expected {}, got {}",
        expected_h_sum,
        policy_h_sum
    );
    assert!(
        (policy_c_sum - expected_c_sum).abs() < 0.01,
        "Policy cell state should be updated. Expected {}, got {}",
        expected_c_sum,
        policy_c_sum
    );

    // Update value states
    manager.update_value_state(new_h, new_c)?;

    // Verify value states were updated
    let (value_h, value_c) = manager.get_value_state();
    let value_h_sum: f32 = value_h.iter().sum();
    let value_c_sum: f32 = value_c.iter().sum();

    assert!(
        (value_h_sum - expected_h_sum).abs() < 0.01,
        "Value hidden state should be updated. Expected {}, got {}",
        expected_h_sum,
        value_h_sum
    );
    assert!(
        (value_c_sum - expected_c_sum).abs() < 0.01,
        "Value cell state should be updated. Expected {}, got {}",
        expected_c_sum,
        value_c_sum
    );

    Ok(())
}

#[test]
fn test_hidden_state_reset_on_done() -> Result<(), Box<dyn std::error::Error>> {
    // Test that states reset when episodes end
    let num_layers = 2;
    let batch_size = 4;
    let hidden_dim = 128;

    let mut manager = HiddenStateManager::with_defaults(num_layers, batch_size, hidden_dim)?;

    let total = num_layers * batch_size * hidden_dim;

    // Set states to non-zero (all ones)
    let ones = vec![1.0_f32; total];
    manager.update_policy_state(ones.clone(), ones.clone())?;
    manager.update_value_state(ones.clone(), ones.clone())?;

    // Create done mask: environments 0 and 2 are done (bool, not u8 tensor)
    let done_mask = vec![true, false, true, false];

    // Reset states for done environments
    manager.reset_on_done(&done_mask)?;

    // Verify policy states
    let (policy_h, policy_c) = manager.get_policy_state();

    // Check that environments 0 and 2 are reset (all zeros)
    // and environments 1 and 3 retain their values (all ones)
    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                let expected = if env == 0 || env == 2 { 0.0 } else { 1.0 };
                assert!(
                    (policy_h[idx] - expected).abs() < 0.01,
                    "Policy hidden state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    policy_h[idx]
                );
                assert!(
                    (policy_c[idx] - expected).abs() < 0.01,
                    "Policy cell state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    policy_c[idx]
                );
            }
        }
    }

    // Verify value states similarly
    let (value_h, value_c) = manager.get_value_state();

    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                let expected = if env == 0 || env == 2 { 0.0 } else { 1.0 };
                assert!(
                    (value_h[idx] - expected).abs() < 0.01,
                    "Value hidden state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    value_h[idx]
                );
                assert!(
                    (value_c[idx] - expected).abs() < 0.01,
                    "Value cell state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    value_c[idx]
                );
            }
        }
    }

    Ok(())
}

#[test]
fn test_hidden_state_batching() -> Result<(), Box<dyn std::error::Error>> {
    // Test that manager handles multiple parallel environments correctly
    let num_layers = 2;
    let batch_size = 8; // 8 parallel environments
    let hidden_dim = 64;

    let mut manager = HiddenStateManager::with_defaults(num_layers, batch_size, hidden_dim)?;

    // Create different values for different environments
    let total = num_layers * batch_size * hidden_dim;
    let mut h_data = vec![0.0_f32; total];
    let mut c_data = vec![0.0_f32; total];

    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                h_data[idx] = (env + 1) as f32; // Environment 0 = 1.0, env 1 = 2.0, etc.
                c_data[idx] = (env + 1) as f32 * 10.0; // Environment 0 = 10.0, env 1 = 20.0, etc.
            }
        }
    }

    // Update states (Vec<f32> directly, no Tensor conversion needed)
    manager.update_policy_state(h_data.clone(), c_data.clone())?;

    // Verify each environment has correct values
    let (policy_h, policy_c) = manager.get_policy_state();

    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                let expected_h = (env + 1) as f32;
                let expected_c = (env + 1) as f32 * 10.0;

                assert!(
                    (policy_h[idx] - expected_h).abs() < 0.01,
                    "Hidden state for env {} should be {}. Got {}",
                    env,
                    expected_h,
                    policy_h[idx]
                );
                assert!(
                    (policy_c[idx] - expected_c).abs() < 0.01,
                    "Cell state for env {} should be {}. Got {}",
                    env,
                    expected_c,
                    policy_c[idx]
                );
            }
        }
    }

    Ok(())
}