[Paper Review] EDGI: Equivariant Diffusion for Planning with Embodied Agents

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

Reinforcement Learning in Robotics4 citations

TL;DR

EDGI introduces a novel equivariant diffusion model for planning with embodied agents that is invariant under the product group of spatial (SE(3)), temporal (Z), and object permutation (S_n) symmetries. By integrating this model into a conditional generative planning framework with classifier guidance, EDGI achieves significantly improved sample efficiency and generalization across symmetry transformations compared to non-equivariant baselines, outperforming them with up to 10× less training data.

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.

Motivation & Objective

• To improve sample efficiency and generalization in model-based reinforcement learning for embodied agents by explicitly encoding spatial, temporal, and permutation symmetries.
• To develop a diffusion-based planning framework that maintains invariance under SE(3)×Z×S_n while allowing flexible, task-specific symmetry breaking at test time.
• To support multiple input representations (e.g., positions, velocities, object labels) within a single equivariant architecture for greater flexibility in real-world robotic tasks.
• To demonstrate that equivariant modeling leads to robust performance under symmetry transformations and low-data regimes.
• To provide a general blueprint for constructing equivariant models on product groups with multi-representation input support.

Proposed method

• EDGI employs a diffusion model trained on offline trajectories, with a denoising network that is equivariant under the full product group SE(3)×Z×S_n.
• The model uses novel temporal, object, and permutation layers that respect the group structure and act consistently across different symmetry types.
• A new embedding mechanism unifies multiple input representations (e.g., 3D coordinates, velocities, object indices) into a shared internal representation while preserving equivariance.
• Planning is performed via conditional sampling from the diffusion model, with task-specific conditioning and classifier guidance to softly break symmetries.
• The framework integrates the equivariant diffusion model into a planning loop that supports both world model learning and trajectory generation.
• The architecture supports multiple data representations without sacrificing equivariance, enabling broader applicability in complex environments.

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

Experimental results

Research questions

• RQ1 Can a diffusion-based planning algorithm achieve improved sample efficiency by explicitly encoding SE(3)×Z×S_n symmetries in the model architecture?
• RQ2 How does equivariance under spatial, temporal, and permutation symmetries affect generalization to unseen symmetry-transformed environments?
• RQ3 To what extent can soft symmetry breaking via conditioning and classifier guidance enable task-specific planning without compromising equivariance?
• RQ4 Does the integration of multi-representation input support enhance performance in complex robotic manipulation and navigation tasks?
• RQ5 Can an equivariant diffusion model maintain high performance with substantially reduced training data compared to non-equivariant baselines?

Key findings

• EDGI matches the performance of the best non-equivariant baseline on navigation and object manipulation tasks when trained on only 10% of the training data.
• In low-data regimes, EDGI significantly outperforms non-equivariant models, demonstrating superior sample efficiency.
• EDGI generalizes robustly to unseen symmetry transformations of the environment, including rotated or permuted object configurations.
• The model maintains strong performance across all three tasks in the Kuka manipulation benchmark, even under symmetry-augmented test conditions.
• The integration of classifier guidance enables effective soft symmetry breaking, allowing the model to solve specific tasks while preserving invariance in the underlying dynamics.
• Empirical results confirm that equivariance leads to improved data efficiency and robustness, validating the core hypothesis of the approach.

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