RSSM基础：变分循环状态空间模型

一句话定位：Dreamer 系列算法的核心世界模型架构，通过同时建模确定性 hidden state 和随机 latent variable ，结合后验-先验对比学习，实现高效的时间序列建模与规划。

前置依赖

Latent Transition（后验 vs 先验结构化设计）
变分推断基础（ELBO、Jensen 不等式）
重参数化技巧（Reparameterization Trick）
循环神经网络基础（GRU/LSTM）
信息论基础（KL 散度、互信息）

核心思想

Recurrent State Space Model（RSSM） 是 Hafner et al. (2019) 在论文”Plearning-Prediction”中提出的世界模型架构，后来成为 Dreamer 系列算法的核心组件。其设计哲学是：将确定性建模的高效性与随机性建模的表达能力结合。

核心观察：

环境的”规律”（physics）是确定性的——给定相同初始条件和动作，物理系统应该产生相同的结果
环境的”噪声”（perception/irreducible noise）是随机的——需要用概率分布建模
时间序列的长期依赖需要高效的信息传递机制

RSSM 的解决方案：

将隐状态空间分解为两个正交组件：

确定性 hidden state ：通过 GRU 传递，承载长期记忆和上下文信息
随机 latent variable ：从后验/先验分布采样，捕获局部不确定性

这种设计的数学好处是：

的更新仅依赖前一时刻的和，梯度路径短且稳定
的采样引入了随机性，但训练时通过重参数化技巧保持梯度流通
先验和后验的对比学习使模型同时学习”物理规律”和”噪声建模”

模型结构图

A. RSSM 完整架构

                    a_t
                     │
         ┌───────────┴───────────┐
         ▼                       ▼
    ┌─────────┐            ┌─────────┐
    │先验网络 │            │后验网络 │
    │p(z_t|  │            │q(z_t|  │
    │h_{t-1},│            │h_t)    │
    │a_{t-1})│            │         │
    └────┬────┘            └────┬────┘
         │                      │
         │     ┌────────────────┘
         │     │  z_t
         ▼     ▼
    ┌──────────────────────────┐
    │   确定性转移函数 (GRU)    │
    │ h_t = f(h_{t-1}, z_t, a_{t-1})│
    └──────────────────────────┘
         │
         ▼
       h_t ──────────┬──────────► 观测解码器 o_hat_t
                     │
                     └──────────► 奖励预测器 r_hat_t

B. 信息流图

时间 t-1:          时间 t:           时间 t+1:
  h_{t-1} ───────► h_t ───────► h_{t+1}    (确定性通道，RNN)
    │                ▲                │
    │                │                │
    ▼                │                ▼
  z_{t-1} ───────► z_t ───────► z_{t+1}    (随机通道，采样)
                     ▲
                     │
                后验 q(z_t|o_{1:T})
                或先验 p(z_t|h_{t-1},a_{t-1})

C. 与原始 SSM 的对比

方面	朴素 SSM	RSSM
隐状态类型	随机	确定性 + 随机
状态转移
不确定性传播	通过序列传播	先验-后验对比
梯度传播	长序列梯度消失	通过高效传播
长期依赖	弱	强

数学推导

A. 完整概率模型定义

RSSM 定义了以下概率分布：

生成模型（世界模型）： $先验确定性转移解码器策略$

变分推断模型（后验）： $后验推断转移$

其中：

：初始 hidden state 先验，通常设为零向量
：初始 hidden state 由观测决定
：确定性推断转移

B. ELBO 的完整推导（核心）

我们的目标是最大化观测对数似然（条件似然，给定动作序列）。

第一步：引入变分分布

使用 Jensen 不等式，引入变分分布：

第二步：应用 Jensen 不等式

由于是凹函数：

第三步：因式分解联合分布

对生成模型进行因式分解：

对变分分布进行因式分解：

第四步：代入并简化

由于（确定性转移，相同函数），这些项在 ELBO 中相互抵消。

最终得到简化的 ELBO：

详细推导：

重构项：

由于是的确定性函数（通过转移函数），我们可以用重参数化技巧计算这个期望。

KL 项：

C. KL 项的深入理解

数学形式：对于高斯分布和：

直观理解： KL 项包含三个组成部分：

方差比： — 鼓励后验方差接近先验方差
均值差： — 惩罚后验均值与先验均值的偏差
方差项： — 额外惩罚方差偏差

在 RSSM 中的作用：

训练时，KL 项迫使后验接近先验
但由于使用的是 -VAE 中的参数（通常），允许后验保留一定的新信息
这使得模型能够学习哪些信息来自观测（后验与先验的差异）

D. 重参数化技巧在 RSSM 中的应用

问题：从后验采样不可导。

解决方案：

梯度计算：

其中（线性），（与相连）。

E. 后验 vs 先验 over latent states

后验（Posterior）：

输入：当前 hidden state （编码了观测的信息）目的：利用当前观测校正对的预测

先验（Prior）：

输入：前一 hidden state 和动作目的：基于”物理规律”预测的分布

对比学习的直观理解：

如果先验已经准确预测了后验（KL ≈ 0），说明可以完全由历史信息确定，没有新信息来自观测
如果后验与先验差异大，说明观测提供了重要的新信息，有效捕获了这些信息

训练细节

A. 完整训练目标

$重构损失（观测）损失$

其中是 KL 权重超参数。

B. 训练流程

输入: 观测序列 o_{1:T}, 动作序列 a_{1:T-1}

1. 初始化:
   h_0 = zeros
   total_loss = 0

2. 对于 t = 1 to T:
   a. 编码观测:
      h_t = encoder(o_t)   # h_t 是 deterministic hidden state
   
   b. 计算先验分布:
      p_mu, p_sigma = prior_net(h_{t-1}, a_{t-1})
      p(z_t) = N(p_mu, p_sigma)
   
   c. 计算后验分布:
      q_mu, q_sigma = posterior_net(h_t)
      q(z_t) = N(q_mu, q_sigma)
   
   d. 重参数化采样（训练时用后验）:
      epsilon ~ N(0, I)
      z_t = q_mu + q_sigma * epsilon
   
   e. 更新 hidden state（确定性转移）:
      h_t = transition_net(h_{t-1}, z_t, a_{t-1})
      # 注意: 这里 h_t 被更新为新的 hidden state
   
   f. 重构观测:
      o_hat_t = decoder(h_t)
      recon_loss = -log p(o_t | h_t)  (或 MSE)
   
   g. 计算 KL 损失:
      kl_loss = D_KL(q(z_t) || p(z_t))
   
   h. 累积损失:
      step_loss = recon_loss + beta * kl_loss
      total_loss += step_loss
   
   i. 更新 h_{t-1} = h_t（为下一步准备）

3. 反向传播:
   total_loss.backward()
   optimizer.step()

C. 关键训练技巧

KL 权重调度（KL Annealing）：

训练初期使用较小的（如 0.1），让模型先学习有效的表征；训练稳定后逐渐增大，加强先验约束。

自由比特（Free Bits）：对每个样本的 KL 项设置下限：

这防止后验分布因过度匹配先验而崩塌。

目标网络（Target Networks）：先验网络和后验网络的参数定期从主网络复制，延迟更新，以稳定训练。

D. 损失函数的分解

损失项	公式	作用
重构损失		确保隐状态保留观测信息
KL 损失		约束后验接近先验，学习物理规律
奖励损失（可选）		额外训练奖励预测，增强规划能力

推理/rollout/规划过程

A. 先验 rollout（想象 rollout）

训练完成后，使用先验分布进行免费 rollout 用于规划：

o_0 = env.reset()
h_0 = encoder(o_0)

对于 t = 1 to T:
    # 1. 从先验采样 z_t（不使用后验，因为没有未来观测）
    z_t ~ p_theta(z_t | h_{t-1}, a_{t-1})
    
    # 2. 确定性转移计算新 hidden state
    h_t = f_theta(h_{t-1}, z_t, a_{t-1})
    
    # 3. 策略网络输出动作
    a_t = pi_theta(h_t)
    
    # 4. (可选) 预测奖励
    r_hat_t = reward_net(h_t)
    
    # 5. 在真实环境中执行动作，获得真实奖励和下一观测
    r_t, o_{t+1} = env.step(a_t)
    
    # 6. 更新隐状态
    h_t = encoder(o_{t+1})

关键：先验 rollout 不需要后验网络，仅依赖先验分布采样，因此可以完全脱离真实观测生成想象轨迹。

B. 后验 rollout（用于重建）

如果要重建真实轨迹，使用后验 rollout：

对于每个时间步 t:
    z_t ~ q_phi(z_t | h_t)   # 后验需要看到观测
    h_{t+1} = f(h_t, z_t, a_t)

后验 rollout 用于训练时的后验推断，确保的采样与真实观测一致。

C. 规划算法

基于梯度的规划：给定累积奖励目标：

动作梯度可显式计算：

其中通过 RSSM 的确定性转移函数计算。

MPC（模型预测控制）：

采样条候选动作序列
使用 RSSM 先验 rollout 预测每条序列的隐状态轨迹
评估累积奖励
选择最优动作序列的第一个动作执行

优点与局限

优点

长期依赖高效：确定性 hidden state 通过 RNN 传递，梯度路径稳定，避免了纯随机 SSM 的梯度消失问题
不确定性显式建模：随机 latent variable 捕获感知噪声和多模态动态
后验-先验对比学习：KL 项结构化地学习物理规律与观测新信息的分离
规划友好：先验 rollout 生成多样化想象轨迹，支持基于梯度的规划和 MPC
端到端训练：整个架构可通过 BPTT 端到端优化

局限

后验推断复杂度：后验需要看到当前观测，无法进行”超前预测”
高斯假设限制：先验和后验都假设高斯分布，无法表达更复杂的隐变量分布
训练不稳定：KL 权重需要仔细调参；过小导致后验崩溃，过大导致先验失效
随机性权衡：需要权衡随机性（多样性与鲁棒性）和确定性（精确预测）

与前后内容的衔接

前置内容：Latent Transition 定义了后验 vs 先验的结构化设计，RSSM 是这一设计的具体实现
后续发展：Dreamer 算法在 RSSM 基础上添加策略学习（actor-critic）和 imagination rollout，实现基于模型的强化学习
核心贡献：RSSM 完成了从理论设计（Latent Transition）到实际算法（Dreamer）的过渡

可复现实现要点

A. PyTorch 实现框架

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
 
class RSSM(nn.Module):
    """Recurrent State Space Model"""
    
    def __init__(self, hidden_dim, action_dim, latent_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.action_dim = action_dim
        self.latent_dim = latent_dim
        
        # 1. 先验网络: p(z_t | h_{t-1}, a_{t-1})
        self.prior_net = nn.Sequential(
            nn.Linear(hidden_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # [mu, log_sigma]
        )
        
        # 2. 后验网络: q(z_t | h_t)
        self.posterior_net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # [mu, log_sigma]
        )
        
        # 3. 确定性转移函数: p(h_t | h_{t-1}, z_t, a_{t-1}) - 用 GRU
        self.gru = nn.GRUCell(hidden_dim + latent_dim + action_dim, hidden_dim)
        
        # 4. 解码器: p(o_t | h_t)
        self.decoder_net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, obs_dim)  # 观测维度
        )
        
        # 5. 奖励预测器: p(r_t | h_t) - 规划用
        self.reward_net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        
        # 6. 策略网络: π(a_t | h_t) - Dreamer 的 actor
        self.policy_net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim * 2)  # [mu, log_sigma] for action
        )
        
    def prior(self, h_prev, action):
        """计算先验分布 p(z_t | h_{t-1}, a_{t-1})"""
        x = torch.cat([h_prev, action], dim=-1)
        stats = self.prior_net(x)
        mu, log_sigma = stats.chunk(2, dim=-1)
        sigma = F.softplus(log_sigma) + 1e-6
        return mu, sigma
    
    def posterior(self, h_t):
        """计算后验分布 q(z_t | h_t)"""
        stats = self.posterior_net(h_t)
        mu, log_sigma = stats.chunk(2, dim=-1)
        sigma = F.softplus(log_sigma) + 1e-6
        return mu, sigma
    
    def forward(self, h_prev, z_t, action):
        """确定性转移函数 h_t = f(h_{t-1}, z_t, a_{t-1})"""
        x = torch.cat([h_prev, z_t, action], dim=-1)
        h_t = self.gru(x, h_prev)
        return h_t
    
    def reparameterize(self, mu, sigma):
        """重参数化技巧"""
        eps = torch.randn_like(mu)
        return mu + sigma * eps
    
    def kl_divergence(self, q_mu, q_sigma, p_mu, p_sigma):
        """计算 KL(q || p) - 单样本版本"""
        # D_KL(N(q_mu, q_sigma) || N(p_mu, p_sigma))
        log_ratio = torch.log(p_sigma / q_sigma + 1e-6)
        return log_ratio + (q_sigma**2 + (q_mu - p_mu)**2) / (2 * p_sigma**2 + 1e-6) - 0.5
    
    def imagination_rollout(self, h_0, action_seq):
        """先验 rollout - 用于规划"""
        T = len(action_seq)
        h_seq = [h_0]
        z_seq = []
        
        h_prev = h_0
        for t in range(T):
            # 1. 先验采样
            p_mu, p_sigma = self.prior(h_prev, action_seq[t])
            z_t = self.reparameterize(p_mu, p_sigma)
            
            # 2. 确定性转移
            h_t = self.forward(h_prev, z_t, action_seq[t])
            
            h_seq.append(h_t)
            z_seq.append(z_t)
            h_prev = h_t
        
        return h_seq, z_seq

B. 完整训练循环

def train_step(model, encoder, optimizer, obs_seq, action_seq, reward_seq=None, beta=1.0):
    """
    RSSM 训练步骤
    obs_seq: (batch, T, obs_dim)
    action_seq: (batch, T-1, action_dim)
    reward_seq: (batch, T-1, 1) - 可选
    """
    T = obs_seq.size(1)
    batch_size = obs_seq.size(0)
    
    # 1. 编码观测序列获取初始 hidden state
    h_prev = encoder(obs_seq[:, 0])  # h_0
    
    total_recon_loss = 0
    total_kl_loss = 0
    
    for t in range(1, T):
        # 2. 计算先验分布 p(z_t | h_{t-1}, a_{t-1})
        p_mu, p_sigma = model.prior(h_prev, action_seq[:, t-1])
        
        # 3. 计算后验分布 q(z_t | h_t) - h_t 是编码后的观测隐状态
        h_t = encoder(obs_seq[:, t])  # 观测编码（简化，实际可能用 RNN）
        q_mu, q_sigma = model.posterior(h_t)
        
        # 4. 重参数化采样（训练时用后验）
        z_t = model.reparameterize(q_mu, q_sigma)
        
        # 5. 确定性转移 h_t = f(h_{t-1}, z_t, a_{t-1})
        h_t = model.forward(h_prev, z_t, action_seq[:, t-1])
        
        # 6. 重构观测
        obs_recon = model.decoder_net(h_t)
        recon_loss = F.mse_loss(obs_recon, obs_seq[:, t], reduction='sum')
        
        # 7. 计算 KL 损失（后验 vs 先验）
        kl_loss = model.kl_divergence(q_mu, q_sigma, p_mu, p_sigma).mean()
        
        total_recon_loss += recon_loss
        total_kl_loss += kl_loss
        
        # 更新 h_prev 为下一步准备
        h_prev = h_t
    
    # 8. 总损失
    loss = total_recon_loss + beta * total_kl_loss
    
    # 9. 反向传播
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    
    return {
        'loss': loss.item(),
        'recon_loss': total_recon_loss.item(),
        'kl_loss': total_kl_loss.item()
    }

C. 关键实现细节

方差保证：

sigma = F.softplus(log_sigma) + 1e-6

softplus 保证正值且可微
小常数防止数值问题

重参数化：

eps = torch.randn_like(mu)
z = mu + sigma * eps

eps 必须保持与计算图连接
torch.randn_like 确保与 mu 同设备同形状

KL 计算：

log_ratio = torch.log(p_sigma / q_sigma + 1e-6)
kl = log_ratio + (q_sigma**2 + (q_mu - p_mu)**2) / (2 * p_sigma**2 + 1e-6) - 0.5

梯度裁剪：

torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

防止梯度爆炸，特别是长序列的 BPTT。

章节摘要

RSSM（Recurrent State Space Model）是 Dreamer 系列算法的核心世界模型架构，通过同时建模确定性 hidden state 和随机 latent variable ，结合后验-先验对比学习，实现高效的时间序列建模与规划。

核心数学：

ELBO 推导：
KL 项的物理含义：后验与先验的差异度量了当前观测带来的新信息量
重参数化：

核心设计思想：

通过 GRU 传递长期记忆，梯度高效传播
通过后验-先验对比学习捕获不确定性
先验 rollout 用于生成想象轨迹，支持规划

RSSM 是连接理论设计（Latent Transition）与实际算法（Dreamer）的桥梁，是世界模型领域最重要的基础架构之一。

关键词

RSSM、变分循环状态空间模型、ELBO 推导、后验分布、先验分布、重参数化技巧、KL 散度、后验-先验对比学习、确定性 hidden state、随机 latent variable、想象 rollout、GRU、信息高速公路、自由比特、KL annealing、模型预测控制、梯度规划、Dreamer

Huajishe Notes

Explorer

4-RSSM基础