fix: re-enable IQN + remove diagnostics for H100 baseline

- iqn_lambda restored to 0.25 (was 0.0 for NaN isolation)
- STEP_DIAG per-step logging removed (was temporary)
- FOXHUNT_GRAD_DIAG env var removed from Argo template

IQN NaN root cause fixed in 8cbabcef5 (f32-as-bf16 type mismatch).
Ready for H100 baseline deployment.

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

config/training/dqn-production.toml

@@ -86,7 +86,7 @@ c51_alpha_max = 0.5
 her_ratio = 0.2
 cql_alpha = 1.0
 curiosity_weight = 0.1
-iqn_lambda = 0.0  # DIAG: disabled to isolate NaN source
+iqn_lambda = 0.25
 spectral_norm_sigma_max = 1.5
 spectral_decoupling_lambda = 0.01
 gradient_clip_norm = 1.0

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

@@ -853,20 +853,6 @@ 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).

docs/superpowers/plans/2026-04-09-h100-nan-fix.md

@@ -0,0 +1,194 @@
+# H100 NaN Fix Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Fix the three remaining H100 training instabilities — IQN type mismatch NaN, IQN gradient amplification, and C51 premature convergence.
+
+**Architecture:** Three independent fixes applied to the CUDA training pipeline. Fix 1 (type mismatch) is the NaN root cause. Fix 2 (gradient amplification) prevents IQN from destabilizing Adam. Fix 3 (alpha cap) keeps MSE gradient alive after C51 converges. Each fix is independently testable.
+
+**Tech Stack:** Rust + CUDA kernels (`.cu`), cuBLAS backward pass, fused DQN training pipeline.
+
+---
+
+### Task 1: Fix f32-as-bf16 type mismatch in apply_iqn_trunk_gradient
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:1082-1135`
+
+The IQN's `d_h_s2_buf` is allocated as `CudaSlice<f32>` (gpu_iqn_head.rs:235) and the backward kernel writes f32 via `atomicAdd`. But `apply_iqn_trunk_gradient` at line 1114-1133 calls `bf16_to_f32_kernel` on this f32 pointer, reinterpreting f32 bits as bf16 — producing garbage/NaN.
+
+- [ ] **Step 1: Remove the bf16→f32 cast and use direct f32 DtoD copy**
+
+In `apply_iqn_trunk_gradient()`, replace the bf16_to_f32_kernel call (lines 1114-1134) with a direct f32→f32 DtoD copy. The IQN d_h_s2 is already f32 — no cast needed.
+
+Replace:
+```rust
+        // ── 2. Cast IQN d_h_s2 (bf16) → bw_d_h_s2 (f32) ─────────────────
+        // IQN head produces bf16 gradient. Cast to f32 for the backward scratch.
+        {
+            let src = iqn_d_h_s2_ptr;
+            let dst = self.ptrs.bw_d_h_s2;
+            let n_elems = (b * sh2) as i32;
+            let blocks = ((b * sh2 + 255) / 256) as u32;
+            unsafe {
+                self.stream
+                    .launch_builder(&self.bf16_to_f32_kernel)
+                    .arg(&src)
+                    .arg(&dst)
+                    .arg(&n_elems)
+                    .launch(LaunchConfig {
+                        grid_dim: (blocks, 1, 1),
+                        block_dim: (256, 1, 1),
+                        shared_mem_bytes: 0,
+                    })
+                    .map_err(|e| MLError::ModelError(format!("IQN d_h_s2 bf16→f32: {e}")))?;
+            }
+        }
+```
+
+With:
+```rust
+        // ── 2. DtoD copy IQN d_h_s2 (f32) → bw_d_h_s2 (f32) ────────────
+        // IQN d_h_s2_buf is f32 (atomicAdd in backward kernel). No cast needed.
+        {
+            let n_bytes = b * sh2 * std::mem::size_of::<f32>();
+            dtod_copy(self.ptrs.bw_d_h_s2, iqn_d_h_s2_ptr, n_bytes, &self.stream, 0, "iqn_d_h_s2_copy")
+                .map_err(|e| MLError::ModelError(format!("IQN d_h_s2 f32 DtoD: {e}")))?;
+        }
+```
+
+- [ ] **Step 2: Build and verify**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml`
+Expected: compiles clean.
+
+- [ ] **Step 3: Run smoke test**
+
+Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::regression --ignored`
+Expected: 4 passed, 0 failed.
+
+- [ ] **Step 4: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+git commit -m "fix(critical): IQN d_h_s2 is f32, not bf16 — remove bogus cast"
+```
+
+---
+
+### Task 2: Fix IQN gradient clipping amplification
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:728-731`
+
+With mean-reduced IQN gradients (~1e-6 magnitude), the IQN Adam's `clip_scale = max_grad_norm / grad_norm = 1.0 / 2.7e-4 ≈ 3700×`. This AMPLIFIES gradients instead of clipping them. The fix: clamp clip_scale to max 1.0 so it only ever reduces, never amplifies.
+
+- [ ] **Step 1: Clamp clip_scale to ≤ 1.0**
+
+In `iqn_adam_kernel`, change line 729-731 from:
+```cuda
+    float clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0.0f)
+        ? max_grad_norm / grad_norm_val : 1.0f;
+```
+
+To:
+```cuda
+    /* Clamp clip_scale to [0, 1]: only reduce, never amplify.
+     * With mean-reduced gradients (inv_batch), grad_norm ≈ 1e-4 which is
+     * below max_grad_norm — clip_scale stays 1.0 (no clipping). */
+    float clip_scale = (grad_norm_val > max_grad_norm)
+        ? max_grad_norm / grad_norm_val : 1.0f;
+```
+
+Note: the existing code already does this correctly (clip_scale is 1.0 when norm ≤ max). The issue is that `max_grad_norm=1.0` is too large relative to the mean-reduced gradient (~2.7e-4). The clipping never fires, and Adam amplifies the tiny gradient via the m_hat/sqrt(v_hat) ratio. The real fix is to not clip at all for IQN (set max_grad_norm=Inf) and let Adam's adaptive LR handle scaling. OR reduce IQN's LR.
+
+Actually, the simplest fix: IQN should use the same LR as the main DQN. Check that `config.lr` is set correctly from the TOML learning_rate.
+
+- [ ] **Step 2: Verify IQN LR matches production config**
+
+In `crates/ml/src/trainers/dqn/fused_training.rs` line 437:
+```rust
+lr: hyperparams.learning_rate as f32,
+```
+This should be `1e-5` from the production TOML. Confirm by searching for the value in the constructor.
+
+- [ ] **Step 3: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
+git commit -m "fix: document IQN gradient clipping behavior with mean-reduced gradients"
+```
+
+---
+
+### Task 3: Set c51_alpha_max=0.5 for production H100
+
+**Files:**
+- Modify: `config/training/dqn-production.toml`
+- Modify: `crates/ml/src/trainers/dqn/config.rs`
+
+The IQN-disabled H100 run showed fold 2 collapsing to grad_norm=0 (C51 converged). With IQN fixed, this may not be the binding constraint, but c51_alpha_max=0.5 is a safety net.
+
+- [ ] **Step 1: Set default c51_alpha_max to 0.5**
+
+In config.rs default constructor, change `c51_alpha_max: 1.0` to `c51_alpha_max: 0.5`.
+In dqn-production.toml, change `c51_alpha_max = 1.0` to `c51_alpha_max = 0.5`.
+
+- [ ] **Step 2: Run full smoke test suite**
+
+Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability --ignored`
+Expected: 6/6 pass.
+
+- [ ] **Step 3: Commit**
+
+```bash
+git add config/training/dqn-production.toml crates/ml/src/trainers/dqn/config.rs
+git commit -m "fix: set c51_alpha_max=0.5 default — prevent C51 gradient convergence"
+```
+
+---
+
+### Task 4: Re-enable IQN, remove diagnostics, deploy baseline
+
+**Files:**
+- Modify: `config/training/dqn-production.toml` — restore `iqn_lambda = 0.25`
+- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — remove STEP_DIAG logging
+- Modify: `infra/k8s/argo/compile-and-train-template.yaml` — remove FOXHUNT_GRAD_DIAG env
+
+- [ ] **Step 1: Restore iqn_lambda, remove diagnostics**
+
+dqn-production.toml: `iqn_lambda = 0.25`
+fused_training.rs: remove the `STEP_DIAG` warn! block (~14 lines)
+compile-and-train-template.yaml: remove `FOXHUNT_GRAD_DIAG` env var
+
+- [ ] **Step 2: Full test suite**
+
+Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability --ignored`
+Expected: 6/6 pass.
+
+- [ ] **Step 3: Commit and push**
+
+```bash
+git add -A
+git commit -m "fix: re-enable IQN + remove diagnostics for H100 baseline"
+git push
+```
+
+- [ ] **Step 4: Deploy H100 baseline (20 epochs)**
+
+```bash
+./scripts/argo-train.sh dqn --baseline --epochs 20 --watch
+```
+
+Expected: grad_norm stable (0.5-1.0) through all 20 epochs, both folds. No NaN.
+
+---
+
+### Task 5: Update memory and documentation
+
+**Files:**
+- Modify: `/home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/followup_gradient_collapse_h100.md`
+
+- [ ] **Step 1: Update memory with final root causes**
+
+Add the IQN f32-as-bf16 type mismatch as root cause #4. Update the verification section with the successful H100 baseline results.

infra/k8s/argo/compile-and-train-template.yaml

@@ -526,8 +526,6 @@ spec:
             value: ":4096:8"
           - name: FOXHUNT_FEATURE_CACHE_DIR
             value: "{{workflow.parameters.feature-cache-dir}}"
-          - name: FOXHUNT_GRAD_DIAG
-            value: "1"
         args:
           - |
             set -e