docs: spec + plan for C51 loss accumulator fix (separate MSE/C51 buffers)

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

docs/superpowers/plans/2026-04-04-c51-loss-accumulator-fix.md

@@ -0,0 +1,167 @@
+# C51 Loss Accumulator Fix
+
+> **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 NaN at step 22 by separating MSE and C51 loss accumulators. Blend at readback time so C51 NaN during warmup (alpha=0) doesn't contaminate the loss scalar.
+
+**Architecture:** Add `mse_loss_buf: CudaSlice<f32>` alongside existing `total_loss_buf` (now C51-only). MSE kernel writes to `mse_loss_buf`. At readback, blend: `loss = (1-α)*mse + α*c51`. CUDA Graph compatible (no conditional kernel launches).
+
+**Tech Stack:** CUDA C, Rust (cudarc), f32 atomicAdd
+
+---
+
+## File Map
+
+| File | Change |
+|------|--------|
+| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Add `mse_loss_buf`, update MSE kernel launch, blend at readback |
+| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | Accept separate output buffer (if currently hardcoded to total_loss_buf) |
+
+---
+
+### Task 1: Add mse_loss_buf and separate the accumulators
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
+- Possibly modify: `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu`
+
+This is a single focused task — all changes are in the loss accumulation path.
+
+- [ ] **Step 1: Read the current code**
+
+Read these sections of `gpu_dqn_trainer.rs`:
+1. `total_loss_buf` field declaration (line 558) — understand the type and allocation
+2. `total_loss_buf` allocation (line 2231) — how it's allocated
+3. `GpuDqnPtrs` struct (line 425) — where `total_loss_buf: u64` raw pointer is stored
+4. `launch_mse_loss()` (line 4689) — find where it passes `total_loss_buf` to the MSE kernel
+5. `launch_c51_loss()` (line 4523) — find where it passes `total_loss_buf` to the C51 kernel
+6. The readback path (around line 3030-3040) — where `total_loss_buf` is read back to host
+7. The memset-zero before each step — where `total_loss_buf` is zeroed
+
+Also read `mse_loss_kernel.cu` to understand if the loss buffer is passed as a kernel argument or if it uses the pointer from `GpuDqnPtrs`.
+
+- [ ] **Step 2: Add `mse_loss_buf` field**
+
+Add next to `total_loss_buf` (line 558):
+```rust
+    pub(crate) mse_loss_buf: CudaSlice<f32>,               // [1] MSE loss accumulator (separate from C51)
+```
+
+If `GpuDqnPtrs` is used to pass pointers to kernels, also add:
+```rust
+    mse_loss_buf: u64,
+```
+
+- [ ] **Step 3: Allocate `mse_loss_buf`**
+
+At line 2231, after `total_loss_buf` allocation:
+```rust
+    let mse_loss_buf = stream.alloc_zeros::<f32>(1)
+        .map_err(|e| MLError::ModelError(format!("alloc mse_loss_f32: {e}")))?;
+```
+
+Wire it into the struct construction.
+
+- [ ] **Step 4: Update MSE loss kernel launch to use `mse_loss_buf`**
+
+In `launch_mse_loss()`, find where `total_loss_buf` (or its raw pointer) is passed as an argument to the MSE kernel. Change it to `mse_loss_buf`.
+
+If the kernel reads the buffer pointer from `GpuDqnPtrs` (a struct uploaded to GPU), you need to set `ptrs.mse_loss_buf = mse_loss_buf.raw_ptr()` instead.
+
+If the kernel takes it as a direct launch argument, change the `.arg()` call.
+
+- [ ] **Step 5: Zero both buffers before each step**
+
+Find where `total_loss_buf` is zeroed (memset) before each training step. Add the same for `mse_loss_buf`:
+```rust
+self.stream.memset_zeros(&mut self.mse_loss_buf)
+    .map_err(|e| MLError::ModelError(format!("zero mse_loss: {e}")))?;
+```
+
+- [ ] **Step 6: Blend at readback**
+
+Find the loss readback (around line 3030-3040). Currently it reads `total_loss_buf[0]`. Change to read BOTH buffers and blend:
+
+```rust
+// Read both loss accumulators
+let mut c51_loss = [0.0_f32; 1];
+let mut mse_loss = [0.0_f32; 1];
+self.stream.memcpy_dtoh(&self.total_loss_buf, &mut c51_loss)
+    .map_err(|e| MLError::ModelError(format!("readback c51_loss: {e}")))?;
+self.stream.memcpy_dtoh(&self.mse_loss_buf, &mut mse_loss)
+    .map_err(|e| MLError::ModelError(format!("readback mse_loss: {e}")))?;
+
+// Blend: (1-α)*MSE + α*C51 — when α=0 (warmup), only MSE contributes
+let alpha = self.c51_alpha;
+let blended_loss = (1.0 - alpha) * mse_loss[0] + alpha * c51_loss[0];
+```
+
+Use `blended_loss` instead of the old `total_loss` value.
+
+**IMPORTANT:** If `alpha * c51_loss[0]` is `0.0 * NaN = NaN`, you need to guard:
+```rust
+let blended_loss = if alpha < 1e-6 {
+    mse_loss[0]
+} else if alpha > 1.0 - 1e-6 {
+    c51_loss[0]
+} else {
+    (1.0 - alpha) * mse_loss[0] + alpha * c51_loss[0]
+};
+```
+
+This prevents `0.0 * NaN = NaN` contamination.
+
+- [ ] **Step 7: Verify compilation**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml
+```
+
+- [ ] **Step 8: Quick NaN test**
+
+```bash
+SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --release --test dqn_training_pipeline_test -- test_dqn_trains_on_es_fut --nocapture --exact 2>&1 | tail -10
+```
+
+If step 22 NaN is gone, the fix works. Expected: test runs past step 22 (may still fail for other reasons, but no NaN).
+
+- [ ] **Step 9: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
+git commit -m "fix: separate MSE/C51 loss accumulators (prevents NaN during warmup)"
+```
+
+---
+
+### Task 2: Run integration tests
+
+- [ ] **Step 1: Pipeline test (fast NaN check)**
+
+```bash
+SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --release --test dqn_training_pipeline_test -- test_dqn_trains_on_es_fut --nocapture --exact
+```
+
+Expected: no NaN, completes or fails for non-NaN reason.
+
+- [ ] **Step 2: Early stopping tests**
+
+```bash
+SQLX_OFFLINE=true cargo test -p ml --release --test dqn_early_stopping_termination_test -- --nocapture
+```
+
+Expected: all 4 pass, <2 min.
+
+- [ ] **Step 3: Full test suite**
+
+```bash
+SQLX_OFFLINE=true cargo test -p ml-core --lib -- gpu::profile
+SQLX_OFFLINE=true cargo test -p ml --lib -- fxcache feature_cache training_profile
+SQLX_OFFLINE=true cargo test -p ml --release --test dqn_action_collapse_fix_test
+```
+
+- [ ] **Step 4: Push**
+
+```bash
+git push origin main
+```

docs/superpowers/specs/2026-04-04-c51-loss-accumulator-fix-design.md

@@ -0,0 +1,58 @@
+# C51 Loss Accumulator Fix
+
+**Goal:** Fix NaN at step 22 caused by C51 loss kernel contaminating the shared loss accumulator during MSE warmup. Separate loss accumulators for MSE and C51, blend at readback time.
+
+**Root cause:** Both MSE and C51 loss kernels always run (CUDA Graph compatibility) and both write to the same `total_loss_buf` via `atomicAdd`. During warmup (`c51_alpha=0`), the C51 loss on random/extreme logits produces NaN which contaminates the shared buffer. The gradient blending is correct (C51 grad × 0 = 0), but the loss scalar is irreversibly NaN.
+
+---
+
+## The Fix
+
+### Separate loss accumulators
+
+Add a second loss buffer `mse_loss_buf: CudaSlice<f32>` to `GpuDqnTrainer`. The existing `total_loss_buf` becomes the C51 accumulator.
+
+| Buffer | Written by | Purpose |
+|--------|-----------|---------|
+| `total_loss_buf` | C51 loss kernel | C51 cross-entropy loss |
+| `mse_loss_buf` | MSE loss kernel | MSE warmup loss |
+
+### Blended readback
+
+At the end of `submit_forward_ops()` (or when the training guard reads the loss), compute:
+
+```
+blended_loss = (1 - alpha) * mse_loss + alpha * c51_loss
+```
+
+When `alpha=0` (warmup): `blended_loss = mse_loss` — C51 NaN is ignored.
+When `alpha=1` (converged): `blended_loss = c51_loss` — MSE is ignored.
+
+This keeps both kernels running unconditionally (CUDA Graph compatible) but prevents NaN contamination.
+
+### Changes
+
+1. **Add `mse_loss_buf`** to `GpuDqnTrainer` struct (allocate alongside `total_loss_buf`)
+2. **Update MSE loss kernel launch** to write to `mse_loss_buf` instead of `total_loss_buf`
+3. **Update loss readback** (in training guard or `process_epoch_boundary`) to blend the two buffers using `alpha`
+4. **Zero both buffers** before each training step
+
+### CUDA Graph compatibility
+
+Both kernels keep running on every step. The conditional logic is ONLY at the scalar readback (host-side blend after GPU→CPU copy). No kernel launches change. Graph stays fixed.
+
+---
+
+## Files touched
+
+| File | Change |
+|------|--------|
+| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Add `mse_loss_buf`, update MSE kernel launch arg, blend at readback |
+| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | May need to accept separate output buffer (check if it already takes a loss_buf arg) |
+
+## What this does NOT change
+
+- C51 loss kernel (still runs, still writes to `total_loss_buf`)
+- Gradient blending (already correct with scale kernel)
+- CUDA Graph structure (no conditional launches)
+- The SAXPY/scale kernel fix (separate bug, already fixed)