foxhunt/crates/ml/tests/dqn_accumulation_convergence_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,
)]
//! Convergence comparison test for gradient accumulation.
//!
//! Compares loss trajectories between:
//! - Training with batch_size=16, accumulation_steps=4 (effective batch=64)
//! - Training with batch_size=64, accumulation_steps=1 (direct batch=64)
//!
//! Both should produce similar loss curves since the effective batch size is identical.

#![allow(unused_crate_dependencies)]

use anyhow::{Context, Result};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use std::path::PathBuf;
use tracing::info;

fn get_6e_fut_data_dir() -> Result<String> {
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .context("Failed to get workspace root")?
        .to_path_buf();
    let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
    if !data_dir.exists() {
        anyhow::bail!("6E.FUT data not found: {}", data_dir.display());
    }
    Ok(data_dir.to_string_lossy().to_string())
}

/// Compare accumulated training (batch=16, accum=4) vs direct training (batch=64).
/// Both have effective batch size 64. Loss trajectories should be within 30%.
#[tokio::test]
#[ignore] // Requires real data + runs two training sessions
async fn test_accumulation_convergence_similar_to_direct() -> Result<()> {
    let data_dir = get_6e_fut_data_dir()?;
    let epochs = 10;

    // --- Run 1: Accumulated training (batch=16, accum=4) ---
    let mut hp_accum = DQNHyperparameters::conservative();
    hp_accum.replay_buffer_vram_fraction = 0.0; // Disable AutoReplaySizer for test determinism
    hp_accum.epochs = epochs;
    hp_accum.batch_size = 16;
    hp_accum.gradient_accumulation_steps = 4;
    hp_accum.learning_rate = 0.0001;
    hp_accum.early_stopping_enabled = false;
    hp_accum.checkpoint_frequency = 100;

    let checkpoint_dir_1 = tempfile::tempdir()?;
    let ckpt_path_1 = checkpoint_dir_1.path().to_path_buf();
    let mut trainer_accum = DQNTrainer::new(hp_accum)?;
    let _metrics_accum = trainer_accum
        .train(&data_dir, "ES.FUT", move |_epoch, data, _is_best| {
            let p = ckpt_path_1.join("accum.safetensors");
            std::fs::write(&p, &data)?;
            Ok(p.to_string_lossy().to_string())
        })
        .await?;
    let loss_accum = trainer_accum.loss_history().to_vec();

    // --- Run 2: Direct training (batch=64, accum=1) ---
    let mut hp_direct = DQNHyperparameters::conservative();
    hp_direct.replay_buffer_vram_fraction = 0.0; // Disable AutoReplaySizer for test determinism
    hp_direct.epochs = epochs;
    hp_direct.batch_size = 64;
    hp_direct.gradient_accumulation_steps = 1;
    hp_direct.learning_rate = 0.0001;
    hp_direct.early_stopping_enabled = false;
    hp_direct.checkpoint_frequency = 100;

    let checkpoint_dir_2 = tempfile::tempdir()?;
    let ckpt_path_2 = checkpoint_dir_2.path().to_path_buf();
    let mut trainer_direct = DQNTrainer::new(hp_direct)?;
    let _metrics_direct = trainer_direct
        .train(&data_dir, "ES.FUT", move |_epoch, data, _is_best| {
            let p = ckpt_path_2.join("direct.safetensors");
            std::fs::write(&p, &data)?;
            Ok(p.to_string_lossy().to_string())
        })
        .await?;
    let loss_direct = trainer_direct.loss_history().to_vec();

    // --- Compare ---
    let min_len = loss_accum.len().min(loss_direct.len());
    assert!(min_len >= 5, "Both runs should complete at least 5 epochs");

    let final_accum = loss_accum.get(min_len - 1).copied().unwrap_or(f64::NAN);
    let final_direct = loss_direct.get(min_len - 1).copied().unwrap_or(f64::NAN);

    assert!(
        final_accum.is_finite() && final_direct.is_finite(),
        "Final losses must be finite: accum={:.6}, direct={:.6}",
        final_accum,
        final_direct
    );

    let avg = (final_accum + final_direct) / 2.0;
    let rel_diff = (final_accum - final_direct).abs() / avg;

    let accum_decrease = 1.0 - final_accum / loss_accum.first().copied().unwrap_or(1.0);
    let direct_decrease = 1.0 - final_direct / loss_direct.first().copied().unwrap_or(1.0);

    info!(
        accum_initial = loss_accum.first().copied().unwrap_or(0.0),
        accum_final = final_accum,
        accum_decrease_pct = accum_decrease * 100.0,
        direct_initial = loss_direct.first().copied().unwrap_or(0.0),
        direct_final = final_direct,
        direct_decrease_pct = direct_decrease * 100.0,
        rel_diff_pct = rel_diff * 100.0,
        "Accumulation convergence test results"
    );

    assert!(
        rel_diff < 0.30,
        "Loss trajectories diverged too much: accum={:.6}, direct={:.6}, rel_diff={:.1}%",
        final_accum,
        final_direct,
        rel_diff * 100.0
    );

    assert!(
        accum_decrease > 0.02,
        "Accumulated training should decrease loss >2% (got {:.1}%)",
        accum_decrease * 100.0
    );
    assert!(
        direct_decrease > 0.02,
        "Direct training should decrease loss >2% (got {:.1}%)",
        direct_decrease * 100.0
    );

    Ok(())
}