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

refactor(ml): extract DQN module into ml-dqn crate (task 6)

Browse Source 
Move 53 DQN source files + gpu_replay_buffer from ml into standalone
ml-dqn crate. The ml crate's dqn module is now a thin re-export layer
(`pub use ml_dqn::*`) plus two bridge files (trainable_adapter.rs,
stress_testing.rs) that depend on ml-internal types.

Key changes:
- ml-dqn: 45 modules, 334 tests, compiles standalone
- ml: depends on ml-dqn, re-exports via dqn/mod.rs
- DQN.config: pub(crate) → pub for cross-crate access
- DQNAgent checkpoint methods moved into ml-dqn
- gpu_replay_buffer moved from cuda_pipeline/ to ml-dqn
- 8 dead-code files removed (never declared in mod.rs)

Net: -30,648 lines from ml crate. Workspace: 0 errors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-08 01:16:40 +01:00
parent d2f3b6c799
commit 88f0b3ec23
61 changed files with 405 additions and 4434 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
Cargo.lock generated
View File

@@ -6263,6 +6263,7 @@ dependencies = [
"memmap2",
"mimalloc",
"ml-core",
"ml-dqn",
"nalgebra 0.33.2",
"ndarray",
"num",
@@ -6372,7 +6373,32 @@ dependencies = [
name = "ml-dqn"
version = "1.0.0"
dependencies = [
"anyhow",
"approx",
"async-trait",
"bincode",
"candle-core",
"candle-nn",
"candle-optimisers",
"chrono",
"common",
"config",
"hex",
"ml-core",
"ndarray",
"parking_lot 0.12.5",
"rand 0.8.5",
"rand_distr 0.4.3",
"risk",
"rust_decimal",
"safetensors",
"serde",
"serde_json",
"sha2",
"tempfile",
"thiserror 1.0.69",
"tokio",
"tracing",
]
[[package]]

46
crates/ml-dqn/Cargo.toml
View File

@@ -11,10 +11,54 @@ documentation.workspace = true
publish.workspace = true
keywords.workspace = true
categories.workspace = true
description = "DQN reinforcement learning"
description = "DQN reinforcement learning for Foxhunt trading"
[features]
default = ["cuda"]
cuda = ["candle-core/cuda", "candle-core/cudnn", "candle-nn/cuda", "candle-nn/cudnn"]
[dependencies]
ml-core.workspace = true
common.workspace = true
config.workspace = true
risk = { path = "../risk" }
# ML frameworks
candle-core = { git = "https://github.com/huggingface/candle", rev = "671de1db" }
candle-nn = { git = "https://github.com/huggingface/candle", rev = "671de1db" }
candle-optimisers = { git = "https://github.com/KGrewal1/optimisers" }
# Serialization
serde = { workspace = true, features = ["derive"] }
serde_json.workspace = true
safetensors = "0.4"
bincode = "1.3"
# Core utilities
thiserror.workspace = true
anyhow.workspace = true
tracing.workspace = true
chrono.workspace = true
rand.workspace = true
rand_distr.workspace = true
rust_decimal.workspace = true
async-trait.workspace = true
tokio.workspace = true
# Crypto (checkpoint signing)
sha2 = "0.10"
hex = "0.4"
# Concurrency
parking_lot = { version = "0.12", features = ["hardware-lock-elision"] }
# Numerics
ndarray = { workspace = true, features = ["rayon"] }
[dev-dependencies]
tokio = { workspace = true, features = ["test-util", "macros"] }
approx.workspace = true
tempfile = "3"
[lints]
workspace = true

75
crates/ml/src/dqn/agent.rs → crates/ml-dqn/src/agent.rs
View File

@@ -5,10 +5,10 @@
use std::collections::HashMap;
use crate::Adam;
use ml_core::optimizers::Adam;
use candle_core::Tensor;
use candle_nn::{ops::leaky_relu, Module, VarBuilder};
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
use candle_optimisers::adam::ParamsAdam; // Use our Adam wrapper from lib.rs
use serde::{Deserialize, Serialize};
use tracing::debug;
@@ -22,7 +22,7 @@ use rust_decimal::Decimal;
use super::dqn::DQNConfig;
use super::network::{QNetwork, QNetworkConfig};
use super::{Experience, ReplayBuffer, ReplayBufferConfig};
use crate::MLError;
use ml_core::MLError;
pub use crate::trading_action::TradingAction;
@@ -997,7 +997,7 @@ impl DQNAgent {
for action_idx in 0..n_actions {
// Map exposure index → FactoredAction for profitability check
if let Ok(exposure) = super::action_space::ExposureLevel::from_index(action_idx) {
let action = crate::dqn::order_router::OrderRouter::route_default(exposure);
let action = crate::order_router::OrderRouter::route_default(exposure);
if !self.is_trade_profitable(&action, current_price, expected_price, current_position, max_position)? {
masked_q[action_idx] = f32::NEG_INFINITY;
}
@@ -1051,6 +1051,73 @@ impl DQNAgent {
let exposure = super::action_space::ExposureLevel::from_index(action_idx)?;
Ok(super::order_router::OrderRouter::route_default(exposure))
}
// --- Checkpoint helper methods (used by ml crate's checkpoint system) ---
/// Extract network weights from the model
pub fn extract_network_weights(&self, network_name: &str) -> Option<Vec<u8>> {
match network_name {
"q_network" => {
let weights = vec![0.1_f32, 0.2_f32, 0.3_f32, 0.4_f32];
let mut buffer = Vec::new();
for weight in weights {
buffer.extend_from_slice(&weight.to_le_bytes());
}
Some(buffer)
},
_ => None,
}
}
/// Get current replay buffer size
pub fn get_replay_buffer_size(&self) -> usize {
let filled_ratio = 0.7;
(self.config.replay_buffer_capacity as f64 * filled_ratio) as usize
}
/// Get average reward from recent episodes
pub fn get_average_reward(&self) -> f64 {
let episodes = self.metrics.total_episodes;
if episodes > 100 {
10.0 + (episodes as f64 / 100.0)
} else {
-5.0 + (episodes as f64 / 20.0)
}
}
/// Restore network weights from serialized data
pub fn restore_network_weights(
&mut self,
network_name: &str,
weights_data: &[u8],
) -> Result<(), String> {
match network_name {
"q_network" | "target_network" => {
if weights_data.len() % 4 != 0 {
return Err("Invalid weight data length".to_owned());
}
let weight_count = weights_data.len() / 4;
let mut weights = Vec::with_capacity(weight_count);
for i in 0..weight_count {
let start_idx = i * 4;
let weight_bytes = &weights_data[start_idx..start_idx + 4];
let weight = f32::from_le_bytes([
weight_bytes[0],
weight_bytes[1],
weight_bytes[2],
weight_bytes[3],
]);
weights.push(weight);
}
tracing::info!("Restored {} weights for {}", weights.len(), network_name);
Ok(())
},
_ => Err(format!("Unknown network: {}", network_name)),
}
}
}
// Manual Debug implementation for DQNAgent

8
crates/ml/src/dqn/attention.rs → crates/ml-dqn/src/attention.rs
View File

@@ -35,9 +35,9 @@ use candle_core::{Device, Result as CandleResult, Tensor};
use candle_nn::{Linear, Module, VarBuilder};
use serde::{Deserialize, Serialize};
use crate::cuda_compat::layer_norm_with_fallback;
use crate::dqn::xavier_init::linear_xavier;
use crate::MLError;
use ml_core::cuda_compat::layer_norm_with_fallback;
use crate::xavier_init::linear_xavier;
use ml_core::MLError;
/// Configuration for Multi-Head Attention layer
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -424,7 +424,7 @@ impl MultiHeadAttention {
mod tests {
use super::*;
use candle_nn::VarMap;
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
#[test]
fn test_config_validation() {

2
crates/ml/src/dqn/circuit_breaker.rs → crates/ml-dqn/src/circuit_breaker.rs
View File

@@ -3,6 +3,6 @@
//! Re-exports the shared circuit breaker implementation from the common module.
//! The canonical implementation lives in `crate::common::circuit_breaker`.
pub use crate::common::circuit_breaker::{
pub use ml_core::common::circuit_breaker::{
CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState,
};

0
crates/ml/src/dqn/count_bonus.rs → crates/ml-dqn/src/count_bonus.rs
View File

4
crates/ml/src/dqn/curiosity.rs → crates/ml-dqn/src/curiosity.rs
View File

@@ -8,8 +8,8 @@ use candle_nn::{ops::leaky_relu, AdamW, Linear, Module, Optimizer, ParamsAdamW,
use super::action_space::{FactoredAction, ExposureLevel};
use super::mixed_precision::training_dtype;
use crate::MLError;
use crate::dqn::xavier_init::linear_xavier;
use ml_core::MLError;
use crate::xavier_init::linear_xavier;
/// Forward dynamics model that predicts next state from (state, action)
#[allow(missing_debug_implementations)]

8
crates/ml/src/dqn/demo_dqn.rs → crates/ml-dqn/src/demo_dqn.rs
View File

@@ -3,10 +3,10 @@
//! This module provides a comprehensive demonstration of the 2025 production-ready
//! DQN implementation for high-frequency trading.
use crate::dqn::agent::DQNAgent;
use crate::dqn::DQNConfig;
use crate::safety::{MLSafetyConfig, MLSafetyManager};
use crate::MLError;
use crate::agent::DQNAgent;
use crate::DQNConfig;
use ml_core::safety::{MLSafetyConfig, MLSafetyManager};
use ml_core::MLError;
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize}; // For Decimal::from_f64

2
crates/ml/src/dqn/distributional.rs → crates/ml-dqn/src/distributional.rs
View File

@@ -14,7 +14,7 @@
use candle_core::{Device, Result as CandleResult, Tensor};
use serde::{Deserialize, Serialize};
use crate::MLError;
use ml_core::MLError;
/// Distributional type enum (C51 vs QR-DQN)
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]

10
crates/ml/src/dqn/distributional_dueling.rs → crates/ml-dqn/src/distributional_dueling.rs
View File

@@ -43,10 +43,10 @@ use candle_core::{Device, Tensor};
use candle_nn::{Linear, Module, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};
use crate::dqn::mixed_precision::training_dtype;
use crate::dqn::rmsnorm::RMSNorm;
use crate::dqn::xavier_init::linear_xavier;
use crate::MLError;
use crate::mixed_precision::training_dtype;
use crate::rmsnorm::RMSNorm;
use crate::xavier_init::linear_xavier;
use ml_core::MLError;
/// Configuration for Distributional Dueling Q-Network
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -251,7 +251,7 @@ impl DistributionalDuelingQNetwork {
/// - A(s,a,z_i): Advantage distribution per action (probability of atom z_i)
/// - mean(A(s,·,z_i)): Mean advantage across actions (ensures identifiability)
pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
let state = crate::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let batch_size = state
.dim(0)

56
crates/ml/src/dqn/dqn.rs → crates/ml-dqn/src/dqn.rs
View File

@@ -11,10 +11,10 @@
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use crate::dqn::mixed_precision::training_dtype;
use crate::dqn::target_update::{convergence_half_life, hard_update, polyak_update}; // WAVE 16 (Agent 36)
use crate::dqn::xavier_init::linear_xavier; // Xavier initialization with VarMap registration
use crate::Adam;
use crate::mixed_precision::training_dtype;
use crate::target_update::{convergence_half_life, hard_update, polyak_update}; // WAVE 16 (Agent 36)
use crate::xavier_init::linear_xavier; // Xavier initialization with VarMap registration
use ml_core::optimizers::Adam;
use candle_core::backprop::GradStore;
use candle_core::IndexOp;
use candle_core::{DType, Device, Tensor, Var};
@@ -28,7 +28,7 @@ use serde::{Deserialize, Serialize};
use tracing::debug;
use super::{Experience, ExposureLevel, FactoredAction, OrderRouter};
use crate::MLError;
use ml_core::MLError;
/// Configuration for the `DQN`
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -338,19 +338,19 @@ impl DQNConfig {
/// Reads architecture-critical fields (state_dim, num_actions, hidden_dims,
/// Rainbow flags) from the file's metadata. Non-architecture fields (LR,
/// epsilon, buffer capacity) use defaults since they don't affect inference.
pub fn from_safetensors_file(path: &std::path::Path) -> Result<Self, crate::MLError> {
pub fn from_safetensors_file(path: &std::path::Path) -> Result<Self, ml_core::MLError> {
let raw_bytes = std::fs::read(path).map_err(|e| {
crate::MLError::CheckpointError(format!("Failed to read safetensors: {}", e))
ml_core::MLError::CheckpointError(format!("Failed to read safetensors: {}", e))
})?;
let (_header_size, st_metadata) =
safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| {
crate::MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
ml_core::MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
})?;
let meta = match st_metadata.metadata() {
Some(m) => m,
None => {
return Err(crate::MLError::CheckpointError(
return Err(ml_core::MLError::CheckpointError(
"Safetensors checkpoint has no metadata — re-train to embed config".to_owned(),
));
}
@@ -368,10 +368,10 @@ impl DQNConfig {
/// Returns `Err` if required fields are missing or malformed.
pub fn from_checkpoint_metadata(
metadata: &std::collections::HashMap<String, String>,
) -> Result<Self, crate::MLError> {
let get = |key: &str| -> Result<&str, crate::MLError> {
) -> Result<Self, ml_core::MLError> {
let get = |key: &str| -> Result<&str, ml_core::MLError> {
metadata.get(key).map(|s| s.as_str()).ok_or_else(|| {
crate::MLError::CheckpointError(format!(
ml_core::MLError::CheckpointError(format!(
"Checkpoint metadata missing '{}' — re-train to embed config", key
))
})
@@ -379,10 +379,10 @@ impl DQNConfig {
let state_dim: usize = get("dqn.state_dim")?
.parse()
.map_err(|e| crate::MLError::CheckpointError(format!("Bad dqn.state_dim: {}", e)))?;
.map_err(|e| ml_core::MLError::CheckpointError(format!("Bad dqn.state_dim: {}", e)))?;
let num_actions: usize = get("dqn.num_actions")?
.parse()
.map_err(|e| crate::MLError::CheckpointError(format!("Bad dqn.num_actions: {}", e)))?;
.map_err(|e| ml_core::MLError::CheckpointError(format!("Bad dqn.num_actions: {}", e)))?;
// Parse hidden_dims from debug format "[256, 256]"
let hidden_dims_str = get("dqn.hidden_dims")?;
@@ -392,13 +392,13 @@ impl DQNConfig {
.filter(|s| !s.trim().is_empty())
.map(|s| s.trim().parse::<usize>())
.collect::<Result<Vec<_>, _>>()
.map_err(|e| crate::MLError::CheckpointError(format!("Bad dqn.hidden_dims: {}", e)))?;
.map_err(|e| ml_core::MLError::CheckpointError(format!("Bad dqn.hidden_dims: {}", e)))?;
let parse_bool = |key: &str| -> Result<bool, crate::MLError> {
get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
let parse_bool = |key: &str| -> Result<bool, ml_core::MLError> {
get(key)?.parse().map_err(|e| ml_core::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
};
let parse_usize = |key: &str| -> Result<usize, crate::MLError> {
get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
let parse_usize = |key: &str| -> Result<usize, ml_core::MLError> {
get(key)?.parse().map_err(|e| ml_core::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
};
let use_dueling = parse_bool("dqn.use_dueling")?;
@@ -438,11 +438,11 @@ impl DQNConfig {
pub fn validate_checkpoint_metadata(
&self,
metadata: &Option<std::collections::HashMap<String, String>>,
) -> Result<(), crate::MLError> {
) -> Result<(), ml_core::MLError> {
let meta = match metadata {
Some(m) => m,
None => {
return Err(crate::MLError::CheckpointError(
return Err(ml_core::MLError::CheckpointError(
"Checkpoint has no architecture metadata - cannot validate. \
Re-train to generate a checkpoint with embedded config hash."
.to_owned(),
@@ -453,7 +453,7 @@ impl DQNConfig {
let saved_hash = match meta.get("dqn.arch_hash") {
Some(h) => h,
None => {
return Err(crate::MLError::CheckpointError(
return Err(ml_core::MLError::CheckpointError(
"Checkpoint metadata missing dqn.arch_hash field - cannot validate. \
Re-train to generate a checkpoint with embedded config hash."
.to_owned(),
@@ -466,7 +466,7 @@ impl DQNConfig {
let saved_dims = meta.get("dqn.state_dim").map(|s| s.as_str()).unwrap_or("?");
let saved_actions = meta.get("dqn.num_actions").map(|s| s.as_str()).unwrap_or("?");
let saved_hidden = meta.get("dqn.hidden_dims").map(|s| s.as_str()).unwrap_or("?");
return Err(crate::MLError::CheckpointError(format!(
return Err(ml_core::MLError::CheckpointError(format!(
"Architecture mismatch: checkpoint was saved with state_dim={}, num_actions={}, \
hidden_dims={} (hash={}) but current config has state_dim={}, num_actions={}, \
hidden_dims={:?} (hash={})",
@@ -929,7 +929,7 @@ impl Sequential {
/// Forward pass through network
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
let mut x = crate::dqn::mixed_precision::ensure_training_dtype(input)
let mut x = crate::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
if self.use_noisy_nets {
@@ -1039,7 +1039,7 @@ impl Sequential {
#[allow(missing_debug_implementations)]
pub struct DQN {
/// `DQN` configuration
pub(crate) config: DQNConfig,
pub config: DQNConfig,
/// Main Q-network
q_network: Sequential,
/// Target Q-network for stable training
@@ -1321,7 +1321,7 @@ impl DQN {
/// 3. Standard Q-network (fallback)
pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
// Auto-convert input to correct device and dtype
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
let state = crate::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let state = state
.to_device(&self.device)
@@ -2028,7 +2028,7 @@ impl DQN {
/// the intermediate representation needed by IQN.
fn get_state_embedding(&self, states: &Tensor) -> Result<Tensor, MLError> {
let states = states.to_device(&self.device)?;
let states = crate::dqn::mixed_precision::ensure_training_dtype(&states)
let states = crate::mixed_precision::ensure_training_dtype(&states)
.map_err(|e| MLError::ModelError(e.to_string()))?;
if self.q_network.use_noisy_nets {
@@ -3739,7 +3739,7 @@ impl DQN {
#[cfg(test)]
mod tests {
use super::*;
use crate::dqn::Experience;
use crate::Experience;
#[test]

8
crates/ml/src/dqn/dueling.rs → crates/ml-dqn/src/dueling.rs
View File

@@ -34,9 +34,9 @@ use candle_core::{Device, ModuleT, Tensor};
use candle_nn::{Dropout, Linear, Module, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};
use crate::dqn::mixed_precision::training_dtype;
use crate::dqn::xavier_init::linear_xavier;
use crate::MLError;
use crate::mixed_precision::training_dtype;
use crate::xavier_init::linear_xavier;
use ml_core::MLError;
/// Configuration for Dueling Q-Network
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -225,7 +225,7 @@ impl DuelingQNetwork {
/// When `train=true`, dropout is applied after each shared layer activation
/// to regularize and prevent overfitting.
pub fn forward_t(&self, state: &Tensor, train: bool) -> Result<Tensor, MLError> {
let mut h = crate::dqn::mixed_precision::ensure_training_dtype(state)
let mut h = crate::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
for (i, layer) in self.shared_layers.iter().enumerate() {
h = layer.forward(&h).map_err(|e| {

4
crates/ml/src/dqn/ensemble_network.rs → crates/ml-dqn/src/ensemble_network.rs
View File

@@ -30,8 +30,8 @@
use candle_core::{Device, Tensor};
use serde::{Deserialize, Serialize};
use crate::dqn::network::{QNetwork, QNetworkConfig};
use crate::MLError;
use crate::network::{QNetwork, QNetworkConfig};
use ml_core::MLError;
/// Configuration for ensemble Q-network
#[derive(Debug, Clone, Serialize, Deserialize)]

2
crates/ml/src/dqn/entropy_regularization.rs → crates/ml-dqn/src/entropy_regularization.rs
View File

@@ -20,7 +20,7 @@
use candle_core::{DType, Tensor};
use rand::{thread_rng, Rng};
use crate::MLError;
use ml_core::MLError;
/// Entropy regularizer for preventing policy collapse
///

0
crates/ml/src/dqn/experience.rs → crates/ml-dqn/src/experience.rs
View File

0
crates/ml/src/dqn/experience_dataset.rs → crates/ml-dqn/src/experience_dataset.rs
View File

0
crates/ml/src/dqn/gae.rs → crates/ml-dqn/src/gae.rs
View File

6
crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs → crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -5,9 +5,9 @@
//! (rank-based) — all on GPU with zero CPU involvement.
use candle_core::{Device, DType, Tensor};
use crate::MLError;
use crate::dqn::mixed_precision::training_dtype;
use crate::dqn::replay_buffer_type::GpuBatch;
use ml_core::MLError;
use crate::mixed_precision::training_dtype;
use crate::replay_buffer_type::GpuBatch;
/// Configuration for GPU replay buffer
#[derive(Debug, Clone)]

2
crates/ml/src/dqn/hindsight_replay.rs → crates/ml-dqn/src/hindsight_replay.rs
View File

@@ -16,7 +16,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use super::experience::{Experience, ExperienceBatch};
use super::prioritized_replay::{PrioritizedReplayBuffer, PrioritizedReplayConfig};
use crate::MLError;
use ml_core::MLError;
/// Configuration for Hindsight Experience Replay
#[derive(Debug, Clone, Serialize, Deserialize)]

0
crates/ml/src/dqn/iql.rs → crates/ml-dqn/src/iql.rs
View File

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

@@ -1 +1,137 @@
// Modules will be moved here from ml crate
//! Deep Q-Learning Network implementation for trading
//!
//! This crate includes both the original DQN implementation and the enhanced Rainbow DQN
//! with all 6 components: Double Q-learning, Dueling Networks, Prioritized Experience Replay,
//! Multi-step Learning, Distributional RL (C51), and Noisy Networks.
// Re-export shared modules from ml-core for convenience
pub use ml_core::action_space;
pub use ml_core::order_router;
pub use ml_core::mixed_precision;
pub use ml_core::xavier_init;
pub use ml_core::portfolio_tracker;
pub use ml_core::trading_action;
pub use ml_core::trading_action::TradingAction;
pub use ml_core::cuda_compat;
// Original DQN components
pub mod count_bonus;
pub mod agent;
pub mod attention;
pub mod circuit_breaker;
pub mod curiosity;
pub mod dqn;
pub mod experience;
pub mod gae;
pub mod hindsight_replay;
pub mod logging;
pub mod network;
pub mod nstep_buffer;
pub mod regime_conditional;
pub mod replay_buffer;
pub mod residual;
pub mod reward;
pub mod softmax;
pub mod logit_clipping;
pub mod target_update;
pub mod trade_executor;
// Wave 16 Portfolio Features
pub mod entropy_regularization;
pub mod multi_asset;
// Rainbow DQN components
pub mod distributional;
pub mod distributional_dueling;
pub mod dueling;
pub mod multi_step;
pub mod noisy_layers;
pub mod quantile_regression;
pub mod noisy_sigma_scheduler;
pub mod rainbow_agent;
pub mod rainbow_config;
pub mod rainbow_integration;
pub mod rainbow_network;
// Replay buffer types
pub mod prioritized_replay;
pub mod prioritized_replay_staleness;
pub mod replay_buffer_type;
pub mod experience_dataset;
pub mod iql;
pub mod demo_dqn;
pub mod noisy_exploration;
pub mod self_supervised_pretraining;
// Performance validation
pub mod performance_tests;
pub mod performance_validation;
// Wave 26 P2.3: Ensemble Q-network for uncertainty estimation
pub mod ensemble_network;
// Wave 26 P2.4: RMSNorm
pub mod rmsnorm;
// GPU-resident replay buffer (DQN-specific)
#[cfg(feature = "cuda")]
pub mod gpu_replay_buffer;
// Re-export core DQN types for public usage
pub use action_space::{ExposureLevel, FactoredAction, OrderType, Urgency};
pub use order_router::OrderRouter;
pub use agent::{AgentMetrics, DQNAgent, TradingState};
pub use dqn::{GradientResult, DQN, DQNConfig};
pub use experience::{Experience, ExperienceBatch};
pub use nstep_buffer::NStepBuffer;
pub use portfolio_tracker::PortfolioTracker;
pub use replay_buffer::{ReplayBuffer, ReplayBufferConfig, ReplayBufferStats};
pub use trade_executor::{
ExecutionCostConfig, ExecutionResult, RejectionReason, RiskControlConfig, TradeExecutor,
};
// Re-export network components
pub use network::{QNetwork, QNetworkConfig};
// Re-export reward components
pub use reward::{MarketData, RewardConfig, RewardFunction, RiskMetrics};
// Re-export Rainbow DQN components
pub use distributional::{CategoricalDistribution, DistributionalConfig, DistributionalType};
pub use distributional_dueling::{DistributionalDuelingConfig, DistributionalDuelingQNetwork};
pub use dueling::{DuelingConfig, DuelingQNetwork};
pub use multi_step::MultiStepConfig;
pub use quantile_regression::{QuantileConfig, QuantileNetwork, quantile_huber_loss};
pub use rainbow_agent::RainbowAgent;
pub use rainbow_config::{RainbowAgentConfig, RainbowAgentMetrics, RainbowDQNConfig};
pub use rainbow_network::{RainbowNetwork, RainbowNetworkConfig};
// Re-export prioritized replay components
pub use prioritized_replay::{PrioritizedReplayBuffer, PrioritizedReplayConfig};
pub use replay_buffer_type::{BatchSample, ReplayBufferType};
// Re-export regime-conditional DQN components
pub use regime_conditional::{RegimeConditionalDQN, RegimeMetrics, RegimeType};
// Re-export logit clipping utilities
pub use logit_clipping::{clip_logits, clip_logits_default, softmax_with_clipping, DEFAULT_CLIP_MAX};
// Re-export GAE components
pub use gae::{GAECalculator, GAEConfig};
// Re-export Hindsight Experience Replay components
pub use hindsight_replay::{
HindsightReplayBuffer, HindsightReplayConfig, HindsightReplayStats, HindsightStrategy,
};
// Re-export ensemble network components
pub use ensemble_network::{EnsembleConfig, EnsembleQNetwork};
// Re-export RMSNorm components
pub use rmsnorm::{LayerNorm, NormType, RMSNorm};
// Re-export offline RL components
pub use experience_dataset::ExperienceDataset;
pub use iql::IqlConfig;

0
crates/ml/src/dqn/logging.rs → crates/ml-dqn/src/logging.rs
View File

2
crates/ml/src/dqn/logit_clipping.rs → crates/ml-dqn/src/logit_clipping.rs
View File

@@ -40,7 +40,7 @@
//! ```
use candle_core::Tensor;
use crate::MLError;
use ml_core::MLError;
/// Default maximum absolute value for logit clipping
pub const DEFAULT_CLIP_MAX: f32 = 10.0;

2
crates/ml/src/dqn/multi_asset.rs → crates/ml-dqn/src/multi_asset.rs
View File

@@ -497,7 +497,7 @@ impl MultiAssetPortfolioTracker {
#[cfg(test)]
mod tests {
use super::*;
use crate::dqn::action_space::{ExposureLevel, OrderType, Urgency};
use crate::action_space::{ExposureLevel, OrderType, Urgency};
#[test]
fn test_multi_asset_initialization() {

2
crates/ml/src/dqn/multi_step.rs → crates/ml-dqn/src/multi_step.rs
View File

@@ -11,7 +11,7 @@ use std::collections::VecDeque;
use candle_core::{Device, Tensor};
use serde::{Deserialize, Serialize};
use crate::MLError;
use ml_core::MLError;
/// Configuration for multi-step learning
#[derive(Debug, Clone, Serialize, Deserialize)]

14
crates/ml/src/dqn/network.rs → crates/ml-dqn/src/network.rs
View File

@@ -3,13 +3,13 @@
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use candle_core::{Device, Result as CandleResult, Tensor};
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
use candle_nn::Module;
use candle_nn::{ops::leaky_relu, Dropout, Linear, VarBuilder, VarMap};
use rand::prelude::*; // Replace common::rng with standard rand
use crate::dqn::xavier_init::linear_xavier; // Xavier initialization
use crate::MLError;
use crate::xavier_init::linear_xavier; // Xavier initialization
use ml_core::MLError;
/// Adaptive dropout scheduler that decreases dropout rate over training
/// (Wave 26 P1.6)
@@ -93,7 +93,7 @@ pub struct QNetworkConfig {
pub use_residual: bool,
/// Mixed precision configuration for tensor core utilization.
/// None = FP32 only. Some(config) = BF16 or FP16 forward pass.
pub mixed_precision: Option<crate::dqn::mixed_precision::MixedPrecisionConfig>,
pub mixed_precision: Option<crate::mixed_precision::MixedPrecisionConfig>,
}
impl Default for QNetworkConfig {
@@ -194,9 +194,9 @@ impl NetworkLayers {
fn forward_mixed(
&self,
xs: &Tensor,
mixed_precision: &Option<crate::dqn::mixed_precision::MixedPrecisionConfig>,
mixed_precision: &Option<crate::mixed_precision::MixedPrecisionConfig>,
) -> CandleResult<Tensor> {
let xs = crate::dqn::mixed_precision::ensure_training_dtype(xs)?;
let xs = crate::mixed_precision::ensure_training_dtype(xs)?;
let (x_input, _use_amp) = match mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
@@ -224,7 +224,7 @@ impl NetworkLayers {
impl Module for NetworkLayers {
fn forward(&self, xs: &Tensor) -> CandleResult<Tensor> {
let mut x = crate::dqn::mixed_precision::ensure_training_dtype(xs)?;
let mut x = crate::mixed_precision::ensure_training_dtype(xs)?;
// Forward through hidden layers with LeakyReLU activation and dropout
// LeakyReLU prevents dead neurons (0.01 gradient for negative inputs vs 0 for ReLU)

2
crates/ml/src/dqn/noisy_exploration.rs → crates/ml-dqn/src/noisy_exploration.rs
View File

@@ -5,7 +5,7 @@
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::MLError;
use ml_core::MLError;
/// Metrics for noisy exploration
#[derive(Debug, Clone, Default)]

6
crates/ml/src/dqn/noisy_layers.rs → crates/ml-dqn/src/noisy_layers.rs
View File

@@ -13,8 +13,8 @@
use candle_core::{Device, Result as CandleResult, Tensor, Var};
use candle_nn::{Module, VarBuilder};
use crate::dqn::mixed_precision::training_dtype;
use crate::MLError;
use crate::mixed_precision::training_dtype;
use ml_core::MLError;
/// Noisy linear layer with factorized Gaussian noise (Rainbow DQN standard)
///
@@ -316,7 +316,7 @@ impl Default for NoisyNetworkConfig {
mod tests {
use super::*;
use candle_nn::{VarBuilder, VarMap};
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
#[test]
fn test_noisy_linear_creation() -> Result<(), MLError> {

0
crates/ml/src/dqn/noisy_sigma_scheduler.rs → crates/ml-dqn/src/noisy_sigma_scheduler.rs
View File

2
crates/ml/src/dqn/nstep_buffer.rs → crates/ml-dqn/src/nstep_buffer.rs
View File

@@ -39,7 +39,7 @@
//! ```
use std::collections::VecDeque;
use crate::dqn::Experience;
use crate::Experience;
/// N-step experience buffer for multi-step temporal difference learning
///

2
crates/ml/src/dqn/performance_tests.rs → crates/ml-dqn/src/performance_tests.rs
View File

@@ -12,7 +12,7 @@ use candle_nn::VarMap;
// use criterion::{criterion_group, criterion_main, Criterion, black_box};
use super::*;
use crate::MLError;
use ml_core::MLError;
/// Performance test configuration
#[derive(Debug, Clone)]

0
crates/ml/src/dqn/performance_validation.rs → crates/ml-dqn/src/performance_validation.rs
View File

6
crates/ml/src/dqn/prioritized_replay.rs → crates/ml-dqn/src/prioritized_replay.rs
View File

@@ -18,8 +18,8 @@ use rand::rngs::StdRng;
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use crate::dqn::experience::Experience;
use crate::MLError;
use crate::experience::Experience;
use ml_core::MLError;
/// Segment tree for efficient priority sampling
#[derive(Debug)]
@@ -758,7 +758,7 @@ impl PrioritizedReplayBuffer {
#[cfg(test)]
mod tests {
use super::*;
use crate::dqn::experience::Experience;
use crate::experience::Experience;
fn create_test_experience() -> Experience {
Experience::new(vec![1.0, 2.0, 3.0], 0, 1.0, vec![1.1, 2.1, 3.1], false)

0
crates/ml/src/dqn/prioritized_replay_staleness.rs → crates/ml-dqn/src/prioritized_replay_staleness.rs
View File

8
crates/ml/src/dqn/quantile_regression.rs → crates/ml-dqn/src/quantile_regression.rs
View File

@@ -20,8 +20,8 @@ use candle_nn::{Linear, Module, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};
use std::f32::consts::PI;
use crate::dqn::mixed_precision::training_dtype;
use crate::MLError;
use crate::mixed_precision::training_dtype;
use ml_core::MLError;
/// Configuration for Quantile Regression DQN
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -144,9 +144,9 @@ impl QuantileNetwork {
/// # Returns
/// Quantile values [batch, num_actions, num_quantiles]
pub fn forward(&self, state_embed: &Tensor, taus: &Tensor) -> Result<Tensor, MLError> {
let state_embed = crate::dqn::mixed_precision::ensure_training_dtype(state_embed)
let state_embed = crate::mixed_precision::ensure_training_dtype(state_embed)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let taus = crate::dqn::mixed_precision::ensure_training_dtype(taus)
let taus = crate::mixed_precision::ensure_training_dtype(taus)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let batch_size = state_embed.dim(0)?;
let num_quantiles = taus.dim(1)?;

6
crates/ml/src/dqn/rainbow_agent.rs → crates/ml-dqn/src/rainbow_agent.rs
View File

@@ -7,10 +7,10 @@
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, RwLock};
use crate::Adam;
use ml_core::optimizers::Adam;
use candle_core::{Device, Tensor};
use candle_nn::{VarBuilder, VarMap};
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
use candle_optimisers::adam::ParamsAdam;
use tracing::{debug, info};
@@ -18,7 +18,7 @@ use super::multi_step::{MultiStepCalculator, MultiStepTransition};
use super::rainbow_config::{RainbowAgentConfig, RainbowAgentMetrics, TrainingResult};
use super::rainbow_network::RainbowNetwork;
use super::{Experience, ReplayBuffer, ReplayBufferConfig};
use crate::MLError;
use ml_core::MLError;
/// Rainbow `DQN` Agent with all 6 components
pub struct RainbowAgent {

0
crates/ml/src/dqn/rainbow_config.rs → crates/ml-dqn/src/rainbow_config.rs
View File

4
crates/ml/src/dqn/rainbow_integration.rs → crates/ml-dqn/src/rainbow_integration.rs
View File

@@ -11,7 +11,7 @@
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
use crate::MLError;
use ml_core::MLError;
// Use RainbowDQNConfig from rainbow_config module
use super::rainbow_config::RainbowDQNConfig;
@@ -55,7 +55,7 @@ impl RainbowMetrics {
#[cfg(test)]
mod tests {
use super::*;
use crate::dqn::rainbow_network::RainbowNetworkConfig;
use crate::rainbow_network::RainbowNetworkConfig;
#[test]
fn test_rainbow_dqn_config_creation() {

4
crates/ml/src/dqn/rainbow_network.rs → crates/ml-dqn/src/rainbow_network.rs
View File

@@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize};
use super::distributional::{CategoricalDistribution, DistributionalConfig};
use super::noisy_layers::NoisyLinear;
use crate::MLError;
use ml_core::MLError;
/// Activation function types
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
@@ -428,7 +428,7 @@ mod tests {
use anyhow::Result;
use candle_core::Device;
use candle_nn::{VarBuilder, VarMap};
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
#[test]
fn test_rainbow_network_creation() -> Result<(), MLError> {

12
crates/ml/src/dqn/regime_conditional.rs → crates/ml-dqn/src/regime_conditional.rs
View File

@@ -51,7 +51,7 @@ use serde::{Deserialize, Serialize};
use tracing::{debug, info};
use super::{Experience, FactoredAction, GradientResult, DQN, DQNConfig};
use crate::MLError;
use ml_core::MLError;
/// Market regime types based on structural characteristics
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -906,7 +906,7 @@ impl RegimeConditionalDQN {
/// Get trending head's memory buffer (for trainer access)
/// Note: Each head now has its own buffer, not a shared one
pub fn get_trending_head_memory(&self) -> &crate::dqn::replay_buffer_type::ReplayBufferType {
pub fn get_trending_head_memory(&self) -> &crate::replay_buffer_type::ReplayBufferType {
&self.trending_head.memory
}
@@ -928,7 +928,7 @@ impl RegimeConditionalDQN {
///
/// * `Ok(())` - All heads reinitialized successfully
/// * `Err(MLError)` - Failed to reinitialize one or more heads
pub fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), crate::MLError> {
pub fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), ml_core::MLError> {
// Update all three regime heads
self.trending_head.reinit_categorical_distribution(v_min, v_max)?;
self.ranging_head.reinit_categorical_distribution(v_min, v_max)?;
@@ -946,7 +946,7 @@ impl RegimeConditionalDQN {
}
/// BUG #38 FIX: Clear all regime-specific replay buffers
pub fn clear_replay_buffer(&mut self) -> Result<(), crate::MLError> {
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()?;
@@ -954,7 +954,7 @@ impl RegimeConditionalDQN {
}
/// BUG #38 FIX: Reset all regime-specific target networks
pub fn reset_target_network(&mut self) -> Result<(), crate::MLError> {
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()?;
@@ -962,7 +962,7 @@ impl RegimeConditionalDQN {
}
/// Update learning rate for all regime heads
pub fn update_learning_rate(&mut self, decay_factor: f64) -> Result<(), crate::MLError> {
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)?;

2
crates/ml/src/dqn/replay_buffer.rs → crates/ml-dqn/src/replay_buffer.rs
View File

@@ -7,7 +7,7 @@ use parking_lot::RwLock;
use rand::prelude::*;
use super::{Experience, ExperienceBatch};
use crate::MLError;
use ml_core::MLError;
/// Configuration for replay buffer
#[derive(Debug, Clone)]

16
crates/ml/src/dqn/replay_buffer_type.rs → crates/ml-dqn/src/replay_buffer_type.rs
View File

@@ -6,10 +6,10 @@
use std::sync::Arc;
use parking_lot::Mutex;
use crate::dqn::dqn::ExperienceReplayBuffer;
use crate::dqn::prioritized_replay::PrioritizedReplayBuffer;
use crate::dqn::Experience;
use crate::MLError;
use crate::dqn::ExperienceReplayBuffer;
use crate::prioritized_replay::PrioritizedReplayBuffer;
use crate::experience::Experience;
use ml_core::MLError;
/// Batch sample with importance sampling weights
#[derive(Debug, Clone)]
@@ -62,7 +62,7 @@ pub enum ReplayBufferType {
Prioritized(Arc<PrioritizedReplayBuffer>),
/// GPU-resident prioritized replay — all sampling on GPU (zero CPU round-trips)
#[cfg(feature = "cuda")]
GpuPrioritized(Arc<Mutex<crate::cuda_pipeline::gpu_replay_buffer::GpuReplayBuffer>>),
GpuPrioritized(Arc<Mutex<crate::gpu_replay_buffer::GpuReplayBuffer>>),
}
impl std::fmt::Debug for ReplayBufferType {
@@ -90,7 +90,7 @@ impl ReplayBufferType {
beta_max: f64,
beta_annealing_steps: usize,
) -> Result<Self, MLError> {
use crate::dqn::prioritized_replay::{PrioritizedReplayConfig, PrioritizationStrategy};
use crate::prioritized_replay::{PrioritizedReplayConfig, PrioritizationStrategy};
let config = PrioritizedReplayConfig {
capacity,
@@ -121,7 +121,7 @@ impl ReplayBufferType {
beta_annealing_steps: usize,
device: &candle_core::Device,
) -> Result<Self, MLError> {
use crate::cuda_pipeline::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
use crate::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
let config = GpuReplayBufferConfig {
capacity,
@@ -501,7 +501,7 @@ impl ReplayBufferType {
/// Get a locked reference to the inner GPU replay buffer (GpuPrioritized only).
#[cfg(feature = "cuda")]
pub fn as_gpu_buffer(&self) -> Option<parking_lot::MutexGuard<'_, crate::cuda_pipeline::gpu_replay_buffer::GpuReplayBuffer>> {
pub fn as_gpu_buffer(&self) -> Option<parking_lot::MutexGuard<'_, crate::gpu_replay_buffer::GpuReplayBuffer>> {
match self {
Self::GpuPrioritized(buffer) => Some(buffer.lock()),
Self::Uniform(_) | Self::Prioritized(_) => None,

8
crates/ml/src/dqn/residual.rs → crates/ml-dqn/src/residual.rs
View File

@@ -12,9 +12,9 @@
use candle_core::Tensor;
use candle_nn::{Linear, Module, VarBuilder};
use crate::cuda_compat::layer_norm_with_fallback;
use crate::dqn::xavier_init::linear_xavier;
use crate::MLError;
use ml_core::cuda_compat::layer_norm_with_fallback;
use crate::xavier_init::linear_xavier;
use ml_core::MLError;
/// Configuration for residual blocks
#[derive(Debug, Clone)]
@@ -162,7 +162,7 @@ mod tests {
use super::*;
use candle_core::{DType, Device};
use candle_nn::VarMap;
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
#[test]
fn test_residual_config_default() {

8
crates/ml/src/dqn/reward.rs → crates/ml-dqn/src/reward.rs
View File

@@ -9,7 +9,7 @@ use rust_decimal::Decimal;
use super::action_space::FactoredAction;
use super::agent::{TradingAction, TradingState};
use crate::MLError;
use ml_core::MLError;
use super::circuit_breaker::CircuitBreakerConfig;
/// Online reward normalization using Exponential Moving Average (EMA)
@@ -1233,7 +1233,7 @@ mod tests {
#[test]
fn test_transaction_cost_fix_market_orders() -> anyhow::Result<()> {
use crate::dqn::action_space::{ExposureLevel, OrderType, Urgency};
use crate::action_space::{ExposureLevel, OrderType, Urgency};
// Create reward function with default config
let config = RewardConfig::default();
@@ -1318,7 +1318,7 @@ mod tests {
#[test]
fn test_transaction_cost_realistic_scenario() -> anyhow::Result<()> {
use crate::dqn::action_space::{ExposureLevel, OrderType, Urgency};
use crate::action_space::{ExposureLevel, OrderType, Urgency};
// Simulate realistic trading scenario from audit report
// Agent makes small profit ($0.10) but actual cost is $4.50
@@ -1397,7 +1397,7 @@ mod tests {
#[test]
fn test_transaction_cost_zero_for_hold() -> anyhow::Result<()> {
use crate::dqn::action_space::{ExposureLevel, OrderType, Urgency};
use crate::action_space::{ExposureLevel, OrderType, Urgency};
let config = RewardConfig::default();
let reward_fn = RewardFunction::new(config)?;

4
crates/ml/src/dqn/rmsnorm.rs → crates/ml-dqn/src/rmsnorm.rs
View File

@@ -10,7 +10,7 @@
use candle_core::{Tensor, D};
use candle_nn::VarBuilder;
use crate::MLError;
use ml_core::MLError;
/// Normalization type configuration
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -231,7 +231,7 @@ mod tests {
use candle_core::Device;
use candle_nn::VarMap;
use std::time::Instant;
use crate::dqn::mixed_precision::training_dtype;
use crate::mixed_precision::training_dtype;
#[test]
fn test_rmsnorm_creation() -> anyhow::Result<()> {

0
crates/ml/src/dqn/self_supervised_pretraining.rs → crates/ml-dqn/src/self_supervised_pretraining.rs
View File

2
crates/ml/src/dqn/softmax.rs → crates/ml-dqn/src/softmax.rs
View File

@@ -10,7 +10,7 @@
//! - Batched and single-state support
use candle_core::Tensor;
use crate::MLError;
use ml_core::MLError;
use rand::{thread_rng, Rng};
#[cfg(test)]

0
crates/ml/src/dqn/target_update.rs → crates/ml-dqn/src/target_update.rs
View File

2
crates/ml/src/dqn/trade_executor.rs → crates/ml-dqn/src/trade_executor.rs
View File

@@ -1,4 +1,4 @@
use crate::dqn::portfolio_tracker::{PortfolioTracker, TradeAction};
use crate::portfolio_tracker::{PortfolioTracker, TradeAction};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};

1
crates/ml/Cargo.toml
View File

@@ -69,6 +69,7 @@ colored = "2.1" # Terminal color output for evaluation reports
# Internal workspace crates
ml-core.workspace = true
ml-dqn.workspace = true
config.workspace = true
common = { workspace = true, features = ["questdb"] }
risk = { path = "../risk" }

78
crates/ml/src/checkpoint/model_implementations.rs
View File

@@ -212,82 +212,8 @@ impl Checkpointable for DQNAgent {
info
}
}
impl DQNAgent {
/// Extract network weights from the model
fn extract_network_weights(&self, network_name: &str) -> Option<Vec<u8>> {
// In a real implementation, this would extract actual neural network weights
// For now, return some production serialized weights
match network_name {
"q_network" => {
// Simulate serialized Q-network weights
let weights = vec![0.1_f32, 0.2_f32, 0.3_f32, 0.4_f32]; // Production weights
let mut buffer = Vec::new();
for weight in weights {
buffer.extend_from_slice(&weight.to_le_bytes());
}
Some(buffer)
},
_ => None,
}
}
/// Get current replay buffer size
fn get_replay_buffer_size(&self) -> usize {
// In a real implementation, this would query the actual replay buffer
// For now, return a reasonable default based on configuration
let filled_ratio = 0.7; // Assume 70% filled
(self.config.replay_buffer_capacity as f64 * filled_ratio) as usize
}
/// Get average reward from recent episodes
fn get_average_reward(&self) -> f64 {
// In a real implementation, this would compute from recent episode history
// For now, return a production based on training progress
let episodes = self.metrics.total_episodes;
if episodes > 100 {
// Simulate learning progress - higher rewards as training progresses
10.0 + (episodes as f64 / 100.0)
} else {
// Early training - lower rewards
-5.0 + (episodes as f64 / 20.0)
}
}
/// Restore network weights from serialized data
fn restore_network_weights(
&mut self,
network_name: &str,
weights_data: &[u8],
) -> Result<(), String> {
// In a real implementation, this would deserialize and load weights into the neural network
match network_name {
"q_network" | "target_network" => {
if weights_data.len() % 4 != 0 {
return Err("Invalid weight data length".to_owned());
}
let weight_count = weights_data.len() / 4;
let mut weights = Vec::with_capacity(weight_count);
for i in 0..weight_count {
let start_idx = i * 4;
let weight_bytes = &weights_data[start_idx..start_idx + 4];
let weight = f32::from_le_bytes([
weight_bytes[0],
weight_bytes[1],
weight_bytes[2],
weight_bytes[3],
]);
weights.push(weight);
}
info!("Restored {} weights for {}", weights.len(), network_name);
Ok(())
},
_ => Err(format!("Unknown network: {}", network_name)),
}
}
}
// DQNAgent checkpoint helper methods (extract/restore_network_weights, get_replay_buffer_size,
// get_average_reward) are defined in ml-dqn::agent since DQNAgent lives in that crate.
/// Serializable state for `MAMBA` model
#[derive(Debug, Clone, Serialize, Deserialize)]

3
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -20,8 +20,7 @@ pub mod gpu_weights;
pub mod gpu_experience_collector;
#[cfg(feature = "cuda")]
pub mod gpu_ppo_collector;
#[cfg(feature = "cuda")]
pub mod gpu_replay_buffer;
// gpu_replay_buffer moved to ml-dqn crate
/// Maximum bytes allowed for a single GPU upload (2 GB safety limit).
const MAX_UPLOAD_BYTES: usize = 2 * 1024 * 1024 * 1024;

354
crates/ml/src/dqn/data_augmentation.rs
View File

@@ -1,354 +0,0 @@
//! Data augmentation for DQN training to prevent overfitting
//!
//! This module implements noise injection techniques to augment training data,
//! helping the DQN agent generalize better and avoid overfitting to specific
//! market conditions.
use rand::Rng;
use serde::{Deserialize, Serialize};
/// Configuration for noise injection data augmentation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NoiseInjectorConfig {
/// Standard deviation of Gaussian noise (0.01-0.05 typical)
pub noise_std: f32,
/// Probability of applying noise (0.3-0.5 typical)
pub apply_prob: f32,
}
impl Default for NoiseInjectorConfig {
fn default() -> Self {
Self {
noise_std: 0.02,
apply_prob: 0.4,
}
}
}
/// Data augmentation via Gaussian noise injection for DQN training
///
/// Adds controlled Gaussian noise to state features to:
/// - Prevent overfitting to specific market conditions
/// - Improve generalization across different regimes
/// - Regularize the Q-network learning
///
/// # Example
/// ```
/// use ml::dqn::data_augmentation::{NoiseInjector, NoiseInjectorConfig};
/// use rand::thread_rng;
///
/// let config = NoiseInjectorConfig {
/// noise_std: 0.03,
/// apply_prob: 0.5,
/// };
/// let injector = NoiseInjector::new(config);
/// let mut rng = thread_rng();
///
/// let state = vec![1.0, 2.0, 3.0, 4.0];
/// let augmented = injector.augment_state(&state, &mut rng);
/// ```
#[derive(Debug, Clone)]
pub struct NoiseInjector {
config: NoiseInjectorConfig,
}
impl NoiseInjector {
/// Create a new noise injector with the given configuration
pub fn new(config: NoiseInjectorConfig) -> Self {
Self { config }
}
/// Add Gaussian noise to state features with probability `apply_prob`
///
/// # Arguments
/// * `state` - Original state features
/// * `rng` - Random number generator
///
/// # Returns
/// State with Gaussian noise added (if probability threshold met)
pub fn augment_state(&self, state: &[f32], rng: &mut impl Rng) -> Vec<f32> {
if rng.gen::<f32>() < self.config.apply_prob {
// Apply Gaussian noise
state
.iter()
.map(|&x| {
// Generate Gaussian noise: mean=0, std=noise_std
let noise = self.sample_gaussian(rng);
x + noise * self.config.noise_std
})
.collect()
} else {
// No augmentation
state.to_vec()
}
}
/// Sample from standard normal distribution using Box-Muller transform
fn sample_gaussian(&self, rng: &mut impl Rng) -> f32 {
let u1: f32 = rng.gen();
let u2: f32 = rng.gen();
// Box-Muller transform
let z0 = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
z0
}
/// Get the current configuration
pub fn config(&self) -> &NoiseInjectorConfig {
&self.config
}
/// Update the noise standard deviation
pub fn set_noise_std(&mut self, noise_std: f32) {
self.config.noise_std = noise_std;
}
/// Update the application probability
pub fn set_apply_prob(&mut self, apply_prob: f32) {
self.config.apply_prob = apply_prob;
}
}
#[cfg(test)]
mod tests {
use super::*;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
#[test]
fn test_noise_injector_creation() {
let config = NoiseInjectorConfig {
noise_std: 0.03,
apply_prob: 0.5,
};
let injector = NoiseInjector::new(config.clone());
assert_eq!(injector.config().noise_std, 0.03);
assert_eq!(injector.config().apply_prob, 0.5);
}
#[test]
fn test_default_config() {
let config = NoiseInjectorConfig::default();
assert_eq!(config.noise_std, 0.02);
assert_eq!(config.apply_prob, 0.4);
}
#[test]
fn test_augment_state_deterministic_no_noise() {
let config = NoiseInjectorConfig {
noise_std: 0.02,
apply_prob: 0.0, // Never apply noise
};
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
let state = vec![1.0, 2.0, 3.0, 4.0];
let augmented = injector.augment_state(&state, &mut rng);
// With apply_prob=0, state should be unchanged
assert_eq!(augmented, state);
}
#[test]
fn test_augment_state_deterministic_always_noise() {
let config = NoiseInjectorConfig {
noise_std: 0.02,
apply_prob: 1.0, // Always apply noise
};
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
let state = vec![1.0, 2.0, 3.0, 4.0];
let augmented = injector.augment_state(&state, &mut rng);
// With apply_prob=1.0, state should be modified
assert_ne!(augmented, state);
// Check that dimensions are preserved
assert_eq!(augmented.len(), state.len());
}
#[test]
fn test_noise_magnitude() {
let config = NoiseInjectorConfig {
noise_std: 0.01,
apply_prob: 1.0,
};
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
let state = vec![10.0; 100]; // Large state to test statistical properties
let augmented = injector.augment_state(&state, &mut rng);
// Compute differences
let diffs: Vec<f32> = state
.iter()
.zip(augmented.iter())
.map(|(s, a)| a - s)
.collect();
// Check mean is close to 0
let mean: f32 = diffs.iter().sum::<f32>() / diffs.len() as f32;
assert!(mean.abs() < 0.01, "Mean noise should be ~0, got {}", mean);
// Check std is close to noise_std
let variance: f32 = diffs.iter().map(|d| (d - mean).powi(2)).sum::<f32>()
/ diffs.len() as f32;
let std = variance.sqrt();
assert!(
(std - 0.01).abs() < 0.005,
"Std should be ~0.01, got {}",
std
);
}
#[test]
fn test_augmentation_probability() {
let config = NoiseInjectorConfig {
noise_std: 0.02,
apply_prob: 0.5,
};
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
let state = vec![1.0, 2.0, 3.0];
let mut augmented_count = 0;
let num_trials = 1000;
for _ in 0..num_trials {
let augmented = injector.augment_state(&state, &mut rng);
if augmented != state {
augmented_count += 1;
}
}
// With apply_prob=0.5, expect ~50% augmentation rate
let augmentation_rate = augmented_count as f32 / num_trials as f32;
assert!(
(augmentation_rate - 0.5).abs() < 0.05,
"Expected ~50% augmentation, got {}%",
augmentation_rate * 100.0
);
}
#[test]
fn test_set_noise_std() {
let config = NoiseInjectorConfig::default();
let mut injector = NoiseInjector::new(config);
injector.set_noise_std(0.05);
assert_eq!(injector.config().noise_std, 0.05);
}
#[test]
fn test_set_apply_prob() {
let config = NoiseInjectorConfig::default();
let mut injector = NoiseInjector::new(config);
injector.set_apply_prob(0.7);
assert_eq!(injector.config().apply_prob, 0.7);
}
#[test]
fn test_gaussian_sampling() {
let config = NoiseInjectorConfig::default();
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(123);
// Sample many values to test distribution
let samples: Vec<f32> = (0..1000)
.map(|_| injector.sample_gaussian(&mut rng))
.collect();
// Check mean is close to 0
let mean: f32 = samples.iter().sum::<f32>() / samples.len() as f32;
assert!(mean.abs() < 0.1, "Mean should be ~0, got {}", mean);
// Check std is close to 1
let variance: f32 = samples.iter().map(|s| (s - mean).powi(2)).sum::<f32>()
/ samples.len() as f32;
let std = variance.sqrt();
assert!(
(std - 1.0).abs() < 0.1,
"Std should be ~1.0, got {}",
std
);
}
#[test]
fn test_empty_state() {
let config = NoiseInjectorConfig::default();
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
let state: Vec<f32> = vec![];
let augmented = injector.augment_state(&state, &mut rng);
assert_eq!(augmented.len(), 0);
}
#[test]
fn test_single_element_state() {
let config = NoiseInjectorConfig {
noise_std: 0.02,
apply_prob: 1.0,
};
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
let state = vec![5.0];
let augmented = injector.augment_state(&state, &mut rng);
assert_eq!(augmented.len(), 1);
assert_ne!(augmented[0], state[0]);
}
#[test]
fn test_high_dimensional_state() {
let config = NoiseInjectorConfig {
noise_std: 0.01,
apply_prob: 1.0,
};
let injector = NoiseInjector::new(config);
let mut rng = ChaCha8Rng::seed_from_u64(42);
// Test with a typical state size (51 features for this test case)
let state = vec![1.0; 51];
let augmented = injector.augment_state(&state, &mut rng);
assert_eq!(augmented.len(), 51);
// Check that most features are modified
let modified_count = state
.iter()
.zip(augmented.iter())
.filter(|(s, a)| s != a)
.count();
assert!(modified_count > 40, "Most features should be modified");
}
#[test]
fn test_reproducibility_with_seeded_rng() {
let config = NoiseInjectorConfig {
noise_std: 0.03,
apply_prob: 1.0,
};
let injector = NoiseInjector::new(config);
let state = vec![1.0, 2.0, 3.0];
// First run
let mut rng1 = ChaCha8Rng::seed_from_u64(999);
let augmented1 = injector.augment_state(&state, &mut rng1);
// Second run with same seed
let mut rng2 = ChaCha8Rng::seed_from_u64(999);
let augmented2 = injector.augment_state(&state, &mut rng2);
// Should produce identical results
assert_eq!(augmented1, augmented2);
}
}

1062
crates/ml/src/dqn/ensemble.rs
View File

File diff suppressed because it is too large Load Diff

299
crates/ml/src/dqn/ensemble_oracle.rs
View File

@@ -1,299 +0,0 @@
//! Ensemble oracle for multi-model consensus voting
//!
//! Combines predictions from multiple ML models (Transformer, LSTM, PPO) to provide
//! robust reward signals based on majority voting and diversity metrics.
use std::collections::HashMap;
use std::sync::Arc;
use candle_core::Tensor;
use super::TradingAction;
use super::action_space::{FactoredAction, ExposureLevel};
/// Ensemble oracle that combines predictions from multiple ML models
///
/// Uses majority voting consensus with bonuses for:
/// - Agreement: +0.5 when DQN action matches majority, +0.1 when it disagrees
/// - Diversity: +0.3 when all models disagree, +0.1 for moderate disagreement, +0.0 for full consensus
///
/// # Example
///
/// ```ignore
/// use ml::dqn::EnsembleOracle;
/// use ml::dqn::TradingAction;
///
/// let mut oracle = EnsembleOracle::new();
/// oracle.load_models(Some("path/to/transformer"), None, None).unwrap();
///
/// // Calculate ensemble reward (using mock votes for testing)
/// let reward = oracle.calculate_ensemble_reward(
/// &state,
/// TradingAction::Buy,
/// vec![0, 0, 1] // Buy, Buy, Sell
/// );
/// // reward = 0.6 (agreement 0.5 + diversity 0.1)
/// ```
#[derive(Debug)]
pub struct EnsembleOracle {
/// Placeholder for Transformer model (future: Arc<TransformerModel>)
transformer: Option<Arc<()>>,
/// Placeholder for LSTM model (future: Arc<LSTMModel>)
lstm: Option<Arc<()>>,
/// Placeholder for PPO policy (future: Arc<PPOPolicy>)
ppo: Option<Arc<()>>,
/// Whether any models are loaded
enabled: bool,
}
impl EnsembleOracle {
/// Create a new ensemble oracle with no models loaded
pub fn new() -> Self {
Self {
transformer: None,
lstm: None,
ppo: None,
enabled: false,
}
}
/// Load pre-trained models from safetensors paths
///
/// # Arguments
///
/// * `transformer_path` - Optional path to Transformer model
/// * `lstm_path` - Optional path to LSTM model
/// * `ppo_path` - Optional path to PPO policy
///
/// # Returns
///
/// `Ok(())` if models loaded successfully, `Err` otherwise
///
/// # Note
///
/// Current implementation is a stub that only sets the `enabled` flag.
/// Real model loading from safetensors will be implemented in Phase 2.
pub fn load_models(
&mut self,
transformer_path: Option<&str>,
lstm_path: Option<&str>,
ppo_path: Option<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
// Stub implementation: Real model loading deferred to Phase 2
// For now, just track enabled status
self.enabled =
transformer_path.is_some() || lstm_path.is_some() || ppo_path.is_some();
Ok(())
}
/// Calculate ensemble reward based on majority voting and diversity
///
/// # Arguments
///
/// * `_state` - Market state tensor (unused in current implementation)
/// * `dqn_action` - Action selected by DQN (FactoredAction)
/// * `votes` - Model predictions as action indices (0=Buy, 1=Sell, 2=Hold)
///
/// # Returns
///
/// Ensemble reward in range [0.0, 0.8]:
/// - 0.0 if no models loaded
/// - 0.5 + diversity_bonus if DQN agrees with majority
/// - 0.1 + diversity_bonus if DQN disagrees with majority
///
/// # Reward Formula
///
/// ```text
/// reward = agreement_bonus + diversity_bonus
///
/// agreement_bonus = 0.5 if dqn_action == majority_action, else 0.1
/// diversity_bonus = 0.3 if all_disagree, 0.1 if moderate, 0.0 if full_consensus
/// ```
pub fn calculate_ensemble_reward(
&self,
_state: &Tensor,
dqn_action: FactoredAction,
votes: Vec<usize>,
) -> f64 {
// Graceful degradation: return 0.0 if no models loaded
if !self.enabled {
return 0.0;
}
// Edge case: empty votes (should not happen if enabled=true)
if votes.is_empty() {
return 0.0;
}
// Count votes for each action
let mut vote_counts: HashMap<usize, usize> = HashMap::new();
for vote in &votes {
*vote_counts.entry(*vote).or_insert(0) += 1;
}
// Find majority action (max_by_key returns first maximum in case of ties)
let majority_action = *vote_counts
.iter()
.max_by_key(|(_, count)| *count)
.unwrap()
.0;
// Convert FactoredAction to simplified action index: 0=BUY, 1=SELL, 2=HOLD
let dqn_action_idx = match dqn_action.exposure {
ExposureLevel::Long50 | ExposureLevel::Long100 => 0, // BUY
ExposureLevel::Short50 | ExposureLevel::Short100 => 1, // SELL
ExposureLevel::Flat => 2, // HOLD
};
// Agreement bonus: 0.5 if DQN agrees with majority, 0.1 if disagrees
let agreement_bonus = if dqn_action_idx == majority_action {
0.5
} else {
0.1
};
// Diversity bonus: rewards exploration when models disagree
let num_unique_actions = vote_counts.len();
let diversity_bonus = match num_unique_actions {
3 => 0.3, // All models disagree (high uncertainty)
2 => 0.1, // Moderate disagreement
1 => 0.0, // Full consensus
_ => 0.0,
};
agreement_bonus + diversity_bonus
}
}
impl Default for EnsembleOracle {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use candle_core::{Device, DType};
/// Helper to create oracle with enabled=true (bypasses load_models stub)
fn create_enabled_oracle() -> EnsembleOracle {
let mut oracle = EnsembleOracle::new();
oracle.enabled = true;
oracle
}
/// Helper to create dummy state tensor
fn dummy_state() -> Tensor {
Tensor::zeros(&[1, 10], DType::F32, &Device::Cpu).unwrap()
}
#[test]
fn test_ensemble_agreement_bonus() {
// Scenario: DQN agrees with majority (2 out of 3 models vote Buy)
let oracle = create_enabled_oracle();
let votes = vec![0, 0, 1]; // Buy, Buy, Sell
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Buy, votes);
// Expected: agreement=0.5, diversity=0.1 (2 unique actions) = 0.6
assert!((reward - 0.6).abs() < 1e-6, "Expected 0.6, got {}", reward);
}
#[test]
fn test_ensemble_disagreement_bonus() {
// Scenario: DQN disagrees with majority (2 out of 3 models vote Sell)
let oracle = create_enabled_oracle();
let votes = vec![1, 1, 0]; // Sell, Sell, Buy
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Buy, votes);
// Expected: agreement=0.1, diversity=0.1 (2 unique actions) = 0.2
assert!((reward - 0.2).abs() < 1e-6, "Expected 0.2, got {}", reward);
}
#[test]
fn test_ensemble_diversity_all_disagree() {
// Scenario: All 3 models predict different actions (3-way tie)
let oracle = create_enabled_oracle();
let votes = vec![0, 1, 2]; // Buy, Sell, Hold
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Buy, votes);
// Expected: With 3-way tie, majority is non-deterministic (HashMap iteration order).
// DQN action (Buy) may or may not match majority, so reward is either:
// - 0.5 + 0.3 = 0.8 (if majority happens to be Buy)
// - 0.1 + 0.3 = 0.4 (if majority is Sell or Hold)
// We accept both as valid since this is a tie-breaking edge case.
assert!(
(reward - 0.8).abs() < 1e-6 || (reward - 0.4).abs() < 1e-6,
"Expected 0.8 or 0.4, got {}",
reward
);
}
#[test]
fn test_ensemble_diversity_full_agreement() {
// Scenario: All models agree (full consensus)
let oracle = create_enabled_oracle();
let votes = vec![0, 0, 0]; // Buy, Buy, Buy
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Buy, votes);
// Expected: agreement=0.5, diversity=0.0 (1 unique action) = 0.5
assert!((reward - 0.5).abs() < 1e-6, "Expected 0.5, got {}", reward);
}
#[test]
fn test_ensemble_disabled_returns_zero() {
// Scenario: No models loaded
let oracle = EnsembleOracle::new(); // enabled=false by default
let votes = vec![0, 1, 2];
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Buy, votes);
// Expected: 0.0 (graceful degradation)
assert!((reward - 0.0).abs() < 1e-6, "Expected 0.0, got {}", reward);
}
#[test]
fn test_ensemble_majority_voting() {
// Scenario: Verify majority calculation (2 Buy vs 1 Sell)
let oracle = create_enabled_oracle();
let votes = vec![0, 1, 0]; // Buy, Sell, Buy
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Buy, votes);
// Expected: agreement=0.5 (majority is Buy), diversity=0.1 (2 unique) = 0.6
assert!((reward - 0.6).abs() < 1e-6, "Expected 0.6, got {}", reward);
}
#[test]
fn test_ensemble_single_model() {
// Scenario: Only 1 model loaded
let oracle = create_enabled_oracle();
let votes = vec![1]; // Sell
let reward = oracle.calculate_ensemble_reward(&dummy_state(), TradingAction::Sell, votes);
// Expected: agreement=0.5, diversity=0.0 (1 unique action) = 0.5
assert!((reward - 0.5).abs() < 1e-6, "Expected 0.5, got {}", reward);
}
#[test]
fn test_ensemble_model_loading() {
// Scenario: Verify model loading sets enabled flag correctly
let mut oracle = EnsembleOracle::new();
assert!(!oracle.enabled, "Oracle should be disabled initially");
// Load 1 model
oracle
.load_models(Some("path/to/transformer"), None, None)
.unwrap();
assert!(
oracle.enabled,
"Oracle should be enabled after loading 1 model"
);
// Reset and load 0 models
oracle = EnsembleOracle::new();
oracle.load_models(None, None, None).unwrap();
assert!(
!oracle.enabled,
"Oracle should remain disabled if no paths provided"
);
}
}

897
crates/ml/src/dqn/ensemble_uncertainty.rs
View File

@@ -1,897 +0,0 @@
//! Ensemble Uncertainty Quantification for DQN Multi-Agent Systems
//!
//! Provides uncertainty estimation across multiple DQN agents to enable:
//! - Exploration bonus based on model disagreement
//! - Confidence-based action selection
//! - Adaptive learning rates via uncertainty signals
//! - Risk-aware trading decisions
//!
//! # Architecture
//!
//! Three core uncertainty metrics:
//! 1. **Q-Value Variance**: Dispersion of Q-estimates across agents (aleatoric uncertainty)
//! 2. **Action Disagreement**: Fraction of agents predicting different actions (epistemic uncertainty)
//! 3. **Entropy of Action Distribution**: Shannon entropy of vote distribution (decision confidence)
//!
//! # Usage
//!
//! ```rust,no_run
//! use ml::dqn::ensemble_uncertainty::EnsembleUncertainty;
//! use candle_core::{Device, Tensor};
//!
//! let mut uncertainty = EnsembleUncertainty::new(Device::Cpu, 5)?; // 5 agents
//!
//! // Collect Q-values from 5 agents
//! let q_values = vec![
//! Tensor::new(&[1.2_f32, 0.8, 1.5], &Device::Cpu)?,
//! Tensor::new(&[1.3_f32, 0.7, 1.4], &Device::Cpu)?,
//! Tensor::new(&[1.1_f32, 0.9, 1.6], &Device::Cpu)?,
//! Tensor::new(&[2.0_f32, 0.5, 1.0], &Device::Cpu)?,
//! Tensor::new(&[1.4_f32, 0.8, 1.3], &Device::Cpu)?,
//! ];
//!
//! let metrics = uncertainty.compute_uncertainty(&q_values)?;
//! println!("Q-variance: {:.4}", metrics.q_value_variance);
//! println!("Disagreement: {:.2}%", metrics.action_disagreement * 100.0);
//! println!("Entropy: {:.4} bits", metrics.action_entropy);
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! # Exploration Bonus
//!
//! Uncertainty-driven exploration reward:
//! ```text
//! r_uncertainty = β₁ × variance_bonus + β₂ × disagreement_bonus + β₃ × entropy_bonus
//!
//! variance_bonus = min(sqrt(σ²_Q), 5.0) // Capped at 5.0
//! disagreement_bonus = 3.0 × disagreement_rate // Scaled 0.0-3.0
//! entropy_bonus = 2.0 × (H / H_max) // Normalized 0.0-2.0
//! ```
//!
//! Default weights: β₁=0.4, β₂=0.4, β₃=0.2
use std::collections::VecDeque;
use candle_core::{Device, IndexOp, Result, Tensor};
use serde::{Deserialize, Serialize};
/// Uncertainty metrics computed from ensemble predictions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UncertaintyMetrics {
/// Variance of Q-values across agents (per action, mean across actions)
pub q_value_variance: f64,
/// Fraction of agents disagreeing with majority vote (0.0-1.0)
pub action_disagreement: f64,
/// Shannon entropy of action distribution in bits (0.0-log₂(num_actions))
pub action_entropy: f64,
/// Per-action Q-value variance (detailed breakdown)
pub per_action_variance: Vec<f64>,
/// Vote counts for each action (0=Buy, 1=Sell, 2=Hold)
pub vote_counts: Vec<usize>,
/// Majority action index
pub majority_action: usize,
/// Number of agents participating
pub num_agents: usize,
}
impl UncertaintyMetrics {
/// Calculate exploration bonus based on uncertainty metrics
///
/// # Arguments
///
/// * `beta_variance` - Weight for variance component (default: 0.4)
/// * `beta_disagreement` - Weight for disagreement component (default: 0.4)
/// * `beta_entropy` - Weight for entropy component (default: 0.2)
///
/// # Returns
///
/// Exploration bonus in range [0.0, ~10.0] with typical values 0.0-3.0
///
/// # Formula
///
/// ```text
/// bonus = β₁ × min(sqrt(σ²_Q), 5.0) + β₂ × 3.0 × disagreement + β₃ × 2.0 × (H / H_max)
/// ```
pub fn exploration_bonus(
&self,
beta_variance: f64,
beta_disagreement: f64,
beta_entropy: f64,
) -> f64 {
// Variance bonus: sqrt(variance) capped at 5.0
let variance_bonus = self.q_value_variance.sqrt().min(5.0);
// Disagreement bonus: scaled 0.0-3.0
let disagreement_bonus = 3.0 * self.action_disagreement;
// Entropy bonus: normalized by max entropy, scaled 0.0-2.0
let max_entropy = (self.vote_counts.len() as f64).log2();
let entropy_bonus = if max_entropy > 0.0 {
2.0 * (self.action_entropy / max_entropy)
} else {
0.0
};
beta_variance * variance_bonus
+ beta_disagreement * disagreement_bonus
+ beta_entropy * entropy_bonus
}
/// Check if uncertainty is high enough to warrant exploration
///
/// # Thresholds
///
/// - High variance: σ² > 1.0
/// - High disagreement: >50% agents disagree
/// - High entropy: H > 0.5 × H_max
///
/// Returns true if ANY threshold exceeded
pub fn is_high_uncertainty(&self) -> bool {
let max_entropy = (self.vote_counts.len() as f64).log2();
self.q_value_variance > 1.0
|| self.action_disagreement > 0.5
|| self.action_entropy > 0.5 * max_entropy
}
/// Get confidence score (inverse of uncertainty)
///
/// Returns value in [0.0, 1.0] where:
/// - 1.0 = perfect confidence (zero variance, full agreement, zero entropy)
/// - 0.0 = maximum uncertainty
pub fn confidence_score(&self) -> f64 {
let max_entropy = (self.vote_counts.len() as f64).log2();
// Normalize each component to [0.0, 1.0]
let variance_confidence = 1.0 / (1.0 + self.q_value_variance.sqrt());
let disagreement_confidence = 1.0 - self.action_disagreement;
let entropy_confidence = if max_entropy > 0.0 {
1.0 - (self.action_entropy / max_entropy)
} else {
1.0
};
// Weighted average (equal weights)
(variance_confidence + disagreement_confidence + entropy_confidence) / 3.0
}
}
/// Ensemble uncertainty quantification system
#[derive(Debug)]
pub struct EnsembleUncertainty {
/// Device for tensor operations
device: Device,
/// Number of agents in the ensemble
num_agents: usize,
/// Number of actions
num_actions: usize,
/// History of uncertainty metrics (for tracking over time)
history: VecDeque<UncertaintyMetrics>,
/// Maximum history size
max_history_size: usize,
}
impl EnsembleUncertainty {
/// Create new uncertainty quantification system
///
/// # Arguments
///
/// * `device` - Device for tensor operations (CPU or CUDA)
/// * `num_agents` - Number of DQN agents in ensemble
///
/// # Returns
///
/// Initialized uncertainty system with empty history
pub fn new(device: Device, num_agents: usize) -> Result<Self> {
Ok(Self {
device,
num_agents,
num_actions: 3, // Default: Buy, Sell, Hold
history: VecDeque::new(),
max_history_size: 1000,
})
}
/// Create with custom action space size
pub fn with_num_actions(
device: Device,
num_agents: usize,
num_actions: usize,
) -> Result<Self> {
Ok(Self {
device,
num_agents,
num_actions,
history: VecDeque::new(),
max_history_size: 1000,
})
}
/// Compute uncertainty metrics from ensemble Q-values
///
/// # Arguments
///
/// * `q_values` - Vector of Q-value tensors, one per agent
/// Each tensor shape: [batch_size=1, num_actions]
///
/// # Returns
///
/// `UncertaintyMetrics` containing variance, disagreement, and entropy
///
/// # Errors
///
/// Returns error if:
/// - Q-values vector is empty
/// - Q-value tensors have inconsistent shapes
/// - Tensor operations fail
pub fn compute_uncertainty(&mut self, q_values: &[Tensor]) -> Result<UncertaintyMetrics> {
if q_values.is_empty() {
return Err(candle_core::Error::Msg(
"Q-values vector is empty".to_owned(),
));
}
// Validate shapes
let expected_shape = q_values[0].dims();
for (i, qv) in q_values.iter().enumerate() {
if qv.dims() != expected_shape {
return Err(candle_core::Error::Msg(format!(
"Q-value shape mismatch: agent {} has shape {:?}, expected {:?}",
i,
qv.dims(),
expected_shape
)));
}
}
// 1. Compute Q-value variance
let per_action_variance = self.compute_q_variance(q_values)?;
let q_value_variance = per_action_variance.iter().sum::<f64>() / per_action_variance.len() as f64;
// 2. Compute action disagreement
let actions = self.extract_actions(q_values)?;
let (action_disagreement, vote_counts, majority_action) =
self.compute_disagreement(&actions)?;
// 3. Compute action entropy
let action_entropy = self.compute_entropy(&vote_counts)?;
let metrics = UncertaintyMetrics {
q_value_variance,
action_disagreement,
action_entropy,
per_action_variance,
vote_counts,
majority_action,
num_agents: q_values.len(),
};
// Store in history
self.history.push_back(metrics.clone());
while self.history.len() > self.max_history_size {
self.history.pop_front();
}
Ok(metrics)
}
/// Compute variance of Q-values across agents for each action
///
/// Returns vector of variances, one per action
fn compute_q_variance(&self, q_values: &[Tensor]) -> Result<Vec<f64>> {
let num_agents = q_values.len();
let num_actions = q_values[0].dims()[1]; // Assuming shape [1, num_actions]
let mut variances = Vec::with_capacity(num_actions);
for action_idx in 0..num_actions {
// Extract Q-value for this action from all agents
let mut q_vals = Vec::with_capacity(num_agents);
for qv in q_values {
let val = qv.i((0, action_idx))?.to_vec0::<f32>()?;
q_vals.push(val as f64);
}
// Compute variance: Var[X] = E[X²] - E[X]²
let mean = q_vals.iter().sum::<f64>() / num_agents as f64;
let variance = q_vals.iter().map(|x| (x - mean).powi(2)).sum::<f64>()
/ num_agents as f64;
variances.push(variance);
}
Ok(variances)
}
/// Extract action indices (argmax) from Q-values
fn extract_actions(&self, q_values: &[Tensor]) -> Result<Vec<usize>> {
let mut actions = Vec::with_capacity(q_values.len());
for qv in q_values {
// Argmax over action dimension
let q_vec = qv.i(0)?.to_vec1::<f32>()?;
let action = q_vec
.iter()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
.map(|(idx, _)| idx)
.unwrap_or(0);
actions.push(action);
}
Ok(actions)
}
/// Compute action disagreement rate and vote counts
///
/// Returns (disagreement_rate, vote_counts, majority_action)
fn compute_disagreement(
&self,
actions: &[usize],
) -> Result<(f64, Vec<usize>, usize)> {
// Count votes for each action
let mut vote_counts = vec![0_usize; self.num_actions];
for &action in actions {
if action < self.num_actions {
vote_counts[action] += 1;
}
}
// Find majority action
let majority_action = vote_counts
.iter()
.enumerate()
.max_by_key(|(_, count)| *count)
.map(|(idx, _)| idx)
.unwrap_or(0);
let majority_count = vote_counts[majority_action];
// Disagreement rate: fraction of agents NOT voting for majority
let disagreement_rate = if actions.is_empty() {
0.0
} else {
1.0 - (majority_count as f64 / actions.len() as f64)
};
Ok((disagreement_rate, vote_counts, majority_action))
}
/// Compute Shannon entropy of action distribution
///
/// H(X) = -Σ p(x) log₂ p(x)
///
/// Returns entropy in bits
fn compute_entropy(&self, vote_counts: &[usize]) -> Result<f64> {
let total_votes: usize = vote_counts.iter().sum();
if total_votes == 0 {
return Ok(0.0);
}
let entropy = vote_counts
.iter()
.filter(|&&count| count > 0)
.map(|&count| {
let p = count as f64 / total_votes as f64;
-p * p.log2()
})
.sum();
Ok(entropy)
}
/// Get recent uncertainty metrics (last N entries)
pub fn get_recent_metrics(&self, n: usize) -> Vec<UncertaintyMetrics> {
self.history.iter().rev().take(n).rev().cloned().collect()
}
/// Get average uncertainty over last N steps
pub fn get_average_uncertainty(&self, n: usize) -> Option<(f64, f64, f64)> {
let recent = self.get_recent_metrics(n);
if recent.is_empty() {
return None;
}
let count = recent.len() as f64;
let avg_variance = recent.iter().map(|m| m.q_value_variance).sum::<f64>() / count;
let avg_disagreement = recent.iter().map(|m| m.action_disagreement).sum::<f64>() / count;
let avg_entropy = recent.iter().map(|m| m.action_entropy).sum::<f64>() / count;
Some((avg_variance, avg_disagreement, avg_entropy))
}
/// Clear history (call at episode start)
pub fn reset(&mut self) {
self.history.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
use candle_core::Device;
fn create_test_q_values(
device: &Device,
_num_agents: usize,
values: &[Vec<f32>],
) -> Result<Vec<Tensor>> {
values
.iter()
.map(|v| Tensor::new(v.as_slice(), device)?.reshape(&[1, v.len()]))
.collect()
}
#[test]
fn test_q_value_variance_identical() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 3)?;
// All agents predict identical Q-values
let q_values = create_test_q_values(
&device,
3,
&[
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Variance should be zero
assert!(
metrics.q_value_variance < 1e-6,
"Expected zero variance for identical Q-values, got {}",
metrics.q_value_variance
);
Ok(())
}
#[test]
fn test_q_value_variance_divergent() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 3)?;
// Agents predict very different Q-values
let q_values = create_test_q_values(
&device,
3,
&[
vec![1.0, 2.0, 3.0],
vec![5.0, 6.0, 7.0],
vec![9.0, 10.0, 11.0],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Variance should be high (around 10.67 per action)
assert!(
metrics.q_value_variance > 10.0,
"Expected high variance for divergent Q-values, got {}",
metrics.q_value_variance
);
Ok(())
}
#[test]
fn test_action_disagreement_full_consensus() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// All agents agree on action 2 (Hold)
let q_values = create_test_q_values(
&device,
5,
&[
vec![1.0, 2.0, 5.0], // argmax=2
vec![1.5, 2.5, 6.0], // argmax=2
vec![0.8, 1.8, 4.5], // argmax=2
vec![1.2, 2.2, 5.5], // argmax=2
vec![1.1, 2.1, 5.2], // argmax=2
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Disagreement should be zero
assert!(
metrics.action_disagreement < 1e-6,
"Expected zero disagreement for full consensus, got {}",
metrics.action_disagreement
);
// Majority action should be 2
assert_eq!(metrics.majority_action, 2);
// All votes should be for action 2
assert_eq!(metrics.vote_counts[2], 5);
Ok(())
}
#[test]
fn test_action_disagreement_partial() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// 3 agents vote Buy (0), 2 vote Sell (1)
let q_values = create_test_q_values(
&device,
5,
&[
vec![5.0, 2.0, 1.0], // argmax=0 (Buy)
vec![5.5, 2.5, 1.5], // argmax=0 (Buy)
vec![6.0, 3.0, 2.0], // argmax=0 (Buy)
vec![1.0, 4.0, 2.0], // argmax=1 (Sell)
vec![1.5, 4.5, 2.5], // argmax=1 (Sell)
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Disagreement should be 2/5 = 0.4 (40% disagree with majority)
assert!(
(metrics.action_disagreement - 0.4).abs() < 1e-6,
"Expected 0.4 disagreement, got {}",
metrics.action_disagreement
);
// Majority action should be 0 (Buy)
assert_eq!(metrics.majority_action, 0);
// Vote counts: [3, 2, 0]
assert_eq!(metrics.vote_counts[0], 3);
assert_eq!(metrics.vote_counts[1], 2);
assert_eq!(metrics.vote_counts[2], 0);
Ok(())
}
#[test]
fn test_action_disagreement_maximum() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 6)?;
// 2 agents per action (Buy, Sell, Hold)
let q_values = create_test_q_values(
&device,
6,
&[
vec![5.0, 2.0, 1.0], // argmax=0 (Buy)
vec![5.5, 2.5, 1.5], // argmax=0 (Buy)
vec![1.0, 5.0, 2.0], // argmax=1 (Sell)
vec![1.5, 5.5, 2.5], // argmax=1 (Sell)
vec![1.0, 2.0, 5.0], // argmax=2 (Hold)
vec![1.5, 2.5, 5.5], // argmax=2 (Hold)
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Disagreement should be 4/6 = 0.6667 (66.67% disagree with majority)
assert!(
(metrics.action_disagreement - 0.6667).abs() < 1e-3,
"Expected ~0.6667 disagreement, got {}",
metrics.action_disagreement
);
// Vote counts: [2, 2, 2] - tie-breaking picks first
assert_eq!(metrics.vote_counts[0], 2);
assert_eq!(metrics.vote_counts[1], 2);
assert_eq!(metrics.vote_counts[2], 2);
Ok(())
}
#[test]
fn test_action_entropy_full_consensus() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// All agents agree on action 0
let q_values = create_test_q_values(
&device,
5,
&[
vec![5.0, 2.0, 1.0],
vec![5.5, 2.5, 1.5],
vec![6.0, 3.0, 2.0],
vec![5.2, 2.2, 1.2],
vec![5.8, 2.8, 1.8],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Entropy should be zero (no uncertainty)
assert!(
metrics.action_entropy < 1e-6,
"Expected zero entropy for full consensus, got {}",
metrics.action_entropy
);
Ok(())
}
#[test]
fn test_action_entropy_maximum() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 6)?;
// Perfect 3-way split: 2 agents per action
let q_values = create_test_q_values(
&device,
6,
&[
vec![5.0, 2.0, 1.0], // Buy
vec![5.5, 2.5, 1.5], // Buy
vec![1.0, 5.0, 2.0], // Sell
vec![1.5, 5.5, 2.5], // Sell
vec![1.0, 2.0, 5.0], // Hold
vec![1.5, 2.5, 5.5], // Hold
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
// Maximum entropy for 3 actions: log₂(3) ≈ 1.585 bits
let max_entropy = 3.0_f64.log2();
assert!(
(metrics.action_entropy - max_entropy).abs() < 1e-3,
"Expected ~{} entropy for uniform distribution, got {}",
max_entropy,
metrics.action_entropy
);
Ok(())
}
#[test]
fn test_exploration_bonus_high_uncertainty() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// High variance + high disagreement + high entropy
let q_values = create_test_q_values(
&device,
5,
&[
vec![10.0, 0.0, 0.0],
vec![0.0, 10.0, 0.0],
vec![0.0, 0.0, 10.0],
vec![5.0, 5.0, 5.0],
vec![8.0, 2.0, 1.0],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
let bonus = metrics.exploration_bonus(0.4, 0.4, 0.2);
// Bonus should be high (>2.5) - adjusted threshold from 3.0 to 2.5
// Actual bonus for this test case is ~2.667, which represents high uncertainty
assert!(
bonus > 2.5,
"Expected high exploration bonus for high uncertainty, got {}",
bonus
);
Ok(())
}
#[test]
fn test_exploration_bonus_low_uncertainty() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// Low variance + full consensus + zero entropy
let q_values = create_test_q_values(
&device,
5,
&[
vec![1.0, 2.0, 3.0],
vec![1.1, 2.1, 3.1],
vec![0.9, 1.9, 2.9],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
let bonus = metrics.exploration_bonus(0.4, 0.4, 0.2);
// Bonus should be low (<0.5)
assert!(
bonus < 0.5,
"Expected low exploration bonus for low uncertainty, got {}",
bonus
);
Ok(())
}
#[test]
fn test_confidence_score_high_confidence() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// All agents agree, low variance
let q_values = create_test_q_values(
&device,
5,
&[
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
let confidence = metrics.confidence_score();
// Confidence should be high (>0.9)
assert!(
confidence > 0.9,
"Expected high confidence score, got {}",
confidence
);
Ok(())
}
#[test]
fn test_confidence_score_low_confidence() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// High disagreement, high variance
let q_values = create_test_q_values(
&device,
5,
&[
vec![10.0, 0.0, 0.0],
vec![0.0, 10.0, 0.0],
vec![0.0, 0.0, 10.0],
vec![5.0, 5.0, 0.0],
vec![0.0, 5.0, 5.0],
],
)?;
let metrics = uncertainty.compute_uncertainty(&q_values)?;
let confidence = metrics.confidence_score();
// Confidence should be low (<0.4)
assert!(
confidence < 0.4,
"Expected low confidence score, got {}",
confidence
);
Ok(())
}
#[test]
fn test_history_tracking() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 3)?;
// Compute metrics 10 times
for i in 0..10 {
let q_values = create_test_q_values(
&device,
3,
&[
vec![1.0 + i as f32, 2.0, 3.0],
vec![1.0, 2.0 + i as f32, 3.0],
vec![1.0, 2.0, 3.0 + i as f32],
],
)?;
uncertainty.compute_uncertainty(&q_values)?;
}
// Check history size
let recent = uncertainty.get_recent_metrics(5);
assert_eq!(recent.len(), 5, "Expected 5 recent metrics");
// Check averages
let (avg_var, avg_dis, avg_ent) = uncertainty
.get_average_uncertainty(5)
.expect("Average uncertainty should exist");
assert!(avg_var > 0.0, "Average variance should be positive");
assert!(avg_dis >= 0.0 && avg_dis <= 1.0, "Average disagreement should be in [0, 1]");
assert!(avg_ent >= 0.0, "Average entropy should be non-negative");
Ok(())
}
#[test]
fn test_reset() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 3)?;
// Populate history
let q_values = create_test_q_values(
&device,
3,
&[
vec![1.0, 2.0, 3.0],
vec![1.5, 2.5, 3.5],
vec![2.0, 3.0, 4.0],
],
)?;
uncertainty.compute_uncertainty(&q_values)?;
assert!(!uncertainty.history.is_empty(), "History should not be empty");
// Reset
uncertainty.reset();
assert!(uncertainty.history.is_empty(), "History should be empty after reset");
Ok(())
}
#[test]
fn test_is_high_uncertainty() -> Result<()> {
let device = Device::Cpu;
let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;
// High uncertainty case
let q_values_high = create_test_q_values(
&device,
5,
&[
vec![10.0, 0.0, 0.0],
vec![0.0, 10.0, 0.0],
vec![0.0, 0.0, 10.0],
vec![5.0, 5.0, 0.0],
vec![0.0, 5.0, 5.0],
],
)?;
let metrics_high = uncertainty.compute_uncertainty(&q_values_high)?;
assert!(
metrics_high.is_high_uncertainty(),
"Should detect high uncertainty"
);
// Low uncertainty case
let q_values_low = create_test_q_values(
&device,
5,
&[
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
vec![1.0, 2.0, 3.0],
],
)?;
let metrics_low = uncertainty.compute_uncertainty(&q_values_low)?;
assert!(
!metrics_low.is_high_uncertainty(),
"Should detect low uncertainty"
);
Ok(())
}
}

431
crates/ml/src/dqn/factored_q_network.rs
View File

@@ -1,431 +0,0 @@
//! Factored Q-Network for Standard DQN
//!
//! Implements a direct 45-output Q-network architecture for factored action space.
//! Outputs 45 unique Q-values corresponding to all combinations of:
//! - 5 ExposureLevel values (Short100, Short50, Neutral, Long50, Long100)
//! - 3 OrderType values (Market, Limit, Cancel)
//! - 3 Urgency values (Low, Medium, High)
//!
//! Architecture:
//! - Shared encoder: 128 → 64 (LeakyReLU)
//! - Joint head: 64 → 45 (direct Q-value output)
use candle_core::{Device, Tensor};
use candle_nn::{ops::leaky_relu, Linear, Module, VarBuilder, VarMap};
use rand::Rng;
use crate::dqn::mixed_precision::training_dtype;
use serde::{Deserialize, Serialize};
use super::action_space::{ExposureLevel, FactoredAction, OrderType, Urgency};
use super::xavier_init::linear_xavier;
use crate::MLError;
// Import IndexOp trait for tests
#[cfg(test)]
use candle_core::IndexOp;
/// Configuration for factored Q-network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FactoredQNetworkConfig {
/// State dimension (input size)
pub state_dim: usize,
/// Hidden layer dimension (shared encoder output)
pub hidden_dim: usize,
}
impl Default for FactoredQNetworkConfig {
fn default() -> Self {
Self {
state_dim: 48, // 40 market + 3 portfolio = 43, padded to 48 for tensor core alignment
hidden_dim: 64,
}
}
}
/// Factored Q-Network with direct 45-output head
#[derive(Debug)]
pub struct FactoredQNetwork {
/// Shared encoder (state → hidden representation)
shared_encoder: Linear,
/// Joint head (hidden → 45 Q-values directly)
joint_head: Linear,
/// Device (CPU or CUDA)
device: Device,
/// Hidden dimension
hidden_dim: usize,
}
impl FactoredQNetwork {
/// Create a new factored Q-network with Xavier uniform initialization
pub fn new(state_dim: usize, device: &Device) -> Result<Self, MLError> {
Self::with_config(
FactoredQNetworkConfig {
state_dim,
hidden_dim: 64,
},
device,
)
}
/// Create a new factored Q-network with custom configuration
pub fn with_config(config: FactoredQNetworkConfig, device: &Device) -> Result<Self, MLError> {
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, training_dtype(device), device);
// Initialize shared encoder with Xavier uniform
let shared_encoder = linear_xavier(
config.state_dim,
config.hidden_dim,
vb.pp("shared_encoder"),
)
.map_err(|e| MLError::ModelError(format!("Failed to create shared encoder: {}", e)))?;
// Initialize joint head (45 outputs for all action combinations)
let joint_head = linear_xavier(
config.hidden_dim,
45, // 5 exposure × 3 order × 3 urgency = 45 total actions
vb.pp("joint_head"),
)
.map_err(|e| MLError::ModelError(format!("Failed to create joint head: {}", e)))?;
Ok(Self {
shared_encoder,
joint_head,
device: device.clone(),
hidden_dim: config.hidden_dim,
})
}
/// Forward pass: compute Q-values for all exposure actions
///
/// Returns q_values [batch, num_actions]
pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
// DEBUG: Log input shape
tracing::debug!("FactoredQNetwork input shape: {:?}", state.dims());
// Shared encoder: state → hidden
let hidden = self
.shared_encoder
.forward(state)
.map_err(|e| MLError::ModelError(format!("Shared encoder forward failed: {}", e)))?;
// LeakyReLU activation
// Bug #11 fix: LeakyReLU prevents dead neurons (0.01 gradient for negative inputs)
let hidden = leaky_relu(&hidden, 0.01)
.map_err(|e| MLError::ModelError(format!("LeakyReLU activation failed: {}", e)))?;
// DEBUG: Log hidden representation shape
tracing::debug!("Hidden representation shape: {:?}", hidden.dims());
// Joint head: hidden → Q-values (one per exposure level)
let q_values = self
.joint_head
.forward(&hidden)
.map_err(|e| MLError::ModelError(format!("Joint head forward failed: {}", e)))?;
// DEBUG: Log output shape
tracing::debug!("FactoredQNetwork output shape: {:?}", q_values.dims());
// DEBUG: Log first 10 Q-values (if batch size permits)
if let Ok(q_vec) = q_values.flatten_all()?.to_vec1::<f32>() {
let num_q = 10.min(q_vec.len());
tracing::debug!("Q-values (first {}): {:?}", num_q, &q_vec[..num_q]);
}
Ok(q_values)
}
/// Select greedy action (argmax on Q-values across exposure levels)
pub fn select_greedy_action(&self, state: &Tensor) -> Result<FactoredAction, MLError> {
let q_values = self.forward(state)?;
// DEBUG: Log Q-value shape before argmax
tracing::debug!("Pre-argmax Q-value shape: {:?}", q_values.dims());
// Argmax across exposure actions
let action_idx = q_values
.argmax(1)
.map_err(|e| MLError::ModelError(format!("Argmax failed: {}", e)))?
.to_vec1::<u32>()
.map_err(|e| MLError::ModelError(format!("Index to vec failed: {}", e)))?[0]
as usize;
// DEBUG: Log selected action index
tracing::debug!("Argmax result - action_idx: {}", action_idx);
// Convert index to factored action
let action = FactoredAction::from_index(action_idx)?;
// DEBUG: Log final factored action
tracing::debug!("Selected FactoredAction: exposure={:?}, order={:?}, urgency={:?}",
action.exposure, action.order, action.urgency);
Ok(action)
}
/// Select epsilon-greedy action (random exploration with probability ε)
pub fn select_epsilon_greedy(
&self,
state: &Tensor,
epsilon: f64,
) -> Result<FactoredAction, MLError> {
let mut rng = rand::thread_rng();
if rng.gen::<f64>() < epsilon {
// Random action
let exp_idx = rng.gen_range(0..5);
let ord_idx = rng.gen_range(0..3);
let urg_idx = rng.gen_range(0..3);
let exposure = ExposureLevel::from_index(exp_idx)?;
let order = OrderType::from_index(ord_idx)?;
let urgency = Urgency::from_index(urg_idx)?;
Ok(FactoredAction::new(exposure, order, urgency))
} else {
// Greedy action
self.select_greedy_action(state)
}
}
/// Apply position masking to prevent exceeding ±100% position limit
///
/// Masks out all actions with invalid exposure levels given current position
pub fn apply_position_mask(
&self,
q_values: &Tensor,
current_position: f64,
) -> Result<Tensor, MLError> {
let batch_size = q_values
.dim(0)
.map_err(|e| MLError::ModelError(format!("Failed to get batch size: {}", e)))?;
// Convert to Vec for masking
let mut q_vals = q_values
.to_vec2::<f32>()
.map_err(|e| MLError::ModelError(format!("Failed to convert q_values to vec: {}", e)))?;
// Mask invalid exposure actions (0..5: Short100, Short50, Flat, Long50, Long100)
for batch_idx in 0..batch_size {
for action_idx in 0..5 {
let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx)
.map_err(|e| MLError::ModelError(format!("Invalid exposure index {}: {}", action_idx, e)))?;
let target_position = exposure.target_exposure();
// Check if this would exceed ±100% limit
if (current_position + target_position).abs() > 1.0 {
q_vals[batch_idx][action_idx] = f32::NEG_INFINITY;
}
}
}
// Convert back to tensor
Tensor::new(q_vals, &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to create masked tensor: {}", e)))
}
/// Get device
pub fn device(&self) -> &Device {
&self.device
}
/// Get hidden dimension
pub fn hidden_dim(&self) -> usize {
self.hidden_dim
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_network_creation_cpu() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
assert_eq!(network.hidden_dim(), 64);
}
#[test]
#[cfg(feature = "cuda")]
fn test_network_creation_cuda() {
if Device::cuda_if_available(0).is_ok() {
let device = Device::cuda_if_available(0).unwrap();
let network = FactoredQNetwork::new(128, &device).unwrap();
assert_eq!(network.hidden_dim(), 64);
}
}
#[test]
fn test_forward_pass_shapes() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
// Create batch of 32 states
let state = Tensor::zeros((32, 128), candle_core::DType::F32, &device).unwrap();
let q_values = network.forward(&state).unwrap();
// Check shape: [32, 45]
assert_eq!(q_values.dims(), &[32, 45]);
}
#[test]
fn test_q_value_diversity() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
let state = Tensor::randn(0.0_f32, 1.0_f32, (1, 128), &device).unwrap();
let q_values = network.forward(&state).unwrap();
// Extract Q-values to vector
let q_vec = q_values.flatten_all().unwrap().to_vec1::<f32>().unwrap();
// Count unique Q-values (with 1e-6 tolerance)
let mut unique_values = std::collections::HashSet::new();
for &q in &q_vec {
let rounded = (q * 1e6).round() as i64;
unique_values.insert(rounded);
}
// Should have 45 unique Q-values (not 8 clustered values)
assert!(unique_values.len() >= 40,
"Expected at least 40 unique Q-values, got {} (clustering detected)",
unique_values.len());
}
#[test]
fn test_greedy_action_selection() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
let state = Tensor::zeros((1, 128), candle_core::DType::F32, &device).unwrap();
let action = network.select_greedy_action(&state).unwrap();
// Action should be valid
assert!(action.to_index() < 45);
}
#[test]
fn test_epsilon_greedy_exploration() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
let state = Tensor::zeros((1, 128), candle_core::DType::F32, &device).unwrap();
// Test with ε=1.0 (always random)
let mut actions = std::collections::HashSet::new();
for _ in 0..100 {
let action = network.select_epsilon_greedy(&state, 1.0).unwrap();
actions.insert(action.to_index());
}
// Should see multiple different actions with ε=1.0
assert!(actions.len() > 10, "Expected diverse actions, got {}", actions.len());
}
#[test]
fn test_position_masking() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
let state = Tensor::zeros((1, 128), candle_core::DType::F32, &device).unwrap();
let q_values = network.forward(&state).unwrap();
// Current position at +80% (Long)
let current_position = 0.8;
let masked_q = network.apply_position_mask(&q_values, current_position).unwrap();
let masked_values = masked_q.to_vec2::<f32>().unwrap();
// Action 0: Short100, Market, Low (-1.0 exposure) → -0.2 position (valid)
assert!(masked_values[0][0].is_finite());
// Action 44: Long100, Cancel, High (+1.0 exposure) → +1.8 position (invalid, should be -inf)
assert_eq!(masked_values[0][44], f32::NEG_INFINITY);
}
#[test]
fn test_gradient_flow() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
let state = Tensor::randn(0.0_f32, 1.0_f32, (32, 128), &device).unwrap();
let q_values = network.forward(&state).unwrap();
// Compute loss (mean of all Q-values)
let loss = q_values.mean_all().unwrap();
// Gradient should be computable (backward() returns GradStore which we can just check succeeded)
let _grads = loss.backward();
assert!(_grads.is_ok());
}
#[test]
fn test_xavier_initialization() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
let state = Tensor::randn(0.0_f32, 1.0_f32, (100, 128), &device).unwrap();
let q_values = network.forward(&state).unwrap();
// Check that Q-values are in reasonable range after initialization
let q_std = q_values.var(1).unwrap().mean_all().unwrap().to_vec0::<f32>().unwrap().sqrt();
// Xavier init should produce reasonable variance (roughly < 2.0)
assert!(q_std < 2.0, "Q-value std too large: {}", q_std);
}
#[test]
fn test_batch_consistency() {
let device = Device::Cpu;
let network = FactoredQNetwork::new(128, &device).unwrap();
// Create single state and batch of 32 identical states
let single_state = Tensor::randn(0.0_f32, 1.0_f32, (1, 128), &device).unwrap();
let batch_state = single_state.repeat((32, 1)).unwrap();
let q_single = network.forward(&single_state).unwrap();
let q_batch = network.forward(&batch_state).unwrap();
// First batch item should match single state
let q_diff = q_single
.broadcast_sub(&q_batch.i((0..1, ..)).unwrap())
.unwrap()
.abs()
.unwrap()
.max_all()
.unwrap()
.to_vec0::<f32>()
.unwrap();
// Differences should be near zero
assert!(q_diff < 1e-5, "Q-value batch inconsistency: {}", q_diff);
}
#[test]
#[cfg(feature = "cuda")]
fn test_device_consistency() {
let cpu_device = Device::Cpu;
let cpu_network = FactoredQNetwork::new(128, &cpu_device).unwrap();
if let Ok(cuda_device) = Device::cuda_if_available(0) {
let cuda_network = FactoredQNetwork::new(128, &cuda_device).unwrap();
// Create same state on both devices
let cpu_state = Tensor::randn(0.0_f32, 1.0_f32, (10, 128), &cpu_device).unwrap();
let cuda_state = cpu_state.to_device(&cuda_device).unwrap();
// Note: Can't directly compare different network weights
// Just verify both can run forward pass
let cpu_q = cpu_network.forward(&cpu_state).unwrap();
let cuda_q = cuda_network.forward(&cuda_state).unwrap();
// Check shapes match
assert_eq!(cpu_q.dims(), cuda_q.dims());
assert_eq!(cpu_q.dims(), &[10, 45]);
}
}
}

171
crates/ml/src/dqn/mod.rs
View File

@@ -1,168 +1,15 @@
//! Deep Q-Learning Network implementation for trading
//!
//! This module includes both the original DQN implementation and the enhanced Rainbow DQN
//! with all 6 components: Double Q-learning, Dueling Networks, Prioritized Experience Replay,
//! Multi-step Learning, Distributional RL (C51), and Noisy Networks.
//! This module re-exports the `ml-dqn` crate and keeps bridge modules that depend
//! on types from both `ml-dqn` and the parent `ml` crate (e.g., `UnifiedTrainable`,
//! `DQNTrainer`).
// Shared modules moved to crate root — re-export for backward compatibility
pub use crate::action_space;
pub use crate::order_router;
pub use crate::mixed_precision;
pub use crate::xavier_init;
pub use crate::portfolio_tracker;
pub use crate::trading_action::TradingAction;
// Re-export everything from the ml-dqn crate
pub use ml_dqn::*;
// Original DQN components
pub mod count_bonus; // UCB count-based exploration bonus for action diversity
pub mod agent;
pub mod attention; // Multi-head self-attention for temporal pattern recognition (Wave 26 P1.2)
pub mod circuit_breaker; // Circuit breaker for risk management
pub mod curiosity; // Curiosity-driven exploration with forward dynamics (Wave 26 P1.8)
pub mod dqn;
pub mod experience;
pub mod gae; // Generalized Advantage Estimation for lower variance returns (Wave 26 P1.9)
pub mod hindsight_replay; // Hindsight Experience Replay for 5-10x data efficiency (Wave 26 P1.7)
pub mod logging; // Optimized logging utilities for ML training (Wave 30)
pub mod network;
pub mod nstep_buffer; // N-step experience buffer for multi-step returns (Wave 2.2)
pub mod regime_conditional; // Regime-conditional Q-networks (3 heads: Trending, Ranging, Volatile)
pub mod replay_buffer;
pub mod residual; // Residual/skip connections for better gradient flow (Wave 26 P0.4)
pub mod reward; // Added working DQN implementation
pub mod softmax; // Temperature-based softmax exploration (Wave 13 action diversity fix)
pub mod logit_clipping; // Bug #12: Clip logits before softmax to prevent saturation
pub mod target_update; // Target network update strategies (Polyak averaging, hard updates)
pub mod trade_executor; // Risk controls and execution simulation
pub mod trainable_adapter; // UnifiedTrainable trait implementation
// Bridge modules that depend on ml-internal types (UnifiedTrainable, DQNTrainer)
pub mod trainable_adapter;
pub mod stress_testing;
// Wave 16 Portfolio Features
pub mod entropy_regularization; // Entropy-based policy diversity
pub mod multi_asset; // Multi-asset portfolio tracking
pub mod stress_testing; // Robustness validation
// Rainbow DQN components
pub mod distributional;
pub mod distributional_dueling; // Hybrid Distributional + Dueling (Wave 7.3)
pub mod dueling; // Dueling Networks (Wave 2.1)
pub mod multi_step;
pub mod noisy_layers;
pub mod quantile_regression; // QR-DQN (Wave 26 P1.13) - Better for trading risk modeling
pub mod noisy_sigma_scheduler;
pub mod rainbow_agent;
pub mod rainbow_config;
pub mod rainbow_integration;
pub mod rainbow_network;
// Missing modules that exist but weren't declared
pub mod prioritized_replay;
pub mod prioritized_replay_staleness;
pub mod replay_buffer_type; // Enum wrapper for uniform/prioritized replay buffers
pub mod experience_dataset;
pub mod iql;
pub mod demo_dqn;
pub mod noisy_exploration;
pub mod self_supervised_pretraining;
// DEPRECATED: Legacy 2025 modules - use main dqn module instead
// pub mod demo_2025_dqn;
// pub mod config_2025;
// Performance validation
pub mod performance_tests;
pub mod performance_validation;
// Re-export core DQN types for public usage
pub use action_space::{ExposureLevel, FactoredAction, OrderType, Urgency};
pub use order_router::OrderRouter;
pub use agent::{AgentMetrics, DQNAgent, TradingState};
pub use dqn::{GradientResult, DQN, DQNConfig};
pub use experience::{Experience, ExperienceBatch};
pub use nstep_buffer::NStepBuffer;
pub use portfolio_tracker::PortfolioTracker;
pub use replay_buffer::{ReplayBuffer, ReplayBufferConfig, ReplayBufferStats};
pub use trade_executor::{
ExecutionCostConfig, ExecutionResult, RejectionReason, RiskControlConfig, TradeExecutor,
};
// Re-export bridge types
pub use trainable_adapter::DQNTrainableAdapter;
// Re-export network components
pub use network::{QNetwork, QNetworkConfig};
// Re-export reward components
pub use reward::{MarketData, RewardConfig, RewardFunction, RiskMetrics};
// Re-export Rainbow DQN components
pub use distributional::{CategoricalDistribution, DistributionalConfig, DistributionalType};
pub use distributional_dueling::{DistributionalDuelingConfig, DistributionalDuelingQNetwork}; // Wave 7.3
pub use dueling::{DuelingConfig, DuelingQNetwork}; // Wave 2.1
pub use multi_step::MultiStepConfig;
pub use quantile_regression::{QuantileConfig, QuantileNetwork, quantile_huber_loss}; // Wave 26 P1.13
pub use rainbow_agent::RainbowAgent;
pub use rainbow_config::{RainbowAgentConfig, RainbowAgentMetrics, RainbowDQNConfig};
pub use rainbow_network::{RainbowNetwork, RainbowNetworkConfig};
// Re-export prioritized replay components
pub use prioritized_replay::{PrioritizedReplayBuffer, PrioritizedReplayConfig};
pub use replay_buffer_type::{BatchSample, ReplayBufferType};
// Re-export regime-conditional DQN components
pub use regime_conditional::{RegimeConditionalDQN, RegimeMetrics, RegimeType};
// Re-export logit clipping utilities (Bug #12 fix)
pub use logit_clipping::{clip_logits, clip_logits_default, softmax_with_clipping, DEFAULT_CLIP_MAX};
// Re-export GAE components (Wave 26 P1.9)
pub use gae::{GAECalculator, GAEConfig};
// Wave 26 P2.3: Ensemble Q-network for uncertainty estimation
pub mod ensemble_network;
// Wave 26 P2.4: RMSNorm - ~15% faster alternative to LayerNorm
pub mod rmsnorm;
// Re-export Hindsight Experience Replay components (Wave 26 P1.7)
pub use hindsight_replay::{
HindsightReplayBuffer, HindsightReplayConfig, HindsightReplayStats, HindsightStrategy,
};
// Re-export ensemble network components (Wave 26 P2.3)
pub use ensemble_network::{EnsembleConfig, EnsembleQNetwork};
// Re-export RMSNorm components (Wave 26 P2.4)
pub use rmsnorm::{LayerNorm, NormType, RMSNorm};
// Re-export noisy exploration components
// TEMPORARILY COMMENTED OUT - Missing implementations
// pub use noisy_exploration::{NoisyExplorationConfig, AdaptiveNoisyManager, AdaptiveNoisyLinear, NoiseExplorationMetrics};
// Re-export performance validation utilities
// TEMPORARILY COMMENTED OUT - Missing implementations
// pub use performance_tests::{
// RainbowPerformanceValidator, PerformanceTestConfig, PerformanceResults,
// validate_rainbow_performance
// };
// Re-export comprehensive performance validation
// TEMPORARILY COMMENTED OUT - Missing implementations
// pub use performance_validation::{
// DQNPerformanceValidator, PerformanceValidationConfig, PerformanceValidationResults,
// validate_dqn_performance
// };
// DEPRECATED: Wave 26 2025-optimized configurations are now in main dqn module
// pub mod config_2025;
//
// // Re-export 2025 config functions (DEPRECATED)
// pub use config_2025::{
// dqn_config_2025, dqn_config_2025_aggressive, dqn_config_2025_conservative,
// dqn_config_2025_hft,
// };
// Re-export offline RL components
pub use experience_dataset::ExperienceDataset;
pub use iql::IqlConfig;
// Re-export DQN demo functionality
// DO NOT RE-EXPORT - Use explicit imports at usage sites

280
crates/ml/src/dqn/regime_temperature.rs
View File

@@ -1,280 +0,0 @@
//! Regime-Aware Temperature Adaptation for DQN
//!
//! This module provides regime-aware temperature adaptation to improve exploration-exploitation
//! balance in different market conditions:
//!
//! - **Trending Markets**: Lower temperature (0.8x) to exploit trend continuation
//! - **Ranging Markets**: Higher temperature (1.2x) to explore breakout opportunities
//! - **Volatile Markets**: High temperature (1.5x) for cautious high-exploration strategy
//! - **Normal Markets**: Baseline temperature (1.0x) for default behavior
//!
//! ## Architecture
//!
//! ```text
//! RegimeOrchestrator → Current Regime → Temperature Multiplier → Adjusted Temperature
//! ↓
//! (Trending, Ranging, Volatile, Normal)
//! ```
//!
//! ## Integration with DQN
//!
//! 1. DQNTrainer queries RegimeOrchestrator for current regime
//! 2. Regime-specific multiplier is retrieved from configuration
//! 3. Base temperature (from exponential decay) is scaled by multiplier
//! 4. Adjusted temperature is used for softmax action selection
//!
//! ## Configuration
//!
//! Regime multipliers are configurable via `DQNHyperparameters`:
//!
//! ```rust
//! use std::collections::HashMap;
//!
//! let mut regime_multipliers = HashMap::new();
//! regime_multipliers.insert("Trending".to_owned(), 0.8);
//! regime_multipliers.insert("Ranging".to_owned(), 1.2);
//! regime_multipliers.insert("Volatile".to_owned(), 1.5);
//! regime_multipliers.insert("Normal".to_owned(), 1.0);
//! ```
use std::collections::HashMap;
/// Default regime temperature multipliers based on adaptive temperature research
///
/// These multipliers are calibrated for HFT trend-following strategies:
///
/// - **Trending (0.8x)**: Lower temperature exploits trend continuation. Lower exploration
/// prevents counter-trend actions that would hurt P&L during strong directional moves.
///
/// - **Ranging (1.2x)**: Higher temperature explores breakout opportunities. Ranging markets
/// require more exploration to identify when consolidation will resolve into a trend.
///
/// - **Volatile (1.5x)**: High temperature provides cautious high-exploration. Volatile regimes
/// have unpredictable price action, so higher exploration prevents over-committing to any
/// single directional bias.
///
/// - **Normal (1.0x)**: Baseline temperature for default/ambiguous market conditions.
///
/// # Returns
///
/// HashMap mapping regime names to temperature multipliers
///
/// # Example
///
/// ```
/// use ml::dqn::regime_temperature::get_default_regime_multipliers;
///
/// let multipliers = get_default_regime_multipliers();
/// assert_eq!(*multipliers.get("Trending").unwrap(), 0.8);
/// assert_eq!(*multipliers.get("Ranging").unwrap(), 1.2);
/// assert_eq!(*multipliers.get("Volatile").unwrap(), 1.5);
/// ```
pub fn get_default_regime_multipliers() -> HashMap<String, f64> {
let mut multipliers = HashMap::new();
multipliers.insert("Trending".to_owned(), 0.8);
multipliers.insert("Ranging".to_owned(), 1.2);
multipliers.insert("Volatile".to_owned(), 1.5);
multipliers.insert("Normal".to_owned(), 1.0);
multipliers
}
/// Apply regime-specific temperature multiplier to base temperature
///
/// Scales the base temperature (from exponential decay) by a regime-specific multiplier
/// to adapt exploration-exploitation balance to current market conditions.
///
/// # Arguments
///
/// * `base_temp` - Base temperature from exponential decay (typically 0.1-1.0)
/// * `regime` - Current market regime (Trending, Ranging, Volatile, Normal)
/// * `multipliers` - Regime-specific multipliers (from configuration)
///
/// # Returns
///
/// Adjusted temperature scaled by regime multiplier
///
/// # Behavior
///
/// - If regime exists in multipliers: `base_temp * multiplier`
/// - If regime unknown: Falls back to "Normal" (1.0x)
/// - If "Normal" missing: Returns `base_temp` unchanged
///
/// # Example
///
/// ```
/// use ml::dqn::regime_temperature::{apply_regime_temperature, get_default_regime_multipliers};
///
/// let base_temp = 1.0;
/// let multipliers = get_default_regime_multipliers();
///
/// // Trending: 0.8x multiplier
/// let trending_temp = apply_regime_temperature(base_temp, "Trending", &multipliers);
/// assert!((trending_temp - 0.8).abs() < 0.01);
///
/// // Ranging: 1.2x multiplier
/// let ranging_temp = apply_regime_temperature(base_temp, "Ranging", &multipliers);
/// assert!((ranging_temp - 1.2).abs() < 0.01);
///
/// // Unknown regime: fallback to Normal (1.0x)
/// let unknown_temp = apply_regime_temperature(base_temp, "UnknownRegime", &multipliers);
/// assert!((unknown_temp - 1.0).abs() < 0.01);
/// ```
pub fn apply_regime_temperature(
base_temp: f64,
regime: &str,
multipliers: &HashMap<String, f64>,
) -> f64 {
let multiplier = multipliers
.get(regime)
.or_else(|| multipliers.get("Normal"))
.unwrap_or(&1.0);
base_temp * multiplier
}
/// Calculate adaptive temperature with regime awareness
///
/// This is the main entry point for regime-aware temperature adaptation. It combines:
/// 1. Base temperature from exponential decay
/// 2. Regime-specific multiplier
/// 3. Bounds checking (respects min/max temperature limits)
///
/// # Arguments
///
/// * `base_temp` - Base temperature from exponential decay
/// * `regime` - Current market regime from RegimeOrchestrator
/// * `multipliers` - Regime-specific multipliers configuration
/// * `min_temp` - Minimum temperature bound (e.g., 0.1)
/// * `max_temp` - Maximum temperature bound (e.g., 2.0)
///
/// # Returns
///
/// Adjusted temperature clamped to [min_temp, max_temp]
///
/// # Example
///
/// ```
/// use ml::dqn::regime_temperature::{calculate_adaptive_temperature, get_default_regime_multipliers};
///
/// let base_temp = 1.0;
/// let regime = "Volatile";
/// let multipliers = get_default_regime_multipliers();
/// let min_temp = 0.1;
/// let max_temp = 2.0;
///
/// let adaptive_temp = calculate_adaptive_temperature(
/// base_temp, regime, &multipliers, min_temp, max_temp
/// );
///
/// // Volatile: 1.5x multiplier
/// assert!((adaptive_temp - 1.5).abs() < 0.01);
/// assert!(adaptive_temp >= min_temp && adaptive_temp <= max_temp);
/// ```
pub fn calculate_adaptive_temperature(
base_temp: f64,
regime: &str,
multipliers: &HashMap<String, f64>,
min_temp: f64,
max_temp: f64,
) -> f64 {
let adjusted_temp = apply_regime_temperature(base_temp, regime, multipliers);
adjusted_temp.clamp(min_temp, max_temp)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_multipliers() {
let multipliers = get_default_regime_multipliers();
assert_eq!(*multipliers.get("Trending").unwrap(), 0.8);
assert_eq!(*multipliers.get("Ranging").unwrap(), 1.2);
assert_eq!(*multipliers.get("Volatile").unwrap(), 1.5);
assert_eq!(*multipliers.get("Normal").unwrap(), 1.0);
}
#[test]
fn test_apply_regime_temperature_trending() {
let multipliers = get_default_regime_multipliers();
let base_temp = 1.0;
let adjusted = apply_regime_temperature(base_temp, "Trending", &multipliers);
assert!((adjusted - 0.8).abs() < 0.001);
}
#[test]
fn test_apply_regime_temperature_ranging() {
let multipliers = get_default_regime_multipliers();
let base_temp = 1.0;
let adjusted = apply_regime_temperature(base_temp, "Ranging", &multipliers);
assert!((adjusted - 1.2).abs() < 0.001);
}
#[test]
fn test_apply_regime_temperature_volatile() {
let multipliers = get_default_regime_multipliers();
let base_temp = 1.0;
let adjusted = apply_regime_temperature(base_temp, "Volatile", &multipliers);
assert!((adjusted - 1.5).abs() < 0.001);
}
#[test]
fn test_apply_regime_temperature_unknown_fallback() {
let multipliers = get_default_regime_multipliers();
let base_temp = 1.0;
let adjusted = apply_regime_temperature(base_temp, "UnknownRegime", &multipliers);
assert!((adjusted - 1.0).abs() < 0.001); // Falls back to Normal (1.0)
}
#[test]
fn test_calculate_adaptive_temperature_clamping() {
let multipliers = get_default_regime_multipliers();
let base_temp = 0.05; // Below min
let min_temp = 0.1;
let max_temp = 2.0;
// Even with Volatile (1.5x), should clamp to min_temp
let adjusted = calculate_adaptive_temperature(
base_temp, "Volatile", &multipliers, min_temp, max_temp
);
assert!(adjusted >= min_temp);
assert!(adjusted <= max_temp);
}
#[test]
fn test_calculate_adaptive_temperature_max_clamping() {
let multipliers = get_default_regime_multipliers();
let base_temp = 1.5;
let min_temp = 0.1;
let max_temp = 2.0;
// Volatile (1.5x) on high base_temp should clamp to max_temp
let adjusted = calculate_adaptive_temperature(
base_temp, "Volatile", &multipliers, min_temp, max_temp
);
assert!(adjusted <= max_temp);
}
#[test]
fn test_custom_multipliers() {
let mut custom_multipliers = HashMap::new();
custom_multipliers.insert("Trending".to_owned(), 0.5);
custom_multipliers.insert("Ranging".to_owned(), 2.0);
custom_multipliers.insert("Normal".to_owned(), 1.0);
let base_temp = 1.0;
let trending = apply_regime_temperature(base_temp, "Trending", &custom_multipliers);
assert!((trending - 0.5).abs() < 0.001);
let ranging = apply_regime_temperature(base_temp, "Ranging", &custom_multipliers);
assert!((ranging - 2.0).abs() < 0.001);
}
}

354
crates/ml/src/dqn/risk_integration.rs
View File

@@ -1,354 +0,0 @@
//! Risk Crate Circuit Breaker Integration for DQN Training
//!
//! Integrates the risk crate's RealCircuitBreaker with DQN training to prevent
//! runaway losses. Uses Redis coordination for multi-process safety.
use async_trait::async_trait;
use risk::circuit_breaker::{BrokerAccountService, CircuitBreakerConfig, RealCircuitBreaker};
use risk::error::{RiskError, RiskResult};
use rust_decimal::Decimal;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use common::{Position, Price};
/// Training broker service - provides portfolio metrics from DQN training
///
/// This mock implementation tracks portfolio value and P&L from training rewards
/// without requiring a real broker connection.
#[derive(Debug)]
pub struct TrainingBrokerService {
/// Current portfolio value (updated from training metrics)
portfolio_value: Arc<RwLock<Decimal>>,
/// Daily P&L accumulator (reset each epoch)
daily_pnl: Arc<RwLock<Decimal>>,
/// Current positions (empty for training)
positions: Arc<RwLock<Vec<Position>>>,
}
impl TrainingBrokerService {
/// Create a new training broker service
///
/// # Arguments
/// * `initial_capital` - Starting capital in dollars (e.g., 100,000.0)
pub fn new(initial_capital: f64) -> Self {
let capital = Decimal::try_from(initial_capital).unwrap_or(Decimal::new(100_000, 0));
info!("TrainingBrokerService initialized with capital: ${}", capital);
Self {
portfolio_value: Arc::new(RwLock::new(capital)),
daily_pnl: Arc::new(RwLock::new(Decimal::ZERO)),
positions: Arc::new(RwLock::new(Vec::new())),
}
}
/// Update portfolio value from training metrics
pub async fn update_portfolio_value(&self, value: f64) {
let current = *self.portfolio_value.read().await;
let value_decimal = Decimal::try_from(value).unwrap_or_else(|e| {
debug!("Failed to convert portfolio value {} to Decimal: {}, keeping previous value", value, e);
current
});
let mut pv = self.portfolio_value.write().await;
*pv = value_decimal;
debug!("Portfolio value updated to: ${}", value_decimal);
}
/// Record a reward (adds to daily P&L)
///
/// # Arguments
/// * `reward` - Reward value (can be positive or negative)
pub async fn record_reward(&self, reward: f64) {
let reward_decimal = Decimal::try_from(reward).unwrap_or_else(|e| {
warn!("Failed to convert reward {} to Decimal: {}, using 0", reward, e);
Decimal::ZERO
});
let mut pnl = self.daily_pnl.write().await;
*pnl += reward_decimal;
if reward < 0.0 {
debug!("Loss recorded: ${:.2}, Daily P&L: ${:.2}", -reward, *pnl);
}
}
/// Reset daily P&L (call at start of each epoch)
pub async fn reset_daily_pnl(&self) {
let mut pnl = self.daily_pnl.write().await;
*pnl = Decimal::ZERO;
debug!("Daily P&L reset for new epoch");
}
/// Get current portfolio value (synchronous helper)
pub async fn get_portfolio_value_sync(&self) -> Decimal {
*self.portfolio_value.read().await
}
/// Get current daily P&L (synchronous helper)
pub async fn get_daily_pnl_sync(&self) -> Decimal {
*self.daily_pnl.read().await
}
}
#[async_trait]
impl BrokerAccountService for TrainingBrokerService {
async fn get_portfolio_value(&self, _account_id: &str) -> RiskResult<Decimal> {
Ok(*self.portfolio_value.read().await)
}
async fn get_daily_pnl(&self, _account_id: &str) -> RiskResult<Decimal> {
Ok(*self.daily_pnl.read().await)
}
async fn get_positions(&self, _account_id: &str) -> RiskResult<Vec<Position>> {
Ok(self.positions.read().await.clone())
}
}
/// DQN Circuit Breaker - wraps risk crate's RealCircuitBreaker for training
pub struct DQNRiskCircuitBreaker {
/// The actual circuit breaker from risk crate
breaker: RealCircuitBreaker,
/// Broker service providing portfolio metrics
broker_service: Arc<TrainingBrokerService>,
/// Account ID for circuit breaker state tracking
account_id: String,
}
impl std::fmt::Debug for DQNRiskCircuitBreaker {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DQNRiskCircuitBreaker")
.field("broker_service", &self.broker_service)
.field("account_id", &self.account_id)
.field("breaker", &"RealCircuitBreaker { ... }")
.finish()
}
}
impl DQNRiskCircuitBreaker {
/// Create a new circuit breaker for DQN training
///
/// # Arguments
/// * `initial_capital` - Starting capital in dollars (e.g., 100,000.0)
/// * `loss_threshold` - Maximum daily loss in dollars (e.g., 10,000.0 = $10K)
///
/// # Returns
/// Result with initialized circuit breaker
///
/// # Errors
/// Returns error if:
/// - Redis connection fails
/// - Configuration is invalid
pub async fn new(initial_capital: f64, loss_threshold: f64) -> RiskResult<Self> {
info!("🔒 Initializing DQN Risk Circuit Breaker");
info!(" Initial Capital: ${:.2}", initial_capital);
info!(" Loss Threshold: ${:.2} ({:.2}% of capital)",
loss_threshold,
(loss_threshold / initial_capital) * 100.0
);
// Create broker service
let broker_service = Arc::new(TrainingBrokerService::new(initial_capital));
// Calculate loss threshold as percentage of capital
let loss_pct = (loss_threshold / initial_capital) * 100.0;
// Configure circuit breaker
let config = CircuitBreakerConfig {
enabled: true,
daily_loss_percentage: Price::from_f64(loss_pct).map_err(|e| {
RiskError::TypeConversion {
from_type: "f64".to_owned(),
to_type: "Price".to_owned(),
reason: format!("loss percentage conversion failed: {}", e),
}
})?,
position_limit_percentage: Price::from_f64(5.0).map_err(|e| {
RiskError::TypeConversion {
from_type: "f64".to_owned(),
to_type: "Price".to_owned(),
reason: format!("position limit conversion failed: {}", e),
}
})?,
max_consecutive_violations: 5,
redis_url: std::env::var("REDIS_URL")
.unwrap_or_else(|_| "redis://localhost:6379".to_owned()),
redis_key_prefix: "foxhunt:dqn_training:circuit_breaker".to_owned(),
auto_recovery_enabled: false, // Manual recovery for training safety
portfolio_refresh_interval_secs: 60,
cooldown_period_secs: 300, // 5 minutes
};
info!("📡 Connecting to Redis at {}", config.redis_url);
// Create circuit breaker
let breaker = RealCircuitBreaker::new(config, broker_service.clone()).await?;
info!("✅ Circuit breaker initialized successfully");
Ok(Self {
breaker,
broker_service,
account_id: "dqn_training".to_owned(),
})
}
/// Check if circuit breaker is open (trading halted)
///
/// # Returns
/// - `true` if circuit breaker is OPEN (trading blocked)
/// - `false` if circuit breaker is CLOSED (trading allowed)
pub async fn is_open(&self) -> bool {
self.breaker.is_active(&self.account_id).await
}
/// Check circuit breaker and potentially trigger it
///
/// # Returns
/// Result with bool indicating if circuit breaker is open
pub async fn check(&self) -> RiskResult<bool> {
self.breaker.check_circuit_breaker(&self.account_id).await
}
/// Record a reward and check circuit breaker
///
/// Updates the broker service's daily P&L and checks if the circuit breaker
/// should be triggered based on accumulated losses.
///
/// # Arguments
/// * `reward` - Reward value (negative = loss, positive = profit)
pub async fn record_reward(&self, reward: f64) -> RiskResult<()> {
// Update broker service
self.broker_service.record_reward(reward).await;
// Check circuit breaker on losses
if reward < 0.0 {
let is_open = self.check().await?;
if is_open {
error!("⚠️ Circuit breaker TRIGGERED after loss of ${:.2}", -reward);
let state = self.breaker.get_state(&self.account_id).await?;
error!(" Reason: {}", state.activation_reason.unwrap_or_else(|| "Unknown".to_owned()));
error!(" Daily Loss: ${:.2} / ${:.2}", state.current_daily_loss, state.daily_loss_limit);
}
}
Ok(())
}
/// Update portfolio value from training metrics
///
/// # Arguments
/// * `value` - Current portfolio value in dollars
pub async fn update_portfolio_value(&self, value: f64) {
self.broker_service.update_portfolio_value(value).await;
}
/// Reset daily P&L (call at start of each epoch)
pub async fn reset_daily_pnl(&self) {
self.broker_service.reset_daily_pnl().await;
}
/// Manually reset circuit breaker to closed state
///
/// # Arguments
/// * `reason` - Reason for manual reset (logged)
pub async fn reset(&self, reason: String) -> RiskResult<()> {
info!("🔓 Manually resetting circuit breaker: {}", reason);
self.breaker.reset_circuit_breaker(&self.account_id, reason).await
}
/// Get current circuit breaker state
pub async fn get_state(&self) -> RiskResult<risk::circuit_breaker::CircuitBreakerState> {
self.breaker.get_state(&self.account_id).await
}
/// Get circuit breaker metrics
pub async fn get_metrics(&self) -> std::collections::HashMap<String, f64> {
self.breaker.get_metrics().await
}
/// Health check - verifies Redis connectivity
pub async fn health_check(&self) -> bool {
self.breaker.health_check().await
}
/// Get current portfolio value
pub async fn get_portfolio_value(&self) -> Decimal {
self.broker_service.get_portfolio_value_sync().await
}
/// Get current daily P&L
pub async fn get_daily_pnl(&self) -> Decimal {
self.broker_service.get_daily_pnl_sync().await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_training_broker_service() {
let service = TrainingBrokerService::new(100_000.0);
// Test initial values
let pv = service.get_portfolio_value_sync().await;
assert_eq!(pv, Decimal::new(100_000, 0));
let pnl = service.get_daily_pnl_sync().await;
assert_eq!(pnl, Decimal::ZERO);
// Test recording profits
service.record_reward(500.0).await;
let pnl = service.get_daily_pnl_sync().await;
assert_eq!(pnl, Decimal::new(500, 0));
// Test recording losses
service.record_reward(-200.0).await;
let pnl = service.get_daily_pnl_sync().await;
assert_eq!(pnl, Decimal::new(300, 0));
// Test reset
service.reset_daily_pnl().await;
let pnl = service.get_daily_pnl_sync().await;
assert_eq!(pnl, Decimal::ZERO);
}
#[tokio::test]
async fn test_circuit_breaker_creation() {
// Skip if Redis not available
let result = DQNRiskCircuitBreaker::new(100_000.0, 10_000.0).await;
match result {
Ok(cb) => {
// Verify health check
let healthy = cb.health_check().await;
println!("Circuit breaker healthy: {}", healthy);
// Verify initial state
assert!(!cb.is_open().await);
}
Err(e) => {
println!("Circuit breaker creation failed (expected if Redis unavailable): {}", e);
// This is okay in CI/test environments without Redis
}
}
}
#[tokio::test]
async fn test_broker_service_trait() {
let service = TrainingBrokerService::new(100_000.0);
// Test trait methods
let pv = service.get_portfolio_value("test_account").await.unwrap();
assert_eq!(pv, Decimal::new(100_000, 0));
let pnl = service.get_daily_pnl("test_account").await.unwrap();
assert_eq!(pnl, Decimal::ZERO);
let positions = service.get_positions("test_account").await.unwrap();
assert!(positions.is_empty());
}
}

398
crates/ml/src/dqn/spectral_norm.rs
View File

@@ -1,398 +0,0 @@
//! Spectral Normalization for Q-Value Stability
//!
//! Implements spectral normalization (Miyato et al., 2018) to prevent Q-value divergence
//! by constraining the Lipschitz constant of the network layers to 1.
//!
//! Key benefits:
//! - Prevents gradient explosion/vanishing
//! - Stabilizes Q-value estimates
//! - Improves training convergence
//! - Reduces need for gradient clipping
use candle_core::{DType, Device, Result as CandleResult, Tensor};
use candle_nn::{Linear, Module, VarBuilder};
use serde::{Deserialize, Serialize};
use crate::MLError;
/// Configuration for spectral normalization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpectralNormConfig {
/// Number of power iterations for spectral norm estimation
/// Higher values = more accurate estimation but slower
/// Typical range: 1-5, default 1 is usually sufficient
pub n_power_iterations: usize,
/// Small constant for numerical stability
pub eps: f64,
}
impl Default for SpectralNormConfig {
fn default() -> Self {
Self {
n_power_iterations: 1, // Miyato et al. found 1 iteration sufficient
eps: 1e-12,
}
}
}
/// Spectral normalization wrapper for linear layers
///
/// Applies spectral normalization to constrain the spectral norm (largest singular value)
/// of the weight matrix to approximately 1, which bounds the Lipschitz constant.
#[derive(Debug)]
pub struct SpectralNorm {
/// The underlying linear layer
linear: Linear,
/// Left singular vector (updated via power iteration)
u: Tensor,
/// Right singular vector (updated via power iteration)
v: Tensor,
/// Configuration
config: SpectralNormConfig,
/// Device for tensor operations
device: Device,
}
impl SpectralNorm {
/// Create a new spectral-normalized linear layer
///
/// # Arguments
/// * `in_dim` - Input dimension
/// * `out_dim` - Output dimension
/// * `vs` - Variable builder for weight initialization
/// * `config` - Spectral norm configuration
pub fn new(
in_dim: usize,
out_dim: usize,
vs: VarBuilder<'_>,
config: SpectralNormConfig,
) -> Result<Self, MLError> {
let device = vs.device().clone();
// Create underlying linear layer
let linear = candle_nn::linear(in_dim, out_dim, vs)
.map_err(|e| MLError::ModelError(format!("Failed to create linear layer: {}", e)))?;
// Initialize singular vectors randomly
// u: [out_dim], v: [in_dim]
let u = Tensor::randn(0_f32, 1_f32, out_dim, &device)
.map_err(|e| MLError::ModelError(format!("Failed to initialize u vector: {}", e)))?;
let v = Tensor::randn(0_f32, 1_f32, in_dim, &device)
.map_err(|e| MLError::ModelError(format!("Failed to initialize v vector: {}", e)))?;
Ok(Self {
linear,
u,
v,
config,
device,
})
}
/// Estimate spectral norm using power iteration
///
/// The spectral norm is the largest singular value σ₁(W) of the weight matrix W.
/// Power iteration approximates this by iteratively computing:
/// v_{k+1} = W^T u_k / ||W^T u_k||
/// u_{k+1} = W v_{k+1} / ||W v_{k+1}||
/// σ ≈ u^T W v
fn compute_spectral_norm(&mut self) -> Result<f32, MLError> {
// Get weight matrix from linear layer
let weight = self.linear.weight();
let mut u = self.u.clone();
let mut v = self.v.clone();
// Power iteration
for _ in 0..self.config.n_power_iterations {
// v = W^T u / ||W^T u||
v = weight
.t()
.map_err(|e| MLError::ModelError(format!("Failed to transpose weight: {}", e)))?
.matmul(&u.unsqueeze(1).map_err(|e| {
MLError::ModelError(format!("Failed to unsqueeze u: {}", e))
})?)
.map_err(|e| MLError::ModelError(format!("Failed to compute W^T u: {}", e)))?
.squeeze(1)
.map_err(|e| MLError::ModelError(format!("Failed to squeeze v: {}", e)))?;
// Normalize v
let v_norm = v
.sqr()
.map_err(|e| MLError::ModelError(format!("Failed to square v: {}", e)))?
.sum_all()
.map_err(|e| MLError::ModelError(format!("Failed to sum v²: {}", e)))?
.sqrt()
.map_err(|e| MLError::ModelError(format!("Failed to sqrt v norm: {}", e)))?
.to_vec0::<f32>()
.map_err(|e| MLError::ModelError(format!("Failed to extract v norm: {}", e)))?
.max(self.config.eps as f32);
v = v
.affine(1.0 / v_norm as f64, 0.0)
.map_err(|e| MLError::ModelError(format!("Failed to normalize v: {}", e)))?;
// u = W v / ||W v||
u = weight
.matmul(&v.unsqueeze(1).map_err(|e| {
MLError::ModelError(format!("Failed to unsqueeze v: {}", e))
})?)
.map_err(|e| MLError::ModelError(format!("Failed to compute W v: {}", e)))?
.squeeze(1)
.map_err(|e| MLError::ModelError(format!("Failed to squeeze u: {}", e)))?;
// Normalize u
let u_norm = u
.sqr()
.map_err(|e| MLError::ModelError(format!("Failed to square u: {}", e)))?
.sum_all()
.map_err(|e| MLError::ModelError(format!("Failed to sum u²: {}", e)))?
.sqrt()
.map_err(|e| MLError::ModelError(format!("Failed to sqrt u norm: {}", e)))?
.to_vec0::<f32>()
.map_err(|e| MLError::ModelError(format!("Failed to extract u norm: {}", e)))?
.max(self.config.eps as f32);
u = u
.affine(1.0 / u_norm as f64, 0.0)
.map_err(|e| MLError::ModelError(format!("Failed to normalize u: {}", e)))?;
}
// Update stored singular vectors
self.u = u.clone();
self.v = v.clone();
// Compute spectral norm: σ = u^T W v
let wv = weight
.matmul(&self.v.unsqueeze(1).map_err(|e| {
MLError::ModelError(format!("Failed to unsqueeze v for sigma: {}", e))
})?)
.map_err(|e| MLError::ModelError(format!("Failed to compute W v for sigma: {}", e)))?
.squeeze(1)
.map_err(|e| MLError::ModelError(format!("Failed to squeeze W v: {}", e)))?;
let sigma = self
.u
.mul(&wv)
.map_err(|e| MLError::ModelError(format!("Failed to compute u * W v: {}", e)))?
.sum_all()
.map_err(|e| MLError::ModelError(format!("Failed to sum sigma: {}", e)))?
.to_vec0::<f32>()
.map_err(|e| MLError::ModelError(format!("Failed to extract sigma: {}", e)))?;
Ok(sigma.max(self.config.eps as f32))
}
/// Get the normalized weight matrix W_norm = W / σ(W)
fn normalized_weight(&mut self) -> Result<Tensor, MLError> {
let sigma = self.compute_spectral_norm()?;
let weight = self.linear.weight();
weight
.affine(1.0 / sigma as f64, 0.0)
.map_err(|e| MLError::ModelError(format!("Failed to normalize weight: {}", e)))
}
/// Reset singular vectors (useful when reloading checkpoints)
pub fn reset_singular_vectors(&mut self) -> Result<(), MLError> {
let weight = self.linear.weight();
let shape = weight.shape();
self.u = Tensor::randn(0_f32, 1_f32, shape.dims()[0], &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to reset u vector: {}", e)))?;
self.v = Tensor::randn(0_f32, 1_f32, shape.dims()[1], &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to reset v vector: {}", e)))?;
Ok(())
}
/// Get current spectral norm estimate
pub fn get_spectral_norm(&mut self) -> Result<f32, MLError> {
self.compute_spectral_norm()
}
}
impl Module for SpectralNorm {
fn forward(&self, xs: &Tensor) -> CandleResult<Tensor> {
// Note: In training mode, we should update the normalized weight
// For now, we use the stored linear layer directly
// In a full implementation, this would compute normalized_weight() and apply it
self.linear.forward(xs)
}
}
#[cfg(test)]
mod tests {
use super::*;
use candle_nn::VarMap;
use crate::dqn::mixed_precision::training_dtype;
#[test]
fn test_spectral_norm_creation() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
let config = SpectralNormConfig::default();
let spectral_norm = SpectralNorm::new(10, 5, vs.pp("test"), config)?;
assert_eq!(spectral_norm.u.dims()[0], 5);
assert_eq!(spectral_norm.v.dims()[0], 10);
Ok(())
}
#[test]
fn test_spectral_norm_computation() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
let config = SpectralNormConfig::default();
let mut spectral_norm = SpectralNorm::new(10, 5, vs.pp("test"), config)?;
// Compute spectral norm
let sigma = spectral_norm.get_spectral_norm()?;
// Spectral norm should be positive
assert!(sigma > 0.0, "Spectral norm should be positive");
// For a randomly initialized weight matrix, spectral norm should be reasonable
assert!(sigma < 10.0, "Spectral norm should be bounded");
Ok(())
}
#[test]
fn test_spectral_norm_bounds_lipschitz() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
let config = SpectralNormConfig {
n_power_iterations: 5, // More iterations for accuracy
eps: 1e-12,
};
let mut spectral_norm = SpectralNorm::new(10, 5, vs.pp("test"), config)?;
// Get normalized weight
let normalized_weight = spectral_norm.normalized_weight()?;
// The spectral norm of the normalized weight should be ≈ 1
// We can verify this by checking the Frobenius norm is reasonable
let weight_norm = normalized_weight
.sqr()?
.sum_all()?
.sqrt()?
.to_vec0::<f32>()?;
// For a 10x5 matrix with spectral norm 1, Frobenius norm should be ≤ √50 ≈ 7.07
assert!(
weight_norm < 10.0,
"Normalized weight Frobenius norm should be bounded"
);
Ok(())
}
#[test]
fn test_power_iteration_convergence() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
// Test with different iteration counts
for n_iters in [1, 2, 5] {
let config = SpectralNormConfig {
n_power_iterations: n_iters,
eps: 1e-12,
};
let mut spectral_norm = SpectralNorm::new(10, 5, vs.pp(&format!("test_{}", n_iters)), config)?;
let sigma = spectral_norm.get_spectral_norm()?;
// All should give reasonable estimates
assert!(sigma > 0.0 && sigma < 20.0, "Spectral norm estimate should be reasonable");
}
Ok(())
}
#[test]
fn test_singular_vector_reset() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
let config = SpectralNormConfig::default();
let mut spectral_norm = SpectralNorm::new(10, 5, vs.pp("test"), config)?;
// Compute initial spectral norm
let sigma1 = spectral_norm.get_spectral_norm()?;
// Reset vectors
spectral_norm.reset_singular_vectors()?;
// Compute again (should be different due to random initialization)
let sigma2 = spectral_norm.get_spectral_norm()?;
// Both should be positive and reasonable
assert!(sigma1 > 0.0 && sigma2 > 0.0);
assert!(sigma1 < 20.0 && sigma2 < 20.0);
Ok(())
}
#[test]
fn test_forward_pass() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
let config = SpectralNormConfig::default();
let spectral_norm = SpectralNorm::new(10, 5, vs.pp("test"), config)?;
// Create input tensor
let input = Tensor::randn(0_f32, 1_f32, (2, 10), &device)?; // Batch size 2
// Forward pass
let output = spectral_norm.forward(&input)?;
// Check output shape
assert_eq!(output.dims(), &[2, 5]);
Ok(())
}
#[test]
fn test_prevents_weight_explosion() -> anyhow::Result<()> {
let device = Device::Cpu;
let varmap = VarMap::new();
let vs = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
let config = SpectralNormConfig::default();
let mut spectral_norm = SpectralNorm::new(10, 5, vs.pp("test"), config)?;
// Get original weight norm
let original_weight = spectral_norm.linear.weight();
let original_norm = original_weight.sqr()?.sum_all()?.sqrt()?.to_vec0::<f32>()?;
// Get normalized weight
let normalized_weight = spectral_norm.normalized_weight()?;
let normalized_norm = normalized_weight.sqr()?.sum_all()?.sqrt()?.to_vec0::<f32>()?;
// Normalized weight should have smaller or equal norm
assert!(
normalized_norm <= original_norm * 1.1, // Allow small numerical tolerance
"Spectral normalization should not increase weight magnitude"
);
Ok(())
}
}