From 2a71f6bf4cc3f81f3c51ef005f7137e2c142a700 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 12 Mar 2026 15:55:45 +0100
Subject: [PATCH] =?UTF-8?q?feat(dqn):=20IQN=20+=20Branching=20coexistence?=
 =?UTF-8?q?=20=E2=80=94=20risk-adjusted=20confidence=20scoring?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two changes enabling IQN and Branching DQN to complement each other:

1. Optimizer fix: IQN vars excluded from optimizer when branching is the
   primary loss path. Previously, IQN weights would silently decay to
   zero via weight_decay with zero gradients.

2. Confidence coexistence in select_action_with_confidence():
   When both use_iqn and use_branching are enabled, branching handles
   action decomposition (per-branch greedy selection) while IQN provides
   distributional risk assessment (CVaR-based confidence scoring).
   This is the "complement" design: branching = what to do, IQN = how
   risky.

Tests: ml-dqn=416, 0 failures

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml-dqn/src/dqn.rs | 52 +++++++++++++++++++++++++++++++---------
 1 file changed, 41 insertions(+), 11 deletions(-)

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(