Stop Regressing: Training Value Functions via Classification for
  Scalable Deep RL
Authors: name: Jesse Farebrother; name: Jordi Orbay; name: Quan Vuong; name: Adrien Ali Taïga; name: Yevgen Chebotar; name: Ted Xiao; name: Alex Irpan; name: Sergey Levine; name: Pablo Samuel Castro; name: Aleksandra Faust; name: Aviral Kumar; name: Rishabh Agarwal
Published: 2024-03-06 18:55:47

Abstract: 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.