foxhunt/docs/superpowers/plans/2026-03-23-stable-distributional-loss.md

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 #defines), add:

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

    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:

    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

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

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

• Step 4: Commit

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:

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

        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

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

• Step 4: Local test — verify loss doesn't explode

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

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:

// 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:

// 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

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

• Step 5: Commit

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:

/// 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:

// 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:

iqn_primary_loss = true

• Step 4: Compile check

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

• Step 5: Commit

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

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)

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)

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

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