From 88f0b3ec231ee66d7dd0bc1e15f8022c552cf963 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 8 Mar 2026 01:16:40 +0100 Subject: [PATCH] refactor(ml): extract DQN module into ml-dqn crate (task 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- Cargo.lock | 26 + crates/ml-dqn/Cargo.toml | 46 +- crates/{ml/src/dqn => ml-dqn/src}/agent.rs | 75 +- .../{ml/src/dqn => ml-dqn/src}/attention.rs | 8 +- .../src/dqn => ml-dqn/src}/circuit_breaker.rs | 2 +- .../{ml/src/dqn => ml-dqn/src}/count_bonus.rs | 0 .../{ml/src/dqn => ml-dqn/src}/curiosity.rs | 4 +- crates/{ml/src/dqn => ml-dqn/src}/demo_dqn.rs | 8 +- .../src/dqn => ml-dqn/src}/distributional.rs | 2 +- .../src}/distributional_dueling.rs | 10 +- crates/{ml/src/dqn => ml-dqn/src}/dqn.rs | 56 +- crates/{ml/src/dqn => ml-dqn/src}/dueling.rs | 8 +- .../dqn => ml-dqn/src}/ensemble_network.rs | 4 +- .../src}/entropy_regularization.rs | 2 +- .../{ml/src/dqn => ml-dqn/src}/experience.rs | 0 .../dqn => ml-dqn/src}/experience_dataset.rs | 0 crates/{ml/src/dqn => ml-dqn/src}/gae.rs | 0 .../src}/gpu_replay_buffer.rs | 6 +- .../dqn => ml-dqn/src}/hindsight_replay.rs | 2 +- crates/{ml/src/dqn => ml-dqn/src}/iql.rs | 0 crates/ml-dqn/src/lib.rs | 138 ++- crates/{ml/src/dqn => ml-dqn/src}/logging.rs | 0 .../src/dqn => ml-dqn/src}/logit_clipping.rs | 2 +- .../{ml/src/dqn => ml-dqn/src}/multi_asset.rs | 2 +- .../{ml/src/dqn => ml-dqn/src}/multi_step.rs | 2 +- crates/{ml/src/dqn => ml-dqn/src}/network.rs | 14 +- .../dqn => ml-dqn/src}/noisy_exploration.rs | 2 +- .../src/dqn => ml-dqn/src}/noisy_layers.rs | 6 +- .../src}/noisy_sigma_scheduler.rs | 0 .../src/dqn => ml-dqn/src}/nstep_buffer.rs | 2 +- .../dqn => ml-dqn/src}/performance_tests.rs | 2 +- .../src}/performance_validation.rs | 0 .../dqn => ml-dqn/src}/prioritized_replay.rs | 6 +- .../src}/prioritized_replay_staleness.rs | 0 .../dqn => ml-dqn/src}/quantile_regression.rs | 8 +- .../src/dqn => ml-dqn/src}/rainbow_agent.rs | 6 +- .../src/dqn => ml-dqn/src}/rainbow_config.rs | 0 .../dqn => ml-dqn/src}/rainbow_integration.rs | 4 +- .../src/dqn => ml-dqn/src}/rainbow_network.rs | 4 +- .../dqn => ml-dqn/src}/regime_conditional.rs | 12 +- .../src/dqn => ml-dqn/src}/replay_buffer.rs | 2 +- .../dqn => ml-dqn/src}/replay_buffer_type.rs | 16 +- crates/{ml/src/dqn => ml-dqn/src}/residual.rs | 8 +- crates/{ml/src/dqn => ml-dqn/src}/reward.rs | 8 +- crates/{ml/src/dqn => ml-dqn/src}/rmsnorm.rs | 4 +- .../src}/self_supervised_pretraining.rs | 0 crates/{ml/src/dqn => ml-dqn/src}/softmax.rs | 2 +- .../src/dqn => ml-dqn/src}/target_update.rs | 0 .../src/dqn => ml-dqn/src}/trade_executor.rs | 2 +- crates/ml/Cargo.toml | 1 + .../src/checkpoint/model_implementations.rs | 78 +- crates/ml/src/cuda_pipeline/mod.rs | 3 +- crates/ml/src/dqn/data_augmentation.rs | 354 ------ crates/ml/src/dqn/ensemble.rs | 1062 ----------------- crates/ml/src/dqn/ensemble_oracle.rs | 299 ----- crates/ml/src/dqn/ensemble_uncertainty.rs | 897 -------------- crates/ml/src/dqn/factored_q_network.rs | 431 ------- crates/ml/src/dqn/mod.rs | 171 +-- crates/ml/src/dqn/regime_temperature.rs | 280 ----- crates/ml/src/dqn/risk_integration.rs | 354 ------ crates/ml/src/dqn/spectral_norm.rs | 398 ------ 61 files changed, 405 insertions(+), 4434 deletions(-) rename crates/{ml/src/dqn => ml-dqn/src}/agent.rs (94%) rename crates/{ml/src/dqn => ml-dqn/src}/attention.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/circuit_breaker.rs (86%) rename crates/{ml/src/dqn => ml-dqn/src}/count_bonus.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/curiosity.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/demo_dqn.rs (96%) rename crates/{ml/src/dqn => ml-dqn/src}/distributional.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/distributional_dueling.rs (98%) rename crates/{ml/src/dqn => ml-dqn/src}/dqn.rs (98%) rename crates/{ml/src/dqn => ml-dqn/src}/dueling.rs (98%) rename crates/{ml/src/dqn => ml-dqn/src}/ensemble_network.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/entropy_regularization.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/experience.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/experience_dataset.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/gae.rs (100%) rename crates/{ml/src/cuda_pipeline => ml-dqn/src}/gpu_replay_buffer.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/hindsight_replay.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/iql.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/logging.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/logit_clipping.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/multi_asset.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/multi_step.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/network.rs (97%) rename crates/{ml/src/dqn => ml-dqn/src}/noisy_exploration.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/noisy_layers.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/noisy_sigma_scheduler.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/nstep_buffer.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/performance_tests.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/performance_validation.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/prioritized_replay.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/prioritized_replay_staleness.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/quantile_regression.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/rainbow_agent.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/rainbow_config.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/rainbow_integration.rs (96%) rename crates/{ml/src/dqn => ml-dqn/src}/rainbow_network.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/regime_conditional.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/replay_buffer.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/replay_buffer_type.rs (97%) rename crates/{ml/src/dqn => ml-dqn/src}/residual.rs (98%) rename crates/{ml/src/dqn => ml-dqn/src}/reward.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/rmsnorm.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/self_supervised_pretraining.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/softmax.rs (99%) rename crates/{ml/src/dqn => ml-dqn/src}/target_update.rs (100%) rename crates/{ml/src/dqn => ml-dqn/src}/trade_executor.rs (99%) delete mode 100644 crates/ml/src/dqn/data_augmentation.rs delete mode 100644 crates/ml/src/dqn/ensemble.rs delete mode 100644 crates/ml/src/dqn/ensemble_oracle.rs delete mode 100644 crates/ml/src/dqn/ensemble_uncertainty.rs delete mode 100644 crates/ml/src/dqn/factored_q_network.rs delete mode 100644 crates/ml/src/dqn/regime_temperature.rs delete mode 100644 crates/ml/src/dqn/risk_integration.rs delete mode 100644 crates/ml/src/dqn/spectral_norm.rs diff --git a/Cargo.lock b/Cargo.lock index 6a0931727..bbbdbf06c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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]] diff --git a/crates/ml-dqn/Cargo.toml b/crates/ml-dqn/Cargo.toml index 63377b5a4..bcd249d76 100644 --- a/crates/ml-dqn/Cargo.toml +++ b/crates/ml-dqn/Cargo.toml @@ -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 diff --git a/crates/ml/src/dqn/agent.rs b/crates/ml-dqn/src/agent.rs similarity index 94% rename from crates/ml/src/dqn/agent.rs rename to crates/ml-dqn/src/agent.rs index 145f9bd48..344900dd2 100644 --- a/crates/ml/src/dqn/agent.rs +++ b/crates/ml-dqn/src/agent.rs @@ -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> { + 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 diff --git a/crates/ml/src/dqn/attention.rs b/crates/ml-dqn/src/attention.rs similarity index 99% rename from crates/ml/src/dqn/attention.rs rename to crates/ml-dqn/src/attention.rs index 43ff77028..834865ca4 100644 --- a/crates/ml/src/dqn/attention.rs +++ b/crates/ml-dqn/src/attention.rs @@ -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() { diff --git a/crates/ml/src/dqn/circuit_breaker.rs b/crates/ml-dqn/src/circuit_breaker.rs similarity index 86% rename from crates/ml/src/dqn/circuit_breaker.rs rename to crates/ml-dqn/src/circuit_breaker.rs index ef9b85cfa..1cd46b62a 100644 --- a/crates/ml/src/dqn/circuit_breaker.rs +++ b/crates/ml-dqn/src/circuit_breaker.rs @@ -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, }; diff --git a/crates/ml/src/dqn/count_bonus.rs b/crates/ml-dqn/src/count_bonus.rs similarity index 100% rename from crates/ml/src/dqn/count_bonus.rs rename to crates/ml-dqn/src/count_bonus.rs diff --git a/crates/ml/src/dqn/curiosity.rs b/crates/ml-dqn/src/curiosity.rs similarity index 99% rename from crates/ml/src/dqn/curiosity.rs rename to crates/ml-dqn/src/curiosity.rs index 18a452e16..0f48d3f49 100644 --- a/crates/ml/src/dqn/curiosity.rs +++ b/crates/ml-dqn/src/curiosity.rs @@ -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)] diff --git a/crates/ml/src/dqn/demo_dqn.rs b/crates/ml-dqn/src/demo_dqn.rs similarity index 96% rename from crates/ml/src/dqn/demo_dqn.rs rename to crates/ml-dqn/src/demo_dqn.rs index d627e4b9c..100bc73b6 100644 --- a/crates/ml/src/dqn/demo_dqn.rs +++ b/crates/ml-dqn/src/demo_dqn.rs @@ -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 diff --git a/crates/ml/src/dqn/distributional.rs b/crates/ml-dqn/src/distributional.rs similarity index 99% rename from crates/ml/src/dqn/distributional.rs rename to crates/ml-dqn/src/distributional.rs index 205ed6a26..a3e3a2b27 100644 --- a/crates/ml/src/dqn/distributional.rs +++ b/crates/ml-dqn/src/distributional.rs @@ -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)] diff --git a/crates/ml/src/dqn/distributional_dueling.rs b/crates/ml-dqn/src/distributional_dueling.rs similarity index 98% rename from crates/ml/src/dqn/distributional_dueling.rs rename to crates/ml-dqn/src/distributional_dueling.rs index 8bd8c465c..1076276d0 100644 --- a/crates/ml/src/dqn/distributional_dueling.rs +++ b/crates/ml-dqn/src/distributional_dueling.rs @@ -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 { - 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) diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml-dqn/src/dqn.rs similarity index 98% rename from crates/ml/src/dqn/dqn.rs rename to crates/ml-dqn/src/dqn.rs index 1b2156c26..b7ddfd7c5 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -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 { + pub fn from_safetensors_file(path: &std::path::Path) -> Result { 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, - ) -> Result { - let get = |key: &str| -> Result<&str, crate::MLError> { + ) -> Result { + 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::()) .collect::, _>>() - .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 { - get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e))) + let parse_bool = |key: &str| -> Result { + get(key)?.parse().map_err(|e| ml_core::MLError::CheckpointError(format!("Bad {}: {}", key, e))) }; - let parse_usize = |key: &str| -> Result { - get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e))) + let parse_usize = |key: &str| -> Result { + 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>, - ) -> 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 { - 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 { // 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 { 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] diff --git a/crates/ml/src/dqn/dueling.rs b/crates/ml-dqn/src/dueling.rs similarity index 98% rename from crates/ml/src/dqn/dueling.rs rename to crates/ml-dqn/src/dueling.rs index fc12570e1..d372d4baf 100644 --- a/crates/ml/src/dqn/dueling.rs +++ b/crates/ml-dqn/src/dueling.rs @@ -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 { - 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| { diff --git a/crates/ml/src/dqn/ensemble_network.rs b/crates/ml-dqn/src/ensemble_network.rs similarity index 99% rename from crates/ml/src/dqn/ensemble_network.rs rename to crates/ml-dqn/src/ensemble_network.rs index 08b45dfa6..6c7f49f8f 100644 --- a/crates/ml/src/dqn/ensemble_network.rs +++ b/crates/ml-dqn/src/ensemble_network.rs @@ -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)] diff --git a/crates/ml/src/dqn/entropy_regularization.rs b/crates/ml-dqn/src/entropy_regularization.rs similarity index 99% rename from crates/ml/src/dqn/entropy_regularization.rs rename to crates/ml-dqn/src/entropy_regularization.rs index 6dd8e822a..86dfd7699 100644 --- a/crates/ml/src/dqn/entropy_regularization.rs +++ b/crates/ml-dqn/src/entropy_regularization.rs @@ -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 /// diff --git a/crates/ml/src/dqn/experience.rs b/crates/ml-dqn/src/experience.rs similarity index 100% rename from crates/ml/src/dqn/experience.rs rename to crates/ml-dqn/src/experience.rs diff --git a/crates/ml/src/dqn/experience_dataset.rs b/crates/ml-dqn/src/experience_dataset.rs similarity index 100% rename from crates/ml/src/dqn/experience_dataset.rs rename to crates/ml-dqn/src/experience_dataset.rs diff --git a/crates/ml/src/dqn/gae.rs b/crates/ml-dqn/src/gae.rs similarity index 100% rename from crates/ml/src/dqn/gae.rs rename to crates/ml-dqn/src/gae.rs diff --git a/crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs similarity index 99% rename from crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs rename to crates/ml-dqn/src/gpu_replay_buffer.rs index 8af98c6e7..07f2fed10 100644 --- a/crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -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)] diff --git a/crates/ml/src/dqn/hindsight_replay.rs b/crates/ml-dqn/src/hindsight_replay.rs similarity index 99% rename from crates/ml/src/dqn/hindsight_replay.rs rename to crates/ml-dqn/src/hindsight_replay.rs index 456e337e8..b8d31bd12 100644 --- a/crates/ml/src/dqn/hindsight_replay.rs +++ b/crates/ml-dqn/src/hindsight_replay.rs @@ -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)] diff --git a/crates/ml/src/dqn/iql.rs b/crates/ml-dqn/src/iql.rs similarity index 100% rename from crates/ml/src/dqn/iql.rs rename to crates/ml-dqn/src/iql.rs diff --git a/crates/ml-dqn/src/lib.rs b/crates/ml-dqn/src/lib.rs index 6ef21459c..b67153144 100644 --- a/crates/ml-dqn/src/lib.rs +++ b/crates/ml-dqn/src/lib.rs @@ -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; diff --git a/crates/ml/src/dqn/logging.rs b/crates/ml-dqn/src/logging.rs similarity index 100% rename from crates/ml/src/dqn/logging.rs rename to crates/ml-dqn/src/logging.rs diff --git a/crates/ml/src/dqn/logit_clipping.rs b/crates/ml-dqn/src/logit_clipping.rs similarity index 99% rename from crates/ml/src/dqn/logit_clipping.rs rename to crates/ml-dqn/src/logit_clipping.rs index 0943893a9..a96f9aa99 100644 --- a/crates/ml/src/dqn/logit_clipping.rs +++ b/crates/ml-dqn/src/logit_clipping.rs @@ -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; diff --git a/crates/ml/src/dqn/multi_asset.rs b/crates/ml-dqn/src/multi_asset.rs similarity index 99% rename from crates/ml/src/dqn/multi_asset.rs rename to crates/ml-dqn/src/multi_asset.rs index 602d3d84b..96b403da1 100644 --- a/crates/ml/src/dqn/multi_asset.rs +++ b/crates/ml-dqn/src/multi_asset.rs @@ -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() { diff --git a/crates/ml/src/dqn/multi_step.rs b/crates/ml-dqn/src/multi_step.rs similarity index 99% rename from crates/ml/src/dqn/multi_step.rs rename to crates/ml-dqn/src/multi_step.rs index 951a1ae88..fcbb94e44 100644 --- a/crates/ml/src/dqn/multi_step.rs +++ b/crates/ml-dqn/src/multi_step.rs @@ -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)] diff --git a/crates/ml/src/dqn/network.rs b/crates/ml-dqn/src/network.rs similarity index 97% rename from crates/ml/src/dqn/network.rs rename to crates/ml-dqn/src/network.rs index 89ee6f1e1..01d07a435 100644 --- a/crates/ml/src/dqn/network.rs +++ b/crates/ml-dqn/src/network.rs @@ -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, + pub mixed_precision: Option, } impl Default for QNetworkConfig { @@ -194,9 +194,9 @@ impl NetworkLayers { fn forward_mixed( &self, xs: &Tensor, - mixed_precision: &Option, + mixed_precision: &Option, ) -> CandleResult { - 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 { - 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) diff --git a/crates/ml/src/dqn/noisy_exploration.rs b/crates/ml-dqn/src/noisy_exploration.rs similarity index 99% rename from crates/ml/src/dqn/noisy_exploration.rs rename to crates/ml-dqn/src/noisy_exploration.rs index 58ae360ee..361ac10f9 100644 --- a/crates/ml/src/dqn/noisy_exploration.rs +++ b/crates/ml-dqn/src/noisy_exploration.rs @@ -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)] diff --git a/crates/ml/src/dqn/noisy_layers.rs b/crates/ml-dqn/src/noisy_layers.rs similarity index 99% rename from crates/ml/src/dqn/noisy_layers.rs rename to crates/ml-dqn/src/noisy_layers.rs index 6617b0851..c1e2e8637 100644 --- a/crates/ml/src/dqn/noisy_layers.rs +++ b/crates/ml-dqn/src/noisy_layers.rs @@ -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> { diff --git a/crates/ml/src/dqn/noisy_sigma_scheduler.rs b/crates/ml-dqn/src/noisy_sigma_scheduler.rs similarity index 100% rename from crates/ml/src/dqn/noisy_sigma_scheduler.rs rename to crates/ml-dqn/src/noisy_sigma_scheduler.rs diff --git a/crates/ml/src/dqn/nstep_buffer.rs b/crates/ml-dqn/src/nstep_buffer.rs similarity index 99% rename from crates/ml/src/dqn/nstep_buffer.rs rename to crates/ml-dqn/src/nstep_buffer.rs index 0adde46e5..c3ea47baa 100644 --- a/crates/ml/src/dqn/nstep_buffer.rs +++ b/crates/ml-dqn/src/nstep_buffer.rs @@ -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 /// diff --git a/crates/ml/src/dqn/performance_tests.rs b/crates/ml-dqn/src/performance_tests.rs similarity index 99% rename from crates/ml/src/dqn/performance_tests.rs rename to crates/ml-dqn/src/performance_tests.rs index 46c867d1c..8c2d4adc5 100644 --- a/crates/ml/src/dqn/performance_tests.rs +++ b/crates/ml-dqn/src/performance_tests.rs @@ -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)] diff --git a/crates/ml/src/dqn/performance_validation.rs b/crates/ml-dqn/src/performance_validation.rs similarity index 100% rename from crates/ml/src/dqn/performance_validation.rs rename to crates/ml-dqn/src/performance_validation.rs diff --git a/crates/ml/src/dqn/prioritized_replay.rs b/crates/ml-dqn/src/prioritized_replay.rs similarity index 99% rename from crates/ml/src/dqn/prioritized_replay.rs rename to crates/ml-dqn/src/prioritized_replay.rs index 54e317c7c..8e443331c 100644 --- a/crates/ml/src/dqn/prioritized_replay.rs +++ b/crates/ml-dqn/src/prioritized_replay.rs @@ -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) diff --git a/crates/ml/src/dqn/prioritized_replay_staleness.rs b/crates/ml-dqn/src/prioritized_replay_staleness.rs similarity index 100% rename from crates/ml/src/dqn/prioritized_replay_staleness.rs rename to crates/ml-dqn/src/prioritized_replay_staleness.rs diff --git a/crates/ml/src/dqn/quantile_regression.rs b/crates/ml-dqn/src/quantile_regression.rs similarity index 99% rename from crates/ml/src/dqn/quantile_regression.rs rename to crates/ml-dqn/src/quantile_regression.rs index 0358ce1df..852e6e1eb 100644 --- a/crates/ml/src/dqn/quantile_regression.rs +++ b/crates/ml-dqn/src/quantile_regression.rs @@ -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 { - 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)?; diff --git a/crates/ml/src/dqn/rainbow_agent.rs b/crates/ml-dqn/src/rainbow_agent.rs similarity index 99% rename from crates/ml/src/dqn/rainbow_agent.rs rename to crates/ml-dqn/src/rainbow_agent.rs index cda68170e..e463781ea 100644 --- a/crates/ml/src/dqn/rainbow_agent.rs +++ b/crates/ml-dqn/src/rainbow_agent.rs @@ -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 { diff --git a/crates/ml/src/dqn/rainbow_config.rs b/crates/ml-dqn/src/rainbow_config.rs similarity index 100% rename from crates/ml/src/dqn/rainbow_config.rs rename to crates/ml-dqn/src/rainbow_config.rs diff --git a/crates/ml/src/dqn/rainbow_integration.rs b/crates/ml-dqn/src/rainbow_integration.rs similarity index 96% rename from crates/ml/src/dqn/rainbow_integration.rs rename to crates/ml-dqn/src/rainbow_integration.rs index 2c0cb09d4..1c8b428ca 100644 --- a/crates/ml/src/dqn/rainbow_integration.rs +++ b/crates/ml-dqn/src/rainbow_integration.rs @@ -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() { diff --git a/crates/ml/src/dqn/rainbow_network.rs b/crates/ml-dqn/src/rainbow_network.rs similarity index 99% rename from crates/ml/src/dqn/rainbow_network.rs rename to crates/ml-dqn/src/rainbow_network.rs index 44e13cc21..f7b23fe4e 100644 --- a/crates/ml/src/dqn/rainbow_network.rs +++ b/crates/ml-dqn/src/rainbow_network.rs @@ -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> { diff --git a/crates/ml/src/dqn/regime_conditional.rs b/crates/ml-dqn/src/regime_conditional.rs similarity index 99% rename from crates/ml/src/dqn/regime_conditional.rs rename to crates/ml-dqn/src/regime_conditional.rs index 64cdac50e..556150742 100644 --- a/crates/ml/src/dqn/regime_conditional.rs +++ b/crates/ml-dqn/src/regime_conditional.rs @@ -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)?; diff --git a/crates/ml/src/dqn/replay_buffer.rs b/crates/ml-dqn/src/replay_buffer.rs similarity index 99% rename from crates/ml/src/dqn/replay_buffer.rs rename to crates/ml-dqn/src/replay_buffer.rs index 9612ede07..9e9a06b9f 100644 --- a/crates/ml/src/dqn/replay_buffer.rs +++ b/crates/ml-dqn/src/replay_buffer.rs @@ -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)] diff --git a/crates/ml/src/dqn/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs similarity index 97% rename from crates/ml/src/dqn/replay_buffer_type.rs rename to crates/ml-dqn/src/replay_buffer_type.rs index 3bb702b77..0d9490316 100644 --- a/crates/ml/src/dqn/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -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), /// GPU-resident prioritized replay — all sampling on GPU (zero CPU round-trips) #[cfg(feature = "cuda")] - GpuPrioritized(Arc>), + GpuPrioritized(Arc>), } impl std::fmt::Debug for ReplayBufferType { @@ -90,7 +90,7 @@ impl ReplayBufferType { beta_max: f64, beta_annealing_steps: usize, ) -> Result { - 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 { - 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> { + pub fn as_gpu_buffer(&self) -> Option> { match self { Self::GpuPrioritized(buffer) => Some(buffer.lock()), Self::Uniform(_) | Self::Prioritized(_) => None, diff --git a/crates/ml/src/dqn/residual.rs b/crates/ml-dqn/src/residual.rs similarity index 98% rename from crates/ml/src/dqn/residual.rs rename to crates/ml-dqn/src/residual.rs index 1a5ee9cbd..c76d138f9 100644 --- a/crates/ml/src/dqn/residual.rs +++ b/crates/ml-dqn/src/residual.rs @@ -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() { diff --git a/crates/ml/src/dqn/reward.rs b/crates/ml-dqn/src/reward.rs similarity index 99% rename from crates/ml/src/dqn/reward.rs rename to crates/ml-dqn/src/reward.rs index ce8b60634..557c7cbf8 100644 --- a/crates/ml/src/dqn/reward.rs +++ b/crates/ml-dqn/src/reward.rs @@ -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)?; diff --git a/crates/ml/src/dqn/rmsnorm.rs b/crates/ml-dqn/src/rmsnorm.rs similarity index 99% rename from crates/ml/src/dqn/rmsnorm.rs rename to crates/ml-dqn/src/rmsnorm.rs index 2bedc5617..99415c493 100644 --- a/crates/ml/src/dqn/rmsnorm.rs +++ b/crates/ml-dqn/src/rmsnorm.rs @@ -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<()> { diff --git a/crates/ml/src/dqn/self_supervised_pretraining.rs b/crates/ml-dqn/src/self_supervised_pretraining.rs similarity index 100% rename from crates/ml/src/dqn/self_supervised_pretraining.rs rename to crates/ml-dqn/src/self_supervised_pretraining.rs diff --git a/crates/ml/src/dqn/softmax.rs b/crates/ml-dqn/src/softmax.rs similarity index 99% rename from crates/ml/src/dqn/softmax.rs rename to crates/ml-dqn/src/softmax.rs index 4da65dad0..eeb02eb4e 100644 --- a/crates/ml/src/dqn/softmax.rs +++ b/crates/ml-dqn/src/softmax.rs @@ -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)] diff --git a/crates/ml/src/dqn/target_update.rs b/crates/ml-dqn/src/target_update.rs similarity index 100% rename from crates/ml/src/dqn/target_update.rs rename to crates/ml-dqn/src/target_update.rs diff --git a/crates/ml/src/dqn/trade_executor.rs b/crates/ml-dqn/src/trade_executor.rs similarity index 99% rename from crates/ml/src/dqn/trade_executor.rs rename to crates/ml-dqn/src/trade_executor.rs index f57052bb7..0ff7c10b3 100644 --- a/crates/ml/src/dqn/trade_executor.rs +++ b/crates/ml-dqn/src/trade_executor.rs @@ -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}; diff --git a/crates/ml/Cargo.toml b/crates/ml/Cargo.toml index 469c0b6ab..c7e93b941 100644 --- a/crates/ml/Cargo.toml +++ b/crates/ml/Cargo.toml @@ -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" } diff --git a/crates/ml/src/checkpoint/model_implementations.rs b/crates/ml/src/checkpoint/model_implementations.rs index ffe9b6bc4..507ff96bf 100644 --- a/crates/ml/src/checkpoint/model_implementations.rs +++ b/crates/ml/src/checkpoint/model_implementations.rs @@ -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> { - // 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)] diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 64d9e5d75..2060262de 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -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; diff --git a/crates/ml/src/dqn/data_augmentation.rs b/crates/ml/src/dqn/data_augmentation.rs deleted file mode 100644 index ed7cfe4a3..000000000 --- a/crates/ml/src/dqn/data_augmentation.rs +++ /dev/null @@ -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 { - if rng.gen::() < 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 = state - .iter() - .zip(augmented.iter()) - .map(|(s, a)| a - s) - .collect(); - - // Check mean is close to 0 - let mean: f32 = diffs.iter().sum::() / 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::() - / 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 = (0..1000) - .map(|_| injector.sample_gaussian(&mut rng)) - .collect(); - - // Check mean is close to 0 - let mean: f32 = samples.iter().sum::() / 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::() - / 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 = 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); - } -} diff --git a/crates/ml/src/dqn/ensemble.rs b/crates/ml/src/dqn/ensemble.rs deleted file mode 100644 index 7b6d40e4b..000000000 --- a/crates/ml/src/dqn/ensemble.rs +++ /dev/null @@ -1,1062 +0,0 @@ -//! DQN Ensemble with Multiple Voting Strategies -//! -//! Multi-agent ensemble system that combines 3-5 DQN agents with diverse architectures -//! to improve decision robustness through consensus voting. Supports 5 voting strategies: -//! -//! 1. **Majority**: Simple majority vote (most common action wins) -//! 2. **QValueWeighted**: Actions weighted by Q-value confidence -//! 3. **Thompson**: Bayesian sampling based on action success rates -//! 4. **MinVariance**: Conservative - chooses action with lowest Q-value variance across agents -//! 5. **MaxVariance**: Aggressive - chooses action with highest variance (explore uncertainty) -//! -//! # Architecture Diversity -//! -//! Each agent has unique characteristics to maximize ensemble diversity: -//! - **Agent 0**: Shallow (2 layers: [256, 128]) - fast convergence, simple patterns -//! - **Agent 1**: Medium (3 layers: [512, 256, 128]) - balanced depth -//! - **Agent 2**: Deep (4 layers: [512, 256, 128, 64]) - complex pattern recognition -//! - **Agent 3**: Wide (3 layers: [1024, 512, 256]) - high capacity -//! - **Agent 4**: Narrow (4 layers: [256, 128, 64, 32]) - regularization via bottleneck -//! -//! # Example -//! -//! ```no_run -//! use ml::dqn::{DQNEnsemble, EnsembleConfig, VotingStrategy, DQNConfig}; -//! -//! // Create ensemble with 3 agents -//! let config = EnsembleConfig::new(3, 128, VotingStrategy::QValueWeighted)?; -//! let mut ensemble = DQNEnsemble::new(config)?; -//! -//! // Select action via ensemble voting -//! let state = vec![0.5; 128]; -//! let action = ensemble.select_action(&state)?; -//! -//! // Train all agents on shared experience -//! let experience = Experience::new(state, action as u8, 0.5, vec![0.6; 128], false); -//! ensemble.train_on_experience(experience)?; -//! # Ok::<(), ml::MLError>(()) -//! ``` - -use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, Mutex}; - -use candle_core::{Device, Tensor}; -use rand::{thread_rng, Rng}; -use serde::{Deserialize, Serialize}; -use tracing::{debug, info}; - -use super::{Experience, TradingAction, DQN, DQNConfig}; -use crate::MLError; - -/// Voting strategy for ensemble decision-making -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum VotingStrategy { - /// Simple majority vote (most common action wins) - Majority, - /// Weight votes by Q-value confidence (action with highest average Q-value) - QValueWeighted, - /// Bayesian Thompson sampling (sample based on historical success rates) - Thompson, - /// Conservative: Choose action with minimum Q-value variance across agents - MinVariance, - /// Aggressive: Choose action with maximum Q-value variance (explore uncertainty) - MaxVariance, -} - -/// Configuration for DQN ensemble -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct EnsembleConfig { - /// Number of agents in ensemble (3-5) - pub num_agents: usize, - /// State dimension (must match feature vector size) - pub state_dim: usize, - /// Voting strategy - pub voting_strategy: VotingStrategy, - /// Base learning rate (will be varied per agent) - pub base_learning_rate: f64, - /// Whether to use shared replay buffer (default: false = separate buffers) - pub shared_replay_buffer: bool, - /// Thompson sampling decay rate (default: 0.99) - pub thompson_decay: f64, - /// Enable diversity penalty (punishes agents for convergent actions) - pub enable_diversity_penalty: bool, - /// Diversity penalty weight (default: 0.1) - pub diversity_penalty_weight: f64, -} - -impl EnsembleConfig { - /// Create new ensemble config with sensible defaults - /// - /// # Arguments - /// - /// * `num_agents` - Number of agents (3-5) - /// * `state_dim` - State dimension - /// * `voting_strategy` - Voting strategy - pub fn new(num_agents: usize, state_dim: usize, voting_strategy: VotingStrategy) -> Result { - if !(3..=5).contains(&num_agents) { - return Err(MLError::ConfigError(format!("num_agents must be 3-5, got {}", num_agents))); - } - - Ok(Self { - num_agents, - state_dim, - voting_strategy, - base_learning_rate: 3e-5, // Wave 7 best parameter - shared_replay_buffer: false, - thompson_decay: 0.99, - enable_diversity_penalty: true, - diversity_penalty_weight: 0.1, - }) - } - - /// Create config for production use (conservative, QValueWeighted) - pub fn production(state_dim: usize) -> Result { - Self::new(3, state_dim, VotingStrategy::QValueWeighted) - } - - /// Create config for exploration (aggressive, MaxVariance) - pub fn exploration(state_dim: usize) -> Result { - Self::new(5, state_dim, VotingStrategy::MaxVariance) - } -} - -/// Thompson sampling statistics for each action -#[derive(Debug, Clone)] -struct ThompsonStats { - /// Number of times action was selected - counts: [u64; 3], - /// Cumulative rewards for each action - rewards: [f64; 3], -} - -impl ThompsonStats { - fn new() -> Self { - Self { - counts: [0; 3], - rewards: [0.0; 3], - } - } - - /// Update statistics after action execution - fn update(&mut self, action: TradingAction, reward: f64) { - let idx = action as usize; - self.counts[idx] += 1; - self.rewards[idx] += reward; - } - - /// Get success rate for action (mean reward) - fn success_rate(&self, action: TradingAction) -> f64 { - let idx = action as usize; - if self.counts[idx] == 0 { - 0.5 // Prior: neutral success rate - } else { - self.rewards[idx] / self.counts[idx] as f64 - } - } -} - -/// DQN Ensemble with multiple voting strategies -#[allow(missing_debug_implementations)] -pub struct DQNEnsemble { - /// Configuration - config: EnsembleConfig, - /// Individual DQN agents with diverse architectures - agents: Vec, - /// Agent configs (stored separately for access) - agent_configs: Vec, - /// Thompson sampling statistics (per-agent) - thompson_stats: Vec, - /// Shared replay buffer (if enabled) - shared_memory: Option>>, - /// Device (CPU or CUDA) - device: Device, - /// Recent actions for diversity monitoring - recent_agent_actions: Vec>, - /// Total training steps - total_steps: u64, -} - -impl DQNEnsemble { - /// Create new DQN ensemble with diverse architectures - pub fn new(config: EnsembleConfig) -> Result { - let device = Device::cuda_if_available(0)?; - - // Create diverse agent architectures - let agents_and_configs: Result, _> = (0..config.num_agents) - .map(|i| Self::create_diverse_agent(i, &config, &device)) - .collect(); - - let agents_and_configs = agents_and_configs?; - let (agents, agent_configs): (Vec<_>, Vec<_>) = agents_and_configs.into_iter().unzip(); - - // Initialize Thompson stats - let thompson_stats = vec![ThompsonStats::new(); config.num_agents]; - - // Create shared replay buffer if enabled - let shared_memory = if config.shared_replay_buffer { - let capacity = 100_000; // Standard replay buffer size - Some(Arc::new(Mutex::new( - super::dqn::ExperienceReplayBuffer::new(capacity), - ))) - } else { - None - }; - - // Initialize diversity tracking - let recent_agent_actions = vec![VecDeque::with_capacity(100); config.num_agents]; - - info!( - "✓ DQN Ensemble initialized: {} agents, {:?} voting, state_dim={}", - config.num_agents, config.voting_strategy, config.state_dim - ); - - Ok(Self { - config, - agents, - agent_configs, - thompson_stats, - shared_memory, - device, - recent_agent_actions, - total_steps: 0, - }) - } - - /// Create agent with architecture diversity - /// - /// Returns (agent, config) tuple for storage - /// - /// Agent architectures (num_agents = 5): - /// - Agent 0: Shallow (2 layers: [256, 128]) - /// - Agent 1: Medium (3 layers: [512, 256, 128]) - /// - Agent 2: Deep (4 layers: [512, 256, 128, 64]) - /// - Agent 3: Wide (3 layers: [1024, 512, 256]) - /// - Agent 4: Narrow (4 layers: [256, 128, 64, 32]) - fn create_diverse_agent( - idx: usize, - config: &EnsembleConfig, - device: &Device, - ) -> Result<(DQN, DQNConfig), MLError> { - let mut agent_config = DQNConfig::emergency_safe_defaults(); - agent_config.state_dim = config.state_dim; - - // Architecture diversity - agent_config.hidden_dims = match idx % 5 { - 0 => vec![256, 128], // Shallow - 1 => vec![512, 256, 128], // Medium (default) - 2 => vec![512, 256, 128, 64], // Deep - 3 => vec![1024, 512, 256], // Wide - 4 => vec![256, 128, 64, 32], // Narrow - _ => unreachable!(), - }; - - // Learning rate diversity (±20% variation) - let lr_multipliers = [0.8, 1.0, 1.2, 0.9, 1.1]; - agent_config.learning_rate = config.base_learning_rate * lr_multipliers[idx % 5]; - - // Gamma diversity (slight variation for different time horizons) - let gammas = [0.95, 0.963, 0.98, 0.93, 0.97]; - agent_config.gamma = gammas[idx % 5]; - - // Epsilon diversity (different exploration rates) - let epsilon_starts = [0.5, 0.7, 0.3, 0.6, 0.4]; - agent_config.epsilon_start = epsilon_starts[idx % 5]; - agent_config.epsilon_decay = 0.995; // Shared decay rate - - // Temperature diversity (softmax action selection) - let temp_starts = [0.8, 1.0, 1.2, 0.9, 1.1]; - agent_config.temperature_start = temp_starts[idx % 5]; - - // Buffer size diversity (memory vs recency trade-off) - let buffer_sizes = [10_000, 20_000, 30_000, 15_000, 25_000]; - agent_config.replay_buffer_capacity = buffer_sizes[idx % 5]; - - // Batch size diversity - let batch_sizes = [128, 256, 192, 224, 160]; - agent_config.batch_size = batch_sizes[idx % 5]; - - // Set unique device seed for each agent - let agent_seed = 42 + (idx as u64) * 1000; - device - .set_seed(agent_seed) - .map_err(|e| MLError::ModelError(format!("Failed to seed agent {}: {}", idx, e)))?; - - debug!( - "Agent {} created: arch={:?}, lr={:.2e}, gamma={:.3}, eps_start={:.2}", - idx, - agent_config.hidden_dims, - agent_config.learning_rate, - agent_config.gamma, - agent_config.epsilon_start - ); - - // Clone config before moving into DQN::new - let agent_config_clone = agent_config.clone(); - let agent = DQN::new(agent_config)?; - - Ok((agent, agent_config_clone)) - } - - /// Select action using ensemble voting - pub fn select_action(&mut self, state: &[f32]) -> Result { - self.total_steps += 1; - - // Collect predictions from all agents - let mut agent_actions = Vec::with_capacity(self.config.num_agents); - let mut agent_q_values = Vec::with_capacity(self.config.num_agents); - - for (idx, agent) in self.agents.iter_mut().enumerate() { - let action = agent.select_action(state)?; - agent_actions.push(action); - - // Get Q-values for voting strategies - let state_tensor = Tensor::from_vec( - state.to_vec(), - (1, self.config.state_dim), - &self.device, - ) - .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?; - - let q_values = agent.forward(&state_tensor)?; - let q_vec = q_values - .flatten_all()? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Failed to extract Q-values: {}", e)))?; - agent_q_values.push(q_vec); - - // Track action for diversity monitoring - self.recent_agent_actions[idx].push_back(action); - if self.recent_agent_actions[idx].len() > 100 { - self.recent_agent_actions[idx].pop_front(); - } - } - - // Apply voting strategy - let action = match self.config.voting_strategy { - VotingStrategy::Majority => self.vote_majority(&agent_actions), - VotingStrategy::QValueWeighted => { - self.vote_q_value_weighted(&agent_actions, &agent_q_values) - } - VotingStrategy::Thompson => self.vote_thompson(&agent_actions), - VotingStrategy::MinVariance => self.vote_min_variance(&agent_q_values), - VotingStrategy::MaxVariance => self.vote_max_variance(&agent_q_values), - }; - - // Log diversity metrics every 1000 steps - if self.total_steps % 1000 == 0 { - self.log_diversity_metrics(); - } - - Ok(action) - } - - /// Majority voting: most common action wins - fn vote_majority(&self, actions: &[TradingAction]) -> TradingAction { - let mut counts = HashMap::new(); - for &action in actions { - *counts.entry(action).or_insert(0) += 1; - } - - // Return action with highest count (ties broken by enum order) - *counts - .iter() - .max_by_key(|(_, count)| *count) - .map(|(action, _)| action) - .unwrap_or(&TradingAction::Hold) - } - - /// Q-value weighted voting: weight by Q-value confidence - fn vote_q_value_weighted( - &self, - actions: &[TradingAction], - q_values: &[Vec], - ) -> TradingAction { - let mut weighted_votes = [0.0_f32; 3]; // BUY, SELL, HOLD - - for (action, q_vals) in actions.iter().zip(q_values.iter()) { - let action_idx = *action as usize; - let q_value = q_vals[action_idx]; - weighted_votes[action_idx] += q_value; - } - - // Return action with highest weighted vote - let max_idx = weighted_votes - .iter() - .enumerate() - .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(idx, _)| idx) - .unwrap_or(2); // Default to HOLD - - TradingAction::from_int(max_idx as u8).unwrap_or(TradingAction::Hold) - } - - /// Thompson sampling: Bayesian action selection based on historical success - fn vote_thompson(&mut self, _actions: &[TradingAction]) -> TradingAction { - // Sample from Beta distribution for each action - let mut rng = thread_rng(); - let mut samples = [0.0; 3]; - - for action in TradingAction::all() { - // Aggregate success rates across all agents - let mut total_success = 0.0; - let mut total_trials = 0.0; - - for stats in &self.thompson_stats { - let rate = stats.success_rate(action); - let count = stats.counts[action as usize]; - total_success += rate * count as f64; - total_trials += count as f64; - } - - // Beta distribution: Beta(α=successes+1, β=failures+1) - let alpha = total_success + 1.0; - let beta = (total_trials - total_success) + 1.0; - - // Simple Beta sampling via rejection (good enough for our use case) - samples[action as usize] = Self::sample_beta(alpha, beta, &mut rng); - } - - // Return action with highest sample - let max_idx = samples - .iter() - .enumerate() - .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(idx, _)| idx) - .unwrap_or(2); - - TradingAction::from_int(max_idx as u8).unwrap_or(TradingAction::Hold) - } - - /// Sample from Beta(α, β) distribution (rejection sampling) - fn sample_beta(alpha: f64, beta: f64, rng: &mut impl Rng) -> f64 { - // Use Gamma ratio method: X ~ Gamma(α), Y ~ Gamma(β), then X/(X+Y) ~ Beta(α, β) - let x = Self::sample_gamma(alpha, rng); - let y = Self::sample_gamma(beta, rng); - x / (x + y) - } - - /// Sample from Gamma(α) distribution (shape parameter) - fn sample_gamma(alpha: f64, rng: &mut impl Rng) -> f64 { - // Marsaglia and Tsang's Method (2000) - simplified for α > 1 - if alpha < 1.0 { - return Self::sample_gamma(alpha + 1.0, rng) * rng.gen::().powf(1.0 / alpha); - } - - let d = alpha - 1.0 / 3.0; - let c = 1.0 / (9.0 * d).sqrt(); - - loop { - let x = Self::sample_normal(rng); // Must be N(0,1), not Uniform - let v = (1.0 + c * x).powi(3); - - if v > 0.0 { - let u = rng.gen::(); - if u < 1.0 - 0.0331 * x.powi(4) { - return d * v; - } - if u.ln() < 0.5 * x.powi(2) + d * (1.0 - v + v.ln()) { - return d * v; - } - } - } - } - - /// Sample from standard normal N(0,1) using Box-Muller transform - fn sample_normal(rng: &mut impl Rng) -> f64 { - let u1 = rng.gen::(); - let u2 = rng.gen::(); - (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos() - } - - /// Min variance voting: choose action with lowest Q-value variance (conservative) - fn vote_min_variance(&self, q_values: &[Vec]) -> TradingAction { - let mut variances = [0.0_f32; 3]; - - for action_idx in 0..3 { - // Extract Q-values for this action across all agents - let q_vals: Vec = q_values.iter().map(|q| q[action_idx]).collect(); - - // Compute variance - let mean = q_vals.iter().sum::() / q_vals.len() as f32; - let variance: f32 = q_vals - .iter() - .map(|q| { - let diff = q - mean; - diff * diff - }) - .sum::() - / q_vals.len() as f32; - - variances[action_idx] = variance; - } - - // Return action with minimum variance (most agent agreement) - let min_idx = variances - .iter() - .enumerate() - .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(idx, _)| idx) - .unwrap_or(2); - - TradingAction::from_int(min_idx as u8).unwrap_or(TradingAction::Hold) - } - - /// Max variance voting: choose action with highest Q-value variance (explore uncertainty) - fn vote_max_variance(&self, q_values: &[Vec]) -> TradingAction { - let mut variances = [0.0_f32; 3]; - - for action_idx in 0..3 { - let q_vals: Vec = q_values.iter().map(|q| q[action_idx]).collect(); - - let mean = q_vals.iter().sum::() / q_vals.len() as f32; - let variance: f32 = q_vals - .iter() - .map(|q| { - let diff = q - mean; - diff * diff - }) - .sum::() - / q_vals.len() as f32; - - variances[action_idx] = variance; - } - - // Return action with maximum variance (most disagreement = exploration opportunity) - let max_idx = variances - .iter() - .enumerate() - .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(idx, _)| idx) - .unwrap_or(2); - - TradingAction::from_int(max_idx as u8).unwrap_or(TradingAction::Hold) - } - - /// Train all agents on shared experience - /// - /// If `shared_replay_buffer` is enabled, experience is stored once and sampled by all agents. - /// Otherwise, each agent maintains its own replay buffer. - pub fn train_on_experience(&mut self, experience: Experience) -> Result, MLError> { - let mut losses = Vec::with_capacity(self.config.num_agents); - - if let Some(ref shared_mem) = self.shared_memory { - // Shared replay buffer: store once, train all agents - shared_mem - .lock() - .map_err(|e| MLError::ConcurrencyError { - operation: format!("lock shared memory: {}", e), - })? - .push(experience.clone()); - - for (agent, cfg) in self.agents.iter_mut().zip(self.agent_configs.iter()) { - // Sample from shared buffer - let buffer = shared_mem.lock().map_err(|e| MLError::ConcurrencyError { - operation: format!("lock shared memory for training: {}", e), - })?; - - if buffer.can_sample(cfg.min_replay_size) { - let batch = buffer.sample(cfg.batch_size)?; - drop(buffer); // Release lock before training - - let (loss, _grad_norm) = agent.train_step(Some(super::replay_buffer_type::BatchSample::uniform(batch)))?; - losses.push(loss); - } else { - losses.push(0.0); // Not enough data yet - } - } - } else { - // Separate replay buffers: each agent stores and trains independently - for agent in &mut self.agents { - agent.store_experience(experience.clone())?; - - if agent.can_train() { - let (loss, _grad_norm) = agent.train_step(None)?; - losses.push(loss); - } else { - losses.push(0.0); - } - } - } - - // Update Thompson statistics (on training, not just selection) - let action = TradingAction::from_int(experience.action).unwrap_or(TradingAction::Hold); - let reward = experience.reward_f32() as f64; - - for stats in &mut self.thompson_stats { - stats.update(action, reward); - } - - // Apply diversity penalty if enabled - if self.config.enable_diversity_penalty { - self.apply_diversity_penalty()?; - } - - Ok(losses) - } - - /// Apply diversity penalty to agents with convergent behavior - /// - /// Penalizes agents whose actions are too similar to the ensemble mean, - /// encouraging diverse exploration strategies. - fn apply_diversity_penalty(&mut self) -> Result<(), MLError> { - // Compute action distribution for each agent - let mut agent_distributions = Vec::with_capacity(self.config.num_agents); - - for agent_actions in &self.recent_agent_actions { - if agent_actions.is_empty() { - agent_distributions.push([0.0; 3]); - continue; - } - - let mut counts = [0.0; 3]; - for &action in agent_actions { - counts[action as usize] += 1.0; - } - - // Normalize to probabilities - let total = counts.iter().sum::(); - if total > 0.0 { - for count in &mut counts { - *count /= total; - } - } - - agent_distributions.push(counts); - } - - // Compute ensemble mean distribution - let mut mean_dist = [0.0; 3]; - for dist in &agent_distributions { - for i in 0..3 { - mean_dist[i] += dist[i]; - } - } - for prob in &mut mean_dist { - *prob /= self.config.num_agents as f32; - } - - // Penalize agents close to mean (KL divergence) - for (idx, dist) in agent_distributions.iter().enumerate() { - let kl_div = Self::kl_divergence(dist, &mean_dist); - - // If KL divergence is low (agent too similar to mean), apply penalty - if kl_div < 0.1 { - // Threshold for "too similar" - debug!( - "Agent {} diversity penalty: KL={:.4} (threshold=0.1)", - idx, kl_div - ); - - // NOTE: Actual penalty application would require modifying agent's loss function - // or adjusting learning rate. This is a monitoring point for now. - } - } - - Ok(()) - } - - /// Compute KL divergence between two distributions - fn kl_divergence(p: &[f32; 3], q: &[f32; 3]) -> f32 { - let mut kl = 0.0; - for i in 0..3 { - if p[i] > 1e-8 && q[i] > 1e-8 { - kl += p[i] * (p[i] / q[i]).ln(); - } - } - kl - } - - /// Log diversity metrics for monitoring - fn log_diversity_metrics(&self) { - let mut total_entropy = 0.0; - - for (idx, actions) in self.recent_agent_actions.iter().enumerate() { - if actions.is_empty() { - continue; - } - - // Compute Shannon entropy for agent's action distribution - let mut counts = [0.0; 3]; - for &action in actions { - counts[action as usize] += 1.0; - } - - let total = counts.iter().sum::(); - let mut entropy = 0.0; - for count in counts { - if count > 0.0 { - let p = count / total; - entropy -= p * p.log2(); - } - } - - total_entropy += entropy; - - debug!( - "Agent {} diversity: entropy={:.3}, actions=[{:.1}%, {:.1}%, {:.1}%]", - idx, - entropy, - counts[0] / total * 100.0, - counts[1] / total * 100.0, - counts[2] / total * 100.0 - ); - } - - let mean_entropy = total_entropy / self.config.num_agents as f32; - info!( - "Ensemble diversity (step {}): mean_entropy={:.3}, voting={:?}", - self.total_steps, mean_entropy, self.config.voting_strategy - ); - } - - /// Update epsilon for all agents (per-epoch decay) - pub fn update_epsilon(&mut self) { - for agent in &mut self.agents { - agent.update_epsilon(); - } - } - - /// Update temperature for all agents (per-epoch decay) - pub fn update_temperature(&mut self) { - for agent in &mut self.agents { - agent.update_temperature(); - } - } - - /// Get current epsilon values for all agents - pub fn get_epsilons(&self) -> Vec { - self.agents.iter().map(|a| a.get_epsilon()).collect() - } - - /// Get current temperature values for all agents - pub fn get_temperatures(&self) -> Vec { - self.agents.iter().map(|a| a.get_temperature()).collect() - } - - /// Get total training steps - pub fn get_total_steps(&self) -> u64 { - self.total_steps - } - - /// Get number of agents - pub fn num_agents(&self) -> usize { - self.config.num_agents - } - - /// Get voting strategy - pub fn voting_strategy(&self) -> VotingStrategy { - self.config.voting_strategy - } - - /// Save ensemble to directory (saves all agent weights with architecture metadata) - pub fn save_to_directory(&self, dir: &str) -> Result<(), MLError> { - std::fs::create_dir_all(dir).map_err(|e| { - MLError::CheckpointError(format!("Failed to create ensemble directory: {}", e)) - })?; - - for (idx, agent) in self.agents.iter().enumerate() { - let path = format!("{}/agent_{}.safetensors", dir, idx); - let vars = agent.get_q_network_vars(); - let vars_data = vars.data().lock().map_err(|e| { - MLError::LockError(format!("Failed to lock vars for ensemble agent {}: {}", idx, e)) - })?; - - let mut tensors: std::collections::HashMap = - std::collections::HashMap::new(); - for (name, var) in vars_data.iter() { - tensors.insert(name.clone(), var.as_tensor().clone()); - } - - let arch_metadata = Some(agent.config.checkpoint_metadata()); - safetensors::serialize_to_file( - &tensors, - &arch_metadata, - std::path::Path::new(&path), - ) - .map_err(|e| { - MLError::CheckpointError(format!("Failed to save ensemble agent {}: {}", idx, e)) - })?; - } - - info!("Ensemble saved to {} ({} agents)", dir, self.config.num_agents); - Ok(()) - } - - /// Load ensemble from directory (loads all agent weights) - pub fn load_from_directory(&mut self, dir: &str) -> Result<(), MLError> { - for (idx, agent) in self.agents.iter_mut().enumerate() { - let path = format!("{}/agent_{}.safetensors", dir, idx); - agent.load_from_safetensors(&path)?; - } - - info!("✓ Ensemble loaded from {} ({} agents)", dir, self.config.num_agents); - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_ensemble_creation() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - let ensemble = DQNEnsemble::new(config)?; - - assert_eq!(ensemble.num_agents(), 3); - assert_eq!(ensemble.voting_strategy(), VotingStrategy::Majority); - Ok(()) - } - - #[test] - fn test_ensemble_config_validation() { - // Too few agents - let result = EnsembleConfig::new(2, 128, VotingStrategy::Majority); - assert!(result.is_err()); - - // Too many agents - let result = EnsembleConfig::new(6, 128, VotingStrategy::Majority); - assert!(result.is_err()); - } - - #[test] - fn test_majority_voting() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - let ensemble = DQNEnsemble::new(config)?; - - // 2 Buy, 1 Sell -> Buy wins - let actions = vec![TradingAction::Buy, TradingAction::Buy, TradingAction::Sell]; - let result = ensemble.vote_majority(&actions); - assert_eq!(result, TradingAction::Buy); - - // 2 Hold, 1 Buy -> Hold wins - let actions = vec![TradingAction::Hold, TradingAction::Hold, TradingAction::Buy]; - let result = ensemble.vote_majority(&actions); - assert_eq!(result, TradingAction::Hold); - - Ok(()) - } - - #[test] - fn test_q_value_weighted_voting() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::QValueWeighted)?; - let ensemble = DQNEnsemble::new(config)?; - - let actions = vec![TradingAction::Buy, TradingAction::Sell, TradingAction::Hold]; - let q_values = vec![ - vec![1.0, 0.5, 0.3], // Agent 0: prefers Buy (Q=1.0) - vec![0.4, 2.0, 0.2], // Agent 1: prefers Sell (Q=2.0) - vec![0.1, 0.2, 0.8], // Agent 2: prefers Hold (Q=0.8) - ]; - - // Weighted votes: Buy=1.0, Sell=2.0, Hold=0.8 -> Sell wins - let result = ensemble.vote_q_value_weighted(&actions, &q_values); - assert_eq!(result, TradingAction::Sell); - - Ok(()) - } - - #[test] - fn test_min_variance_voting() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::MinVariance)?; - let ensemble = DQNEnsemble::new(config)?; - - let q_values = vec![ - vec![1.0, 0.5, 0.9], // Agent 0 - vec![1.1, 2.0, 0.95], // Agent 1 - vec![0.9, 0.3, 1.0], // Agent 2 - ]; - - // Q-value variances: Buy=[1.0, 1.1, 0.9] var=0.0067, Sell=[0.5, 2.0, 0.3] var=0.62, Hold=[0.9, 0.95, 1.0] var=0.0017 - // Min variance is Hold (agents agree most on Hold) - let result = ensemble.vote_min_variance(&q_values); - assert_eq!(result, TradingAction::Hold); - - Ok(()) - } - - #[test] - fn test_max_variance_voting() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::MaxVariance)?; - let ensemble = DQNEnsemble::new(config)?; - - let q_values = vec![ - vec![1.0, 0.5, 0.9], // Agent 0 - vec![1.1, 2.0, 0.95], // Agent 1 - vec![0.9, 0.3, 1.0], // Agent 2 - ]; - - // Max variance is Sell (agents disagree most on Sell) - let result = ensemble.vote_max_variance(&q_values); - assert_eq!(result, TradingAction::Sell); - - Ok(()) - } - - #[test] - fn test_action_selection_consistency() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - let mut ensemble = DQNEnsemble::new(config)?; - - // Select action 10 times (should not crash) - let state = vec![0.5; 128]; - for _ in 0..10 { - let action = ensemble.select_action(&state)?; - assert!(matches!( - action, - TradingAction::Buy | TradingAction::Sell | TradingAction::Hold - )); - } - - Ok(()) - } - - #[test] - fn test_training_with_shared_buffer() -> anyhow::Result<()> { - let mut config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - config.shared_replay_buffer = true; - - let mut ensemble = DQNEnsemble::new(config)?; - - // Create experience - let experience = Experience::new( - vec![0.5; 128], - TradingAction::Buy as u8, - 0.8, - vec![0.6; 128], - false, - ); - - // Train (may not have enough data yet, but should not crash) - let result = ensemble.train_on_experience(experience); - assert!(result.is_ok()); - - Ok(()) - } - - #[test] - fn test_training_with_separate_buffers() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - let mut ensemble = DQNEnsemble::new(config)?; - - let experience = Experience::new( - vec![0.5; 128], - TradingAction::Buy as u8, - 0.8, - vec![0.6; 128], - false, - ); - - let result = ensemble.train_on_experience(experience); - assert!(result.is_ok()); - - Ok(()) - } - - #[test] - fn test_epsilon_update() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - let mut ensemble = DQNEnsemble::new(config)?; - - let initial_epsilons = ensemble.get_epsilons(); - ensemble.update_epsilon(); - let updated_epsilons = ensemble.get_epsilons(); - - // All agents should have decayed epsilon - for (initial, updated) in initial_epsilons.iter().zip(updated_epsilons.iter()) { - assert!(updated <= initial, "Epsilon should decay"); - } - - Ok(()) - } - - #[test] - fn test_temperature_update() -> anyhow::Result<()> { - let config = EnsembleConfig::new(3, 128, VotingStrategy::Majority)?; - let mut ensemble = DQNEnsemble::new(config)?; - - let initial_temps = ensemble.get_temperatures(); - ensemble.update_temperature(); - let updated_temps = ensemble.get_temperatures(); - - // All agents should have decayed temperature - for (initial, updated) in initial_temps.iter().zip(updated_temps.iter()) { - assert!(updated <= initial, "Temperature should decay"); - } - - Ok(()) - } - - #[test] - fn test_production_config() -> anyhow::Result<()> { - let config = EnsembleConfig::production(128)?; - assert_eq!(config.num_agents, 3); - assert_eq!(config.voting_strategy, VotingStrategy::QValueWeighted); - assert!(!config.shared_replay_buffer); - Ok(()) - } - - #[test] - fn test_exploration_config() -> anyhow::Result<()> { - let config = EnsembleConfig::exploration(128)?; - assert_eq!(config.num_agents, 5); - assert_eq!(config.voting_strategy, VotingStrategy::MaxVariance); - Ok(()) - } - - #[test] - fn test_kl_divergence() { - let p = [0.33, 0.33, 0.34]; - let q = [0.5, 0.3, 0.2]; - - let kl = DQNEnsemble::kl_divergence(&p, &q); - assert!(kl > 0.0, "KL divergence should be positive"); - assert!(kl < 1.0, "KL divergence should be reasonable"); - - // Self KL divergence should be 0 - let kl_self = DQNEnsemble::kl_divergence(&p, &p); - assert!(kl_self < 1e-6, "Self KL divergence should be ~0"); - } - - #[test] - fn test_thompson_stats_update() { - let mut stats = ThompsonStats::new(); - - stats.update(TradingAction::Buy, 0.5); - stats.update(TradingAction::Buy, 0.8); - stats.update(TradingAction::Sell, -0.3); - - assert_eq!(stats.counts[0], 2); // Buy count - assert_eq!(stats.counts[1], 1); // Sell count - assert_eq!(stats.counts[2], 0); // Hold count - - let buy_rate = stats.success_rate(TradingAction::Buy); - assert!((buy_rate - 0.65).abs() < 1e-6); // (0.5 + 0.8) / 2 = 0.65 - - let hold_rate = stats.success_rate(TradingAction::Hold); - assert!((hold_rate - 0.5).abs() < 1e-6); // Prior = 0.5 - } - - #[test] - fn test_architectural_diversity() -> anyhow::Result<()> { - let config = EnsembleConfig::new(5, 128, VotingStrategy::Majority)?; - let ensemble = DQNEnsemble::new(config)?; - - // Verify agents have different architectures - let architectures: Vec<_> = ensemble - .agent_configs - .iter() - .map(|cfg| cfg.hidden_dims.clone()) - .collect(); - - // Check that we have at least 3 unique architectures - let unique_archs: std::collections::HashSet<_> = architectures.iter().collect(); - assert!( - unique_archs.len() >= 3, - "Expected diverse architectures, got {} unique", - unique_archs.len() - ); - - Ok(()) - } -} diff --git a/crates/ml/src/dqn/ensemble_oracle.rs b/crates/ml/src/dqn/ensemble_oracle.rs deleted file mode 100644 index 9a1d26fb0..000000000 --- a/crates/ml/src/dqn/ensemble_oracle.rs +++ /dev/null @@ -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) - transformer: Option>, - /// Placeholder for LSTM model (future: Arc) - lstm: Option>, - /// Placeholder for PPO policy (future: Arc) - ppo: Option>, - /// 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> { - // 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, - ) -> 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 = 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" - ); - } -} diff --git a/crates/ml/src/dqn/ensemble_uncertainty.rs b/crates/ml/src/dqn/ensemble_uncertainty.rs deleted file mode 100644 index 32ca71f98..000000000 --- a/crates/ml/src/dqn/ensemble_uncertainty.rs +++ /dev/null @@ -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>(()) -//! ``` -//! -//! # 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, - - /// Vote counts for each action (0=Buy, 1=Sell, 2=Hold) - pub vote_counts: Vec, - - /// 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, - - /// 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 { - 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 { - 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 { - 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::() / 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> { - 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::()?; - q_vals.push(val as f64); - } - - // Compute variance: Var[X] = E[X²] - E[X]² - let mean = q_vals.iter().sum::() / num_agents as f64; - let variance = q_vals.iter().map(|x| (x - mean).powi(2)).sum::() - / num_agents as f64; - - variances.push(variance); - } - - Ok(variances) - } - - /// Extract action indices (argmax) from Q-values - fn extract_actions(&self, q_values: &[Tensor]) -> Result> { - 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::()?; - 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)> { - // 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 { - 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 { - 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::() / count; - let avg_disagreement = recent.iter().map(|m| m.action_disagreement).sum::() / count; - let avg_entropy = recent.iter().map(|m| m.action_entropy).sum::() / 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], - ) -> Result> { - 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(()) - } -} diff --git a/crates/ml/src/dqn/factored_q_network.rs b/crates/ml/src/dqn/factored_q_network.rs deleted file mode 100644 index da5afba50..000000000 --- a/crates/ml/src/dqn/factored_q_network.rs +++ /dev/null @@ -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::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 { - 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 { - // 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::() { - 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 { - 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::() - .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 { - let mut rng = rand::thread_rng(); - - if rng.gen::() < 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 { - 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::() - .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::().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::().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::().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::() - .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]); - } - } -} diff --git a/crates/ml/src/dqn/mod.rs b/crates/ml/src/dqn/mod.rs index 80f61f07f..59a9c58c8 100644 --- a/crates/ml/src/dqn/mod.rs +++ b/crates/ml/src/dqn/mod.rs @@ -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 diff --git a/crates/ml/src/dqn/regime_temperature.rs b/crates/ml/src/dqn/regime_temperature.rs deleted file mode 100644 index 07210cc90..000000000 --- a/crates/ml/src/dqn/regime_temperature.rs +++ /dev/null @@ -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 { - 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, -) -> 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, - 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); - } -} diff --git a/crates/ml/src/dqn/risk_integration.rs b/crates/ml/src/dqn/risk_integration.rs deleted file mode 100644 index 585576a3f..000000000 --- a/crates/ml/src/dqn/risk_integration.rs +++ /dev/null @@ -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>, - /// Daily P&L accumulator (reset each epoch) - daily_pnl: Arc>, - /// Current positions (empty for training) - positions: Arc>>, -} - -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 { - Ok(*self.portfolio_value.read().await) - } - - async fn get_daily_pnl(&self, _account_id: &str) -> RiskResult { - Ok(*self.daily_pnl.read().await) - } - - async fn get_positions(&self, _account_id: &str) -> RiskResult> { - 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, - /// 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 { - 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 { - 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 { - self.breaker.get_state(&self.account_id).await - } - - /// Get circuit breaker metrics - pub async fn get_metrics(&self) -> std::collections::HashMap { - 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()); - } -} diff --git a/crates/ml/src/dqn/spectral_norm.rs b/crates/ml/src/dqn/spectral_norm.rs deleted file mode 100644 index 3f7767988..000000000 --- a/crates/ml/src/dqn/spectral_norm.rs +++ /dev/null @@ -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 { - 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 { - // 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::() - .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::() - .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::() - .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 { - 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 { - self.compute_spectral_norm() - } -} - -impl Module for SpectralNorm { - fn forward(&self, xs: &Tensor) -> CandleResult { - // 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::()?; - - // 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::()?; - - // Get normalized weight - let normalized_weight = spectral_norm.normalized_weight()?; - let normalized_norm = normalized_weight.sqr()?.sum_all()?.sqrt()?.to_vec0::()?; - - // 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(()) - } -}