变分自编码器 (VAE)

机器人经常需要生成东西，而不只是给输入打标签：想象下一帧相机画面会是什么样（用来在“脑子里”试动作），为一个没见过的物体提出几个可行的抓取，或者在“从左边绕”和“从右边绕”都对的时候，给出一个具体的动作而不是两者的平均。这些都要求一个能从中采样的概率模型 \(p(\mathbf x)\)，而且 \(\mathbf x\) 往往是高维的：一张图像、一段动作序列。

18.2 节的自编码器把图像压成一个低维编码 \(\mathbf z\) 再还原，学到的编码很有用，但它没有告诉我们编码服从什么分布：编码散落在隐空间里形状任意、中间有大片空洞，随便取一个 \(\mathbf z\) 解码，多半不像任何真实图像。VAE（Kingma 与 Welling，2014；Rezende 等，2014 同时提出）的做法是把整件事写成一个隐变量模型：先从一个简单的先验 \(p(\mathbf z)=\mathcal N(\mathbf 0,I)\) 采样，再由解码网络生成 \(\mathbf x\)。训练这个模型需要反推“哪个 \(\mathbf z\) 生成了这个 \(\mathbf x\)”（后验），这在高维里算不出来，于是再训练一个编码网络去近似它，并用一个可计算的下界 ELBO 代替算不出的似然。

这个框架在机器人学习里随处可见：ACT（22.2 节）用条件 VAE 生成动作块；World Models 用 VAE 把游戏画面压成隐向量再在隐空间里学控制；Dreamer（24.3 节）的世界模型是一个按时间展开的隐变量模型，训练目标就是一个 ELBO；扩散模型（19.6 节）也可以看成层数很多、编码器固定的 VAE。本节把这些背后的共同部件一个个拆开。

先看一个我们完全知道答案的生成模型。一台平面二连杆机械臂（连杆长 0.62 m 和 0.50 m，7.3 节的正运动学），它的关节角由一个二维隐变量决定：\(\hl{prior}{\mathbf z\sim\mathcal N(\mathbf 0,I)}\)，\(q_1=\tfrac{\pi}{2}+0.7z_1\)，\(q_2=z_2\)。我们看不到关节角，只看到手的位置，而且带噪声：\(\hl{lik}{\mathbf x\sim\mathcal N(f(\mathbf z),\sigma_x^2I)}\)，\(f\) 是正运动学。这就是一个隐变量模型：\(\mathbf z\) 是原因，\(\mathbf x\) 是观测，“解码器” \(f\) 是已知的。

从这个模型可以问两个问题。生成：手会出现在哪里？把很多 \(\mathbf z\) 推过去，得到下图左边的青色点云，它的密度就是边缘似然 \(\hl{px}{p(\mathbf x)=\int p(\mathbf z)\,p(\mathbf x|\mathbf z)\,d\mathbf z}\)。推断：看到手在某处，关节角可能是什么？这是贝叶斯公式（4.2 节）给出的后验 \(\hl{post}{p(\mathbf z|\mathbf x)}=p(\mathbf x|\mathbf z)p(\mathbf z)/p(\mathbf x)\)。因为 \(\mathbf z\) 只有两维，我们可以在一张 181×181 的网格上把它直接算出来，画在右边。拖动琥珀色的目标点，看后验怎样变。

现在换成真正的 VAE：数据是机械臂的图像，没有人告诉网络关节角。我们按 FIG 1 的机械臂在 \(q_1\in[\tfrac\pi2-1.1,\tfrac\pi2+1.1]\)、\(q_2\in[-1.7,1.7]\) 里均匀随机取 500 个姿态，各画一张 14×10 像素的灰度图（像素值在 0 到 1 之间，140 维），另取 100 张作测试集。模型有两个网络：

• 编码器 \(q_\phi(\mathbf z|\mathbf x)=\mathcal N(\boldsymbol\mu_\phi(\mathbf x),\operatorname{diag}\boldsymbol\sigma^2_\phi(\mathbf x))\)：MLP 140 → 48（ReLU）→ 4，输出 2 个均值和 2 个 \(\log\sigma^2\)。
• 解码器 \(p_\theta(\mathbf x|\mathbf z)\)：MLP 2 → 48（ReLU）→ 140，每个像素输出一个 logit，按伯努利分布计算似然，于是重建项就是 140 个像素的二元交叉熵之和（4.11 节）。

每一步取 16 张图，对每张：编码得到 \(\boldsymbol\mu,\boldsymbol\sigma\)，采一个 \(\boldsymbol\epsilon\sim\mathcal N(\mathbf 0,I)\)，令 \(\mathbf z=\boldsymbol\mu+\boldsymbol\sigma\odot\boldsymbol\epsilon\)（重参数化，第 4 节），解码，损失为

KL 项有闭式（4.11 节高斯 KL 公式的逐维求和），不用采样。\(\beta=1\) 时 \(\mathcal L\) 恰好是负的 ELBO（第 5 节推导），这就是标准 VAE；\(\beta\) 是你可以拉的旋钮。优化器是 Adam（学习率 0.002，6.3 节），全部由 nn.js 在你的浏览器里一步步算出，每步的随机数都有固定种子，所以下面引用的数字可以复现。

回到 FIG 1 的机械臂：在那里 \(\log p(\mathbf x)\) 和真实后验都能算，所以可以亲眼看到 VAE 的核心等式（第 5 节推导）：对任意分布 \(q(\mathbf z)\)，

左边与 \(q\) 无关，所以提高 ELBO 就等于缩小缺口，即让 \(q\) 贴近真实后验；缺口为零当且仅当 \(q\) 就是后验。下图中 \(q=\mathcal N(\mathbf m,\operatorname{diag}(s_1^2,s_2^2))\) 是对角高斯，和 VAE 编码器的输出形式一样。你可以手动拖它，也可以按“拟合 q”，用重参数化的随机梯度上升（每步 8 个样本，Adam，学习率 0.03，共 600 步）去最大化 ELBO。ELBO 用 24×24 点的高斯-埃尔米特求积精确计算，不含采样噪声。

这正是 4.11 节讲的反向 KL 的“抓峰”行为：缺口是 \(D_{\mathrm{KL}}(q\|p)\)，\(q\) 放在 \(p\) 很小的地方会被重罚，所以它宁可只蹲在一个峰上。对 VAE 而言，这意味着 ELBO 总是低估 \(\log p(\mathbf x)\)，低估多少取决于后验离“对角高斯”有多远。缺口有两个来源：近似族不够灵活（本图：对角高斯盖不住双峰和弯曲），以及摊销：VAE 不为每个 \(\mathbf x\) 单独优化 \(q\)（像本图这样），而是用一个编码网络一次性输出所有 \(\mathbf x\) 的 \(q\)，网络不完美又会多出一块“摊销缺口”。

训练要对 \(\mathbb E_{q_\phi}[f(\mathbf z)]\) 关于编码器参数 \(\phi\) 求梯度，可“采样”这个操作本身没有导数。有两条路。得分函数（score function，也叫 REINFORCE、似然比）估计：\(\nabla_\mu\mathbb E[f]=\mathbb E\big[f(z)\,\nabla_\mu\log q(z)\big]\)，只需要能算 \(f\) 的值，不需要 \(f\) 可导，20.3 节的策略梯度用的就是它。重参数化估计：把随机性挪到一个与参数无关的噪声上，\(z=\mu+\sigma\hl{eps}{\epsilon}\)，\(\epsilon\sim\mathcal N(0,1)\)，于是 \(\nabla_\mu\mathbb E[f]=\mathbb E\big[f'(\mu+\sigma\epsilon)\big]\)，梯度直接沿网络反传（16.3 节）。两者的期望相同，方差却天差地别。下面在一维里比较：\(z\) 是抓取的横向偏移，\(f(z)=(z-1.5)^2\) 是离目标偏移 1.5 的平方误差，我们要 \(\nabla_\mu\mathbb E_{z\sim\mathcal N(\mu,\sigma^2)}[f(z)]\)，真值是 \(2(\mu-1.5)\)。

重参数化之所以好，是因为它用到了 \(f\) 的导数：它知道“往哪边挪 \(z\) 代价会降”，而得分函数只知道“这个样本代价高还是低”，要靠大量样本的相关性才能平均出方向。代价是要求 \(\mathbf z\) 连续、\(f\) 可微、分布能写成“参数的可微函数 + 固定噪声”（高斯、拉普拉斯、logistic 等可以；离散的类别变量不行，要用 Gumbel-softmax 之类的松弛，或者退回得分函数）。VAE 的解码器是可微网络，正好满足条件，这是 VAE 能在 2014 年训练起来的关键。

训练好的 VAE 给了我们一张“隐空间地图”：因为 KL 项把每个编码都拉向 \(\mathcal N(\mathbf 0,I)\)、又让它们带着噪声，相邻的 \(\mathbf z\) 必须解码成相似的图像，先验覆盖的整个区域都被“填满”。下图用 FIG 2 的模型（两图共享，在这里也可以继续训练）把 \([-2.4,2.4]^2\) 上 9×9 个 \(\mathbf z\) 解码成图像，并比较两种插值：在隐空间里走直线 \((1-t)\mathbf z_A+t\mathbf z_B\) 再解码，和在像素空间里直接混合 \((1-t)\mathbf x_A+t\mathbf x_B\)。

生成模型是 \(p_\theta(\mathbf x,\mathbf z)=\hl{prior}{p(\mathbf z)}\,\hl{lik}{p_\theta(\mathbf x|\mathbf z)}\)，先验通常固定为 \(\mathcal N(\mathbf 0,I)\)，解码器是参数为 \(\theta\) 的网络。给定数据集 \(\{\mathbf x^{(n)}\}\)，最大似然（4.6 节）要最大化 \(\sum_n\log p_\theta(\mathbf x^{(n)})\)，其中

这个积分在高维里没法算（FIG 1 的核心洞见）。用蒙特卡洛从先验采样 \(\frac1K\sum_k p_\theta(\mathbf x|\mathbf z_k)\) 理论上无偏，实际上几乎所有先验样本都解码成与 \(\mathbf x\) 无关的图像，估计的方差大到没法用：有用的 \(\mathbf z\) 集中在后验附近，而后验正是我们不知道的。

引入任意一个分布 \(\hl{q}{q_\phi(\mathbf z|\mathbf x)}\)（在 \(p\) 非零处非零），对 \(\log p_\theta(\mathbf x)\) 做恒等变形。因为 \(\log p_\theta(\mathbf x)\) 不依赖 \(\mathbf z\)，可以对 \(q\) 取期望，再把 \(p_\theta(\mathbf x)=p_\theta(\mathbf x,\mathbf z)/p_\theta(\mathbf z|\mathbf x)\) 代入、上下同乘 \(q\)：

因为 KL 非负（4.11 节），\(\mathrm{ELBO}\le\log p_\theta(\mathbf x)\)，这就是“证据下界”名字的由来（证据 evidence 是 \(p(\mathbf x)\) 的别名）。也可以用 Jensen 不等式一步得到：\(\log\mathbb E_q[p(\mathbf x,\mathbf z)/q]\ge\mathbb E_q[\log(p(\mathbf x,\mathbf z)/q)]\)，但上面的恒等式多告诉了我们缺口是什么。于是同时对 \(\theta\) 和 \(\phi\) 最大化 ELBO 一举两得：对 \(\theta\) 是在（近似地）做最大似然，对 \(\phi\) 是在缩小缺口，让 \(q\) 逼近真实后验（FIG 3）。

最后一行的两项正是 FIG 2 的两条曲线。重建项要求“从 \(q\) 采出的编码能把 \(\mathbf x\) 解码回来”；KL 项要求“每个 \(\mathbf x\) 的编码分布别离先验太远”。对数据集平均后，KL 项还可以写成 \(I_q(\mathbf x;\mathbf z)+D_{\mathrm{KL}}(q(\mathbf z)\|p(\mathbf z))\)，其中 \(q(\mathbf z)=\mathbb E_{\mathbf x}[q_\phi(\mathbf z|\mathbf x)]\) 是“聚合后验”：它既限制了编码携带的信息量，又要求所有编码合起来像先验，后者正是“从先验采样能得到像样的图像”的保证。

KL 项：\(q=\mathcal N(\boldsymbol\mu,\operatorname{diag}\boldsymbol\sigma^2)\)、\(p=\mathcal N(\mathbf 0,I)\) 时，把 4.11 节的一维公式逐维相加：

每一维在 \(\mu_i=0,\sigma_i=1\) 处为零：一维“坍塌”就是指它停在了这里。编码器输出 \(\log\sigma^2\) 而不是 \(\sigma\)，这样不需要约束正负，数值也稳定（本页把它截在 \([-8,4]\)）。

重建项：对每张图只采一个 \(\mathbf z\) 做蒙特卡洛估计，它的形式取决于你给解码器选的似然。伯努利解码器（本页）：\(-\log p_\theta(\mathbf x|\mathbf z)=\sum_k \mathrm{BCE}(\ell_k,x_k)\)，\(\ell_k\) 是 logit。高斯解码器 \(\mathcal N(\hat{\mathbf x}_\theta(\mathbf z),s^2I)\)：\(-\log p=\frac{1}{2s^2}\|\mathbf x-\hat{\mathbf x}\|^2+\frac D2\log(2\pi s^2)\)。于是“均方误差 + KL”这个常见写法其实默认了 \(s^2=\tfrac12\)；用 \(\|\mathbf x-\hat{\mathbf x}\|^2+\beta D_{\mathrm{KL}}\) 等价于 \(s^2=\beta/2\)：β 和解码器噪声方差是同一个旋钮。

写 \(\mathbf z=g_\phi(\boldsymbol\epsilon,\mathbf x)=\boldsymbol\mu_\phi(\mathbf x)+\boldsymbol\sigma_\phi(\mathbf x)\odot\hl{eps}{\boldsymbol\epsilon}\)，\(\boldsymbol\epsilon\sim\mathcal N(\mathbf 0,I)\) 与 \(\phi\) 无关，于是期望可以和梯度交换：

这三行就是 FIG 2 每一步的反传：解码器的 backward 返回 \(\partial\mathcal L/\partial\mathbf z\)，乘上 \(1\) 和 \(\tfrac12\epsilon_i\sigma_i\) 送进编码器输出，再加上 KL 的闭式梯度。得分函数估计则是 \(\nabla_\mu\mathbb E[f]=\mathbb E[f(z)\,(z-\mu)/\sigma^2]\)。对 FIG 4 的 \(f(z)=(z-a)^2\)，记 \(d=\mu-a\)，按高斯矩 \(\mathbb E\epsilon^2=1,\ \mathbb E\epsilon^4=3,\ \mathbb E\epsilon^6=15\) 展开，可得三个单样本估计的方差：重参数化 \(\hat g_{\mathrm{rp}}\)、得分函数 \(\hat g_{\mathrm{sf}}\)，以及减去基线 \(b=\mathbb E_q f\) 的得分函数 \(\hat g_{\mathrm{sf}}^{\,b}\)：

\(\sigma\to0\) 时重参数化的方差趋于 0，得分函数的却像 \(d^4/\sigma^2\) 那样发散；基线去掉了发散项，但仍留下与 \(\sigma\) 无关的 \(8d^2\)。例如 \(\mu=0,\sigma=0.2\)（\(d=-1.5\)）：两者标准差之比约 31，加基线后约 10.7，和 FIG 4 实测一致。

Higgins 等（2017）的 β-VAE 最小化 \(-\mathbb E_q[\log p_\theta(\mathbf x|\mathbf z)]+\beta D_{\mathrm{KL}}(q\|p)\)。β > 1 逼编码器用更少的信息，常被用来换取更“解耦”的隐维度；β < 1 换取更好的重建。只有 β = 1 时目标才是 ELBO，此时比较不同模型的 ELBO 才有意义。后验坍塌是指 KL 项赢了：\(q_\phi(\mathbf z|\mathbf x)\approx p(\mathbf z)\) 对所有 \(\mathbf x\) 成立，\(\mathbf z\) 不再携带任何信息，解码器只好对所有输入输出数据的平均（FIG 2 里 β = 10 时重建项升到约 27 nats，β = 1 时还不到 15）。它在两种情况下最常见：KL 权重相对太大（大 β，或者重建项用了“对像素取平均”而不是求和，见第 7 节误区 3）；以及解码器强到不需要 \(\mathbf z\) 也能把 \(\mathbf x\) 建模好，例如自回归解码器（19.4 节，Bowman 等 2016 在文本 VAE 上首先报告）。常见对策：训练初期把 β 从 0 慢慢升到 1（KL 退火），“free bits”（每维 KL 低于某个阈值时不罚），或者削弱解码器。

机器人模仿学习最常见的难题之一：示教者面对同一个场景会做出不同但都正确的动作。下面的任务里，夹爪要从 (0, −1) 走到 (0, 1)，中间 \(y=0\) 处有一个半径 0.28 m 的障碍物，横向位置 \(o\) 每次不同。示教者从障碍物左边或右边 0.5 m 处绕过（加 0.06 m 的噪声），障碍物越靠右越倾向于走左边（走左边的概率 \(1/(1+e^{-4o})\)）。动作 \(a\) 就是路径在 \(y=0\) 处的横向位置，路径是过起点、\((a,0)\) 和终点的二次贝塞尔曲线。共 400 条示教。

用均方误差回归 \(a=\pi(o)\)（22.1 节的行为克隆）会学到条件均值 \(\mathbb E[a|o]\)，两个模式的平均，在 \(o\approx0\) 时正好穿过障碍物。条件 VAE（CVAE，Sohn 等 2015）把所有分布都以条件 \(\mathbf c=o\) 为条件：编码器 \(q_\phi(z|a,o)\)，解码器 \(p_\theta(a|z,o)=\mathcal N(\hat a_\theta(z,o),0.08^2)\)，先验仍用固定的 \(\mathcal N(0,1)\)，ELBO 变成

直觉上，一维的 \(\hl{cz}{z}\) 学会了编码“这次走哪边”，这正是光看 \(o\) 无法确定、只有看到动作才知道的那部分信息。测试时没有动作可编码，就从先验采 \(z\) 再解码，每个样本都是一个完整的、在某一边的动作。两个网络都是小型 tanh MLP（编码器 2 → 32 → 2，解码器 2 → 32 → 32 → 1；MSE 策略 1 → 32 → 32 → 1），每步 32 条示教，Adam，学习率 0.003。

注意那 10% 到 15% 的碰撞样本：解码器是连续函数，\(z\) 从一边的模式连续变到另一边时必然经过中间，而先验在 \(z\approx0\) 附近恰好密度最大。所以“取先验均值 \(z=0\)”这个看似保险的做法，在两个模式差不多一样可能时（\(o\approx0\)），反而正落在缝里（读数“z = 0 的路径”）。这是连续隐变量模型对付多模态的一个结构性弱点，也是扩散模型和 flow matching（19.6、19.7 节）在动作生成上受欢迎的原因之一；22.2 节会把 ACT 与 Diffusion Policy 放在一起比较。