diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index edd8465f5..eb502b169 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -177,24 +177,6 @@ pub struct DQNConfig { /// Default: 3. Flows to `CountBonus` and `BranchingConfig`. pub num_urgency_levels: usize, - /// Regime-conditional importance sampling for the branching loss path. - /// Feature index for ADX(14) in the market feature vector. - /// Default: 40 (position in the 42-dim feature layout). - /// Must match the feature extraction pipeline. - pub regime_adx_idx: usize, - - /// Feature index for CUSUM direction in the market feature vector. - /// Default: 41 (position in the 42-dim feature layout). - pub regime_cusum_idx: usize, - - /// ADX threshold for trending regime classification (normalized 0-1). - /// Default: 0.25 (= raw ADX > 25). - pub regime_adx_threshold: f32, - - /// CUSUM direction threshold for volatile regime classification. - /// Default: 0.7 (|CUSUM dir| > 0.7 → volatile). - pub regime_cusum_threshold: f32, - /// Hidden layer width for curiosity forward dynamics model. /// Default: 128. Must match `CUR_HIDDEN` in CUDA kernel. pub curiosity_hidden_dim: usize, @@ -279,11 +261,6 @@ impl Default for DQNConfig { num_order_types: 3, num_urgency_levels: 3, - regime_adx_idx: 40, - regime_cusum_idx: 41, - regime_adx_threshold: 0.25, - regime_cusum_threshold: 0.7, - // Curiosity model dimensions curiosity_hidden_dim: 128, curiosity_market_dim: 42, @@ -689,10 +666,6 @@ impl DQNConfig { branch_hidden_dim: 128, num_order_types: 3, num_urgency_levels: 3, - regime_adx_idx: 40, - regime_cusum_idx: 41, - regime_adx_threshold: 0.25, - regime_cusum_threshold: 0.7, curiosity_hidden_dim: 128, curiosity_market_dim: 42, cvar_alpha: 0.05, @@ -1991,6 +1964,52 @@ impl DQN { (exp, ord, urg) } + /// Per-branch Q-values for the GPU action selector. + /// + /// Returns `Some((exposure, order, urgency))` as raw `CudaSlice` buffers + /// when the branching network is initialised; `None` otherwise. Each buffer is + /// `[batch_size, branch_size]` flattened (exposure = 4 actions, order = 3, + /// urgency = 3). + pub fn batch_branching_q_values( + &self, + states: &GpuTensor, + ) -> Result, cudarc::driver::CudaSlice, cudarc::driver::CudaSlice)>, MLError> { + let Some(branching_net) = self.branching_q_network.as_ref() else { + return Ok(None); + }; + + let n = states.shape().first().copied().unwrap_or(0); + if n == 0 { + return Err(MLError::ModelError("Empty batch for batch_branching_q_values".into())); + } + + let output = branching_net.forward_branches_eval(states)?; + if output.advantages.len() < 3 { + return Err(MLError::ModelError(format!( + "batch_branching_q_values: branching network produced {} branches, expected 3", + output.advantages.len() + ))); + } + + let exp_slice = output.advantages.first() + .ok_or_else(|| MLError::ModelError("missing branch 0 (exposure)".into()))? + .clone() + .into_parts() + .0; + let ord_slice = output.advantages.get(1) + .ok_or_else(|| MLError::ModelError("missing branch 1 (order)".into()))? + .clone() + .into_parts() + .0; + let urg_slice = output.advantages.get(2) + .ok_or_else(|| MLError::ModelError("missing branch 2 (urgency)".into()))? + .clone() + .into_parts() + .0; + + Ok(Some((exp_slice, ord_slice, urg_slice))) + } + /// Task 0.6 — NoisyNets σ mean for a specific action branch /// (0 = direction, 1 = magnitude, 2 = order, 3 = urgency). Mean is over /// the branch's fc + out NoisyLinear layers' weight_sigma + bias_sigma @@ -2113,11 +2132,10 @@ impl DQN { /// ([`serialize_model`](../../ml/src/trainers/dqn/trainer/mod.rs)) emits one /// merged safetensors file with three regime-prefixed copies of the branching /// weight manifest: `trending__`, `ranging__`, `volatile__`. The - /// single-head `DQN` struct only owns one branching network, so this loader - /// restores the `trending__` copy by default (matching - /// [`RegimeConditionalDQN::load_from_merged_safetensors`]). For checkpoints - /// saved without a prefix (e.g., the per-crate round-trip tests), the loader - /// falls back to plain names. + /// `DQN` owns one branching network. Checkpoints saved by the Phase 4+ + /// serializer use plain names (no prefix). For legacy checkpoints saved + /// with a `trending__` prefix (pre-Phase 4), the loader falls back to + /// plain names after stripping the prefix. /// /// # Arguments /// diff --git a/crates/ml-dqn/src/lib.rs b/crates/ml-dqn/src/lib.rs index 59ea83c36..da128296a 100644 --- a/crates/ml-dqn/src/lib.rs +++ b/crates/ml-dqn/src/lib.rs @@ -35,7 +35,7 @@ pub mod gae; pub mod logging; pub mod network; pub mod nstep_buffer; -pub mod regime_conditional; +pub mod regime_classifier; pub mod residual; pub mod reward; pub mod softmax; @@ -125,8 +125,8 @@ pub use rainbow_network::{RainbowNetwork, RainbowNetworkConfig}; pub use replay_buffer_type::{GpuBatch, StagedGpuBuffer}; pub use gpu_replay_buffer::GpuBatchPtrs; -// Re-export regime-conditional DQN components -pub use regime_conditional::{RegimeClassConfig, RegimeConditionalDQN, RegimeMetrics, RegimeType}; +// Re-export regime classifier types (used by validation layer and ml-regime-detection) +pub use regime_classifier::{RegimeClassConfig, RegimeType}; // Re-export logit clipping utilities pub use logit_clipping::{clip_logits, clip_logits_default, softmax_with_clipping, DEFAULT_CLIP_MAX}; diff --git a/crates/ml-dqn/src/moe.rs b/crates/ml-dqn/src/moe.rs index e6dda58c2..3568856b9 100644 --- a/crates/ml-dqn/src/moe.rs +++ b/crates/ml-dqn/src/moe.rs @@ -1,8 +1,9 @@ //! Mixture-of-Experts components for the DQN policy network. //! -//! Replaces the vestigial RegimeConditionalDQN with a learned-gate -//! K=8 expert mixture per the spec at +//! K=8 learned-gate expert mixture per the spec at //! `docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md`. +//! Phase 4 removed `RegimeConditionalDQN`; MoE is now the sole +//! regime-conditioning mechanism. use std::sync::Arc; diff --git a/crates/ml-dqn/src/regime_classifier.rs b/crates/ml-dqn/src/regime_classifier.rs new file mode 100644 index 000000000..d2067e023 --- /dev/null +++ b/crates/ml-dqn/src/regime_classifier.rs @@ -0,0 +1,144 @@ +//! Regime classifier types: feature-index config and three-class enum. +//! +//! These types are used by the validation layer and offline analysis to classify +//! market states from feature vectors. The three-head `RegimeConditionalDQN` +//! routing logic was deleted in Phase 4; regime conditioning now lives exclusively +//! in the MoE gate inside `DQN`. + +use std::sync::Arc; +use cudarc::driver::CudaStream; +use ml_core::cuda_autograd::GpuTensor; +use ml_core::MLError; +use serde::{Deserialize, Serialize}; +use tracing::debug; + +/// Dynamic regime classification parameters. +/// +/// Feature indices and thresholds for CPU-side regime classification used by the +/// validation layer. The `DQN` MoE gate learns its own decomposition from the full +/// state vector — these fields are no longer consumed by the training path. +/// +/// ## Default Feature Layout +/// +/// OHLCV(0-4) + Technical(5-9) + PricePatterns(10-15) + Volume(16-21) +/// + Time(22-26) + Statistical(27-39) + Regime(40-41) +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct RegimeClassConfig { + /// Feature index for ADX(14) in the market feature vector. + pub adx_idx: usize, + /// Feature index for CUSUM direction in the market feature vector. + pub cusum_idx: usize, + /// ADX threshold for trending classification (normalized 0-1). + pub adx_threshold: f32, + /// CUSUM direction threshold for volatile classification. + pub cusum_threshold: f32, +} + +impl Default for RegimeClassConfig { + fn default() -> Self { + Self { + adx_idx: 40, + cusum_idx: 41, + adx_threshold: 0.25, + cusum_threshold: 0.7, + } + } +} + +/// Market regime types based on structural characteristics. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum RegimeType { + /// Strong directional trend (ADX > 25) + Trending, + /// Mean-reverting market (ADX <= 25, Entropy <= 0.7) + Ranging, + /// High volatility/uncertainty (ADX <= 25, Entropy > 0.7) + Volatile, +} + +impl RegimeType { + /// Classify regime from market features (CPU path). + /// + /// Uses ADX and CUSUM direction from the state vector at positions specified + /// by `cfg`. Returns `Ranging` when the feature vector is shorter than required. + pub fn classify_from_features(features: &[f32], cfg: &RegimeClassConfig) -> Self { + if features.len() <= cfg.cusum_idx { + debug!( + "Feature vector too small ({} elements), returning Ranging regime as default", + features.len() + ); + return Self::Ranging; + } + + let adx = features.get(cfg.adx_idx).copied().unwrap_or(0.0); + let cusum_direction = features.get(cfg.cusum_idx).copied().unwrap_or(0.0); + + if adx > cfg.adx_threshold { + Self::Trending + } else if cusum_direction.abs() > cfg.cusum_threshold { + Self::Volatile + } else { + Self::Ranging + } + } + + /// Classify regime for a batch of states on GPU using tensor ops. + /// + /// Returns 3 binary mask tensors (trending, ranging, volatile) each of shape + /// `[batch_size]`. Zero CPU roundtrip — all operations are GPU tensor ops. + pub fn classify_regime_masks_gpu( + states: &GpuTensor, + cfg: &RegimeClassConfig, + stream: &Arc, + ) -> Result<(GpuTensor, GpuTensor, GpuTensor), MLError> { + let shape = states.shape(); + let batch_size = shape.first().copied().ok_or_else(|| { + MLError::ModelError( + "classify_regime_masks_gpu: states tensor has no dimensions".to_owned(), + ) + })?; + let state_dim = shape.get(1).copied().ok_or_else(|| { + MLError::ModelError( + "classify_regime_masks_gpu: states tensor needs 2 dimensions".to_owned(), + ) + })?; + + let host = states.to_host(stream)?; + + let mut trending = Vec::with_capacity(batch_size); + let mut ranging = Vec::with_capacity(batch_size); + let mut volatile = Vec::with_capacity(batch_size); + + for i in 0..batch_size { + let offset = i * state_dim; + let row = host.get(offset..offset + state_dim).ok_or_else(|| { + MLError::ModelError(format!( + "classify_regime_masks_gpu: row {i} out of bounds" + )) + })?; + let regime = Self::classify_from_features(row, cfg); + match regime { + RegimeType::Trending => { + trending.push(1.0_f32); + ranging.push(0.0); + volatile.push(0.0); + } + RegimeType::Ranging => { + trending.push(0.0_f32); + ranging.push(1.0); + volatile.push(0.0); + } + RegimeType::Volatile => { + trending.push(0.0_f32); + ranging.push(0.0); + volatile.push(1.0); + } + } + } + + let t = GpuTensor::from_host(&trending, vec![batch_size], stream)?; + let r = GpuTensor::from_host(&ranging, vec![batch_size], stream)?; + let v = GpuTensor::from_host(&volatile, vec![batch_size], stream)?; + Ok((t, r, v)) + } +} diff --git a/crates/ml-dqn/src/regime_conditional.rs b/crates/ml-dqn/src/regime_conditional.rs deleted file mode 100644 index 9cd624266..000000000 --- a/crates/ml-dqn/src/regime_conditional.rs +++ /dev/null @@ -1,876 +0,0 @@ -//! Regime-Conditional Deep Q-Network -//! -//! Implements separate Q-network heads for different market regimes (Trending, Ranging, Volatile). -//! Routes actions and training through regime-specific networks for improved performance. -//! -//! ## Architecture -//! -//! ```text -//! State [46+ dims] -> Regime Classifier -> Regime-Specific Q-Head -> Action [45 dims] -//! (or 54 core features) | | -//! ADX + Entropy Trending / Ranging / Volatile -//! ``` -//! -//! ## Key Features -//! -//! - **3 Independent Heads**: Separate networks for trending, ranging, volatile regimes -//! - **Automatic Routing**: Classifies regime from market features (ADX, Entropy) -//! - **Shared Experience Replay**: All regimes contribute to shared buffer -//! - **Per-Regime Metrics**: Track training progress per regime -//! - **Regime-Specific Epsilon**: Independent exploration strategies -//! - **Checkpoint Support**: Save/load all 3 heads simultaneously -//! -//! ## Regime Classification -//! -//! - **Trending**: ADX > 25 (strong directional trend) -//! - **Volatile**: ADX <= 25 AND Entropy > 0.7 (high uncertainty) -//! - **Ranging**: ADX <= 25 AND Entropy <= 0.7 (mean-reverting) - -use std::collections::HashMap; - -use std::sync::Arc; -use cudarc::driver::{CudaSlice, CudaStream}; -use ml_core::cuda_autograd::GpuTensor; -use ml_core::device::MlDevice; -use serde::{Deserialize, Serialize}; -use tracing::{debug, info}; - -use ml_core::MLError; -use super::dqn::{DQN, DQNConfig}; -use crate::FactoredAction; - -/// Dynamic regime classification parameters. -/// -/// Feature indices and thresholds for regime classification, derived from `DQNConfig` -/// rather than hardcoded. This allows hyperopt to tune thresholds and supports -/// feature layouts where regime features are at non-standard positions. -/// -/// ## Default Feature Layout -/// -/// OHLCV(0-4) + Technical(5-9) + PricePatterns(10-15) + Volume(16-21) -/// + Time(22-26) + Statistical(27-39) + Regime(40-41) -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] -pub struct RegimeClassConfig { - /// Feature index for ADX(14) in the market feature vector. - pub adx_idx: usize, - /// Feature index for CUSUM direction in the market feature vector. - pub cusum_idx: usize, - /// ADX threshold for trending classification (normalized 0-1). - pub adx_threshold: f32, - /// CUSUM direction threshold for volatile classification. - pub cusum_threshold: f32, -} - -impl Default for RegimeClassConfig { - fn default() -> Self { - Self { - adx_idx: 40, - cusum_idx: 41, - adx_threshold: 0.25, - cusum_threshold: 0.7, - } - } -} - -impl RegimeClassConfig { - /// Build from `DQNConfig` fields. - pub const fn from_dqn_config(config: &DQNConfig) -> Self { - Self { - adx_idx: config.regime_adx_idx, - cusum_idx: config.regime_cusum_idx, - adx_threshold: config.regime_adx_threshold, - cusum_threshold: config.regime_cusum_threshold, - } - } -} - -/// Market regime types based on structural characteristics -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub enum RegimeType { - /// Strong directional trend (ADX > 25) - Trending, - /// Mean-reverting market (ADX <= 25, Entropy <= 0.7) - Ranging, - /// High volatility/uncertainty (ADX <= 25, Entropy > 0.7) - Volatile, -} - -impl RegimeType { - /// Classify regime from market features (CPU path). - /// - /// Uses ADX and CUSUM direction from the state vector at positions - /// specified by `cfg`. These features are produced by `FeatureExtractor` - /// and are always present in the 42-dim (or 45/48-dim) state vector. - /// - /// # Classification Rules - /// - /// - **Trending**: ADX > `adx_threshold` - /// - **Volatile**: NOT trending AND |CUSUM direction| > `cusum_threshold` - /// - **Ranging**: NOT trending AND NOT volatile - /// - /// # Fallback - /// - /// Returns Ranging if the feature vector is too short (e.g. tests). - pub fn classify_from_features(features: &[f32], cfg: &RegimeClassConfig) -> Self { - if features.len() <= cfg.cusum_idx { - debug!( - "Feature vector too small ({} elements), returning Ranging regime as default", - features.len() - ); - return Self::Ranging; - } - - let adx = features.get(cfg.adx_idx).copied().unwrap_or(0.0); - let cusum_direction = features.get(cfg.cusum_idx).copied().unwrap_or(0.0); - - if adx > cfg.adx_threshold { - Self::Trending - } else if cusum_direction.abs() > cfg.cusum_threshold { - Self::Volatile - } else { - Self::Ranging - } - } - - /// Classify regime for a batch of states entirely on GPU using tensor ops. - /// - /// Returns 3 binary mask tensors (trending, ranging, volatile) each of shape [`batch_size`]. - /// Feature indices and thresholds come from `cfg` (derived from `DQNConfig`). - /// - /// Zero CPU roundtrip -- all operations are GPU tensor ops dispatched on device. - pub fn classify_regime_masks_gpu( - states: &GpuTensor, - cfg: &RegimeClassConfig, - stream: &Arc, - ) -> Result<(GpuTensor, GpuTensor, GpuTensor), MLError> { - // Host-side classification: download states, classify per sample, upload masks. - // The hot path uses the fused CUDA trainer which bypasses this entirely. - let shape = states.shape(); - let batch_size = shape.first().copied().ok_or_else(|| { - MLError::ModelError("classify_regime_masks_gpu: states tensor has no dimensions".to_owned()) - })?; - let state_dim = shape.get(1).copied().ok_or_else(|| { - MLError::ModelError("classify_regime_masks_gpu: states tensor needs 2 dimensions".to_owned()) - })?; - - let host = states.to_host(stream)?; - - let mut trending = Vec::with_capacity(batch_size); - let mut ranging = Vec::with_capacity(batch_size); - let mut volatile = Vec::with_capacity(batch_size); - - for i in 0..batch_size { - let offset = i * state_dim; - let row = host.get(offset..offset + state_dim).ok_or_else(|| { - MLError::ModelError(format!("classify_regime_masks_gpu: row {i} out of bounds")) - })?; - let regime = Self::classify_from_features(row, cfg); - match regime { - RegimeType::Trending => { trending.push(1.0_f32); ranging.push(0.0); volatile.push(0.0); } - RegimeType::Ranging => { trending.push(0.0_f32); ranging.push(1.0); volatile.push(0.0); } - RegimeType::Volatile => { trending.push(0.0_f32); ranging.push(0.0); volatile.push(1.0); } - } - } - - let t = GpuTensor::from_host(&trending, vec![batch_size], stream)?; - let r = GpuTensor::from_host(&ranging, vec![batch_size], stream)?; - let v = GpuTensor::from_host(&volatile, vec![batch_size], stream)?; - Ok((t, r, v)) - } - - /// Get regime-specific reward scaling factor - /// - /// Adjusts rewards based on regime characteristics: - /// - Trending: Amplify trend-following rewards (1.2x) - /// - Ranging: Penalize volatility (0.8x) - /// - Volatile: Reduce magnitude to prevent overreaction (0.6x) - pub const fn reward_scale_factor(&self) -> f32 { - match self { - Self::Trending => 1.2, - Self::Ranging => 0.8, - Self::Volatile => 0.6, - } - } -} - -/// Per-regime training metrics -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct RegimeMetrics { - /// Number of training steps for this regime - pub training_steps: u64, - /// Cumulative loss - pub cumulative_loss: f64, - /// Average gradient norm - pub avg_grad_norm: f64, - /// Number of actions selected in this regime - pub action_count: u64, -} - -/// Regime-Conditional Deep Q-Network -/// -/// Maintains 3 independent Q-network heads (trending, ranging, volatile) and routes -/// actions/training based on detected market regime. -#[allow(missing_debug_implementations)] -pub struct RegimeConditionalDQN { - /// Trending regime Q-network head - trending_head: DQN, - /// Ranging regime Q-network head - ranging_head: DQN, - /// Volatile regime Q-network head - volatile_head: DQN, - - /// Per-regime training metrics - metrics: HashMap, - - /// Regime classification config (feature indices + thresholds) - regime_config: RegimeClassConfig, - - /// CUDA stream for GPU operations - stream: Arc, - - /// Gradient collapse counter (consecutive epochs with `grad_norm` below threshold) - gradient_collapse_counter: usize, -} - -impl RegimeConditionalDQN { - /// Create new regime-conditional DQN with 3 independent heads - pub fn new(config: DQNConfig) -> Result { - let device = MlDevice::cuda(0)?; - Self::new_on_device(config, device) - } - - /// Create regime-conditional DQN on a specific device. - pub fn new_on_device(config: DQNConfig, device: MlDevice) -> Result { - let stream = Arc::clone(device.cuda_stream()?); - let regime_config = RegimeClassConfig::from_dqn_config(&config); - - let trending_config = config.clone(); - let ranging_config = config.clone(); - let volatile_config = config; - - let trending_head = DQN::new_on_device(trending_config, device.clone())?; - let ranging_head = DQN::new_on_device(ranging_config, device.clone())?; - let volatile_head = DQN::new_on_device(volatile_config, device)?; - - let mut metrics = HashMap::new(); - metrics.insert(RegimeType::Trending, RegimeMetrics::default()); - metrics.insert(RegimeType::Ranging, RegimeMetrics::default()); - metrics.insert(RegimeType::Volatile, RegimeMetrics::default()); - - info!("RegimeConditionalDQN created with 3 independent heads"); - - Ok(Self { - trending_head, - ranging_head, - volatile_head, - metrics, - regime_config, - stream, - gradient_collapse_counter: 0, - }) - } - - /// Get a reference to the primary (trending) DQN head. - pub const fn primary_head(&self) -> &DQN { - &self.trending_head - } - - /// Get a mutable reference to the primary (trending) DQN head. - pub const fn primary_head_mut(&mut self) -> &mut DQN { - &mut self.trending_head - } - - /// Total training steps across all regime heads. - fn total_training_steps(&self) -> u64 { - self.metrics.values().map(|m| m.training_steps).sum() - } - - /// Gradient collapse detection at epoch boundary. - pub fn log_diagnostics(&mut self, grad_norm: f32) -> Result<(), MLError> { - let config = &self.trending_head.config; - let warmup_steps = (config.replay_buffer_capacity as f64 * 0.2) as u64; - let past_warmup = self.total_training_steps() > warmup_steps; - - let threshold = config.learning_rate as f32 * config.gradient_collapse_multiplier as f32; - let patience = config.gradient_collapse_patience; - - if past_warmup && grad_norm < threshold { - self.gradient_collapse_counter += 1; - tracing::warn!( - "GRADIENT COLLAPSE (regime-conditional): norm={:.6} (threshold: {:.6}) at total_steps {} (consecutive: {}/{})", - grad_norm, threshold, self.total_training_steps(), - self.gradient_collapse_counter, patience, - ); - - if self.gradient_collapse_counter >= patience { - return Err(MLError::TrainingError(format!( - "Training stopped: Gradient collapse detected for {} consecutive epochs \ - (threshold: {:.6}, actual: {:.6}). \ - This indicates the model has stopped learning. Consider:\n\ - 1. Increasing learning rate (current: {:.2e})\n\ - 2. Reducing batch size (current: {})\n\ - 3. Checking for data quality issues\n\ - 4. Adjusting network architecture", - self.gradient_collapse_counter, threshold, grad_norm, - config.learning_rate, config.batch_size, - ))); - } - } else if self.gradient_collapse_counter > 0 { - tracing::info!( - "Gradients recovered (regime-conditional): norm={:.6} (threshold: {:.6}), \ - resetting collapse counter (was: {})", - grad_norm, threshold, self.gradient_collapse_counter, - ); - self.gradient_collapse_counter = 0; - } else { - // Gradient norm healthy and no prior collapse - } - - Ok(()) - } - - /// Gradient collapse check without dead neuron detection (zero GPU->CPU sync). - pub fn check_gradient_collapse(&mut self, grad_norm: f32) -> Result<(), MLError> { - let config = &self.trending_head.config; - let warmup_steps = (config.collapse_warmup_capacity as f64 * 0.2) as u64; - let past_warmup = self.total_training_steps() > warmup_steps; - - let threshold = config.learning_rate as f32 * config.gradient_collapse_multiplier as f32; - let patience = config.gradient_collapse_patience; - - if past_warmup && grad_norm < threshold { - self.gradient_collapse_counter += 1; - if self.gradient_collapse_counter >= patience { - return Err(MLError::TrainingError(format!( - "Gradient collapse detected (regime-conditional): norm={:.6} < threshold={:.6} \ - for {} consecutive checks", - grad_norm, threshold, self.gradient_collapse_counter, - ))); - } - } else if self.gradient_collapse_counter > 0 { - self.gradient_collapse_counter = 0; - } else { - // Gradient norm healthy and no prior collapse - } - - Ok(()) - } - - /// Forward pass through regime-specific Q-network head - pub fn forward(&self, state: &GpuTensor, regime: RegimeType) -> Result { - match regime { - RegimeType::Trending => self.trending_head.forward(state), - RegimeType::Ranging => self.ranging_head.forward(state), - RegimeType::Volatile => self.volatile_head.forward(state), - } - } - - /// Select action using regime-specific head - pub fn select_action(&mut self, state: &[f32]) -> Result { - let regime = RegimeType::classify_from_features(state, &self.regime_config); - - let action = match regime { - RegimeType::Trending => self.trending_head.select_action(state)?, - RegimeType::Ranging => self.ranging_head.select_action(state)?, - RegimeType::Volatile => self.volatile_head.select_action(state)?, - }; - - if let Some(metrics) = self.metrics.get_mut(®ime) { - metrics.action_count += 1; - } - - debug!("Action selected via {:?} head: {:?}", regime, action); - - Ok(action) - } - - /// Batch greedy action selection across regime heads. - pub fn batch_greedy_actions(&self, states: &GpuTensor) -> Result, MLError> { - let n = states.shape().first().copied().unwrap_or(0); - let q_values = self.batch_q_values(states)?; - let reductions = ml_core::cuda_autograd::ReductionKernels::new(&self.stream)?; - reductions.argmax_rows(&q_values, n, self.config().num_actions) - } - - /// Batch Q-value computation across regime heads -- fully GPU-resident. - /// - /// Returns raw `CudaSlice` of shape `[batch_size, num_actions]` (flattened). - /// Callers know the shape from context — no GpuTensor wrapper needed at the boundary. - pub fn batch_q_values(&self, states: &GpuTensor) -> Result, MLError> { - let n = states.shape().first().copied().unwrap_or(0); - if n == 0 { - return Err(MLError::ModelError("Empty batch for batch_q_values".into())); - } - - let (trending_mask, ranging_mask, volatile_mask) = - RegimeType::classify_regime_masks_gpu(states, &self.regime_config, &self.stream)?; - - let trending_q = self.trending_head.q_values_for_batch(states)?; - let ranging_q = self.ranging_head.q_values_for_batch(states)?; - let volatile_q = self.volatile_head.q_values_for_batch(states)?; - - let trending_mask = trending_mask.unsqueeze(1, &self.stream).map_err(|e| { - MLError::ModelError(format!("trending unsqueeze: {e}")) - })?; - let ranging_mask = ranging_mask.unsqueeze(1, &self.stream).map_err(|e| { - MLError::ModelError(format!("ranging unsqueeze: {e}")) - })?; - let volatile_mask = volatile_mask.unsqueeze(1, &self.stream).map_err(|e| { - MLError::ModelError(format!("volatile unsqueeze: {e}")) - })?; - - let blended = trending_q.broadcast_mul(&trending_mask, &self.stream).map_err(|e| { - MLError::ModelError(format!("trending mul: {e}")) - })?; - let ranging_product = ranging_q.broadcast_mul(&ranging_mask, &self.stream).map_err(|e| { - MLError::ModelError(format!("ranging mul: {e}")) - })?; - let blended = blended.add(&ranging_product, &self.stream).map_err(|e| { - MLError::ModelError(format!("ranging add: {e}")) - })?; - let volatile_product = volatile_q.broadcast_mul(&volatile_mask, &self.stream).map_err(|e| { - MLError::ModelError(format!("volatile mul: {e}")) - })?; - let blended = blended.add(&volatile_product, &self.stream).map_err(|e| { - MLError::ModelError(format!("volatile add: {e}")) - })?; - - Ok(blended.into_parts().0) - } - - /// Per-branch Q-values blended across regime heads -- for branching DQN training path. - /// - /// Returns `Some((exposure, order, urgency))` as raw `CudaSlice` buffers. - /// Each buffer is `[batch_size, branch_size]` flattened. - pub fn batch_branching_q_values( - &self, - states: &GpuTensor, - ) -> Result, CudaSlice, CudaSlice)>, MLError> { - // use_branching is always on — unconditionally proceed - - let n = states.shape().first().copied().unwrap_or(0); - if n == 0 { - return Err(MLError::ModelError("Empty batch for batch_branching_q_values".into())); - } - - let (trending_mask, ranging_mask, volatile_mask) = - RegimeType::classify_regime_masks_gpu(states, &self.regime_config, &self.stream)?; - - let trending_mask = trending_mask.unsqueeze(1, &self.stream).map_err(|e| { - MLError::ModelError(format!("trending unsqueeze: {e}")) - })?; - let ranging_mask = ranging_mask.unsqueeze(1, &self.stream).map_err(|e| { - MLError::ModelError(format!("ranging unsqueeze: {e}")) - })?; - let volatile_mask = volatile_mask.unsqueeze(1, &self.stream).map_err(|e| { - MLError::ModelError(format!("volatile unsqueeze: {e}")) - })?; - - let mut branch_accumulators: Option<(GpuTensor, GpuTensor, GpuTensor)> = None; - - for (head, mask, label) in [ - (&self.trending_head, &trending_mask, "trending"), - (&self.ranging_head, &ranging_mask, "ranging"), - (&self.volatile_head, &volatile_mask, "volatile"), - ] { - let Some(branching_net) = head.branching_q_network.as_ref() else { - return Ok(None); - }; - - let output = branching_net.forward_branches_eval(states)?; - if output.advantages.len() < 3 { - return Err(MLError::ModelError(format!( - "{label} head branching produced {} branches, expected 3", - output.advantages.len() - ))); - } - - let masked_exp = output.advantages.first() - .ok_or_else(|| MLError::ModelError(format!("{label}: missing branch 0")))? - .broadcast_mul(mask, &self.stream).map_err(|e| { - MLError::ModelError(format!("{label} exp mask mul: {e}")) - })?; - let masked_ord = output.advantages.get(1) - .ok_or_else(|| MLError::ModelError(format!("{label}: missing branch 1")))? - .broadcast_mul(mask, &self.stream).map_err(|e| { - MLError::ModelError(format!("{label} ord mask mul: {e}")) - })?; - let masked_urg = output.advantages.get(2) - .ok_or_else(|| MLError::ModelError(format!("{label}: missing branch 2")))? - .broadcast_mul(mask, &self.stream).map_err(|e| { - MLError::ModelError(format!("{label} urg mask mul: {e}")) - })?; - - branch_accumulators = Some(match branch_accumulators { - None => (masked_exp, masked_ord, masked_urg), - Some((acc_e, acc_o, acc_u)) => ( - acc_e.add(&masked_exp, &self.stream).map_err(|e| { - MLError::ModelError(format!("{label} exp add: {e}")) - })?, - acc_o.add(&masked_ord, &self.stream).map_err(|e| { - MLError::ModelError(format!("{label} ord add: {e}")) - })?, - acc_u.add(&masked_urg, &self.stream).map_err(|e| { - MLError::ModelError(format!("{label} urg add: {e}")) - })?, - ), - }); - } - - Ok(branch_accumulators.map(|(e, o, u)| { - (e.into_parts().0, o.into_parts().0, u.into_parts().0) - })) - } - - /// Batch softmax action selection across regime heads (Gumbel-max, GPU-resident). - pub fn batch_softmax_actions( - &self, - states: &GpuTensor, - temperature: f64, - ) -> Result, MLError> { - let batch_size = states.shape().first().copied().unwrap_or(0); - let num_actions = self.config().num_actions; - let q_values = self.batch_q_values(states)?; - let temp = temperature.max(1e-6) as f32; - let shape = vec![batch_size, num_actions]; - - let q_host = self.stream.clone_dtoh(&q_values).map_err(|e| { - MLError::ModelError(format!("batch_softmax dtoh: {e}")) - })?; - let scaled: Vec = q_host.iter().map(|&v| v / temp).collect(); - let scaled_tensor = GpuTensor::from_host(&scaled, shape.clone(), &self.stream)?; - - let n = scaled.len(); - let mut gumbel_vals = Vec::with_capacity(n); - let mut rng = rand::thread_rng(); - use rand::Rng; - for _ in 0..n { - let u: f32 = rng.gen_range(0.001_f32..0.999_f32); - gumbel_vals.push(-(-u.ln()).ln()); - } - let gumbel = GpuTensor::from_host(&gumbel_vals, shape, &self.stream)?; - - let perturbed = scaled_tensor.add(&gumbel, &self.stream) - .map_err(|e| MLError::ModelError(format!("Gumbel perturbation failed: {}", e)))?; - perturbed - .argmax(1, &self.stream) - .map_err(|e| MLError::ModelError(format!("Gumbel-max argmax failed: {}", e))) - } - - /// GPU-resident hierarchical factored softmax. - pub fn batch_hierarchical_softmax_actions( - &self, - states: &GpuTensor, - temperature: f64, - ) -> Result, MLError> { - self.batch_softmax_actions(states, temperature) - } - - pub fn set_noise_sigma_scale(&mut self, scale: f64) { - self.trending_head.set_noise_sigma_scale(scale); - self.ranging_head.set_noise_sigma_scale(scale); - self.volatile_head.set_noise_sigma_scale(scale); - } - - pub fn reset_count_bonus(&mut self) { - self.trending_head.reset_count_bonus(); - self.ranging_head.reset_count_bonus(); - self.volatile_head.reset_count_bonus(); - } - - /// Get count bonuses for all actions (UCB exploration). - pub fn get_count_bonuses(&self) -> Vec { - self.trending_head.get_count_bonuses() - } - - /// Get shared DQN config (all heads share the same config). - pub const fn config(&self) -> &super::DQNConfig { - &self.trending_head.config - } - - /// Update epsilon for specific regime head - pub fn update_epsilon(&mut self, regime: RegimeType) { - match regime { - RegimeType::Trending => self.trending_head.update_epsilon(), - RegimeType::Ranging => self.ranging_head.update_epsilon(), - RegimeType::Volatile => self.volatile_head.update_epsilon(), - } - } - - /// Set epsilon for all regime heads - pub const fn set_epsilon_all(&mut self, epsilon: f64) { - self.trending_head.set_epsilon(epsilon); - self.ranging_head.set_epsilon(epsilon); - self.volatile_head.set_epsilon(epsilon); - } - - /// Get epsilon for specific regime head - pub const fn get_epsilon(&self, regime: RegimeType) -> f32 { - match regime { - RegimeType::Trending => self.trending_head.get_epsilon(), - RegimeType::Ranging => self.ranging_head.get_epsilon(), - RegimeType::Volatile => self.volatile_head.get_epsilon(), - } - } - - /// Update target networks for all heads - pub const fn update_target_networks(&mut self) -> Result<(), MLError> { - Ok(()) - } - - /// Get replay buffer size - pub fn get_replay_buffer_size(&self) -> Result { - self.trending_head.get_replay_buffer_size() - } - - /// Get per-regime training metrics - pub const fn get_regime_metrics(&self) -> &HashMap { - &self.metrics - } - - /// Save checkpoint for all 3 heads - pub fn save_checkpoint(&self, path: &str) -> Result<(), MLError> { - let trending_path = format!("{}_trending.safetensors", path); - let ranging_path = format!("{}_ranging.safetensors", path); - let volatile_path = format!("{}_volatile.safetensors", path); - - let save_head = |_head: &super::dqn::DQN, head_path: &str, label: &str| -> Result<(), MLError> { - // Checkpoint save is handled by the flat buffer path in GpuDqnTrainer. - let _ = head_path; - tracing::debug!("Would save {label} head checkpoint (flat buffer path)"); - Ok(()) - }; - - save_head(&self.trending_head, &trending_path, "trending")?; - save_head(&self.ranging_head, &ranging_path, "ranging")?; - save_head(&self.volatile_head, &volatile_path, "volatile")?; - - info!("RegimeConditionalDQN checkpoint saved: {}", path); - Ok(()) - } - - /// Load checkpoint for all 3 heads - pub fn load_checkpoint(&mut self, path: &str) -> Result<(), MLError> { - let trending_path = format!("{}_trending.safetensors", path); - let ranging_path = format!("{}_ranging.safetensors", path); - let volatile_path = format!("{}_volatile.safetensors", path); - - self.trending_head.load_from_safetensors(&trending_path)?; - self.ranging_head.load_from_safetensors(&ranging_path)?; - self.volatile_head.load_from_safetensors(&volatile_path)?; - - info!("RegimeConditionalDQN checkpoint loaded: {}", path); - Ok(()) - } - - /// Load all 3 heads from a single merged safetensors file. - pub fn load_from_merged_safetensors(&mut self, path: &str) -> Result<(), MLError> { - self.trending_head.load_from_safetensors(path)?; - info!( - "RegimeConditionalDQN loaded from merged checkpoint (trending head only): {}", - path - ); - Ok(()) - } - - /// Scale reward based on regime type - pub fn scale_reward(&self, reward: f32, regime: RegimeType) -> f32 { - reward * regime.reward_scale_factor() - } - - /// Get trending head (for testing and var access) - pub const fn get_trending_head(&self) -> Option<&DQN> { - Some(&self.trending_head) - } - - /// Get ranging head (for testing and var access) - pub const fn get_ranging_head(&self) -> Option<&DQN> { - Some(&self.ranging_head) - } - - /// Get volatile head (for testing and var access) - pub const fn get_volatile_head(&self) -> Option<&DQN> { - Some(&self.volatile_head) - } - - /// Get trending head's memory buffer (for trainer access) - pub const fn get_trending_head_memory(&self) -> &crate::replay_buffer_type::StagedGpuBuffer { - &self.trending_head.memory - } - - /// Get mutable trending head's memory buffer (for sampling). - pub fn get_trending_head_memory_mut(&mut self) -> &mut crate::replay_buffer_type::StagedGpuBuffer { - &mut self.trending_head.memory - } - - /// Get CUDA stream (for trainer access) - pub const fn get_stream(&self) -> &Arc { - &self.stream - } - - /// Get device (for trainer access) - pub fn get_device(&self) -> MlDevice { - MlDevice::Cuda { - context: self.stream.context().clone(), - stream: Arc::clone(&self.stream), - } - } - - /// Reinitialize categorical distribution with new bounds (adaptive C51) - pub fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), ml_core::MLError> { - self.trending_head.reinit_categorical_distribution(v_min, v_max)?; - self.ranging_head.reinit_categorical_distribution(v_min, v_max)?; - self.volatile_head.reinit_categorical_distribution(v_min, v_max)?; - Ok(()) - } - - /// Get state dimension from configuration - pub const fn get_state_dim(&self) -> usize { - self.trending_head.get_state_dim() - } - - /// Clear all regime-specific replay buffers - pub fn clear_replay_buffer(&mut self) -> Result<(), ml_core::MLError> { - self.trending_head.clear_replay_buffer()?; - self.ranging_head.clear_replay_buffer()?; - self.volatile_head.clear_replay_buffer()?; - Ok(()) - } - - /// Reset all regime-specific target networks - pub fn reset_target_network(&mut self) -> Result<(), ml_core::MLError> { - self.trending_head.reset_target_network()?; - self.ranging_head.reset_target_network()?; - self.volatile_head.reset_target_network()?; - Ok(()) - } - - /// Update learning rate for all regime heads - pub fn update_learning_rate(&mut self, decay_factor: f64) -> Result<(), ml_core::MLError> { - self.trending_head.update_learning_rate(decay_factor)?; - self.ranging_head.update_learning_rate(decay_factor)?; - self.volatile_head.update_learning_rate(decay_factor)?; - Ok(()) - } -} - -#[allow(clippy::items_after_test_module)] -#[cfg(test)] -#[allow(clippy::identity_op, clippy::erasing_op)] -mod tests { - use super::*; - - fn test_stream() -> Arc { - let ctx = cudarc::driver::CudaContext::new(0).expect("CUDA context required"); - ctx.new_stream().expect("CUDA stream required") - } - - #[test] - fn test_regime_classification_trending() { - let cfg = RegimeClassConfig::default(); - let mut features = vec![0.0_f32; 48]; - if let Some(v) = features.get_mut(cfg.adx_idx) { *v = 0.35; } - if let Some(v) = features.get_mut(cfg.cusum_idx) { *v = 0.0; } - assert_eq!(RegimeType::classify_from_features(&features, &cfg), RegimeType::Trending); - - if let Some(v) = features.get_mut(cfg.cusum_idx) { *v = 0.9; } - assert_eq!(RegimeType::classify_from_features(&features, &cfg), RegimeType::Trending); - } - - #[test] - fn test_regime_classification_volatile() { - let cfg = RegimeClassConfig::default(); - let mut features = vec![0.0_f32; 48]; - if let Some(v) = features.get_mut(cfg.adx_idx) { *v = 0.15; } - if let Some(v) = features.get_mut(cfg.cusum_idx) { *v = 0.85; } - assert_eq!(RegimeType::classify_from_features(&features, &cfg), RegimeType::Volatile); - - if let Some(v) = features.get_mut(cfg.cusum_idx) { *v = -0.9; } - assert_eq!(RegimeType::classify_from_features(&features, &cfg), RegimeType::Volatile); - } - - #[test] - fn test_regime_classification_ranging() { - let cfg = RegimeClassConfig::default(); - let mut features = vec![0.0_f32; 48]; - if let Some(v) = features.get_mut(cfg.adx_idx) { *v = 0.10; } - if let Some(v) = features.get_mut(cfg.cusum_idx) { *v = 0.3; } - assert_eq!(RegimeType::classify_from_features(&features, &cfg), RegimeType::Ranging); - - let features_zero = vec![0.0_f32; 48]; - assert_eq!(RegimeType::classify_from_features(&features_zero, &cfg), RegimeType::Ranging); - } - - #[test] - fn test_regime_classification_fallback() { - let cfg = RegimeClassConfig::default(); - let short = vec![0.0_f32; 10]; - assert_eq!(RegimeType::classify_from_features(&short, &cfg), RegimeType::Ranging); - - let empty: Vec = vec![]; - assert_eq!(RegimeType::classify_from_features(&empty, &cfg), RegimeType::Ranging); - } - - #[test] - fn test_regime_gpu_masks() { - let cfg = RegimeClassConfig::default(); - let stream = test_stream(); - let batch_size = 4; - let state_dim = 48; - - let mut data = vec![0.0_f32; batch_size * state_dim]; - // Row 0: Trending - if let Some(v) = data.get_mut(0 * state_dim + cfg.adx_idx) { *v = 0.5; } - if let Some(v) = data.get_mut(0 * state_dim + cfg.cusum_idx) { *v = 0.0; } - // Row 1: Volatile - if let Some(v) = data.get_mut(1 * state_dim + cfg.adx_idx) { *v = 0.1; } - if let Some(v) = data.get_mut(1 * state_dim + cfg.cusum_idx) { *v = 0.9; } - // Row 2: Ranging - if let Some(v) = data.get_mut(2 * state_dim + cfg.adx_idx) { *v = 0.1; } - if let Some(v) = data.get_mut(2 * state_dim + cfg.cusum_idx) { *v = 0.2; } - // Row 3: Trending - if let Some(v) = data.get_mut(3 * state_dim + cfg.adx_idx) { *v = 0.3; } - if let Some(v) = data.get_mut(3 * state_dim + cfg.cusum_idx) { *v = -0.8; } - - let states = GpuTensor::from_host(&data, vec![batch_size, state_dim], &stream).unwrap(); - let (trending, ranging, volatile) = RegimeType::classify_regime_masks_gpu(&states, &cfg, &stream).unwrap(); - - let t_host = trending.to_host(&stream).unwrap(); - let r_host = ranging.to_host(&stream).unwrap(); - let v_host = volatile.to_host(&stream).unwrap(); - - assert!((t_host.get(0).copied().unwrap_or(0.0) - 1.0).abs() < 1e-6); - assert!((t_host.get(1).copied().unwrap_or(0.0) - 0.0).abs() < 1e-6); - assert!((t_host.get(2).copied().unwrap_or(0.0) - 0.0).abs() < 1e-6); - assert!((t_host.get(3).copied().unwrap_or(0.0) - 1.0).abs() < 1e-6); - - assert!((r_host.get(2).copied().unwrap_or(0.0) - 1.0).abs() < 1e-6); - assert!((v_host.get(1).copied().unwrap_or(0.0) - 1.0).abs() < 1e-6); - - for i in 0..batch_size { - let row_sum = t_host.get(i).copied().unwrap_or(0.0) - + r_host.get(i).copied().unwrap_or(0.0) - + v_host.get(i).copied().unwrap_or(0.0); - assert!((row_sum - 1.0).abs() < 1e-6, "Row {i} mask sum: {row_sum}"); - } - } - - #[test] - fn test_reward_scaling() { - assert_eq!(RegimeType::Trending.reward_scale_factor(), 1.2); - assert_eq!(RegimeType::Ranging.reward_scale_factor(), 0.8); - assert_eq!(RegimeType::Volatile.reward_scale_factor(), 0.6); - } -} - -impl std::fmt::Debug for RegimeConditionalDQN { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("RegimeConditionalDQN") - .field("trending_head", &"DQN { ... }") - .field("ranging_head", &"DQN { ... }") - .field("volatile_head", &"DQN { ... }") - .field("metrics", &self.metrics) - .field("stream", &"Arc") - .finish() - } -} diff --git a/crates/ml/src/trainers/curriculum.rs b/crates/ml/src/trainers/curriculum.rs index d5dfbae02..7a39934e2 100644 --- a/crates/ml/src/trainers/curriculum.rs +++ b/crates/ml/src/trainers/curriculum.rs @@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize}; use crate::dqn::action_space::get_valid_action_mask; -use crate::dqn::regime_conditional::RegimeType; +use crate::dqn::RegimeType; /// Number of DQN exposure actions (ShortSmall..LongFull + Flat). const NUM_ACTIONS: usize = 7; diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 52c8a8102..3a47fe8da 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -4,25 +4,30 @@ //! Contains all training configuration including learning rates, //! batch sizes, exploration parameters, and risk management settings. +use std::sync::Arc; + use cudarc::driver::CudaSlice; use ml_core::device::MlDevice; use ml_core::cuda_autograd::GpuTensor; use crate::dqn::action_space::FactoredAction; use crate::dqn::dqn::DQN; -use crate::dqn::regime_conditional::{RegimeConditionalDQN, RegimeMetrics, RegimeType}; use crate::MLError; -/// DQN agent wrapper around `RegimeConditionalDQN` (3 heads: Trending, Ranging, Volatile). +/// DQN agent wrapper around `DQN`. /// /// Provides the unified API surface that the training loop, fused CUDA trainer, -/// and checkpoint code call. All methods delegate to `self.agent`. +/// and checkpoint code call. Phase 4 of the MoE redesign collapsed the 3-head +/// `RegimeConditionalDQN` wrapper down to a single `DQN` since the regime +/// conditioning infrastructure was vestigial (only `trending_head` ever trained). +/// MoE in the underlying `DQN` provides the regime-conditioned behavior via a +/// learned gate (Phase 3). pub struct DQNAgentType { - pub(crate) agent: RegimeConditionalDQN, + pub(crate) agent: DQN, } impl DQNAgentType { - pub fn new(agent: RegimeConditionalDQN) -> Self { + pub fn new(agent: DQN) -> Self { Self { agent } } } @@ -34,26 +39,29 @@ impl std::fmt::Debug for DQNAgentType { } impl DQNAgentType { - /// Select action using regime-conditional agent + /// Select action using the DQN agent. pub fn select_action(&mut self, state: &[f32]) -> Result { self.agent.select_action(state) } - /// Batch greedy action selection — single forward pass per architecture head. + /// Batch greedy action selection — single forward pass. /// - /// Returns `Vec` action indices directly (cold path, CPU readback). + /// Returns `Vec` action indices (cold path, CPU readback from GPU argmax). pub fn batch_greedy_actions(&self, states: &GpuTensor) -> Result, MLError> { - self.agent.batch_greedy_actions(states) + let tensor = self.agent.batch_greedy_actions(states)?; + let host = tensor.to_host(self.agent.cuda_stream()) + .map_err(|e| MLError::ModelError(format!("batch_greedy_actions DtoH: {e}")))?; + Ok(host.iter().map(|&v| v as u32).collect()) } /// Batch Q-value computation — returns raw `CudaSlice` of shape `[batch, num_actions]`. /// /// Used by `GpuBacktestEvaluator` for GPU-side argmax. pub fn batch_q_values(&self, states: &GpuTensor) -> Result, MLError> { - self.agent.batch_q_values(states) + self.agent.q_values_for_batch(states).map(|t| t.into_parts().0) } - /// Per-branch Q-values for branching DQN — regime-blended. + /// Per-branch Q-values for branching DQN. /// /// Returns `Some((exposure, order, urgency))` as raw `CudaSlice` buffers when /// branching is active, `None` otherwise. Used by the GPU action selector's @@ -67,239 +75,214 @@ impl DQNAgentType { /// Batch softmax (Boltzmann) action selection. /// - /// Returns `Vec` action indices directly (cold path, CPU readback). + /// Returns `GpuTensor` of action indices (GPU-resident, shape `[batch]`). pub fn batch_softmax_actions( &self, states: &GpuTensor, temperature: f64, - ) -> Result, MLError> { + ) -> Result { self.agent.batch_softmax_actions(states, temperature) } - /// Hierarchical factored softmax -- exposure-first, then order/urgency within. + /// Hierarchical factored softmax — exposure-first, then order/urgency within. /// - /// Returns `Vec` action indices directly (cold path, CPU readback). + /// Returns `GpuTensor` of action indices (GPU-resident, shape `[batch]`). pub fn batch_hierarchical_softmax_actions( &self, states: &GpuTensor, temperature: f64, - ) -> Result, MLError> { + ) -> Result { self.agent.batch_hierarchical_softmax_actions(states, temperature) } - /// Task 0.6 — NoisyNets σ mean for a specific action branch + /// NoisyNets σ mean for a specific action branch /// (0 = direction, 1 = magnitude, 2 = order, 3 = urgency). - /// Delegates through primary_head → branching network. Returns 0.0 if - /// branching not initialized. H7 detection signal for HEALTH_DIAG. + /// Returns 0.0 if branching network not initialized. H7 detection signal for HEALTH_DIAG. pub fn branch_noisy_sigma_mean(&self, branch_idx: usize) -> f32 { - self.agent.primary_head().branch_noisy_sigma_mean(branch_idx) + self.agent.branch_noisy_sigma_mean(branch_idx) } - /// Set noise sigma scale for annealed noisy exploration + /// Set noise sigma scale for annealed noisy exploration. pub fn set_noise_sigma_scale(&mut self, scale: f64) { self.agent.set_noise_sigma_scale(scale); } - /// Reset count-based exploration bonus (call at epoch boundary) + /// Reset count-based exploration bonus (call at epoch boundary). pub fn reset_count_bonus(&mut self) { self.agent.reset_count_bonus(); } /// Get count bonuses for all actions (UCB exploration). - /// Returns a Vec of length num_actions with exploration bonuses. + /// Returns a Vec of length `num_actions` with exploration bonuses. pub fn get_count_bonuses(&self) -> Vec { self.agent.get_count_bonuses() } /// Get per-branch count bonuses for branching DQN (direction [4], order [3], urgency [3]). - /// RegimeConditional uses flat bonuses, so this always returns None. - pub fn get_count_bonuses_branched(&self) -> Option<([f32; 4], [f32; 3], [f32; 3])> { - None + /// + /// Delegates directly to `DQN::get_count_bonuses_branched`. UCB count bonuses now + /// reach the GPU action selector — the ghost feature from the Phase 0 deferral is + /// fixed in this commit. + pub fn get_count_bonuses_branched(&self) -> ([f32; 4], [f32; 3], [f32; 3]) { + self.agent.get_count_bonuses_branched() } - - /// Number of discrete actions in the action space. pub fn num_actions(&self) -> usize { - self.agent.config().num_actions + self.agent.config.num_actions } /// Branch sizes for branching DQN: `(direction, magnitude, order_type, urgency)`. /// /// Returns the actual branch sizes from the DQN config. pub fn branch_sizes(&self) -> (usize, usize, usize, usize) { - let cfg = self.agent.config(); + let cfg = &self.agent.config; (cfg.num_actions, 3, cfg.num_order_types, cfg.num_urgency_levels) } - /// Update epsilon for exploration decay (all regime heads) + /// Update epsilon for exploration decay. pub fn update_epsilon(&mut self) { - self.agent.update_epsilon(RegimeType::Trending); - self.agent.update_epsilon(RegimeType::Ranging); - self.agent.update_epsilon(RegimeType::Volatile); + self.agent.update_epsilon(); } - /// Get current epsilon value (trending head as representative) + /// Get current epsilon value. pub fn get_epsilon(&self) -> f32 { - self.agent.get_epsilon(RegimeType::Trending) + self.agent.get_epsilon() } /// Get the effective epsilon for exploration, respecting noisy_epsilon_floor. /// The batch training path must use this instead of get_epsilon() to avoid /// zero exploration when noisy nets are enabled. pub fn get_effective_epsilon(&self) -> f32 { - self.agent.primary_head().get_effective_epsilon() + self.agent.get_effective_epsilon() } - /// Set epsilon value for exploration (all regime heads) + /// Set epsilon value for exploration. pub fn set_epsilon(&mut self, epsilon: f64) { - self.agent.set_epsilon_all(epsilon); + self.agent.set_epsilon(epsilon); } - /// Save checkpoint to disk with architecture metadata. + /// Get regime-specific metrics. /// - /// Downloads GpuVarStore parameters to host and serializes via safetensors. - pub fn save_checkpoint(&self, path: &str) -> Result<(), MLError> { - self.agent.save_checkpoint(path) + /// Always returns `None` — regime metrics were part of the deleted + /// `RegimeConditionalDQN`. MoE gate utilization is visible via HEALTH_DIAG + /// `aux_moe` line instead. + pub fn get_regime_metrics(&self) -> Option<&std::collections::HashMap> { + None } - /// Get regime-specific metrics - pub fn get_regime_metrics(&self) -> Option<&std::collections::HashMap> { - Some(self.agent.get_regime_metrics()) - } - - /// Get replay buffer size + /// Get replay buffer size. pub fn get_replay_buffer_size(&self) -> Result { self.agent.get_replay_buffer_size() } - /// Forward pass through Q-network. + /// Forward pass through Q-network (branching path). /// - /// Uses GPU regime classification masks to blend Q-values from all 3 heads. /// Returns raw `CudaSlice` of shape `[batch, num_actions]`. pub fn forward(&self, state: &GpuTensor) -> Result, MLError> { - self.agent.batch_q_values(state) + self.agent.q_values_for_batch(state).map(|t| t.into_parts().0) } /// Track action for diversity monitoring. - /// - /// Regime-conditional agent tracks actions via regime metrics, so this is a no-op. - pub fn track_action(&mut self, _action: FactoredAction) { - // Regime-conditional agent doesn't use track_action - // Action tracking happens via regime metrics instead + pub fn track_action(&mut self, action: FactoredAction) { + self.agent.track_action(action); } /// Track a batch of actions for diversity monitoring. - /// - /// Regime-conditional agent tracks actions via regime metrics, so this is a no-op. - pub fn track_actions_batch(&mut self, _actions: &[FactoredAction]) { - // Regime-conditional agent doesn't use track_action + pub fn track_actions_batch(&mut self, actions: &[FactoredAction]) { + self.agent.track_actions_batch(actions); } // Legacy insert_batch_tensors DELETED — use insert_batch directly // from the training loop with f32 CudaSlice states (Task #30). - /// Check if agent can train (has enough replay buffer samples) + /// Check if agent can train (has enough replay buffer samples). pub fn can_train(&self) -> bool { - // Check if replay buffer has minimum samples - // Use 100 as min_replay_size (same as in regime_conditional.rs train_step) - if let Ok(buffer) = self.agent.get_replay_buffer_size() { - buffer >= 100 - } else { - false - } + self.agent.can_train() } - /// Get reference to replay buffer memory + /// Get reference to replay buffer memory. pub fn memory(&self) -> &crate::dqn::replay_buffer_type::StagedGpuBuffer { - self.agent.get_trending_head_memory() + &self.agent.memory } /// Get mutable reference to replay buffer memory (for sampling). pub fn memory_mut(&mut self) -> &mut crate::dqn::replay_buffer_type::StagedGpuBuffer { - self.agent.get_trending_head_memory_mut() + &mut self.agent.memory } - /// WAVE 23 P0: Log diagnostics and check for gradient collapse (early stopping) - /// Returns Err if gradient collapse detected for consecutive epochs + /// Log diagnostics and check for gradient collapse (early stopping). + /// Returns Err if gradient collapse detected for consecutive epochs. pub fn log_diagnostics(&mut self, grad_norm: f32) -> Result<(), MLError> { self.agent.log_diagnostics(grad_norm) } - /// Gradient collapse check only -- no dead neuron detection (zero GPU->CPU sync). + /// Gradient collapse check only — no dead neuron detection (zero GPU→CPU sync). /// Use in CUDA guard hot path; dead neuron detection runs at epoch boundary. pub fn check_gradient_collapse(&mut self, grad_norm: f32) -> Result<(), MLError> { self.agent.check_gradient_collapse(grad_norm) } /// Suppress or enable forward() monitoring (GPU→CPU sync avoidance in hot path). - /// - /// No-op for regime-conditional agent. - pub fn set_training_forward_active(&mut self, _active: bool) { - // Regime-conditional agent doesn't use training_forward_active + pub fn set_training_forward_active(&mut self, active: bool) { + self.agent.set_training_forward_active(active); } - /// WAVE 23 P0: Log Q-values and check for Q-value divergence (early stopping) - /// Returns Err if Q-value divergence detected for consecutive checks - pub fn log_q_values(&mut self, _states_tensor: &GpuTensor) -> Result<(), MLError> { - // Regime-conditional doesn't implement Q-value divergence detection yet - Ok(()) + /// Log Q-values and check for Q-value divergence (early stopping). + /// Returns Err if Q-value divergence detected for consecutive checks. + pub fn log_q_values(&mut self, states_tensor: &GpuTensor) -> Result<(), MLError> { + self.agent.log_q_values(states_tensor) } /// Log Q-values from pre-computed GPU statistics (zero tensor readback). /// Returns Err if Q-value divergence detected for consecutive checks. pub fn log_q_values_from_stats( &mut self, - _q_min: f32, - _q_max: f32, - _q_mean: f32, - _q_variance: f32, - _n_actions: usize, + q_min: f32, + q_max: f32, + q_mean: f32, + q_variance: f32, + n_actions: usize, ) -> Result<(), MLError> { - // Regime-conditional doesn't implement Q-value divergence detection yet - Ok(()) + self.agent.log_q_values_from_stats(q_min, q_max, q_mean, q_variance, n_actions) } - /// Get device (CPU or CUDA) + /// Get CUDA device. pub fn device(&self) -> MlDevice { - self.agent.get_device() - } - - - /// Build safetensors checkpoint metadata from the agent's config. - pub fn checkpoint_metadata(&self) -> std::collections::HashMap { - // Use trending head config as representative for regime-conditional - match self.agent.get_trending_head() { - Some(head) => head.config.checkpoint_metadata(), - None => std::collections::HashMap::new(), + MlDevice::Cuda { + context: self.agent.cuda_stream().context().clone(), + stream: Arc::clone(self.agent.cuda_stream()), } } - /// Get immutable reference to the primary DQN head. + /// Build safetensors checkpoint metadata from the agent's config. + pub fn checkpoint_metadata(&self) -> std::collections::HashMap { + self.agent.config.checkpoint_metadata() + } + + /// Get immutable reference to the underlying DQN. pub fn primary_dqn(&self) -> &DQN { - self.agent.primary_head() + &self.agent } - /// Get mutable reference to the primary DQN head. + /// Get mutable reference to the underlying DQN. pub fn primary_dqn_mut(&mut self) -> &mut DQN { - self.agent.primary_head_mut() + &mut self.agent } - /// Get state dimension from agent configuration + /// Get state dimension from agent configuration. /// - /// WAVE 10.4: Added to fix hardcoded STATE_DIM bug - /// Returns the actual state dimension (72 without OFI, 80 with OFI) - /// Layout: 42 market + 8 portfolio + 16 multi-timeframe + (8 OFI), 8-aligned + /// Returns the actual state dimension from the canonical STATE_DIM constant. pub fn get_state_dim(&self) -> usize { self.agent.get_state_dim() } - /// BUG #38 FIX: Clear replay buffer(s) + /// Clear replay buffer. pub fn clear_replay_buffer(&mut self) -> Result<(), MLError> { self.agent.clear_replay_buffer() } - /// BUG #38 FIX: Reset target network(s) to match main Q-network(s) + /// Reset target network to match main Q-network. pub fn reset_target_network(&mut self) -> Result<(), MLError> { self.agent.reset_target_network() } @@ -365,39 +348,29 @@ impl DQNAgentType { self.primary_dqn().memory.capacity() } - /// Update learning rate for the optimizer(s) by applying a decay factor. - /// - /// Updates all three regime head optimizers. + /// Update learning rate for the optimizer by applying a decay factor. pub fn update_learning_rate(&mut self, decay_factor: f64) -> Result<(), crate::MLError> { self.agent.update_learning_rate(decay_factor) } - /// Whether the active agent architecture uses Branching DQN (3 independent heads). + /// Whether the active agent architecture uses Branching DQN. /// /// When `true`, weight sets are extracted via `weight_sets_from_branching()` /// which provides zero-copy pointer views into the branching network's live weights. pub fn is_using_branching(&self) -> bool { - self.agent.get_trending_head() - .is_some_and(|h| h.branching_q_network.is_some()) + self.agent.branching_q_network.is_some() } /// Network dimensions for pure-CUDA forward kernel: `(shared_h1, shared_h2, value_h, adv_h)`. /// /// These are injected as `#define` constants into the NVRTC-compiled forward kernel. - /// The CUDA kernel assumes exactly 2 shared layers (matching DQN's `hidden_dims[0..2]`). pub fn network_dims(&self) -> (usize, usize, usize, usize) { - // Use trending head config as representative - match self.agent.get_trending_head() { - Some(head) => { - let h = &head.config.hidden_dims; - let shared_h1 = h.first().copied().unwrap_or(256); - let shared_h2 = h.get(1).copied().unwrap_or(shared_h1); - let value_h = head.config.dueling_hidden_dim; - let adv_h = head.config.branch_hidden_dim; - (shared_h1, shared_h2, value_h, adv_h) - } - None => (256, 256, 128, 128), // safe defaults - } + let h = &self.agent.config.hidden_dims; + let shared_h1 = h.first().copied().unwrap_or(256); + let shared_h2 = h.get(1).copied().unwrap_or(shared_h1); + let value_h = self.agent.config.dueling_hidden_dim; + let adv_h = self.agent.config.branch_hidden_dim; + (shared_h1, shared_h2, value_h, adv_h) } } diff --git a/crates/ml/src/trainers/dqn/trainer/action.rs b/crates/ml/src/trainers/dqn/trainer/action.rs index a88817c7c..e6a688724 100644 --- a/crates/ml/src/trainers/dqn/trainer/action.rs +++ b/crates/ml/src/trainers/dqn/trainer/action.rs @@ -107,13 +107,12 @@ impl DQNTrainer { let coeff = self.hyperparams.count_bonus_coefficient.unwrap_or(0.0); if coeff > 0.0 { let agent_r = self.agent.read().await; - if let Some((mut be, mut bo, mut bu)) = agent_r.get_count_bonuses_branched() { - let coeff_f32 = coeff as f32; - for x in &mut be { *x *= coeff_f32; } - for x in &mut bo { *x *= coeff_f32; } - for x in &mut bu { *x *= coeff_f32; } - let _ = selector.set_count_bonuses(&be, &bo, &bu); - } + let (mut be, mut bo, mut bu) = agent_r.get_count_bonuses_branched(); + let coeff_f32 = coeff as f32; + for x in &mut be { *x *= coeff_f32; } + for x in &mut bo { *x *= coeff_f32; } + for x in &mut bu { *x *= coeff_f32; } + let _ = selector.set_count_bonuses(&be, &bo, &bu); drop(agent_r); } diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 6ea0eb786..e2e77da07 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -19,7 +19,7 @@ use crate::dqn::curiosity::CuriosityModule; use crate::dqn::dqn::DQNConfig; use crate::dqn::logging::{LoggingConfig, MetricsAggregator}; use crate::dqn::portfolio_tracker::PortfolioTracker; -use crate::dqn::regime_conditional::RegimeConditionalDQN; +use crate::dqn::dqn::DQN; use crate::dqn::reward::{RewardConfig, RewardFunction}; use crate::features::microstructure_features::*; use crate::labeling::triple_barrier::TripleBarrierEngine; @@ -322,13 +322,13 @@ impl DQNTrainer { } else { device.clone() }; - // Regime-conditional Q-network is always active - 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, agent_device.clone()) - .map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?; - let agent = DQNAgentType::new(regime_agent); + // Single DQN with MoE gate (Phase 3+). Regime conditioning via learned + // MoE gate that reads ADX (idx 40) + CUSUM (idx 41) as part of the full + // state vector — strictly subsuming the deleted 3-head threshold classifier. + info!("Creating DQN with MoE gate (8 experts, learned regime conditioning)"); + let dqn = DQN::new_on_device(config, agent_device.clone()) + .map_err(|e| anyhow::anyhow!("Failed to create DQN: {}", e))?; + let agent = DQNAgentType::new(dqn); // Initialize portfolio tracker with $100k starting capital and 1 basis point spread diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index e0fc8eef3..2d00e6442 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -1310,42 +1310,25 @@ impl DQNTrainer { /// 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. + /// Serializes the single DQN (MoE + branching heads) into a safetensors file. + /// The Phase 3 MoE gate and 8 expert weights are included via the branching + /// network's `named_weight_slices`. Phase 4 removed the 3-head + /// `RegimeConditionalDQN` wrapper; checkpoint format is now a single flat + /// namespace (no `trending__` / `ranging__` / `volatile__` prefixes). pub async fn serialize_model(&self) -> Result> { let agent = self.agent.read().await; let tensors: std::collections::HashMap = { - let regime = &agent.agent; - { - 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) - })?; - // Extract named weight slices directly from branching network. - let br = if let Some(ref br) = head.branching_q_network { - br - } else { - continue; - }; - for (name, slice, shape) in br.named_weight_slices() { - if let Ok(tensor) = GpuTensor::new(slice.clone(), shape) { - all_tensors.insert( - format!("{}{}", prefix, name), - tensor, - ); - } + let dqn = &agent.agent; + let mut all_tensors = std::collections::HashMap::new(); + if let Some(ref br) = dqn.branching_q_network { + for (name, slice, shape) in br.named_weight_slices() { + if let Ok(tensor) = GpuTensor::new(slice.clone(), shape) { + all_tensors.insert(name, tensor); } } - all_tensors } + all_tensors }; // Embed architecture metadata in safetensors header diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 33962db83..92d35aa2e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -832,3 +832,25 @@ loads cubin, exposes test_mixture_forward using mapped pinned buffers exclusively (no HtoD/HtoH per `feedback_no_htod_htoh_only_mapped_pinned.md`). Unit test in `crates/ml/tests/moe_kernels_test.rs` verifies CPU reference match within 1e-6 for B=4, K=8, C=256. + +RegimeConditionalDQN deletion (2026-04-27): atomic removal per +feedback_no_partial_refactor.md. Deleted regime_adx_idx, regime_cusum_idx, +regime_adx_threshold, regime_cusum_threshold fields from DQNConfig (both +struct definition and Default + aggressive() builder entries). DQNAgentType +rewritten as thin wrapper over single DQN directly: no per-regime delegation +methods, no multi-head epsilon/epsilon_all calls, no get_trending_head / +primary_head indirection, no get_trending_head_memory accessor. The ghost +feature get_count_bonuses_branched (previously hardcoded None) now delegates +to DQN::get_count_bonuses_branched() — UCB count bonuses reach the GPU +action selector for the first time. action.rs caller simplified to direct +destructure of fixed arrays (no Option match). serialize_model in mod.rs +rewritten to serialize the single DQN branching network directly (no +trending__/ranging__/volatile__ prefix namespace). Constructor drops +RegimeConditionalDQN::new_on_device and calls DQN::new_on_device directly. +curriculum.rs import fixed from crate::dqn::regime_conditional::RegimeType +to crate::dqn::RegimeType (re-exported from regime_classifier). Phase 3 MoE +(commit a52d99613) provides the regime-conditioned behavior the legacy 3-head +architecture pretended to do. Workspace clean: 0 errors across all crates, tests, and examples. +Smoke validates no regression: 3/3 folds passed (405s), gate utilization +preserved (util=0.119,0.119,0.169,...), HEALTH_DIAG aux_moe line emits, +all fold checkpoints written.