# IQL Full Integration Design

## Goal

Wire IQL's V(s) output into every component that can benefit from it: PER priorities, C51 atom support, gradient scaling, and exploration. Eliminate all hardcoded constants — everything derives from the data.

## Architecture

IQL trains V(s) via expectile regression on Q(s, a_taken). A second IQL instance at τ=0.3 provides the expectile gap for uncertainty estimation. The five integration points are:

1. **PER priority modulation with staleness correction** — advantage-weighted replay with age-decay
2. **Per-sample C51 atom support** — each sample gets `[v_min, v_max, delta_z]` centered on V(s)
3. **Per-branch advantage decomposition** — dynamic gradient scaling per action branch
4. **Expectile gap exploration** — state-dependent epsilon from dual-tau uncertainty
5. **Dead code cleanup** — remove batch-level v_range adaptation infrastructure

## Component 1: PER Priority Modulation + Staleness

### Kernel: `iql_modulate_td_errors`

Located in `iql_value_kernel.cu`. Modifies `td_errors_buf [B]` in-place before the seg_tree update.

```
td_errors[b] *= clamp(adv_weights[b], 1/K, K) * exp(-λ * age[b] / τ_stale)
```

- `K = exp(β * 3 * σ_adv)` — dynamic clamp derived from advantage temperature and observed advantage standard deviation
- `σ_adv` — running EMA of advantage standard deviation, updated each step from the batch
- `age[b] = current_step - step_collected[b]` — transition age in training steps
- `λ` — staleness decay rate, config param `iql_staleness_lambda` (default 0.001)
- `τ_stale` — age normalizer, config param `iql_staleness_tau` (default 10000.0)

### Advantage Variance Tracking

A reduce kernel `iql_adv_variance_reduce` computes `mean(adv_weights)` and `var(adv_weights)` from the batch. Output goes to a pinned device-mapped `σ_adv_ema` scalar (host-side EMA update, graph-safe).

### Data Requirements

- `age[b]` is derived from the replay buffer's circular index: `age[b] = (buffer_write_pos - indices[b]) % buffer_capacity`. Higher distance = older transition. `indices [B]` comes from `GpuBatch::indices_ptr`. `buffer_write_pos` is `GpuReplayBuffer::write_cursor` (host-side `usize`) — passed to the kernel as a scalar `i32` parameter. `buffer_capacity` is also a scalar from `GpuReplayBuffer::capacity()`. No replay buffer schema change needed.

### Config Additions

```rust
// DQNHyperparameters
pub iql_staleness_lambda: f32,  // default 0.001
pub iql_staleness_tau: f32,     // default 10000.0
```

### Files Touched

- `iql_value_kernel.cu` — new kernels: `iql_modulate_td_errors`, `iql_adv_variance_reduce`
- `gpu_iql_trainer.rs` — new buffers: `adv_sigma_ema` (pinned scalar), new methods: `modulate_td_errors`, `update_adv_sigma`
- `fused_training.rs` — call `modulate_td_errors` between IQL step and PER update
- `config.rs` — two new hyperparams

## Component 2: Per-Sample C51 Atom Support

### Design

Each sample gets its own distributional support centered on V(s), with width derived from the Q-value spread around V(s):

```
q_spread(b) = max(|Q_max(s_b) - V(s_b)|, |Q_min(s_b) - V(s_b)|)
half_width(b) = q_spread(b) * (1 + gamma)
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)
```

The `(1 + gamma)` headroom factor accounts for the Bellman backup shifting the support.

When `delta_z(b) < 1e-7` (all Q-values identical for this sample — point mass distribution), the sample's loss contribution is skipped. This is the only guard — derived from f32 machine epsilon, not a tuning knob.

### Precomputation Kernel: `iql_compute_per_sample_support`

- Reads: `v_out_buf [B]`, `q_out_buf [B, total_actions]`, `gamma` (scalar)
- Writes: `per_sample_support [B, 3]` — packed `[v_min, v_max, delta_z]` per sample
- Launch: `grid=ceil(B/256), block=256` — one thread per sample, each thread scans `total_actions` Q-values
- Lives in `iql_value_kernel.cu`

### C51 Loss Kernel Changes (`c51_loss_batched`)

- Replace parameter `const float* __restrict__ v_range_buf` (`[2]`) with `const float* __restrict__ per_sample_support` (`[B, 3]`)
- Remove shared `v_min/v_max` read from `v_range_buf`
- Per-sample register loads:
  ```cuda
  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];
  ```
- Guard: `if (delta_z < 1e-7f) { per_sample_loss[sample_id] = 0.0f; td_errors[sample_id] = 0.0f; return; }`
- `shmem_support` array computed in-kernel: `shmem_support[j] = v_min + j * delta_z`
- Bellman projection uses per-sample `v_min/v_max/delta_z` — same math, per-sample values

### Dead Code Removal

With per-sample support, the following become dead code and are deleted:

**`GpuDqnTrainer` (gpu_dqn_trainer.rs):**
- `adapt_v_range_full`, `adapt_v_range`, `adapt_v_range_with_gap` methods
- `v_range` (the `[f32; 2]` getter), `v_range_buf_ptr` accessor
- `v_range_pinned` field and its allocation
- `v_range_dev_ptr` field
- `reward_std_ema` field
- The `v_range_buf` allocation in `GpuDqnTrainer::new`
- All non-graph `c51_loss_batched` launches that pass `v_range_dev_ptr` (lines ~1728, ~4171, ~4278, ~4430) — update to pass `per_sample_support_buf` pointer

**`GpuExperienceCollector` (gpu_experience_collector.rs):**
- `v_range_ptr: u64` field
- `set_v_range_ptr` method
- All `compute_expected_q` launches that pass `v_range_ptr` — these need `per_sample_support` or a simplified v_range for experience collection (experience collector computes Q-values for action selection, not C51 loss — it can use a batch-level estimate from the IQL trainer's mean V(s))

**`FusedTrainingCtx` (fused_training.rs):**
- `trainer_v_range_buf_ptr` wrapper
- `adapt_v_range`, `adapt_v_range_with_gap`, `adapt_v_range_full` wrappers
- `v_range` wrapper

**`training_loop.rs`:**
- All `adapt_v_range*` calls and Q-stat plumbing
- `collector.set_v_range_ptr(fused_ctx.trainer_v_range_buf_ptr())` calls

### Files Touched

- `iql_value_kernel.cu` — new kernel: `iql_compute_per_sample_support`
- `c51_loss_kernel.cu` — change `v_range_buf [2]` to `per_sample_support [B, 3]`, per-sample support logic
- `gpu_dqn_trainer.rs` — delete v_range infrastructure, add `per_sample_support_buf [B, 3]` allocation, update C51 launch to pass new buffer
- `gpu_iql_trainer.rs` — new method: `compute_per_sample_support`
- `fused_training.rs` — call `compute_per_sample_support` before graph_forward, remove v_range adaptation calls
- `training_loop.rs` — remove `adapt_v_range*` calls and Q-stat plumbing for v_range

### Note on Graph Capture and One-Step Lag

`per_sample_support [B, 3]` depends on `v_out_buf` (from IQL forward) and `q_out_buf` (from `populate_q_out`), which both require graph_forward to have already run. Therefore the C51 kernel inside graph_forward uses `per_sample_support` from the PREVIOUS step (one-step lag). This is the same pattern as the target network — the support adapts with a one-step delay, which provides natural smoothing.

**Step 0 initialization:** Before the first graph_forward replay, `per_sample_support` must contain valid values. At construction time, `GpuIqlTrainer` seeds the buffer with a uniform default: `v_min = -1.0, v_max = 1.0, delta_z = 2.0 / (num_atoms - 1)` for all B samples. This is a safe initial support — Xavier V(s) and initial Q-values are near zero, so the `[-1, 1]` range covers them. After step 0 completes, `compute_per_sample_support` overwrites with data-driven values and the seed is never used again.

The buffer pointer is stable (pre-allocated `CudaSlice`), so the CUDA Graph captures the pointer and the per-step content changes are picked up on replay. Same pattern as `states_buf`.

## Component 3: Per-Branch Advantage Decomposition

### Kernel: `iql_per_branch_advantage`

Computes per-branch advantage magnitudes by marginalizing Q over the other 3 branches:

```
For each sample b, for each branch d (0..3):
  Q_marginal(b, d, a_d) = mean over other 3 branches of Q(s_b, a) where a has branch d = a_d
  A_branch(b, d) = |Q_marginal(b, d, a_taken_d) - V(s_b)|

branch_scale(b, d) = A_branch(b, d) / max_d(A_branch(b, d))
```

- Reads: `q_out_buf [B, total_actions]`, `v_out_buf [B]`, `actions_buf [B]`
- Writes: `branch_scales_buf [B, 4]`
- Launch: `grid=B, block=256` — one block per sample, threads cooperate on marginalization

For 4 branches of size 3 each (81 total actions), each branch marginalization averages over 27 actions. The block's 256 threads handle this cooperatively — 64 threads per branch, each thread handles ceil(27/64) actions.

The `max_d` normalization ensures the highest-advantage branch gets scale 1.0, others get proportionally less. No hardcoded zeros — if magnitude genuinely matters for a state, it gets gradient.

### C51 Grad Kernel Changes (`c51_grad_kernel`)

- New parameter: `const float* __restrict__ branch_scales` — `[B, 4]`
- Line 72: delete `float branch_scale = (d == 1) ? 0.0f : 1.0f;`
- Replace with: `float branch_scale = branch_scales[b * 4 + d];`
- **Entropy interaction:** The entropy regularization term (lines 60-67) uses per-branch `ent_scale` (d==1 gets 5.0f, d==2 gets 3.0f) and accumulates into `d_val_sum` BEFORE `branch_scale` is applied. With dynamic `branch_scale`, the entropy gradient for magnitude is no longer silently zeroed. The `ent_scale` multipliers become redundant — replace with `1.0f` for all branches since `branch_scales` already captures per-branch importance. Delete the `ent_scale` conditional block entirely.

### Files Touched

- `iql_value_kernel.cu` — new kernel: `iql_per_branch_advantage`
- `c51_grad_kernel.cu` — replace hardcoded branch_scale with per-sample buffer read
- `gpu_iql_trainer.rs` — new buffer: `branch_scales_buf [B, 4]`, new method: `compute_branch_scales`
- `gpu_dqn_trainer.rs` — update C51 grad kernel launch to pass `branch_scales_buf`
- `fused_training.rs` — call `compute_branch_scales` in `submit_aux_ops`

## Component 4: Expectile Gap Exploration

### Architecture

A second `GpuIqlTrainer` instance at τ=0.3. Both trainers share the same `q_taken_buf` (gathered Q-values), but train separate V(s) networks with different expectile parameters.

The gap `V_0.7(s) - V_0.3(s)` measures advantage distribution spread — high gap = high uncertainty about the best action = explore more.

### Gap Kernel: `iql_expectile_gap`

- Reads: `v_out_high [B]` (τ=0.7 trainer), `v_out_low [B]` (τ=0.3 trainer)
- Writes: `expectile_gap_buf [B]`, `gap_mean_buf [1]` (batch mean)
- Two-phase: element-wise subtraction (`grid=ceil(B/256), block=256`) + sequential mean reduce (`grid=1, block=1`)

### Per-Sample Epsilon Kernel: `iql_compute_per_sample_epsilon`

Converts gap to state-dependent exploration rate:

```
epsilon(b) = base_epsilon * sigmoid(gap[b] / gap_mean - 1.0)
```

- When `gap[b] == gap_mean`: sigmoid(0) = 0.5 → half the base epsilon
- When `gap[b] >> gap_mean`: sigmoid → 1 → full exploration
- When `gap[b] << gap_mean`: sigmoid → 0 → exploit confidently
- `gap_mean` normalizer is relative — no hardcoded thresholds

- Reads: `expectile_gap_buf [B]`, `gap_mean_buf [1]`, `base_epsilon` (scalar)
- Writes: `per_sample_epsilon_buf [B]`
- Launch: `grid=ceil(B/256), block=256`

### Epsilon-Greedy Kernel Changes (`epsilon_greedy_kernel.cu`)

- The `branching_action_select` kernel currently takes three scalar epsilons: `epsilon_exp`, `epsilon_ord`, `epsilon_urg` (one per non-direction branch)
- Replace all three with a single `const float* __restrict__ per_sample_epsilon` — `[B]`
- Each thread reads `per_sample_epsilon[b]` and uses it for all three branch epsilons
- The per-branch epsilon differentiation was a manual heuristic — the expectile gap provides a principled per-state signal that supersedes per-branch tuning

### Experience Collector Changes

- `GpuExperienceCollector` gains a `per_sample_epsilon_buf: CudaSlice<f32>` input
- `FusedTrainingCtx` passes the buffer from the IQL low trainer's gap computation

### Second IQL Trainer Initialization

In `FusedTrainingCtx::new`:
```rust
let iql_low_config = GpuIqlConfig {
    expectile_tau: 1.0 - hyperparams.iql_expectile_tau, // mirror: 0.7 → 0.3
    ..iql_config.clone()
};
let gpu_iql_low = GpuIqlTrainer::new(stream.clone(), iql_low_config)?;
```

The low-tau trainer uses the same architecture, same lr, same batch size. Only the expectile parameter differs. The mirroring `1 - tau` ensures the gap is symmetric around the median.

### Files Touched

- `iql_value_kernel.cu` — new kernels: `iql_expectile_gap`, `iql_compute_per_sample_epsilon`
- `epsilon_greedy_kernel.cu` — per-sample epsilon parameter
- `gpu_iql_trainer.rs` — new buffers: `expectile_gap_buf [B]`, `gap_mean_buf [1]`, `per_sample_epsilon_buf [B]`, new methods
- `gpu_action_selector.rs` — pass per-sample epsilon buffer
- `gpu_experience_collector.rs` — accept and forward per-sample epsilon
- `fused_training.rs` — initialize second IQL trainer, run both, compute gap + epsilon

## Component 5: Dead Code Cleanup

All batch-level v_range infrastructure is removed (listed in Component 2). Additionally:

- Delete the `IQL_EXPECTILE_TAU` compile-time define from the old kernel (now a runtime parameter)
- Delete the `total_loss` parameter from the old `iql_forward_loss_kernel` signature (already done in earlier fix)
- Remove `use_iql` config flag — IQL is mandatory (both τ=0.7 and τ=0.3 always active)
- Remove `Option<GpuIqlTrainer>` — make it a plain `GpuIqlTrainer` field
- Delete `if let Some(ref mut iql)` conditionals — replace with direct field access

## Testing Strategy

All features ship together (big bang). Validation:

1. **Compilation**: `SQLX_OFFLINE=true cargo check --workspace`
2. **Smoke test**: `test_no_nan_after_graph_capture` — verifies graph capture + replay with all new kernels
3. **Stability test**: `test_production_training_stability` — 10-epoch training with all features active
4. **Determinism test**: Two identical runs, compare EPOCH_DIAG values — should be bit-identical within-process (zero atomicAdd maintained)
5. **Metrics monitoring**: Track `atom_entropy`, `atom_utilization` (should improve with per-sample support), action diversity per branch (should shift with dynamic branch scaling), and exploration rate distribution (should vary with expectile gap)

## Summary of New CUDA Kernels

| Kernel | File | Launch | Purpose |
|--------|------|--------|---------|
| `iql_modulate_td_errors` | `iql_value_kernel.cu` | `ceil(B/256), 256` | PER priority × advantage × staleness |
| `iql_adv_variance_reduce` | `iql_value_kernel.cu` | `1, 1` | Batch advantage σ for dynamic clamp |
| `iql_compute_per_sample_support` | `iql_value_kernel.cu` | `ceil(B/256), 256` | V(s)-centered `[v_min, v_max, delta_z]` per sample |
| `iql_per_branch_advantage` | `iql_value_kernel.cu` | `B, 256` | Q-marginal → branch scales |
| `iql_expectile_gap` | `iql_value_kernel.cu` | `ceil(B/256), 256` + `1, 1` | V_high - V_low + batch mean |
| `iql_compute_per_sample_epsilon` | `iql_value_kernel.cu` | `ceil(B/256), 256` | Gap → state-dependent epsilon |

## Summary of New GPU Buffers

| Buffer | Shape | Owner | Purpose |
|--------|-------|-------|---------|
| `per_sample_support` | `[B, 3]` | `GpuIqlTrainer` | Per-sample C51 atom support (pointer passed to `GpuDqnTrainer` for C51 launch) |
| `branch_scales_buf` | `[B, 4]` | `GpuIqlTrainer` | Dynamic per-branch gradient scale |
| `expectile_gap_buf` | `[B]` | `GpuIqlTrainer` | V_high - V_low per sample |
| `gap_mean_buf` | `[1]` | `GpuIqlTrainer` | Batch mean of expectile gap |
| `per_sample_epsilon_buf` | `[B]` | `GpuIqlTrainer` | State-dependent exploration rate |
| `adv_sigma_ema` | scalar (pinned) | `GpuIqlTrainer` | Running σ of advantage weights |

## Config Changes

| Field | Type | Default | Purpose |
|-------|------|---------|---------|
| `iql_staleness_lambda` | `f32` | `0.001` | Age-decay rate for PER staleness |
| `iql_staleness_tau` | `f32` | `10000.0` | Age normalizer (steps) |

`use_iql` removed — IQL is mandatory. `iql_expectile_tau` stays (controls high-tau trainer). Low-tau trainer derives from `1 - iql_expectile_tau`.
