foxhunt/docs/plans/2026-02-20-dqn-algorithm-fix-design.md

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():

   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)

// 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)

// 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:

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