From c26f0ae45119cbcfbf8a99de0eae49e697327481 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 9 Apr 2026 16:35:26 +0200
Subject: [PATCH] diag: add per-step STEP_DIAG logging (zero-cost, uses
 existing readback)

Logs grad_norm, loss, c51_alpha every 50 steps via existing async readback
scalars. No stream sync or extra GPU work. Works with mega-graph path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml/src/trainers/dqn/fused_training.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs
index b9cd644ce..e3274fa20 100644
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -853,6 +853,20 @@ impl FusedTrainingCtx {
             .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;
 
         self.steps_since_varmap_sync += 1;
+
+        // Per-step gradient diagnostic — zero cost (uses existing readback, no sync).
+        // Log every 50 steps to see gradient trajectory during collapse.
+        if self.steps_since_varmap_sync % 50 == 1 {
+            let alpha = self.trainer.c51_alpha();
+            tracing::warn!(
+                step = self.steps_since_varmap_sync,
+                grad_norm = fused_result.grad_norm,
+                loss = fused_result.total_loss,
+                c51_alpha = alpha,
+                batch_size = self.batch_size,
+                "STEP_DIAG: per-step gradient readback"
+            );
+        }
         self.last_combined_norm = fused_result.grad_norm;
 
         // Capture mega-graph at step 2 (all sub-trainers initialized).