循环世界模型(Recurrent World Models)——真实世界建模的强化学习利器

世界模型

智能体可以在它们头脑中的世界进行学习吗？

  世界模型（World Model）是NIPS 2018获得口头报告的论文之一，也是谷歌AI和Deepmind 2019年力作深度规划网络 PlaNet的主要参考文献之一。这篇博文是对World Model论文的翻译和学习笔记。原文地址 https://arxiv.org/abs/1803.10122

摘要

  我们探索建立通用强化学习环境的生成神经网络模型。我们的世界模型能够以一种快速的无监督学习方式，学习一个压缩的空间和时间的环境表示（Environment Representation）。通过使用从世界模型中提取的特征作为智能体的输入，我们可以训练出一个非常紧凑、简单的策略来解决所需的任务。我们甚至可以将智能体完全训练在由它的世界模型所产生的幻想的环境（own dream environment）中，并将这一策略转移到真实世界的环境中。

1. Introduction

  人类以有限的感官所能感知到的事物为基础，形成了一个心理世界模型（mental model of the world）。我们所做的决定和行动都是基于这个内部模型。系统动力学之父杰伊·赖特(Jay Wright)将心理模型描述为：

  “我们脑中所承载的周围世界的形象，只是一个模型。他脑子里没有办法能把世界、政府或国家都搞清楚。他只选择了（他理解的世界政府国家的）概念和它们之间的关系，并使用这些概念来表示真实的系统。”

  为了处理流经我们日常生活的大量信息，我们的大脑学会了信息的空间域和时域的抽象表示。我们能够观察一个场景并记住其中的抽象描述[4，5]。证据还表明，我们在任何时刻所感知的，都是由我们的大脑根据我们的内部心理模型对未来的预测所决定的。（比如下图，我们的大脑认为左图的两个正方形并不标准，认为右图的黑白点在不停波动，这些都是大脑根据我们的视觉系统建立了自己的模型，而后做出的判断，但真实世界并非如此。）

  有一种理解大脑内部预测模型的方法是，internal model可能不仅仅是预测未来，而且根据我们当前的运动行为来预测未来的感官数据。我们能够在这种预测模型上采取行动，并在我们面临危险时表现出快速的行为[13]，而不需要有意识地规划一个行动路线。以棒球为例，一个棒球运动员只有毫秒级的时间来决定如何挥动球棍，这个时间甚至比视觉信号从眼球传到大脑的时间还要短。对专业运动员来说，这个动作几乎是下意识的，他们的肌肉在正确的时间和地点挥动球棍，与他们内部模型的预测一致[7]。他们可以根据他们对未来的预测迅速采取行动，而不需要有意识地推出可行的击球计划。

  在许多强化学习问题中，智能体既需要一个对过去和现在状态的良好描述，还需要一个优秀的模型来预测未来的状态，最好是一个如递归神经网络RNN一样的强大的预测模型。

  大参数量的RNN模型是具有高度表示能力的模型（强大的表征能力可能代表强大的决策能力），可以学习丰富的数据空间和时间表示（Spatial and temporal representation）。然而，许多文献中的model-free的强化学习算法往往只采用参数较少的小型神经网络。传统RL算法受限于credit assignment问题。

  PS：credit assignment问题是说，在许多RL任务中，reward都是很稀疏的，往往在一串序列结束后才给出，而在这之前agent已经采取了很多action，credit assignment 问题就是要解决究竟是哪一个action对最后的reward的影响最大，对final reward，哪些action起了有益的作用，哪些action起了负面作用。credit assignment问题使得传统的RL算法很难训练数百万个参数的大模型，因此在实际应用中，由于小网络在训练过程中更快，反而有比较好的策略。（我的直观理解就是，由于sparse reward，大型网络的参数很不好训练，所以在性能上输给了小型网络，但作者认为采用小型网络只是暂时的成功，未来要用的还是大型网络）

  理想情况下，我们希望能够训练大参数量的RNN模型，比如我们已经有了训练大规模网络的反向传播算法。因此这篇论文的工作就是探讨如何训练一个大型网络来解决强化学习的任务。在文中我们讨论的大型网络是指参数量在$10^7$量级的网络，相比$10^3-10^6$量级的小型网络，这是一个进步，但是和深度学习领域最成功的巨型网络相比还是太小，后者具有高达$10^8-10^9$量级的参数规模。

  我们把agent模型拆分成两部分，一部分是大型的world model，另一部分是小型的controller model。首先采用无监督学习的方式训练world model，作为对真实世界的一个表征，或者叫agent本身对真实世界的一个理解，也就是前文所述的mental model。然后再训练controller model利用学到的world model去完成切实的RL任务。controller model使训练算法在一个较小的搜索空间上处理credit assignment问题，而不是牺牲capacity和表达能力，在更大的世界模型上处理credit assignment问题。通过使controller model从world model的视角下进行训练，我们发现它可以学习到一种高度紧凑的策略来执行它的任务。

  虽然目前有大量的研究有关model-based强化学习算法，但本文并不是要做一个这方面的综述。相反，本文的目标是从1990-2015年一系列论文中凝练出几个关键的概念，这些概念是融合了RNN-based world model和controller model 领域的。我们还将讨论其他相关的工作，（相关工作指，学习一个世界模型和然后使用该模型训练controller等类似思想的论文）。

  在本文中，我们还提出了一个简化框架，用于在实验中展示论文提到的若干核心观点。同时提出了一些前瞻性见解，以便将这些idea应用于不同的RL算法框架。本文的RL领域的学术术语和标记与On Learning to Think:Algorithmic Information Theory for Novel Combinations of RL Controllers and RNN World Models该文章保持一致。

  我们搭建了一个基于OpenAI Gym的概率生成模型。我们采用从真实实验环境中收集到的观察数据训练RNN-based world model。当world model训练好之后，可以用来模拟真实环境，或者训练controller。

2. 智能体模型 Agent Model

  受我们人类的认知系统启发，本文提出了一个简单的智能体模型。模型包括一个视觉感知组件（Visual Sensory Component），它将看到的图像压缩成一个低维的表征向量（Representative code）作为模型的输入。模型还包括一个记忆组件（memory component），记忆组件基于历史信息，对未来的表征向量做出预测。最后，我们的模型还有一个决策组件（decision-making component），它仅仅基于视觉组件和记忆组件的表征向量决定采取的动作。

  如上图，视觉组件和记忆组件组成了World Model，决策组件就是前文所述的controller model。每个时间步骤下，视觉组件将高维的原始图像编码成一个低维的潜层向量，记忆组件接收视觉组件编码的潜层向量，并基于其本身的历史信息，对未来状态做出预测。决策/控制组件同时接受视觉组件（V）和记忆组件（M）的输出，采取动作，与真实环境进行交互。

2.1 变分自编码器（VAE）模型——即视觉（V）组件

  在每个时间步长，真实环境给智能体的输入是一个高维图像观测信息（文中叫high dimensional input observation，所谓高维信息是指图像作为一个三维矩阵，尺寸比较大）。视觉组件的输入是一个二维的图像帧，通常是视频序列的一部分。视觉组件的任务就是学习一个抽象的、压缩的表示（Representation）来描述每一帧的输入。

  论文中使用了一个简单的变分自编码器作为视觉组件，VAE能够将原始观察图像通过编码成一个低维潜层向量$z$，再经过解码器，解码成重构图像。向量$z$就是我们需要的低维潜层向量（small latent vector）。

2.2 MDN-RNN模型——记忆组件

  视觉组件的任务是将输入图像进行压缩，我们同时也希望将时域上的信息也进行压缩。为实现这个目的，记忆组件的任务就是预测未来。详细地说，记忆组件需要预测出视觉组件在未来可能输出的潜层向量$z$，换句话说，记忆组件需要预见将要发生的游戏画面（的潜层表征）。由于自然环境中许多复杂的模型都是随机性的，因此我们训练RNN模型去输出$z$的概率密度函数$p(z)$，而不是输出一个确定的向量$z$。

  在论文在中，我们用混合高斯分布（a mixture of Gaussian Distribution）来估计$p(z)$。给定当前信息和历史信息，我们训练RNN来输出下一个潜层向量$z_{t+1}$的概率分布（如何输出概率密度？）。

  确切地说，RNN建模的是这样一个条件概率：$P(z_{t+1}|a_t,z_t,h_t)$，式中，$a_t$是$t$时刻智能体采取的动作，$h_t$是RNN模型在$t$时刻的隐层状态。在采样过程中，我们通过控制温度参数$\tau$来控制模型的不确定性（具体是怎样控制的？），而且实验发现$\tau$值对接下来的controller model的训练是有帮助的。

  生成概率密度的方法被称作是MDN-RNN模型，MDN是混合密度网络（Mixture Density Network）的简称，由大名鼎鼎的《模式识别与机器学习》的作者CM.Bishop发明。MDN-RNN模型先前被应用于序列生成问题，比如手写字或简笔画的生成。下图是SketchRNN的例子，图中竖线左边是人画的第一笔，模型可以根据这一笔线条预测之后画的内容。在本文中，我们并不用它来预测简笔画，而是预测下一步视觉组件要输出的潜层向量$z_t$。

2.3 Controller ( C ) Model

  controller model负责确定要采取的action，以使一轮rollout的期望回报最大。在本文的实验中，我们故意使controller model的参数规模尽可能的小，并且和视觉组件、记忆组件分开训练，这样就能保证智能体的多样性大部分都集中在世界模型中（也就是V和M模型）。

  Controller Model是一个简单的单层线性模型，将$z_t$和$h_t$直接映射到动作空间$a_t$：
$$a_t=W_c[z_t{h}_t]+b_c$$

  在上式中，$W_c$是权重矩阵，$b_c$是偏置向量，controller model首先将$z_t$和$h_t$向量拼接在一起，然后进行线性映射，得到动作$a_t$。

2.4 将V,C,M模型进行整合

  下图展示了V,M,C模型是如何整合在一起，并与环境进行交互的。在每个时刻$t$，首先由V模型处理原始图像，得到$z_t$，随后，C模型将$z_t$向量与M模块的$h_t$向量进行拼接，并输出动作向量$a_t$，完成与环境的交互。最后将$z_t$和$a_t$输入M模型，更新隐状态$h_{t+1}$。

  下面给出在OpenAI Gym环境下的整体模型的伪代码。给定一个controller model，在这个函数下运行，就能得到一个episode的累计回报值。

def rollout(controller):
	'''
	env,rnn,vae are global variables
	'''
	obs = env.reset() # 初始化游戏环境
	h = rnn.initial_state() # 获取初始的RNN隐状态
	done = False # 游戏结束标志
	cumulative_reward = 0 # 累计回报
	while not done:
		z = vae.encode(obs) # VAE模型对原始图像进行编码
		a = controller.action([z, h]) # controller model根据当前状态z和之前隐状态h采取动作
		obs, reward, done = env.step(a) # agent与环境交互，得到下一个时刻的环境状态和reward
		cumulative_reward += reward # 更新累计回报
		h = rnn.forward([a, z, h]) # 更新隐状态
	return cumulative_reward

  controller model的minimal design（我的理解是，minimal design是将controller model的参数量设计得比较小的这一做法）也具有非常重要的实践意义。深度学习领域的研究进展给我们提供了高效训练大规模复杂神经网络的工具，也提供了设计一个可微的、性能良好的损失函数的可能。我们的V模型和M模型就是在多GPU加速的反向传播算法上训练得到的，因此我们希望模型的大部分复杂性以及模型的参数都驻留在V模型和M模型中（不理解原文中这句话前后的逻辑在哪里，设计一个小controller和较大的V,M模型是有助于将（从真实环境学来的）模型复杂性留存在VM模型中，但这和反向传播算法有什么关系？