feat(rl): IQN complementary Q-head — quantile embedding + Huber loss

Phase 2 of the Q-learning improvements spec. Adds an Implicit
Quantile Network head alongside the existing C51 head:

- rl_iqn_forward.cu: quantile embedding φ(τ) = ReLU(W × cos(iπτ)),
  element-wise h_t ⊙ φ(τ), action-value projection. Plus
  rl_iqn_expected_q for tau-mean reduction.
- rl_iqn_loss.cu: quantile Huber loss ρ_τ(δ) = |τ-1(δ<0)| × Huber(δ).
  Block tree-reduce per batch (no atomicAdd).
- rl/iqn.rs: IqnHead struct with online + target weights, Xavier init,
  forward/forward_target/expected_q/compute_loss methods.

ISV slots: 543 N_TAU (32), 544 ensemble_alpha (0.5), 545 LR (1e-3).

Not yet wired into the trainer — head is constructible and kernels
compile. Ensemble integration is the next step.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/ml-alpha/build.rs

@@ -91,6 +91,8 @@ const KERNELS: &[&str] = &[
     "rl_min_hold_check",             // hard minimum hold time — override close actions to Hold before N steps
     "rl_asymmetric_trail_decay",     // auto-tighten losers, auto-widen winners — structural P&L asymmetry
     "rl_session_risk_check",         // session-level loss limit circuit breaker
+    "rl_iqn_forward",                // IQN distributional Q-head: quantile embedding + action-value projection; complementary to C51
+    "rl_iqn_loss",                   // IQN quantile Huber loss: ρ_τ(δ) = |τ - 1(δ<0)| × Huber(δ, κ=1.0); forward + backward
 ];
 
 // Cache bust v31 — five new reduce / derive kernels populate the input

crates/ml-alpha/cuda/rl_iqn_forward.cu

@@ -0,0 +1,143 @@
+// rl_iqn_forward.cu — Implicit Quantile Network (IQN) forward pass.
+//
+// Complementary distributional Q-head running alongside C51. The IQN
+// head models the full return distribution via learned quantile
+// functions rather than a fixed atom support.
+//
+// Forward pass:
+//   1. Quantile embedding: phi(tau) = ReLU(W_embed × cos(i×π×τ) + b_embed)
+//      where i=1..EMBED_DIM (64). Output: [B, N_TAU, HIDDEN_DIM]
+//   2. Element-wise: combined = h_t ⊙ phi(tau) → [B, N_TAU, HIDDEN_DIM]
+//   3. Action value: Q = W_out × combined + b_out → [B, N_TAU, N_ACTIONS]
+//
+// Expected Q for ensemble action selection: mean over the tau dimension.
+//
+// Block layout:
+//   rl_iqn_forward: grid = (B, N_TAU, 1); block = (HIDDEN_DIM, 1, 1).
+//     One block per (batch, tau) pair. HIDDEN_DIM threads cooperate on
+//     the quantile embedding and then compute action values via a
+//     strided inner loop over N_ACTIONS.
+//
+//   rl_iqn_expected_q: grid = (B, 1, 1); block = (N_ACTIONS, 1, 1).
+//     One block per batch. Each thread (one per action) reduces across
+//     N_TAU quantile samples to compute E[Q(s,a)] = mean_tau Q(s,tau,a).
+//     Uses block tree-reduce pattern (no atomicAdd per
+//     `feedback_no_atomicadd.md`).
+
+#define HIDDEN_DIM   128
+#define N_ACTIONS    11
+#define N_TAU_MAX    64
+#define EMBED_DIM    64
+#define PI_F         3.14159265f
+
+// ─────────────────────────────────────────────────────────────────────
+// rl_iqn_forward:
+//   Compute Q(s, tau, a) for all (batch, tau, action) triples.
+//
+// Inputs:
+//   h_t         [B, HIDDEN_DIM]           — encoder hidden state
+//   tau         [B, N_TAU]                — quantile samples from U(0,1)
+//   w_embed     [EMBED_DIM, HIDDEN_DIM]   — quantile embedding weight
+//   b_embed     [HIDDEN_DIM]              — quantile embedding bias
+//   w_out       [HIDDEN_DIM, N_ACTIONS]   — output projection weight
+//   b_out       [N_ACTIONS]               — output projection bias
+//   B                                     — batch size
+//   N_TAU                                 — number of quantile samples
+// Outputs:
+//   q_values    [B, N_TAU, N_ACTIONS]     — quantile action values
+// ─────────────────────────────────────────────────────────────────────
+extern "C" __global__ void rl_iqn_forward(
+    const float* __restrict__ h_t,         // [B, HIDDEN_DIM]
+    const float* __restrict__ tau,          // [B, N_TAU]
+    const float* __restrict__ w_embed,      // [EMBED_DIM, HIDDEN_DIM]
+    const float* __restrict__ b_embed,      // [HIDDEN_DIM]
+    const float* __restrict__ w_out,        // [HIDDEN_DIM, N_ACTIONS]
+    const float* __restrict__ b_out,        // [N_ACTIONS]
+    int          B,
+    int          N_TAU,
+    float*       __restrict__ q_values      // [B, N_TAU, N_ACTIONS]
+) {
+    const int batch = blockIdx.x;
+    const int tau_idx = blockIdx.y;
+    const int c = threadIdx.x;  // hidden dim index
+    if (batch >= B) return;
+    if (tau_idx >= N_TAU) return;
+    if (c >= HIDDEN_DIM) return;
+
+    // Step 1: Compute quantile embedding phi(tau)[c]
+    // phi(tau) = ReLU( sum_{i=1..EMBED_DIM} cos(i * pi * tau) * W_embed[i, c]  + b_embed[c] )
+    const float tau_val = tau[batch * N_TAU + tau_idx];
+    float embed_acc = b_embed[c];
+    #pragma unroll
+    for (int i = 0; i < EMBED_DIM; ++i) {
+        float cos_feat = cosf((float)(i + 1) * PI_F * tau_val);
+        embed_acc += cos_feat * w_embed[i * HIDDEN_DIM + c];
+    }
+    // ReLU activation
+    float phi_c = fmaxf(embed_acc, 0.0f);
+
+    // Step 2: Element-wise product with encoder hidden state
+    float combined_c = h_t[batch * HIDDEN_DIM + c] * phi_c;
+
+    // Store combined in shared memory for the matmul reduction
+    __shared__ float s_combined[HIDDEN_DIM];
+    s_combined[c] = combined_c;
+    __syncthreads();
+
+    // Step 3: Compute Q(s, tau, a) = W_out^T × combined + b_out
+    // Each thread computes one action's contribution from its hidden dim,
+    // then we need a reduction across hidden dims. Instead, thread c
+    // contributes to all actions via the output projection.
+    // With HIDDEN_DIM threads and N_ACTIONS outputs, each thread iterates
+    // over actions and contributes its portion.
+    //
+    // Output: q_values[batch, tau_idx, a] = sum_c(W_out[c, a] * combined[c]) + b_out[a]
+    // Since HIDDEN_DIM=128 and N_ACTIONS=11, we assign each thread to
+    // compute partial sums and use warp-shuffle reduction.
+    //
+    // Strategy: thread c writes combined[c] to shared mem (done above).
+    // First N_ACTIONS threads each compute one full dot product.
+    if (c < N_ACTIONS) {
+        float q_acc = b_out[c];
+        #pragma unroll
+        for (int k = 0; k < HIDDEN_DIM; ++k) {
+            q_acc += w_out[k * N_ACTIONS + c] * s_combined[k];
+        }
+        q_values[batch * N_TAU * N_ACTIONS + tau_idx * N_ACTIONS + c] = q_acc;
+    }
+}
+
+// ─────────────────────────────────────────────────────────────────────
+// rl_iqn_expected_q:
+//   Compute E[Q(s, a)] = (1/N_TAU) × Σ_{tau} Q(s, tau, a)
+//   for ensemble action selection.
+//
+// Inputs:
+//   q_values    [B, N_TAU, N_ACTIONS]     — full quantile Q tensor
+//   B                                     — batch size
+//   N_TAU                                 — number of quantile samples
+// Outputs:
+//   expected_q  [B, N_ACTIONS]            — mean Q per action
+//
+// Block layout: grid = (B, 1, 1); block = (N_ACTIONS, 1, 1).
+//   One thread per action; each thread sums over N_TAU quantiles.
+//   No cross-thread reduction needed (each action is independent).
+// ─────────────────────────────────────────────────────────────────────
+extern "C" __global__ void rl_iqn_expected_q(
+    const float* __restrict__ q_values,    // [B, N_TAU, N_ACTIONS]
+    int          B,
+    int          N_TAU,
+    float*       __restrict__ expected_q    // [B, N_ACTIONS]
+) {
+    const int batch = blockIdx.x;
+    const int act = threadIdx.x;
+    if (batch >= B) return;
+    if (act >= N_ACTIONS) return;
+
+    float sum = 0.0f;
+    const int base = batch * N_TAU * N_ACTIONS;
+    for (int t = 0; t < N_TAU; ++t) {
+        sum += q_values[base + t * N_ACTIONS + act];
+    }
+    expected_q[batch * N_ACTIONS + act] = sum / (float)N_TAU;
+}

crates/ml-alpha/cuda/rl_iqn_loss.cu

@@ -0,0 +1,130 @@
+// rl_iqn_loss.cu — IQN quantile Huber loss (forward + backward).
+//
+// Computes the quantile regression loss between online and target IQN
+// outputs for the action taken:
+//
+//   δ_{ij} = target_q[b, j, a*] - online_q[b, i, a*]
+//   ρ_{τ_i}(δ_{ij}) = |τ_i - 1(δ_{ij} < 0)| × Huber(δ_{ij}, κ=1.0)
+//   L_b = (1 / N_TAU_ONLINE) × (1 / N_TAU_TARGET) × Σ_i Σ_j ρ_{τ_i}(δ_{ij})
+//
+// where:
+//   Huber(δ, κ) = 0.5 × δ² / κ       if |δ| ≤ κ
+//               = |δ| - 0.5 × κ       otherwise
+//
+// Backward: ∂L/∂online_q[b, i, a*] is computed and written to grad_out.
+//
+// Block layout:
+//   grid = (B, 1, 1); block = (N_TAU_ONLINE_MAX, 1, 1).
+//   One block per batch entry. Each thread handles one online quantile
+//   index i, reducing across all target quantile indices j.
+//   Cross-batch loss aggregation uses block tree-reduce into shared
+//   memory (no atomicAdd per `feedback_no_atomicadd.md`), with block 0
+//   writing the final scalar loss.
+//
+// κ (Huber threshold) = 1.0 is standard for IQN (Dabney et al. 2018).
+
+#define HIDDEN_DIM       128
+#define N_ACTIONS        11
+#define N_TAU_MAX        64
+#define KAPPA            1.0f
+
+// ─────────────────────────────────────────────────────────────────────
+// rl_iqn_loss_fwd:
+//   Compute per-batch quantile Huber loss and per-online-quantile gradients.
+//
+// Inputs:
+//   online_q        [B, N_TAU_ONLINE, N_ACTIONS] — online IQN Q values
+//   target_q        [B, N_TAU_TARGET, N_ACTIONS] — target IQN Q values
+//   tau_online      [B, N_TAU_ONLINE]            — online quantile fractions
+//   actions_taken   [B]                          — action index per batch
+//   B                                           — batch size
+//   N_TAU_ONLINE                                — number of online quantiles
+//   N_TAU_TARGET                                — number of target quantiles
+// Outputs:
+//   loss_per_batch  [B]                         — per-batch scalar loss
+//   grad_online_q   [B, N_TAU_ONLINE, N_ACTIONS] — gradient w.r.t. online Q
+//                   (only the action-taken slice is non-zero)
+// ─────────────────────────────────────────────────────────────────────
+extern "C" __global__ void rl_iqn_loss_fwd(
+    const float* __restrict__ online_q,       // [B, N_TAU_ONLINE, N_ACTIONS]
+    const float* __restrict__ target_q,       // [B, N_TAU_TARGET, N_ACTIONS]
+    const float* __restrict__ tau_online,     // [B, N_TAU_ONLINE]
+    const int*   __restrict__ actions_taken,  // [B]
+    int          B,
+    int          N_TAU_ONLINE,
+    int          N_TAU_TARGET,
+    float*       __restrict__ loss_per_batch, // [B]
+    float*       __restrict__ grad_online_q   // [B, N_TAU_ONLINE, N_ACTIONS]
+) {
+    const int batch = blockIdx.x;
+    const int i = threadIdx.x;  // online quantile index
+    if (batch >= B) return;
+    if (i >= N_TAU_ONLINE) return;
+
+    const int act = actions_taken[batch];
+
+    // Online Q value for this quantile and action
+    const float q_online_i = online_q[batch * N_TAU_ONLINE * N_ACTIONS + i * N_ACTIONS + act];
+    const float tau_i = tau_online[batch * N_TAU_ONLINE + i];
+
+    float loss_sum = 0.0f;
+    float grad_sum = 0.0f;
+
+    // Sum over all target quantiles j
+    for (int j = 0; j < N_TAU_TARGET; ++j) {
+        const float q_target_j = target_q[batch * N_TAU_TARGET * N_ACTIONS + j * N_ACTIONS + act];
+        const float delta = q_target_j - q_online_i;
+        const float abs_delta = fabsf(delta);
+
+        // Huber loss
+        float huber;
+        float huber_grad;  // dHuber/d(online_q) = -dHuber/d(delta)
+        if (abs_delta <= KAPPA) {
+            huber = 0.5f * delta * delta / KAPPA;
+            huber_grad = -delta / KAPPA;
+        } else {
+            huber = abs_delta - 0.5f * KAPPA;
+            huber_grad = (delta < 0.0f) ? 1.0f : -1.0f;
+        }
+
+        // Quantile weight: |tau_i - 1(delta < 0)|
+        float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
+        float weight = fabsf(tau_i - indicator);
+
+        loss_sum += weight * huber;
+        // Gradient: d/d(online_q) of weight * huber
+        // weight doesn't depend on online_q (it depends on sign of delta
+        // which is a subgradient boundary). Standard IQN gradient:
+        grad_sum += weight * huber_grad;
+    }
+
+    // Normalize by N_TAU_TARGET (the N_TAU_ONLINE normalization is applied
+    // when writing to loss_per_batch after the shared-mem reduce below)
+    float loss_i = loss_sum / (float)N_TAU_TARGET;
+    float grad_i = grad_sum / ((float)N_TAU_ONLINE * (float)N_TAU_TARGET);
+
+    // Write gradient for this online quantile's action slot
+    // Zero-init the non-taken-action gradient slots
+    const int grad_base = batch * N_TAU_ONLINE * N_ACTIONS + i * N_ACTIONS;
+    for (int a = 0; a < N_ACTIONS; ++a) {
+        grad_online_q[grad_base + a] = (a == act) ? grad_i : 0.0f;
+    }
+
+    // Block tree-reduce to get per-batch loss (sum over online quantiles,
+    // then divide by N_TAU_ONLINE)
+    __shared__ float s_loss[N_TAU_MAX];
+    s_loss[i] = loss_i;
+    __syncthreads();
+
+    // Tree reduce within the block
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (i < stride && (i + stride) < N_TAU_ONLINE) {
+            s_loss[i] += s_loss[i + stride];
+        }
+        __syncthreads();
+    }
+
+    if (i == 0) {
+        loss_per_batch[batch] = s_loss[0] / (float)N_TAU_ONLINE;
+    }
+}

crates/ml-alpha/src/rl/iqn.rs

@@ -0,0 +1,434 @@
+//! IQN (Implicit Quantile Network) distributional Q-head for the
+//! integrated RL trainer — complementary to the existing C51 head.
+//!
+//! Sits on the shared Mamba2+CfC encoder's `h_t` output alongside
+//! the C51 head. Produces per-action quantile values of shape
+//! `[B, N_TAU, N_ACTIONS]` for `N_TAU` quantile fractions sampled
+//! from U(0,1). The ensemble action selection combines both heads:
+//!
+//! ```text
+//!   E_Q = α × E_C51 + (1-α) × E_IQN
+//! ```
+//!
+//! where α is read from `ISV[RL_IQN_ENSEMBLE_ALPHA_INDEX=544]`.
+//!
+//! ## Architecture
+//!
+//! ```text
+//!   1. Quantile embedding:
+//!      cos_feats[i] = cos((i+1) × π × τ)   for i=0..EMBED_DIM-1
+//!      phi(τ) = ReLU(W_embed × cos_feats + b_embed)   → [HIDDEN_DIM]
+//!
+//!   2. Element-wise product:
+//!      combined = h_t ⊙ phi(τ)   → [HIDDEN_DIM]
+//!
+//!   3. Action value projection:
+//!      Q(s, τ, a) = W_out × combined + b_out   → [N_ACTIONS]
+//! ```
+//!
+//! Expected Q for action selection: `E_IQN[a] = mean_τ Q(s, τ, a)`.
+//!
+//! ## Loss
+//!
+//! Quantile Huber loss (Dabney et al. 2018):
+//!
+//! ```text
+//!   δ_{ij} = target_q[j,a*] - online_q[i,a*]
+//!   ρ_{τ_i}(δ_{ij}) = |τ_i - 1(δ<0)| × Huber(δ, κ=1.0)
+//!   L = (1/N_TAU²) Σ_i Σ_j ρ_{τ_i}(δ_{ij})
+//! ```
+//!
+//! ## Initialisation
+//!
+//! Per `pearl_scoped_init_seed_for_reproducibility`, `IqnHead::new`
+//! installs a `scoped_init_seed` guard before drawing Xavier samples.
+//! Xavier uniform scaled by 0.01 (same as C51 head) so initial Q
+//! outputs sit near zero.
+//!
+//! ## Constraints honoured
+//!
+//!   * `feedback_no_atomicadd.md` — loss kernel uses block tree-reduce.
+//!   * `feedback_no_htod_htoh_only_mapped_pinned.md` — weight uploads
+//!     stage through `MappedF32Buffer`.
+//!   * `feedback_no_nvrtc.md` — pre-compiled cubins via `build.rs`.
+//!   * `feedback_isv_for_adaptive_bounds.md` — N_TAU, ensemble α, and LR
+//!     live in ISV slots 543-545.
+//!   * `feedback_cpu_is_read_only.md` — forward/backward are GPU kernels.
+
+use std::sync::Arc;
+
+use anyhow::{Context, Result};
+use cudarc::driver::{
+    CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
+};
+use ml_core::cuda_autograd::init::scoped_init_seed;
+use ml_core::device::MlDevice;
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha8Rng;
+
+use crate::heads::HIDDEN_DIM;
+use crate::pinned_mem::MappedF32Buffer;
+use crate::rl::common::N_ACTIONS;
+
+/// Quantile embedding dimension — number of cosine basis functions.
+/// Matches the `EMBED_DIM` define in `cuda/rl_iqn_forward.cu`.
+const EMBED_DIM: usize = 64;
+
+/// Default number of quantile samples per forward pass. ISV-overridable
+/// via `RL_IQN_N_TAU_INDEX`. Matches the `N_TAU` default in the kernel.
+const DEFAULT_N_TAU: usize = 32;
+
+const IQN_FWD_CUBIN: &[u8] = include_bytes!(concat!(
+    env!("OUT_DIR"),
+    "/rl_iqn_forward.cubin"
+));
+const IQN_LOSS_CUBIN: &[u8] = include_bytes!(concat!(
+    env!("OUT_DIR"),
+    "/rl_iqn_loss.cubin"
+));
+
+/// Construction config for [`IqnHead`].
+#[derive(Clone, Debug)]
+pub struct IqnHeadConfig {
+    /// Encoder hidden dimension `h_t [B, HIDDEN_DIM]` feeding the head.
+    pub hidden_dim: usize,
+    /// Random seed for Xavier init. Reproducibility-critical per
+    /// `pearl_scoped_init_seed_for_reproducibility`.
+    pub seed: u64,
+    /// Number of quantile samples per forward pass. Defaults to
+    /// `DEFAULT_N_TAU` (32); overridden at runtime by
+    /// `ISV[RL_IQN_N_TAU_INDEX]` in the training loop.
+    pub n_tau: usize,
+}
+
+impl Default for IqnHeadConfig {
+    fn default() -> Self {
+        Self {
+            hidden_dim: HIDDEN_DIM,
+            seed: 0x19A,
+            n_tau: DEFAULT_N_TAU,
+        }
+    }
+}
+
+/// IQN distributional Q-head. Owns device weights for the quantile
+/// embedding (`w_embed`, `b_embed`) and output projection (`w_out`,
+/// `b_out`), plus a target network copy for Bellman bootstrapping.
+///
+/// The head produces `Q(s, τ, a)` for sampled τ ∈ U(0,1). Expected Q
+/// for action selection is `E[Q(s,a)] = mean_τ Q(s,τ,a)`, combined
+/// with C51's E[Q] via the ensemble α weight from ISV.
+pub struct IqnHead {
+    cfg: IqnHeadConfig,
+    stream: Arc<CudaStream>,
+
+    _fwd_module: Arc<CudaModule>,
+    /// Forward kernel: quantile embedding + action-value projection.
+    pub fwd_fn: CudaFunction,
+    /// Expected-Q reduction kernel: mean over tau dimension.
+    pub expected_q_fn: CudaFunction,
+
+    _loss_module: Arc<CudaModule>,
+    /// Loss kernel: quantile Huber loss forward + backward.
+    pub loss_fn: CudaFunction,
+
+    /// Quantile embedding weight `[EMBED_DIM, HIDDEN_DIM]`, row-major.
+    pub w_embed_d: CudaSlice<f32>,
+    /// Quantile embedding bias `[HIDDEN_DIM]`.
+    pub b_embed_d: CudaSlice<f32>,
+    /// Output projection weight `[HIDDEN_DIM, N_ACTIONS]`, row-major.
+    pub w_out_d: CudaSlice<f32>,
+    /// Output projection bias `[N_ACTIONS]`.
+    pub b_out_d: CudaSlice<f32>,
+
+    /// Target-network weights — same shapes as online. Soft-updated
+    /// using the same τ from `ISV[RL_TARGET_TAU_INDEX]` as C51.
+    pub w_embed_target_d: CudaSlice<f32>,
+    pub b_embed_target_d: CudaSlice<f32>,
+    pub w_out_target_d: CudaSlice<f32>,
+    pub b_out_target_d: CudaSlice<f32>,
+}
+
+impl IqnHead {
+    /// Allocate device weights, load cubins, and cache kernel handles.
+    /// Online and target networks start with IDENTICAL Xavier draws so
+    /// the first Bellman backup sees zero-divergence (same bootstrap
+    /// principle as C51).
+    pub fn new(dev: &MlDevice, cfg: IqnHeadConfig) -> Result<Self> {
+        let stream: Arc<CudaStream> = dev.cuda_stream().context("iqn_head stream")?.clone();
+        let ctx = dev.cuda_context().context("iqn_head ctx")?;
+
+        let fwd_module = ctx
+            .load_cubin(IQN_FWD_CUBIN.to_vec())
+            .context("load rl_iqn_forward cubin")?;
+        let fwd_fn = fwd_module
+            .load_function("rl_iqn_forward")
+            .context("load rl_iqn_forward fn")?;
+        let expected_q_fn = fwd_module
+            .load_function("rl_iqn_expected_q")
+            .context("load rl_iqn_expected_q fn")?;
+
+        let loss_module = ctx
+            .load_cubin(IQN_LOSS_CUBIN.to_vec())
+            .context("load rl_iqn_loss cubin")?;
+        let loss_fn = loss_module
+            .load_function("rl_iqn_loss_fwd")
+            .context("load rl_iqn_loss_fwd fn")?;
+
+        // Per pearl_scoped_init_seed_for_reproducibility: install the
+        // scoped seed guard BEFORE drawing any Xavier samples.
+        let _seed_guard = scoped_init_seed(cfg.seed);
+        let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed);
+
+        // Xavier uniform scaled by 0.01 — initial Q outputs near zero.
+        let embed_scale =
+            0.01_f32 * (6.0_f32 / (EMBED_DIM + cfg.hidden_dim) as f32).sqrt();
+        let w_embed_host: Vec<f32> = (0..EMBED_DIM * cfg.hidden_dim)
+            .map(|_| rng.gen_range(-embed_scale..embed_scale))
+            .collect();
+        let b_embed_host: Vec<f32> = vec![0.0_f32; cfg.hidden_dim];
+
+        let out_scale =
+            0.01_f32 * (6.0_f32 / (cfg.hidden_dim + N_ACTIONS) as f32).sqrt();
+        let w_out_host: Vec<f32> = (0..cfg.hidden_dim * N_ACTIONS)
+            .map(|_| rng.gen_range(-out_scale..out_scale))
+            .collect();
+        let b_out_host: Vec<f32> = vec![0.0_f32; N_ACTIONS];
+
+        let w_embed_d = upload(&stream, &w_embed_host)?;
+        let b_embed_d = upload(&stream, &b_embed_host)?;
+        let w_out_d = upload(&stream, &w_out_host)?;
+        let b_out_d = upload(&stream, &b_out_host)?;
+
+        // Target network: identical init for zero-divergence bootstrap.
+        let w_embed_target_d = upload(&stream, &w_embed_host)?;
+        let b_embed_target_d = upload(&stream, &b_embed_host)?;
+        let w_out_target_d = upload(&stream, &w_out_host)?;
+        let b_out_target_d = upload(&stream, &b_out_host)?;
+
+        Ok(Self {
+            cfg,
+            stream,
+            _fwd_module: fwd_module,
+            fwd_fn,
+            expected_q_fn,
+            _loss_module: loss_module,
+            loss_fn,
+            w_embed_d,
+            b_embed_d,
+            w_out_d,
+            b_out_d,
+            w_embed_target_d,
+            b_embed_target_d,
+            w_out_target_d,
+            b_out_target_d,
+        })
+    }
+
+    /// Forward pass: compute Q(s, τ, a) for all (batch, tau, action)
+    /// triples. Caller provides pre-sampled `tau [B, N_TAU]` from U(0,1).
+    ///
+    /// Output: `q_values [B, N_TAU, N_ACTIONS]`.
+    pub fn forward(
+        &self,
+        h_t: &CudaSlice<f32>,
+        tau: &CudaSlice<f32>,
+        b_size: usize,
+        n_tau: usize,
+        q_values_out: &mut CudaSlice<f32>,
+    ) -> Result<()> {
+        debug_assert_eq!(h_t.len(), b_size * self.cfg.hidden_dim);
+        debug_assert_eq!(tau.len(), b_size * n_tau);
+        debug_assert_eq!(q_values_out.len(), b_size * n_tau * N_ACTIONS);
+
+        let b_i = b_size as i32;
+        let n_tau_i = n_tau as i32;
+        let launch_cfg = LaunchConfig {
+            grid_dim: (b_size as u32, n_tau as u32, 1),
+            block_dim: (HIDDEN_DIM as u32, 1, 1),
+            shared_mem_bytes: (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
+        };
+        let mut launch = self.stream.launch_builder(&self.fwd_fn);
+        launch
+            .arg(h_t)
+            .arg(tau)
+            .arg(&self.w_embed_d)
+            .arg(&self.b_embed_d)
+            .arg(&self.w_out_d)
+            .arg(&self.b_out_d)
+            .arg(&b_i)
+            .arg(&n_tau_i)
+            .arg(q_values_out);
+        unsafe {
+            launch.launch(launch_cfg).context("rl_iqn_forward launch")?;
+        }
+        Ok(())
+    }
+
+    /// Target-network forward: same kernel as online forward, but reads
+    /// from the target weights. Used for the Bellman bootstrap.
+    pub fn forward_target(
+        &self,
+        h_t: &CudaSlice<f32>,
+        tau: &CudaSlice<f32>,
+        b_size: usize,
+        n_tau: usize,
+        q_values_out: &mut CudaSlice<f32>,
+    ) -> Result<()> {
+        debug_assert_eq!(h_t.len(), b_size * self.cfg.hidden_dim);
+        debug_assert_eq!(tau.len(), b_size * n_tau);
+        debug_assert_eq!(q_values_out.len(), b_size * n_tau * N_ACTIONS);
+
+        let b_i = b_size as i32;
+        let n_tau_i = n_tau as i32;
+        let launch_cfg = LaunchConfig {
+            grid_dim: (b_size as u32, n_tau as u32, 1),
+            block_dim: (HIDDEN_DIM as u32, 1, 1),
+            shared_mem_bytes: (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
+        };
+        let mut launch = self.stream.launch_builder(&self.fwd_fn);
+        launch
+            .arg(h_t)
+            .arg(tau)
+            .arg(&self.w_embed_target_d)
+            .arg(&self.b_embed_target_d)
+            .arg(&self.w_out_target_d)
+            .arg(&self.b_out_target_d)
+            .arg(&b_i)
+            .arg(&n_tau_i)
+            .arg(q_values_out);
+        unsafe {
+            launch
+                .launch(launch_cfg)
+                .context("rl_iqn_forward (target) launch")?;
+        }
+        Ok(())
+    }
+
+    /// Compute expected Q values: `E_IQN[b,a] = mean_τ Q(s,τ,a)`.
+    ///
+    /// Input: `q_values [B, N_TAU, N_ACTIONS]` from [`forward`].
+    /// Output: `expected_q [B, N_ACTIONS]`.
+    pub fn expected_q(
+        &self,
+        q_values: &CudaSlice<f32>,
+        b_size: usize,
+        n_tau: usize,
+        expected_q_out: &mut CudaSlice<f32>,
+    ) -> Result<()> {
+        debug_assert_eq!(q_values.len(), b_size * n_tau * N_ACTIONS);
+        debug_assert_eq!(expected_q_out.len(), b_size * N_ACTIONS);
+
+        let b_i = b_size as i32;
+        let n_tau_i = n_tau as i32;
+        let launch_cfg = LaunchConfig {
+            grid_dim: (b_size as u32, 1, 1),
+            block_dim: (N_ACTIONS as u32, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let mut launch = self.stream.launch_builder(&self.expected_q_fn);
+        launch
+            .arg(q_values)
+            .arg(&b_i)
+            .arg(&n_tau_i)
+            .arg(expected_q_out);
+        unsafe {
+            launch
+                .launch(launch_cfg)
+                .context("rl_iqn_expected_q launch")?;
+        }
+        Ok(())
+    }
+
+    /// Quantile Huber loss forward + backward for the taken action.
+    ///
+    /// Computes per-batch loss and gradient w.r.t. the online Q values.
+    /// The gradient is only non-zero at the taken action's quantile
+    /// slice — all other actions receive zero gradient.
+    #[allow(clippy::too_many_arguments)]
+    pub fn compute_loss(
+        &self,
+        online_q: &CudaSlice<f32>,
+        target_q: &CudaSlice<f32>,
+        tau_online: &CudaSlice<f32>,
+        actions_taken: &CudaSlice<i32>,
+        b_size: usize,
+        n_tau_online: usize,
+        n_tau_target: usize,
+        loss_per_batch: &mut CudaSlice<f32>,
+        grad_online_q: &mut CudaSlice<f32>,
+    ) -> Result<()> {
+        debug_assert_eq!(online_q.len(), b_size * n_tau_online * N_ACTIONS);
+        debug_assert_eq!(target_q.len(), b_size * n_tau_target * N_ACTIONS);
+        debug_assert_eq!(tau_online.len(), b_size * n_tau_online);
+        debug_assert_eq!(actions_taken.len(), b_size);
+        debug_assert_eq!(loss_per_batch.len(), b_size);
+        debug_assert_eq!(grad_online_q.len(), b_size * n_tau_online * N_ACTIONS);
+
+        let b_i = b_size as i32;
+        let n_tau_online_i = n_tau_online as i32;
+        let n_tau_target_i = n_tau_target as i32;
+        let launch_cfg = LaunchConfig {
+            grid_dim: (b_size as u32, 1, 1),
+            block_dim: (n_tau_online as u32, 1, 1),
+            // Shared memory for the tree-reduce over online quantiles.
+            shared_mem_bytes: (n_tau_online * std::mem::size_of::<f32>()) as u32,
+        };
+        let mut launch = self.stream.launch_builder(&self.loss_fn);
+        launch
+            .arg(online_q)
+            .arg(target_q)
+            .arg(tau_online)
+            .arg(actions_taken)
+            .arg(&b_i)
+            .arg(&n_tau_online_i)
+            .arg(&n_tau_target_i)
+            .arg(loss_per_batch)
+            .arg(grad_online_q);
+        unsafe {
+            launch.launch(launch_cfg).context("rl_iqn_loss_fwd launch")?;
+        }
+        Ok(())
+    }
+
+    /// Stream used to launch all kernels owned by this head.
+    pub fn stream(&self) -> &Arc<CudaStream> {
+        &self.stream
+    }
+
+    /// Encoder hidden dim this head was built for.
+    pub fn hidden_dim(&self) -> usize {
+        self.cfg.hidden_dim
+    }
+
+    /// Number of quantile samples configured at construction.
+    pub fn n_tau(&self) -> usize {
+        self.cfg.n_tau
+    }
+}
+
+// ── pinned-staging upload helper (mirrors dqn.rs::upload) ──────────
+
+fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
+    let n = host.len();
+    let staging = unsafe { MappedF32Buffer::new(n) }
+        .map_err(|e| anyhow::anyhow!("iqn_head upload staging: {e}"))?;
+    staging.write_from_slice(host);
+    let mut dst = stream
+        .alloc_zeros::<f32>(n)
+        .context("iqn_head upload alloc")?;
+    if n > 0 {
+        let nbytes = n * std::mem::size_of::<f32>();
+        unsafe {
+            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
+            cudarc::driver::result::memcpy_dtod_async(
+                dst_ptr,
+                staging.dev_ptr,
+                nbytes,
+                stream.cu_stream(),
+            )
+            .context("iqn_head upload DtoD")?;
+        }
+    }
+    Ok(dst)
+}

crates/ml-alpha/src/rl/isv_slots.rs

@@ -49,7 +49,9 @@
 //! | 494-497     | SP20 P5 trail-stop bounds (MIN/MAX/K_INIT/ADJUST_RATE)  | rl_isv_write (seed) + rl_trail_mutate + rl_unit_state_update |
 //! | 498-503     | SP20 P3 FRD head (λ/LR/h1-3 ticks/bucket-range σ)       | rl_isv_write (seed) + rl_frd_fwd + rl_frd_bwd + loader label gen |
 //!
-//! Total: 104 slots; `RL_SLOTS_END = 504`.
+//! | 543-545     | 3 IQN head slots (N_TAU + ensemble α + LR)              | IqnHead + ensemble action selector |
+//!
+//! Total: 146 slots; `RL_SLOTS_END = 546`.
 
 /// Discount factor γ. Controller input: mean trade duration / anchor.
 /// Bootstrap 0.99.
@@ -955,6 +957,22 @@ pub const RL_SESSION_EMA_ALPHA_INDEX: usize = 541;
 /// instead of r + γQ(s'). Default 10 (2.5 seconds at 4 events/sec).
 pub const RL_N_STEP_INDEX: usize = 542;
 
+/// IQN quantile count (N_TAU). Number of quantile fractions sampled
+/// from U(0,1) per forward pass. Higher values improve distributional
+/// coverage at the cost of more compute per step. Default 32.
+pub const RL_IQN_N_TAU_INDEX: usize = 543;
+
+/// IQN ensemble alpha: weight of C51 in the combined expected Q.
+/// `E_Q = alpha × E_C51 + (1 - alpha) × E_IQN`. Default 0.5
+/// (equal weighting). Higher values favor C51's fixed-support
+/// estimate; lower values favor IQN's implicit quantile estimate.
+pub const RL_IQN_ENSEMBLE_ALPHA_INDEX: usize = 544;
+
+/// IQN head learning rate. Default 1e-3 (canonical Adam starting LR).
+/// Subject to the same plateau-decay controller as other heads once
+/// wired into the training loop.
+pub const RL_IQN_LR_INDEX: usize = 545;
+
 /// Last RL-allocated slot index (exclusive). The integrated trainer
 /// extends `ISV_TOTAL_DIM` to at least this value at trainer init time.
-pub const RL_SLOTS_END: usize = 543;
+pub const RL_SLOTS_END: usize = 546;

crates/ml-alpha/src/rl/mod.rs

@@ -18,6 +18,7 @@
 pub mod common;
 pub mod dqn;
 pub mod frd;
+pub mod iqn;
 pub mod isv_slots;
 pub mod loss_balance;
 pub mod ppo;