PAPER DEEP DIVE
PAC-DP:PAC贝叶斯扩散策略学习
提出PAC贝叶斯扩散策略学习框架,将PAC贝叶斯理论与扩散策略结合,为机器人操作策略学习提供理论保证与泛化能力。
1. 论文概览：PAC-贝叶斯扩散策略学习
PAC-DP 由 Mohammad Hasan Yeganegi、Dian Yu、Andrea Del Prete、Majid Khadiv 和 Matteo Saveriano 于 2026 年 7 月提出，将扩散策略建模为贝叶斯神经网络，并基于 PAC-贝叶斯学习理论推导出一种新的训练目标——在标准去噪损失上增加先验与后验参数分布之间的 KL 散度正则化项。扩散策略虽在复杂操作任务中表现出色，但在低数据场景下仅最小化经验去噪目标并不能保证泛化：最优得分函数可能使生成样本紧绑训练数据，甚至在有限数据集上出现记忆现象。PAC-DP 从统计学习理论出发提供了有原则的正则化机制，实验表明其在去噪精度、变分负对数似然和下游机器人任务成功率上均优于标准 DP，尤其在低数据和高难度任务中改进最显著。
2. 核心问题：扩散策略的泛化缺失
扩散策略通过迭代去噪建模多模态动作分布，已成为机器人模仿学习的主导范式之一。然而其泛化性质几乎未被理论探索——尤其是在收集专家演示昂贵且耗时的低数据场景下。近期关于扩散模型的理论和实验研究表明，在低数据训练条件下仅最小化经验去噪目标不足以保证泛化：Li 等证明最小化去噪目标得到的最优得分函数会使 DM 生成与训练数据紧绑的样本；有限数据集或正则化不足时 DM 会记忆训练样本。这些发现同样适用于扩散策略。现有改进方法主要依赖架构归纳偏置、结构化条件或更丰富的数据集，间接提升泛化，但对训练目标本身缺乏理论指导。
3. 方法：PAC-DP 框架
3.1 扩散策略预备
DP 将模仿学习建模为条件 DDPM：动作序列 $x_0 \in \mathbb{R}^{N \times T_a}$ 以最近 $T_o$ 个观测 $o$ 为条件。前向扩散逐步加噪：$x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$，$\epsilon \sim \mathcal{N}(0, I)$，训练网络 $\epsilon_\theta(x_t, t, o)$ 预测注入噪声，最小化去噪目标：
$$\mathcal{L}_{\text{DP}}(\theta) = \mathbb{E}_{x_0, t, \epsilon}\left[\|\epsilon - \epsilon_\theta(x_t, t, o)\|_2^2\right] \tag{2}$$推理时从高斯噪声迭代去噪生成动作，采用 DDIM 加速采样。
3.2 PAC-贝叶斯泛化界
PAC-贝叶斯理论为随机预测器提供数据相关的期望泛化误差上界。给定训练集 $\mathcal{S}=\{z_i\}_{i=1}^N$、先验 $P(w)$ 和学到的后验 $Q(w)$，经验风险与总体风险定义为：
$$\hat{R}(Q) = \mathbb{E}_{w \sim Q}\left[\frac{1}{N}\sum_{i=1}^N \ell(w, z_i)\right], \quad R(Q) = \mathbb{E}_{w \sim Q}\left[\mathbb{E}_{z \sim \mathcal{D}}[\ell(w, z)]\right] \tag{3-4}$$由 PAC-Bayes-kl 不等式与 Refined Pinsker 不等式松弛，以至少 $1-\delta$ 概率成立：
$$R(Q) \leq \hat{R}(Q) + \sqrt{\frac{2\hat{R}(Q)(D_{\text{KL}}(Q\|P) + \ln\frac{2\sqrt{N}}{\delta})}{N}} + \frac{2(D_{\text{KL}}(Q\|P) + \ln\frac{2\sqrt{N}}{\delta})}{N} \tag{5}$$该界揭示核心原理：期望总体风险由经验风险和后验-先验散度联合控制。最小化由此导出的代理目标自然平衡数据拟合与后验复杂度。
3.3 PAC-DP 目标
将噪声预测网络建模为贝叶斯神经网络，确定性权重 $\theta$ 替换为从后验 $Q(w)$ 采样的随机权重 $w$。先验与后验均参数化为分解高斯 $Q(w) = \mathcal{N}(\mu, \text{diag}(\sigma^2))$。权重用重参数化技巧采样 $w = \mu + \sigma \odot \xi$，标准差参数化为 $\sigma = \log(1+\exp(\rho))$ 以保证非负。贝叶斯层的采样权重通过以下重参数化公式获得： $$w = \mu + \sigma \odot \xi, \quad \xi \sim \mathcal{N}(0, I) \tag{7}$$
去噪目标的后验期望 $\hat{R}(Q) = \mathbb{E}_{w \sim Q}[\mathcal{L}_{\text{DP}}(w)]$，代入 PAC-Bayes 界得 PAC-DP 目标：
$$\mathcal{L}_{\text{PAC-DP}} = \hat{R}(Q) + \sqrt{\frac{2\hat{R}(Q)(\lambda D_{\text{KL}}(Q\|P) + \ln\frac{2\sqrt{N}}{\delta})}{N}} + \frac{2(\lambda D_{\text{KL}}(Q\|P) + \ln\frac{2\sqrt{N}}{\delta})}{N} \tag{10}$$其中 $\lambda$ 控制正则化强度。因先验后验均为分解高斯，KL 散度有闭式解，可逐贝叶斯层独立计算：
$$D_{\text{KL}}(Q\|P) = \sum_{l=1}^L D_{\text{KL}}(Q_l \| P_l) \tag{11}$$其中 $L$ 为贝叶斯层数。推理时可用后验均值 $\mu$ 获确定性策略或采样保留随机性。
4. 实验
4.1 PushT 优化动态
在 PushT 任务上用 200 演示（70%训练/30%验证）训练 $10^6$ 步。DP 和 BDP 随训练进行 NPE 逐渐升高（泛化退化），而 PAC-DP 全程保持低 NPE，收敛至 $0.060 \pm 0.002$，相比 DP 的 $0.184 \pm 0.008$ 降低约 67%。变分 NLL 方面，DP/BDP 快速饱和于 $-3.24$（随机高斯去噪器水平），PAC-DP 维持 $-8.275 \pm 0.142$。下游成功率上 PAC-DP 收敛至 $93.7\% \pm 1.0\%$，DP 降至 $86.1\%$，BDP 为 $90.2\%$。
| 方法 | NPE ↓ | 变分 NLL ↓ | 成功率 SR ↑ |
|---|---|---|---|
| DP | 0.184 | -3.24 | 86.1% |
| BDP | 0.191 | -3.24 | 90.2% |
| PAC-DP | 0.060 | -8.275 | 93.7% |
4.2 Robomimic 基准与低数据评估
在 Lift、Can、Square、Transport、ToolHang 五项任务上用 200 演示评估。PAC-DP 在所有任务上 NPE 最低、NLL 最紧。成功率方面，在 ToolHang 等高难度任务上 PAC-DP 达 73.8%，远超 DP 的 40.1% 和 BDP 的 54.4%。低数据评估中（20/40/80/200 演示），PAC-DP 在 20 演示时优势最大——ToolHang 从 64.3%(DP) 提升至 85.8%(PAC-DP, 200 演示)。
| 任务 | #演示 | DP SR | BDP SR | PAC-DP SR |
|---|---|---|---|---|
| PushT | 20 | 48.1% | 49.5% | 54.9% |
| PushT | 200 | 90.1% | 92.6% | 95.6% |
| Can | 20 | 78.1% | 86.5% | 90.6% |
| ToolHang | 200 | 64.3% | 72.0% | 85.8% |
消融实验表明 BDP（贝叶斯参数化但无 PAC-Bayes 正则化）在成功率上略优于 DP，但在 NPE/NLL 上无改善——只有 PAC-DP 在所有指标上全面领先，说明贝叶斯参数化与 PAC-Bayes 正则化两者结合才能最大化性能。超参数敏感性分析显示 $\lambda \in \{10^{-5}, 10^{-6}\}$、$\rho \in \{-3.5, -4.0, -4.5, -5.0\}$ 范围内 PAC-DP 表现稳定，不依赖精细调参。
5. 局限性
- 先验选择：当前使用固定的高斯先验，未来需研究数据相关先验以获得更紧的 PAC-Bayes 界。
- 损失有界假设：PAC-Bayes-kl 不等式假设损失有界，而去噪损失理论上无界；实践中虽快速降至1以下使裁剪损失与原始近似，但理论严谨性仍有缺口。
- 计算开销：贝叶斯参数化增加了每次前向传播的采样开销，虽训练时间增加不显著，但在实时推理场景下仍有影响。
6. 总结
PAC-DP 将扩散策略建模为贝叶斯神经网络，从 PAC-Bayes 泛化界推导出耦合经验去噪损失与 KL 散度正则化的训练目标，在统计学习理论与扩散策略优化之间建立了有原则的联系。实验表明 PAC-DP 在去噪精度、变分 NLL 和任务成功率上均优于标准 DP，尤其在低数据和高难度任务中改进最显著——ToolHang 200 演示时成功率从 64.3% 提升至 85.8%。消融证明增益主要来自 PAC-Bayes 正则化而非贝叶斯参数化本身。真机 Franka 验证了仿真到真实的迁移。核心洞见是：扩散策略的泛化不能仅靠去噪目标保证——PAC-Bayes 界告诉我们，好的策略不仅要在训练数据上去噪准确，其参数分布还必须保持简洁，与先验的偏离应当被数据充分支撑。
flowchart TD
A["训练数据 S = {z_i}"] --> B["前向扩散: x_t = sqrt(a_t) x_0 + sqrt(1-a_t) epsilon"]
B --> C["贝叶斯噪声预测网络 epsilon_w(x_t, t, o)"]
D["先验 P(w) = N(0, sigma^2)"] --> E["KL 散度 D_KL(Q || P)"]
F["后验 Q(w) = N(mu, diag(sigma^2))"] --> G["重参数化采样: w = mu + sigma * xi"]
G --> C
C --> H["经验去噪风险 R_hat(Q)"]
H --> I["PAC-DP 目标: R_hat + KL 正则化"]
E --> I
I --> J["SGD 更新后验参数 mu, rho"]
J --> F
K["推理: 均值 mu 或采样"] --> L["DDIM 迭代去噪生成动作"]
