使用离线强化学习对通才代理进行预训练

1725522943903.jpg

强化学习 (RL) 算法可以学习技能来解决决策任务,例如玩游戏、使机器人拾取物体,甚至优化微芯片设计。但是,在现实世界中运行 RL 算法需要昂贵的主动数据收集。事实证明,对不同的数据集进行预训练可以针对自然语言处理(NLP) 和视觉问题中的各个下游任务进行数据高效的微调。就像BERT或GPT-3模型为 NLP 提供通用初始化一样,大型 RL 预训练模型可以为决策提供通用初始化。因此,我们要问一个问题:我们能否启用类似的预训练来加速 RL 方法并为跨各种任务的高效 RL 创建通用的“主干”?

在即将于ICLR 2023上发表的 “离线 Q 学习在多样化多任务数据上既可扩展又具有泛化能力”中,我们讨论了如何扩展离线 RL,它可用于在之前收集的静态数据集上训练价值函数,以提供这种通用的预训练方法。我们证明,使用多样化数据集的缩放 Q 学习足以学习有助于快速迁移到新任务的表示,以及在任务的新变体上快速在线学习,这比现有的表示学习方法甚至使用更大模型的基于 Transformer 的方法都有显著改进。

缩放 Q 学习:使用保守 Q 学习进行多任务预训练

为了提供通用的预训练方法,离线 RL 需要具有可扩展性,使我们能够对不同任务的数据进行预训练,并利用富有表现力的神经网络模型来获取强大的预训练主干,专门用于各个下游任务。我们的离线 RL 预训练方法基于保守 Q 学习(CQL),这是一种简单的离线 RL 方法,它将标准Q 学习更新与额外的正则化器相结合,以最小化未见动作的价值。对于离散动作,CQL 正则化器相当于标准交叉熵损失,这是对标准深度 Q 学习的简单一行修改。一些关键的设计决策使这成为可能:

神经网络大小:我们发现多游戏 Q 学习需要大型神经网络架构。虽然之前的方法通常使用相对较浅的卷积网络,但我们发现像ResNet 101这样的大型模型比较小的模型有显著的改进。

神经网络架构:为了学习可用于新游戏的预训练主干,我们的最终架构使用共享神经网络主干,其中单独的 1 层头输出每个游戏的 Q 值。这种设计避免了预训练期间游戏之间的干扰,同时仍提供足够的数据共享来学习单个共享表示。我们的共享视觉主干还利用了学习位置嵌入(类似于 Transformer 模型)来跟踪游戏中的空间信息。

表征正则化:最近的研究发现,Q 学习容易受到表征 崩溃 问题的影响,即使是大型神经网络也可能无法学习有效的表征。为了解决这个问题,我们利用之前的工作来规范化 Q 网络共享部分的最后一层特征。此外,我们在 Q 学习中使用了分类分布 RL损失,众所周知,它可以提供更丰富的表征,从而提高下游任务的性能。

多任务 Atari 基准测试

我们在一系列Atari 游戏上评估了我们的可扩展离线 RL 方法,其目标是使用来自低质量(即次优)玩家的异构数据训练单个 RL 代理玩一系列游戏,然后使用生成的网络主干快速学习预训练游戏的新变体或全新的游戏。即使使用标准的在线深度 RL 方法,训练一个可以玩许多不同 Atari 游戏的单一策略也已经足够困难,因为每个游戏都需要不同的策略和不同的表示。在离线设置中,一些先前的研究,如多游戏决策变压器,提出完全放弃 RL,而是利用条件模仿学习,尝试使用大型神经网络架构(如变压器)进行扩展。然而,在这项工作中,我们表明,这种多游戏预训练可以通过 RL 有效地完成,方法是结合使用 CQL 和一些精心的设计决策,我们将在下面描述。

训练游戏的可扩展性

我们使用两种数据组合来评估 Scaled Q-Learning 方法的性能和可扩展性:(1) 接近最优的数据,包括之前 RL 运行的重放缓冲区中出现的所有训练数据,以及 (2) 低质量数据,包括重放缓冲区中前 20% 次试验的数据(即,仅来自极次优策略的数据)。在下面的结果中,我们将具有 8000 万参数模型的 Scaled Q-Learning 与具有 4000 万或 8000 万参数模型的多游戏决策转换器 (DT) 以及行为克隆(模仿学习)基线 (BC) 进行了比较。我们观察到,Scaled Q-Learning 是唯一一种比离线数据有所改进的方法,可达到人类标准化性能的约 80%。

此外,如下所示,Scaled Q-Learning 在性能方面有所改进,但它也具有良好的扩展特性:就像预训练语言和视觉模型的性能随着网络规模的扩大而提高一样,享受通常所说的“幂律扩展”,我们表明 Scaled Q-Learning 的性能具有类似的扩展特性。虽然这可能并不令人惊讶,但这种扩展在 RL 中一直难以实现,性能通常会随着模型规模的增大而下降。这表明,Scaled Q-Learning 与上述设计选择相结合,可以更好地释放离线 RL 利用大型模型的能力。

对新游戏和新变化进行微调

为了评估从这种离线初始化进行的微调,我们考虑了两种设置:(1)使用来自该游戏的少量离线数据(对应于 2M 次游戏转换)对一款全新的、完全未见过的游戏进行微调,以及(2)对具有在线互动的游戏的新变体进行微调。从离线游戏数据进行微调如下所示。请注意,这种情况通常更有利于模仿式方法、决策转换器和行为克隆,因为新游戏的离线数据质量相对较高。尽管如此,我们看到在大多数情况下,缩放 Q 学习比其他方法(平均提高 80%)以及专用的表示学习方法(如MAE或CPC )有更好的效果,这些方法仅使用离线数据来学习视觉表示而不是价值函数。

在线设置中,我们发现使用 Scaled Q-Learning 进行预训练可以带来更大的改进。在这种情况下,MAE 等表征学习方法在在线强化学习期间的改进微乎其微,而 Scaled Q-Learning 可以成功整合有关预训练游戏的先验知识,在 20k 个在线交互步骤后显著提高最终得分。

这些结果表明,使用多任务离线强化学习预训练通用价值函数主干可以显著提高强化学习在下游任务上的性能,无论是在离线还是在线模式下。请注意,这些微调任务非常困难:各种 Atari 游戏,甚至同一游戏的不同变体,在外观和动态上都存在很大差异。例如,Breakout 中的目标块在游戏变体中消失(如下所示),使得控制变得困难。然而,Scaled Q 学习的成功,尤其是与 MAE 和 CPC 等视觉表征学习技术相比,表明该模型实际上是在学习游戏动态的一些表征,而不仅仅是提供更好的视觉特征。

使用在线 RL 对 Freeway、Hero 和 Breakout 游戏变体进行微调。微调中使用的新变体显示在每张图的下一行,预训练中看到的原始游戏显示在上一行。Scaled Q-Learning 的微调明显优于 MAE(一种视觉表示学习方法)和使用单游戏 DQN 从头开始学习。

结论和要点

我们介绍了一种基于CQL算法 的可扩展离线 RL 预训练方法 Scaled Q-Learning ,并展示了它如何实现高效的离线 RL 多任务训练。这项工作在实现更实用的真实世界 RL 代理训练方面取得了初步进展,可以替代昂贵而复杂的基于模拟的流程或大规模实验。也许从长远来看,类似的工作将带来具有普遍能力的预训练 RL 代理,这些代理可以通过大规模离线预训练开发出广泛适用的探索和交互技能。在机器人技术(参见一些初步结果)和 NLP 等领域的更广泛、更现实的任务上验证这些结果是未来研究的重要方向。离线 RL 预训练具有很大的潜力,我们期待在未来的工作中看到该领域的许多进展。

致谢

这项工作由 Aviral Kumar、Rishabh Agarwal、Xinyang Geng、George Tucker 和 Sergey Levine 完成。特别感谢 Sherry Yang、Ofir Nachum 和 Kuang-Huei Lee 对多游戏决策转换器代码库(用于评估)和多游戏 Atari 基准的帮助,以及 Tom Small 提供的插图和动画。

版权声明

本文仅代表作者观点,不代表本站立场。
本文系作者授权发表,未经许可,不得转载。

评论