From f0cda76229eebafc33628047eabeb94c5bb9602a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 13 Apr 2026 14:56:27 +0200 Subject: [PATCH] =?UTF-8?q?docs:=20IQL=20full=20integration=20implementati?= =?UTF-8?q?on=20plan=20=E2=80=94=209=20tasks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 6 new CUDA kernels, C51/grad/epsilon kernel mods, dual-tau IQL, v_range dead code deletion, per-sample everything. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-04-13-iql-full-integration.md | 1248 +++++++++++++++++ 1 file changed, 1248 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-13-iql-full-integration.md diff --git a/docs/superpowers/plans/2026-04-13-iql-full-integration.md b/docs/superpowers/plans/2026-04-13-iql-full-integration.md new file mode 100644 index 000000000..28a1d1077 --- /dev/null +++ b/docs/superpowers/plans/2026-04-13-iql-full-integration.md @@ -0,0 +1,1248 @@ +# 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, // [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, // [B, 3] v_min, v_max, delta_z + + // ── Per-branch advantage ──────────────────────────────────────── + branch_scales_buf: CudaSlice, // [B, 4] gradient scale per branch + + // ── Expectile gap exploration ─────────────────────────────────── + expectile_gap_buf: CudaSlice, // [B] V_high - V_low + gap_mean_buf: CudaSlice, // [1] batch mean of gap + per_sample_epsilon_buf: CudaSlice, // [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 = (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, + indices: &CudaSlice, + 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, + ) -> 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, + actions_buf: &CudaSlice, + ) -> 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 { + &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`** + +Change field: +```rust + pub(crate) gpu_iql: Option, +``` +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` 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" +```