diff --git a/docs/superpowers/specs/2026-04-13-iql-full-integration-design.md b/docs/superpowers/specs/2026-04-13-iql-full-integration-design.md new file mode 100644 index 000000000..d847829f0 --- /dev/null +++ b/docs/superpowers/specs/2026-04-13-iql-full-integration-design.md @@ -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` 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` — 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`.