随机动力学建模

一句话定位：在确定性动力学的基础上引入随机隐变量，显式建模环境动态的不确定性和多模态性，使模型能够同时学习环境的”规律”与”噪声”。

前置依赖

确定性动力学建模（World Models 基础）
变分推断基础（VAE 原理）
贝叶斯深度学习基本概念
重参数化技巧（Reparameterization Trick）
KL 散度与 ELBO

核心思想

真实世界的物理系统并非完全确定：传感器存在噪声、执行器存在误差、环境的物理规律本身也可能有内在随机性（如布朗运动）。确定性动力学模型无法表达这种不确定性，而随机动力学建模通过引入随机隐变量 来解决这一问题。

核心假设：状态转移可以分解为两部分——确定性的”规律”和随机性的”噪声”。

数学表达为：

其中是从先验分布中采样的随机变量。

这一设计的直观理解是：

是确定性转移函数，给定隐变量后输出确定性的下一状态
是隐变量先验，通常设为标准正态分布
对求积分（边缘化）得到边缘转移分布，该分布可以是多峰的、非高斯的

随机动力学的核心优势：

表达多模态动态：同一状态-动作对可导致多种结果（如冰面、干燥路面）
建模感知不确定性：传感器噪声被吸收到随机变量中
防止误差累积爆炸：随机性可抵消部分单步误差，避免误差指数级增长

模型结构图

                    a_t           h_t
                     │             │
                     ▼             ▼
              ┌─────────────────────────┐
              │    随机转移函数         │
              │ p(h_{t+1}|h_t,a_t,z_t) │
              └─────────────────────────┘
                     │   ▲
                     │   │ 采样
                     ▼   │
              ┌─────────────────────────┐
              │   隐变量先验 p(z_t)     │
              │   (通常为 N(0,I))       │
              └─────────────────────────┘
                     │
                     ▼
                  h_{t+1} ~ q(h_{t+1}|...)

与确定性动力学的对比：

方面	确定性动力学	随机动力学
转移函数
随机性来源	无	隐变量采样
输出分布	Dirac	边缘化后的复杂分布
多模态表达	不支持	支持（通过混合）
训练难度	简单	需要变分推断

数学推导

A. 边缘转移分布的构建

给定当前隐状态和动作，我们希望建模状态转移。在随机动力学中，我们引入隐变量并假设：

$确定性转移先验$

其中：

：标准高斯先验
：给定隐变量的确定性输出分布（通常假设方差固定）

由于是确定性的（给定后输出点估计），上述积分将高斯先验”传输”到输出空间，使仍然保持高斯分布。

B. 变分下界（ELBO）推导

直接最大化观测的对数似然是困难的，因为隐变量的后验无法解析计算。引入变分近似，我们可以推导 ELBO：

完整推导：

从贝叶斯定理出发：

引入变分分布：

使用 Jensen 不等式（是凹函数，）：

$Extra close brace or missing open brace= \underbrace{\int q_\phi(z_{1:T} \mid o_{1:T}) \log p_\theta(o_{1:T}, z_{1:T}) dz_{1:T}}_{\text重构 + 隐变量先验}} - \underbrace{\int q_\phi(z_{1:T} \mid o_{1:T}) \log q_\phi(z_{1:T} \mid o_{1:T}) dz_{1:T}}_{H(q_\phi)}$

将联合分布分解：

时间因式分解：利用马尔可夫假设：

C. 重建项与KL项的直观理解

重构项（第一项）：

该期望表示：给定变分后验采样的隐变量，通过确定性转移函数计算，再通过解码器重构观测。期望操作允许服从后验分布，引入不确定性。

先验匹配项（第三项 KL 项）：

该 KL 项度量变分后验与先验之间的距离。最小化该 KL 项使后验分布向先验分布靠拢，防止后验分布过拟合观测数据。这实际上是正则化项。

D. 重参数化技巧

直接从中采样进行梯度下降是不可行的，因为采样操作通常是不可微的。重参数化技巧将随机性转移到另一个独立的随机变量：

典型的重参数化形式：

其中和是神经网络输出的均值和标准差，是逐元素乘法。

梯度计算：

其中是确定性的（因为是和的显式函数）。

E. posterior 与 prior 的对比

先验：

通常设为标准高斯
不依赖观测数据，是固定的正则化目标
作用：约束后验分布不会过度偏离先验

后验：

由神经网络参数化的变分分布
依赖当前观测序列，能够根据数据调整
在训练时用于采样，在推理时可用先验替代

训练时，我们希望后验分布能够”吸收”观测中的信息，因此后验应该比先验有更小的方差（更确定）；推理时，我们仅使用先验采样，使模型能够在未见过的状态上生成多样化的预测。

训练细节

A. 训练目标

完整的随机动力学训练目标（ELBO）： $重构损失正则化$

其中是 KL 项权重，超参数平衡重构与正则化。

B. 训练流程

对于每个训练 batch (o_{1:T}, a_{1:T-1}):

1. 编码观测序列:
   h_t = Encoder(o_t)  for t = 1 to T

2. 推断后验分布:
   q_phi(z_t) = AmortizedInference(h_t, h_{t-1})
   # 输出均值 mu_phi 和方差 sigma_phi^2

3. 重参数化采样:
   z_t = mu_phi + sigma_phi * epsilon,  epsilon ~ N(0,I)

4. 随机转移:
   h'_t = f_theta(h_{t-1}, a_{t-1}, z_t)
   # 基于采样隐变量计算新的隐状态

5. 重构观测:
   o'_t = Decoder(h'_t)

6. 计算 ELBO:
   recon_loss = MSE(o'_t, o_t)
   kl_loss = KL(q_phi(z_t) || N(0,I))
   loss = recon_loss - beta * kl_loss

7. 反向传播更新:
   loss.backward()
   optimizer.step()

C. 关键训练技巧

KL annealing：在训练初期逐渐增加（从 0.1 到 1.0），防止后验分布崩塌
自由比特（Free Bits）：当某个样本的 KL 项过小时，不对该样本的 KL 项进行惩罚，允许后验保留多样性
目标网络：变分分布的参数定期从主网络复制，延迟更新以稳定训练
重参数化梯度估计：使用 1-5 个样本的 MC 估计即可（方差较高时增加样本数）

推理/rollout/规划过程

A. 免费 rollout（Free Rollout）

在随机动力学模型的推理过程中，我们不使用后验分布，而是直接从先验采样：

h_0 = Encoder(o_0)
对于 t = 1 to T:
    z_t ~ N(0, I)                    # 从先验采样（非后验）
    h_{t+1} = f_theta(h_t, a_t, z_t) # 确定性转移
    a_{t+1} = policy(h_{t+1})         # 策略网络输出动作

这种设计使 rollout 过程具有随机性，能够生成多样化的想象轨迹，模拟真实环境中的不确定性。

B. 后验 vs 先验 Rollout

方面	先验 rollout	后验 rollout
来源
目的	生成多样化的想象轨迹	推理真实轨迹的隐变量
随机性	高（完全随机）	低（由观测决定）
应用场景	规划、策略学习	训练时的后验推断

C. 规划与模型预测控制

随机动力学模型的规划与确定性模型类似，但每次 rollout 具有随机性：

初始化当前隐状态
采样条独立 rollout 轨迹
计算每条轨迹的累积回报
选择期望回报最高的动作（或者使用 MPC 选最优动作）
执行实际动作，从真实环境获得奖励和下一观测

随机性使得规划结果具有鲁棒性——如果某条 trajectory 的高回报依赖于特定的随机实现，实际执行时可能无法复现。

优点与局限

优点

表达不确定性：显式建模感知噪声和动态随机性
多模态动态：通过混合高斯表达同一状态-动作对的不同结果
防止误差累积爆炸：随机性抵消部分单步误差
与变分框架结合：有坚实的统计推断理论基础
支持后验推断：可用于视频预测等需要推理隐变量的任务

局限

训练复杂：需要变分推断、KL 项计算、重参数化梯度估计
后验崩溃风险：KL 项权重过高时，后验分布可能退化为先验，隐变量失去意义
先验假设限制：高斯先验可能不足以表达复杂的隐变量分布
梯度估计高方差：MC 采样和重参数化引入额外的梯度方差
计算开销：每次前向传播需要采样和 KL 计算

与前后内容的衔接

前置内容：确定性动力学提供了随机动力学的基线；确定性 GRU 是随机转移函数的参考设计
后续发展：RSSM 在随机动力学基础上引入后验 vs 先验的结构化设计，同时建模确定性 hidden state 和随机 latent variable
核心区别：
- 随机动力学引入隐变量但没有结构化的时间转移
- RSSM 显式建模（确定性）和（随机性）的交互，实现后验-先验对比学习

可复现实现要点

A. PyTorch 实现框架

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class StochasticDynamics(nn.Module):
    def __init__(self, hidden_dim, action_dim, latent_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim
        
        # 确定性转移函数
        self.transition_net = nn.Sequential(
            nn.Linear(hidden_dim + action_dim + latent_dim, hidden_dim),
            nn.GRU(hidden_dim, hidden_dim)
        )
        
        # 变分推断网络（编码器）
        self.posterior_net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # 均值 + 对数方差
        )
        
    def forward(self, h_prev, action, z_t):
        """确定性转移函数 p(h_{t+1} | h_t, a_t, z_t)"""
        x = torch.cat([h_prev, action, z_t], dim=-1)
        h_next = self.transition_net(x)
        return h_next
    
    def infer_posterior(self, h_t):
        """推断后验分布 q(z_t | o_{1:T}) ~ N(mu, sigma)"""
        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 reparameterize(self, mu, sigma):
        """重参数化技巧 z = mu + sigma * epsilon"""
        epsilon = torch.randn_like(mu)
        z = mu + sigma * epsilon
        return z
    
    def kl_divergence(self, mu, sigma):
        """计算 KL(q(z) || N(0,I))"""
        # KL(N(mu, sigma) || N(0, I)) = -0.5 * (1 + log(sigma^2) - mu^2 - sigma^2)
        return -0.5 * torch.sum(
            1 + torch.log(sigma**2 + 1e-6) - mu**2 - sigma**2,
            dim=-1
        )

B. 完整训练循环

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)
    
    # 1. 编码观测序列
    h_seq = encoder(obs_seq)  # (batch, T, hidden_dim)
    
    total_recon_loss = 0
    total_kl_loss = 0
    
    # 2. 免费 rollout（训练时使用后验）
    for t in range(1, T):
        h_prev = h_seq[:, t-1]
        a_t = action_seq[:, t-1]
        
        # 推断后验（编码未来观测信息）
        mu, sigma = model.infer_posterior(h_seq[:, t])
        
        # 重参数化采样
        z_t = model.reparameterize(mu, sigma)
        
        # 随机转移
        h_pred = model.forward(h_prev, a_t, z_t)
        
        # 重构观测
        obs_recon = decoder(h_pred)
        recon_loss = F.mse_loss(obs_recon, obs_seq[:, t], reduction='sum')
        
        # KL 损失
        kl_loss = model.kl_divergence(mu, sigma)
        
        total_recon_loss += recon_loss
        total_kl_loss += kl_loss.mean()
    
    # 3. 总损失
    loss = total_recon_loss - beta * total_kl_loss
    
    # 4. 反向传播
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    
    return {
        'recon_loss': total_recon_loss.item(),
        'kl_loss': total_kl_loss.item(),
        'loss': loss.item()
    }

C. 关键实现细节

方差保证：使用 softplus 或 exp 将神经网络的 log 方差输出转换为正值的方差
数值稳定性：KL 计算中加上小常数防止
重参数化：epsilon 必须与计算图相连以支持梯度反向传播
采样次数：训练时每个时间步通常采样 1 次 MC 估计即可

章节摘要

随机动力学建模在确定性动力学的基础上引入随机隐变量，通过变分推断框架实现对环境动态不确定性的显式建模。

核心数学形式：

边缘转移：
ELBO 目标：
重参数化：

核心问题：随机动力学解决了确定性模型无法表达多模态动态和不确定性的问题，但训练需要变分推断和 MC 梯度估计。这引出了 RSSM 的结构化设计——同时建模确定性 hidden state 和随机 latent variable，并引入后验-先验对比学习机制。

关键词

随机动力学、变分推断、ELBO、重参数化技巧、KL 散度、后验分布、先验分布、隐变量模型、多模态动态、自由比特、KL annealing、模型预测控制、边缘化、梯度估计

Huajishe Notes

Explorer

2-随机动力学