fix: fused training supports RegimeConditional + GpuTensor::cat dim>0

Three fixes:
1. FusedTrainingCtx::new now accepts RegimeConditionalDQN by using
   primary_head() — the old code rejected it, causing silent fallback
   to non-fused train_step() which has action-space mismatch with
   branching DQN (5-action Q-table vs 45-action factored indices →
   CUDA_ERROR_ILLEGAL_ADDRESS buffer overrun)

2. ensure_fused_ctx() returns Result instead of () — fused init
   failure is now a hard error (no silent CPU fallback)

3. GpuTensor::cat now supports dim>0 concatenation (was unimplemented,
   caused "dim=1 > 0 not yet implemented" error in validation path)

Root cause chain:
  RegimeConditional rejected by fused init
  → silent fallback to DQN::train_step()
  → compute_loss_internal() uses num_actions=5 but actions are 0-44
  → gather with out-of-bounds offsets → CUDA_ERROR_ILLEGAL_ADDRESS
  → async error poisons CUDA context
  → next stream.synchronize() deadlocks forever

Remaining: curiosity kernel crash also poisons the stream. The
curiosity training runs inside collect_gpu_experiences() and its
async error blocks the PER insert. This needs separate investigation
of curiosity_training_kernel.cu.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/ml-core/src/cuda_autograd/gpu_tensor.rs

@@ -248,10 +248,27 @@ impl GpuTensor {
             }
             return Self::from_host(&flat, new_shape, stream);
         }
-        // General case: not yet optimized
-        Err(MLError::ModelError(format!(
-            "cat: dim={dim} > 0 not yet implemented for GpuTensor"
-        )))
+        // General dim>0: interleave rows along the concat dimension
+        let ref_shape = tensors.first().map_or(&[] as &[usize], |t| t.shape());
+        // outer = product of dims before `dim`, inner = product of dims after `dim`
+        let outer: usize = ref_shape[..dim].iter().product();
+        let inner: usize = if dim + 1 < ref_shape.len() { ref_shape[dim + 1..].iter().product() } else { 1 };
+        let mut flat = Vec::with_capacity(outer * total_dim * inner);
+        for o in 0..outer {
+            for (ti, t) in tensors.iter().enumerate() {
+                let d = t.shape().get(dim).copied().unwrap_or(0);
+                let src = &all_data[ti];
+                let src_stride = d * inner;
+                let start = o * src_stride;
+                let end = start + src_stride;
+                if end <= src.len() {
+                    flat.extend_from_slice(&src[start..end]);
+                }
+            }
+        }
+        let mut new_shape = ref_shape.to_vec();
+        new_shape[dim] = total_dim;
+        Self::from_host(&flat, new_shape, stream)
     }
 
     /// Stack tensors along a new dimension 0.

crates/ml/src/trainers/dqn/fused_training.rs

@@ -98,11 +98,7 @@ impl FusedTrainingCtx {
     ) -> Result<Self> {
         let dqn = match agent {
             DQNAgentType::Standard(d) => d,
-            DQNAgentType::RegimeConditional(_) => {
-                return Err(anyhow::anyhow!(
-                    "Fused CUDA training requires Standard DQN agent (not RegimeConditional)"
-                ));
-            }
+            DQNAgentType::RegimeConditional(regime) => regime.primary_head(),
         };
 
         let branching_net = dqn.branching_q_network.as_ref().ok_or_else(|| {

crates/ml/src/trainers/dqn/trainer/train_step.rs

@@ -634,18 +634,17 @@ impl DQNTrainer {
 
     /// Lazy-init fused CUDA training context (Standard DQN only).
     /// Recreate if batch_size changed (OOM recovery).
-    pub(crate) async fn ensure_fused_ctx(&mut self) {
+    pub(crate) async fn ensure_fused_ctx(&mut self) -> anyhow::Result<()> {
         let needs_init = match &self.fused_ctx {
             None => self.device.is_cuda(),
             Some(ctx) => ctx.batch_size() != self.current_batch_size,
         };
         if !needs_init {
-            return;
+            return Ok(());
         }
-        // Fused context requires the unified forked CudaStream
         let stream = match self.cuda_stream {
             Some(ref s) => std::sync::Arc::clone(s),
-            None => return,
+            None => return Ok(()),
         };
         if self.fused_ctx.is_some() {
             info!("Fused CUDA context: batch_size changed, recreating");
@@ -660,8 +659,12 @@ impl DQNTrainer {
                 self.fused_ctx = Some(ctx);
             }
             Err(e) => {
-                warn!("Fused CUDA training init failed, using Candle path: {e}");
+                return Err(anyhow::anyhow!(
+                    "Fused CUDA training init FAILED (no fallback — non-fused path \
+                     has action-space mismatch with branching DQN): {e}"
+                ));
             }
         }
+        Ok(())
     }
 }

crates/ml/src/trainers/dqn/trainer/training_loop.rs

@@ -885,7 +885,7 @@ impl DQNTrainer {
         // Lazy-init fused CUDA training context (Standard DQN only).
         // Recreate if batch_size changed (OOM recovery).
         if num_training_steps > 0 {
-            self.ensure_fused_ctx().await;
+            self.ensure_fused_ctx().await?;
         }
 
         // Lazy-init + reset guard accumulators for this epoch