docs: Add DQN algorithm fix & 2026 modernization design
Research-validated design for fixing critical DQN issues: - CQL regularization for offline RL training on historical data - IQN integration replacing broken C51 (Candle scatter_add bug) - CVaR risk-aware action selection - State dimension consolidation (51 is canonical) - NaN panic fix, Adam epsilon fix per Rainbow paper Verified against 20+ papers and 2026 SOTA. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
246
docs/plans/2026-02-20-dqn-algorithm-fix-design.md
Normal file
246
docs/plans/2026-02-20-dqn-algorithm-fix-design.md
Normal file
@@ -0,0 +1,246 @@
|
||||
# DQN Algorithm Fix & 2026 Modernization Design
|
||||
|
||||
**Date:** 2026-02-20
|
||||
**Branch:** feature/codebase-cleanup (then new branch for implementation)
|
||||
**Status:** Design approved, pending implementation plan
|
||||
|
||||
## Problem Statement
|
||||
|
||||
The foxhunt DQN implementation has three critical issues making it non-functional:
|
||||
|
||||
1. **C51 distributional RL is disabled** — Candle's `scatter_add` breaks gradient flow (BUG #36), removing the most impactful Rainbow component
|
||||
2. **No offline RL regularization** — Training on historical data without CQL causes Q-value overestimation and divergence
|
||||
3. **State dimension chaos** — 5 different state_dim values (51/54/52/32/64) across configs, causing silent tensor mismatches
|
||||
|
||||
Additionally, several bugs and hyperparameter deviations from the Rainbow paper reduce training stability.
|
||||
|
||||
## Research Validation (2026 SOTA)
|
||||
|
||||
Verified against 20+ papers and reference implementations (Feb 2026):
|
||||
|
||||
| Component | Paper/Source | Our Status | Verdict |
|
||||
|-----------|-------------|------------|---------|
|
||||
| Rainbow DQN architecture | Hessel et al. 2018 | 5/6 correct | SOUND — still mainstream in 2026 |
|
||||
| Double Q-learning | van Hasselt 2016 | Correct | No changes needed |
|
||||
| Dueling Networks | Wang 2016 | Correct | No changes needed |
|
||||
| Prioritized Replay | Schaul 2016 | Correct | No changes needed |
|
||||
| Multi-step Returns | Sutton 1988 | Correct | No changes needed |
|
||||
| Noisy Networks | Fortunato 2018 | Correct | No changes needed |
|
||||
| C51 (scatter_add) | Bellemare 2017 | Algorithm correct, framework broken | Replace with IQN |
|
||||
| QR-DQN/IQN | Dabney 2018a/b | Architecture exists, not integrated | Upgrade & integrate |
|
||||
| CQL offline RL | Kumar et al. 2020 | MISSING | Add — critical for historical training |
|
||||
| CVaR risk-aware policy | RPADiRL 2026, CFA 2025 | compute_cvar() exists, not wired | Wire into action selection |
|
||||
|
||||
### Key Literature Findings
|
||||
|
||||
- **Rainbow DQN remains viable** for discrete trading (2025 SSRN: 25.58% annualized, 7.41% max drawdown)
|
||||
- **Distributional RL is the most impactful** Rainbow component (Ceron & Castro 2021)
|
||||
- **CQL is critical for offline training** — without it, Q-values overestimate on out-of-distribution actions
|
||||
- **IQN > QR-DQN > C51** for trading (adaptive quantiles, direct CVaR, no fixed bins)
|
||||
- **FQF not worth the complexity** — 20% slower for marginal gains over IQN
|
||||
- **Candle is viable** for DQN-class training; monitor Burn framework for future
|
||||
|
||||
## Design
|
||||
|
||||
### Section 1: CQL Regularization (P0)
|
||||
|
||||
**Problem:** Training on historical market data is offline RL. Without conservative Q-value estimation, the agent overestimates Q-values for state-action pairs underrepresented in the training data. This is the likely root cause of Q-value divergence (BUG #37) and training instability.
|
||||
|
||||
**Solution:** Add CQL regularization term to `DQN::train_step()`.
|
||||
|
||||
**Formula (Kumar et al. 2020):**
|
||||
```
|
||||
CQL_loss = td_loss + α * (E_a[logsumexp(Q(s,a))] - E_{a~data}[Q(s, a_data)])
|
||||
```
|
||||
|
||||
This penalizes high Q-values for actions NOT taken in the training data, pushing the learned Q-function to be a lower bound of the true Q-function.
|
||||
|
||||
**Implementation:**
|
||||
- Add `cql_alpha: f64` to `DQNConfig` (default: 1.0, tunable via hyperopt)
|
||||
- Add `use_cql: bool` to `DQNConfig` (default: true for offline training)
|
||||
- In `train_step()`, after computing `current_q_values` (line ~1227):
|
||||
1. Compute `logsumexp(Q(s, all_actions))` — the soft maximum over all Q-values
|
||||
2. Get `Q(s, a_data)` — Q-value for the action actually taken in the dataset
|
||||
3. CQL penalty = `logsumexp - Q_data` (averaged over batch)
|
||||
4. Total loss = `td_loss + cql_alpha * cql_penalty`
|
||||
- The `logsumexp` operation uses standard tensor ops (`exp`, `sum`, `log`) — NO scatter_add
|
||||
|
||||
**Files modified:**
|
||||
- `ml/src/dqn/dqn.rs` — Add CQL config fields, CQL loss computation in train_step()
|
||||
|
||||
**Candle operations needed:** `exp()`, `sum()`, `log()`, `mean_all()` — all battle-tested, no gradient flow risk.
|
||||
|
||||
### Section 2: IQN Integration (P0)
|
||||
|
||||
**Problem:** C51 is broken due to Candle's scatter_add gradient bug. Without distributional RL, the agent cannot model return distributions or tail risk.
|
||||
|
||||
**Solution:** Replace C51 with IQN (Implicit Quantile Networks). The existing `QuantileNetwork` in `quantile_regression.rs` already uses IQN's cosine embedding architecture.
|
||||
|
||||
**Key insight:** Our `QuantileNetwork` IS architecturally IQN — it uses cosine embedding for τ (lines 141-179 of quantile_regression.rs). The upgrade from QR-DQN to true IQN requires only changing fixed τ fractions to random τ sampling during training.
|
||||
|
||||
**Implementation:**
|
||||
|
||||
1. **Expand QuantileNetwork output** from single-action `[batch, num_quantiles]` to multi-action `[batch, num_actions, num_quantiles]`:
|
||||
- Change `output_layer` from `linear(state_dim, 1)` to `linear(state_dim, num_actions)`
|
||||
- Reshape output: `[batch, num_quantiles, num_actions]` → transpose → `[batch, num_actions, num_quantiles]`
|
||||
|
||||
2. **Add IQN to DQN struct:**
|
||||
- New field: `iqn_network: Option<QuantileNetwork>`
|
||||
- New field: `iqn_target_network: Option<QuantileNetwork>`
|
||||
- Created when `distributional_type == DistributionalType::QRDQN` (rename to IQN)
|
||||
|
||||
3. **Add IQN loss path in train_step():**
|
||||
- Sample τ ~ Uniform(0,1) with shape `[batch, num_quantiles]`
|
||||
- Forward: get quantile values for current states `Z(s, a, τ)` → `[batch, num_quantiles]`
|
||||
- Target: get target quantile values with detached target network
|
||||
- Loss: `quantile_huber_loss(predicted, target, τ, κ)` — already implemented
|
||||
- NO scatter_add in the entire path
|
||||
|
||||
4. **IQN action selection in select_action():**
|
||||
- Sample fixed τ fractions (for deterministic action selection at inference)
|
||||
- Compute Q(s,a) = mean over quantiles of Z(s,a,τ)
|
||||
- Argmax over actions
|
||||
- Risk-aware mode: use CVaR instead of mean (Section 3)
|
||||
|
||||
5. **Random τ sampling during training** (the QR-DQN → IQN upgrade):
|
||||
- Training: `τ ~ Uniform(0,1)` with shape `[batch, num_quantiles]` (random per sample)
|
||||
- Inference: `τ_i = (i + 0.5) / N` (fixed, deterministic)
|
||||
- This one change unlocks arbitrary-precision risk metrics at inference time
|
||||
|
||||
**Files modified:**
|
||||
- `ml/src/dqn/quantile_regression.rs` — Expand output to multi-action, add random τ sampling
|
||||
- `ml/src/dqn/dqn.rs` — Add IQN fields to DQN struct, IQN loss path, IQN action selection
|
||||
- `ml/src/dqn/mod.rs` — Update re-exports if needed
|
||||
|
||||
**Candle operations needed:** `matmul`, `broadcast_as`, `cos`, `mul`, `abs`, `where_cond`, `mean` — all standard, no scatter_add.
|
||||
|
||||
### Section 3: CVaR Action Selection (P1)
|
||||
|
||||
**Problem:** Even with distributional RL, using `argmax(mean Q)` ignores tail risk. The agent optimizes for average returns, not worst-case protection.
|
||||
|
||||
**Solution:** Add risk-aware action selection using CVaR (Conditional Value at Risk).
|
||||
|
||||
**Implementation:**
|
||||
- Add `use_cvar_action_selection: bool` to DQNConfig (default: false)
|
||||
- Add `cvar_alpha: f32` to DQNConfig (default: 0.05 = worst 5%)
|
||||
- In `select_action()` when IQN is active:
|
||||
- If `use_cvar_action_selection`: score each action by `CVaR_α(Z(s,a))` instead of `mean(Z(s,a))`
|
||||
- `compute_cvar()` already exists in `quantile_regression.rs:236-246`
|
||||
- Select `argmax_a CVaR_α(Z(s,a))` — the action with the best worst-case outcome
|
||||
|
||||
**Files modified:**
|
||||
- `ml/src/dqn/dqn.rs` — Add CVaR config, modify select_action()
|
||||
|
||||
### Section 4: State Dimension Consolidation (P1)
|
||||
|
||||
**Problem:** 5 different state_dim values across configs cause silent tensor mismatches.
|
||||
|
||||
**The correct value:** 51 (Wave 23: 45 market features + 6 portfolio features)
|
||||
|
||||
**Changes:**
|
||||
1. `dqn::dqn::DQNConfig::default()` — state_dim: 54 → **51**, num_actions: 45 → **45** (already correct)
|
||||
2. `dqn::agent::DQNConfig::default()` — state_dim: 52 → **51**
|
||||
3. Leave `aggressive()/conservative()/emergency()` at 32 (test configs with num_actions: 3)
|
||||
4. Add runtime validation in `DQN::new()`:
|
||||
```rust
|
||||
if config.state_dim == 0 {
|
||||
return Err(MLError::ConfigError("state_dim must be > 0".into()));
|
||||
}
|
||||
```
|
||||
|
||||
**Files modified:**
|
||||
- `ml/src/dqn/dqn.rs` — Fix default state_dim
|
||||
- `ml/src/dqn/agent.rs` — Fix default state_dim
|
||||
|
||||
### Section 5: Bug Fixes (P1)
|
||||
|
||||
**5a. NaN Panic Fix (dqn.rs:1716)**
|
||||
|
||||
```rust
|
||||
// Before (panics on NaN):
|
||||
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
|
||||
// After (safe):
|
||||
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
```
|
||||
|
||||
**5b. Adam Epsilon Fix (rainbow_agent_impl.rs:81)**
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
eps: 1e-8,
|
||||
|
||||
// After (per Rainbow paper, Hessel et al. 2018):
|
||||
eps: 1.5e-4,
|
||||
```
|
||||
|
||||
**5c. Remove Duplicate DistributionalType Enum**
|
||||
|
||||
- Keep the enum in `distributional.rs:21`
|
||||
- Remove duplicate from `quantile_regression.rs:314`
|
||||
- Import from distributional module instead
|
||||
|
||||
**Files modified:**
|
||||
- `ml/src/dqn/dqn.rs` — NaN sort fix
|
||||
- `ml/src/dqn/rainbow_agent_impl.rs` — Adam epsilon
|
||||
- `ml/src/dqn/quantile_regression.rs` — Remove duplicate enum
|
||||
|
||||
## Architecture After Changes
|
||||
|
||||
```
|
||||
DQN::train_step()
|
||||
├── CQL regularization (NEW — offline RL penalty)
|
||||
│ └── logsumexp(Q) - Q(data) → no scatter_add
|
||||
├── IQN quantile Huber loss (NEW — replaces C51)
|
||||
│ ├── Random τ sampling → Uniform(0,1)
|
||||
│ ├── QuantileNetwork forward → [batch, num_actions, num_quantiles]
|
||||
│ └── quantile_huber_loss() → standard tensor ops
|
||||
├── Scalar MSE/Huber loss (existing — standard DQN path)
|
||||
│ └── Works correctly, no changes needed
|
||||
└── PER priority updates (existing — no changes needed)
|
||||
|
||||
DQN::select_action()
|
||||
├── IQN mode (NEW):
|
||||
│ ├── Fixed τ fractions for deterministic selection
|
||||
│ ├── Mean Q-values = mean over quantiles
|
||||
│ └── CVaR mode: argmax CVaR_α(Z(s,a)) for risk-aware selection
|
||||
├── Standard mode (existing): argmax Q(s,a)
|
||||
└── Noisy Networks mode (existing): no epsilon needed
|
||||
```
|
||||
|
||||
## What We're NOT Doing (YAGNI)
|
||||
|
||||
- **FQF** — 20% slower than IQN for marginal gains. Not worth it.
|
||||
- **Decision Transformer** — Requires GPT-2 weights + LoRA infrastructure. Not available in Candle.
|
||||
- **SAC** — Good for continuous control but foxhunt already has PPO for that.
|
||||
- **Framework switch** — Candle works for DQN-class models. Monitor Burn for future.
|
||||
- **Multi-agent RL** — Ensemble voting (already implemented) is simpler and proven.
|
||||
- **LLM features** — Feature engineering question, not RL architecture.
|
||||
|
||||
## Safety Gates
|
||||
|
||||
Every task must pass:
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check --workspace # Zero errors
|
||||
SQLX_OFFLINE=true cargo test -p ml # All tests pass
|
||||
```
|
||||
|
||||
## Files Changed Summary
|
||||
|
||||
| File | Changes |
|
||||
|------|---------|
|
||||
| `ml/src/dqn/dqn.rs` | CQL config+loss, IQN fields, IQN loss path, CVaR action selection, state_dim fix, NaN fix |
|
||||
| `ml/src/dqn/quantile_regression.rs` | Multi-action output, random τ sampling, remove duplicate enum |
|
||||
| `ml/src/dqn/agent.rs` | state_dim: 52 → 51 |
|
||||
| `ml/src/dqn/rainbow_agent_impl.rs` | Adam eps: 1e-8 → 1.5e-4 |
|
||||
| `ml/src/dqn/mod.rs` | Update re-exports if needed |
|
||||
|
||||
## References
|
||||
|
||||
- Hessel et al. (2018) "Rainbow: Combining Improvements in Deep RL" — AAAI
|
||||
- Kumar et al. (2020) "Conservative Q-Learning for Offline RL" — NeurIPS
|
||||
- Dabney et al. (2018a) "Distributional RL with Quantile Regression" — AAAI
|
||||
- Dabney et al. (2018b) "Implicit Quantile Networks" — ICML
|
||||
- Ceron & Castro (2021) "Revisiting Rainbow" — ablation study
|
||||
- RPADiRL (2026) "Personalized Risk-Sensitive DiRL for Stock Trading" — Applied Soft Computing
|
||||
- CFA Institute (2025) "RL and Inverse RL Practitioner's Guide" — Chapter 6
|
||||
Reference in New Issue
Block a user