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

feat(cuda): fused DQN training kernel + trainer split

Browse Source 
Replace 7k-line monolithic trainer.rs with modular trainer/ directory:
  action.rs, constructor.rs, metrics.rs, mod.rs, state.rs,
  tests.rs, training_loop.rs, train_step.rs (6048 lines total)

New fused CUDA training pipeline:
  - dqn_training_kernel.cu: single-kernel forward+loss+backward
  - gpu_dqn_trainer.rs: host-side fused training orchestration
  - fused_training.rs: Rust-side fused training integration

Eliminates per-step CPU↔GPU synchronization in DQN training loop.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-15 23:50:29 +01:00
parent 6b3a00a0ac
commit 3ac51679a1
14 changed files with 9399 additions and 7125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1281
crates/ml/src/cuda_pipeline/dqn_training_kernel.cu Normal file
View File

File diff suppressed because it is too large Load Diff

1595
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs Normal file
View File

File diff suppressed because it is too large Load Diff

34
crates/ml/src/trainers/dqn/config.rs
View File

@@ -67,7 +67,7 @@ impl DQNAgentType {
/// Batch greedy action selection — single forward pass per architecture head.
///
/// Dispatches to the underlying DQN or RegimeConditionalDQN batch method.
pub fn batch_greedy_actions(&self, states: &Tensor) -> Result<Vec<usize>, MLError> {
pub fn batch_greedy_actions(&self, states: &Tensor) -> Result<Tensor, MLError> {
match self {
Self::Standard(agent) => agent.batch_greedy_actions(states),
Self::RegimeConditional(agent) => agent.batch_greedy_actions(states),
@@ -123,7 +123,7 @@ impl DQNAgentType {
&self,
states: &Tensor,
temperature: f64,
) -> Result<Vec<usize>, MLError> {
) -> Result<Tensor, MLError> {
match self {
Self::Standard(agent) => agent.batch_softmax_actions(states, temperature),
Self::RegimeConditional(agent) => agent.batch_softmax_actions(states, temperature),
@@ -135,7 +135,7 @@ impl DQNAgentType {
&self,
states: &Tensor,
temperature: f64,
) -> Result<Vec<usize>, MLError> {
) -> Result<Tensor, MLError> {
match self {
Self::Standard(agent) => {
agent.batch_hierarchical_softmax_actions(states, temperature)
@@ -592,6 +592,34 @@ impl DQNAgentType {
self.memory().step();
}
/// Post-step bookkeeping for the fused CUDA training path.
///
/// Delegates to `DQN::fused_post_step` -- increments training_steps, updates
/// PER priorities, steps beta annealing, and runs Polyak target update.
/// Only supported for Standard agents (not RegimeConditional).
pub fn fused_post_step(&mut self, td_errors: &[f32], indices: &[usize]) -> Result<(), crate::MLError> {
match self {
Self::Standard(agent) => agent.fused_post_step(td_errors, indices),
Self::RegimeConditional(_) => Err(crate::MLError::ModelError(
"Fused CUDA training not supported for RegimeConditional agent".into(),
)),
}
}
/// Post-step bookkeeping WITHOUT target EMA (for GPU-native EMA path).
///
/// Delegates to `DQN::fused_post_step_no_ema` -- increments training_steps,
/// updates PER priorities, steps beta annealing. Target EMA is done by the
/// GPU EMA kernel in `FusedTrainingCtx::run_full_step`.
pub fn fused_post_step_no_ema(&mut self, td_errors: &[f32], indices: &[usize]) -> Result<(), crate::MLError> {
match self {
Self::Standard(agent) => agent.fused_post_step_no_ema(td_errors, indices),
Self::RegimeConditional(_) => Err(crate::MLError::ModelError(
"Fused CUDA training not supported for RegimeConditional agent".into(),
)),
}
}
/// Flush GPU-accumulated max priority to CPU (single scalar readback per epoch).
///
/// Call once at epoch boundary after all `update_priorities_gpu` calls in the

443
crates/ml/src/trainers/dqn/fused_training.rs Normal file
View File

@@ -0,0 +1,443 @@
//! Fused CUDA Training Module
//!
//! High-performance H100-optimized training path that replaces 2,100+ Candle kernel
//! dispatches per batch with 3 fused CUDA kernels captured in a CUDA Graph:
//!
//! 1. **Forward + Loss kernel** -- shared trunk, 3 branching advantage heads, C51 distributional loss
//! 2. **Backward kernel** -- full backprop through all layers with gradient clipping
//! 3. **Adam optimizer kernel** -- fused parameter update with weight decay
//!
//! After the first batch, all 3 kernels are captured into a CUDA Graph and replayed
//! on subsequent steps with zero launch overhead.
//!
//! ## Weight Sync Architecture (GPU-native EMA)
//!
//! Device pointers in `CudaSlice` buffers are stable across in-place updates, so
//! the CUDA Graph remains valid across Polyak EMA target updates:
//!
//! 1. **After fused step**: online `CudaSlice` weights are already updated by the Adam kernel
//! 2. **GPU Polyak EMA**: `target[i] = (1-tau)*target[i] + tau*online[i]` via EMA kernel
//! 3. **No reverse/forward sync per step** -- VarMap sync deferred to epoch boundary
//!
//! This eliminates ~120 D2D copies + 120 Candle flatten/contiguous ops per step
//! (~13 GB wasted PCIe per epoch).
use std::sync::Arc;
use anyhow::Result;
use candle_core::{Device, Tensor};
use tracing::info;
use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer};
use crate::cuda_pipeline::gpu_weights::{
self, BranchingWeightSet, DuelingWeightSet,
};
use crate::dqn::dqn::GpuTrainResult;
use crate::dqn::mixed_precision::training_dtype;
use crate::dqn::replay_buffer_type::BatchSample;
use super::config::DQNAgentType;
use super::DQNHyperparameters;
/// Fused CUDA training context -- owns the `GpuDqnTrainer` and extracted weight sets.
///
/// Weight sets are extracted from the Candle VarMap at initialization. Per-step
/// training operates entirely on GPU `CudaSlice` buffers: the Adam kernel updates
/// online weights, then the EMA kernel blends online into target weights in-place.
///
/// VarMap sync is deferred to epoch boundary to avoid 120 D2D copies per step.
/// Device pointers are stable across in-place updates, so the CUDA Graph stays valid.
pub(crate) struct FusedTrainingCtx {
trainer: GpuDqnTrainer,
online_dueling: DuelingWeightSet,
online_branching: BranchingWeightSet,
target_dueling: DuelingWeightSet,
target_branching: BranchingWeightSet,
stream: Arc<candle_core::cuda_backend::cudarc::driver::CudaStream>,
/// Batch size at creation time -- must match `current_batch_size` to reuse CUDA Graph.
batch_size: usize,
/// Steps since last VarMap sync (deferred to epoch boundary).
steps_since_varmap_sync: usize,
}
impl FusedTrainingCtx {
/// Create a new fused training context.
///
/// Extracts online + target weight sets from the DQN's Candle VarMaps into
/// `CudaSlice` buffers, compiles the 3 fused kernels + EMA kernel, and
/// allocates all pre-allocated buffers for CUDA Graph capture.
///
/// Only valid when `device.is_cuda() && agent.is_using_branching()`.
pub(crate) fn new(
device: &Device,
agent: &DQNAgentType,
hyperparams: &DQNHyperparameters,
batch_size: usize,
) -> Result<Self> {
let dqn = match agent {
DQNAgentType::Standard(d) => d,
DQNAgentType::RegimeConditional(_) => {
return Err(anyhow::anyhow!(
"Fused CUDA training requires Standard DQN agent (not RegimeConditional)"
));
}
};
let branching_net = dqn.branching_q_network.as_ref().ok_or_else(|| {
anyhow::anyhow!("Fused CUDA training requires branching Q-network")
})?;
let branching_target = dqn.branching_target_network.as_ref().ok_or_else(|| {
anyhow::anyhow!("Fused CUDA training requires branching target network")
})?;
// Get CudaStream from device
let cuda_dev = match device {
Device::Cuda(dev) => dev,
_ => return Err(anyhow::anyhow!("Fused CUDA training requires CUDA device")),
};
let stream = cuda_dev.cuda_stream();
// Build config from DQN network dimensions
let (shared_h1, shared_h2, value_h, adv_h) = agent.network_dims();
let config = GpuDqnTrainConfig {
state_dim: dqn.config.state_dim,
shared_h1,
shared_h2,
value_h,
adv_h,
num_atoms: dqn.config.num_atoms,
v_min: dqn.config.v_min,
v_max: dqn.config.v_max,
branch_0_size: dqn.config.num_actions,
branch_1_size: dqn.config.num_order_types,
branch_2_size: dqn.config.num_urgency_levels,
batch_size,
gamma: hyperparams.gamma as f32,
lr: hyperparams.learning_rate as f32,
beta1: 0.9,
beta2: 0.999,
epsilon: 1e-8,
weight_decay: 1e-5,
max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32,
};
// Extract weight sets from VarMaps (online + target)
let online_vars = branching_net.vars();
let target_vars = branching_target.vars();
let online_dueling =
gpu_weights::extract_dueling_weights_branching(online_vars, &stream)
.map_err(|e| anyhow::anyhow!("Extract online dueling weights: {e}"))?;
let online_branching =
gpu_weights::extract_branching_weights(online_vars, &stream)
.map_err(|e| anyhow::anyhow!("Extract online branching weights: {e}"))?;
let target_dueling =
gpu_weights::extract_dueling_weights_branching(target_vars, &stream)
.map_err(|e| anyhow::anyhow!("Extract target dueling weights: {e}"))?;
let target_branching =
gpu_weights::extract_branching_weights(target_vars, &stream)
.map_err(|e| anyhow::anyhow!("Extract target branching weights: {e}"))?;
// Create the fused trainer (compiles kernels, allocates buffers)
let trainer = GpuDqnTrainer::new(stream.clone(), config)
.map_err(|e| anyhow::anyhow!("GpuDqnTrainer init: {e}"))?;
info!(
batch_size,
"Fused CUDA training initialized: 4 kernels + EMA compiled, \
~291K params, CUDA Graph will capture on first step"
);
Ok(Self {
trainer,
online_dueling,
online_branching,
target_dueling,
target_branching,
stream,
batch_size,
steps_since_varmap_sync: 0,
})
}
/// Batch size this context was created for.
pub(crate) fn batch_size(&self) -> usize {
self.batch_size
}
/// Steps since last VarMap sync.
pub(crate) fn steps_since_varmap_sync(&self) -> usize {
self.steps_since_varmap_sync
}
/// Run one full fused training step.
///
/// Executes the complete training cycle with GPU-native EMA:
/// 1. Extract flat arrays from `BatchSample`
/// 2. Forward + loss + backward + Adam via fused CUDA kernels (or CUDA Graph replay)
/// 3. GPU EMA: `target[i] = (1-tau)*target[i] + tau*online[i]` (20 kernel launches)
/// 4. PER priority update + beta annealing + training_steps increment
/// 5. Return `GpuTrainResult` with GPU scalar tensors for monitoring
///
/// No VarMap sync per step -- deferred to epoch boundary via `sync_to_varmap()`.
/// This eliminates ~120 D2D copies + 120 Candle ops per step (~13 GB/epoch saved).
pub(crate) fn run_full_step(
&mut self,
batch: &BatchSample,
agent: &mut DQNAgentType,
device: &Device,
) -> Result<GpuTrainResult> {
let state_dim = agent.get_state_dim();
let (states, next_states, actions, rewards, dones, is_weights) =
extract_batch_arrays(batch, state_dim);
// Step 1: Fused forward + loss + backward + Adam
// The Adam kernel updates online CudaSlice weights in-place.
// Device pointers are stable -- CUDA Graph stays valid.
let fused_result = self.trainer.train_step(
&states, &next_states, &actions, &rewards, &dones, &is_weights,
&self.online_dueling, &self.online_branching,
&self.target_dueling, &self.target_branching,
).map_err(|e| anyhow::anyhow!("Fused train_step: {e}"))?;
// Step 2: GPU-native Polyak EMA target update
// Computes cosine-annealed tau and applies target[i] = (1-tau)*target[i] + tau*online[i]
// entirely on GPU. No VarMap round-trip.
{
let dqn = agent.as_standard_mut().ok_or_else(|| {
anyhow::anyhow!("Fused training requires Standard DQN agent")
})?;
if dqn.config.use_soft_updates {
let tau = compute_cosine_annealed_tau(
dqn.get_training_steps(),
dqn.config.tau,
dqn.config.tau_final,
dqn.config.tau_anneal_steps,
);
self.trainer.target_ema_update(
&self.online_dueling, &self.online_branching,
&self.target_dueling, &self.target_branching,
tau as f32,
).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
}
}
// Step 3: PER priority update + training_steps++ + beta annealing
// (no target EMA -- already done by GPU kernel above)
// td_errors from fused kernel are already f32 on CPU.
agent.fused_post_step_no_ema(&fused_result.td_errors, &batch.indices)
.map_err(|e| anyhow::anyhow!("Fused post_step_no_ema: {e}"))?;
self.steps_since_varmap_sync += 1;
// Step 4: Create GPU scalar tensors from fused results for monitoring code
Ok(GpuTrainResult {
loss_gpu: Tensor::new(fused_result.total_loss, device)
.map_err(|e| anyhow::anyhow!("Fused loss->Tensor: {e}"))?,
grad_norm_gpu: Tensor::new(fused_result.grad_norm, device)
.map_err(|e| anyhow::anyhow!("Fused grad_norm->Tensor: {e}"))?,
})
}
/// Sync CudaSlice weights back to VarMap (deferred -- called at epoch boundary).
///
/// Copies:
/// - Online `CudaSlice -> VarMap` (reverse sync: Adam-updated weights)
/// - Target `CudaSlice -> VarMap` (reverse sync: EMA-blended weights)
///
/// This keeps the Candle VarMap in sync for checkpointing, Q-value estimation
/// via `forward()`, and any non-fused code paths.
pub(crate) fn sync_to_varmap(&mut self, agent: &mut DQNAgentType) -> Result<()> {
let dqn = agent.as_standard_mut().ok_or_else(|| {
anyhow::anyhow!("VarMap sync requires Standard DQN agent")
})?;
// Reverse sync online weights: CudaSlice -> VarMap
let online_vars = dqn.branching_q_network.as_ref()
.ok_or_else(|| anyhow::anyhow!("Missing branching Q-network"))?
.vars();
gpu_weights::reverse_sync_dueling_weights_branching(
online_vars, &self.online_dueling, &self.stream,
).map_err(|e| anyhow::anyhow!("Reverse sync online dueling: {e}"))?;
gpu_weights::reverse_sync_branching_weights(
online_vars, &self.online_branching, &self.stream,
).map_err(|e| anyhow::anyhow!("Reverse sync online branching: {e}"))?;
// Reverse sync target weights: CudaSlice -> VarMap
let target_vars = dqn.branching_target_network.as_ref()
.ok_or_else(|| anyhow::anyhow!("Missing branching target network"))?
.vars();
gpu_weights::reverse_sync_dueling_weights_branching(
target_vars, &self.target_dueling, &self.stream,
).map_err(|e| anyhow::anyhow!("Reverse sync target dueling: {e}"))?;
gpu_weights::reverse_sync_branching_weights(
target_vars, &self.target_branching, &self.stream,
).map_err(|e| anyhow::anyhow!("Reverse sync target branching: {e}"))?;
self.steps_since_varmap_sync = 0;
Ok(())
}
}
/// Cosine-annealed Polyak EMA coefficient (BYOL/MoCo v3 schedule).
///
/// `tau(t) = tau_final - (tau_final - tau_base) * (cos(pi*t/T) + 1) / 2`
///
/// - Early training: `tau ~ tau_base` (fast target adaptation)
/// - Late training: `tau ~ tau_final` (stability)
/// - `anneal_steps == 0`: fixed `tau_base` (no annealing)
fn compute_cosine_annealed_tau(
training_steps: u64,
tau_base: f64,
tau_final: f64,
anneal_steps: u64,
) -> f64 {
if anneal_steps > 0 {
let progress = (training_steps as f64 / anneal_steps as f64).min(1.0);
let cosine_factor = (std::f64::consts::PI * progress).cos();
tau_final - (tau_final - tau_base) * (cosine_factor + 1.0) / 2.0
} else {
tau_base
}
}
/// Extract flat `f32`/`i32` arrays from a `BatchSample` for the fused CUDA trainer.
///
/// Returns `(states, next_states, actions, rewards, dones, is_weights)`.
/// Rewards are converted from fixed-point (`i32 / 1_000_000`) to `f32`.
fn extract_batch_arrays(
batch: &BatchSample,
state_dim: usize,
) -> (Vec<f32>, Vec<f32>, Vec<i32>, Vec<f32>, Vec<f32>, Vec<f32>) {
let b = batch.experiences.len();
let mut states = Vec::with_capacity(b * state_dim);
let mut next_states = Vec::with_capacity(b * state_dim);
let mut actions = Vec::with_capacity(b);
let mut rewards = Vec::with_capacity(b);
let mut dones = Vec::with_capacity(b);
for exp in &batch.experiences {
states.extend_from_slice(&exp.state);
next_states.extend_from_slice(&exp.next_state);
actions.push(exp.action as i32);
rewards.push(exp.reward_f32());
dones.push(if exp.done { 1.0_f32 } else { 0.0_f32 });
}
(states, next_states, actions, rewards, dones, batch.weights.clone())
}
/// GPU Q-value estimation -- called every 50 training steps for monitoring.
///
/// Samples 10 experiences from the replay buffer, runs a forward pass through
/// the branching Q-network, and uses GPU-native reduction kernels for:
/// - Q-value divergence check (early stopping on runaway Q-values)
/// - Q-value statistics (min/max/mean/variance)
/// - Welford running mean accumulation (zero CPU sync)
///
/// Shared by both `train_step_single_batch` and `train_step_with_accumulation`.
///
/// Returns `Ok(avg_q)` or an error if GPU Q-value accumulation fails.
#[allow(clippy::indexing_slicing)]
pub(super) fn gpu_q_value_estimation(
agent: &mut DQNAgentType,
training_guard: &mut crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard,
device: &Device,
) -> Result<f64> {
use candle_core::IndexOp;
let buffer = agent.memory();
if buffer.len() == 0 {
return Err(anyhow::anyhow!(
"GPU Q-value estimation requires non-empty replay buffer"
));
}
let sample_size = buffer.len().min(10);
let batch_sample = buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?;
let state_dim = agent.get_state_dim();
let mut batch_tensor_opt: Option<Tensor> = None;
// GPU PER path: use gpu_batch.states directly
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = Some(
gpu.states
.to_dtype(training_dtype(agent.device()))
.map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?,
);
}
// CPU fallback: build tensor from experiences
if batch_tensor_opt.is_none() {
let mut state_data = Vec::with_capacity(sample_size * state_dim);
for exp in &batch_sample.experiences {
state_data.extend_from_slice(&exp.state);
}
if !state_data.is_empty() {
let tensor = Tensor::from_vec(
state_data,
(sample_size, state_dim),
device,
)
.map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))?
.to_dtype(training_dtype(device))
.map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?;
batch_tensor_opt = Some(tensor);
}
}
let batch_tensor = batch_tensor_opt.ok_or_else(|| {
anyhow::anyhow!("GPU Q-value estimation: no tensor built (empty batch?)")
})?;
// Suppress forward() monitoring to avoid to_vec2 GPU->CPU sync
agent.set_training_forward_active(true);
let batch_q_values = agent
.forward(&batch_tensor)
.map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?;
agent.set_training_forward_active(false);
let num_actions = batch_q_values.dims().get(1).copied().unwrap_or(5);
// Divergence check on first sample
let first_q = batch_q_values
.i(0)
.map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?;
let div_result = training_guard
.qvalue_divergence(&first_q, num_actions, 10000.0)
.map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?;
agent
.log_q_values_from_stats(
div_result.q_min,
div_result.q_max,
div_result.q_mean,
div_result.q_variance,
num_actions,
)
.map_err(|e| {
tracing::info!("Early stopping (Q-value divergence): {}", e);
anyhow::anyhow!("Early stopping: {}", e)
})?;
// Batch average via GPU reduction (one-step delay due to double-buffering)
let stats = training_guard
.qvalue_stats(&batch_q_values, sample_size, num_actions)
.map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?;
let cached_avg_q = stats.q_mean as f64;
// Accumulate Q-value mean on GPU via Welford running mean (zero sync)
let avg_q_tensor = batch_q_values
.max(1)
.map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))?
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?;
training_guard
.accumulate_q_value(&avg_q_tensor)
.map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?;
Ok(cached_avg_q)
}

1
crates/ml/src/trainers/dqn/mod.rs
View File

@@ -23,6 +23,7 @@ mod data_loading;
mod early_stopping;
pub(crate) mod financials;
mod features;
mod fused_training;
pub mod lr_scheduler;
mod monitoring;
mod risk;

7122
crates/ml/src/trainers/dqn/trainer.rs
View File

File diff suppressed because it is too large Load Diff

396
crates/ml/src/trainers/dqn/trainer/action.rs Normal file
View File

@@ -0,0 +1,396 @@
//! DQN Trainer — Action selection, routing, and fill simulation
use anyhow::Result;
use candle_core::Tensor;
use tracing::{debug, info};
use super::DQNTrainer;
use crate::dqn::action_space::{ExposureLevel, FactoredAction};
use crate::dqn::TradingState;
use crate::dqn::mixed_precision::training_dtype;
use crate::dqn::order_router::OrderRouter;
use ml_core::fill_simulator::FillResult;
impl DQNTrainer {
/// Select action using epsilon-greedy
pub(crate) async fn select_action(&self, state: &TradingState) -> Result<FactoredAction> {
let _agent = self.agent.read().await;
// Convert state to tensor with tensor core alignment padding
let state_vec = state.to_vector();
let raw_dim = state_vec.len();
let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device);
let padded: Vec<f32> = if aligned_dim > raw_dim {
let mut v = state_vec.to_vec();
v.resize(aligned_dim, 0.0);
v
} else {
state_vec.to_vec()
};
let state_tensor = Tensor::new(&*padded, &self.device)
.map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))?
.unsqueeze(0)?; // Add batch dimension
// Get Q-values (epsilon-greedy handled by agent internally)
let action_idx = self.epsilon_greedy_action(&state_tensor).await?;
let exposure = ExposureLevel::from_index(action_idx)
.map_err(|e| anyhow::anyhow!("Invalid action index {}: {}", action_idx, e))?;
// Phase C: Use smart routing with trainer's spread/vol EMAs
Ok(self.route_action(exposure, self.hyperparams.avg_spread as f32))
}
/// Phase C: Route an exposure-level action using smart order routing.
///
/// Uses the trainer's running spread/volatility EMAs to determine
/// optimal order type (Market/Limit/IoC) and urgency (Patient/Normal/Aggressive).
/// The DQN selects exposure; OrderRouter selects execution strategy.
pub(crate) fn route_action(&self, exposure: ExposureLevel, current_spread: f32) -> FactoredAction {
OrderRouter::route(
exposure,
current_spread,
self.hyperparams.avg_spread as f32,
self.vol_ema as f32,
self.median_vol as f32,
)
}
/// Phase C: Simulate order fill and return result.
///
/// Returns (action, fill_result) — if not filled, action is overridden to Flat
/// so the agent learns that limit orders in certain conditions don't execute.
pub(crate) fn simulate_fill(
&self,
action: FactoredAction,
step: usize,
) -> (FactoredAction, FillResult) {
let normalized_vol = if self.median_vol > 0.0 {
(self.vol_ema / self.median_vol) as f32
} else {
1.0
};
let spread_bps = self.hyperparams.avg_spread * 10000.0; // fractional → bps
let fill_result = self.fill_simulator.simulate_fill(
action.order,
action.urgency,
normalized_vol,
spread_bps,
step,
action.exposure as usize,
);
if fill_result.filled {
(action, fill_result)
} else {
// Order didn't fill — position stays unchanged (Flat action, no trade)
(OrderRouter::route_default(ExposureLevel::Flat), fill_result)
}
}
/// Select actions for a batch of states (GPU-optimized)
///
/// This method reduces GPU kernel launches by batching all action selections
/// into a single forward pass. Provides 125× reduction in kernel launches
/// compared to sequential select_action() calls.
///
/// # Performance Impact
/// - Single GPU kernel launch for entire batch (vs. one per sample)
/// - Reduced CPU-GPU synchronization overhead
/// - Better GPU utilization through larger batch sizes
///
/// # Arguments
/// * `states` - Slice of TradingState objects to process
///
/// # Returns
/// Vector of TradingAction decisions (same order as input states)
pub(crate) async fn select_actions_batch(&mut self, states: &[TradingState]) -> Result<Vec<FactoredAction>> {
if states.is_empty() {
return Ok(Vec::new());
}
let agent = self.agent.read().await;
let batch_size = states.len();
// Get state dimension from first state, then align for tensor cores
let first_vec = states
.first()
.map(|s| s.to_vector())
.ok_or_else(|| anyhow::anyhow!("Empty states slice"))?;
let raw_state_dim = first_vec.len();
let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &self.device);
let pad = aligned_dim - raw_state_dim;
// Pre-allocate flat buffer, zero-padding each state to aligned dimension
let mut flat_states = Vec::with_capacity(batch_size * aligned_dim);
flat_states.extend_from_slice(&first_vec);
flat_states.extend(std::iter::repeat_n(0.0_f32, pad));
for (i, state) in states.iter().enumerate().skip(1) {
let vec = state.to_vector();
if vec.len() != raw_state_dim {
return Err(anyhow::anyhow!(
"State {} dimension mismatch: expected {}, got {}",
i,
raw_state_dim,
vec.len()
));
}
flat_states.extend_from_slice(&vec);
flat_states.extend(std::iter::repeat_n(0.0_f32, pad));
}
// Create batched tensor directly from flat buffer
let batch_tensor = Tensor::from_vec(flat_states, (batch_size, aligned_dim), &self.device) .map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))?
.to_dtype(training_dtype(&self.device))
.map_err(|e| anyhow::anyhow!("Failed to cast batched state tensor to training dtype: {}", e))?;
// FIX: Use get_effective_epsilon() which respects noisy_epsilon_floor.
// Previously used get_epsilon() which returns the decayed epsilon (0.0 with noisy nets),
// making the batch path have ZERO random exploration — root cause of action collapse.
let base_epsilon = agent.get_effective_epsilon() as f64;
let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon);
let epsilon = adjusted_epsilon as f32;
debug!("Epsilon: base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon);
// Single forward pass for all samples (GPU-optimized).
// For RegimeConditional, forward() now blends all 3 heads via regime masks.
let batch_q_values = agent
.forward(&batch_tensor)
.map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?;
// Branching DQN: get per-branch Q-values while agent lock is held.
// Returns (exposure [batch,5], order [batch,3], urgency [batch,3]).
let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> =
if self.hyperparams.use_branching {
agent
.batch_branching_q_values(&batch_tensor)
.map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))?
} else {
None
};
drop(agent); // Release lock early
// Fused GPU epsilon-greedy: argmax + RNG in a single CUDA kernel launch.
// Eliminates the intermediate GPU→CPU sync from Candle's argmax().
// Lazy-init the GPU action selector on first call
{
if self.gpu_action_selector.is_none() && self.device.is_cuda() {
let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new(
&self.device,
self.hyperparams.batch_size.max(batch_size).max(8192),
0xDEAD_BEEF_CAFE_u64,
).map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?;
info!("GPU action selector initialized for select_actions_batch");
self.gpu_action_selector = Some(selector);
}
let selector = self.gpu_action_selector.as_mut()
.ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?;
// Branching path: per-branch epsilon-greedy → factored indices (0-44)
let factored_tensor = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors {
selector
.select_actions_branching(q_exp, q_ord, q_urg, epsilon)
.map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))?
} else {
// Non-branching: exposure-only epsilon-greedy → GPU route to factored
let exposure_tensor = selector
.select_actions(&batch_q_values, epsilon, batch_size, 5)
.map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?;
selector.route_exposure_to_factored(
&exposure_tensor,
batch_size,
self.hyperparams.avg_spread as f32,
self.hyperparams.avg_spread as f32,
self.vol_ema as f32,
self.median_vol as f32,
).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))?
};
// Per-element scalar readback (no to_vec1)
// narrow(0,i,1) → shape [1]; squeeze(0) → scalar [] for to_scalar
let mut actions = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let idx = factored_tensor.narrow(0, i, 1)
.and_then(|t| t.squeeze(0))
.and_then(|t| t.to_scalar::<u32>())
.map_err(|e| anyhow::anyhow!("Factored index readback [{i}]: {e}"))?;
let action = FactoredAction::from_index(idx as usize)
.map_err(|e| anyhow::anyhow!("Invalid factored index {idx}: {e}"))?;
actions.push(action);
}
Ok(actions)
}
}
/// GPU-optimized batch action selection using pre-built state tensor.
///
/// Skips the state→Vec→flatten→Tensor pipeline (~130 allocs per batch).
/// GPU-batched action selection with optional fused routing + fill simulation.
///
/// Returns `(actions, gpu_handled_fill)`:
/// - `gpu_handled_fill = true`: actions are post-fill (routed + fill-checked by GPU kernel).
/// Caller must NOT apply CPU `route_action()` or `simulate_fill()` — already done.
/// - `gpu_handled_fill = false`: actions have basic routing only. Caller should apply
/// CPU routing + fill as before.
///
/// The fused kernel (`epsilon_greedy_routed`) is used when:
/// 1. GPU action selector is available (CUDA device)
/// 2. Not using branching DQN (branching learns order type via network heads)
/// 3. Median volatility > 0 (fill simulation requires vol context)
pub(crate) async fn select_actions_batch_gpu(
&mut self,
batch_tensor: &Tensor,
batch_start: usize,
) -> Result<(Vec<FactoredAction>, bool)> {
let batch_size = batch_tensor.dims()[0];
if batch_size == 0 {
return Ok((Vec::new(), false));
}
let agent = self.agent.read().await;
// FIX: Use get_effective_epsilon() which respects noisy_epsilon_floor.
// Same fix as select_actions_batch() — previously used get_epsilon() which
// returned 0.0 with noisy nets, causing zero exploration in GPU batch path.
let base_epsilon = agent.get_effective_epsilon() as f64;
let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon);
let epsilon = adjusted_epsilon as f32;
debug!("Epsilon (GPU path): base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon);
// Single forward pass — tensor already on GPU, no construction needed
let batch_q_values = agent
.forward(batch_tensor)
.map_err(|e| anyhow::anyhow!("GPU batched forward pass failed: {}", e))?;
// C2 FIX: Count bonus removed from Q-value computation.
// Noisy nets are the sole exploration mechanism during training.
// Count bonus kept for diversity metrics only (record_action tracking).
// Branching DQN: get per-branch Q-values if branching mode is active.
// Uses unified dispatch that supports both Standard and RegimeConditional.
// For RC, regime classification masks blend per-branch Q-values from all 3 heads.
let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = if self.hyperparams.use_branching {
agent
.batch_branching_q_values(batch_tensor)
.map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))?
} else {
None
};
drop(agent);
// Fused GPU epsilon-greedy: argmax + RNG + routing in CUDA kernels.
// Lazy-init GPU action selector on first call.
if self.gpu_action_selector.is_none() && self.device.is_cuda() {
let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new(
&self.device,
self.hyperparams.batch_size.max(batch_size).max(8192),
0xDEAD_BEEF_CAFE_u64,
).map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?;
info!("GPU action selector initialized for select_actions_batch_gpu");
self.gpu_action_selector = Some(selector);
}
let selector = self.gpu_action_selector.as_mut()
.ok_or_else(|| anyhow::anyhow!("GPU action selector not initialized on non-CUDA device"))?;
// Use fused routing+fill kernel when conditions allow:
// - Not branching (branching learns order type via network heads)
// - Median vol > 0 (fill simulation needs vol context)
let use_routed = !self.hyperparams.use_branching && self.median_vol > 0.0;
// Select actions via fused kernel — one launch, no intermediate GPU→CPU sync
let factored_tensor = if self.hyperparams.use_branching {
if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors {
// Branching: 3-head epsilon-greedy → factored (0-44) directly
selector
.select_actions_branching(q_exp, q_ord, q_urg, epsilon)
.map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))?
} else {
// Branching Q forward failed — exposure-only → route on GPU
let exposure_tensor = selector
.select_actions(&batch_q_values, epsilon, batch_size, 5)
.map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?;
selector.route_exposure_to_factored(
&exposure_tensor, batch_size,
self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32,
self.vol_ema as f32, self.median_vol as f32,
).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))?
}
} else if use_routed {
// Fused: epsilon-greedy + routing + fill simulation in one kernel
let spread = self.hyperparams.avg_spread as f32;
let spread_bps = (self.hyperparams.avg_spread * 10000.0) as f32;
selector
.select_actions_routed(
&batch_q_values, epsilon, batch_size, 5,
batch_start as i32,
spread, spread,
self.vol_ema as f32, self.median_vol as f32,
spread_bps, 0.85, 0.30, 0.80, 0.50, 0.50,
)
.map_err(|e| anyhow::anyhow!("GPU routed action selection failed: {e}"))?
} else {
// Exposure-only → route on GPU
let exposure_tensor = selector
.select_actions(&batch_q_values, epsilon, batch_size, 5)
.map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?;
selector.route_exposure_to_factored(
&exposure_tensor, batch_size,
self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32,
self.vol_ema as f32, self.median_vol as f32,
).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))?
};
// Per-element scalar readback (no to_vec1)
// narrow(0,i,1) → shape [1]; squeeze(0) → scalar [] for to_scalar
let mut actions = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let idx = factored_tensor.narrow(0, i, 1)
.and_then(|t| t.squeeze(0))
.and_then(|t| t.to_scalar::<u32>())
.map_err(|e| anyhow::anyhow!("Factored index readback [{i}]: {e}"))?;
let action = FactoredAction::from_index(idx as usize)
.map_err(|e| anyhow::anyhow!("Invalid factored index {idx}: {e}"))?;
actions.push(action);
}
Ok((actions, use_routed))
}
/// Epsilon-greedy action selection for single-step inference.
///
/// This is batch_size=1 — the overhead of a CUDA kernel launch (~5us) exceeds
/// the benefit of fusing argmax+RNG for a single element. Candle's argmax +
/// to_scalar is already minimal for this path. Keep on CPU.
pub(crate) async fn epsilon_greedy_action(&self, state: &Tensor) -> Result<usize> {
use rand::Rng;
let epsilon = self.get_epsilon().await? as f32;
let mut rng = rand::thread_rng();
if rng.gen::<f32>() < epsilon {
// Random action (exploration) over 5 exposure levels
Ok(rng.gen_range(0..5))
} else {
// Greedy action (exploitation) - use actual Q-network
let agent = self.agent.read().await;
let q_values = agent.forward(state)?;
// GPU-native argmax — single u32 scalar transfer instead of full Q-value vector
let best_action = q_values
.argmax(1)
.and_then(|t| t.squeeze(0))
.and_then(|t| t.to_scalar::<u32>()) .map(|v| v as usize)
.ok()
.unwrap_or(2); // Default to HOLD (index 2) on error
Ok(best_action)
}
}
}

732
crates/ml/src/trainers/dqn/trainer/constructor.rs Normal file
View File

@@ -0,0 +1,732 @@
//! DQN Trainer constructor — `new_internal` and init helpers.
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::Duration;
use anyhow::Result;
use candle_core::Device;
use ml_core::fill_simulator::FillSimulator;
use risk::drawdown_monitor::DrawdownMonitor;
use risk::safety::position_limiter::HybridPositionLimiter;
use risk::safety::PositionLimiterConfig;
use rust_decimal::Decimal;
use tokio::sync::RwLock;
use tracing::info;
use crate::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use crate::dqn::curiosity::CuriosityModule;
use crate::dqn::dqn::{DQN, DQNConfig};
use crate::dqn::logging::{LoggingConfig, MetricsAggregator};
use crate::dqn::portfolio_tracker::PortfolioTracker;
use crate::dqn::regime_conditional::RegimeConditionalDQN;
use crate::dqn::reward::{RewardConfig, RewardFunction};
use crate::features::microstructure_features::*;
use crate::labeling::triple_barrier::TripleBarrierEngine;
use crate::memory_optimization::auto_batch_size::{AutoBatchSizer, BatchSizeConfig};
use crate::trainers::TargetUpdateMode;
use crate::TrainingMetrics;
use super::super::config::{DQNAgentType, DQNHyperparameters};
use super::DQNTrainer;
impl DQNTrainer {
pub(crate) fn new_internal(mut hyperparams: DQNHyperparameters, debug_logging: bool, override_device: Option<Device>) -> Result<Self> {
// Validate batch size is non-zero
if hyperparams.batch_size == 0 {
return Err(anyhow::anyhow!(
"Batch size must be greater than 0, got: {}",
hyperparams.batch_size
));
}
// WAVE 26 P2.2: Validate gradient_accumulation_steps > 0
if hyperparams.gradient_accumulation_steps == 0 {
return Err(anyhow::anyhow!(
"gradient_accumulation_steps must be greater than 0, got: {}",
hyperparams.gradient_accumulation_steps
));
}
// Pre-compute hidden dims to get accurate model size for batch sizing.
// Align input_dim to 8 so the log matches the actual model dimensions.
// (device not yet created, so use the formula directly — CUDA always aligns)
let ofi_pre = hyperparams.mbp10_data_dir.is_some();
let input_dim: usize = if ofi_pre { 56 } else { 48 }; // (53+7)&!7=56, (45+7)&!7=48
let output_dim: usize = 5;
let hidden_dims: Vec<usize> = match hyperparams.hidden_dim_base {
Some(base) => {
let b = crate::cuda_pipeline::align_to_tensor_cores(base);
vec![b, b] // Constant-width: no tapering, no silently-discarded narrow layer
}
None => {
let caps = crate::gpu::capabilities::cached_capabilities();
let base = crate::gpu::memory_profile::resolve_hidden_dim_base(
caps.free_vram_mb,
);
let b = crate::cuda_pipeline::align_to_tensor_cores(base);
vec![b, b] // Constant-width: no tapering
}
};
// Compute accurate model size from actual network dimensions
let full_dims: Vec<usize> = std::iter::once(input_dim)
.chain(hidden_dims.iter().copied())
.chain(std::iter::once(output_dim))
.collect();
let param_count = crate::gpu::memory_profile::network_param_count(&full_dims);
// FP32 params + AdamW state (2x for momentum/variance) + target network copy = ~4x
let model_overhead_mb = (param_count as f64 * 4.0 * 4.0) / (1024.0 * 1024.0);
info!(
"DQN network: {:?} → {} params, {:.1} MB overhead",
full_dims, param_count, model_overhead_mb
);
// Dynamic batch sizing: scale UP for larger GPUs, cap DOWN for smaller ones.
// Uses HardwareBudget for consistent sizing across DQN/PPO.
const STATIC_MAX_BATCH_SIZE: usize = 8192;
let max_safe_batch = match AutoBatchSizer::new() {
Ok(sizer) => {
let config = BatchSizeConfig {
model_memory_mb: model_overhead_mb,
safety_margin: 0.15,
..BatchSizeConfig::default()
};
let safe = sizer.max_safe_batch_size(&config);
info!(
"AutoBatchSizer: GPU VRAM ceiling = {} (configured: {})",
safe, hyperparams.batch_size
);
safe
}
Err(e) => {
info!(
"AutoBatchSizer unavailable ({}), using static cap: {}",
e, STATIC_MAX_BATCH_SIZE
);
STATIC_MAX_BATCH_SIZE
}
};
// Cap to VRAM ceiling from AutoBatchSizer (no separate scale-UP —
// AutoBatchSizer already accounts for model size and available VRAM)
if hyperparams.batch_size > max_safe_batch {
info!(
"DQN batch_size capped from {} → {} (VRAM ceiling)",
hyperparams.batch_size, max_safe_batch
);
hyperparams.batch_size = max_safe_batch;
}
// Use override device if provided (hyperopt shares one CUDA context),
// otherwise auto-detect GPU
let device = if let Some(dev) = override_device {
dev
} else {
Device::cuda_if_available(0)
.map_err(|e| anyhow::anyhow!("Failed to initialize device: {}", e))?
};
// Dynamic replay buffer sizing: scale replay capacity to available VRAM.
// Only activates when replay_buffer_vram_fraction > 0 and GPU is detected.
// Also computes the PER memory budget from actual VRAM — no hardcoded caps.
let original_buffer_size = hyperparams.buffer_size; // Save before AutoReplaySizer mutates it
let mut per_max_memory_bytes: usize = 4 * 1024 * 1024 * 1024; // CPU fallback: 4 GB
if hyperparams.replay_buffer_vram_fraction > 0.0 && device.is_cuda() {
use ml_core::memory_optimization::detect_gpu_hardware;
match detect_gpu_hardware() {
Ok(hw) => {
let raw_sd = if hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 };
let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &device);
let replay_cfg = hw.optimal_replay_config(
aligned_sd,
hyperparams.replay_buffer_vram_fraction,
);
per_max_memory_bytes = replay_cfg.per_max_buffer_bytes;
if replay_cfg.capacity != hyperparams.buffer_size {
info!(
"AutoReplaySizer: replay buffer {} -> {} (VRAM={:.0}MB, fraction={:.0}%, PER budget={:.0}MB)",
hyperparams.buffer_size,
replay_cfg.capacity,
hw.free_memory_mb,
hyperparams.replay_buffer_vram_fraction * 100.0,
per_max_memory_bytes as f64 / (1024.0 * 1024.0),
);
hyperparams.buffer_size = replay_cfg.capacity;
}
}
Err(e) => {
info!(
"AutoReplaySizer unavailable ({}), using static buffer_size: {}",
e, hyperparams.buffer_size
);
}
}
} else if device.is_cuda() {
// No auto-sizer, but still compute PER budget from VRAM
use ml_core::memory_optimization::detect_gpu_hardware;
if let Ok(hw) = detect_gpu_hardware() {
per_max_memory_bytes = hw.per_max_buffer_bytes();
} else {
// GPU detection failed — keep CPU default (4 GB)
}
} else {
// CPU device — keep default 4 GB PER budget
}
info!(
"Initializing DQN trainer on device: {:?}, using 5 exposure actions + OrderRouter",
if device.is_cuda() { "CUDA GPU" } else { "CPU" },
);
// Auto-detect mixed precision capability based on GPU architecture
let mixed_precision_detected = if device.is_cuda() {
match crate::memory_optimization::auto_batch_size::detect_gpu_memory() {
Ok((_total, _free, ref name)) => {
let detected = crate::dqn::mixed_precision::detect_from_gpu_name(name);
match &detected {
Some(c) => info!("GPU mixed precision: {:?} enabled (GPU: {})", c.dtype, name),
None => info!("GPU mixed precision: disabled (GPU: {})", name),
}
detected
}
Err(_) => None,
}
} else {
None
};
// Create DQN configuration
// 42-feature architecture: OHLCV, technical, patterns, volume, time, statistical, regime
// Portfolio features (3) are populated via PortfolioTracker → 45 total state_dim
// With MBP-10 OFI features: +8 OFI features → 53 total
//
// Tensor core alignment: state_dim is rounded up to the next multiple of 8
// (53→56, 45→48) so that cuBLAS dispatches BF16 HMMA instructions instead
// of falling back to scalar FMA. The extra columns are zero-padded at the
// data pipeline boundaries (GpuPreloadedData and train_batch CPU path).
let ofi_enabled = hyperparams.mbp10_data_dir.is_some();
let raw_state_dim = if ofi_enabled { 53 } else { 45 };
let state_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &device);
let config = DQNConfig {
state_dim,
num_actions: 5, // 5 exposure levels (Short100, Short50, Flat, Long50, Long100)
hidden_dims,
learning_rate: hyperparams.learning_rate,
gamma: hyperparams.gamma as f32,
epsilon_start: hyperparams.epsilon_start as f32,
epsilon_end: hyperparams.epsilon_end as f32,
epsilon_decay: hyperparams.epsilon_decay as f32,
replay_buffer_capacity: hyperparams.buffer_size,
collapse_warmup_capacity: original_buffer_size,
batch_size: hyperparams.batch_size,
min_replay_size: hyperparams.min_replay_size.min(hyperparams.buffer_size), // Cap at buffer_size to prevent deadlock
target_update_freq: hyperparams.target_update_frequency, // Use hyperparameter instead of hardcoded 1000
use_double_dqn: true,
use_huber_loss: hyperparams.use_huber_loss,
huber_delta: hyperparams.huber_delta as f32,
leaky_relu_alpha: 0.01, // Standard LeakyReLU alpha (prevents dead neurons)
gradient_clip_norm: hyperparams.gradient_clip_norm.unwrap_or(10.0), // Wave 11 Bug #1 fix: Dynamic clipping
// WAVE 16 (Agent 36): Target update configuration
tau: hyperparams.tau,
tau_final: hyperparams.tau * 0.1, // Anneal to 10% of base tau
tau_anneal_steps: 100_000,
use_soft_updates: matches!(hyperparams.target_update_mode, TargetUpdateMode::Soft),
// Rainbow DQN warmup period
warmup_steps: hyperparams.warmup_steps,
// PER configuration
initial_capital: hyperparams.initial_capital as f64,
use_per: hyperparams.use_per,
use_gpu_replay_buffer: hyperparams.use_gpu_replay_buffer,
per_alpha: hyperparams.per_alpha,
per_beta_start: hyperparams.per_beta_start,
per_beta_max: 1.0,
per_beta_annealing_steps: hyperparams.epochs * 2000, // ~2000 steps/epoch (130k bars / ~64 batch_size)
per_max_memory_bytes,
// Wave 2.1: Dueling Networks (ENABLED BY DEFAULT - Wave 6.4)
use_dueling: hyperparams.use_dueling,
dueling_hidden_dim: hyperparams.dueling_hidden_dim,
// Wave 2.2: Multi-Step Returns (N-step TD) (ENABLED BY DEFAULT - Wave 6.4)
n_steps: hyperparams.n_steps, // Default: 3 (Rainbow DQN standard)
// Wave 2.3: Distributional RL (C51) (ENABLED BY DEFAULT - Wave 6.4)
use_distributional: hyperparams.use_distributional, // Default: enabled (C51 distributional RL)
num_atoms: hyperparams.num_atoms, // Rainbow DQN standard: 51 atoms
v_min: hyperparams.v_min as f32, // Minimum value for distribution support
v_max: hyperparams.v_max as f32, // Maximum value for distribution support
// Wave 2.4: Noisy Networks for Exploration (ENABLED BY DEFAULT - Wave 6.4)
use_noisy_nets: hyperparams.use_noisy_nets, // Default: enabled (replaces epsilon-greedy)
noisy_sigma_init: hyperparams.noisy_sigma_init, // Rainbow DQN standard: 0.5
// BUG #37 FIX: Q-value clipping (prevents step-level explosions)
enable_q_value_clipping: true,
q_value_clip_min: -500.0,
q_value_clip_max: 500.0,
// WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold (from hyperparams)
gradient_collapse_multiplier: hyperparams.gradient_collapse_multiplier,
gradient_collapse_patience: hyperparams.gradient_collapse_patience,
use_cql: hyperparams.use_cql,
cql_alpha: hyperparams.cql_alpha,
use_iqn: hyperparams.use_qr_dqn, // Controlled by hyperopt
iqn_num_quantiles: hyperparams.num_quantiles, // Controlled by hyperopt
iqn_kappa: hyperparams.qr_kappa as f32, // Controlled by hyperopt (f64→f32)
iqn_embedding_dim: 64, // Fixed (not in search space)
use_branching: hyperparams.use_branching,
branch_hidden_dim: hyperparams.branch_hidden_dim,
use_regime_conditioning: true, // Always enable per-regime IS weights for branching loss
use_cvar_action_selection: false,
cvar_alpha: 0.05,
#[allow(clippy::cast_possible_truncation)]
minimum_profit_factor: hyperparams.minimum_profit_factor as f32,
weight_decay: hyperparams.weight_decay,
dropout_rate: if hyperparams.enable_dropout_scheduler { hyperparams.dropout_initial } else { 0.0 },
mixed_precision: hyperparams.mixed_precision.clone().or(mixed_precision_detected),
entropy_coefficient: hyperparams.entropy_coefficient.unwrap_or(0.01),
noisy_epsilon_floor: hyperparams.noisy_epsilon_floor.unwrap_or(0.0) as f32, // C2: NoisyNet handles exploration
use_count_bonus: hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, // C3 FIX: enable when coefficient > 0
count_bonus_coefficient: hyperparams.count_bonus_coefficient.unwrap_or(0.0),
..DQNConfig::default()
};
// Extract curiosity dims before config is moved into the agent
let curiosity_market_dim = config.curiosity_market_dim;
let curiosity_hidden_dim = config.curiosity_hidden_dim;
// Create DQN agent
let agent = if hyperparams.enable_regime_qnetwork {
info!("Creating regime-conditional DQN with 3 heads (Trending, Ranging, Volatile)");
info!(" - Regime detection: ADX (raw index 40) + CUSUM direction (raw index 41)");
info!(" - Classification: Trending (ADX>0.25), Volatile (ADX≤0.25 & |CUSUM|>0.7), Ranging (otherwise)");
let regime_agent = RegimeConditionalDQN::new_on_device(config, device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?;
DQNAgentType::RegimeConditional(regime_agent)
} else {
info!("Creating standard DQN with single Q-network head");
let standard_agent = DQN::new_on_device(config, device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create DQN agent: {}", e))?;
DQNAgentType::Standard(standard_agent)
};
// Initialize portfolio tracker with $100k starting capital and 1 basis point spread
// Bug #2 fix: Portfolio features were hardcoded as [0.0, 0.0, 0.0] at line 1528
let portfolio_tracker = PortfolioTracker::new(
hyperparams.initial_capital, // P2-A: Configurable capital
0.0001, // 1 basis point spread (0.01%)
hyperparams.cash_reserve_percent, // Cash reserve requirement
);
// Initialize reward function with hyperparameter-driven configuration
// WAVE 10-A9 FIX: Wire hold_penalty_weight from hyperparameters to RewardConfig
// BUG #17 FIX: Add normalization and percentage-based P&L (enabled by default)
let reward_config = RewardConfig {
pnl_weight: Decimal::ONE,
risk_weight: Decimal::try_from(0.1).unwrap_or(Decimal::ZERO),
cost_weight: Decimal::ONE, // Bug #2 fix: 100% transaction cost weight (was 0.05, 20x too low)
hold_reward: Decimal::ZERO, // Flat position = no edge = zero reward (was +0.001, 20x trade PnL)
movement_threshold: Decimal::try_from(hyperparams.movement_threshold)
.unwrap_or(Decimal::ZERO),
hold_penalty_weight: Decimal::try_from(hyperparams.hold_penalty_weight)
.unwrap_or(Decimal::ZERO), // CRITICAL FIX
diversity_weight: Decimal::try_from(-0.1).unwrap_or(Decimal::ZERO),
enable_normalization: !hyperparams.use_dsr, // DSR replaces EMA normalizer
use_percentage_pnl: true, // Bug #17: Use percentage returns for scale-invariance
circuit_breaker_config: CircuitBreakerConfig::default(),
triple_barrier_profit_bonus: Decimal::try_from(0.5).unwrap_or(Decimal::ZERO),
triple_barrier_stop_penalty: Decimal::try_from(0.5).unwrap_or(Decimal::ZERO),
sharpe_weight: Decimal::ZERO, // WAVE 26 P1.3: Disabled by default
sharpe_window: 20, // WAVE 26 P1.3: Standard 20-period window
use_dsr: hyperparams.use_dsr,
dsr_eta: hyperparams.dsr_eta,
initial_capital: hyperparams.initial_capital as f64,
};
let reward_fn = RewardFunction::new_with_debug(reward_config, debug_logging)?;
// WAVE 1.1: Initialize triple barrier engine (max 1000 active trackers)
let triple_barrier = Arc::new(RwLock::new(TripleBarrierEngine::new(1000)));
info!("Triple barrier engine initialized with 1000 max trackers");
// WAVE 16S: Initialize Kelly optimizer if enabled
let kelly_optimizer = if hyperparams.enable_kelly_sizing {
use crate::risk::kelly_optimizer::{KellyCriterionOptimizer, KellyOptimizerConfig};
let kelly_config = KellyOptimizerConfig {
max_fraction: hyperparams.kelly_max_fraction,
min_fraction: 0.01,
lookback_period: 252,
confidence_threshold: 0.6,
volatility_adjustment: true,
drawdown_protection: true,
};
let optimizer = KellyCriterionOptimizer::new(kelly_config)
.map_err(|e| anyhow::anyhow!("Failed to create Kelly optimizer: {}", e))?;
info!("Kelly optimizer enabled (fractional={}, max={})",
hyperparams.kelly_fractional, hyperparams.kelly_max_fraction);
Some(Arc::new(optimizer))
} else {
None
};
// Wave 16 Portfolio Features: Initialize action masking, entropy regularization, and stress testing
let enable_action_masking = hyperparams.enable_action_masking;
let max_position = hyperparams.max_position_absolute; // BLOCKER #2: Use hyperopt-tunable position limit
// Entropy regularization: SAC-style computed directly on Q-values in DQN::compute_loss_internal
if hyperparams.enable_entropy_regularization {
let coeff = hyperparams.entropy_coefficient.unwrap_or(0.01);
info!("Entropy regularization enabled (coefficient={coeff:.4}, applied to Q-value softmax in loss)");
}
// Multi-asset portfolio tracking (disabled -- single-asset is the current production mode)
//
// When expanding to multi-asset trading:
// 1. Add `enable_multi_asset: bool` to DQNHyperparams (default false).
// 2. Initialize MultiAssetPortfolioTracker here when the flag is set.
// 3. Wire portfolio state into the DQN observation: expand state_dim to
// include per-asset position, PnL, and correlation features so the
// agent can learn cross-asset hedging and allocation.
let multi_asset_portfolio: Option<Arc<crate::dqn::multi_asset::MultiAssetPortfolioTracker>> = None;
// Stress testing for robustness validation
// Initialized as None here; call init_stress_tester() after construction
// to resolve the circular dependency (DQNStressTester needs a DQNTrainer).
let stress_tester: Option<Arc<crate::dqn::stress_testing::DQNStressTester>> = None;
if enable_action_masking {
info!(
"Action masking enabled (max_position=±{:.1}, 30-50% filtering expected)",
max_position
);
} else {
info!("Action masking disabled (all 5 exposure levels available)");
}
// Wave 16 Core Risk Features: Initialize drawdown monitor, position limiter, circuit breaker
// These are ALWAYS enabled by default for production safety
// 1. Drawdown Monitor (15% max drawdown, alerts at 10%, 12.5%, 15%)
let drawdown_monitor = {
// DrawdownMonitor will be configured in first training step
// Config will be applied via async configure_alerts() in train_epoch
info!("Drawdown monitor enabled (thresholds: 10%, 12.5%, 15%)");
Some(Arc::new(DrawdownMonitor::new()))
};
// 2. Position Limiter (3-tier limits: ±10.0 absolute, 1M notional, 10% concentration)
let position_limiter = {
let config = PositionLimiterConfig {
enabled: true,
cache_ttl: Duration::from_secs(60),
rpc_check_threshold_percent: 0.8,
max_position_per_symbol: 10.0, // ±10.0 absolute position limit
max_order_value: 1_000_000.0, // $1M notional limit
max_daily_loss: 0.10, // 10% concentration limit
};
let limiter = HybridPositionLimiter::new(config);
info!("Position limiter enabled (abs=±10.0, notional=$1M, concentration=10%)");
Some(Arc::new(limiter))
};
// 3. Circuit Breaker (5 consecutive failures, 60s cooldown)
let circuit_breaker = {
let config = CircuitBreakerConfig {
failure_threshold: 5,
success_threshold: 3,
timeout_duration: Duration::from_secs(60),
half_open_max_calls: 2,
};
let breaker = CircuitBreaker::new(config);
info!("Circuit breaker enabled (threshold=5 failures, cooldown=60s)");
Some(Arc::new(breaker))
};
// WAVE 24: Capture patience before hyperparams is moved
let early_stopping_patience = hyperparams.gradient_collapse_patience;
// WAVE 26 P1: Initialize advanced DQN features
// P1.3: Sharpe ratio reward component
let sharpe_weight = hyperparams.sharpe_weight;
let sharpe_window = hyperparams.sharpe_window;
// P1.6: Adaptive dropout scheduler
let dropout_scheduler = hyperparams.enable_dropout_scheduler.then(|| {
use crate::dqn::network::DropoutScheduler;
info!("Dropout scheduler enabled (initial={}, final={}, steps={})",
hyperparams.dropout_initial, hyperparams.dropout_final, hyperparams.dropout_anneal_steps);
DropoutScheduler::new(
hyperparams.dropout_initial,
hyperparams.dropout_final,
hyperparams.dropout_anneal_steps,
)
});
// P1.7: Hindsight Experience Replay (HER)
let her_buffer = (hyperparams.her_ratio > 0.0)
.then(|| {
use crate::dqn::hindsight_replay::{HindsightReplayBuffer, HindsightReplayConfig, HindsightStrategy};
use crate::dqn::prioritized_replay::PrioritizedReplayConfig;
let her_strategy = match hyperparams.her_strategy.as_str() {
"final" => HindsightStrategy::Final,
_ => HindsightStrategy::Future, // Default to Future
};
let config = HindsightReplayConfig {
base_config: PrioritizedReplayConfig {
capacity: hyperparams.buffer_size,
..Default::default()
},
her_ratio: hyperparams.her_ratio,
her_strategy,
goal_dim: 1, // Single goal dimension for trading (target return)
k_future: 4, // Sample 4 future goals for Future strategy
batch_size: hyperparams.batch_size,
};
info!("HER buffer enabled (ratio={}, strategy={:?}, capacity={})",
hyperparams.her_ratio, her_strategy, hyperparams.buffer_size);
HindsightReplayBuffer::new(config)
.map(Arc::new)
.map_err(|e| anyhow::anyhow!("Failed to create HER buffer: {}", e))
})
.transpose()?;
// P1.9: Generalized Advantage Estimation (GAE)
let gae_calculator = hyperparams.enable_gae.then(|| {
use crate::dqn::gae::GAECalculator;
info!("GAE calculator enabled (lambda={}, gamma={})",
hyperparams.gae_lambda, hyperparams.gamma);
GAECalculator::new(hyperparams.gae_lambda, hyperparams.gamma)
});
// P1.11: Noisy network sigma scheduling
let noisy_sigma_scheduler = hyperparams.enable_noisy_sigma_scheduler.then(|| {
use crate::dqn::noisy_sigma_scheduler::NoisySigmaScheduler;
info!("Noisy sigma scheduler enabled (initial={}, final={}, steps={})",
hyperparams.noisy_sigma_initial, hyperparams.noisy_sigma_final, hyperparams.noisy_sigma_anneal_steps);
NoisySigmaScheduler::new(
hyperparams.noisy_sigma_initial,
hyperparams.noisy_sigma_final,
hyperparams.noisy_sigma_anneal_steps,
)
});
// WAVE 26 P1.8: Initialize curiosity module if curiosity_weight > 0
// Must be created BEFORE hyperparams and device are moved
let curiosity_module = (hyperparams.curiosity_weight > 0.0)
.then(|| {
CuriosityModule::new(
device.clone(),
0.001, // Forward model learning rate
2.0, // Max curiosity reward (clip to prevent noise exploitation)
curiosity_market_dim,
curiosity_hidden_dim,
3, // Action categories (Short/Flat/Long)
).map_err(|e| anyhow::anyhow!("Failed to create curiosity module: {}", e))
})
.transpose()?;
// WAVE 26 P0.6: Initialize learning rate scheduler with warmup
// Must be created BEFORE hyperparams is moved
let lr_scheduler = {
use super::super::lr_scheduler::LRScheduler;
LRScheduler::new(
hyperparams.learning_rate,
hyperparams.warmup_steps,
hyperparams.lr_decay_type,
)
};
// WAVE 44: Initialize n-step buffer if n_steps > 1
let nstep_buffer = (hyperparams.n_steps > 1).then(|| {
info!("🎯 Multi-step returns ENABLED: n_steps={}, gamma={}",
hyperparams.n_steps, hyperparams.gamma);
crate::dqn::nstep_buffer::NStepBuffer::new(
hyperparams.n_steps,
hyperparams.gamma
)
});
// Capture values before hyperparams is moved into Self
let initial_batch_size = hyperparams.batch_size;
let base_tau = hyperparams.tau;
// GPU pipeline: pre-allocate staging buffers on CUDA devices
let buffer_pool = device.is_cuda().then(|| {
info!("GpuBufferPool: pre-allocated staging buffers (100k bars, 42 features, 4 targets)");
crate::cuda_pipeline::GpuBufferPool::new(100_000, 42, 4)
});
// GPU pipeline: double-buffered loader for zero-downtime fold transitions
let double_buffer = device.is_cuda().then(|| {
crate::cuda_pipeline::double_buffer::DoubleBufferedLoader::new(device.clone())
});
// Multi-GPU: auto-detect if multiple CUDA devices are available
let multi_gpu = crate::cuda_pipeline::multi_gpu::MultiGpuConfig::detect()
.unwrap_or(None);
if let Some(ref mg) = multi_gpu {
info!("Multi-GPU: {} devices detected, data parallelism enabled", mg.world_size);
}
Ok(Self {
agent: Arc::new(RwLock::new(agent)),
hyperparams,
device,
metrics: Arc::new(RwLock::new(TrainingMetrics::new())),
loss_history: Vec::new(),
q_value_history: Vec::new(),
best_val_loss: f64::INFINITY, // Start with worst possible loss
val_data: Vec::new(),
val_loss_history: Vec::new(),
sharpe_history: Vec::new(),
best_sharpe: f64::NEG_INFINITY, // C4: Start with worst possible Sharpe
best_epoch: 0,
gradient_logging_step: 0,
collapse_warmup_buffer_size: original_buffer_size,
portfolio_tracker,
feature_stats: None, // WAVE 3 FIX #2: Start with None, collect stats in epochs 0-10
recent_actions: VecDeque::with_capacity(100),
reward_fn,
// WAVE 16S: Adaptive risk management
kelly_optimizer,
trade_history: VecDeque::with_capacity(500),
volatility_returns: VecDeque::with_capacity(20), // Use default instead of moved hyperparams
pnl_history: VecDeque::with_capacity(1000),
// Wave 16 Portfolio Features
enable_action_masking,
max_position,
multi_asset_portfolio,
stress_tester,
// Wave 16 Core Risk Features
drawdown_monitor,
position_limiter,
circuit_breaker,
// WAVE 3.10: Microstructure feature calculators
micro_high_low_spread: HighLowSpread::default(),
micro_vw_spread: VolumeWeightedSpread::default(),
micro_tick_count: TickCount::default(),
micro_inter_arrival: InterArrivalTime::default(),
micro_buy_sell_imbalance: BuySellImbalance::default(),
micro_kyle_lambda: KyleLambda::default(),
micro_price_impact: PriceImpact::default(),
micro_variance_ratio: VarianceRatio::default(),
last_timestamp_ns: 0,
last_close: 0.0,
// WAVE 1.1: Triple barrier integration
triple_barrier,
active_position_tracker: None,
previous_simulated_position: 0.0, // WAVE P3: Start with flat position
// WAVE 1.2: Safety Infrastructure Integration (8 Systems)
safety_loss_history: VecDeque::with_capacity(30),
safety_loss_plateau_counter: 0,
safety_action_counts: std::collections::HashMap::new(),
safety_memory_manager: Arc::new(RwLock::new(
crate::safety::memory_manager::SafeMemoryManager::new(
&crate::safety::MLSafetyConfig::default()
)
)),
safety_level: crate::safety::SafetyLevel::Normal, // Default to Normal mode
safety_step_counter: 0,
feature_cache_dir: None,
prev_epoch_q_mean: 0.0,
adaptive_tau: base_tau,
// WAVE 24 (Agent 17): Initialize patience-based early stopping
// Use gradient_collapse_patience from hyperparams for consistency
// Set min_delta to 0.001 (0.1% improvement threshold)
early_stopping: super::super::early_stopping::EarlyStopping::new(
early_stopping_patience, // Reuse patience parameter (default: 5)
0.001, // 0.1% minimum improvement
),
// WAVE 26 P0.6: Use pre-initialized learning rate scheduler
lr_scheduler,
// WAVE 26 P1.8: Use pre-initialized curiosity module
curiosity_module,
// WAVE 26 P1: Advanced DQN Features
// P1.3: Sharpe ratio reward
returns_history: VecDeque::with_capacity(sharpe_window),
sharpe_weight,
// P1.6: Adaptive dropout
dropout_scheduler,
// P1.7: Hindsight Experience Replay
her_buffer,
// P1.9: Generalized Advantage Estimation
gae_calculator,
// P1.11: Noisy sigma scheduler
noisy_sigma_scheduler,
// WAVE 30: Structured logging integration
logging_config: LoggingConfig::default(),
metrics_aggregator: MetricsAggregator::new(),
// WAVE 44: Multi-step returns
nstep_buffer,
// OOM recovery: track effective batch size
current_batch_size: initial_batch_size,
// Q-value estimation: periodic (every 50 steps) instead of every step
cached_avg_q: 0.0,
q_estimation_counter: 0,
// GPU pipeline: pre-uploaded training data (initialized lazily at first epoch)
gpu_data: None,
gpu_portfolio_sim: None,
targets_raw_cuda: None,
features_raw_cuda: None,
gpu_experience_collector: None,
gpu_action_selector: None,
training_guard: None,
gpu_monitoring: None,
cached_n_episodes: None,
// GPU pipeline: staging buffer pool (auto-initialized on CUDA devices)
buffer_pool,
// GPU pipeline: double-buffered loader for fold transitions
double_buffer,
// GPU walk-forward: initialized lazily when enable_gpu_walk_forward=true
gpu_walk_forward: None,
// Multi-GPU: auto-detected data parallelism
multi_gpu,
// OFI features: populated during data loading when MBP-10 data is available
ofi_features: None,
ofi_val_offset: 0,
val_features_gpu: None,
val_closes_gpu: None,
val_ofi_gpu: None,
// Phase C: Fill simulation and smart order routing
fill_simulator: FillSimulator::default(),
vol_ema: 0.01, // Initial volatility estimate (1% daily)
median_vol: 0.01, // Slowly adapting baseline
// Fused CUDA training: lazy-initialized on first training step
fused_ctx: None,
})
}
}

782
crates/ml/src/trainers/dqn/trainer/metrics.rs Normal file
View File

@@ -0,0 +1,782 @@
//! DQN Trainer — Training metrics, Q-value diagnostics, and validation
use anyhow::Result;
use candle_core::Tensor;
use super::DQNTrainer;
use crate::dqn::TradingState;
use crate::dqn::mixed_precision::training_dtype;
use crate::TrainingMetrics;
use super::super::config::DQNAgentType;
use super::super::statistics::QValueStats;
impl DQNTrainer {
/// Full training loop with existing logic (Wave 12 Group 3)
/// Calculate average metrics for an epoch
pub(crate) fn calculate_epoch_metrics(
epoch_loss: f64,
epoch_q_value: f64,
epoch_gradient_norm: f64,
samples_processed: usize,
) -> (f64, f64, f64) {
if samples_processed > 0 {
let count = samples_processed as f64;
(
epoch_loss / count,
epoch_q_value / count,
epoch_gradient_norm / count,
)
} else {
(0.0, 0.0, 0.0)
}
}
/// Compute validation loss on held-out data
/// WAVE 10.6: Batched validation for 5-10x speedup
/// Collect Q-value statistics from replay buffer
///
/// Samples experiences from the replay buffer and computes Q-value statistics
/// (min, max, mean, std) for adaptive C51 bounds calculation.
///
/// # Returns
///
/// QValueStats with min/max/mean/std of Q-values
pub(crate) async fn collect_qvalue_statistics(&self) -> Result<QValueStats, crate::MLError> {
let agent = self.agent.read().await;
// Determine sample size (min of buffer size or 1000)
let buffer_size = agent.get_replay_buffer_size()?;
let sample_size = buffer_size.min(1000);
if sample_size == 0 {
return Err(crate::MLError::TrainingError(
"Replay buffer is empty, cannot collect Q-value statistics".to_owned()
));
}
// Sample experiences from replay buffer
let batch_sample = agent.memory().sample(sample_size)?;
// GPU PER path: use gpu_batch.states directly (always active in CUDA builds)
let gpu_batch = batch_sample.gpu_batch.as_ref()
.ok_or_else(|| crate::MLError::TrainingError(
"GPU PER must be active — gpu_batch is None".to_owned()
))?;
let batch_tensor = gpu_batch.states.to_dtype(training_dtype(agent.device()))
.map_err(|e| crate::MLError::ModelError(format!("GPU Q-stat states dtype cast: {}", e)))?;
// Forward pass to get Q-values [batch_size, num_actions]
let q_values = agent.forward(&batch_tensor)?;
// GPU-side statistics: flatten Q-values and compute min/max/mean/std on device.
// Only 4 scalar readbacks (16 bytes) instead of downloading the entire tensor.
let q_f32 = q_values
.to_dtype(candle_core::DType::F32)
.map_err(|e| crate::MLError::ModelError(format!("Q-value F32 cast: {}", e)))?;
let q_flat = q_f32
.flatten_all()
.map_err(|e| crate::MLError::ModelError(format!("Q-value flatten: {}", e)))?;
let count = q_flat.elem_count();
// Compute all stats on GPU, single batched readback (4 floats in 1 DMA)
let min_t = q_flat.min(0)
.map_err(|e| crate::MLError::ModelError(format!("Q-value min: {}", e)))?;
let max_t = q_flat.max(0)
.map_err(|e| crate::MLError::ModelError(format!("Q-value max: {}", e)))?;
let mean_t = q_flat.mean_all()
.map_err(|e| crate::MLError::ModelError(format!("Q-value mean: {}", e)))?;
let var_t = q_flat.broadcast_sub(&mean_t)
.and_then(|d| d.sqr())
.and_then(|sq| sq.mean_all())
.map_err(|e| crate::MLError::ModelError(format!("Q-value variance: {}", e)))?;
let stats = Tensor::cat(
&[&min_t.unsqueeze(0)?, &max_t.unsqueeze(0)?, &mean_t.unsqueeze(0)?, &var_t.unsqueeze(0)?], 0
).and_then(|t| t.to_dtype(candle_core::DType::F32))
.and_then(|t| t.to_vec1::<f32>())
.map_err(|e| crate::MLError::ModelError(format!("Q-value stats readback: {}", e)))?;
Ok(QValueStats {
min: *stats.first().unwrap_or(&0.0) as f64,
max: *stats.get(1).unwrap_or(&0.0) as f64,
mean: *stats.get(2).unwrap_or(&0.0) as f64,
std: (*stats.get(3).unwrap_or(&0.0) as f64).sqrt(),
sample_count: count,
})
}
/// Create final training metrics
pub(crate) async fn create_final_metrics(
&self,
total_loss: f64,
total_q_value: f64,
total_gradient_norm: f64,
total_reward: f64,
num_epochs: usize,
training_duration: std::time::Duration,
early_stopped: bool,
total_action_counts: [usize; 5], // 5 exposure levels
total_factored_action_counts: [usize; 45], // 45 factored actions
) -> Result<TrainingMetrics> {
let final_loss = total_loss / num_epochs as f64;
let avg_q_value_final = total_q_value / num_epochs as f64;
let avg_grad_norm_final = total_gradient_norm / num_epochs as f64;
let avg_episode_reward = total_reward / num_epochs as f64;
let mut metrics = TrainingMetrics {
loss: final_loss,
accuracy: 0.0,
precision: 0.0,
recall: 0.0,
f1_score: 0.0,
training_time_seconds: training_duration.as_secs_f64(),
epochs_trained: num_epochs as u32,
convergence_achieved: final_loss < 1.0,
additional_metrics: std::collections::HashMap::new(),
};
metrics.add_metric("avg_q_value", avg_q_value_final);
metrics.add_metric("avg_gradient_norm", avg_grad_norm_final);
metrics.add_metric("final_epsilon", self.get_epsilon().await.unwrap_or(0.1));
metrics.add_metric("avg_episode_reward", avg_episode_reward);
// Action metrics: use factored 45-action space if branching, else 5 exposure levels
let total_factored: usize = total_factored_action_counts.iter().sum();
let total_exposure: usize = total_action_counts.iter().sum();
let total_actions = total_factored.max(total_exposure);
if total_actions > 0 {
// Factored 45-action diversity (primary metric when branching)
if total_factored > 0 {
let unique_factored = total_factored_action_counts.iter()
.filter(|&&count| count > 0).count();
let factored_diversity = (unique_factored as f64 / 45.0) * 100.0;
metrics.add_metric("action_diversity", factored_diversity);
metrics.add_metric("factored_unique_actions", unique_factored as f64);
metrics.add_metric("action_space_size", 45.0);
// Active factored actions (used >0.5% of the time)
let active_threshold = (total_factored as f64 * 0.005).max(1.0);
let active_count = total_factored_action_counts.iter()
.filter(|&&count| count as f64 >= active_threshold)
.count();
let active_diversity_pct = (active_count as f64 / 45.0) * 100.0;
metrics.add_metric("active_actions_count", active_count as f64);
metrics.add_metric("active_diversity_pct", active_diversity_pct);
// Top factored actions
let mut sorted_actions: Vec<(usize, usize)> = total_factored_action_counts.iter()
.enumerate()
.map(|(idx, &count)| (idx, count))
.collect();
sorted_actions.sort_by(|a, b| b.1.cmp(&a.1));
if let Some((top1_idx, top1_count)) = sorted_actions.first() {
let top1_pct = (*top1_count as f64 / total_factored as f64) * 100.0;
metrics.add_metric("top1_action_idx", *top1_idx as f64);
metrics.add_metric("top1_action_count", *top1_count as f64);
metrics.add_metric("top1_action_pct", top1_pct);
}
let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum();
let top5_coverage_pct = (top5_count as f64 / total_factored as f64) * 100.0;
metrics.add_metric("top5_coverage_pct", top5_coverage_pct);
} else {
// Fallback: 5 exposure levels
let unique_actions = total_action_counts.iter()
.filter(|&&count| count > 0).count();
let action_diversity = (unique_actions as f64 / 5.0) * 100.0;
metrics.add_metric("action_diversity", action_diversity);
metrics.add_metric("action_space_size", 5.0);
let active_threshold = (total_exposure as f64 * 0.005).max(1.0);
let active_count = total_action_counts.iter()
.filter(|&&count| count as f64 >= active_threshold)
.count();
let active_diversity_pct = (active_count as f64 / 5.0) * 100.0;
metrics.add_metric("active_actions_count", active_count as f64);
metrics.add_metric("active_diversity_pct", active_diversity_pct);
let mut sorted_actions: Vec<(usize, usize)> = total_action_counts.iter()
.enumerate()
.map(|(idx, &count)| (idx, count))
.collect();
sorted_actions.sort_by(|a, b| b.1.cmp(&a.1));
if let Some((top1_idx, top1_count)) = sorted_actions.first() {
let top1_pct = (*top1_count as f64 / total_exposure as f64) * 100.0;
metrics.add_metric("top1_action_idx", *top1_idx as f64);
metrics.add_metric("top1_action_count", *top1_count as f64);
metrics.add_metric("top1_action_pct", top1_pct);
}
let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum();
let top5_coverage_pct = (top5_count as f64 / total_exposure as f64) * 100.0;
metrics.add_metric("top5_coverage_pct", top5_coverage_pct);
}
metrics.add_metric("total_actions", total_actions as f64);
// Buy/sell/hold always from 5-exposure space (meaningful for P&L)
// 0=Short100, 1=Short50, 2=Flat, 3=Long50, 4=Long100
let sell_count: usize = total_action_counts.get(0).copied().unwrap_or(0)
+ total_action_counts.get(1).copied().unwrap_or(0);
let hold_count: usize = total_action_counts.get(2).copied().unwrap_or(0);
let buy_count: usize = total_action_counts.get(3).copied().unwrap_or(0)
+ total_action_counts.get(4).copied().unwrap_or(0);
metrics.add_metric("buy_count", buy_count as f64);
metrics.add_metric("sell_count", sell_count as f64);
metrics.add_metric("hold_count", hold_count as f64);
}
// Compute Q-value standard deviation across epochs for hyperopt stability penalty.
// self.q_value_history stores per-epoch average Q-values; their std measures
// how much Q-values fluctuate during training (volatility indicator).
if self.q_value_history.len() >= 2 {
let n = self.q_value_history.len() as f64;
let mean = self.q_value_history.iter().sum::<f64>() / n;
let variance = self
.q_value_history
.iter()
.map(|&q| (q - mean).powi(2))
.sum::<f64>()
/ n;
let std_dev = variance.sqrt();
metrics.add_metric("q_value_std", std_dev);
} else {
// Not enough data points to compute std; default to 0.0 (no volatility signal)
metrics.add_metric("q_value_std", 0.0);
}
if early_stopped {
metrics.add_metric("early_stopped", 1.0);
}
Ok(metrics)
}
/// Epoch-end Q-value diagnostics: gap analysis + per-action averages.
///
/// Merges the former `compute_q_gap_for_epoch` and `compute_per_action_q_values`
/// into a single forward pass + readback, eliminating one redundant buffer sample,
/// forward pass, and `to_vec2` GPU-CPU transfer per epoch.
///
/// Returns (gap_stats, per_action_avgs) where:
/// - gap_stats: (mean_gap, min_gap, max_gap) of Q_best - Q_second_best
/// - per_action_avgs: `[f64; 5]` averages (one per exposure action)
pub(crate) async fn compute_epoch_q_diagnostics(&self) -> Option<(
(f64, f64, f64),
[f64; 5],
)> {
let agent = self.agent.read().await;
let buffer = agent.memory();
if buffer.len() < 10 {
return None;
}
// Use the larger sample size (200) for both diagnostics
let sample_size = buffer.len().min(200);
let batch_sample = match buffer.sample(sample_size) {
Ok(s) => s,
Err(_) => return None,
};
let state_dim = agent.get_state_dim();
// GPU PER path: use gpu_batch.states directly (experiences vec is empty)
#[allow(unused_mut)]
let mut batch_tensor_opt: Option<Tensor> = None;
{
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = gpu.states
.to_dtype(training_dtype(agent.device()))
.ok();
}
}
let batch_tensor = if let Some(t) = batch_tensor_opt {
t
} else {
let batched_states: Vec<f32> = batch_sample
.experiences
.iter()
.flat_map(|exp| {
let mut s = exp.state.clone();
s.resize(state_dim, 0.0);
s
})
.collect();
let t = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) { Ok(t) => t,
Err(_) => return None,
};
match t.to_dtype(training_dtype(agent.device())) {
Ok(t) => t,
Err(_) => return None,
}
};
// Single forward pass for both gap and per-action diagnostics
let batch_q_values = match agent.forward(&batch_tensor) {
Ok(q) => q,
Err(_) => return None,
};
// All paths use GPU-native diagnostics — no to_vec2 readback
if self.device.is_cuda() {
return Self::compute_q_diagnostics_gpu(&batch_q_values).ok();
}
// Non-CUDA: compute diagnostics via tensor ops (no to_vec2)
// Sort Q-values descending per row, gap = sorted[0] - sorted[1]
let sorted = match batch_q_values.sort_last_dim(false) {
Ok((s, _)) => s,
Err(_) => return None,
};
let n_actions = batch_q_values.dims().get(1).copied().unwrap_or(5);
if n_actions < 2 { return None; }
let best = match sorted.narrow(1, 0, 1) {
Ok(t) => t.flatten_all().unwrap_or(sorted.clone()),
Err(_) => return None,
};
let second_best = match sorted.narrow(1, 1, 1) {
Ok(t) => t.flatten_all().unwrap_or(sorted.clone()),
Err(_) => return None,
};
let gaps_tensor = match best.sub(&second_best) {
Ok(t) => t,
Err(_) => return None,
};
let gap_mean = gaps_tensor.mean_all().ok()?.to_scalar::<f32>().ok()? as f64;
let gap_min = gaps_tensor.min(0).ok()?.to_scalar::<f32>().ok()? as f64;
let gap_max = gaps_tensor.max(0).ok()?.to_scalar::<f32>().ok()? as f64;
// Per-action average Q-values via mean(dim=0)
let per_action = match batch_q_values.mean(0) {
Ok(t) => t,
Err(_) => return None,
};
let mut avgs = [0.0_f64; 5];
for i in 0..5_usize.min(n_actions) {
if let Ok(v) = per_action.narrow(0, i, 1).and_then(|t| t.to_scalar::<f32>()) {
avgs[i] = v as f64;
}
}
Some(((gap_mean, gap_min, gap_max), avgs))
}
/// Compute Q-value gap and per-action averages on GPU.
/// Returns (mean_gap, min_gap, max_gap, per_action_avgs[5]).
/// Single 8-float readback at epoch end.
fn compute_q_diagnostics_gpu(
q_values: &Tensor, // [batch, 5]
) -> candle_core::Result<((f64, f64, f64), [f64; 5])> {
// sort_last_dim returns (sorted_values, indices) — destructure the tuple
let (sorted, _indices) = q_values.sort_last_dim(true)?; // descending
let best = sorted.narrow(1, 0, 1)?;
let second = sorted.narrow(1, 1, 1)?;
let gaps = best.sub(&second)?;
// Compute gap stats on GPU, batch into one tensor for single readback
let gaps_flat = gaps.flatten_all()?;
let mean_gap = gaps_flat.mean_all()?;
let min_gap = gaps_flat.min(0)?;
let max_gap = gaps_flat.max(0)?;
// Per-action means: mean along batch dim [5]
let per_action = q_values.mean(0)?;
// Single readback: cat [mean, min, max, per_action_0..4] → 8 floats in one DMA
let gap_vec = Tensor::cat(
&[&mean_gap.unsqueeze(0)?, &min_gap.unsqueeze(0)?, &max_gap.unsqueeze(0)?], 0
)?;
let all_stats = Tensor::cat(&[&gap_vec, &per_action.flatten_all()?], 0)?
.to_dtype(candle_core::DType::F32)?
.to_vec1::<f32>()?;
let mean_g = *all_stats.first().unwrap_or(&0.0) as f64;
let min_g = *all_stats.get(1).unwrap_or(&0.0) as f64;
let max_g = *all_stats.get(2).unwrap_or(&0.0) as f64;
let mut avgs = [0.0_f64; 5];
for i in 0..5_usize {
avgs[i] = *all_stats.get(3 + i).unwrap_or(&0.0) as f64;
}
Ok(((mean_g, min_g, max_g), avgs))
}
/// Get current training metrics
pub async fn get_metrics(&self) -> TrainingMetrics {
self.metrics.read().await.clone()
}
/// Get Q-values for a given state
pub(crate) async fn get_q_values(&self, state: &TradingState) -> Result<Vec<f64>> {
let agent = self.agent.read().await;
let state_vec = state.to_vector();
let raw_dim = state_vec.len();
let aligned = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device);
let padded: Vec<f32> = if aligned > raw_dim {
let mut v = state_vec.to_vec();
v.resize(aligned, 0.0);
v
} else {
state_vec.to_vec()
};
let state_tensor = Tensor::new(&*padded, &self.device)?.unsqueeze(0)?; // Add batch dimension
let q_values_tensor = agent.forward(&state_tensor)?.squeeze(0)?;
// Single readback: download entire Q-value vector in one DMA
let q_f32 = q_values_tensor
.to_dtype(candle_core::DType::F32)?
.to_vec1::<f32>()?;
Ok(q_f32.into_iter().map(|v| v as f64).collect())
}
/// Check if early stopping criteria are met
pub(crate) fn check_early_stopping(&self, avg_q_value: f64, epoch: usize) -> Option<String> {
if !self.hyperparams.early_stopping_enabled
|| epoch + 1 < self.hyperparams.min_epochs_before_stopping
{
return None;
}
// Criterion 1: Q-value floor check
if avg_q_value < self.hyperparams.q_value_floor {
return Some(format!(
"Q-value {:.4} below floor threshold {:.4}",
avg_q_value, self.hyperparams.q_value_floor
));
}
// C4 FIX: Criterion 2 — Sharpe plateau check (was val-loss plateau).
// Sharpe directly measures trading quality. A plateau means the model
// has stopped improving its trading strategy, even if TD-loss still moves.
if self.sharpe_history.len() >= self.hyperparams.plateau_window {
let window = self.hyperparams.plateau_window;
let recent_sharpes: Vec<f64> = self
.sharpe_history
.iter()
.rev()
.take(window)
.copied()
.collect();
if let (Some(&newest), Some(&oldest)) = (recent_sharpes.first(), recent_sharpes.last()) {
// For Sharpe, improvement = newest - oldest (higher is better)
let improvement = newest - oldest;
if improvement < 0.01 {
let msg = if improvement < -0.01 {
format!(
"Sharpe worsening detected (delta: {:.4}, window: {})",
improvement,
window
)
} else {
format!(
"Sharpe plateau detected (improvement: {:.4}, window: {})",
improvement,
window
)
};
return Some(msg);
}
}
}
None
}
#[allow(unused_variables, unreachable_code, unused_mut)]
pub(crate) async fn compute_validation_loss(&mut self) -> Result<f64> {
if self.val_data.is_empty() {
return Ok(0.0);
}
// Save current epsilon and force to 0 for deterministic evaluation
let original_epsilon = self.get_epsilon().await?;
self.set_epsilon(0.0).await?; // Pure greedy selection
let sample_size = self.val_data.len().min(1000); // Sample up to 1000 for speed
let aligned_state_dim = {
let agent = self.agent.read().await;
agent.get_state_dim()
};
// ── GPU-resident validation data (lazy init, uploaded once) ──
// Pre-upload val features [sample_size, 42], close prices, and OFI features to GPU.
// Subsequent epochs reuse the same GPU tensors — zero CPU loop.
if self.val_features_gpu.is_none() {
let mut flat_features = Vec::with_capacity(sample_size * 42);
let mut current_closes = Vec::with_capacity(sample_size);
let mut next_closes = Vec::with_capacity(sample_size);
for (feature_vec, target) in self.val_data.iter().take(sample_size) {
for &v in feature_vec.iter() {
flat_features.push(v as f32);
}
let cur = if target.len() >= 2 { target[0] } else { feature_vec[3] };
let nxt = if target.len() >= 2 { target[1] } else { cur };
current_closes.push(cur as f32);
next_closes.push(nxt as f32);
}
self.val_features_gpu = Some(
Tensor::from_vec(flat_features, (sample_size, 42), &self.device)
.map_err(|e| anyhow::anyhow!("GPU val features upload: {e}"))?
);
let cur_t = Tensor::from_vec(current_closes, &[sample_size], &self.device)
.map_err(|e| anyhow::anyhow!("GPU val current_closes upload: {e}"))?;
let nxt_t = Tensor::from_vec(next_closes, &[sample_size], &self.device)
.map_err(|e| anyhow::anyhow!("GPU val next_closes upload: {e}"))?;
self.val_closes_gpu = Some((cur_t, nxt_t));
// Upload OFI features for validation range if available
if let Some(ref ofi) = self.ofi_features {
let mut flat_ofi = Vec::with_capacity(sample_size * 8);
for i in 0..sample_size {
let idx = self.ofi_val_offset + i;
if let Some(row) = ofi.get(idx) {
for &v in row.iter() {
flat_ofi.push(v as f32);
}
} else {
flat_ofi.extend_from_slice(&[0.0_f32; 8]);
}
}
self.val_ofi_gpu = Some(
Tensor::from_vec(flat_ofi, (sample_size, 8), &self.device)
.map_err(|e| anyhow::anyhow!("GPU val OFI upload: {e}"))?
);
}
}
// ── Build state tensor on GPU: cat [features, portfolio, ofi] ──
let features_gpu = self.val_features_gpu.as_ref()
.ok_or_else(|| anyhow::anyhow!("val_features_gpu must be initialized"))?;
// Portfolio features: 3 scalars from current tracker, broadcast to [sample_size, 3]
// Use first validation sample's close price for portfolio evaluation
let val_price = self.val_data.first()
.map(|(fv, tgt)| if tgt.len() >= 2 { tgt[0] as f32 } else { fv[3] as f32 })
.unwrap_or(0.0);
let portfolio_f = self.portfolio_tracker.get_portfolio_features(val_price);
let portfolio_gpu = Tensor::new(&portfolio_f[..], &self.device)
.map_err(|e| anyhow::anyhow!("GPU portfolio tensor: {e}"))?
.unsqueeze(0)
.map_err(|e| anyhow::anyhow!("GPU portfolio unsqueeze: {e}"))?
.broadcast_left(sample_size)
.map_err(|e| anyhow::anyhow!("GPU portfolio broadcast: {e}"))?;
// Concat features + portfolio (+ OFI if enabled) → [sample_size, raw_state_dim]
let raw_state_dim = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 };
let state_gpu = if let Some(ref ofi_gpu) = self.val_ofi_gpu {
Tensor::cat(&[features_gpu, &portfolio_gpu, ofi_gpu], 1)
.map_err(|e| anyhow::anyhow!("GPU val state cat (with OFI): {e}"))?
} else {
Tensor::cat(&[features_gpu, &portfolio_gpu], 1)
.map_err(|e| anyhow::anyhow!("GPU val state cat: {e}"))?
};
// Pad to aligned dim if needed (trailing zeros for tensor core alignment)
let batch_tensor = if aligned_state_dim > raw_state_dim {
let pad_width = aligned_state_dim - raw_state_dim;
let pad = Tensor::zeros((sample_size, pad_width), candle_core::DType::F32, &self.device)
.map_err(|e| anyhow::anyhow!("GPU val pad zeros: {e}"))?;
Tensor::cat(&[&state_gpu, &pad], 1)
.map_err(|e| anyhow::anyhow!("GPU val state pad: {e}"))?
.to_dtype(training_dtype(&self.device))
.map_err(|e| anyhow::anyhow!("GPU val state dtype: {e}"))?
} else {
state_gpu
.to_dtype(training_dtype(&self.device))
.map_err(|e| anyhow::anyhow!("GPU val state dtype: {e}"))?
};
// Close prices already on GPU from lazy init
let (ref val_current_closes_t, ref val_next_closes_t) = self.val_closes_gpu.as_ref()
.ok_or_else(|| anyhow::anyhow!("val_closes_gpu must be initialized"))?;
let agent = self.agent.read().await;
let batch_q_values = agent.forward(&batch_tensor)
.map_err(|e| anyhow::anyhow!("Batched validation forward pass failed: {}", e))?;
// Get branching Q-values if branching is enabled
let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> =
if self.hyperparams.use_branching {
agent
.batch_branching_q_values(&batch_tensor)
.map_err(|e| anyhow::anyhow!("Validation branching Q-values failed: {e}"))?
} else {
None
};
drop(agent); // Release lock early
// Fused GPU greedy action selection (epsilon=0.0) + GPU routing.
{
if self.gpu_action_selector.is_none() && self.device.is_cuda() {
let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new(
&self.device,
self.hyperparams.batch_size.max(sample_size).max(8192),
0xDEAD_BEEF_CAFE_u64,
).map_err(|e| anyhow::anyhow!("Validation GPU action selector init failed: {e}"))?;
self.gpu_action_selector = Some(selector);
}
let selector = self.gpu_action_selector.as_mut()
.ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?;
let factored_tensor = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors {
selector
.select_actions_branching(q_exp, q_ord, q_urg, 0.0)
.map_err(|e| anyhow::anyhow!("Validation GPU branching select failed: {e}"))?
} else {
let exposure_tensor = selector
.select_actions(&batch_q_values, 0.0, sample_size, 5)
.map_err(|e| anyhow::anyhow!("Validation GPU greedy select failed: {e}"))?;
selector.route_exposure_to_factored(
&exposure_tensor, sample_size,
self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32,
self.vol_ema as f32, self.median_vol as f32,
).map_err(|e| anyhow::anyhow!("Validation GPU route failed: {e}"))?
};
// --- GPU PnL-based reward + Sharpe computation ---
// Instead of reading factored indices back to CPU and looping through
// the scalar reward function, compute PnL rewards entirely on GPU:
// reward_i = direction_i * (next_close_i - current_close_i) / current_close_i
// Then compute Sharpe = mean(rewards) / std(rewards) * sqrt(252) on device
// with a single scalar readback for the final value.
// 1) Close prices already on GPU (pre-uploaded)
let current_closes_t: &Tensor = val_current_closes_t;
let next_closes_t: &Tensor = val_next_closes_t;
// 2) Extract exposure index from factored tensor:
// factored_index = exposure * 9 + order * 3 + urgency
// => exposure_idx = factored_index / 9
let factored_f32 = factored_tensor
.to_dtype(candle_core::DType::F32)
.map_err(|e| anyhow::anyhow!("GPU val factored to f32: {e}"))?;
let nine = Tensor::new(&[9.0_f32], &self.device)
.map_err(|e| anyhow::anyhow!("GPU val nine const: {e}"))?;
let exposure_idx_f32 = factored_f32
.broadcast_div(&nine)
.map_err(|e| anyhow::anyhow!("GPU val exposure div: {e}"))?
.floor()
.map_err(|e| anyhow::anyhow!("GPU val exposure floor: {e}"))?;
// 3) Map exposure index to direction multiplier via lookup table:
// [0=Short100→-1.0, 1=Short50→-0.5, 2=Flat→0.0, 3=Long50→0.5, 4=Long100→1.0]
let direction_lut = Tensor::new(
&[-1.0_f32, -0.5, 0.0, 0.5, 1.0], &self.device,
).map_err(|e| anyhow::anyhow!("GPU val direction LUT: {e}"))?;
let exposure_idx_u32 = exposure_idx_f32
.to_dtype(candle_core::DType::U32)
.map_err(|e| anyhow::anyhow!("GPU val exposure to u32: {e}"))?;
let directions = direction_lut
.index_select(&exposure_idx_u32, 0)
.map_err(|e| anyhow::anyhow!("GPU val direction gather: {e}"))?;
// 4) Compute PnL-based rewards: direction * (next - current) / current
let price_returns = next_closes_t
.sub(&current_closes_t)
.map_err(|e| anyhow::anyhow!("GPU val price diff: {e}"))?
.broadcast_div(&current_closes_t)
.map_err(|e| anyhow::anyhow!("GPU val price returns div: {e}"))?;
let rewards = directions
.mul(&price_returns)
.map_err(|e| anyhow::anyhow!("GPU val rewards mul: {e}"))?;
// 5) Compute Sharpe on GPU: mean / std * sqrt(252) — single batched readback
let mean_t = rewards.mean_all()
.map_err(|e| anyhow::anyhow!("GPU val rewards mean: {e}"))?;
let var_t = rewards
.broadcast_sub(&mean_t)
.map_err(|e| anyhow::anyhow!("GPU val rewards center: {e}"))?
.sqr()
.map_err(|e| anyhow::anyhow!("GPU val rewards sqr: {e}"))?
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU val rewards var: {e}"))?;
let stats = Tensor::cat(&[&mean_t.unsqueeze(0)?, &var_t.unsqueeze(0)?], 0)
.and_then(|t| t.to_dtype(candle_core::DType::F32))
.and_then(|t| t.to_vec1::<f32>())
.map_err(|e| anyhow::anyhow!("GPU val Sharpe stats readback: {e}"))?;
let mean_scalar = *stats.first().unwrap_or(&0.0) as f64;
let var_scalar = *stats.get(1).unwrap_or(&0.0) as f64;
let std_val = var_scalar.sqrt();
let val_sharpe = if std_val > 1e-10 {
(mean_scalar / std_val) * (252.0_f64).sqrt()
} else {
0.0
};
// Restore original epsilon after evaluation
self.set_epsilon(original_epsilon).await?;
// Return negative Sharpe as the "loss" (lower = better Sharpe)
return Ok(-val_sharpe);
}
}
/// Estimate average Q-value from replay buffer samples for monitoring
///
/// WAVE 23 P0: Now includes Q-value divergence check (early stopping)
/// OPTIMIZATION: Batched Q-value estimation for 10× speedup via GPU parallelization
pub(crate) async fn estimate_avg_q_value_with_early_stopping(&self, agent: &mut DQNAgentType) -> Result<f64> {
// Get a few samples from the replay buffer to estimate Q-values
let buffer = agent.memory();
if buffer.len() == 0 {
return Ok(0.0);
}
// Sample up to 10 experiences for Q-value estimation
let sample_size = buffer.len().min(10);
let batch_sample = buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Failed to sample experiences: {}", e))?;
// GPU PER path: use gpu_batch.states directly (always active in CUDA builds)
let gpu_batch = batch_sample.gpu_batch.as_ref()
.ok_or_else(|| anyhow::anyhow!("GPU PER must be active — gpu_batch is None"))?;
let batch_tensor = gpu_batch.states.to_dtype(training_dtype(agent.device()))
.map_err(|e| anyhow::anyhow!("GPU Q-est states dtype cast: {}", e))?;
// WAVE 23 P0 Fix: Check for Q-value divergence (early stopping)
// This calls log_q_values() which returns Err if divergence detected for consecutive checks
agent.log_q_values(&batch_tensor)
.map_err(|e| {
tracing::info!("🛑 Early stopping triggered (Q-value divergence): {}", e);
anyhow::anyhow!("Early stopping: {}", e)
})?;
// Single forward pass for all samples (10× faster than sequential)
let batch_q_values = agent
.forward(&batch_tensor)
.map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?;
// Get max Q-value per sample across action dimension
let max_q_values = batch_q_values
.max(1)
.map_err(|e| anyhow::anyhow!("Failed to compute max Q-values: {}", e))?;
// Compute average across batch
let avg_q = max_q_values
.mean_all()
.map_err(|e| anyhow::anyhow!("Failed to compute mean Q-value: {}", e))?
.to_scalar::<f32>() .map_err(|e| anyhow::anyhow!("Failed to extract average Q-value: {}", e))?
as f64;
Ok(avg_q)
}
}

872
crates/ml/src/trainers/dqn/trainer/mod.rs Normal file
View File

@@ -0,0 +1,872 @@
//! DQN Trainer Implementation
//!
//! Main training loop and execution logic for Deep Q-Network.
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Context, Result};
use candle_core::{Device, Tensor};
use crate::cuda_pipeline::DqnGpuData;
use ml_core::fill_simulator::FillSimulator;
use risk::drawdown_monitor::DrawdownMonitor;
use risk::safety::position_limiter::HybridPositionLimiter;
use tokio::sync::RwLock;
use tracing::info;
use uuid::Uuid;
use crate::dqn::action_space::FactoredAction;
use crate::dqn::circuit_breaker::CircuitBreaker;
use crate::dqn::curiosity::CuriosityModule;
use crate::dqn::logging::{LoggingConfig, MetricsAggregator};
use crate::dqn::portfolio_tracker::PortfolioTracker;
use crate::dqn::reward::RewardFunction;
use crate::TrainingMetrics;
use crate::features::extraction::FeatureVector;
use crate::features::microstructure_features::*;
use crate::labeling::triple_barrier::TripleBarrierEngine;
use super::config::{DQNAgentType, DQNHyperparameters};
use super::statistics::{FeatureStatistics, QValueStats};
pub(super) use super::EPISODE_LENGTH;
mod action;
mod metrics;
mod state;
mod constructor;
mod training_loop;
mod train_step;
#[cfg(test)]
mod tests;
pub struct DQNTrainer {
/// DQN agent
pub(crate) agent: Arc<RwLock<DQNAgentType>>,
/// Training hyperparameters
pub(crate) hyperparams: DQNHyperparameters,
/// Device (GPU or CPU)
pub(crate) device: Device,
/// Training metrics
pub(crate) metrics: Arc<RwLock<TrainingMetrics>>,
/// Loss history for plateau detection
pub(crate) loss_history: Vec<f64>,
/// Q-value history for floor detection
pub(crate) q_value_history: Vec<f64>,
/// Best validation loss achieved so far
pub(crate) best_val_loss: f64,
/// Validation data for computing validation loss
pub(crate) val_data: Vec<(FeatureVector, Vec<f64>)>,
/// Validation loss history for early stopping
pub(crate) val_loss_history: Vec<f64>,
/// C4: Sharpe history for Sharpe-based early stopping
pub(crate) sharpe_history: Vec<f64>,
/// C4: Best Sharpe ratio achieved so far (higher = better)
pub(crate) best_sharpe: f64,
/// Epoch with best validation loss
pub(crate) best_epoch: usize,
/// Step counter for gradient logging (logs every 10 steps)
pub(crate) gradient_logging_step: usize,
/// Original buffer_size before AutoReplaySizer (for gradient collapse warmup)
pub(crate) collapse_warmup_buffer_size: usize,
/// Portfolio state tracker for P&L-based rewards (Bug #2 fix)
pub portfolio_tracker: PortfolioTracker,
/// Feature normalization statistics (WAVE 3 FIX #2)
/// None during stats collection phase (epochs 0-10), Some during normalization phase (epochs 11+)
pub feature_stats: Option<FeatureStatistics>,
/// Sliding window of recent actions for reward calculation (max 100)
pub(crate) recent_actions: VecDeque<FactoredAction>,
/// Reward function for calculating rewards with recent actions
pub(crate) reward_fn: RewardFunction,
// WAVE 16S: Adaptive Risk Management Components
/// Kelly criterion optimizer for position sizing (None if disabled)
pub(crate) kelly_optimizer: Option<Arc<crate::risk::kelly_optimizer::KellyCriterionOptimizer>>,
/// Trade history for Kelly calculation (wins/losses)
pub(crate) trade_history: VecDeque<f64>,
/// Volatility tracker for epsilon adjustment (None if disabled)
pub(crate) volatility_returns: VecDeque<f64>,
/// PnL history for Sharpe calculation (max 1000 entries)
pub(crate) pnl_history: VecDeque<f64>,
// Wave 16 Portfolio Features
/// Enable action masking (filters invalid actions before Q-value computation)
pub enable_action_masking: bool,
/// Maximum position size for action masking (default: 2.0)
pub max_position: f64,
/// Entropy regularizer for preventing policy collapse (None if disabled)
// entropy_regularizer removed — SAC-style entropy is computed directly on Q-value tensors in DQN::compute_loss_internal
/// Multi-asset portfolio tracker (None if single-asset mode)
pub multi_asset_portfolio: Option<Arc<crate::dqn::multi_asset::MultiAssetPortfolioTracker>>,
/// Stress tester for robustness validation (None if disabled)
pub stress_tester: Option<Arc<crate::dqn::stress_testing::DQNStressTester>>,
// Wave 16 Core Risk Features Integration
/// Drawdown monitor for tracking portfolio drawdowns (15% max drawdown)
pub drawdown_monitor: Option<Arc<DrawdownMonitor>>,
/// Position limiter with 3-tier limits (±10.0 absolute, 1M notional, 10% concentration)
pub position_limiter: Option<Arc<HybridPositionLimiter>>,
/// Circuit breaker for stopping training on consecutive failures
pub circuit_breaker: Option<Arc<CircuitBreaker>>,
// WAVE 3.10: Microstructure Feature Calculators (12 features)
pub(crate) micro_high_low_spread: HighLowSpread,
pub(crate) micro_vw_spread: VolumeWeightedSpread,
pub(crate) micro_tick_count: TickCount,
pub(crate) micro_inter_arrival: InterArrivalTime,
pub(crate) micro_buy_sell_imbalance: BuySellImbalance,
pub(crate) micro_kyle_lambda: KyleLambda,
pub(crate) micro_price_impact: PriceImpact,
pub(crate) micro_variance_ratio: VarianceRatio,
// Note: Roll Measure, Corwin-Schultz, Amihud, VPIN already exist in ml/src/microstructure/
// We'll integrate those in the update logic
/// Track last timestamp for inter-arrival time calculation
pub(crate) last_timestamp_ns: u64,
/// Track last close price for microstructure calculations
pub(crate) last_close: f64,
// WAVE 1.1: Triple Barrier Integration
/// Triple barrier engine for position exit labeling
pub(crate) triple_barrier: Arc<RwLock<TripleBarrierEngine>>,
/// Active position tracker ID (None = no active position)
pub(crate) active_position_tracker: Option<Uuid>,
/// WAVE P3: Track previous simulated position for barrier tracking continuity
pub(crate) previous_simulated_position: f32,
// WAVE 1.2: Safety Infrastructure Integration (8 Systems)
/// Loss history window for spike detection (size: 30)
pub(crate) safety_loss_history: VecDeque<f32>,
/// Loss plateau counter for anomaly detection
pub(crate) safety_loss_plateau_counter: usize,
/// Action counts for diversity monitoring (5 exposure actions)
pub(crate) safety_action_counts: std::collections::HashMap<usize, usize>,
/// Memory manager for GPU OOM risk monitoring
pub(crate) safety_memory_manager: Arc<RwLock<crate::safety::memory_manager::SafeMemoryManager>>,
/// Safety enforcement level (Strict/Normal/Permissive)
pub(crate) safety_level: crate::safety::SafetyLevel,
/// Step counter for periodic safety checks
pub(crate) safety_step_counter: usize,
/// Optional path to feature cache directory for faster hyperopt
pub(crate) feature_cache_dir: Option<PathBuf>,
/// Previous epoch's mean Q-value for overestimation detection
pub(crate) prev_epoch_q_mean: f64,
/// Current effective tau (may be temporarily increased if Q-values drift)
pub(crate) adaptive_tau: f64,
/// WAVE 24 (Agent 17): Patience-based early stopping for anti-overfitting
pub(crate) early_stopping: super::early_stopping::EarlyStopping,
/// WAVE 26 P0.6: Learning rate scheduler with warmup
pub(crate) lr_scheduler: super::lr_scheduler::LRScheduler,
/// WAVE 26 P1.8: Curiosity module for intrinsic rewards (None if curiosity_weight = 0.0)
pub(crate) curiosity_module: Option<CuriosityModule>,
// WAVE 26 P1: Advanced DQN Features Integration
// P1.3: Sharpe Ratio Reward Component
/// Rolling buffer of returns for Sharpe ratio calculation (max: sharpe_window)
pub(crate) returns_history: VecDeque<f64>,
/// Sharpe reward weight (0.0 = disabled)
pub(crate) sharpe_weight: f64,
// P1.6: Adaptive Dropout Scheduling
/// Optional dropout scheduler (None if disabled)
pub(crate) dropout_scheduler: Option<crate::dqn::network::DropoutScheduler>,
// P1.7: Hindsight Experience Replay (HER)
/// Optional HER buffer (None if her_ratio = 0.0)
pub(crate) her_buffer: Option<Arc<crate::dqn::hindsight_replay::HindsightReplayBuffer>>,
// P1.9: Generalized Advantage Estimation (GAE)
/// Optional GAE calculator (None if disabled)
pub(crate) gae_calculator: Option<crate::dqn::gae::GAECalculator>,
// P1.11: Noisy Network Sigma Scheduling
/// Optional noisy sigma scheduler (None if disabled)
pub(crate) noisy_sigma_scheduler: Option<crate::dqn::noisy_sigma_scheduler::NoisySigmaScheduler>,
// WAVE 30: Structured Logging Integration
/// Logging configuration for training metrics
pub(crate) logging_config: LoggingConfig,
/// Metrics aggregator for windowed training statistics
pub(crate) metrics_aggregator: MetricsAggregator,
// WAVE 44: Multi-step returns integration
/// N-step buffer for multi-step TD learning (None if n_steps=1)
pub(crate) nstep_buffer: Option<crate::dqn::nstep_buffer::NStepBuffer>,
/// Current effective batch size (may be reduced by OOM recovery)
pub(crate) current_batch_size: usize,
/// Cached Q-value estimate for periodic monitoring (avoids extra forward pass every step)
pub(crate) cached_avg_q: f64,
/// Counter for Q-value estimation frequency (estimate every N training steps)
pub(crate) q_estimation_counter: u64,
/// Pre-uploaded GPU training data (set once, reused across epochs)
pub(crate) gpu_data: Option<DqnGpuData>,
/// GPU portfolio simulator for CUDA-accelerated experience collection
pub(crate) gpu_portfolio_sim: Option<crate::cuda_pipeline::gpu_portfolio::GpuPortfolioSimulator>,
/// Raw cudarc targets buffer for CUDA kernel (parallel to candle Tensor in gpu_data)
pub(crate) targets_raw_cuda: Option<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>>,
/// Raw cudarc features buffer for CUDA experience kernel [num_bars * 42]
pub(crate) features_raw_cuda: Option<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>>,
/// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b)
pub(crate) gpu_experience_collector: Option<crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector>,
/// GPU-fused epsilon-greedy action selector (eliminates argmax GPU->CPU sync barrier)
pub(crate) gpu_action_selector: Option<crate::cuda_pipeline::gpu_action_selector::GpuActionSelector>,
/// GPU training guard for zero-sync safety checks (loss clip, NaN, grad collapse)
pub(crate) training_guard: Option<crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard>,
/// GPU monitoring reducer — accumulates reward/action stats across kernel launches
pub(crate) gpu_monitoring: Option<crate::cuda_pipeline::gpu_monitoring::GpuMonitoringReducer>,
/// Reusable GPU staging buffers for zero-alloc fold transitions
pub(crate) buffer_pool: Option<crate::cuda_pipeline::GpuBufferPool>,
/// Double-buffered GPU data for zero-downtime fold transitions
pub(crate) double_buffer: Option<crate::cuda_pipeline::double_buffer::DoubleBufferedLoader>,
/// GPU-resident walk-forward data (entire dataset on GPU, per-fold views via index ranges)
pub(crate) gpu_walk_forward: Option<crate::cuda_pipeline::gpu_walk_forward::GpuWalkForwardData>,
/// Multi-GPU configuration for data-parallel training (None = single GPU)
pub(crate) multi_gpu: Option<crate::cuda_pipeline::multi_gpu::MultiGpuConfig>,
/// Cached GPU n_episodes (computed once from nvidia-smi, reused across epochs)
/// Avoids forking nvidia-smi subprocess every epoch (~5-10ms per fork).
pub(crate) cached_n_episodes: Option<i32>,
/// Pre-computed OFI features per bar (indexed by global bar position).
/// Populated during data loading when MBP-10 order book data is available.
/// Passed as `regime_features` in `TradingState::from_normalized()`.
/// Arc-shared to avoid 2.68 GB copy per hyperopt trial (41.9M × 8 × 8 bytes).
pub(crate) ofi_features: Option<Arc<[[f64; 8]]>>,
/// Number of training bars (OFI offset for validation data).
/// val_data[i] corresponds to ofi_features[ofi_val_offset + i].
pub(crate) ofi_val_offset: usize,
/// GPU-resident validation features [sample_size, 42] — pre-uploaded once, reused per epoch
pub(crate) val_features_gpu: Option<Tensor>,
/// GPU-resident validation close prices (current, next) — pre-uploaded once
pub(crate) val_closes_gpu: Option<(Tensor, Tensor)>,
/// GPU-resident validation OFI features [sample_size, 8] — pre-uploaded once
pub(crate) val_ofi_gpu: Option<Tensor>,
// Phase C: Fill simulation and smart order routing
/// Fill simulator for order type-dependent execution modeling
pub(crate) fill_simulator: FillSimulator,
/// EMA of bar volatility (|close log return|) for OrderRouter routing decisions
pub(crate) vol_ema: f64,
/// Running median volatility estimate (slowly adapting EMA)
pub(crate) median_vol: f64,
/// Fused CUDA training context: pre-allocated batch buffers + single entry point
/// for CUDA Graph capture. Only active for Standard DQN on CUDA devices.
/// Lazy-initialized on first training step; dropped and recreated if batch_size changes.
pub(crate) fused_ctx: Option<super::fused_training::FusedTrainingCtx>,
}
impl std::fmt::Debug for DQNTrainer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DQNTrainer")
.field("hyperparams", &self.hyperparams)
.finish_non_exhaustive()
}
}
impl DQNTrainer {
/// Create new DQN trainer with hyperparameters and debug logging disabled
pub fn new(hyperparams: DQNHyperparameters) -> Result<Self> {
Self::new_with_debug(hyperparams, false)
}
/// Create new DQN trainer with a specific compute device.
/// Used by hyperopt to share a single CUDA context across parallel trials.
pub fn new_with_device(hyperparams: DQNHyperparameters, device: Device) -> Result<Self> {
Self::new_internal(hyperparams, false, Some(device))
}
/// Create new DQN trainer with hyperparameters and configurable debug logging
///
/// # Arguments
/// * `hyperparams` - DQN training hyperparameters
/// * `debug_logging` - Enable debug logging (REWARD_DEBUG, gradient norms, etc.)
pub fn new_with_debug(hyperparams: DQNHyperparameters, debug_logging: bool) -> Result<Self> {
Self::new_internal(hyperparams, debug_logging, None)
}
/// Two-phase stress tester initialization.
///
/// Call this after constructing `DQNTrainer` to resolve the circular dependency:
/// `DQNStressTester::new()` requires a `DQNTrainer`, so the tester cannot be
/// created *during* trainer construction. This method builds a lightweight
/// inner trainer (with stress testing itself disabled to avoid recursion) and
/// hands it to `DQNStressTester::new()`.
///
/// No-op if `hyperparams.enable_stress_testing` is `false`.
pub fn init_stress_tester(&mut self) -> Result<()> {
if !self.hyperparams.enable_stress_testing {
return Ok(());
}
// Clone hyperparams with stress testing disabled so the inner trainer
// does not recursively try to initialise its own stress tester.
// Also disable GPU-heavy features that the stress tester doesn't need —
// otherwise we double the VRAM usage (3 extra regime heads, 3 extra GPU PER buffers,
// experience collector, etc.) which causes OOM on 4GB GPUs.
let mut inner_hp = self.hyperparams.clone();
inner_hp.enable_stress_testing = false;
inner_hp.enable_regime_qnetwork = false;
inner_hp.use_per = false;
inner_hp.buffer_size = 1;
inner_hp.enable_gpu_experience_collector = false;
let inner_trainer = Self::new_with_device(inner_hp, self.device.clone())
.context("Failed to create inner DQNTrainer for stress tester")?;
let tester = crate::dqn::stress_testing::DQNStressTester::new(inner_trainer)?;
self.stress_tester = Some(Arc::new(tester));
info!("Stress testing enabled (8 scenarios)");
Ok(())
}
/// Get a reference to the double-buffered loader, if GPU is active.
pub fn double_buffer(&self) -> Option<&crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
self.double_buffer.as_ref()
}
/// Get a mutable reference to the double-buffered loader, if GPU is active.
pub fn double_buffer_mut(&mut self) -> Option<&mut crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
self.double_buffer.as_mut()
}
/// Set feature cache directory for faster hyperopt
///
/// Enables loading pre-computed features from disk instead of recomputing them
pub fn with_feature_cache(mut self, cache_dir: PathBuf) -> Self {
self.feature_cache_dir = Some(cache_dir);
self
}
/// Train DQN on market data from DBN files
///
/// # Arguments
///
/// * `dbn_data_dir` - Directory containing DBN files (e.g., "test_data/real/databento/ml_training/")
/// * `checkpoint_callback` - Callback for saving checkpoints (epoch, model_data, is_final) -> `Result<String>`
///
/// # Returns
///
/// Training metrics (loss, accuracy, gradient norms, Q-values)
pub async fn train<F>(
&mut self,
dbn_data_dir: &str,
checkpoint_callback: F,
) -> Result<TrainingMetrics>
where
F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
{
info!(
"Starting DQN training for {} epochs with batch size {}",
self.hyperparams.epochs, self.hyperparams.batch_size
);
// Load market data from DBN files (ALL data for walk-forward or single-pass)
let (training_data, val_data) = self.load_training_data(dbn_data_dir).await?;
info!(
"Loaded {} training samples, {} validation samples",
training_data.len(),
val_data.len()
);
// GPU walk-forward: upload ALL data to GPU, run expanding-window folds
if self.hyperparams.enable_gpu_walk_forward && self.device.is_cuda() {
// Merge train+val into a single dataset for walk-forward splitting
let mut all_data = training_data;
all_data.extend(val_data);
info!(
"GPU walk-forward enabled: {} total bars, uploading to VRAM",
all_data.len(),
);
return self.train_walk_forward(&all_data, checkpoint_callback).await;
}
// Standard single-pass training
self.ofi_val_offset = training_data.len();
self.val_data = val_data;
self.val_features_gpu = None;
self.val_closes_gpu = None;
self.val_ofi_gpu = None;
self.train_with_data_full_loop(&training_data, checkpoint_callback)
.await
}
/// Train with preloaded data (skips disk I/O and feature extraction).
///
/// Accepts pre-split training and validation data that was loaded once and
/// cached across hyperopt trials. This avoids re-reading 36 `.dbn.zst` files
/// and re-extracting 42 features on every trial, eliminating minutes of GPU
/// idle time at each trial boundary.
///
/// # Arguments
///
/// * `training_data` - Pre-extracted (features, targets) for training split
/// * `val_data` - Pre-extracted (features, targets) for validation split
/// * `checkpoint_callback` - Checkpoint save callback
///
/// # Returns
///
/// Training metrics from the completed run
///
/// # Errors
///
/// Returns error if the training loop fails
pub async fn train_with_preloaded_data<F>(
&mut self,
training_data: Vec<(FeatureVector, Vec<f64>)>,
val_data: Vec<(FeatureVector, Vec<f64>)>,
checkpoint_callback: F,
) -> Result<TrainingMetrics>
where
F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
{
info!(
"Starting DQN training with preloaded data: {} train, {} val samples",
training_data.len(),
val_data.len()
);
// Store validation data for loss computation
self.ofi_val_offset = training_data.len();
self.val_data = val_data;
self.val_features_gpu = None;
self.val_closes_gpu = None;
self.val_ofi_gpu = None;
// Use the common training loop (Wave 12 Group 3 refactor)
self.train_with_data_full_loop(&training_data, checkpoint_callback)
.await
}
/// Train with shared preloaded data (zero-copy for hyperopt).
///
/// Same as [`train_with_preloaded_data`] but accepts `Arc`-wrapped data,
/// avoiding a ~150 MB deep clone per hyperopt trial.
pub async fn train_with_shared_data<F>(
&mut self,
training_data: &[(FeatureVector, Vec<f64>)],
val_data: Vec<(FeatureVector, Vec<f64>)>,
checkpoint_callback: F,
) -> Result<TrainingMetrics>
where
F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
{
info!(
"Starting DQN training with shared data: {} train, {} val samples",
training_data.len(),
val_data.len()
);
self.ofi_val_offset = training_data.len();
self.val_data = val_data;
self.val_features_gpu = None;
self.val_closes_gpu = None;
self.val_ofi_gpu = None;
self.train_with_data_full_loop(training_data, checkpoint_callback)
.await
}
/// Train with GPU-resident walk-forward cross-validation.
///
/// Uploads the ENTIRE dataset to GPU VRAM once, then runs expanding-window
/// walk-forward: each fold trains on [0..T], validates on [T..V], tests on
/// [V..E]. Fold transitions are zero-copy (index range changes only).
///
/// Returns the metrics from the LAST fold (most data, most representative).
pub async fn train_walk_forward<F>(
&mut self,
training_data: &[(FeatureVector, Vec<f64>)],
mut checkpoint_callback: F,
) -> Result<TrainingMetrics>
where
F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
{
use crate::cuda_pipeline::gpu_walk_forward::{GpuWalkForwardConfig, GpuWalkForwardData};
let wf_config = GpuWalkForwardConfig {
initial_train_fraction: self.hyperparams.wf_initial_train_fraction,
val_fraction: self.hyperparams.wf_val_fraction,
test_fraction: self.hyperparams.wf_test_fraction,
step_fraction: self.hyperparams.wf_step_fraction,
};
// Upload ALL data to GPU once
let wf_data = GpuWalkForwardData::upload(
training_data,
self.ofi_features.as_deref(),
&wf_config,
&self.device,
).map_err(|e| anyhow::anyhow!("GPU walk-forward upload: {e}"))?;
let num_folds = wf_data.num_folds();
if num_folds == 0 {
return Err(anyhow::anyhow!(
"Insufficient data for walk-forward: {} bars, need at least {} for one fold",
training_data.len(),
((wf_config.initial_train_fraction + wf_config.val_fraction + wf_config.test_fraction) * training_data.len() as f64) as usize,
));
}
info!(
"GPU walk-forward: {} folds, {:.1} MB VRAM, {} total bars",
num_folds, wf_data.vram_bytes as f64 / 1_048_576.0, wf_data.total_bars,
);
// Store GPU walk-forward data and cudarc buffers for the experience collector
self.features_raw_cuda = Some(wf_data.features);
self.targets_raw_cuda = Some(wf_data.targets);
let mut last_metrics = TrainingMetrics::new();
for fold_idx in 0..num_folds {
let fold = wf_data.folds.get(fold_idx).ok_or_else(|| {
anyhow::anyhow!("Fold {fold_idx} out of range")
})?;
info!(
"=== Walk-Forward Fold {}/{} === train: {} bars, val: {} bars, test: {} bars",
fold_idx + 1, num_folds, fold.train_len(), fold.val_len(), fold.test_len(),
);
// Split training_data into fold's train and val slices (for CPU-side data)
let fold_train = &training_data[fold.train_start..fold.train_end];
let fold_val: Vec<(FeatureVector, Vec<f64>)> =
training_data[fold.val_start..fold.val_end].to_vec();
// Store validation data for this fold
self.ofi_val_offset = fold.train_end;
self.val_data = fold_val;
// Reset training state for new fold
self.gpu_data = None; // Force re-upload via DqnGpuData for the fold's range
self.best_sharpe = f64::NEG_INFINITY;
self.best_val_loss = f64::INFINITY;
self.loss_history.clear();
self.q_value_history.clear();
self.val_loss_history.clear();
self.sharpe_history.clear();
// Run training loop on this fold's data
last_metrics = self
.train_with_data_full_loop(fold_train, &mut checkpoint_callback)
.await?;
info!(
"Fold {}/{} complete: loss={:.6}, epochs={}",
fold_idx + 1, num_folds,
last_metrics.loss,
last_metrics.epochs_trained,
);
}
// Clean up GPU walk-forward buffers (features/targets already stored in self)
self.gpu_walk_forward = None;
Ok(last_metrics)
}
/// Calculate adaptive bounds with margin
///
/// # Arguments
///
/// * `stats` - Q-value statistics from Phase 1
/// * `margin` - Safety margin as fraction (e.g., 0.3 = 30%)
///
/// # Returns
///
/// Tuple of (v_min, v_max) with safety margin applied
fn calculate_adaptive_bounds(stats: &QValueStats, margin: f64) -> (f64, f64) {
let range = stats.max - stats.min;
let v_min = stats.min - range * margin;
let v_max = stats.max + range * margin;
// Cap at ±10,000 to prevent explosion
(v_min.max(-10000.0), v_max.min(10000.0))
}
/// Reinitialize categorical distribution with new bounds
async fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), crate::MLError> {
let mut agent = self.agent.write().await;
match &mut *agent {
DQNAgentType::Standard(agent) => agent.reinit_categorical_distribution(v_min, v_max)?,
DQNAgentType::RegimeConditional(agent) => agent.reinit_categorical_distribution(v_min, v_max)?,
}
Ok(())
}
/// Calculate reward based on price movement
///
/// # Arguments
/// * `current_close` - Current bar's close price
/// * `next_close` - Next bar's close price (target)
///
/// # Returns
/// Normalized reward in [-1.0, 1.0] based on price change
fn calculate_reward(&self, current_close: f64, next_close: f64) -> f32 {
let price_change = next_close - current_close;
// Normalize by 10.0 for ES futures typical moves (±10 points)
// Clamp to [-1.0, 1.0] to prevent extreme rewards
(price_change / 10.0).clamp(-1.0, 1.0) as f32
}
/// Check if we can train (buffer has enough samples)
async fn can_train(&self) -> Result<bool> {
let agent = self.agent.read().await;
Ok(agent.can_train())
}
/// Get current epsilon value
async fn get_epsilon(&self) -> Result<f64> {
let agent = self.agent.read().await;
Ok(agent.get_epsilon() as f64)
}
/// Set epsilon value (used for deterministic evaluation)
async fn set_epsilon(&self, epsilon: f64) -> Result<()> {
let mut agent = self.agent.write().await;
agent.set_epsilon(epsilon);
Ok(())
}
/// Get best validation loss achieved during training
///
/// Returns the lowest validation loss seen across all epochs.
/// Used by hyperopt adapter to optimize for generalization.
pub fn get_best_val_loss(&self) -> f64 {
self.best_val_loss
}
/// Get epoch number where best validation loss was achieved
///
/// Returns the 1-indexed epoch number with the best validation loss.
pub fn get_best_epoch(&self) -> usize {
self.best_epoch
}
/// Get validation data for backtest integration
///
/// Returns a reference to the validation dataset for hyperopt backtest evaluation.
/// Each entry contains a FeatureVector (42 market + 3 portfolio = 45 dims) and the corresponding target values.
/// Used by hyperopt adapter to run backtests on unseen data after training.
pub fn get_val_data(&self) -> &[(FeatureVector, Vec<f64>)] {
&self.val_data
}
/// Update portfolio tracker to reflect current position from backtest engine.
///
/// Called between chunks so the next chunk's portfolio features
/// accurately reflect the current position (direction, value, exposure).
pub fn set_portfolio_for_backtest(
&mut self,
position_size: f32,
entry_price: f32,
current_price: f32,
) {
self.portfolio_tracker = PortfolioTracker::new(
self.portfolio_tracker.initial_capital(),
self.portfolio_tracker.spread(),
0.0,
);
if position_size.abs() > f32::EPSILON {
self.portfolio_tracker
.set_position_direct(position_size, entry_price, current_price);
}
}
/// Get the device used by this trainer
pub fn device(&self) -> &candle_core::Device {
&self.device
}
/// Get access to the DQN agent
///
/// Returns a reference to the Arc<RwLock<DQNAgentType>> for checkpoint saving.
/// Used by hyperopt adapter to save model weights after training.
pub fn get_agent(&self) -> &Arc<RwLock<DQNAgentType>> {
&self.agent
}
/// Get reference to training hyperparameters
///
/// Returns a reference to the DQN hyperparameters used for this trainer.
/// Used by tests to validate configuration.
pub fn hyperparams(&self) -> &DQNHyperparameters {
&self.hyperparams
}
/// Get current learning rate from scheduler
///
/// Returns the current learning rate after applying warmup and decay.
/// Used by tests and monitoring to track LR schedule.
pub fn get_current_lr(&self) -> f64 {
self.lr_scheduler.get_lr()
}
/// Serialize model to bytes with architecture metadata embedded in safetensors header.
///
/// For RegimeConditional agents, serializes ALL 3 heads (trending, ranging,
/// volatile) into a single safetensors file using prefixed tensor names
/// (`trending__`, `ranging__`, `volatile__`). This ensures walk-forward
/// checkpoint restore loads all heads, not just the trending head.
pub async fn serialize_model(&self) -> Result<Vec<u8>> {
let agent = self.agent.read().await;
let tensors: std::collections::HashMap<String, candle_core::Tensor> = match &*agent {
crate::trainers::dqn::DQNAgentType::RegimeConditional(regime) => {
let mut all_tensors = std::collections::HashMap::new();
for (prefix, head_opt) in [
("trending__", regime.get_trending_head()),
("ranging__", regime.get_ranging_head()),
("volatile__", regime.get_volatile_head()),
] {
let head = head_opt.ok_or_else(|| {
anyhow::anyhow!("Missing {} head for serialization", prefix)
})?;
let vars = head.get_q_network_vars();
let vars_data = vars.data().lock().map_err(|_| {
anyhow::anyhow!("Failed to lock VarMap for {} head", prefix)
})?;
for (name, var) in vars_data.iter() {
all_tensors.insert(
format!("{}{}", prefix, name),
var.as_tensor().clone(),
);
}
}
all_tensors
}
_ => {
let vars = agent.get_q_network_vars();
let vars_data = vars.data().lock().map_err(|_| {
anyhow::anyhow!("Failed to lock VarMap for serialization")
})?;
vars_data
.iter()
.map(|(name, var)| (name.clone(), var.as_tensor().clone()))
.collect()
}
};
// Embed architecture metadata in safetensors header
let arch_metadata = Some(agent.checkpoint_metadata());
let data = safetensors::serialize(&tensors, &arch_metadata)
.map_err(|e| anyhow::anyhow!("Failed to serialize safetensors: {}", e))?;
Ok(data)
}
/// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`).
///
/// The trainer's `train_epoch` lazily uploads data on first call.
/// Use this to provide data that was uploaded in advance by a
/// `DoubleBufferedLoader`, skipping the per-fold upload latency.
pub fn set_gpu_data(&mut self, data: DqnGpuData) {
info!(
"DqnTrainer: injected pre-uploaded GPU data ({} bars, {:.1} MB)",
data.num_bars,
data.vram_bytes() as f64 / 1_048_576.0,
);
self.gpu_data = Some(data);
}
/// Drop cached GPU data, freeing VRAM for the next fold.
pub fn clear_gpu_data(&mut self) {
if self.gpu_data.is_some() {
info!("DqnTrainer: cleared GPU data (VRAM freed)");
self.gpu_data = None;
}
}
/// BUG #38 FIX: Clear replay buffer of contaminated experiences
pub async fn clear_replay_buffer(&mut self) -> Result<()> {
let mut agent = self.agent.write().await;
agent.clear_replay_buffer().map_err(|e| {
anyhow::anyhow!("Failed to clear replay buffer: {}", e)
})?;
let buffer_size = agent.get_replay_buffer_size().unwrap_or(0);
info!("Replay buffer cleared successfully. Current size: {}", buffer_size);
Ok(())
}
/// BUG #38 FIX: Reset target network to match current network
pub async fn reset_target_network(&mut self) -> Result<()> {
let mut agent = self.agent.write().await;
agent.reset_target_network().map_err(|e| {
anyhow::anyhow!("Failed to reset target network: {}", e)
})?;
info!("Target network reset successfully");
Ok(())
}
/// Get per-epoch training loss history (for smoke test verification)
pub fn loss_history(&self) -> &[f64] {
&self.loss_history
}
/// Get per-epoch validation loss history
pub fn val_loss_history(&self) -> &[f64] {
&self.val_loss_history
}
/// Get current epsilon from the DQN agent
pub async fn get_agent_epsilon(&self) -> f32 {
let agent_lock = self.agent.read().await;
agent_lock.get_epsilon()
}
}
// ---------------------------------------------------------------------------
// GPU Q-value diagnostics (Task 6)
// ---------------------------------------------------------------------------

179
crates/ml/src/trainers/dqn/trainer/state.rs Normal file
View File

@@ -0,0 +1,179 @@
//! DQN Trainer — State/feature vector conversion
use anyhow::{Context, Result};
use candle_core::Tensor;
use common::CommonError;
use num_traits::ToPrimitive;
use super::DQNTrainer;
use crate::dqn::TradingState;
use crate::features::extraction::FeatureVector;
impl DQNTrainer {
/// Convert feature vector to TradingState (42 market features → 45-dim state with portfolio)
///
/// CRITICAL BUG FIX: Features 0-3 are LOG RETURNS (signed), not raw prices.
/// Using .abs() destroys directional information (bullish vs bearish).
/// We now use TradingState::from_normalized() to preserve sign information.
///
/// Feature mapping:
/// - Features 0-3: OHLC log returns → price_features (signed, normalized)
/// - Features 4-224: All other features → technical_indicators (221 features including Wave D)
///
/// # Arguments
///
/// * `feature_vec` - 42-dimensional FeatureVector from extraction pipeline
/// * `close_price` - Current close price for portfolio feature calculation (optional)
///
/// # Bug #4 Fix
///
/// Added close_price parameter to enable portfolio feature population from PortfolioTracker.
pub(crate) fn feature_vector_to_state(
&self,
feature_vec: &FeatureVector,
close_price: Option<rust_decimal::Decimal>,
) -> Result<TradingState> {
self.feature_vector_to_state_with_ofi(feature_vec, close_price, None)
}
pub(crate) fn feature_vector_to_state_with_ofi(
&self,
feature_vec: &FeatureVector,
close_price: Option<rust_decimal::Decimal>,
ofi_index: Option<usize>,
) -> Result<TradingState> {
// States are pre-normalized during data loading
let normalized_features: Vec<f32> = feature_vec.iter().map(|&v| v as f32).collect();
// Features 0-3 are LOG RETURNS - preserve sign information for price direction
let price_features: Vec<f32> = vec![
normalized_features[0], // open log return (can be negative)
normalized_features[1], // high log return (can be negative)
normalized_features[2], // low log return (can be negative)
normalized_features[3], // close log return (can be negative)
];
// 42-FEATURE ARCHITECTURE: Extract market features (indices 4-41)
assert_eq!(
normalized_features.len(),
42,
"Expected 42 market features (got {})",
normalized_features.len()
);
let market_features: Vec<f32> = normalized_features[4..42]
.iter()
.map(|&x| x as f32)
.collect();
// Legacy technical_indicators (empty for 42-feature architecture)
let technical_indicators = vec![];
// BUG #36 FIX: Use NORMALIZED portfolio features to prevent Q-value explosion
let portfolio_features = if let Some(price) = close_price {
let price_f32 = price.to_f32().unwrap_or(0.0);
self.portfolio_tracker
.get_portfolio_features(price_f32)
.to_vec()
} else {
vec![0.0, 0.0, 0.0] // Fallback if no price provided
};
// OFI regime features: 8 features from MBP-10 order book data.
// When OFI is enabled (mbp10_data_dir set), always return 8 features
// (zeros if data didn't load) to match state_dim=53.
let ofi_enabled = self.hyperparams.mbp10_data_dir.is_some();
let regime_features: Vec<f32> = if let (Some(ofi), Some(idx)) = (&self.ofi_features, ofi_index) {
ofi.get(idx)
.map(|f| f.iter().map(|&v| v as f32).collect())
.unwrap_or_else(|| vec![0.0; 8])
} else if ofi_enabled {
vec![0.0; 8]
} else {
vec![]
};
// Use from_normalized() to preserve sign information
Ok(TradingState::from_normalized(
price_features,
technical_indicators,
market_features,
portfolio_features,
regime_features,
))
}
/// Convert feature vector to state tensor for action selection
///
/// Public wrapper around internal state conversion for hyperopt backtest integration.
/// Converts a 42-dimensional feature vector to a 45-dimensional state tensor
/// suitable for DQN agent's select_action method.
///
/// # Arguments
///
/// * `feature_vec` - 42-dimensional market feature vector
/// * `close_price` - Current close price for portfolio feature calculation
///
/// # Returns
///
/// Result containing the 45-dimensional state tensor ready for model inference.
/// Portfolio features (last 3 dimensions) are populated via PortfolioTracker.
pub fn convert_to_state(
&self,
feature_vec: &FeatureVector,
close_price: f64,
) -> Result<Tensor> {
let close = rust_decimal::Decimal::try_from(close_price)
.map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?;
// Use internal conversion method (returns TradingState)
let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?;
// Convert TradingState to flat vector, pad for tensor core alignment
let state_vec = trading_state.to_vector();
let raw_dim = state_vec.len();
let aligned = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device);
let padded: Vec<f32> = if aligned > raw_dim {
let mut v = state_vec.to_vec();
v.resize(aligned, 0.0);
v
} else {
state_vec.to_vec()
};
// Convert to Tensor using trainer's device (GPU or CPU)
Tensor::new(padded.as_slice(), &self.device)
.context("Failed to create state tensor from TradingState")
}
/// Convert feature vector to flat state Vec<f32> (CPU only, no GPU tensor).
///
/// Same as `convert_to_state` but returns the raw vector instead of a GPU tensor.
/// Used by chunked batch inference to avoid per-bar GPU allocations.
pub fn convert_to_state_vec(
&self,
feature_vec: &FeatureVector,
close_price: f64,
) -> Result<Vec<f32>> {
let close = rust_decimal::Decimal::try_from(close_price)
.map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?;
let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?;
Ok(trading_state.to_vector())
}
/// Convert feature vector to flat state Vec<f32> with OFI features at the given index.
///
/// Same as `convert_to_state_vec` but injects OFI features from the preloaded
/// array at `ofi_index`, preventing train/eval feature mismatch.
pub fn convert_to_state_vec_with_ofi(
&self,
feature_vec: &FeatureVector,
close_price: f64,
ofi_index: usize,
) -> Result<Vec<f32>> {
let close = rust_decimal::Decimal::try_from(close_price)
.map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?;
let trading_state = self.feature_vector_to_state_with_ofi(feature_vec, Some(close), Some(ofi_index))?;
Ok(trading_state.to_vector())
}
}

609
crates/ml/src/trainers/dqn/trainer/tests.rs Normal file
View File

@@ -0,0 +1,609 @@
use super::*;
use crate::hyperopt::ParameterSpace;
use std::sync::OnceLock;
/// Shared CUDA device across all trainer tests.
///
/// Root cause fix for CUBLAS_STATUS_NOT_INITIALIZED cascades: each
/// `Device::cuda_if_available(0)` creates a new cuBLAS handle. With
/// 400+ tests doing this in rapid succession (even with --test-threads=1),
/// the driver's internal handle pool is exhausted. Sharing one device
/// eliminates the churn entirely.
static SHARED_DEVICE: OnceLock<Device> = OnceLock::new();
fn shared_cuda_device() -> Device {
// Initialize tracing so kernel compilation/launch logs are visible.
static TRACING_INIT: std::sync::Once = std::sync::Once::new();
TRACING_INIT.call_once(|| {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
tracing_subscriber::fmt()
.with_env_filter(filter)
.with_test_writer()
.try_init()
.ok();
});
SHARED_DEVICE
.get_or_init(|| {
Device::cuda_if_available(0).unwrap_or(Device::Cpu)
})
.clone()
}
// Helper function to create test hyperparameters
// Uses conservative defaults suitable for testing
fn create_test_params() -> DQNHyperparameters {
let mut params = DQNHyperparameters::conservative();
// Production default: branching DQN (3 heads: exposure, order, urgency).
// Always enabled — the warp-cooperative kernel on H100 requires it.
params.use_branching = true;
params.hidden_dim_base = Some(32); // Small for fast test iterations
params.buffer_size = 10_000;
params
}
fn create_test_trainer() -> Result<DQNTrainer> {
DQNTrainer::new_with_device(create_test_params(), shared_cuda_device())
}
fn create_test_trainer_with(params: DQNHyperparameters) -> Result<DQNTrainer> {
DQNTrainer::new_with_device(params, shared_cuda_device())
}
/// Pad a TradingState's regime_features so that `state.dimension()` matches the
/// trainer's aligned state_dim (e.g. 45→48 on CUDA due to tensor core alignment).
fn pad_state_to_aligned(state: &mut TradingState, trainer: &DQNTrainer) {
let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(
state.dimension(),
&trainer.device,
);
let pad = aligned_dim.saturating_sub(state.dimension());
if pad > 0 {
state.regime_features.extend(vec![0.0_f32; pad]);
}
}
#[tokio::test]
async fn test_dqn_trainer_creation() {
let hyperparams = create_test_params();
let trainer = create_test_trainer_with(hyperparams);
assert!(
trainer.is_ok(),
"Failed to create DQN trainer: {:?}",
trainer.err()
);
}
#[tokio::test]
async fn test_batch_size_validation() {
let mut hyperparams = create_test_params();
hyperparams.batch_size = 500;
// VRAM ceiling clamps if needed, never rejects
let trainer = create_test_trainer_with(hyperparams);
assert!(
trainer.is_ok(),
"Should clamp oversized batch, not reject: {:?}",
trainer.err()
);
}
#[tokio::test]
async fn test_feature_vector_to_state() {
let hyperparams = create_test_params();
let trainer = create_test_trainer_with(hyperparams).unwrap();
// Create a synthetic 42-dim feature vector (42 market features)
let mut feature_vec = [0.0; 42];
feature_vec[0] = 4000.0; // open
feature_vec[1] = 4010.0; // high
feature_vec[2] = 3990.0; // low
feature_vec[3] = 4005.0; // close
feature_vec[4] = 1000.0; // volume
// Fill remaining features with synthetic data
for i in 5..42 {
feature_vec[i] = (i as f64) * 0.1;
}
let close_price =
rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price));
assert!(
state.is_ok(),
"Failed to convert feature vector: {:?}",
state.err()
);
let state = state.unwrap();
// State dimension is 45 (42 market + 3 portfolio)
// - Market features: 0-41 (42 features)
// - Portfolio features: 42-44 (3 features, populated by PortfolioTracker)
assert_eq!(
state.dimension(),
45,
"State dimension should be 45 (42 market + 3 portfolio features)"
);
}
#[tokio::test]
async fn test_batched_action_selection() {
let hyperparams = create_test_params();
let mut trainer = create_test_trainer_with(hyperparams).unwrap();
// Create multiple synthetic states for batched action selection
let batch_size = 10;
let mut states = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let mut feature_vec = [0.0; 42]; // 42 market features
// Create varied states for testing
feature_vec[0] = 4000.0 + (i as f64 * 10.0); // open
feature_vec[1] = 4010.0 + (i as f64 * 10.0); // high
feature_vec[2] = 3990.0 + (i as f64 * 10.0); // low
feature_vec[3] = 4005.0 + (i as f64 * 10.0); // close
feature_vec[4] = 1000.0 + (i as f64 * 100.0); // volume
// Fill remaining features
for j in 5..42 {
feature_vec[j] = (j as f64 + i as f64) * 0.1;
}
let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
.unwrap_or(rust_decimal::Decimal::ZERO);
let mut state = trainer
.feature_vector_to_state(&feature_vec, Some(close_price))
.unwrap();
pad_state_to_aligned(&mut state, &trainer);
states.push(state);
}
// Test batched action selection
let actions_result = trainer.select_actions_batch(&states).await;
assert!(
actions_result.is_ok(),
"Batched action selection failed: {:?}",
actions_result.err()
);
let actions = actions_result.unwrap();
assert_eq!(
actions.len(),
batch_size,
"Expected {} actions, got {}",
batch_size,
actions.len()
);
// Verify all actions have valid exposure indices (0-4)
for (i, action) in actions.iter().enumerate() {
let exp_idx = action.exposure as usize;
assert!(
exp_idx < 5,
"Action {} has invalid exposure index {}: {:?}",
i,
exp_idx,
action
);
}
}
#[tokio::test]
async fn test_batched_vs_sequential_action_selection_consistency() {
let hyperparams = create_test_params();
let mut trainer = create_test_trainer_with(hyperparams).unwrap();
// Create test states
let batch_size = 5;
let mut states = Vec::with_capacity(batch_size);
for i in 0..batch_size {
let mut feature_vec = [0.0; 42]; // 42 market features
feature_vec[0] = 4000.0 + (i as f64 * 50.0);
feature_vec[1] = 4050.0 + (i as f64 * 50.0);
feature_vec[2] = 3950.0 + (i as f64 * 50.0);
feature_vec[3] = 4025.0 + (i as f64 * 50.0);
feature_vec[4] = 5000.0 + (i as f64 * 500.0);
for j in 5..42 {
feature_vec[j] = (j as f64) * 0.5 + (i as f64);
}
let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
.unwrap_or(rust_decimal::Decimal::ZERO);
let mut state = trainer
.feature_vector_to_state(&feature_vec, Some(close_price))
.unwrap();
pad_state_to_aligned(&mut state, &trainer);
states.push(state);
}
// Get batched actions (GPU-optimized)
let batched_actions = trainer.select_actions_batch(&states).await.unwrap();
// Both should return valid actions
assert_eq!(
batched_actions.len(),
batch_size,
"Batched action count mismatch"
);
// Verify all actions have valid exposure indices (0-4)
for action in &batched_actions {
let exp_idx = action.exposure as usize;
assert!(exp_idx < 5, "Invalid exposure index {}: {:?}", exp_idx, action);
}
}
#[tokio::test]
async fn test_empty_batch_handling() {
let hyperparams = create_test_params();
let mut trainer = create_test_trainer_with(hyperparams).unwrap();
let empty_states: Vec<TradingState> = Vec::new();
let result = trainer.select_actions_batch(&empty_states).await;
assert!(result.is_ok(), "Empty batch should be handled gracefully");
assert_eq!(
result.unwrap().len(),
0,
"Empty batch should return empty actions"
);
}
#[tokio::test]
async fn test_zero_batch_size_handling() {
// Test DQN rejects zero batch size
let mut hyperparams = create_test_params();
hyperparams.batch_size = 0;
let result = create_test_trainer_with(hyperparams);
// Should fail with descriptive error
assert!(
result.is_err(),
"DQN should reject zero batch size, but got: {:?}",
result
);
// Error message should mention batch size
let error_msg = result.unwrap_err().to_string();
assert!(
error_msg.to_lowercase().contains("batch"),
"Error message should mention batch size, got: {}",
error_msg
);
}
// ===== Agent 23 Test #6: Batch Size Mismatch Validation Tests =====
/// Production-critical test: Verify trainer handles batch smaller than configured
#[tokio::test]
async fn test_batch_size_mismatch_smaller_than_configured() {
let mut hyperparams = create_test_params();
hyperparams.batch_size = 32;
let mut trainer = create_test_trainer_with(hyperparams).unwrap();
// Create batch with 16 states (half of configured 32)
let mut feature_vec = [0.0; 42]; // 42 market features
for i in 0..4 {
feature_vec[i] = 4000.0 + (i as f64 * 10.0);
}
for i in 5..42 {
feature_vec[i] = (i as f64) * 0.1;
}
let close_price =
rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
let mut state = trainer
.feature_vector_to_state(&feature_vec, Some(close_price))
.unwrap();
pad_state_to_aligned(&mut state, &trainer);
let smaller_batch = vec![state.clone(); 16];
let result = trainer.select_actions_batch(&smaller_batch).await;
assert!(
result.is_ok(),
"DQN should handle smaller batches: {:?}",
result.err()
);
assert_eq!(
result.unwrap().len(),
16,
"Should return action for each state"
);
}
/// Production-critical test: Verify trainer handles batch larger than configured
#[tokio::test]
async fn test_batch_size_mismatch_larger_than_configured() {
let mut hyperparams = create_test_params();
hyperparams.batch_size = 16;
let mut trainer = create_test_trainer_with(hyperparams).unwrap();
// Create batch with 64 states (4x configured 16)
let mut feature_vec = [0.0; 42]; // 42 market features
for i in 0..4 {
feature_vec[i] = 4000.0 + (i as f64 * 10.0);
}
for i in 5..42 {
feature_vec[i] = (i as f64) * 0.1;
}
let close_price =
rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
let mut state = trainer
.feature_vector_to_state(&feature_vec, Some(close_price))
.unwrap();
pad_state_to_aligned(&mut state, &trainer);
let larger_batch = vec![state.clone(); 64];
let result = trainer.select_actions_batch(&larger_batch).await;
assert!(
result.is_ok(),
"DQN should handle larger batches: {:?}",
result.err()
);
assert_eq!(
result.unwrap().len(),
64,
"Should return action for each state"
);
}
/// Production-critical test: Verify empty batch handling
#[tokio::test]
async fn test_empty_batch_returns_empty_actions() {
let mut trainer = create_test_trainer().unwrap();
let empty_batch: Vec<TradingState> = vec![];
let result = trainer.select_actions_batch(&empty_batch).await;
assert!(result.is_ok(), "Should handle empty batch gracefully");
assert_eq!(
result.unwrap().len(),
0,
"Empty batch should return empty actions"
);
}
/// Production-critical test: Verify single-sample batch handling
#[tokio::test]
async fn test_single_sample_batch() {
let mut hyperparams = create_test_params();
hyperparams.batch_size = 32;
let mut trainer = create_test_trainer_with(hyperparams).unwrap();
let mut feature_vec = [0.0; 42]; // 42 market features
for i in 0..4 {
feature_vec[i] = 4000.0;
}
for i in 5..42 {
feature_vec[i] = (i as f64) * 0.1;
}
let close_price =
rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO);
let mut state = trainer
.feature_vector_to_state(&feature_vec, Some(close_price))
.unwrap();
pad_state_to_aligned(&mut state, &trainer);
let single_batch = vec![state];
let result = trainer.select_actions_batch(&single_batch).await;
assert!(
result.is_ok(),
"Should handle single-sample batch: {:?}",
result.err()
);
assert_eq!(result.unwrap().len(), 1, "Should return exactly one action");
}
/// Large batch sizes are accepted (VRAM ceiling is the only cap)
#[test]
fn test_large_batch_size_accepted() {
let mut hyperparams = create_test_params();
hyperparams.batch_size = 2048;
let result = create_test_trainer_with(hyperparams);
assert!(
result.is_ok(),
"Should accept large batch sizes within VRAM ceiling: {:?}",
result.err()
);
}
/// Production-critical test: Non-power-of-2 batch sizes
#[tokio::test]
async fn test_non_power_of_two_batch_size() {
let mut hyperparams = create_test_params();
hyperparams.batch_size = 13; // Not a power of 2
let result = create_test_trainer_with(hyperparams);
assert!(
result.is_ok(),
"Should accept non-power-of-2 batch sizes: {:?}",
result.err()
);
}
/// Production-critical test: Train with empty dataset doesn't crash
#[tokio::test]
async fn test_train_with_empty_data_completes_gracefully() {
let mut params = create_test_params();
params.epochs = 5; // Short run — just checking it doesn't panic
params.early_stopping_enabled = false;
params.gradient_collapse_patience = 1000;
params.buffer_size = 1000;
let device = candle_core::Device::new_cuda(0).expect("CUDA device required");
let mut trainer = DQNTrainer::new_with_device(params, device).unwrap();
let empty_data: Vec<(FeatureVector, Vec<f64>)> = vec![];
let checkpoint_callback = |_, _, _| Ok(String::new());
let result = trainer
.train_with_data_full_loop(&empty_data, checkpoint_callback)
.await;
assert!(
result.is_err(),
"Training with empty data should return an error (no CPU fallback)"
);
}
/// Test reward function calculates actual price changes correctly
#[test]
fn test_reward_function_price_changes() {
let trainer = create_test_trainer().unwrap();
// Test upward price move (+14.25 points, should clamp to +1.0)
let reward_up = trainer.calculate_reward(5900.0, 5914.25);
assert!(
(reward_up - 1.0).abs() < 1e-6,
"Upward move should return +1.0 (clamped), got: {}",
reward_up
);
// Test downward price move (-14.25 points, should clamp to -1.0)
let reward_down = trainer.calculate_reward(5914.25, 5900.0);
assert!(
(reward_down - (-1.0)).abs() < 1e-6,
"Downward move should return -1.0 (clamped), got: {}",
reward_down
);
// Test flat market (0 points, should return 0.0)
let reward_flat = trainer.calculate_reward(5900.0, 5900.0);
assert!(
reward_flat.abs() < 1e-6,
"Flat market should return 0.0, got: {}",
reward_flat
);
// Test small upward move (+5 points, should return +0.5)
let reward_small_up = trainer.calculate_reward(5900.0, 5905.0);
assert!(
(reward_small_up - 0.5).abs() < 1e-6,
"Small upward move (+5) should return +0.5, got: {}",
reward_small_up
);
// Test small downward move (-5 points, should return -0.5)
let reward_small_down = trainer.calculate_reward(5905.0, 5900.0);
assert!(
(reward_small_down - (-0.5)).abs() < 1e-6,
"Small downward move (-5) should return -0.5, got: {}",
reward_small_down
);
// Test unclamped move (+3 points, should return +0.3)
let reward_unclamped = trainer.calculate_reward(5900.0, 5903.0);
assert!(
(reward_unclamped - 0.3).abs() < 1e-6,
"Move of +3 points should return +0.3, got: {}",
reward_unclamped
);
}
#[test]
fn test_dynamic_batch_size_l4() {
// L4 has 24GB VRAM — HardwareBudget should allow batch_size >> 230
let budget = crate::hyperopt::HardwareBudget {
gpu_memory_mb: 24_000,
gpu_name: "NVIDIA L4".to_string(),
};
let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0);
assert!(batch.unwrap_or(0.0) > 230.0, "L4 should support DQN batch > 230, got {:?}", batch);
}
#[test]
fn test_dynamic_batch_size_h100() {
// H100 has 80GB VRAM — should hit the 8192 ceiling
let budget = crate::hyperopt::HardwareBudget {
gpu_memory_mb: 81_920,
gpu_name: "NVIDIA H100".to_string(),
};
let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0);
assert!((batch.unwrap_or(0.0) - 8192.0).abs() < 1.0, "H100 should hit 8192 ceiling, got {:?}", batch);
}
// ── C2 Overhaul Smoke Tests ─────────────────────────────────────────
/// Verify DQN action space is 5 exposure levels (not 45 factored actions).
#[test]
fn test_c2_dqn_default_num_actions_is_5() {
let config = crate::dqn::DQNConfig::default();
assert_eq!(config.num_actions, 5, "DQN default must be 5 exposure-level actions");
}
/// Verify 5 exposure indices produce 5 distinct exposure levels.
#[test]
fn test_c2_five_actions_produce_distinct_exposures() {
use crate::dqn::action_space::ExposureLevel;
use crate::dqn::order_router::OrderRouter;
let actions: Vec<_> = (0..5)
.filter_map(|idx| ExposureLevel::from_index(idx).ok())
.map(|e| OrderRouter::route_default(e))
.collect();
assert_eq!(actions.len(), 5);
let unique: std::collections::HashSet<_> = actions.iter().map(|a| a.exposure).collect();
assert_eq!(unique.len(), 5, "All 5 exposure levels must be distinct");
}
/// Verify hyperopt search space is 29D (C4: sharpe_weight, L2: branch_hidden_dim).
#[test]
fn test_c3_search_space_is_27d() {
let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds();
assert_eq!(bounds.len(), 30, "Search space must be 30D (C6: gradient_accumulation_steps added)");
let names = crate::hyperopt::adapters::dqn::DQNParams::param_names();
assert_eq!(names.len(), 30);
assert!(names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient must be in search space (C3)");
assert!(names.contains(&"sharpe_weight"), "sharpe_weight must be in search space (C4)");
assert!(names.contains(&"branch_hidden_dim"), "branch_hidden_dim must be in search space (L2)");
assert!(!names.contains(&"curiosity_weight"), "curiosity_weight must not be in search space");
assert!(!names.contains(&"noisy_epsilon_floor"), "noisy_epsilon_floor must not be in search space");
}
/// Verify noisy_epsilon_floor is fixed to 0.10 (prevents action collapse).
#[test]
fn test_noisy_epsilon_floor_fixed() {
let params = crate::hyperopt::adapters::dqn::DQNParams::default();
assert!(
(params.noisy_epsilon_floor - 0.10).abs() < 1e-6,
"noisy_epsilon_floor must default to 0.10 (prevents action collapse)"
);
}
/// Verify exploration params are fixed after C2 cleanup.
#[test]
fn test_c2_exploration_params_fixed() {
let params = crate::hyperopt::adapters::dqn::DQNParams::default();
assert!(
params.curiosity_weight.abs() < f64::EPSILON,
"curiosity_weight must be fixed at 0.0"
);
assert!(
(params.noisy_epsilon_floor - 0.10).abs() < 1e-6,
"noisy_epsilon_floor must be fixed at 0.10"
);
assert!(
params.count_bonus_coefficient.abs() < f64::EPSILON,
"count_bonus_coefficient must be fixed at 0.0"
);
// Roundtrip through from_continuous should preserve fixed values
let continuous = params.to_continuous();
let recovered = crate::hyperopt::adapters::dqn::DQNParams::from_continuous(&continuous).unwrap();
assert!(
recovered.curiosity_weight.abs() < f64::EPSILON,
"curiosity_weight must remain 0.0 after roundtrip"
);
assert!(
recovered.count_bonus_coefficient.abs() < f64::EPSILON,
"count_bonus_coefficient must remain 0.0 after roundtrip"
);
}

631
crates/ml/src/trainers/dqn/trainer/train_step.rs Normal file
View File

@@ -0,0 +1,631 @@
//! DQN training step methods — single-batch and gradient-accumulation paths.
use anyhow::Result;
use candle_core::{IndexOp, Tensor};
use tracing::{debug, info, warn};
use crate::dqn::mixed_precision::training_dtype;
use super::DQNTrainer;
impl DQNTrainer {
/// Perform one training step using real DQN algorithm
///
/// This method implements the core Deep Q-Learning algorithm:
/// 1. Sample batch from experience replay buffer
/// 2. Compute current Q-values: Q(s, a)
/// 3. Compute target Q-values: r + γ * max_a' Q_target(s', a')
/// 4. Calculate TD-error and MSE loss
/// 5. Backpropagate gradients and update Q-network
/// 6. Periodically update target network
///
/// WAVE 26 P2.2: Now supports gradient accumulation for larger effective batch sizes
///
/// Returns: (loss, avg_q_value, grad_norm)
async fn train_step(&mut self) -> Result<(f64, f64, f64)> {
let accumulation_steps = self.hyperparams.gradient_accumulation_steps;
// OOM recovery loop: retry up to 3 times with halved batch size
const MAX_OOM_RETRIES: usize = 3;
for retry in 0..=MAX_OOM_RETRIES {
let result = if accumulation_steps > 1 {
self.train_step_with_accumulation().await
} else {
self.train_step_single_batch().await
};
match result {
Ok(metrics) => return Ok(metrics),
Err(e) => {
// Check if this is an OOM error by inspecting the error chain
let err_str = format!("{:?}", e).to_lowercase();
let is_oom = err_str.contains("out of memory")
|| err_str.contains("oom")
|| err_str.contains("cuda error 2")
|| err_str.contains("cudamalloc")
|| err_str.contains("failed to allocate");
if is_oom && retry < MAX_OOM_RETRIES {
let old_batch = self.current_batch_size;
self.current_batch_size = (old_batch / 2).max(1);
warn!(
"OOM detected (retry {}/{}): reducing batch size {} -> {}",
retry + 1,
MAX_OOM_RETRIES,
old_batch,
self.current_batch_size
);
// Continue to next retry
} else {
return Err(e);
}
}
}
}
Err(anyhow::anyhow!(
"Training failed after {} OOM retries",
MAX_OOM_RETRIES
))
}
/// Standard single-batch training step (no gradient accumulation)
///
/// Pre-samples from PER buffer using READ lock before acquiring WRITE lock
/// for GPU training. This preserves PER IS-weights and indices for proper
/// importance sampling correction and priority updates.
async fn train_step_single_batch(&mut self) -> Result<(f64, f64, f64)> {
// Pre-sample batch OUTSIDE the write lock using read-only access to the buffer.
// PER IS-weights and indices are preserved for correct importance sampling.
// The write lock is only held during GPU forward/backward + optimizer step.
let explicit_batch = {
let agent = self.agent.read().await;
let buffer = agent.memory();
let sample_size = self.current_batch_size;
buffer.can_sample(sample_size).then(|| {
buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Failed to pre-sample batch: {}", e))
}).transpose()?
}; // READ lock released here
let mut agent = self.agent.write().await;
// train_step returns GpuTrainResult with GPU-resident scalar tensors.
#[allow(unused_variables)]
let gpu_result = agent
.train_step(explicit_batch)
.map_err(|e| anyhow::anyhow!("Training step failed: {}", e))?;
// GPU training guard: on-device NaN/loss-clip/grad-collapse checks.
// Zero cudaStreamSynchronize — kernel writes halt flags to pinned host memory.
let (loss_clipped, grad_norm) = {
// Lazy-init training guard on first call
if self.training_guard.is_none() && self.device.is_cuda() {
match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::new(&self.device) {
Ok(guard) => {
info!("GPU training guard initialized");
self.training_guard = Some(guard);
}
Err(e) => {
return Err(anyhow::anyhow!("GPU training guard init FAILED (no CPU fallback): {e}"));
}
}
}
if let Some(ref mut guard) = self.training_guard {
let grad_collapse_threshold =
self.hyperparams.learning_rate as f32
* self.hyperparams.gradient_collapse_multiplier as f32;
// Use original buffer_size (before AutoReplaySizer) for warmup guard
let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64;
let past_warmup = self.gradient_logging_step as u64 > warmup_steps;
let result = guard
.check_and_accumulate(
&gpu_result.loss_gpu,
&gpu_result.grad_norm_gpu,
1e6_f32, // loss clip threshold
grad_collapse_threshold,
!past_warmup,
)
.map_err(|e| anyhow::anyhow!("GPU guard check: {e}"))?;
// Handle halt conditions
if result.halt_nan {
return Err(anyhow::anyhow!(
"NaN/Inf detected in loss ({}) or grad_norm ({})",
result.raw_loss,
result.raw_grad_norm
));
}
if result.halt_loss_clip {
warn!(
"Loss clipped from {:.2e} to 1.0e6 (TD error explosion, epoch {})",
result.raw_loss,
self.loss_history.len() + 1
);
}
// GPU guard path: collapse check only (no detect_dead_neurons GPU->CPU sync).
// Dead neuron detection runs at epoch boundary via log_diagnostics().
agent.check_gradient_collapse(result.raw_grad_norm).map_err(|e| {
tracing::info!("Early stopping triggered (gradient collapse): {}", e);
anyhow::anyhow!("Early stopping: {}", e)
})?;
(result.clipped_loss as f64, result.raw_grad_norm as f64)
} else {
return Err(anyhow::anyhow!(
"GPU training guard not initialized — CUDA device required for DQN training"
));
}
};
// Unreachable in non-cuda mode (return above), but Rust still name-checks.
// Q-value estimation: periodic (every 50 steps)
self.q_estimation_counter += 1;
if self.q_estimation_counter % 50 == 1 {
// GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback)
let mut gpu_q_done = false;
{
if let Some(ref mut guard) = self.training_guard {
let buffer = agent.memory();
if buffer.len() > 0 {
let sample_size = buffer.len().min(10);
let batch_sample = buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?;
let state_dim = agent.get_state_dim();
let mut batch_tensor_opt: Option<Tensor> = None;
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = Some(
gpu.states
.to_dtype(training_dtype(agent.device()))
.map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?,
);
}
if batch_tensor_opt.is_none() {
let mut state_data =
Vec::with_capacity(sample_size * state_dim);
for exp in &batch_sample.experiences {
state_data.extend_from_slice(&exp.state);
}
if !state_data.is_empty() {
let tensor = Tensor::from_vec( state_data,
(sample_size, state_dim),
&self.device,
)
.map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))?
.to_dtype(training_dtype(&self.device))
.map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?;
batch_tensor_opt = Some(tensor);
}
}
if let Some(ref batch_tensor) = batch_tensor_opt {
// Suppress forward() monitoring to avoid to_vec2 GPU→CPU sync
agent.set_training_forward_active(true);
let batch_q_values = agent
.forward(batch_tensor)
.map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?;
agent.set_training_forward_active(false);
let num_actions =
batch_q_values.dims().get(1).copied().unwrap_or(5);
// Divergence check on first sample
let first_q = batch_q_values
.i(0)
.map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?;
let div_result = guard
.qvalue_divergence(&first_q, num_actions, 10000.0)
.map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?;
agent
.log_q_values_from_stats(
div_result.q_min,
div_result.q_max,
div_result.q_mean,
div_result.q_variance,
num_actions,
)
.map_err(|e| {
tracing::info!(
"Early stopping (Q-value divergence): {}",
e
);
anyhow::anyhow!("Early stopping: {}", e)
})?;
// Batch average via GPU reduction (one-step delay due to double-buffering)
let stats = guard
.qvalue_stats(&batch_q_values, sample_size, num_actions)
.map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?;
self.cached_avg_q = stats.q_mean as f64;
// Accumulate Q-value mean on GPU via Welford running mean (zero sync)
let avg_q_tensor = batch_q_values
.max(1)
.map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))?
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?;
guard
.accumulate_q_value(&avg_q_tensor)
.map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?;
gpu_q_done = true;
}
}
}
}
// CUDA: GPU Q-value accumulation is mandatory — no CPU fallback.
if !gpu_q_done {
return Err(anyhow::anyhow!(
"GPU Q-value accumulation FAILED (no CPU fallback). \
Check GpuTrainingGuard initialization."
));
}
}
let avg_q_value = self.cached_avg_q;
debug!("Gradient norm after clip (actual): {:.4}", grad_norm);
self.gradient_logging_step += 1;
if self.gradient_logging_step % 10 == 0 {
debug!(
"Step {}: grad={:.4}, loss={:.4}",
self.gradient_logging_step, grad_norm, loss_clipped
);
}
Ok((loss_clipped, avg_q_value, grad_norm))
}
/// Training step with true gradient accumulation across N mini-batches.
///
/// Unlike the previous implementation which ran N independent optimizer
/// steps, this version computes gradients for each mini-batch, accumulates
/// them, averages, and then applies a **single** optimizer step. This
/// simulates training with an effective batch size of
/// `accumulation_steps * batch_size` while keeping memory usage at
/// `batch_size`.
///
/// Returns: (avg_loss, avg_q_value, final_grad_norm)
async fn train_step_with_accumulation(&mut self) -> Result<(f64, f64, f64)> {
let accumulation_steps = self.hyperparams.gradient_accumulation_steps;
debug!(
"Starting true gradient accumulation with {} steps (effective batch: {})",
accumulation_steps,
self.current_batch_size * accumulation_steps
);
// Pre-sample ALL mini-batches using READ lock (no GPU contention).
let pre_sampled: Vec<Option<crate::dqn::replay_buffer_type::BatchSample>> = {
let agent = self.agent.read().await;
let buffer = agent.memory();
let sample_size = self.current_batch_size;
let mut batches = Vec::with_capacity(accumulation_steps);
for step_idx in 0..accumulation_steps {
batches.push(
buffer.can_sample(sample_size).then(|| {
buffer.sample(sample_size).map_err(|e| {
anyhow::anyhow!(
"Failed to pre-sample batch (accum step {}): {}",
step_idx,
e
)
})
}).transpose()?,
);
}
batches
}; // READ lock released
let mut agent = self.agent.write().await;
// === Phase 1: Accumulate gradients across N mini-batches ===
let mut accumulated_grads: Option<candle_core::backprop::GradStore> = None;
// Used by non-CUDA fallback and CUDA empty-tensor fallback paths.
#[allow(unused_mut, unused_assignments, unused_variables)]
let mut total_loss = 0.0_f64;
let mut all_td_errors = Vec::new();
let mut all_indices = Vec::new();
#[allow(unused_mut, unused_assignments, unused_variables)]
let mut final_grad_norm = 0.0_f32;
let mut gpu_td_errors: Vec<candle_core::Tensor> = Vec::new();
let mut gpu_indices: Vec<candle_core::Tensor> = Vec::new();
let mut gpu_loss_tensors: Vec<candle_core::Tensor> = Vec::new();
let mut gpu_grad_tensors: Vec<candle_core::Tensor> = Vec::new();
for (step, batch) in pre_sampled.into_iter().enumerate() {
// Compute forward pass + backward WITHOUT optimizer step
let result = agent
.compute_gradients(batch)
.map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?;
// Get vars for accumulation. Var is an Arc wrapper so cloning is cheap.
let vars: Vec<candle_core::Var> = agent
.optimizer_vars()
.map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?;
crate::gradient_accumulation::accumulate_grads(
&mut accumulated_grads,
result.grads,
&vars,
)
.map_err(|e| anyhow::anyhow!("Gradient accumulation step {} failed: {}", step, e))?;
all_td_errors.extend(result.td_errors);
all_indices.extend(result.indices);
// GPU guard: check + accumulate loss/grad for this sub-step (borrows
// tensors before the move into gpu_*_tensors below).
{
if let (Some(ref loss_gpu), Some(ref gn_gpu)) =
(&result.loss_tensor_gpu, &result.grad_norm_gpu)
{
if let Some(ref mut guard) = self.training_guard {
let grad_collapse_threshold =
self.hyperparams.learning_rate as f32
* self.hyperparams.gradient_collapse_multiplier as f32;
// Use original buffer_size (before AutoReplaySizer) for warmup guard
let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64;
let past_warmup = self.gradient_logging_step as u64 > warmup_steps;
let guard_result = guard.check_and_accumulate(
loss_gpu,
gn_gpu,
1e6_f32,
grad_collapse_threshold,
!past_warmup,
).map_err(|e| anyhow::anyhow!("GPU guard sub-step {}: {e}", step))?;
if guard_result.halt_nan {
return Err(anyhow::anyhow!(
"NaN/Inf at accumulation sub-step {}: loss={}, grad={}",
step, guard_result.raw_loss, guard_result.raw_grad_norm
));
}
if guard_result.halt_grad_collapse {
agent.check_gradient_collapse(guard_result.raw_grad_norm).map_err(|e| {
tracing::info!("Early stopping (gradient collapse): {}", e);
anyhow::anyhow!("Early stopping: {}", e)
})?;
}
}
}
}
{
if let Some(td_gpu) = result.td_errors_gpu {
gpu_td_errors.push(td_gpu);
}
if let Some(idx_gpu) = result.indices_gpu {
gpu_indices.push(idx_gpu);
}
if let Some(loss_gpu) = result.loss_tensor_gpu {
gpu_loss_tensors.push(loss_gpu);
}
if let Some(gn_gpu) = result.grad_norm_gpu {
gpu_grad_tensors.push(gn_gpu);
}
}
// CPU sentinel fallback (non-CUDA only)
}
// === Phase 2: Average and apply gradients (single optimizer step) ===
if let Some(ref mut grads) = accumulated_grads {
let vars: Vec<candle_core::Var> = agent
.optimizer_vars()
.map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?;
crate::gradient_accumulation::scale_grads(
grads,
&vars,
1.0 / accumulation_steps as f64,
)
.map_err(|e| anyhow::anyhow!("Gradient scaling failed: {}", e))?;
let guard_active = self.training_guard.is_some();
crate::gradient_accumulation::check_gradients_finite_guarded(
grads,
&vars,
guard_active,
).map_err(|e| anyhow::anyhow!("Training halted: {}", e))?;
agent
.apply_accumulated_gradients(grads)
.map_err(|e| anyhow::anyhow!("Apply accumulated gradients failed: {}", e))?;
}
// === Phase 3: Bookkeeping ===
{
// GPU PER path: concatenate GPU tensors and update in one shot
if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() {
let td_cat = candle_core::Tensor::cat(&gpu_td_errors, 0)
.map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?;
let idx_cat = candle_core::Tensor::cat(&gpu_indices, 0)
.map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?;
agent
.update_priorities_gpu(&idx_cat, &td_cat)
.map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?;
} else if !all_indices.is_empty() {
agent
.update_priorities(&all_indices, &all_td_errors)
.map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?;
} else {
// No priority updates needed (uniform buffer or empty batch)
}
}
agent.step_replay_buffer();
// Single readback at accumulation boundary — prefer GPU guard accumulators
// (zero extra sync), fall back to cat+mean+to_scalar if guard absent.
let (avg_loss, final_grad_norm_f64) = {
if let Some(ref mut guard) = self.training_guard {
let (avg_l, avg_gn) = guard.read_accumulators()
.map_err(|e| anyhow::anyhow!("GPU guard read_accumulators: {e}"))?;
guard.reset_accumulators()
.map_err(|e| anyhow::anyhow!("GPU guard reset: {e}"))?;
(avg_l, avg_gn)
} else {
return Err(anyhow::anyhow!(
"GPU training guard not initialized — CUDA device required"
));
}
};
// Gradient collapse detection (early stopping) -- no dead neuron GPU->CPU sync.
// Dead neuron detection runs at epoch boundary via log_diagnostics().
agent
.check_gradient_collapse(final_grad_norm_f64 as f32)
.map_err(|e| {
tracing::info!("Early stopping triggered (gradient collapse): {}", e);
anyhow::anyhow!("Early stopping: {}", e)
})?;
// Clip averaged loss
let loss_clipped = if avg_loss > 1e6 {
warn!("Averaged loss clipped from {:.2e} to 1.0e6", avg_loss);
1e6
} else {
avg_loss
};
// Q-value estimation: periodic (every 50 steps)
self.q_estimation_counter += 1;
if self.q_estimation_counter % 50 == 1 {
// GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback)
let mut gpu_q_done = false;
{
if let Some(ref mut guard) = self.training_guard {
let buffer = agent.memory();
if buffer.len() > 0 {
let sample_size = buffer.len().min(10);
let batch_sample = buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?;
let state_dim = agent.get_state_dim();
let mut batch_tensor_opt: Option<Tensor> = None;
if let Some(ref gpu) = batch_sample.gpu_batch {
batch_tensor_opt = Some(
gpu.states
.to_dtype(training_dtype(agent.device()))
.map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?,
);
}
if batch_tensor_opt.is_none() {
let mut state_data =
Vec::with_capacity(sample_size * state_dim);
for exp in &batch_sample.experiences {
state_data.extend_from_slice(&exp.state);
}
if !state_data.is_empty() {
let tensor = Tensor::from_vec( state_data,
(sample_size, state_dim),
&self.device,
)
.map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))?
.to_dtype(training_dtype(&self.device))
.map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?;
batch_tensor_opt = Some(tensor);
}
}
if let Some(ref batch_tensor) = batch_tensor_opt {
// Suppress forward() monitoring to avoid to_vec2 GPU→CPU sync
agent.set_training_forward_active(true);
let batch_q_values = agent
.forward(batch_tensor)
.map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?;
agent.set_training_forward_active(false);
let num_actions =
batch_q_values.dims().get(1).copied().unwrap_or(5);
// Divergence check on first sample
let first_q = batch_q_values
.i(0)
.map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?;
let div_result = guard
.qvalue_divergence(&first_q, num_actions, 10000.0)
.map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?;
agent
.log_q_values_from_stats(
div_result.q_min,
div_result.q_max,
div_result.q_mean,
div_result.q_variance,
num_actions,
)
.map_err(|e| {
tracing::info!(
"Early stopping (Q-value divergence): {}",
e
);
anyhow::anyhow!("Early stopping: {}", e)
})?;
// Batch average via GPU reduction (one-step delay due to double-buffering)
let stats = guard
.qvalue_stats(&batch_q_values, sample_size, num_actions)
.map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?;
self.cached_avg_q = stats.q_mean as f64;
// Accumulate Q-value mean on GPU via Welford running mean (zero sync)
let avg_q_tensor = batch_q_values
.max(1)
.map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))?
.mean_all()
.map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?;
guard
.accumulate_q_value(&avg_q_tensor)
.map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?;
gpu_q_done = true;
}
}
}
}
// CUDA: GPU Q-value accumulation is mandatory — no CPU fallback.
if !gpu_q_done {
return Err(anyhow::anyhow!(
"GPU Q-value accumulation FAILED (no CPU fallback). \
Check GpuTrainingGuard initialization."
));
}
}
let avg_q_value = self.cached_avg_q;
Ok((loss_clipped, avg_q_value, final_grad_norm_f64))
}
/// Lazy-init fused CUDA training context (Standard DQN only).
/// Recreate if batch_size changed (OOM recovery).
pub(crate) async fn ensure_fused_ctx(&mut self) {
let needs_init = match &self.fused_ctx {
None => self.device.is_cuda(),
Some(ctx) => ctx.batch_size() != self.current_batch_size,
};
if !needs_init {
return;
}
if self.fused_ctx.is_some() {
info!("Fused CUDA context: batch_size changed, recreating");
self.fused_ctx = None;
}
let agent = self.agent.read().await;
match super::super::fused_training::FusedTrainingCtx::new(
&self.device, &*agent, &self.hyperparams, self.current_batch_size,
) {
Ok(ctx) => {
info!("Fused CUDA training initialized (batch_size={})", self.current_batch_size);
self.fused_ctx = Some(ctx);
}
Err(e) => {
warn!("Fused CUDA training init failed, using Candle path: {e}");
}
}
}
}

1847
crates/ml/src/trainers/dqn/trainer/training_loop.rs Normal file
View File

File diff suppressed because it is too large Load Diff