From 3e5db0b377fbc8176ddaeea74f7e39e6c8e34c12 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 20 Feb 2026 14:51:48 +0100
Subject: [PATCH] fix(dqn): NaN sort panic, Adam eps per Rainbow paper, dedup
 DistributionalType
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Fix NaN panic at dqn.rs:1716 with unwrap_or(Ordering::Equal)
- Fix Adam epsilon 1e-8 → 1.5e-4 per Hessel et al. 2018
- Remove duplicate DistributionalType enum from quantile_regression.rs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ml/src/dqn/dqn.rs                 |  2 +-
 ml/src/dqn/quantile_regression.rs | 23 ++---------------------
 ml/src/dqn/rainbow_agent_impl.rs  |  2 +-
 3 files changed, 4 insertions(+), 23 deletions(-)

diff --git a/ml/src/dqn/dqn.rs b/ml/src/dqn/dqn.rs
index 070d2f29b..761b64f99 100644
--- a/ml/src/dqn/dqn.rs
+++ b/ml/src/dqn/dqn.rs
@@ -1713,7 +1713,7 @@ impl DQN {
         // Log top-5 best actions (shows which actions network prefers)
         let mut indexed: Vec<(usize, f32)> = q_vec.iter().enumerate()
             .map(|(i, &v)| (i, v)).collect();
-        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
 
         tracing::info!(
             "Top-5 actions: {}",
diff --git a/ml/src/dqn/quantile_regression.rs b/ml/src/dqn/quantile_regression.rs
index 69935500c..105146b72 100644
--- a/ml/src/dqn/quantile_regression.rs
+++ b/ml/src/dqn/quantile_regression.rs
@@ -309,21 +309,8 @@ pub fn quantile_huber_loss(
     quantile_loss?.mean_all()
 }
 
-/// Distributional type enum (C51 vs QR-DQN)
-#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
-pub enum DistributionalType {
-    /// Categorical DQN (C51) - fixed bins
-    C51,
-    /// Quantile Regression DQN - adaptive quantiles
-    QRDQN,
-}
-
-impl Default for DistributionalType {
-    fn default() -> Self {
-        // Default to QR-DQN for better trading risk modeling
-        Self::QRDQN
-    }
-}
+// Re-export from distributional module (single source of truth)
+pub use super::distributional::DistributionalType;
 
 #[cfg(test)]
 mod tests {
@@ -337,12 +324,6 @@ mod tests {
         assert_eq!(config.kappa, 1.0);
     }
 
-    #[test]
-    fn test_distributional_type_default() {
-        let dist_type = DistributionalType::default();
-        assert_eq!(dist_type, DistributionalType::QRDQN);
-    }
-
     #[test]
     fn test_sample_uniform_quantiles() -> Result<(), MLError> {
         let config = QuantileConfig::default();
diff --git a/ml/src/dqn/rainbow_agent_impl.rs b/ml/src/dqn/rainbow_agent_impl.rs
index 5418947e4..9a3179eda 100644
--- a/ml/src/dqn/rainbow_agent_impl.rs
+++ b/ml/src/dqn/rainbow_agent_impl.rs
@@ -78,7 +78,7 @@ impl RainbowAgent {
             lr: config.learning_rate,
             beta_1: 0.9,
             beta_2: 0.999,
-            eps: 1e-8,
+            eps: 1.5e-4, // Rainbow paper (Hessel et al. 2018) standard for distributional stability
             weight_decay: None,
             amsgrad: false,
         };