foxhunt/crates/ml/tests/rainbow_loss_shape_test.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,
)]
//! Test for Rainbow DQN Loss Computation Shape Mismatch
//!
//! This test reproduces the shape mismatch bug in compute_rainbow_loss:
//! `shape mismatch in mul, lhs: [32, 1], rhs: [32]`
//!
//! Migrated from Candle Tensor to GpuTensor (native cudarc).
#![allow(unused_crate_dependencies)]

// candle eliminated — test uses native GpuTensor APIs
use std::sync::OnceLock;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::device::MlDevice;
use ml_core::MLError;

static SHARED_CUDA: OnceLock<MlDevice> = OnceLock::new();

fn cuda_device() -> MlDevice {
    SHARED_CUDA.get_or_init(|| MlDevice::cuda(0).expect("CUDA required")).clone()
}

/// Test: Reproduce shape mismatch in target Q-value computation
///
/// This test simulates the exact tensor operations in `compute_rainbow_loss`
/// that cause the shape mismatch between target_q [32, 1] and gamma_tensor [32]
#[test]
fn test_target_q_value_shape_mismatch() {
    let device = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");
    let batch_size = 32;

    // Simulate next_q_values shape [32, 4, 1] (batch, actions, 1) from get_q_values
    let next_q_values = GpuTensor::randn(&[batch_size, 4, 1], 1.0, stream).unwrap();

    // Simulate next_actions [32] from argmax — all zeros (action index 0)
    let next_actions: Vec<u32> = vec![0_u32; batch_size];

    // Use index_select on dim=1 to pick Q-values for selected actions
    // index_select on dim=1 with single index gives [32, 1, 1]
    let gathered = next_q_values.index_select(1, &next_actions, stream).unwrap();

    // squeeze(1) produces [32, 1] - THIS IS THE BUG
    let target_q = gathered.squeeze(1, stream).unwrap();

    // Verify shape is [32, 1] (this is the problematic shape)
    assert_eq!(target_q.shape(), &[32, 1]);

    // Create gamma_tensor [32]
    let gamma_tensor = GpuTensor::full(&[batch_size], 0.99, stream).unwrap();

    // Verify shape is [32]
    assert_eq!(gamma_tensor.shape(), &[32]);

    // This multiplication SHOULD FAIL with shape mismatch [32, 1] vs [32]
    let result = target_q.mul(&gamma_tensor, stream);

    // The test should fail here showing the shape mismatch
    match result {
        Ok(_) => panic!("Expected shape mismatch error but operation succeeded!"),
        Err(e) => {
            let error_msg = format!("{:?}", e);
            assert!(
                error_msg.contains("mismatch") || error_msg.contains("Mismatch") || error_msg.contains("incompatible"),
                "Expected shape mismatch error, got: {}",
                error_msg
            );
        },
    }
}

/// Test: Correct shape handling with squeeze
///
/// This test shows the FIX - we need to squeeze both dimensions after gather
#[test]
fn test_target_q_value_shape_fix() {
    let device = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");
    let batch_size = 32;

    // Simulate next_q_values shape [32, 4, 1] (batch, actions, 1) from get_q_values
    let next_q_values = GpuTensor::randn(&[batch_size, 4, 1], 1.0, stream).unwrap();

    // Simulate next_actions [32] from argmax — all zeros
    let next_actions: Vec<u32> = vec![0_u32; batch_size];

    // index_select on dim=1 gives [32, 1, 1]
    let gathered = next_q_values.index_select(1, &next_actions, stream).unwrap();

    // FIX: squeeze BOTH dimensions to get [32]
    let target_q = gathered
        .squeeze(1, stream).unwrap()
        .squeeze(1, stream).unwrap();

    // Verify shape is [32] (fixed!)
    assert_eq!(target_q.shape(), &[32]);

    // Create gamma_tensor [32]
    let gamma_tensor = GpuTensor::full(&[batch_size], 0.99, stream).unwrap();

    // Verify shape is [32]
    assert_eq!(gamma_tensor.shape(), &[32]);

    // This multiplication should now work!
    let result = target_q.mul(&gamma_tensor, stream);
    assert!(
        result.is_ok(),
        "Multiplication should succeed with matching shapes"
    );

    let product = result.unwrap();
    assert_eq!(product.shape(), &[32]);
}

/// Test: Current action Q-values shape handling
///
/// Verifies the same issue exists for current_action_q computation
#[test]
fn test_current_action_q_shape_mismatch() {
    let device = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");
    let batch_size = 32;

    // Simulate current_q_values shape [32, 4, 1] from get_q_values
    let current_q_values = GpuTensor::randn(&[batch_size, 4, 1], 1.0, stream).unwrap();

    // Simulate actions [32] — cycling through 0..3
    let actions: Vec<u32> = (0..batch_size).map(|i| (i % 4) as u32).collect();

    // index_select on dim=1 gives [32, 1, 1]
    let gathered = current_q_values.index_select(1, &actions, stream).unwrap();

    // squeeze(1) produces [32, 1] - same bug
    let current_action_q = gathered.squeeze(1, stream).unwrap();

    // Verify shape is [32, 1]
    assert_eq!(current_action_q.shape(), &[32, 1]);

    // Create target_values [32]
    let target_values = GpuTensor::randn(&[batch_size], 1.0, stream).unwrap();

    // Verify shape is [32]
    assert_eq!(target_values.shape(), &[32]);

    // Subtraction should fail with shape mismatch
    let result = current_action_q.sub(&target_values, stream);

    match result {
        Ok(_) => panic!("Expected shape mismatch error but operation succeeded!"),
        Err(e) => {
            let error_msg = format!("{:?}", e);
            assert!(
                error_msg.contains("mismatch") || error_msg.contains("Mismatch") || error_msg.contains("incompatible"),
                "Expected shape mismatch error, got: {}",
                error_msg
            );
        },
    }
}

/// Test: Current action Q-values shape fix
///
/// Verifies the fix works for current_action_q computation
#[test]
fn test_current_action_q_shape_fix() {
    let device = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");
    let batch_size = 32;

    // Simulate current_q_values shape [32, 4, 1] from get_q_values
    let current_q_values = GpuTensor::randn(&[batch_size, 4, 1], 1.0, stream).unwrap();

    // Simulate actions [32] — cycling through 0..3
    let actions: Vec<u32> = (0..batch_size).map(|i| (i % 4) as u32).collect();

    // index_select on dim=1 gives [32, 1, 1]
    let gathered = current_q_values.index_select(1, &actions, stream).unwrap();

    // FIX: squeeze BOTH dimensions to get [32]
    let current_action_q = gathered
        .squeeze(1, stream).unwrap()
        .squeeze(1, stream).unwrap();

    // Verify shape is [32]
    assert_eq!(current_action_q.shape(), &[32]);

    // Create target_values [32]
    let target_values = GpuTensor::randn(&[batch_size], 1.0, stream).unwrap();

    // Verify shape is [32]
    assert_eq!(target_values.shape(), &[32]);

    // Subtraction should now work!
    let result = current_action_q.sub(&target_values, stream);
    assert!(
        result.is_ok(),
        "Subtraction should succeed with matching shapes"
    );

    let diff = result.unwrap();
    assert_eq!(diff.shape(), &[32]);
}