spec: Q-learning improvements — n-step + IQN ensemble + noisy nets
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
231
docs/superpowers/specs/2026-05-25-q-learning-improvements.md
Normal file
231
docs/superpowers/specs/2026-05-25-q-learning-improvements.md
Normal file
@@ -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<f32>, // 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<f32>, // 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<NStepEntry>,
|
||||
n: usize, // from ISV
|
||||
}
|
||||
struct NStepEntry {
|
||||
h_t: CudaSlice<f32>,
|
||||
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
|
||||
Reference in New Issue
Block a user