Прекратить регрессию: обучение функций ценности с помощью классификации для масштабируемого глубокого обучения с подкреплением

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

March 6, 2024

Авторы: 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

cs.AI

Аннотация

Функции ценности являются центральным компонентом глубокого обучения с подкреплением (RL). Эти функции, параметризованные нейронными сетями, обучаются с использованием среднеквадратичной ошибки регрессии для соответствия бутстрапированным целевым значениям. Однако масштабирование методов RL на основе значений, использующих регрессию к большим сетям, таким как высокоемкие трансформеры, оказалось сложной задачей. Эта сложность находится в явном противоречии с обучением с учителем: используя потери классификации кросс-энтропии, методы обучения с учителем надежно масштабируются к массивным сетям. Наблюдая это расхождение, в данной статье мы исследуем, может ли масштабируемость глубокого RL также быть улучшена просто путем использования классификации вместо регрессии для обучения функций ценности. Мы демонстрируем, что функции ценности, обученные с использованием категориальной кросс-энтропии, значительно улучшают производительность и масштабируемость в различных областях. Среди них: однозадачное RL на играх Atari 2600 с использованием SoftMoEs, многозадачное RL на Atari с крупномасштабными ResNets, робототехническое манипулирование с Q-трансформерами, игра в шахматы без поиска и задача языкового агента Wordle с высокоемкими трансформерами, достигая передовых результатов в этих областях. Через тщательный анализ мы показываем, что преимущества категориальной кросс-энтропии в основном происходят от ее способности смягчать проблемы, присущие RL на основе значений, такие как шумные цели и нестационарность. В целом, мы утверждаем, что простое переключение на обучение функций ценности с использованием категориальной кросс-энтропии может привести к существенным улучшениям в масштабируемости глубокого RL при минимальных или отсутствующих затратах.

English

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.