c645e6222d
CRITICAL FINDINGS from 3-trial validation: - 85,120 gradient clipping warnings (81.6% of logs) - REGRESSION - Rainbow features DISABLED: use_dueling=false, use_distributional=false, use_noisy_nets=false - Negative Q-values confirmed: HOLD -1000 to -3250 - Performance: Sharpe 0.29 (target 0.77) Changes: - Fixed N-Step compilation (7/7 tests passing) - Fixed Distributional compilation (6/6 tests passing) - Fixed Dueling CUDA errors (10/10 tests passing) - Added TDD validation for state_dim=225 - Total: 23/23 Wave 11 tests passing (100%) Issues requiring investigation: 1. Why are Dueling/Distributional/Noisy disabled in hyperopt? 2. Why gradient explosion despite previous fixes? 3. Test coverage gaps - unit tests pass but integration fails 🤖 Generated with Claude Code Co-Authored-By: Claude <noreply@anthropic.com>
834 lines
31 KiB
Rust
834 lines
31 KiB
Rust
//! DQN Training with Full Rainbow Architecture (DEFAULT)
|
||
//!
|
||
//! This implementation includes ALL Rainbow DQN components by default:
|
||
//! - Double DQN (always enabled) - reduces overestimation bias
|
||
//! - Dueling Networks (separate value/advantage streams) - better credit assignment
|
||
//! - Prioritized Experience Replay (PER) - samples high TD-error transitions
|
||
//! - Multi-Step Returns (n=3) - balances bias vs variance
|
||
//! - Distributional RL (C51 with 51 atoms) - models full return distribution
|
||
//! - Noisy Networks (learnable exploration) - replaces epsilon-greedy
|
||
//!
|
||
//! # Usage
|
||
//!
|
||
//! ```bash
|
||
//! # Train with default parameters (100 epochs, ALL Rainbow features enabled)
|
||
//! cargo run -p ml --example train_dqn --release --features cuda
|
||
//!
|
||
//! # Disable specific components (opt-out)
|
||
//! cargo run -p ml --example train_dqn --release --features cuda -- \
|
||
//! --no-dueling --no-distributional --no-noisy-nets
|
||
//!
|
||
//! # Train vanilla DQN (all Rainbow features disabled)
|
||
//! cargo run -p ml --example train_dqn --release --features cuda -- \
|
||
//! --no-dueling --no-distributional --no-noisy-nets --n-steps 1
|
||
//!
|
||
//! # Custom hyperparameters with Rainbow
|
||
//! cargo run -p ml --example train_dqn --release --features cuda -- \
|
||
//! --epochs 500 \
|
||
//! --n-steps 5 \
|
||
//! --num-atoms 101 \
|
||
//! --dueling-hidden-dim 256
|
||
//! ```
|
||
|
||
// Use mimalloc allocator for 10-25% performance improvement
|
||
#[cfg(feature = "mimalloc-allocator")]
|
||
use mimalloc::MiMalloc;
|
||
#[cfg(feature = "mimalloc-allocator")]
|
||
#[global_allocator]
|
||
static GLOBAL: MiMalloc = MiMalloc;
|
||
|
||
use anyhow::{Context, Result};
|
||
use clap::Parser;
|
||
use std::path::PathBuf;
|
||
use std::sync::atomic::{AtomicBool, Ordering};
|
||
use std::sync::Arc;
|
||
use tokio::signal;
|
||
use tracing::{info, warn};
|
||
use tracing_subscriber::FmtSubscriber;
|
||
|
||
use ml::checkpoint::{CheckpointConfig, CheckpointManager};
|
||
use ml::data_loaders::BarSamplingMethod;
|
||
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
||
use ml::trainers::TargetUpdateMode;
|
||
|
||
/// Train DQN model on market data
|
||
#[derive(Debug, Parser)]
|
||
#[command(name = "train_dqn", about = "Train DQN model on market data")]
|
||
struct Opts {
|
||
/// Number of training epochs
|
||
#[arg(long, default_value = "100")]
|
||
epochs: usize,
|
||
|
||
/// Learning rate
|
||
/// BUG #18 FIX (Wave 16S-V17): Reduced 10× to 0.00001 to prevent Q-value explosion
|
||
/// after Bug #16 fix increased reward scale (raw portfolio values vs normalized).
|
||
/// Previous: 0.0001 caused gradient collapse (Q-values hit 1000.0 clamp, grad_norm → 0).
|
||
#[arg(long, default_value = "0.00001")]
|
||
learning_rate: f64,
|
||
|
||
/// Batch size (max 230 for RTX 3050 Ti 4GB)
|
||
/// Optimal value from hyperopt (42 trials, 2025-10-31): 32
|
||
#[arg(long, default_value = "32")]
|
||
batch_size: usize,
|
||
|
||
/// Discount factor (gamma)
|
||
/// Optimal value from hyperopt (42 trials, 2025-10-31): 0.9626
|
||
#[arg(long, default_value = "0.9626")]
|
||
gamma: f64,
|
||
|
||
/// Checkpoint save frequency (epochs)
|
||
#[arg(long, default_value = "10")]
|
||
checkpoint_frequency: usize,
|
||
|
||
/// Output directory for trained model
|
||
#[arg(long, default_value = "ml/trained_models")]
|
||
output_dir: String,
|
||
|
||
/// Data directory containing DBN files
|
||
#[arg(long, default_value = "test_data/real/databento/ml_training")]
|
||
data_dir: String,
|
||
|
||
/// Parquet file path (overrides data_dir if specified)
|
||
#[arg(long)]
|
||
parquet_file: Option<String>,
|
||
|
||
/// Verbose logging
|
||
#[arg(short, long)]
|
||
verbose: bool,
|
||
|
||
/// Enable early stopping (recommended, use --no-early-stopping to disable)
|
||
#[arg(long)]
|
||
early_stopping: bool,
|
||
|
||
/// Disable early stopping
|
||
#[arg(long)]
|
||
no_early_stopping: bool,
|
||
|
||
/// Q-value floor threshold for early stopping
|
||
#[arg(long, default_value = "0.5")]
|
||
q_value_floor: f64,
|
||
|
||
/// Minimum loss improvement percentage for plateau detection
|
||
#[arg(long, default_value = "2.0")]
|
||
min_loss_improvement: f64,
|
||
|
||
/// Plateau detection window size (epochs)
|
||
/// Optimal value from hyperopt: 5
|
||
#[arg(long, default_value = "5")]
|
||
plateau_window: usize,
|
||
|
||
/// Minimum epochs before early stopping can trigger
|
||
/// Updated to 50 to prevent premature stopping (was 10)
|
||
#[arg(long, default_value = "50")]
|
||
min_epochs_before_stopping: usize,
|
||
|
||
/// Initial exploration rate (epsilon start)
|
||
/// Updated to 0.3 for more initial exploration (was 1.0)
|
||
#[arg(long, default_value = "0.3")]
|
||
epsilon_start: f64,
|
||
|
||
/// Final exploration rate (epsilon end)
|
||
/// Updated to 0.05 to maintain exploration (was 0.01)
|
||
#[arg(long, default_value = "0.05")]
|
||
epsilon_end: f64,
|
||
|
||
/// Exploration decay rate
|
||
/// Updated to 0.995 for slower decay (was 0.9968)
|
||
#[arg(long, default_value = "0.995")]
|
||
epsilon_decay: f64,
|
||
|
||
/// Replay buffer capacity
|
||
/// Optimal value from hyperopt: 104346
|
||
#[arg(long, default_value = "104346")]
|
||
buffer_size: usize,
|
||
|
||
/// Minimum replay buffer size before training starts
|
||
/// Updated to 500 for more diverse experiences (was auto-calculated as batch_size * 2 = 64)
|
||
#[arg(long, default_value = "500")]
|
||
min_replay_size: usize,
|
||
|
||
/// Checkpoint directory (overrides output_dir for checkpoints)
|
||
#[arg(long)]
|
||
checkpoint_dir: Option<String>,
|
||
|
||
/// Alternative bar sampling method (time, tick, volume, dollar, imbalance, run)
|
||
#[arg(long, default_value = "time")]
|
||
bar_method: String,
|
||
|
||
/// HOLD penalty weight (higher = stronger penalty for holding)
|
||
#[arg(long, default_value = "0.01")]
|
||
hold_penalty_weight: f64,
|
||
|
||
/// Price movement threshold for dynamic HOLD reward (as fraction, e.g., 0.01 = 1%)
|
||
#[arg(long, default_value = "0.01")]
|
||
bar_threshold: Option<f64>,
|
||
|
||
/// Disable preprocessing (use raw non-stationary prices - NOT RECOMMENDED)
|
||
#[arg(long)]
|
||
no_preprocessing: bool,
|
||
|
||
/// Preprocessing window size (default: 50 bars)
|
||
#[arg(long, default_value = "50")]
|
||
preprocessing_window: i64,
|
||
|
||
/// Preprocessing clip sigma (default: 5.0σ)
|
||
#[arg(long, default_value = "5.0")]
|
||
preprocessing_clip_sigma: f64,
|
||
|
||
/// Warmup steps for random exploration before training (Rainbow DQN: 80K)
|
||
/// Default: ADAPTIVE (0 for <200K steps, scaled 200K-1M, 80K for >1M)
|
||
/// Set explicitly to override adaptive behavior
|
||
#[arg(long)]
|
||
warmup_steps: Option<usize>,
|
||
|
||
/// Initial capital for portfolio trading (default: $100,000, minimum: $1,000)
|
||
#[arg(long, default_value = "100000.0")]
|
||
initial_capital: f32,
|
||
|
||
/// Cash reserve requirement as a percentage of portfolio value (0.0-100.0)
|
||
#[arg(long, default_value = "0.0")]
|
||
cash_reserve_percent: f64,
|
||
|
||
/// Polyak averaging coefficient (tau) for soft target updates (default: 0.001 = soft updates)
|
||
/// Rainbow DQN standard: tau=0.001 (693-step convergence half-life)
|
||
/// Lower values = slower convergence, higher values = faster convergence
|
||
#[arg(long, default_value = "0.001")]
|
||
tau: f64,
|
||
|
||
/// Disable soft target updates (Polyak averaging) and use hard updates instead
|
||
/// Soft updates blend target network gradually with main network (default)
|
||
/// Default: soft updates (tau=0.001, gradual blending every step)
|
||
#[arg(long)]
|
||
no_soft_updates: bool,
|
||
|
||
/// Maximum absolute position size for action masking (1.0-10.0 contracts)
|
||
/// Default: 10.0 (matches hyperopt Trial #26)
|
||
#[arg(long, default_value = "10.0")]
|
||
max_position: f64,
|
||
|
||
/// Entropy regularization coefficient (0.0-0.1)
|
||
/// Controls exploration diversity via action entropy bonus
|
||
/// Default: 0.01 (1% entropy bonus from Wave 9-13)
|
||
#[arg(long, default_value = "0.01")]
|
||
entropy_coefficient: f64,
|
||
|
||
/// Transaction cost multiplier (0.5-2.0)
|
||
/// Multiplies base transaction fees: Market 0.15%, LimitMaker 0.05%, IoC 0.10%
|
||
/// Default: 1.0 (100% of base fees, production default)
|
||
#[arg(long, default_value = "1.0")]
|
||
transaction_cost_multiplier: f64,
|
||
|
||
/// Huber loss delta parameter (0.1-2.0)
|
||
/// Controls transition from quadratic (MSE) to linear (MAE) loss
|
||
#[arg(long, default_value = "1.0")]
|
||
huber_delta: f64,
|
||
|
||
// Wave 3 (Phase 2): Prioritized Experience Replay (PER) Arguments
|
||
/// Enable Prioritized Experience Replay (PER)
|
||
/// Samples high TD-error transitions more frequently for faster convergence
|
||
/// Expected improvement: 25-40% fewer epochs to reach target performance
|
||
#[arg(long)]
|
||
use_per: bool,
|
||
|
||
/// PER alpha parameter: prioritization exponent (0.0-1.0)
|
||
/// Controls how much prioritization to use (0.0 = uniform, 1.0 = full priority)
|
||
/// Rainbow DQN standard: 0.6 (balanced prioritization)
|
||
#[arg(long, default_value = "0.6")]
|
||
per_alpha: f64,
|
||
|
||
/// PER beta start: importance sampling correction exponent (0.0-1.0)
|
||
/// Anneals from beta_start to 1.0 over training to remove bias
|
||
/// Rainbow DQN standard: 0.4 (start) → 1.0 (end)
|
||
#[arg(long, default_value = "0.4")]
|
||
per_beta_start: f64,
|
||
|
||
// Wave 3 (Phase 2): Triple Barrier Method Arguments
|
||
/// Enable triple barrier method for multi-step reward labeling
|
||
#[arg(long)]
|
||
enable_triple_barrier: bool,
|
||
|
||
/// Triple barrier profit target in basis points (e.g., 100 = 1%)
|
||
#[arg(long, default_value = "100")]
|
||
profit_target: u32,
|
||
|
||
/// Triple barrier stop loss in basis points (e.g., 50 = 0.5%)
|
||
#[arg(long, default_value = "50")]
|
||
stop_loss: u32,
|
||
|
||
/// Triple barrier time limit in seconds (e.g., 3600 = 1 hour)
|
||
#[arg(long, default_value = "3600")]
|
||
time_limit: u64,
|
||
|
||
// Regime-Conditional DQN Arguments
|
||
/// Enable regime-conditional Q-network (3 heads: Trending, Ranging, Volatile)
|
||
/// When enabled, routes actions and training through regime-specific Q-networks
|
||
/// Expected improvement: +10-15% Sharpe ratio via regime-adaptive strategies
|
||
#[arg(long)]
|
||
enable_regime_qnetwork: bool,
|
||
|
||
// Wave 6.4: Rainbow DQN Opt-Out Flags (ALL ENABLED BY DEFAULT)
|
||
/// Disable dueling networks (enabled by default)
|
||
/// Dueling networks separate value and advantage streams for better credit assignment
|
||
#[arg(long)]
|
||
no_dueling: bool,
|
||
|
||
/// Disable distributional RL / C51 algorithm (enabled by default)
|
||
/// Distributional RL models full return distribution instead of scalar Q-values
|
||
#[arg(long)]
|
||
no_distributional: bool,
|
||
|
||
/// Disable noisy networks (enabled by default)
|
||
/// Noisy networks add learnable noise to network parameters for exploration
|
||
#[arg(long)]
|
||
no_noisy_nets: bool,
|
||
|
||
/// Set n_steps for multi-step returns (default: 3)
|
||
/// Higher values reduce bias but increase variance
|
||
#[arg(long, default_value = "3")]
|
||
n_steps: usize,
|
||
|
||
/// Number of atoms for C51 distributional RL (default: 51)
|
||
#[arg(long, default_value = "51")]
|
||
num_atoms: usize,
|
||
|
||
/// Minimum value for C51 distribution support (default: -1000.0)
|
||
#[arg(long, default_value = "-1000.0")]
|
||
v_min: f64,
|
||
|
||
/// Maximum value for C51 distribution support (default: 1000.0)
|
||
#[arg(long, default_value = "1000.0")]
|
||
v_max: f64,
|
||
|
||
/// Noisy network sigma init (default: 0.5, Rainbow DQN standard)
|
||
#[arg(long, default_value = "0.5")]
|
||
noisy_sigma_init: f64,
|
||
|
||
/// Dueling hidden dimension (default: 128)
|
||
#[arg(long, default_value = "128")]
|
||
dueling_hidden_dim: usize,
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn main() -> Result<()> {
|
||
// Parse CLI options
|
||
let opts = Opts::parse();
|
||
|
||
// Setup logging
|
||
let level = if opts.verbose {
|
||
tracing::Level::DEBUG
|
||
} else {
|
||
tracing::Level::INFO
|
||
};
|
||
|
||
let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
|
||
tracing::subscriber::set_global_default(subscriber)
|
||
.context("Failed to set tracing subscriber")?;
|
||
|
||
#[cfg(feature = "mimalloc-allocator")]
|
||
info!("🚀 Using mimalloc allocator for improved performance");
|
||
#[cfg(not(feature = "mimalloc-allocator"))]
|
||
info!("ℹ️ Using system allocator (consider --features mimalloc-allocator for 10-25% speedup)");
|
||
info!("🚀 Starting DQN Training");
|
||
info!("Configuration:");
|
||
info!(" • Epochs: {}", opts.epochs);
|
||
info!(" • Learning rate: {}", opts.learning_rate);
|
||
info!(" • Batch size: {}", opts.batch_size);
|
||
info!(" • Gamma: {}", opts.gamma);
|
||
info!(
|
||
" • Checkpoint frequency: {} epochs",
|
||
opts.checkpoint_frequency
|
||
);
|
||
info!(" • Output directory: {}", opts.output_dir);
|
||
info!(" • Data directory: {}", opts.data_dir);
|
||
info!(" • Bar sampling method: {}", opts.bar_method);
|
||
if let Some(threshold) = opts.bar_threshold {
|
||
info!(" • Bar threshold: {}", threshold);
|
||
}
|
||
info!(" • Epsilon start: {}", opts.epsilon_start);
|
||
info!(" • Epsilon end: {}", opts.epsilon_end);
|
||
info!(" • Epsilon decay: {}", opts.epsilon_decay);
|
||
info!(" • Buffer size: {}", opts.buffer_size);
|
||
info!(" • Min replay size: {}", opts.min_replay_size);
|
||
|
||
info!(" • Initial capital: ${:.2}", opts.initial_capital);
|
||
info!(" • Cash reserve: {}%", opts.cash_reserve_percent);
|
||
|
||
// Log target update configuration
|
||
if opts.no_soft_updates {
|
||
info!(" • Target update mode: Hard (complete replacement every 10K steps)");
|
||
info!(" • Tau (τ): {} (no blending, hard copy)", opts.tau);
|
||
} else {
|
||
info!(" • Target update mode: Soft (Polyak averaging) [DEFAULT]");
|
||
info!(
|
||
" • Tau (τ): {} (convergence half-life: {:.0} steps)",
|
||
opts.tau,
|
||
(-0.5_f64.ln()) / (-(1.0 - opts.tau).ln())
|
||
);
|
||
}
|
||
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
// ADAPTIVE WARMUP CALCULATION
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
// Estimate total training steps for adaptive warmup
|
||
let avg_steps_per_epoch = 1392; // Empirical: ES_FUT_180d.parquet has 1,392 steps/epoch
|
||
let total_steps_estimate = opts.epochs * avg_steps_per_epoch;
|
||
|
||
info!(
|
||
"📊 Training length estimate: {}K steps ({} epochs × {} steps/epoch)",
|
||
total_steps_estimate / 1000,
|
||
opts.epochs,
|
||
avg_steps_per_epoch
|
||
);
|
||
|
||
// Adaptive warmup calculation
|
||
let effective_warmup = if let Some(explicit_warmup) = opts.warmup_steps {
|
||
// User explicitly set warmup - respect it
|
||
info!(
|
||
"🎯 Using EXPLICIT warmup: {}K steps (user override)",
|
||
explicit_warmup / 1000
|
||
);
|
||
explicit_warmup
|
||
} else {
|
||
// Adaptive warmup based on training length
|
||
let adaptive_warmup = match total_steps_estimate {
|
||
0..=200_000 => {
|
||
info!("⚡ ADAPTIVE WARMUP: 0 steps (short training <200K steps)");
|
||
0
|
||
},
|
||
200_001..=500_000 => {
|
||
let warmup = total_steps_estimate / 20; // 5% warmup
|
||
info!(
|
||
"⚡ ADAPTIVE WARMUP: {}K steps (~5% of {}K total)",
|
||
warmup / 1000,
|
||
total_steps_estimate / 1000
|
||
);
|
||
warmup
|
||
},
|
||
500_001..=1_000_000 => {
|
||
let warmup = total_steps_estimate / 12; // ~8% warmup
|
||
info!(
|
||
"⚡ ADAPTIVE WARMUP: {}K steps (~8% of {}K total)",
|
||
warmup / 1000,
|
||
total_steps_estimate / 1000
|
||
);
|
||
warmup
|
||
},
|
||
_ => {
|
||
info!("⚡ ADAPTIVE WARMUP: 80K steps (Rainbow DQN standard for >1M steps)");
|
||
80_000 // Full Rainbow warmup for very long runs
|
||
},
|
||
};
|
||
adaptive_warmup
|
||
};
|
||
|
||
// Warning: warmup consuming too much of training
|
||
if effective_warmup > 0 {
|
||
let warmup_ratio = effective_warmup as f64 / total_steps_estimate as f64;
|
||
if warmup_ratio > 0.10 {
|
||
warn!(
|
||
"⚠️ Warmup period ({}K steps) is {:.1}% of total training ({}K steps)",
|
||
effective_warmup / 1000,
|
||
warmup_ratio * 100.0,
|
||
total_steps_estimate / 1000
|
||
);
|
||
warn!("⚠️ This may significantly delay learning. Consider:");
|
||
warn!(
|
||
" • Increase --epochs to lengthen training (recommended: {}+)",
|
||
(effective_warmup * 10) / avg_steps_per_epoch
|
||
);
|
||
warn!(
|
||
" • Reduce --warmup-steps to {} (10% of total)",
|
||
total_steps_estimate / 10
|
||
);
|
||
warn!(" • Set --warmup-steps 0 to disable warmup entirely");
|
||
}
|
||
|
||
info!(
|
||
" • Warmup steps: {}K ({:.1}% of training, Rainbow DQN random exploration)",
|
||
effective_warmup / 1000,
|
||
warmup_ratio * 100.0
|
||
);
|
||
} else {
|
||
info!(" • Warmup steps: 0 (disabled for short training runs)");
|
||
}
|
||
|
||
// Validate initial capital
|
||
if opts.initial_capital < 1000.0 {
|
||
eprintln!("❌ Error: initial_capital must be >= $1,000 (got: ${:.2})", opts.initial_capital);
|
||
eprintln!(" Use --initial-capital <amount> to specify a valid amount");
|
||
std::process::exit(1);
|
||
}
|
||
|
||
// Validate cash reserve percent
|
||
if !(0.0..=100.0).contains(&opts.cash_reserve_percent) {
|
||
eprintln!(
|
||
"❌ Error: cash_reserve_percent must be between 0.0 and 100.0 (got: {})",
|
||
opts.cash_reserve_percent
|
||
);
|
||
eprintln!(" Use --cash-reserve-percent <percentage> to specify a valid amount");
|
||
std::process::exit(1);
|
||
}
|
||
|
||
// Setup graceful shutdown handler for containerized environments (RunPod, Docker, K8s)
|
||
let shutdown_flag = Arc::new(AtomicBool::new(false));
|
||
let shutdown_clone = shutdown_flag.clone();
|
||
|
||
tokio::spawn(async move {
|
||
let ctrl_c = signal::ctrl_c();
|
||
|
||
#[cfg(unix)]
|
||
{
|
||
use tokio::signal::unix::{signal, SignalKind};
|
||
let mut sigterm =
|
||
signal(SignalKind::terminate()).expect("Failed to setup SIGTERM handler");
|
||
|
||
tokio::select! {
|
||
_ = ctrl_c => {
|
||
info!("🛑 Received Ctrl+C, initiating graceful shutdown...");
|
||
}
|
||
_ = sigterm.recv() => {
|
||
info!("🛑 Received SIGTERM, initiating graceful shutdown...");
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(not(unix))]
|
||
{
|
||
ctrl_c.await.expect("Failed to listen for Ctrl+C");
|
||
info!("🛑 Received Ctrl+C, initiating graceful shutdown...");
|
||
}
|
||
|
||
shutdown_clone.store(true, Ordering::Relaxed);
|
||
});
|
||
|
||
info!("✅ Graceful shutdown handler registered (Ctrl+C / SIGTERM)");
|
||
|
||
// Determine early stopping (enabled by default, unless --no-early-stopping is specified)
|
||
let early_stopping_enabled = !opts.no_early_stopping;
|
||
info!(
|
||
" • Early stopping: {}",
|
||
if early_stopping_enabled {
|
||
"enabled"
|
||
} else {
|
||
"disabled"
|
||
}
|
||
);
|
||
if early_stopping_enabled {
|
||
info!(" - Q-value floor: {}", opts.q_value_floor);
|
||
info!(" - Min loss improvement: {}%", opts.min_loss_improvement);
|
||
info!(" - Plateau window: {} epochs", opts.plateau_window);
|
||
info!(
|
||
" - Min epochs before stopping: {}",
|
||
opts.min_epochs_before_stopping
|
||
);
|
||
}
|
||
|
||
// Create output and checkpoint directories
|
||
let output_path = PathBuf::from(&opts.output_dir);
|
||
let checkpoint_path = if let Some(ref dir) = opts.checkpoint_dir {
|
||
PathBuf::from(dir)
|
||
} else {
|
||
output_path.clone()
|
||
};
|
||
|
||
if !output_path.exists() {
|
||
std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
|
||
info!("✅ Created output directory: {}", opts.output_dir);
|
||
}
|
||
|
||
if !checkpoint_path.exists() && checkpoint_path != output_path {
|
||
std::fs::create_dir_all(&checkpoint_path)
|
||
.context("Failed to create checkpoint directory")?;
|
||
info!(
|
||
"✅ Created checkpoint directory: {}",
|
||
checkpoint_path.display()
|
||
);
|
||
}
|
||
|
||
if opts.checkpoint_dir.is_some() {
|
||
info!(" • Checkpoint directory: {}", checkpoint_path.display());
|
||
}
|
||
|
||
// Configure DQN hyperparameters with optimal values from hyperopt (42 trials, 2025-10-31)
|
||
let hyperparams = DQNHyperparameters {
|
||
learning_rate: opts.learning_rate,
|
||
batch_size: opts.batch_size,
|
||
gamma: opts.gamma,
|
||
epsilon_start: opts.epsilon_start,
|
||
epsilon_end: opts.epsilon_end,
|
||
epsilon_decay: opts.epsilon_decay,
|
||
buffer_size: opts.buffer_size,
|
||
min_replay_size: opts.min_replay_size, // Configurable min replay size
|
||
epochs: opts.epochs,
|
||
checkpoint_frequency: opts.checkpoint_frequency,
|
||
early_stopping_enabled,
|
||
q_value_floor: opts.q_value_floor,
|
||
min_loss_improvement_pct: opts.min_loss_improvement,
|
||
plateau_window: opts.plateau_window,
|
||
min_epochs_before_stopping: opts.min_epochs_before_stopping, // NOW CONFIGURABLE!
|
||
hold_penalty: -0.001,
|
||
// Bug #3 fix: Enable gradient clipping to prevent Q-value collapse
|
||
gradient_clip_norm: Some(10.0), // Conservative clipping at max_norm=10.0
|
||
// Bug #3 fix: Enable Huber loss for robustness to outliers
|
||
use_huber_loss: true,
|
||
huber_delta: opts.huber_delta,
|
||
// Enable Double DQN to reduce overestimation bias
|
||
use_double_dqn: true,
|
||
// HOLD penalty weight (Bug #3 fix)
|
||
hold_penalty_weight: opts.hold_penalty_weight, // Configurable via CLI
|
||
// Price movement threshold for HOLD penalty (2% price movement)
|
||
movement_threshold: 0.02,
|
||
// Wave 14 Agent 32: Preprocessing configuration
|
||
enable_preprocessing: !opts.no_preprocessing, // Enabled by default, disable with --no-preprocessing
|
||
preprocessing_window: opts.preprocessing_window,
|
||
preprocessing_clip_sigma: opts.preprocessing_clip_sigma,
|
||
|
||
// Target update configuration (soft updates enabled by default)
|
||
tau: opts.tau, // CLI-configurable (default: 0.001 = soft updates)
|
||
target_update_mode: if opts.no_soft_updates {
|
||
TargetUpdateMode::Hard
|
||
} else {
|
||
TargetUpdateMode::Soft
|
||
},
|
||
|
||
// P2-B Enhancement: Cash reserve requirement
|
||
cash_reserve_percent: opts.cash_reserve_percent, // Configurable via CLI
|
||
target_update_frequency: 10000, // Hard update frequency (every 10K steps)
|
||
|
||
// Rainbow DQN warmup
|
||
warmup_steps: effective_warmup, // Adaptive warmup (0 for <200K, scaled 200K-1M, 80K for >1M)
|
||
|
||
// P2-A Enhancement
|
||
initial_capital: opts.initial_capital,
|
||
|
||
// WAVE 16S: Adaptive Risk Management (ALL ENABLED BY DEFAULT)
|
||
enable_kelly_sizing: true,
|
||
enable_volatility_epsilon: true,
|
||
enable_risk_adjusted_rewards: true,
|
||
kelly_fractional: 0.5,
|
||
kelly_max_fraction: 0.25,
|
||
kelly_min_trades: 20,
|
||
volatility_window: 20,
|
||
|
||
// WAVE 35: Advanced Features
|
||
enable_regime_qnetwork: opts.enable_regime_qnetwork, // CLI-configurable via --enable-regime-qnetwork
|
||
enable_compliance: true, // Enabled by default
|
||
|
||
// WAVE 16: Core Risk Management (ALL ENABLED BY DEFAULT)
|
||
enable_drawdown_monitoring: true,
|
||
enable_position_limits: true,
|
||
enable_circuit_breaker: true,
|
||
|
||
// Wave 16 Portfolio Features (ALL ENABLED BY DEFAULT)
|
||
enable_action_masking: true,
|
||
enable_entropy_regularization: true,
|
||
enable_stress_testing: true,
|
||
max_position_absolute: opts.max_position,
|
||
|
||
// Wave 17: Hyperopt 9D Search Space Extensions
|
||
entropy_coefficient: Some(opts.entropy_coefficient),
|
||
transaction_cost_multiplier: opts.transaction_cost_multiplier,
|
||
|
||
// Wave 3 (Phase 2): Prioritized Experience Replay Configuration
|
||
use_per: opts.use_per,
|
||
per_alpha: opts.per_alpha,
|
||
per_beta_start: opts.per_beta_start,
|
||
|
||
// Wave 3 (Phase 2): Triple Barrier Method Configuration
|
||
enable_triple_barrier: opts.enable_triple_barrier,
|
||
triple_barrier_profit_target_bps: opts.profit_target,
|
||
triple_barrier_stop_loss_bps: opts.stop_loss,
|
||
triple_barrier_max_holding_seconds: opts.time_limit,
|
||
|
||
// Wave 6.4: Rainbow DQN Features (ALL ENABLED BY DEFAULT)
|
||
// Wave 2.1: Dueling Networks
|
||
use_dueling: !opts.no_dueling, // Default: enabled, opt-out with --no-dueling
|
||
dueling_hidden_dim: opts.dueling_hidden_dim, // Default: 128
|
||
|
||
// Wave 2.2: Multi-Step Returns
|
||
n_steps: opts.n_steps, // Default: 3
|
||
|
||
// Wave 2.3: Distributional RL (C51)
|
||
use_distributional: !opts.no_distributional, // Default: enabled, opt-out with --no-distributional
|
||
num_atoms: opts.num_atoms, // Default: 51
|
||
v_min: opts.v_min, // Default: -1000.0
|
||
v_max: opts.v_max, // Default: 1000.0
|
||
|
||
// Wave 2.4: Noisy Networks
|
||
use_noisy_nets: !opts.no_noisy_nets, // Default: enabled, opt-out with --no-noisy-nets
|
||
noisy_sigma_init: opts.noisy_sigma_init, // Default: 0.5
|
||
};
|
||
|
||
// Configure alternative bar sampling (Wave B)
|
||
let bar_sampling = match opts.bar_method.as_str() {
|
||
"tick" => BarSamplingMethod::TickBars(opts.bar_threshold.unwrap_or(100.0) as usize),
|
||
"volume" => BarSamplingMethod::VolumeBars(opts.bar_threshold.unwrap_or(10000.0)),
|
||
"dollar" => BarSamplingMethod::DollarBars(opts.bar_threshold.unwrap_or(2_000_000.0)),
|
||
"imbalance" => BarSamplingMethod::ImbalanceBars(opts.bar_threshold.unwrap_or(1000.0)),
|
||
"run" => BarSamplingMethod::RunBars(opts.bar_threshold.unwrap_or(50.0) as usize),
|
||
_ => BarSamplingMethod::TimeBars,
|
||
};
|
||
|
||
info!("✅ Bar sampling configured: {:?}", bar_sampling);
|
||
|
||
// Create DQN trainer
|
||
let mut trainer = DQNTrainer::new(hyperparams).context("Failed to create DQN trainer")?;
|
||
|
||
// Note: DQN trainer will need to accept bar_sampling parameter
|
||
// This requires updating DQNTrainer to use DbnSequenceLoader
|
||
info!("✅ DQN trainer initialized");
|
||
|
||
// Setup checkpoint manager
|
||
let checkpoint_config = CheckpointConfig {
|
||
base_dir: output_path.clone(),
|
||
max_checkpoints_per_model: 10,
|
||
auto_cleanup: true,
|
||
validate_checksums: true,
|
||
..Default::default()
|
||
};
|
||
|
||
let _checkpoint_manager =
|
||
CheckpointManager::new(checkpoint_config).context("Failed to create checkpoint manager")?;
|
||
|
||
info!("✅ Checkpoint manager initialized (max 10 checkpoints, auto-cleanup enabled)");
|
||
|
||
// Create checkpoint callback with interruption handling
|
||
let checkpoint_dir_for_callback = opts
|
||
.checkpoint_dir
|
||
.clone()
|
||
.unwrap_or_else(|| opts.output_dir.clone());
|
||
let shutdown_check = shutdown_flag.clone();
|
||
|
||
let checkpoint_callback =
|
||
move |epoch: usize, model_data: Vec<u8>, is_best: bool| -> Result<String> {
|
||
// Check if shutdown was requested
|
||
let interrupted = shutdown_check.load(Ordering::Relaxed);
|
||
|
||
let filename = if is_best {
|
||
// Best model checkpoint (overwrites previous best)
|
||
"dqn_best_model.safetensors".to_string()
|
||
} else if interrupted {
|
||
format!("dqn_interrupted_epoch{}.safetensors", epoch)
|
||
} else {
|
||
// Periodic checkpoint
|
||
format!("dqn_epoch_{}.safetensors", epoch)
|
||
};
|
||
|
||
let checkpoint_path = PathBuf::from(&checkpoint_dir_for_callback).join(filename);
|
||
|
||
// Save checkpoint to disk
|
||
std::fs::write(&checkpoint_path, &model_data)
|
||
.context(format!("Failed to save checkpoint: {:?}", checkpoint_path))?;
|
||
|
||
let checkpoint_type = if is_best {
|
||
"🎉 BEST"
|
||
} else if interrupted {
|
||
"⚠️ INTERRUPTED"
|
||
} else {
|
||
"💾 PERIODIC"
|
||
};
|
||
|
||
info!(
|
||
"{} Checkpoint saved: {} ({} bytes)",
|
||
checkpoint_type,
|
||
checkpoint_path.display(),
|
||
model_data.len()
|
||
);
|
||
|
||
Ok(checkpoint_path.to_string_lossy().to_string())
|
||
};
|
||
|
||
// Train the model
|
||
info!("\n🏋️ Starting training...\n");
|
||
let start_time = std::time::Instant::now();
|
||
|
||
let metrics = if let Some(ref parquet_path) = opts.parquet_file {
|
||
info!("Using Parquet file: {}", parquet_path);
|
||
trainer
|
||
.train_from_parquet(parquet_path, checkpoint_callback)
|
||
.await
|
||
.context("Training from Parquet failed")?
|
||
} else {
|
||
info!("Using DBN directory: {}", opts.data_dir);
|
||
trainer
|
||
.train(&opts.data_dir, checkpoint_callback)
|
||
.await
|
||
.context("Training failed")?
|
||
};
|
||
|
||
let training_duration = start_time.elapsed();
|
||
|
||
// Check if training was interrupted
|
||
if shutdown_flag.load(Ordering::Relaxed) {
|
||
info!("\n⚠️ Training was interrupted by shutdown signal");
|
||
info!("💾 Interrupted checkpoint saved, safe to terminate");
|
||
info!("📊 Partial training metrics:");
|
||
info!(" • Epochs completed: {}", metrics.epochs_trained);
|
||
info!(
|
||
" • Training time: {:.1}s ({:.1} min)",
|
||
metrics.training_time_seconds,
|
||
metrics.training_time_seconds / 60.0
|
||
);
|
||
return Ok(());
|
||
}
|
||
|
||
// Print final metrics
|
||
info!("\n✅ Training completed successfully!");
|
||
info!("\n📊 Final Metrics:");
|
||
info!(" • Final loss: {:.6}", metrics.loss);
|
||
info!(" • Epochs trained: {}", metrics.epochs_trained);
|
||
info!(
|
||
" • Training time: {:.1}s ({:.1} min)",
|
||
metrics.training_time_seconds,
|
||
metrics.training_time_seconds / 60.0
|
||
);
|
||
info!(
|
||
" • Actual elapsed time: {:.1}s (includes data loading + overhead)",
|
||
training_duration.as_secs_f64()
|
||
);
|
||
info!(
|
||
" • Convergence: {}",
|
||
if metrics.convergence_achieved {
|
||
"✅ Yes"
|
||
} else {
|
||
"❌ No"
|
||
}
|
||
);
|
||
|
||
// Additional metrics from training
|
||
if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
|
||
info!(" • Average Q-value: {:.4}", avg_q_value);
|
||
}
|
||
if let Some(final_epsilon) = metrics.additional_metrics.get("final_epsilon") {
|
||
info!(" • Final epsilon: {:.4}", final_epsilon);
|
||
}
|
||
if let Some(grad_norm) = metrics.additional_metrics.get("avg_gradient_norm") {
|
||
info!(" • Average gradient norm: {:.6}", grad_norm);
|
||
}
|
||
|
||
// Save final model
|
||
let final_model_path = output_path.join(format!("dqn_final_epoch{}.safetensors", opts.epochs));
|
||
info!("\n💾 Saving final model to: {}", final_model_path.display());
|
||
|
||
// Get final model state
|
||
let final_checkpoint_data = trainer
|
||
.serialize_model()
|
||
.await
|
||
.context("Failed to serialize final model")?;
|
||
|
||
std::fs::write(&final_model_path, &final_checkpoint_data)
|
||
.context("Failed to save final model")?;
|
||
|
||
info!(
|
||
"✅ Final model saved: {} ({} bytes)",
|
||
final_model_path.display(),
|
||
final_checkpoint_data.len()
|
||
);
|
||
|
||
info!("\n🎉 DQN training complete!");
|
||
info!("📁 Model files saved to: {}", opts.output_dir);
|
||
|
||
Ok(())
|
||
}
|