变分推断在 RL 中的应用

0. 本讲导读

上一讲我们推出了证据下界（Evidence Lower Bound, ELBO）：对任意分布 $q_i(z)$，都有 $\log p(x_i) \ge \mathcal{L}_i(p, q_i)$，于是「最大化难算的 $\log p(x_i)$」被替换成「最大化好算的 $\mathcal{L}_i$」。 但那还只是一个数学恒等式，不是一个算法。真正落地时马上会撞上三个问题：

• $q_i$ 是每个数据点各配一份的分布，参数量随数据集线性增长，而且来了一个新样本还得重新跑一轮优化；
• 就算只优化一个 $q_i$，$\nabla_\phi \E_{z\sim q_\phi}[\cdot]$ 这个「对期望的分布本身求导」到底怎么估？
• 这套东西跟强化学习有什么关系？

本讲按三段回答。第一段把 $q_i$ 换成一张网络 $q_\phi(z|x)$——摊销变分推断（amortized variational inference）， 并给出两种梯度估计器：策略梯度（score function / REINFORCE）与重参数化技巧（reparameterization trick）， 逐条对比它们的适用范围与方差。第二段把这套工具组装成生成模型：VAE、条件 VAE、 以及把 VAE 铺到时间轴上得到的状态空间模型（state space model）——这正是后面 model-based RL 里 「从像素学潜动力学」的骨架；顺带说清 diffusion / flow matching 为什么可以看成一个层次化 VAE。 第三段换一个方向：不是「用变分推断做 RL 里的某个模块」，而是把整个最优控制问题本身 写成一个概率图模型上的推断问题——控制即推断（control as inference）。这一段会推出 soft 值迭代、 $\pi = \exp(A)$ 的玻尔兹曼策略、以及前向/后向消息相交产生的「橄榄球形」状态边缘分布。

1. 从变分下界到摊销推断

1.1 先把上一讲的下界重新摆一遍

设生成模型是 $p_\theta(x,z)=p_\theta(x|z)p(z)$，其中 $z$ 是隐变量（latent variable），$p(z)$ 通常取 $\mathcal{N}(0,I)$。 我们想做的是极大似然：$\theta \leftarrow \argmax_\theta \frac{1}{N}\sum_i \log p_\theta(x_i)$。 但 $\log p_\theta(x_i)=\log\int p_\theta(x_i|z)p(z)\,dz$ 里的积分对神经网络解码器是算不出来的。

引入任意分布 $q_i(z)$，用一次乘除再配 Jensen 不等式：

把最后一项拆开，就是幻灯片上那个写法：

其中 $\mathcal{H}(q)=-\E_{z\sim q}[\log q(z)]$ 是熵。这个不等式松了多少？把差值算出来：

注意上图算法里对 $\theta$ 的梯度为什么可以只保留一项：$\mathcal{L}_i$ 中只有 $\log p_\theta(x_i|z)$ 含 $\theta$， $\log p(z)$ 和 $\mathcal{H}(q_i)$ 都与 $\theta$ 无关，而且期望的分布 $q_i(z)$ 也不含 $\theta$， 所以 $\nabla_\theta\mathcal{L}_i=\E_{z\sim q_i}[\nabla_\theta\log p_\theta(x_i|z)]\approx\nabla_\theta\log p_\theta(x_i|z)$，单样本无偏。 麻烦全在 $q_i$ 那一侧。

1.2 「每个样本一份 $q_i$」为什么不可行

把这个参数量算算实数。假设数据集是 $N=10^6$ 张图，隐变量维度 $|z|=32$，那么每个样本要存 $\mu_i\in\R^{32}$ 和 $\sigma_i\in\R^{32}$， 一共 $10^6\times 64=6.4\times 10^7$ 个额外参数——已经和一个中等规模的解码器相当。更糟的是三件事：

1. 每个参数只被一个数据点的梯度更新。$\mu_i$ 只出现在 $\mathcal{L}_i$ 里，跑完一个 epoch 它只被更新了一次（或几次）， 统计效率极差；解码器 $\theta$ 却被 $N$ 个样本共同训练。二者的学习速度严重失配。
2. 数据集变大就崩。参数量 $O(N)$ 意味着换个更大的数据集就得重新分配显存。
3. 测试时无法推断。来了一个从没见过的 $x_{\text{new}}$，你没有对应的 $\mu_{\text{new}},\sigma_{\text{new}}$， 必须现场跑一个内层优化循环把它优化出来。对需要在线感知的机器人来说这是致命的。

1.3 摊销后的算法长什么样

形式上只是把下标 $i$ 换成了下标 $\phi$，但含义变了：现在 $\phi$ 是所有样本共享的一组参数， 每个 mini-batch 的每个样本都在更新它。总参数量变成 $|\theta|+|\phi|$，与 $N$ 无关。

目标函数写全：

$\theta$ 的梯度还是老样子。真正新的是 $\nabla_\phi$：$\phi$ 同时出现在被积函数外面（期望所依赖的分布）和 熵项里面。熵项好办——高斯熵有闭式：

直接对 $\sigma_\phi(x)$ 求导即可，无需采样。难的是第一项 $J(\phi)=\E_{z\sim q_\phi(z|x_i)}[r(x_i,z)]$， 其中我们把 $$ r(x_i,z)\;\triangleq\;\log p_\theta(x_i|z)+\log p(z) $$ 记成一个「奖励」。这个记号不是巧合——它和 RL 里 $\E_{a\sim\pi_\phi(a|s)}[r(s,a)]$ 的形状一模一样： $x_i$ 对应状态 $s$，$z$ 对应动作 $a$，$q_\phi(z|x)$ 对应策略 $\pi_\phi(a|s)$。所以我们已经会算它了。

2. 两种梯度估计器：策略梯度 vs 重参数化

2.1 直接搬策略梯度

推导就是第 4 讲的 log-derivative trick，一步不落地写出来：

这个估计器是无偏的，而且对离散 $z$ 一样成立（把积分换成求和即可）。那「What's wrong」在哪？方差。

2.2 重参数化技巧

为什么可以换基？因为 $\mathcal{N}(\mu,\sigma^2)$ 是位置-尺度族（location-scale family）： 若 $\epsilon\sim\mathcal{N}(0,1)$，则 $\mu+\sigma\epsilon\sim\mathcal{N}(\mu,\sigma^2)$。 于是对任意可测函数 $r$，$\E_{z\sim\mathcal{N}(\mu_\phi,\sigma_\phi^2)}[r(z)]=\E_{\epsilon\sim\mathcal{N}(0,1)}[r(\mu_\phi+\sigma_\phi\epsilon)]$。 这是恒等变形，不是近似。

换基之后，$\phi$ 只出现在被积函数里，普通链式法则就够了：

关键差别一目了然：这里出现了 $\partial r/\partial z$。估计器知道把 $z$ 往哪挪能让 $r$ 变大， 而不是像策略梯度那样只知道「这次采到的 $z$ 好不好」。这种估计器叫路径导数（pathwise derivative）估计器。 $r$ 里那个 $-5000$ 的公共偏移量在求导后直接消失了——它是常数，导数为零。这就是低方差的根源。

2.3 换个写法：重构项 + KL 项

这个写法在工程上比原写法优越，原因是方差进一步降低：KL 项从「用采样估的两项之差」变成了精确解析值， 只有重构项还带蒙特卡洛噪声。对 $q=\mathcal{N}(\mu,\diag(\sigma^2))$、$p=\mathcal{N}(0,I)$，逐维展开得到那个人人都会背的式子：

两项的解释也很清楚：$\E_{q}[\log p_\theta(x_i|z)]$ 是重构（reconstruction）项，逼着 $z$ 保留足够信息把 $x_i$ 还原； $-D_{\mathrm{KL}}(q_\phi(z|x_i)\|p(z))$ 是正则项，逼着每个样本的后验都别离先验太远。两者拉锯： 重构项想让不同 $x_i$ 的 $z$ 互相分开、$\sigma$ 尽量小（信息多）；KL 项想让它们全都缩回 $\mathcal{N}(0,I)$（信息少）。 最终得到的隐空间是被填满的、连续的——这正是采样能工作的前提，见 §3.2。

2.4 两种估计器的正面对比

2.5 最小实现：同一个下界的两种梯度

下面这段代码把两个估计器并排实现，用的是同一个编码器/解码器。跑起来你会看到： elbo_reparam 的梯度范数稳定，elbo_score_function 的梯度范数在批与批之间跳好几个数量级。

import math
import torch, torch.nn as nn, torch.nn.functional as F

D_X, D_Z, H = 784, 16, 400

class Encoder(nn.Module):      # q_phi(z|x) = N(mu_phi(x), diag(sigma_phi(x)^2))
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(nn.Linear(D_X, H), nn.ReLU(),
                                  nn.Linear(H, H),   nn.ReLU())
        self.mu, self.logstd = nn.Linear(H, D_Z), nn.Linear(H, D_Z)
    def forward(self, x):
        h = self.body(x)
        return self.mu(h), self.logstd(h).clamp(-5.0, 2.0)   # clamp 防止 sigma 爆炸/塌缩

class Decoder(nn.Module):      # p_theta(x|z)：伯努利像素，log p(x|z) 就是负 BCE
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(D_Z, H), nn.ReLU(),
                                 nn.Linear(H, H),   nn.ReLU(),
                                 nn.Linear(H, D_X))
    def forward(self, z):
        return self.net(z)                                    # 返回 logits

enc, dec = Encoder(), Decoder()
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=1e-3)

def log_px_given_z(x, z):
    logits = dec(z)
    return -F.binary_cross_entropy_with_logits(logits, x, reduction='none').sum(-1)

# ---------- 写法 A：重参数化 + 解析 KL（就是 VAE） ----------
def elbo_reparam(x):
    mu, logstd = enc(x)
    std = logstd.exp()
    eps = torch.randn_like(std)          # eps 与 phi 无关，梯度可以直接穿过去
    z   = mu + eps * std                 # <-- 重参数化
    rec = log_px_given_z(x, z)
    kl  = 0.5 * (std.pow(2) + mu.pow(2) - 1.0 - 2.0 * logstd).sum(-1)   # 解析式
    return (rec - kl).mean()

# ---------- 写法 B：同一个下界，phi 的梯度改用 policy gradient ----------
def surrogate_score_function(x):
    """返回值本身不是 ELBO，但它对 (theta, phi) 的梯度 = REINFORCE 估计器。"""
    mu, logstd = enc(x)
    std = logstd.exp()
    with torch.no_grad():                        # 采样这一步不回传：这正是问题所在
        z = mu + torch.randn_like(std) * std
    # log q_phi(z|x)，z 视作常数，梯度只经由 mu / logstd 流向 phi
    logq = (-0.5 * ((z - mu) / std).pow(2) - logstd
            - 0.5 * math.log(2 * math.pi)).sum(-1)
    rec  = log_px_given_z(x, z)                  # theta 的梯度从这里来
    r    = (rec - 0.5 * z.pow(2).sum(-1)).detach()   # r = log p(x|z) + log p(z)，当常数
    ent  = (logstd + 0.5 * math.log(2 * math.pi * math.e)).sum(-1)  # 高斯熵，解析可导
    return (rec + logq * r + ent).mean()

def grad_norm(loss):
    opt.zero_grad(); loss.backward()
    g = torch.cat([p.grad.flatten() for p in enc.parameters()])
    return g.norm().item()

# 训练循环（用写法 A）
for x in dataloader:                     # x: [B, 784]，取值 0/1
    loss = -elbo_reparam(x)
    opt.zero_grad(); loss.backward(); opt.step()

两点实现细节值得强调。第一，网络输出的是 logstd 而不是 std， 这样 std = logstd.exp() 天然为正，不需要额外约束；同时 KL 的解析式里 $\log\sigma_j^2=2\cdot\texttt{logstd}$ 直接可用。 第二，写法 B 里的 with torch.no_grad() 不是可选项——如果不 detach，PyTorch 会因为 z = mu + randn * std 而自动给你算出重参数化梯度，那就不是 score function 估计器了。 这也从侧面说明重参数化在现代框架里有多「免费」：你什么都不用做，它就发生了。

3. 变分自编码器与表示学习

3.1 VAE：把上一节的零件拼起来

没有新东西了——VAE 就是「摊销变分推断 + 重参数化 + 解析 KL」这三件事的名字。 唯一值得单独说的是解码器分布的选择：

• $p_\theta(x|z)=\mathcal{N}(\mu_\theta(z),\sigma^2 I)$ 且 $\sigma$ 固定 $\Rightarrow$ $\log p_\theta(x|z)=-\frac{1}{2\sigma^2}\|x-\mu_\theta(z)\|^2+\text{const}$， 重构项退化成加权的 L2 损失。此时 $\sigma$ 起的作用是调节重构项与 KL 项的相对权重：$\sigma$ 越小，重构越被看重。
• $p_\theta(x|z)$ 取伯努利（二值图像）$\Rightarrow$ 重构项是二元交叉熵。
• 幻灯片上写的是 $\mathcal{N}(\mu_\theta(z),\sigma_\theta(z))$，即方差也由网络输出。这更灵活，但训练容易不稳： 网络会发现把 $\sigma_\theta$ 推到 0 能让某些点的对数似然趋于 $+\infty$。实践中要给 $\log\sigma_\theta$ 加 clamp。

3.2 用 VAE：为什么祖先采样能工作

这一页的问题很值得认真答，因为它解释了 VAE 与普通自编码器的本质区别。

更严格一点的说法：把所有样本的后验平均起来得到聚合后验（aggregate posterior） $\bar{q}(z)=\frac{1}{N}\sum_i q_\phi(z|x_i)$。可以证明

其中 $I_{\bar q}(x;z)\ge 0$ 是数据与隐变量在联合分布 $\frac{1}{N}\sum_i\delta_{x_i}(x)q_\phi(z|x)$ 下的互信息。 所以 ELBO 里的 KL 惩罚同时压 $D_{\mathrm{KL}}(\bar q\|p)$（让聚合后验对齐先验，采样才有效） 和 $I(x;z)$（限制 $z$ 携带的信息量）。这两个作用的拉扯就是 VAE 的两大经典毛病： 「先验空洞」（prior hole，$\bar q$ 没铺满 $p$，采样落到空区）与「后验塌缩」（posterior collapse，$I(x;z)\to0$，$z$ 被忽略）。

3.3 第一个 RL 用法：把 $z$ 当状态

「Why is this a good idea?」这个问题 Levine 留给听众，答案至少有四条：

1. 降维。$Q(s,a)$ 直接吃 $84\times84\times3$ 的像素，需要一个大 CNN，而这个 CNN 的唯一训练信号是 TD 误差——一个标量、而且在稀疏奖励任务里几乎全是零。换成 $Q(z,a)$ 且 $z\in\R^{32}$，值函数网络可以只是个小 MLP。
2. 密集的无监督信号。重构损失对每个像素都有梯度，是极其密集的监督。它替 TD 误差把「视觉特征」这部分活干了。 这是「用自监督预训练替代端到端」在 RL 里的最早形态之一。
3. 光滑性与泛化。VAE 的隐空间是连续的、局部线性的（KL 正则的副产品）， 视觉上相似的状态在 $z$ 空间里也相近，值函数在其上更容易插值，样本效率更高。 左下那张图展示的正是这一点：沿单个隐维度移动，椅子的朝向 / 人脸的姿态连续且解耦地变化。
4. 能用先验数据（黄框那句）。训 VAE 只需要 状态，不需要动作、不需要奖励、不需要是本任务的数据。 所以你可以拿一大堆无标注的视频、别的任务的 replay buffer 先把表示学好，再上任务。

4. 条件模型与多模态模仿学习

4.1 条件 VAE 的下界

推导与无条件情形逐字对应，只要把每一处概率都加上 $|x_i$：

同样可以合并成重构 + KL 的形式：

两个容易搞错的地方：

4.2 为什么模仿学习非要多模态不可

把这件事说透：设两个模式各占一半，动作是一维的转向角，左绕 $a=-1$、右绕 $a=+1$。 用一个高斯 $\mathcal{N}(\mu,\sigma^2)$ 做最大似然拟合，闭式解是 $\mu=0$、$\sigma=1$。 $\mu=0$ 就是「直行撞树」——而且这个策略的似然还挺高，因为它把两个模式都覆盖住了。 最大似然对多模态数据的这种「取平均」行为不是训练不充分，而是模型类本身表达不了。

加了隐变量之后，模型类变成了混合分布： $p(a|s)=\int p(a|s,z)p(z)\,dz$。虽然每个 $p(a|s,z)$ 还是单峰高斯，但对 $z$ 积分之后 可以逼近任意复杂的多峰分布（这就是无限混合高斯）。解码器要学的是「给定 $z$ 的这个取值，我该走哪条路」， 而 $z$ 扮演的是没被观测到的意图 / 风格变量：示范者当时打算左绕还是右绕，这个信息在数据里是缺失的， 所以它天然应该是隐变量。

ACT 把两件事叠在一起，值得分开理解：

• 动作分块（action chunking）：一次预测未来 $k$ 步动作而不是 1 步。这缓解的是模仿学习里的复合误差 和人类示范的时间抖动；从概率角度看，它把建模对象从 $p(a_t|s_t)$ 换成了 $p(a_{t:t+k}|s_t)$， 后者的多模态性更强（整段轨迹的模式比单步动作的模式更明显），所以更需要隐变量。
• CVAE 的隐变量 $z$：吸收「这次示范是哪种风格 / 哪个意图」的信息。训练时编码器能看到真实动作序列， 所以它能把「这一段属于哪个模式」编进 $z$，解码器于是不必在模式间平均。 推理时直接取 $z=0$（先验均值）或从 $\mathcal{N}(0,I)$ 采样。

4.3 最小实现：带隐变量的多模态 BC 策略

import torch, torch.nn as nn

D_S, D_A, D_Z, H = 32, 4, 8, 256

class LatentPolicy(nn.Module):
    """条件 VAE 形式的行为克隆策略：pi(a|s) = int p(a|s,z) p(z) dz"""
    def __init__(self):
        super().__init__()
        # q_phi(z | s, a)：编码器必须看到动作，否则无法分辨模式
        self.enc = nn.Sequential(nn.Linear(D_S + D_A, H), nn.ReLU(),
                                 nn.Linear(H, 2 * D_Z))
        # p_theta(a | s, z)：解码器只看 s 和 z
        self.dec = nn.Sequential(nn.Linear(D_S + D_Z, H), nn.ReLU(),
                                 nn.Linear(H, 2 * D_A))

    def loss(self, s, a, beta=1.0):
        mu, logstd = self.enc(torch.cat([s, a], -1)).chunk(2, -1)
        logstd = logstd.clamp(-5.0, 2.0)
        z = mu + torch.randn_like(mu) * logstd.exp()          # 重参数化
        a_mu, a_logstd = self.dec(torch.cat([s, z], -1)).chunk(2, -1)
        a_logstd = a_logstd.clamp(-5.0, 2.0)
        # 重构项：高斯对数似然
        rec = (-0.5 * ((a - a_mu) / a_logstd.exp()).pow(2) - a_logstd).sum(-1)
        # KL(q || N(0,I)) 解析式
        kl  = 0.5 * ((2 * logstd).exp() + mu.pow(2) - 1.0 - 2 * logstd).sum(-1)
        return (-rec + beta * kl).mean()                      # 最小化 -ELBO

    @torch.no_grad()
    def act(self, s, stochastic=True):
        z = torch.randn(s.shape[0], D_Z) if stochastic else torch.zeros(s.shape[0], D_Z)
        a_mu, a_logstd = self.dec(torch.cat([s, z], -1)).chunk(2, -1)
        return a_mu + (torch.randn_like(a_mu) * a_logstd.clamp(-5.0, 2.0).exp()
                       if stochastic else 0.0)

注意 act 里的两次随机：一次是采 $z$（选模式：左绕还是右绕）， 一次是在选定模式内采动作（模式内的抖动）。这两级随机性的分工正是隐变量策略的价值所在。 如果只想要「确定地选一个模式」，把 $z$ 固定成 0 即可——ACT 在真机上就是这么做的。 $\beta$ 是 KL 权重（$\beta$-VAE 的那个 $\beta$）：$\beta$ 太大会后验塌缩、退回单峰策略；太小则 $z$ 会把动作原样记住， 而测试时你从先验采的 $z$ 与训练时的分布对不上，输出乱掉。

5. 与 diffusion / flow matching 的关系

把这张图翻译成一句话：diffusion 就是隐变量为一整条加噪轨迹 $a_{\Delta\tau:1}$ 的条件 VAE， 而且它的编码器不用学。逐条对比：

关于「层次化 VAE」这个说法：普通 VAE 是 $z\to x$ 两层；层次化 VAE 是 $z_L\to z_{L-1}\to\dots\to z_1\to x$， 每一层都是一个条件生成步，ELBO 里有 $L$ 个 KL 项。diffusion 正是这个结构， 只不过 $L$ 特别大（上千）、各层维度相同、且推断方向 $q$ 被写死。至于 flow matching， 它把离散的去噪步换成一个连续时间的常微分方程 $\frac{da_\tau}{d\tau}=v_\theta(a_\tau,\tau,s)$， 训练目标不再是显式的 ELBO 而是回归一个已知的条件速度场； 但「用一条固定的、把数据搅成噪声的路径来定义隐变量」这个骨架是一样的。 幻灯片右下那条注记提醒的是索引约定：DDPM 习惯 $\tau=1$ 是噪声、$\tau=0$ 是数据， flow matching 文献常反过来，读论文时别弄混。

6. 状态空间模型：把 VAE 铺到时间轴上

6.1 三个组件的角色

先统一记号。$o_t$ 是观测（可能只是环境的一部分，比如一帧图像），$z_t$ 是隐状态 （我们希望它满足马尔可夫性，扮演真正的 $s_t$），$a_t$ 是动作。整个模型是

6.2 序列 ELBO 的完整推导

直接把 VAE 的下界套上去，然后利用两个乘积形式的因子分解拆开。

这个式子每一项都有直白的含义：

• 第一项：$z_t$ 必须留住足够信息把 $o_t$ 画回来。防止 $z$ 退化成常数。
• 第三项：编码器在 $t+1$ 时刻推断出的状态，必须跟动力学从 $z_t$ 预测出来的一致。 这一项同时训练了两边：既逼动力学学准（不然预测不上编码器给的答案）， 又逼编码器输出可预测的表示（不然 KL 罚不下去）。这是整个目标里最重要的一项， 也是「学出来的 $z$ 为什么会近似满足马尔可夫性」的原因。
• 第二项：起始状态对齐标准正态，只是个锚。

6.3 编码器的因子分解：filtering、smoothing 还是单帧？

幻灯片给的是 $q_\phi(z|o)=\prod_t q_\phi(z_t|o_{1:t})$，这叫滤波（filtering）后验。但这不是唯一选择， 不同选择在「精度 / 计算量 / 能否在线用」之间取舍：

为什么滤波是甜点？因为你在控制时本来就只有 $o_{1:t}$。如果训练时用了平滑后验， 部署时就得再单独搞一个滤波器，两者不一致会掉性能。而滤波形式的 $q_\phi(z_t|o_{1:t})$ 用一个 RNN 实现： 维护隐藏态 $h_t=\mathrm{RNN}(h_{t-1},[\,\mathrm{enc}(o_t),a_{t-1}\,])$，然后 $q_\phi(z_t|o_{1:t})=\mathcal{N}(\mu_\phi(h_t),\sigma_\phi(h_t))$。 一次前向就能采出整条 $z_{1:T}$，重参数化照常穿过去。

6.4 最小实现：一个能跑的潜动力学模型

import torch, torch.nn as nn, torch.distributions as D

D_O, D_A, D_Z, H = 64, 4, 32, 256

def gauss(params, lo=-5.0, hi=2.0):
    mu, logstd = params.chunk(2, dim=-1)
    return D.Normal(mu, logstd.clamp(lo, hi).exp())

class LatentDynamics(nn.Module):
    def __init__(self):
        super().__init__()
        self.obs_enc = nn.Sequential(nn.Linear(D_O, H), nn.ReLU(), nn.Linear(H, H))
        self.rnn     = nn.GRUCell(H + D_A, H)                 # 把 o_{1:t}, a_{1:t-1} 累积进 h_t
        self.q_head  = nn.Linear(H, 2 * D_Z)                  # 编码器 q(z_t | o_{1:t})  <- 信念状态
        self.dyn     = nn.Sequential(nn.Linear(D_Z + D_A, H), nn.ReLU(),
                                     nn.Linear(H, 2 * D_Z))   # 先验 p(z_{t+1} | z_t, a_t) <- 动力学
        self.dec     = nn.Sequential(nn.Linear(D_Z, H), nn.ReLU(),
                                     nn.Linear(H, 2 * D_O))   # 解码器 p(o_t | z_t)

    def elbo(self, obs, act, free_nats=1.0):
        """obs: [B,T,D_O]；act: [B,T,D_A]，act[:,t] 是在 t 时刻采取的动作。"""
        B, T, _ = obs.shape
        h = obs.new_zeros(B, H)
        z_prev, a_prev = None, obs.new_zeros(B, D_A)
        rec_total = obs.new_zeros(B)
        kl_total  = obs.new_zeros(B)

        for t in range(T):
            # --- 编码器：把 o_t 和上一步动作并进 RNN，得到 q(z_t | o_{1:t}) ---
            h   = self.rnn(torch.cat([self.obs_enc(obs[:, t]), a_prev], dim=-1), h)
            q_t = gauss(self.q_head(h))
            z_t = q_t.rsample()                         # rsample = 重参数化采样

            # --- 解码器：重构 o_t ---
            rec_total = rec_total + gauss(self.dec(z_t)).log_prob(obs[:, t]).sum(-1)

            # --- 先验：t=0 用 N(0,I)，之后用学出来的转移 ---
            if z_prev is None:
                p_t = D.Normal(torch.zeros_like(z_t), torch.ones_like(z_t))
            else:
                p_t = gauss(self.dyn(torch.cat([z_prev, a_prev], dim=-1)))

            kl_t = D.kl_divergence(q_t, p_t).sum(-1)
            kl_total = kl_total + kl_t.clamp(min=free_nats)   # free bits，防后验塌缩

            z_prev, a_prev = z_t, act[:, t]

        return (rec_total - kl_total).mean()

model = LatentDynamics()
opt   = torch.optim.Adam(model.parameters(), lr=3e-4)

for obs, act in dataloader:                      # obs:[B,T,D_O]  act:[B,T,D_A]
    loss = -model.elbo(obs, act)
    opt.zero_grad(); loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 100.0)
    opt.step()

# ---- 部署：用编码器在线维护信念，然后只靠 dyn 在潜空间里“想象” ----
@torch.no_grad()
def imagine(model, h, a_seq):
    """h 由在线跑 RNN 得到（已吸收 o_{1:t}）；a_seq: [B,K,D_A] 是候选动作序列。"""
    z = gauss(model.q_head(h)).mean
    traj = []
    for k in range(a_seq.shape[1]):
        z = gauss(model.dyn(torch.cat([z, a_seq[:, k]], dim=-1))).mean
        traj.append(z)
    return torch.stack(traj, dim=1)              # [B,K,D_Z]，喂给 reward 头或 Q 网络

几处工程要点：（1）rsample() 而不是 sample()——前者走重参数化、可回传梯度， 后者切断梯度，写错了模型会一动不动；（2）kl_t.clamp(min=free_nats) 就是 free bits， 让 KL 在低于 1 nat 时不再被压，给编码器留出携带信息的余地； （3）clip_grad_norm_ 是必须的，序列模型的 BPTT 梯度尖峰很常见； （4）imagine 里完全不需要解码器和真实观测——这就是「在潜空间里做规划」的含义， 把它接上 CEM / MPPI 就是 PlaNet，接上在潜空间里学的 actor-critic 就是 Dreamer。

7. 控制即推断：动机与概率图模型

前面六节是「用变分推断做 RL 里的某个模块」。现在换一个更激进的问题： 能不能把最优控制问题本身，整个写成一次概率推断？ 如果能，我们就可以把 HMM、卡尔曼滤波、消息传递这些成熟的推断算法直接搬过来解控制问题。

7.1 出发点：用最优控制解释人的行为

这一页的立场是：人和动物的行为看起来是在优化某个目标，所以我们可以假设他们是最优控制器， 然后去拟合那个让观测到的行为「最合理」的奖励函数 $r$。这就是逆强化学习（inverse RL）的基本设定。 但立刻就有一个大问题。

7.2 数据不是最优的

严格的 $\argmax$ 模型无法解释这张图。它给「最优轨迹」概率 1，给其余所有轨迹概率 0， 于是只要数据里出现任何一点偏差，模型的似然就是 $-\infty$。可现实是：

• 人的行为是随机的——同一个任务做十次，十条轨迹都不一样；
• 但这些随机不是均匀的：在路径中段绕一点弯几乎无所谓（反正最后能到），在终点抓偏 5 厘米就完全失败。 某些错误比另一些更要紧；
• 而且好行为依然是最可能的——不是唯一可能，是最可能。

我们需要的是一个模型，它给轨迹赋予的概率随奖励单调递增，好轨迹概率高、差轨迹概率低但非零， 而且这个「衰减速度」应该由奖励差距决定。

7.3 决策的概率图模型：最优性变量

这个构造的巧妙之处：$\mathcal{O}_t\in\{0,1\}$ 是一个人造的二值变量，读作「第 $t$ 步的行为是最优的」。 它的取值我们永远设为 1 并当作已观测，于是「求最优行为」就变成了标准的推断问题 「已知 $\mathcal{O}_{1:T}=1$，问 $\tau$ 的后验是什么」。

现在检查它是否满足 §7.2 的三条要求。假设动力学是确定的，那么 $p(\tau)$ 在所有可行轨迹上是常数 （每条可行轨迹的转移概率都是 1，不可行的是 0），于是

这是一个玻尔兹曼分布（Boltzmann distribution），能量就是负的总回报。三条要求全部满足：

7.4 这个模型为什么值得做

幻灯片列了三条理由，每条都对应本课后面的一整块内容：

7.5 推断 = 规划：三个要问的问题

这三个量的含义值得先建立直觉：

• $\beta_t(s_t,a_t)=p(\mathcal{O}_{t:T}|s_t,a_t)$：在 $s_t$ 做了 $a_t$ 之后，从 $t$ 到 $T$ 全都最优的概率。 它只看未来，不看过去——所以它是「$Q$ 函数式」的量。
• $\pi(a_t|s_t)=p(a_t|s_t,\mathcal{O}_{1:T})$：已知全程最优，在 $s_t$ 时会做什么动作。这就是我们要的策略。
• $\alpha_t(s_t)=p(s_t|\mathcal{O}_{1:t-1})$：假设前 $t-1$ 步都最优，现在会在哪些状态。它只看过去。 两者相乘就得到完整的状态边缘 $p(s_t|\mathcal{O}_{1:T})$。

这个「后向消息 × 前向消息」的结构和 HMM 的前向-后向算法（forward-backward）、 卡尔曼滤波与平滑是完全同构的——这正是本讲最后一页要强调的「control = inference (similar to HMM, EKF, etc.)」。

8. 后向消息：soft 值迭代的由来

8.1 后向递推的推导

8.2 取对数：soft 值迭代现身

对两条递推取对数。第二条（$\beta_t(s_t)=\int\beta_t(s_t,a_t)p(a_t|s_t)da_t$，均匀先验下 $p(a|s)$ 是常数）给出

这就是 log-sum-exp，也叫 soft max（注意不是 softmax 函数，是「软化的 max」）。为什么它近似 $\max$？

第一条递推取对数（用 $p(\mathcal{O}_t|s_t,a_t)=\exp r$）给出

和标准值迭代 $Q(s,a)\leftarrow r(s,a)+\gamma\E[V(s')]$ 并排一看，差别就两处： $\max_a$ 变成了 $\log\int\exp$，$\E_{s'}[V]$ 变成了 $\log\E_{s'}[\exp V]$。第一处是我们想要的（它带来了随机策略）， 第二处是个大问题。

8.3 「乐观转移」：这个模型的真正缺陷

Levine 在幻灯片上直接写了 “not a good idea!”，并给出两种目前可以接受的处理：

• 确定性动力学：此时 $\E_{s'}[\exp V]=\exp V(s_{t+1})$，$\log$ 与 $\exp$ 抵消，得到 $Q_t(s_t,a_t)=r(s_t,a_t)+V_{t+1}(s_{t+1})$——和标准 Bellman 备份完全一致。 所以在确定性环境里，control-as-inference 除了 $V$ 的 soft max 之外没有任何偏差。
• 随机动力学：留到下一讲用变分推断来修——办法是不去求精确后验 $p(\tau|\mathcal{O}_{1:T})$， 而是在一个约束了转移必须等于真实转移的分布族里做变分近似 $q(\tau)=p(s_1)\prod_t p(s_{t+1}|s_t,a_t)q(a_t|s_t)$。这样一来智能体就不能通过「改变对动力学的信念」来作弊， 只能改策略。最大化对应的 ELBO 得到的正是最大熵 RL 的目标 $\sum_t\E[r(s_t,a_t)+\mathcal{H}(q(a_t|s_t))]$，也就是 Soft Actor-Critic 的出发点。

注意这里的语义是概率而不是回报：$\beta_t\in[0,1]$（在 $r\le0$ 的约定下）， 它的对数才是我们熟悉的 $Q$。这解释了为什么值函数会以 $\log$ 的形式出现—— 「累加奖励」在概率空间里是「连乘 $\exp(r)$」，取对数就变回累加。

9. 动作先验：为什么可以放心假设均匀

推导只有一行代数，但结论很有分量。带一般先验时，$V$ 的更新是

令 $\tilde{Q}(s_t,a_t)\triangleq Q(s_t,a_t)+\log p(a_t|s_t)=\underbrace{r(s_t,a_t)+\log p(a_t|s_t)}_{\tilde r(s_t,a_t)}+\log\E[\exp V(s_{t+1})]$， 则 $V(s_t)=\log\int\exp(\tilde Q(s_t,a_t))\,da_t$——形式与均匀先验时一模一样， 只是奖励换成了 $\tilde r=r+\log p(a|s)$。所以：任何带动作先验的问题，都等价于一个奖励被修改过、 动作先验为均匀的问题。

10. 策略计算：$\pi=\exp(A)$

10.1 从后验中读出策略

这个式子有非常朴素的解释：分子是「做了 $a_t$ 之后全程最优的概率」， 分母是「在 $s_t$ 什么都还没决定时全程最优的概率」，比值就是这个动作对『还能最优』的相对贡献。 贡献大的动作被选中的概率高。

10.2 换成值函数写法

代换一步即可：$\beta_t(s_t,a_t)=\exp(Q_t)$、$\beta_t(s_t)=\exp(V_t)$，所以

顺带验证它自动归一化：$\int\exp(Q_t-V_t)da_t=e^{-V_t}\int e^{Q_t}da_t=e^{-V_t}\cdot e^{V_t}=1$。 这不是巧合——$V_t=\log\int\exp Q_t$ 这个 soft max 的定义恰恰就是那个归一化常数（配分函数）的对数。

10.3 温度：从贪心到均匀的一根旋钮

温度是从哪来的？把最优性变量的定义改成 $p(\mathcal{O}_t|s_t,a_t)=\exp\big(r(s_t,a_t)/\alpha\big)$， 所有推导原样成立，只是各处的 $r$ 换成 $r/\alpha$，于是 $Q,V$ 都被 $1/\alpha$ 缩放。 换句话说 $\alpha$ 就是「你认为示范者/智能体有多接近最优」的刻度。

四条性质逐条对上：更好的动作概率更大（$\exp$ 单调）；随机打破平局 （两个 $Q$ 完全相等时各 50%，而 $\argmax$ 要靠实现细节任意挑一个）； 与玻尔兹曼探索同构（$\pi\propto\exp(Q/\alpha)$ 正是 Boltzmann / softmax 探索的公式， 只不过这里它是从推断里推出来的，而不是手工加的探索启发式）； 温度降低趋于贪心（上表第三行）。

10.4 最小实现：表格式 soft 值迭代

import numpy as np
from scipy.special import logsumexp

def soft_value_iteration(r, P, T, alpha=1.0):
    """r: [S,A] 奖励；P: [S,A,S] 转移概率；T: 时域长度；alpha: 温度。
    返回每个时刻的 Q、V 和策略 pi。"""
    S, A = r.shape
    V  = np.zeros((T + 1, S))          # V[T] = 0 作为边界条件
    Q  = np.zeros((T, S, A))
    logP = np.log(P + 1e-300)          # 防 log(0)
    for t in range(T - 1, -1, -1):
        # log E_{s'}[exp V_{t+1}(s')] = logsumexp_{s'}( log P(s'|s,a) + V_{t+1}(s') )
        soft_next = logsumexp(logP + V[t + 1][None, None, :], axis=2)   # [S,A]
        Q[t] = r / alpha + soft_next                 # <- “乐观”转移就藏在 soft_next 里
        V[t] = logsumexp(Q[t], axis=1)               # V_t(s) = log sum_a exp Q_t(s,a)
    pi = np.exp(Q - V[:T, :, None])                  # pi(a|s) = exp(Q - V) = exp(A)
    return Q, V[:T], pi

def hard_value_iteration(r, P, T, gamma=1.0):
    """对照组：标准（hard）有限时域值迭代。"""
    S, A = r.shape
    V = np.zeros((T + 1, S)); Q = np.zeros((T, S, A))
    for t in range(T - 1, -1, -1):
        Q[t] = r + gamma * P @ V[t + 1]              # E_{s'}[V]，不是 log E[exp V]
        V[t] = Q[t].max(axis=1)
    return Q, V[:T]

把 alpha 从 1 调到 0.05 再跑 soft_value_iteration，你会看到 pi 逐渐变成 one-hot， 而 Q * alpha 收敛到 hard_value_iteration 的 Q——但只在确定性转移下。 在随机转移的 MDP 上，即使 $\alpha\to0$，soft_next 也会因为 $\log\E[\exp]$ 而收敛到 $\max_{s'}V(s')$ 而不是 $\E_{s'}[V(s')]$，两者的差就是 §8.3 那个乐观偏差。 自己造一个「有 1% 概率通向大奖」的动作跑一下，能非常直观地看到这个模型的病灶。

11. 前向消息与状态边缘分布

11.1 前向递推

11.2 两条消息相乘 = 状态边缘

最后一步是把两个方向的信息合起来：

第一处用的是条件独立：给定 $s_t$，未来的最优性 $\mathcal{O}_{t:T}$ 与过去的 $\mathcal{O}_{1:t-1}$ 独立。 最后 $p(\mathcal{O}_{1:t-1})$ 与 $s_t$ 无关，作为常数丢掉。

为什么是橄榄球（或者说梭形）？

• $\alpha_t(s_t)$ 只看过去：从一个确定的起点出发，随着时间推移，不确定性单调增长——一个向右张开的锥。
• $\beta_t(s_t)$ 只看未来：离终点越近，「还能救回来」的状态集合越小，越接近终点约束越紧——一个向左张开的锥。
• 二者相乘：起点附近被 $\alpha$ 卡死（你就在那儿），终点附近被 $\beta$ 卡死（必须能到目标）， 中间没人管，所以方差最大。

11.3 最小实现：前向消息与状态边缘

import numpy as np
from scipy.special import logsumexp

def forward_messages(p1, P, r, T):
    """alpha[t, s] = p(s_t | O_{1:t-1})，与 soft_value_iteration 用同一套 (P, r)。"""
    S, A = r.shape
    alpha = np.zeros((T, S))
    alpha[0] = p1                                        # alpha_1(s) = p(s_1)
    # p(a|s,O) 正比于 exp(r(s,a)) * p(a|s)，均匀先验下就是对 a 做 softmax
    pa_given_sO = np.exp(r - logsumexp(r, axis=1, keepdims=True))     # [S,A]
    for t in range(1, T):
        # alpha_t(s') = sum_{s,a} P(s'|s,a) p(a|s,O) alpha_{t-1}(s)
        alpha[t] = np.einsum('s,sa,sax->x', alpha[t - 1], pa_given_sO, P)
        alpha[t] /= alpha[t].sum()                       # 归一化（前面推导里丢掉的常数）
    return alpha

def state_marginals(alpha, V):
    """p(s_t | O_{1:T}) 正比于 beta_t(s_t) * alpha_t(s_t)，其中 beta_t(s) = exp(V_t(s))。"""
    m = alpha * np.exp(V - V.max(axis=1, keepdims=True))  # 减最大值防溢出
    return m / m.sum(axis=1, keepdims=True)

# 用法：先后向、再前向、最后相乘
Q, V, pi = soft_value_iteration(r, P, T, alpha=1.0)
al       = forward_messages(p1, P, r, T)
marg     = state_marginals(al, V)
# marg[t] 就是那串绿色椭圆；把 marg 的熵按 t 画出来，会看到中间高、两端低
ent = -(marg * np.log(marg + 1e-12)).sum(axis=1)

把最后那个 ent（每个时刻状态边缘分布的熵）画成曲线，就是 Li & Todorov 那条实验曲线的 理论对应物：起点熵接近 0（初始状态确定），终点熵很低（目标约束紧），中间最高。

本讲小结

速查表：本讲所有关键公式

要点清单

• 摊销把逐样本变分参数换成一张编码器网络，参数量与数据量脱钩，测试时一次前向即可推断； 代价是引入摊销间隙。
• 重参数化之所以低方差，是因为它用了 $\partial r/\partial z$ 这个一阶信息， 而且 $r$ 的常数偏移会被求导消掉。策略梯度只用零阶信息，$r$ 的绝对量级会直接放大方差。 连续隐变量优先用重参数化，离散隐变量只能用策略梯度或连续松弛。
• VAE 能采样的原因是 KL 项让每个样本在隐空间占一团有体积的云、并把这些云拉向先验， 从而把隐空间连续地填满。这也带来两个孪生病症：先验空洞与后验塌缩。
• 隐变量策略是模仿学习处理多模态的标准手段：$p(a|s)=\int p(a|s,z)p(z)dz$， $z$ 承担「示范者当时的意图」这个缺失变量。不用它，最大似然会在模式之间取平均——直接撞树。
• diffusion 是编码器写死的层次化 VAE：省掉了学 $q$ 的全部麻烦，KL 项解析化， 代价是采样要跑很多步。
• 状态空间模型把 VAE 铺到时间轴：先验成了要学的动力学，编码器成了信念状态。 序列 ELBO 里那个逐步 KL 是核心——它同时训练动力学的准确性和表示的可预测性。 编码器用滤波形式 $q(z_t|o_{1:t})$ 是因为控制时只能看到过去。
• control as inference 用最优性变量把「求最优策略」变成「求后验」。 它自然地建模了次优行为、给出了随机策略的第一性解释、并把推断算法引入控制。
• 后向传递就是 soft 值迭代：$\max_a\to\log\int\exp$。 但 $\E_{s'}[V]\to\log\E_{s'}[\exp V]$ 是个 bug 不是 feature——它让智能体变成风险偏好者。 确定性动力学下没这个问题；随机动力学下要靠下一讲的变分近似来修。
• $\pi=\exp(A)$ 就是玻尔兹曼探索，只是这里它是推出来的而非手工加的。 温度 $\alpha$ 在贪心与均匀之间连续插值。
• 前向 × 后向 = 状态边缘，形状是两端窄中间宽的橄榄球， 与 Li & Todorov 2006 测到的人类运动方差曲线吻合。

延伸阅读

变分推断与 VAE 的基础

• Auto-Encoding Variational Bayes (Kingma & Welling, 2013) — 本讲第一、二段的原始论文，摊销推断 + 重参数化的出处，短且好读，值得逐行推一遍。
• Stochastic Backpropagation and Approximate Inference in Deep Generative Models (Rezende et al., 2014) — 与上一篇同期独立提出同一套方法，对梯度估计器的讨论更细致。
• Neural Variational Inference and Learning in Belief Networks (Mnih & Gregor, 2014) — 离散隐变量只能用 score function 时，怎么做方差缩减；理解「为什么大家都想用重参数化」的反面教材。
• Categorical Reparameterization with Gumbel-Softmax (Jang et al., 2016) 与 The Concrete Distribution (Maddison et al., 2016) — 把重参数化推广到离散隐变量的两篇同期工作。
• Variational Inference with Normalizing Flows (Rezende & Mohamed, 2015) — 当对角高斯的 $q$ 表达力不够时怎么办：用可逆变换堆出更灵活的后验。
• β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework (Higgins et al., 2017) — 幻灯片里椅子与人脸那组隐维度遍历图的来源，也是「调大 KL 权重换解耦表示」这个做法的出处。

表示学习与隐变量策略

• Learning Latent Plans from Play (Lynch et al., 2019) — 幻灯片右侧那张框架图的论文：plan recognition = $q_\phi$，plan proposal = $p(z|x)$，一个条件 VAE 撑起一个通用策略。
• Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware (Zhao et al., 2023) — 即 ACT / ALOHA。CVAE + 动作分块 + transformer，是当前机器人模仿学习最常被复现的基线之一。
• Diffusion Policy (Chi et al., 2023) — 把第 5 节讲的 diffusion 直接当成策略的动作分布，多模态表达力显著强于 CVAE。
• Denoising Diffusion Probabilistic Models (Ho et al., 2020) — 把 diffusion 的变分下界化简成 $\|\epsilon-\epsilon_\theta\|^2$ 的关键一步，正好对应第 5 节那张表。
• Flow Matching for Generative Modeling (Lipman et al., 2022) — 连续时间版本，训练目标是回归条件速度场；理解 slides 上那条「DDPM 与 flow matching 时间轴方向相反」的注记。

状态空间模型 / 潜动力学

• Embed to Control (Watter et al., 2015) — 最早把「学一个隐空间使动力学在其中局部线性」这件事做通的工作，之后可以直接上 iLQR。
• Deep Kalman Filters (Krishnan et al., 2015) — 序列 VAE 里编码器因子分解（filtering / smoothing）的系统讨论。
• Learning Latent Dynamics for Planning from Pixels (Hafner et al., 2018) — 即 PlaNet。第 6 节代码的完整工业级版本，含 RSSM 的「确定性 + 随机」双通道设计。
• Dream to Control (Hafner et al., 2019) 与 Mastering Diverse Domains through World Models (Hafner et al., 2023) — Dreamer / DreamerV3：在学出来的隐空间里直接跑 actor-critic，是「imagine 函数」的终点形态。
• Stochastic Latent Actor-Critic (Lee et al., 2019) — 把本讲两半（潜状态模型 + 最大熵 RL）真正拼在一起的工作。
• SOLAR: Deep Structured Representations for Model-Based RL (Zhang et al., 2018) — 在隐空间里学线性高斯动力学，从而可以套用 LQR 类的高效控制器。

控制即推断

• Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review (Levine, 2018) — 本讲第三段的完整书面版，包括本讲略过的随机动力学变分修正、与最大熵 RL 的等价性证明。 如果只读一篇，读这篇。
• Maximum Entropy Inverse Reinforcement Learning (Ziebart et al., 2008) — $p(\tau)\propto\exp(\sum_t r)$ 这个模型在逆 RL 里的首次系统使用，幻灯片上出租车路线图的来源。
• An Iterative Optimal Control and Estimation Design for Nonlinear Stochastic Systems (Li & Todorov, 2006) — 橄榄球形状态边缘那条实验曲线的来源，也是「最小干预原则」的代表工作。
• Reinforcement Learning with Deep Energy-Based Policies (Haarnoja et al., 2017) — soft Q-learning：把本讲的 soft 值迭代变成能在连续动作、深度网络上跑的算法。
• Soft Actor-Critic (Haarnoja et al., 2018) — 下一讲的终点。最大熵目标 + actor-critic，目前连续控制上最稳的 off-policy 算法之一。