RigidFormer：基于Transformer的刚体动力学学习 RigidFormer: Learning Rigid Dynamics using Transformers

RigidFormer的方法流程可分为五个步骤。步骤一：输入特征构造。对每个物体的每个顶点 $v$，构造12维特征 $h_{t}^{(i,v)} = [d_{t}^{(i,v)}; v_{t}^{(i,v)}; r_{t}^{(i,v)}; \phi^{(i)}]$，包含到最近邻物体的位移向量 $d$、位置增量（离散速度代理）$v$、参考偏移 $r$（相对序列首帧）和广播的物理参数（质量、摩擦、恢复系数）。步骤二：物体级编码。用层次化PointNet编码器 $\text{Enc}_\theta$ 将每个物体的所有点特征聚合成单个 $D$ 维物体token $o_t^{(i)} \in \mathbb{R}^D$。编码器在全局和三个下采样层级提取几何特征后拼接融合，共享权重以促进跨几何形状的泛化。步骤三：物体级交互。将 $M$ 个物体token与 $N_r=16$ 个可学习register token拼接，输入 $L$ 层Transformer。每层依次应用自注意力、step-size FiLM条件（$c=(s, s^2)$，$s \propto \Delta t$）、门控注意力（$\sigma(G(Q)) \odot \text{Attn}(Q,K,V)$，门控值抑制噪声读取以稳定长时域滚动）和FFN。步骤四：锚点级状态推进。用FPS从每个物体采 $N_a=4$ 个锚点，每个锚点提取位置特征并通过MLP投影得到查询 $Q_t \in \mathbb{R}^{(MN_a) \times D}$。每个查询通过交叉注意力从物体token中获取上下文，再与AVP模块聚合的局部接触特征 $u_t^{(i,k)} \in \mathbb{R}^{256}$（使用可学习距离核 $w = \exp(-\|x^{(i,v)} - q^{(i,k)}\|/\sigma)$ 加权聚合）拼接，预测每个锚点的加速度 $a_t^{(i,k)}$。用Verlet积分推进：$\hat{q}_{t+1}^{(i,k)} = a_t^{(i,k)} \Delta t^2 + 2q_t^{(i,k)} - q_{t-1}^{(i,k)}$。步骤五：刚性投影与状态输出。对每个物体，用Kabsch对齐参考锚点 $q_{\text{ref}}^{(i,k)}$ 和预测锚点 $\hat{q}_{t+1}^{(i,k)}$，得到刚体变换 $(R^{(i)}, t^{(i)}) \in SE(3)$：计算 $H = \sum_k (q_{\text{ref}}^{(i,k)} - \bar{q}_{\text{ref}}^{(i)})(\hat{q}_{t+1}^{(i,k)} - \bar{\hat{q}}^{(i)})^\top$，SVD分解 $H=U\Sigma V^\top$，取 $R^{(i)} = V\text{diag}(1,1,\det(VU^\top))U^\top$。最终全分辨率顶点更新为 $x_{t+1}^{(i,v)} = R^{(i)} x_{\text{ref}}^{(i,v)} + t^{(i)}$。Kabsch对齐通过RoMa实现以保证SVD梯度在退化奇异值附近的稳定性。

RigidFormer的技术新颖性体现在五个层面的方法论贡献。首先是问题重定义：与MGN/FIGNet/HopNet等将仿真视为'逐顶点图推理'问题不同，RigidFormer将其重新定义为'物体级注意力+锚点级SE(3)流形投影'问题，这种重新定义为无网格建模提供了理论支撑。其次是ARoPE的数学性质保证：本文严格证明ARoPE满足物体token排列的置换等变性（permutation equivariance）和锚点重排的不变性（anchor-indexing invariance），填补了'无序集合作为Transformer输入'的位置编码理论空白。第三是归纳偏置设计：通过FPS锚点+Kabsch刚性投影，将归纳偏置深度融入架构而非仅靠数据驱动，减少了达到相同性能所需的数据量。第四是计算复杂度优化：交互成本从顶点级 $O((MN_v)^2)$ 降到物体级 $O(M^2)$ 和锚点级 $O((MN_a)^2)$，$N_a=4 \ll N_v$ 是关键。第五是可扩展性验证：在WreckingBall数据集上扩展到216个物体仍保持稳定（位置误差仅0.130m），这是首次在如此规模下验证点云神经仿真器的可行性。训练目标为位置+加速度的Smooth L1损失的组合：$\mathcal{L} = \lambda_{\text{pos}}(\mathcal{L}_{\text{raw}}^{\text{pos}} + \mathcal{L}_{\text{rigid}}^{\text{pos}}) + \lambda_{\text{acc}}(\mathcal{L}_{\text{raw}}^{\text{acc}} + \mathcal{L}_{\text{rigid}}^{\text{acc}})$，其中'raw'和'rigid'分别表示Kabsch对齐前后的损失，$\lambda_{\text{pos}}=10, \lambda_{\text{acc}}=1$，训练300个epoch，AdamW，$\beta_1=0.9, \beta_2=0.999$，基础学习率 $10^{-4}$。

尽管RigidFormer在多个维度取得突破，仍存在几个值得深入分析的弱点。弱点一：物体标签的强依赖。这是最显著的实用性瓶颈——模型将每个点云分段输入并需要预分割。在感知-物理一体化的趋势下，这相当于把分割的困难外包给上游pipeline。改进方向：引入端到端的点-物体关联机制，例如用类似Mask3D的分割头，或用slot attention让模型自动发现物体。弱点二：FPS锚点的稳定性问题。FPS是确定性的，对点云的微小扰动可能产生完全不同的锚点集合，影响推理稳定性。虽然Table 9的随机化FPS实验显示模型对锚点选择鲁棒，但实际部署中输入点云噪声可能导致锚点漂移。改进方向：引入可学习锚点初始化（如用K-means中心或聚类中心替代FPS），或设计锚点去噪模块。弱点三：自回归滚动的误差累积本质。即使有Kabsch投影，每步预测误差仍会累积——MOVi-B 100步位置误差0.161m是50步0.050m的3.2倍，呈现明显累积趋势。改进方向：训练时加入随机扰动（denoising score matching思想）、引入观测校正模块（类似Kalman滤波），或将长滚动分解为多尺度预测。弱点四：FiLM条件编码的简化。当前$c=(s, s^2)$仅编码step size标量，忽略了系统状态相关的自适应步长。改进方向：让模型自适应决定下一时刻的$\Delta t$（如基于预测不确定性）。弱点五：register token数量的固定。$N_r=16$个register token是固定值，对不同物体数量的场景可能不是最优。改进方向：基于物体数量动态调整register数量。弱点六：缺少物理一致性约束。损失函数仅监督位置和加速度，未包含能量守恒、动量守恒等物理规律。改进方向：引入物理信息损失（physics-informed loss）。弱点七：实验评估的benchmark范围有限。MOVi数据集相对简单，未在接触更复杂、形状更丰富的真实数据集上验证。改进方向：在PartNet-Mobility、ShapeNet、CMU-Mocap等更具挑战性的数据集上评估。

作者提出了几个有前景的未来方向，结合本文成果可延伸出更多研究方向。第一，感知-物理一体化：作者指出'当前方法需要物体标签'是关键限制，未来可探索从原始RGB-D或点云流中联合推断物体身份和物理状态，这需要结合3D场景图、实例分割或开放词汇检测的前沿工作。第二，混合刚体-可形变场景：现实世界中很多场景包含刚体和软体（如机器人抓取软物体、布料与刚体交互）。本文的物体级框架可通过引入可形变物体的参数化（如神经辐射场、皮肤模型）扩展到此类场景。第三，自适应时间步长：当前仅支持离散step size，未来可探索连续$\Delta t$的傅里叶特征编码，或让模型基于状态不确定性自适应选择步长。第四，inverse dynamics：从观测推断物理参数$\phi=[m, \mu, \epsilon]$，这对机器人交互（如估计物体摩擦系数）有重要价值。第五，更复杂的铰接体控制：当前仅展示方向控制，未来可扩展到全身运动控制、接触式操作（contact-rich manipulation）。第六，从仿真到真实（sim-to-real）：在仿真器中预训练，再迁移到真实场景。本文的无网格特性使其天然适合sim-to-real，但需要处理真实点云的噪声和稀疏性。第七，多模态输入融合：结合RGB图像、深度、法向等信息增强感知鲁棒性。第八，作为世界模型：在强化学习中，将RigidFormer作为环境模型用于planning——其快速推理和长时域稳定性使其成为MBRL的理想选择。第九，时空生成式建模：将RigidFormer与扩散模型结合，从噪声生成物理合理的轨迹，而非仅做自回归预测。

从复现角度看，本文提供了较为详细的信息但仍存在挑战。架构细节：编码器为层次化PointNet（在全局+3个下采样层级提取特征并融合），$L$层Transformer block，$N_r=16$ 个register token，$N_a=4$ 个FPS锚点，$D=256$ 维特征。训练细节：300 epochs，AdamW（$\beta_1=0.9, \beta_2=0.999$，weight decay 0.01），基础学习率 $10^{-4}$，10 epoch线性warmup（从 $10^{-5}$） + cosine衰减到 $10^{-6}$，梯度裁剪norm=1.0。损失：$\mathcal{L}=\lambda_{\text{pos}}(\mathcal{L}_{\text{raw}}^{\text{pos}}+\mathcal{L}_{\text{rigid}}^{\text{pos}})+\lambda_{\text{acc}}(\mathcal{L}_{\text{raw}}^{\text{acc}}+\mathcal{L}_{\text{rigid}}^{\text{acc}})$，$\lambda_{\text{pos}}=10, \lambda_{\text{acc}}=1$，Smooth L1损失。算力需求：NVIDIA A100 GPU训练（未明确GPU数量，附录C中应有详细说明）。推理环境：RTX 5080，50步滚动平均10次迭代。代码与数据：作者承诺'Code will be released upon publication'，但截至目前（2026年6月）尚未明确开源仓库链接。MOVi数据集来自Kubric项目，可公开获取。第三方依赖：RoMa（用于鲁棒可微Kabsch对齐）。复现难度评估：中等。算法思想清晰，架构组件明确，损失函数简单；但完整复现需要实现：(1)层次化PointNet编码器；(2)门控注意力+FiLM条件的Transformer；(3)FPS锚点+AVP模块；(4)Verlet积分+Kabsch对齐的联合优化；(5)数据预处理流水线（物体点云提取+物理参数注入）。整体而言，对熟悉物理仿真和Transformer的研究者，约1-2个月可复现核心结果。