Research map / Reinforcement learning / Policy gradient and actor-critic
Stop Regressing: Training Value Functions via Classification for Scalable Deep RL
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that replacing MSE regression with categorical cross-entropy, especially HL-Gauss target distributions, improves performance and scalability of value-based deep RL across Atari, multi-game Atari, Wordle, chess distillation, and robotic manipulation.
Problem
Value-based deep RL is trained via MSE regression to bootstrapped targets and has struggled to scale to large networks, unlike supervised learning which scales with cross-entropy classification.
Approach
Recasts TD regression as classification: a softmax over fixed bins is trained with cross-entropy against a target distribution built from the scalar TD target, with HL-Gauss (Gaussian-smoothed histogram targets) as the most effective choice; the loss is a drop-in replacement for MSE.
Key results
- HL-Gauss substantially outperforms Two-Hot and MSE in online Atari DQN over 60 games with 200M frames, and also improves over C51.
- HL-Gauss gives about 67% higher peak performance than the regression baseline on large-scale robotic manipulation with Q-Transformer.
- HL-Gauss yields about 45% improvement over the best prior multi-game Atari result with ResNet-101 (80M parameters) by IQM human-normalized score.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| HL-Gauss cross-entropy loss replaces MSE regression loss in DQN | DQN+Adam · Atari 60 games online | IQM human-normalized score: substantially outperforms MSE | Helped |
| Two-Hot cross-entropy loss replaces MSE in DQN | DQN+Adam · Atari 60 games online | IQM human-normalized score | Mixed |
| HL-Gauss loss replaces MSE in offline CQL | CQL · Atari 10% DQN replay offline | behavior-policy normalized score | Helped |
| HL-Gauss loss replaces MSE in SoftMoE DQN across expert counts | SoftMoE DQN · Atari 20 games, 1-8 experts | performance vs MSE: consistent constant-factor improvement (~30% per text) | Helped |
| HL-Gauss critic loss replaces MSE in IMPALA with scaled ResNets | IMPALA multi-task · Asteroids and Space Invaders multi-task | performance with scaling | Helped |
| HL-Gauss replaces C51 loss in multi-game offline setup | ResNet-101 multi-game · 40-game offline Atari | IQM human-normalized score: about 45% over best prior result | Helped |
| HL-Gauss targets replace MSE in Wordle offline Q-learning | GPT-like Transformer 125M · Wordle | performance: 40% better per intro | Helped |
| HL-Gauss targets replace 1-Hot for chess distillation | Causal Transformer 270M · Chess puzzles without search | puzzle accuracy: 70% improvement per intro | Helped |
| HL-Gauss replaces MSE regression loss in Q-Transformer | Q-Transformer 60M · Robotic manipulation | peak performance: 67% higher peak performance | Helped |
| Softmax parameterization of Q-values trained with MSE (no cross-entropy) (ablation) | DQN MSE · Atari online and offline RL | performance | No clear effect |
| Increase HL-Gauss smoothing ratio sigma/bin width across wide range vs Two-Hot (ablation) | HL-Gauss DQN · Atari 13-game subset, online | IQM score | Helped |
| HL-Gauss loss replaces MSE under noisy rewards (ablation) | CQL · Offline Atari with noisy rewards (eta 0.1,0.3,1.0) | performance degradation | Helped |
Surprising findings
- Two-Hot cross-entropy, despite being a cross-entropy loss, underperforms other classification losses and sometimes MSE because it does not spread probability to neighboring classes.
- Softmax parameterization alone with MSE gives no gains; the cross-entropy loss accounts for most of the improvement.
Limitations
- Cross-entropy alone does not fully resolve the representation, stability, or non-stationarity problems it is shown to partially address.
- Multi-game setup was run with a single seed due to prohibitive compute.
- Evaluation limited to specific settings; pre-training, fine-tuning, and continual RL not explored.
Experimental setup
Datasets: Arcade Learning Environment (60 Atari games) Atari DQN replay dataset (10%) Wordle offline dataset Lichess puzzles Stockfish 16 annotated chess games simulated robotic manipulation dataset
Benchmarks: Atari 2600 (60 games) multi-task Atari Asteroids and Space Invaders multi-game Atari (40 games) Wordle chess puzzles (Lichess) large-scale robotic manipulation
Metrics: IQM human-normalized score optimality gap behavior-normalized score puzzle accuracy success rate MSE linear probing performance
Baselines: DQN (MSE) C51 Two-Hot SoftMoE + MSE IMPALA MSE Q-Transformer MSE AlphaZero with MCTS 1-Hot classification
Builds on: DQN C51 CQL IMPALA Q-Transformer HL-Gauss
Closest papers by meaning
- Bigger, Regularized, Categorical: High-Capacity Value Functions are Efficient Multi-Task Learners (NeurIPS 2025 · similarity 0.88)
- Compute-Optimal Scaling for Value-Based Deep RL (NeurIPS 2025 · similarity 0.88)
- Value-Based Deep RL Scales Predictably (ICML 2025 · similarity 0.87)
- Symmetric Reinforcement Learning Loss for Robust Learning on Diverse Tasks and Model Scales (ICML 2025 · similarity 0.87)
- Is Value Learning Really the Main Bottleneck in Offline RL? (NeurIPS 2024 · similarity 0.87)
- Mixtures of Experts Unlock Parameter Scaling for Deep RL (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.