foxhunt/docs/codebase-cleanup/batch_diversity_implementation.md

WAVE 26 P0.2: Batch Diversity Enforcement Implementation

Overview

Added batch diversity enforcement to prevent sampling the same experience twice in consecutive batches.

Implementation Changes

1. Added HashSet for Tracking Recently Sampled Indices

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/prioritized_replay.rs

Imports

use std::collections::HashSet;

Struct Fields

pub struct PrioritizedReplayBuffer {
    // ... existing fields ...
    /// Tracks recently sampled indices to prevent duplicates
    recently_sampled: Arc<Mutex<HashSet<usize>>>,
    /// Counter for cooldown management (cleared every 50 batches)
    sample_counter: AtomicUsize,
}

Constructor

pub fn new(config: PrioritizedReplayConfig) -> Result<Self, MLError> {
    // ... existing initialization ...
    Ok(Self {
        // ... existing fields ...
        recently_sampled: Arc::new(Mutex::new(HashSet::new())),
        sample_counter: AtomicUsize::new(0),
        config,
    })
}

2. Modified sample() Method

Rejection Sampling with Fallback:

pub fn sample(&self, batch_size: usize) -> Result<(Vec<Experience>, Vec<f32>, Vec<usize>), MLError> {
    // ... existing code ...

    let mut recently_sampled = self.recently_sampled.lock();

    for _ in 0..batch_size {
        // Try sampling with diversity enforcement (max 100 attempts)
        let mut attempts = 0;
        let max_attempts = 100;
        let idx = loop {
            let value = rng.gen::<f32>() * total_priority;
            let sampled_idx = tree.sample(value)?;

            // Check if this index was recently sampled
            if !recently_sampled.contains(&sampled_idx) {
                break sampled_idx;
            }

            attempts += 1;
            if attempts >= max_attempts {
                // Fallback: clear cooldown and use this sample
                recently_sampled.clear();
                break sampled_idx;
            }
        };

        // ... process experience ...
        recently_sampled.insert(idx);
    }

    // Increment sample counter and clear cooldown every 50 batches
    let sample_count = self.sample_counter.fetch_add(1, Ordering::Relaxed);
    if sample_count % 50 == 49 {
        recently_sampled.clear();
    }

    // ... return results ...
}

3. Updated clear() Method

pub fn clear(&self) {
    // ... existing clearing code ...

    self.sample_counter.store(0, Ordering::Release);

    // Clear diversity tracking
    {
        let mut recently_sampled = self.recently_sampled.lock();
        recently_sampled.clear();
    }

    // ... reset metrics ...
}

Test Coverage

Test 1: Basic Diversity Enforcement

#[test]
fn test_batch_diversity_no_duplicates() {
    // Tests that consecutive batches don't share indices
    // Verifies HashSet prevents duplicate sampling
}

Test 2: Cooldown After 50 Batches

#[test]
fn test_batch_diversity_cooldown_cleared_after_50_batches() {
    // Verifies cooldown is cleared after 50 batches
    // Ensures long-running training doesn't accumulate all indices
}

Test 3: Fallback on Exhaustion

#[test]
fn test_batch_diversity_fallback_on_exhaustion() {
    // Tests fallback when buffer size < 2*batch_size
    // Ensures graceful degradation instead of infinite loop
}

Performance Characteristics

• Space: O(batch_size) for recently_sampled HashSet
• Time: O(1) expected per sample (rejection sampling), O(100) worst case
• Fallback: Clears cooldown after 100 failed attempts (prevents infinite loops)
• Memory: Cleared every 50 batches (prevents unbounded growth)

Benefits

1. No duplicate experiences in consecutive batches
2. Better gradient diversity - each update uses unique experiences
3. Graceful fallback - handles small buffers without failing
4. Automatic cleanup - cooldown reset every 50 batches
5. Minimal overhead - HashSet lookups are O(1)

Edge Cases Handled

1. Small buffers: Fallback clears cooldown if can't find unique samples
2. Long training: Automatic cooldown reset every 50 batches
3. Thread safety: Mutex protects HashSet from concurrent access
4. Clear buffer: Resets both counter and cooldown set