From 33a3ee99eaa45987401329bb7050ddf780276334 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 20 Feb 2026 15:17:29 +0100
Subject: [PATCH] feat(dqn): add IQN network fields to DQN struct with
 initialization

Add iqn_network and iqn_target_network as Option<QuantileNetwork> fields
to the DQN struct. When use_iqn is enabled, DQN::new() creates both networks
with QuantileConfig derived from DQNConfig (embed_dim from last hidden layer).
Each network gets its own VarMap for independent parameter tracking.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ml/src/dqn/dqn.rs | 54 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 53 insertions(+), 1 deletion(-)

diff --git a/ml/src/dqn/dqn.rs b/ml/src/dqn/dqn.rs
index 33fed992c..1f0d4577b 100644
--- a/ml/src/dqn/dqn.rs
+++ b/ml/src/dqn/dqn.rs
@@ -730,6 +730,11 @@ pub struct DQN {
     /// N-step buffer for multi-step TD learning (Wave 11.5)
     nstep_buffer: Option<Arc<Mutex<super::NStepBuffer>>>,
 
+    /// IQN (Implicit Quantile Network) for distributional RL (Dabney et al. 2018b)
+    iqn_network: Option<super::quantile_regression::QuantileNetwork>,
+    /// IQN target network (frozen copy for stable target computation)
+    iqn_target_network: Option<super::quantile_regression::QuantileNetwork>,
+
     // WAVE 23 P0 Fix #2: Early stopping on gradient collapse
     /// Consecutive epochs with gradient collapse (resets when gradients recover)
     gradient_collapse_counter: usize,
@@ -868,6 +873,34 @@ impl DQN {
             None
         };
 
+        // IQN initialization (replaces broken C51 — no scatter_add needed)
+        let (iqn_network, iqn_target_network) = if config.use_iqn {
+            let iqn_config = super::quantile_regression::QuantileConfig {
+                num_quantiles: config.iqn_num_quantiles,
+                quantile_embedding_dim: config.iqn_embedding_dim,
+                kappa: config.iqn_kappa,
+                num_actions: config.num_actions,
+            };
+
+            let embed_dim = config.hidden_dims.last().copied().unwrap_or(config.state_dim);
+
+            let iqn_vars = VarMap::new();
+            let iqn_vb = VarBuilder::from_varmap(&iqn_vars, DType::F32, &device);
+            let iqn_net = super::quantile_regression::QuantileNetwork::new(
+                &iqn_config, embed_dim, iqn_vb
+            )?;
+
+            let iqn_target_vars = VarMap::new();
+            let iqn_target_vb = VarBuilder::from_varmap(&iqn_target_vars, DType::F32, &device);
+            let iqn_target = super::quantile_regression::QuantileNetwork::new(
+                &iqn_config, embed_dim, iqn_target_vb
+            )?;
+
+            (Some(iqn_net), Some(iqn_target))
+        } else {
+            (None, None)
+        };
+
         Ok(Self {
             epsilon: config.epsilon_start,
             q_network,
@@ -886,7 +919,8 @@ impl DQN {
             device,
             recent_actions: VecDeque::with_capacity(100),
             nstep_buffer,
-            // WAVE 23 P0 Fix #2/#4: Initialize early stopping counters
+            iqn_network,
+            iqn_target_network,
             gradient_collapse_counter: 0,
             q_value_divergence_counter: 0,
         })
@@ -2522,4 +2556,22 @@ mod tests {
         assert!(loss.is_finite(), "CQL loss should be finite, got {}", loss);
         Ok(())
     }
+
+    #[test]
+    fn test_dqn_with_iqn_creation() {
+        let mut config = DQNConfig::default();
+        config.state_dim = 8;
+        config.num_actions = 3;
+        config.hidden_dims = vec![16, 16];
+        config.use_iqn = true;
+        config.iqn_num_quantiles = 32;
+        config.use_distributional = false;
+        config.use_dueling = false;
+        let dqn = DQN::new(config);
+        assert!(
+            dqn.is_ok(),
+            "DQN with IQN should create successfully: {:?}",
+            dqn.err()
+        );
+    }
 }