foxhunt/docs/ADR-001-dqn-refactoring.md

ADR-001: DQN Trainer Modularization

Status

PROPOSED

Context

The DQN trainer module (ml/src/trainers/dqn.rs) has grown to 4,975 lines, making it:

• Unmaintainable: Changes require navigating massive single file
• Untestable: Unit testing individual components is difficult
• Unreadable: Logical boundaries obscured by file size
• Risky: High chance of introducing bugs during modifications

Current Dependencies

• 19 test files depend on trainers::dqn::{DQNHyperparameters, DQNTrainer}
• Hyperopt adapter (hyperopt/adapters/dqn.rs, 3,162 lines) imports trainer types
• Public API exports: pub use dqn::{DQNHyperparameters, DQNTrainer}

Decision

Refactor into module structure:

ml/src/trainers/dqn/
├── mod.rs                    # Public API, re-exports (~50 lines)
├── config.rs                 # Hyperparameters, constants (~800 lines)
├── agent_wrapper.rs          # DQNAgentType enum (~260 lines)
├── training_monitor.rs       # Validation, metrics (~270 lines)
├── trainer_core.rs           # Struct, constructors (~600 lines)
├── training_loop.rs          # Main train() methods (~1200 lines)
├── data_loading.rs           # OHLCV extraction (~600 lines)
└── checkpointing.rs          # Save/load logic (~200 lines)

Module Boundaries

config.rs (Lines 48-747)

Responsibility: Configuration and normalization

pub const EPISODE_LENGTH: usize = 200;
pub type FeatureVector = [f64; 54];
pub type FeatureVector51 = [f64; 51];
pub struct FeatureStatistics { /* Welford's algorithm */ }
pub struct DQNHyperparameters { /* 60+ fields */ }

agent_wrapper.rs (Lines 164-424)

Responsibility: Unified agent API (Standard vs Regime-Conditional)

pub enum DQNAgentType {
    Standard(WorkingDQN),
    RegimeConditional(RegimeConditionalDQN),
}
pub struct QValueStats { /* C51 bounds */ }

training_monitor.rs (Lines 748-1012)

Responsibility: Training validation and diagnostics

struct TrainingMonitor {
    // Validates: reward diversity, action diversity, Q-value balance
}

trainer_core.rs (Lines 1013-1500)

Responsibility: DQNTrainer struct and initialization

pub struct DQNTrainer {
    agent: Arc<RwLock<DQNAgentType>>,
    hyperparams: DQNHyperparameters,
    // ... 20+ fields for risk management, microstructure, etc.
}
impl DQNTrainer {
    pub fn new(hyperparams: DQNHyperparameters) -> Result<Self> { /* ... */ }
}

training_loop.rs (Lines 1464-2980)

Responsibility: Main training algorithms

impl DQNTrainer {
    pub async fn train<F>(&mut self, ...) -> Result<TrainingMetrics> { /* ... */ }
    pub async fn train_from_parquet<F>(&mut self, ...) -> Result<TrainingMetrics> { /* ... */ }
}

data_loading.rs (Lines 3482-4600)

Responsibility: OHLCV extraction, feature engineering

impl DQNTrainer {
    pub fn extract_ohlcv_bars_from_dbn(&self, ...) -> Result<Vec<OHLCVBar>> { /* ... */ }
    fn create_features(&self, ...) -> Result<Vec<f32>> { /* ... */ }
}

checkpointing.rs (Lines scattered)

Responsibility: Model save/load

impl DQNTrainer {
    fn save_checkpoint(&self, ...) -> Result<()> { /* ... */ }
}

Public API Preservation

Before (trainers/mod.rs):

pub mod dqn;
pub use dqn::{DQNHyperparameters, DQNTrainer};

After (trainers/dqn/mod.rs):

// Internal modules
mod config;
mod agent_wrapper;
mod training_monitor;
mod trainer_core;
mod training_loop;
mod data_loading;
mod checkpointing;

// Public re-exports (maintain API compatibility)
pub use config::{
    DQNHyperparameters,
    FeatureStatistics,
    FeatureVector,
    FeatureVector51,
    EPISODE_LENGTH,
};
pub use agent_wrapper::{DQNAgentType, QValueStats};
pub use training_monitor::TrainingMonitor;
pub use trainer_core::DQNTrainer;

External code unchanged:

// Tests, hyperopt adapter, CLI - all continue working
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

Consequences

Benefits

1. Maintainability: Each module <1,000 lines, easier to understand
2. Testability: Unit test individual components (config validation, agent wrapper)
3. Readability: Clear separation of concerns
4. Safety: Smaller diffs, easier code review
5. Performance: No runtime overhead (zero-cost abstraction)

Risks

1. Breaking changes: If public API not preserved correctly
   • Mitigation: Comprehensive re-export testing
2. Build failures: Complex dependency graph
   • Mitigation: Incremental extraction with cargo check after each step
3. Test failures: Imports may need adjustment
   • Mitigation: Run full test suite after refactoring

Implementation Strategy

Phase 1: Extract Non-Dependent Modules (Safe)

1. Create dqn/config.rs (self-contained)
2. Create dqn/training_monitor.rs (depends only on config)
3. Create dqn/agent_wrapper.rs (depends on external dqn:: modules)
4. Verify: cargo check --package ml

Phase 2: Split Trainer Implementation

5. Create dqn/trainer_core.rs (struct + constructors)
6. Create dqn/data_loading.rs (data methods)
7. Create dqn/training_loop.rs (train methods)
8. Create dqn/checkpointing.rs (checkpoint methods)
9. Verify: cargo check --package ml

Phase 3: Integration

10. Create dqn/mod.rs with comprehensive re-exports
11. Delete original dqn.rs (keep .backup)
12. Update trainers/mod.rs to use pub mod dqn; (already correct)
13. Verify: cargo test --package ml --lib

Phase 4: Validation

14. Run full test suite (19 DQN tests)
15. Verify hyperopt adapter still compiles
16. Check all public API consumers

Rollback Plan

If refactoring causes issues:

rm -rf ml/src/trainers/dqn/
mv ml/src/trainers/dqn.rs.backup ml/src/trainers/dqn.rs
cargo check --package ml

Alternatives Considered

Alternative 1: Keep as Single File

Rejected: File already unmanageable, will only grow

Alternative 2: Split by Feature (Rainbow extensions, risk management, etc.)

Rejected: Cross-cutting concerns make this difficult

Alternative 3: Extract to Separate Crate

Rejected: Too aggressive, increases build complexity

References

• Original issue: SWARM_TASK swarm_1764253799645_zlazqh589
• Target: All trainer files <1,000 lines
• Current: dqn.rs (4,975), tft.rs (2,915), mamba2.rs (3,247), hyperopt/adapters/dqn.rs (3,162)