foxhunt/ml/tests/gradient_accumulation_tests.rs

//! TDD Tests for Gradient Accumulation (WAVE 26 P2.2)
//!
//! Gradient accumulation allows larger effective batch sizes by accumulating
//! gradients over multiple forward/backward passes before applying optimizer step.
//!
//! Benefits:
//! - Larger effective batch size without GPU memory limits
//! - More stable gradient estimates (reduced variance)
//! - Better convergence in low-data regimes
//!
//! Example: batch_size=64, accumulation_steps=4 -> effective_batch_size=256

use ml::trainers::dqn::DQNHyperparameters;
use ml::trainers::dqn::DQNTrainer;

/// Helper to create minimal hyperparameters for testing
fn create_test_hyperparams() -> DQNHyperparameters {
    let mut params = DQNHyperparameters::conservative();
    params.epochs = 2; // Minimal epochs for fast testing
    params.batch_size = 32;
    params.buffer_size = 1000;
    params.min_replay_size = 100;
    params.learning_rate = 0.001;
    params.gradient_accumulation_steps = 1; // Default: no accumulation
    params
}

/// TEST 1: Verify default configuration (no accumulation)
#[test]
fn test_default_accumulation_steps_is_one() {
    let params = create_test_hyperparams();
    assert_eq!(
        params.gradient_accumulation_steps, 1,
        "Default accumulation_steps should be 1 (no accumulation)"
    );
}

/// TEST 2: Verify accumulation_steps field exists in config
#[test]
fn test_accumulation_steps_field_in_config() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 4;
    assert_eq!(
        params.gradient_accumulation_steps, 4,
        "Should be able to set gradient_accumulation_steps"
    );
}

/// TEST 3: Verify trainer accepts gradient_accumulation_steps config
#[test]
fn test_trainer_accepts_accumulation_config() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 2;

    let trainer = DQNTrainer::new(params);
    assert!(
        trainer.is_ok(),
        "Trainer should accept gradient_accumulation_steps config"
    );
}

/// TEST 4: Verify gradient accumulation field is stored in trainer
#[test]
fn test_trainer_stores_accumulation_steps() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 4;

    let trainer = DQNTrainer::new(params).expect("Failed to create trainer");

    // Trainer should store the accumulation_steps value
    // This will be accessible via trainer.hyperparams().gradient_accumulation_steps
    assert_eq!(
        trainer.hyperparams().gradient_accumulation_steps, 4,
        "Trainer should store accumulation_steps from config"
    );
}

/// TEST 5: Validate accumulation_steps > 0
#[test]
fn test_reject_zero_accumulation_steps() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 0; // Invalid

    let trainer = DQNTrainer::new(params);
    assert!(
        trainer.is_err(),
        "Should reject gradient_accumulation_steps = 0"
    );

    let err_msg = trainer.unwrap_err().to_string();
    assert!(
        err_msg.contains("gradient_accumulation_steps") || err_msg.contains("greater than 0"),
        "Error message should mention gradient_accumulation_steps validation"
    );
}

/// TEST 6: Verify effective batch size calculation
#[test]
fn test_effective_batch_size_calculation() {
    let mut params = create_test_hyperparams();
    params.batch_size = 64;
    params.gradient_accumulation_steps = 4;

    let trainer = DQNTrainer::new(params).expect("Failed to create trainer");

    // Effective batch size = batch_size * accumulation_steps
    let expected_effective_batch = 64 * 4; // 256

    // This validates that the trainer understands the concept
    assert_eq!(
        trainer.hyperparams().batch_size * trainer.hyperparams().gradient_accumulation_steps,
        expected_effective_batch,
        "Effective batch size should be batch_size * accumulation_steps"
    );
}

/// TEST 7: Verify gradient accumulation doesn't exceed replay buffer size
#[test]
fn test_accumulation_bounded_by_replay_buffer() {
    let mut params = create_test_hyperparams();
    params.batch_size = 64;
    params.gradient_accumulation_steps = 4; // Effective batch = 256
    params.buffer_size = 1000; // Buffer has enough samples

    let trainer = DQNTrainer::new(params);
    assert!(
        trainer.is_ok(),
        "Should accept accumulation when buffer is large enough"
    );

    // But if effective batch exceeds buffer...
    let mut params2 = create_test_hyperparams();
    params2.batch_size = 64;
    params2.gradient_accumulation_steps = 100; // Effective batch = 6400 > buffer_size=1000
    params2.buffer_size = 1000;

    let trainer2 = DQNTrainer::new(params2);
    // Trainer creation should succeed (validation happens at runtime)
    assert!(
        trainer2.is_ok(),
        "Trainer creation should succeed (runtime will validate)"
    );
}

/// TEST 8: Verify gradient accumulation preserves loss scaling
/// Loss should be scaled by 1/accumulation_steps to maintain gradient magnitude
#[tokio::test]
async fn test_loss_scaling_during_accumulation() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 4;
    params.batch_size = 32;
    params.min_replay_size = 50; // Low threshold for faster testing

    let trainer = DQNTrainer::new(params).expect("Failed to create trainer");

    // When training with accumulation, each mini-batch loss should be divided by accumulation_steps
    // This test verifies the concept exists (actual implementation will test functional behavior)

    // The scaled loss prevents gradient explosion when accumulating
    let accumulation_steps = trainer.hyperparams().gradient_accumulation_steps;
    let scale_factor = 1.0 / accumulation_steps as f64;

    assert!(
        scale_factor > 0.0 && scale_factor <= 1.0,
        "Loss scale factor should be in range (0, 1] for accumulation_steps >= 1"
    );
}

/// TEST 9: Verify optimizer step called once per accumulation cycle
/// With accumulation_steps=4, optimizer.step() should be called every 4 mini-batches
#[test]
fn test_optimizer_step_frequency() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 4;

    let trainer = DQNTrainer::new(params).expect("Failed to create trainer");

    // Verify the trainer has the accumulation_steps configured
    assert_eq!(
        trainer.hyperparams().gradient_accumulation_steps, 4,
        "Trainer should have accumulation_steps=4"
    );

    // The implementation will ensure optimizer.step() is called after every 4 mini-batches
    // This test validates the configuration exists
}

/// TEST 10: Verify gradient accumulation resets after optimizer step
#[tokio::test]
async fn test_gradients_reset_after_optimizer_step() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 2;
    params.batch_size = 32;

    let trainer = DQNTrainer::new(params).expect("Failed to create trainer");

    // After accumulating gradients and calling optimizer.step(),
    // gradients should be zeroed via optimizer.zero_grad()
    // This test verifies the configuration (implementation will test functional behavior)

    assert_eq!(
        trainer.hyperparams().gradient_accumulation_steps, 2,
        "Should use accumulation_steps=2"
    );
}

/// TEST 11: Verify accumulation works with various batch sizes
#[test]
fn test_accumulation_with_different_batch_sizes() {
    // Test small batch + accumulation
    let mut params1 = create_test_hyperparams();
    params1.batch_size = 16;
    params1.gradient_accumulation_steps = 8; // Effective: 128
    assert!(DQNTrainer::new(params1).is_ok());

    // Test medium batch + accumulation
    let mut params2 = create_test_hyperparams();
    params2.batch_size = 64;
    params2.gradient_accumulation_steps = 4; // Effective: 256
    assert!(DQNTrainer::new(params2).is_ok());

    // Test large batch + no accumulation
    let mut params3 = create_test_hyperparams();
    params3.batch_size = 128;
    params3.gradient_accumulation_steps = 1; // Effective: 128
    assert!(DQNTrainer::new(params3).is_ok());
}

/// TEST 12: Verify accumulation respects GPU memory limits
/// Even with accumulation, each mini-batch must fit in GPU memory
#[test]
fn test_accumulation_respects_gpu_memory_limit() {
    let mut params = create_test_hyperparams();
    params.batch_size = 250; // Exceeds GPU limit (230)
    params.gradient_accumulation_steps = 4;

    let trainer = DQNTrainer::new(params);
    assert!(
        trainer.is_err(),
        "Should reject batch_size > 230 even with accumulation"
    );

    let err_msg = trainer.unwrap_err().to_string();
    assert!(
        err_msg.contains("230") || err_msg.contains("GPU memory"),
        "Error should mention GPU memory limit"
    );
}

/// TEST 13: Verify accumulation_steps=1 behaves like normal training
#[tokio::test]
async fn test_no_accumulation_is_backward_compatible() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 1; // No accumulation
    params.batch_size = 32;

    let trainer = DQNTrainer::new(params).expect("Failed to create trainer");

    // With accumulation_steps=1, behavior should be identical to normal training
    // (optimizer.step() called after every batch, no gradient accumulation)
    assert_eq!(
        trainer.hyperparams().gradient_accumulation_steps, 1,
        "accumulation_steps=1 should behave like normal training"
    );
}

/// TEST 14: Verify large accumulation steps (stress test)
#[test]
fn test_large_accumulation_steps() {
    let mut params = create_test_hyperparams();
    params.batch_size = 32;
    params.gradient_accumulation_steps = 16; // Effective batch: 512
    params.buffer_size = 10000; // Large buffer to support effective batch

    let trainer = DQNTrainer::new(params);
    assert!(
        trainer.is_ok(),
        "Should accept large accumulation_steps with sufficient buffer"
    );
}

/// TEST 15: Verify accumulation works with all DQN features
#[test]
fn test_accumulation_with_rainbow_dqn_features() {
    let mut params = create_test_hyperparams();
    params.gradient_accumulation_steps = 4;
    params.use_double_dqn = true;
    params.use_dueling = true;
    params.use_distributional = true;
    params.use_per = true;
    params.use_noisy_nets = true;

    let trainer = DQNTrainer::new(params);
    assert!(
        trainer.is_ok(),
        "Gradient accumulation should work with all Rainbow DQN features"
    );
}