世界模型——代理可以在它们的梦中学习吗？

摘要：我们在最近很有名的强化学习环境中构建生成式神经网络模型。我们的“世界模型”可以在非监督条件下快速训练学习到周边环境的时空表示。通过将从世界模型中提取出来的特征输入到一个代理中，我们可以训练得到一个简要的策略来解决简单的任务。我们甚至可以完全在世界模型中生成的幻觉梦境中训练代理，然后将生成的策略迁移回现实环境。

导语

人类基于有限的感知能力所感知到的事物在脑中构建一个对世界的意识模型，并且以此为依据来做出选择和行动。系统动力学之父Jay Wright Forrester对意识模型的定义如下：“我们的头脑中对周边的世界的映像，就是一个模型。没有人能在他的脑中得到世界、政府或国家的映像，人们只是选择了一些概念和这些概念彼此之间的关系来代表真正的系统。” (Forrester, 1971)

图1 世界观（Scott McCloud的关于认知的漫画）

为了处理日常生活中接触的大量信息，我们的大脑学会如何提取这些信息的空间和时间的关键信息。我们可以观察到一个景象后简要地记住相应的描述，也有证据表明我们在任何时候感知到的事物是被基于大脑中的世界模型而做出的预测影响的。

一种理解我们大脑中的预测模型的方式是，不要将它理解为对未来的总体上的预测，而是基于当下的行为对未来的感知信息的预测。有了这种预测模式，我们就可以在面临危险时快速地条件式反射做出反应，不需要有意识地计划出一系列行动。

图2 我们所看见的基于我们大脑对未来的预判

以棒球为例，棒球手只有几毫秒的时间来决定是否要挥起球棒（这个时间比视觉信息到达大脑所用的时间都短）。他们之所以可以击中球速高达115mph的球，就是因为他们可以用直觉预测到球在什么时候会飞到哪里。对于职业棒球手来说，击球的整个过程都是发生在潜意识中的。他们的肌肉条件反射性地收缩舒张来操控球棒在对的时间和位置击到球，与他们的内在模型的预判一致。他们才可以基于对未来的预判做出行动，而不需要有意识地进行各种对未来可能的场景计划。

图3 棒球手的例子

在许多强化学习（RL）的问题中，如果对过去和现在的状态，以及对未来的预测有一个很好的表示形式，人工代理也会有很好的表现，最好是强大的预测模型在一般的计算机上实现，如递归神经网络（RNN）。

如果在一系列步骤结束时才给予（积极或消极的）反馈，这就需要解决奖励分配问题，即哪些步骤应得到奖励？哪些步骤应得到惩罚？

大型RNN是具有高度表达性的模型，可以学习丰富的空间和时间数据的表示。然而，文献中的许多无模型RL方法通常只使用参数很少的小型神经网络。 RL算法的瓶颈常常是奖惩分配问题，这使得传统的RL算法难以学习大型模型的数百万个权重，因此在实践中使用较小的网络，因为它们在训练期间更快地迭代到良好的策略。

理想情况下，我们希望能够高效地训练基于大型的RNN的代理。反向传播算法可以用来有效地训练大型神经网络。在这篇论文的研究工作中，我们将代理分成大型世界模型和小型的控制模型，从而训练一个大型神经网络来处理RL任务。我们首先以无监督的方式训练一个大的神经网络，学习代理的世界模型，然后训练较小的控制器模型，学习使用这个世界模型执行任务。通过一个小型控制器让训练算法专注于解决小型搜索空间上的奖惩分配问题，同时不会牺牲较大的世界模型的容量和表现能力。从世界模型的角度，我们发现所训练的代理可以得到一个高度简要的策略来完成它的任务。

本文的目标是从1990-2015年关于结合基于RNN的世界模型和控制器的一系列论文中提炼几个关键概念，进一步提出如何在不同的增强学习环境中应用这些想法。我们建立了OpenAI Gym环境中的概率生成模型。这个基于递归神经网络的世界模型是通过训练实际游戏环境中录制的游戏场景得到的。之后，我们训练代理使用这些模型来模拟完整的环境。

代理模型

受人类本身认知系统的启发，我们提出了一个简单的代理模型。在这个模型中，代理有一个将它所看见的事物压缩到一段代表性的编码的视觉感官组件。它同时也有基于历史信息对未来代码做出预测的记忆组件。最后，我们的代理有一个基于其视觉和记忆组件做出行动决定的决策组件。

图4 世界模型的视觉组件（V）记忆组件（M）和控制器（C）

恒变自动编码VAE(V)模型

图5 恒变自动编码器工作流程图

在我们的模型中，使用恒变自动编码器作为视觉V组件。它将每一步接收到的图像压缩成一个低维的隐向量，这个向量还可以被反编码从而还原原来的图像。

混合密度-递归神经网络记忆模型MDN-RNN(M)

除了压缩视觉信息，我们还想压缩记录下来的事情。为了达到这个目的，记忆M模型通过产生一个z向量来预测未来。因为自然环境中很多环境都是随机的,我们训练递归神经网络生成一个概率密度函数p(z)而不是一个确定性的预测z。

图6 MDN-RNN(M)模型示意图

在我们的M模型中，有混合密度网络输出层的递归神经网络，基于现有信息，输出面向下一个隐向量的混合高斯模型的参数。具体来说，RNN会建立一个模型，at是在某一时刻t所采取的行动，ht是RNN彼时的隐状态。在随机抽样阶段，我们可以调整一个温度参数τ来控制模型不确定程度。

控制器（C）模型

控制器模型负责在特定环境中决定一系列动作来最大化累计的奖励。在我们的实验中，分别从视觉模型V和记忆模型M中训练，并且让它尽可能地简单、尽可能地小，因此我们代理恩典复杂度就只是存在于世界模型（V和M）中了。

C是一层简单的线性模型，将每一步的预测Zt和隐状态ht直接映射到动作at上：

图7 代理模型的流程图

将V，M，和C合并

上图是我们的代理模型的流程图，观察所得的原始图像先被每一步骤的视觉组件V处理生成zt，控制器C的输入数据是和记忆组件M生成的隐状态ht合并的向量，控制器C会接着输出一个控制马达的动作向量at，会接着影响环境，记忆组件M就会记录现在的Zt和动作at作为输入来更新它自己的隐状态。下面是我们的模型在open AI Gym 环境中的伪代码。

用一个特定的控制器C来运行这段代码会得到一轮中的累计奖励。控制器C的最小化设计也带来了一些实际的好处。只要定义一个区分度好的损失函数，就能通过深度学习提升训练大型、复杂的模型的效率。我们的视觉和记忆模型可以在GPU加速器的帮助下用反向传播算法高效地进行训练，所以我们希望模型的大部分复杂参数都留在视觉模型V和记忆模型M中。控制器是一个线性模型，其参数数量已经是最少的了。这就让我们可以尝试用不同的方式来训练控制器C，例如，可以使用演化策略来解决增强学习中有挑战性的奖惩分配问题。

为了优化控制器C的参数，我们选择了协方差适应演化策略（CMA-ES）作为优化算法，同时并行地在一台有多个CPU的机器上演化控制器C的参数。

赛车实验

我们通过之前介绍的代理模型来解决赛车任务。据我们所知，我们的模型是第一个达到得分要求的方法。

预测性世界模型可以帮助我们提取出有用的时间和空间代表的特征信息。通过使用这些特征信息作为控制器的输入数据，可以训练出一个实现控制任务的最小化控制器。

在这个环境中，赛道在每次比赛中都是随机生成的，我们的代理小汽车要在尽可能短的时间内跑过尽可能长距离的跑道。代理小汽车控制这三种连续的动作：向左/右转、加速、刹车。

我们先用这个环境中10000个随机比赛来训练视觉模型。我们先让代理小汽车自由随机地探索周边的环境，训练的VAE恒变自动编码模型将观察到的每一帧图像编码为低维的音向量z；然后使用训练好的V视觉模型来预处理每一时刻传入记忆模型中的图像，这样记忆模型就可以训练得到混合高斯概率模型。

在这次实验中，世界模型（V和M）对实际环境中的奖励信号一无所知，它的任务就是压缩和预测一系列观察到的图像。只有控制器C可以得到环境中的奖励反馈。因为线性控制器模型只有867个参数，CMA-ES优化算法也很适合这个优化任务。

图8 模型参数数量

实验结果

一开始实验只是使用了视觉模型，尽管我们的代理能够在这种情况下找到跑道，但是它会在急转弯处偏离跑道，在100次随机实验中平均分是632±251；接着我们训练了记忆模型M来做出预测。我们发现，代理小汽车同时有了视觉模型产生的低维向量和记忆模型产生的隐状态，可以显著地提升其行驶性能，行驶更为稳定，并且代理小汽车可以像我们先前提到的棒球手一样凭直觉预测什么时候、在哪里行进，在100次随机实验中拿到了906±21的历史最高分。之前用Deep RL方法的只得到了平均591-652的分数。

图9 同时加入视觉模型和记忆模型的代理小汽车行驶更为平稳

“幻梦”中的赛车比赛

由于我们的世界模型可以对未来加农，我们可以让模型自己“幻想”假设的赛车情景。我们可以让它基于当前的状态，产生接下来的概率分布，在这个概率分布中采样，当作是真实的观察。我们重新将控制器C放回记忆模型M产生的幻想环境中。这向我们展示了训练出一个可以完全在它自己模拟的梦境世界中执行任务的代理的可能性。论文还对VizDoom进行了一系列梦境内学习的相关试验，对迭代式训练过程有更详尽的介绍，感兴趣的朋友可以点击下方链接阅读交互式的论文原文。

“大数据与数据科学家”公众号

主编：王宏志

特邀副主编：朱劼

副主编：丁小欧

责任编辑：齐志鑫，宋扬，万晓珑、魏龑，张荣恩

编辑： 陶颖安

-精彩内容，记得分享到朋友圈-