foxhunt/docs/superpowers/plans/2026-03-24-c51-loss-convergence.md

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:

  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:

  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.