diff --git a/crates/ml-dqn/src/branching.rs b/crates/ml-dqn/src/branching.rs index 0a76b0ebe..6d26a3e9a 100644 --- a/crates/ml-dqn/src/branching.rs +++ b/crates/ml-dqn/src/branching.rs @@ -36,7 +36,7 @@ //! - **`NoisyNet`**: Factorized Gaussian noise in value/branch heads for learned exploration. //! - **State Dim Alignment**: Auto-pad `state_dim` to multiples of 8 for tensor core HMMA. -use candle_core::{DType, Device, ModuleT, Tensor}; +use candle_core::{DType, Device, ModuleT, Tensor, Var}; use candle_nn::{Dropout, Linear, Module, VarBuilder, VarMap}; use serde::{Deserialize, Serialize}; @@ -218,6 +218,32 @@ impl MaybeNoisyLinear { } Ok(()) } + + /// Collect only sigma (noise std dev) `Var`s from `NoisyLinear` layers. + /// + /// When mu vars are registered in `VarMap`, this avoids double-counting them + /// in `all_trainable_vars()` and `noisy_vars_ordered()`. + fn noisy_sigma_vars(&self) -> Vec { + match self { + Self::Standard(_) => Vec::new(), + Self::Noisy(n) => n.sigma_vars().iter().map(|v| (*v).clone()).collect(), + } + } + + /// Register mu (weight/bias) `Var`s in `VarMap` under `{name}.weight` / `{name}.bias`. + /// + /// The GPU experience collector looks up weights by name from `VarMap`. `NoisyLinear` + /// creates standalone `Var`s not in `VarMap`, so the collector fails with + /// "Missing weight: `value_fc.weight`". Registering mu vars fixes this. + fn register_mu_in_varmap(&self, varmap: &VarMap, name: &str) { + if let Self::Noisy(n) = self { + let [w_mu, b_mu] = n.mu_vars(); + if let Ok(mut data) = varmap.data().lock() { + data.insert(format!("{name}.weight"), w_mu.clone()); + data.insert(format!("{name}.bias"), b_mu.clone()); + } + } + } } impl std::fmt::Debug for MaybeNoisyLinear { @@ -359,6 +385,18 @@ impl BranchingDuelingQNetwork { ) }).transpose()?; + // Register NoisyLinear mu weights in VarMap so the GPU experience collector + // can find them by name (e.g. "value_fc.weight"). Without this, the collector + // fails with "Missing weight: value_fc.weight" and falls back to CPU. + value_fc.register_mu_in_varmap(&vars, "value_fc"); + value_out.register_mu_in_varmap(&vars, "value_out"); + for (d, fc) in branch_fcs.iter().enumerate() { + fc.register_mu_in_varmap(&vars, &format!("branch_{d}_fc")); + } + for (d, out) in branch_outs.iter().enumerate() { + out.register_mu_in_varmap(&vars, &format!("branch_{d}_out")); + } + Ok(Self { shared_layers, value_fc, @@ -900,6 +938,46 @@ impl BranchingDuelingQNetwork { &self.vars } + /// Collect ALL trainable `Var`s: `VarMap` (shared encoder + mu) + `NoisyLinear` sigma. + /// + /// Mu vars (`weight_mu`, `bias_mu`) are registered in `VarMap` at construction + /// time for GPU experience collector compatibility. Sigma vars (`weight_sigma`, + /// `bias_sigma`) remain standalone. This method collects both without duplication. + pub fn all_trainable_vars(&self) -> Vec { + let mut vars = self.vars.all_vars(); // shared encoder + NoisyLinear mu vars + // Only sigma vars — mu already in VarMap + vars.extend(self.value_fc.noisy_sigma_vars()); + vars.extend(self.value_out.noisy_sigma_vars()); + for fc in &self.branch_fcs { + vars.extend(fc.noisy_sigma_vars()); + } + for out in &self.branch_outs { + vars.extend(out.noisy_sigma_vars()); + } + vars + } + + /// Collect only the `NoisyLinear` sigma `Var`s (for target network Polyak updates). + /// + /// Mu vars are registered in `VarMap`, so `polyak_update()` on `VarMap` handles them. + /// This method returns only sigma vars for `polyak_update_var_pairs()`. + /// + /// Returns vars in a deterministic order: `value_fc`, `value_out`, then + /// `branch_fcs[0..D]`, `branch_outs[0..D]`. Both online and target networks + /// produce the same order, so vars can be zipped for Polyak update. + pub fn noisy_vars_ordered(&self) -> Vec { + let mut vars = Vec::new(); + vars.extend(self.value_fc.noisy_sigma_vars()); + vars.extend(self.value_out.noisy_sigma_vars()); + for fc in &self.branch_fcs { + vars.extend(fc.noisy_sigma_vars()); + } + for out in &self.branch_outs { + vars.extend(out.noisy_sigma_vars()); + } + vars + } + /// Get device. pub const fn device(&self) -> &Device { &self.device @@ -911,24 +989,59 @@ impl BranchingDuelingQNetwork { } /// Copy weights from another branching network (target network sync). + /// + /// Copies both `VarMap` vars (shared encoder) AND `NoisyLinear` head vars. pub fn copy_weights_from(&mut self, other: &BranchingDuelingQNetwork) -> 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), - })?; + // 1. Copy VarMap vars (shared encoder layers) + { + 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) { - self_var.set(other_var.as_tensor()).map_err(|e| { - MLError::ModelError(format!("Copy weight {}: {}", name, e)) - })?; + for (name, self_var) in self_vars.iter() { + if let Some(other_var) = other_vars.get(name) { + self_var.set(other_var.as_tensor()).map_err(|e| { + MLError::ModelError(format!("Copy weight {}: {}", name, e)) + })?; + } } } + // 2. Copy NoisyLinear head vars (not in VarMap — standalone Vars) + Self::copy_noisy_layer(&mut self.value_fc, &other.value_fc, "value_fc")?; + Self::copy_noisy_layer(&mut self.value_out, &other.value_out, "value_out")?; + for (d, (self_fc, other_fc)) in self.branch_fcs.iter_mut().zip(other.branch_fcs.iter()).enumerate() { + Self::copy_noisy_layer(self_fc, other_fc, &format!("branch_{}_fc", d))?; + } + for (d, (self_out, other_out)) in self.branch_outs.iter_mut().zip(other.branch_outs.iter()).enumerate() { + Self::copy_noisy_layer(self_out, other_out, &format!("branch_{}_out", d))?; + } + Ok(()) } + + /// Copy `NoisyLinear` weights between matching layers (no-op for `Standard` layers). + fn copy_noisy_layer(dst: &mut MaybeNoisyLinear, src: &MaybeNoisyLinear, label: &str) -> Result<(), MLError> { + match (dst, src) { + (MaybeNoisyLinear::Noisy(d), MaybeNoisyLinear::Noisy(s)) => { + let dst_vars = d.vars(); + let src_vars = s.vars(); + for (dv, sv) in dst_vars.iter().zip(src_vars.iter()) { + dv.set(sv.as_tensor()).map_err(|e| { + MLError::ModelError(format!("Copy noisy weight {}: {}", label, e)) + })?; + } + Ok(()) + } + (MaybeNoisyLinear::Standard(_), MaybeNoisyLinear::Standard(_)) => Ok(()), + _ => Err(MLError::ModelError(format!( + "Layer type mismatch copying {}", label + ))), + } + } } impl std::fmt::Debug for BranchingDuelingQNetwork { @@ -1810,6 +1923,89 @@ mod tests { Ok(()) } + #[test] + fn test_all_trainable_vars_includes_noisy() -> anyhow::Result<()> { + // NoisyLinear vars must be in all_trainable_vars() for the optimizer to update them + let mut config = trading_config_distributional(8); + config.use_noisy = true; + config.noisy_sigma_init = 0.5; + + let net = BranchingDuelingQNetwork::new(config, Device::Cpu)?; + + let varmap_only = net.vars().all_vars(); + let all_vars = net.all_trainable_vars(); + let sigma_only = net.noisy_vars_ordered(); + + // VarMap has shared encoder (4) + NoisyLinear mu vars (8 layers × 2 = 16) = 20 + assert_eq!( + varmap_only.len(), 20, + "VarMap should have shared encoder (4) + mu weights (16)" + ); + // noisy_vars_ordered returns only sigma vars: 8 layers × 2 = 16 + assert_eq!( + sigma_only.len(), 16, + "8 NoisyLinear layers × 2 sigma vars each" + ); + // all_trainable_vars = VarMap (shared + mu) + sigma + assert_eq!( + all_vars.len(), + varmap_only.len() + sigma_only.len(), + "all_trainable_vars should combine VarMap ({}) + sigma ({})", + varmap_only.len(), + sigma_only.len() + ); + // Total: 20 + 16 = 36 (4 shared + 8×4 NoisyLinear params) + assert_eq!(all_vars.len(), 36, "4 shared + 32 NoisyLinear = 36 total"); + + Ok(()) + } + + #[test] + fn test_noisy_weight_copy() -> anyhow::Result<()> { + // Verify copy_weights_from syncs ALL vars (VarMap mu + standalone sigma) + let mut config = trading_config_distributional(8); + config.use_noisy = true; + config.noisy_sigma_init = 0.5; + + let net1 = BranchingDuelingQNetwork::new(config.clone(), Device::Cpu)?; + let mut net2 = BranchingDuelingQNetwork::new(config, Device::Cpu)?; + + let state = Tensor::ones((1, 8), DType::F32, &Device::Cpu)?; + + // Before copy: outputs should differ (random mu init) + let out1 = net1.forward_branches_eval(&state)?; + let out2 = net2.forward_branches_eval(&state)?; + let v1: Vec = out1.value.flatten_all()?.to_vec1()?; + let v2: Vec = out2.value.flatten_all()?.to_vec1()?; + let diff_before: f32 = v1.iter().zip(v2.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum(); + assert!(diff_before > 1e-6, "Before copy, outputs should differ: {}", diff_before); + + net2.copy_weights_from(&net1)?; + + // After copy: outputs should match (all weights synced) + let out1a = net1.forward_branches_eval(&state)?; + let out2a = net2.forward_branches_eval(&state)?; + let v1a: Vec = out1a.value.flatten_all()?.to_vec1()?; + let v2a: Vec = out2a.value.flatten_all()?.to_vec1()?; + let diff_after: f32 = v1a.iter().zip(v2a.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum(); + assert!(diff_after < 1e-6, "After copy, outputs should match: {}", diff_after); + + // Also verify sigma vars were copied (ordered, so zip is deterministic) + let s1 = net1.noisy_vars_ordered(); + let s2 = net2.noisy_vars_ordered(); + assert_eq!(s1.len(), s2.len()); + for (a, b) in s1.iter().zip(s2.iter()) { + let d = a.as_tensor().sub(b.as_tensor())?.sqr()?.sum_all()?.to_scalar::()?; + assert!(d < 1e-10, "Sigma var mismatch: {}", d); + } + + Ok(()) + } + #[test] fn test_config_serde_defaults() -> anyhow::Result<()> { // Verify serde defaults for new fields diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 1822ff8d0..8a20f1dbc 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -2268,9 +2268,14 @@ impl DQN { // Wave 11.6: Fix Wave 10.3 optimizer issue - use correct network parameters // Priority: branching > IQN+base > hybrid > dueling > standard let mut vars = if self.config.use_branching { - // Branching DQN: use branching network parameters exclusively + // Branching DQN: use ALL trainable parameters (VarMap + NoisyLinear heads). + // NoisyLinear creates standalone Vars via Var::from_tensor() — NOT + // registered in the VarMap. Using only vars().all_vars() leaves + // the head parameters frozen and can cause device mismatch in + // clip_grad_norm when backward() produces gradients the optimizer + // doesn't know about. if let Some(ref branching_net) = self.branching_q_network { - branching_net.vars().all_vars() + branching_net.all_trainable_vars() } else { self.q_network.vars().all_vars() } @@ -3314,9 +3319,30 @@ impl DQN { grad_norm_tensor }; + // Device guard: ensure returned tensors are on the network's device. + // Prevents device mismatch in RegimeConditionalDQN's accumulation loop. + let loss_gpu = if loss_clamped.device().location() != self.device.location() { + tracing::warn!( + "train_step: loss on {:?}, expected {:?} — migrating", + loss_clamped.device(), self.device + ); + loss_clamped.to_device(&self.device)? + } else { + loss_clamped + }; + let grad_norm_gpu = if grad_norm_scalar.device().location() != self.device.location() { + tracing::warn!( + "train_step: grad_norm on {:?}, expected {:?} — migrating", + grad_norm_scalar.device(), self.device + ); + grad_norm_scalar.to_device(&self.device)? + } else { + grad_norm_scalar + }; + Ok(GpuTrainResult { - loss_gpu: loss_clamped, - grad_norm_gpu: grad_norm_scalar, + loss_gpu, + grad_norm_gpu, }) } @@ -3530,6 +3556,7 @@ impl DQN { } // Also update Branching target network if present + // Two-phase: (1) VarMap vars (shared encoder), (2) NoisyLinear head vars if let (Some(ref branching_net), Some(ref branching_target)) = (&self.branching_q_network, &self.branching_target_network) { @@ -3537,6 +3564,14 @@ impl DQN { .map_err(|e| { MLError::TrainingError(format!("Branching EMA update failed: {}", e)) })?; + // Polyak-update NoisyLinear head vars (not in VarMap) + let online_noisy = branching_net.noisy_vars_ordered(); + let target_noisy = branching_target.noisy_vars_ordered(); + super::target_update::polyak_update_var_pairs( + &online_noisy, &target_noisy, current_tau, + ).map_err(|e| { + MLError::TrainingError(format!("Branching NoisyLinear EMA update failed: {}", e)) + })?; } // Log EMA update every 1000 steps diff --git a/crates/ml-dqn/src/noisy_layers.rs b/crates/ml-dqn/src/noisy_layers.rs index 574da4c03..99ac9704b 100644 --- a/crates/ml-dqn/src/noisy_layers.rs +++ b/crates/ml-dqn/src/noisy_layers.rs @@ -269,6 +269,22 @@ impl NoisyLinear { vec![&self.weight_mu, &self.bias_mu, &self.weight_sigma, &self.bias_sigma] } + /// Get only mu (mean) parameters — `weight_mu`, `bias_mu`. + /// + /// Used when mu vars are registered in `VarMap` for GPU experience collector + /// compatibility; sigma vars are managed separately. + pub const fn mu_vars(&self) -> [&Var; 2] { + [&self.weight_mu, &self.bias_mu] + } + + /// Get only sigma (noise std dev) parameters — `weight_sigma`, `bias_sigma`. + /// + /// Used alongside `VarMap` vars in `all_trainable_vars()`: mu vars live in + /// `VarMap` (for GPU weight extraction), sigma vars are standalone. + pub const fn sigma_vars(&self) -> [&Var; 2] { + [&self.weight_sigma, &self.bias_sigma] + } + /// Disable noise for evaluation (use mean parameters only) pub fn disable_noise(&mut self) -> Result<(), MLError> { // Set epsilon buffers to zero in training dtype (effectively uses μ only) diff --git a/crates/ml-dqn/src/target_update.rs b/crates/ml-dqn/src/target_update.rs index f33c99cfb..3249d5a58 100644 --- a/crates/ml-dqn/src/target_update.rs +++ b/crates/ml-dqn/src/target_update.rs @@ -70,6 +70,30 @@ pub fn polyak_update(online_vars: &VarMap, target_vars: &VarMap, tau: f64) -> Ca Ok(()) } +/// Polyak (EMA) update on raw `Var` pairs — for `NoisyLinear` vars not in a `VarMap`. +/// +/// `θ_target` = (1-τ) × `θ_target` + τ × `θ_online` +pub fn polyak_update_var_pairs( + online: &[candle_core::Var], + target: &[candle_core::Var], + tau: f64, +) -> CandleResult<()> { + debug_assert_eq!( + online.len(), + target.len(), + "polyak_update_var_pairs: online ({}) and target ({}) var counts must match", + online.len(), + target.len(), + ); + for (o, t) in online.iter().zip(target.iter()) { + let online_t = o.as_tensor(); + let target_t = t.as_tensor(); + let new_target = ((target_t * (1.0 - tau))? + (online_t * tau)?)?; + t.set(&new_target)?; + } + Ok(()) +} + /// Hard update (copy all weights) /// /// Used for: diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index eac9697cb..0f5eb2b3a 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -910,6 +910,11 @@ pub struct DQNTrainer { epochs: usize, buffer_size_max: usize, runtime_handle: Option, + /// Owned multi-thread Tokio runtime — keeps the handle valid for the adapter lifetime. + /// Without this, each trial creates a single-threaded `Runtime::new()` that pins all + /// async work (training, backtest, data loading) to one OS thread. + #[allow(dead_code)] + _owned_runtime: Option>, training_paths: TrainingPaths, device: candle_core::Device, // Initialize CUDA early like MAMBA-2 /// Pool of GPU devices for multi-GPU trial parallelism. @@ -1032,15 +1037,24 @@ impl DQNTrainer { .map_err(|e| MLError::ConfigError(format!("CUDA GPU required for DQN hyperopt: {}", e)))?; info!(" Device: CUDA GPU"); - // Try to reuse existing Tokio runtime, create new one if needed - let runtime_handle = match tokio::runtime::Handle::try_current() { + // Reuse existing runtime or create a multi-threaded one. + // Runtime::new() creates a current-thread scheduler → 1 OS thread → CPU bottleneck. + // Multi-thread allows async overlap: GPU kernel dispatch + CPU backtest + data I/O. + let (runtime_handle, _owned_runtime) = match tokio::runtime::Handle::try_current() { Ok(handle) => { info!(" Runtime: Reusing existing Tokio runtime"); - Some(handle) + (Some(handle), None) }, Err(_) => { - info!(" Runtime: Will create new Tokio runtime per trial"); - None + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .map_err(|e| MLError::ConfigError(format!("Failed to create Tokio runtime: {}", e)))?; + let handle = rt.handle().clone(); + let rt_arc = Arc::new(rt); + info!(" Runtime: Created multi-thread Tokio runtime (4 workers)"); + (Some(handle), Some(rt_arc)) }, }; @@ -1052,6 +1066,7 @@ impl DQNTrainer { epochs, buffer_size_max, runtime_handle, + _owned_runtime, training_paths, device_pool: vec![device.clone()], device, @@ -1258,15 +1273,12 @@ impl DQNTrainer { } // Load data using the internal trainer's async method - let (train_data, val_data) = if let Some(handle) = &self.runtime_handle { - handle.block_on(loader.load_training_data(data_path_str)) - } else { - let runtime = tokio::runtime::Runtime::new().map_err(|e| { - MLError::TrainingError(format!("Failed to create runtime for preload: {}", e)) - })?; - runtime.block_on(loader.load_training_data(data_path_str)) - } - .map_err(|e| MLError::TrainingError(format!("Failed to preload data: {}", e)))?; + let preload_handle = self.runtime_handle.as_ref().ok_or_else(|| { + MLError::ConfigError("BUG: runtime_handle is None — DQNTrainer::new() should always set it".to_owned()) + })?; + let (train_data, val_data) = preload_handle + .block_on(loader.load_training_data(data_path_str)) + .map_err(|e| MLError::TrainingError(format!("Failed to preload data: {}", e)))?; // Load OFI features from MBP-10 data (separate from OHLCV loading). // The internal trainer's load_training_data() only loads OHLCV — OFI must be loaded @@ -2676,39 +2688,22 @@ impl HyperparameterOptimizable for DQNTrainer { // Borrow preloaded data (zero-copy) or fall back to loading from disk. // Val data is always cloned into the trainer (small: ~20% of total). // Training data is borrowed via train_with_shared_data (large: ~80%). - let training_metrics = if let Some(handle) = &self.runtime_handle { - if use_preloaded { - let train_ref = cached_train.as_ref().ok_or_else(|| { - MLError::TrainingError("cached_train is None but use_preloaded is true".to_owned()) - })?; - let val_data = cached_val.as_ref().map(|a| a.as_ref().clone()).unwrap_or_default(); - info!("Training DQN with shared data ({} train, {} val samples)", - train_ref.len(), val_data.len()); - handle.block_on(internal_trainer.train_with_shared_data( - train_ref, val_data, checkpoint_callback, - )) - } else { - info!("Training DQN with DBN directory: {}", data_path_str); - handle.block_on(internal_trainer.train(data_path_str, checkpoint_callback)) - } - } else { - let runtime = tokio::runtime::Runtime::new().map_err(|e| { - MLError::TrainingError(format!("Failed to create runtime: {}", e)) + let handle = self.runtime_handle.as_ref().ok_or_else(|| { + MLError::ConfigError("BUG: runtime_handle is None — DQNTrainer::new() should always set it".to_owned()) + })?; + let training_metrics = if use_preloaded { + let train_ref = cached_train.as_ref().ok_or_else(|| { + MLError::TrainingError("cached_train is None but use_preloaded is true".to_owned()) })?; - if use_preloaded { - let train_ref = cached_train.as_ref().ok_or_else(|| { - MLError::TrainingError("cached_train is None but use_preloaded is true".to_owned()) - })?; - let val_data = cached_val.as_ref().map(|a| a.as_ref().clone()).unwrap_or_default(); - info!("Training DQN with shared data ({} train, {} val samples)", - train_ref.len(), val_data.len()); - runtime.block_on(internal_trainer.train_with_shared_data( - train_ref, val_data, checkpoint_callback, - )) - } else { - info!("Training DQN with DBN directory: {}", data_path_str); - runtime.block_on(internal_trainer.train(data_path_str, checkpoint_callback)) - } + let val_data = cached_val.as_ref().map(|a| a.as_ref().clone()).unwrap_or_default(); + info!("Training DQN with shared data ({} train, {} val samples)", + train_ref.len(), val_data.len()); + handle.block_on(internal_trainer.train_with_shared_data( + train_ref, val_data, checkpoint_callback, + )) + } else { + info!("Training DQN with DBN directory: {}", data_path_str); + handle.block_on(internal_trainer.train(data_path_str, checkpoint_callback)) } .map_err(|e| MLError::TrainingError(format!("DQN training failed: {}", e)))?; @@ -2938,10 +2933,10 @@ impl HyperparameterOptimizable for DQNTrainer { ); // Acquire read lock once — batch_greedy_actions is &self - let runtime = tokio::runtime::Runtime::new().map_err(|e| { - MLError::TrainingError(format!("Failed to create runtime for backtest: {}", e)) + let bt_handle = self.runtime_handle.as_ref().ok_or_else(|| { + MLError::ConfigError("BUG: runtime_handle is None — DQNTrainer::new() should always set it".to_owned()) })?; - let agent_guard = runtime.block_on(agent_arc.read()); + let agent_guard = bt_handle.block_on(agent_arc.read()); // Collect per-window metrics let mut window_metrics: Vec = Vec::with_capacity(window_count);