From 91af2b933481f9fe44e3dfbc6f465a2c7b348c7d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 20 Feb 2026 14:39:58 +0100 Subject: [PATCH] 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 --- .../2026-02-20-dqn-algorithm-fix-design.md | 246 ++++++++++++++++++ 1 file changed, 246 insertions(+) create mode 100644 docs/plans/2026-02-20-dqn-algorithm-fix-design.md diff --git a/docs/plans/2026-02-20-dqn-algorithm-fix-design.md b/docs/plans/2026-02-20-dqn-algorithm-fix-design.md new file mode 100644 index 000000000..ae14a422a --- /dev/null +++ b/docs/plans/2026-02-20-dqn-algorithm-fix-design.md @@ -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` + - New field: `iqn_target_network: Option` + - 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