foxhunt/crates/ml/tests/dqn_logit_clipping_bug12_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,
)]
//! Bug #12: Clip logits before softmax to prevent saturation
//!
//! This test suite validates that logit clipping prevents:
//! - Numerical overflow: exp(large_value) -> inf
//! - Probability saturation: exp(-large_value) -> 0.0
//! - Gradient vanishing: d/dx softmax(saturated) ~ 0.0
//!
//! Expected behavior:
//! - Logits clipped to [-10, 10] before softmax
//! - All softmax probabilities > 0.0001 (no saturation)
//! - Gradients remain non-zero after backward pass
//!
//! Post-migration: uses `ml::dqn::clip_logits` and `ml::dqn::softmax_with_clipping`
//! which operate on `&[f32]` slices (CPU-side). The GPU tensor paths are tested
//! separately via the DQN trainer integration tests.

use anyhow::Result;
use ml::dqn::{clip_logits, softmax_with_clipping};

/// Compute softmax on a 1-D slice.
fn softmax_1d(logits: &[f32]) -> Vec<f32> {
    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = logits.iter().map(|&x| (x - max).exp()).collect();
    let sum: f32 = exps.iter().sum();
    exps.iter().map(|&e| e / sum).collect()
}

/// Compute softmax for a 2-D batch (row-wise) stored in a flat slice.
fn softmax_2d(logits: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut result = Vec::with_capacity(rows * cols);
    for r in 0..rows {
        let row = &logits[r * cols..(r + 1) * cols];
        result.extend(softmax_1d(row));
    }
    result
}

#[test]
fn test_logits_clipped_to_range() -> Result<()> {
    // Create logits spanning full range
    let logits = vec![-44.0f32, -10.0, 0.0, 10.0, 44.0];

    // Apply clipping
    let clipped = clip_logits(&logits, -10.0, 10.0);

    // Verify all values in [-10, 10]
    for (i, &val) in clipped.iter().enumerate() {
        assert!(
            val >= -10.0 && val <= 10.0,
            "Value {} at index {} outside [-10, 10]",
            val,
            i
        );
    }

    Ok(())
}

#[test]
fn test_extreme_positive_logits_clipped() -> Result<()> {
    // Create extreme logits: [44.0, -44.0, 0.0]
    let logits = vec![44.0f32, -44.0f32, 0.0f32];

    // Apply clipping
    let clipped = clip_logits(&logits, -10.0, 10.0);

    // Verify clipping
    assert!(
        (clipped[0] - 10.0).abs() < 1e-6,
        "44.0 should clip to 10.0, got {}",
        clipped[0]
    );
    assert!(
        (clipped[1] + 10.0).abs() < 1e-6,
        "-44.0 should clip to -10.0, got {}",
        clipped[1]
    );
    assert!(
        (clipped[2] - 0.0).abs() < 1e-6,
        "0.0 should remain unchanged, got {}",
        clipped[2]
    );

    Ok(())
}

#[test]
fn test_extreme_negative_logits_clipped() -> Result<()> {
    // Create extreme negative logits
    let logits = vec![-100.0f32, -44.0f32, -10.0f32];

    // Apply clipping
    let clipped = clip_logits(&logits, -10.0, 10.0);

    // All should be clipped to -10.0
    for (i, &val) in clipped.iter().enumerate() {
        assert!(
            (val + 10.0).abs() < 1e-6,
            "Value {} at index {} should be clipped to -10.0, got {}",
            logits[i],
            i,
            val
        );
    }

    Ok(())
}

#[test]
fn test_softmax_no_saturation() -> Result<()> {
    // Extreme logits that would saturate without clipping
    let logits = vec![50.0f32, -50.0f32, 0.0f32];

    // Clip then softmax
    let clipped = clip_logits(&logits, -10.0, 10.0);
    let probs = softmax_1d(&clipped);

    // Verify no complete saturation (all probs > 0.0)
    // After clipping [50, -50, 0] -> [10, -10, 0], softmax gives:
    // - exp(10)/(exp(10) + exp(-10) + exp(0)) ~ 0.9999546
    // - exp(-10)/(exp(10) + exp(-10) + exp(0)) ~ 2.06e-9  (very small but NON-ZERO)
    // - exp(0)/(exp(10) + exp(-10) + exp(0)) ~ 4.54e-5
    // This is MUCH better than exp(-50) ~ 0.0 (complete saturation)
    for (i, &prob) in probs.iter().enumerate() {
        assert!(
            prob > 1e-9,
            "Probability {} should be > 1e-9 (got {}), complete saturation would be ~0.0",
            i,
            prob
        );
    }

    // Verify sum to 1.0
    let sum: f32 = probs.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Probabilities should sum to 1.0, got {}",
        sum
    );

    Ok(())
}

#[test]
fn test_gradients_nonzero_after_clipping() -> Result<()> {
    // Create extreme logits
    let logits = vec![40.0f32, -40.0f32, 0.0f32];

    // Clip and apply softmax
    let clipped = clip_logits(&logits, -10.0, 10.0);
    let probs = softmax_1d(&clipped);

    // Compute a simple loss (cross-entropy with target [1, 0, 0])
    let target = vec![1.0f32, 0.0f32, 0.0f32];
    let loss: f32 = target
        .iter()
        .zip(probs.iter())
        .map(|(&t, &p)| -t * p.ln())
        .sum();

    // Verify loss is finite (not NaN or Inf)
    assert!(
        loss.is_finite(),
        "Loss should be finite, got {}",
        loss
    );

    // Verify loss is positive
    assert!(loss > 0.0, "Loss should be positive, got {}", loss);

    Ok(())
}

#[test]
fn test_action_diversity_maintained() -> Result<()> {
    // Simulate Q-values from multiple states
    let batch_size = 5;
    let num_actions = 3;

    // Create extreme logits for batch
    let logits_vec: Vec<f32> = (0..batch_size)
        .flat_map(|_| vec![50.0f32, -50.0f32, 0.0f32])
        .collect();

    // Clip batch
    let clipped: Vec<f32> = logits_vec
        .chunks(num_actions)
        .flat_map(|chunk| clip_logits(chunk, -10.0, 10.0))
        .collect();

    // Softmax per row
    let probs = softmax_2d(&clipped, batch_size, num_actions);

    // Argmax per row -> actions
    let actions: Vec<usize> = (0..batch_size)
        .map(|b| {
            let row = &probs[b * num_actions..(b + 1) * num_actions];
            row.iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .map(|(i, _)| i)
                .unwrap()
        })
        .collect();

    // Verify all actions are in valid range
    for (i, &action) in actions.iter().enumerate() {
        assert!(
            action < num_actions,
            "Action {} at index {} exceeds num_actions {}",
            action,
            i,
            num_actions
        );
    }

    // After clipping [50, -50, 0] -> [10, -10, 0], softmax ~ [0.9999, 0.0, 0.0]
    // So all actions will be 0 (highest Q-value)
    // This is EXPECTED behavior - diversity is preserved when Q-values are similar

    // Let's test with more balanced Q-values
    let balanced_logits = vec![5.0f32, 3.0f32, 4.0f32];
    let clipped_balanced = clip_logits(&balanced_logits, -10.0, 10.0);
    let probs_balanced = softmax_1d(&clipped_balanced);

    // All probabilities should be > 0.01 (no single action dominates completely)
    let min_prob = probs_balanced.iter().copied().fold(f32::INFINITY, f32::min);
    assert!(
        min_prob > 0.01,
        "Minimum probability should be > 0.01 for balanced Q-values, got {}",
        min_prob
    );

    Ok(())
}

#[test]
fn test_batch_logit_clipping() -> Result<()> {
    // Batch of logits: [batch_size=2, num_actions=3]
    let logits = vec![
        44.0f32, -44.0f32, 0.0f32,
        20.0f32, -20.0f32, 5.0f32,
    ];

    let rows = 2;
    let cols = 3;

    // Clip
    let clipped: Vec<f32> = logits
        .chunks(cols)
        .flat_map(|chunk| clip_logits(chunk, -10.0, 10.0))
        .collect();

    // Verify shape preserved
    assert_eq!(clipped.len(), rows * cols);

    // Verify all values in range
    for (i, &val) in clipped.iter().enumerate() {
        let r = i / cols;
        let c = i % cols;
        assert!(
            val >= -10.0 && val <= 10.0,
            "Value {} at [{}, {}] outside [-10, 10]",
            val,
            r,
            c
        );
    }

    Ok(())
}

#[test]
fn test_distributional_atom_clipping() -> Result<()> {
    // Simulate distributional RL: [batch=2, actions=3, atoms=51]
    let batch_size = 2;
    let num_actions = 3;
    let num_atoms = 51;

    // Create extreme atom logits
    let logits_vec: Vec<f32> = (0..batch_size * num_actions * num_atoms)
        .map(|i| {
            if i % 10 == 0 {
                100.0 // Extreme positive
            } else if i % 10 == 5 {
                -100.0 // Extreme negative
            } else {
                (i as f32) / 10.0 // Normal range
            }
        })
        .collect();

    // Clip all values
    let clipped: Vec<f32> = logits_vec
        .iter()
        .map(|&v| v.clamp(-10.0, 10.0))
        .collect();

    // Apply softmax over atoms for each (batch, action) pair
    // Reshape to [batch*actions, atoms] and softmax per row
    let total_rows = batch_size * num_actions;
    let probs = softmax_2d(&clipped, total_rows, num_atoms);

    // Verify all probabilities > 0 (no saturation) and rows sum to 1
    for i in 0..total_rows {
        let row = &probs[i * num_atoms..(i + 1) * num_atoms];
        let row_sum: f32 = row.iter().sum();
        assert!(
            (row_sum - 1.0).abs() < 1e-4,
            "Row {} probabilities should sum to 1.0, got {}",
            i,
            row_sum
        );

        // Check for saturation (no probability should be exactly 0.0 or 1.0)
        let min_prob = row.iter().copied().fold(f32::INFINITY, f32::min);
        let max_prob = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);

        assert!(
            min_prob > 0.0,
            "Row {} has saturated probability (min={})",
            i,
            min_prob
        );
        assert!(
            max_prob < 1.0,
            "Row {} has saturated probability (max={})",
            i,
            max_prob
        );
    }

    Ok(())
}