From 40febb006164945d0f0d62e34368709554e941a0 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sat, 28 Feb 2026 17:01:01 +0100
Subject: [PATCH] fix(cuda): enable GPU experience collection for
 regime-conditional DQN

The GPU experience collector only supported DQNAgentType::Standard, but
enable_regime_qnetwork defaults to true, creating RegimeConditional
agents. This silently fell back to CPU with a debug! message invisible
at INFO log level.

- Add primary_head() accessor to RegimeConditionalDQN (returns trending head)
- Extract weights from primary head for both collector init and weight sync
- Upgrade debug! to warn! for missing GPU collection prerequisites
- Add warn! to PPO for silent GPU skip when raw market data not uploaded

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml/src/dqn/regime_conditional.rs |  8 ++++++++
 crates/ml/src/trainers/dqn/trainer.rs   | 21 ++++++++++++++++-----
 crates/ml/src/trainers/ppo.rs           | 12 ++++++++++++
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/crates/ml/src/dqn/regime_conditional.rs b/crates/ml/src/dqn/regime_conditional.rs
index 4c5ad73c8..52922ee8b 100644
--- a/crates/ml/src/dqn/regime_conditional.rs
+++ b/crates/ml/src/dqn/regime_conditional.rs
@@ -235,6 +235,14 @@ impl RegimeConditionalDQN {
         })
     }
 
+    /// Get a reference to the primary (trending) DQN head.
+    ///
+    /// Used by GPU experience collector for weight extraction — the trending head
+    /// is the most common regime and provides representative network weights.
+    pub fn primary_head(&self) -> &DQN {
+        &self.trending_head
+    }
+
     /// Forward pass through regime-specific Q-network head
     ///
     /// # Arguments
diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs
index d3c36a071..00f6aa502 100644
--- a/crates/ml/src/trainers/dqn/trainer.rs
+++ b/crates/ml/src/trainers/dqn/trainer.rs
@@ -1412,9 +1412,16 @@ impl DQNTrainer {
                     use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector;
                     let stream = cuda_dev.cuda_stream();
 
-                    // Read-lock agent to check for dueling networks
+                    // Read-lock agent to extract dueling network weights for GPU collector.
+                    // Works with both Standard and RegimeConditional agents — for regime-conditional,
+                    // we use the primary (trending) head's weights since the GPU kernel runs a single
+                    // Q-network for fast experience collection.
                     let agent = self.agent.read().await;
-                    let init_result = if let DQNAgentType::Standard(ref dqn) = &*agent {
+                    let dqn_ref: Option<&crate::dqn::DQN> = match &*agent {
+                        DQNAgentType::Standard(ref dqn) => Some(dqn),
+                        DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()),
+                    };
+                    let init_result = if let Some(dqn) = dqn_ref {
                         match (
                             dqn.dueling_q_network.as_ref(),
                             dqn.dueling_target_network.as_ref(),
@@ -1432,12 +1439,12 @@ impl DQNTrainer {
                                 ))
                             }
                             _ => {
-                                debug!("GPU experience collector skipped: requires dueling networks + curiosity module");
+                                warn!("GPU experience collector skipped: requires dueling networks + curiosity module (have dueling_q={}, dueling_target={}, curiosity={})",
+                                    dqn.dueling_q_network.is_some(), dqn.dueling_target_network.is_some(), self.curiosity_module.is_some());
                                 None
                             }
                         }
                     } else {
-                        debug!("GPU experience collector skipped: regime-conditional DQN not supported");
                         None
                     };
                     drop(agent);
@@ -2254,7 +2261,11 @@ impl DQNTrainer {
             #[cfg(feature = "cuda")]
             if let Some(ref mut collector) = self.gpu_experience_collector {
                 let agent = self.agent.read().await;
-                if let DQNAgentType::Standard(ref dqn) = &*agent {
+                let dqn_ref: Option<&crate::dqn::DQN> = match &*agent {
+                    DQNAgentType::Standard(ref dqn) => Some(dqn),
+                    DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()),
+                };
+                if let Some(dqn) = dqn_ref {
                     if let Some(ref online) = dqn.dueling_q_network {
                         if let Err(e) = collector.sync_online_weights(online.vars()) {
                             warn!("GPU online weight sync failed: {}", e);
diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs
index 7d608caf6..0692f5b8e 100644
--- a/crates/ml/src/trainers/ppo.rs
+++ b/crates/ml/src/trainers/ppo.rs
@@ -502,6 +502,18 @@ impl PpoTrainer {
                     }
                 }
             } else {
+                // Log why GPU collection is not available so it doesn't silently fall back to CPU
+                #[cfg(feature = "cuda")]
+                {
+                    if gpu_ppo_collector.is_none() {
+                        warn!("PPO GPU experience collection skipped: collector not initialized");
+                    } else if self.features_raw_cuda.is_none() || self.targets_raw_cuda.is_none() {
+                        warn!("PPO GPU experience collection skipped: raw market data not uploaded (call set_raw_market_data() before train())");
+                    } else {
+                        // All conditions met but pattern match still failed — shouldn't happen
+                        warn!("PPO GPU experience collection skipped: unexpected state");
+                    }
+                }
                 None
             };