隐状态转移函数（Latent Transition）
一句话定位：系统化定义隐状态空间中状态转移的数学形式，包括后验 vs 先验分布的结构化设计，以及时间维度的信息传播机制。
前置依赖
- 随机动力学建模（变分推断、ELBO、重参数化）
- 马尔可夫假设与条件独立性
- 贝叶斯网络与因子图基础
- 信息论基础（KL 散度、互信息）
核心思想
在前两节（确定性动力学和随机动力学）中，我们分别讨论了无随机性和有随机性的动力学建模。在本节中，我们将系统化地定义隐状态转移函数的数学形式，为 RSSM（Recurrent State Space Model）奠定理论基础。
核心问题：如何设计隐状态空间的转移结构，使其既能高效地进行时间序列推断，又能保留表达不确定性所需的随机性？
解决方案：将隐状态空间分解为两个相互作用的组件——确定性隐藏状态（deterministic hidden state）
| 组件 | 特点 | 作用 |
|---|---|---|
| 确定性，可长期记忆 | 传递长期上下文信息，梯度平稳传播 | |
| 随机性，捕捉局部不确定性 | 建模观测噪声、多模态动态 |
这种分离设计的直觉来源于：环境的”规律”（physics）是确定性的，而”噪声”（noise）是随机的。确定性部分可以被 RNN 记住并在时间上传播；随机部分只需要在局部建模。
模型结构图
A. 朴素状态空间模型（Naive SSM）
o_t o_{t+1}
│ │
▼ ▼
┌────────┐ ┌────────┐
│编码器 │ │编码器 │
└───┬────┘ └──┬────┘
│ │
▼ ▼
z_t ~~~> z_{t+1}
▲ ▲
│ │
└── a_t ────┘
朴素 SSM 将整个隐状态建模为随机变量
B. 混合状态空间模型（Hybrid SSM / RSSM）
a_t
│
o_t──────────────┼──────────────┐
│ │ │
▼ ▼ ▼
┌─────────┐ ┌─────────────┐ ┌─────────┐
│后验q(z_t│h_{t-1},a_{t-1})│ │先验p(z_t│h_{t-1},a_{t-1})
└────┬────┘ └──────┬──────┘ └────┬────┘
│ │ │
▼ ▼ ▼
z_t ───────────────> z_{t+1}
│ │ │
▼ ▼ ▼
┌─────────────────────────────────────────────┐
│ 确定性转移函数 p(h_{t+1}|h_t,z_t,a_t) │
└─────────────────────────────────────────────┘
│
▼
h_{t+1}
这种结构的关键特点是：
- 随机变量
仍然存在，但其分布由两部分组成：先验（预测）和后验（观测校正） - 确定性 hidden state
作为”信息高速公路”，高效传递长期依赖 - 后验-先验对比学习：训练时通过 KL 项让后验靠近先验，但保持一定距离以保留信息
C. 时序信息流
时间步 t-1 时间步 t 时间步 t+1
│ │ │
▼ ▼ ▼
┌──────┐ ┌──────┐ ┌──────┐
│ h_t-1│ ────────► │ h_t │ ────────► │h_t+1 │ (确定性信息流)
└──────┘ └──────┘ └──────┘
│ ▲ │
│ │ │
▼ ▼ ▼
┌──────┐ ┌──────┐ ┌──────┐
│ z_t-1│ ────────► │ z_t │ ────────► │z_t+1 │ (随机信息流)
└──────┘ └──────┘ └──────┘
信息在两个通道中同时传播：
数学推导
A. 状态空间模型的因子分解
给定观测序列
生成模型（第一因子分解）：
各因子的含义：
| 因子 | 分布 | 说明 |
|---|---|---|
| 先验 | 初始 hidden state，通常设为固定值或零向量 | |
| 先验 | 基于历史预测 latent variable 先验分布 | |
| 转移函数 | 确定性的 hidden state 转移 | |
| 解码器/似然 | 基于 hidden state 重构观测 |
推断模型（第二因子分解）：
其中
B. 后验分布的分解
利用马尔可夫假设和条件独立性，我们有：
这个分解说明：后验分布 = 先验 × 似然，即先验分布被当前观测
更具体地：
- 先验
：基于历史信息对当前 latent variable 的预测 - 似然
：给定 latent variable 时观测的条件分布
在实践中，我们通常用神经网络参数化：
C. 先验分布的设计
标准高斯先验：
条件先验（RSSM 中使用）：
先验网络的输入是前一时刻的 hidden state
先验的设计动机：
- 先验网络能够学习”物理规律”——给定当前状态和动作，预测可能的下一状态分布
- 动作影响先验分布，使得不同动作下的预测有不同的不确定性水平
D. ELBO 的完整推导
目标：最大化观测对数似然的下界
从 Jensen 不等式出发：
引入变分分布
应用 Jensen 不等式（对
将联合分布因式分解代入：
对每个时间步
代入 ELBO：
简化后的 ELBO：
解释：
- 重构项：给定后验采样的 latent variable，通过确定性 hidden state 预测观测
- KL 项：让后验分布
接近先验分布
E. 后验-先验对比学习的直观理解
训练目标（KL 项）：
最小化这个 KL 项，使
关键洞察：在训练完成后，
为什么 latent variable 有用：
- 先验
：基于历史预测下一状态分布 - 后验
：基于历史 + 当前观测校正预测
两者之间的差异度量了当前观测带来的新信息量。如果差异大，说明观测提供了关于当前状态的重要信息；如果差异小，说明预测已经足够准确。
训练细节
A. 训练目标
其中单步 ELBO 为：
B. 训练算法
对于每个训练 batch (o_{1:T}, a_{1:T-1}):
1. 编码观测序列:
h_0 = zeros
对于 t = 1 to T:
h_t = Encoder(o_t) # 或使用 RNN 逐步编码
2. 初始化累积损失:
total_loss = 0
3. 对于 t = 1 to T:
a. 计算先验分布:
p_theta(z_t | h_{t-1}, a_{t-1}) = N(mu_prior, sigma_prior)
mu_prior, sigma_prior = PriorNet(h_{t-1}, a_{t-1})
b. 计算后验分布:
q_phi(z_t | o_{1:T}) = N(mu_post, sigma_post)
mu_post, sigma_post = PosteriorNet(h_t)
c. 重参数化采样:
z_t = mu_post + sigma_post * epsilon, epsilon ~ N(0,I)
d. 计算确定性 hidden state 转移:
h_t = TransitionNet(h_{t-1}, z_t, a_{t-1})
e. 重构观测:
o_hat_t = Decoder(h_t)
f. 计算单步损失:
recon_loss = -log p_theta(o_t | h_t) (或 MSE)
kl_loss = D_KL(q_phi || p_theta)
step_loss = recon_loss + beta * kl_loss
g. 累积损失:
total_loss += step_loss
4. 反向传播更新:
total_loss.backward()
optimizer.step()
C. 关键训练技巧
- 延迟后验（Delayed Posterior）：后验网络输入不仅是当前观测
，还可以包含未来观测 ，以获得更好的全局信息 - 自由比特（Free Bits）：避免后验分布崩塌，允许每个样本保留一定的 KL”预算”
- KL 权重调度：
从小到大 schedule，防止后验过度依赖先验
推理/rollout/规划过程
A. 先验 rollout（训练后用于规划）
训练完成后，使用先验分布进行免费 rollout：
h_0 = encoder(o_0)
对于 t = 1 to T:
z_t ~ p_theta(z_t | h_{t-1}, a_{t-1}) # 从先验采样
h_t = f_theta(h_{t-1}, z_t, a_{t-1}) # 确定性转移
a_t = policy(h_t) # 策略网络
B. 后验 rollout（用于视频预测/重建）
如果要生成与真实轨迹一致的预测，可以使用后验 rollout：
h_0 = encoder(o_0)
对于 t = 1 to T:
z_t ~ q_phi(z_t | o_{1:T}) # 从后验采样
h_t = f_theta(h_{t-1}, z_t, a_{t-1}) # 确定性转移
后验 rollout 保证生成轨迹与真实观测一致（因为
C. 与确定性 RNN 的对比
| 方面 | 确定性 RNN | Latent Transition (RSSM) |
|---|---|---|
| 状态更新 | ||
| 随机性 | 无 | 有（通过 |
| 后验推断 | 不支持 | 支持（ |
| 不确定性传播 | 无 | 有（先验-后验对比） |
| 多模态建模 | 不支持 | 支持 |
优点与局限
优点
- 长期依赖高效传播：确定性 hidden state
作为信息高速公路，避免了随机变量序列的梯度消失问题 - 不确定性显式建模：随机 latent variable
显式捕获动态不确定性 - 后验-先验对比学习：通过 KL 项结构化地学习哪些信息来自观测，哪些来自先验预测
- 多模态表达：latent variable 的混合分布能够表达同一状态-动作对的不同结果
- 规划友好：先验 rollout 生成多样化的想象轨迹用于规划
局限
- 后验推断复杂度：后验
需要看到整个观测序列，计算成本高 - 训练不稳定：KL 项权重需要仔细调参，
过小导致后验崩溃， 过大导致先验失效 - 隐变量维度设计：隐变量维度需要权衡表达能力和计算效率
- 先验假设限制：高斯先验假设可能不足以捕获复杂的隐变量分布
与前后内容的衔接
- 前置内容：确定性动力学提供了
的确定性转移机制；随机动力学引入了 的变分推断框架 - 后续内容：RSSM 是本节 Latent Transition 概念的完整实现，定义了具体网络架构和训练流程
- 核心贡献：本节明确了后验 vs 先验的结构化设计，为 RSSM 的 KL 对比学习提供理论基础
可复现实现要点
A. 核心模块实现
import torch
import torch.nn as nn
import torch.nn.functional as F
class LatentTransition(nn.Module):
"""Latent Transition 模块：后验-先验对比结构"""
def __init__(self, hidden_dim, action_dim, latent_dim):
super().__init__()
self.hidden_dim = hidden_dim
self.latent_dim = latent_dim
# 先验网络：基于 h_{t-1} 和 a_{t-1} 预测 latent 分布
self.prior_net = nn.Sequential(
nn.Linear(hidden_dim + action_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, latent_dim * 2) # 均值 + 对数方差
)
# 后验网络：基于 h_t（编码了 o_t 信息）推断后验分布
self.posterior_net = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, latent_dim * 2)
)
# 确定性转移函数：p(h_t | h_{t-1}, z_t, a_{t-1})
self.transition_net = nn.Sequential(
nn.Linear(hidden_dim + latent_dim + action_dim, hidden_dim),
nn.GRU(hidden_dim, hidden_dim)
)
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):
"""确定性转移函数 p(h_t | h_{t-1}, z_t, a_{t-1})"""
x = torch.cat([h_prev, z_t, action], dim=-1)
h_t = self.transition_net(x, h_prev)
return h_t
def reparameterize(self, mu, sigma):
"""重参数化技巧"""
epsilon = torch.randn_like(mu)
return mu + sigma * epsilon
def kl_divergence(self, q_mu, q_sigma, p_mu, p_sigma):
"""计算 D_KL(q || p)"""
# KL(N(q_mu, q_sigma) || N(p_mu, p_sigma))
# = log(p_sigma/q_sigma) + (q_sigma^2 + (q_mu-p_mu)^2)/(2*p_sigma^2) - 0.5
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.5B. 训练循环
def train_step(model, encoder, decoder, optimizer, obs_seq, action_seq, beta=1.0):
T = obs_seq.size(1)
batch_size = obs_seq.size(0)
# 初始化 hidden state
h_prev = torch.zeros(batch_size, model.hidden_dim, device=obs_seq.device)
total_recon_loss = 0
total_kl_loss = 0
for t in range(1, T):
# 编码当前观测获取 h_t（简化版，实际可能用 RNN）
h_t = encoder(obs_seq[:, t])
# 1. 先验分布
p_mu, p_sigma = model.prior(h_prev, action_seq[:, t-1])
# 2. 后验分布
q_mu, q_sigma = model.posterior(h_t)
# 3. 重参数化采样（训练时用后验）
z_t = model.reparameterize(q_mu, q_sigma)
# 4. 确定性转移
h_t = model.forward(h_prev, z_t, action_seq[:, t-1])
# 5. 重构观测
obs_recon = decoder(h_t)
recon_loss = F.mse_loss(obs_recon, obs_seq[:, t], reduction='sum')
# 6. KL 损失
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_t
# 总损失
loss = total_recon_loss + beta * total_kl_loss
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
return loss.item(), total_recon_loss.item(), total_kl_loss.item()C. 关键实现细节
- 先验/后验共享部分统计量：在某些实现中，后验网络接收额外信息（如未来观测），以提升后验质量
- KL 权重
：使用 标准的 KL 权重，或使用 的 KL annealing - 多步 KL 分解：在长序列训练中，可将 KL 项分解到每一步，而非使用全局后验
章节摘要
Latent Transition 是对隐状态转移函数的系统化设计，通过将隐状态空间分解为确定性 hidden state
- 高效的时间信息传播：
通过 RNN 传递长期依赖 - 显式的不确定性建模：
捕获局部随机性 - 后验-先验对比学习：通过 KL 项结构化地学习预测与观测的差异
核心数学：
- 生成模型：
- ELBO：
- KL 项度量后验与先验的差异，反映观测带来的新信息量
这一设计为 RSSM 提供了理论基础，是 Dreamer 系列算法的核心组件。
关键词
Latent Transition、后验分布、先验分布、状态空间模型、后验-先验对比、KL 散度、马尔可夫假设、条件独立性、变分推断、重参数化技巧、隐变量模型、信息高速公路、自由比特、KL annealing、确定性 hidden state、随机 latent variable