From bafe784dbfe7f01fcd904eb75b8f9c0d0db2b97e Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 20 Feb 2026 15:30:15 +0100
Subject: [PATCH] feat(dqn): add IQN action selection with optional CVaR
 risk-aware mode

Wire IQN into select_action() greedy branch. When use_iqn is enabled,
computes quantile-based Q-values via uniform quantile sampling. CVaR mode
(use_cvar_action_selection) optimizes for worst-case outcomes instead of
mean, providing risk-averse action selection for conservative trading.

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

diff --git a/ml/src/dqn/dqn.rs b/ml/src/dqn/dqn.rs
index 0177b41e5..116d345b9 100644
--- a/ml/src/dqn/dqn.rs
+++ b/ml/src/dqn/dqn.rs
@@ -1116,14 +1116,47 @@ impl DQN {
             )
             .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?;
 
-            let q_values = self.forward(&state_tensor)?;
-            let best_action_idx = q_values
-                .argmax(1)?
-                .get(0)?
-                .to_scalar::<u32>()
-                .map_err(|e| MLError::ModelError(format!("Failed to get best action: {}", e)))?;
+            if self.config.use_iqn && self.iqn_network.is_some() {
+                // IQN: Compute quantile-based Q-values
+                let iqn_net = self.iqn_network.as_ref().unwrap();
+                let state_embed = self.get_state_embedding(&state_tensor)?;
 
-            FactoredAction::from_index(best_action_idx as usize)?
+                // Use fixed uniform quantiles for deterministic action selection
+                let taus = iqn_net.sample_uniform_quantiles(1, &self.device)
+                    .map_err(|e| MLError::ModelError(format!("Failed to sample taus: {}", e)))?;
+
+                // Forward: [1, num_actions, num_quantiles]
+                let all_quantiles = iqn_net.forward(&state_embed, &taus)?;
+
+                // Score each action
+                let action_scores = if self.config.use_cvar_action_selection {
+                    // CVaR: Optimize for worst-case outcomes (risk-averse)
+                    iqn_net.compute_cvar(&all_quantiles, self.config.cvar_alpha)
+                        .map_err(|e| MLError::ModelError(format!("CVaR computation failed: {}", e)))?
+                } else {
+                    // Standard: Mean over quantiles (risk-neutral)
+                    iqn_net.to_expected_q(&all_quantiles)
+                        .map_err(|e| MLError::ModelError(format!("Expected Q computation failed: {}", e)))?
+                };
+
+                // Argmax over actions
+                let best_action_idx = action_scores
+                    .argmax(1)?
+                    .get(0)?
+                    .to_scalar::<u32>()?;
+
+                FactoredAction::from_index(best_action_idx as usize)?
+            } else {
+                // Standard Q-network action selection
+                let q_values = self.forward(&state_tensor)?;
+                let best_action_idx = q_values
+                    .argmax(1)?
+                    .get(0)?
+                    .to_scalar::<u32>()
+                    .map_err(|e| MLError::ModelError(format!("Failed to get best action: {}", e)))?;
+
+                FactoredAction::from_index(best_action_idx as usize)?
+            }
         };
 
         // Track action for entropy penalty calculation
@@ -2704,6 +2737,50 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_iqn_action_selection() {
+        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 = 8;
+        config.use_distributional = false;
+        config.use_dueling = false;
+        config.epsilon_start = 0.0; // Force greedy for testing
+        config.use_noisy_nets = false;
+        config.warmup_steps = 0;
+
+        let mut dqn = DQN::new(config).unwrap();
+
+        let state = vec![0.5f32; 8];
+        let action = dqn.select_action(&state);
+        assert!(action.is_ok(), "IQN action selection should succeed: {:?}", action.err());
+    }
+
+    #[test]
+    fn test_iqn_cvar_action_selection() {
+        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 = 8;
+        config.use_distributional = false;
+        config.use_dueling = false;
+        config.epsilon_start = 0.0;
+        config.use_noisy_nets = false;
+        config.warmup_steps = 0;
+        config.use_cvar_action_selection = true;
+        config.cvar_alpha = 0.05;
+
+        let mut dqn = DQN::new(config).unwrap();
+
+        let state = vec![0.5f32; 8];
+        let action = dqn.select_action(&state);
+        assert!(action.is_ok(), "IQN CVaR action selection should succeed: {:?}", action.err());
+    }
+
     #[test]
     fn test_dqn_with_iqn_creation() {
         let mut config = DQNConfig::default();