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

refactor(dqn): delete vestigial RegimeConditionalDQN — MoE replaces it

Browse Source 
Atomic cleanup per feedback_no_partial_refactor.md:

- Delete crates/ml-dqn/src/regime_conditional.rs and all
  RegimeConditional* exports from lib.rs. regime_classifier.rs
  (RegimeType, RegimeClassConfig) is kept — used by validation layer
  and ml-regime-detection crate.
- Rewrite DQNAgentType to wrap DQN directly (no RegimeConditionalDQN
  field, no 3-head delegation gymnastics).
- DQNAgentType::get_count_bonuses_branched no longer returns hardcoded
  None — it delegates to DQN::get_count_bonuses_branched() and UCB
  count bonuses reach the GPU action selector for the first time
  (ghost feature from Phase 0 deferral). action.rs caller simplified
  to direct destructure of fixed arrays, no Option matching.
- Drop 4 regime-threshold fields from DQNConfig (regime_adx_idx,
  regime_cusum_idx, regime_adx_threshold, regime_cusum_threshold) +
  matching Default and aggressive() builder entries.
- serialize_model rewritten to serialize single DQN branching network
  directly (no trending__/ranging__/volatile__ prefix namespace).
- curriculum.rs import fixed: crate::dqn::regime_conditional::RegimeType
  → crate::dqn::RegimeType (re-exported from regime_classifier).
- Constructor drops RegimeConditionalDQN::new_on_device, calls
  DQN::new_on_device directly.
- Stale doc comments in dqn.rs, moe.rs, regime_classifier.rs updated.
- Wire-up audit updated.

Phase 3 MoE (commit a52d99613) provides the regime-conditioned behavior
the legacy 3-head architecture pretended to do — gate sees ADX (40)
and CUSUM (41) as part of the full 128-dim state vector and learns its
own decomposition, strictly subsuming the threshold classifier.

Smoke: 3/3 folds passed (405s), gate util=0.119,0.119,0.169,... preserved,
all fold checkpoints written. Workspace: 0 errors across all crates,
tests, and examples.

Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §7.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-27 20:21:09 +02:00
parent a52d996135
commit ff00af68a3
11 changed files with 353 additions and 1089 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
crates/ml-dqn/src/dqn.rs
View File

@@ -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<f32>` 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<Option<(cudarc::driver::CudaSlice<f32>, cudarc::driver::CudaSlice<f32>, cudarc::driver::CudaSlice<f32>)>, 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
///

6
crates/ml-dqn/src/lib.rs
View File

@@ -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};

5
crates/ml-dqn/src/moe.rs
View File

@@ -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;

144
crates/ml-dqn/src/regime_classifier.rs Normal file
View File

@@ -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<CudaStream>,
) -> 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))
}
}

876
crates/ml-dqn/src/regime_conditional.rs
View File

@@ -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<CudaStream>,
) -> 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<RegimeType, RegimeMetrics>,
/// Regime classification config (feature indices + thresholds)
regime_config: RegimeClassConfig,
/// CUDA stream for GPU operations
stream: Arc<CudaStream>,
/// 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<Self, MLError> {
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<Self, MLError> {
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<GpuTensor, MLError> {
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<FactoredAction, MLError> {
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(&regime) {
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<Vec<u32>, 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<f32>` 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<CudaSlice<f32>, 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<f32>` buffers.
/// Each buffer is `[batch_size, branch_size]` flattened.
pub fn batch_branching_q_values(
&self,
states: &GpuTensor,
) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, 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<Vec<u32>, 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<f32> = 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<Vec<u32>, 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<f32> {
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<usize, MLError> {
self.trending_head.get_replay_buffer_size()
}
/// Get per-regime training metrics
pub const fn get_regime_metrics(&self) -> &HashMap<RegimeType, RegimeMetrics> {
&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<CudaStream> {
&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<CudaStream> {
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<f32> = 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<CudaStream>")
.finish()
}
}

2
crates/ml/src/trainers/curriculum.rs
View File

@@ -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;

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

@@ -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<FactoredAction, MLError> {
self.agent.select_action(state)
}
/// Batch greedy action selection — single forward pass per architecture head.
/// Batch greedy action selection — single forward pass.
///
/// Returns `Vec<u32>` action indices directly (cold path, CPU readback).
/// Returns `Vec<u32>` action indices (cold path, CPU readback from GPU argmax).
pub fn batch_greedy_actions(&self, states: &GpuTensor) -> Result<Vec<u32>, 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<f32>` of shape `[batch, num_actions]`.
///
/// Used by `GpuBacktestEvaluator` for GPU-side argmax.
pub fn batch_q_values(&self, states: &GpuTensor) -> Result<CudaSlice<f32>, 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<f32>` 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<u32>` 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<Vec<u32>, MLError> {
) -> Result<GpuTensor, MLError> {
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<u32>` 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<Vec<u32>, MLError> {
) -> Result<GpuTensor, MLError> {
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<f32> of length num_actions with exploration bonuses.
/// Returns a Vec<f32> of length `num_actions` with exploration bonuses.
pub fn get_count_bonuses(&self) -> Vec<f32> {
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<crate::dqn::RegimeType, crate::validation::RegimeMetrics>> {
None
}
/// Get regime-specific metrics
pub fn get_regime_metrics(&self) -> Option<&std::collections::HashMap<RegimeType, RegimeMetrics>> {
Some(self.agent.get_regime_metrics())
}
/// Get replay buffer size
/// Get replay buffer size.
pub fn get_replay_buffer_size(&self) -> Result<usize, MLError> {
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<f32>` of shape `[batch, num_actions]`.
pub fn forward(&self, state: &GpuTensor) -> Result<CudaSlice<f32>, 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<String, String> {
// 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<String, String> {
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)
}
}

13
crates/ml/src/trainers/dqn/trainer/action.rs
View File

@@ -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);
}

16
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

@@ -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

41
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -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<Vec<u8>> {
let agent = self.agent.read().await;
let tensors: std::collections::HashMap<String, GpuTensor> = {
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

22
docs/dqn-wire-up-audit.md
View File

@@ -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.