用视觉模型改善强化学习评价设计

文 / 研究工程师 Mohammad Babaeizadeh 和研究员 Dumitru Erhan,Google Research

无模型强化学习已在 机器人游戏 和无人驾驶汽车等一系列领域得到成功验证。这些系统仅通过简单的试错法进行学习,因此在解决给定任务之前还需要大量尝试。

相比之下,基于模型的强化学习 (Model-based Reinforcement Learning) 学习的是环境模型(通常被称为世界模型动力学模型),智能体 (Agent) 能够预测潜在行动的结果,从而减少解决任务所需的环境交互量。

原则上,规划严格意义上的需要的只是预测未来奖励 (Rewards) 可以用于选择接近最优的未来行动。尽管如此,最近的许多方法,如 DreamerPlaNetSimPLe,都额外利用了预测未来图像的训练信号。但预测未来图像是否真的有必要或者能起到作用?视觉 MBRL 算法实际上还能从预测未来图像中获得什么好处?预测整个图像的计算和表示成本相当大,因此了解其是否真的有用对 MBRL 研究具有深远的意义。

在“模型、像素与奖励:在基于视觉模型的强化学习中评估设计权衡 (Models, Pixels, and Rewards:Evaluating Design Trade-offs in Visual Model-Based Reinforcement Learning)”中,我们证明了预测未来图像可以带来巨大收益,事实上,它是训练成功的视觉 MBRL 智能体的关键要素。我们开发了一个新的开源库,世界模型库,它使我们能够严格评估各种世界模型设计,以确定图像预测对每个模型返回奖励的相对影响。

世界模型库

世界模型库专为视觉 MBRL 训练和评估而设计,能够就每个设计决策对智能体在多个任务中的最终性能的影响进行大规模实证研究。该库引入了与平台无关的视觉 MBRL 仿真循环和 API,以无缝定义新的世界模型、规划器和任务,或从现有目录中进行选择,其中包括智能体(如 PlaNet)、视频模型(如 SV2P)以及各种 DeepMind Control 任务和规划器,如 CEM 和 MPPI。

利用该库,开发者可以研究 MBRL 中各种因素(如模型设计或表示空间)对智能体在一组任务上的性能的影响。该库支持从头开始或在预先收集的一组轨迹上训练智能体,并支持在给定任务上评估预训练的智能体。模型、规划算法和任务可以轻松混合和匹配到任何需要的组合。

为了向用户提供最大的灵活性,该库使用 NumPy 接口构建,可以在 TensorFlow、Pytorch 或 JAX 中实现不同的组件。此 Colab 为快速介绍。

图像预测的影响

我们使用世界模型库训练了多个具有不同图像预测水平的世界模型。这些模型都使用相同的输入(先前观察到的图像)来预测图像和奖励,但它们在预测图像的百分比上有所不同。随着智能体预测的图像像素数量的增加,由真实奖励衡量的智能体性能通常有所提高。

模型的输入是固定的(先前观察到的图像),但是预测图像的分数是变化的。如右图所示,增加预测像素的数量可以显著提高模型性能

值得注意的是,奖励预测准确率与智能体性能之间的相关性并不强,在某些情况下,更准确的奖励预测甚至会导致智能体性能降低。同时,图像重建误差与智能体性能之间也有很强的相关性。

图像/奖励预测的准确率(x 轴)和任务性能(y 轴)之间的相关性。这张图清楚表明图像预测准确率和任务性能之间有较强的相关性

这种现象与 探索 (Exploration) 直接相关, 即智能体为了收集更多关于环境中未知选项的信息而尝试风险更大和潜在奖励更低的行动。这可以通过在 离线环境 (offline setup) 中测试和比较模型来证明(即从预先收集的数据集中学习策略,与在线 (online) RL 相反,后者通过与环境交互学习策略)。离线环境可以确保不发生探索并且所有模型都在相同的数据上训练。我们发现,在离线环境中,更适合数据的模型通常表现更好。不仅如此,这些模型可能并不是从头学习和探索时表现得最好的模型。

不同视觉 MBRL 模型在不同任务中取得的分数。图表上半部分和下半部分分别显示了每个任务在在线和离线环境下训练取得的分数。每种颜色对应一种模型。在预先收集的数据(离线环境)上训练时,在线环境中表现较差的模型通常会获得高分,反之亦然

结论

我们已经通过实验证明,与仅预测预期奖励的模型相比,预测图像可以大幅提高任务性能。我们还证明,图像预测的准确率与这些模型的最终任务性能密切相关。这些发现可以推动更好的模型设计,并且格外适合数据收集成本较高的高维输入空间的未来环境。

如果您想开发自己的模型和实验,访问我们的 GitHub 仓库Colab 即可获取有关如何重现这项工作以及使用或扩展世界模型库的说明。

致谢

我们要特别感谢 Google Brain 团队的多位研究员和论文的合著者:Mohammad Taghi Saffar、Danijar Hafner、Harini Kannan、Chelsea Finn 和 Sergey Levine。

原文:Evaluating Design Trade-offs in Visual Model-Based Reinforcement Learning
中文:谷歌开发者公众号