[论文解读] EDGI: Equivariant Diffusion for Planning with Embodied Agents

Johann Brehmer, Joey Bose+2 more|arXiv (Cornell University)|Mar 22, 2023

Reinforcement Learning in Robotics被引用 4

一句话总结

EDGI 提出了一种新颖的等变扩散模型，用于具身智能体的规划，该模型在空间（SE(3)）、时间（Z）和物体置换（S_n）对称性的乘积群下保持不变。通过将该模型集成到带有分类器引导的条件生成规划框架中，EDGI 相较于非等变基线模型，在对称变换下的样本效率和泛化能力显著提升，仅需最多 10 倍少的训练数据即可超越基线模型。

ABSTRACT

Embodied agents operate in a structured world, often solving tasks with spatial, temporal, and permutation symmetries. Most algorithms for planning and model-based reinforcement learning (MBRL) do not take this rich geometric structure into account, leading to sample inefficiency and poor generalization. We introduce the Equivariant Diffuser for Generating Interactions (EDGI), an algorithm for MBRL and planning that is equivariant with respect to the product of the spatial symmetry group SE(3), the discrete-time translation group Z, and the object permutation group Sn. EDGI follows the Diffuser framework (Janner et al., 2022) in treating both learning a world model and planning in it as a conditional generative modeling problem, training a diffusion model on an offline trajectory dataset. We introduce a new SE(3)xZxSn-equivariant diffusion model that supports multiple representations. We integrate this model in a planning loop, where conditioning and classifier guidance let us softly break the symmetry for specific tasks as needed. On object manipulation and navigation tasks, EDGI is substantially more sample efficient and generalizes better across the symmetry group than non-equivariant models.

研究动机与目标

• 通过在模型架构中显式编码空间、时间与置换对称性，提升具身智能体在基于模型的强化学习中的样本效率与泛化能力。
• 开发一种在 SE(3)×Z×S_n 下保持不变的基于扩散的规划框架，同时在推理阶段允许灵活的任务特定对称性破缺。
• 在单一等变架构中支持多种输入表示（如位置、速度、物体标签），以提升真实机器人任务中的灵活性。
• 证明等变建模可实现对称变换下鲁棒的性能表现以及低数据场景下的优异性能。
• 提供一个通用蓝图，用于在乘积群上构建支持多表示输入的等变模型。

提出的方法

• EDGI 使用在离线轨迹上训练的扩散模型，其去噪网络在完整乘积群 SE(3)×Z×S_n 下保持等变性。
• 模型引入了新颖的时间、物体与置换层，尊重群结构，并在不同对称类型间保持一致的行为。
• 一种新型嵌入机制将多种输入表示（如三维坐标、速度、物体索引）统一映射为共享的内部表示，同时保持等变性。
• 通过从扩散模型进行条件采样实现规划，结合任务特定的条件输入与分类器引导，实现对称性的软破缺。
• 该框架将等变扩散模型集成到规划循环中，支持世界模型学习与轨迹生成。
• 该架构在不牺牲等变性的前提下支持多种数据表示，从而在复杂环境中具备更广泛的应用潜力。

Figure 1: Schematic of EDGI in a navigation task, where the agent (red square) plans the next actions (red dots) to reach the goal (green star) without touching obstacles (grey circles). Top : planning as conditional sampling from a diffusion model. Bottom : effect of a group action. Equivariance re

实验结果

研究问题

• RQ1 通过在模型架构中显式编码 SE(3)×Z×S_n 对称性，基于扩散的规划算法是否能实现更高的样本效率？
• RQ2 在空间、时间与置换对称性下的等变性，如何影响模型在未见过的对称变换环境中的泛化能力？
• RQ3 通过条件化与分类器引导实现的软对称性破缺，在多大程度上可实现任务特定规划，同时不损害等变性？
• RQ4 多表示输入支持的集成是否能提升复杂机器人操作与导航任务中的性能？
• RQ5 与非等变基线相比，等变扩散模型是否能在显著减少训练数据的情况下仍保持高性能？

主要发现

• 当仅使用 10% 的训练数据时，EDGI 在导航与物体操作任务上的性能与最佳非等变基线模型相当。
• 在低数据场景下，EDGI 显著优于非等变模型，展现出卓越的样本效率。
• EDGI 对环境的未见对称变换（包括旋转或置换的物体构型）具有鲁棒的泛化能力。
• 在 Kuka 操作基准测试的全部三项任务中，EDGI 即使在对称性增强的测试条件下，仍保持了强劲的性能表现。
• 分类器引导的集成实现了有效的软对称性破缺，使模型能够在保持底层动力学不变性的同时解决特定任务。
• 实证结果证实，等变性可提升数据效率与鲁棒性，验证了该方法的核心假设。

Figure 2: Architecture of our $\mathrm{SE(3)}\times\mathbb{Z}\times\mathrm{S}_{n}$ -equivariant denoising network. Input trajectories (top left), which consist of features in different representations of the symmetry group, are first transformed into a single internal representation (green block). T