论文笔记（八十三）STACKGEN: Generating Stable Structures from Silhouettes via Diffusion

文章概括

引用：

@article{sun2024stackgen,title={StackGen: Generating Stable Structures from Silhouettes via Diffusion},author={Sun, Luzhe and Yoneda, Takuma and Wheeler, Samuel W and Jiang, Tianchong and Walter, Matthew R},journal={arXiv preprint arXiv:2409.18098},year={2024}
}

Sun, L., Yoneda, T., Wheeler, S.W., Jiang, T. and Walter, M.R., 2024. StackGen: Generating Stable Structures from Silhouettes via Diffusion. arXiv preprint arXiv:2409.18098.

主页：https://ripl.github.io/StackGen/
原文：https://ripl.github.io/StackGen/static/StackGen.pdf
代码、数据和视频：

系列文章：
请在 $《$文章$》$ 专栏中查找

---

宇宙声明！

引用解析部分属于自我理解补充，如有错误可以评论讨论然后改正！

---

摘要

人类通过观察和与世界互动，自然而然地获得对刚性物体相互作用及稳定性的直觉。正是这种直觉支配了我们在环境中配置物体的方式，使我们能够利用简单的日常物体构建复杂结构。另一方面，机器人代理传统上需要一个包含每个物体的详细几何信息和环境动态分析模型的显式世界模型，而这种模型难以扩展且无法实现泛化。相反，机器人将受益于对直觉物理的认知，使其能够以类似方式推理环境中物体的稳定相互作用。为实现这一目标，我们提出了 STACKGEN——一种扩散模型，用于生成与目标轮廓匹配的多样化稳定积木配置。 为了展示该方法的能力，我们在模拟环境中对其进行了评估，并在真实环境中部署了机械臂以组装模型生成的结构。我们的代码可在 https://ripl.github.io/StackGen 获取。

I. INTRODUCTION

理解场景物理原理是执行许多物理任务（例如堆叠、（拆）装以及移动物体）的先决条件。人类可以通过视觉线索、力反馈和经验知识的结合直观地评估和预测结构的稳定性。另一方面，机器人缺乏自然的多模态传感整合和对直觉物理学的理解。传统上，机器人依赖一个世界模型，该模型包括对环境中物体详细几何形状的表示以及支配它们相互作用的动力学解析模型。这种依赖性给在未准备环境中部署机器人代理带来了重大挑战。

将多种积木组合成稳定结构的能力作为研究智能体对物体组合与交互理解的测试平台，有着悠久的历史 [1, 2, 3, 4]。虽然看似原始，但这一能力具有许多实际意义，例如机器人辅助建造 [5, 6, 7, 8, 9]，并可作为智能体处理复杂真实物体集合的下游应用的支柱。

基于对物理的直觉理解而构建3D结构的当代方法，将场景预测的前向动力学作为规划器的一部分，用于将积木组合成目标结构。这通常包括首先训练一个作为直觉物理引擎的前向动力学模型，然后通过某种拒绝采样的方式，利用该模型模拟候选物体放置的行为。这种方法成本很高，因为它需要在大范围的坐标空间中搜索，并为每种可能的积木放置建模动力学。

我们不训练前向动力学模型，而是考虑学习并生成组成稳定3D结构的物体的$\text{SE(3)}$姿态的联合分布（图1）。我们以用户提供的结构规范为条件，对该分布进行约束，使他们能够在测试时控制生成。

图1：STACKGEN 包含一个扩散模型，该模型以目标结构轮廓和可用积木形状列表作为输入。然后，模型生成一组积木姿态 { p ^ 1 , … , p ^ k } \{\hat p_1,\dots,\hat p_k\} {p^​1​,…,p^​k​}，用于构造与目标轮廓一致的稳定结构。最后，可以使用机器人手臂来搭建所得到的结构。

受其在计算机视觉 [10, 11, 12, 13, 14] 领域以及近年来机器人领域 [15, 16, 17] 取得的成功启发，我们采用条件扩散模型 [18]——一种在各种生成领域中表现出色的生成模型家族——在我们的工作中用于生成稳定的6-DoF物体姿态。与那些通过草图或轮廓等空间信息控制图像生成的方法 [19] 在精神上类似，我们要求用户提供一个模糊描述所需结构的轮廓，并将其用作条件信号。不同于 Zhang 等人 [19]，我们仅训练一个基于 Transformer 架构的条件扩散模型。需要注意的是，与标准图像生成不同，我们的方法生成的是一组姿态。我们旨在生成能够形成物理上稳定结构的姿态。

我们的模型（STACKGEN）在不同积木的6-DoF姿态上进行推理，以实现其组合作为与不同用户提供的目标规范相一致的稳定3D结构。在接下来的章节中，我们将描述支撑我们扩散模型的基于 Transformer 的架构，以及我们用于生成用于训练和评估的稳定积木配置的过程。我们通过基线比较以及真实场景实验（展示其在使用 UR5 机械臂进行真实场景生成中的优势）来评估 STACKGEN 的能力。

II. 相关工作

A. 从直觉物理学学习稳定性

与我们的工作类似，一些研究 [2, 3, 15] 考虑了相对简单物体的交互，以探讨直觉物理学的概念。在考虑用于评估稳定性的视觉信号时，ShapeStacks [20] 成功地在单链堆叠场景中学习了凸对象的物理特性。具体做法是垂直堆叠物体、计算其质心（CoM），并扩充数据集，通过监督学习训练视觉模型，从而在堆叠之前实现稳定性预测。然而，对于多链堆叠中的组合，计算质心要复杂得多。另一种直觉物理学形式涉及预测一组物体状态随时间演变的能力，其中包括连续性和对象永久性的概念。这促使人们开发了衡量模型在此类任务上能力的基准，称为违反期望（VoE） [21, 22, 23, 24]。在机器人领域，Agrawal 等人 [25] 收集了机器人机械臂拨动物体的视频序列。利用该数据集，他们从像素输入训练了前向和逆向动力学模型，并演示了该模型能够让机器人推理出一系列合适的拨动动作以实现目标图像。
其他研究也采取了类似方法 [26, 27]。

B. 用于姿态生成的扩散模型

鉴于扩散模型在学习多模态分布方面的惊人能力，许多研究 [28] 利用扩散模型来学习支持机器人规划的$SE(3)$姿态分布 [29, 30, 31]。Urain 等人 [29] 使用条件扩散模型来预测机器人操作中基于目标物体形状的合理末端执行器位置。Simeonov 等人 [30] 通过对场景中物体与其环境之间的空间关系建模，使用扩散模型来预测物体的最优放置位置，以识别架设、堆叠或悬挂等任务的目标姿态。他们的方法将3D点云重建作为上下文信息纳入，以确保预测姿态在真实场景中既功能性强又可行。Liu 等人 [32] 和 Xu 等人 [33] 将大型语言模型与组合扩散模型相结合，以分析用户指令并生成所需物体放置的基于图的表示。他们随后通过优化联合目标来预测物体排列模式，有效地将语言理解与空间推理融合。

C. 自动化顺序装配

与我们的数据生成过程相关，Tian 等人 [34] 提出了一种依赖于逆向拆解过程的装配方法（ASAP），其中每个组件被放置在唯一位置以保证物理可行性。然而，由于该装配场景假设组件与特定位置之间存在一对一映射，这种方法并未考虑多种组合可能引发的结构不稳定性。相比之下，我们的工作通过一对多映射方法，仅使用2D轮廓来解决保持重力稳定结构的问题。该方法确保在过渡到3D环境时结构稳定性得以保留并准确再现。我们的重点在于生成和验证结构稳定的积木配置，而非优化装配顺序。与 ASAP 方法类似，该方法生成在重力作用下中间配置保持稳定的逐步装配序列，我们提出了一种“通过拆解来构建”的方法，该方法通过预测多样的稳定配置，实现可扩展的数据生成，无需依赖预定义的装配路径。在本节中，我们描述了用于生成共同构成与用户提供的场景规范一致的稳定结构的积木$SE(3)$姿态的基于扩散的框架。然后，我们讨论模型训练的流程，包括生成包含多样稳定积木配置的训练集的方法。

III. 方法

A. 用于$SE(3)$积木姿态生成的扩散模型

我们的模型（图2）生成创建一个既符合给定条件（例如轮廓）又稳定的3D结构所需的$\text{SE(3)}$积木姿态。我们框架的基础是一个基于Transformer的扩散模型，该模型表示稳定六自由度姿态的分布，而无需显式指定其组成积木的数量、类型或位置。通过这种方式，模型利用反向扩散过程生成能够共同构成稳定结构的积木姿态。我们另外训练了一个卷积神经网络（CNN），根据目标轮廓预测构建所需的积木数量和类型。在测试时，我们首先使用CNN预测积木列表，然后将该列表和目标轮廓一起作为输入提供给扩散模型。扩散模型随后对可能构成稳定结构的积木姿态进行采样。

图2：STACKGEN基于Transformer的架构可视化。我们首先将噪声姿态 { p ~ 1 t , … , p ~ k t } \{\tilde p_1^t,\dots,\tilde p_k^t\} {p~​1t​,…,p~​kt​}、形状嵌入 { s 1 , … , s k } \{s_1,\dots,s_k\} {s1​,…,sk​}和时间步$t$映射到相同的维度$d$，然后相加以构造对象令牌，$\in {\mathbb{R}}^{k\times d}$。随后，我们对目标轮廓进行分块（patchify），并将每个分块嵌入到维度$d$，得到 { i 1 , … , i c } \{i_1,\dots,i_c\} {i1​,…,ic​}，再与位置编码${\mathrm{Pos}}_{1:c}$相加，以构造轮廓令牌，$\in {\mathbb{R}}^{c\times d}$。接着，我们将对象令牌、填充令牌和轮廓令牌一起输入图中的去噪Transformer $D_{\theta }$（式(2)），以预测噪声 { ϵ ^ 1 t , … , ϵ ^ k t } \{\hat\epsilon_1^t,\dots,\hat\epsilon_k^t\} {ϵ^1t​,…,ϵ^kt​}。最后，通过运行式(4)得到下一个反向状态 { p ~ 1 t − 1 , … , p ~ k t − 1 } \{\tilde p_1^{t-1},\dots,\tilde p_k^{t-1}\} {p~​1t−1​,…,p~​kt−1​}。

这个流程其实就是把“带噪乱序的物体姿态”＋“目标轮廓”一起丢进一个 Transformer
去噪，然后按照扩散模型的公式一步步还原出整齐的排列。我们可以把它拆成四个阶段来看：

---

1. 对象令牌（Object Tokens）构建

1. 输入

   • 当前扩散步 $t$ 下的所有物体姿态： { p ~ 1 t , … , p ~ k t } \{\tilde p_1^t,\dots,\tilde p_k^t\} {p~​1t​,…,p~​kt​}，每个 ${\tilde{p}}_i^t\in {\mathbb{R}}^m$（比如一个物体的中心点、朝向等参数）。
   • 每个物体对应的形状嵌入： { s 1 , … , s k } \{s_1,\dots,s_k\} {s1​,…,sk​}，每个 $s_i\in {\mathbb{R}}^d$。
   • 当前的扩散时间步标记 $t$（标量），也会被映射到同样的维度。

2. 线性投影

   • 把 ${\tilde{p}}_i^t$、$s_i$ 和 $t$ 各自通过三条不同的全连接（FC）层，投影到同一维度 $d$。
   • 得到三组向量：${\mathrm{Proj}}_p({\tilde{p}}_i^t)\in {\mathbb{R}}^d$、${\mathrm{Proj}}_s(s_i)\in {\mathbb{R}}^d$、${\mathrm{Proj}}_t(t)\in {\mathbb{R}}^d$。

3. 相加融合 $${\text{object\_token}}_i={\mathrm{Proj}}_p({\tilde{p}}_i^t)+{\mathrm{Proj}}_s(s_i)+{\mathrm{Proj}}_t(t)\in {\mathbb{R}}^d,i=1,\dots ,k$$ 这就得到了 $k$ 个“对象令牌”（$\in {\mathbb{R}}^{k\times d}$），每个令牌里同时编码了物体的当前位置、形状信息和当前是第几步。

---

2. 轮廓令牌（Silhouette Tokens）构建

1. Patchify

   • 把目标轮廓（二值化的 silhouette）划分成 $c$ 个小块（patch），常见做法是把图像切成 $N\times N$ 的网格，每块大小比如 $16\times 16$ 像素。

2. Patch 嵌入

   • 对每个小块做一次线性投影（Patch Projection），把它展平成一个向量，再投到维度 $d$。
   • 记作 { i 1 , … , i c } \{i_1,\dots,i_c\} {i1​,…,ic​}，每个 $i_j\in {\mathbb{R}}^d$。

3. 位置编码

   • 给每个 patch 加上自己在网格中的位置编码 ${\mathrm{Pos}}_j\in {\mathbb{R}}^d$。
   • 令：
     $${\text{silh\_token}}_j=i_j+{\mathrm{Pos}}_j,j=1,\dots ,c.$$
   • 得到 $c$ 个“轮廓令牌”（$\in {\mathbb{R}}^{c\times d}$）。

---

3. 拼接并送入去噪 Transformer

1. 序列拼接

   • 把上面得到的三部分令牌在“序列长度”维度上串起来：

     [ object_token_1, …, object_token_k,padding_token (若 k+c < L 则用若干),silh_token_1, …, silh_token_c ]
     

   • 这里通常会固定一个最大长度 $L$，多余位置用特殊的 padding_token 占位（Transformer 不会对它们进行预测）。

2. 去噪 Transformer $D_{\theta }$

   • 这个 Transformer 接受长度为 $L$、每个维度为 $d$ 的 token 序列，经过多层自注意力和前馈层，输出同样维度的序列。
   • 最终在前 $k$ 个位置读取输出，得到每个物体的噪声预测 { ϵ ^ 1 t , … , ϵ ^ k t } \{\hat\epsilon_1^t,\dots,\hat\epsilon_k^t\} {ϵ^1t​,…,ϵ^kt​}。

---

4. 反向一步更新姿态（Eq. (4)）

根据扩散模型的反向采样公式（图中标注的 Eq.(4)）： $${\tilde{p}}_i^{t-1}=\frac{1}{\sqrt{{\alpha }_t}}({\tilde{p}}_i^t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\widehat{ϵ}}_i^t)+{\sigma }_{t}z,z\sim \mathcal{N}(0,I),$$ 把预测到的噪声 ${\widehat{ϵ}}_i^t$ 和当前的 ${\tilde{p}}_i^t$ 结合，就能得到更新后、稍微“干净”一些的姿态 ${\tilde{p}}_i^{t-1}$。然后迭代：把它当新的输入，继续做第 $t-1$、$t-2$ …… 直到 $t=0$，就得到了最终物体的排列。

---

整体逻辑小结

1. 初始化：所有物体姿态从纯噪声开始；给定目标轮廓（silhouette）。
2. 迭代去噪：在每个扩散步 $t$，把「当前带噪姿态 + 形状 + 时间步」和「目标轮廓」一起送进 Transformer，预测噪声残差。
3. 更新姿态：用扩散反向公式消一点噪，${\tilde{p}}^t\to {\tilde{p}}^{t-1}$。
4. 终止：最后得到的 ${\tilde{p}}^0$ 就是我们希望的、与轮廓匹配的无噪声姿态排布。

这样，你就把目标外形（轮廓）条件化到了一个多物体的扩散模型里，让它一步步把乱序噪声“吸”到符合轮廓的那个排列上。

我们采用去噪扩散概率模型（DDPM）[28]作为STACKGEN的核心框架。概率扩散模型[35]是一种由正向扩散过程和反向扩散过程组成的生成模型。正向过程是一个一阶马尔可夫链，将噪声逐步加入从数据分布$q(p^1)$中抽取的样本$p^1\sim q(p^1)$中，而反向过程则是一个马尔可夫链，迭代地对噪声输入$p^T\sim \mathcal{N}(0,I)$进行去噪。模型通过学习如何加入和去除噪声，从而能够从目标数据分布$q(p^1)$中生成样本。

在STACKGEN中， p 1 = { p 1 1 , p 2 1 , … , p k 1 } p^1=\{p^1_1,p_2^1,\dots,p_k^1\} p1={p11​,p21​,…,pk1​}是物体姿态的集合，其中$p_i^1\in {\mathbb{R}}^6$是第$i$个积木的$\text{SE(3)}$姿态，而$q(p^1)$是所有积木姿态构成稳定堆叠的分布。（我们在应用扩散框架前对姿态进行归一化；在推理时将生成的姿态逆归一化。）遵循DDPM，STACKGEN使用以下正向扩散过程注入噪声：

$${\tilde{p}}_i^t=\sqrt{{\bar{\alpha }}_t}{p}_i^1+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_i,{ϵ}_i\sim \mathcal{N}(\text{0},\text{I}),\text{\hspace{1em}(1)}$$

其中 t ∼ U n i f { 1 , 2 , … , T } t\sim\mathrm{Unif}\{1,2,\dots,T\} t∼Unif{1,2,…,T} 是定义噪声规模的扩散时间步，${\bar{\alpha }}_t$是由噪声调度决定的系数，${ϵ}_i$是每个步骤注入的噪声。$p_i^1$是第 $i$ 个积木的稳定姿态，${\tilde{p}}_i^t$是时间步$t$下给$p_i^1$加入噪声之后的样子。STACKGEN将噪声姿态${\tilde{p}}_i^t(i=1,2,\dots ,k)$连同扩散时间步$t$、形状嵌入 { s 1 , … , s k } \{s_1,\dots,s_k\} {s1​,…,sk​}和积木轮廓$S$一起提供给我们的去噪网络（图2中的Transformer）$D_{\theta }$，得到

$${\widehat{ϵ}}_{1:k}=D_{\theta }({\tilde{p}}_{1:k},t,s_{1:k},S),\text{\hspace{1em}(2)}$$

其中记号$X_{1:k}$等价于 { X 1 , … , X k } \{X_1,\dots,X_k\} {X1​,…,Xk​}，${\widehat{ϵ}}_i$是对第$i$个积木的预测噪声。有了这些预测噪声，单样本的训练目标为
$$\frac{1}{k}\sum _{i=1}^k\Vert {ϵ}_i-{\widehat{ϵ}}_i{\Vert }^2.\text{\hspace{1em}(3)}$$
在每个训练步骤中，我们从 { 1 , 2 , … , T } \{1,2,\dots,T\} {1,2,…,T}中均匀随机采样扩散时间步$t$。

一旦去噪网络训练完成，采样过程从高斯分布中采样噪声姿态${\tilde{p}}_i^T\sim \mathcal{N}(\text{0},\text{I})$开始。然后从$t=T$迭代到1执行以下步骤：
$${\tilde{p}}_i^{t-1}=\frac{1}{\sqrt{{\alpha }_t}}({\tilde{p}}_i^t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\widehat{ϵ}}_i^t)+{\sigma }_{t}z,z\sim \mathcal{N}(\text{0},\text{I}),\text{\hspace{1em}(4)}$$
其中${\widehat{ϵ}}_i$由式(2)给出，$z\sim \mathcal{N}(\text{0},\text{I})$当$t>1$时，否则$z=0$。最终得到的${\tilde{p}}_{1:k}^1$便是生成的稳定姿态。

B. 模型架构

任务的挑战性需求来自于：1）模型必须能够处理可变数量的积木姿态，因为不同的堆叠使用不同数量和形状的积木；2）模型必须处理来自不同模态的输入，包括具有空间信息的姿态、形状和轮廓。

我们的模型（图2）基于Transformer架构[36]构建，可以处理可能来源于不同模态的输入令牌。为了初始化过程，我们使用卷积神经网络（CNN）从结构的轮廓预测积木列表，如图1所示。该CNN将单通道掩码输入转换为十通道，然后经过三十个残差层、最大池化和两个全连接层。

在训练时，我们使用整数索引对结构中立方体、矩形、长矩形和三角形的数量进行唯一编码，并通过用参数θ对CNN $C_{\theta }$进行训练，以使用交叉熵损失模拟积木数量的联合分布，该联合分布由每个索引的类别概率表示：
$$L(D,\theta )=\frac{1}{|D|}\sum _{(S_i,y_i)\in D}-\log \frac{\exp (C_{\theta }(y_i\mid S_i))}{{\sum }_{y_k}\exp (C_{\theta }(y_k\mid S_i))}\text{\hspace{1em}(5)}$$.
其中$D=\{(S_i,y_i)\}$是我们的有标签训练集，$S_i$是结构的轮廓，$y_i$是对应于该结构积木列表的索引。该预测的积木列表作为模型后续步骤的输入之一。

这段话其实是在讲，第一步我们用一个 CNN（记作 $C_{\theta }$）把“结构的轮廓”映射到一个“积木列表”的预测上，而这一步的训练目标就是让 CNN 在所有可能的积木列表中，给出真正那个列表的高概率。公式 (5) 就是经典的交叉熵损失（cross‐entropy loss），它的含义可以拆开来看：

---

1. 数据与目标

• 我们有一个有标签训练集 $$D=\{(S_i,y_i)\},$$ 其中
  • $S_i$ 是第 $i$ 个样本的“轮廓”输入（单通道掩码图像）。
  • $y_i$ 是对应的“积木列表”的整数索引（这里他们事先把“立方体、矩形、长矩形、三角形”四种积木的数量编码成一个唯一的整数标签）。

---

2. CNN 的输出与 Softmax

• CNN $C_{\theta }$ 对每个可能的标签 $y_k$ 都打出一个“logit”分数： $$C_{\theta }(y_k\mid S_i)\in \mathbb{R}.$$
• 然后通过 Softmax 变成概率分布： $$P_{\theta }(y_k\mid S_i)=\frac{\exp (C_{\theta }(y_k\mid S_i))}{{\sum }_j\exp (C_{\theta }(y_j\mid S_i))}.$$

---

3. 交叉熵损失

• 对于一个样本 $(S_i,y_i)$，我们希望模型给出“真实标签” $y_i$ 的概率 $P_{\theta }(y_i\mid S_i)$ 尽可能高。一个自然的度量就是取负对数概率： $$-\log P_{\theta }(y_i\mid S_i)=-\log \frac{\exp (C_{\theta }(y_i\mid S_i))}{{\sum }_k\exp (C_{\theta }(y_k\mid S_i))}.$$
• 公式 (5) 把所有样本的这一项加起来，再除以样本数 $|D|$，就是平均的交叉熵损失： $$L(D,\theta )=\frac{1}{|D|}\sum _{(S_i,y_i)\in D}-\log \frac{\exp (C_{\theta }(y_i\mid S_i))}{{\sum }_{y_k}\exp (C_{\theta }(y_k\mid S_i))}.$$

---

4. 换句话说

• 训练时：给定一个轮廓 $S_i$，CNN 会输出一个对每个可能“积木组合索引” 的打分向量。
• 损失评估：对真正对应的那个索引 $y_i$，我们看模型给了多大概率（Softmax 后的输出），取它的负对数就是这个样本的损失。
• 优化目标：让所有样本的平均负对数概率尽可能小，也就是让 CNN 学会在轮廓和对应积木列表之间建立正确的映射。

这就是公式 (5) 的全部含义——标准的多分类交叉熵。它保证模型会把真实标签预测成最高的概率。

对于包含$k$个积木的稳定堆叠场景，我们提取它们的姿态列表$p_{1:k}\in {\mathbb{R}}^{k\times 6}$和形状嵌入$s_{1:k}\in {\mathbb{R}}^{k\times d}$，其中$d=512$（图2）。对于姿态$p_i$，我们使用由平移的笛卡尔坐标和旋转的指数坐标组成的六维表示。形状嵌入$s_i$从一个码本中检索，该码本为每种形状存储了唯一的可训练嵌入。姿态通过一个分别应用于每个对象的多层感知机（MLP）投射到$d$维空间。同样地，我们将扩散时间步$t\in [1,T]$投射到$d$维嵌入空间。然后对每个对象，将姿态、形状和时间步嵌入相加，得到$k$个对象令牌。为了处理可变数量的积木，我们向Transformer编码器输入固定数量$N$个对象令牌，并根据需要用零向量填充剩余的$N-k$个令牌。

1. 输入：姿态与形状的表示

• 场景里有 $k$ 块积木，我们要预测它们的三维“姿态”（平移+旋转）和已知的“形状类别”。

• 姿态 $p_i\in {\mathbb{R}}^6$：

  • 前 3 维是笛卡尔坐标 $(x,y,z)$（物体中心的位置）。
  • 后 3 维是旋转的指数坐标（用一个 3 维向量表示在空间中绕某个轴旋转的“角度”）。

• 形状嵌入 $s_i\in {\mathbb{R}}^{512}$：

  • 事先为“立方体、长方体、三角形……”等每种形状准备了一条 512 维的可训练向量，储存在一个“形状码本”里。

---

2. 把姿态、形状、时间步都映射到同一个空间

1. 姿态投影

   • 对每个 $p_i\in {\mathbb{R}}^6$，用一个小型的多层感知机（MLP）把它映射到 512 维：

     $${\mathrm{Proj}}_p(p_i)\in {\mathbb{R}}^{512}.$$

2. 形状嵌入

   • 从码本里直接查到 $s_i\in {\mathbb{R}}^{512}$。

3. 时间步嵌入

   • 扩散模型在第 $t$ 步运行时，把标量 $t$ 通过另一条 MLP 映射到 512 维：

     $$ \mathrm{Proj}_t(t);\in;\mathbb R^{512}.

4. 三者相加

   • 得到每个积木的对象令牌（object token）：

     $${\text{obj}}_i={\mathrm{Proj}}_p(p_i)+s_i+{\mathrm{Proj}}_t(t)\in {\mathbb{R}}^{512}.$$

   • 一共 $k$ 个这样的向量，组成矩阵 ${\mathbb{R}}^{k\times 512}$。

---

3. 支持可变数量的积木：填充或截断

• 模型在硬件上预设能“看见”最多 $N=10$ 个对象令牌。
• 如果真实场景里 $k<10$，就用全零向量把剩下的 $10-k$ 个槽位填满（padding）； 如果 $k>10$，一般在数据预处理时保证不会超过这个上限。

积木结构的轮廓$S$表示为大小为$64\times 64$的二值图像$I$。按照Dosovitskiy等人[37]的方法，我们将二值图像$I$分成$c=16$个大小为$16\times 16$的块，并分别使用两层MLP对每个块进行编码，以得到轮廓令牌$i_{1:c}\in {\mathbb{R}}^{c\times d}$。我们向轮廓令牌中添加正弦位置编码${\mathrm{Pos}}_{1:c}\in {\mathbb{R}}^{c\times d}$以保留空间信息。然后我们将块嵌入与位置编码相加，得到$c$个轮廓令牌。

4. 轮廓令牌（Silhouette Tokens）：用 ViT 风格处理轮廓图

• 输入轮廓：把积木堆叠渲染成一张 $64\times 64$ 的二值图像 $I$（前景1 / 背景0）。

• Patchify：按 $16\times 16$ 划分成 $c=16$ 块（每块是 $16\times 16\times 1$ 通道）。

• 块嵌入：对每个小块，用两层 MLP 投影到 512 维，得到 $i_j\in {\mathbb{R}}^{512}$，共 $j=1,\dots ,16$。

• 位置编码：加上正弦-余弦位置编码 ${\mathrm{Pos}}_j\in {\mathbb{R}}^{512}$，保留块在图中的空间顺序

  $${\text{silh}}_j=i_j+{\mathrm{Pos}}_j(j=1,\dots ,16).$$

STACKGEN随后将对象令牌、填充令牌和轮廓令牌连接，并输入一个六层的Transformer编码器。该Transformer使用隐藏维度$d=512$，并通过维度$d_{\text{ff}}=2048$的前馈网络扩展表示。在自注意力层中，我们使用$h=8$个自注意力头，每个头的维度为$\frac{d}{h}=64$。

5. 序列拼接：对象令牌 + 填充令牌 + 轮廓令牌

• 最终把三部分拼成一条长度为 $N+c=10+16=26$ 的序列，送入 Transformer 编码器。

---

6. Transformer 编码器参数

• 层数：6 层 Encoder。
• 隐藏维度 $d=512$。
• 前馈网络：每层有一个内层维度 $d_{\mathrm{ff}}=2048$ 的两层 MLP。
• 多头自注意力：头数 $h=8$，每头维度 $512/8=64$。

在编码器的最后一层，STACKGEN将每个上下文化的积木令牌线性投射回姿态空间${\mathbb{R}}^6$，并与对应姿态中加入的原始噪声计算均方误差（MSE）损失，以按照DDPM框架进行监督学习。图2总结了该过程和架构。

在论文报道的实验中，STACKGEN使用$N=10$、$T=50$，以及基于超参数调优的线性噪声调度，从$[10^{-4},0.189]$。

7. 输出与训练损失

1. 线性投影回姿态空间

   • 在最顶层 Transformer 输出中，取前 $k$ 个“真实”对象令牌（不含填充），每个 512 维，再用一条线性层映射回 6
     维：

     $${\widehat{ϵ}}_i\in {\mathbb{R}}^6.$$

2. DDPM 框架下的监督

   • 训练时，我们对真实姿态 $p_i^0$ 加噪得到 ${\tilde{p}}_i^t$，网络预测噪声残差
     ${\widehat{ϵ}}_i$。 * 用均方误差（MSE）衡量预测噪声和真实噪声 ${ϵ}_i$ 之间的差：

     $$\mathrm{MSE}=\Vert {\widehat{ϵ}}_i-{ϵ}_i{\Vert }^2.$$ * 所有对象、所有时间步的 MSE 平均起来，就是最终的训练目标。

---

8. 关键超参数

• 对象槽数 $N=10$。
• 扩散总步数 $T=50$。
• 线性噪声调度：噪声强度从 $10^{-4}$ 线性增加到 $0.189$（通过超参数调优确定）。

---

整体流程小结

1. 准备令牌

   • 对每块积木：把姿态、形状标签、当前扩散步映射到 512 维相加 → 对象令牌。 * 对轮廓图：划分 Patch → MLP 投影 → 加位置编码 → 轮廓令牌。

2. 拼接序列 → 送入 6 层、512 维、8 头的 Transformer。

3. 输出预测 → 线性映射回 6D 噪声残差 → 用 MSE 在 DDPM 框架下监督去噪学习。

4. 迭代 $t=1\dots 50$ → 最终恢复出符合轮廓的、稳定堆叠的积木姿态。

这样，你就能看到 STACKGEN 如何把“任意数量的积木姿态”“形状信息”“轮廓图像”三条信息流，统一在同一个 Transformer 里去降噪、融合、预测，最终得出稳固的3D堆叠方案。

C. 数据生成

为了训练能够生成多样化稳定积木姿态的模型，数据集的质量和多样性至关重要。我们希望拥有一种算法，能够以合成方式大规模采样各种稳定的积木配置以生成此类数据集。如果我们过分强调积木堆叠的多样性，一种简单而通用的方法是：在仿真中随机生成并释放一个随机选择的形状于随机姿态，等待其稳定落下，然后重复此过程，直到在场景中构建出一个有意义的堆叠（例如通过检查其高度或积木之间的碰撞）。如果该过程未能形成堆叠，我们可以将其拒绝并重新开始，重复此过程。这种方法有可能产生一个非常通用且极其多样化的积木堆叠数据集，然而已被证明效率低下且不切实际。

一、为什么要合成数据？

• 目标是训练一个模型，能够“从轮廓”生成多样化且物理稳定的积木姿态。
• 要做到这一点，就需要一个足够大、足够多样的“真实”或“近似真实”数据集。
• 直接从现实世界收集数据成本高；用合成仿真自动生成就显得尤为必要。

---

二、直接随机释放法（效率低下）

1. 思路

   • 随机在空中生成一块随机形状（立方体、长方体、三角形……），随即给予一个随机的姿态（位置+朝向）。
   • 在物理仿真中“放下”，等它静止后，再继续放下一块；重复若干次，如果整体堆叠形成了“有意义的结构”（比如高度足够高、没有严重重叠），就把它当一条样本。
   • 如果堆叠失败（掉落散开、倒塌），就扔掉重来。

2. 问题

   • 因为随机性太强，大多数尝试看出来都不稳定或形状凌乱，需要反复丢弃和重做。
   • 仿真器调用很多次，非常耗时，效率极低，不适合生成百万级数据。

作为替代，我们采用一种“通过拆解来构建”的方法，首先从由不同积木形状组成的密集结构开始，然后进行积木移除过程，迭代地从堆叠中移除积木，直到结构变得不稳定。虽然初始结构由预定义的网格指导，但我们发现随机的水平位移和积木移除过程可以生成一组多样且非平凡的结构。

具体而言，我们考虑一个 $4\times 4$ 的网格，作为积木堆叠设计的支架。我们从底部开始构建初始密集结构，尝试在当前行中放置一个随机选择的积木（仅在顶行放三角形（在本文中，我们考虑了四种不同的形状：{三角形、立方体、矩形、长矩形}。））而不超过最大宽度四个单元格。一旦某行至少有三个单元格被占用，我们便移动到下一层。这样就得到了积木堆叠的初始模板。然后我们将模板转换为对应积木的 $\text{SE(3)}$ 姿态集合，并在它们的水平位置上加入少量噪声。接着，我们使用仿真器验证在重力作用下堆叠的稳定性，渲染其正面轮廓，并将该姿态集合及轮廓加入到数据集中。如果堆叠倒塌，我们则直接拒绝该设计。值得注意的是，最终数据集中包含了围绕垂直轴轻微旋转的积木，如图3右侧所示。这是由于物理引擎的不准确性，当我们运行前向动力学并等待堆叠其他部分稳定时，积木会不断轻微滑动并随机旋转。尽管这并非有意，但我们仍将其保留在数据集中，认为这种随机性有助于增加积木姿态的多样性。

图3：左侧展示了我们生成多样化稳定堆叠的策略。将设计网格填充形状后，我们在模拟器中验证堆叠的稳定性，然后开始逐个移除积木，并保存所有稳定的堆叠。右侧展示了数据集中一些具有挑战性的示例。

对于上述生成的每个积木堆叠，我们通过积木移除生成额外的数据点，即移除缺失后不会导致结构塌陷的积木。从图3所示的初始堆叠开始，我们尝试移除每个积木，并模拟剩余积木的影响。如果堆叠仍然稳定，我们就将移除后的姿态集合和轮廓加入数据集，然后对下一个积木重复此过程。我们对每个稳定配置递归应用该过程，最多移除四个积木。注意，堆叠顶部的积木不被移除，因此数据样本始终具有四层立方体的高度。按照此过程，我们生成了191k个稳定积木堆叠实例，并以9:1的比例将其划分为训练集和测试集。

三、“拆解式”生成法（高效多样）

• 为了避开上面方法的低效率，作者提出先构建一个“密集模板”——再逐步拆解的策略：

3.1 初始模板构建

1. 网格框架

   • 在水平面上画一个 $4\times 4$ 的网格，作为最多放四列积木的“框架”。

2. 一行一行填充

   • 从底层开始：在当前行的四个格子里，随机挑选一种积木（立方体、矩形、长矩形、三角形），放入一个或多个相邻格子，注意最多占用 4 个格子宽度。
   • 只要这一行“至少放满 3 个格子”（即占用数量 ≥3），就算这一层“合格”，然后往上一层继续。
   • 对顶层（第四层）仅允许放三角形。

3. 把模板变成实际的 3D 堆叠

   • 每个“格子”对应一个积木中心的大致位置，在这些位置上放置对应形状，并给它们加上少量随机的水平偏移和旋转噪声。
   • 这样就有了一个“初始密集堆叠”。

3.2 稳定性检测与渲染

1. 仿真器验证

   • 把这个初始堆叠放到物理仿真里，加上重力。
   • 等待若干秒，检查所有积木是否保持“整体不倒塌”。

2. 渲染轮廓

   • 如果稳定，就对其正面做一个二值化的轮廓渲染，用作数据集中的“条件输入”之一。
   • 如果不稳定，直接丢弃这一设计。

小细节：

• 在仿真过程中，由于物理引擎的微小误差，积木会绕垂直轴发生少量随机旋转和滑动。
• 这些无意中的随机性反而增加了“姿态多样性”，所以作者保留了这种现象。

---

四、积木移除：递归生成更多样本

仅靠上述“初始密集堆叠”还不够多样，还要再拆解式地生成更多稳定子结构：

1. 尝试移除一个积木

   • 对于当前的稳定堆叠，依次尝试“拿掉”其中的每一块积木。
   • 在仿真里验证“剩下的积木”是否仍然稳定。

2. 接受或拒绝

   • 如果移除某块后整个结构依旧不倒，就把“移除后的结构”及其轮廓渲染作为一个新的样本加入数据集。
   • 如果不稳定，就把该移除操作丢弃，保留原结构不变。

3. 递归进行

   • 对每个新加入的数据点，最多再继续“移除”多达四次（即一个链条里最多拆四块积木）。
   • 保证顶层那一层的四块积木不被移除，这样所有样本都至少有四层的高度。

---

五、最终数据结果

• 这样一来，从一个“网格密集模板”出发，通过一次“仿真挑选”＋多次“移除尝试”＋“仿真验证”，可以高效地产生成百上千、乃至上万的稳定且多样化堆叠实例。
• 实际上，作者一共生成了 191,000 个稳定堆叠样本。
• 最后按 9:1 的比例，分成训练集和测试集，用来训练和评估模型。

---

整体流程小结

1. 构建“填满大半行”的初始模板 →
2. 仿真中验证稳定性 →
3. 渲染并收集样本 →
4. 对每个已采样结构，递归地尝试移除积木，再次验证 →
5. 收集所有通过测试的子结构 →
6. 得到一个大规模、多样且都保证物理稳定的数据集。

IV. 实验

我们评估模型生成与参考输入一致的稳定物体配置的能力，该参考输入可以是一个积木结构示例或所需结构的草图（图 1）。随后，我们展示了使用 UR5 机械臂构建不同结构的真实场景结果。

A. 仿真评估

我们使用保留的测试数据集来评估模型。图 4 可视化了 STACKGEN 针对单一轮廓生成的一组多样化堆叠，展示了其对多模态分布学习的能力。

图4：（左）参考（即真实）堆叠及其轮廓；（右）本模型根据该轮廓生成的一组多样化结构。

我们使用两个指标来评估方法：

1. 由本方法生成的积木配置中结构上稳定的比例；
2. 生成的堆叠与目标轮廓的一致性。

由于从给定积木列表和轮廓生成稳定结构的问题可能有多种解，我们的评估方法对测试集中每对轮廓与积木列表采样三组积木姿态进行测试。我们将本方法与四个基线进行比较：两种启发式基线（暴力搜索基线和贪婪随机基线）以及两种基于学习的基线（Transformer-回归基线和 Transformer-VAE 基线）。

1. 暴力搜索基线：
对于给定的轮廓和可用积木集，该算法通过最大化由轮廓交集衡量的对齐得分并最小化积木间碰撞惩罚，搜索每个积木的潜在放置姿态。为获得高对齐得分，对每个积木我们从 $[-3,3]$ 和 { 1 , 3 , 5 , 7 } \{1,3,5,7\} {1,3,5,7} 中分别均匀采样 20 个 $(x_i,0,z_i)$ 坐标，然后从每个点沿 $x$ 轴正反两个方向进行 20 次线性搜索，以找到最优对齐与碰撞组合。该算法在 $k\times 400$ 个候选位置的离散化空间中进行暴力搜索，虽然不属于 NP-困难问题，但计算量极大。

2. 贪婪随机基线：
该方法采用自左向右、自下而上的算法，基于结构轮廓放置积木。从最低层到最高层，每层分配固定高度。算法先测量从左到右最长连续像素行的长度，然后在当前积木列表中选取宽度小于此长度的所有积木，贪婪地放置最长者。由于该算法是确定性的，我们引入置换机制以增加多样性：以概率 $\sigma$，算法会将同一层中相邻的两个立方体与轮廓中其它位置的一个矩形互换（因为两立方体与一矩形长度相等）。通过控制 $\sigma$，可调节生成配置的多样性。该置换机制也应用于暴力搜索基线，以控制其可变性。

3. Transformer-回归基线：
我们加入一个基于 Transformer 的回归模型（图 5 右）作为非生成式学习方法。此基线用于评估生成模型在捕捉条件于轮廓的多模态积木姿态分布方面的优势。为保证公平，我们训练的 Transformer 在层数、头数、前馈维度和隐藏维度上均与 STACKGEN 中的 Transformer 一致。

图5：我们的基于学习的基线包括：（左）Tran-VAE：一个基于Transformer的VAE模型，具有两层编码器和四层解码器；（右）Tran-Reg：一个基于Transformer的回归模型。

4. Transformer-VAE 基线：
Tran-VAE 基线（图 5 左）采用一个 Transformer 变分自编码器，包含两层编码器和四层解码器，总层数与 STACKGEN 的扩散模型（六层）相同。包括 STACKGEN 在内的所有三种学习型模型均使用相同的批量大小和训练周期，以确保公平比较。

Transformer 基线（Tran-Reg 和 Tran-VAE）在形状嵌入中加入了位置编码，而 STACKGEN 则没有，因为扩散模型能够通过各自不同的噪声姿态区分同一积木的不同实例。

为量化预测姿态的多样性，我们按层从左到右提取每个预测结果的积木列表。若两种构建的分层列表不同，则视为不同。多样性得分定义为针对给定输入生成的不同姿态数量除以采样总次数。例如，图 4 中两个场景的平均多样性得分为 $\frac{5}{6}=83.33\%$。

在稳定性评估中，我们根据生成的姿态放置积木，使用仿真器观察其后续行为，并检查是否有积木跌落到其初始层以下，若是则判定该样本为不稳定。为评估轮廓一致性，我们在运行前向动力学后提取生成结构的轮廓，计算来自三个视图（正面、侧面和顶部，见图 6 左）的交并比（IoU），并取其平均值。坍塌结构的 IoU 记为零。

图6：柱状图比较了STACKGEN与基线方法在（左）IoU和（右）稳定性方面的表现。所有方法均使用CNN预测的积木列表。

我们在 500 个场景上评估了所有五种模型，使用真实标签和预训练 CNN 积木列表预测器，每个场景生成三次样本。由于 STACKGEN 实现了 60.47% 的多样性水平，我们在暴力搜索和贪婪随机基线中将 $\sigma =0.6$ 以匹配此多样性水平。如表 I 和图 6 所示，STACKGEN 在稳定性和 IoU 两个指标上均显著优于各基线。

表 I：不同方法在稳定性、多样性和IoU方面的比较。所有结果均以百分比（%）呈现。标注 $†$ 的模型使用了真实积木列表；未标注†的模型使用了CNN预测的积木列表。

由于扩散模型能够表示多模态分布，STACKGEN 自然生成满足轮廓约束的多样稳定设计，而启发式和学习型基线在堆叠稳定性、与草图一致性（即 IoU）以及（除 Tran-VAE 基线外）生成多样设计以匹配输入轮廓的能力方面表现明显较差。与此同时，尽管暴力搜索方法取得了较高的 IoU 分数，但它需要评估巨量的姿态组合，推理时间超过六小时（而 STACKGEN 仅需数秒），其计算复杂度使其在真实机器人应用中不可行。

B. 预测积木列表 vs. 真实积木列表

为了评估我们基于 CNN 的积木列表预测的影响，我们将其与使用真实积木列表的方法进行比较。使用测试集中的 500 个场景，我们为每个场景生成三次样本，并分别使用 CNN 预测的积木列表 ${\hat{s}}_1,\dots ,{\hat{s}}_k$ 和真实积木列表 $s_1,\dots ,s_k$ 运行每个模型。如表 I 所示（标注 $†$ 的模型使用了真实积木列表），STACKGEN 在稳定性和 IoU 两个指标上的表现，在使用 CNN 预测积木列表和使用真实积木列表两种情况下的差异均不超过 2%，这验证了基于 CNN 的预测对 STACKGEN 的实用价值。
然而，我们注意到其他基于学习的方法对积木列表的变化非常敏感，即便该列表已足以生成堆叠。

C. 真实环境中的积木堆叠

为了证明我们的方法在真实环境中表现良好，我们使用玩具积木和 UR5 机械臂进行了实验。我们的目标是构建一个如下运作的流程：首先，用户通过展示玩具积木的参考堆叠或绘制所需结构的草图来提供一个轮廓。从堆叠或草图中提取轮廓后，我们的模型生成与所提供轮廓相匹配的稳定积木配置。最后，UR5 机械臂使用真实积木在桌面上组装生成的堆叠。

图1：STACKGEN 包含一个扩散模型，该模型以目标结构轮廓和可用积木形状列表作为输入。然后，模型生成一组积木姿态 { p ^ 1 , … , p ^ k } \{\hat p_1,\dots,\hat p_k\} {p^​1​,…,p^​k​}，用于构造与目标轮廓一致的稳定结构。最后，可以使用机器人手臂来搭建所得到的结构。

1. 堆叠→堆叠：在此场景中，如图 1 所示，通过由 RGBD 相机（Realsense 435D）、玩具积木和白色背景组成的简单装置，从堆叠中提取轮廓。该装置拍摄用户搭建的积木堆叠照片，然后通过利用深度读数过滤掉背景像素，并应用中值滤波对轮廓进行平滑，去除任何残留的白色像素，最后将结果调整大小并粘贴到$64\times 64$画布上，生成二值轮廓。

2. 草图→堆叠：在这种情况下，我们使用相机捕捉用户的手绘草图（图 1）。将其转换为二值图像，使用中值滤波进行平滑，并在其周围放置一个带有$4\times 4$网格的边界框。然后，我们计算占据情况，以确定每个网格单元是完全占据还是部分占据（例如，三角形）。

利用提取的轮廓，我们使用预训练的 CNN 来预测积木列表（对于草图→堆叠示例，我们采用启发式方法而不是CNN来识别积木列表）。扩散模型随后生成候选积木姿态。对于生成姿态集合中的每个积木，UR5 机械臂执行拾取并放置操作，将积木放置到其对应的姿态位置。执行顺序采用自左向右、自下而上的贪婪方式设定。

在我们测试的八种案例中，该流程成功地稳定构建了所有堆叠，且与原始轮廓仅存在细微差异（考虑到积木初始位置的误差）。然而，我们注意到这并不意味着我们的系统是完美无缺的。正如在第 IV-A 节中讨论的，模型有时会生成不稳定的积木配置。尽管如此，在这些真实环境实验中，成功率表明该模型足够鲁棒，能够有效应对潜在的分布外轮廓。

图7：展示了由 UR5 机械臂根据以 (a) 图像 和 (b) 目标结构草图 形式给出的目标规范构建的各种稳定 3D 结构示例。请注意，STACKGEN 旨在匹配输入的轮廓，因此单个积木的颜色和类型可能与参考输入不同。

V. 结论

在本文中，我们提出了一种新方法，使机器人能够对物体的六自由度姿态进行推理，以实现稳定的三维结构。在给定稳定结构数据集的情况下，STACKGEN 学习不同物体原语的 $\text{SE(3)}$ 姿态分布，并以用户提供的目标结构轮廓为条件。在推理阶段，STACKGEN 生成一组多样化的候选组合，这些组合既与轮廓对齐，又保证物理可行性。我们在仿真环境中进行了实验，展示了我们的方法能够在不显式建模物理的情况下，有效生成符合用户提供轮廓的稳定结构。此外，我们将该方法部署于真实环境，证明了该方法能够以数据驱动的方式，有效且可靠地生成稳定且有效的积木结构，成功地搭建了从视觉设计输入到物理构建的桥梁。