diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index 0a8f1fbb4..e97d51ee5 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -530,11 +530,13 @@ fn evaluate_ppo_fold( num_actions: args.num_actions, policy_hidden_dims: { let base = hp_usize(hp, "hidden_dim_base").unwrap_or(128); - vec![base, base / 2] + let align = |x: usize| x.div_ceil(8) * 8; // tensor-core alignment (must match PpoTrainer) + vec![align(base), align(base / 2)] }, value_hidden_dims: { let base = hp_usize(hp, "hidden_dim_base").unwrap_or(128); - vec![base * 4, base * 3, base * 2, base, base / 2] // 5-layer to match trainer + let align = |x: usize| x.div_ceil(8) * 8; // tensor-core alignment (must match PpoTrainer) + vec![align(base * 4), align(base * 3), align(base * 2), align(base), align(base / 2)] }, policy_learning_rate: 3e-4, value_learning_rate: 1e-3, diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index ea376f0f3..92ac68f35 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -320,9 +320,112 @@ impl DQNConfig { meta.insert("dqn.use_distributional".to_owned(), self.use_distributional.to_string()); meta.insert("dqn.use_noisy_nets".to_owned(), self.use_noisy_nets.to_string()); meta.insert("dqn.use_iqn".to_owned(), self.use_iqn.to_string()); + meta.insert("dqn.dueling_hidden_dim".to_owned(), self.dueling_hidden_dim.to_string()); + meta.insert("dqn.num_atoms".to_owned(), self.num_atoms.to_string()); + meta.insert("dqn.iqn_embedding_dim".to_owned(), self.iqn_embedding_dim.to_string()); + meta.insert("dqn.iqn_num_quantiles".to_owned(), self.iqn_num_quantiles.to_string()); + meta.insert("dqn.gamma".to_owned(), self.gamma.to_string()); meta } + /// Load DQNConfig from a safetensors checkpoint file. + /// + /// 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 { + let raw_bytes = std::fs::read(path).map_err(|e| { + crate::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)) + })?; + + let meta = match st_metadata.metadata() { + Some(m) => m, + None => { + return Err(crate::MLError::CheckpointError( + "Safetensors checkpoint has no metadata — re-train to embed config".to_owned(), + )); + } + }; + + Self::from_checkpoint_metadata(meta) + } + + /// Reconstruct a DQNConfig from safetensors checkpoint metadata. + /// + /// Reads architecture-critical fields (state_dim, num_actions, hidden_dims, + /// Rainbow flags) from the metadata map. Non-architecture fields (LR, epsilon, + /// buffer capacity) use defaults since they don't affect inference. + /// + /// 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> { + metadata.get(key).map(|s| s.as_str()).ok_or_else(|| { + crate::MLError::CheckpointError(format!( + "Checkpoint metadata missing '{}' — re-train to embed config", key + )) + }) + }; + + let state_dim: usize = get("dqn.state_dim")? + .parse() + .map_err(|e| crate::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)))?; + + // Parse hidden_dims from debug format "[256, 256]" + let hidden_dims_str = get("dqn.hidden_dims")?; + let hidden_dims: Vec = hidden_dims_str + .trim_matches(|c| c == '[' || c == ']') + .split(',') + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim().parse::()) + .collect::, _>>() + .map_err(|e| crate::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_usize = |key: &str| -> Result { + get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e))) + }; + + let use_dueling = parse_bool("dqn.use_dueling")?; + let use_distributional = parse_bool("dqn.use_distributional")?; + let use_noisy_nets = parse_bool("dqn.use_noisy_nets")?; + let use_iqn = parse_bool("dqn.use_iqn")?; + let dueling_hidden_dim = parse_usize("dqn.dueling_hidden_dim")?; + let num_atoms = parse_usize("dqn.num_atoms")?; + let iqn_embedding_dim = parse_usize("dqn.iqn_embedding_dim")?; + let iqn_num_quantiles = parse_usize("dqn.iqn_num_quantiles")?; + + let gamma: f32 = metadata.get("dqn.gamma") + .and_then(|s| s.parse().ok()) + .unwrap_or(0.95); + + Ok(Self { + state_dim, + num_actions, + hidden_dims, + use_dueling, + dueling_hidden_dim, + use_distributional, + num_atoms, + use_noisy_nets, + use_iqn, + iqn_embedding_dim, + iqn_num_quantiles, + gamma, + ..Self::default() + }) + } + /// Validate that a checkpoint's metadata matches this config's architecture. /// /// Returns `Ok(())` if the hash matches. Returns an error if metadata is @@ -1904,8 +2007,9 @@ impl DQN { .gather(&next_actions_for_gather, 1)? .squeeze(1)?; // [batch, num_quantiles] - // Compute target quantiles: r + γ * (1 - done) * Z_target - let gamma = self.config.gamma; + // Compute target quantiles: r + γ^n * (1 - done) * Z_target + // Must use gamma^n for n-step returns (NStepBuffer accumulates n-step rewards) + let gamma = self.config.gamma.powi(self.config.n_steps as i32); let rewards_broadcast = rewards_tensor .unsqueeze(1)? .broadcast_as((batch_size, num_quantiles))?; @@ -2020,12 +2124,14 @@ impl DQN { }; // Apply Bellman operator to get target distribution + // Must use gamma^n for n-step returns (NStepBuffer accumulates n-step rewards) + let gamma_n = self.config.gamma.powi(self.config.n_steps as i32); let target_dists = if let Some(ref cat_dist) = self.categorical_dist { cat_dist.apply_bellman_operator( &rewards_tensor, &next_dists, &dones_tensor, - self.config.gamma, + gamma_n, )? } else { return Err(MLError::TrainingError( diff --git a/crates/ml/src/ppo/ppo.rs b/crates/ml/src/ppo/ppo.rs index 78191264b..f85a87489 100644 --- a/crates/ml/src/ppo/ppo.rs +++ b/crates/ml/src/ppo/ppo.rs @@ -1029,6 +1029,14 @@ impl PPO { /// /// Returns the average value loss across all mini-batches and epochs. pub fn update_value_only(&mut self, batch: &mut TrajectoryBatch) -> Result { + // Circuit breaker guard (consistent with update()) + if let Some(ref circuit_breaker) = self.circuit_breaker { + if !circuit_breaker.allow_request() { + warn!("Circuit breaker is open - skipping value-only update"); + return Ok(0.0); + } + } + // Initialize optimizers (we only use value_optimizer, but init both for consistency) self.init_optimizers()?; diff --git a/services/trading_service/src/services/enhanced_ml.rs b/services/trading_service/src/services/enhanced_ml.rs index c366b1f17..f9cf28a0d 100644 --- a/services/trading_service/src/services/enhanced_ml.rs +++ b/services/trading_service/src/services/enhanced_ml.rs @@ -1444,31 +1444,7 @@ impl RealDQNModel { use ml::dqn::agent::DQNAgent; use ml::dqn::DQNConfig; - // TODO: Load architecture config from checkpoint metadata instead of hardcoding - // DQN configuration matching production training defaults - // (51 market + 3 portfolio features, 5 exposure x 3 order x 3 urgency factored actions) - let config = DQNConfig { - state_dim: 54, - num_actions: 45, - hidden_dims: vec![128, 64, 32], - learning_rate: 0.0001, - gamma: 0.99, - epsilon_start: 0.1, - epsilon_end: 0.01, - epsilon_decay: 0.995, - replay_buffer_capacity: 100000, - batch_size: 128, - target_update_freq: 1000, - minimum_profit_factor: 1.5, - tau: 0.005, - weight_decay: 1e-4, - ..DQNConfig::default() - }; - - let mut agent = DQNAgent::new(config) - .map_err(|e| ml::MLError::ModelError(format!("Failed to create DQN agent: {}", e)))?; - - // Try safetensors first, fall back to JSON checkpoint + // Resolve safetensors path let safetensors_path = if checkpoint_path .extension() .is_some_and(|ext| ext == "safetensors") @@ -1483,6 +1459,37 @@ impl RealDQNModel { } }; + // Load architecture config from checkpoint metadata when available. + // This ensures hidden_dims, Rainbow flags, etc. match exactly what was + // trained, regardless of hyperopt-selected architecture. + let config = if let Some(ref st_path) = safetensors_path { + match DQNConfig::from_safetensors_file(st_path) { + Ok(mut cfg) => { + // Override inference-only settings + cfg.epsilon_start = 0.0; // No exploration during inference + cfg.epsilon_end = 0.0; + tracing::info!( + "DQN config loaded from checkpoint metadata: state_dim={}, num_actions={}, hidden_dims={:?}", + cfg.state_dim, cfg.num_actions, cfg.hidden_dims + ); + cfg + } + Err(e) => { + tracing::warn!( + "Failed to read DQN config from checkpoint metadata, using defaults: {}", + e + ); + Self::fallback_dqn_config() + } + } + } else { + // JSON checkpoint — no metadata available + Self::fallback_dqn_config() + }; + + let mut agent = DQNAgent::new(config) + .map_err(|e| ml::MLError::ModelError(format!("Failed to create DQN agent: {}", e)))?; + if let Some(st_path) = safetensors_path { match agent.load_from_safetensors(&st_path) { Ok(()) => { @@ -1518,6 +1525,26 @@ impl RealDQNModel { feature_count: 54, }) } + + /// Fallback DQN config for legacy JSON checkpoints without metadata. + fn fallback_dqn_config() -> ml::dqn::DQNConfig { + ml::dqn::DQNConfig { + state_dim: 54, + num_actions: 45, + hidden_dims: vec![128, 64, 32], + learning_rate: 0.0001, + gamma: 0.95, + epsilon_start: 0.0, + epsilon_end: 0.0, + replay_buffer_capacity: 100_000, + batch_size: 128, + target_update_freq: 1000, + minimum_profit_factor: 1.5, + tau: 0.005, + weight_decay: 1e-4, + ..ml::dqn::DQNConfig::default() + } + } } #[async_trait::async_trait]