# 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)
