diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 9fbb6b62a..9ffd07028 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -2042,18 +2042,44 @@ impl DQN { })?; let output = branching_net.forward_branches_eval(&state_tensor)?; - // Per-branch confidence: average of per-branch softmax max probs - let mut total_conf = 0.0_f32; - for adv in &output.advantages { - let branch_conf = Self::softmax_confidence(adv)?; - total_conf += branch_conf; - } - let conf = (total_conf / output.advantages.len() as f32).clamp(0.5, 0.95); - let branch_actions = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output)?; let exposure = ExposureLevel::from_index(branch_actions.first().copied().unwrap_or(2) as usize)?; let order = OrderType::from_index(branch_actions.get(1).copied().unwrap_or(0) as usize)?; let urgency = Urgency::from_index(branch_actions.get(2).copied().unwrap_or(1) as usize)?; + + // Confidence: IQN CVaR when available (risk-sensitive), else branching softmax. + // IQN + Branching coexistence: branching handles action decomposition, + // IQN provides distributional risk assessment for confidence scoring. + let conf = if self.config.use_iqn { + if let Some(ref iqn_net) = self.iqn_network { + let state_embed = self.get_state_embedding(&state_tensor)?; + let taus = iqn_net + .sample_uniform_quantiles(1, &self.device) + .map_err(|e| MLError::ModelError(format!("IQN taus: {}", e)))?; + let all_quantiles = iqn_net.forward(&state_embed, &taus)?; + let cvar_scores = iqn_net + .compute_cvar(&all_quantiles, self.config.cvar_alpha) + .map_err(|e| MLError::ModelError(format!("IQN CVaR: {}", e)))?; + // CVaR-based confidence: softmax over CVaR scores of exposure actions + Self::softmax_confidence(&cvar_scores)? + } else { + // IQN enabled but network not initialized — fall through to branching + let mut total_conf = 0.0_f32; + for adv in &output.advantages { + total_conf += Self::softmax_confidence(adv)?; + } + total_conf / output.advantages.len() as f32 + } + } else { + // Per-branch confidence: average of per-branch softmax max probs + let mut total_conf = 0.0_f32; + for adv in &output.advantages { + total_conf += Self::softmax_confidence(adv)?; + } + total_conf / output.advantages.len() as f32 + }; + let conf = conf.clamp(0.5, 0.95); + return Ok((FactoredAction { exposure, order, urgency }, conf)); } @@ -2422,9 +2448,13 @@ impl DQN { self.q_network.vars().all_vars() }; - // IQN: Add IQN network vars (loss flows through both IQN and base q_network) - if let Some(ref iqn_net) = self.iqn_network { - vars.extend(iqn_net.vars().all_vars()); + // IQN: Add IQN network vars only when IQN is the primary loss path. + // When branching is active, IQN gets zero gradients but weight_decay + // would silently decay its weights to zero — skip to prevent that. + if !self.config.use_branching { + if let Some(ref iqn_net) = self.iqn_network { + vars.extend(iqn_net.vars().all_vars()); + } } self.optimizer = Some(