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

2024-03-07

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.

Machine Learning,Artificial Intelligence

What problem does this paper attempt to address?

This paper mainly discusses the problem of value function training in deep reinforcement learning (RL). Traditional value function training uses mean squared error regression loss, but this method encounters difficulties when scaling to large networks such as high-capacity Transformers. The authors propose an alternative approach, which is to train the value function using classification loss (cross-entropy), in order to improve performance and scalability.

The paper points out that classification methods have been successfully applied in large-scale networks in supervised learning, so the research team explores whether this idea can be applied to value function training in RL. They found that by converting continuous target values into discrete categories and using classification cross-entropy loss, performance can be significantly improved in various domains, including Atari games, multi-task RL, robot manipulation, search-free chess games, and language agent tasks (Wordle).

The experimental results show that classification cross-entropy loss outperforms traditional regression methods in these domains, especially in high-capacity networks like Transformers. Classification methods are better able to handle noisy targets and non-stationarity issues, thereby improving the stability and generalization ability of RL algorithms.

In conclusion, this paper proposes the strategy of transforming regression problems into classification problems to improve value function training in deep reinforcement learning. It proves the effectiveness of this approach in multiple tasks and network architectures, providing a new perspective for the future development of more efficient RL algorithms.