foxhunt/crates/ml/tests/dqn_gradient_collapse_root_cause_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,
)]
// CRITICAL TEST: Isolate DQN gradient collapse root cause
// Purpose: Replicate production failure with dtype conversion hypothesis
// Created: 2025-11-21 - Test-Driven Development Campaign
// Updated: 2026-03-04 - Account for BF16 mixed precision on CUDA (Ampere+)
// Updated: 2026-03-19 - Migrated from Candle to GpuTensor native API
//
// Based on Zen analysis: Line 1087 dtype conversion likely breaks gradient flow

use ml_core::cuda_autograd::GpuTensor;
use ml_core::device::MlDevice;
use ml_core::native_types::NativeDType;
use ml::MLError;
use std::sync::Arc;

/// Test: Dtype conversion preserves values (CRITICAL for loss computation)
///
/// Production code converts BF16 network output to F32 for loss computation.
/// This test verifies the conversion preserves values within acceptable tolerance.
/// (BF16 has ~3 decimal digits of precision)
///
/// Note: GpuTensor operates in F32 on GPU; dtype is metadata-only. The round-trip
/// test validates that to_dtype (a no-op clone for F32) does not corrupt data.
#[test]
fn test_dtype_conversion_preserves_values() -> Result<(), MLError> {
    let device = MlDevice::new_cuda(0).expect("CUDA required");
    let stream = device.cuda_stream()?.clone();

    // Create F32 tensor with known random values
    let original = GpuTensor::randn(&[32, 51], 1.0, &stream)?;

    // Round-trip: F32 -> "BF16" -> F32 (GpuTensor is always F32; to_dtype is a clone)
    let as_bf16 = original.to_dtype(NativeDType::BF16, &stream)?;
    let back_to_f32 = as_bf16.to_dtype(NativeDType::F32, &stream)?;

    // Compute max absolute element-wise difference on host
    let orig_host = original.to_host(&stream)?;
    let rt_host = back_to_f32.to_host(&stream)?;

    let max_elem_diff: f32 = orig_host
        .iter()
        .zip(rt_host.iter())
        .map(|(a, b)| (a - b).abs())
        .fold(0.0_f32, f32::max);

    // GpuTensor to_dtype is a clone (always F32), so diff should be exactly 0.
    // Allow a tiny tolerance for floating-point DtoH round-trip edge cases.
    assert!(
        max_elem_diff < 0.02,
        "Round-trip max element-wise diff = {} exceeds 0.02",
        max_elem_diff
    );

    // Verify mean is preserved
    let orig_mean = original.mean_all(&stream)?;
    let rt_mean = back_to_f32.mean_all(&stream)?;
    let mean_diff = (orig_mean - rt_mean).abs();
    assert!(
        mean_diff < 0.01,
        "Round-trip changed mean: {} -> {} (diff={})",
        orig_mean, rt_mean, mean_diff
    );

    Ok(())
}

/// Test: Network output dtype on CUDA uses BF16 mixed precision
/// On Ampere+ GPUs, networks output BF16 (not F32) -- conversion needed in loss path
///
/// Note: With GpuTensor, all data is F32 on GPU. This test validates that
/// DistributionalDuelingQNetwork produces valid finite output.
#[test]
fn test_network_output_dtype() -> Result<(), MLError> {
    use ml::dqn::{DistributionalDuelingConfig, DistributionalDuelingQNetwork};

    let device = MlDevice::new_cuda(0).expect("CUDA required");
    let stream = device.cuda_stream()?.clone();

    let config = DistributionalDuelingConfig {
        state_dim: 54,
        num_actions: 45,
        num_atoms: 51,
        shared_hidden_dims: vec![128, 64],
        value_hidden_dim: 64,
        advantage_hidden_dim: 64,
        leaky_relu_alpha: 0.01,
    };

    let network = DistributionalDuelingQNetwork::new(config, stream.clone())?;

    // Forward pass: create GpuTensor input, run forward, read back to host
    let input_host: Vec<f32> = (0..32 * 54).map(|i| (i as f32 * 0.01).sin()).collect();
    let input_gpu = GpuTensor::from_host(&input_host, vec![32, 54], &stream)?;
    let output = network.forward(&input_gpu)?;

    // Read output to host for validation
    let output_host = output.to_host(&stream)?;

    // GpuTensor is always F32 -- verify output is finite
    assert!(
        output_host.iter().all(|v| v.is_finite()),
        "Network output contains NaN/Inf"
    );
    assert!(
        !output_host.is_empty(),
        "Network output is empty"
    );

    Ok(())
}

/// Test: Gather operation preserves dtype
/// Goal: Verify indexing does not corrupt values
///
/// Note: With GpuTensor, all data is F32. This test validates that
/// host-side gather logic works correctly.
#[test]
fn test_gather_preserves_dtype() -> Result<(), MLError> {
    let device = MlDevice::new_cuda(0).expect("CUDA required");
    let stream = device.cuda_stream()?.clone();

    // Create a [32, 45, 51] tensor on GPU
    let n = 32 * 45 * 51;
    let host_data: Vec<f32> = (0..n).map(|i| (i as f32 * 0.001).sin()).collect();
    let tensor = GpuTensor::from_host(&host_data, vec![32, 45, 51], &stream)?;

    // Download to host and manually gather action=0 for each batch element
    let flat = tensor.to_host(&stream)?;

    // Gather: for each batch element b, pick action 0, get 51 atoms
    // Index: flat[b * 45 * 51 + 0 * 51 .. b * 45 * 51 + 1 * 51]
    let mut gathered = Vec::with_capacity(32 * 51);
    for b in 0..32 {
        let start = b * 45 * 51; // action 0
        for atom in 0..51 {
            gathered.push(flat[start + atom]);
        }
    }

    // Verify gathered values match expected
    for b in 0..32 {
        for atom in 0..51 {
            let expected = host_data[b * 45 * 51 + atom];
            let actual = gathered[b * 51 + atom];
            assert_eq!(
                expected, actual,
                "Gather mismatch at batch={b}, atom={atom}: expected {expected}, got {actual}"
            );
        }
    }

    Ok(())
}