很多同学在学习深度强化学习课程时往往会被“变分推断”这一讲卡住。它不像策略梯度那样有一个明确的“采样-估计-更新”回路也不像 Q-learning 那样直接维护一张价值表而是同时牵扯概率图模型、最优化理论和神经网络训练。尤其当课程进入第 12 讲前面已经消化了 DQN、PPO、SAC 这些主流算法突然切换到“推断”视角会让人产生一种“我们还是不是在做强化学习”的恍惚感。本文把这部分内容拆开重讲一遍先说变分推断解决什么问题再推导 ELBO 和重参数化技巧然后回到强化学习本身解释“控制即推断”、最大熵强化学习与变分推断的关系最后给出一个可以直接运行的状态表示学习小项目。无论你是想把概念弄懂还是想在实践中快速上手这篇文章都值得收藏备用。1. 背景与核心概念1.1 为什么强化学习需要变分推断强化学习本质上是在学一个策略给定状态 (s)输出动作 (a)目标是最大化长期累积回报。这个描述听起来和概率推断没什么关系但一旦我们把“智能体是理性的”这个假设转化为概率语言就会发现策略学习其实可以被看成一个后验推断问题。举个例子。假设我们观测到了一个智能体的一系列行为并且知道这些行为带来了较高的回报。那么我们可以问在“这个智能体表现得不错”这个条件下它更可能采取什么样的策略这就是一个条件概率 (p(\pi \mid \text{高回报})) 的问题。可惜这个后验分布通常没有解析表达式而且状态动作空间稍微大一点精确计算就是指数级的复杂度。这个时候变分推断就派上了用场。在深度学习时代变分推断的应用更直接。基于模型的强化学习需要学习环境的状态转移模型而真实转移往往是未知且有随机性的需要用变分方法做近似后验估计。部分可观测环境POMDP中智能体要从观测序列中推断隐藏状态这本身就是变分推断擅长的问题。最大熵策略优化中的熵正则项也可以从变分目标中自然推导出来。离线强化学习、多智能体协作学习里变分推断也常被用来处理奖励稀疏、联合动作空间过大等问题。所以“强化学习中的变分推断”并不是为了给课程凑章节而是很多前沿算法背后的共同数学骨架。1.2 变分推断的基本问题变分推断要解决的核心问题可以描述成一句话我们想计算一个复杂的后验分布 (p(\theta \mid x))但这个分布算不出来于是用一个简单的分布 (q(\theta)) 去近似它。具体到强化学习中这个“后验分布”可能是给定一批轨迹数据时策略参数 (\theta) 的后验分布给定观测序列时隐状态 (z) 的后验分布给定离线数据集时Q 函数或奖励模型的后验分布。变分推断的思路不是直接去求后验而是构建一个带参数的分布族 ({q_\phi})然后寻找一个参数 (\phi)使得 (q_\phi(\theta)) 和真实后验 (p(\theta \mid x)) 尽可能接近。两个分布的接近程度一般用 KL 散度衡量。于是问题就变成了[ \min_{\phi} \operatorname{KL}(q_\phi(\theta) | p(\theta \mid x)) ]由于 (p(\theta \mid x)) 中包含难以计算的归一化常数直接最小化这个 KL 散度仍然行不通。变分推断的做法是对目标函数做一个等价变形得到一个可以优化的下界也就是 ELBOEvidence Lower Bound证据下界。1.3 与 MCMC 等方法的关系很多人在学习时容易把变分推断和 MCMC马尔可夫链蒙特卡洛搞混。两者都是近似推断工具但定位完全不同。MCMC 通过采样的方式逼近后验理论保证好但计算量大收敛慢不适合深度神经网络的端到端训练。变分推断把推断问题变成优化问题速度快可扩展性强但得到的是近似解而且可能陷入局部最优。在强化学习中我们通常面对的是高维连续状态动作空间需要和神经网络一起做梯度更新因此变分推断方法更常用。MCMC 更多用于小规模、对精度要求极高的离线分析。2. 数学基础ELBO 与重参数化技巧2.1 ELBO 是怎么推导出来的我们回到一般形式。假设有观测数据 (x)隐变量 (z)联合分布 (p(x, z) p(z) p(x \mid z))。我们希望得到后验 (p(z \mid x))但它往往不可解。引入变分分布 (q(z))看它对数似然的下界[ \log p(x) \log \int p(x, z) dz ]这个积分很难算所以我们对 (\log p(x)) 做一个技巧性变换[ \log p(x) \log p(x) \int q(z) dz \int q(z) \log p(x) dz ]因为 (q(z)) 是一个概率分布所以 (\int q(z) dz 1)。接下来把 (p(x)) 写成 (p(x, z) / p(z \mid x))继续展开[ \log p(x) \int q(z) \log \frac{p(x, z)}{p(z \mid x)} dz ]把 (\log) 拆开并与 (q(z)) 组合成 KL 散度的形式[ \log p(x) \int q(z) \log \frac{p(x, z)}{q(z)} dz\operatorname{KL}(q(z) | p(z \mid x)) ]第一项就是 ELBO[ \text{ELBO}(q) \mathbb{E}_{q(z)}[\log p(x, z) - \log q(z)] ]第二项是变分分布与真实后验之间的 KL 散度它一定大于等于 0。因此[ \log p(x) \ge \text{ELBO}(q) ]当且仅当 (q(z) p(z \mid x)) 时等号成立。2.2 把 ELBO 拆成两项ELBO 还可以继续展开成两种直观写法。第一种写法[ \text{ELBO}(q) \mathbb{E}_{q(z)}[\log p(x \mid z)] - \operatorname{KL}(q(z) | p(z)) ]这个形式可以理解为最大化 ELBO 就是同时做两件事——让重构概率 (\log p(x \mid z)) 尽可能大即隐变量要能很好解释观测数据让变分分布 (q(z)) 尽量接近先验 (p(z))即隐变量不能偏离先验太远。在 VAE 中第一项是重构损失第二项是正则项。在强化学习的变分模型学习中第一项对应状态预测误差第二项对应隐变量正则化。第二种写法[ \text{ELBO}(q) \mathbb{E}_{q(z)}[\log p(x, z)] H(q(z)) ]这里 (H(q(z))) 是变分分布的熵。这个形式告诉我们变分推断鼓励在保持与数据似然一致的同时让分布尽可能“模糊”。这也是为什么最大熵强化学习与变分推断有天然的联系——熵正则项会从变分目标中自然浮现。2.3 重参数化技巧变分推断要用于神经网络训练必须能对参数 (\phi) 求梯度。但 ELBO 中的期望是 (\mathbb{E}{q\phi(z)}[\cdot])对分布参数直接求梯度会非常困难因为采样过程不可导。Kingma 等人提出的重参数化技巧解决了这个问题。以高斯分布为例假设[ z \sim \mathcal{N}(\mu_\phi, \sigma_\phi^2) ]我们可以把随机性从参数中剥离出来。先从一个标准正态分布中采样[ \epsilon \sim \mathcal{N}(0, I) ]然后通过变换得到 (z)[ z \mu_\phi \sigma_\phi \odot \epsilon ]这样一来(z) 关于 (\mu_\phi) 和 (\sigma_\phi) 的梯度就可以正常反向传播而随机性全部集中在与参数无关的 (\epsilon) 上。这个技巧几乎是所有深度变分方法的基石后面的实战代码会再次用到。2.4 均值场近似与分布族选择变分推断中(q(z)) 的分布族选择决定了表达能力和计算复杂度。最常用的是均值场近似[ q(z) \prod_i q(z_i) ]即假设各个隐变量之间相互独立。这显然会损失变量之间的相关性但好处是优化复杂度大大降低。在强化学习中如果状态表示的各维度之间相关性并不重要均值场假设往往够用。另一个常见选择是各向同性高斯分布即协方差矩阵为对角阵。这时候我们需要优化的参数只有每个维度的均值和方差。再进一步为了数值稳定性通常让网络输出 (\log \sigma^2) 而不是直接输出 (\sigma)这样能避免方差为负的问题。3. 强化学习中的变分推断3.1 控制即推断视角“控制即推断”Control as Inference是理解强化学习与变分推断关系的重要视角。它把强化学习过程建模成一个概率图模型并引入“最优性变量” (O_t) 表示时刻 (t) 的决策是否达到最优。通常定义[ p(O_t 1 \mid s_t, a_t) \exp(r(s_t, a_t)) ]奖励越大该时刻被判定为“最优”的概率越高。智能体的目标是在给定整条轨迹“都是最优”的条件下推断出每一时刻应该采取什么动作即计算[ p(a_t \mid s_t, O_{1:T} 1) ]这个后验分布才是“合理策略”的完整描述。精确计算同样困难于是可以找一个变分分布 (q_\theta(a_t \mid s_t)) 来近似它。对这个变分目标做推导会发现优化过程不只是最大化累积回报还天然带上了熵正则项。换句话说变分推断视角下的策略优化不是把奖励直接变成概率而是在“追求高回报”和“保持动作多样性”之间做权衡。3.2 最大熵强化学习与变分目标Soft Actor-CriticSAC这类最大熵强化学习算法目标函数是[ J(\pi) \sum_t \mathbb{E}{(s_t, a_t) \sim \rho\pi}[ r(s_t, a_t) \alpha H(\pi(\cdot \mid s_t)) ] ]其中 (\alpha) 是温度参数控制熵项的权重。从控制即推断的视角来看SAC 的熵正则项并不是额外加进去的“技巧”而是变分推断目标展开之后的必然产物。这一点对理解 SAC 非常重要如果你只把 SAC 看成“PPO 加了个熵奖励”就很难理解为什么要动态调节 (\alpha)。如果你从变分推断的角度看(\alpha) 其实是在控制变分分布与最优后验之间的偏离程度是算法内在的一部分。3.3 基于模型强化学习中的变分表示学习基于模型的强化学习经常需要学习一个环境模型。真实环境的转移往往是高维、非线性的直接预测原始像素级状态非常困难因此在很多方法中会先学习一个低维隐状态表示然后在这个隐空间里做转移预测和规划。这就是变分推断发挥作用的地方。以 Dreamer 这类算法为例它的世界模型包含表示模型从观测 (x_t) 编码出隐状态 (z_t)转移模型从 (z_t) 和动作 (a_t) 预测下一时刻隐状态 (z_{t1})奖励模型从隐状态和动作预测当前奖励。训练这些模型时目标函数通常就是 ELBO 的变体既要让隐状态能重构出原始观测又要让隐状态的预测变化尽量符合先验。这与 VAE 的优化目标是同构的。可以这样说如果你理解了 VAE再去看 Dreamer 的世界模型会发现很多概念都是相通的。3.4 离线与多智能体场景中的应用在离线强化学习中智能体只能使用已有的固定数据集无法与环境交互。这时一个核心问题是策略分布与数据分布之间可能存在偏差导致 Q 函数被过高估计。变分推断在这里的一个思路是把离线策略优化建模为在数据分布约束下的后验推断。常见的做法是对策略分布与行为策略分布之间施加 KL 约束或者对 Q 函数更新施加保守项。这些约束在数学上都可以统一到变分推断的目标中。在多智能体强化学习中联合动作空间随智能体数量指数增长。变分推断可以把联合策略分解为各个智能体策略的近似分布或者用来推断智能体之间隐藏的通信意图。例如在协作场景中每个智能体不仅要知道自己的动作还要推断队友可能采取的行为这种推断天然带有不确定性变分方法能够给出分布式的近似解。4. 实战用 PyTorch 实现变分推断辅助的状态表示学习4.1 环境准备下面的示例在常见的深度学习环境下即可运行。本文不绑定具体版本建议使用Python 3.8 及以上PyTorch 1.10 及以上NumPy 1.20 及以上如果使用 GPU安装对应版本的 CUDA 版 PyTorch没有 GPU 也能运行只是速度略慢。可以用 pip 快速确认依赖pip install torch numpy matplotlib4.2 项目结构说明为了让代码更贴近真实项目这里按以下结构组织variational_rl_demo/ ├── main.py # 训练入口 ├── models.py # 变分自编码器 重参数化 └── utils.py # 数据加载与可视化接下来我们逐个文件实现。4.3 实现重参数化采样重参数化模块是整个变分推断的基石。它把“从高斯分布采样”转换成“从标准正态采样再加线性变换”保证梯度可以自由传播。文件路径models.pyimport torch import torch.nn as nn import torch.nn.functional as F class ReparameterizedSampler(nn.Module): 重参数化高斯采样模块。 输入编码网络的输出返回采样得到的隐变量 z。 def __init__(self, latent_dim: int): super().__init__() self.latent_dim latent_dim def forward(self, mu: torch.Tensor, log_var: torch.Tensor): # 在推理时通常直接使用均值不做随机采样 if not self.training: return mu # 标准正态噪声 eps torch.randn_like(mu) # 从 log_var 还原为标准差并保证数值稳定 std torch.exp(0.5 * log_var) # 重参数化z mu std * eps z mu std * eps return z这里的关键点是使用log_var而不是直接使用方差原因有两个神经网络可以输出任意实数不需要额外加约束来保证方差非负训练过程中方差的尺度变化范围很大用对数形式可以让梯度更稳定。4.4 搭建变分自编码器下面实现一个简单的 VAE。它接收二维坐标点作为输入通过 MLP 编码得到隐变量的均值和对数方差然后用重参数化采样得到隐变量再通过解码器重构输入。文件路径models.pyclass VAE(nn.Module): 在简单二维数据上训练的最小变分自编码器。 输入输出维度都是 input_dim隐空间维度是 latent_dim。 def __init__(self, input_dim: int 2, latent_dim: int 2): super().__init__() self.input_dim input_dim self.latent_dim latent_dim # 编码器输入 - 均值和对数方差 self.encoder nn.Sequential( nn.Linear(input_dim, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU(), ) self.mu_head nn.Linear(32, latent_dim) self.log_var_head nn.Linear(32, latent_dim) # 重参数化采样器 self.sampler ReparameterizedSampler(latent_dim) # 解码器隐变量 - 重构输入 self.decoder nn.Sequential( nn.Linear(latent_dim, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, input_dim), ) def encode(self, x: torch.Tensor): h self.encoder(x) mu self.mu_head(h) log_var self.log_var_head(h) return mu, log_var def decode(self, z: torch.Tensor): return self.decoder(z) def forward(self, x: torch.Tensor): mu, log_var self.encode(x) z self.sampler(mu, log_var) x_recon self.decode(z) return x_recon, mu, log_var4.5 训练循环与 ELBO 损失训练时我们用 ELBO 作为损失函数它由两部分组成重构误差这里使用 MSE衡量生成结果与原始输入之间的差距KL 散度衡量编码分布与标准正态先验之间的差距。在高斯假设下(q_\phi(z|x) \mathcal{N}(\mu, \sigma^2))先验 (p(z)\mathcal{N}(0, I))两者之间的 KL 散度有解析表达式[ \operatorname{KL}(q_\phi(z|x) | p(z)) -\frac{1}{2} \sum_{j1}^{d} \left(1 \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\right) ]在代码中我们不需要自己手算这个式子可以直接调用F.kl_div的变体也可以按上述公式自行实现。这里给出自行实现的完整损失函数。文件路径main.pyimport torch import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset from models import VAE def vae_loss(x, x_recon, mu, log_var): VAE 的 ELBO 损失返回 (loss, recon_loss, kl_loss)。 # 重构损失MSE recon_loss F.mse_loss(x_recon, x, reductionmean) # KL 散度解析形式 kl_loss -0.5 * (1 log_var - mu.pow(2) - log_var.exp()).sum(dim1).mean() total_loss recon_loss kl_loss return total_loss, recon_loss, kl_loss为了能快速看到训练效果我们使用一个简单的二维高斯混合分布作为数据集。这样既不需要下载 MNIST也能直观看见变分自动编码器学习到的状态表示。文件路径main.pydef generate_toy_data(n_samples2000, seed42): 生成四个簇的二维高斯混合数据模拟强化学习中的状态分布。 torch.manual_seed(seed) centers torch.tensor([ [2.0, 2.0], [-2.0, 2.0], [2.0, -2.0], [-2.0, -2.0], ]) cluster_idx torch.randint(0, 4, (n_samples,)) base torch.randn(n_samples, 2) * 0.3 data base centers[cluster_idx] return data.float()接下来是训练函数。注意这里训练过程中 VAE 处于train()模式因此ReparameterizedSampler会进行随机采样。文件路径main.pydef train_vae(model, data_loader, epochs60, lr1e-3): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0.0 total_recon 0.0 total_kl 0.0 n_batches 0 for batch in data_loader: x batch[0] optimizer.zero_grad() x_recon, mu, log_var model(x) loss, recon_loss, kl_loss vae_loss(x, x_recon, mu, log_var) loss.backward() optimizer.step() total_loss loss.item() total_recon recon_loss.item() total_kl kl_loss.item() n_batches 1 if (epoch 1) % 10 0: print( fEpoch {epoch 1:3d} | fLoss {total_loss / n_batches:.4f} | fRecon {total_recon / n_batches:.4f} | fKL {total_kl / n_batches:.4f} )文件路径main.pyif __name__ __main__: data generate_toy_data() dataset TensorDataset(data) loader DataLoader(dataset, batch_size64, shuffleTrue) model VAE(input_dim2, latent_dim2) train_vae(model, loader, epochs60) # 训练后前向一次观察隐变量分布 model.eval() with torch.no_grad(): x data[:16] x_recon, mu, log_var model(x) print(\n原始数据前 4 个点) print(x[:4]) print(\n重构数据前 4 个点) print(x_recon[:4]) print(\n隐变量均值前 4 个点) print(mu[:4])4.6 运行与验证运行命令python main.py预期输出风格如下Epoch 10 | Loss 1.5321 | Recon 1.0312 | KL 0.5009 Epoch 20 | Loss 1.2015 | Recon 0.7643 | KL 0.4372 Epoch 30 | Loss 1.0124 | Recon 0.6128 | KL 0.3996 Epoch 40 | Loss 0.8734 | Recon 0.4912 | KL 0.3822 Epoch 50 | Loss 0.7712 | Recon 0.4123 | KL 0.3589 Epoch 60 | Loss 0.6902 | Recon 0.3421 | KL 0.3481从输出可以看到随着训练进行重构损失和 KL 损失同时下降说明变分自编码器在拟合数据的同时隐变量的后验分布也在向标准正态先验靠拢。这正是 ELBO 优化过程的直观体现。在实际强化学习项目中这个 VAE 可以换成更复杂的卷积结构输入改为图像帧隐变量则作为策略网络或世界模型的状态表示。5. 常见问题与排查思路问题现象常见原因解决思路KL 损失迅速变为 0重构损失却很大后验坍塌posterior collapse解码器忽略了隐变量减小 KL 项权重或者使用 KL 退火检查解码器是否过强训练不稳定loss 出现 NaN对数方差输出数值过大导致 exp 后溢出限制 log_var 的边界例如 clamp 到 [-10, 10]重构效果好但隐变量没有语义KL 项权重过小模型退化成普通自编码器适当增大 KL 权重或使用 β-VAE 的方式逐步调整采样得到的隐变量和训练时差距较大训练和推理阶段没有区分采样模式在推理时使用均值 μ而不是随机采样检查是否固定了随机种子强化学习策略训练时 entropy 项不稳定温度参数 α 固定且过大/过小使用自动调节温度例如 SAC 中的目标熵机制基于模型的状态预测发散隐状态先验与真实转移不匹配在转移模型中引入不确定性建模用变分分布而非点估计6. 最佳实践与工程建议6.1 数值稳定性是第一位变分推断中几乎到处都是对数、指数和除法稍不注意就会出现 NaN。常见的处理建议网络输出层使用log_var并对其做clamp避免极端值计算标准差时使用torch.exp(0.5 * log_var)不要在模型中手动计算方差再开方如果使用F.kl_div要确认输入是对数概率还是概率传错维度会得到非常奇怪的结果训练早期可以先把 KL 项权重设小等重构损失稳定后再逐渐增大。6.2 可视化隐空间比看 loss 更重要训练变分模型时Loss 曲线只能说明目标函数在优化并不能说明隐表示是否真的有用。建议每次训练后把隐变量投影到二维平面观察不同轨迹或不同簇的分布。如果隐空间中的点能明显分开说明变分编码器学到了有区分度的表示如果所有点挤在一起很可能是后验坍塌或者先验太强。在强化学习中也可以用类似手段把训练过程中收集到的经验数据编码到隐空间然后观察策略更新前后的分布变化这能帮助你判断策略是在探索新的状态区域还是在同一个区域内反复打转。6.3 与策略网络结合时的熵调节策略当把变分推断引入策略优化时熵项权重不能拍脑袋定死。工程上推荐的做法是先固定一个较大的熵权重让策略充分探索观察策略的熵值变化设定一个目标熵例如动作维度的负数使用类似 SAC 的温度自调节机制让 α 自动更新。这样做的好处是策略不会因为熵权过大而一直随机游走也不会因为熵权过小而过早陷入局部最优。7. 总结与学习路线这一讲的内容表面上看是在讲数学实际上是在给强化学习换一个视角。掌握变分推断之后你会发现很多算法之间的内在联系VAE 的 ELBO 与基于模型强化学习的模型训练目标同源SAC 的熵正则可以从控制即推断中自然推出离线强化学习中的保守约束也能放在变分框架下理解。因此建议的学习顺序是先把 ELBO 推导和重参数化技巧弄明白再实现一个 VAE然后回到 SAC 和 Dreamer 这类算法尝试找到它们目标函数中的变分结构。当你下次在论文摘要里看到“variational bound”时就不会再下意识跳过了而是会想这里到底在近似哪个后验分布如果本文对你有帮助可以收藏备用遇到变分推断相关的算法或代码时再翻出来对照。也欢迎在评论区交流你踩过的坑特别是 KL 坍塌和训练不稳定这两类问题每个人遇到的具体表现可能都不一样互相参考能少走很多弯路。