# C51 Loss Convergence 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 C51 distributional loss divergence so the DQN model converges on real ES futures data and learns profitable trading behavior.

**Architecture:** Two-phase approach: (1) fix the entropy coefficient dominance that poisons the gradient signal, (2) add a Huber/MSE warmup phase before C51 kicks in, matching the pre-GPU Candle path that already worked.

**Tech Stack:** CUDA kernels (C51 loss, utility kernels), Rust GPU trainer (`gpu_dqn_trainer.rs`, `fused_training.rs`), DQN config.

---

## Root Cause Analysis

### The Problem

50-epoch local training on RTX 3050 with real ES futures data shows:
- **Loss**: 4.0 → 40.0 (diverges, never recovers)
- **Gradient norm (pre-clip)**: 883 → 5.4M (grows linearly with epoch count)
- **Q-values with ±2.0 v_range**: oscillate in [0.15, 1.57] — no longer saturate v_max ✓
- **Q-values with ±0.5 v_range**: saturate at v_max=0.5 from epoch 10 onward (fixed by widening)
- **Sharpe**: 0.00 for all 50 epochs — model never learns to trade profitably
- **Actions**: converge to ~77% Flat from epoch 8 onward

### Why C51 Diverges with Reward v4

The C51 cross-entropy loss computes `CE = -Σ projected_j * log(online_j)` per branch per sample.

With reward v4's per-bar returns (~0.001), the Bellman target is:
```
target = r + γ^n * Q_target = 0.001 + 0.876 * 0.3 = 0.264
```
This projects onto the C51 support atoms as a VERY narrow distribution (most mass on 1-2 atoms). The online distribution is similar. The cross-entropy between two nearly-identical distributions is their KL divergence, which is tiny (~0.001).

**But the entropy regularization gradient is 10x larger:**
```
entropy_grad = 0.01 * (1 + log_prob) ≈ 0.01 * (1 + (-5)) = -0.04 per atom
reward_signal = 0.001 per step
entropy / reward = 40x
```

The entropy gradient DOMINATES the reward signal. The model learns to maximize entropy (spread its distribution) rather than maximize returns.

### Why the Pre-GPU Huber Path Worked

The Candle CPU path at `dqn.rs:2484` uses Huber loss on expected Q-values:
```
td_error = Q(s,a) - (r + γ * max_a' Q_target(s',a'))
loss = Huber(td_error, δ=10.0)
```

This is a SCALAR loss on the expected value — no distributional complexity, no cross-entropy between 51-atom distributions. The gradient is simply `dL/dQ = sign(td_error)` (for |td_error| > δ) or `dL/dQ = td_error` (for small errors). This directly pushes Q(s,a) toward the Bellman target.

### The Fix Strategy

**Phase 1 (epochs 0..warmup_epochs):** Use Huber/MSE loss on expected Q-values.
- The C51 forward pass still runs (needed for action selection)
- But the LOSS is computed from `E[Q] = Σ p_j * z_j` vs Bellman target
- This gives a clean, strong gradient signal proportional to the reward

**Phase 2 (epochs warmup_epochs..end):** Switch to full C51 cross-entropy.
- By now Q-values have converged to the right neighborhood
- C51 refines the VALUE DISTRIBUTION (not just the mean)
- Entropy regularization is helpful here (prevents distribution collapse)

**Plus:** Reduce entropy coefficient from 0.01 to 0.001 regardless of phase.

---

## File Structure

| File | Action | Responsibility |
|------|--------|---------------|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | Add `mse_loss_kernel` + `mse_grad_kernel`, add `warmup_epochs` config field, switch loss mode |
| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | Modify | Add MSE/Huber loss + gradient CUDA kernels (inline) |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Modify | Pass `warmup_epochs` and current epoch to trainer, wire loss mode switch |
| `crates/ml/src/trainers/dqn/config.rs` | Modify | Add `c51_warmup_epochs` field to DQNHyperparameters |
| `crates/ml/src/hyperopt/adapters/dqn.rs` | Modify | Add `c51_warmup_epochs` to hyperopt search space |
| `crates/ml-dqn/src/dqn.rs` | Modify | Reduce `entropy_coefficient` defaults |
| `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` | Modify | Update convergence test to validate loss decreases |

---

## Task 1: Fix Entropy Coefficient (5 min)

**Files:**
- Modify: `crates/ml-dqn/src/dqn.rs:304` — DQNConfig default
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:192` — GpuDqnTrainConfig default

The entropy coefficient must be ≤ the reward magnitude. With per-bar returns ~0.001, entropy_coeff=0.01 is 10x too large. Set to 0.001.

- [ ] **Step 1:** Change `entropy_coefficient: 0.1` → `0.001` in `dqn.rs:304` (DQNConfig::default)
- [ ] **Step 2:** Change `entropy_coefficient: 0.01` → `0.001` in `gpu_dqn_trainer.rs:192` (GpuDqnTrainConfig::default)
- [ ] **Step 3:** Run `SQLX_OFFLINE=true cargo check -p ml -p ml-dqn` — verify clean build
- [ ] **Step 4:** Commit: "fix: entropy coefficient 0.01→0.001 to match reward v4 scale"

---

## Task 2: Add MSE Loss CUDA Kernel (15 min)

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — add `compile_mse_loss_kernel()` and `compile_mse_grad_kernel()`

The MSE/Huber loss operates on expected Q-values (scalars), not C51 distributions:
```
expected_Q = Σ_j softmax(logits_j) * z_j  (per branch, per sample)
target_Q = r + γ^n * max_a' expected_Q_target(s', a')
td_error = expected_Q(s, a) - target_Q
loss = 0.5 * td_error^2   (or Huber)
grad = td_error * d(expected_Q)/d(logits)
     = td_error * p_j * (z_j - expected_Q)   (softmax gradient)
```

- [ ] **Step 1:** Add `mse_loss_kernel` inline CUDA source in `gpu_dqn_trainer.rs` — computes expected Q from logits + support, then MSE vs Bellman target. Writes per-sample loss and td_errors.
- [ ] **Step 2:** Add `mse_grad_kernel` inline CUDA source — backprop through softmax expectation into logit gradients. Routes through dueling architecture (same as c51_grad).
- [ ] **Step 3:** Add `compile_mse_loss_kernel()` and `compile_mse_grad_kernel()` functions.
- [ ] **Step 4:** Add `mse_loss_kernel` and `mse_grad_kernel` fields to `GpuDqnTrainer` struct. Compile in `new()`.
- [ ] **Step 5:** Run `SQLX_OFFLINE=true cargo check -p ml` — verify clean build.
- [ ] **Step 6:** Commit: "feat: add MSE loss CUDA kernels for C51 warmup phase"

---

## Task 3: Add Warmup Config + Loss Mode Switch (10 min)

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — add `set_loss_mode()` method
- Modify: `crates/ml/src/trainers/dqn/config.rs` — add `c51_warmup_epochs` field
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — wire warmup config

- [ ] **Step 1:** Add `c51_warmup_epochs: usize` to `DQNHyperparameters` with default `5`.
- [ ] **Step 2:** Add `loss_mode: LossMode` enum (`Mse`, `C51`) to `GpuDqnTrainer`. Default `Mse`.
- [ ] **Step 3:** Add `pub fn set_loss_mode(&mut self, mode: LossMode)` — switches which kernels are used in `launch_loss()` and `launch_loss_grad()`.
- [ ] **Step 4:** In `fused_training.rs`, pass `c51_warmup_epochs` to the trainer config. In the training loop, switch loss mode at the epoch boundary:
  ```rust
  if epoch >= hyperparams.c51_warmup_epochs {
      fused.set_loss_mode(LossMode::C51);
  }
  ```
- [ ] **Step 5:** Run `SQLX_OFFLINE=true cargo check -p ml` — verify clean build.
- [ ] **Step 6:** Commit: "feat: C51 warmup — MSE loss for first N epochs, then switch to C51"

---

## Task 4: Validate Convergence Locally (10 min)

**Files:**
- Modify: `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs`

- [ ] **Step 1:** Run the 50-epoch convergence test:
  ```bash
  FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_50_epoch_convergence --ignored --nocapture
  ```
- [ ] **Step 2:** Verify loss trajectory: should DECREASE in MSE phase (epochs 0-4), then stabilize or slowly increase when C51 kicks in (epochs 5+).
- [ ] **Step 3:** Verify Q-values stay within v_range ±2.0 throughout all 50 epochs.
- [ ] **Step 4:** Verify gradient norm stays bounded (should be much lower during MSE phase).
- [ ] **Step 5:** Update the smoke test assertions based on observed metrics.
- [ ] **Step 6:** Commit: "test: validate C51 warmup convergence on real ES data"

---

## Task 5: Wire into Hyperopt Search Space (5 min)

**Files:**
- Modify: `crates/ml/src/hyperopt/adapters/dqn.rs`

- [ ] **Step 1:** Add `c51_warmup_epochs` to `DQNParams` struct (default 5).
- [ ] **Step 2:** Add to `continuous_bounds()` with range [0, 20] (dim 40).
- [ ] **Step 3:** Add to `from_continuous()` — round to nearest integer.
- [ ] **Step 4:** Add to `to_continuous()` and `param_names()`.
- [ ] **Step 5:** Wire through `build_hyperparams()` at line ~3039.
- [ ] **Step 6:** Update tests for 41D search space.
- [ ] **Step 7:** Commit: "feat: add c51_warmup_epochs to hyperopt search space"

---

## Key Design Decisions

### Why MSE Warmup Instead of Just Reducing Entropy?

Reducing entropy from 0.01 to 0.001 helps but doesn't fix the fundamental issue: C51 cross-entropy between nearly-identical distributions produces near-zero gradient signal for the VALUE (mean Q). The C51 gradient tells the network "make your distribution look like the target distribution" — but when both distributions are almost identical, the gradient is noise.

MSE directly says "make E[Q(s,a)] closer to the Bellman target" — a clear, strong signal.

### Why Not Remove C51 Entirely?

C51 distributional RL provides richer value representation than scalar Q-values. Once Q-values are in the right neighborhood (after MSE warmup), C51 refines the SHAPE of the value distribution — capturing uncertainty, multi-modality, and tail risk. This is valuable for trading where the return distribution is heavy-tailed.

### Why 5 Epochs Default?

With 2000+ steps per epoch on H100, 5 epochs = 10,000+ training steps. That's enough for MSE to push Q-values from random initialization to the correct neighborhood (~0.02 for ES futures). After that, C51 has a meaningful target to refine.

On the smoke test (50 steps/epoch), 5 epochs = 250 steps — still enough for Q-values to move in the right direction.
