gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(ml): wire train_baseline_rl DQN to DQNTrainer — enables all GPU optimizations

Browse Source 
Replace raw DQN::new() + manual training loop in the walk-forward
training binary with DQNTrainer, which automatically activates:
- Mixed precision (BF16/F16 auto-detected from GPU)
- Dynamic batch sizing (AutoBatchSizer + HardwareBudget)
- Gradient accumulation
- Full Rainbow DQN (PER, dueling, C51, noisy nets, n-step)
- Regime-conditional Q-networks
- Portfolio tracking, Kelly sizing, entropy regularization

The walk-forward fold structure (data loading, feature extraction,
window generation, normalization) stays in the binary — only per-fold
training delegates to DQNTrainer::train_with_preloaded_data().

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-01 20:21:36 +01:00
parent 69a9808107
commit 95dfd4ad1a
1 changed files with 67 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files

304
crates/ml/examples/train_baseline_rl.rs
View File

@@ -29,7 +29,7 @@ use rand::Rng;
use serde_json::Value;
use tracing::{error, info, warn};
use ml::dqn::{DQNConfig, Experience, DQN};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
#[allow(unreachable_pub)]
mod baseline_common;
@@ -218,10 +218,31 @@ fn compute_reward(close_current: f64, close_next: f64, action_idx: u8, total_cos
// DQN Training
// ---------------------------------------------------------------------------
/// Train a DQN model on a single walk-forward fold.
/// Convert pre-aligned features and bars into the format expected by
/// `DQNTrainer::train_with_preloaded_data`: `Vec<([f64; 51], Vec<f64>)>`.
///
/// The `Vec<f64>` target contains OHLCV prices for the trainer's internal
/// reward computation and portfolio simulation.
fn features_to_trainer_format(
features: &[[f64; 51]],
bars: &[OHLCVBar],
) -> Vec<([f64; 51], Vec<f64>)> {
features
.iter()
.zip(bars.iter())
.map(|(feat, bar)| {
(*feat, vec![bar.open, bar.high, bar.low, bar.close])
})
.collect()
}
/// Train a DQN model on a single walk-forward fold using `DQNTrainer`.
///
/// Delegates all GPU optimizations (mixed precision, dynamic batching, gradient
/// accumulation, Rainbow DQN components) to the production trainer. The walk-forward
/// fold structure stays in this binary; only per-fold training delegates to `DQNTrainer`.
///
/// Returns the best validation loss achieved.
#[allow(clippy::cognitive_complexity)]
fn train_dqn_fold(
fold: usize,
train_features: &[[f64; 51]],
@@ -233,7 +254,6 @@ fn train_dqn_fold(
hp: &Option<Value>,
) -> Result<f64> {
// Caller must pre-align bars to features (skip warmup period before calling).
// If these differ, bar-index reward computation will silently use wrong prices.
debug_assert_eq!(
train_bars.len(),
train_features.len(),
@@ -250,254 +270,64 @@ fn train_dqn_fold(
);
info!(" [DQN] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len());
// Configure DQN -- hyperopt params override compiled defaults, CLI args override both
let config = DQNConfig {
state_dim: args.feature_dim,
num_actions: args.num_actions,
hidden_dims: {
let base = hp_usize(hp, "hidden_dim_base").unwrap_or(128);
if base > 128 {
vec![base, base / 2, base / 4] // GPU-scaled: 3-layer network
} else {
vec![128, 64] // Default small network
}
},
// Build DQNHyperparameters -- hyperopt JSON overrides defaults, CLI args override both.
// `..DQNHyperparameters::default()` enables all Rainbow components (PER, dueling,
// distributional C51, noisy nets, gradient accumulation, mixed precision auto-detect).
let hyperparams = DQNHyperparameters {
learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
gamma: hp_f64(hp, "gamma").unwrap_or(0.95) as f32,
batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
gamma: hp_f64(hp, "gamma").unwrap_or(0.95),
epsilon_start: 1.0,
epsilon_end: 0.05,
epsilon_decay: 0.995,
replay_buffer_capacity: hp_usize(hp, "buffer_size").unwrap_or(50_000),
batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
buffer_size: hp_usize(hp, "buffer_size").unwrap_or(50_000),
min_replay_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
target_update_freq: 200,
epochs: args.epochs,
checkpoint_frequency: 10,
hidden_dim_base: hp_usize(hp, "hidden_dim_base"),
warmup_steps: 0,
use_double_dqn: true,
use_huber_loss: true,
use_per: false,
use_dueling: false,
use_distributional: false,
use_noisy_nets: false,
use_cql: false,
use_iqn: false,
use_cvar_action_selection: false,
..DQNConfig::default()
// Early stopping is managed by DQNTrainer internally
early_stopping_enabled: true,
..DQNHyperparameters::default()
};
let mut dqn = DQN::new(config).context("Failed to create DQN model")?;
// Create DQNTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing
let mut trainer = DQNTrainer::new(hyperparams)
.context("Failed to create DQNTrainer")?;
let mut best_val_loss = f64::MAX;
let mut epochs_without_improvement = 0_usize;
let mut rng = rand::thread_rng();
let fold_str = fold.to_string();
// Convert features + bars to trainer format
let training_data = features_to_trainer_format(train_features, train_bars);
let val_data = features_to_trainer_format(val_features, val_bars);
for epoch in 0..args.epochs {
let epoch_start = std::time::Instant::now();
// --- Training pass ---
let mut epoch_loss = 0.0_f64;
let mut epoch_steps = 0_usize;
// Checkpoint callback: save best model to output directory
let output_dir_owned = output_dir.to_path_buf();
let checkpoint_callback = move |epoch: usize, data: Vec<u8>, is_best: bool| -> Result<String> {
let suffix = if is_best { "best" } else { &format!("epoch{}", epoch) };
let ckpt_path = output_dir_owned.join(format!("dqn_fold{}_{}.safetensors", fold, suffix));
std::fs::write(&ckpt_path, &data)
.with_context(|| format!("Failed to write checkpoint: {}", ckpt_path.display()))?;
info!(" [DQN] Fold {} saved checkpoint: {}", fold, ckpt_path.display());
Ok(ckpt_path.to_string_lossy().into_owned())
};
// Process training data sequentially
let n_train = train_features.len();
if n_train < 2 {
warn!(" [DQN] Fold {} -- insufficient training features ({})", fold, n_train);
break;
}
// Run the async training loop via a tokio runtime.
// The trainer handles epochs, early stopping, epsilon decay, validation, and
// all Rainbow DQN components internally.
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("Failed to create tokio runtime for DQN training")?;
let step_limit = if args.max_steps_per_epoch > 0 {
args.max_steps_per_epoch.min(n_train.saturating_sub(1))
} else {
n_train.saturating_sub(1)
};
// Features skip warmup period, so feature[i] corresponds to
// train_bars[i + warmup_offset]. We must use the same offset for
// reward computation to avoid misalignment.
let warmup_offset = train_bars.len().saturating_sub(n_train);
for i in 0..step_limit {
let Some(state_f64) = train_features.get(i) else {
continue;
};
let Some(next_f64) = train_features.get(i + 1) else {
continue;
};
let metrics = rt.block_on(
trainer.train_with_preloaded_data(training_data, val_data, checkpoint_callback)
).context("DQNTrainer training failed")?;
let state: Vec<f32> = state_f64.iter().map(|&v| v as f32).collect();
let next_state: Vec<f32> = next_f64.iter().map(|&v| v as f32).collect();
// Select action with epsilon-greedy
let action = dqn.select_action(&state)
.map(|fa| fa.to_index().min(args.num_actions.saturating_sub(1)) as u8)
.unwrap_or_else(|_| rng.gen_range(0..args.num_actions as u8));
// Compute reward from bar prices (offset by warmup to align with features)
let bar_idx = i + warmup_offset;
let close_cur = train_bars.get(bar_idx).map(|b| b.close).unwrap_or(0.0);
let close_next = train_bars.get(bar_idx + 1).map(|b| b.close).unwrap_or(close_cur);
let total_cost = args.tx_cost_bps + spread_cost_bps(close_cur, args.tick_size, args.spread_ticks);
let reward = compute_reward(close_cur, close_next, action, total_cost);
let done = i + 2 >= n_train;
// Store experience
let exp = Experience::new(state, action, reward, next_state, done);
if let Err(e) = dqn.store_experience(exp) {
// Non-fatal: buffer may not be ready
if epoch == 0 && i < 5 {
info!(" [DQN] store_experience: {}", e);
}
}
// Train step (returns (loss, grad_norm))
if let Ok((loss, _grad_norm)) = dqn.train_step(None) {
if loss.is_nan() || loss.is_infinite() {
metrics::record_nan_detected("dqn", &fold_str);
}
epoch_loss += loss as f64;
epoch_steps += 1;
}
}
let avg_train_loss = if epoch_steps > 0 {
epoch_loss / epoch_steps as f64
} else {
f64::MAX
};
// --- Validation pass ---
let val_loss = evaluate_dqn_validation(&mut dqn, val_features, val_bars, args);
// Prometheus metrics
let epoch_elapsed = epoch_start.elapsed();
metrics::set_epoch("dqn", &fold_str, (epoch + 1) as f64);
metrics::set_epoch_loss("dqn", &fold_str, avg_train_loss);
metrics::set_validation_loss("dqn", &fold_str, val_loss);
metrics::set_batches_processed("dqn", &fold_str, epoch_steps as f64);
metrics::set_iteration_seconds("dqn", &fold_str, epoch_elapsed.as_secs_f64());
if epoch_elapsed.as_secs_f64() > 0.001 {
metrics::set_batches_per_second(
"dqn",
&fold_str,
epoch_steps as f64 / epoch_elapsed.as_secs_f64(),
);
}
// Decay epsilon at epoch level
let current_eps = dqn.get_epsilon();
let new_eps = (current_eps * 0.995_f32).max(0.05);
dqn.set_epsilon(new_eps as f64);
info!(
" [DQN] Fold {} Epoch {}/{} -- train_loss={:.6} val_loss={:.6} eps={:.4}",
fold,
epoch + 1,
args.epochs,
avg_train_loss,
val_loss,
dqn.get_epsilon()
);
// Early stopping check
if val_loss < best_val_loss {
best_val_loss = val_loss;
epochs_without_improvement = 0;
// Save best checkpoint
let ckpt_path = output_dir.join(format!("dqn_fold{}_best.safetensors", fold));
let ckpt_start = std::time::Instant::now();
if let Err(e) = dqn.get_q_network_vars().save(ckpt_path.to_string_lossy().as_ref()) {
metrics::record_checkpoint_failure("dqn", &fold_str);
warn!(" [DQN] Failed to save checkpoint: {}", e);
} else {
let ckpt_size = ckpt_path.metadata().map(|m| m.len() as f64).unwrap_or(0.0);
metrics::record_checkpoint_save(
"dqn",
&fold_str,
ckpt_start.elapsed().as_secs_f64(),
ckpt_size,
);
info!(" [DQN] Saved best checkpoint: {}", ckpt_path.display());
}
} else {
epochs_without_improvement += 1;
if epochs_without_improvement >= args.patience {
info!(
" [DQN] Early stopping at epoch {} (patience {} exhausted)",
epoch + 1,
args.patience
);
break;
}
}
}
Ok(best_val_loss)
}
/// Evaluate DQN on validation features and return average loss proxy.
///
/// Since `DQN.train_step` uses replay buffer internally, we estimate validation
/// performance via average absolute reward (lower is closer to zero = better).
fn evaluate_dqn_validation(
dqn: &mut DQN,
val_features: &[[f64; 51]],
val_bars: &[OHLCVBar],
args: &Args,
) -> f64 {
// Bars must be pre-aligned to features by the caller.
debug_assert_eq!(
val_bars.len(),
val_features.len(),
"val_bars ({}) must be pre-aligned to val_features ({})",
val_bars.len(),
val_features.len(),
info!(
" [DQN] Fold {} complete -- loss={:.6} epochs_trained={} converged={}",
fold, metrics.loss, metrics.epochs_trained, metrics.convergence_achieved
);
// Run greedy inference on validation set and compute negative average reward.
// Lower values = better model (early stopping minimizes this).
let n = val_features.len();
if n < 2 {
return f64::MAX;
}
// Save epsilon and force greedy for validation.
// Note: select_action increments total_steps (side effect), but this is benign
// because warmup_steps=0 in our config, so the warmup branch is never taken.
let saved_epsilon = dqn.get_epsilon();
dqn.set_epsilon(0.0);
let warmup_offset = val_bars.len().saturating_sub(n);
let mut total_reward = 0.0_f64;
let mut count = 0_usize;
for i in 0..n.saturating_sub(1) {
let state: Vec<f32> = match val_features.get(i) {
Some(f) => f.iter().map(|&v| v as f32).collect(),
None => break,
};
// Greedy action (epsilon=0)
let action = match dqn.select_action(&state) {
Ok(fa) => fa.to_index().min(2) as u8,
Err(_) => 2, // HOLD on error
};
let bar_idx = i + warmup_offset;
let close_cur = val_bars.get(bar_idx).map(|b| b.close).unwrap_or(0.0);
let close_next = val_bars.get(bar_idx + 1).map(|b| b.close).unwrap_or(close_cur);
let total_cost = args.tx_cost_bps + spread_cost_bps(close_cur, args.tick_size, args.spread_ticks);
let reward = compute_reward(close_cur, close_next, action, total_cost);
total_reward += reward as f64;
count += 1;
}
// Restore epsilon
dqn.set_epsilon(saved_epsilon as f64);
if count > 0 {
// Return negative mean reward: lower = better model
-(total_reward / count as f64)
} else {
f64::MAX
}
Ok(metrics.loss)
}
// ---------------------------------------------------------------------------