【模型式强化学习】（五）：PlaNet——从 RSSM 潜在状态建模到 MPC 动作规划

前言

本系列将从模型式强化学习的基本原理出发，逐步介绍经典方法与现代世界模型，结合核心公式、算法架构与核心代码，讨论智能体如何在学习行动的同时预测行动的后果，并利用预测结果指导决策
上一期我们介绍了世界模型的开山之作 World Models（2018） 。该方法将整体架构划分为 V、M、C 三个相对独立的模块：
- V 使用 VAE 将视觉观测压缩为低维潜在表示
- M 使用 MDN-RNN 对潜在状态的时序动态进行建模，并通过混合密度网络（MDN）预测未来潜在状态的概率分布；
- C 则采用协方差矩阵自适应进化策略（CMA-ES），优化控制器在学习到的世界模型中的行为表现。
- 三个模块分别承担视觉表征、环境动态建模与策略优化的任务，可以分阶段训练。训练完成后，控制器既可以在真实环境中运行，也可以在学习到的世界模型中进行闭环优化与评估。
本期我们要讲解的是 PlaNet（2019）。与 World Models 将世界模型与控制器分开训练、再通过 CMA-ES 优化控制器不同，PlaNet 将学习到的潜在动力学模型与模型预测控制（MPC）相结合，使智能体能够直接在潜在空间中预测未来状态、评估动作序列并进行在线规划。
最终，我们加载公开的预训练权重，在本地验证 Walker 任务的潜在空间规划闭环，效果如下：

文章目录

- 前言
- [1 回顾](#1 回顾)
- - - [1-1 PETS](#1-1 PETS)
    - [1-2 PETS 与 PlaNet 对比](#1-2 PETS 与 PlaNet 对比)
- [2 PlaNet](#2 PlaNet)
- - - [2-1 介绍](#2-1 介绍)
    - [2-2 结构](#2-2 结构)
    - - [2-2-1 总览](#2-2-1 总览)
      - [2-2-2 整体流程](#2-2-2 整体流程)
      - [2-2-3 三个要点](#2-2-3 三个要点)
    - [2-3 Recurrent State-Space Model（RSSM，循环状态空间模型）](#2-3 Recurrent State-Space Model（RSSM，循环状态空间模型）)
    - - [2-3-1 RSSM 模型含义](#2-3-1 RSSM 模型含义)
      - [2-3-2 先验与后验](#2-3-2 先验与后验)
      - [2-3-2-1 先验：预测当前状态](#2-3-2-1 先验：预测当前状态)
        
        [2-3-2-2 后验：结合观测推断当前状态](#2-3-2-2 后验：结合观测推断当前状态)
        
        [2-3-2-3 约束](#2-3-2-3 约束)
      - [2-3-3 RSSM 结构](#2-3-3 RSSM 结构)
      - [2-3-4 GRU](#2-3-4 GRU)
      - [2-3-5 先验分布构建](#2-3-5 先验分布构建)
      - [2-3-6 CNN](#2-3-6 CNN)
      - [2-3-6-1 结构](#2-3-6-1 结构)
      - [2-3-6-2 后验拼接](#2-3-6-2 后验拼接)
      - [2-3-7 后验估计](#2-3-7 后验估计)
      - [2-3-8 图像解码](#2-3-8 图像解码)
      - [2-3-9 奖励预测](#2-3-9 奖励预测)
    - [2-4 多步预测与Latent Overshooting（潜在超前预测）](#2-4 多步预测与Latent Overshooting（潜在超前预测）)
    - - [2-4-1 向前滚动 k 步](#2-4-1 向前滚动 k 步)
      - [2-4-2 代码里的滚动开关](#2-4-2 代码里的滚动开关)
      - [2-4-3 Latent Overshooting](#2-4-3 Latent Overshooting)
    - [2-5 Planning in Latent Spaces（潜在空间规划）](#2-5 Planning in Latent Spaces（潜在空间规划）)
    - - [2-5-1 CEM：如何搜索较优动作序列？](#2-5-1 CEM：如何搜索较优动作序列？)
      - [2-5-2 滚动时域规划：为什么只执行第一个动作？](#2-5-2 滚动时域规划：为什么只执行第一个动作？)
    - [2-6 数据流串联](#2-6 数据流串联)
    - [2-7 训练](#2-7 训练)
    - [2-8 部署](#2-8 部署)
    - [2-9 对比](#2-9 对比)
    - - [2-9-1 实验介绍](#2-9-1 实验介绍)
      - [2-9-2 实验结果](#2-9-2 实验结果)
- [3 Walker 实验](#3 Walker 实验)
- - - [3-1 任务与数据](#3-1 任务与数据)
    - - [3-1-1 walker-walk 任务](#3-1-1 walker-walk 任务)
      - [3-1-2 观测来源](#3-1-2 观测来源)
      - [3-1-3 奖励函数](#3-1-3 奖励函数)
      - [3-1-4 数据采集](#3-1-4 数据采集)
      - [3-1-5 序列采样](#3-1-5 序列采样)
    - [3-2 实验步骤](#3-2 实验步骤)
    - - [3-2-1 权重加载](#3-2-1 权重加载)
      - [3-2-2 评测流程](#3-2-2 评测流程)
    - [3-3 实验结果](#3-3 实验结果)
- 总结

1 回顾

PlaNet 与 PETS 都属于 MBRL（Model-Based Reinforcement Learning，模型式强化学习），并且都采用 MPC（Model Predictive Control，模型预测控制）与 CEM（Cross-Entropy Method，交叉熵方法）进行动作规划。本节先快速回顾 PETS 的核心思想，再分析 PlaNet 在此基础上的改进

1-1 PETS

PETS 的核心可以概括为 PE + TS + MPC ：
- PE（Probabilistic Ensembles，概率模型集成）：训练多个概率神经网络，通过 Bootstrap 采样构建模型集成。各个模型分别预测下一状态的概率分布，从而同时刻画数据本身的随机性与不同模型之间的预测分歧
- TS（Trajectory Sampling，轨迹采样）：在多步预测过程中处理模型不确定性，主要包括 TS1 和 TS∞ 两种方式。TS1 在每个预测时刻重新选择集成模型；TS∞ 则在一条预测轨迹中固定使用同一个模型
- MPC（Model Predictive Control，模型预测控制）：利用学习到的动态模型评估候选动作序列，通过 CEM 搜索较优方案，执行当前动作后，再根据新状态重新规划
也就是说，PETS 通过概率模型集成刻画环境动态的不确定性，利用轨迹采样预测未来状态，并结合 MPC 与 CEM 在模型中进行滚动动作规划。

1-2 PETS 与 PlaNet 对比

下表中的 RSSM 指循环状态空间模型（Recurrent State-Space Model），其结构在 2-3 节展开

对比维度	PETS	PlaNet
状态输入	通常是低维状态向量	从图像学习潜在状态
动态模型	概率神经网络集成	RSSM 潜在动态模型
不确定性	概率预测与模型集成	潜在状态中的随机成分
规划方式	MPC + CEM	MPC + CEM
核心难点	动态预测误差与不确定性	视觉表征与潜在动态学习

可以看到，PlaNet 并没有抛弃 PETS 的规划思路，而是将其拓展到视觉输入场景：先从图像中学习潜在状态，再通过 RSSM 建模潜在动态，最后利用 MPC + CEM 在潜在空间中规划动作
那么，PlaNet 是如何把图像转换成适合预测和规划的潜在状态的？这就要从它的核心模型 RSSM 开始讲起

2 PlaNet

2-1 介绍

PlaNet（Learning Latent Dynamics for Planning from Pixels）是 Hafner 等人于 2019 年提出的一种基于潜在动态模型的模型式强化学习方法
它的核心思想是：从图像观测中学习紧凑的潜在状态表示，在潜在空间中预测环境未来的动态，并通过模型预测控制（MPC）搜索动作序列，从而完成决策
论文进一步提出了两项关键设计：
1. RSSM（Recurrent State-Space Model，循环状态空间模型），通过确定性隐状态与随机潜在状态相结合，建模环境的长期动态与不确定性；
2. Latent Overshooting（潜在空间多步预测约束），利用多步潜在状态预测来改善模型的长期预测能力。
实验表明，PlaNet 能够仅凭像素观测解决包含接触动力学、部分可观测性和稀疏奖励的连续控制任务，并以更少的环境交互达到具有竞争力的控制性能

2-2 结构

2-2-1 总览

#mermaid-svg-6ZYOwQaA0t54m3VF{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-6ZYOwQaA0t54m3VF .error-icon{fill:#552222;}#mermaid-svg-6ZYOwQaA0t54m3VF .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-6ZYOwQaA0t54m3VF .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-6ZYOwQaA0t54m3VF .marker{fill:#333333;stroke:#333333;}#mermaid-svg-6ZYOwQaA0t54m3VF .marker.cross{stroke:#333333;}#mermaid-svg-6ZYOwQaA0t54m3VF svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-6ZYOwQaA0t54m3VF p{margin:0;}#mermaid-svg-6ZYOwQaA0t54m3VF .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF .cluster-label text{fill:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF .cluster-label span{color:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF .cluster-label span p{background-color:transparent;}#mermaid-svg-6ZYOwQaA0t54m3VF .label text,#mermaid-svg-6ZYOwQaA0t54m3VF span{fill:#333;color:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF .node rect,#mermaid-svg-6ZYOwQaA0t54m3VF .node circle,#mermaid-svg-6ZYOwQaA0t54m3VF .node ellipse,#mermaid-svg-6ZYOwQaA0t54m3VF .node polygon,#mermaid-svg-6ZYOwQaA0t54m3VF .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-6ZYOwQaA0t54m3VF .rough-node .label text,#mermaid-svg-6ZYOwQaA0t54m3VF .node .label text,#mermaid-svg-6ZYOwQaA0t54m3VF .image-shape .label,#mermaid-svg-6ZYOwQaA0t54m3VF .icon-shape .label{text-anchor:middle;}#mermaid-svg-6ZYOwQaA0t54m3VF .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-6ZYOwQaA0t54m3VF .rough-node .label,#mermaid-svg-6ZYOwQaA0t54m3VF .node .label,#mermaid-svg-6ZYOwQaA0t54m3VF .image-shape .label,#mermaid-svg-6ZYOwQaA0t54m3VF .icon-shape .label{text-align:center;}#mermaid-svg-6ZYOwQaA0t54m3VF .node.clickable{cursor:pointer;}#mermaid-svg-6ZYOwQaA0t54m3VF .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-6ZYOwQaA0t54m3VF .arrowheadPath{fill:#333333;}#mermaid-svg-6ZYOwQaA0t54m3VF .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-6ZYOwQaA0t54m3VF .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-6ZYOwQaA0t54m3VF .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-6ZYOwQaA0t54m3VF .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-6ZYOwQaA0t54m3VF .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-6ZYOwQaA0t54m3VF .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-6ZYOwQaA0t54m3VF .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-6ZYOwQaA0t54m3VF .cluster text{fill:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF .cluster span{color:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-6ZYOwQaA0t54m3VF .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-6ZYOwQaA0t54m3VF rect.text{fill:none;stroke-width:0;}#mermaid-svg-6ZYOwQaA0t54m3VF .icon-shape,#mermaid-svg-6ZYOwQaA0t54m3VF .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-6ZYOwQaA0t54m3VF .icon-shape p,#mermaid-svg-6ZYOwQaA0t54m3VF .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-6ZYOwQaA0t54m3VF .icon-shape .label rect,#mermaid-svg-6ZYOwQaA0t54m3VF .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-6ZYOwQaA0t54m3VF .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-6ZYOwQaA0t54m3VF .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-6ZYOwQaA0t54m3VF :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-6ZYOwQaA0t54m3VF .input>*{fill:#DCEBFF!important;stroke:#3973AC!important;color:#17365D!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .input span{fill:#DCEBFF!important;stroke:#3973AC!important;color:#17365D!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .input tspan{fill:#17365D!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .encoder>*{fill:#D9F2E6!important;stroke:#27966A!important;color:#124D37!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .encoder span{fill:#D9F2E6!important;stroke:#27966A!important;color:#124D37!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .encoder tspan{fill:#124D37!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .rssm>*{fill:#FFF0CC!important;stroke:#D99A16!important;color:#704900!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .rssm span{fill:#FFF0CC!important;stroke:#D99A16!important;color:#704900!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .rssm tspan{fill:#704900!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .latent>*{fill:#E9DFFF!important;stroke:#8756B5!important;color:#43256A!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .latent span{fill:#E9DFFF!important;stroke:#8756B5!important;color:#43256A!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .latent tspan{fill:#43256A!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .output>*{fill:#FFE0E0!important;stroke:#C94F55!important;color:#75272B!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .output span{fill:#FFE0E0!important;stroke:#C94F55!important;color:#75272B!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .output tspan{fill:#75272B!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .planner>*{fill:#D7F0F0!important;stroke:#258A8A!important;color:#124F4F!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .planner span{fill:#D7F0F0!important;stroke:#258A8A!important;color:#124F4F!important;stroke-width:2px!important;}#mermaid-svg-6ZYOwQaA0t54m3VF .planner tspan{fill:#124F4F!important;} MPC：潜在空间规划与动作执行
世界模型：状态推断与学习
KL 散度约束
提供未来奖励预测
控制信号
当前观测

图像 oₜ
CNN 视觉编码器

eₜ
后验分布

q(zₜ | hₜ, eₜ)
上一时刻动作 aₜ₋₁
GRU 状态更新

hₜ = f(hₜ₋₁, zₜ₋₁, aₜ₋₁)
上一时刻随机状态 zₜ₋₁
上一时刻确定性状态 hₜ₋₁
先验分布

p(zₜ | hₜ)
后验采样得到 zₜ
潜在状态

(hₜ, zₜ)
奖励预测器

预测 rₜ
图像解码器

重建观测 ôₜ
当前潜在状态

(hₜ, zₜ)
CEM 采样与优化

候选动作序列
RSSM 潜在动态展开

预测未来状态与奖励
累计预测奖励

评估候选序列
选出当前较优动作序列
只执行第一个动作

aₜ
真实环境执行动作
获得新图像 oₜ₊₁
CNN + RSSM

更新当前潜在状态

PlaNet 的整体架构可以分为两个相互配合的部分：世界模型（World Model）负责从图像中学习环境的潜在动态，模型预测控制（MPC）则利用学习到的动态模型，在潜在空间中预测未来并规划动作
PlaNet 以高维图像作为输入，首先通过 CNN 将图像编码为特征向量，再利用循环状态空间模型（RSSM）维护环境的潜在状态。在此基础上，模型可以预测未来状态及其对应的奖励，并通过 MPC 搜索较优的动作序列

2-2-2 整体流程

整体流程如下：
- 首先，CNN 将图像观测 o t o_t ot 编码为特征 e t e_t et，RSSM 结合历史状态与动作，得到由确定性状态 h t h_t ht 和随机状态 z t z_t zt 构成的潜在状态
- 世界模型通过先验预测、后验推断学习环境动态，并利用图像重建与奖励预测进行训练；论文另外提出的 Latent Overshooting 约束多步预测，但最终智能体没有采用
- 完成训练后，MPC 利用 CEM 在潜在空间中搜索动作序列，通过 RSSM 预测未来状态、奖励头评估累计回报，只执行较优序列的第一个动作，再根据新观测更新状态并重新规划
后面的小节按这条数据流展开：先说明 RSSM 如何得到 ( h t , z t ) (h_t,z_t) (ht,zt)，再说明 Latent Overshooting 如何改善多步预测，最后说明如何用这套模型在潜在空间中做规划

2-2-3 三个要点

PlaNet 的核心可以围绕三个方面展开：
- Recurrent State-Space Model（循环状态空间模型，RSSM）：结合确定性隐状态与随机潜在状态，利用循环结构积累历史信息，并通过先验预测与后验推断建模环境动态，使模型能够处理部分可观测性和环境中的随机性
- Latent Overshooting（潜在空间多步预测约束）：训练时不只约束相邻时刻的状态转移，还进一步约束模型对多个未来时刻的潜在状态预测，使多步预测分布与真实观测条件下的后验分布保持一致，从而改善世界模型的长期动态预测能力
- Planning in Latent Spaces（潜在空间规划）：将高维图像观测编码为紧凑的潜在状态，在潜在空间中预测未来状态与奖励，并通过模型预测控制（MPC）搜索较优的动作序列，避免规划过程中反复生成完整图像

需要说明的是，PlaNet 各个模块之间耦合度较高，单独拆分讲解容易造成内容重复。因此，本文将从 RSSM 入手，围绕潜在空间规划、RSSM 和 Latent Overshooting 三个核心要点展开分析。

当然，这样的组织思路也与论文作者的讲解脉络相呼应。

2-3 Recurrent State-Space Model（RSSM，循环状态空间模型）

2-3-1 RSSM 模型含义

RSSM（Recurrent State-Space Model，循环状态空间模型）是 PlaNet 用于学习环境潜在动态的核心模型。它将循环神经网络与随机状态空间模型结合起来，通过历史状态、动作和当前观测，学习环境状态随时间变化的规律。
RSSM 同时维护两种状态：
- 确定性状态（Deterministic State）： 由 GRU 递归更新，结合历史状态与动作，持续积累过去的信息。
- 随机状态（Stochastic State）： 通过概率分布采样得到，用于描述当前潜在状态中的随机性与不确定性。
两种状态共同构成模型的潜在状态表示： s t = ( h t , z t ) s_t=(h_t,z_t) st=(ht,zt) 其中：
- h t h_t ht ：确定性状态
- z t z_t zt：随机状态。
在每个时间步，RSSM 首先根据上一时刻的状态和动作更新确定性状态，再通过先验分布预测随机状态；当获得当前观测后，还可以结合观测特征，通过后验分布进一步推断随机状态。

说人话：RSSM 用 GRU 记住过去发生了什么，再用随机变量描述当前潜在状态。预测未来时依靠先验分布，获得真实观测后则通过后验分布修正状态。

2-3-2 先验与后验

RSSM 的一个重要设计是同时引入先验分布（Prior）与后验分布（Posterior），分别用于预测和推断随机状态 z t z_t zt。
以防你忘记概率论基础：
- 先验概率（Prior）： 在获得新的观测数据之前，根据已有信息对某个变量作出的概率判断。
- 后验概率（Posterior）： 在获得新的观测数据之后，结合已有信息与新观测，对该变量重新进行概率判断。

放到 RSSM 中，两者分别对应：
- 先验分布：根据历史信息 h t h_t ht 预测当前随机状态 z t z_t zt，不依赖当前观测。
- 后验分布：在历史信息 h t h_t ht 的基础上，进一步结合当前观测特征 e t e_t et，推断当前随机状态 z t z_t zt。

也就是说，区别于World Models 的 LSTM + MDN 更像是一个基于历史信息预测下一步潜变量的时序模型；PlaNet 的 GRU + Prior/Posterior 则把状态估计和未来预测明确区分开了。

2-3-2-1 先验：预测当前状态

首先，模型根据上一时刻的确定性状态、随机状态以及动作，更新当前确定性状态 h t h_t ht，随后通过先验分布预测当前随机状态： p θ ( z t ∣ h t ) = N ( μ θ ( h t ) , σ θ ( h t ) ) p_\theta(z_t\mid h_t) =\mathcal{N}\left(\mu_\theta(h_t),\sigma_\theta(h_t)\right) pθ(zt∣ht)=N(μθ(ht),σθ(ht))先验分布只依赖确定性状态 h t h_t ht，不需要当前时刻的真实观测，因此可以用于未来状态预测。具体的网络头与采样在 2-3-5 节展开。

2-3-2-2 后验：结合观测推断当前状态

但仅依靠历史信息进行预测，可能无法准确反映当前环境的实际情况。因此，RSSM 进一步引入当前观测的编码特征 e t e_t et，构建后验分布： q ϕ ( z t ∣ h t , e t ) = N ( μ ϕ ( h t , e t ) , σ ϕ ( h t , e t ) ) q_\phi(z_t\mid h_t,e_t) =\mathcal{N}\left(\mu_\phi(h_t,e_t),\sigma_\phi(h_t,e_t)\right) qϕ(zt∣ht,et)=N(μϕ(ht,et),σϕ(ht,et)) 后验分布同时利用历史信息和当前观测，对随机状态进行推断，可以理解为对先验预测的观测校正。拼接与网络头分别在 2-3-6-2 与 2-3-7 节展开。

2-3-2-3 约束

训练时，模型通过 KL 散度约束先验分布与后验分布，使先验能够逐渐学习到后验所提供的状态信息：
L K L = D K L ( q ϕ ( z t ∣ h t , e t ) ∥ p θ ( z t ∣ h t ) ) \mathcal{L}{\mathrm{KL}}=D{KL}\left(q_\phi(z_t\mid h_t,e_t)\,\|\,p_\theta(z_t\mid h_t)\right) LKL=DKL(qϕ(zt∣ht,et)∥pθ(zt∣ht))

python 复制代码

# 对应 2-3-2-3：先验与后验之间的 KL，方向为 KL(q || p)，与公式一致
kl = kl_divergence(Normal(posterior_means, posterior_std_devs),
                   Normal(prior_means, prior_std_devs)).sum(dim=2)
# free nats 以内的偏差不再产生梯度，避免先验被强行拉到后验的细节上
kl_loss = torch.max(kl, free_nats).mean(dim=(0, 1))

这种设计使 RSSM 同时具备基于历史信息预测未来和利用真实观测推断当前状态的能力。训练时可以利用后验分布获得更准确的状态表示，而在规划阶段，由于未来观测未知，模型只能依靠先验分布向前预测。

说人话：先验相当于根据过去推测现在可能是什么样，后验则是看到当前图像后重新判断现在是什么样。训练时让两者尽可能一致，这样模型以后即使看不到未来图像，也能依靠先验预测未来状态。

2-3-3 RSSM 结构

有了上述关于潜在状态 s t = ( h t , z t ) s_t=(h_t,z_t) st=(ht,zt) 以及先验与后验分布的介绍，我们再通过整体结构图梳理 RSSM 的工作流程：

#mermaid-svg-xRBMg9CmWWMBZNnC{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-xRBMg9CmWWMBZNnC .error-icon{fill:#552222;}#mermaid-svg-xRBMg9CmWWMBZNnC .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-xRBMg9CmWWMBZNnC .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-xRBMg9CmWWMBZNnC .marker{fill:#333333;stroke:#333333;}#mermaid-svg-xRBMg9CmWWMBZNnC .marker.cross{stroke:#333333;}#mermaid-svg-xRBMg9CmWWMBZNnC svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-xRBMg9CmWWMBZNnC p{margin:0;}#mermaid-svg-xRBMg9CmWWMBZNnC .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC .cluster-label text{fill:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC .cluster-label span{color:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC .cluster-label span p{background-color:transparent;}#mermaid-svg-xRBMg9CmWWMBZNnC .label text,#mermaid-svg-xRBMg9CmWWMBZNnC span{fill:#333;color:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC .node rect,#mermaid-svg-xRBMg9CmWWMBZNnC .node circle,#mermaid-svg-xRBMg9CmWWMBZNnC .node ellipse,#mermaid-svg-xRBMg9CmWWMBZNnC .node polygon,#mermaid-svg-xRBMg9CmWWMBZNnC .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-xRBMg9CmWWMBZNnC .rough-node .label text,#mermaid-svg-xRBMg9CmWWMBZNnC .node .label text,#mermaid-svg-xRBMg9CmWWMBZNnC .image-shape .label,#mermaid-svg-xRBMg9CmWWMBZNnC .icon-shape .label{text-anchor:middle;}#mermaid-svg-xRBMg9CmWWMBZNnC .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-xRBMg9CmWWMBZNnC .rough-node .label,#mermaid-svg-xRBMg9CmWWMBZNnC .node .label,#mermaid-svg-xRBMg9CmWWMBZNnC .image-shape .label,#mermaid-svg-xRBMg9CmWWMBZNnC .icon-shape .label{text-align:center;}#mermaid-svg-xRBMg9CmWWMBZNnC .node.clickable{cursor:pointer;}#mermaid-svg-xRBMg9CmWWMBZNnC .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-xRBMg9CmWWMBZNnC .arrowheadPath{fill:#333333;}#mermaid-svg-xRBMg9CmWWMBZNnC .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-xRBMg9CmWWMBZNnC .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-xRBMg9CmWWMBZNnC .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xRBMg9CmWWMBZNnC .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-xRBMg9CmWWMBZNnC .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xRBMg9CmWWMBZNnC .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-xRBMg9CmWWMBZNnC .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-xRBMg9CmWWMBZNnC .cluster text{fill:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC .cluster span{color:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-xRBMg9CmWWMBZNnC .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-xRBMg9CmWWMBZNnC rect.text{fill:none;stroke-width:0;}#mermaid-svg-xRBMg9CmWWMBZNnC .icon-shape,#mermaid-svg-xRBMg9CmWWMBZNnC .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xRBMg9CmWWMBZNnC .icon-shape p,#mermaid-svg-xRBMg9CmWWMBZNnC .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-xRBMg9CmWWMBZNnC .icon-shape .label rect,#mermaid-svg-xRBMg9CmWWMBZNnC .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xRBMg9CmWWMBZNnC .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-xRBMg9CmWWMBZNnC .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-xRBMg9CmWWMBZNnC :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-xRBMg9CmWWMBZNnC .input>*{fill:#DCEBFF!important;stroke:#3973AC!important;color:#17365D!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .input span{fill:#DCEBFF!important;stroke:#3973AC!important;color:#17365D!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .input tspan{fill:#17365D!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .encoder>*{fill:#D9F2E6!important;stroke:#27966A!important;color:#124D37!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .encoder span{fill:#D9F2E6!important;stroke:#27966A!important;color:#124D37!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .encoder tspan{fill:#124D37!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .core>*{fill:#FFF0CC!important;stroke:#D99A16!important;color:#704900!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .core span{fill:#FFF0CC!important;stroke:#D99A16!important;color:#704900!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .core tspan{fill:#704900!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .latent>*{fill:#E9DFFF!important;stroke:#8756B5!important;color:#43256A!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .latent span{fill:#E9DFFF!important;stroke:#8756B5!important;color:#43256A!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .latent tspan{fill:#43256A!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .output>*{fill:#FFE0E0!important;stroke:#C94F55!important;color:#75272B!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .output span{fill:#FFE0E0!important;stroke:#C94F55!important;color:#75272B!important;stroke-width:2px!important;}#mermaid-svg-xRBMg9CmWWMBZNnC .output tspan{fill:#75272B!important;} RSSM：循环状态空间模型
KL 散度约束
当前观测图像

oₜ
CNN 视觉编码器

eₜ = Encoder(oₜ)
上一时刻动作

aₜ₋₁
确定性状态更新

hₜ = f(hₜ₋₁, zₜ₋₁, aₜ₋₁)
上一时刻确定性状态

hₜ₋₁
上一时刻随机状态

zₜ₋₁
先验分布

p(zₜ | hₜ)
后验分布

q(zₜ | hₜ, eₜ)
采样随机状态

zₜ
潜在状态

(hₜ, zₜ)
奖励预测器

预测奖励
图像解码器

重建观测

从整体结构来看，可以分为三个部分：
- 确定性状态更新： 上一时刻的确定性状态 h t − 1 h_{t-1} ht−1、随机状态 z t − 1 z_{t-1} zt−1 与动作 a t − 1 a_{t-1} at−1 输入 GRU，递归得到当前确定性状态 h t h_t ht。
- 随机状态推断： RSSM 根据 h t h_t ht 得到先验分布 p ( z t ∣ h t ) p(z_t\mid h_t) p(zt∣ht)；同时，CNN 将当前观测编码为特征 e t e_t et，与 h t h_t ht 一起得到后验分布 q ( z t ∣ h t , e t ) q(z_t\mid h_t,e_t) q(zt∣ht,et)。训练时通过 KL 散度约束两个分布的一致性，并从后验分布中采样得到 z t z_t zt。
- 预测与重建： 将确定性状态与随机状态组合成潜在状态 ( h t , z t ) (h_t,z_t) (ht,zt)，分别输入奖励预测器和图像解码器，用于预测奖励和重建当前观测。
需要注意的是，训练阶段可以利用当前观测，通过后验分布推断随机状态；而在规划未来时，由于无法提前获得未来观测，只能依靠先验分布进行预测。

2-3-4 GRU

先验要建立在 h t h_t ht 之上，所以这一节先把 h t h_t ht 讲清楚：在获得当前观测之前，我们只能根据上一时刻留下的状态与动作估计当前环境处于什么状态，这份估计就是确定性状态
确定性状态由上一时刻的随机状态、动作与自身递推得到：
h t = f θ ( h t − 1 , z t − 1 , a t − 1 ) h_t=f_\theta\left(h_{t-1},z_{t-1},a_{t-1}\right) ht=fθ(ht−1,zt−1,at−1)
式中的 h t − 1 h_{t-1} ht−1 不是被喂进去的普通输入，它就是循环网络的隐状态，每一步更新完再传给下一步，因此 h t h_t ht 携带的是从回合开始到现在累积的整段历史
而确定性状态在RSSM中由GRU实现
关于GRU这里不过多介绍，感兴趣的朋友可以去看：
- 【深入理解 LSTM】从时间反向传播（BPTT）到长时记忆：门控机制与实验验证
- 【RNN GRU LSTM】从 Vanilla RNN 到 LSTM：公式推导、手写 cell 与三组不训练实验
代码中 f θ f_\theta fθ 由 GRUCell 实现，先拼接随机状态与动作，再递推隐状态：

python 复制代码

# 对应确定性状态更新：拼接上一时刻的随机状态与动作，再递推 belief
hidden = self.act_fn(self.fc_embed_state_action(torch.cat([_state, actions[t]], dim=1)))
beliefs[t + 1] = self.rnn(hidden, beliefs[t])

fc_embed_state_action 把 30 维随机状态与 6 维动作拼成 36 维后映射到 200 维，GRUCell 在此基础上维护跨时间步的隐状态
确定性通路的必要性来自记忆长度：纯随机状态模型每一步都重新采样，难以在多个时间步上保留同一份信息

说人话：GRU 负责"记住"，随机状态负责"表达不确定"

2-3-5 先验分布构建

h t h_t ht 是确定性向量，不能直接当随机状态用：同一段历史之后可能有多种走向，应当用一个分布把这种可能性表示出来
先验分支是一个两层 MLP：第一层把 h t h_t ht 映射到 200 维隐藏单元，第二层输出 60 维参数，对半切开后，一半作为均值，另一半经 softplus 得到标准差
两层线性变换的定义如下：

python 复制代码

# 先验头：只读取确定性状态，输出 2 * state_size 个参数
self.fc_embed_belief_prior = nn.Linear(belief_size, hidden_size)   # 输入 200、输出 200
self.fc_state_prior = nn.Linear(hidden_size, 2 * state_size)       # 输入 200、输出 60

数据在各层之间的形状变化如下：

层	输入	输出	说明
fc_embed_belief_prior 加激活	200	200	把 h t h_t ht 映射到隐藏单元
fc_state_prior	200	60	输出均值与标准差的原始值
chunk	60	30 与 30	切分为 μ p \mu_p μp 与 σ p \sigma_p σp

写成公式，先验是一个均值与标准差都由 h t h_t ht 决定的对角高斯：
μ p , σ ~ p = s p l i t ( W 2 R e L U ( W 1 h t + b 1 ) + b 2 ) , σ p = s o f t p l u s ( σ ~ p ) + 0.1 \mu_p,\tilde\sigma_p=\mathrm{split}\big(W_2\,\mathrm{ReLU}(W_1h_t+b_1)+b_2\big),\qquad \sigma_p=\mathrm{softplus}(\tilde\sigma_p)+0.1 μp,σ~p=split(W2ReLU(W1ht+b1)+b2),σp=softplus(σ~p)+0.1
p θ ( z t ∣ h t ) = N ( z t ; μ p , d i a g ( σ p 2 ) ) p_\theta(z_t\mid h_t)=\mathcal{N}\big(z_t;\mu_p,\mathrm{diag}(\sigma_p^2)\big) pθ(zt∣ht)=N(zt;μp,diag(σp2))
对应的分支代码只有四行：

python 复制代码

# 先验分支：只使用确定性状态
hidden = self.act_fn(self.fc_embed_belief_prior(beliefs[t + 1]))
prior_means[t + 1], _prior_std_dev = torch.chunk(self.fc_state_prior(hidden), 2, dim=1)
prior_std_devs[t + 1] = F.softplus(_prior_std_dev) + self.min_std_dev

标准差那一侧多做一步 σ = s o f t p l u s ( ⋅ ) + 0.1 \sigma=\mathrm{softplus}(\cdot)+0.1 σ=softplus(⋅)+0.1：softplus 保证取值为正，下限 0.1 则防止方差塌到零
- 下限的必要性在 KL 里：摊开的高斯散度含有 ( μ q − μ p ) 2 / 2 σ p 2 (\mu_q-\mu_p)^2/2\sigma_p^2 (μq−μp)2/2σp2 这样的项， σ p \sigma_p σp 趋近 0 时分母把这一项推向无穷，训练早期就被打断
- 标准差没有上界，先验因此可以表达"完全不知道"，这给探索留出了空间
采样用重参数化，把随机性单独放进一个标准正态噪声，梯度仍然能穿过均值与标准差回传：
z t = μ θ ( h t ) + σ θ ( h t ) ⊙ ε , ε ∼ N ( 0 , I ) z_t=\mu_\theta\left(h_t\right)+\sigma_\theta\left(h_t\right)\odot\varepsilon,\qquad \varepsilon\sim\mathcal{N}(0,I) zt=μθ(ht)+σθ(ht)⊙ε,ε∼N(0,I)
对应的实现只有一行：

python 复制代码

prior_states[t + 1] = prior_means[t + 1] + prior_std_devs[t + 1] * torch.randn_like(prior_means[t + 1])

说人话： h t h_t ht 说的是"历史走到哪了"，先验负责把它翻译成"下一步可能是哪几种状态、各有几分可能"

2-3-6 CNN

2-3-6-1 结构

PlaNet 没有像 World Models（2018）那样先单独预训练 VAE，而是把 CNN 编码器、RSSM 与解码器放进同一个序列变分下界里联合训练
- World Models 分两个阶段：先用重建损失单独训练 VAE 得到每帧的潜变量，再冻结潜变量去训练时序模型
- 重建项与 KL 项同时作用于卷积层与 RSSM，没有单独的表示学习阶段
CNN 编码器仍然逐帧提取当前图像特征 e t e_t et，随机状态则由后验结合历史确定性状态 h t h_t ht 与 e t e_t et 共同推断，构成滤波后验 q ϕ ( z t ∣ h t , e t ) q_\phi(z_t\mid h_t,e_t) qϕ(zt∣ht,et)
PlaNet 论文没有专门论证这一取舍，附录 A 只说明卷积与反卷积结构沿用 Ha & Schmidhuber（2018）；两阶段与联合训练的差别由作者在后续的 Dreamer（2020）中明确指出：World Models 以两阶段方式学习潜在动力学，PlaNet 则联合学习
- 这一差别落在潜变量的用途上：单独预训练时 VAE 只对重建负责，潜变量只需还原画面；联合训练时重建与 KL 由同一个变分下界权衡，编码器必须给出既能还原观测、又能被先验预测的状态，规划用到的正是后者

说人话：图像重建得好，不代表潜变量就适合预测未来。PlaNet 在训练编码器时，不仅要求潜变量能够还原当前画面，还要求 RSSM 能够利用历史信息预测这些潜变量。这样学到的表示才真正服务于后续的动态建模与规划。

这点可以在我们的另一篇文章中得到启发：【POMDP】从单帧图像到时序感知：CNN、VAE、GRU 与 PPO 的视觉控制实验
编码器是卷积神经网络（CNN，Convolutional Neural Network），把 64×64×3 的图像压缩成 1024 维嵌入：
e t = E n c ϕ ( o t ) e_t=\mathrm{Enc}_\phi\left(o_t\right) et=Encϕ(ot)
结构为四层步长为 2 的卷积，通道数依次为 32、64、128、256，展平后得到 1024 维向量

python 复制代码

self.conv1 = nn.Conv2d(3, 32, 4, stride=2)
self.conv2 = nn.Conv2d(32, 64, 4, stride=2)
self.conv3 = nn.Conv2d(64, 128, 4, stride=2)
self.conv4 = nn.Conv2d(128, 256, 4, stride=2)
hidden = hidden.view(-1, 1024)

嵌入 e t e_t et 只有一个去处：作为后验分布的附加输入（2-3-2 节）
解码器的结构与训练用途放在 2-3-8 节单独说明
以本文的实验为例，输入为 walker 在 t = 10 s 的 64×64 观测，训练好的编码器前三层输出的特征图如下：
浅层特征图保留了躯干与双腿的位置与边缘，通道之间的差异主要体现在关注的区域；越往深层，空间分辨率越低，单个单元覆盖的画面范围越大，编码的内容逐渐从局部纹理变成整体姿态

2-3-6-2 后验拼接

编码器的输出 e t e_t et 只有一个去处：和确定性状态 h t h_t ht 拼在一起，作为后验网络的输入（完整的后验分支在 2-3-7 节）
两者维度并不相同， h t h_t ht 为 200 维， e t e_t et 为 1024 维，拼接后是 1224 维：
z t ∼ q ϕ ( z t ∣ h t , e t ) , $h t ; e t$ ∈ R 1224 z_t\sim q_\phi\left(z_t\mid h_t,e_t\right),\qquad \left $h_t;e_t\\right$ \in\mathbb{R}^{1224} zt∼qϕ(zt∣ht,et), $ht;et$ ∈R1224
拼接本身只有一行，沿特征维把两支向量接起来：

python 复制代码

# belief 200 维 + 观测嵌入 1024 维，沿特征维拼接
hidden = self.act_fn(self.fc_embed_belief_posterior(torch.cat([beliefs[t + 1], observations[t]], dim=1)))

融合方式选拼接而不是相加： h t h_t ht 记的是历史， e t e_t et 只描述当前帧，两者的维度与含义都不同，拼接把信息原样交给全连接层，由网络自己学权重
观测必须先穿过采样步骤进入 z t z_t zt，再由 z t z_t zt 参与重建；否则编码器可以把 e t e_t et 直接交给解码器，形成一条绕过随机状态的确定性捷径

2-3-7 后验估计

后验分支与先验分支的结构完全对称，区别只在输入：后验在 h t h_t ht 之外再拼接当前观测的嵌入 e t e_t et，第一层的输入因此从 200 维变成 1224 维
两层线性变换的定义如下：

python 复制代码

# 后验头：输入为确定性状态与观测嵌入的拼接
self.fc_embed_belief_posterior = nn.Linear(belief_size + embedding_size, hidden_size)  # 输入 1224、输出 200
self.fc_state_posterior = nn.Linear(hidden_size, 2 * state_size)                      # 输入 200、输出 60

拼接之后，各层与切分对应的形状变化如下：

层	输入	输出	说明
fc_embed_belief_posterior 加激活	1224	200	压回隐藏单元
fc_state_posterior	200	60	输出均值与标准差的原始值
chunk	60	30 与 30	切分为 μ q \mu_q μq 与 σ q \sigma_q σq

写成公式，后验同样是对角高斯，只是均值与标准差改由 h t h_t ht 与 e t e_t et 共同决定：
μ q , σ ~ q = s p l i t ( W 2 R e L U ( W 1 $h t ; e t$ + b 1 ) + b 2 ) , σ q = s o f t p l u s ( σ ~ q ) + 0.1 \mu_q,\tilde\sigma_q=\mathrm{split}\big(W_2\,\mathrm{ReLU}(W_1 $h_t;e_t$ +b_1)+b_2\big),\qquad \sigma_q=\mathrm{softplus}(\tilde\sigma_q)+0.1 μq,σ~q=split(W2ReLU(W1 $ht;et$ +b1)+b2),σq=softplus(σ~q)+0.1
q ϕ ( z t ∣ h t , e t ) = N ( z t ; μ q , d i a g ( σ q 2 ) ) q_\phi(z_t\mid h_t,e_t)=\mathcal{N}\big(z_t;\mu_q,\mathrm{diag}(\sigma_q^2)\big) qϕ(zt∣ht,et)=N(zt;μq,diag(σq2))
完整的后验分支如下：

python 复制代码

# 后验分支：额外拼接当前观测的嵌入
hidden = self.act_fn(self.fc_embed_belief_posterior(torch.cat([beliefs[t + 1], observations[t]], dim=1)))
posterior_means[t + 1], _posterior_std = torch.chunk(self.fc_state_posterior(hidden), 2, dim=1)
posterior_std_devs[t + 1] = F.softplus(_posterior_std) + self.min_std_dev
posterior_states[t + 1] = posterior_means[t + 1] + posterior_std_devs[t + 1] * torch.randn_like(posterior_means[t + 1])

两条分支除了输入不同，其余完全一致：同样的 200 维隐藏单元、同样的 60 维输出、同样的 softplus 下限与重参数化采样
两条分支在数值上没有区别，区别只在信息：先验看不到当前帧，后验看得到，所以后验更贴近当前的真实状态
后验要求观测已经到达。部署时当前这一步有真实的 o t o_t ot，仍然可以用后验校正状态；但规划阶段沿候选动作向前滚动的那 12 步没有观测，只能退回先验
后验给出的是滤波估计而不是平滑：它只用到 t t t 时刻为止的观测与动作，不会回头修改 t − 1 t-1 t−1 的状态

说人话：先验是"闭着眼睛猜现在"，后验是"睁开眼看一眼再修正"；滚动到未来时眼睛闭着，只能靠先验

2-3-8 图像解码

那么同样的，有编码就要有解码。在图像通过CNN输入RSSM后，我们可以得到
解码器负责将潜在状态重新映射到像素空间，通过重建误差为潜在表示提供监督信号：
o ^ t = D e c θ ( h t , z t ) \hat{o}t=\mathrm{Dec}\theta\left(h_t,z_t\right) o^t=Decθ(ht,zt)
解码器与编码器近似镜像：先把 h t h_t ht 与 z t z_t zt 拼接后映射到嵌入维，再用四层转置卷积放大回 64×64×3

python 复制代码

hidden = self.fc1(torch.cat([belief, state], dim=1))
hidden = hidden.view(-1, self.embedding_size, 1, 1)
hidden = self.act_fn(self.conv1(hidden))
hidden = self.act_fn(self.conv2(hidden))
hidden = self.act_fn(self.conv3(hidden))
observation = self.conv4(hidden)

四层转置卷积的通道数依次为 128、64、32、3，卷积核大小依次为 5、5、6、6，步长与编码器对称
重建损失为像素均方误差，只在训练阶段提供监督；规划与部署都不需要生成图像，这是潜在空间规划计算量较小的直接原因
- 重建还决定了潜在状态里留下什么：只靠奖励头优化时，与回报无关的画面信息可以被随意丢弃；有了重建约束，潜在状态必须保留足够还原画面的环境结构，规划时才有可用的几何与接触信息

说人话：解码器只负责教模型"看清楚"，真正决策时用不到它

以本文的实验为例，输入为回放中四个时刻的观测，解码后的图像如下：
重建保住了躯干与双腿的姿态与空间位置，纹理与反光则明显模糊，说明解码器提供的是结构级的监督信号

2-3-9 奖励预测

前面介绍了图像解码器，它负责将潜在状态重建为图像，为世界模型提供观测重建监督。但在真正进行规划时，模型还需要知道预测出来的未来状态究竟有多好，这就需要奖励预测器（Reward Predictor）。

为什么 PlaNet 需要单独设计一个奖励预测器？主要与它的规划方式有关：
- PETS 等方法可以直接利用任务定义的奖励函数，对预测轨迹进行评价。
- 但 PlaNet 主要处理视觉输入，模型内部规划时使用的是潜在状态 ( h t , z t ) (h_t,z_t) (ht,zt)，而不是真实环境状态。
- 由于潜在状态是模型从图像中学习得到的隐式表示，真实环境中的奖励函数通常无法直接作用于这个潜在表示。因此，PlaNet 需要额外学习从潜在状态到奖励的映射： r ^ t = g θ ( h t , z t ) \hat r_t=g_\theta(h_t,z_t) r^t=gθ(ht,zt)
- 这样，在规划阶段，RSSM 可以利用先验分布向前预测未来潜在状态，再由奖励预测器评估每一步的预测奖励，最终计算候选动作序列的累计回报，为 MPC 提供评价依据。

说人话：PlaNet 在潜在空间里规划，但潜在状态并不是真实环境状态，原本的奖励函数没办法直接拿来使用。因此，它专门训练了一个奖励预测器，让模型能够根据自己预测出来的未来状态判断这条行动路线好不好。

潜在空间规划绕开了这个困难：奖励头从少数几次拿到奖励的片段里学出"哪种状态值钱"，规划时把 H H H 步的预测奖励直接相加来给候选序列排序，既不需要价值函数，也不需要把稀薄的成功信号反传到每一个动作上
- 这也解释了论文为什么强调模型必须能在多个时间步上准确预测奖励：CEM 比的是一整条序列的累计回报，只把下一步预测准没有用
奖励头以潜在状态为输入，输出标量奖励：
r ^ t = g θ ( h t , z t ) \hat{r}t=g\theta\left(h_t,z_t\right) r^t=gθ(ht,zt)
实现为两层隐藏层的全连接网络，隐藏维度为 200：

python 复制代码

hidden = self.act_fn(self.fc1(torch.cat([belief, state], dim=1)))
hidden = self.act_fn(self.fc2(hidden))
reward = self.fc3(hidden).squeeze(dim=1)

评估一条候选序列时，规划器先把当前潜在状态复制 J J J 份，再沿动作序列滚动，由奖励头逐时刻打分并求和：

python 复制代码

# 把当前潜在状态复制 J 份，所有候选在批次维上并行展开
# B 为并行环境数，self.candidates 即候选数 J
belief = belief.unsqueeze(dim=1).expand(B, self.candidates, belief.size(1)).reshape(-1, belief.size(1))
state = state.unsqueeze(dim=1).expand(B, self.candidates, Z).reshape(-1, Z)

# actions 的形状为 (H, B*J, A)
beliefs, states, _, _ = self.transition_model(state, actions, belief)
returns = self.reward_model(beliefs.view(-1, H), states.view(-1, Z)).view(H, -1).sum(dim=0)

beliefs 与 states 的形状分别为 ( H , B ⋅ J , 200 ) (H,B\cdot J,200) (H,B⋅J,200) 与 ( H , B ⋅ J , 30 ) (H,B\cdot J,30) (H,B⋅J,30)，view 把时间维并入批次后一次算完所有时刻的奖励
实现中不做折扣，直接累加预测奖励；评估过程不产生真实环境交互，也不更新参数，所有候选序列都从同一个当前潜在状态出发
- 候选动作序列的排序完全由奖励头决定，因此多步一致性比单步精度更关键
- 奖励损失与重建损失、KL 损失联合优化，三者共享编码器与 RSSM 的梯度

2-4 多步预测与Latent Overshooting（潜在超前预测）

2-4-1 向前滚动 k 步

到目前为止 RSSM 都在做滤波：每一步都有真实观测 e t e_t et 进后验，把状态往真实值上拉。规划要做的是另一件事，给定一条候选动作序列，在没有观测的情况下推出后果
做法是把后验那条回边断开，让模型自己喂自己：
h t + 1 = f θ ( h t , z t , a t ) , z t + 1 ∼ p θ ( z t + 1 ∣ h t + 1 ) , r ^ t + 1 = g θ ( h t + 1 , z t + 1 ) h_{t+1}=f_\theta\left(h_t,z_t,a_t\right),\qquad z_{t+1}\sim p_\theta\left(z_{t+1}\mid h_{t+1}\right),\qquad \hat{r}{t+1}=g\theta\left(h_{t+1},z_{t+1}\right) ht+1=fθ(ht,zt,at),zt+1∼pθ(zt+1∣ht+1),r^t+1=gθ(ht+1,zt+1)
这三行重复 k k k 次，就得到未来 k k k 步的潜在轨迹与逐步奖励；整个过程不需要解码成图像，也不需要与真实环境交互

说人话：训练时，RSSM 每走一步都能看到真实图像，用它来修正自己的状态；但规划时，未来的图像还没出现，模型只能根据当前状态和候选动作，一步步预测下一个状态和奖励，再把预测出的状态作为下一步的输入。重复这个过程，就能在潜在空间里提前推演未来 k k k 步，而不需要真的让机器人执行动作。

能把这条递推跑起来，靠的是 RSSM 结构上的三点：
- h h h 是确定性递推：每一步都用同一个函数 f θ f_\theta fθ，历史信息沿着隐状态一路带下去，不必每一步都从噪声里重新回忆
- z z z 每一步从先验重采样：同一条动作序列在不同采样下得到的结果略有差异，不确定性被显式传播，而不是压成一个点估计
- 全程没有校正：第 t + 1 t+1 t+1 步的先验只依赖 h t + 1 h_{t+1} ht+1，而 h t + 1 h_{t+1} ht+1 只依赖历史，一旦方向偏了，后面每一步都带着这个偏差
这也解释了规划时域为什么只取 12 步：越往后，推演的成分越多于证据

2-4-2 代码里的滚动开关

滤波与预测共用同一个 TransitionModel.forward，唯一的差别是 observations 传不传：

python 复制代码

# 有观测：每一帧都用来校正，得到后验状态（训练拟合与在线更新状态走这条）
beliefs, prior_states, prior_means, prior_std_devs, posterior_states, posterior_means, posterior_std_devs = transition_model(
    init_state, actions[:-1], init_belief, bottle(encoder, (observations[1:], )), nonterminals[:-1])
# 无观测：从 t 时刻的后验状态出发，只用动作向前滚 d 步，得到每一步的先验（规划走这条）
beliefs, prior_states, prior_means, prior_std_devs = transition_model(
    posterior_states[t], actions[t:d], beliefs[t], None, nonterminals[t:d])

循环体里那一行就是全部的分叉：

python 复制代码

_state = prior_states[t] if observations is None else posterior_states[t]

预测模式下后验分支整段跳过，每一步只剩三件事：把 z z z 与 a a a 拼起来映射到隐状态维度、递推 h h h、对 h h h 求先验并采样新的 z z z
- 回传的 z z z 是自己上一步采样的结果，因此这条路径是自由滚动： k k k 步里没有任何一次来自真实环境的修正
规划器复用的就是这段代码（2-5 节）：CEM 的候选序列不是另写的一套网络，而是一批并行的 observations=None 滚动，把 J J J 份候选放进批次维一次算完

2-4-3 Latent Overshooting

训练时序列变分下界里的 KL 只比较同一时刻的先验与后验，也就是只把单步预测训到位
论文在 Limited capacity （神经网络的容量有限）一段里把这层关系讲得更直接：如果模型能把单步预测做到完美，多步预测自然完美；但容量有限，只按单步目标训到收敛，得到的模型一般不是多步预测最好的那一个

说人话：单步预测准确，不代表连续向前预测很多步也准确。

Latent Overshooting 的思路是让训练时的滚动与规划时的滚动走同一条路：从 t t t 时刻的后验状态出发，喂入真实动作、不喂观测，滚 d d d 步，再把每一步的先验与该时刻真实观测给出的后验比较
L o v e r ( θ ) = 1 D ∑ d = 1 D β d E $K L ( q ϕ ( z t + d ∣ h t + d , e t + d ) ∥ p θ ( d ) ( z t + d ∣ h t , a t : t + d - 1 ) )$ \mathcal{L}{\mathrm{over}}(\theta)=\frac{1}{D}\sum{d=1}^{D}\beta_d\,\mathbb{E}\left $\\mathrm{KL}\\left(q_\\phi\\left(z_{t+d}\\mid h_{t+d},e_{t+d}\\right)\\,\\\|\\,p_\\theta\^{(d)}\\left(z_{t+d}\\mid h_t,a_{t:t+d-1}\\right)\\right)\\right$ Lover(θ)=D1d=1∑DβdE $KL(qϕ(zt+d∣ht+d,et+d)∥pθ(d)(zt+d∣ht,at:t+d-1))$
p θ ( d ) p_\theta^{(d)} pθ(d) 就是从 t t t 出发滚 d d d 步得到的先验， d = 1 d=1 d=1 时这一项退回标准变分下界里的 KL， d > 1 d>1 d>1 的部分才真正约束多步预测
实现上后验一侧停止梯度，让多步先验去追已经被观测校正过的后验，而不是让后验迁就模型自己的预测：

python 复制代码

# 后验统计量先 detach，再与滚动得到的先验比较
kl_loss += (1 / args.overshooting_distance) * args.overshooting_kl_beta * torch.max(
    (kl_divergence(Normal(post_mean.detach(), post_std.detach()), Normal(prior_mean, prior_std)) * seq_mask).sum(dim=2),
    free_nats).mean(dim=(0, 1)) * (args.chunk_size - 1)

说人话：训练时让 RSSM 从某个时刻出发，连续向前预测多步，并对这些多步预测结果计算 KL Loss。

结果值得记一笔：论文附录 D 的对照里，这一项对 DRNN 这类模型提升明显，对 RSSM 反而略有下降，最终采用的智能体没有用它，仓库实现也把 overshooting_kl_beta 默认设成 0
论文没有解释这个反差，一个自然的读法是 RSSM 的确定性通路本身已经承担了多步一致性，再叠加显式的多步约束，等于在同一个方向上重复施力

2-5 Planning in Latent Spaces（潜在空间规划）

2-5-1 CEM：如何搜索较优动作序列？

规划直接在 2-3 节的潜在状态 ( h t , z t ) (h_t,z_t) (ht,zt) 上展开：候选序列推演出的后果全部由 RSSM 沿先验给出，不需要解码成图像
与 PETS 一样，环境模型只负责"给定一条动作序列，沿先验滚动 H H H 步并累加预测奖励"，它不会直接告诉我们该执行哪个动作
- 动作空间是连续的，长度为 H H H 的序列无法穷举；PlaNet 沿用 PETS 的做法，用 CEM 在这个回报函数上迭代筛选候选，筛出的较优序列再交给滚动时域策略执行（2-5-2 节）
CEM（Cross-Entropy Method，交叉熵方法）：一种基于采样和迭代更新的优化方法，用于搜索高回报的动作序列
记号约定： B B B 为并行环境数， J J J 为每轮候选序列数， K K K 为精英数， H H H 为规划时域， I I I 为 CEM 迭代次数， Z Z Z 与 A A A 分别为随机状态维度与动作维度
候选序列的分数由潜在动力学模型与奖励预测器给出，即时域内预测奖励之和：
R ( a t : t + H − 1 ) = ∑ τ = t t + H − 1 r ^ τ , r ^ τ = g θ ( h τ + 1 , z τ + 1 ) R\left(a_{t:t+H-1}\right)=\sum_{\tau=t}^{t+H-1}\hat r_\tau,\qquad\hat r_\tau=g_\theta\left(h_{\tau+1},z_{\tau+1}\right) R(at:t+H−1)=τ=t∑t+H−1r^τ,r^τ=gθ(hτ+1,zτ+1)
CEM 维护一个关于动作序列的因子化高斯分布，初值为标准正态：
q ( a t : t + H − 1 ) = N ( μ , diag ⁡ ( σ 2 ) ) , μ 0 = 0 , σ 0 = 1 q\left(a_{t:t+H-1}\right)=\mathcal{N}\left(\mu,\operatorname{diag}\left(\sigma^{2}\right)\right),\qquad\mu_0=0,\quad\sigma_0=1 q(at:t+H−1)=N(μ,diag(σ2)),μ0=0,σ0=1
第 i i i 轮从当前分布采样 J J J 条候选，逐条按 R R R 打分，再保留分数最高的 K K K 条，把它们的编号组成精英集合：
E i = TopK ⁡ ( { R ( a ( j ) ) } j = 1 J ) , ∣ E i ∣ = K \mathcal{E}i=\operatorname{TopK}\left(\left\{R\left(a^{(j)}\right)\right\}{j=1}^{J}\right),\qquad\left|\mathcal{E}_i\right|=K Ei=TopK({R(a(j))}j=1J),∣Ei∣=K
分布更新取精英样本的均值与标准差：
μ i + 1 = 1 K ∑ j ∈ E i a ( j ) , σ i + 1 = 1 K ∑ j ∈ E i ( a ( j ) − μ i + 1 ) 2 \mu_{i+1}=\frac{1}{K}\sum_{j\in\mathcal{E}i}a^{(j)},\qquad\sigma{i+1}=\sqrt{\frac{1}{K}\sum_{j\in\mathcal{E}i}\left(a^{(j)}-\mu{i+1}\right)^{2}} μi+1=K1j∈Ei∑a(j),σi+1=K1j∈Ei∑(a(j)−μi+1)2

说人话：CEM 就是先随机生成一批动作序列，让世界模型逐一打分，再挑出表现最好的 K 条，根据它们的分布调整下一轮采样范围。如此反复迭代，让候选动作逐渐集中到高回报区域，最终找到较优的动作序列。

均值与标准差按元素计算，每个时刻的每个动作维各自维护一组参数，因此待优化的分布是时间与动作维上的独立高斯乘积
标准差用有偏估计（除以 K K K 而不是 K − 1 K-1 K−1），与代码里的 unbiased=False 一致；采样后先裁剪到动作区间，保证候选始终合法
对应的迭代过程如下：

python 复制代码

# 对应 2-5-1 的 CEM：采样、评估、选精英、更新分布
# H、Z、A 分别为规划时域、随机状态维度与动作维度
action_mean = torch.zeros(H, B, 1, A)
action_std_dev = torch.ones(H, B, 1, A)
for _ in range(self.optimisation_iters):
    actions = (action_mean + action_std_dev * torch.randn(H, B, self.candidates, A)).view(H, B * self.candidates, A)
    actions.clamp_(min=self.min_action, max=self.max_action)
    beliefs, states, _, _ = self.transition_model(state, actions, belief)
    returns = self.reward_model(beliefs.view(-1, H), states.view(-1, Z)).view(H, -1).sum(dim=0)
    _, topk = returns.reshape(B, self.candidates).topk(self.top_candidates, dim=1, largest=True, sorted=False)
    topk += self.candidates * torch.arange(0, B).unsqueeze(dim=1)  # 展平后按批次偏移
    best_actions = actions[:, topk.view(-1)].reshape(H, B, self.top_candidates, A)
    action_mean = best_actions.mean(dim=2, keepdim=True)
    action_std_dev = best_actions.std(dim=2, unbiased=False, keepdim=True)

每轮迭代都把采样范围收缩到上一轮的高回报区域，因此 CEM 是"用迭代换精度"：候选数不变，分布越来越窄
与随机射击相比，CEM 多了分布更新这一步，同样的候选预算下更容易收敛到较优序列附近；这一差异对应 2-9 节消融表里的"随机射击替代 CEM"一行

2-5-2 滚动时域规划：为什么只执行第一个动作？

经过 CEM 多轮采样与筛选后，规划器得到一条预测回报较高的动作序列。但 PlaNet 并不会一次性执行整条序列，而是采用滚动时域规划
- 滚动时域规划（Receding-Horizon Planning）：每次规划未来一段时间的动作序列，但只执行当前较优序列的第一个动作
经过 I I I 轮迭代后，规划器返回第一段的均值动作，其余动作全部丢弃：
a t = μ I , t a_t=\mu_{I,t} at=μI,t
环境执行该动作并返回新观测，后验结合观测更新潜在状态，下一控制步从新的状态重新规划：
h t + 1 = f θ ( h t , z t , a t ) , z t + 1 ∼ q ϕ ( z ∣ h t + 1 , e t + 1 ) , e t + 1 = E n c ϕ ( o t + 1 ) h_{t+1}=f_\theta\left(h_t,z_t,a_t\right),\qquad z_{t+1}\sim q_\phi\left(z\mid h_{t+1},e_{t+1}\right),\qquad e_{t+1}=\mathrm{Enc}\phi\left(o{t+1}\right) ht+1=fθ(ht,zt,at),zt+1∼qϕ(z∣ht+1,et+1),et+1=Encϕ(ot+1)
对应的闭环代码先更新后验，再调用规划器，最后只把第一段动作交给环境：

python 复制代码

# 对应 2-5-2：更新后验、重新规划、只执行第一个动作
belief, _, _, _, posterior_state, _, _ = transition_model(
    posterior_state, action.unsqueeze(dim=0), belief, encoder(observation).unsqueeze(dim=0))
belief, posterior_state = belief.squeeze(dim=0), posterior_state.squeeze(dim=0)
action = planner(belief, posterior_state)
action.clamp_(min=min_action, max=max_action)
observation, reward, done = env.step(action[0].cpu())

控制步与仿真步通过动作重复对应：一次规划出的动作在仿真中重复执行 2 步，之后重新感知与规划
- 重新规划把模型误差限制在单个控制步之内：时域再长，真正承担风险的也只有执行的那一步
采集阶段会在规划结果上叠加探索噪声，部署阶段不加；这一点与 CEM 本身无关，属于数据收集策略
每个控制步都要重新求解一次 CEM，计算开销集中在规划而不是策略前向，参数与耗时在 2-8 节展开

说人话：每次都只走一步就重新看路，长时域预测只用来比较候选，不用来执行

2-6 数据流串联

图像 oₜ
CNN 视觉编码器

eₜ
后验分布

q(zₜ | hₜ, eₜ)
上一时刻动作 aₜ₋₁
GRU 状态更新

hₜ = f(hₜ₋₁, zₜ₋₁, aₜ₋₁)
上一时刻随机状态 zₜ₋₁
上一时刻确定性状态 hₜ₋₁
先验分布

p(zₜ | hₜ)
后验采样得到 zₜ
潜在状态

(hₜ, zₜ)
奖励预测器

预测 rₜ
图像解码器

重建观测 ôₜ
当前潜在状态

(hₜ, zₜ)
CEM 采样与优化

候选动作序列
RSSM 潜在动态展开

预测未来状态与奖励
累计预测奖励

评估候选序列
选出当前较优动作序列
只执行第一个动作

aₜ
真实环境执行动作
获得新图像 oₜ₊₁
CNN + RSSM

更新当前潜在状态

我们回到一开始的整体结构：前面几节把编码器、GRU、先验、后验、解码器与奖励头拆开讲了一遍，这一节把它们拼回去，看一次前向里数据在每一站变成什么形状
形状约定：序列统一按时间维在前存放，一次采样的观测是 ( L , B , 3 , 64 , 64 ) (L,B,3,64,64) (L,B,3,64,64)、动作是 ( L , B , 6 ) (L,B,6) (L,B,6)、奖励是 ( L , B ) (L,B) (L,B)，其中 L = 50 L=50 L=50 为片段长度， B = 50 B=50 B=50 为批大小， 6 6 6 为 walker 的动作维
编码器、解码器与奖励头都只处理单个时刻，所以由 bottle 先把时间维并进批次维，算完再拆回去：

python 复制代码

def bottle(f, x_tuple):
  x_sizes = tuple(map(lambda x: x.size(), x_tuple))
  y = f(*map(lambda x: x[0].view(x[1][0] * x[1][1], *x[1][2:]), zip(x_tuple, x_sizes)))
  y_size = y.size()
  return y.view(x_sizes[0][0], x_sizes[0][1], *y_size[1:])

于是整段前向里，编码器一次就吞下 49 × 50 = 2450 49\times50=2450 49×50=2450 张图，解码器与奖励头同理，时序被摊平之后梯度沿批次维自然累加
单步之内，数据的形状依次变化如下：

环节	输入	输出
拼接 $z t - 1 ; a t - 1$ $z_{t-1};a_{t-1}$ $zt-1;at-1$	30 + 6 30+6 30+6	36 36 36
状态动作嵌入	36 36 36	200 200 200
GRU 递推	200 200 200 输入 + 200 200 200 隐状态	h t h_t ht： 200 200 200
先验头	200 → 200 200\to200 200→200	μ θ \mu_\theta μθ 与 σ θ \sigma_\theta σθ 各 30 30 30
CNN编码器	3 × 64 × 64 3\times64\times64 3×64×64	e t e_t et： 1024 1024 1024
拼接 $h t ; e t$ $h_t;e_t$ $ht;et$	200 + 1024 200+1024 200+1024	1224 1224 1224
后验头	1224 → 200 1224\to200 1224→200	μ ϕ \mu_\phi μϕ 与 σ ϕ \sigma_\phi σϕ 各 30 30 30
重参数化采样	均值与标准差各 30 30 30	z t z_t zt： 30 30 30
解码器	200 + 30 → 1024 200+30\to1024 200+30→1024	64 × 64 × 3 64\times64\times3 64×64×3
奖励头	200 + 30 → 200 → 200 200+30\to200\to200 200+30→200→200	标量奖励

解码器一路把空间维从 1 × 1 1\times1 1×1 放回来： 1024 1024 1024 嵌入先变成 1 × 1 1\times1 1×1 的张量，四层转置卷积依次经过 5 × 5 5\times5 5×5、 13 × 13 13\times13 13×13、 30 × 30 30\times30 30×30，最后一层回到 64 × 64 × 3 64\times64\times3 64×64×3
整段序列一次算完，靠的是 transition_model 里那个按时间展开的循环，动作、观测与终止标志都按时间维输入：

python 复制代码

# 一段 50 帧的片段只做 49 次转移：第 t 步用上一拍的 z、上一拍的动作与当前帧 o_t
beliefs, prior_states, prior_means, prior_std_devs, posterior_states, posterior_means, posterior_std_devs = transition_model(
    init_state, actions[:-1], init_belief, bottle(encoder, (observations[1:], )), nonterminals[:-1])

输入为什么都错开一格：第 t t t 步的转移要用 ( h t − 1 , z t − 1 ) (h_{t-1},z_{t-1}) (ht−1,zt−1) 与 a t − 1 a_{t-1} at−1 才能算出 h t h_t ht，再与当前帧 o t o_t ot 一起得到 z t z_t zt，所以 o t o_t ot 配 a t − 1 a_{t-1} at−1，这就是 actions[:-1] 与 observations[1:] 的来源
- 错开也保证了因果性：后验只用 t t t 时刻及以前的观测，不会提前看到下一帧
终止标志负责把片段断开：遇到终止时，上一拍的随机状态先被清零再喂给 GRU，避免一个回合的结尾污染下一个回合的开头
拿到的整段潜在状态送入两个只看单拍的模块，再算三项损失：

python 复制代码

# 重建与奖励都用后验状态，时序由 bottle 处理
observation_loss = F.mse_loss(bottle(observation_model, (beliefs, posterior_states)),
                              observations[1:], reduction='none').sum(dim=(2, 3, 4)).mean(dim=(0, 1))
reward_loss = F.mse_loss(bottle(reward_model, (beliefs, posterior_states)),
                         rewards[:-1], reduction='none').mean(dim=(0, 1))
# KL 项的构造见 2-3-2-3，这里只保留训练时的写法
kl_loss = torch.max(kl_divergence(Normal(posterior_means, posterior_std_devs),
                                  Normal(prior_means, prior_std_devs)).sum(dim=2),
                    free_nats).mean(dim=(0, 1))

三项损失共享同一套梯度：重建与奖励从后验状态往回传，依次经过解码器或奖励头、后验头、GRU 与编码器；KL 项同时牵动先验头与后验头，把先验往被观测校正过的后验上拉
至此一次前向的完整链路就闭合了：像素进来，潜在状态出去，奖励与重建作为监督，规划阶段则只用到中间那段 ( h t , z t ) (h_t,z_t) (ht,zt)

2-7 训练

总目标为三项损失之和：
L = L r e c o n + L r e w a r d + β L K L \mathcal{L}=\mathcal{L}{\mathrm{recon}}+\mathcal{L}{\mathrm{reward}}+\beta\,\mathcal{L}_{\mathrm{KL}} L=Lrecon+Lreward+βLKL

python 复制代码

# 三项损失相加后一次反传，三项各自的计算见 2-6 节
loss = observation_loss + reward_loss + kl_loss
optimiser.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(param_list, args.grad_clip_norm, norm_type=2)
optimiser.step()

实现里 β \beta β 取 1，三项直接相加，不再单独给 KL 加权
KL 项设 free nats 为 3，低于该阈值的散度不再产生梯度，避免后验过早塌缩到先验
优化器为 Adam，主要超参数如下：

参数	取值	含义
回合预算	1000	每个任务采集的回合数
种子回合	5	初始随机数据
片段长度	50	每次拟合的序列步数
批大小	50	每次拟合的序列条数
动作重复	2	一个控制步对应的仿真步数
学习率	1e-3	Adam 步长
梯度裁剪	1000	全局范数上限

训练初期模型只见过随机动作产生的数据，规划质量有限；随着回放池扩大，先验逐渐准确，规划带来的数据分布改善又会反过来提升模型
这一循环是 PlaNet 与纯监督式世界模型的主要区别：数据不是固定数据集，而是由规划策略持续生成
一个训练的模板就是

python 复制代码

# 1. 用 S = 5 个随机动作的种子回合填满回放池
D = ReplayBuffer(capacity=experience_size)
for _ in range(seed_episodes):
    D.append_episode(rollout(env, policy=random_actions))

# 2. 之后每一轮：先拟合 C = 100 次，再用当前规划器采集一幕写回 D
for episode in range(seed_episodes, episodes):
    for s in range(collect_interval):
        observations, actions, rewards, nonterminals = D.sample(batch_size, chunk_size)
        # 一次前向得到整段的信念、先验与后验，再算三项损失
        loss = observation_loss + reward_loss + kl_loss
        optimiser.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(param_list, grad_clip_norm, norm_type=2)
        optimiser.step()

    # 采集：每一步先用真实观测更新后验，再由 CEM 选动作，加探索噪声后执行
    observation, belief, posterior_state = env.reset(), zeros(belief_size), zeros(state_size)
    for t in range(max_episode_length // action_repeat):
        belief, posterior_state, action, observation, reward, done = update_belief_and_act(
            planner, transition_model, encoder, belief, posterior_state, action, observation, explore=True)
        D.append(observation, action, reward, done)
        if done:
            break

代码里的 collect_interval 就是论文 Algorithm 1 中的采集间隔 C C C，这里取 100：每拟合 100 次，才用新的模型采一幕数据

2-8 部署

部署阶段只保留编码器、RSSM、奖励头与 CEM 规划器，解码器不参与推理
每个控制步的执行流程如下：
1. 用编码器把当前观测编码为 e t e_t et
2. 结合上一时刻的 ( h t − 1 , z t − 1 ) (h_{t-1},z_{t-1}) (ht−1,zt−1) 与动作，用后验更新得到 ( h t , z t ) (h_t,z_t) (ht,zt)
3. 从 ( h t , z t ) (h_t,z_t) (ht,zt) 出发，用 CEM 在潜在空间搜索长度为 H H H 的动作序列
4. 只执行序列中的第一个动作，环境返回新观测后回到第 1 步
这四步在同一个循环里完成，控制回路的全部内容只有五行：

python 复制代码

# 部署循环：滤波一拍、规划一拍、执行第一个动作
for step in range(max_episode_length // action_repeat):
    belief, _, _, _, posterior_state, _, _ = transition_model(
        posterior_state, action.unsqueeze(0), belief, encoder(observation).unsqueeze(0))
    action = planner(belief, posterior_state)
    action.clamp_(env.action_range[0], env.action_range[1])
    observation, reward, done = env.step(action[0].cpu())

这里给 transition_model 传了一步动作与一步观测，所以走的是 2-4-2 节里的后验分支，返回的 posterior_state 就是这一拍的 z t z_t zt，与 belief 一起构成 ( h t , z t ) (h_t,z_t) (ht,zt)
规划器拿到 ( h t , z t ) (h_t,z_t) (ht,zt) 之后，内部再做 observations=None 的 H H H 步滚动（2-4-2 节），所以同一段代码同时承担了滤波与预测
初始时刻 belief、posterior_state 与 action 都置零：还没有历史，第一拍之后才由观测与动作逐步积累信息
动作要夹到环境允许的范围，CEM 采样出来的是高斯样本，本身不保证落在界内
整个推理包在 torch.no_grad() 里，不产生梯度，也不更新任何参数
权重里仍保存着解码器，但它只在需要生成重建对比图时被调用一次，控制回路用不到它
规划参数与训练阶段保持一致：

参数	取值	含义
规划时域 H H H	12	每次向前预测的步数
CEM 迭代次数	10	分布更新轮数
候选序列数	1000	每轮评估的序列条数
精英样本数	100	参与分布更新的序列条数

每执行一个动作就重新规划，新的观测会修正潜在状态，避免整段回合始终开环预测，减轻预测误差持续累积的影响
与 MBPO 不同，PlaNet 部署时仍然依赖模型，计算开销集中在规划而不是策略前向

2-9 对比

2-9-1 实验介绍

六个任务全部来自 DeepMind Control Suite，观测只有 64 × 64 × 3 64\times64\times3 64×64×3 的第三人称相机图像，动作是连续控制量，回报由各任务自己定义
论文把这些环境摆在同一张图里，从左到右依次是下面六个：
六个任务卡住的地方各不相同：

任务	画面	难点	动作重复
Cartpole Swing Up	可摆动的小车	相机固定，杆摆起后小车会跑出视野，要求长时域规划与记忆	8
Reacher Easy	二连杆手臂与目标球	手与目标区域重叠才有奖励，稀疏	4
Cheetah Run	四足奔跑	状态与动作空间更大，接触频繁	4
Finger Spin	手指拨动方块	手指与物体之间的接触动力学	2
Cup Catch	杯子接住滚来的球	只有接住的那一刻有奖励，稀疏	4
Walker Walk	双足机器人	先站起再行走，倒地时的碰撞难以预测	2

动作重复是一个控制步里执行多少次仿真步，数值取自论文附录；它决定规划时域 12 步对应多长的真实时间
六个任务合起来覆盖了四类难点：部分可观测、接触动力学、稀疏奖励，以及更大的状态与动作空间，后面的三组对照正是在这些场景上做的

2-9-2 实验结果

第一组对照把 PlaNet、两个无模型基线与两种更简单的模型设计画在一起，六张子图共用一套坐标系，横轴是采集的回合数，纵轴是测试回报，实线取中位数，阴影是 5% 到 95% 分位，结果如下图：
看图之前先看横轴：只到 1000 个回合，而两条虚线是 A3C 与 D4PG 训练 10 万回合之后的最终水平，蓝线用的数据量只有基线的百分之一
在回合预算远小于两个无模型基线的情况下，PlaNet 在前 100 个回合内就超过了 A3C，并在约 500 个回合后接近 D4PG，差距主要体现在手指旋转任务
同一张图里还叠了模型设计的消融：绿线是纯随机状态模型（SSM），红线是纯确定性模型（GRU），六张子图上两者都明显低于蓝线
- Cartpole 与 Walker 上最直观：纯 GRU 的曲线几乎平在 200 以下，没有随机状态，模型就学不到"同一段历史之后可能有多种走向"
- Reacher 与 Cup Catch 的阴影格外宽，论文在图注里点明是稀疏奖励造成的：只有少数几次撞上奖励，分位带就跟着大起大落
第二组对照检查智能体自身的设计，论文单独给了一张图：
红线是 random shooting：每一步从 1000 条候选序列里挑最好的一条，不做 CEM 的迭代精修，在 Cartpole 与 Finger 上都明显低于蓝线
绿线是 random collection：用随机动作采集数据再训练，Cheetah 与 Finger 上最终只到 400 上下，Walker 上的差距更明显
三组对照合起来看，设计上的每一处取舍都对应一段曲线差距：

消融对象	对照设置	观察
动力学模型	纯 GRU、纯随机状态模型	两类状态都必要
数据来源	随机动作采集	在线规划采集更有效
规划方式	随机射击替代 CEM	CEM 的迭代优化带来提升

两张图分别出自论文 Figure 4 与 Figure 5；纵轴是各任务自身的回报，不同任务之间不能直接比高低，只能比较同一张子图里的几条曲线
论文还验证了单个智能体在六个任务上共享模型的多任务设置，说明潜在动力学可以跨任务复用

3 Walker 实验

3-1 任务与数据

3-1-1 walker-walk 任务

本章在 walker-walk 任务上验证二维平面内的双足行走，动作为 6 维关节力矩，观测为 64 × 64 × 3 64\times64\times3 64×64×3 的第三人称图像，动作重复 2，回合长度 1000 个仿真步，对应 500 个控制步
walker-walk 的一帧画面如下，画面中机器人正处在迈步过程中：
该任务的难点在于智能体需要先从倒地姿态起身，再维持稳定步态，接触变化难以精确预测
walker 不设失败终止条件，回合固定跑满时间上限，因此各集回报可以直接比较

3-1-2 观测来源

基于像素的世界模型不依赖离线数据集，训练数据全部来自智能体与仿真环境的交互
- 观测由 DeepMind Control Suite 的 walker 场景实时渲染：固定使用 0 号相机，画面缩放至 64 × 64 64\times64 64×64，按 5 bit 量化后以 uint8 存储
量化用于压缩回放池占用，训练取样时将量化误差还原为均匀噪声，相当于用连续分布近似离散图像的似然：

python 复制代码

def _images_to_observation(images, bit_depth):
  images = torch.tensor(cv2.resize(images, (64, 64), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1), dtype=torch.float32)
  preprocess_observation_(images, bit_depth)  # 量化到 5 bit、居中到 [-0.5, 0.5] 并加回噪声
  return images.unsqueeze(dim=0)

动作重复 2 表示智能体每执行一个决策，仿真器内部前进两个 0.025 s 的物理步，因此 1000 个仿真步对应 500 次决策

3-1-3 奖励函数

奖励由环境给出，不依赖人工标注：walker 每一步的奖励由站立项与移动项相乘得到
r = 3 s t a n d + u p r i g h t 4 ⋅ 5 m o v e + 1 6 r=\frac{3\,\mathrm{stand}+\mathrm{upright}}{4}\cdot\frac{5\,\mathrm{move}+1}{6} r=43stand+upright⋅65move+1
- 站立项 s t a n d \mathrm{stand} stand 是躯干高度对 1.2 m 的软阈值，高于 1.2 m 取 1，低于该高度按线性衰减； u p r i g h t \mathrm{upright} upright 是躯干朝向的归一化结果，二者共同要求机器人保持直立
- 移动项 m o v e \mathrm{move} move 是水平速度对 1 m/s 的软阈值，达到 1 m/s 取 1，速度不足时衰减，使奖励偏向前进而非原地站立
单个仿真步的奖励落在 $0 , 1$ $0,1$ $0,1$ ，一个控制步内累加两次，因此单控制步的奖励上限为 2，一回合的回报上限为 1000：

python 复制代码

reward = 0
for k in range(self.action_repeat):
  state = self._env.step(action)
  reward += state.reward
  self.t += 1
  done = state.last() or self.t == self.max_episode_length

该奖励在三个环节被使用：训练时由奖励头回归，规划时 CEM 以奖励头的预测值评估候选动作序列，评测时将真实奖励逐控制步累加得到回合回报
- 与 reacher、cup 等稀疏奖励任务不同，walker 的奖励每一步都有信号，能为奖励头提供更密集的监督；奖励头的结构不变，仍是 $h t ; z t$ $h_t;z_t$ $ht;zt$ 上的多层感知机，230 维输入经两层 200 维隐藏单元输出一个标量，训练目标仍是 2-3-9 节的均方误差，这也是本章选择该任务验证闭环的原因之一

3-1-4 数据采集

训练使用的回放池由智能体自身写入，分为两段：前 5 个回合以随机动作采集种子数据，之后每完成 100 次拟合，使用当前规划器采集一个回合写入回放池
- 采集时动作叠加标准差为 0.3 的高斯噪声，否则回放池只包含当前策略的确定性轨迹，模型容易过拟合到单一分布
- 每条记录为（观测、动作、奖励、终止标志）四元组，写入容量 100 万的环形缓冲，写满后覆盖最旧的数据；1000 个回合合计约 50 万条，缓冲不会写满
种子回合的采集代码如下：

python 复制代码

for s in range(1, seed_episodes + 1):
  observation, done, t = env.reset(), False, 0
  while not done:
    action = env.sample_random_action()
    next_observation, reward, done = env.step(action)
    D.append(observation, action, reward, done)
    observation = next_observation
    t += 1

种子数据使模型在首次拟合前即可看到足够多样的画面，否则回放池规模过小，先验与后验均缺少学习依据

3-1-5 序列采样

拟合时按序列抽样而非逐条抽样：每次独立抽取 50 条 50 帧的连续片段，单条片段不能跨过环形缓冲的写入游标
代码通过取模将片段绕成环形，并检查片段是否包含写入游标：

python 复制代码

def _sample_idx(self, L):
  valid_idx = False
  while not valid_idx:
    idx = np.random.randint(0, self.size if self.full else self.idx - L)
    idxs = np.arange(idx, idx + L) % self.size
    valid_idx = not self.idx in idxs[1:]  # 保证片段不跨过写入游标
  return idxs

def sample(self, n, L):
  idxs = np.asarray([self._sample_idx(L) for _ in range(n)])
  return self._retrieve_batch(idxs, n, L)

抽样得到的片段按 2-6 节的方式对齐：第 t 步使用第 t 帧与第 t-1 个动作，一条 50 帧片段产生 49 个转移
- 该批数据用于三项损失：重建损失监督编码器与解码器，奖励损失监督奖励头，KL 项约束先验与后验；规划阶段不读取回放池，仅使用当前潜在状态

3-2 实验步骤

3-2-1 权重加载

为节省计算开销，本文不重新训练模型，而是加载与论文同等回合预算下训练得到的公开权重，验证从图像到动作的完整闭环
- 权重包含四个模块的参数：过渡模型（GRU 与先验、后验头）、观测模型（解码器）、奖励模型与编码器
加载后将四个模块切换至推理模式：

python 复制代码

sd = torch.load(models, map_location=device)
transition_model.load_state_dict(sd['transition_model'])
observation_model.load_state_dict(sd['observation_model'])
reward_model.load_state_dict(sd['reward_model'])
encoder.load_state_dict(sd['encoder'])
transition_model.eval(), observation_model.eval(), reward_model.eval(), encoder.eval()

评测阶段不需要回放池与优化器，二者仅用于训练

3-2-2 评测流程

评测按回合依次进行，每个控制步的顺序固定为：编码当前帧、以观测更新潜在状态后验、CEM 规划、执行第一个动作
与训练时的数据采集相比有两处差异：动作不叠加探索噪声，数据不写回回放池

python 复制代码

with torch.no_grad():
  for episode in range(episodes):
    observation = env.reset()
    belief = torch.zeros(1, belief_size, device=device)
    posterior_state = torch.zeros(1, state_size, device=device)
    action = torch.zeros(1, action_size, device=device)
    total = 0.0
    for step in range(max_episode_length // action_repeat):
      belief, _, _, _, posterior_state, _, _ = transition_model(
          posterior_state, action.unsqueeze(0), belief, encoder(observation).unsqueeze(0))
      belief, posterior_state = belief.squeeze(0), posterior_state.squeeze(0)
      action = planner(belief, posterior_state)
      action.clamp_(action_range[0], action_range[1])
      observation, reward, done = env.step(action[0].cpu())
      total += reward
      if done:
        break
    returns.append(total)

每集开始时信念、潜在状态与动作均置零，第一拍之后由观测与动作逐步积累历史信息
重建质量单独评估：在若干固定控制步上，用解码器将当前 ( h t , z t ) (h_t,z_t) (ht,zt) 还原为图像，与真实观测并排保存；规划过程不使用解码器
评测共 10 集，随机种子固定为 1，动作不叠加噪声，同一权重在同一台机器上的评测轨迹可以复现

3-3 实验结果

实验内容：在 walker-walk 上评估预训练权重的行为表现与重建质量
实验实现：每个控制步用编码器更新后验，由 CEM 在潜在空间规划动作并只执行第一个动作；同时记录真实观测与解码器重建
下图给出一个回合内的行走过程：
智能体从倒地姿态起身，随后保持直立并交替迈步，说明编码、潜在推断与规划三段链路均已闭合
下图对比同一回合中四个时刻的真实观测与重建：
重建保留了躯干与双腿的姿态，但细节明显模糊；这是像素均方误差训练的典型表现，也说明潜在状态编码的是用于规划的结构信息，而非逐像素复现
十集回报的均值为 750.6，标准差 121.6，最低 423.5，最高 882.1；按每回合上限 1000 折算，平均每个控制步获得约 1.5 的奖励
- 论文的 walker 曲线收敛于 950 附近，本地结果略低
- 该差距不来自链路错误，逐集评测与官方测试入口给出的均值完全一致
- 更可能的原因是仿真器版本差异：任务模型文件多年未变，但物理引擎从 2.x 升级到 3.x 后接触求解行为发生变化，固定策略对这类差异较为敏感

总结

本文从 PETS 的规划思路出发，介绍了 PlaNet 的 RSSM、latent overshooting 与潜在空间规划，并在 Walker 任务上用预训练权重验证了完整闭环
核心要点回顾：
- 确定性状态与随机状态分工：GRU 维护长期记忆，随机状态表达不确定性
- 先验与后验的 KL 约束：训练时用观测校正，规划时沿先验外推
- 多步预测约束：latent overshooting 用多跨度的 KL 直接约束多步先验，但消融显示它对 DRNN 提升明显，对 RSSM 反而略有下降，最终智能体没有启用
- 潜在空间规划：CEM 在低维状态上评估动作序列，只执行第一个动作
- 模型在部署时仍在使用：这一点与 MBPO 不同，计算开销集中在规划
如有错误，欢迎指出！感谢观看！