PAD 阅读笔记
Prediction with Action: Visual Policy Learning via Joint Denoising Process — Yanjiang Guo et al., NeurIPS 2024
摘要
PAD 做了一件之前没人真正做成的事:把"预测未来会发生什么"和"决定机器人该做什么"塞进同一个扩散去噪过程。乍看像是硬把两个功能拼在一起,但论文用数据证明它们不只是兼容,而是互相增强。图像预测迫使模型学会物理规律,这种物理理解反过来让动作预测更准。
PAD 用一个 Diffusion Transformer (DiT) 处理多种模态(RGB图像、机器人姿态、深度图、文本指令),把它们全部编码成 token 后拼在一起联合去噪。关键设计是 attention mask:缺失的模态直接屏蔽掉,同一个模型可以同时吃有动作标注的机器人数据和没有标注的纯视频数据。
在 Metaworld 50 个任务上,一个文本条件策略达到了 72.5% 的平均成功率,比之前最强的 GR-1(57.4%)高出 26.3%。在真实世界中面对从没见过的物体(水果玩具替代方块),成功率也提升了 28%。但最让人吃惊的消融结果是:去掉图像预测后性能暴跌 28.9 个百分点。也就是说,推理时被扔掉的图像预测,在训练时反而贡献了绝大部分性能。
一、Motivation
你在教一个机械臂从桌上抓起一个方块放进碗里。传统 Diffusion Policy 的做法是:看一帧画面,直接输出"手该往哪移"。它学到的是一种从像素到动作的映射,但对于"方块被推会滑多远""手指碰到碗边会怎样"这些物理规律,它只有隐式理解,全压缩在网络权重里,看不见摸不着。
但如果你同时让模型预测"做完这个动作后画面会变成什么样",它就不得不学会这些物理细节。预测下一帧图像相当于额外加了一门"物理考试":不只要答对"该做什么",还要答对"做了之后世界会变成什么样"。两门考试共用一套理解,互相交叉验证。
在 PAD 之前已经有人试过这个思路,但都有问题:
- 两阶段方案(SuSIE, UniPi):先用一个模型预测目标图像,再用另一个模型从图像反推动作。问题是预测图像的时候根本不知道什么信息对动作有用,大量算力花在还原跟控制无关的像素细节上,真正重要的特征反而被稀释了。
- 自回归交替(GR-1):用 GPT 风格的 transformer 一帧一帧地交替预测图像和动作。问题是自回归天然会误差累积,一旦前一帧预测偏了,后面就全跟着偏。
PAD 的切入点很直接:既然图像生成和动作生成在技术上都是扩散去噪(从噪声逐步恢复干净信号),那就别分两步或交替做了,直接把它们合并成一个联合去噪过程。同一个 attention 层同时看到图像 token 和动作 token,两边的梯度信号在每一步去噪里都互相流动。
二、现存问题
- 两阶段信息瓶颈:SuSIE、UniPi 先预测目标图像再推动作,但预测阶段对"哪些像素对控制重要"一无所知,预测质量和动作质量只能分别优化。两个模型之间传递的是一整张图像,高维、冗余,也无法保证保留了控制所需的关键信息。
- 自回归误差累积:GR-1 逐帧预测,第 1 帧偏了 1%,第 2 帧在此基础上偏 1%……到第 10 帧可能已经完全偏离轨道。在需要精确操作的长序列任务中(比如插销、篮球投篮),这种漂移致命。
- 纯动作监督容易过拟合:Diffusion Policy 只用动作做监督信号,而机器人姿态只有 4-7 维(末端位置 + 夹爪),信息量极其稀疏。50 条示范轨迹 × 4 维动作,模型很容易记住训练数据而非学会泛化。
- 多模态扩展困难:现有框架加一种新模态(深度图、触觉),往往需要重新设计网络结构,不是即插即用的。
四、方法详解
Figure 3:PAD 框架全景。左:当前观测经编码器变成条件 latent,与未来帧的噪声 latent 拼接后送入 DiT 联合去噪。右:DiT 内部结构,多模态 token 通过 self-attention 互相看到。
4.1 整体思路:k 步联合预测
给定当前观测(RGB 图像 $c_I$、机器人姿态 $c_A$、可选的深度图 $c_E$、文本指令 $l$),PAD 同时预测未来 $k$ 步的图像、动作和深度图。每次推理后只执行第一个预测动作 $x_A^1$,然后重新观测、重新预测,形成闭环控制。论文中 $k=3$,帧间隔 $i=4$。
4.2 输入编码:每种模态一个编码器
- RGB 图像:通过冻结的 VAE 编码器压成 latent(256×256 输入 → 32×32×4 latent),再 patchify 成 256 个 token。
- 机器人姿态:通过 MLP 映射成与图像 token 同维度的 1 个 token。7 维姿态(3 位置 + 3 旋转 + 1 夹爪)被压缩成一个 embedding。
- 深度图(可选):下采样到 32×32×1,patchify 成 16 个 token。
- 文本指令:通过冻结的 CLIP 编码器提取嵌入向量 $c_l$。
4.3 联合去噪:信息双向流动
关键操作是 channel-wise 拼接:把条件 latent(当前帧的编码,干净的)和噪声 latent(未来帧的带噪版本)在 channel 维度拼起来。比如当前图像编码形状是 $c \times d \times d$,未来 $k$ 帧噪声 latent 形状是 $kc \times d \times d$,拼完就是 $(k+1)c \times d \times d$。
拼好后切成 token 序列,连同动作 token、深度 token 一起送进多层 DiT block。每个 block 里所有 token 做 full self-attention:图像 token 能看到动作 token,动作 token 也能看到图像 token。这就是 PAD 和两阶段方法最本质的区别:信息在去噪的每一步都是双向流动的。
4.4 缺失模态处理:attention mask
当某种模态缺失时(比如纯视频数据没有动作标注),PAD 用 attention mask 把对应位置遮掉,padding token 不参与去噪计算。这意味着同一个模型、同一个训练循环里可以混合两种数据:有标注的机器人数据贡献图像+动作损失,纯视频数据只贡献图像预测损失。
4.5 训练目标
标准 DDPM 噪声预测损失,同时对所有模态计算:
$$\mathcal{L}(\theta) = \lambda_I \cdot \mathcal{L}_{\text{diff}}^I + \lambda_A \cdot \mathcal{L}_{\text{diff}}^A + \lambda_E \cdot \mathcal{L}_{\text{diff}}^E$$其中每一项的形式相同:
$$\mathcal{L}_{\text{diff}}^\delta(\theta) = \mathbb{E}_{\epsilon^\delta \sim \mathcal{N}(0,1), t, C, l} \left[ \|\epsilon^\delta - \epsilon_\theta^\delta(z_t^\delta, t, C, l)\|_2^2 \right]$$4.6 训练策略:先学物理,再学控制
权重初始化自 ImageNet 预训练的 DiT(原标签嵌入层丢弃,文本嵌入层和动作编解码器零初始化)。训练分两步:
- 预训练:在 BridgeData-v2(60,000 条机器人轨迹)上训练 200k 步。
- 适配:在目标域(Metaworld 或真实世界)训练 100k 步。
适配阶段有个关键细节:$\lambda_I$ 始终保持 1.0,而 $\lambda_A$ 和 $\lambda_E$ 从 0 线性增加到 2.0(100k 步内完成)。换句话说,模型先专注学好图像预测,把物理理解的"地基"打牢,再逐步加入动作学习的压力。最终 $\lambda_A = 2.0 > \lambda_I = 1.0$,动作损失的权重反超图像,毕竟最终目的还是控制。
整个训练在 4 张 NVIDIA A100 上大约需要 3 天(预训练 2 天 + 适配 1 天)。推理用 75 步 DDIM 采样。
五、实验结果
Figure 1:PAD 在 Metaworld(50 任务)、真实世界已见任务和未见任务上的表现对比。红色 PAD 全面领先。
5.1 Metaworld 50 任务
用一个文本条件策略覆盖全部 50 个任务,每个任务仅有 50 条示范轨迹。PAD-XL/2(661M 参数,119.1 GFLOPs)达到 72.5% 的平均成功率。
5.2 真实世界:已见任务
在 Franka Panda 平台上测试 6 类任务(按钮、走线、抓取、放置、开抽屉、关抽屉),每个任务 200 条遥操作示范、50 次评估。PAD 平均成功率 72%,加上深度模态后提升至 78%。作为对比,RT-2* 为 69%、SuSIE 为 49%、Diffusion Policy 仅 38%。
5.3 真实世界:未见任务泛化
训练数据只有彩色方块和盘子,测试时换成从没见过的水果蔬菜玩具(草莓、胡萝卜、茄子)。难度分三级:Easy(1-4 干扰物)、Medium(5-15 干扰物)、Hard(未见物体 + 5-15 干扰或未见背景)。PAD 在所有级别上都显著领先,其他 baseline 在 Hard 级别几乎全军覆没。
Figure 5:泛化测试。黄色框标出目标位置。PAD 在面对从未见过的物体和干扰时依然能完成任务。
5.4 消融实验:图像预测是关键
这是整篇论文最有说服力的实验:
- 去掉图像预测(PAD w/o img):72.5% → 43.6%,暴跌 28.9 个百分点。这是最关键的发现:仅靠 4 维姿态监督,50 条示范根本不够让模型学会泛化,图像预测提供的高维像素级监督是性能的核心来源。
- 去掉协同训练(PAD w/o co-train):72.5% → 59.2%,下降 13.3 个百分点。大规模视频数据中的物理先验贡献显著。论文还展示了图像质量的对比:不协同训练时预测画面中方块直接消失了。
- 加入深度模态(PAD-Depth):真实世界从 72% 升到 78%,证明框架的多模态扩展性是实打实的。深度信息帮助模型判断距离,在精细操作中特别有用。
5.5 Scaling 特性
PAD 展示了清晰的计算量-性能正相关:
Figure 10:计算量(GFLOPs)与成功率的关系。对数刻度下接近线性,暗示 scaling law。
从 PAD-B/2(128M, 22.5 GFLOPs, 62.4%)到 PAD-XL/2(661M, 119.1 GFLOPs, 72.5%),计算量每增加一截都能换来稳定提升。但 PAD-XL/8(同样 661M 参数,只是 patch size 更大,GFLOPs 仅 7.7)只有 48.2%。这说明关键不在参数量,而在实际的计算密度,也就是 token 的空间细粒度。
六、总结
PAD 的思路很简洁:既然图像预测和动作生成都是扩散去噪,就让它们在同一个 DiT 里共享 attention、互相增强。最重要的贡献不在数字本身,而在于打通了一条路:机器人策略可以直接从海量互联网视频中学到物理理解,不需要每条视频都有动作标注。
它打开的可能性:任何新模态(触觉、点云、力传感器)都可以通过加一个编码器 + token 化接入同一个去噪过程。
留下的问题:推理速度不乐观(75 步 DDIM 联合去噪图像和动作),只验证了三种模态,真实世界实验的规模也还比较小。
七、Insight
训练时的"想象力"比推理时的"想象力"更有价值。
消融实验揭示了一个反直觉的事实:PAD 推理时只用第一个动作预测,图像预测直接扔掉;但一旦去掉图像预测的训练,性能暴跌 28.9 个百分点。这说明图像预测的价值不在于"推理时看到未来画面",而在于"训练时被迫建立物理表征"。它本质上是一种极其高效的正则化手段,用像素级的高维监督信号防止模型在稀疏的动作标签上过拟合。这个发现后来被 Fast-WAM、GigaWorld-Policy 等后续工作进一步验证,并发展成"训练时想象、推理时跳过"的设计范式。
八、关键引用
[Section 1] "Although the diffusion generative model and diffusion policy exhibit distinct capabilities—image prediction and robotic action, respectively—they technically follow a similar denoising process."
扩散生成模型和扩散策略看似能力不同(一个预测图像,一个生成动作),但技术上走的是同一条去噪路径。
[Section 4.4] "The absence of image prediction compromises the robot's ability to utilize the physical knowledge encoded in the image modalities, which may be crucial for robotic control."
去掉图像预测后,机器人就失去了利用图像模态中编码的物理知识的能力,而这些知识可能对控制至关重要。
[Section 4.4] "Predicting solely the robot pose provides only low-dimensional supervision signals, potentially leading to overfitting of the training data."
只预测机器人姿态提供的监督信号太低维了,容易导致模型在训练数据上过拟合。
[Section 4.5] "Our findings reveal a strong correlation between computational allocation (measured as transformer Gflops) and the success rate of the learned policy."
计算量(用 transformer GFLOPs 衡量)和策略成功率之间存在强正相关,暗示这个方向可能有 scaling law。
九、Q&A
Diffusion Policy 只对动作做扩散去噪,观测作为条件输入。PAD 则把图像和动作都当作需要去噪的目标,放在同一个 DiT 里联合处理。关键区别在于动作 token 和图像 token 在 self-attention 里互相可见:动作生成会受到图像预测的梯度信号影响,反过来也一样。这种双向信息流是 PAD 性能提升的核心来源。
两阶段方案(SuSIE)的瓶颈在于预测阶段不知道什么信息对动作有用,可能花大量精力还原光照细节、阴影变化等跟控制无关的像素,而物体的精确位置和接触状态(真正影响动作选择的信息)反而被淹没了。PAD 的联合去噪让动作侧的梯度信号回传到图像侧,引导图像预测把注意力放在对控制真正重要的特征上。
如果一上来就同时训练图像和动作,两个目标会互相拉扯,容易不收敛。所以先让模型专注学好图像预测($\lambda_A = 0$ 时模型只做图像去噪),等物理理解的"地基"打好,再逐步加入动作学习的压力。$\lambda_A$ 最终升到 2.0(超过 $\lambda_I = 1.0$),是因为最终目的是控制,后期需要让动作损失主导梯度方向。
通过 attention mask。纯视频数据只有图像,没有动作和深度,PAD 就把对应位置的 token 用 mask 屏蔽掉,不参与去噪计算。这样同一个训练循环里就能混合两种数据:有标注的机器人数据贡献图像 + 动作损失,纯视频只贡献图像预测损失。BridgeData-v2 的 60,000 条轨迹就是这样混合使用的。
核心原因是监督信号的维度差异。机器人姿态只有 4-7 维,50 条轨迹 × 4 维,每个任务大约只有 200 个标量的动作监督。而图像预测提供的是 32×32×4 = 4096 维的像素级监督,信息量高出三个数量级。图像预测强迫模型理解每个动作对世界的完整影响,这种理解会转化为更好的内部表征,最终反哺动作预测的准确性。
PAD-XL/8(patch size 8)只有 48.2%,PAD-XL/2(patch size 2)达到 72.5%,同样 661M 参数。patch size 越小,每张图像产生的 token 越多(patch 2 产生 256 个 token,patch 8 只产生 16 个),self-attention 能看到更细粒度的空间信息。机器人操作需要精确的空间推理(物体在哪、手指和物体的相对位置),粗粒度的 16 个 token 根本不够。GFLOPs 从 7.7 涨到 119.1,但成功率也从 48.2% 涨到 72.5%。
论文承认 "control frequency is not very high" 但没报具体数字。考虑到 75 步 DDIM 采样、661M 参数的 DiT、同时去噪图像和动作,单张 A100 上估计低于 5Hz。后续工作 DreamZero 在 14B 模型上通过系统优化做到了 7Hz,PAD 在 661M 规模上应该也可以通过减少采样步数等手段达到类似水平,只是论文没有做这方面的优化。
Figure 10 显示 GFLOPs 和成功率在对数刻度下接近线性关系,而 72.5% 远未饱和(理论上限是 100%),这暗示继续加大计算量(更大模型、更小 patch size)应该还能涨。不过论文的 backbone 是 ImageNet 预训练的 DiT,并非视频扩散模型。后续工作(DreamZero 用 14B Wan2.1、Cosmos Policy 用 2B Cosmos-Predict2)进一步证实了 scaling 的有效性。