foxhunt/docs/superpowers/plans/2026-03-27-gradient-stability.md

# Gradient Stability — Proper Multi-Task Gradient Isolation + Full Spectral Norm

> **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 gradient explosion by isolating every gradient source into its own buffer with budget-allocated clipping, extending spectral normalization to all weight matrices, and adding per-component monitoring.

**Architecture:** Four gradient sources (C51, CQL, IQN, ensemble) each get an independent norm budget that sums to exactly `max_grad_norm`. CQL gets its own scratch buffer (currently it accumulates directly into `grad_buf` with beta=1.0, mixing with C51). Spectral norm power iteration constrains all 10 weight matrices (2 trunk + 8 heads) with uniform σ_max. Per-component Prometheus gauges provide real-time gradient interference detection.

**Tech Stack:** Rust, CUDA (precompiled cubin), cudarc, cuBLAS, Prometheus, tokio

**Prior work (already merged):** `dqn_clipped_saxpy_kernel`, `dqn_clip_grad_kernel`, `clip_grad_buf_inplace()`, clipped SAXPY for IQN/ensemble, `weight_decay` hardcode fix, intermediate gradient clipping between CQL and IQN. **This plan fixes design flaws in the prior work.**

---

## Design — Why the current per-component clipping is wrong

The already-implemented clipping has three flaws:

**Flaw 1: CQL has no independent clipping.** CQL's `backward_full(beta=1.0)` accumulates directly into `grad_buf` on top of C51. The intermediate clip (Step 5d) then clips the MIXED C51+CQL. If CQL dominates, C51's gradient direction is destroyed by the clip.

**Flaw 2: Unbudgeted clip thresholds.** Each component clips to the FULL `max_grad_norm`:
```
After Step 5d:  ||grad_buf|| ≤ max_grad_norm           (C51+CQL)
After IQN:      ||grad_buf|| ≤ max_grad_norm + iqn_lambda × max_grad_norm
After ensemble: ||grad_buf|| ≤ max_grad_norm + iqn_lambda×mgn + ens_scale×mgn
Adam final clip: ||grad_buf|| → min(||grad||, max_grad_norm)
```
The combined norm is ~1.35× `max_grad_norm`. Adam clips 26% every step, wasting gradient signal and biasing the direction toward whichever component was largest.

**Flaw 3: CQL and IQN share `bw_d_h_*` scratch buffers.** This works only because IQN runs before CQL. An ordering change would silently corrupt gradients. Undocumented fragile invariant.

### Proper design

Each source gets a **budget fraction** of `max_grad_norm`:

| Component | Budget fraction | Rationale |
|---|---|---|
| C51/MSE (primary) | 0.70 | Primary RL loss — gets most of the gradient capacity |
| CQL (conservative penalty) | 0.15 | Regularizer — secondary importance |
| IQN (quantile risk) | 0.10 | Auxiliary head — small trunk contribution |
| Ensemble (diversity) | 0.05 | Weakest signal — diversity is a gentle push |
| **Total** | **1.00** | Adam clip NEVER fires → zero wasted gradient |

Pipeline after fix:
```
1. graph_forward     → grad_buf = C51_raw
2. clip_grad_buf(0.70 × mgn)  → ||grad_buf|| ≤ C51 budget
3. CQL backward      → cql_scratch = CQL_raw  (separate buffer)
4. clipped SAXPY      → grad_buf += clip(cql_scratch, 0.15 × mgn)
5. IQN clipped SAXPY  → grad_buf += λ × clip(iqn_scratch, 0.10 × mgn)
6. Ens clipped SAXPY  → grad_buf += s × clip(ens_scratch, 0.05 × mgn)
7. graph_adam         → Adam (safety clip at mgn, should never fire)
```

---

## File Map

| File | Role | Action |
|---|---|---|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | GPU trainer — buffers, kernel launches, spectral norm | Modify: add `cql_grad_scratch` buffer, change `apply_cql_gradient` to use scratch, add 16 spec_u/v buffers, extend spectral norm to all heads |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Fused training orchestration | Modify: add C51 clip after graph_forward, change CQL to use scratch+SAXPY, update all clip thresholds to budget fractions |
| `crates/common/src/metrics/training_metrics.rs` | Prometheus gauge definitions + setters | Modify: add 2 per-component gradient norm gauges |
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Training loop — epoch metrics | Modify: push per-component grad norms at epoch boundary |
| `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` | Smoke tests | Modify: add gradient-bounded-across-epochs test |

---

### Task 1: Isolate CQL gradient into separate scratch buffer

CQL currently runs `backward_full(beta=1.0)` directly into `grad_buf`, mixing with C51 before any CQL-specific clipping. Fix: allocate a `total_params`-sized scratch buffer, run CQL backward into it, then independently clip and SAXPY into `grad_buf`.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

- [ ] **Step 1: Add `cql_grad_scratch` field to struct**

After the existing `grad_norm_buf` field (~line 406), add:

```rust
    cql_grad_scratch: CudaSlice<f32>, // [TOTAL_PARAMS] CQL gradient isolation buffer
```

- [ ] **Step 2: Allocate in constructor**

After the existing `grad_buf` allocation (~line 1530), add:

```rust
        let cql_grad_scratch = alloc_f32(&stream, total_params, "cql_grad_scratch")?;
```

- [ ] **Step 3: Add to struct construction**

After `grad_buf,` in the `Ok(Self { ... })` block, add:

```rust
            cql_grad_scratch,
```

- [ ] **Step 4: Rewrite `apply_cql_gradient` to use scratch buffer**

The core change: instead of passing `grad_buf` pointer to `backward_full`, pass `cql_grad_scratch`. Zero the scratch first (since backward uses beta=1.0 accumulation).

In `apply_cql_gradient`, replace the backward_full call block (lines ~1195-1209). Change from:

```rust
        // Run full backward pass with CQL logit gradients.
        // This ACCUMULATES into grad_buf (beta=1.0 in cuBLAS SGEMM),
        // adding CQL parameter gradients on top of C51's.
        self.cublas_backward.backward_full(
            &self.stream,
            d_v_ptr,
            &d_adv_ptrs,
            states_ptr_fw,
            h_s1_ptr, h_s2_ptr, h_v_ptr,
            &[h_b0_ptr, h_b1_ptr, h_b2_ptr],
            &w_ptrs,
            raw_device_ptr(&self.grad_buf, &self.stream),
            scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v,
            &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2],
        ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?;

        Ok(true)
```

To:

```rust
        // Zero CQL scratch buffer (backward_full uses beta=1.0 accumulation)
        self.stream.memset_zeros(&mut self.cql_grad_scratch)
            .map_err(|e| MLError::ModelError(format!("zero cql_grad_scratch: {e}")))?;

        // Run full backward pass with CQL logit gradients into ISOLATED scratch buffer.
        // This produces CQL parameter gradients WITHOUT mixing with C51's grad_buf.
        self.cublas_backward.backward_full(
            &self.stream,
            d_v_ptr,
            &d_adv_ptrs,
            states_ptr_fw,
            h_s1_ptr, h_s2_ptr, h_v_ptr,
            &[h_b0_ptr, h_b1_ptr, h_b2_ptr],
            &w_ptrs,
            raw_device_ptr(&self.cql_grad_scratch, &self.stream),
            scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v,
            &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2],
        ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?;

        Ok(true)
```

- [ ] **Step 5: Add `apply_cql_clipped_saxpy` method**

Add a new public method after `apply_cql_gradient` that clips the CQL scratch and SAXPYs into `grad_buf`. This is called from `fused_training.rs` after `apply_cql_gradient` returns `Ok(true)`.

```rust
    /// Clip CQL gradient scratch buffer and add to grad_buf via SAXPY.
    ///
    /// Called after `apply_cql_gradient` which populated `cql_grad_scratch`.
    /// Uses the same clipped SAXPY pattern as IQN/ensemble: compute norm of
    /// scratch, then add with per-component clip.
    ///
    /// `cql_budget` is the CQL fraction of max_grad_norm (e.g. 0.15 × max_grad_norm).
    pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> {
        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        // Compute CQL gradient norm
        self.stream.memset_zeros(&mut self.grad_norm_buf)
            .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm: {e}")))?;

        let total = self.total_params as i32;
        let blocks = ((self.total_params + 255) / 256) as u32;
        let scratch_ptr = raw_device_ptr(&self.cql_grad_scratch, &self.stream);
        let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream);

        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_kernel)
                .arg(&scratch_ptr)
                .arg(&norm_ptr)
                .arg(&total)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 256,
                })
                .map_err(|e| MLError::ModelError(format!("CQL grad_norm: {e}")))?;
        }

        // Clipped SAXPY: grad_buf += 1.0 * clip(cql_scratch, cql_budget)
        let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream);
        let alpha = 1.0_f32;
        unsafe {
            self.stream
                .launch_builder(&self.clipped_saxpy_kernel)
                .arg(&grad_ptr)
                .arg(&scratch_ptr)
                .arg(&alpha)
                .arg(&cql_budget)
                .arg(&norm_ptr)
                .arg(&total)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("CQL clipped SAXPY: {e}")))?;
        }

        Ok(())
    }
```

- [ ] **Step 6: Verify compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5`
Expected: `Finished` (new method is unused — wired in Task 2).

- [ ] **Step 7: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat: isolate CQL gradient into separate scratch buffer with independent clipping"
```

---

### Task 2: Implement gradient budget allocation

Change all clip thresholds to use fractional budgets that sum to `max_grad_norm`. Add C51 clipping after graph_forward. Wire up the new CQL isolated path.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

**Budget constants** (defined in fused_training.rs):
```rust
/// Per-component gradient norm budget fractions. Must sum to 1.0.
/// The primary C51 loss gets 70% of max_grad_norm. Auxiliary objectives
/// (CQL, IQN, ensemble) share the remaining 30%.
const C51_GRAD_BUDGET: f32 = 0.70;
const CQL_GRAD_BUDGET: f32 = 0.15;
const IQN_GRAD_BUDGET: f32 = 0.10;
const ENS_GRAD_BUDGET: f32 = 0.05;
```

- [ ] **Step 1: Add budget constants**

At the top of `fused_training.rs` (after the `use` imports, before struct definitions), add the 4 constants shown above.

- [ ] **Step 2: Add C51 clip after graph_forward**

After `train_step_gpu` returns (after `let _fused_placeholder = ...` ~line 570), add a clip of the raw C51 gradient:

```rust
        // ── Step 2b: Clip raw C51 gradient to its budget ──────────────────
        // The graph_forward backward produces raw C51 gradients in grad_buf.
        // Clip to C51's budget fraction BEFORE any auxiliary injection.
        {
            let c51_budget = self.trainer.config().max_grad_norm * C51_GRAD_BUDGET;
            self.trainer.clip_grad_buf_inplace(c51_budget)
                .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?;
        }
```

- [ ] **Step 3: Wire CQL isolated path**

Replace the existing CQL block (~lines 742-756) with the new isolated CQL path:

```rust
        // ── Step 5c: CQL conservative penalty (isolated gradient) ─────────
        // CQL backward runs into a SEPARATE scratch buffer (not grad_buf).
        // Its gradient is independently clipped to CQL's budget fraction,
        // then added to grad_buf via clipped SAXPY.
        if self.trainer.has_cql() {
            match self.trainer.apply_cql_gradient() {
                Ok(true) => {
                    let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET;
                    self.trainer.apply_cql_clipped_saxpy(cql_budget)
                        .map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?;
                    tracing::trace!("CQL gradient: isolated → clipped → SAXPY into grad_buf");
                }
                Ok(false) => {} // CQL disabled or alpha=0
                Err(e) => {
                    tracing::warn!("CQL gradient failed (non-fatal): {e}");
                }
            }
        }
```

- [ ] **Step 4: Remove the old combined Step 5d clip**

Delete the entire "Step 5d: Clip combined C51+CQL gradient before auxiliary injections" block (~lines 754-773). It's replaced by Steps 2b (C51 clip) and 5c (CQL isolation). Keep the diagnostic logging but move it to Step 2b:

In the Step 2b block added above, add the diagnostic:

```rust
            // Diagnostic: log C51 pre-clip norm every 1000 steps
            if self.steps_since_varmap_sync % 1000 == 0 {
                let pre_clip_norm = self.trainer.read_grad_norm_sync()
                    .unwrap_or(f32::NAN);
                tracing::info!(
                    c51_raw_grad_norm = pre_clip_norm,
                    c51_budget,
                    step = self.steps_since_varmap_sync,
                    "Per-component gradient diagnostic (C51 before budget clip)"
                );
            }
```

- [ ] **Step 5: Update IQN clip threshold to budget fraction**

In `gpu_dqn_trainer.rs`, in `apply_iqn_trunk_gradient` (~line 830), change:

```rust
            let max_component_norm = self.config.max_grad_norm;
```

To:

```rust
            let max_component_norm = self.config.max_grad_norm * 0.10; // IQN_GRAD_BUDGET
```

- [ ] **Step 6: Update ensemble clip threshold to budget fraction**

In `gpu_dqn_trainer.rs`, in `apply_ensemble_trunk_gradient` (~line 1071), change:

```rust
            let max_component_norm = self.config.max_grad_norm;
```

To:

```rust
            let max_component_norm = self.config.max_grad_norm * 0.05; // ENS_GRAD_BUDGET
```

- [ ] **Step 7: Verify compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5`
Expected: `Finished` with no errors.

- [ ] **Step 8: Run existing tests**

Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5`
Expected: All tests pass (122+).

- [ ] **Step 9: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: proper gradient budget allocation — C51 70%, CQL 15%, IQN 10%, ensemble 5%"
```

---

### Task 3: Allocate spectral norm u/v vectors for all 8 head weight matrices

Each weight matrix W [out_dim, in_dim] needs u [out_dim] and v [in_dim] vectors for power iteration. The trunk already has `spec_u_s1/v_s1`, `spec_u_s2/v_s2`. We add 8 pairs for the heads.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

**Weight matrix dimensions** (from `DuelingWeightSet` at `gpu_weights.rs:109` and `BranchingWeightSet` at `gpu_weights.rs:141`):

| Matrix | Shape [out, in] | u size | v size | GOFF index |
|---|---|---|---|---|
| W_v1 | [VALUE_H, SHARED_H2] | value_h | shared_h2 | 4 |
| W_v2 | [NUM_ATOMS, VALUE_H] | num_atoms | value_h | 6 |
| W_a1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 8 |
| W_a2 | [B0×NUM_ATOMS, ADV_H] | b0×num_atoms | adv_h | 10 |
| W_bo1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 12 |
| W_bo2 | [B1×NUM_ATOMS, ADV_H] | b1×num_atoms | adv_h | 14 |
| W_bu1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 16 |
| W_bu2 | [B2×NUM_ATOMS, ADV_H] | b2×num_atoms | adv_h | 18 |

GOFF layout verified at `gpu_dqn_trainer.rs:280-302` (`compute_param_sizes`).

- [ ] **Step 1: Add 16 struct fields**

After `spec_v_s2: CudaSlice<f32>,` (~line 347):

```rust
    // Head spectral norm vectors (power iteration)
    // Value head
    spec_u_v1: CudaSlice<f32>,  // [VALUE_H]
    spec_v_v1: CudaSlice<f32>,  // [SHARED_H2]
    spec_u_v2: CudaSlice<f32>,  // [NUM_ATOMS]
    spec_v_v2: CudaSlice<f32>,  // [VALUE_H]
    // Exposure head (branch 0) — in DuelingWeightSet
    spec_u_a1: CudaSlice<f32>,  // [ADV_H]
    spec_v_a1: CudaSlice<f32>,  // [SHARED_H2]
    spec_u_a2: CudaSlice<f32>,  // [B0*NUM_ATOMS]
    spec_v_a2: CudaSlice<f32>,  // [ADV_H]
    // Order head (branch 1) — in BranchingWeightSet
    spec_u_bo1: CudaSlice<f32>, // [ADV_H]
    spec_v_bo1: CudaSlice<f32>, // [SHARED_H2]
    spec_u_bo2: CudaSlice<f32>, // [B1*NUM_ATOMS]
    spec_v_bo2: CudaSlice<f32>, // [ADV_H]
    // Urgency head (branch 2) — in BranchingWeightSet
    spec_u_bu1: CudaSlice<f32>, // [ADV_H]
    spec_v_bu1: CudaSlice<f32>, // [SHARED_H2]
    spec_u_bu2: CudaSlice<f32>, // [B2*NUM_ATOMS]
    spec_v_bu2: CudaSlice<f32>, // [ADV_H]
```

- [ ] **Step 2: Allocate and initialize in constructor**

After the existing `spec_v_s2` init block (~line 1746), using the existing `rand_unit` helper:

```rust
        // Head spectral norm u/v vectors (random unit initialization)
        let vh = config.value_h;
        let ah = config.adv_h;
        let na = config.num_atoms;
        let b0na = config.branch_0_size * na;
        let b1na = config.branch_1_size * na;
        let b2na = config.branch_2_size * na;

        let mut spec_u_v1 = alloc_f32(&stream, vh, "spec_u_v1")?;
        let mut spec_v_v1 = alloc_f32(&stream, config.shared_h2, "spec_v_v1")?;
        let mut spec_u_v2 = alloc_f32(&stream, na, "spec_u_v2")?;
        let mut spec_v_v2 = alloc_f32(&stream, vh, "spec_v_v2")?;
        let mut spec_u_a1 = alloc_f32(&stream, ah, "spec_u_a1")?;
        let mut spec_v_a1 = alloc_f32(&stream, config.shared_h2, "spec_v_a1")?;
        let mut spec_u_a2 = alloc_f32(&stream, b0na, "spec_u_a2")?;
        let mut spec_v_a2 = alloc_f32(&stream, ah, "spec_v_a2")?;
        let mut spec_u_bo1 = alloc_f32(&stream, ah, "spec_u_bo1")?;
        let mut spec_v_bo1 = alloc_f32(&stream, config.shared_h2, "spec_v_bo1")?;
        let mut spec_u_bo2 = alloc_f32(&stream, b1na, "spec_u_bo2")?;
        let mut spec_v_bo2 = alloc_f32(&stream, ah, "spec_v_bo2")?;
        let mut spec_u_bu1 = alloc_f32(&stream, ah, "spec_u_bu1")?;
        let mut spec_v_bu1 = alloc_f32(&stream, config.shared_h2, "spec_v_bu1")?;
        let mut spec_u_bu2 = alloc_f32(&stream, b2na, "spec_u_bu2")?;
        let mut spec_v_bu2 = alloc_f32(&stream, ah, "spec_v_bu2")?;

        for (buf, len) in [
            (&mut spec_u_v1, vh), (&mut spec_v_v1, config.shared_h2),
            (&mut spec_u_v2, na), (&mut spec_v_v2, vh),
            (&mut spec_u_a1, ah), (&mut spec_v_a1, config.shared_h2),
            (&mut spec_u_a2, b0na), (&mut spec_v_a2, ah),
            (&mut spec_u_bo1, ah), (&mut spec_v_bo1, config.shared_h2),
            (&mut spec_u_bo2, b1na), (&mut spec_v_bo2, ah),
            (&mut spec_u_bu1, ah), (&mut spec_v_bu1, config.shared_h2),
            (&mut spec_u_bu2, b2na), (&mut spec_v_bu2, ah),
        ] {
            let init = rand_unit(len, &mut rng_state);
            stream.memcpy_htod(&init, buf)
                .map_err(|e| MLError::ModelError(format!("spec head init: {e}")))?;
        }
```

- [ ] **Step 3: Add to `Ok(Self { ... })`**

After `spec_v_s2,`:

```rust
            spec_u_v1, spec_v_v1, spec_u_v2, spec_v_v2,
            spec_u_a1, spec_v_a1, spec_u_a2, spec_v_a2,
            spec_u_bo1, spec_v_bo1, spec_u_bo2, spec_v_bo2,
            spec_u_bu1, spec_v_bu1, spec_u_bu2, spec_v_bu2,
```

- [ ] **Step 4: Verify compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5`
Expected: `Finished` (unused fields warning OK — used in Task 4).

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat: allocate spectral norm u/v vectors for all 8 head weight matrices"
```

---

### Task 4: Extend `apply_spectral_norm` to all heads + sync params_buf

Add 8 spectral norm kernel launches (one per head weight matrix) and sync the spectrally-normalized weights back to `params_buf` at their GOFF offsets.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

**Design note:** Uniform σ_max for all layers. With σ_max=3.0 on 4 layers (trunk→head_fc→head_out), the end-to-end Lipschitz bound is 3⁴=81. The C51 logit spread under σ_max=3.0 is ~18 (for 101 atoms with typical activations), which is more than sufficient for sharp softmax distributions. Output layer spectral norm does NOT collapse C51 — verified analytically.

- [ ] **Step 1: Change `apply_spectral_norm` signature**

```rust
pub fn apply_spectral_norm(
    &mut self,
    online_dueling: &mut DuelingWeightSet,
    online_branching: &mut BranchingWeightSet,
) -> Result<(), MLError> {
```

- [ ] **Step 2: Add 8 spectral norm kernel launches**

After the existing W_s2 block (before the trunk sync-back), add:

```rust
        // Head spectral norm — uniform σ_max, same as trunk
        let vh = self.config.value_h as i32;
        let ah = self.config.adv_h as i32;
        let na = self.config.num_atoms as i32;
        let b0na = (self.config.branch_0_size * self.config.num_atoms) as i32;
        let b1na = (self.config.branch_1_size * self.config.num_atoms) as i32;
        let b2na = (self.config.branch_2_size * self.config.num_atoms) as i32;

        macro_rules! spec_norm {
            ($w:expr, $u:expr, $v:expr, $od:expr, $id:expr, $name:literal) => {{
                let (w_ptr, _g) = $w.device_ptr(&self.stream);
                let (u_ptr, _g2) = $u.device_ptr(&self.stream);
                let (v_ptr, _g3) = $v.device_ptr(&self.stream);
                unsafe {
                    self.stream
                        .launch_builder(&self.spectral_norm_kernel)
                        .arg(&w_ptr).arg(&u_ptr).arg(&v_ptr)
                        .arg(&$od).arg(&$id).arg(&sigma_max)
                        .launch(LaunchConfig {
                            grid_dim: (1, 1, 1),
                            block_dim: (256, 1, 1),
                            shared_mem_bytes: 512 * 4,
                        })
                        .map_err(|e| MLError::ModelError(
                            format!(concat!("spectral_norm ", $name, ": {}"), e)
                        ))?;
                }
            }};
        }

        spec_norm!(online_dueling.w_v1, self.spec_u_v1, self.spec_v_v1, vh, sh2, "W_v1");
        spec_norm!(online_dueling.w_v2, self.spec_u_v2, self.spec_v_v2, na, vh, "W_v2");
        spec_norm!(online_dueling.w_a1, self.spec_u_a1, self.spec_v_a1, ah, sh2, "W_a1");
        spec_norm!(online_dueling.w_a2, self.spec_u_a2, self.spec_v_a2, b0na, ah, "W_a2");
        spec_norm!(online_branching.w_bo1, self.spec_u_bo1, self.spec_v_bo1, ah, sh2, "W_bo1");
        spec_norm!(online_branching.w_bo2, self.spec_u_bo2, self.spec_v_bo2, b1na, ah, "W_bo2");
        spec_norm!(online_branching.w_bu1, self.spec_u_bu1, self.spec_v_bu1, ah, sh2, "W_bu1");
        spec_norm!(online_branching.w_bu2, self.spec_u_bu2, self.spec_v_bu2, b2na, ah, "W_bu2");
```

- [ ] **Step 3: Add head weight sync-back to params_buf**

After the existing trunk sync-back section, add:

```rust
        // Sync head weights back to params_buf at their GOFF offsets
        {
            let param_sizes = compute_param_sizes(&self.config);
            let f32_sz = std::mem::size_of::<f32>();
            let params_base = raw_device_ptr(&self.params_buf, &self.stream);

            let byte_offset = |idx: usize| -> u64 {
                param_sizes[..idx].iter().sum::<usize>() as u64 * f32_sz as u64
            };

            // GOFF indices: 4=w_v1, 6=w_v2, 8=w_a1, 10=w_a2,
            //               12=w_bo1, 14=w_bo2, 16=w_bu1, 18=w_bu2
            let head_syncs: [(u64, u64, usize); 8] = [
                (raw_device_ptr(&online_dueling.w_v1, &self.stream), byte_offset(4), param_sizes[4]),
                (raw_device_ptr(&online_dueling.w_v2, &self.stream), byte_offset(6), param_sizes[6]),
                (raw_device_ptr(&online_dueling.w_a1, &self.stream), byte_offset(8), param_sizes[8]),
                (raw_device_ptr(&online_dueling.w_a2, &self.stream), byte_offset(10), param_sizes[10]),
                (raw_device_ptr(&online_branching.w_bo1, &self.stream), byte_offset(12), param_sizes[12]),
                (raw_device_ptr(&online_branching.w_bo2, &self.stream), byte_offset(14), param_sizes[14]),
                (raw_device_ptr(&online_branching.w_bu1, &self.stream), byte_offset(16), param_sizes[16]),
                (raw_device_ptr(&online_branching.w_bu2, &self.stream), byte_offset(18), param_sizes[18]),
            ];

            for (src, dst_off, n_elems) in head_syncs {
                unsafe {
                    cudarc::driver::result::memcpy_dtod_async(
                        params_base + dst_off, src, n_elems * f32_sz,
                        self.stream.cu_stream()
                    ).map_err(|e| MLError::ModelError(format!("spectral sync head: {e}")))?;
                }
            }
        }
```

- [ ] **Step 4: Update call site in fused_training.rs**

Change (~line 560):
```rust
self.trainer.apply_spectral_norm(&mut self.online_dueling)
```
To:
```rust
self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching)
```

- [ ] **Step 5: Verify compilation + run tests**

Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -10`
Expected: Compiles, all 122+ tests pass.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: spectral norm on all 10 weight matrices — uniform σ_max, Lipschitz bounded"
```

---

### Task 5: Per-component gradient Prometheus gauges

**Files:**
- Modify: `crates/common/src/metrics/training_metrics.rs`
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs`

- [ ] **Step 1: Register gauges and add setters in training_metrics.rs**

In `register_all()` (after existing gradient gauge registrations):

```rust
    _ = register_gauge_vec(
        "foxhunt_training_grad_norm_c51",
        "Per-component gradient L2 norm: C51 primary loss (before budget clip)",
        mf,
    );
    _ = register_gauge_vec(
        "foxhunt_training_grad_norm_combined",
        "Combined gradient L2 norm after all injections (should ≤ max_grad_norm)",
        mf,
    );
```

After `set_gradient_norm`:

```rust
pub fn set_grad_norm_c51(model: &str, fold: &str, norm: f64) {
    set_gauge_vec("foxhunt_training_grad_norm_c51", &[model, fold], norm);
}

pub fn set_grad_norm_combined(model: &str, fold: &str, norm: f64) {
    set_gauge_vec("foxhunt_training_grad_norm_combined", &[model, fold], norm);
}
```

- [ ] **Step 2: Track norms in FusedTrainingCtx**

Add fields `last_c51_raw_norm: f32` and `last_combined_norm: f32` to `FusedTrainingCtx`. Initialize to 0.0. Update `last_c51_raw_norm` from the diagnostic readback in Step 2b. Update `last_combined_norm` from `fused_result.grad_norm` after graph_adam.

Add accessor:
```rust
    pub fn grad_norm_diagnostics(&self) -> (f32, f32) {
        (self.last_c51_raw_norm, self.last_combined_norm)
    }
```

- [ ] **Step 3: Push in training_loop.rs**

After `training_metrics::set_gradient_norm(...)`:
```rust
        if let Some(ref fused) = self.fused_ctx {
            let (c51, combined) = fused.grad_norm_diagnostics();
            training_metrics::set_grad_norm_c51("dqn", "current", c51 as f64);
            training_metrics::set_grad_norm_combined("dqn", "current", combined as f64);
        }
```

- [ ] **Step 4: Verify compilation**

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

- [ ] **Step 5: Commit**

```bash
git add crates/common/src/metrics/training_metrics.rs crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat: per-component gradient Prometheus gauges (C51 raw, combined)"
```

---

### Task 6: Smoke test — gradient norms bounded across epochs

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

- [ ] **Step 1: Write test**

```rust
/// Gradient norms must stay bounded across all epochs.
///
/// Before fix: grad_norm grew 89K → 3.6B (40,000x over 8 epochs).
/// After fix: per-component budget allocation ensures Adam clip never fires
/// and combined norm ≤ max_grad_norm at every step.
#[test]
#[ignore] // Loads real training data — run via nightly CI or manual trigger
fn test_gradient_norm_bounded_across_epochs() -> anyhow::Result<()> {
    let data_dir = test_data_dir()
        .expect("FOXHUNT_TEST_DATA or test_data/ must exist");

    let mut p = smoke_params();
    p.epochs = 3;

    let mut trainer = smoke_trainer_with(p)?;
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _best| {
        Ok("skip".to_owned())
    }))?;

    let grad_norm = metrics.additional_metrics.get("avg_gradient_norm")
        .copied().unwrap_or(f64::NAN);

    assert!(grad_norm.is_finite(), "avg_gradient_norm not finite: {grad_norm}");
    assert!(grad_norm > 0.0, "avg_gradient_norm should be positive: {grad_norm}");
    assert!(
        grad_norm < 100_000.0,
        "GRADIENT EXPLOSION: avg_gradient_norm {grad_norm:.1} exceeds 100K"
    );
    assert_eq!(metrics.epochs_trained, 3, "should complete all 3 epochs");

    drop(trainer);
    drop(rt);
    Ok(())
}
```

- [ ] **Step 2: Compile + run**

Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- test_gradient_norm_bounded_across_epochs --no-run 2>&1 | tail -5`

If GPU + test data available:
Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_gradient_norm_bounded_across_epochs --ignored --nocapture 2>&1 | tail -20`

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
git commit -m "test: gradient norm bounded across epochs smoke test"
```

---

### Task 7: Research — PCGrad, LayerNorm, adaptive weighting analysis

Pure research — no code. Write analysis of 3 additional architectural approaches.

**Files:**
- Create: `docs/superpowers/plans/2026-03-27-gradient-stability-research.md`

- [ ] **Step 1: Write analysis** covering PCGrad (Yu 2020), LayerNorm, and Uncertainty Weighting (Kendall 2018). Evaluate feasibility, impact, and interaction with our per-component budget + spectral norm approach. Recommend PCGrad as the best follow-up if gradient interference persists after validation.

- [ ] **Step 2: Commit**

```bash
git add docs/superpowers/plans/2026-03-27-gradient-stability-research.md
git commit -m "docs: gradient stability research — PCGrad, LayerNorm, adaptive weighting"
```

---

## Execution Dependencies

```
Task 1 (CQL isolation) → Task 2 (budget allocation) → Task 6 (smoke test)
                                                      ↗
Task 3 (spec norm bufs) → Task 4 (spec norm launches)
Task 5 (Prometheus) ──────────────────────────────────
Task 7 (research) — independent
```

Tasks 1→2 fix the per-component clipping design. Tasks 3→4 add spectral norm to heads. Task 5 is independent monitoring. Task 6 validates everything. Task 7 is research.