世界模型：代理能在自己的梦中学习吗？（上）

来源：worldmodel

编译：weakish

编者按：Google Brain hardmaru的新作。探索了为流行的强化学习环境构建生成式神经网络模型。其世界模型（world model）可以通过无监督方式快速训练，以学习一个压缩的环境的时空表示。使用提取自世界模型的特征作为代理的输入，训练一个非常紧凑和简单的策略解决所需任务。甚至可以让代理完全在自己的幻梦（由世界模型生成）中训练，然后迁移策略至实际环境。

介绍

一个世界模型，取自Scott McCloud的Understanding Comics

人类基于其有限的感官接受到的信息发展出了关于世界的心智模型。人类基于这一内在模型做出决策和行动。系统动力学之父Jay Wright Forrester如此定义心智模型：

我们脑海中的周围世界的图像，只是一个模型。没人在脑中想象整个世界，政府、国家，等等。只有一些挑选出来的概念，以及概念之间的关系。我们用这些来表示真实系统。

为了处理日常生活中的海量信息，我们的大脑学习信息的时间和空间的抽象表示。我们能够观察一个场景，然后记忆关于该场景的一个抽象描述。有证据表明，我们在任何给定时刻接收到的信息，受到我们大脑对未来的预测所控制，该预测基于我们的内部模型作出。

我们大脑中的预测模型能够基于当前的运动信息预测未来的感官数据。这让我们在遇到危险时快速反应，根据直觉行动。例如，棒球手的反应时间以毫秒计——短于视觉信号从眼睛传到大脑的时间。因此棒球手的反应基于对明天（４月１日愚人节）何时来球，球来自何方的预测，这完全是下意识的。棒球手的肌肉基于内部模型的预测作出反应，在合适的时间和位置击球。

在许多强化学习（RL）问题中，人工代理同样从具备过去和当前状态的良好表示以及预测未来的良好模型中受益，比如循环神经网络（RNN）模型。

大型RNN是具有高度表达力的模型，可以学习丰富的时空表示。然而，许多RL方法常常使用参数很少（103到106）的小型网络。这是因为RL算法的瓶颈经常在奖励分配，这让传统RL算法难以学习大型模型数以百万计的权重。所以，实践中一般使用较小的网络，这样在训练中能更快地迭代，以得到良好的策略。

理想情况下，我们希望能高效地训练基于大型RNN的代理。在本文中，我们探索了训练大型神经网络（参数量级：107）来处理RL问题。我们将代理分为一个较大的世界模型和一个较小的控制者模型，然后训练较小的控制者模型使用较大的世界模型学习执行任务。较小的控制者模型让训练算法在较小的搜索空间里解决奖励分配问题，而较大的世界模型则避免了牺牲表达能力。基于世界模型训练代理，可以让代理学习到执行任务的高度紧凑的表示。

代理模型

我们的模型受我们自身的认知系统的启发。我们的代理有一个视觉感官组件将所见压缩为较小的具有代表性的编码。还有一个记忆组件，基于历史信息预测未来编码。最后，一个决策组件，仅基于视觉组件和记忆组件创建的表示做出决策。

VAE（V）模型

环境给代理提供的是每时步的高维输入，通常是视频序列中的2D图像。V模型的角色是学习输入帧的抽象、压缩表示。

我们使用变分自动编码器（VAE），将每时步t的视频帧编码为低维的潜向量zt。这一压缩表示可用于重建图像。

MDN-RNN (M)模型

M模型用于预测V未来会生成的z向量。我们使用了循环神经网络（RNN），搭配混合密度网络（MDN）作为输出层，以输出z的概率分布（混合高斯分布）。

控制者（C）模型

控制者模型负责决定行动序列以最大化期望奖励。我们特意让C尽可能地简单，并且单独训练，这样我们的代理模型的大部分复杂性都在世界模型（V和M）之中。

C是一个简单的线性模型：

at= Wc[ztht] + bc

将一切整合起来

下图展示了V、M、C是如何和环境交互的：

用伪代码表示上面的流程图：

C的极简设计在实践中额外提供了重要的益处。我们的V和M模型可以在现代GPU上基于反向传播高效地训练，如前所述，我们的模型的大部分复杂度（参数）在V和M中。相比之下，C作为线性模型，参数是极简的。这让我们得以探索非传统的训练C的方法——比如，我们甚至可以使用演化策略（ES）来应对更具挑战性的RL任务。事实上，我们使用了协方差矩阵自适应演化策略（CMA-ES）来优化C的参数，因为一般认为这一优化算法在几千个参数的解空间内表现优异。

赛车实验：世界模型用于特征提取

一个预测性的世界模型可以帮助我们提取有用的时空表示。将这些特征作为控制者的输入，我们可以训练一个紧凑、极简的控制者执行持续的控制任务，例如在一个基于像素输入的自顶向下的赛车环境中学习驾驶。本节将描述我们如何训练之前描述的代理模型来完成这一赛车任务。就我们所知，我们的代理是该任务的第一个已知解答（达到所需分数）。

在这一环境中，每次测试随机生成赛道，代理控制三种持续的行动：左/右转、加速、刹车。

我们首先让代理随机行动以探索环境，并记录采取的随机行动at以及对环境的观测所得。基于这一数据集，我们训练V学习观测到的帧的潜空间。然后使用训练过的V模型预处理每一时刻t的帧为zt，进而使用这些数据来训练我们的M模型。

理论上，我们能够以端到端的方式同时训练V和M。不过在实践中，我们发现分别训练V和M更容易操作，并且可以取得满意的结果。在单张NVIDIA P100 GPU上，V和M分别只需要不到1小时的训练时间。

在这一实验中，世界模型（V和M）对环境中的奖励信号一无所知。它的任务不过是压缩和预测观测到的图像序列。只有控制者（C）模型才知道环境中的奖励信息。由于线性控制者模型只有867个参数，因此像CMA-ES这样的演化算法很适合这一优化任务。

下图比较了代理的实际观测和世界模型捕捉的观测：

模型参数数量

赛车实验结果

仅使用V模型

如果我们具备了观测的良好表示，那么训练代理驾驶就不是一项困难的任务。之前的研究显示，基于手工提取的观测信息，例如LIDAR信息、角度、位置、风速，很容易就能训练一个小型前馈网络，输出满意的导航策略。因此，我们打算首先给C设置一些障碍，让它只能访问V模型，而不能访问M模型。

在这一设定下，尽管代理仍能通行于赛道，我们注意到它游移不定，并在大转弯处脱离赛道。这个残障代理在100次随机测试中的平均成绩是632 ± 251，和OpenAI Gym排行榜及A3C之类传统深度RL方法的表现相当。给C加上一个隐藏层后，成绩提升至788 ± 141.

全世界模型（V和M）

结合了V和M之后，控制者模型C具备了一个良好的表示，该表示包括当前的观测和对未来的期望。

我们看到，驾驶更平稳了，并且代理能够有效地处理大转弯。

基于全世界模型，我们的代理能够在100次测试中取得906 ± 21的平均成绩，解决了这一问题（达到900分以上），并且得到了新的最先进表现。传统深度RL方法的平均成绩在591至652之间，而排行榜上报告的最高成绩是838 ± 11. 另外，传统深度RL算法经常需要预处理每帧，例如施加边缘检测，还需要堆叠最近的一些帧。而我们的世界模型接受原始RGB像素图像作为输入，直接学习时空表示。

赛车梦

因为我们的世界模型能够建模未来，我们能让它自行产生假想的赛车场景的幻觉。我们可以给定当前状态，让它生成概率分布zt+1，取样一个zt+1，然后将这个样本用作真实观测。我们可以在这一M生成的幻觉场景中运行训练过的C模型。

我们发现，看起来在真实环境中学习到的策略某种程度上能在梦中起作用。这提出了一个问题——我们能否训练代理在自己的梦中学习，然后将策略迁移回实际环境？我们将在下篇回答这个问题，敬请期待。