foxhunt/ml/examples/diagnose_q_value_diversity.rs

//! Diagnostic Script: Analyze Q-value Diversity in FactoredQNetwork
//!
//! This script investigates why only 8 out of 45 factored actions are being selected
//! during DQN training with factored actions enabled.
//!
//! HYPOTHESIS:
//! Additive Q-value factorization (Q(e,o,u) = Q_exp[e] + Q_ord[o] + Q_urg[u])
//! may create duplicate Q-values due to commutative addition, resulting in
//! clustering where multiple actions have identical Q-values.
//!
//! TEST METHODOLOGY:
//! 1. Initialize FactoredQNetwork with random weights (Xavier uniform)
//! 2. Generate random state tensor
//! 3. Forward pass to get 3 Q-value heads
//! 4. Compute joint Q-values using additive factorization
//! 5. Count unique Q-values (with float tolerance ε=1e-6)
//! 6. Analyze distribution and clustering patterns
//! 7. Test with multiple random seeds and states
//!
//! EXPECTED OUTCOMES:
//! - If 45 unique Q-values: Factorization works correctly, issue is in action selection logic
//! - If <45 unique Q-values: Additive factorization creates duplicates (FIX: multiplicative or different architecture)
//! - If exactly 8 unique Q-values: Confirms hypothesis from training logs

use candle_core::{Device, Tensor};
use ml::dqn::action_space::FactoredAction;
use ml::dqn::factored_q_network::FactoredQNetwork;
use std::collections::HashMap;

const FLOAT_TOLERANCE: f32 = 1e-6;

/// Round float to nearest multiple of tolerance for uniqueness checking
fn round_to_tolerance(value: f32, tolerance: f32) -> i64 {
    (value / tolerance).round() as i64
}

/// Analyze Q-value diversity from factored network output
fn analyze_q_diversity(q_values: &[f32]) -> (usize, HashMap<i64, Vec<usize>>) {
    let mut clusters: HashMap<i64, Vec<usize>> = HashMap::new();

    for (idx, &q_val) in q_values.iter().enumerate() {
        let key = round_to_tolerance(q_val, FLOAT_TOLERANCE);
        clusters.entry(key).or_insert_with(Vec::new).push(idx);
    }

    (clusters.len(), clusters)
}

/// Run diagnostic test with given seed
fn run_diagnostic_test(seed: u64, device: &Device) -> anyhow::Result<()> {
    println!("\n{}", "=".repeat(80));
    println!("TEST RUN: Seed {}", seed);
    println!("{}", "=".repeat(80));

    // 1. Initialize FactoredQNetwork
    let network = FactoredQNetwork::new(128, device)?;

    // 2. Generate random state (using seed for reproducibility)
    let state = if seed == 0 {
        Tensor::zeros((1, 128), candle_core::DType::F32, device)?
    } else {
        Tensor::randn(seed as f32, 1.0f32, (1, 128), device)?
    };

    println!("State shape: {:?}", state.dims());

    // 3. Forward pass
    let (q_exposure, q_order, q_urgency) = network.forward(&state)?;

    println!("Q-value head shapes:");
    println!("  - Exposure: {:?}", q_exposure.dims());
    println!("  - Order: {:?}", q_order.dims());
    println!("  - Urgency: {:?}", q_urgency.dims());

    // 4. Compute joint Q-values
    let joint_q = network.compute_joint_q(&q_exposure, &q_order, &q_urgency)?;

    println!("Joint Q-values shape: {:?}", joint_q.dims());

    // 5. Extract Q-values
    let q_vec = joint_q.squeeze(0)?.to_vec1::<f32>()?;

    println!("Extracted {} Q-values", q_vec.len());

    // 6. Analyze diversity
    let (num_unique, clusters) = analyze_q_diversity(&q_vec);

    println!("\n{}", "-".repeat(80));
    println!("DIVERSITY ANALYSIS");
    println!("{}", "-".repeat(80));
    println!("Total actions: 45");
    println!("Unique Q-values: {}", num_unique);
    println!(
        "Duplicate rate: {:.1}%",
        100.0 * (45 - num_unique) as f32 / 45.0
    );

    // 7. Show Q-value statistics
    let min_q = q_vec.iter().copied().fold(f32::INFINITY, f32::min);
    let max_q = q_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let mean_q = q_vec.iter().sum::<f32>() / q_vec.len() as f32;
    let variance = q_vec.iter().map(|&x| (x - mean_q).powi(2)).sum::<f32>() / q_vec.len() as f32;
    let std_q = variance.sqrt();

    println!("\nQ-VALUE STATISTICS:");
    println!("  Min: {:.6}", min_q);
    println!("  Max: {:.6}", max_q);
    println!("  Mean: {:.6}", mean_q);
    println!("  Std: {:.6}", std_q);
    println!("  Range: {:.6}", max_q - min_q);

    // 8. Show cluster distribution
    println!("\nCLUSTER DISTRIBUTION:");
    let mut cluster_sizes: Vec<_> = clusters.values().map(|v| v.len()).collect();
    cluster_sizes.sort_unstable_by(|a, b| b.cmp(a)); // Descending

    for (i, &size) in cluster_sizes.iter().enumerate().take(10) {
        println!(
            "  Cluster {}: {} actions ({:.1}%)",
            i + 1,
            size,
            100.0 * size as f32 / 45.0
        );
    }

    // 9. Show example clusters (if duplicates exist)
    if num_unique < 45 {
        println!("\nEXAMPLE DUPLICATE CLUSTERS (showing first 3):");
        let mut sorted_clusters: Vec<_> = clusters
            .iter()
            .filter(|(_, actions)| actions.len() > 1)
            .collect();
        sorted_clusters.sort_by(|a, b| b.1.len().cmp(&a.1.len()));

        for (cluster_idx, (q_key, action_indices)) in sorted_clusters.iter().take(3).enumerate() {
            let q_value = **q_key as f32 * FLOAT_TOLERANCE;
            println!(
                "\n  Cluster #{} (Q={:.6}, {} actions):",
                cluster_idx + 1,
                q_value,
                action_indices.len()
            );

            for &idx in action_indices.iter().take(5) {
                let action = FactoredAction::from_index(idx)?;
                println!("    - Action {}: {:?}", idx, action);
            }

            if action_indices.len() > 5 {
                println!("    ... and {} more", action_indices.len() - 5);
            }
        }
    }

    // 10. Analyze head contributions
    println!("\n{}", "-".repeat(80));
    println!("HEAD CONTRIBUTION ANALYSIS");
    println!("{}", "-".repeat(80));

    let exp_vec = q_exposure.squeeze(0)?.to_vec1::<f32>()?;
    let ord_vec = q_order.squeeze(0)?.to_vec1::<f32>()?;
    let urg_vec = q_urgency.squeeze(0)?.to_vec1::<f32>()?;

    let exp_range = exp_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max)
        - exp_vec.iter().copied().fold(f32::INFINITY, f32::min);
    let ord_range = ord_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max)
        - ord_vec.iter().copied().fold(f32::INFINITY, f32::min);
    let urg_range = urg_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max)
        - urg_vec.iter().copied().fold(f32::INFINITY, f32::min);

    println!("Exposure head (5 values):");
    println!("  Values: {:?}", exp_vec);
    println!("  Range: {:.6}", exp_range);

    println!("\nOrder head (3 values):");
    println!("  Values: {:?}", ord_vec);
    println!("  Range: {:.6}", ord_range);

    println!("\nUrgency head (3 values):");
    println!("  Values: {:?}", urg_vec);
    println!("  Range: {:.6}", urg_range);

    let total_range = exp_range + ord_range + urg_range;
    println!("\nRelative Contributions:");
    println!(
        "  Exposure: {:.1}% ({:.6} / {:.6})",
        100.0 * exp_range / total_range,
        exp_range,
        total_range
    );
    println!(
        "  Order: {:.1}% ({:.6} / {:.6})",
        100.0 * ord_range / total_range,
        ord_range,
        total_range
    );
    println!(
        "  Urgency: {:.1}% ({:.6} / {:.6})",
        100.0 * urg_range / total_range,
        urg_range,
        total_range
    );

    // 11. Show argmax behavior
    println!("\n{}", "-".repeat(80));
    println!("ARGMAX BEHAVIOR");
    println!("{}", "-".repeat(80));

    let argmax_idx = q_vec
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
        .map(|(idx, _)| idx)
        .unwrap();

    let max_q_value = q_vec[argmax_idx];
    let action = FactoredAction::from_index(argmax_idx)?;

    println!("Greedy action: {} (Q={:.6})", argmax_idx, max_q_value);
    println!("  {:?}", action);

    // Count how many actions share this Q-value
    let tied_actions: Vec<_> = q_vec
        .iter()
        .enumerate()
        .filter(|(_, &q)| (q - max_q_value).abs() < FLOAT_TOLERANCE)
        .map(|(idx, _)| idx)
        .collect();

    if tied_actions.len() > 1 {
        println!(
            "\nWARNING: {} actions tied for max Q-value!",
            tied_actions.len()
        );
        println!("Tied actions: {:?}", tied_actions);
        println!(
            "Argmax will deterministically select action {} (first occurrence)",
            argmax_idx
        );
    }

    Ok(())
}

fn main() -> anyhow::Result<()> {
    // Enable info logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("\n{}", "=".repeat(80));
    println!("FACTORED Q-NETWORK DIVERSITY DIAGNOSTIC");
    println!("{}", "=".repeat(80));
    println!("\nThis script analyzes whether additive Q-value factorization");
    println!("Q(e,o,u) = Q_exposure[e] + Q_order[o] + Q_urgency[u]");
    println!("creates duplicate Q-values across the 45-action space.");
    println!("\nExpected: 45 unique Q-values");
    println!("Observed in training: Only 8 unique actions selected");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("\nDevice: {:?}", device);

    // Run multiple tests with different seeds
    let seeds = vec![0, 1, 42, 123, 999];

    let mut all_unique_counts = Vec::new();

    for seed in &seeds {
        run_diagnostic_test(*seed, &device)?;

        // Re-run to get unique count for statistics
        let network = FactoredQNetwork::new(128, &device)?;
        let state = if *seed == 0 {
            Tensor::zeros((1, 128), candle_core::DType::F32, &device)?
        } else {
            Tensor::randn(*seed as f32, 1.0f32, (1, 128), &device)?
        };
        let (q_exposure, q_order, q_urgency) = network.forward(&state)?;
        let joint_q = network.compute_joint_q(&q_exposure, &q_order, &q_urgency)?;
        let q_vec = joint_q.squeeze(0)?.to_vec1::<f32>()?;
        let (num_unique, _) = analyze_q_diversity(&q_vec);
        all_unique_counts.push(num_unique);
    }

    // Summary statistics
    println!("\n\n{}", "=".repeat(80));
    println!("SUMMARY ACROSS {} RANDOM INITIALIZATIONS", seeds.len());
    println!("{}", "=".repeat(80));

    let min_unique = *all_unique_counts.iter().min().unwrap();
    let max_unique = *all_unique_counts.iter().max().unwrap();
    let mean_unique =
        all_unique_counts.iter().sum::<usize>() as f32 / all_unique_counts.len() as f32;

    println!("Unique Q-values per run: {:?}", all_unique_counts);
    println!("  Min: {}", min_unique);
    println!("  Max: {}", max_unique);
    println!("  Mean: {:.1}", mean_unique);

    if max_unique < 45 {
        println!("\n{}", "=".repeat(80));
        println!("ROOT CAUSE IDENTIFIED");
        println!("{}", "=".repeat(80));
        println!("Additive factorization Q(e,o,u) = Q_exp[e] + Q_ord[o] + Q_urg[u]");
        println!("creates DUPLICATE Q-values due to commutative addition.");
        println!("\nExpected: 45 unique Q-values (5 × 3 × 3 combinations)");
        println!(
            "Actual: {}-{} unique Q-values ({:.1}% duplication)",
            min_unique,
            max_unique,
            100.0 * (45.0 - mean_unique) / 45.0
        );
        println!("\n{}", "=".repeat(80));
        println!("RECOMMENDED FIX");
        println!("{}", "=".repeat(80));
        println!("Option 1: Multiplicative factorization");
        println!("  Q(e,o,u) = Q_exp[e] × Q_ord[o] × Q_urg[u]");
        println!("  Pros: Preserves 45 unique values, non-commutative");
        println!("  Cons: Requires careful initialization (avoid zeros)");
        println!("\nOption 2: Concatenation + MLP");
        println!("  hidden = MLP([Q_exp, Q_ord, Q_urg])");
        println!("  Q(e,o,u) = Linear(hidden)");
        println!("  Pros: Learnable interactions, no duplicates");
        println!("  Cons: More parameters, slower inference");
        println!("\nOption 3: Weighted sum with learnable weights");
        println!("  Q(e,o,u) = w_e * Q_exp[e] + w_o * Q_ord[o] + w_u * Q_urg[u]");
        println!("  Pros: Simple, learnable weights");
        println!("  Cons: Still commutative, may still cluster");
    } else {
        println!("\nSURPRISE: Factorization produces 45 unique Q-values!");
        println!("Root cause is NOT in Q-value computation.");
        println!("Check action selection logic (epsilon_greedy, argmax, mapping).");
    }

    Ok(())
}