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>
153 lines
4.2 KiB
Markdown
153 lines
4.2 KiB
Markdown
# 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
|
|
```rust
|
|
use std::collections::HashSet;
|
|
```
|
|
|
|
#### Struct Fields
|
|
```rust
|
|
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
|
|
```rust
|
|
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**:
|
|
```rust
|
|
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
|
|
|
|
```rust
|
|
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
|
|
```rust
|
|
#[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
|
|
```rust
|
|
#[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
|
|
```rust
|
|
#[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
|