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

📄 arXiv: 2403.03950v1 📥 PDF

作者: 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.LG, cs.AI, stat.ML

发布日期: 2024-03-06

💡 一句话要点

通过分类训练价值函数以提升深度强化学习的可扩展性

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 深度强化学习 价值函数 分类训练 可扩展性 噪声目标 非平稳性 Transformer 游戏AI

📋 核心要点

  1. 现有的价值基础强化学习方法在大规模网络上扩展时面临困难，尤其是使用回归训练时。
  2. 本文提出通过使用分类交叉熵替代均方误差回归来训练价值函数，以提高深度强化学习的可扩展性。
  3. 实验结果表明，使用分类交叉熵训练的价值函数在多个任务中表现出显著的性能提升，达到了最先进的水平。

📝 摘要（中文）

价值函数是深度强化学习（RL）的核心组成部分，通常通过均方误差回归目标进行训练。然而，使用回归的价值基础RL方法在大规模网络（如高容量Transformer）上扩展时面临挑战。本文探讨了通过将回归替换为分类来改善深度RL的可扩展性。我们展示了使用分类交叉熵训练的价值函数在多个领域显著提升了性能和可扩展性，包括Atari 2600游戏的单任务RL、多任务RL、机器人操作、无搜索的国际象棋以及高容量Transformer的语言代理Wordle任务，均取得了最先进的结果。通过仔细分析，我们表明分类交叉熵的优势主要源于其缓解了价值基础RL固有的噪声目标和非平稳性问题。

🔬 方法详解

问题定义：本文旨在解决深度强化学习中价值函数训练的可扩展性问题。现有方法使用均方误差回归进行训练，导致在大规模网络上性能不佳，尤其是在处理噪声目标和非平稳性时。

核心思路：论文提出通过使用分类交叉熵损失函数替代传统的回归损失，来训练价值函数。分类方法在处理大规模数据时表现更为稳定，能够有效缓解价值基础RL中的一些固有问题。

技术框架：整体架构包括使用深度神经网络（如SoftMoEs和ResNets）作为价值函数的参数化模型，通过分类交叉熵损失进行训练。模型在多个任务上进行评估，包括单任务和多任务的强化学习。

关键创新：最重要的创新在于将分类方法引入价值函数训练，显著提升了深度强化学习的可扩展性和性能。这一方法与传统的回归训练在本质上有所区别，后者在处理复杂任务时常常受到限制。

关键设计：在实验中，使用了分类交叉熵作为损失函数，结合高容量的Transformer架构。具体的参数设置和网络结构设计经过优化，以确保在不同任务中都能实现最佳性能。实验还考虑了噪声目标和非平稳性对训练的影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，使用分类交叉熵训练的价值函数在Atari 2600游戏、机器人操作和Wordle任务中均取得了显著的性能提升，达到了最先进的结果。例如，在Atari游戏中，性能提升幅度超过了20%，在无搜索的国际象棋任务中也表现出色。

🎯 应用场景

该研究的潜在应用领域包括游戏AI、机器人操作和自然语言处理等。通过提升深度强化学习的可扩展性，能够在更复杂的环境中实现更高效的决策和控制，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.