From 02393a4cd1ea958a9eb42b1eb3ec47c6289ca2ef Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 20 Feb 2026 15:09:29 +0100
Subject: [PATCH] feat(dqn): add CQL offline RL regularization to train_step
 (Kumar et al. 2020)

CQL penalty = logsumexp(Q(s, all_a)) - Q(s, a_data)
- Numerically stable logsumexp with max subtraction
- Controlled by use_cql and cql_alpha config fields
- Periodic logging every 100 steps
- test_cql_regularization passes

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

diff --git a/ml/src/dqn/dqn.rs b/ml/src/dqn/dqn.rs
index acc568876..33fed992c 100644
--- a/ml/src/dqn/dqn.rs
+++ b/ml/src/dqn/dqn.rs
@@ -1636,7 +1636,37 @@ impl DQN {
         let entropy_penalty = self.calculate_entropy_penalty()?;
         let entropy_weight = 0.1; // Match Wave 5-A1 diversity_weight
         let entropy_term = (entropy_penalty * entropy_weight)?;
-        let loss = loss_value.add(&entropy_term)?;
+        let loss_with_entropy = loss_value.add(&entropy_term)?;
+
+        // CQL regularization for offline RL (Kumar et al. 2020)
+        // Penalizes high Q-values for actions not taken in the training data
+        let loss = if self.config.use_cql {
+            // logsumexp(Q(s, all_actions)) — numerically stable soft maximum
+            let q_max = current_q_values.max(1)?;
+            let q_max_broadcast = q_max.unsqueeze(1)?
+                .broadcast_as(current_q_values.shape())?;
+            let q_shifted = (&current_q_values - &q_max_broadcast)?;
+            let logsumexp = (q_shifted.exp()?.sum(1)?.log()? + q_max)?;
+
+            // Q(s, a_data) — Q-values for actions actually taken
+            let q_data = &state_action_values;
+
+            // CQL penalty = mean(logsumexp - Q_data)
+            let cql_penalty = (logsumexp - q_data)?.mean_all()?;
+            let cql_term = (cql_penalty * self.config.cql_alpha)?;
+
+            if self.training_steps % 100 == 0 {
+                let cql_val: f32 = cql_term.to_scalar().unwrap_or(0.0);
+                tracing::debug!(
+                    "CQL penalty at step {}: {:.4} (alpha: {:.2})",
+                    self.training_steps, cql_val, self.config.cql_alpha
+                );
+            }
+
+            loss_with_entropy.add(&cql_term)?
+        } else {
+            loss_with_entropy
+        };
 
         // Extract loss value before backward pass
         let loss_value = loss
@@ -2457,4 +2487,39 @@ mod tests {
         assert!(result.is_ok());
         Ok(())
     }
+
+    #[test]
+    fn test_cql_regularization() -> anyhow::Result<()> {
+        let mut config = DQNConfig::emergency_safe_defaults();
+        config.state_dim = 8;
+        config.num_actions = 3;
+        config.hidden_dims = vec![16, 16];
+        config.use_cql = true;
+        config.cql_alpha = 1.0;
+        config.use_iqn = false;
+        config.use_distributional = false;
+        config.use_dueling = false;
+        config.batch_size = 4;
+        config.min_replay_size = 4;
+
+        let mut dqn = DQN::new(config)?;
+
+        for i in 0..10 {
+            let exp = Experience::new(
+                vec![0.1 * i as f32; 8],
+                (i % 3) as u8,
+                0.5,
+                vec![0.2 * i as f32; 8],
+                false,
+            );
+            dqn.store_experience(exp)?;
+        }
+
+        let result = dqn.train_step(None);
+        assert!(result.is_ok(), "Training with CQL should succeed: {:?}", result.err());
+
+        let (loss, _grad_norm) = result.unwrap();
+        assert!(loss.is_finite(), "CQL loss should be finite, got {}", loss);
+        Ok(())
+    }
 }