Learning Latent Dynamics for Planning from Pixels

Danijar Hafner^{1 2} Timothy Lillicrap^{3} Ian Fischer^{4} Ruben Villegas^{1 5}
David Ha^{1} Honglak Lee^{1} James Davidson^{1}

^{1}Google Brain ^{2}University of Toronto ^{3}DeepMind ^{4}Google Research ^{5}University of Michigan. Correspondence to: Danijar Hafner mail@danijar.com.

Proceedings of the 36^{th} International Conference on Machine Learning, Long Beach, California, PMLR 97, 2019. Copyright 2019 by the author(s).

Abstract

摘要

Planning has been very successful for control tasks with known environment dynamics. To leverage planning in unknown environments, the agent needs to learn the dynamics from interactions with the world. However, learning dynamics models that are accurate enough for planning has been a long-standing challenge, especially in image-based domains. We propose the Deep Planning Network (PlaNet), a purely model-based agent that learns the environment dynamics from images and chooses actions through fast online planning in latent space. To achieve high performance, the dynamics model must accurately predict the rewards ahead for multiple time steps. We approach this using a latent dynamics model with both deterministic and stochastic transition components. Moreover, we propose a multi-step variational inference objective that we name latent overshooting. Using only pixel observations, our agent solves continuous control tasks with contact dynamics, partial observability, and sparse rewards, which exceed the difficulty of tasks that were previously solved by planning with learned models. PlaNet uses substantially fewer episodes and reaches final performance close to and sometimes higher than strong model-free algorithms.

规划在环境动力学已知的控制任务上非常成功。要把规划用到未知环境里，智能体需要从与世界的交互中学习动力学。然而，学到准到足以规划的动力学模型，一直是长期挑战，在基于图像的领域尤其如此。我们提出深度规划网络（Deep Planning Network，PlaNet）：一个纯基于模型的智能体，从图像学习环境动力学，并通过在潜空间里做快速在线规划来选择动作。要达到高水平表现，动力学模型必须能准确预测未来多个时间步的奖励。我们用同时带有确定性与随机转移分量的潜变量动力学模型来处理这一点。此外，我们提出一种多步变分推断目标，并把它叫做潜空间过冲（latent overshooting）。只使用像素观测，我们的智能体就能解决带有接触力学、部分可观测性和稀疏奖励的连续控制任务，这些任务的难度超过了此前用所学模型做规划所能解决的任务。PlaNet 使用的回合数明显更少，最终表现接近、有时还高于强模型自由算法。

1. Introduction

1. 引言

Planning is a natural and powerful approach to decision making problems with known dynamics, such as game playing and simulated robot control (Tassa et al., 2012; Silver et al., 2017; Moravčík et al., 2017). To plan in unknown environments, the agent needs to learn the dynamics from experience. Learning dynamics models that are accurate enough for planning has been a long-standing challenge. Key difficulties include model inaccuracies, accumulating errors of multi-step predictions, failure to capture multiple possible futures, and overconfident predictions outside of the training distribution.

规划是处理已知动力学决策问题的一种自然而有力的方法，例如博弈和仿真机器人控制（Tassa et al., 2012; Silver et al., 2017; Moravčík et al., 2017）。要在未知环境中规划，智能体需要从经验中学习动力学。学到准到足以规划的动力学模型，一直是长期挑战。主要困难包括模型不准确、多步预测误差累积、无法刻画多种可能的未来，以及在训练分布之外做出过于自信的预测。

Planning using learned models offers several benefits over model-free reinforcement learning. First, model-based planning can be more data efficient because it leverages a richer training signal and does not require propagating rewards through Bellman backups. Moreover, planning carries the promise of increasing performance just by increasing the computational budget for searching for actions, as shown by Silver et al. (2017). Finally, learned dynamics can be independent of any specific task and thus have the potential to transfer well to other tasks in the environment.

用所学模型做规划，相对模型自由强化学习有几项好处。首先，基于模型的规划可以更省数据，因为它利用更丰富的训练信号，并且不必把奖励沿贝尔曼备份回传。此外，规划还带来这样一种可能：只要增加搜索动作的计算预算，就能提高表现，Silver et al.（2017）已经展示了这一点。最后，学到的动力学可以独立于任何特定任务，因而有潜力迁移到同一环境中的其他任务。

Recent work has shown promise in learning the dynamics of simple low-dimensional environments (Deisenroth & Rasmussen, 2011; Gal et al., 2016; Amos et al., 2018; Chua et al., 2018; Henaff et al., 2018). However, these approaches typically assume access to the underlying state of the world and the reward function, which may not be available in practice. In high-dimensional environments, we would like to learn the dynamics in a compact latent space to enable fast planning. The success of such latent models has previously been limited to simple tasks such as balancing cartpoles and controlling 2-link arms from dense rewards (Watter et al., 2015; Banijamali et al., 2017).

近期工作表明，在简单的低维环境中学习动力学是有希望的（Deisenroth & Rasmussen, 2011; Gal et al., 2016; Amos et al., 2018; Chua et al., 2018; Henaff et al., 2018）。然而这些方法通常假定能够获得世界的底层状态和奖励函数，实践中未必做得到。在高维环境里，我们希望在紧凑的潜空间中学习动力学，以便快速规划。此类潜模型此前的成功还局限于简单任务，例如用稠密奖励做倒立摆平衡和控制两连杆手臂（Watter et al., 2015; Banijamali et al., 2017）。

In this paper, we propose the Deep Planning Network (PlaNet), a model-based agent that learns the environment dynamics from pixels and chooses actions through online planning in a compact latent space. To learn the dynamics, we use a transition model with both stochastic and deterministic components. Moreover, we experiment with a novel generalized variational objective that encourages multi-step predictions. PlaNet solves continuous control tasks from pixels that are more difficult than those previously solved by planning with learned models.

本文提出深度规划网络（PlaNet）：一个基于模型的智能体，从像素学习环境动力学，并通过在紧凑潜空间中做在线规划来选择动作。为了学习动力学，我们使用同时带有随机与确定性分量的转移模型。此外，我们试验了一种新的广义变分目标，以鼓励多步预测。PlaNet 能从像素解决连续控制任务，这些任务比此前用所学模型做规划所解决的任务更难。

Key contributions of this work are summarized as follows:

本工作的主要贡献概括如下：

• Planning in latent spaces We solve a variety of tasks from the DeepMind control suite, shown in Figure 1, by learning a dynamics model and efficiently planning in its latent space. Our agent substantially outperforms the model-free A3C and in some cases D4PG algorithm in final performance, with on average 200× less environment interaction and similar computation time.

• 在潜空间中规划 我们通过学习动力学模型并在其潜空间中高效规划，解决了 DeepMind 控制套件中的多项任务，如图 1 所示。我们的智能体在最终表现上显著超过模型自由的 A3C，在部分情形下也超过 D4PG，平均只需约 200 倍更少的环境交互，计算时间则相近。

• Recurrent state space model We design a latent dynamics model with both deterministic and stochastic components (Buesing et al., 2018; Chung et al., 2015). Our experiments indicate having both components to be crucial for high planning performance.

• 循环状态空间模型 我们设计了一种同时带有确定性与随机分量的潜变量动力学模型（Buesing et al., 2018; Chung et al., 2015）。实验表明，这两个分量对高水平的规划表现都至关重要。

• Latent overshooting We generalize the standard variational bound to include multi-step predictions. Using only terms in latent space results in a fast regularizer that can improve long-term predictions and is compatible with any latent sequence model.

• 潜空间过冲 我们把标准变分界推广到包含多步预测。只使用潜空间中的项，会得到一种快速的正则，可以改进长期预测，并且与任何潜变量序列模型兼容。

Figure 1: Image-based control domains used in our experiments. The images show agent observations before downscaling to $64 \times 64 \times 3$ pixels. (a) The cartpole swingup task has a fixed camera so the cart can move out of sight. (b) The reacher task has only a sparse reward. (c) The cheetah running task includes both contacts and a larger number of joints. (d) The finger spinning task includes contacts between the finger and the object. (e) The cup task has a sparse reward that is only given once the ball is caught. (f) The walker task requires balance and predicting difficult interactions with the ground when the robot is lying down.

图 1：实验所用的基于图像的控制领域。图中给出的是下采样到 $64 \times 64 \times 3$ 像素之前的智能体观测。(a) 倒立摆起摆任务使用固定相机，因此小车可能移出视野。(b) 伸手任务只有稀疏奖励。(c) 猎豹奔跑任务既包含接触，关节数也更多。(d) 手指旋转任务包含手指与物体之间的接触。(e) 杯子任务只有在球被接住时才给出稀疏奖励。(f) 步行者任务需要保持平衡，并在机器人躺在地上时预测与地面的困难交互。

2. Latent Space Planning

2. 潜空间规划

To solve unknown environments via planning, we need to model the environment dynamics from experience. PlaNet does so by iteratively collecting data using planning and training the dynamics model on the gathered data. In this section, we introduce notation for the environment and describe the general implementation of our model-based agent. In this section, we assume access to a learned dynamics model. Our design and training objective for this model are detailed in Section 3.

要通过规划来解决未知环境，我们需要从经验中为环境动力学建模。PlaNet 的做法是：迭代地用规划收集数据，并在收集到的数据上训练动力学模型。本节引入环境的记号，并描述我们基于模型的智能体的一般实现。在本节中，我们假定已经有一个学到的动力学模型。该模型的设计和训练目标将在第 3 节详细给出。

Problem setup Since individual image observations generally do not reveal the full state of the environment, we consider a partially observable Markov decision process (POMDP). We define a discrete time step $t$, hidden states $s_t$, image observations $o_t$, continuous action vectors $a_t$, and scalar rewards $r_t$, that follow the stochastic dynamics

问题设置 由于单张图像观测通常并不能揭示环境的全部状态，我们考虑部分可观测马尔可夫决策过程（POMDP）。我们定义离散时间步 $t$、隐状态 $s_t$、图像观测 $o_t$、连续动作向量 $a_t$ 以及标量奖励 $r_t$，它们服从如下随机动力学

$$\begin{aligned}
\text{Transition function:} &\quad s_t \sim \mathrm{p}(s_t \mid s_{t-1}, a_{t-1}) \
\text{Observation function:} &\quad o_t \sim \mathrm{p}(o_t \mid s_t) \
\text{Reward function:} &\quad r_t \sim \mathrm{p}(r_t \mid s_t) \
\text{Policy:} &\quad a_t \sim \mathrm{p}(a_t \mid o_{\le t}, a_{<t}),
\end{aligned} \tag{1}$$

where we assume a fixed initial state $s_0$ without loss of generality. The goal is to implement a policy $\mathrm{p}(a_t \mid o_{\le t}, a_{<t})$ that maximizes the expected sum of rewards $\mathrm{E}_{\mathrm{p}}[\sum_{t=1}^{T} r_t]$, where the expectation is over the distributions of the environment and the policy.

其中我们在不损失一般性的前提下假定初始状态 $s_0$ 固定。目标是实现策略 $\mathrm{p}(a_t \mid o_{\le t}, a_{<t})$，使期望累计奖励 $\mathrm{E}_{\mathrm{p}}[\sum_{t=1}^{T} r_t]$ 最大，期望是对环境分布和策略分布取的。

Algorithm 1: Deep Planning Network (PlaNet)

Input:

$R$	Action repeat	$p(s_t \mid s_{t-1}, a_{t-1})$	Transition model
$S$	Seed episodes	$p(o_t \mid s_t)$	Observation model
$C$	Collect interval	$p(r_t \mid s_t)$	Reward model
$B$	Batch size	$q(s_t \mid o_{\le t}, a_{<t})$	Encoder
$L$	Chunk length	$p(\epsilon)$	Exploration noise
$\alpha$	Learning rate

1. Initialize dataset $\mathcal{D}$ with $S$ random seed episodes.
2. Initialize model parameters $\theta$ randomly.
3. while not converged do
   • // Model fitting
   4. for update step $s = 1..C$ do
      5. Draw sequence chunks $\{(o_t, a_t, r_t)_{t=k}^{L+k}\}_{i=1}^{B} \sim \mathcal{D}$ uniformly at random from the dataset.
      6. Compute loss $\mathcal{L}(\theta)$ from Equation 3.
      7. Update model parameters $\theta \leftarrow \theta - \alpha \nabla_\theta \mathcal{L}(\theta)$.
   • // Data collection
   8. $o_1 \leftarrow \texttt{env.reset()}$
   9. for time step $t = 1..\lceil T/R \rceil$ do
      10. Infer belief over current state $q(s_t \mid o_{\le t}, a_{<t})$ from the history.
      11. $a_t \leftarrow \texttt{planner}(q(s_t \mid o_{\le t}, a_{<t}), p)$, see Algorithm 2 in the appendix for details.
      12. Add exploration noise $\epsilon \sim p(\epsilon)$ to the action.
      13. for action repeat $k = 1..R$ do
          14. $r_t^k, o_{t+1}^k \leftarrow \texttt{env.step}(a_t)$
      14. $r_t, o_{t+1} \leftarrow \sum_{k=1}^{R} r_t^k,\; o_{t+1}^R$
   10. $\mathcal{D} \leftarrow \mathcal{D} \cup \{(o_t, a_t, r_t)_{t=1}^{T}\}$

算法 1： 深度规划网络（PlaNet）

输入：

$R$	动作重复	$p(s_t \mid s_{t-1}, a_{t-1})$	转移模型
$S$	种子回合	$p(o_t \mid s_t)$	观测模型
$C$	收集间隔	$p(r_t \mid s_t)$	奖励模型
$B$	批量大小	$q(s_t \mid o_{\le t}, a_{<t})$	编码器
$L$	片段长度	$p(\epsilon)$	探索噪声
$\alpha$	学习率

1. 用 $S$ 条随机种子回合初始化数据集 $\mathcal{D}$。
2. 随机初始化模型参数 $\theta$。
3. while 尚未收敛 do
   • // 模型拟合
   4. for 更新步 $s = 1..C$ do
      5. 从数据集中均匀随机抽取序列片段 $\{(o_t, a_t, r_t)_{t=k}^{L+k}\}_{i=1}^{B} \sim \mathcal{D}$。
      6. 按方程 3 计算损失 $\mathcal{L}(\theta)$。
      7. 更新模型参数 $\theta \leftarrow \theta - \alpha \nabla_\theta \mathcal{L}(\theta)$。
   • // 数据收集
   8. $o_1 \leftarrow \texttt{env.reset()}$
   9. for 时间步 $t = 1..\lceil T/R \rceil$ do
      10. 根据历史推断当前状态的信念 $q(s_t \mid o_{\le t}, a_{<t})$。
      11. $a_t \leftarrow \texttt{planner}(q(s_t \mid o_{\le t}, a_{<t}), p)$，细节见附录中的算法 2。
      12. 向动作加入探索噪声 $\epsilon \sim p(\epsilon)$。
      13. for 动作重复 $k = 1..R$ do
          14. $r_t^k, o_{t+1}^k \leftarrow \texttt{env.step}(a_t)$
      14. $r_t, o_{t+1} \leftarrow \sum_{k=1}^{R} r_t^k,\; o_{t+1}^R$
   10. $\mathcal{D} \leftarrow \mathcal{D} \cup \{(o_t, a_t, r_t)_{t=1}^{T}\}$

Model-based planning PlaNet learns a transition model $p(s_t \mid s_{t-1}, a_{t-1})$, observation model $p(o_t \mid s_t)$, and reward model $p(r_t \mid s_t)$ from previously experienced episodes (note italic letters for the model compared to upright letters for the true dynamics). The observation model provides a rich training signal but is not used for planning. We also learn an encoder $q(s_t \mid o_{\le t}, a_{<t})$ to infer an approximate belief over the current hidden state from the history using filtering. Given these components, we implement the policy as a planning algorithm that searches for the best sequence of future actions. We use model-predictive control (MPC; Richards, 2005) to allow the agent to adapt its plan based on new observations, meaning we replan at each step. In contrast to model-free and hybrid reinforcement learning algorithms, we do not use a policy or value network.

基于模型的规划 PlaNet 从先前经历过的回合中学习转移模型 $p(s_t \mid s_{t-1}, a_{t-1})$、观测模型 $p(o_t \mid s_t)$ 和奖励模型 $p(r_t \mid s_t)$（注意模型用斜体字母，真实动力学用正体字母）。观测模型提供丰富的训练信号，但规划时并不使用它。我们还学习一个编码器 $q(s_t \mid o_{\le t}, a_{<t})$，用滤波从历史中推断当前隐状态的近似信念。有了这些组件，我们把策略实现成一个搜索最佳未来动作序列的规划算法。我们使用模型预测控制（MPC；Richards, 2005），让智能体能够根据新观测调整计划，也就是每一步都重新规划。与模型自由和混合强化学习算法不同，我们不使用策略网络或价值网络。

Experience collection Since the agent may not initially visit all parts of the environment, we need to iteratively collect new experience and refine the dynamics model. We do so by planning with the partially trained model, as shown in Algorithm 1. Starting from a small amount of $S$ seed episodes collected under random actions, we train the model and add one additional episode to the data set every $C$ update steps. When collecting episodes for the data set, we add small Gaussian exploration noise to the action. To reduce the planning horizon and provide a clearer learning signal to the model, we repeat each action $R$ times, as common in reinforcement learning (Mnih et al., 2015; 2016).

经验收集 由于智能体一开始未必会访问到环境的所有部分，我们需要迭代地收集新经验并改进动力学模型。做法是用部分训练好的模型做规划，如算法 1 所示。从少量在随机动作下收集的 $S$ 条种子回合出发，我们训练模型，并每隔 $C$ 次更新向数据集再加入一条回合。为数据集收集回合时，我们给动作加上较小的高斯探索噪声。为了缩短规划视野，并向模型提供更清晰的学习信号，我们把每个动作重复 $R$ 次，这在强化学习中很常见（Mnih et al., 2015; 2016）。

Planning algorithm We use the cross entropy method (CEM; Rubinstein, 1997; Chua et al., 2018) to search for the best action sequence under the model, as outlined in Algorithm 2. We decided on this algorithm because of its robustness and because it solved all considered tasks when given the true dynamics for planning. CEM is a population-based optimization algorithm that infers a distribution over action sequences that maximize the objective. As detailed in Algorithm 2 in the appendix, we initialize a time-dependent diagonal Gaussian belief over optimal action sequences $a_{t:t+H} \sim \mathrm{Normal}(\mu_{t:t+H}, \sigma_{t:t+H}^2 I)$, where $t$ is the current time step of the agent and $H$ is the length of the planning horizon. Starting from zero mean and unit variance, we repeatedly sample $J$ candidate action sequences, evaluate them under the model, and re-fit the belief to the top $K$ action sequences. After $I$ iterations, the planner returns the mean of the belief for the current time step, $\mu_t$. Importantly, after receiving the next observation, the belief over action sequences starts from zero mean and unit variance again to avoid local optima.

规划算法 我们使用交叉熵方法（CEM；Rubinstein, 1997; Chua et al., 2018）在模型下搜索最佳动作序列，概要见算法 2。我们选择这一算法，是因为它稳健，而且在给定真实动力学做规划时能够解决所考虑的全部任务。CEM 是一种基于种群的优化算法，它推断出能最大化目标的动作序列分布。如附录中算法 2 所述，我们在最优动作序列上初始化一个随时间变化的对角高斯信念 $a_{t:t+H} \sim \mathrm{Normal}(\mu_{t:t+H}, \sigma_{t:t+H}^2 I)$，其中 $t$ 是智能体的当前时间步，$H$ 是规划视野长度。从零均值、单位方差出发，我们反复采样 $J$ 条候选动作序列，在模型下评估它们，再把信念重新拟合到最好的 $K$ 条动作序列上。经过 $I$ 次迭代后，规划器返回当前时间步信念的均值 $\mu_t$。重要的是，在收到下一个观测之后，动作序列上的信念会再次从零均值、单位方差开始，以避免局部最优。

To evaluate a candidate action sequence under the learned model, we sample a state trajectory starting from the current state belief, and sum the mean rewards predicted along the sequence. Since we use a population-based optimizer, we found it sufficient to consider a single trajectory per action sequence and thus focus the computational budget on evaluating a larger number of different sequences. Because the reward is modeled as a function of the latent state, the planner can operate purely in latent space without generating images, which allows for fast evaluation of large batches of action sequences. The next section introduces the latent dynamics model that the planner uses.

为了在所学模型下评估一条候选动作序列，我们从当前状态信念出发采样一条状态轨迹，并把沿该序列预测的平均奖励加起来。由于我们使用基于种群的优化器，发现对每条动作序列只考虑一条轨迹就足够了，于是可以把计算预算集中在评估更多不同的序列上。因为奖励被建模成潜状态的函数，规划器可以纯粹在潜空间中运作而不生成图像，从而能够快速评估大批量的动作序列。下一节介绍规划器所使用的潜变量动力学模型。

3. Recurrent State Space Model

3. 循环状态空间模型

For planning, we need to evaluate thousands of action sequences at every time step of the agent. Therefore, we use a recurrent state-space model (RSSM) that can predict forward purely in latent space, similar to recently proposed models (Karl et al., 2016; Buesing et al., 2018; Doerr et al., 2018). This model can be thought of as a non-linear Kalman filter or sequential VAE. Instead of an extensive comparison to prior architectures, we highlight two findings that can guide future designs of dynamics models: our experiments show that both stochastic and deterministic paths in the transition model are crucial for successful planning. In this section, we remind the reader of latent state-space models and then describe our dynamics model.

为了规划，我们需要在智能体的每一个时间步评估数千条动作序列。因此，我们使用可以纯粹在潜空间中向前预测的循环状态空间模型（RSSM），这与近期提出的一些模型类似（Karl et al., 2016; Buesing et al., 2018; Doerr et al., 2018）。可以把这个模型想成非线性卡尔曼滤波或序列 VAE。我们不做与先前架构的广泛比较，而是强调两项可以指导未来动力学模型设计的发现：实验表明，转移模型中的随机路径和确定性路径对成功规划都至关重要。本节先回顾潜变量状态空间模型，再描述我们的动力学模型。

Latent dynamics We consider sequences $\{o_t, a_t, r_t\}_{t=1}^T$ with discrete time step $t$, image observations $o_t$, continuous action vectors $a_t$, and scalar rewards $r_t$. A typical latent state-space model is shown in Figure 2b and resembles the structure of a partially observable Markov decision process. It defines the generative process of the images and rewards using a hidden state sequence $\{s_t\}_{t=1}^T$,

潜动力学 我们考虑序列 $\{o_t, a_t, r_t\}_{t=1}^T$，其中离散时间步为 $t$，图像观测为 $o_t$，连续动作向量为 $a_t$，标量奖励为 $r_t$。一个典型的潜变量状态空间模型如图 2b 所示，其结构类似于部分可观测马尔可夫决策过程。它用隐状态序列 $\{s_t\}_{t=1}^T$ 定义图像和奖励的生成过程，

$$\begin{aligned}
\text{Transition model:} &\quad s_t \sim p(s_t \mid s_{t-1}, a_{t-1}) \
\text{Observation model:} &\quad o_t \sim p(o_t \mid s_t) \
\text{Reward model:} &\quad r_t \sim p(r_t \mid s_t),
\end{aligned} \tag{2}$$

where we assume a fixed initial state $s_0$ without loss of generality. The transition model is Gaussian with mean and variance parameterized by a feed-forward neural network, the observation model is Gaussian with mean parameterized by a deconvolutional neural network and identity covariance, and the reward model is a scalar Gaussian with mean parameterized by a feed-forward neural network and unit variance. Note that the log-likelihood under a Gaussian distribution with unit variance equals the mean squared error up to a constant.

其中我们在不损失一般性的前提下假定初始状态 $s_0$ 固定。转移模型是高斯的，均值和方差由前馈神经网络参数化；观测模型是高斯的，均值由反卷积神经网络参数化，协方差为单位阵；奖励模型是标量高斯，均值由前馈神经网络参数化，方差为 1。注意，单位方差高斯分布下的对数似然，在相差一个常数的意义下等于均方误差。

Variational encoder Since the model is non-linear, we cannot directly compute the state posteriors that are needed for parameter learning. Instead, we use an encoder $q(s_{1:T} \mid o_{1:T}, a_{1:T}) = \prod_{t=1}^{T} q(s_t \mid s_{t-1}, a_{t-1}, o_t)$ to infer approximate state posteriors from past observations and actions, where $q(s_t \mid s_{t-1}, a_{t-1}, o_t)$ is a diagonal Gaussian with mean and variance parameterized by a convolutional neural network followed by a feed-forward neural network. We use the filtering posterior that conditions on past observations since we are ultimately interested in using the model for planning, but one may also use the full smoothing posterior during training (Babaeizadeh et al., 2017; Gregor & Besse, 2018).

变分编码器 由于模型是非线性的，我们无法直接计算参数学习所需的状态后验。于是我们使用编码器 $q(s_{1:T} \mid o_{1:T}, a_{1:T}) = \prod_{t=1}^{T} q(s_t \mid s_{t-1}, a_{t-1}, o_t)$，从过去的观测和动作推断近似状态后验，其中 $q(s_t \mid s_{t-1}, a_{t-1}, o_t)$ 是对角高斯，均值和方差由卷积神经网络再接前馈神经网络来参数化。我们使用以过去观测为条件的滤波后验，因为我们最终关心的是用该模型做规划；不过在训练时也可以使用完整的平滑后验（Babaeizadeh et al., 2017; Gregor & Besse, 2018）。

Training objective Using the encoder, we construct a variational bound on the data log-likelihood. For simplicity, we write losses for predicting only the observations — the reward losses follow by analogy. The variational bound obtained using Jensen’s inequality is

训练目标 我们用编码器构造数据对数似然的变分下界。为简单起见，我们只写出预测观测的损失——奖励损失可类推。用 Jensen 不等式得到的变分界是

$$\begin{aligned}
\ln p(o_{1:T} \mid a_{1:T}) &\triangleq \ln \int \prod_{t} p(s_t \mid s_{t-1}, a_{t-1}) p(o_t \mid s_t), ds_{1:T} \
&\ge \sum_{t=1}^{T} \Big(
\underbrace{\mathrm{E}{q(s_t \mid o{\le t}, a_{<t})}[\ln p(o_t \mid s_t)]}{\text{reconstruction}} \
&\quad - \underbrace{\mathrm{E}{q(s_{t-1} \mid o_{\le t-1}, a_{<t-1})}[\mathrm{KL}[q(s_t \mid o_{\le t}, a_{<t}) \parallel p(s_t \mid s_{t-1}, a_{t-1})]]}_{\text{complexity}}
\Big).
\end{aligned} \tag{3}$$

For the derivation, please see Equation 8 in the appendix. Estimating the outer expectations using a single reparameterized sample yields an efficient objective for inference and learning in non-linear latent variable models that can be optimized using gradient ascent (Kingma & Welling, 2013; Rezende et al., 2014; Krishnan et al., 2017).

推导见附录中的方程 8。用单个重参数化样本估计外层期望，会得到一个可用于非线性潜变量模型中推断与学习的高效目标，并可用梯度上升来优化（Kingma & Welling, 2013; Rezende et al., 2014; Krishnan et al., 2017）。

Deterministic path Despite its generality, the purely stochastic transitions make it difficult for the transition model to reliably remember information for multiple time steps. In theory, this model could learn to set the variance to zero for some state components, but the optimization procedure may not find this solution. This motivates including a deterministic sequence of activation vectors $\{h_t\}_{t=1}^T$ that allow the model to access not just the last state but all previous states deterministically (Chung et al., 2015; Buesing et al., 2018). We use such a model, shown in Figure 2c, that we name recurrent state-space model (RSSM),

确定性路径 尽管纯粹随机的转移具有一般性，但它使转移模型难以可靠地记住多个时间步的信息。理论上，这个模型可以学会把某些状态分量的方差设为零，但优化过程未必找得到这个解。这促使我们加入一串确定性的激活向量 $\{h_t\}_{t=1}^T$，让模型不仅能访问上一个状态，还能以确定性的方式访问所有先前状态（Chung et al., 2015; Buesing et al., 2018）。我们使用如图 2c 所示的这种模型，并把它叫做循环状态空间模型（RSSM），

$$\begin{aligned}
\text{Deterministic state model:} &\quad h_t = f(h_{t-1}, s_{t-1}, a_{t-1}) \
\text{Stochastic state model:} &\quad s_t \sim p(s_t \mid h_t) \
\text{Observation model:} &\quad o_t \sim p(o_t \mid h_t, s_t) \
\text{Reward model:} &\quad r_t \sim p(r_t \mid h_t, s_t),
\end{aligned} \tag{4}$$

where $f(h_{t-1}, s_{t-1}, a_{t-1})$ is implemented as a recurrent neural network (RNN). Intuitively, we can understand this model as splitting the state into a stochastic part $s_t$ and a deterministic part $h_t$, which depend on the stochastic and deterministic parts at the previous time step through the RNN. We use the encoder $q(s_{1:T} \mid o_{1:T}, a_{1:T}) = \prod_{t=1}^{T} q(s_t \mid h_t, o_t)$ to parameterize the approximate state posteriors. Importantly, all information about the observations must pass through the sampling step of the encoder to avoid a deterministic shortcut from inputs to reconstructions.

其中 $f(h_{t-1}, s_{t-1}, a_{t-1})$ 实现为循环神经网络（RNN）。直观上，可以把这个模型理解成把状态拆成随机部分 $s_t$ 和确定性部分 $h_t$，它们通过 RNN 依赖于上一时间步的随机部分和确定性部分。我们用编码器 $q(s_{1:T} \mid o_{1:T}, a_{1:T}) = \prod_{t=1}^{T} q(s_t \mid h_t, o_t)$ 来参数化近似状态后验。重要的是，关于观测的全部信息都必须经过编码器的采样步骤，以避免从输入到重构的确定性捷径。

In the next section, we identify a limitation of the standard objective for latent sequence models and propose a generalization of it that improves long-term predictions.

下一节中，我们指出潜变量序列模型标准目标的一个局限，并提出它的一种推广，以改进长期预测。

Figure 2: Latent dynamics model designs. In this example, the model observes the first two time steps and predicts the third. Circles represent stochastic variables and squares deterministic variables. Solid lines denote the generative process and dashed lines the inference model. (a) Transitions in a recurrent neural network are purely deterministic. This prevents the model from capturing multiple futures and makes it easy for the planner to exploit inaccuracies. (b) Transitions in a state-space model are purely stochastic. This makes it difficult to remember information over multiple time steps. (c) We split the state into stochastic and deterministic parts, allowing the model to robustly learn to predict multiple futures.

图 2：潜变量动力学模型的设计。在这个例子中，模型观测前两个时间步并预测第三个。圆圈表示随机变量，方块表示确定性变量。实线表示生成过程，虚线表示推断模型。(a) 循环神经网络中的转移是纯粹确定性的。这使模型无法刻画多种未来，也让规划器容易利用模型的不准确之处。(b) 状态空间模型中的转移是纯粹随机的。这使得在多个时间步上记住信息变得困难。(c) 我们把状态拆成随机部分和确定性部分，让模型能够稳健地学会预测多种未来。

4. Latent Overshooting

4. 潜空间过冲

In the previous section, we derived the typical variational bound for learning and inference in latent sequence models (Equation 3). As show in Figure 3a, this objective function contains reconstruction terms for the observations and KL-divergence regularizers for the approximate posteriors. A limitation of this objective is that the stochastic path of the transition function $p(s_t \mid s_{t-1}, a_{t-1})$ is only trained via the KL-divergence regularizers for one-step predictions: the gradient flows through $p(s_t \mid s_{t-1}, a_{t-1})$ directly into $q(s_{t-1})$ but never traverses a chain of multiple $p(s_t \mid s_{t-1}, a_{t-1})$. In this section, we generalize this variational bound to latent overshooting, which trains all multi-step predictions in latent space. We found that several dynamics models benefit from latent overshooting, although our final agent using the RSSM model does not require it (see Appendix D).

上一节我们推导了潜变量序列模型中用于学习与推断的典型变分界（方程 3）。如图 3a 所示，这一目标函数包含观测的重构项，以及对近似后验的 KL 散度正则。这一目标的局限在于，转移函数 $p(s_t \mid s_{t-1}, a_{t-1})$ 的随机路径只通过一步预测的 KL 散度正则来训练：梯度经由 $p(s_t \mid s_{t-1}, a_{t-1})$ 直接流入 $q(s_{t-1})$，但从未穿过一串多个 $p(s_t \mid s_{t-1}, a_{t-1})$。本节把这一变分界推广为潜空间过冲，在潜空间中训练所有多步预测。我们发现若干动力学模型都能从潜空间过冲中受益，尽管我们最终使用 RSSM 模型的智能体并不需要它（见附录 D）。

Figure 3: Unrolling schemes. The labels $s_{i|j}$ are short for the state at time $i$ conditioned on observations up to time $j$. Arrows pointing at shaded circles indicate log-likelihood loss terms. Wavy arrows indicate KL-divergence loss terms. (a) The standard variational objectives decodes the posterior at every step to compute the reconstruction loss. It also places a KL on the prior and posterior at every step, which trains the transition function for one-step predictions. (b) Observation overshooting (Amos et al., 2018) decodes all multi-step predictions to apply additional reconstruction losses. This is typically too expensive in image domains. (c) Latent overshooting predicts all multi-step priors. These state beliefs are trained towards their corresponding posteriors in latent space to encourage accurate multi-step predictions.

图 3：展开方案。标记 $s_{i|j}$ 是时刻 $i$、以截至时刻 $j$ 的观测为条件的状态的简写。指向着色圆圈的箭头表示对数似然损失项。波浪箭头表示 KL 散度损失项。(a) 标准变分目标在每一步解码后验以计算重构损失。它还在每一步对先验和后验施加 KL，从而按一步预测训练转移函数。(b) 观测过冲（Amos et al., 2018）解码所有多步预测，以施加额外的重构损失。这在图像领域通常过于昂贵。(c) 潜空间过冲预测所有多步先验。这些状态信念在潜空间中被训练去靠近对应的后验，以鼓励准确的多步预测。

Limited capacity If we could train our model to make perfect one-step predictions, it would also make perfect multi-step predictions, so this would not be a problem. However, when using a model with limited capacity and restricted distributional family, training the model only on one-step predictions until convergence does in general not coincide with the model that is best at multi-step predictions. For successful planning, we need accurate multi-step predictions. Therefore, we take inspiration from Amos et al. (2018) and earlier related ideas (Krishnan et al., 2015; Lamb et al., 2016; Chiappa et al., 2017), and train the model on multi-step predictions of all distances. We develop this idea for latent sequence models, showing that multi-step predictions can be improved by a loss in latent space, without having to generate additional images.

有限容量 如果我们能把模型训练到一步预测是完美的，那它的多步预测也会是完美的，也就不会有这个问题。然而，当模型容量有限、分布族也受限制时，只在一步预测上训练到收敛，一般并不等于在多步预测上最好的那个模型。要成功规划，我们需要准确的多步预测。因此，我们从 Amos et al.（2018）以及更早的相关想法（Krishnan et al., 2015; Lamb et al., 2016; Chiappa et al., 2017）中获得启发，在所有距离的多步预测上训练模型。我们把这一想法发展用于潜变量序列模型，并表明可以通过潜空间中的损失来改进多步预测，而不必再生成额外的图像。

Multi-step prediction We start by generalizing the standard variational bound (Equation 3) from training one-step predictions to training multi-step predictions of a fixed distance $d$. For ease of notation, we omit actions in the conditioning set here; every distribution over $s_t$ is conditioned upon $a_{<t}$. We first define multi-step predictions, which are computed by repeatedly applying the transition model and integrating out the intermediate states,

多步预测 我们先把标准变分界（方程 3）从训练一步预测，推广到训练固定距离 $d$ 的多步预测。为便于记号，这里在条件集合中省略动作；每一个关于 $s_t$ 的分布都以 $a_{<t}$ 为条件。我们先定义多步预测：它通过反复应用转移模型并积掉中间状态来计算，

$$\begin{aligned}
p(s_t \mid s_{t-d}) &\triangleq \int \prod_{\tau=t-d+1}^{t} p(s_\tau \mid s_{\tau-1}), ds_{t-d+1:t-1} \
&= \mathrm{E}{p(s{t-1} \mid s_{t-d})}[p(s_t \mid s_{t-1})].
\end{aligned} \tag{5}$$

The case $d=1$ recovers the one-step transitions used in the original model. Given this definition of a multi-step prediction, we generalize Equation 3 to the variational bound on the multi-step predictive distribution $p_d$,

$d=1$ 的情形恢复为原模型中使用的一步转移。有了这一多步预测的定义，我们把方程 3 推广为多步预测分布 $p_d$ 上的变分界，

$$\begin{aligned}
\ln p_d(o_{1:T}) &\triangleq \ln \int \prod_{t=1}^{T} p(s_t \mid s_{t-d}) p(o_t \mid s_t), ds_{1:T} \
&\ge \sum_{t=1}^{T} \Big(
\underbrace{\mathrm{E}{q(s_t \mid o{\le t})}[\ln p(o_t \mid s_t)]}{\text{reconstruction}} \
&\quad - \underbrace{\mathrm{E}{p(s_{t-1} \mid s_{t-d}) q(s_{t-d} \mid o_{\le t-d})}[\mathrm{KL}[q(s_t \mid o_{\le t}) \parallel p(s_t \mid s_{t-1})]]}_{\text{multi-step prediction}}
\Big).
\end{aligned} \tag{6}$$

For the derivation, please see Equation 9 in the appendix. Maximizing this objective trains the multi-step predictive distribution. This reflects the fact that during planning, the model makes predictions without having access to all the preceding observations.

推导见附录中的方程 9。最大化这一目标，就是在训练多步预测分布。这反映了规划时模型做预测却拿不到此前全部观测这一事实。

We conjecture that Equation 6 is also a lower bound on $\ln p(o_{1:T})$ based on the data processing inequality. Since the latent state sequence is Markovian, for $d \ge 1$ we have $\mathrm{I}(s_t; s_{t-d}) \le \mathrm{I}(s_t; s_{t-1})$ and thus $\mathrm{E}[\ln p_d(o_{1:T})] \le \mathrm{E}[\ln p(o_{1:T})]$. Hence, every bound on the multi-step predictive distribution is also a bound on the one-step predictive distribution in expectation over the data set. For details, please see Equation 10 in the appendix. In the next paragraph, we alleviate the limitation that a particular $p_d$ only trains predictions of one distance and arrive at our final objective.

我们猜想，基于数据处理不等式，方程 6 也是 $\ln p(o_{1:T})$ 的一个下界。由于潜状态序列是马尔可夫的，对 $d \ge 1$ 有 $\mathrm{I}(s_t; s_{t-d}) \le \mathrm{I}(s_t; s_{t-1})$，因此 $\mathrm{E}[\ln p_d(o_{1:T})] \le \mathrm{E}[\ln p(o_{1:T})]$。于是，在数据集上取期望时，多步预测分布上的每一个界，也是一步预测分布上的界。细节见附录中的方程 10。下一段中，我们消除“某一个 $p_d$ 只训练一种距离的预测”这一局限，得到最终目标。

Latent overshooting We introduced a bound on predictions of a given distance $d$. However, for planning we need accurate predictions not just for a fixed distance but for all distances up to the planning horizon. We introduce latent overshooting for this, an objective function for latent sequence models that generalizes the standard variational bound (Equation 3) to train the model on multi-step predictions of all distances $1 \le d \le D$,

潜空间过冲 我们引入了给定距离 $d$ 上预测的一个界。然而，规划需要的不只是某一固定距离上的准确预测，而是直到规划视野的所有距离。为此我们引入潜空间过冲：这是潜变量序列模型的一个目标函数，它把标准变分界（方程 3）推广为在所有距离 $1 \le d \le D$ 的多步预测上训练模型，

$$\begin{aligned}
\frac{1}{D}\sum_{d=1}^{D} \ln p_d(o_{1:T})
&\ge \sum_{t=1}^{T} \Big(
\underbrace{\mathrm{E}{q(s_t \mid o{\le t})}[\ln p(o_t \mid s_t)]}{\text{reconstruction}} \
&\quad - \underbrace{\frac{1}{D}\sum{d=1}^{D} \beta_d, \mathrm{E}{p(s{t-1} \mid s_{t-d}) q(s_{t-d} \mid o_{\le t-d})}[\mathrm{KL}[q(s_t \mid o_{\le t}) \parallel p(s_t \mid s_{t-1})]]}_{\text{latent overshooting}}
\Big).
\end{aligned} \tag{7}$$

Latent overshooting can be interpreted as a regularizer in latent space that encourages consistency between one-step and multi-step predictions, which we know should be equivalent in expectation over the data set. We include weighting factors \{eta_d\}_{d=1}^{D} analogously to the $\beta$-VAE (Higgins et al., 2016). While we set all $\beta_{>1}$ to the same value for simplicity, they could be chosen to let the model focus more on long-term or short-term predictions. In practice, we stop gradients of the posterior distributions for overshooting distances $d>1$, so that the multi-step predictions are trained towards the informed posteriors, but not the other way around.

潜空间过冲可以解释成潜空间中的一种正则，它鼓励一步预测与多步预测之间的一致性——我们知道二者在数据集上取期望时应当等价。我们加入权重因子 \{eta_d\}_{d=1}^{D}，做法类似于 $\beta$-VAE（Higgins et al., 2016）。虽然为简单起见我们把所有 $\beta_{>1}$ 设成同一个值，但也可以把它们选成让模型更关注长期或短期预测。实践中，我们对过冲距离 $d>1$ 停止后验分布的梯度，从而使多步预测被训练去靠近有信息的后验，而不是反过来。

5. Experiments

5. 实验

We evaluate PlaNet on six continuous control tasks from pixels. We explore multiple design axes of the agent: the stochastic and deterministic paths in the dynamics model, iterative planning, and online experience collection. We refer to the appendix for hyper parameters (Appendix A) and additional experiments (Appendices C to E). Besides the action repeat, we use the same hyper parameters for all tasks. Within less than one hundredth the episodes, PlaNet outperforms A3C (Mnih et al., 2016) and achieves similar performance to the top model-free algorithm D4PG (Barth-Maron et al., 2018). The training time of 10 to 20 hours (depending on the task) on a single Nvidia V100 GPU compares favorably to that of A3C and D4PG. Our implementation uses TensorFlow Probability (Dillon et al., 2017). Please visit https://danijar.com/planet for access to the code and videos of the trained agent.

我们从像素出发，在六项连续控制任务上评估 PlaNet。我们探索智能体的多个设计轴：动力学模型中的随机与确定性路径、迭代规划，以及在线经验收集。超参数见附录 A，补充实验见附录 C 至 E。除动作重复外，所有任务使用相同的超参数。在不到百分之一的回合数内，PlaNet 超过了 A3C（Mnih et al., 2016），并达到与顶尖模型自由算法 D4PG（Barth-Maron et al., 2018）相近的表现。在单块 Nvidia V100 GPU 上，训练时间为 10 到 20 小时（视任务而定），与 A3C 和 D4PG 相比并不逊色。我们的实现使用 TensorFlow Probability（Dillon et al., 2017）。代码和训练后智能体的视频见 https://danijar.com/planet 。

For our evaluation, we consider six image-based continuous control tasks of the DeepMind control suite (Tassa et al., 2018), shown in Figure 1. These environments provide qualitatively different challenges. The cartpole swingup task requires a long planning horizon and to memorize the cart when it is out of view, reacher has a sparse reward given when the hand and goal area overlap, finger spinning includes contact dynamics between the finger and the object, cheetah exhibits larger state and action spaces, the cup task only has a sparse reward for when the ball is caught, and the walker is challenging because the robot first has to stand up and then walk, resulting in collisions with the ground that are difficult to predict. In all tasks, the only observations are third-person camera images of size $64 \times 64 \times 3$ pixels.

评估中，我们考虑 DeepMind 控制套件（Tassa et al., 2018）里的六项基于图像的连续控制任务，如图 1 所示。这些环境提供了性质不同的挑战。倒立摆起摆任务需要长规划视野，并在小车移出视野时记住它；伸手任务在手与目标区域重叠时给出稀疏奖励；手指旋转包含手指与物体之间的接触力学；猎豹具有更大的状态和动作空间；杯子任务只在球被接住时给出稀疏奖励；步行者则更具挑战，因为机器人必须先站起来再走路，从而与地面发生难以预测的碰撞。在所有任务中，唯一的观测是大小为 $64 \times 64 \times 3$ 像素的第三人称相机图像。

Comparison to model-free methods Figure 4 compares the performance of PlaNet to the model-free algorithms reported by Tassa et al. (2018). Within 100 episodes, PlaNet outperforms the policy-gradient method A3C trained from proprioceptive states for 100,000 episodes, on all tasks. After 500 episodes, it achieves performance similar to D4PG, trained from images for 100,000 episodes, except for the finger task. PlaNet surpasses the final performance of D4PG with a relative improvement of 26% on the cheetah running task. We refer to Table 1 for numerical results, which also includes the performance of CEM planning with the true dynamics of the simulator.

与模型自由方法的比较 图 4 把 PlaNet 的表现与 Tassa et al.（2018）报告的模型自由算法做了比较。在 100 个回合之内，PlaNet 在全部任务上都超过了用本体感受状态训练 100,000 回合的策略梯度方法 A3C。到 500 个回合时，除手指任务外，它已达到与用图像训练 100,000 回合的 D4PG 相近的表现。在猎豹奔跑任务上，PlaNet 以 26% 的相对提升超过了 D4PG 的最终表现。数值结果见表 1，其中还包括用仿真器真实动力学做 CEM 规划的表现。

Model designs Figure 4 additionally compares design choices of the dynamics model. We train PlaNet using our recurrent state-space model (RSSM), as well as versions with purely deterministic GRU (Cho et al., 2014), and purely stochastic state-space model (SSM). We observe the importance of both stochastic and deterministic elements in the transition function on all tasks. The deterministic part allows the model to remember information over many time steps. The stochastic component is even more important — the agent does not learn without it. This could be because the tasks are stochastic from the agent's perspective due to partial observability of the initial states. The noise might also add a safety margin to the planning objective that results in more robust action sequences.

模型设计 图 4 还比较了动力学模型的设计选择。我们用循环状态空间模型（RSSM）训练 PlaNet，同时也训练了纯粹确定性的 GRU 版本（Cho et al., 2014），以及纯粹随机的状态空间模型（SSM）版本。我们在全部任务上都观察到，转移函数中的随机元素和确定性元素都很重要。确定性部分让模型能够在许多时间步上记住信息。随机分量甚至更重要——没有它，智能体就学不会。这可能是因为，由于初始状态只是部分可观测的，这些任务从智能体的视角看是随机的。噪声也许还为规划目标增加了一道安全裕度，从而得到更稳健的动作序列。

Agent designs Figure 5 compares PlaNet, a version collecting episodes under random actions rather than by planning, and a version that at each environment step selects the best action out of 1000 sequences rather than iteratively refining plans via CEM. We observe that online data collection helps for all tasks and is necessary for the cartpole, finger, and walker tasks. Iterative search for action sequences using CEM improves performance on all tasks.

智能体设计 图 5 比较了 PlaNet、一个不靠规划而是在随机动作下收集回合的版本，以及一个在每个环境步上从 1000 条序列中选出最佳动作、而不用 CEM 迭代改进计划的版本。我们观察到，在线数据收集对所有任务都有帮助，并且对倒立摆、手指和步行者任务是必要的。用 CEM 迭代搜索动作序列，在全部任务上都提高了表现。

One agent all tasks Figure 7 in the appendix shows the performance of a single agent trained on all six tasks. The agent is not told which task it is facing; it needs to infer this from the image observations. We pad the action spaces with unused elements to make them compatible and adapt Algorithm 1 to collect one episode of each task every $C$ update steps. We use the same hyper parameters as for the main experiments above. The agent solves all tasks while learning slower compared to individually trained agents. This indicates that the model can learn to predict multiple domains, regardless of the conceptually different visuals.

一个智能体应对全部任务 附录中的图 7 展示了在全部六项任务上训练的单个智能体的表现。智能体并不知道自己面对的是哪一项任务，需要从图像观测中推断出来。我们用未使用的分量把动作空间补齐以使其兼容，并把算法 1 改成每隔 $C$ 次更新为每项任务各收集一条回合。超参数与上面的主要实验相同。该智能体解决了全部任务，但学习比分别训练的智能体更慢。这表明，尽管各领域的视觉在概念上不同，模型仍能学会预测多个领域。

Figure 4: Comparison of PlaNet to model-free algorithms and other model designs. Plots show test performance over the number of collected episodes. We compare PlaNet using our RSSM (Section 3) to purely deterministic (GRU) and purely stochastic models (SSM). The RNN does not use latent overshooting, as it does not have stochastic latents. The lines show medians and the areas show percentiles 5 to 95 over 5 seeds and 10 trajectories. The shaded areas are large on two of the tasks due to the sparse rewards.

图 4：PlaNet 与模型自由算法以及其他模型设计的比较。图中给出测试表现随所收集回合数的变化。我们把使用 RSSM（第 3 节）的 PlaNet，与纯粹确定性（GRU）和纯粹随机（SSM）的模型相比较。RNN 没有随机潜变量，因此不使用潜空间过冲。折线表示中位数，阴影区域表示 5 个随机种子和 10 条轨迹上的第 5 到第 95 百分位。由于稀疏奖励，其中两项任务上的阴影区域较大。

Table 1: Comparison of PlaNet to the model-free algorithms A3C and D4PG reported by Tassa et al. (2018). The training curves for these are shown as orange lines in Figure 4 and as solid green lines in Figure 6 in their paper. From these, we estimate the number of episodes that D4PG takes to achieve the final performance of PlaNet to estimate the data efficiency gain. We further include CEM planning ($H = 12$, $I = 10$, $J = 1000$, $K = 100$) with the true simulator instead of learned dynamics as an estimated upper bound on performance. Numbers indicate mean final performance over 5 seeds and 10 trajectories.

Method	Modality	Episodes	Cartpole Swing Up	Reacher Easy	Cheetah Run	Finger Spin	Cup Catch	Walker Walk
A3C	proprioceptive	100,000	558	285	214	129	105	311
D4PG	pixels	100,000	862	967	524	985	980	968
PlaNet (ours)	pixels	1,000	821	832	662	700	930	951
CEM + true simulator	simulator state	0	850	964	656	825	993	994
Data efficiency gain PlaNet over D4PG (factor)			250	40	500+	300	100	90

表 1：PlaNet 与 Tassa et al.（2018）报告的模型自由算法 A3C 和 D4PG 的比较。后者的训练曲线在图 4 中用橙色线表示，在他们论文的图 6 中用实心绿线表示。据此，我们估计 D4PG 达到 PlaNet 最终表现所需的回合数，以估计数据效率增益。我们还加入了用真实仿真器而非所学动力学、参数为 $H = 12$、$I = 10$、$J = 1000$、$K = 100$ 的 CEM 规划，作为表现的估计上界。数字为 5 个随机种子和 10 条轨迹上的平均最终表现。

方法	模态	回合数	倒立摆起摆	伸手（简单）	猎豹奔跑	手指旋转	杯子接球	步行者走路
A3C	本体感受	100,000	558	285	214	129	105	311
D4PG	像素	100,000	862	967	524	985	980	968
PlaNet（本文）	像素	1,000	821	832	662	700	930	951
CEM + 真实仿真器	仿真器状态	0	850	964	656	825	993	994
PlaNet 相对 D4PG 的数据效率增益（倍数）			250	40	500+	300	100	90

Figure 5: Comparison of agent designs. Plots show test performance over the number of collected episodes. We compare PlaNet, a version that collects data under random actions (random collection), and a version that chooses the best action out of 1000 sequences at each environment step (random shooting) without iteratively refining plans via CEM. The lines show medians and the areas show percentiles 5 to 95 over 5 seeds and 10 trajectories.

图 5：智能体设计的比较。图中给出测试表现随所收集回合数的变化。我们比较 PlaNet、一个在随机动作下收集数据的版本（随机收集），以及一个在每个环境步从 1000 条序列中选出最佳动作、而不用 CEM 迭代改进计划的版本（随机射击）。折线表示中位数，阴影区域表示 5 个随机种子和 10 条轨迹上的第 5 到第 95 百分位。

6. Related Work

6. 相关工作

Previous work in model-based reinforcement learning has focused on planning in low-dimensional state spaces (Gal et al., 2016; Higuera et al., 2018; Henaff et al., 2018; Chua et al., 2018), combining the benefits of model-based and model-free approaches (Kalweit & Boedecker, 2017; Nagabandi et al., 2017; Weber et al., 2017; Kurutach et al., 2018; Buckman et al., 2018; Ha & Schmidhuber, 2018; Wayne et al., 2018; Igl et al., 2018; Srinivas et al., 2018), and pure video prediction without planning (Oh et al., 2015; Krishnan et al., 2015; Karl et al., 2016; Chiappa et al., 2017; Babaeizadeh et al., 2017; Gemici et al., 2017; Denton & Fergus, 2018; Buesing et al., 2018; Doerr et al., 2018; Gregor & Besse, 2018). Appendix G reviews these orthogonal research directions in more detail.

此前基于模型的强化学习工作，主要集中在低维状态空间中的规划（Gal et al., 2016; Higuera et al., 2018; Henaff et al., 2018; Chua et al., 2018），结合基于模型与模型自由方法的长处（Kalweit & Boedecker, 2017; Nagabandi et al., 2017; Weber et al., 2017; Kurutach et al., 2018; Buckman et al., 2018; Ha & Schmidhuber, 2018; Wayne et al., 2018; Igl et al., 2018; Srinivas et al., 2018），以及不做规划的纯视频预测（Oh et al., 2015; Krishnan et al., 2015; Karl et al., 2016; Chiappa et al., 2017; Babaeizadeh et al., 2017; Gemici et al., 2017; Denton & Fergus, 2018; Buesing et al., 2018; Doerr et al., 2018; Gregor & Besse, 2018）。附录 G 对这些正交的研究方向有更详细的评述。

Relatively few works have demonstrated successful planning from pixels using learned dynamics models. The robotics community focuses on video prediction models for planning (Agrawal et al., 2016; Finn & Levine, 2017; Ebert et al., 2018; Zhang et al., 2018) that deal with the visual complexity of the real world and solve tasks with a simple gripper, such as grasping or pushing objects. In comparison, we focus on simulated environments, where we leverage latent planning to scale to larger state and action spaces, longer planning horizons, as well as sparse reward tasks. E2C (Watter et al., 2015) and RCE (Banijamali et al., 2017) embed images into a latent space, where they learn local-linear latent transitions and plan for actions using LQR. These methods balance simulated cartpoles and control 2-link arms from images, but have been difficult to scale up. We lift the Markov assumption of these models, making our method applicable under partial observability, and present results on more challenging environments that include longer planning horizons, contact dynamics, and sparse rewards.

用所学动力学模型从像素成功做规划的工作相对较少。机器人社区关注用于规划的视频预测模型（Agrawal et al., 2016; Finn & Levine, 2017; Ebert et al., 2018; Zhang et al., 2018），它们要应对真实世界的视觉复杂性，并用简单夹爪解决抓取或推动物体等任务。相比之下，我们关注仿真环境，借助潜空间规划扩展到更大的状态和动作空间、更长的规划视野，以及稀疏奖励任务。E2C（Watter et al., 2015）和 RCE（Banijamali et al., 2017）把图像嵌入潜空间，在其中学习局部线性的潜转移，并用 LQR 规划动作。这些方法能从图像平衡仿真倒立摆、控制两连杆手臂，但一直难以扩大规模。我们取消了这些模型的马尔可夫假定，使方法适用于部分可观测情形，并在包含更长规划视野、接触力学和稀疏奖励的更有挑战的环境上给出结果。

7. Discussion

7. 讨论

We present PlaNet, a model-based agent that learns a latent dynamics model from image observations and chooses actions by fast planning in latent space. To enable accurate long-term predictions, we design a model with both stochastic and deterministic paths. We show that our agent succeeds at several continuous control tasks from image observations, reaching performance that is comparable to the best model-free algorithms while using $200\times$ fewer episodes and similar or less computation time. The results show that learning latent dynamics models for planning in image domains is a promising approach.

我们提出 PlaNet：一个基于模型的智能体，从图像观测学习潜变量动力学模型，并通过在潜空间中快速规划来选择动作。为了实现准确的长期预测，我们设计了同时带有随机路径和确定性路径的模型。我们表明，该智能体能够从图像观测成功完成多项连续控制任务，表现可与最好的模型自由算法相比，同时只用 $200\times$ 更少的回合，计算时间相近或更少。这些结果表明，学习潜变量动力学模型并在图像领域做规划，是一条有希望的路。

Directions for future work include learning temporal abstraction instead of using a fixed action repeat, possibly through hierarchical models. To further improve final performance, one could learn a value function to approximate the sum of rewards beyond the planning horizon. Moreover, gradient-based planning could increase the computational efficiency of the agent and learning representations without reconstruction could help to solve tasks with higher visual diversity. Our work provides a starting point for multi-task control by sharing the dynamics model.

未来工作的方向包括学习时间抽象，而不是使用固定的动作重复，这或许可以通过层次模型来做。为进一步提高最终表现，可以学习一个价值函数，来近似规划视野之外的奖励之和。此外，基于梯度的规划可以提高智能体的计算效率；不借助重构来学习表示，则可能有助于解决视觉多样性更高的任务。我们的工作通过共享动力学模型，为多任务控制提供了一个起点。

Acknowledgements

We thank Jacob Buckman, Nicolas Heess, John Schulman, Rishabh Agarwal, Silviu Pitis, Mohammad Norouzi, George Tucker, David Duvenaud, Shane Gu, Chelsea Finn, Steven Bohez, Jimmy Ba, Stephanie Chan, and Jenny Liu for helpful discussions.

我们感谢 Jacob Buckman、Nicolas Heess、John Schulman、Rishabh Agarwal、Silviu Pitis、Mohammad Norouzi、George Tucker、David Duvenaud、Shane Gu、Chelsea Finn、Steven Bohez、Jimmy Ba、Stephanie Chan 和 Jenny Liu 的有益讨论。

References

Agrawal, P., Nair, A. V., Abbeel, P., Malik, J., and Levine, S. Learning to poke by poking: Experiential learning of intuitive physics. In Advances in Neural Information Processing Systems, pp. 5074–5082, 2016.

Amos, B., Dinh, L., Cabi, S., Rothörl, T., Muldal, A., Erez, T., Tassa, Y., de Freitas, N., and Denil, M. Learning awareness models. In International Conference on Learning Representations, 2018.

Babaeizadeh, M., Finn, C., Erhan, D., Campbell, R. H., and Levine, S. Stochastic variational video prediction. arXiv preprint arXiv:1710.11252, 2017.

Banijamali, E., Shu, R., Ghavamzadeh, M., Bui, H., and Ghodsi, A. Robust locally-linear controllable embedding. arXiv preprint arXiv:1710.05373, 2017.

Barth-Maron, G., Hoffman, M. W., Budden, D., Dabney, W., Horgan, D., Muldal, A., Heess, N., and Lillicrap, T. Distributed distributional deterministic policy gradients. arXiv preprint arXiv:1804.08617, 2018.

Bengio, S., Vinyals, O., Jaitly, N., and Shazeer, N. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Neural Information Processing Systems, pp. 1171–1179, 2015.

Buckman, J., Hafner, D., Tucker, G., Brevdo, E., and Lee, H. Sample-efficient reinforcement learning with stochastic ensemble value expansion. arXiv preprint arXiv:1807.01675, 2018.

Buesing, L., Weber, T., Racaniere, S., Eslami, S., Rezende, D., Reichert, D. P., Viola, F., Besse, F., Gregor, K., Hassabis, D., et al. Learning and querying fast generative models for reinforcement learning. arXiv preprint arXiv:1802.03006, 2018.

Chiappa, S., Racaniere, S., Wierstra, D., and Mohamed, S. Recurrent environment simulators. arXiv preprint arXiv:1704.02254, 2017.

Cho, K., Van Merriënboer, B., Gulcehre, C., Bahdanau, D., Bougares, F., Schwenk, H., and Bengio, Y. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078, 2014.

Chua, K., Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. arXiv preprint arXiv:1805.12114, 2018.

Chung, J., Kastner, K., Dinh, L., Goel, K., Courville, A. C., and Bengio, Y. A recurrent latent variable model for sequential data. In Advances in neural information processing systems, pp. 2980–2988, 2015.

Clevert, D.-A., Unterthiner, T., and Hochreiter, S. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289, 2015.

Deisenroth, M. and Rasmussen, C. E. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11), pp. 465–472, 2011.

Denton, E. and Fergus, R. Stochastic video generation with a learned prior. arXiv preprint arXiv:1802.07687, 2018.

Dillon, J. V., Langmore, I., Tran, D., Brevdo, E., Vasudevan, S., Moore, D., Patton, B., Alemi, A., Hoffman, M., and Saurous, R. A. Tensorflow distributions. arXiv preprint arXiv:1711.10604, 2017.

Doerr, A., Daniel, C., Schiegg, M., Nguyen-Tuong, D., Schaal, S., Toussaint, M., and Trimpe, S. Probabilistic recurrent state-space models. arXiv preprint arXiv:1801.10395, 2018.

Ebert, F., Finn, C., Dasari, S., Xie, A., Lee, A., and Levine, S. Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv:1812.00568, 2018.

Finn, C. and Levine, S. Deep visual foresight for planning robot motion. In Robotics and Automation (ICRA), 2017 IEEE International Conference on, pp. 2786–2793. IEEE, 2017.

Gal, Y., McAllister, R., and Rasmussen, C. E. Improving pilco with bayesian neural network dynamics models. In Data-Efficient Machine Learning workshop, ICML, 2016.

Gemici, M., Hung, C.-C., Santoro, A., Wayne, G., Mohamed, S., Rezende, D. J., Amos, D., and Lillicrap, T. Generative temporal models with memory. arXiv preprint arXiv:1702.04649, 2017.

Gregor, K. and Besse, F. Temporal difference variational auto-encoder. arXiv preprint arXiv:1806.03107, 2018.

Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv:1803.10122, 2018.

Henaff, M., Whitney, W. F., and LeCun, Y. Model-based planning with discrete and continuous actions. arXiv preprint arXiv:1705.07177, 2018.

Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. beta-vae: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2016.

Higuera, J. C. G., Meger, D., and Dudek, G. Synthesizing neural network controllers with probabilistic model based reinforcement learning. arXiv preprint arXiv:1803.02291, 2018.

Igl, M., Zintgraf, L., Le, T. A., Wood, F., and Whiteson, S. Deep variational reinforcement learning for pomdps. arXiv preprint arXiv:1806.02426, 2018.

Kalchbrenner, N., Oord, A. v. d., Simonyan, K., Danihelka, I., Vinyals, O., Graves, A., and Kavukcuoglu, K. Video pixel networks. arXiv preprint arXiv:1610.00527, 2016.

Kalweit, G. and Boedecker, J. Uncertainty-driven imagination for continuous deep reinforcement learning. In Conference on Robot Learning, pp. 195–206, 2017.

Karl, M., Soelch, M., Bayer, J., and van der Smagt, P. Deep variational bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint arXiv:1605.06432, 2016.

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.

Kingma, D. P. and Dhariwal, P. Glow: Generative flow with invertible 1x1 convolutions. arXiv preprint arXiv:1807.03039, 2018.

Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.

Krishnan, R. G., Shalit, U., and Sontag, D. Deep kalman filters. arXiv preprint arXiv:1511.05121, 2015.

Krishnan, R. G., Shalit, U., and Sontag, D. Structured inference networks for nonlinear state space models. In AAAI, pp. 2101–2109, 2017.

Kurutach, T., Clavera, I., Duan, Y., Tamar, A., and Abbeel, P. Model-ensemble trust-region policy optimization. arXiv preprint arXiv:1802.10592, 2018.

Lamb, A. M., GOYAL, A. G. A. P., Zhang, Y., Zhang, S., Courville, A. C., and Bengio, Y. Professor forcing: A new algorithm for training recurrent networks. In Advances In Neural Information Processing Systems, pp. 4601–4609, 2016.

Mathieu, M., Couprie, C., and LeCun, Y. Deep multi-scale video prediction beyond mean square error. arXiv preprint arXiv:1511.05440, 2015.

Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. Nature, 518(7540): 529, 2015.

Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., and Kavukcuoglu, K. Asynchronous methods for deep reinforcement learning. In International Conference on Machine Learning, pp. 1928–1937, 2016.

Moerland, T. M., Broekens, J., and Jonker, C. M. Learning multimodal transition dynamics for model-based reinforcement learning. arXiv preprint arXiv:1705.00470, 2017.

Moravčík, M., Schmid, M., Burch, N., Lisý, V., Morrill, D., Bard, N., Davis, T., Waugh, K., Johanson, M., and Bowling, M. Deepstack: Expert-level artificial intelligence in heads-up no-limit poker. Science, 356(6337):508–513, 2017.

Nagabandi, A., Kahn, G., Fearing, R. S., and Levine, S. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. arXiv preprint arXiv:1708.02596, 2017.

Nair, V. and Hinton, G. E. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference on machine learning (ICML-10), pp. 807–814, 2010.

Oh, J., Guo, X., Lee, H., Lewis, R. L., and Singh, S. Action-conditional video prediction using deep networks in atari games. In Advances in Neural Information Processing Systems, pp. 2863–2871, 2015.

Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint arXiv:1401.4082, 2014.

Richards, A. G. Robust constrained model predictive control. PhD thesis, Massachusetts Institute of Technology, 2005.

Rubinstein, R. Y. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1):89–112, 1997.

Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. Mastering the game of go without human knowledge. Nature, 550(7676):354, 2017.

Srinivas, A., Jabri, A., Abbeel, P., Levine, S., and Finn, C. Universal planning networks. arXiv preprint arXiv:1804.00645, 2018.

Talvitie, E. Model regularization for stable sample rollouts. In UAI, pp. 780–789, 2014.

Tassa, Y., Erez, T., and Todorov, E. Synthesis and stabilization of complex behaviors through online trajectory optimization. In Intelligent Robots and Systems (IROS), 2012 IEEE/RSJ International Conference on, pp. 4906–4913. IEEE, 2012.

Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., et al. Deepmind control suite. arXiv preprint arXiv:1801.00690, 2018.

van den Oord, A., Vinyals, O., et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, pp. 6309–6318, 2017.

Venkatraman, A., Hebert, M., and Bagnell, J. A. Improving multi-step prediction of learned time series models. In AAAI, pp. 3024–3030, 2015.

Vondrick, C., Pirsiavash, H., and Torralba, A. Generating videos with scene dynamics. In Advances In Neural Information Processing Systems, 2016.

Watter, M., Springenberg, J., Boedecker, J., and Riedmiller, M. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, pp. 2746–2754, 2015.

Wayne, G., Hung, C.-C., Amos, D., Mirza, M., Ahuja, A., Grabska-Barwinska, A., Rae, J., Mirowski, P., Leibo, J. Z., Santoro, A., et al. Unsupervised predictive memory in a goal-directed agent. arXiv preprint arXiv:1803.10760, 2018.

Weber, T., Racanière, S., Reichert, D. P., Buesing, L., Guez, A., Rezende, D. J., Badia, A. P., Vinyals, O., Heess, N., Li, Y., et al. Imagination-augmented agents for deep reinforcement learning. arXiv preprint arXiv:1707.06203, 2017.

Zhang, M., Vikram, S., Smith, L., Abbeel, P., Johnson, M., and Levine, S. SOLAR: deep structured representations for model-based reinforcement learning. arXiv preprint arXiv:1808.09105, 2018.