From 6192bd4eb650f88b8cc3f61c41a92d2ea6a4c454 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 25 May 2026 20:07:10 +0200 Subject: [PATCH] =?UTF-8?q?spec:=20Q-learning=20improvements=20=E2=80=94?= =?UTF-8?q?=20n-step=20+=20IQN=20ensemble=20+=20noisy=20nets?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three-phase plan to push Q convergence past profitability: 1. N-step returns (n=10): 10× more direct reward signal 2. IQN complementary head alongside C51: ensemble action selection 3. Noisy linear layers: state-dependent exploration C51 stays for the confidence gate's distributional LCB. IQN adds flexible quantile estimation. Ensemble combines both for action selection. All ISV-driven. Co-Authored-By: Claude Opus 4.7 --- .../2026-05-25-q-learning-improvements.md | 231 ++++++++++++++++++ 1 file changed, 231 insertions(+) create mode 100644 docs/superpowers/specs/2026-05-25-q-learning-improvements.md diff --git a/docs/superpowers/specs/2026-05-25-q-learning-improvements.md b/docs/superpowers/specs/2026-05-25-q-learning-improvements.md new file mode 100644 index 000000000..34d4cc144 --- /dev/null +++ b/docs/superpowers/specs/2026-05-25-q-learning-improvements.md @@ -0,0 +1,231 @@ +# Q-Learning Improvements: N-Step + IQN Ensemble + Noisy Nets + +> **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:** Push Q convergence deeper and break the profitability barrier by adding n-step returns, an IQN complementary head alongside C51, and noisy exploration. + +**Architecture:** Keep C51 (categorical, 21 atoms) for risk-aware gating (confidence gate reads its distributional LCB). Add IQN (implicit quantile) as a second Q-head that shares the encoder but has its own weights. Ensemble action selection: `E_Q = α × E_C51 + (1-α) × E_IQN` with ISV-driven α. N-step returns accelerate both heads. Noisy linear layers replace the P_MIN floor for state-dependent exploration. + +**Tech Stack:** CUDA kernels via cudarc 0.19, Rust 1.85, ISV-driven parameters, PER=32768 + +**Current state:** C51 with 21 atoms [-0.5,+0.5], 1-step Bellman, γ=0.995, l_q best=1.65, win rate 49.4%, approaching break-even. + +--- + +## Context + +### Current Q path in `step_synthetic` (integrated.rs) + +``` +1. dqn_head.forward(h_t) → q_logits_d [B × N_ACTIONS × Q_N_ATOMS] +2. dqn_head.forward_target(h_tp1) → q_target_full_d +3. argmax_expected_q(q_logits_tp1) → next_actions_d (Double DQN) +4. select_action_atoms(target, next_actions) → q_target_action_d [B × Q_N_ATOMS] +5. project_bellman_target(target_atoms, rewards, dones, γ) → target_dist_d +6. C51 cross-entropy loss: KL(target_dist || softmax(q_logits_at_action)) +7. backward through q_logits → grad_h_t +``` + +### Key files +- `crates/ml-alpha/src/rl/dqn.rs` — DqnHead struct, forward/backward/bellman +- `crates/ml-alpha/src/rl/common.rs` — Transition, N_ACTIONS, Q_N_ATOMS +- `crates/ml-alpha/src/rl/replay.rs` — PER buffer +- `crates/ml-alpha/src/trainer/integrated.rs` — step pipeline, push/sample replay +- `crates/ml-alpha/cuda/bellman_target_projection.cu` — C51 Bellman projection +- `crates/ml-alpha/cuda/rl_confidence_gate.cu` — reads C51 distributional LCB + +--- + +## Phase 1: N-Step Returns + +### What changes + +The Bellman target currently uses 1-step: `r + γ × Q(s')`. + +N-step uses: `R_n + γⁿ × Q(s_{t+n})` where `R_n = Σ_{k=0}^{n-1} γᵏ r_{t+k}`. + +At γ=0.995 with n=10, the agent sees 10 real rewards (~2.5 seconds) before bootstrapping from Q. This dramatically reduces bootstrap error and accelerates convergence. + +### ISV slots +- `RL_N_STEP_INDEX` (default 10) — ISV-driven n + +### Transition changes +```rust +pub struct Transition { + pub h_t: CudaSlice, // state at step t (unchanged) + pub action: u32, // action at step t (unchanged) + pub reward: f32, // BECOMES: n-step discounted return R_n + pub raw_reward: f32, // BECOMES: raw (unscaled) R_n + pub next_h_t: CudaSlice, // BECOMES: state at step t+n (not t+1) + pub done: bool, // BECOMES: any done in [t, t+n) + pub n_step_gamma: f32, // NEW: γⁿ (effective discount for bootstrap) + pub log_pi_old: f32, // unchanged + pub q_value_old: [f32; N_ACTIONS], // unchanged +} +``` + +### N-step ring buffer (Rust side, in IntegratedTrainer) +```rust +struct NStepBuffer { + entries: VecDeque, + n: usize, // from ISV +} +struct NStepEntry { + h_t: CudaSlice, + action: u32, + raw_reward: f32, + scaled_reward: f32, + done: bool, + log_pi_old: f32, + gamma: f32, +} +``` + +On each step: push entry to ring. When ring.len() == n, pop the oldest, compute R_n, push to PER with `h_t` from oldest entry and `next_h_t` from current step. + +### Bellman kernel change +`bellman_target_projection.cu` line 151: +```cuda +// OLD: const float t_z = r + gamma_eff * atom_value; +// NEW: const float t_z = r + n_step_gamma * gamma_eff * atom_value; +``` +Add `n_step_gamma` as kernel parameter (read from the sampled transition's `n_step_gamma` field, uploaded per batch). + +### Testing +- GPU oracle: push 10 transitions with known rewards, verify R_n computation +- Smoke: 1k steps, verify l_q drops faster than 1-step baseline + +--- + +## Phase 2: IQN Complementary Head + +### What changes + +Add an IQN head ALONGSIDE C51. Both share the encoder output `h_t [B, HIDDEN_DIM]`. Each has its own weights and loss. Action selection ensembles both. + +### IQN architecture + +``` +Input: h_t [B, HIDDEN_DIM], τ [B, N_TAU] ~ U(0,1) + +1. Quantile embedding: + φ(τ) = ReLU(Linear(cos(i × π × τ for i=1..64))) → [B, N_TAU, HIDDEN_DIM] + +2. Element-wise combine: + combined = h_t.unsqueeze(1) ⊙ φ(τ) → [B, N_TAU, HIDDEN_DIM] + +3. Action-value: + Q(s, a, τ) = Linear(combined) → [B, N_TAU, N_ACTIONS] + +4. Expected Q for action selection: + E_Q[a] = mean_τ(Q(s, a, τ)) +``` + +### Loss: Quantile Huber + +``` +For each pair (τ_i, τ_j) of sample and target quantiles: + δ = target_Q(τ_j) - online_Q(τ_i) + ρ_τ(δ) = |τ_i - 1(δ < 0)| × Huber(δ, κ=1.0) + L = (1/N_TAU²) Σ_i Σ_j ρ_{τ_i}(δ_{ij}) +``` + +### Ensemble action selection + +``` +E_ensemble[a] = α × E_C51[a] + (1 - α) × E_IQN[a] +``` + +ISV-driven α (default 0.5). The confidence gate continues reading C51's distributional LCB (unchanged). The policy π receives gradient from the ensemble value. + +### New ISV slots +- `RL_IQN_N_TAU_INDEX` (default 32) — quantile samples per forward +- `RL_IQN_ENSEMBLE_ALPHA_INDEX` (default 0.5) — C51 weight in ensemble +- `RL_IQN_LR_INDEX` — separate learning rate for IQN head + +### New files +- `crates/ml-alpha/cuda/rl_iqn_forward.cu` — quantile embedding + forward +- `crates/ml-alpha/cuda/rl_iqn_loss.cu` — quantile Huber loss +- `crates/ml-alpha/cuda/rl_iqn_backward.cu` — gradients through IQN +- `crates/ml-alpha/cuda/rl_ensemble_action_value.cu` — combine C51 + IQN +- `crates/ml-alpha/src/rl/iqn.rs` — IqnHead struct, weight management + +### Testing +- GPU oracle: verify quantile embedding produces correct shape +- GPU oracle: verify quantile Huber loss matches analytical formula +- Smoke: ensemble action selection produces valid actions + +--- + +## Phase 3: Noisy Linear Layers + +### What changes + +Replace the P_MIN probability floor (brute-force 1.5% per action) with learnable parameter noise. The network learns WHERE to explore — more noise in unfamiliar states, less in confident ones. + +### Noisy linear formula + +``` +y = (μ_w + σ_w ⊙ ε_w) × x + (μ_b + σ_b ⊙ ε_b) +``` + +Where: +- μ_w, μ_b: learned mean weights (same as standard linear) +- σ_w, σ_b: learned noise scale (initialized small, e.g., 0.5/√fan_in) +- ε_w, ε_b: noise samples ~ N(0,1), resampled each forward pass + +### Factored noise (efficient) + +For a layer with fan_in=p, fan_out=q: +- Instead of p×q independent noise samples, use factored: + `ε_w[i,j] = f(ε_i) × f(ε_j)` where `f(x) = sign(x) × √|x|` +- Only p + q noise samples needed (not p×q) + +### Integration points + +Replace the final linear layer in: +1. DQN head (Q values) — `dqn_fwd.cu` weight multiply +2. Policy head (π logits) — `rl_pi_action_kernel.cu` upstream +3. IQN head (if Phase 2 is done) — action-value linear + +Remove P_MIN floor from `rl_pi_action_kernel.cu` (no longer needed). + +### New ISV slots +- `RL_NOISY_SIGMA_INIT_INDEX` (default 0.5) — initial noise scale + +### New files +- `crates/ml-alpha/cuda/rl_noisy_linear_forward.cu` +- `crates/ml-alpha/cuda/rl_noisy_linear_backward.cu` +- `crates/ml-alpha/src/rl/noisy.rs` — NoisyLinear struct + +### Testing +- GPU oracle: noisy forward with fixed seed produces deterministic output +- GPU oracle: noise disabled (σ=0) matches standard linear +- Smoke: verify exploration without P_MIN floor + +--- + +## Dependency graph + +``` +Phase 1 (N-step) ─────────── independent, highest ROI, do first + │ +Phase 2 (IQN ensemble) ────── depends on nothing, but benefits from Phase 1 + │ +Phase 3 (Noisy nets) ──────── depends on nothing, remove P_MIN after +``` + +## Kill criteria + +- **Phase 1 success:** Q-loss drops below 1.5 at 100k steps (vs 2.14 current) +- **Phase 2 success:** ensemble Q converges faster than C51-alone +- **Phase 3 success:** exploration diversity (action entropy) maintained without P_MIN floor + +## Implementation order + +1. Phase 1 first (2 hours) — n-step ring buffer + modified Bellman kernel +2. Validate with 200k cluster run +3. Phase 2 (6 hours) — IQN head + ensemble +4. Validate with 200k cluster run +5. Phase 3 (3 hours) — noisy linear layers +6. Final 1M run with all three