From 9c0be855ddfbbb688b0571c0b300a451d42e7bc3 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 25 May 2026 20:35:06 +0200 Subject: [PATCH] =?UTF-8?q?feat(rl):=20IQN=20complementary=20Q-head=20?= =?UTF-8?q?=E2=80=94=20quantile=20embedding=20+=20Huber=20loss?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 2 + crates/ml-alpha/cuda/rl_iqn_forward.cu | 143 ++++++++ crates/ml-alpha/cuda/rl_iqn_loss.cu | 130 ++++++++ crates/ml-alpha/src/rl/iqn.rs | 434 +++++++++++++++++++++++++ crates/ml-alpha/src/rl/isv_slots.rs | 22 +- crates/ml-alpha/src/rl/mod.rs | 1 + 6 files changed, 730 insertions(+), 2 deletions(-) create mode 100644 crates/ml-alpha/cuda/rl_iqn_forward.cu create mode 100644 crates/ml-alpha/cuda/rl_iqn_loss.cu create mode 100644 crates/ml-alpha/src/rl/iqn.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index ea7164b72..e060aa5f1 100644 --- a/crates/ml-alpha/build.rs +++ b/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 diff --git a/crates/ml-alpha/cuda/rl_iqn_forward.cu b/crates/ml-alpha/cuda/rl_iqn_forward.cu new file mode 100644 index 000000000..f036f017c --- /dev/null +++ b/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; +} diff --git a/crates/ml-alpha/cuda/rl_iqn_loss.cu b/crates/ml-alpha/cuda/rl_iqn_loss.cu new file mode 100644 index 000000000..b72e79ad2 --- /dev/null +++ b/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; + } +} diff --git a/crates/ml-alpha/src/rl/iqn.rs b/crates/ml-alpha/src/rl/iqn.rs new file mode 100644 index 000000000..e6ac70d2c --- /dev/null +++ b/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, + + _fwd_module: Arc, + /// 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, + /// 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, + /// Quantile embedding bias `[HIDDEN_DIM]`. + pub b_embed_d: CudaSlice, + /// Output projection weight `[HIDDEN_DIM, N_ACTIONS]`, row-major. + pub w_out_d: CudaSlice, + /// Output projection bias `[N_ACTIONS]`. + pub b_out_d: CudaSlice, + + /// 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, + pub b_embed_target_d: CudaSlice, + pub w_out_target_d: CudaSlice, + pub b_out_target_d: CudaSlice, +} + +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 { + let stream: Arc = 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 = (0..EMBED_DIM * cfg.hidden_dim) + .map(|_| rng.gen_range(-embed_scale..embed_scale)) + .collect(); + let b_embed_host: Vec = 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 = (0..cfg.hidden_dim * N_ACTIONS) + .map(|_| rng.gen_range(-out_scale..out_scale)) + .collect(); + let b_out_host: Vec = 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, + tau: &CudaSlice, + b_size: usize, + n_tau: usize, + q_values_out: &mut CudaSlice, + ) -> 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::()) 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, + tau: &CudaSlice, + b_size: usize, + n_tau: usize, + q_values_out: &mut CudaSlice, + ) -> 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::()) 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, + b_size: usize, + n_tau: usize, + expected_q_out: &mut CudaSlice, + ) -> 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, + target_q: &CudaSlice, + tau_online: &CudaSlice, + actions_taken: &CudaSlice, + b_size: usize, + n_tau_online: usize, + n_tau_target: usize, + loss_per_batch: &mut CudaSlice, + grad_online_q: &mut CudaSlice, + ) -> 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::()) 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 { + &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, host: &[f32]) -> Result> { + 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::(n) + .context("iqn_head upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + 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) +} diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index c0e15d6f8..730f40843 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/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; diff --git a/crates/ml-alpha/src/rl/mod.rs b/crates/ml-alpha/src/rl/mod.rs index 6113c29b3..0068d054e 100644 --- a/crates/ml-alpha/src/rl/mod.rs +++ b/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;