From 2df6a17ee7dfe1e098b8d3ebd35c2354e273c52c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 10 Mar 2026 13:06:19 +0100 Subject: [PATCH] feat(cuda): wire branching DQN GPU action selection into trainer Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/src/trainers/dqn/trainer.rs | 836 +++++++++++++++++--------- 1 file changed, 537 insertions(+), 299 deletions(-) diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 45a684500..f5304cd72 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -28,7 +28,7 @@ use ml_core::fill_simulator::{FillSimulator, FillResult}; use crate::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; use crate::dqn::curiosity::CuriosityModule; use crate::dqn::dqn::{DQN, DQNConfig}; -use crate::dqn::logging::{LoggingConfig, MetricsAggregator, log_epoch_start, log_epoch_end, log_training_progress, log_gradient_stats}; +use crate::dqn::logging::{LoggingConfig, MetricsAggregator, log_epoch_start, log_epoch_end, log_training_progress}; use crate::dqn::portfolio_tracker::PortfolioTracker; use crate::dqn::regime_conditional::RegimeConditionalDQN; use crate::dqn::reward::{RewardConfig, RewardFunction}; @@ -234,6 +234,10 @@ pub struct DQNTrainer { #[cfg(feature = "cuda")] gpu_experience_collector: Option, + /// GPU-fused epsilon-greedy action selector (eliminates argmax GPU->CPU sync barrier) + #[cfg(feature = "cuda")] + gpu_action_selector: Option, + /// Reusable GPU staging buffers for zero-alloc fold transitions buffer_pool: Option, @@ -896,6 +900,8 @@ impl DQNTrainer { features_raw_cuda: None, #[cfg(feature = "cuda")] gpu_experience_collector: None, + #[cfg(feature = "cuda")] + gpu_action_selector: None, // GPU pipeline: staging buffer pool (auto-initialized on CUDA devices) buffer_pool, @@ -2070,6 +2076,15 @@ impl DQNTrainer { num_atoms: self.hyperparams.num_atoms as i32, v_min: self.hyperparams.v_min as f32, v_max: self.hyperparams.v_max as f32, + // Fill simulation: order routing + stochastic fill check + fill_median_spread: self.hyperparams.avg_spread as f32, + fill_median_vol: self.median_vol as f32, + fill_ioc_fill_prob: 0.85, + fill_limit_fill_min: 0.30, + fill_limit_fill_max: 0.80, + fill_spread_cost_frac: 0.50, + fill_spread_capture_frac: 0.50, + fill_simulation_enabled: self.median_vol > 0.0, ..Default::default() }; @@ -2233,17 +2248,20 @@ impl DQNTrainer { }; // Batched action selection — GPU tensor path skips state→Vec→flatten→Tensor ✅ - let actions = if let Some(ref bt) = batch_tensor { - self.select_actions_batch_gpu(bt).await? + // gpu_handled_fill: when true, GPU kernel already did routing + fill simulation. + // Inner loop skips CPU route_action() + simulate_fill() to avoid double-work. + let (actions, gpu_handled_fill) = if let Some(ref bt) = batch_tensor { + self.select_actions_batch_gpu(bt, batch_start).await? } else { - self.select_actions_batch(&states).await? + (self.select_actions_batch(&states).await?, false) }; - // GPU portfolio simulation for the batch (if CUDA available) + // GPU portfolio simulation for the batch (if CUDA available). + // When gpu_handled_fill=true, actions are post-fill so GPU sim PnL is correct. + // When gpu_handled_fill=false, actions are pre-fill — GPU sim would simulate + // unfilled trades, so we fall back to CPU portfolio tracking. #[cfg(feature = "cuda")] - // GPU portfolio sim results: computed for metrics/validation, CPU inner loop remains source of truth. - // Phase 2b will wire these into the inner loop after CPU/GPU result equivalence is validated. - let _gpu_sim_results = { + let gpu_sim_results = if gpu_handled_fill { if let (Some(ref mut sim), Some(ref targets_buf)) = (&mut self.gpu_portfolio_sim, &self.targets_raw_cuda) { @@ -2258,29 +2276,61 @@ impl DQNTrainer { } else { None } + } else { + None }; + #[cfg(not(feature = "cuda"))] + let gpu_sim_results: Option<()> = None; + // Store experiences with batched actions for (idx_in_batch, &i) in batch_indices.iter().enumerate() { let state = &states[idx_in_batch]; let raw_action = actions[idx_in_batch]; - // Phase C/C+: Order routing - // Branching DQN: network heads already selected order type + urgency - // Standard DQN: deterministic routing based on spread/vol context - let current_spread = { - let features = &training_data[i].0; - // HL spread estimate: ln(high/low) ≈ features[1] - features[2] - (features[1] - features[2]).abs() as f32 - }; - let routed_action = if self.hyperparams.use_branching { - raw_action // Branching: order/urgency learned by network - } else { - self.route_action(raw_action.exposure, current_spread) - }; + // Extract GPU sim data for this sample (reward, portfolio features, done flag). + // Available when gpu_handled_fill=true and GPU portfolio sim succeeded. + // When Some, skips CPU portfolio execution and RewardFunction::calculate_reward(). + #[cfg(feature = "cuda")] + let gpu_sim_sample: Option<(f32, [f32; 3], bool)> = + gpu_sim_results.as_ref().and_then(|r| { + Some(( + *r.rewards.get(idx_in_batch)?, + [ + *r.portfolio_features.get(idx_in_batch * 3)?, + *r.portfolio_features.get(idx_in_batch * 3 + 1)?, + *r.portfolio_features.get(idx_in_batch * 3 + 2)?, + ], + *r.done_flags.get(idx_in_batch)? != 0, + )) + }); + #[cfg(not(feature = "cuda"))] + let gpu_sim_sample: Option<(f32, [f32; 3], bool)> = None; - // Phase C: Fill simulation gate — check if order fills - let (action, _fill_result) = self.simulate_fill(routed_action, i); + // Phase C/C+: Order routing + fill simulation + // When gpu_handled_fill=true, the GPU fused kernel already did: + // 1. epsilon-greedy action selection + // 2. order routing (spread/vol → order type + urgency) + // 3. fill simulation (splitmix64 hash → fill/no-fill) + // 4. unfilled → Flat override + // Skip CPU routing + fill to avoid double-work. + let action = if gpu_handled_fill { + raw_action // Already routed + fill-checked by GPU kernel + } else { + let current_spread = { + let features = &training_data[i].0; + // HL spread estimate: ln(high/low) ≈ features[1] - features[2] + (features[1] - features[2]).abs() as f32 + }; + let routed_action = if self.hyperparams.use_branching { + raw_action // Branching: order/urgency learned by network + } else { + self.route_action(raw_action.exposure, current_spread) + }; + // Phase C: Fill simulation gate — check if order fills + let (fill_action, _fill_result) = self.simulate_fill(routed_action, i); + fill_action + }; // Update volatility EMA from close log return (features[3]) let vol_sample = training_data[i].0[3].abs(); @@ -2441,70 +2491,59 @@ impl DQNTrainer { } } - // BUG #42 FIX: Execute portfolio action to enable reward calculation - // Root Cause: Portfolio tracker never updated during experience collection, - // causing state.portfolio_features[0] to always equal initial_capital, - // which results in 100% zero rewards since (V - V) / V = 0. - // Solution: Execute trade on portfolio tracker BEFORE reward calculation - // to update portfolio value, enabling P&L-based rewards. - // - // Flow: - // 1. state.portfolio_features[0] = current portfolio value (from tracker) - // 2. Execute action A at current_price → portfolio updates - // 3. Market moves to next_price → position marked-to-market - // 4. next_state.portfolio_features[0] = new portfolio value (after price move) - // 5. reward = (next_value - current_value) / current_value (percentage P&L) + // Portfolio execution + next_state update. + // GPU path: GPU sim already computed PnL, portfolio features, and done flags. + // CPU path: execute trade on portfolio tracker (BUG #42 FIX). + if let Some((_, ref gpu_features, _)) = gpu_sim_sample { + // GPU sim computed portfolio features for this sample. + // Skip CPU portfolio tracker execution — GPU kernel already did + // the full position management, mark-to-market, and PnL calculation. + if next_state.portfolio_features.len() >= 3 { + next_state.portfolio_features[0] = gpu_features[0]; // Normalized value + next_state.portfolio_features[1] = gpu_features[1]; // Normalized position + // portfolio_features[2] is spread, leave unchanged (already set) + } + } else { + // CPU portfolio execution (BUG #42 FIX) + let position_before = self.portfolio_tracker.current_position(); + let target_exposure = action.target_exposure() as f32; - // ========== DIAGNOSTIC: Capture position BEFORE action execution ========== - let position_before = self.portfolio_tracker.current_position(); - let target_exposure = action.target_exposure() as f32; + let current_price_f32 = current_close as f32; + let max_position_f32 = self.max_position as f32; + self.portfolio_tracker.execute_action(action, current_price_f32, max_position_f32); - // Execute trade on portfolio tracker - let current_price_f32 = current_close as f32; - let max_position_f32 = self.max_position as f32; - self.portfolio_tracker.execute_action(action, current_price_f32, max_position_f32); + let position_after = self.portfolio_tracker.current_position(); - // ========== DIAGNOSTIC: Capture position AFTER action execution ========== - let position_after = self.portfolio_tracker.current_position(); + if i < 100 { + let action_idx = action.exposure as usize; + debug!( + "SIGN_DIAG step={} action={} target_exp={:+.3} pos_before={:+.4} pos_after={:+.4}", + i, action_idx, target_exposure, position_before, position_after + ); + } - // SIGN INVERSION DIAGNOSTIC (Bug #42, now fixed) - // Only emitted at RUST_LOG=debug level - if i < 100 { - let action_idx = action.exposure as usize; - debug!( - "SIGN_DIAG step={} action={} target_exp={:+.3} pos_before={:+.4} pos_after={:+.4}", - i, action_idx, target_exposure, position_before, position_after - ); - } + let next_price_f32 = next_close as f32; + let updated_portfolio_features = self.portfolio_tracker.get_portfolio_features(next_price_f32); - // Update next_state portfolio features to reflect trade + price movement - // This is critical for reward calculation which compares current_state vs next_state - // CRITICAL: Use get_portfolio_features() (NORMALIZED) not total_value() (RAW) - // to match the normalization used in state.portfolio_features - let next_price_f32 = next_close as f32; - let updated_portfolio_features = self.portfolio_tracker.get_portfolio_features(next_price_f32); + if next_state.portfolio_features.len() >= 3 { + next_state.portfolio_features[0] = updated_portfolio_features[0]; + next_state.portfolio_features[1] = updated_portfolio_features[1]; + } - // Override next_state portfolio features with NORMALIZED post-trade values - if next_state.portfolio_features.len() >= 3 { - next_state.portfolio_features[0] = updated_portfolio_features[0]; // Normalized value - next_state.portfolio_features[1] = updated_portfolio_features[1]; // Normalized position - // portfolio_features[2] is spread, leave unchanged (already set) - } - - // Debug logging to verify trade execution and reward signal - if i % 100 == 0 || i < 10 { - let current_value = state.portfolio_features.get(0).unwrap_or(&1.0); - let new_value = updated_portfolio_features[0]; - let new_position = updated_portfolio_features[1]; - let reward_signal = if *current_value > 0.0 { - ((new_value - current_value) / current_value) * 100.0 - } else { - 0.0 - }; - debug!( - "BUG42_FIX: step={}, action={:?}, current_value={:.6}, new_value={:.6}, position={:.4}, reward_signal={:.4}%", - i, action, current_value, new_value, new_position, reward_signal - ); + if i % 100 == 0 || i < 10 { + let current_value = state.portfolio_features.get(0).unwrap_or(&1.0); + let new_value = updated_portfolio_features[0]; + let new_position = updated_portfolio_features[1]; + let reward_signal = if *current_value > 0.0 { + ((new_value - current_value) / current_value) * 100.0 + } else { + 0.0 + }; + debug!( + "BUG42_FIX: step={}, action={:?}, current_value={:.6}, new_value={:.6}, position={:.4}, reward_signal={:.4}%", + i, action, current_value, new_value, new_position, reward_signal + ); + } } // Track action for diversity penalty @@ -2566,82 +2605,88 @@ impl DQNTrainer { } } - // Calculate reward using RewardFunction (portfolio tracking, diversity penalty, movement threshold) - let recent_actions_vec: Vec = - self.recent_actions.iter().copied().collect(); - let reward_decimal = self.reward_fn.calculate_reward( - action, - state, - &next_state, - &recent_actions_vec, - )?; - let raw_reward = reward_decimal.to_f64().unwrap_or(0.0); - - // WAVE P2: Apply triple barrier scaling if we have a label (updated for Option) - let barrier_scaled_reward = if let Some(label) = barrier_label { - let scaled = self.reward_fn.apply_triple_barrier_scaling(reward_decimal, label); - scaled.to_f64().unwrap_or(0.0) + // Calculate reward: GPU path uses pre-computed PnL, CPU path uses RewardFunction. + // Both paths apply barrier scaling and risk adjustment on CPU (lightweight scalar ops). + let (raw_reward, reward) = if let Some((gpu_reward, _, _)) = gpu_sim_sample { + // GPU portfolio sim already computed PnL reward with risk penalty. + // Apply barrier scaling + DSR on top (CPU scalar ops). + let raw = gpu_reward as f64; + let barrier_scaled = if let Some(label) = barrier_label { + let reward_dec = rust_decimal::Decimal::try_from(raw) + .unwrap_or(rust_decimal::Decimal::ZERO); + let scaled = self.reward_fn.apply_triple_barrier_scaling(reward_dec, label); + scaled.to_f64().unwrap_or(0.0) + } else { + raw + }; + let risk_adjusted = if self.hyperparams.use_dsr { + barrier_scaled + } else { + self.calculate_risk_adjusted_reward(barrier_scaled) + }; + (raw, risk_adjusted as f32) } else { - raw_reward - }; - - // When DSR is active, it already handles risk-adjustment internally. - // Legacy path: divide by PnL volatility (Sharpe-like normalization) - let risk_adjusted_reward = if self.hyperparams.use_dsr { - barrier_scaled_reward // DSR output is already risk-adjusted - } else { - self.calculate_risk_adjusted_reward(barrier_scaled_reward) - }; - - // WAVE 26 P1.8: Add intrinsic curiosity reward if enabled - let total_reward = if let Some(ref mut curiosity) = self.curiosity_module { - // Convert state and next_state to tensors [1, 54] (batch_size=1) - let state_vec = state.to_vector(); - let state_tensor = Tensor::from_vec( - state_vec.clone(), - (1, state_vec.len()), - &self.device - ).map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Failed to cast state tensor to training dtype: {}", e))?; - - let next_state_vec = next_state.to_vector(); - let next_state_tensor = Tensor::from_vec( - next_state_vec, - (1, state_vec.len()), // Use same length as state_vec - &self.device - ).map_err(|e| anyhow::anyhow!("Failed to create next_state tensor: {}", e))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Failed to cast next_state tensor to training dtype: {}", e))?; - - // C1 FIX: Calculate curiosity for logging only. - // Intrinsic reward is NOT added to replay buffer — Q-values - // must learn pure trading P&L, not novelty-seeking. - let _intrinsic_reward = curiosity.calculate_curiosity_reward( - &state_tensor, + // CPU reward path: full RewardFunction pipeline + let recent_actions_vec: Vec = + self.recent_actions.iter().copied().collect(); + let reward_decimal = self.reward_fn.calculate_reward( action, - &next_state_tensor, - ).map_err(|e| anyhow::anyhow!("Curiosity reward calculation failed: {}", e))?; + state, + &next_state, + &recent_actions_vec, + )?; + let raw = reward_decimal.to_f64().unwrap_or(0.0); - // C1: Store ONLY extrinsic (trading) reward - risk_adjusted_reward - } else { - risk_adjusted_reward + let barrier_scaled = if let Some(label) = barrier_label { + let scaled = self.reward_fn.apply_triple_barrier_scaling(reward_decimal, label); + scaled.to_f64().unwrap_or(0.0) + } else { + raw + }; + + let risk_adjusted = if self.hyperparams.use_dsr { + barrier_scaled + } else { + self.calculate_risk_adjusted_reward(barrier_scaled) + }; + + // Curiosity forward model training (C1: reward not used, only weight update) + // Uses train_forward_model() — zero GPU→CPU roundtrip (no scalar readback) + if let Some(ref mut curiosity) = self.curiosity_module { + let state_vec = state.to_vector(); + let state_tensor = Tensor::from_vec( + state_vec.clone(), + (1, state_vec.len()), + &self.device + ).map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Failed to cast state tensor to training dtype: {}", e))?; + + let next_state_vec = next_state.to_vector(); + let next_state_tensor = Tensor::from_vec( + next_state_vec, + (1, state_vec.len()), + &self.device + ).map_err(|e| anyhow::anyhow!("Failed to create next_state tensor: {}", e))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Failed to cast next_state tensor to training dtype: {}", e))?; + + curiosity.train_forward_model( + &state_tensor, + action, + &next_state_tensor, + ).map_err(|e| anyhow::anyhow!("Curiosity forward model training failed: {}", e))?; + } + + (raw, risk_adjusted as f32) }; // Calculate price return for volatility tracking let price_return = (next_close - current_close) / current_close; - // BUG #17 FIX (Wave 16S-V16): Break positive feedback loop - // CRITICAL: Use raw_reward (NOT risk_adjusted_reward) to update trackers - // Otherwise: amplified rewards → inflated Sharpe → even larger amplification → explosion - // Example: raw=1.0 → Sharpe=10 → adjusted=10.0 → stored=10.0 → next Sharpe=100 → ... + // BUG #17 FIX: Use raw_reward to update trackers (prevents feedback loop) self.update_risk_trackers(raw_reward, price_return); - // Convert back to f32 for experience storage - // C1: total_reward now contains only extrinsic reward (trading P&L) - let reward = total_reward as f32; - // WAVE 16S: Log adaptive features periodically if i % 100 == 0 && i > 0 { let kelly_frac = self.get_kelly_fraction(); @@ -2675,9 +2720,10 @@ impl DQNTrainer { // 2. Fixed episode length boundary (fallback for compatibility) // 3. Data boundary reached let barrier_done = barrier_label.is_some(); // Any barrier event (Some(1/0/-1)) + let gpu_sim_done = gpu_sim_sample.map(|(_, _, d)| d).unwrap_or(false); let time_done = (i + 1) % EPISODE_LENGTH == 0; let data_done = i + 1 >= training_data.len(); - let done = barrier_done || time_done || data_done; + let done = barrier_done || gpu_sim_done || time_done || data_done; // Log barrier-driven terminations for analysis if barrier_done { @@ -2696,9 +2742,10 @@ impl DQNTrainer { } // P1 FIX: Reset portfolio at episode boundaries (updated for barrier termination) - // Note: Portfolio already reset in barrier detection block (lines 1717-1718) - // This handles time/data boundary resets - if done && !barrier_done { + // Note: Portfolio already reset in barrier detection block + // GPU sim done: GPU kernel already reset its internal portfolio state. + // This handles time/data boundary resets on the CPU tracker. + if done && !barrier_done && !gpu_sim_done { self.portfolio_tracker.reset(); self.reward_fn.reset_dsr(); @@ -2912,129 +2959,126 @@ impl DQNTrainer { monitor.track_q_value_range(avg_q); } } else { - // ═══ CPU readback path (non-GPU-PER buffers) ═══ + // ═══ GPU-accumulation path for non-GPU-PER buffers ═══ + // Same zero-sync pattern as GPU-PER path: accumulate loss/grad on GPU, + // single readback at epoch boundary. Safety checks deferred to epoch end + // (gradient collapse detected within ~8 training steps, acceptable delay). + use candle_core::Tensor; + + let mut loss_acc = Tensor::new(0.0_f32, &self.device) + .map_err(|e| anyhow::anyhow!("loss acc init: {e}"))?; + let mut grad_acc = Tensor::new(0.0_f32, &self.device) + .map_err(|e| anyhow::anyhow!("grad acc init: {e}"))?; + for _ in 0..num_training_steps { - match self.train_step().await { - Ok((loss, q_value, grad_norm)) => { - // WAVE 1.2 SAFETY #1: NaN/Inf Detection (catch corrupted values EARLY) - if !loss.is_finite() || !q_value.is_finite() || !grad_norm.is_finite() { - let msg = format!( - "SAFETY: NaN/Inf detected at epoch {} - loss={:.6}, q_value={:.6}, grad_norm={:.6}", - epoch, loss, q_value, grad_norm - ); - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{} (stopping training)", msg)); - }, - crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - debug!("{} (continuing training)", msg); - continue; // Skip this step - }, + // Pre-sample batch (READ lock) + let explicit_batch = { + let agent = self.agent.read().await; + let buffer = agent.memory(); + buffer.can_sample(self.current_batch_size).then(|| { + buffer + .sample(self.current_batch_size) + .map_err(|e| anyhow::anyhow!("Pre-sample: {e}")) + }).transpose()? + }; + + // GPU train step (WRITE lock) — zero CPU sync + let gpu_result = { + let mut agent = self.agent.write().await; + match agent.train_step(explicit_batch) { + Ok(r) => r, + Err(e) => { + let msg = e.to_string(); + if msg.contains("Early stopping") || msg.contains("Gradient collapse") { + return Err(anyhow::anyhow!("{}", msg)); } + warn!("GPU train step failed: {e}, continuing..."); + continue; } + } + }; - // WAVE 1.2 SAFETY #2: Gradient Monitor (explosion >10K or vanishing <1e-6) - if grad_norm > 10000.0 { - let msg = format!( - "SAFETY: Gradient explosion detected at epoch {} - norm={:.2e} > 10K threshold", - epoch, grad_norm - ); - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{} (stopping training)", msg)); - }, - crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - debug!("{}", msg); - }, - } - } else if grad_norm < 1e-6 && epoch > 10 { - debug!( - "SAFETY: Gradient vanishing detected at epoch {} - norm={:.2e} < 1e-6 threshold", - epoch, grad_norm - ); - } else { - // Gradient norm within healthy range - } - - // WAVE 1.2 SAFETY #3: Q-Value Bounds Check (±1M limits) - if q_value < -1_000_000.0 || q_value > 1_000_000.0 { - debug!( - "SAFETY: Q-value out of bounds at epoch {}: {:.2e} (bounds: ±1M)", - epoch, q_value - ); - } - - // WAVE 1.2 SAFETY #4: Loss Spike Detector (>50% jump) - if !self.safety_loss_history.is_empty() { - let Some(&prev_loss_raw) = self.safety_loss_history.back() else { continue; }; - let prev_loss = prev_loss_raw as f64; - // FIX: Use absolute difference instead of percentage (robust for negative/small values) - let loss_diff = (loss - prev_loss).abs(); - - if loss_diff > 0.5 { // Threshold: 0.5 absolute change - let msg = format!( - "SAFETY: Loss spike detected at epoch {}: {:.6} → {:.6} (Δ={:+.4})", - epoch, prev_loss, loss, loss_diff - ); - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{} (stopping training)", msg)); - }, - crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - debug!("{}", msg); - }, - } - } - } - - // Update loss history - self.safety_loss_history.push_back(loss as f32); - if self.safety_loss_history.len() > 30 { - self.safety_loss_history.pop_front(); - } - - epoch_loss += loss; - epoch_q_value += q_value; - epoch_gradient_norm += grad_norm; - train_step_count += 1; - - // WAVE 30: Record metrics in aggregator - let current_avg_reward = if !monitor.reward_history.is_empty() { - monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 - } else { - 0.0 - }; - self.metrics_aggregator.record(loss as f32, q_value as f32, current_avg_reward); - self.metrics_aggregator.record_gradient(grad_norm as f32); - - // WAVE 30: Log training progress at configured intervals - if self.metrics_aggregator.should_log(&self.logging_config) { - let aggregated = self.metrics_aggregator.aggregate_and_clear(); - log_training_progress(&aggregated); - } - - // WAVE 30: Log gradient statistics at configured intervals - if self.metrics_aggregator.batch_count() % self.logging_config.gradient_log_interval == 0 { - log_gradient_stats(grad_norm as f32, grad_norm as f32, 10.0); - } - - // WAVE 9-11: Track Q-value range for production monitoring - monitor.track_q_value_range(q_value); - }, - Err(e) => { - // WAVE 23 P0 FIX: Propagate early stopping errors - let error_msg = e.to_string(); - if error_msg.contains("Early stopping") || - error_msg.contains("Gradient collapse") || - error_msg.contains("Q-value divergence") { - tracing::error!("TERMINATING: Early stopping triggered - {}", e); - return Err(e.into()); - } - warn!("Training step failed: {}, continuing...", e); - }, - } + // GPU-side accumulation — no readback + loss_acc = (&loss_acc + &gpu_result.loss_gpu) + .map_err(|e| anyhow::anyhow!("loss acc: {e}"))?; + grad_acc = (&grad_acc + &gpu_result.grad_norm_gpu) + .map_err(|e| anyhow::anyhow!("grad acc: {e}"))?; + train_step_count += 1; } - } // end CPU path + + // ═══ Single epoch-boundary readback ═══ + if train_step_count > 0 { + let n = train_step_count as f64; + let n_tensor = Tensor::new(train_step_count as f32, &self.device) + .map_err(|e| anyhow::anyhow!("n tensor: {e}"))?; + let avg_loss = loss_acc.broadcast_div(&n_tensor) + .map_err(|e| anyhow::anyhow!("avg loss: {e}"))? + .to_scalar::() + .map_err(|e| anyhow::anyhow!("loss readback: {e}"))?; + let avg_grad = grad_acc.broadcast_div(&n_tensor) + .map_err(|e| anyhow::anyhow!("avg grad: {e}"))? + .to_scalar::() + .map_err(|e| anyhow::anyhow!("grad readback: {e}"))?; + + // Q-value estimation at epoch boundary + let avg_q = { + let mut agent = self.agent.write().await; + self.estimate_avg_q_value_with_early_stopping(&mut agent).await? + }; + + // Epoch-level safety checks (same as GPU-PER path) + if !avg_loss.is_finite() || !avg_grad.is_finite() { + let msg = format!( + "SAFETY: NaN/Inf in epoch {} avg — loss={:.6}, grad={:.6}", + epoch, avg_loss, avg_grad + ); + match self.safety_level { + crate::safety::SafetyLevel::Strict => { + return Err(anyhow::anyhow!("{} (stopping training)", msg)); + }, + crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { + debug!("{} (continuing training)", msg); + }, + } + } + + // Loss history (epoch-level average) + self.safety_loss_history.push_back(avg_loss); + if self.safety_loss_history.len() > 30 { + self.safety_loss_history.pop_front(); + } + + epoch_loss = avg_loss as f64 * n; + epoch_q_value = avg_q * n; + epoch_gradient_norm = avg_grad as f64 * n; + + // Metrics aggregator + let current_avg_reward = if !monitor.reward_history.is_empty() { + monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 + } else { + 0.0 + }; + self.metrics_aggregator.record(avg_loss, avg_q as f32, current_avg_reward); + self.metrics_aggregator.record_gradient(avg_grad); + + if self.metrics_aggregator.should_log(&self.logging_config) { + let aggregated = self.metrics_aggregator.aggregate_and_clear(); + log_training_progress(&aggregated); + } + + // Gradient diagnostics at epoch boundary + { + let mut agent = self.agent.write().await; + agent.log_diagnostics(avg_grad) + .map_err(|e| { + tracing::info!("Early stopping (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + } + + monitor.track_q_value_range(avg_q); + } + } // end non-GPU-PER path // Phase 2b: Sync GPU weight copies after training updates #[cfg(feature = "cuda")] @@ -3972,7 +4016,7 @@ impl DQNTrainer { /// /// # Returns /// Vector of TradingAction decisions (same order as input states) - async fn select_actions_batch(&self, states: &[TradingState]) -> Result> { + async fn select_actions_batch(&mut self, states: &[TradingState]) -> Result> { if states.is_empty() { return Ok(Vec::new()); } @@ -4034,31 +4078,70 @@ impl DQNTrainer { drop(agent); // Release lock early - // Extract Q-values and select actions (epsilon-greedy) + // Fused GPU epsilon-greedy: argmax + RNG in a single kernel launch, + // eliminating the intermediate GPU->CPU sync from Candle's argmax(). + // Falls back to the two-step path when CUDA kernel is unavailable. + #[cfg(feature = "cuda")] + { + // Lazy-init the GPU action selector on first call + if self.gpu_action_selector.is_none() && self.device.is_cuda() { + match crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( + &self.device, + self.hyperparams.batch_size.max(batch_size).max(8192), + 0xDEAD_BEEF_CAFE_u64, + ) { + Ok(selector) => { + info!("GPU action selector initialized for select_actions_batch"); + self.gpu_action_selector = Some(selector); + } + Err(e) => { + warn!("GPU action selector init failed, using CPU fallback: {e}"); + } + } + } + + if let Some(ref mut selector) = self.gpu_action_selector { + let action_indices_tensor = selector + .select_actions(&batch_q_values, epsilon, batch_size, 5) + .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))?; + // Single GPU->CPU readback for route_action() dispatch + let action_indices = action_indices_tensor + .to_vec1::() + .map_err(|e| anyhow::anyhow!("Failed to readback action indices: {}", e))?; + + let mut actions = Vec::with_capacity(batch_size); + for idx in &action_indices { + let action_idx = *idx as usize; + let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) + .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; + let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); + actions.push(action); + } + return Ok(actions); + } + } + + // CPU fallback: Candle argmax + per-sample RNG let mut actions = Vec::with_capacity(batch_size); let mut rng = rand::thread_rng(); - // WAVE 10.6: GPU-optimized argmax - compute argmax on GPU, only pull indices to CPU let greedy_action_indices = batch_q_values .argmax(1) .map_err(|e| anyhow::anyhow!("Failed to compute argmax on GPU: {}", e))? .to_vec1::() .map_err(|e| anyhow::anyhow!("Failed to transfer argmax results to CPU: {}", e))?; - for i in 0..batch_size { + for idx in &greedy_action_indices { use rand::Rng; let action_idx = if rng.gen::() < epsilon { - // Random exploration over 5 exposure levels rng.gen_range(0..5) } else { - // Greedy exploitation: use precomputed argmax from GPU - greedy_action_indices[i] as usize + *idx as usize }; let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; - // Phase C: Use smart routing with trainer's spread/vol EMAs let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); actions.push(action); @@ -4070,10 +4153,26 @@ impl DQNTrainer { /// GPU-optimized batch action selection using pre-built state tensor. /// /// Skips the state→Vec→flatten→Tensor pipeline (~130 allocs per batch). - async fn select_actions_batch_gpu(&self, batch_tensor: &Tensor) -> Result> { + /// 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) + async fn select_actions_batch_gpu( + &mut self, + batch_tensor: &Tensor, + batch_start: usize, + ) -> Result<(Vec, bool)> { let batch_size = batch_tensor.dims()[0]; if batch_size == 0 { - return Ok(Vec::new()); + return Ok((Vec::new(), false)); } let agent = self.agent.read().await; @@ -4095,9 +4194,137 @@ impl DQNTrainer { // Noisy nets are the sole exploration mechanism during training. // Count bonus kept for diversity metrics only (record_action tracking). + // Branching DQN: get per-head Q-values if branching mode is active. + // Must be done while agent lock is held (branching_q_network is on agent). + // branching_q_network lives on the DQN struct, only accessible via Standard variant. + #[cfg(feature = "cuda")] + let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = if self.hyperparams.use_branching { + let dqn_ref: Option<&crate::dqn::DQN> = match &*agent { + DQNAgentType::Standard(ref dqn) => Some(dqn), + DQNAgentType::RegimeConditional(_) => None, + }; + if let Some(dqn) = dqn_ref { + if let Some(ref branching_net) = dqn.branching_q_network { + let branch_output = branching_net + .forward_branches(batch_tensor, false) + .map_err(|e| anyhow::anyhow!("Branching forward_branches failed: {}", e))?; + if branch_output.advantages.len() >= 3 { + Some(( + branch_output.advantages[0].clone(), + branch_output.advantages[1].clone(), + branch_output.advantages[2].clone(), + )) + } else { + warn!("Branching forward produced {} heads, expected 3", branch_output.advantages.len()); + None + } + } else { + None + } + } else { + None + } + } else { + None + }; + drop(agent); - // GPU argmax + epsilon-greedy (same as CPU path) + // Fused GPU epsilon-greedy: argmax + RNG in a single kernel launch. + // Same fusion as select_actions_batch() — see comments there. + #[cfg(feature = "cuda")] + { + if self.gpu_action_selector.is_none() && self.device.is_cuda() { + match crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( + &self.device, + self.hyperparams.batch_size.max(batch_size).max(8192), + 0xDEAD_BEEF_CAFE_u64, + ) { + Ok(selector) => { + info!("GPU action selector initialized for select_actions_batch_gpu"); + self.gpu_action_selector = Some(selector); + } + Err(e) => { + warn!("GPU action selector init failed, using CPU fallback: {e}"); + } + } + } + + if let Some(ref mut selector) = self.gpu_action_selector { + // 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; + + let action_indices_tensor = if self.hyperparams.use_branching { + 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 { + // Fallback: use standard path if branching forward failed + selector + .select_actions(&batch_q_values, epsilon, batch_size, 5) + .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))? + } + } else if use_routed { + 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, // DQN exposure actions + batch_start as i32, + spread, // current spread (batch-level avg) + spread, // median spread (same as avg for batch) + self.vol_ema as f32, + self.median_vol as f32, + spread_bps, + 0.85, // ioc_fill_prob + 0.30, // limit_fill_min + 0.80, // limit_fill_max + 0.50, // spread_cost_frac + 0.50, // spread_capture_frac + ) + .map_err(|e| anyhow::anyhow!("GPU routed action selection failed: {}", e))? + } else { + selector + .select_actions(&batch_q_values, epsilon, batch_size, 5) + .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))? + }; + + let action_indices = action_indices_tensor + .to_vec1::() + .map_err(|e| anyhow::anyhow!("Failed to readback action indices: {}", e))?; + + let mut actions = Vec::with_capacity(batch_size); + if self.hyperparams.use_branching { + // Branching: factored indices 0-44 + for &idx in &action_indices { + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; + actions.push(action); + } + } else { + // Standard: exposure indices 0-4 + for &idx in &action_indices { + let action_idx = idx as usize; + let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) + .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; + // When routed kernel was used, exposure is already post-fill + // (unfilled → Flat override done on GPU). Route to get the + // FactoredAction struct but the exposure level is final. + let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); + actions.push(action); + } + } + return Ok((actions, use_routed)); + } + } + + // CPU fallback: Candle argmax + per-sample RNG let greedy_action_indices = batch_q_values .argmax(1) .map_err(|e| anyhow::anyhow!("Failed to compute argmax on GPU: {}", e))? @@ -4106,23 +4333,26 @@ impl DQNTrainer { let mut actions = Vec::with_capacity(batch_size); let mut rng = rand::thread_rng(); - for i in 0..batch_size { + for idx in &greedy_action_indices { use rand::Rng; let action_idx = if rng.gen::() < epsilon { rng.gen_range(0..5) } else { - greedy_action_indices[i] as usize + *idx as usize }; let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; - // Phase C: Use smart routing with trainer's spread/vol EMAs let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); actions.push(action); } - Ok(actions) + Ok((actions, false)) } - /// Epsilon-greedy action selection + /// 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. async fn epsilon_greedy_action(&self, state: &Tensor) -> Result { use rand::Rng; @@ -4279,17 +4509,25 @@ impl DQNTrainer { let mut agent = self.agent.write().await; - // train_step returns GpuTrainResult — single readback here (CPU path) + // train_step returns GpuTrainResult with GPU-resident scalar tensors. let gpu_result = agent .train_step(explicit_batch) .map_err(|e| anyhow::anyhow!("Training step failed: {}", e))?; - let loss_f32 = gpu_result.loss_gpu - .to_scalar::() - .map_err(|e| anyhow::anyhow!("Loss readback: {e}"))?; - let grad_norm_f32 = gpu_result.grad_norm_gpu - .to_scalar::() - .map_err(|e| anyhow::anyhow!("Grad norm readback: {e}"))?; + // Batch both GPU scalars into a single [2]-element tensor for ONE DMA + // readback instead of two separate GPU sync barriers. Each to_scalar() + // triggers a cudaMemcpy + stream synchronize; by stacking first we pay + // that cost only once. + let loss_unsqueezed = gpu_result.loss_gpu.unsqueeze(0) + .map_err(|e| anyhow::anyhow!("Loss unsqueeze: {e}"))?; + let grad_unsqueezed = gpu_result.grad_norm_gpu.unsqueeze(0) + .map_err(|e| anyhow::anyhow!("Grad norm unsqueeze: {e}"))?; + let stacked = candle_core::Tensor::cat(&[&loss_unsqueezed, &grad_unsqueezed], 0) + .map_err(|e| anyhow::anyhow!("Loss+grad stack: {e}"))?; + let readback = stacked.to_vec1::() + .map_err(|e| anyhow::anyhow!("Batched readback: {e}"))?; + let loss_f32 = readback.first().copied().unwrap_or(0.0); + let grad_norm_f32 = readback.get(1).copied().unwrap_or(0.0); // Check for gradient collapse (early stopping) agent.log_diagnostics(grad_norm_f32) @@ -5132,7 +5370,7 @@ mod tests { #[tokio::test] async fn test_batched_action_selection() { let hyperparams = create_test_params(); - let trainer = DQNTrainer::new(hyperparams).unwrap(); + let mut trainer = DQNTrainer::new(hyperparams).unwrap(); // Create multiple synthetic states for batched action selection let batch_size = 10; @@ -5195,7 +5433,7 @@ mod tests { #[tokio::test] async fn test_batched_vs_sequential_action_selection_consistency() { let hyperparams = create_test_params(); - let trainer = DQNTrainer::new(hyperparams).unwrap(); + let mut trainer = DQNTrainer::new(hyperparams).unwrap(); // Create test states let batch_size = 5; @@ -5242,7 +5480,7 @@ mod tests { #[tokio::test] async fn test_empty_batch_handling() { let hyperparams = create_test_params(); - let trainer = DQNTrainer::new(hyperparams).unwrap(); + let mut trainer = DQNTrainer::new(hyperparams).unwrap(); let empty_states: Vec = Vec::new(); let result = trainer.select_actions_batch(&empty_states).await; @@ -5286,7 +5524,7 @@ mod tests { async fn test_batch_size_mismatch_smaller_than_configured() { let mut hyperparams = create_test_params(); hyperparams.batch_size = 32; - let trainer = DQNTrainer::new(hyperparams).unwrap(); + let mut trainer = DQNTrainer::new(hyperparams).unwrap(); // Create batch with 16 states (half of configured 32) let mut feature_vec = [0.0; 42]; // 42 market features @@ -5323,7 +5561,7 @@ mod tests { async fn test_batch_size_mismatch_larger_than_configured() { let mut hyperparams = create_test_params(); hyperparams.batch_size = 16; - let trainer = DQNTrainer::new(hyperparams).unwrap(); + let mut trainer = DQNTrainer::new(hyperparams).unwrap(); // Create batch with 64 states (4x configured 16) let mut feature_vec = [0.0; 42]; // 42 market features @@ -5358,7 +5596,7 @@ mod tests { /// Production-critical test: Verify empty batch handling #[tokio::test] async fn test_empty_batch_returns_empty_actions() { - let trainer = DQNTrainer::new(create_test_params()).unwrap(); + let mut trainer = DQNTrainer::new(create_test_params()).unwrap(); let empty_batch: Vec = vec![]; let result = trainer.select_actions_batch(&empty_batch).await; @@ -5375,7 +5613,7 @@ mod tests { async fn test_single_sample_batch() { let mut hyperparams = create_test_params(); hyperparams.batch_size = 32; - let trainer = DQNTrainer::new(hyperparams).unwrap(); + let mut trainer = DQNTrainer::new(hyperparams).unwrap(); let mut feature_vec = [0.0; 42]; // 42 market features for i in 0..4 {