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:

    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:

        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:

            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:

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

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

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

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

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

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

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

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

            let max_component_norm = self.config.max_grad_norm;

To:

            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:

            let max_component_norm = self.config.max_grad_norm;

To:

            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

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

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

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

            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

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

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:

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

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

self.trainer.apply_spectral_norm(&mut self.online_dueling)

To:

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

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

    _ = 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:

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:

    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(...):

        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

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

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

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

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.