Stop Regressing: Training Value Functions via Classification for Scalable Deep RL

Jesse Farebrother; Jordi Orbay; Quan Vuong; Adrien Ali Taïga; Yevgen Chebotar; Ted Xiao; Alex Irpan; Sergey Levine; Pablo Samuel Castro; Aleksandra Faust; Aviral Kumar; Rishabh Agarwal

Source-linked AI summary

Stop Regressing: Training Value Functions via Classification for Scalable Deep RL

Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taïga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, Aviral Kumar, Rishabh Agarwal

arXiv:2403.03950v1Mar 6, 2024cs.LGcs.AIstat.ML

TL;DR

Deep RL commonly trains value functions by MSE regression, but scaling this approach to large networks remains difficult. The paper replaces regression with categorical cross-entropy, finding broad performance and scalability gains across RL domains while linking them to improved handling of noisy and non-stationary targets.

Problem
Value-based deep RL relies mainly on regression, yet regression-based methods have been difficult to scale to large networks such as high-capacity Transformers.
Method
The paper evaluates categorical cross-entropy methods, especially HL-Gauss, as replacements for MSE value-function training across online and offline RL settings.
Results
Categorical cross-entropy substantially improves performance and scalability across Atari, Wordle, chess, and robotic manipulation, with HL-Gauss reporting gains from 30% to 70% across highlighted domains.
Takeaways & Limitations
Cross-entropy can serve as a broadly useful replacement for MSE in value-based RL, with benefits attributed to better handling of noisy and non-stationary targets.
Takeaways & Limitations
Further evaluation is needed in RL settings involving pre-training, fine-tuning, or continual learning.

Abstract

from arXiv · show

Value functions are a central component of deep reinforcement learning (RL). These functions, parameterized by neural networks, are trained using a mean squared error regression objective to match bootstrapped target values. However, scaling value-based RL methods that use regression to large networks, such as high-capacity Transformers, has proven challenging. This difficulty is in stark contrast to supervised learning: by leveraging a cross-entropy classification loss, supervised methods have scaled reliably to massive networks. Observing this discrepancy, in this paper, we investigate whether the scalability of deep RL can also be improved simply by using classification in place of regression for training value functions. We demonstrate that value functions trained with categorical cross-entropy significantly improves performance and scalability in a variety of domains. These include: single-task RL on Atari 2600 games with SoftMoEs, multi-task RL on Atari with large-scale ResNets, robotic manipulation with Q-transformers, playing Chess without search, and a language-agent Wordle task with high-capacity Transformers, achieving state-of-the-art results on these domains. Through careful analysis, we show that the benefits of categorical cross-entropy primarily stem from its ability to mitigate issues inherent to value-based RL, such as noisy targets and non-stationarity. Overall, we argue that a simple shift to training value functions with categorical cross-entropy can yield substantial improvements in the scalability of deep RL at little-to-no cost.

1. Introduction

The paper asks whether replacing value-function regression with classification can improve deep RL scalability, motivated by classification’s success with large neural networks. It finds that categorical cross-entropy, especially HL-Gauss, improves performance, robustness, and scaling across diverse architectures and domains.

Motivation: Value-based deep RL mainly uses MSE regression, yet scaling regression-based methods to high-capacity networks such as Transformers has been challenging.This contrasts with supervised learning, where classification has supported effective scaling to large neural networks.
Findings: 30% better performance with Mixture-of-Experts on single-task Atari, 1.8–2.1× in multi-task Atari, 40% on Wordle, 70% in search-free chess, and 67% in robotic manipulation.These gains are reported for HL-Gauss relative to traditional regression-based approaches across different domains and architectures.
Approach: The study evaluates methods for deriving classification labels and uses categorical cross-entropy to train value functions instead of MSE regression.The approach is presented as a potential drop-in replacement for the MSE loss.
Analysis: Diagnostic experiments suggest that cross-entropy mitigates noisy-target problems and helps networks use capacity to fit non-stationary targets.The analysis is intended to explain the empirical gains over mean-squared regression.

2. Preliminaries and Background

Deep RL learns policies from discounted returns and commonly trains neural value functions by regressing toward bootstrapped Bellman targets. The paper reframes this regression as learning a categorical target distribution whose expectation recovers the scalar value, optimized with cross-entropy.

Regression as Classification: Regression as classification models the target as a conditional distribution, learns a parameterized distribution, and recovers the prediction as its expectation.The target distribution is constructed so that its mean equals the scalar regression target, while the predicted distribution is trained toward it with cross-entropy.
Categorical Representation: The categorical representation restricts predictions to m evenly spaced support locations between v_min and v_max, with probabilities assigned to the corresponding classes.The target distribution must be constructed and projected onto this categorical support for RL.
Reinforcement Learning: A value function estimates expected discounted return, while DQN approximates the optimal action-value function with a neural network.The return is the discounted sum of future rewards, and the action-value function conditions that return on a state and action.
Value-Based RL: DQN trains on transitions using a temporal-difference target produced by a slowly updated target network and a sample Bellman optimality operator.Most deep RL value-learning algorithms use variations of this target-network regression recipe.
Paper Objective: The paper replaces the squared TD-error objective with a classification-style cross-entropy objective for value-based and actor-critic methods in online and offline RL.This is the central methodological shift developed from the preceding RL setup.

3. Value-Based RL with Classification

The paper casts value-function TD learning as classification by representing action values with categorical distributions and training them using cross-entropy. It develops scalar-target and distributional constructions that preserve target information or distribute probability across neighboring value locations.

Value-Based RL with Classification: Value-function TD learning replaces squared-error regression between scalar Q-values and TD targets with cross-entropy between categorical distributions.The predicted action value is represented as the expectation of a categorical distribution whose probabilities come from softmax logits.
Value-Based RL with Classification: The categorical representation uses m evenly spaced classes between v_min and v_max, with probabilities assigned to those value locations.The target distribution must use the same categorical support as the prediction so the cross-entropy loss can be computed directly.
3.1. Constructing Categorical Distributions from Scalars: Two-Hot represents a scalar TD target exactly by placing probability on its two bounding locations, avoiding the lossy errors of one-hot discretization.All other locations receive zero probability, but the method does not exploit the natural ordering among classes.
3.1. Constructing Categorical Distributions from Scalars: HL-Gauss converts a Gaussian centered on the scalar TD target into histogram probabilities spanning neighboring bins, with σ² controlling label smoothing.This distributes mass beyond the two nearest locations and uses the target distribution's variance as a smoothing hyperparameter.
3.2. Modelling the Categorical Return Distribution: Categorical distributional RL directly models future-return distributions, shifting and scaling categorical locations before projecting mass proportionally onto neighboring support points.The stochastic distributional Bellman operator motivates the categorical projection used by CDRL.

4. Evaluating Classification Losses in RL

Across online and offline Atari tasks, categorical losses—especially HL-Gauss—outperform MSE and improve scaling across larger networks, multi-task settings, and language-based offline RL. The results also show greater stability during prolonged offline training and substantial gains over prior methods.

4.1. Single-Task RL on Atari Games: HL-Gauss consistently outperforms MSE across online Atari, offline Atari, multi-task Atari, and Wordle evaluations.It also surpasses C51 in both online and offline Atari settings.
4.1. Single-Task RL on Atari Games: Cross-entropy losses remain stable during prolonged offline training, whereas MSE performance degrades; Two-Hot is more stable than MSE but underperforms other classification methods.This pattern appears in the Atari offline RL comparison.
4.2. Scaling Value-based RL to Large Networks: ≈30% IQM improvement is provided by HL-Gauss with SoftMoE, while SoftMoE mitigates negative scaling observed with MSE alone.HL-Gauss outperforms all regression configurations even with a single expert.
4.2.1. Scaling with Mixture-of-Experts: HL-Gauss scales reliably better than MSE on multi-task Asteroids, with larger networks improving HL-Gauss performance while MSE degrades.The comparison uses normalized scores relative to a baseline IMPALA agent with MSE loss.
4.2.2. Training Generalist Policies with ResNets: ≈45% improvement over the best prior multi-game result is achieved with ResNet-101, measured by IQM human normalized score.HL-Gauss also continues improving with model capacity while MSE typically plateaus beyond ResNet-34.
4.3.1. Language Agent: Wordle: HL-Gauss yields substantially higher Wordle success rates than MSE across varying strengths of behavior regularization.The task evaluates one-turn word-guessing success from partially played games using an offline-trained Transformer.

5. Why Does Classification Benefit RL?

The experiments attribute classification’s gains primarily to cross-entropy itself, distributed targets, improved representations, and greater robustness to noisy and non-stationary value targets.

Categorical representations: Cross-entropy, rather than softmax parameterization alone, provides the bulk of the performance improvement over MSE.Adding softmax to MSE produced no gains in online or offline Atari RL.
Target distributions: HL-Gauss outperforms Two-Hot because distributing probability across neighboring bins reduces overfitting and exploits ordinal structure.The best smoothing scale appears independent of the number of bins.
Robustness to noisy targets: HL-Gauss degrades more gracefully than MSE as reward noise increases, indicating greater robustness to noisy targets.The comparison varies additive reward noise in offline Atari RL.
Robustness to noisy targets: With stochastic dynamics, HL-Gauss outperforms MSE, whereas under deterministic dynamics they perform comparably and both outperform C51.Stochasticity is introduced through sticky actions in Atari.
Representations and non-stationarity: Cross-entropy-trained value networks learn more expressive representations and better handle non-stationary targets than regression-based networks.Linear probing evaluates whether frozen representations retain information needed to relearn policies.

6. Related Work

The paper builds on prior regression-as-classification and categorical distributional RL work while isolating cross-entropy’s role in scalable value-based RL.

Regression as classification: HL-Gauss extends prior regression-as-classification work by enabling scalable value-based RL across Atari, robotic manipulation, chess, and Wordle.Earlier HL-Gauss results focused on small-scale supervised regression tasks outside RL.
Categorical losses: Two-Hot is reported to perform worse than other cross-entropy losses and MSE because it does not distribute probability effectively across neighboring classes.This contrasts with C51 and HL-Gauss, which distribute probability across nearby values.
Categorical distributional RL: The paper distinguishes its cross-entropy findings from distributional RL theory, where statistical benefits are orthogonal to categorical representation or cross-entropy objectives.Related work also includes analyses of categorical one-step distributional RL and scaling in offline value-based RL.
Representations: Linear probing addresses a gap in prior representation analyses by comparing representations learned with cross-entropy losses against those learned with MSE.Earlier studies focused primarily on MSE-trained TD-learning representations.

7. Conclusion

The paper concludes that replacing MSE with categorical cross-entropy substantially improves value-based RL across tasks and architectures by improving representations and handling noise and non-stationarity better.

Conclusion: Framing regression as classification and minimizing categorical cross-entropy yields large performance and scalability improvements across diverse value-based RL tasks and architectures.The conclusion presents this as a small change with a substantial effect.
Conclusion: The gains stem specifically from cross-entropy enabling more expressive representations and better handling noise and non-stationarity.The loss does not fully alleviate these problems, but produces a substantial difference.
Future directions: The findings suggest classification may provide a smoother path for transferring advances in value-based RL to Transformer architectures.The authors frame this as an implication for future algorithm design.
Future directions: Further evaluation is needed for settings involving pre-training, fine-tuning, or continual reinforcement learning.These settings were not covered by the experiments discussed in the conclusion.

A. Reference Implementations

The reference implementations define HL-Gauss by transforming scalar targets into categorical probabilities and recovering scalar predictions from bin probabilities.

JAX implementation: The target transform computes normalized bin probabilities from error-function evaluations over evenly spaced support points.The normalization uses the difference between endpoint CDF evaluations.
JAX implementation: The inverse transform recovers a scalar prediction as the probability-weighted sum of adjacent-bin centers.Bin centers are computed from neighboring support points.
PyTorch implementation: HL-Gauss converts logits into categorical probabilities using cross-entropy against target probabilities derived from a histogram transform.The implementation exposes separate forward and target-to-probability transformations.

B. Experimental Methodology

The experiments evaluate regression and classification value-learning methods across online, offline, multi-task, multi-game, and chess settings using established RL architectures and task protocols.

Experiments compare DQN+Adam regression with C51, Two-Hot, and HL-Gauss classification implementations across the reported domains.The regression baselines use Dopamine's Jax DQN+Adam implementation, while classification methods build on its C51 implementation.
Atari: The Atari online and offline studies use the standard DQN Nature architecture, with three convolutional layers followed by one nonlinear fully connected layer.Online results cover 60 Atari games with five seeds per game; offline results cover 17 games with three seeds per game.
SoftMoE: SoftMoE experiments replace DQN's penultimate layer, use Impala ResNets, reuse 20 games, and run each configuration for five seeds per game.Classification methods reuse the reported C51, Two-Hot, and HL-Gauss parameter settings.
Multi-task and multi-game: Multi-task and multi-game experiments follow prior protocols, using five seeds per game for multi-task runs and one seed per configuration for computationally expensive multi-game runs.The multi-task studies reuse HL-Gauss hyperparameters from Table B.3.
Chess: The chess setup projects Stockfish action-values into a categorical distribution with HL-Gauss using 128 bins over [0, 1] and smoothing ratio σ/ζ = 0.75.Each parameter configuration trains a single agent.

B.3. Robotic manipulation experiments.

The robotic manipulation study uses a vision-based mobile manipulator and a Q-Transformer, while a synthetic experiment tests whether classification remains effective for non-stationary, increasing-magnitude targets.

Robotic manipulation: The robotic domain is a simulated tabletop with randomized objects, transformed by RetinaGAN to make visual inputs more realistic.The study uses a mobile manipulator with seven degrees of freedom.
Synthetic robustness: The synthetic network takes CIFAR-10 images as inputs and predicts a scalar, with randomly sampled target parameters matching the convolutional architecture.The construction is intended to test non-stationary targets whose magnitude increases without making regression intrinsically more difficult.
Synthetic robustness: The synthetic setup models TD learning by increasing target biases through b ∈ {0, 8, 16, 24, 32} and resampling target parameters at each stage.Each stage uses 5,000 gradient steps with batch size 512 and Adam at learning rate 10^-3.
Synthetic robustness: The synthetic experiment finds that L2 regression loses the ability to rapidly fit targets of increasing magnitude, whereas Two-Hot and HL-Gauss are evaluated as classification alternatives.MSE is tracked across training and averaged over 30 seeds, using 101 bins over [-40, 40] for both classification methods.

C. Per-Game Atari Results

Across 60 Atari games, HL-Gauss generally outperforms MSE, with per-game summaries showing broad gains and substantial advantages on many games.

Training setup: Training curves cover all 60 Atari games over 200M frames for DQN(Adam), C51, Two-Hot, and HL-Gauss.The comparison includes both regression and categorical value-learning methods.
Per-game comparison: HL-Gauss outperforms MSE in approximately three-quarters of Atari games and scores at least 10% higher in half of them.Figure C.2 summarizes relative final performance and IQM-normalized training curves.
Per-game comparison: The per-game analysis compares HL-Gauss against MSE using relative final performance and IQM-normalized training curves.These summaries aggregate the training curves shown in Figure C.1.

D. Additional Results

Additional Atari scaling results consistently favor HL-Gauss over MSE as architectures grow, and aggregate offline results show faster improvement than C51 with increasing parameter count.

Multi-task Atari: On Space Invaders with 29 concurrently trained variants, HL-Gauss scales better than MSE for every evaluated architecture.The result is shown for each architecture's training curves.
Multi-task Atari: In Space Invaders scaling experiments with IMPALA ResNets, HL-Gauss outperforms MSE for all models.Scores are normalized against a baseline IMPALA agent with MSE loss because human scores are unavailable for the variants.
Multi-task Atari: On Asteroids with 63 concurrently trained variants, HL-Gauss keeps improving beyond 1300M frames while MSE performance drops for architectures larger than ResNet-18.The larger architectures also reach higher peak IQM scores with HL-Gauss.
Offline RL: In multi-task offline RL, HL-Gauss improves faster than C51 as the number of parameters scales up under DQN normalization.The comparison is reported using aggregate normalized results.