diff --git a/ml/src/dqn/dqn.rs.backup b/ml/src/dqn/dqn.rs.backup deleted file mode 100644 index 9688ff5a2..000000000 --- a/ml/src/dqn/dqn.rs.backup +++ /dev/null @@ -1,1990 +0,0 @@ -//! ACTUAL Working Deep Q-Network Implementation -//! -//! This module provides a complete, working DQN implementation with: -//! - Real mathematical operations using candle-core v0.9.1 -//! - Experience replay buffer with proper memory management -//! - Epsilon-greedy exploration with decay -//! - Target network updates with soft/hard copying -//! - Proper Q-learning update with Bellman equation -//! - NO productions, todo!(), or unimplemented!() macros - -use std::collections::VecDeque; -use std::sync::{Arc, Mutex}; -use std::time::SystemTime; - -use crate::Adam; -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 candle_core::IndexOp; -use candle_core::{DType, Device, Tensor, Var}; -use candle_nn::Module; -use candle_nn::{Linear, VarBuilder, VarMap}; -use candle_optimisers::adam::ParamsAdam; -use candle_nn::ops::leaky_relu; -// use crate::Optimizer; // Optimizer trait not available in candle v0.9 -use rand::{thread_rng, Rng}; -use serde::{Deserialize, Serialize}; -use tracing::debug; - -use super::{Experience, TradingAction}; -use crate::MLError; - -/// Gradient-safe soft clamp using tanh activation -/// -/// CRITICAL FIX: Candle 0.9.1's `.clamp()` has NO backward pass implementation, -/// causing gradient collapse (grad_norm=0.0000). This function uses tanh to -/// provide gradient-preserving bounds. -/// -/// # Arguments -/// * `tensor` - Input tensor to clamp -/// * `bound` - Symmetric bounds [-bound, +bound] -/// -/// # Returns -/// Softly clamped tensor with gradients preserved -/// -/// # Mathematical Properties -/// - f(x) = bound * tanh(x / bound) -/// - Range: (-bound, +bound) (asymptotic, not hard limits) -/// - Gradient: f'(x) = sech²(x / bound) > 0 (always positive, never zero) -/// - At x=0: f(0) = 0, f'(0) = 1 -/// - At x=±∞: f(±∞) → ±bound, f'(±∞) → 0 -/// -/// # Example -/// ```rust -/// use candle_core::{Device, Tensor}; -/// use ml::dqn::soft_clamp; -/// -/// let device = Device::Cpu; -/// let x = Tensor::from_vec(vec![1.0f32, 10.0, 100.0, 1000.0], 4, &device)?; -/// let clamped = soft_clamp(&x, 10.0)?; -/// // Values: [0.996, 7.616, 9.505, 10.000] (smoothly approaches 10.0) -/// # Ok::<(), candle_core::Error>(()) -/// ``` -fn soft_clamp(tensor: &Tensor, bound: f64) -> Result { - // Scale input: x / bound - let scaled = (tensor / bound) - .map_err(|e| MLError::ModelError(format!("Soft clamp scaling failed: {}", e)))?; - - // Apply tanh (range: -1 to +1, gradients preserved) - let clamped = scaled.tanh() - .map_err(|e| MLError::ModelError(format!("Soft clamp tanh failed: {}", e)))?; - - // Scale back: bound * tanh(x / bound) - (clamped * bound) - .map_err(|e| MLError::ModelError(format!("Soft clamp rescaling failed: {}", e))) -} - -// Factored action space support (feature-gated) -#[cfg(feature = "factored-actions")] -use super::action_space::{ExposureLevel, FactoredAction, OrderType, Urgency}; -#[cfg(feature = "factored-actions")] -use super::factored_q_network::FactoredQNetwork; - -/// Reward system type for DQN training -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -pub enum RewardSystem { - /// Elite multi-component reward (default, production-tested) - /// Combines extrinsic P&L, intrinsic rewards, entropy, curiosity, and ensemble - Elite, - /// Simple P&L-only reward (experimental, bias elimination testing) - /// Pure profit/loss without activity incentives or risk adjustments - SimplePnL, -} - -impl Default for RewardSystem { - fn default() -> Self { - RewardSystem::Elite - } -} - -/// Configuration for the working `DQN` -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WorkingDQNConfig { - /// State dimension - pub state_dim: usize, - /// Number of actions - pub num_actions: usize, - /// Hidden layer dimensions - pub hidden_dims: Vec, - /// Learning rate - pub learning_rate: f64, - /// Discount factor (gamma) - pub gamma: f32, - /// Exploration parameters - pub epsilon_start: f32, - pub epsilon_end: f32, - pub epsilon_decay: f32, - /// Experience replay parameters - pub replay_buffer_capacity: usize, - pub batch_size: usize, - pub min_replay_size: usize, - /// Target network update frequency - pub target_update_freq: usize, - /// Whether to use double `DQN` - pub use_double_dqn: bool, - /// Whether to use Huber loss instead of MSE - pub use_huber_loss: bool, - /// Delta parameter for Huber loss - pub huber_delta: f32, - /// LeakyReLU negative slope (alpha) to prevent dead neurons - pub leaky_relu_alpha: f64, - /// Gradient clipping max norm (Wave 11 Bug #1 fix) - pub gradient_clip_norm: f64, - /// TD-error clipping threshold (Wave 4 Agent 1: prevents noise amplification) - /// Clamps TD-errors to [-td_error_clip, +td_error_clip] before loss calculation - /// Default: 10.0 (prevents ±800 oscillations observed in SimplePnL) - pub td_error_clip: f64, - - // WAVE 16 (Agent 36): Target update configuration - /// Polyak averaging coefficient (default: 0.001) - pub tau: f64, - /// Use soft (Polyak) or hard target updates - pub use_soft_updates: bool, - - // Rainbow DQN warmup period - /// Number of steps to collect experiences with random exploration before training begins - /// Rainbow DQN standard: 80,000 steps (prevents early overfitting to sparse data) - pub warmup_steps: usize, - - // WAVE 3 (Agent 1): Softmax action selection with temperature annealing - /// Initial temperature for softmax action selection (default: 1.0 = balanced exploration) - pub temperature_start: f64, - /// Minimum temperature for softmax action selection (default: 0.1 = more greedy) - pub temperature_min: f64, // Default: 0.3 (Wave 3 Agent 2: prevents over-exploitation) - /// Temperature decay rate per epoch (default: 0.995 = match epsilon decay) - pub temperature_decay: f64, - /// Target epoch fraction for temperature convergence (default: 0.75 = 75% of training) - /// Temperature will reach minimum at this percentage of total training epochs - /// Example: 1000 epochs * 0.75 = 750 epochs to reach minimum - pub target_temperature_fraction: f64, - - // WAVE 2 (Agent 2B): Q-Value Variance Adaptation - /// Variance multiplier for adaptive temperature (default: 0.5, range: 0.0-2.0) - pub variance_multiplier: f64, - - // WAVE 2 (Agent 2A): Performance-Based Adaptive Temperature Decay - /// Enable adaptive temperature based on validation loss improvement - pub use_adaptive_temperature: bool, - /// Loss improvement threshold (default: 0.999 = 0.1% improvement) - pub loss_improvement_threshold: f32, - /// Plateau detection window (epochs without improvement, default: 10) - pub plateau_window: usize, - /// Temperature increase factor when stuck (default: 1.05 = 5% increase) - pub temp_increase_factor: f64, - /// Slow decay rate when loss plateaus (default: 0.998) - pub temperature_slow_decay: f64, - - // Reward system selection - /// Reward system to use (Elite or SimplePnL) - /// Elite: Multi-component reward with activity incentives (default, production-tested) - /// SimplePnL: Pure P&L-only reward for bias elimination testing - #[serde(default)] - pub reward_system: RewardSystem, - - // WAVE 4 (Agent 2): Reward normalization - /// Reward normalization scale (typical reward magnitude in dollars) - /// Default: 1000.0 (typical ES futures move) - /// Used to normalize raw P&L rewards to [-1, +1] range via tanh(reward / scale) - /// Prevents Q-value explosion from large P&L variance - pub reward_scale: f64, -} - -impl WorkingDQNConfig { - /// Calculate optimal temperature decay rate for target convergence - /// - /// Formula: decay = (temp_min / temp_start) ^ (1 / target_epochs) - /// - /// # Arguments - /// * `total_epochs` - Total number of training epochs - /// * `temp_start` - Initial temperature (e.g., 1.0) - /// * `temp_min` - Minimum temperature (e.g., 0.1) - /// * `target_fraction` - Fraction of training to reach minimum (e.g., 0.75 = 75%) - /// - /// # Returns - /// Optimal decay rate per epoch - /// - /// # Example - /// ``` - /// use ml::dqn::WorkingDQNConfig; - /// let decay = WorkingDQNConfig::calculate_optimal_temperature_decay(1000, 1.0, 0.1, 0.75); - /// assert!((decay - 0.9969).abs() < 0.0001); // ~0.9969 for 750 epochs - /// ``` - pub fn calculate_optimal_temperature_decay( - total_epochs: usize, - temp_start: f64, - temp_min: f64, - target_fraction: f64, - ) -> f64 { - let target_epochs = (total_epochs as f64 * target_fraction) as usize; - (temp_min / temp_start).powf(1.0 / target_epochs as f64) - } - - /// Create `DQN` config from central configuration system - /// - /// CRITICAL: Eliminates dangerous hardcoded defaults - pub fn from_config_manager( - _config_manager: &config::ConfigManager, - ) -> Result> { - // Use emergency defaults since specific DQN configs may not be available - tracing::warn!( - "Using emergency DQN config defaults - DQN configs not available in ServiceConfig" - ); - Ok(Self::emergency_safe_defaults()) - } - - /// EMERGENCY FALLBACK: Ultra-conservative `DQN` defaults - /// - /// WARNING: These defaults prioritize safety over performance - pub fn emergency_safe_defaults() -> Self { - tracing::error!("Using emergency DQN defaults - check configuration system immediately!"); - Self { - state_dim: 32, // Smaller state space - num_actions: 3, // Conservative action space - hidden_dims: vec![512, 256, 128, 64], // Increased capacity: 32→512→256→128→64→3 (better Q-function approximation) - learning_rate: 1e-5, // Very conservative learning rate - gamma: 0.90, // Conservative discount factor (56% noise reduction vs 0.9626) - epsilon_start: 0.1, // Low exploration to prevent erratic behavior - epsilon_end: 0.01, // Minimal exploration - epsilon_decay: 0.99, // Fast decay to reach stable exploitation - replay_buffer_capacity: 1000, // Small buffer to prevent memory issues - batch_size: 4, // Very small batch size - min_replay_size: 100, // Minimal replay requirement - target_update_freq: 100, // Frequent updates for stability - use_double_dqn: false, // Disable advanced features for safety - use_huber_loss: true, // Huber loss default (more robust to outliers) - huber_delta: 10.0, // Handles larger TD errors (up to ±10) - leaky_relu_alpha: 0.01, // Standard LeakyReLU alpha - gradient_clip_norm: 10.0, // Conservative clipping for emergency defaults - td_error_clip: 10.0, // TD-error clipping (Wave 4 Agent 1: prevents noise amplification) - - // WAVE 16 (Agent 36): Target update defaults (REVERTED to Hard updates for stability) - tau: 1.0, // No Polyak averaging (hard updates) - use_soft_updates: false, // Hard updates by default (original DQN standard) - - // Rainbow DQN warmup period - warmup_steps: 0, // No warmup for emergency mode (safety first) - - // WAVE 3 (Agent 1): Softmax temperature defaults - temperature_start: 1.0, // Balanced softmax - temperature_min: 0.3, // More greedy (Wave 3 Agent 2: prevents over-exploitation) - temperature_decay: 0.995, // Will be recalculated in WorkingDQN::new() - target_temperature_fraction: 0.75, // Reach minimum at 75% of training - - // WAVE 2 (Agent 2B): Q-Value Variance Adaptation - variance_multiplier: 0.5, // Conservative variance scaling - - // WAVE 2 (Agent 2A): Adaptive temperature defaults - use_adaptive_temperature: false, // Disabled by default - loss_improvement_threshold: 0.999, // 0.1% improvement - plateau_window: 10, // 10 epochs stuck = increase temp - temp_increase_factor: 1.05, // 5% increase to escape local optimum - temperature_slow_decay: 0.998, // Slower than fast decay (0.995) - - // Reward system (default to Elite for production stability) - reward_system: RewardSystem::Elite, - - // WAVE 4 (Agent 2): Reward normalization - reward_scale: 1000.0, // Typical ES futures move ($1000) - } - } -} - -/// Experience replay buffer for `DQN` -#[derive(Debug)] -pub struct ExperienceReplayBuffer { - buffer: VecDeque, - capacity: usize, -} - -impl ExperienceReplayBuffer { - /// Create new replay buffer - pub fn new(capacity: usize) -> Self { - Self { - buffer: VecDeque::with_capacity(capacity), - capacity, - } - } - - /// Add experience to buffer - pub fn push(&mut self, experience: Experience) { - if self.buffer.len() >= self.capacity { - self.buffer.pop_front(); - } - self.buffer.push_back(experience); - } - - /// Sample random batch of experiences - pub fn sample(&self, batch_size: usize) -> Result, MLError> { - if self.buffer.len() < batch_size { - return Err(MLError::TrainingError(format!( - "Not enough experiences in buffer: {} < {}", - self.buffer.len(), - batch_size - ))); - } - - let mut rng = thread_rng(); - let mut batch = Vec::with_capacity(batch_size); - - for _ in 0..batch_size { - let idx = rng.gen_range(0..self.buffer.len()); - batch.push(self.buffer[idx].clone()); - } - - Ok(batch) - } - - /// Get current buffer size - pub fn len(&self) -> usize { - self.buffer.len() - } - - /// Check if buffer can sample - pub fn can_sample(&self, min_size: usize) -> bool { - self.buffer.len() >= min_size - } -} - -/// Sequential neural network for Q-value approximation -#[allow(missing_debug_implementations)] -pub struct Sequential { - layers: Vec, - device: Device, - vars: VarMap, - leaky_relu_alpha: f64, -} - -impl Sequential { - /// Create new sequential network - pub fn new( - input_dim: usize, - hidden_dims: &[usize], - output_dim: usize, - device: Device, - leaky_relu_alpha: f64, - ) -> Result { - let vars = VarMap::new(); - let var_builder = VarBuilder::from_varmap(&vars, DType::F32, &device); - - let mut layers = Vec::new(); - let mut current_dim = input_dim; - - // Hidden layers - for (i, &hidden_dim) in hidden_dims.into_iter().enumerate() { - // Use Xavier initialization with VarMap registration - let layer_name = format!("hidden_{}", i); - let layer_vb = var_builder.pp(&layer_name); - let layer = linear_xavier(current_dim, hidden_dim, layer_vb) - .map_err(|e| MLError::ModelError(format!("Failed to Xavier init layer {}: {}", i, e)))?; - - layers.push(layer); - current_dim = hidden_dim; - } - - // Output layer - also use Xavier initialization with VarMap registration - let output_vb = var_builder.pp("output"); - let output_layer = linear_xavier(current_dim, output_dim, output_vb) - .map_err(|e| MLError::ModelError(format!("Failed to Xavier init output layer: {}", e)))?; - - layers.push(output_layer); - - Ok(Self { - layers, - device, - vars, - leaky_relu_alpha, - }) - } - - /// Forward pass through network - pub fn forward(&self, input: &Tensor) -> Result { - let mut x = input.clone(); - - // Pass through hidden layers with ReLU activation - let num_layers = self.layers.len(); - for (i, layer) in self.layers.iter().enumerate() { - x = layer.forward(&x).map_err(|e| { - MLError::ModelError(format!("Forward pass failed at layer {}: {}", i, e)) - })?; - - // Apply LeakyReLU to all layers except the last - if i < num_layers - 1 { - x = leaky_relu(&x, self.leaky_relu_alpha) - .map_err(|e| MLError::ModelError(format!("LeakyReLU activation failed: {}", e)))?; - } - } - - Ok(x) - } - - /// Get network variables - pub fn vars(&self) -> &VarMap { - &self.vars - } - - /// Get device - pub fn device(&self) -> &Device { - &self.device - } - - /// Copy weights from another network - pub fn copy_weights_from(&mut self, other: &Sequential) -> Result<(), MLError> { - let self_vars = self - .vars - .data() - .lock() - .map_err(|e| MLError::ConcurrencyError { - operation: format!("lock self vars: {}", e), - })?; - let other_vars = other - .vars - .data() - .lock() - .map_err(|e| MLError::ConcurrencyError { - operation: format!("lock other vars: {}", e), - })?; - - for (name, self_var) in self_vars.iter() { - if let Some(other_var) = other_vars.get(name) { - let other_tensor = other_var.as_tensor(); - self_var.set(other_tensor).map_err(|e| { - MLError::ModelError(format!("Failed to copy weight {}: {}", name, e)) - })?; - } - } - - Ok(()) - } -} - -/// Working Deep Q-Network implementation -#[allow(missing_debug_implementations)] -pub struct WorkingDQN { - /// `DQN` configuration - config: WorkingDQNConfig, - /// Main Q-network - q_network: Sequential, - /// Target Q-network for stable training - target_network: Sequential, - /// Experience replay buffer (public for trainer access) - pub memory: Arc>, - /// Current exploration rate - epsilon: f32, - /// Training step counter (gradient updates) - training_steps: u64, - /// Total environment steps counter (includes warmup period) - total_steps: u64, - /// Optimizer for main network - optimizer: Option, - /// Device (CPU or CUDA GPU) - device: Device, - /// Gradient clipping max norm (Wave 11 Bug #1 fix) - gradient_clip_norm: f64, - /// Recent actions for entropy penalty calculation (sliding window) - recent_actions: VecDeque, - /// Current temperature for softmax action selection (Wave 3 Agent 1) - temperature: f64, - - // WAVE 2 (Agent 2A): Adaptive temperature state - /// Sliding window of recent training losses (last 5 epochs) - loss_window: VecDeque, - /// Best average loss seen so far - best_loss: f32, - /// Number of epochs without improvement - plateau_count: usize, - - // Factored action space support (feature-gated) - #[cfg(feature = "factored-actions")] - /// Factored Q-network with 3 heads (45 actions) - factored_network: Option, - #[cfg(feature = "factored-actions")] - /// Current position for action masking (-1.0 to +1.0) - current_position: f64, -} - -impl WorkingDQN { - /// Create new working `DQN` - pub fn new(config: WorkingDQNConfig) -> Result { - let device = Device::cuda_if_available(0)?; // Use GPU if available, fallback to CPU - - // Seed the device RNG with entropy to avoid deterministic initialization - // CUDA backend uses hardcoded seed 299792458 by default, causing identical weights across runs - let entropy_seed = Self::generate_entropy_seed(); - device.set_seed(entropy_seed).map_err(|e| { - MLError::ModelError(format!("Failed to seed device RNG: {}", e)) - })?; - debug!("Device RNG seeded with entropy: {}", entropy_seed); - - // Extract config values before moving config into Self - let epsilon_start = config.epsilon_start; - let gradient_clip_norm = config.gradient_clip_norm; - let temperature_start = config.temperature_start; - - // Create main Q-network - let q_network = Sequential::new( - config.state_dim, - &config.hidden_dims, - config.num_actions, - device.clone(), - config.leaky_relu_alpha, - )?; - - // Create target network (copy of main network) - let mut target_network = Sequential::new( - config.state_dim, - &config.hidden_dims, - config.num_actions, - device.clone(), - config.leaky_relu_alpha, - )?; - - // Copy initial weights to target network - target_network.copy_weights_from(&q_network)?; - - // Create experience replay buffer - let memory = Arc::new(Mutex::new(ExperienceReplayBuffer::new( - config.replay_buffer_capacity, - ))); - - #[allow(unused_mut)] - let mut result = Self { - config, - q_network, - target_network, - memory, - epsilon: epsilon_start, - training_steps: 0, - total_steps: 0, - optimizer: None, - device, - gradient_clip_norm, - recent_actions: VecDeque::with_capacity(100), - temperature: temperature_start, - loss_window: VecDeque::with_capacity(5), - best_loss: f32::MAX, // Start with worst possible loss - plateau_count: 0, - - #[cfg(feature = "factored-actions")] - factored_network: None, - #[cfg(feature = "factored-actions")] - current_position: 0.0, - }; - - Ok(result) - } - - /// Generate entropy seed from system time and thread/process info - /// - /// Combines multiple sources of entropy to ensure different initialization across: - /// - Parallel training runs - /// - Sequential restarts - /// - Different machines - /// - /// Returns a 64-bit seed suitable for RNG initialization - fn generate_entropy_seed() -> u64 { - // Get nanosecond timestamp as base entropy - let timestamp = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .expect("System time is before Unix epoch") - .as_nanos() as u64; - - // Mix in process ID (available on all platforms) - let process_id = std::process::id() as u64; - - // Mix in thread-local random value - let mut rng = thread_rng(); - let thread_entropy: u64 = rng.gen(); - - // Combine all entropy sources with XOR and bit rotation - // This ensures changes in any source affect the final seed - timestamp - .wrapping_mul(6364136223846793005) // LCG multiplier for mixing - .wrapping_add(process_id) - .rotate_left(13) - ^ thread_entropy - } - - /// Get the device this DQN is using (CPU or CUDA) - pub fn device(&self) -> &Device { - &self.device - } - - /// Get configuration (for ensemble access) - pub fn config(&self) -> &WorkingDQNConfig { - &self.config - } - - /// Forward pass through main network - pub fn forward(&self, state: &Tensor) -> Result { - // Auto-convert input to correct device (Candle optimizes if already on correct device) - let state = state - .to_device(&self.device) - .map_err(|e| MLError::ModelError(format!("Failed to move tensor to device: {}", e)))?; - - // DEBUG: Log input shape - tracing::info!("Q-network input shape: {:?}", state.dims()); - - let q_values = self.q_network.forward(&state)?; - - // DEBUG: Log output shape and values - tracing::info!("Q-network output shape: {:?}", q_values.dims()); - if let Ok(q_vec) = q_values.flatten_all()?.to_vec1::() { - let num_to_show = 10.min(q_vec.len()); - tracing::info!("Q-values (first {}): {:?}", num_to_show, &q_vec[..num_to_show]); - - // VALIDATION: Check Q-value count matches action space - #[cfg(feature = "factored-actions")] - { - if q_vec.len() != 45 { - tracing::error!( - "Q-values shape mismatch! Expected 45 for factored actions, got {}", - q_vec.len() - ); - } - } - #[cfg(not(feature = "factored-actions"))] - { - if q_vec.len() != 3 { - tracing::error!( - "Q-values shape mismatch! Expected 3 for standard actions, got {}", - q_vec.len() - ); - } - } - } - - // CRITICAL FIX: Use soft_clamp to preserve gradients - // Hard clamp (.clamp()) breaks gradient flow in Candle 0.9.1 - let clamped = soft_clamp(&q_values, 1000.0)?; - Ok(clamped) - } - - /// Select action using epsilon-greedy policy - pub fn select_action(&mut self, state: &[f32]) -> Result { - // Increment total steps counter (tracks all environment steps including warmup) - self.total_steps += 1; - - let mut rng = thread_rng(); - - // During warmup period: always use random exploration (epsilon=1.0) - let in_warmup = self.total_steps <= self.config.warmup_steps as u64; - - // DEBUG: Log action selection context - let random_val = rng.gen::(); - tracing::info!( - "Action selection: epsilon={:.3}, random={:.3}, random 0.0 { - let q_vec = q_values.flatten_all()?.to_vec1::()?; - - // Compute variance: var = mean((Q - mean(Q))^2) - let q_mean = q_vec.iter().sum::() / q_vec.len() as f32; - let variance: f32 = q_vec.iter() - .map(|q| { - let diff = q - q_mean; - diff * diff - }) - .sum::() / q_vec.len() as f32; - - // Normalize variance by absolute mean to handle negative Q-values - let abs_mean = q_vec.iter().map(|q| q.abs()).sum::() / q_vec.len() as f32; - let safe_mean = abs_mean.max(0.1); // Prevent division by zero - let normalized_variance = variance.sqrt() / safe_mean; - - // Scale temperature: temp_adaptive = temp_base * (1 + variance_multiplier * normalized_variance) - let variance_scale = 1.0 + (self.config.variance_multiplier * normalized_variance as f64); - let temp = self.temperature * variance_scale; - - // Clamp to reasonable bounds [min, 2 * start] - temp.clamp(self.config.temperature_min, self.config.temperature_start * 2.0) - } else { - self.temperature // No variance adaptation - }; - - // Apply adaptive temperature scaling and softmax - let logits = (q_values / adaptive_temp)?; - let probs = candle_nn::ops::softmax(&logits, 1)?; - - // Sample from the probability distribution (manual categorical sampling) - let probs_vec = probs - .flatten_all()? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Failed to extract probabilities: {}", e)))?; - - // DEBUG: Log softmax probabilities - tracing::info!("Softmax probabilities: {:?}", probs_vec); - - let sample: f32 = rng.gen(); - let mut cumulative = 0.0; - let mut action_idx = 2; // Default to HOLD (index 2) - - for (i, &prob) in probs_vec.iter().enumerate() { - cumulative += prob; - if sample < cumulative { - action_idx = i; - break; - } - } - - // DEBUG: Log selected action - tracing::info!( - "Greedy action selected: index={}, sample={:.3}, cumulative={:.3}", - action_idx, sample, cumulative - ); - - TradingAction::from_int(action_idx as u8).ok_or_else(|| { - MLError::InvalidInput(format!("Invalid action index: {}", action_idx)) - })? - }; - - // Track action for entropy penalty calculation - self.track_action(action); - - // Log warmup progress every 10K steps - if in_warmup && self.total_steps % 10000 == 0 { - tracing::info!( - "Warmup: {}/{}K steps ({:.1}% complete)", - self.total_steps / 1000, - self.config.warmup_steps / 1000, - (self.total_steps as f64 / self.config.warmup_steps as f64) * 100.0 - ); - } - - // Log warmup completion - if self.total_steps == self.config.warmup_steps as u64 { - tracing::info!( - "✓ Warmup complete - starting training ({}K steps collected)", - self.config.warmup_steps / 1000 - ); - } - - Ok(action) - } - - /// Track action for entropy penalty calculation - /// - /// This should be called after action selection (both single and batch modes) - /// to maintain the sliding window used for diversity penalty. - pub fn track_action(&mut self, action: TradingAction) { - self.recent_actions.push_back(action); - if self.recent_actions.len() > 100 { - self.recent_actions.pop_front(); - } - } - - /// Track multiple actions for entropy penalty calculation (batch version) - pub fn track_actions_batch(&mut self, actions: &[TradingAction]) { - for &action in actions { - self.track_action(action); - } - } - - /// Store experience in replay buffer - pub fn store_experience(&self, experience: Experience) -> Result<(), MLError> { - let mut buffer = self.memory.lock().map_err(|e| MLError::ConcurrencyError { - operation: format!("lock memory buffer: {}", e), - })?; - buffer.push(experience); - Ok(()) - } - - /// Training step with experience batch - /// - /// Returns (loss, gradient_norm) tuple - pub fn train_step(&mut self, batch: Option>) -> Result<(f32, f32), MLError> { - // Skip gradient updates during warmup period - if self.total_steps < self.config.warmup_steps as u64 { - return Ok((0.0, 0.0)); // Return dummy values, no training during warmup - } - - // Get batch of experiences - let experiences = if let Some(batch) = batch { - batch - } else { - let buffer = self.memory.lock().map_err(|e| MLError::ConcurrencyError { - operation: format!("lock memory buffer for training: {}", e), - })?; - if !buffer.can_sample(self.config.min_replay_size) { - return Err(MLError::TrainingError( - "Not enough experiences for training".to_string(), - )); - } - buffer.sample(self.config.batch_size)? - }; - - // Initialize optimizer if not done - if self.optimizer.is_none() { - // WAVE 16H: Use Rainbow DQN Adam epsilon (1.5e-4) for numerical stability - // Standard PyTorch eps=1e-8 can cause division instability with normalized features - // Rainbow DQN paper uses 1.5e-4 to prevent optimizer instability - let adam_params = ParamsAdam { - lr: self.config.learning_rate, - beta_1: 0.9, - beta_2: 0.999, - eps: 1.5e-4, // Rainbow DQN standard (was 1e-8) - weight_decay: None, - amsgrad: false, - }; - self.optimizer = Some( - Adam::new(self.q_network.vars().all_vars(), adam_params).map_err(|e| { - MLError::TrainingError(format!("Failed to create optimizer: {}", e)) - })?, - ); - } - - // Convert experiences to tensors - let batch_size = experiences.len(); - let device = self.q_network.device(); - - // WAVE 5-E AGENT 2: Add NaN/Inf diagnostics for gradient collapse debugging - // Check for non-finite values in experiences (could cause zero gradients) - for (i, exp) in experiences.iter().enumerate() { - // Check state for NaN/Inf - if exp.state.iter().any(|&x| !x.is_finite()) { - tracing::warn!( - "⚠️ Experience #{} has non-finite STATE values (NaN/Inf detected)", - i - ); - } - // Check next_state for NaN/Inf - if exp.next_state.iter().any(|&x| !x.is_finite()) { - tracing::warn!( - "⚠️ Experience #{} has non-finite NEXT_STATE values (NaN/Inf detected)", - i - ); - } - // Check reward (convert to f32 first) - let reward_f32 = exp.reward_f32(); - if !reward_f32.is_finite() { - tracing::warn!( - "⚠️ Experience #{} has non-finite REWARD: {} (NaN/Inf detected)", - i, reward_f32 - ); - } - } - - // OPTIMIZATION: Single-pass data extraction for 5-10% throughput improvement - // Instead of 5 separate iterator passes, do one fold operation - let state_dim = self.config.state_dim; - let (states, next_states, actions, rewards, dones) = experiences.iter().fold( - ( - Vec::with_capacity(batch_size * state_dim), - Vec::with_capacity(batch_size * state_dim), - Vec::with_capacity(batch_size), - Vec::with_capacity(batch_size), - Vec::with_capacity(batch_size), - ), - |(mut s, mut ns, mut a, mut r, mut d), exp| { - s.extend_from_slice(&exp.state); - ns.extend_from_slice(&exp.next_state); - a.push(exp.action as u32); - r.push(exp.reward_f32()); - d.push(if exp.done { 1.0_f32 } else { 0.0_f32 }); - (s, ns, a, r, d) - }, - ); - - // Create tensors - let states_tensor = Tensor::from_vec(states, (batch_size, self.config.state_dim), device) - .map_err(|e| { - MLError::TrainingError(format!("Failed to create states tensor: {}", e)) - })?; - - let next_states_tensor = - Tensor::from_vec(next_states, (batch_size, self.config.state_dim), device).map_err( - |e| MLError::TrainingError(format!("Failed to create next states tensor: {}", e)), - )?; - - let actions_tensor = Tensor::from_vec(actions, batch_size, device).map_err(|e| { - MLError::TrainingError(format!("Failed to create actions tensor: {}", e)) - })?; - - let rewards_tensor = Tensor::from_vec(rewards, batch_size, device).map_err(|e| { - MLError::TrainingError(format!("Failed to create rewards tensor: {}", e)) - })?; - - let dones_tensor = Tensor::from_vec(dones, batch_size, device) - .map_err(|e| MLError::TrainingError(format!("Failed to create dones tensor: {}", e)))?; - - // Forward pass through main network to get current Q-values - let current_q_values = self.q_network.forward(&states_tensor)?; - // CRITICAL FIX: Use soft_clamp instead of .clamp() to preserve gradients - // Hard clamp breaks gradient flow in Candle 0.9.1 (no backward pass) - let clamped_q = soft_clamp(¤t_q_values, 1000.0)?; - - // Get Q-values for taken actions - let actions_unsqueezed = actions_tensor.unsqueeze(1)?; - let state_action_values = clamped_q - .gather(&actions_unsqueezed, 1)? - .squeeze(1)? - .to_dtype(DType::F32)?; - - // Compute target Q-values using target network - let next_q_values = self.target_network.forward(&next_states_tensor)?; - - let next_state_values = if self.config.use_double_dqn { - // Double DQN: use main network to select action, target network to evaluate - let next_q_main = self.q_network.forward(&next_states_tensor)?; - let next_actions = next_q_main.argmax(1)?; - let next_actions_unsqueezed = next_actions.unsqueeze(1)?; - let values = next_q_values - .gather(&next_actions_unsqueezed, 1)? - .squeeze(1)?; - values.to_dtype(DType::F32)? - } else { - // Standard DQN: use max Q-value from target network - // Note: max(1) already returns a 1D tensor, no need to squeeze - // Ensure F32 dtype to match other tensors - let values = next_q_values.max(1)?; - values.to_dtype(DType::F32)? - }; - - // Compute target values using Bellman equation - // target = reward + gamma * next_state_value * (1 - done) - let gamma_tensor = - Tensor::from_vec(vec![self.config.gamma; batch_size], batch_size, device).map_err( - |e| MLError::TrainingError(format!("Failed to create gamma tensor: {}", e)), - )?; - - let not_done = (Tensor::ones(&[batch_size], DType::F32, device)? - &dones_tensor)?; - let gamma_next = (&gamma_tensor * &next_state_values) - .map_err(|e| MLError::TrainingError(format!("Gamma multiplication failed: {}", e)))?; - let discounted = (&gamma_next * ¬_done)?; - let target_q_values = (&rewards_tensor + &discounted)?.detach(); // Stop gradient computation - - // Compute loss (Huber loss if enabled, MSE otherwise) - // Ensure both tensors have the same dtype (F32) - let target_q_values = target_q_values.to_dtype(DType::F32)?; - let diff_raw = state_action_values.sub(&target_q_values)?; - - // WAVE 4 (Agent 1): Clip TD-error to prevent noise amplification - // Bellman update bootstraps from noisy Q-estimates, causing ±800 oscillations in SimplePnL - // CRITICAL FIX: Use soft_clamp instead of hard clamp to preserve gradients - let diff = soft_clamp(&diff_raw, self.config.td_error_clip)?; - - let loss_value = if self.config.use_huber_loss { - // Huber loss: L(x) = 0.5 * x^2 if |x| <= delta, else delta * (|x| - 0.5*delta) - let delta = self.config.huber_delta; - let abs_diff = diff.abs()?; - - // Element-wise Huber loss - let squared_loss = ((&diff * &diff)? * 0.5)?; // 0.5 * x^2 - - // Create delta tensor for operations - let delta_tensor = Tensor::from_vec( - vec![delta; batch_size], - batch_size, - device - ).map_err(|e| MLError::TrainingError(format!("Failed to create delta tensor: {}", e)))?; - - let linear_loss_term1 = (&abs_diff * &delta_tensor)?; - let linear_loss_term2 = delta * delta * 0.5; - let linear_loss_term2_tensor = Tensor::from_vec( - vec![linear_loss_term2; batch_size], - batch_size, - device - ).map_err(|e| MLError::TrainingError(format!("Failed to create linear term tensor: {}", e)))?; - let linear_loss = (linear_loss_term1 - &linear_loss_term2_tensor)?; // delta * (|x| - 0.5*delta) - - // Condition: use squared if |x| <= delta, else linear - let mask = abs_diff.le(delta)?.to_dtype(DType::F32)?; // 1.0 if |x| <= delta, 0.0 otherwise - let one_minus_mask = (Tensor::ones(mask.shape(), DType::F32, device)? - &mask)?; - let huber_loss = ((&squared_loss * &mask)? + (&linear_loss * &one_minus_mask)?)?; - huber_loss.mean_all()? - } else { - // MSE fallback - (&diff * &diff)?.mean_all()? - }; - - // BUG FIX (Wave 5-E Agent 2): REMOVE entropy penalty from loss computation - // The entropy penalty was a detached constant tensor (created via Tensor::from_vec) - // which broke gradient flow when added to the loss. This caused gradient collapse - // (grad_norm=0.0000) in production training despite Q-values being non-zero. - // - // Root Cause: Tensor::from_vec() creates a DETACHED constant with no gradient tracking. - // When added to loss via .add(), it detaches the entire computation graph. - // - // Solution: Use loss_value directly WITHOUT adding entropy penalty. - // Entropy regularization should be applied at reward level (already done in reward.rs), - // NOT as a loss term (which requires proper gradient tracking). - let loss = loss_value; - - // Extract loss value AFTER backward pass (preserve computation graph) - // Backward pass with gradient monitoring (Adam provides natural stabilization) - let grad_norm = if let Some(ref mut optimizer) = self.optimizer { - let norm = optimizer - .backward_step_with_monitoring(&loss, self.gradient_clip_norm) - .map_err(|e| MLError::TrainingError(format!("Backward step with monitoring failed: {}", e)))?; - - tracing::debug!("Gradient norm: {:.4}", norm); - norm as f32 - } else { - return Err(MLError::TrainingError("Optimizer not initialized".to_string())); - }; - - // Extract loss value after backward pass (safe to detach now) - let loss_value = loss - .to_scalar::() - .map_err(|e| MLError::TrainingError(format!("Failed to extract loss: {}", e)))?; - - // Update training steps (epsilon decay moved to epoch-level in trainer) - self.training_steps += 1; - - // WAVE 10-A4: Real-time diagnostic monitoring - // Q-value monitoring every 10 steps - if self.training_steps % 10 == 0 { - self.log_q_values(&states_tensor)?; - } - // Dead neuron detection every 100 steps - if self.training_steps % 100 == 0 { - self.log_diagnostics(grad_norm)?; - } - - // WAVE 16 (Agent 36): Update target network with Polyak averaging or hard updates - if self.config.use_soft_updates { - // Polyak averaging: Update every step with tau coefficient - polyak_update(self.q_network.vars(), self.target_network.vars(), self.config.tau) - .map_err(|e| MLError::TrainingError(format!("Polyak update failed: {}", e)))?; - - // Log soft update every 1000 steps - if self.training_steps % 1000 == 0 { - let half_life = convergence_half_life(self.config.tau); - debug!("Soft target update at step {} (τ={}, half-life={:.0} steps)", - self.training_steps, self.config.tau, half_life); - } - } else { - // Hard update: Full copy every N steps (legacy mode) - if self.training_steps % self.config.target_update_freq as u64 == 0 { - hard_update(self.q_network.vars(), self.target_network.vars()) - .map_err(|e| MLError::TrainingError(format!("Hard update failed: {}", e)))?; - debug!("Hard target update at step {} (every {} steps)", - self.training_steps, self.config.target_update_freq); - } - } - - Ok((loss_value, grad_norm)) - } - - /// Log Q-values for the first state in batch (Wave 10-A4 diagnostic monitoring) - fn log_q_values(&self, states_tensor: &Tensor) -> Result<(), MLError> { - // Get Q-values for first state in batch - let first_state = states_tensor.i(0)?; - let first_state = first_state.unsqueeze(0)?; // Add batch dimension - let q_values = self.q_network.forward(&first_state)?; - - // Extract Q-values for each action - let q_buy = q_values.i((0, 0))?.to_scalar::()?; - let q_sell = q_values.i((0, 1))?.to_scalar::()?; - let q_hold = q_values.i((0, 2))?.to_scalar::()?; - - tracing::info!( - "Step {} Q-values: BUY={:.6}, SELL={:.6}, HOLD={:.6}", - self.training_steps, q_buy, q_sell, q_hold - ); - - // Alert if Q-value collapse detected (all Q-values near zero) - if q_buy.abs() < 0.0001 && q_sell.abs() < 0.0001 && q_hold.abs() < 0.0001 { - tracing::warn!( - "⚠️ Q-VALUE COLLAPSE DETECTED at step {}: BUY={:.6}, SELL={:.6}, HOLD={:.6}", - self.training_steps, q_buy, q_sell, q_hold - ); - } - - Ok(()) - } - - /// Detect dead neurons and log comprehensive diagnostics (Wave 10-A4) - fn log_diagnostics(&self, grad_norm: f32) -> Result<(), MLError> { - let dead_pct = self.detect_dead_neurons()?; - - tracing::info!( - "Step {} Diagnostics: grad_norm={:.2}, dead_neurons={:.2}%", - self.training_steps, grad_norm, dead_pct - ); - - // Alert if gradient collapse detected - if grad_norm < 1.0 { - tracing::warn!( - "⚠️ GRADIENT COLLAPSE: norm={:.6} at step {}", - grad_norm, self.training_steps - ); - } - - Ok(()) - } - - /// Detect dead ReLU neurons (weights stuck at zero) - fn detect_dead_neurons(&self) -> Result { - let mut dead_count = 0; - let mut total_count = 0; - - // Lock VarMap to inspect weights - let vars_data = self.q_network.vars().data().lock().map_err(|e| { - MLError::ConcurrencyError { - operation: format!("lock VarMap for dead neuron detection: {}", e), - } - })?; - - // Check each layer's weights - for (_name, var) in vars_data.iter() { - let tensor = var.as_tensor(); - let values = tensor.flatten_all()?.to_vec1::()?; - - for &val in values.iter() { - total_count += 1; - if val.abs() < 1e-6 { - dead_count += 1; - } - } - } - - Ok((dead_count as f32 / total_count as f32) * 100.0) - } - - /// Calculate entropy-based diversity penalty from recent actions - /// Returns a penalty tensor (negative entropy encourages diversity) - fn calculate_entropy_penalty(&self) -> Result { - if self.recent_actions.is_empty() { - return Tensor::zeros(&[], DType::F32, &self.device) - .map_err(|e| MLError::ModelError(format!("Failed to create zero penalty: {}", e))); - } - - // Count action frequencies - let mut counts = [0, 0, 0]; // BUY, SELL, HOLD - for action in &self.recent_actions { - counts[*action as usize] += 1; - } - - // Calculate Shannon entropy: H = -Σ(p_i * log2(p_i)) - let total = self.recent_actions.len() as f64; - let mut entropy = 0.0_f64; - for &count in &counts { - if count > 0 { - let p = count as f64 / total; - entropy -= p * p.log2(); - } - } - - // Return negative entropy as penalty (lower entropy = higher penalty) - // This encourages the agent to maximize entropy (balanced actions) - let penalty = -entropy as f32; - Tensor::from_vec(vec![penalty], &[], &self.device) - .map_err(|e| MLError::ModelError(format!("Failed to create entropy penalty tensor: {}", e))) - } - - /// Update exploration epsilon (called once per epoch by trainer) - pub fn update_epsilon(&mut self) { - self.epsilon = (self.epsilon * self.config.epsilon_decay).max(self.config.epsilon_end); - } - - /// Update softmax temperature (called once per epoch by trainer) (Wave 3 Agent 1) - pub fn update_temperature(&mut self) { - self.temperature = (self.temperature * self.config.temperature_decay).max(self.config.temperature_min); - } - - /// Update temperature adaptively based on training loss (Wave 2 Agent 2A) - /// - /// # Arguments - /// * `current_loss` - Training loss from current epoch - /// - /// # Behavior - /// - If loss improves >0.1%: Decay temperature faster (0.99 instead of 0.995) - /// - If loss plateaus <0.1% change: Slow decay (0.998 instead of 0.995) - /// - If loss plateaus >10 epochs: Increase temperature 5% (escape local optimum) - /// - /// When `use_adaptive_temperature` is false, falls back to fixed decay. - pub fn update_temperature_adaptive(&mut self, current_loss: f32) { - // If adaptive temperature disabled, use fixed decay - if !self.config.use_adaptive_temperature { - self.update_temperature(); - return; - } - - // Add current loss to sliding window (5-epoch average reduces noise) - self.loss_window.push_back(current_loss); - if self.loss_window.len() > 5 { - self.loss_window.pop_front(); - } - - // Compute average loss over window - let avg_loss = self.loss_window.iter().sum::() / self.loss_window.len() as f32; - - // Initialize best_loss on first call - if self.best_loss == f32::MAX { - self.best_loss = avg_loss; - tracing::debug!("Initializing best_loss to {:.6}", self.best_loss); - return; // Skip first update to establish baseline - } - - // Check if loss improved significantly (>0.1% improvement) - let improvement_threshold = self.best_loss * self.config.loss_improvement_threshold; - - if avg_loss < improvement_threshold { - // Loss improved → decay temperature faster (exploit learned policy) - self.temperature = (self.temperature * self.config.temperature_decay) - .max(self.config.temperature_min); - - // Update best loss and reset plateau counter - self.best_loss = avg_loss; - self.plateau_count = 0; - - tracing::debug!( - "Loss improved ({:.6} → {:.6}), decaying temperature to {:.4}", - self.best_loss, avg_loss, self.temperature - ); - } else { - // Loss plateaued → check if stuck in local optimum - self.plateau_count += 1; - - if self.plateau_count > self.config.plateau_window { - // Stuck for too long → increase temperature (escape local optimum) - self.temperature = (self.temperature * self.config.temp_increase_factor) - .min(self.config.temperature_start); - - // Reset plateau counter after temperature increase - self.plateau_count = 0; - - tracing::warn!( - "Loss plateaued for {} epochs, increasing temperature to {:.4}", - self.config.plateau_window, self.temperature - ); - } else { - // Slower decay (give more time to improve) - self.temperature = (self.temperature * self.config.temperature_slow_decay) - .max(self.config.temperature_min); - - tracing::debug!( - "Loss plateau {} epochs, slow decay to {:.4}", - self.plateau_count, self.temperature - ); - } - } - } - - /// Get plateau count (for testing) - pub fn get_plateau_count(&self) -> usize { - self.plateau_count - } - - /// Get best loss (for testing) - pub fn get_best_loss(&self) -> f32 { - self.best_loss - } - - /// Get current temperature value - pub fn get_temperature(&self) -> f64 { - self.temperature - } - - /// Set temperature value (used for evaluation or manual control) - pub fn set_temperature(&mut self, temperature: f64) { - self.temperature = temperature.max(0.01); // Prevent division by zero - } - - /// Update target network by copying weights from main network - fn update_target_network(&mut self) -> Result<(), MLError> { - self.target_network.copy_weights_from(&self.q_network)?; - Ok(()) - } - - /// Get target network Q-values for next state (for Bellman TD error computation) - /// - /// Used in validation loss to compute proper TD targets: reward + gamma * max_Q(s',a') - pub fn forward_target(&self, state: &Tensor) -> Result { - // Auto-convert input to correct device - let state = state - .to_device(&self.device) - .map_err(|e| MLError::ModelError(format!("Failed to move tensor to device: {}", e)))?; - - let q_values = self.target_network.forward(&state)?; - - // CRITICAL FIX: Use soft_clamp to preserve gradients - // Hard clamp (.clamp()) breaks gradient flow in Candle 0.9.1 - let clamped = soft_clamp(&q_values, 1000.0)?; - Ok(clamped) - } - - /// Get current epsilon value - pub fn get_epsilon(&self) -> f32 { - self.epsilon - } - - /// Set epsilon value (used for deterministic evaluation) - pub fn set_epsilon(&mut self, epsilon: f64) { - self.epsilon = epsilon.clamp(0.0, 1.0) as f32; - } - - /// Get training steps (gradient updates only, excludes warmup) - pub fn get_training_steps(&self) -> u64 { - self.training_steps - } - - /// Get total environment steps (includes warmup period) - pub fn get_total_steps(&self) -> u64 { - self.total_steps - } - - /// Check if in warmup period - pub fn is_in_warmup(&self) -> bool { - self.total_steps < self.config.warmup_steps as u64 - } - - /// Get warmup steps configured - pub fn get_warmup_steps(&self) -> usize { - self.config.warmup_steps - } - - /// Get replay buffer size - pub fn get_replay_buffer_size(&self) -> Result { - let buffer = self.memory.lock().map_err(|e| MLError::ConcurrencyError { - operation: format!("lock memory buffer for size check: {}", e), - })?; - Ok(buffer.len()) - } - - /// Get Q-network variables for serialization - pub fn get_q_network_vars(&self) -> &VarMap { - self.q_network.vars() - } - - /// Get target network variables for testing - pub fn get_target_network_vars(&self) -> &VarMap { - self.target_network.vars() - } - - /// Load model weights from safetensors checkpoint - /// - /// Loads pre-trained weights from a safetensors file and updates both - /// the Q-network and target network. Follows the MAMBA2 pattern for - /// checkpoint loading. - /// - /// # Arguments - /// - /// * `path` - Path to the safetensors file (with or without .safetensors extension) - /// - /// # Returns - /// - /// * `Ok(())` - Checkpoint loaded successfully - /// * `Err(MLError::CheckpointError)` - File not found or invalid format - /// * `Err(MLError::LockError)` - Failed to acquire VarMap lock - /// - /// # Example - /// - /// ```no_run - /// use ml::dqn::{WorkingDQN, WorkingDQNConfig}; - /// - /// let config = WorkingDQNConfig::emergency_safe_defaults(); - /// let mut dqn = WorkingDQN::new(config)?; - /// dqn.load_from_safetensors("/path/to/checkpoint")?; - /// # Ok::<(), ml::MLError>(()) - /// ``` - pub fn load_from_safetensors(&mut self, path: &str) -> Result<(), MLError> { - // Add .safetensors extension if not present - let safetensors_path = if !path.ends_with(".safetensors") { - format!("{}.safetensors", path) - } else { - path.to_string() - }; - - // Verify checkpoint file exists - if !std::path::Path::new(&safetensors_path).exists() { - return Err(MLError::CheckpointError(format!( - "Checkpoint file not found: {}", - safetensors_path - ))); - } - - // Load tensors from safetensors - let tensors = candle_core::safetensors::load(&safetensors_path, &self.device).map_err( - |e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)), - )?; - - // Populate VarMap with loaded tensors - let mut vars_data = self.q_network.vars().data().lock().map_err(|e| { - MLError::LockError(format!("Failed to lock VarMap for checkpoint load: {}", e)) - })?; - - for (name, tensor) in tensors.iter() { - // Create new Var from loaded tensor - let var = Var::from_tensor(tensor)?; - vars_data.insert(name.clone(), var); - } - - // Release lock before updating target network - drop(vars_data); - - // Update target network to match loaded weights - self.update_target_network()?; - - debug!( - "✓ DQN checkpoint loaded successfully: {} ({} tensors)", - safetensors_path, - tensors.len() - ); - - Ok(()) - } - - /// Check if ready for training - pub fn can_train(&self) -> bool { - match self.memory.lock() { - Ok(buffer) => buffer.can_sample(self.config.min_replay_size), - Err(_) => false, // If we can't lock, assume we can't train - } - } - - // ========== FACTORED ACTION SPACE SUPPORT (feature-gated) ========== - - #[cfg(feature = "factored-actions")] - /// Initialize factored Q-network with 3 heads (45 actions) - /// - /// Creates a factored Q-network with exposure, order, and urgency heads. - /// Must be called before using factored action selection. - pub fn init_factored_network(&mut self) -> Result<(), MLError> { - let factored_net = FactoredQNetwork::new(self.config.state_dim, &self.device)?; - self.factored_network = Some(factored_net); - tracing::info!("Factored Q-network initialized with 45 actions (5 exposure × 3 order × 3 urgency)"); - Ok(()) - } - - #[cfg(feature = "factored-actions")] - /// Update current position for action masking - /// - /// Call this before each action selection to ensure invalid actions are masked - pub fn set_current_position(&mut self, position: f64) { - self.current_position = position.clamp(-1.0, 1.0); - } - - #[cfg(feature = "factored-actions")] - /// Get current position for action masking - pub fn get_current_position(&self) -> f64 { - self.current_position - } - - #[cfg(feature = "factored-actions")] - /// Select action using factored Q-network with epsilon-greedy policy - /// - /// Uses position masking to prevent invalid actions (exceeding ±100% limit) - pub fn select_factored_action(&mut self, state: &[f32]) -> Result { - let factored_net = self.factored_network.as_ref().ok_or_else(|| { - MLError::ModelError("Factored network not initialized. Call init_factored_network() first.".to_string()) - })?; - - // Increment total steps counter - self.total_steps += 1; - - let mut rng = thread_rng(); - let in_warmup = self.total_steps <= self.config.warmup_steps as u64; - - // Convert state to tensor - let state_tensor = Tensor::from_vec( - state.to_vec(), - (1, self.config.state_dim), - factored_net.device(), - ) - .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?; - - // Epsilon-greedy action selection - let action = if in_warmup || rng.gen::() < self.epsilon { - // Random exploration - factored_net.select_epsilon_greedy(&state_tensor, 1.0)? - } else { - // Greedy exploitation with masking - let (q_exp, q_ord, q_urg) = factored_net.forward(&state_tensor)?; - - // Apply position masking to prevent invalid actions - let masked_q_exp = factored_net.apply_position_mask(&q_exp, self.current_position)?; - - // Select best action from masked Q-values - let exp_idx = masked_q_exp - .argmax(1) - .map_err(|e| MLError::ModelError(format!("Exposure argmax failed: {}", e)))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Exposure index to vec failed: {}", e)))?[0] - as usize; - - let ord_idx = q_ord - .argmax(1) - .map_err(|e| MLError::ModelError(format!("Order argmax failed: {}", e)))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Order index to vec failed: {}", e)))?[0] - as usize; - - let urg_idx = q_urg - .argmax(1) - .map_err(|e| MLError::ModelError(format!("Urgency argmax failed: {}", e)))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Urgency index to vec failed: {}", e)))?[0] - as usize; - - let exposure = ExposureLevel::from_index(exp_idx)?; - let order = OrderType::from_index(ord_idx)?; - let urgency = Urgency::from_index(urg_idx)?; - - FactoredAction::new(exposure, order, urgency) - }; - - Ok(action) - } - - #[cfg(feature = "factored-actions")] - /// Check if factored network is initialized - pub fn has_factored_network(&self) -> bool { - self.factored_network.is_some() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::dqn::Experience; - // use crate::safe_operations; // DISABLED - module not found - - /// Unit test: Verify soft_clamp gradient flow - /// - /// This test ensures that soft_clamp preserves gradients, unlike hard .clamp() - /// which has NO backward pass in Candle 0.9.1. - #[test] - fn test_soft_clamp_gradient_flow() -> anyhow::Result<()> { - let device = Device::Cpu; - - // Create input tensor with gradient tracking enabled - let x = Tensor::from_vec(vec![1.0f32, 10.0, 100.0, 1000.0], 4, &device)?; - - // Apply soft_clamp - let clamped = soft_clamp(&x, 10.0)?; - - // Verify output values are bounded - let values = clamped.to_vec1::()?; - for (i, &val) in values.iter().enumerate() { - assert!( - val.abs() <= 10.0, - "Value {} at index {} exceeds bound 10.0", - val, i - ); - } - - // Verify asymptotic behavior (large inputs approach bound) - assert!( - (values[3] - 10.0).abs() < 0.1, - "Large input (1000.0) should approach bound 10.0, got {}", - values[3] - ); - - // Verify smooth scaling (small inputs nearly unchanged) - assert!( - (values[0] - 1.0).abs() < 0.1, - "Small input (1.0) should be nearly unchanged, got {}", - values[0] - ); - - println!("✅ Soft clamp values: {:?}", values); - println!(" Expected: [~1.0, ~7.6, ~9.5, ~10.0]"); - - Ok(()) - } - - /// Unit test: Verify soft_clamp preserves computation graph - /// - /// This test ensures gradients can flow through soft_clamp for backpropagation. - #[test] - fn test_soft_clamp_preserves_computation_graph() -> anyhow::Result<()> { - let device = Device::Cpu; - - // Create VarMap and Variable for gradient tracking - let varmap = VarMap::new(); - let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); - - // Create trainable variable (initialized via VarBuilder) - let x_var_tensor = vb.get((4,), "x")?; - - // Get underlying Var and set initial values - let vars_data = varmap.data().lock().unwrap(); - if let Some(x_var) = vars_data.get("x") { - x_var.set(&Tensor::from_vec(vec![1.0f32, 10.0, 100.0, 1000.0], 4, &device)?)?; - } - drop(vars_data); // Release lock - - // Apply soft_clamp on the Variable tensor - let clamped = soft_clamp(&x_var_tensor, 10.0)?; - - // Compute loss (sum of squared clamped values) - let loss = clamped.sqr()?.sum_all()?; - - // Backward pass - this will FAIL if soft_clamp breaks gradient flow - let grads = loss.backward()?; - - // Verify gradients exist for the variable - let grad = grads.get(&x_var_tensor).ok_or_else(|| { - anyhow::anyhow!("No gradient computed for variable 'x' - gradient flow broken!") - })?; - - // Verify gradients are non-zero (at least for some elements) - let grad_values = grad.to_vec1::()?; - let non_zero_grads = grad_values.iter().filter(|&&g| g.abs() > 1e-6).count(); - - assert!( - non_zero_grads > 0, - "GRADIENT COLLAPSE: All gradients are zero! soft_clamp broke gradient flow.\nGradients: {:?}", - grad_values - ); - - println!("✅ Gradients computed successfully: {:?}", grad_values); - println!(" Non-zero gradients: {}/4", non_zero_grads); - - Ok(()) - } - - /// Unit test: Compare soft_clamp vs hard clamp gradient behavior - /// - /// Demonstrates that hard clamp breaks gradients while soft_clamp preserves them. - #[test] - fn test_soft_vs_hard_clamp_gradient_comparison() -> anyhow::Result<()> { - let device = Device::Cpu; - - // Test input: mix of small and large values - let test_values = vec![1.0f32, 5.0, 10.0, 50.0, 100.0]; - let bound = 10.0; - - // Test soft_clamp - let x_soft = Tensor::from_vec(test_values.clone(), 5, &device)?; - let clamped_soft = soft_clamp(&x_soft, bound)?; - let soft_values = clamped_soft.to_vec1::()?; - - // Note: We cannot test hard clamp's gradient behavior directly because - // Candle's .clamp() has no backward pass. This test documents expected behavior. - - // Verify soft_clamp properties - println!("Soft clamp output:"); - for (input, output) in test_values.iter().zip(soft_values.iter()) { - let ratio = output / input; - println!(" input={:6.1} → output={:6.3} (ratio={:.3})", input, output, ratio); - - // Verify output is bounded - assert!( - output.abs() <= bound as f32 + 0.1, - "Output {} exceeds bound {}", - output, bound - ); - } - - // Verify small inputs are nearly preserved (linear region) - assert!( - (soft_values[0] / test_values[0] - 1.0).abs() < 0.05, - "Small input should be ~linearly scaled, got ratio {}", - soft_values[0] / test_values[0] - ); - - // Verify large inputs are compressed (saturation region) - assert!( - soft_values[4] < test_values[4] * 0.2, - "Large input should be compressed, got {} from {}", - soft_values[4], test_values[4] - ); - - Ok(()) - } - - #[test] - fn test_working_dqn_creation() -> anyhow::Result<()> { - // Test DQN creation concepts - let initial_epsilon = 1.0; - let training_steps = 0; - - assert_eq!(initial_epsilon, 1.0); - assert_eq!(training_steps, 0); - Ok(()) - } - - #[test] - fn test_action_selection() -> anyhow::Result<()> { - // Test action selection concepts - let num_actions = 3; - let selected_action = 1; // Sample action - assert!(selected_action < num_actions); - Ok(()) - } - - #[test] - fn test_experience_storage() -> anyhow::Result<()> { - // Test experience storage concepts - let replay_buffer_size = 1; - let experience_count = 1; - - assert_eq!(experience_count, replay_buffer_size); - Ok(()) - } - - #[test] - fn test_training_update() -> anyhow::Result<()> { - // Test training update concepts - let batch_size = 32; - // SAFETY: Learning rate must come from configuration, not hardcoded - let config = WorkingDQNConfig::emergency_safe_defaults(); - let learning_rate = config.learning_rate; - - assert!(batch_size > 0); - assert!(learning_rate > 0.0); - Ok(()) - } - - #[test] - fn test_training_step_without_enough_data() -> anyhow::Result<()> { - let config = WorkingDQNConfig::emergency_safe_defaults(); - let mut dqn = WorkingDQN::new(config)?; - - // Try training without enough experiences - let result = dqn.train_step(None); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn test_training_step_with_data() -> anyhow::Result<()> { - let mut config = WorkingDQNConfig::emergency_safe_defaults(); - config.min_replay_size = 4; - config.batch_size = 4; - config.state_dim = 52; // Match the state vector size used in test data (4 prices + 16 technical + 16 microstructure + 16 portfolio) - let mut dqn = WorkingDQN::new(config)?; - - // Add enough experiences - for i in 0..10 { - let experience = Experience::new( - vec![i as f32 * 0.1; 52], - (i % 3) as u8, - i as f32, - vec![(i + 1) as f32 * 0.1; 52], - i == 9, - ); - dqn.store_experience(experience)?; - } - - // Training should work now - let result = dqn.train_step(None); - if let Err(ref e) = result { - eprintln!("train_step error: {:?}", e); - } - assert!(result.is_ok(), "train_step failed: {:?}", result.err()); - - let (loss, grad_norm) = result?; - assert!(loss >= 0.0); // Loss should be non-negative - assert!(grad_norm >= 0.0); // Gradient norm should be non-negative - Ok(()) - } - - #[test] - fn test_epsilon_decay() -> anyhow::Result<()> { - let mut config = WorkingDQNConfig::emergency_safe_defaults(); - config.epsilon_start = 1.0; - config.epsilon_decay = 0.9; - config.epsilon_end = 0.1; - let mut dqn = WorkingDQN::new(config)?; - - let initial_epsilon = dqn.get_epsilon(); - dqn.update_epsilon(); - let new_epsilon = dqn.get_epsilon(); - - assert!(new_epsilon < initial_epsilon); - assert!(new_epsilon >= 0.1); // Should not go below epsilon_end - Ok(()) - } - - #[test] - fn test_target_network_update() -> anyhow::Result<()> { - let config = WorkingDQNConfig::emergency_safe_defaults(); - let mut dqn = WorkingDQN::new(config)?; - - let result = dqn.update_target_network(); - assert!(result.is_ok()); - Ok(()) - } - - /// Wave 5-E Agent 1: TDD Gradient Flow Test - /// - /// This test is designed to FAIL with Bug #2 (gradient collapse). - /// It verifies that gradients flow correctly through the Q-value computation path - /// during training, specifically testing: - /// 1. .clamp() preserves gradients - /// 2. .gather() supports backward pass - /// 3. .to_dtype() doesn't detach computation graph - /// - /// Expected behavior: - /// - FAIL with current code (grad_norm = 0.0) - /// - PASS after Bug #2 fix (grad_norm > 0.0 and < 100.0) - /// - /// Root cause (suspected): - /// - .clamp() may not preserve gradients in Candle 0.9.1 - /// - .gather() gradient support might be broken - /// - .to_dtype() conversion might detach computation graph - #[test] - fn test_gradient_flow_not_zero() -> anyhow::Result<()> { - // TEST ON CUDA (not CPU) to match production environment - let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); - println!("✅ Testing gradient flow on device: {:?}", device); - - // Create minimal DQN with simple config - let mut config = WorkingDQNConfig::emergency_safe_defaults(); - config.state_dim = 225; // Production state dimension - config.num_actions = 3; // BUY, SELL, HOLD - config.batch_size = 8; // Small batch for quick test - config.min_replay_size = 8; // Match batch size - config.replay_buffer_capacity = 100; // Small buffer - config.learning_rate = 0.0001; // Standard learning rate - config.gamma = 0.99; - config.gradient_clip_norm = 10.0; // Wave 11 Bug #1 fix - - // Save gradient_clip_norm before moving config - let gradient_clip_norm = config.gradient_clip_norm; - - let mut dqn = WorkingDQN::new(config)?; - - // Create sample experiences (8 transitions with diverse actions) - let experiences = vec![ - Experience::new( - vec![0.5f32; 225], - 0, // BUY - 1.0, - vec![0.6f32; 225], - false, - ), - Experience::new( - vec![0.6f32; 225], - 1, // SELL - -0.5, - vec![0.4f32; 225], - false, - ), - Experience::new( - vec![0.4f32; 225], - 2, // HOLD - 0.1, - vec![0.45f32; 225], - false, - ), - Experience::new( - vec![0.45f32; 225], - 0, // BUY - 0.8, - vec![0.55f32; 225], - false, - ), - Experience::new( - vec![0.55f32; 225], - 1, // SELL - -0.3, - vec![0.5f32; 225], - false, - ), - Experience::new( - vec![0.5f32; 225], - 2, // HOLD - 0.05, - vec![0.52f32; 225], - false, - ), - Experience::new( - vec![0.52f32; 225], - 0, // BUY - 1.2, - vec![0.65f32; 225], - false, - ), - Experience::new( - vec![0.65f32; 225], - 1, // SELL - -0.7, - vec![0.3f32; 225], - true, // Terminal state - ), - ]; - - // Store experiences in replay buffer - for exp in experiences.clone() { - dqn.store_experience(exp)?; - } - - // Perform one training step - let result = dqn.train_step(None); - - // Validate training succeeded - assert!( - result.is_ok(), - "Training step should succeed, got error: {:?}", - result.err() - ); - - let (loss, grad_norm) = result?; - - // CRITICAL ASSERTION #1: Gradient norm should be non-zero - // This assertion will FAIL with Bug #2 (gradient collapse) - assert!( - grad_norm > 0.0, - "GRADIENT COLLAPSE DETECTED: Gradient norm should be non-zero after training step, got: {:.8}. \ - This indicates gradients are not flowing through the Q-value computation path. \ - Suspected causes: .clamp(), .gather(), or .to_dtype() breaking gradient flow.", - grad_norm - ); - - // CRITICAL ASSERTION #2: Gradients should not explode - assert!( - grad_norm < 100.0, - "GRADIENT EXPLOSION DETECTED: Gradients should not explode, got: {:.4}. \ - Gradient clipping (max_norm={:.1}) may be insufficient.", - grad_norm, - gradient_clip_norm - ); - - // ASSERTION #3: Loss should be finite and reasonable - assert!( - loss.is_finite(), - "Loss should be finite, got: {}", - loss - ); - - assert!( - loss >= 0.0, - "Loss should be non-negative (Huber/MSE), got: {}", - loss - ); - - // ASSERTION #4: Loss should be reasonable magnitude (not collapsed) - assert!( - loss < 1000.0, - "Loss too large, possible training instability: {}", - loss - ); - - // SUCCESS: If we reach here, gradients are flowing correctly - eprintln!("✅ Gradient flow test PASSED:"); - eprintln!(" - Gradient norm: {:.6}", grad_norm); - eprintln!(" - Loss: {:.6}", loss); - eprintln!(" - Training steps: {}", dqn.training_steps); - - Ok(()) - } -} diff --git a/ml/src/dqn/factored_q_network.rs.backup b/ml/src/dqn/factored_q_network.rs.backup deleted file mode 100644 index 5a4b9cfac..000000000 --- a/ml/src/dqn/factored_q_network.rs.backup +++ /dev/null @@ -1,564 +0,0 @@ -//! Factored Q-Network for Standard DQN -//! -//! Implements a factored Q-network architecture with 3 separate heads for exposure, -//! order type, and urgency sub-actions. Uses additive Q-value factorization: -//! Q(s,a) = Q_exposure(s,a_exp) + Q_order(s,a_ord) + Q_urgency(s,a_urg) -//! -//! Architecture: -//! - Shared encoder: 128 → 64 (ReLU) -//! - Exposure head: 64 → 5 (ExposureLevel) -//! - Order head: 64 → 3 (OrderType) -//! - Urgency head: 64 → 3 (Urgency) - -use candle_core::{Device, Tensor}; -use candle_nn::{Linear, Module, VarBuilder, VarMap}; -use rand::Rng; -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: 128, - hidden_dim: 64, - } - } -} - -/// Factored Q-Network with 3 heads for exposure, order, and urgency -#[derive(Debug)] -pub struct FactoredQNetwork { - /// Shared encoder (state → hidden representation) - shared_encoder: Linear, - /// Exposure head (hidden → 5 Q-values) - exposure_head: Linear, - /// Order type head (hidden → 3 Q-values) - order_head: Linear, - /// Urgency head (hidden → 3 Q-values) - urgency_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, candle_core::DType::F32, 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 exposure head (5 outputs) - let exposure_head = linear_xavier( - config.hidden_dim, - 5, // ExposureLevel has 5 values - vb.pp("exposure_head"), - ) - .map_err(|e| MLError::ModelError(format!("Failed to create exposure head: {}", e)))?; - - // Initialize order head (3 outputs) - let order_head = linear_xavier( - config.hidden_dim, - 3, // OrderType has 3 values - vb.pp("order_head"), - ) - .map_err(|e| MLError::ModelError(format!("Failed to create order head: {}", e)))?; - - // Initialize urgency head (3 outputs) - let urgency_head = linear_xavier( - config.hidden_dim, - 3, // Urgency has 3 values - vb.pp("urgency_head"), - ) - .map_err(|e| MLError::ModelError(format!("Failed to create urgency head: {}", e)))?; - - Ok(Self { - shared_encoder, - exposure_head, - order_head, - urgency_head, - device: device.clone(), - hidden_dim: config.hidden_dim, - }) - } - - /// Forward pass: compute Q-values for all 3 heads - /// - /// Returns (q_exposure [batch, 5], q_order [batch, 3], q_urgency [batch, 3]) - pub fn forward(&self, state: &Tensor) -> Result<(Tensor, Tensor, Tensor), MLError> { - // DEBUG: Log input shape - tracing::info!("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)))?; - - // ReLU activation - let hidden = hidden - .relu() - .map_err(|e| MLError::ModelError(format!("ReLU activation failed: {}", e)))?; - - // DEBUG: Log hidden representation shape - tracing::info!("Hidden representation shape: {:?}", hidden.dims()); - - // Exposure head - let q_exposure = self - .exposure_head - .forward(&hidden) - .map_err(|e| MLError::ModelError(format!("Exposure head forward failed: {}", e)))?; - - // Order head - let q_order = self - .order_head - .forward(&hidden) - .map_err(|e| MLError::ModelError(format!("Order head forward failed: {}", e)))?; - - // Urgency head - let q_urgency = self - .urgency_head - .forward(&hidden) - .map_err(|e| MLError::ModelError(format!("Urgency head forward failed: {}", e)))?; - - // DEBUG: Log output shapes - tracing::info!( - "FactoredQNetwork output shapes - exposure: {:?}, order: {:?}, urgency: {:?}", - q_exposure.dims(), - q_order.dims(), - q_urgency.dims() - ); - - // DEBUG: Log first 5 Q-values from each head (if batch size permits) - if let Ok(exp_vec) = q_exposure.flatten_all()?.to_vec1::() { - let num_exp = 5.min(exp_vec.len()); - tracing::info!("Exposure Q-values (first {}): {:?}", num_exp, &exp_vec[..num_exp]); - } - if let Ok(ord_vec) = q_order.flatten_all()?.to_vec1::() { - let num_ord = 3.min(ord_vec.len()); - tracing::info!("Order Q-values (first {}): {:?}", num_ord, &ord_vec[..num_ord]); - } - if let Ok(urg_vec) = q_urgency.flatten_all()?.to_vec1::() { - let num_urg = 3.min(urg_vec.len()); - tracing::info!("Urgency Q-values (first {}): {:?}", num_urg, &urg_vec[..num_urg]); - } - - Ok((q_exposure, q_order, q_urgency)) - } - - /// Compute joint Q-values using additive factorization - /// - /// Q(s,a) = Q_exposure(s,a_exp) + Q_order(s,a_ord) + Q_urgency(s,a_urg) - /// - /// Returns [batch, 45] tensor of joint Q-values - pub fn compute_joint_q( - &self, - q_exposure: &Tensor, - q_order: &Tensor, - q_urgency: &Tensor, - ) -> Result { - let batch_size = q_exposure - .dim(0) - .map_err(|e| MLError::ModelError(format!("Failed to get batch size: {}", e)))?; - - // Reshape to [batch, 5, 1, 1] for broadcasting - let q_exp = q_exposure - .reshape((batch_size, 5, 1, 1)) - .map_err(|e| MLError::ModelError(format!("Failed to reshape q_exposure: {}", e)))?; - - // Reshape to [batch, 1, 3, 1] for broadcasting - let q_ord = q_order - .reshape((batch_size, 1, 3, 1)) - .map_err(|e| MLError::ModelError(format!("Failed to reshape q_order: {}", e)))?; - - // Reshape to [batch, 1, 1, 3] for broadcasting - let q_urg = q_urgency - .reshape((batch_size, 1, 1, 3)) - .map_err(|e| MLError::ModelError(format!("Failed to reshape q_urgency: {}", e)))?; - - // Broadcast and sum: [batch, 5, 3, 3] - let joint_q = q_exp - .broadcast_add(&q_ord) - .map_err(|e| MLError::ModelError(format!("Failed to add q_exposure + q_order: {}", e)))?; - - let joint_q = joint_q - .broadcast_add(&q_urg) - .map_err(|e| MLError::ModelError(format!("Failed to add q_urgency: {}", e)))?; - - // Flatten to [batch, 45] - let joint_q = joint_q - .reshape((batch_size, 45)) - .map_err(|e| MLError::ModelError(format!("Failed to flatten joint Q-values: {}", e)))?; - - Ok(joint_q) - } - - /// Select greedy action (argmax per head) - pub fn select_greedy_action(&self, state: &Tensor) -> Result { - let (q_exposure, q_order, q_urgency) = self.forward(state)?; - - // DEBUG: Log Q-value shapes before argmax - tracing::debug!("Pre-argmax Q-value shapes - exposure: {:?}, order: {:?}, urgency: {:?}", - q_exposure.dims(), q_order.dims(), q_urgency.dims()); - - // Argmax per head - let exp_idx = q_exposure - .argmax(1) - .map_err(|e| MLError::ModelError(format!("Exposure argmax failed: {}", e)))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Exposure index to vec failed: {}", e)))?[0] - as usize; - - let ord_idx = q_order - .argmax(1) - .map_err(|e| MLError::ModelError(format!("Order argmax failed: {}", e)))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Order index to vec failed: {}", e)))?[0] - as usize; - - let urg_idx = q_urgency - .argmax(1) - .map_err(|e| MLError::ModelError(format!("Urgency argmax failed: {}", e)))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Urgency index to vec failed: {}", e)))?[0] - as usize; - - // DEBUG: Log selected indices - tracing::info!("Argmax results - exposure_idx: {}, order_idx: {}, urgency_idx: {}", - exp_idx, ord_idx, urg_idx); - - // Convert indices to action - let exposure = ExposureLevel::from_index(exp_idx)?; - let order = OrderType::from_index(ord_idx)?; - let urgency = Urgency::from_index(urg_idx)?; - - // DEBUG: Log final factored action - tracing::info!("Selected FactoredAction: exposure={:?}, order={:?}, urgency={:?}", - exposure, order, urgency); - - Ok(FactoredAction::new(exposure, order, urgency)) - } - - /// 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 exposure levels that would exceed the limit given current position - pub fn apply_position_mask( - &self, - q_exposure: &Tensor, - current_position: f64, - ) -> Result { - let batch_size = q_exposure - .dim(0) - .map_err(|e| MLError::ModelError(format!("Failed to get batch size: {}", e)))?; - - // Convert to Vec for masking - let mut q_values = q_exposure - .to_vec2::() - .map_err(|e| MLError::ModelError(format!("Failed to convert q_exposure to vec: {}", e)))?; - - // Mask invalid actions - for batch_idx in 0..batch_size { - for exp_idx in 0..5 { - let exposure = ExposureLevel::from_index(exp_idx)?; - let target_position = exposure.target_exposure(); - - // Check if this would exceed ±100% limit - if (current_position + target_position).abs() > 1.0 { - q_values[batch_idx][exp_idx] = f32::NEG_INFINITY; - } - } - } - - // Convert back to tensor - Tensor::new(q_values, &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_exp, q_ord, q_urg) = network.forward(&state).unwrap(); - - // Check shapes - assert_eq!(q_exp.dims(), &[32, 5]); - assert_eq!(q_ord.dims(), &[32, 3]); - assert_eq!(q_urg.dims(), &[32, 3]); - } - - #[test] - fn test_compute_joint_q_shape() { - let device = Device::Cpu; - let network = FactoredQNetwork::new(128, &device).unwrap(); - - let state = Tensor::zeros((32, 128), candle_core::DType::F32, &device).unwrap(); - let (q_exp, q_ord, q_urg) = network.forward(&state).unwrap(); - - let joint_q = network.compute_joint_q(&q_exp, &q_ord, &q_urg).unwrap(); - - // Check shape: [32, 45] - assert_eq!(joint_q.dims(), &[32, 45]); - } - - #[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_exp, _, _) = network.forward(&state).unwrap(); - - // Current position at +80% (Long) - let current_position = 0.8; - let masked_q = network.apply_position_mask(&q_exp, current_position).unwrap(); - - let masked_values = masked_q.to_vec2::().unwrap(); - - // Short100 (-1.0) would result in -0.2 (valid) - assert!(masked_values[0][0].is_finite()); - - // Long100 (+1.0) would result in +1.8 (invalid, should be -inf) - assert_eq!(masked_values[0][4], f32::NEG_INFINITY); - } - - #[test] - fn test_gradient_flow() { - let device = Device::Cpu; - let network = FactoredQNetwork::new(128, &device).unwrap(); - - let state = Tensor::randn(0.0f32, 1.0f32, (32, 128), &device).unwrap(); - let (q_exp, q_ord, q_urg) = network.forward(&state).unwrap(); - - // Compute loss (mean of all Q-values) - let loss = q_exp - .mean_all() - .unwrap() - .broadcast_add(&q_ord.mean_all().unwrap()) - .unwrap() - .broadcast_add(&q_urg.mean_all().unwrap()) - .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.0f32, 1.0f32, (100, 128), &device).unwrap(); - let (q_exp, q_ord, q_urg) = network.forward(&state).unwrap(); - - // Check that Q-values are in reasonable range after initialization - let exp_std = q_exp.var(1).unwrap().mean_all().unwrap().to_vec0::().unwrap().sqrt(); - let ord_std = q_ord.var(1).unwrap().mean_all().unwrap().to_vec0::().unwrap().sqrt(); - let urg_std = q_urg.var(1).unwrap().mean_all().unwrap().to_vec0::().unwrap().sqrt(); - - // Xavier init should produce reasonable variance (roughly < 2.0) - assert!(exp_std < 2.0, "Exposure std too large: {}", exp_std); - assert!(ord_std < 2.0, "Order std too large: {}", ord_std); - assert!(urg_std < 2.0, "Urgency std too large: {}", urg_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.0f32, 1.0f32, (1, 128), &device).unwrap(); - let batch_state = single_state.repeat((32, 1)).unwrap(); - - let (q_exp_single, q_ord_single, q_urg_single) = network.forward(&single_state).unwrap(); - let (q_exp_batch, q_ord_batch, q_urg_batch) = network.forward(&batch_state).unwrap(); - - // First batch item should match single state - let exp_diff = q_exp_single - .broadcast_sub(&q_exp_batch.i((0..1, ..)).unwrap()) - .unwrap() - .abs() - .unwrap() - .max_all() - .unwrap() - .to_vec0::() - .unwrap(); - - let ord_diff = q_ord_single - .broadcast_sub(&q_ord_batch.i((0..1, ..)).unwrap()) - .unwrap() - .abs() - .unwrap() - .max_all() - .unwrap() - .to_vec0::() - .unwrap(); - - let urg_diff = q_urg_single - .broadcast_sub(&q_urg_batch.i((0..1, ..)).unwrap()) - .unwrap() - .abs() - .unwrap() - .max_all() - .unwrap() - .to_vec0::() - .unwrap(); - - // Differences should be near zero - assert!(exp_diff < 1e-5, "Exposure batch inconsistency: {}", exp_diff); - assert!(ord_diff < 1e-5, "Order batch inconsistency: {}", ord_diff); - assert!(urg_diff < 1e-5, "Urgency batch inconsistency: {}", urg_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.0f32, 1.0f32, (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_exp, cpu_ord, cpu_urg) = cpu_network.forward(&cpu_state).unwrap(); - let (cuda_exp, cuda_ord, cuda_urg) = cuda_network.forward(&cuda_state).unwrap(); - - // Check shapes match - assert_eq!(cpu_exp.dims(), cuda_exp.dims()); - assert_eq!(cpu_ord.dims(), cuda_ord.dims()); - assert_eq!(cpu_urg.dims(), cuda_urg.dims()); - } - } -} diff --git a/ml/src/hyperopt/adapters/dqn.rs b/ml/src/hyperopt/adapters/dqn.rs index 0799bc7cf..3d5c2c661 100644 --- a/ml/src/hyperopt/adapters/dqn.rs +++ b/ml/src/hyperopt/adapters/dqn.rs @@ -232,7 +232,7 @@ impl Default for DQNParams { batch_size: 128, gamma: 0.99, buffer_size: 100_000, - hold_penalty_weight: 2.0, // User-discovered optimal value + hold_penalty_weight: 0.01, // WAVE 10 Bug Fix: Align with CLI/RewardConfig defaults (was 2.0, causing 200x mismatch) max_position_absolute: 2.0, // BLOCKER #2: Default matches production (±2.0) huber_delta: 1.0, entropy_coefficient: 0.01, @@ -2474,6 +2474,7 @@ mod tests { 0.6, 0.4, // per_alpha, per_beta_start -1000.0, 1000.0, 0.5_f64.ln(), // v_min, v_max, noisy_sigma_init (Wave 6.4) 256.0, 3.0, 101.0, // dueling_hidden_dim, n_steps, num_atoms (Wave 6.4) + 1.5, // minimum_profit_factor (mid-point of 1.1-2.0 range, BUG #7) // WAVE 11: use_dueling, use_distributional, use_noisy_nets REMOVED (always TRUE) ]; @@ -2492,6 +2493,7 @@ mod tests { 0.4, 0.2, // per_alpha min, per_beta_start min -1000.0, 1000.0, 0.5_f64.ln(), // v_min, v_max, noisy_sigma_init (Wave 6.4) 128.0, 1.0, 51.0, // dueling_hidden_dim min, n_steps min, num_atoms min (Wave 6.4) + 1.1, // minimum_profit_factor min (BUG #7) // WAVE 11: All Rainbow booleans always TRUE ]; let params_min = DQNParams::from_continuous(&continuous_min).unwrap(); @@ -2503,6 +2505,7 @@ mod tests { 0.8, 0.6, // per_alpha max, per_beta_start max -1000.0, 1000.0, 0.5_f64.ln(), // v_min, v_max, noisy_sigma_init (Wave 6.4) 512.0, 5.0, 201.0, // dueling_hidden_dim max, n_steps max, num_atoms max (Wave 6.4) + 2.0, // minimum_profit_factor max (BUG #7) // WAVE 11: All Rainbow booleans always TRUE ]; let params_max = DQNParams::from_continuous(&continuous_max).unwrap(); diff --git a/ml/trained_models/dqn_final_epoch10.safetensors b/ml/trained_models/dqn_final_epoch10.safetensors index 02dceffe1..714dcd408 100644 Binary files a/ml/trained_models/dqn_final_epoch10.safetensors and b/ml/trained_models/dqn_final_epoch10.safetensors differ