BREAKING CHANGES: - Removed orphaned dqn.rs monolithic trainer (4,975 lines) - Removed orphaned dqn_ensemble.rs module (816 lines) - Removed orphaned tft.rs and tft_complete_int8_integration_test.rs - TFT trainer split into modular directory structure DQN Module Refactoring: - Split trainers/dqn.rs into modular structure (config.rs, statistics.rs, trainer.rs) - Fixed hyperopt 39D search space (continuous params only) - Boolean flags (use_dueling, use_double_dqn, use_per, use_noisy_nets) are now FIXED architectural decisions - use_distributional defaults to false (Candle BUG #36 - scatter_add gradient issues) Clean Module Structure: - ml/src/trainers/dqn/ directory with proper mod.rs exports - ml/src/trainers/tft/ directory with config.rs, types.rs, model.rs, trainer.rs, tests.rs - All P0 features validated: TD-error clamping, batch diversity, LR scheduler, priority staleness Documentation: - Added comprehensive docs in docs/codebase-cleanup/ - ADR-001 for DQN refactoring decisions - Rainbow DQN component matrix and quick reference guides Build Status: Compiles with zero errors 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
205 lines
5.8 KiB
Markdown
205 lines
5.8 KiB
Markdown
# WAVE 26 P0.2: Batch Diversity Enforcement - Implementation Summary
|
|
|
|
## Objective
|
|
Prevent sampling the same experience twice in consecutive batches to improve gradient diversity and training stability.
|
|
|
|
## Changes Made
|
|
|
|
### 1. Added HashSet Tracking (`/home/jgrusewski/Work/foxhunt/ml/src/dqn/prioritized_replay.rs`)
|
|
|
|
**Imports:**
|
|
```rust
|
|
use std::collections::HashSet;
|
|
```
|
|
|
|
**Struct Fields:**
|
|
```rust
|
|
pub struct PrioritizedReplayBuffer {
|
|
// ... existing fields ...
|
|
|
|
/// Tracks recently sampled indices to prevent duplicates (WAVE 26 P0.2)
|
|
recently_sampled: Arc<Mutex<HashSet<usize>>>,
|
|
|
|
/// Counter for cooldown management - cleared every 50 batches (WAVE 26 P0.2)
|
|
sample_counter: AtomicUsize,
|
|
}
|
|
```
|
|
|
|
**Constructor:**
|
|
```rust
|
|
pub fn new(config: PrioritizedReplayConfig) -> Result<Self, MLError> {
|
|
// ...
|
|
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:**
|
|
```rust
|
|
let mut recently_sampled = self.recently_sampled.lock();
|
|
|
|
for _ in 0..batch_size {
|
|
// WAVE 26 P0.2: Batch diversity enforcement
|
|
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);
|
|
}
|
|
|
|
// Clear cooldown every 50 batches
|
|
let sample_count = self.sample_counter.fetch_add(1, Ordering::Relaxed);
|
|
if sample_count % 50 == 49 {
|
|
recently_sampled.clear();
|
|
}
|
|
```
|
|
|
|
### 3. Updated `clear()` Method
|
|
|
|
**Reset Diversity Tracking:**
|
|
```rust
|
|
pub fn clear(&self) {
|
|
// ... existing clearing code ...
|
|
|
|
self.sample_counter.store(0, Ordering::Release);
|
|
|
|
// WAVE 26 P0.2: Clear diversity tracking
|
|
{
|
|
let mut recently_sampled = self.recently_sampled.lock();
|
|
recently_sampled.clear();
|
|
}
|
|
|
|
// ... reset metrics ...
|
|
}
|
|
```
|
|
|
|
## Test Coverage
|
|
|
|
### Test 1: `test_batch_diversity_no_duplicates`
|
|
**Purpose:** Verify consecutive batches don't share indices
|
|
**Setup:**
|
|
- Buffer capacity: 200
|
|
- Add 100 experiences
|
|
- Sample 2 batches of 32 each
|
|
|
|
**Validation:**
|
|
```rust
|
|
let set1: HashSet<_> = batch1_indices.iter().copied().collect();
|
|
let set2: HashSet<_> = batch2_indices.iter().copied().collect();
|
|
assert!(set1.is_disjoint(&set2));
|
|
```
|
|
|
|
### Test 2: `test_batch_diversity_cooldown_cleared_after_50_batches`
|
|
**Purpose:** Verify cooldown resets after 50 batches
|
|
**Setup:**
|
|
- Buffer capacity: 500
|
|
- Add 500 experiences
|
|
- Sample 51 batches of 10 each
|
|
|
|
**Validation:**
|
|
- First 50 batches: all indices must be unique
|
|
- 51st batch: can reuse indices (cooldown cleared)
|
|
|
|
### Test 3: `test_batch_diversity_fallback_on_exhaustion`
|
|
**Purpose:** Verify graceful fallback when buffer too small
|
|
**Setup:**
|
|
- Buffer capacity: 100
|
|
- Add only 40 experiences
|
|
- Sample 2 batches of 32 each
|
|
|
|
**Validation:**
|
|
- Both batches should succeed (no infinite loop)
|
|
- Overlap is expected and acceptable
|
|
- Fallback clears cooldown after 100 failed attempts
|
|
|
|
## Algorithm Details
|
|
|
|
### Rejection Sampling
|
|
1. Sample index from prioritized distribution
|
|
2. Check if index in `recently_sampled` HashSet
|
|
3. If yes, retry (max 100 attempts)
|
|
4. If max attempts reached, clear cooldown and use current sample
|
|
5. Add sampled index to `recently_sampled`
|
|
|
|
### Cooldown Management
|
|
- Counter incremented after each `sample()` call
|
|
- Every 50 batches (`sample_count % 50 == 49`), clear `recently_sampled`
|
|
- Prevents unbounded HashSet growth during long training runs
|
|
|
|
### Edge Cases Handled
|
|
|
|
1. **Small buffers:** Fallback prevents infinite loops
|
|
2. **Long training:** Automatic cooldown every 50 batches
|
|
3. **Thread safety:** Mutex protects HashSet
|
|
4. **Buffer clear:** Resets both counter and HashSet
|
|
|
|
## Performance Characteristics
|
|
|
|
- **Space Complexity:** O(batch_size) for recently_sampled HashSet
|
|
- **Time Complexity:** O(1) expected per sample, O(100) worst case
|
|
- **Memory Overhead:** Minimal - HashSet cleared every 50 batches
|
|
- **Fallback Safety:** Prevents infinite loops in edge cases
|
|
|
|
## Benefits
|
|
|
|
1. **No Duplicate Experiences:** Consecutive batches are guaranteed disjoint (up to 50 batches)
|
|
2. **Better Gradient Diversity:** Each training update uses unique experiences
|
|
3. **Graceful Degradation:** Handles small buffers without failing
|
|
4. **Automatic Cleanup:** Prevents memory growth over long runs
|
|
5. **Minimal Overhead:** O(1) HashSet lookups per sample
|
|
|
|
## Integration Notes
|
|
|
|
- **No API Changes:** Existing `sample()` signature unchanged
|
|
- **Backward Compatible:** Feature automatically enabled
|
|
- **Configuration:** Hardcoded to 50-batch cooldown (can be parameterized later)
|
|
- **Thread Safe:** Uses `Arc<Mutex<HashSet<usize>>>` for concurrent access
|
|
|
|
## Next Steps (Future Enhancements)
|
|
|
|
1. Make cooldown period configurable via `PrioritizedReplayConfig`
|
|
2. Add metrics to track fallback frequency
|
|
3. Consider adaptive cooldown based on buffer utilization
|
|
4. Expose diversity statistics in `PrioritizedReplayMetrics`
|
|
|
|
## Files Modified
|
|
|
|
- `/home/jgrusewski/Work/foxhunt/ml/src/dqn/prioritized_replay.rs`
|
|
- Added `HashSet` import
|
|
- Added `recently_sampled` and `sample_counter` fields
|
|
- Modified `new()`, `sample()`, and `clear()` methods
|
|
- Added 3 comprehensive tests
|
|
|
|
## Documentation
|
|
|
|
- Implementation details: `/home/jgrusewski/Work/foxhunt/docs/codebase-cleanup/batch_diversity_implementation.md`
|
|
- Summary: This file
|
|
|
|
---
|
|
|
|
**Status:** ✅ Implementation complete, tests added, awaiting test results
|
|
**Wave:** WAVE 26 P0.2
|
|
**Date:** 2025-11-27
|