docs: IQL full integration design spec — 5 components
Covers: PER advantage-weighted staleness, per-sample C51 atom support centered on V(s), per-branch advantage decomposition, dual-tau expectile gap exploration, and v_range dead code removal. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
274
docs/superpowers/specs/2026-04-13-iql-full-integration-design.md
Normal file
274
docs/superpowers/specs/2026-04-13-iql-full-integration-design.md
Normal file
@@ -0,0 +1,274 @@
|
||||
# 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. No replay buffer schema change needed — `indices [B]` and `buffer_write_pos` are already available in the GPU batch.
|
||||
|
||||
### 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::adapt_v_range_full`
|
||||
- `GpuDqnTrainer::adapt_v_range`
|
||||
- `GpuDqnTrainer::adapt_v_range_with_gap`
|
||||
- `GpuDqnTrainer::v_range` (the `[f32; 2]` getter)
|
||||
- `GpuDqnTrainer::v_range_buf_ptr`
|
||||
- `v_range_pinned` field and its allocation
|
||||
- `v_range_dev_ptr` field
|
||||
- `reward_std_ema` field
|
||||
- All callers in `training_loop.rs` that pass Q-stats to `adapt_v_range*`
|
||||
- The `v_range_buf` allocation in `GpuDqnTrainer::new`
|
||||
|
||||
### 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. On step 0, `per_sample_support` is initialized from Xavier V(s) output and initial Q-values (both small but non-zero).
|
||||
|
||||
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];`
|
||||
|
||||
### 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 a scalar `epsilon`
|
||||
- Replace with `const float* __restrict__ per_sample_epsilon` — `[B]`
|
||||
- Each thread reads `per_sample_epsilon[b]` instead of the uniform scalar
|
||||
|
||||
### 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`.
|
||||
Reference in New Issue
Block a user