# IQL Full Integration 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:** Wire IQL V(s) into PER priorities, C51 atom support, gradient scaling, and exploration — eliminating all hardcoded constants.

**Architecture:** Six new CUDA kernels in `iql_value_kernel.cu` + modifications to `c51_loss_kernel.cu`, `c51_grad_kernel.cu`, and `epsilon_greedy_kernel.cu`. A second `GpuIqlTrainer` (τ=0.3) provides expectile gap uncertainty. All v_range batch-level infrastructure is deleted.

**Tech Stack:** CUDA C (kernels), Rust (cudarc launcher), f32 throughout, zero atomicAdd.

**Spec:** `docs/superpowers/specs/2026-04-13-iql-full-integration-design.md`

---

### Task 1: New CUDA Kernels in `iql_value_kernel.cu`

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/iql_value_kernel.cu`

Add 6 new kernels after the existing `iql_compute_advantage_weights` kernel. These are all element-wise or reduction kernels — no shared memory complexity.

- [ ] **Step 1: Add `iql_modulate_td_errors` kernel**

Append after the `iql_compute_advantage_weights` kernel (after line ~593):

```cuda
/* ------------------------------------------------------------------ */
/*  PER Priority Modulation with Staleness Correction                  */
/* ------------------------------------------------------------------ */
/**
 * Modulate td_errors in-place with advantage weights and age-decay.
 *
 * td_errors[b] *= clamp(adv_weights[b], 1/K, K) * exp(-lambda * age / tau)
 *
 * K = exp(beta * 3 * sigma_adv) — dynamic clamp from advantage distribution.
 * age = (write_pos - indices[b]) % capacity — circular buffer distance.
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_modulate_td_errors(
    float* __restrict__ td_errors,             /* [B] in-place */
    const float* __restrict__ adv_weights,     /* [B] */
    const int* __restrict__ indices,           /* [B] replay buffer indices */
    float sigma_adv,                           /* EMA of advantage std */
    float beta,                                /* advantage temperature */
    float staleness_lambda,
    float staleness_tau,
    int write_pos,                             /* current replay buffer write cursor */
    int capacity,                              /* replay buffer capacity */
    int batch_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    /* Dynamic clamp bounds from advantage distribution */
    float K = expf(beta * 3.0f * fmaxf(sigma_adv, 1e-6f));
    K = fmaxf(K, 1.1f); /* ensure clamp range is at least [0.91, 1.1] */
    float inv_K = 1.0f / K;

    float w = fminf(fmaxf(adv_weights[b], inv_K), K);

    /* Staleness decay from circular buffer distance */
    int age = (write_pos - indices[b] + capacity) % capacity;
    float decay = expf(-staleness_lambda * (float)age / fmaxf(staleness_tau, 1.0f));

    td_errors[b] *= w * decay;
}
```

- [ ] **Step 2: Add `iql_adv_variance_reduce` kernel**

```cuda
/* ------------------------------------------------------------------ */
/*  Advantage Variance Reduce (sequential for determinism)             */
/* ------------------------------------------------------------------ */
/**
 * Compute mean and variance of advantage weights.
 * Output: stats[0] = mean, stats[1] = variance.
 *
 * Sequential (grid=1, block=1) for full determinism.
 */
extern "C" __global__
void iql_adv_variance_reduce(
    const float* __restrict__ adv_weights,   /* [B] */
    float* __restrict__ stats,               /* [2]: mean, variance */
    int batch_size
)
{
    float sum = 0.0f;
    for (int b = 0; b < batch_size; b++)
        sum += adv_weights[b];
    float mean = sum / (float)batch_size;

    float var_sum = 0.0f;
    for (int b = 0; b < batch_size; b++) {
        float d = adv_weights[b] - mean;
        var_sum += d * d;
    }
    stats[0] = mean;
    stats[1] = var_sum / (float)batch_size;
}
```

- [ ] **Step 3: Add `iql_compute_per_sample_support` kernel**

```cuda
/* ------------------------------------------------------------------ */
/*  Per-Sample C51 Atom Support                                        */
/* ------------------------------------------------------------------ */
/**
 * Compute per-sample distributional support centered on V(s).
 *
 * v_min(b) = V(s_b) - half_width(b)
 * v_max(b) = V(s_b) + half_width(b)
 * delta_z(b) = (v_max(b) - v_min(b)) / (num_atoms - 1)
 *
 * half_width = q_spread * (1 + gamma), where q_spread = max distance
 * from V(s) to any Q-value for this sample.
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_compute_per_sample_support(
    const float* __restrict__ v_out,         /* [B] V(s) from IQL */
    const float* __restrict__ q_out,         /* [B, total_actions] */
    float* __restrict__ per_sample_support,  /* [B, 3]: v_min, v_max, delta_z */
    float gamma,
    int batch_size,
    int total_actions,
    int num_atoms
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float v = v_out[b];
    const float* q = q_out + b * total_actions;

    /* Find max distance from V(s) to any Q-value */
    float spread = 0.0f;
    for (int a = 0; a < total_actions; a++) {
        float dist = fabsf(q[a] - v);
        spread = fmaxf(spread, dist);
    }

    float half_w = spread * (1.0f + gamma);
    float v_min = v - half_w;
    float v_max = v + half_w;
    float delta_z = (v_max - v_min) / (float)(num_atoms - 1);

    per_sample_support[b * 3 + 0] = v_min;
    per_sample_support[b * 3 + 1] = v_max;
    per_sample_support[b * 3 + 2] = delta_z;
}
```

- [ ] **Step 4: Add `iql_per_branch_advantage` kernel**

```cuda
/* ------------------------------------------------------------------ */
/*  Per-Branch Advantage Decomposition                                 */
/* ------------------------------------------------------------------ */
/**
 * Compute per-branch advantage magnitudes via Q-value marginalization.
 *
 * For each branch d, average Q over the other 3 branches for the taken
 * action in branch d, then measure distance from V(s).
 *
 * branch_scale(b,d) = A_branch(b,d) / max_d(A_branch(b,d))
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_per_branch_advantage(
    const float* __restrict__ q_out,         /* [B, total_actions] */
    const float* __restrict__ v_out,         /* [B] */
    const int* __restrict__ actions,         /* [B] factored action index */
    float* __restrict__ branch_scales,       /* [B, 4] */
    int batch_size,
    int total_actions,
    int b0_size, int b1_size, int b2_size, int b3_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float v = v_out[b];
    const float* q = q_out + b * total_actions;

    /* Decode factored action */
    int factored = actions[b];
    int max_act = b0_size * b1_size * b2_size * b3_size;
    if (factored < 0 || factored >= max_act) factored = 0;
    int a0 = factored / (b1_size * b2_size * b3_size);
    int a1 = (factored / (b2_size * b3_size)) % b1_size;
    int a2 = (factored / b3_size) % b2_size;
    int a3 = factored % b3_size;
    int branch_actions[4] = { a0, a1, a2, a3 };
    int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };

    /* For each branch, marginalize Q over other branches for the taken sub-action */
    float a_branch[4];
    float max_a = 0.0f;

    for (int d = 0; d < 4; d++) {
        int a_d = branch_actions[d];
        float q_sum = 0.0f;
        int count = 0;

        /* Iterate all joint actions where branch d == a_d */
        for (int joint = 0; joint < total_actions; joint++) {
            /* Extract branch d's sub-action from the joint action */
            int sub_d;
            if (d == 0) sub_d = joint / (b1_size * b2_size * b3_size);
            else if (d == 1) sub_d = (joint / (b2_size * b3_size)) % b1_size;
            else if (d == 2) sub_d = (joint / b3_size) % b2_size;
            else sub_d = joint % b3_size;

            if (sub_d == a_d) {
                q_sum += q[joint];
                count++;
            }
        }

        float q_marginal = (count > 0) ? (q_sum / (float)count) : v;
        a_branch[d] = fabsf(q_marginal - v);
        max_a = fmaxf(max_a, a_branch[d]);
    }

    /* Normalize: highest-advantage branch gets 1.0, others proportionally less */
    float inv_max = (max_a > 1e-8f) ? (1.0f / max_a) : 1.0f;
    for (int d = 0; d < 4; d++) {
        branch_scales[b * 4 + d] = a_branch[d] * inv_max;
    }
}
```

- [ ] **Step 5: Add `iql_expectile_gap` kernel**

```cuda
/* ------------------------------------------------------------------ */
/*  Expectile Gap (V_high - V_low)                                     */
/* ------------------------------------------------------------------ */
/**
 * Compute expectile gap and its batch mean.
 *
 * gap[b] = v_high[b] - v_low[b]
 * gap_mean = mean(gap)
 *
 * Two-phase: element-wise (grid=ceil(B/256)) + sequential mean (grid=1).
 */
extern "C" __global__
void iql_expectile_gap(
    const float* __restrict__ v_high,   /* [B] V(s) at tau_high */
    const float* __restrict__ v_low,    /* [B] V(s) at tau_low */
    float* __restrict__ gap,            /* [B] output */
    int batch_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;
    gap[b] = v_high[b] - v_low[b];
}

extern "C" __global__
void iql_gap_mean_reduce(
    const float* __restrict__ gap,      /* [B] */
    float* __restrict__ gap_mean,       /* [1] */
    int batch_size
)
{
    float sum = 0.0f;
    for (int b = 0; b < batch_size; b++)
        sum += gap[b];
    gap_mean[0] = sum / (float)batch_size;
}
```

- [ ] **Step 6: Add `iql_compute_per_sample_epsilon` kernel**

```cuda
/* ------------------------------------------------------------------ */
/*  Per-Sample Epsilon from Expectile Gap                              */
/* ------------------------------------------------------------------ */
/**
 * epsilon(b) = base_epsilon * sigmoid(gap[b] / gap_mean - 1.0)
 *
 * High gap (uncertainty) → more exploration.
 * Low gap (confidence) → exploit.
 * gap_mean normalizer keeps it relative — no hardcoded thresholds.
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_compute_per_sample_epsilon(
    const float* __restrict__ gap,          /* [B] */
    const float* __restrict__ gap_mean,     /* [1] */
    float* __restrict__ per_sample_eps,     /* [B] */
    float base_epsilon,
    int batch_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float gm = fmaxf(gap_mean[0], 1e-8f);
    float x = gap[b] / gm - 1.0f;
    float sig = 1.0f / (1.0f + expf(-x));
    per_sample_eps[b] = base_epsilon * sig;
}
```

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

Run: `SQLX_OFFLINE=true cargo build -p ml 2>&1 | head -5`
Expected: build.rs compiles `iql_value_kernel.cu` to cubin without errors.

- [ ] **Step 8: Commit**

```bash
git add crates/ml/src/cuda_pipeline/iql_value_kernel.cu
git commit -m "feat(iql): add 6 integration kernels — PER modulation, per-sample support, branch advantage, expectile gap"
```

---

### Task 2: Modify `c51_loss_kernel.cu` for Per-Sample Support

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu`

- [ ] **Step 1: Change `c51_loss_batched` signature — replace `v_range_buf` with `per_sample_support`**

In the kernel signature (line ~205), replace:
```cuda
    const float* __restrict__ v_range_buf,        /* [2] adaptive C51 z-support: [v_min, v_max] */
```
with:
```cuda
    const float* __restrict__ per_sample_support, /* [B, 3] per-sample: [v_min, v_max, delta_z] */
```

- [ ] **Step 2: Replace shared v_range read with per-sample register loads**

Replace lines 232-234:
```cuda
    /* Read adaptive v_range from device buffer (graph-safe, L1 cached) */
    float v_min = v_range_buf[0];
    float v_max = v_range_buf[1];
```
with:
```cuda
    /* Per-sample C51 support centered on V(s) */
    float v_min = per_sample_support[sample_id * 3];
    float v_max = per_sample_support[sample_id * 3 + 1];
    float delta_z = per_sample_support[sample_id * 3 + 2];

    /* Skip sample if support is degenerate (all Q-values identical) */
    if (delta_z < 1e-7f) {
        if (tid == 0) {
            per_sample_loss[sample_id] = 0.0f;
            td_errors[sample_id] = 0.0f;
        }
        return;
    }
```

- [ ] **Step 3: Remove the old delta_z computation**

Delete line 257:
```cuda
    float delta_z = (v_max - v_min) / (float)(num_atoms - 1);
```

This is now loaded from `per_sample_support` in Step 2.

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

Run: `SQLX_OFFLINE=true cargo build -p ml 2>&1 | head -5`
Expected: cubin compiles successfully.

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
git commit -m "feat(c51): per-sample atom support — V(s)-centered distributional RL"
```

---

### Task 3: Modify `c51_grad_kernel.cu` for Dynamic Branch Scales

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu`

- [ ] **Step 1: Add `branch_scales` parameter to kernel signature**

In the signature (line 14), after the `entropy_coeff` parameter, add:
```cuda
    const float* __restrict__ branch_scales    /* [B, 4] per-sample per-branch gradient scale */
```

Don't forget to add a comma after `float entropy_coeff,` (currently last param has no comma).

- [ ] **Step 2: Replace hardcoded branch_scale with dynamic per-sample read**

Replace line 72:
```cuda
        float branch_scale = (d == 1) ? 0.0f : 1.0f;
```
with:
```cuda
        float branch_scale = branch_scales[b * 4 + d];
```

- [ ] **Step 3: Remove hardcoded entropy ent_scale — branch_scales subsumes it**

Replace lines 60-67:
```cuda
        /* Entropy regularization */
        if (entropy_coeff > 0.0f) {
            float ent_scale;
            if (d == 1) ent_scale = 5.0f;
            else if (d == 2) ent_scale = 3.0f;
            else ent_scale = 1.0f;
            float lp_clamped = fmaxf(lp, -10.0f);
            d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
        }
```
with:
```cuda
        /* Entropy regularization — scale derived from branch_scales (no hardcoded per-branch weights) */
        if (entropy_coeff > 0.0f) {
            float lp_clamped = fmaxf(lp, -10.0f);
            d_combined += entropy_coeff * (1.0f + lp_clamped);
        }
```

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

Run: `SQLX_OFFLINE=true cargo build -p ml 2>&1 | head -5`
Expected: cubin compiles.

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
git commit -m "feat(c51): dynamic per-branch gradient scaling from IQL advantage decomposition"
```

---

### Task 4: Modify `epsilon_greedy_kernel.cu` for Per-Sample Epsilon

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu`

- [ ] **Step 1: Replace 3 scalar epsilon params with per-sample buffer**

In the `branching_action_select` signature (line 191), replace:
```cuda
    const float epsilon_exp,
    const float epsilon_ord,
    const float epsilon_urg,
```
with:
```cuda
    const float* __restrict__ per_sample_epsilon, /* [B] state-dependent epsilon */
```

- [ ] **Step 2: Update epsilon reads in kernel body**

Replace lines 210-212:
```cuda
    float eps_exp_bf = bf16(epsilon_exp);
    float eps_ord_bf = bf16(epsilon_ord);
    float eps_urg_bf = bf16(epsilon_urg);
```
with:
```cuda
    float eps = per_sample_epsilon[idx];
```

Then replace all 3 uses of the old per-branch epsilons:
- Line 217: `if (r1 < eps_exp_bf)` → `if (r1 < eps)`
- Line 247: `if (r2 < eps_ord_bf)` → `if (r2 < eps)`
- Line 282: `if (r3 < eps_urg_bf)` → `if (r3 < eps)`

- [ ] **Step 3: Verify build**

Run: `SQLX_OFFLINE=true cargo build -p ml 2>&1 | head -5`

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu
git commit -m "feat(explore): per-sample epsilon from IQL expectile gap"
```

---

### Task 5: Rust — `GpuIqlTrainer` New Buffers and Methods

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

- [ ] **Step 1: Add new fields to `GpuIqlTrainer` struct**

After the existing `advantage_weights_buf` field, add:

```rust
    // ── PER modulation buffers ──────────────────────────────────────
    adv_stats_buf: CudaSlice<f32>,       // [2] mean, variance of adv weights
    adv_sigma_ema: f32,                  // host-side EMA of advantage std

    // ── Per-sample C51 support ──────────────────────────────────────
    per_sample_support_buf: CudaSlice<f32>, // [B, 3] v_min, v_max, delta_z

    // ── Per-branch advantage ────────────────────────────────────────
    branch_scales_buf: CudaSlice<f32>,   // [B, 4] gradient scale per branch

    // ── Expectile gap exploration ───────────────────────────────────
    expectile_gap_buf: CudaSlice<f32>,   // [B] V_high - V_low
    gap_mean_buf: CudaSlice<f32>,        // [1] batch mean of gap
    per_sample_epsilon_buf: CudaSlice<f32>, // [B] state-dependent epsilon
```

- [ ] **Step 2: Add new kernel fields**

After `advantage_weight_kernel`:

```rust
    modulate_td_kernel: CudaFunction,
    adv_variance_kernel: CudaFunction,
    per_sample_support_kernel: CudaFunction,
    branch_advantage_kernel: CudaFunction,
    expectile_gap_kernel: CudaFunction,
    gap_mean_kernel: CudaFunction,
    per_sample_epsilon_kernel: CudaFunction,
```

- [ ] **Step 3: Update `IqlKernels` struct and `compile_iql_kernels`**

Add the 7 new kernel fields to `IqlKernels` and load them from the cubin:

```rust
struct IqlKernels {
    // ... existing 10 fields ...
    modulate_td: CudaFunction,
    adv_variance: CudaFunction,
    per_sample_support: CudaFunction,
    branch_advantage: CudaFunction,
    expectile_gap: CudaFunction,
    gap_mean: CudaFunction,
    per_sample_epsilon: CudaFunction,
}
```

In `compile_iql_kernels`, add:
```rust
    modulate_td: load("iql_modulate_td_errors")?,
    adv_variance: load("iql_adv_variance_reduce")?,
    per_sample_support: load("iql_compute_per_sample_support")?,
    branch_advantage: load("iql_per_branch_advantage")?,
    expectile_gap: load("iql_expectile_gap")?,
    gap_mean: load("iql_gap_mean_reduce")?,
    per_sample_epsilon: load("iql_compute_per_sample_epsilon")?,
```

- [ ] **Step 4: Allocate new buffers in `GpuIqlTrainer::new`**

After the existing `advantage_weights_buf` allocation:

```rust
    let adv_stats_buf = alloc_f32(&stream, 2, "iql_adv_stats")?;
    let per_sample_support_buf = alloc_f32(&stream, b * 3, "iql_per_sample_support")?;
    let branch_scales_buf = alloc_f32(&stream, b * 4, "iql_branch_scales")?;
    let expectile_gap_buf = alloc_f32(&stream, b, "iql_expectile_gap")?;
    let gap_mean_buf = alloc_f32(&stream, 1, "iql_gap_mean")?;
    let per_sample_epsilon_buf = alloc_f32(&stream, b, "iql_per_sample_epsilon")?;
```

Seed `per_sample_support_buf` with safe defaults for step 0:

```rust
    // Seed per-sample support with [-1, 1] default for step 0 (before first IQL forward)
    let default_support: Vec<f32> = (0..b).flat_map(|_| vec![-1.0_f32, 1.0, 2.0 / (config.batch_size.max(2) as f32 - 1.0)]).collect();
    // Actually delta_z should use num_atoms, not batch_size. We pass num_atoms via a new config field.
```

Wait — `GpuIqlConfig` doesn't have `num_atoms`. We need to add it. Add `pub num_atoms: usize` to `GpuIqlConfig` with default 51, and use it:

```rust
    let na = config.num_atoms.max(2) as f32;
    let default_delta_z = 2.0 / (na - 1.0);
    let mut default_support = vec![0.0_f32; b * 3];
    for i in 0..b {
        default_support[i * 3]     = -1.0;  // v_min
        default_support[i * 3 + 1] =  1.0;  // v_max
        default_support[i * 3 + 2] = default_delta_z;
    }
    super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?;
```

- [ ] **Step 5: Add `num_atoms` and `total_actions` to `GpuIqlConfig`**

```rust
pub struct GpuIqlConfig {
    // ... existing fields ...
    /// Number of C51 atoms.
    pub num_atoms: usize,
    /// Total factored actions (b0*b1*b2*b3, e.g. 81).
    pub total_actions: usize,
    /// Branch sizes for advantage decomposition.
    pub branch_sizes: [usize; 4],
    /// Discount factor for Bellman headroom in per-sample support.
    pub gamma: f32,
    /// Staleness decay rate for PER modulation.
    pub staleness_lambda: f32,
    /// Staleness age normalizer (steps).
    pub staleness_tau: f32,
}
```

Update `Default`:
```rust
    num_atoms: 51,
    total_actions: 81,
    branch_sizes: [3, 3, 3, 3],
    gamma: 0.99,
    staleness_lambda: 0.001,
    staleness_tau: 10000.0,
```

- [ ] **Step 6: Add launcher methods**

Add these methods to `impl GpuIqlTrainer`:

```rust
    /// Modulate td_errors in-place with advantage weights and staleness decay.
    pub fn modulate_td_errors(
        &mut self,
        td_errors: &mut CudaSlice<f32>,
        indices: &CudaSlice<i32>,
        write_pos: i32,
        capacity: i32,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let batch_i32 = b as i32;
        let beta = self.config.advantage_temperature;
        let sigma = self.adv_sigma_ema;
        let lambda = self.config.staleness_lambda;
        let tau = self.config.staleness_tau;
        let blocks = (b + 255) / 256;

        unsafe {
            self.stream
                .launch_builder(&self.modulate_td_kernel)
                .arg(td_errors)
                .arg(&self.advantage_weights_buf)
                .arg(indices)
                .arg(&sigma)
                .arg(&beta)
                .arg(&lambda)
                .arg(&tau)
                .arg(&write_pos)
                .arg(&capacity)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL modulate_td_errors: {e}")))?;
        }
        Ok(())
    }

    /// Update the EMA of advantage standard deviation.
    pub fn update_adv_sigma(&mut self) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let batch_i32 = b as i32;

        unsafe {
            self.stream
                .launch_builder(&self.adv_variance_kernel)
                .arg(&self.advantage_weights_buf)
                .arg(&mut self.adv_stats_buf)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL adv_variance_reduce: {e}")))?;
        }

        // Read back variance and update EMA on host
        let mut stats = [0.0_f32; 2];
        self.stream.memcpy_dtoh(&self.adv_stats_buf, &mut stats)
            .map_err(|e| MLError::ModelError(format!("IQL adv_stats DtoH: {e}")))?;
        let sigma = stats[1].max(0.0).sqrt();
        const EMA_BETA: f32 = 0.99;
        if self.adv_sigma_ema < 1e-8 {
            self.adv_sigma_ema = sigma;
        } else {
            self.adv_sigma_ema = EMA_BETA * self.adv_sigma_ema + (1.0 - EMA_BETA) * sigma;
        }
        Ok(())
    }

    /// Compute per-sample C51 atom support centered on V(s).
    pub fn compute_per_sample_support(
        &mut self,
        q_out_buf: &CudaSlice<f32>,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let batch_i32 = b as i32;
        let total_actions_i32 = self.config.total_actions as i32;
        let num_atoms_i32 = self.config.num_atoms as i32;
        let gamma = self.config.gamma;
        let blocks = (b + 255) / 256;

        unsafe {
            self.stream
                .launch_builder(&self.per_sample_support_kernel)
                .arg(&self.v_out_buf)
                .arg(q_out_buf)
                .arg(&mut self.per_sample_support_buf)
                .arg(&gamma)
                .arg(&batch_i32)
                .arg(&total_actions_i32)
                .arg(&num_atoms_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL per_sample_support: {e}")))?;
        }
        Ok(())
    }

    /// Compute per-branch gradient scales from Q-value marginalization.
    pub fn compute_branch_scales(
        &mut self,
        q_out_buf: &CudaSlice<f32>,
        actions_buf: &CudaSlice<i32>,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let batch_i32 = b as i32;
        let ta = self.config.total_actions as i32;
        let bs = self.config.branch_sizes;
        let b0 = bs[0] as i32;
        let b1 = bs[1] as i32;
        let b2 = bs[2] as i32;
        let b3 = bs[3] as i32;
        let blocks = (b + 255) / 256;

        unsafe {
            self.stream
                .launch_builder(&self.branch_advantage_kernel)
                .arg(q_out_buf)
                .arg(&self.v_out_buf)
                .arg(actions_buf)
                .arg(&mut self.branch_scales_buf)
                .arg(&batch_i32)
                .arg(&ta)
                .arg(&b0)
                .arg(&b1)
                .arg(&b2)
                .arg(&b3)
                .launch(LaunchConfig {
                    grid_dim: (blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL branch_advantage: {e}")))?;
        }
        Ok(())
    }

    /// Compute expectile gap between two V(s) estimates.
    pub fn compute_expectile_gap(
        &mut self,
        v_low_ptr: u64,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let batch_i32 = b as i32;
        let blocks = (b + 255) / 256;

        // Phase 1: element-wise gap
        unsafe {
            self.stream
                .launch_builder(&self.expectile_gap_kernel)
                .arg(&self.v_out_buf)     // v_high (this trainer, tau=0.7)
                .arg(&v_low_ptr)          // v_low (other trainer, tau=0.3)
                .arg(&mut self.expectile_gap_buf)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL expectile_gap: {e}")))?;
        }

        // Phase 2: sequential mean
        unsafe {
            self.stream
                .launch_builder(&self.gap_mean_kernel)
                .arg(&self.expectile_gap_buf)
                .arg(&mut self.gap_mean_buf)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL gap_mean_reduce: {e}")))?;
        }
        Ok(())
    }

    /// Compute per-sample epsilon from expectile gap.
    pub fn compute_per_sample_epsilon(
        &mut self,
        base_epsilon: f32,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let batch_i32 = b as i32;
        let blocks = (b + 255) / 256;

        unsafe {
            self.stream
                .launch_builder(&self.per_sample_epsilon_kernel)
                .arg(&self.expectile_gap_buf)
                .arg(&self.gap_mean_buf)
                .arg(&mut self.per_sample_epsilon_buf)
                .arg(&base_epsilon)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQL per_sample_epsilon: {e}")))?;
        }
        Ok(())
    }

    /// Raw pointer to per-sample support buffer [B, 3].
    pub fn per_sample_support_ptr(&self) -> u64 {
        self.per_sample_support_buf.raw_ptr()
    }

    /// Raw pointer to branch scales buffer [B, 4].
    pub fn branch_scales_ptr(&self) -> u64 {
        self.branch_scales_buf.raw_ptr()
    }

    /// Raw pointer to per-sample epsilon buffer [B].
    pub fn per_sample_epsilon_ptr(&self) -> u64 {
        self.per_sample_epsilon_buf.raw_ptr()
    }

    /// Reference to per-sample epsilon buffer.
    pub fn per_sample_epsilon_buf(&self) -> &CudaSlice<f32> {
        &self.per_sample_epsilon_buf
    }
```

- [ ] **Step 7: Update struct initialization in `new`**

Wire all new fields into the `Ok(Self { ... })` return.

- [ ] **Step 8: Verify build**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5`
Expected: compiles with possible unused warnings (callers not updated yet).

- [ ] **Step 9: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
git commit -m "feat(iql): new buffers and launcher methods for full integration"
```

---

### Task 6: Rust — Update `GpuDqnTrainer` (C51 + Grad Kernel Launches + Dead Code)

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

- [ ] **Step 1: Add `per_sample_support_ptr` and `branch_scales_ptr` fields**

Add to `GpuDqnTrainer` struct:
```rust
    /// Per-sample C51 support [B, 3] — pointer set by IQL trainer.
    per_sample_support_ptr: u64,
    /// Per-branch gradient scales [B, 4] — pointer set by IQL trainer.
    branch_scales_ptr: u64,
```

Initialize both to 0 in `GpuDqnTrainer::new`.

Add setters:
```rust
    pub fn set_per_sample_support_ptr(&mut self, ptr: u64) {
        self.per_sample_support_ptr = ptr;
    }
    pub fn set_branch_scales_ptr(&mut self, ptr: u64) {
        self.branch_scales_ptr = ptr;
    }
```

- [ ] **Step 2: Update `launch_c51_loss` — replace `v_range_dev_ptr` with `per_sample_support_ptr`**

In `launch_c51_loss` (line ~5383), replace:
```rust
                .arg(&self.v_range_dev_ptr)
```
with:
```rust
                .arg(&self.per_sample_support_ptr)
```

- [ ] **Step 3: Update `launch_c51_grad` — add `branch_scales_ptr` parameter**

In `launch_c51_grad` (line ~5508), after `.arg(&entropy_coeff)`, add:
```rust
                .arg(&self.branch_scales_ptr)
```

- [ ] **Step 4: Delete v_range infrastructure**

Delete the following from `GpuDqnTrainer`:
- Fields: `v_range_pinned`, `v_range_dev_ptr`, `reward_std_ema`
- Methods: `adapt_v_range_full`, `adapt_v_range`, `adapt_v_range_with_gap`, `v_range`, `v_range_buf_ptr`
- The v_range allocation code in `GpuDqnTrainer::new`

Update all other `v_range_dev_ptr` usages in the file (non-graph C51 launches at lines ~1728, ~4171, ~4278, ~4430) to use `self.per_sample_support_ptr`.

- [ ] **Step 5: Verify build**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10`
Expected: errors about missing callers in fused_training.rs and training_loop.rs (expected — fixed in Task 7).

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat(c51): per-sample support + branch scales in kernel launches, delete v_range"
```

---

### Task 7: Rust — Wire Everything in `fused_training.rs`

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

- [ ] **Step 1: Make IQL mandatory — remove `Option<GpuIqlTrainer>`**

Change field:
```rust
    pub(crate) gpu_iql: Option<GpuIqlTrainer>,
```
to:
```rust
    pub(crate) gpu_iql: GpuIqlTrainer,
    pub(crate) gpu_iql_low: GpuIqlTrainer,
```

- [ ] **Step 2: Update `FusedTrainingCtx::new` — always create both IQL trainers**

Replace the `if hyperparams.use_iql { ... }` block with unconditional initialization:

```rust
        let iql_config = GpuIqlConfig {
            expectile_tau: hyperparams.iql_expectile_tau,
            advantage_temperature: hyperparams.iql_advantage_temperature,
            value_hidden_dim: 128,
            state_dim: dqn.config.state_dim,
            batch_size,
            lr: hyperparams.learning_rate as f32,
            max_grad_norm: resolved_grad_norm as f32,
            num_atoms: dqn.config.num_atoms,
            total_actions: dqn.config.branch_0_size * dqn.config.branch_1_size
                * dqn.config.branch_2_size * dqn.config.branch_3_size,
            branch_sizes: [
                dqn.config.branch_0_size,
                dqn.config.branch_1_size,
                dqn.config.branch_2_size,
                dqn.config.branch_3_size,
            ],
            gamma: dqn.config.gamma as f32,
            staleness_lambda: hyperparams.iql_staleness_lambda,
            staleness_tau: hyperparams.iql_staleness_tau,
            ..GpuIqlConfig::default()
        };

        let gpu_iql = GpuIqlTrainer::new(stream.clone(), iql_config.clone())?;

        let iql_low_config = GpuIqlConfig {
            expectile_tau: 1.0 - hyperparams.iql_expectile_tau,
            ..iql_config
        };
        let gpu_iql_low = GpuIqlTrainer::new(stream.clone(), iql_low_config)?;
```

- [ ] **Step 3: Set per_sample_support_ptr and branch_scales_ptr on trainer**

After creating the `GpuDqnTrainer` (the `trainer` variable), set the IQL buffer pointers:

```rust
        trainer.set_per_sample_support_ptr(gpu_iql.per_sample_support_ptr());
        trainer.set_branch_scales_ptr(gpu_iql.branch_scales_ptr());
```

- [ ] **Step 4: Update `submit_aux_ops` — full IQL pipeline**

Replace the existing IQL block (`if let Some(ref mut iql) = self.gpu_iql { ... }`) with the full pipeline:

```rust
        // ── IQL pipeline: gather → train both → advantage weights → support → branch scales → gap → epsilon → PER modulation ──
        {
            // 1. Populate q_out_buf from current logits.
            self.trainer.populate_q_out(self.batch_size)
                .map_err(|e| anyhow::anyhow!("IQL populate_q_out: {e}"))?;

            // 2. Gather Q(s, a_taken) for both IQL trainers.
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            let total_actions = self.trainer.total_actions();
            self.gpu_iql.gather_q_taken(q_out, actions, total_actions)
                .map_err(|e| anyhow::anyhow!("IQL high gather_q_taken: {e}"))?;
            self.gpu_iql_low.gather_q_taken(q_out, actions, total_actions)
                .map_err(|e| anyhow::anyhow!("IQL low gather_q_taken: {e}"))?;

            // 3. Train V(s) — both tau=0.7 and tau=0.3.
            let states_f32 = self.trainer.states_buf();
            self.gpu_iql.train_value_step(states_f32)
                .map_err(|e| anyhow::anyhow!("IQL high train: {e}"))?;
            let states_f32 = self.trainer.states_buf();
            self.gpu_iql_low.train_value_step(states_f32)
                .map_err(|e| anyhow::anyhow!("IQL low train: {e}"))?;

            // 4. Compute advantage weights (from high-tau V(s)).
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            self.gpu_iql.compute_advantage_weights(q_out, actions, total_actions)
                .map_err(|e| anyhow::anyhow!("IQL advantage_weights: {e}"))?;

            // 5. Update advantage sigma EMA (for dynamic PER clamp bounds).
            self.gpu_iql.update_adv_sigma()
                .map_err(|e| anyhow::anyhow!("IQL adv_sigma: {e}"))?;

            // 6. Compute per-sample C51 support.
            let q_out = self.trainer.q_out_buf();
            self.gpu_iql.compute_per_sample_support(q_out)
                .map_err(|e| anyhow::anyhow!("IQL per_sample_support: {e}"))?;

            // 7. Compute per-branch gradient scales.
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            self.gpu_iql.compute_branch_scales(q_out, actions)
                .map_err(|e| anyhow::anyhow!("IQL branch_scales: {e}"))?;

            // 8. Compute expectile gap and per-sample epsilon.
            let v_low_ptr = self.gpu_iql_low.v_out_ptr();
            self.gpu_iql.compute_expectile_gap(v_low_ptr)
                .map_err(|e| anyhow::anyhow!("IQL expectile_gap: {e}"))?;

            let base_epsilon = agent.primary_dqn().config.epsilon_end as f32;
            self.gpu_iql.compute_per_sample_epsilon(base_epsilon)
                .map_err(|e| anyhow::anyhow!("IQL per_sample_epsilon: {e}"))?;
        }
```

- [ ] **Step 5: Wire PER modulation — modulate td_errors before seg_tree update**

In `run_full_step`, before the `agent.update_priorities_from_td(...)` call (line ~1039), add:

```rust
        // Modulate td_errors with IQL advantage weights + staleness decay.
        {
            let indices = &gpu_batch.indices_ptr;
            // Get write_pos and capacity from the replay buffer
            let write_pos = agent.memory().gpu_write_cursor() as i32;
            let capacity = agent.memory().gpu_capacity() as i32;
            self.gpu_iql.modulate_td_errors(
                self.trainer.td_errors_buf_mut(),
                // indices is a raw pointer — need to pass appropriately
                // This requires adding a method to get the CudaSlice reference
            ).map_err(|e| anyhow::anyhow!("IQL modulate_td: {e}"))?;
        }
```

Note: The exact integration depends on how `indices` is surfaced from `GpuBatch`. The `GpuBatch` has `indices_ptr: u64` (raw pointer). The `modulate_td_errors` method needs a `&CudaSlice<i32>` or raw pointer. Adapt the method signature to take `indices_ptr: u64` as a raw pointer, or use the batch's slice directly.

- [ ] **Step 6: Delete v_range wrappers**

Delete from `FusedTrainingCtx`:
- `trainer_v_range_buf_ptr`
- `adapt_v_range`, `adapt_v_range_with_gap`, `adapt_v_range_full`
- `v_range`

- [ ] **Step 7: Update `pre_replay_state_update`**

Change `if let Some(ref mut iql)` to direct access:
```rust
        self.gpu_iql.increment_adam_step();
        self.gpu_iql_low.increment_adam_step();
```

- [ ] **Step 8: Verify build**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -20`

- [ ] **Step 9: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat(iql): full pipeline wired — PER modulation, per-sample support, branch scales, gap exploration"
```

---

### Task 8: Rust — Update Config + Training Loop + Experience Collector

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs`
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_action_selector.rs`

- [ ] **Step 1: Add config params to `DQNHyperparameters`**

Add after `use_iql`:
```rust
    pub iql_staleness_lambda: f32,
    pub iql_staleness_tau: f32,
```

Set defaults:
```rust
    iql_staleness_lambda: 0.001,
    iql_staleness_tau: 10000.0,
```

Remove `use_iql` field entirely.

- [ ] **Step 2: Update `training_loop.rs` — remove v_range calls**

Delete all calls to:
- `fused_ctx.adapt_v_range(...)`
- `fused_ctx.adapt_v_range_with_gap(...)`
- `fused_ctx.adapt_v_range_full(...)`
- `collector.set_v_range_ptr(...)`

These are dead code with per-sample support.

- [ ] **Step 3: Update `gpu_experience_collector.rs` — remove v_range_ptr**

Delete the `v_range_ptr: u64` field and `set_v_range_ptr` method. For the experience collector's `compute_expected_q` calls that still need a v_range: use a fixed `[-100, 100]` range (experience collection only needs Q-values for action selection argmax — atom support precision doesn't matter for argmax). Replace:
```rust
    .arg(&self.v_range_ptr)
```
with a local constant:
```rust
    let exp_v_range = [-100.0_f32, 100.0_f32];
    // upload to a small device buffer or use pinned memory
```

- [ ] **Step 4: Update `gpu_action_selector.rs` — pass per-sample epsilon buffer**

Update the `branching_action_select` launcher to pass a `per_sample_epsilon_ptr: u64` instead of 3 scalar epsilons. The `GpuActionSelector::select_branching` method signature changes from:
```rust
    epsilon_exp: f32, epsilon_ord: f32, epsilon_urg: f32
```
to:
```rust
    per_sample_epsilon_ptr: u64
```

Update the kernel launch accordingly.

- [ ] **Step 5: Verify full build**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10`
Expected: clean compile.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/trainers/dqn/config.rs \
       crates/ml/src/trainers/dqn/trainer/training_loop.rs \
       crates/ml/src/cuda_pipeline/gpu_experience_collector.rs \
       crates/ml/src/cuda_pipeline/gpu_action_selector.rs
git commit -m "feat(iql): config + training loop + experience collector integration"
```

---

### Task 9: Full Integration Test

**Files:**
- No new files — run existing tests.

- [ ] **Step 1: Full workspace check**

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

- [ ] **Step 2: Smoke test — graph capture + NaN regression**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_no_nan_after_graph_capture --ignored --nocapture 2>&1 | tail -20`
Expected: `test result: ok. 1 passed`

- [ ] **Step 3: Stability test — 10 epochs**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_production_training_stability --ignored --nocapture 2>&1 | tail -30`
Expected: `test result: ok. 1 passed`, loss converges, no NaN/panic.

- [ ] **Step 4: Determinism test**

Run the same test twice and compare EPOCH_DIAG:
```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_walk_forward_oos_metrics --ignored --nocapture 2>&1 | grep "EPOCH_DIAG" > /tmp/run_a.txt
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_walk_forward_oos_metrics --ignored --nocapture 2>&1 | grep "EPOCH_DIAG" > /tmp/run_b.txt
diff /tmp/run_a.txt /tmp/run_b.txt
```
Expected: within-process determinism maintained (no new atomicAdd sources).

- [ ] **Step 5: Run full smoke test suite**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -5`
Expected: all smoke tests pass.

- [ ] **Step 6: Commit integration verification**

```bash
git commit --allow-empty -m "test: IQL full integration verified — smoke + stability + determinism"
```
