精读：PAD——Prediction with Action

论文：Prediction with Action: Visual Policy Learning via Joint Denoising Process
方法名：PAD，Prediction with Action Diffuser
发表：NeurIPS 2024

这篇论文和你刚读的 GR-1 非常接近：二者都认为，机器人策略不应该只预测动作，还应该学习“接下来世界会变成什么样”。区别在于：

• GR-1：用自回归 Transformer 分别预测未来图像 token 和动作；

• PAD：把未来图像和动作都看作需要从噪声中恢复的变量，放进同一个 DiT 联合去噪过程。

PAD 的一句话概括是：

机器人一边“想象未来画面”，一边生成动作，而且这两件事在同一个扩散模型里共同完成。

论文认为，未来图像预测和动作生成都受到相同物理动力学支配，所以不应该像传统方法一样分开学习。

1. 论文试图解决什么问题？

1.1 传统 Diffusion Policy 只预测动作

标准 Diffusion Policy 学习的是：

其中：

• $o_t$：当前视觉观测；

• $l$：语言指令；

• $a_{t:t+H}$：未来动作序列。

它的优势是能够拟合多模态动作分布。例如，同一个物体可以从左边抓，也可以从右边抓，扩散模型可以保留多种可能性。

但它有一个问题：

动作是低维信号，单纯预测动作并不强迫模型真正理解场景将如何变化。

例如，机械臂执行“把红色方块放进碗里”：

• 只预测动作：输出机械臂轨迹；

• 同时预测未来图像：模型还必须知道红色方块将移动到碗中、机械臂位置如何变化、遮挡关系如何变化。

后者提供了更丰富的监督信号。

1.2 两阶段方法割裂了预测与控制

UniPi、SuSIE 等方法通常采用：

即先用生成模型想象目标图像，再用另一个策略根据目标图像生成动作。

问题在于：

1. 图像生成器和动作策略分开训练；

2. 图像成为两个模型之间的信息瓶颈；

3. 图像生成模型内部丰富的 latent feature 没有直接用于动作预测；

4. 图像生成错误会继续传递给低层策略。

PAD 则希望直接建模：

也就是：

• $x^I$：未来图像；

• $x^A$：未来机器人动作或位姿；

• $c^I_t$：当前图像；

• $c^A_t$：当前机器人状态；

• $l$：语言指令。

图像和动作不再是串联关系，而是联合生成关系。

2. PAD 的核心思想

PAD 建模如下联合分布：

其中：

模型不是只预测一步，而是同时预测未来 $k$ 步：

但真正执行时只执行第一个动作：

然后获取新的观测，再重新预测。这就是类似 MPC 的 receding-horizon control，滚动时域控制。

3. 整体架构

PAD 的核心架构可以写成：

随后：

也就是说，一个 DiT 同时预测不同模态中加入的噪声。

论文的 Figure 3 是理解方法最关键的一张图：不同模态先进入各自编码器，再形成统一 token 序列，在同一组 DiT block 内通过 self-attention 交互，最后同时恢复未来 RGB、动作和深度。

4. 输入是怎样构造的？

这是 PAD 最容易看混的部分。

4.1 当前观测作为条件

对于图像模态，当前图像经过冻结的 VAE encoder：

未来 $k$ 帧图像对应的 latent 被加入扩散噪声：

作者将当前条件 latent 和未来的 noisy latent 在 channel 维度拼接：

假设单张图像 latent 的尺寸为：

预测 $k$ 帧时，未来 noisy latent 的尺寸为：

所以拼接后的输入为：

直观上相当于：

模型的任务是保留当前观测作为条件，把后面的噪声逐渐还原为未来图像。

动作和深度采用相同思想：

这种设计类似条件图像生成中的 channel concatenation。

4.2 各模态怎样变成 token？

RGB 图像

输入分辨率是 $256\times256$。

冻结 VAE 将其压缩为：

在 PAD-XL/2 中，使用 $2\times2$ patch：

所以图像一共产生 256 个 token。

动作

动作序列先拼成一个向量，再经过 MLP 和线性投影。

值得注意的是：

PAD 把整段机器人动作序列压缩为 1 个 action token。

因此默认模型的 token 构成大致是：

深度模态加入后还会增加 16 个 depth tokens。

这会带来一个值得思考的问题：

图像在 self-attention 中拥有 256 个 token，而动作只有 1 个 token，图像的信息容量和计算占比显著更大。

作者通过动作损失权重来部分平衡这种模态不对称。

5. 为什么选择 DiT？

PAD 没有继续使用常见的 U-Net diffusion backbone，而是使用 Diffusion Transformer。

原因主要有三个。

5.1 不同模态可以直接拼成 token 序列

通过 self-attention，图像 token 和动作 token 可以直接交互：

$$

\operatorname{Attention}(Q,K,V)

\operatorname{softmax}
\left(
\frac{QK^\top}{\sqrt d}
\right)V

\text{先图像预测，再动作预测}

\text{图像预测}\leftrightarrow\text{动作预测}

D_{\text{video}}={I_{1:T}}

对于没有动作的视频，PAD：

1. 将动作位置填充为 padding；

2. 使用 attention mask 屏蔽这些位置；

3. 只计算有效模态的输出和损失。

因此同一个 DiT 可以处理：

• RGB-only 视频；

• RGB + action；

• RGB + action + depth。

这是 PAD 能够进行异构数据联合训练的关键。

5.3 可以从图像生成 DiT 初始化

PAD 使用 ImageNet 上预训练的 DiT 权重初始化。

但原 DiT 只有图像生成能力，没有语言、动作、深度模块，因此作者：

• 删除原来的类别 label embedding；

• 新增 text embedding 层并置零初始化；

• 动作 encoder、decoder 置零初始化；

• 对图像 tokenizer 权重复制 $k+1$ 次，以支持当前帧和未来 $k$ 帧的 channel stacking。

置零初始化的意义是：

新增加的动作和文本分支在训练初期不会严重扰乱已经学到的图像生成能力。

6. 联合扩散过程

先回顾标准扩散模型。

给定干净 latent $z_0$，前向加噪为：

神经网络学习预测噪声：

其损失为：

$$

\mathcal L_{\mathrm{diff}}

\mathbb E
\left[
\left|
\epsilon-
\epsilon_\theta(z_t,t,c)
\right|_2^2
\right]

其中：

最终联合目标为：

$$

\boxed{
\mathcal L(\theta)

\lambda_I\mathcal L^I_{\mathrm{diff}}
+
\lambda_A\mathcal L^A_{\mathrm{diff}}
+
\lambda_E\mathcal L^E_{\mathrm{diff}}
}

a_t\in\mathbb R^4

a_t\in\mathbb R^7

x_{t+1}\in\mathbb R^{H\times W\times3}

s_{t+1}=f(s_t,a_t)

p(o_{t+1}\mid o_t,a_t)

p(a_t\mid o_t,l)

p(o_{t+1},a_t\mid o_t,l)

\text{robot video pretraining}
\rightarrow
\text{robot policy improvement}

\text{general web video}
\rightarrow
\text{robot control improvement}

\lambda_I=1

\lambda_A,\lambda_E:0\rightarrow2

k=3

i=4

t+4,\quad t+8,\quad t+12

(\hat I_{t+1:t+k},\hat a_{t:t+k-1})
\sim
p_\theta(\cdot\mid I_t,s_t,l)

a_t=\hat a_t

I_{t+1},s_{t+1}
\leftarrow
\operatorname{EnvironmentStep}(a_t)

72.5-57.4=15.1

\frac{72.5-57.4}{57.4}
\approx26.3%

78%-72%=6%

72.5-43.6=28.9

\text{image prediction auxiliary task is useful}

\text{world dynamics understanding}

\text{representation regularization}

72.5-59.2=13.3

\begin{aligned}
\text{no image prediction}&:43.6%\
\text{image prediction without video pretraining}&:59.2%\
\text{image prediction with video pretraining}&:72.5%
\end{aligned}

48.2%\quad\text{vs.}\quad72.5%

24.3

17\quad\text{vs.}\quad257

\text{GR-1: }
p(x,a)=\prod_i p(y_i\mid y_{<i})

\text{PAD: }
(x,a)\sim p_\theta(x,a\mid c,l)

a_t

\hat I_{t+1}

\text{latent world prediction}
\rightarrow
\text{action head}

\text{pixel generation}
+
\text{action generation}

\text{image prediction quality}
\rightarrow
\text{action success}

k=3

T_I\leftrightarrow T_A

T_I^{(j)}
\rightarrow
a_t

z_j\leftarrow0

z_j\leftarrow z_j’

\Delta a=
\left|
a(z)-a(z_{-j})
\right|

\operatorname{CausalImportance}(z_j\rightarrow a)

m_t=\operatorname{Select}
\left(
z_{t+1:t+k}
\right)

a_t=
\pi
\left(
o_t,
m_t
\right)

\mathcal M=
{(z_{\text{context}},z_{\text{causal}},a,r)}

m_t=\operatorname{Retrieve}(\mathcal M,o_t,l)

p(x_I,x_A\mid c_I,c_A,l,m_t)

最值得质疑的地方

论文把性能提升解释为“学习到了物理知识”，但目前实验更直接证明的是：

至于模型是否真正学到了可组合、可迁移、可用于规划的物理规律，证据仍然不足。

18. 最终总结

PAD 的整体计算链可以浓缩为：

将未来目标加噪：

输入统一 DiT：

联合优化：

$$

\mathcal L

\lambda_I\mathcal L_I
+
\lambda_A\mathcal L_A
+
\lambda_E\mathcal L_E

(\hat x_I^{1:k},\hat x_A^{1:k})
\sim
p_\theta(\cdot\mid c_I,c_A,l)

\hat x_A^1