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