# Stable Distributional Loss Architecture — 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:** Eliminate C51 cross-entropy loss explosion (4.2→242K→796K) via three-layer defense: per-sample loss clamping, label smoothing, and promoting IQN to primary loss with fixed τ (QR-DQN style, CUDA Graph compatible).

**Architecture:** Layer 1 breaks the PER feedback loop (loss → priority → IS weight → amplified loss). Layer 2 prevents target distribution collapse that triggers log(near-zero). Layer 3 makes the Huber quantile loss (bounded by construction) the primary gradient source, with C51 demoted to auxiliary. Fixed τ values replace random sampling, enabling full CUDA Graph capture.

**Tech Stack:** CUDA C (kernels), Rust (orchestration), cudarc (driver API)

---

## File Structure

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` | Modify | Layer 1: loss clamp. Layer 2: label smoothing after Bellman projection |
| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Modify | Layer 3: replace random τ with fixed midpoints, expose per_sample_loss for PER |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Modify | Layer 3: swap loss weighting (IQN primary, C51 auxiliary), use IQN loss for PER |
| `crates/ml/src/trainers/dqn/config.rs` | Modify | Add `iqn_is_primary: bool` config flag |
| `config/training/dqn-production.toml` | Modify | Set IQN as primary |

---

### Task 1: Layer 1 — Per-Sample C51 Loss Clamping

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu:592-601`

- [ ] **Step 1: Add MAX_PER_SAMPLE_CE constant**

At the top of the file (near other `#define`s), add:
```cuda
#ifndef MAX_PER_SAMPLE_CE
#define MAX_PER_SAMPLE_CE 50.0f  /* 10x converged loss — breaks PER feedback loop */
#endif
```

- [ ] **Step 2: Clamp avg_ce before IS weight multiplication**

Replace lines 596-601:
```cuda
    if (tid == 0) {
        float weighted_loss = avg_ce * is_weight;
        per_sample_loss[sample_id] = weighted_loss;
        td_errors[sample_id]       = avg_ce;
        atomicAdd(total_loss, weighted_loss / (float)batch_size);
    }
```

With:
```cuda
    if (tid == 0) {
        /* Layer 1: clamp per-sample loss to prevent PER feedback loop.
         * High loss → high PER priority → high IS weight → amplified loss → repeat.
         * Clamping at 50.0 (10x converged loss) breaks the loop. */
        float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
        float weighted_loss = clamped_ce * is_weight;
        per_sample_loss[sample_id] = weighted_loss;
        td_errors[sample_id]       = clamped_ce;  /* PER sees clamped too */
        atomicAdd(total_loss, weighted_loss / (float)batch_size);
    }
```

- [ ] **Step 3: Compile check**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

Expected: clean (CUDA compiled at runtime via NVRTC).

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
git commit -m "fix: Layer 1 — per-sample C51 loss clamp breaks PER feedback loop"
```

---

### Task 2: Layer 2 — Label Smoothing on Bellman Target

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu:568-575`

- [ ] **Step 1: Add label smoothing constant**

Near MAX_PER_SAMPLE_CE:
```cuda
#ifndef LABEL_SMOOTHING_EPS
#define LABEL_SMOOTHING_EPS 0.01f  /* 1% uniform mix — prevents target collapse */
#endif
```

- [ ] **Step 2: Apply label smoothing after Bellman projection**

After the `block_bellman_project` call (line 568-569) and before saving projected (line 573), add:

```cuda
        block_bellman_project(shmem_proj, reward, done, gamma,
                              shmem_lp, shmem_support, shmem_reduce, tid);
        __syncthreads();

        /* Layer 2: Label smoothing — mix 1% uniform into projected target.
         * Prevents any atom from having zero probability, which would cause
         * log(0) = -inf in the cross-entropy. Standard technique from
         * Szegedy et al. (2016), adapted for C51. */
        {
            const float uniform = LABEL_SMOOTHING_EPS / (float)NUM_ATOMS;
            const float smooth_scale = 1.0f - LABEL_SMOOTHING_EPS;
            for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
                shmem_lp[j] = smooth_scale * shmem_lp[j] + uniform;
            __syncthreads();
        }
```

Note: `shmem_lp` holds the projected target after `block_bellman_project`. The smoothing happens IN-PLACE before it's saved to `save_projected` and used in the cross-entropy.

- [ ] **Step 3: Compile check**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 4: Local test — verify loss doesn't explode**

```bash
SQLX_OFFLINE=true timeout 180 ./target/release/examples/hyperopt_baseline_rl \
  --model dqn --phase fast --trials 1 --n-initial 1 --epochs 3 \
  --data-dir test_data/futures-baseline --symbol ES.FUT \
  --output /tmp/layer12_test.json --seed 42
```

Expected: train_loss stays < 100 (was 242K at epoch 10 without these fixes).

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
git commit -m "fix: Layer 2 — label smoothing on C51 Bellman target prevents log(0)"
```

---

### Task 3: Layer 3a — Fixed τ Values (QR-DQN Style, CUDA Graph Compatible)

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs:320-350`

- [ ] **Step 1: Pre-compute fixed τ midpoints in the constructor**

In the `new()` method, after allocating `online_taus` and `target_taus`, initialize them with fixed midpoint values instead of leaving them zero:

```rust
// Fixed τ midpoints: τ_i = (2i - 1) / (2N) for i = 1..N
// This makes IQN equivalent to QR-DQN — deterministic, CUDA Graph compatible.
let fixed_taus: Vec<f32> = (0..n)
    .map(|i| (2.0 * (i as f32) + 1.0) / (2.0 * n as f32))
    .collect();
// Tile across batch: each sample gets the same fixed τ values
let mut tiled_taus = Vec::with_capacity(b * n);
for _ in 0..b {
    tiled_taus.extend_from_slice(&fixed_taus);
}
stream.memcpy_htod(&tiled_taus, &mut online_taus)
    .map_err(|e| MLError::ModelError(format!("IQN fixed tau upload: {e}")))?;
stream.memcpy_htod(&tiled_taus, &mut target_taus)
    .map_err(|e| MLError::ModelError(format!("IQN fixed tau upload: {e}")))?;
```

- [ ] **Step 2: Remove random τ sampling from train_iqn_step_gpu**

In `train_iqn_step_gpu()` (lines 323-350), remove the two `sample_taus_kernel` launches:

```rust
// REMOVED: random τ sampling — now using fixed midpoints (QR-DQN style)
// The online_taus and target_taus buffers contain pre-computed fixed values.
// No per-step random sampling → fully deterministic → CUDA Graph compatible.
```

Delete lines 323-350 (the Philox PRNG kernel launches for online_taus and target_taus).

- [ ] **Step 3: Also remove random sampling from compute_cvar_scales**

In `compute_cvar_scales()`, remove the `sample_taus_kernel` launch there too. The fixed τ values in `self.online_taus` are already correct.

- [ ] **Step 4: Compile check**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
git commit -m "feat: Layer 3a — fixed τ midpoints (QR-DQN style, CUDA Graph compatible)"
```

---

### Task 4: Layer 3b — Promote IQN to Primary Loss

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs:429-448`
- Modify: `crates/ml/src/trainers/dqn/config.rs`

- [ ] **Step 1: Add iqn_primary config flag**

In `config.rs`, add to DQNHyperparameters:
```rust
/// When true, IQN's Huber quantile loss is the primary PER signal
/// and C51 cross-entropy is auxiliary. Default: true (IQN is bounded).
pub iqn_primary_loss: bool,
```

Default to `true` in `conservative()`.

- [ ] **Step 2: Use IQN per_sample_loss for PER when primary**

In `fused_training.rs`, after the IQN training step (line 427), add:

```rust
// Layer 3b: When IQN is primary, use its per_sample_loss for PER priorities.
// IQN's Huber loss is bounded → no PER feedback explosion.
if hyperparams.iqn_primary_loss {
    if let Some(ref iqn) = self.gpu_iqn {
        // Overwrite the C51 td_errors with IQN per_sample_loss
        // so PER uses the bounded IQN loss for priority updates.
        let iqn_loss_buf = iqn.per_sample_loss();
        let dst = self.trainer.td_errors_buf_mut();
        let n_bytes = effective_gpu.batch_size * std::mem::size_of::<f32>();
        let (src_ptr, _s) = iqn_loss_buf.device_ptr(&self.stream);
        let (dst_ptr, _d) = dst.device_ptr(&self.stream);
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, n_bytes, self.stream.cu_stream()
            ).map_err(|e| anyhow::anyhow!("IQN→PER loss DtoD: {e}"))?;
        }
    }
}
```

- [ ] **Step 3: Update TOML configs**

`config/training/dqn-production.toml`:
```toml
iqn_primary_loss = true
```

- [ ] **Step 4: Compile check**

```bash
SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/config.rs config/training/dqn-production.toml
git commit -m "feat: Layer 3b — IQN primary loss for PER (bounded Huber replaces unbounded CE)"
```

---

### Task 5: Integration Test + Local Verification

- [ ] **Step 1: Rebuild release**

```bash
SQLX_OFFLINE=true cargo build --release -p ml --example hyperopt_baseline_rl 2>&1 | tail -3
```

- [ ] **Step 2: Local test — seed 42 (previously NaN'd)**

```bash
SQLX_OFFLINE=true timeout 180 ./target/release/examples/hyperopt_baseline_rl \
  --model dqn --phase fast --trials 1 --n-initial 1 --epochs 3 \
  --data-dir test_data/futures-baseline --symbol ES.FUT \
  --output /tmp/stable_loss_test.json --seed 42
```

Expected: NO NaN, train_loss < 50, grad_norm < 5000.

- [ ] **Step 3: Local test — seed 99 (different params)**

```bash
SQLX_OFFLINE=true timeout 180 ./target/release/examples/hyperopt_baseline_rl \
  --model dqn --phase fast --trials 1 --n-initial 1 --epochs 3 \
  --data-dir test_data/futures-baseline --symbol ES.FUT \
  --output /tmp/stable_loss_test2.json --seed 99
```

Expected: NO NaN, train_loss decreasing.

- [ ] **Step 4: Push and submit H100**

```bash
git push origin main
```

- [ ] **Step 5: H100 hyperopt — verify no loss explosion across 20 trials**

Submit workflow with `hyperopt-epochs=50, hyperopt-trials=20`. Monitor:
- Trial 1 epoch 10: train_loss should be < 50 (was 242K before)
- No TRIAL_SUMMARY with objective=1000000 from NaN
- Q-gap should survive > 0.5 through epoch 50
