PaperlayerView exact revision on arXiv

Source-linked AI summary

EDGI: Equivariant Diffusion for Planning with Embodied Agents

Johann Brehmer, Joey Bose, Pim de Haan, Taco Cohen

arXiv:2303.12410v2cs.LGcs.ROstat.ML

TL;DR

Many planning and model-based reinforcement-learning algorithms overlook the spatial, temporal, and permutation symmetries of embodied tasks, a setting associated with sample inefficiency and limited robustness. EDGI uses an equivariant diffusion model for planning and reports stronger low-data performance and robustness to symmetry transformations.

  • Problem

    Many planning and model-based reinforcement-learning algorithms do not account for embodied tasks' spatial, temporal, and permutation symmetries, and can be sample-inefficient and lack robustness to environmental changes.

    Source: §1 · p. 1

  • Method

    EDGI trains a diffusion model on offline trajectories that is equivariant to SE(3) × Z × Sn and uses test-time conditioning and classifier guidance for soft symmetry breaking.

    Source: §1 · p. 2, §3.2 · p. 6, §3.4 · p. 7

  • Results

    Across 3D navigation and robotic object manipulation, EDGI matches the best non-equivariant baseline with an order of magnitude less training data and generalizes well to unseen symmetry configurations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    EDGI provides a blueprint for equivariant networks supporting multiple data representations while modeling an invariant base density and solving non-invariant tasks through task-specific soft symmetry breaking.

    Source: §6 · p. 10

    Training and planning are currently expensive.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Embodied tasks exhibit spatial, temporal, and object-permutation symmetries, but many RL and planning algorithms do not account for them and can be sample-inefficient and lack robustness. EDGI builds these symmetries into a diffusion-based planning model and reports improved low-data performance and robustness to symmetry transformations.

1 Introduction

  • Embodied tasks often have spatial, temporal, or permutation symmetries, while many RL and planning algorithms do not account for this structure.The introduction identifies sample inefficiency and lack of robustness to environmental changes as weaknesses of these algorithms.Source: §1 · p. 1
  • EDGI is an equivariant algorithm for model-based reinforcement learning and planning, with a base component equivariant to SE(3) × Z × Sn.The model supports multiple representations and allows soft symmetry breaking at test time for specific tasks.Source: §1 · p. 2
  • EDGI follows Diffuser by treating dynamics learning and planning as generative modeling, training a diffusion model on offline state-action trajectories.Planning samples trajectories conditionally on the current state and uses classifier guidance to maximize reward.Source: §1 · p. 2
  • Its diffusion model supports multiple data representations and combines novel temporal, object, and permutation layers for the corresponding symmetries.Task specifications can softly break symmetry at test time through conditioning and classifier guidance.Source: §1 · p. 2
  • EDGI matches the best non-equivariant baseline with an order of magnitude less training data and generalizes well to unseen symmetry configurations.The reported experiments cover 3D navigation and robotic object manipulation.Source: §1 · p. 2

2 Background

The background defines equivariance and invariance for group actions, then reviews diffusion models and trajectory optimization. It frames diffusion-based planning as conditional trajectory sampling shaped by task constraints or reward.

2 Background

  • A function is G-equivariant when applying a group action to its input corresponds to applying that action to its output; it is G-invariant when the output action is trivial.The discussion focuses on linear representations, including rotation and permutation matrices.Source: §2 · pp. 2–3
  • For generative modeling, an equivariant map from a G-invariant prior can produce a G-invariant target density, as established in cited work.The paper applies this principle to trajectory distributions symmetric under SE(3) × Z × Sn.Source: §2 · p. 3
  • Diffusion models add noise to data over forward steps and generate samples by iteratively denoising noise toward the empirical data distribution.The forward transition is a Gaussian with a prescribed variance schedule, while the reverse transition is learned.Source: §2 · p. 3
  • Classifier guidance modifies the diffusion score using the gradient of a label log-likelihood, with the guidance weight controlling its contribution.The background also notes noise prediction as a practical parameterization choice.Source: §2 · p. 3
  • Trajectory optimization seeks an action sequence maximizing an objective that sums per-timestep rewards over a planning horizon.The trajectory contains states and actions, and the dynamics specify the next state from the current state and action.Source: §2 · p. 3
  • Diffuser-style planning trains a diffusion model on offline state-action trajectories and samples conditionally from a task-weighted trajectory distribution.Constraints encode trajectory requirements; conditioning and reward guidance shape test-time sampling.Source: §2 · p. 3

3 Equivariant diffuser for generating interactions (EDGI)

EDGI builds an SE(3) × Z × Sn-equivariant diffusion model for trajectories and uses invariant or task-conditioned sampling for planning.

3.1 Symmetry and representations

  • The SE(3) × Z × Sn group captures spatial, temporal, and object-permutation symmetries, with features organized as scalar/vector and object/global representations.Gravity or distinguishable objects may partially break the group; referencing positions to a point makes translation equivariance reducible to SO(3) equivariance.Source: §3.1 · pp. 3–4

3.2 Equivariant diffusion model

  • EDGI transforms heterogeneous trajectory features into one scalar–vector internal representation, processes them with temporal convolutions, object attention, and geometric layers, then maps outputs back.The internal representation pairs one scalar and one vector; the authors note that multiple copies may be beneficial when scalar or vector quantities predominate.Source: §3.2 · pp. 5–6
  • EDGI trains its denoiser on offline trajectories without reward information, predicting the noise added at a sampled diffusion step.Training minimizes a simplified variational lower bound expressed as a noise-prediction loss.Source: §3.2 · p. 6

3.3 Invariant sampling and symmetry breaking

  • Unconditional sampling and sampling guided by an invariant classifier or reward remain SE(3) × Z × Sn-invariant; non-invariant guides can break that symmetry.The invariant-guidance result applies because the product-group representations are orthogonal.Source: §3.3 · p. 7

3.4 Planning with equivariant diffusion

  • For planning, EDGI conditions trajectories to begin at the current state and uses classifier-based reward guidance; task-specific conditions or rewards can softly break symmetry at test time.Classifier-free guidance would require jointly modeling equivariant unconditional and non-equivariant task-conditional distributions, which the authors say is generally difficult.Source: §3.4 · pp. 7–8

4 Experiments

Experiments evaluate EDGI on obstacle navigation and Kuka manipulation, comparing full-data performance, data efficiency, and generalization under symmetry shifts. Across these tests, EDGI matches or exceeds Diffuser with full data and performs better in low-data and shifted settings.

Experimental setup

  • The experiments use obstacle navigation and three Kuka block tasks, with roughly 10^5 navigation and 10^4 manipulation offline trajectories.EDGI plans through conditional sampling; classifier guidance is used for conditional and rearrangement Kuka tasks. Baselines include Diffuser, augmented Diffuser, BCQ, CQL, and partial-equivariance variants.Source: §4 · p. 8

Task performance

  • Across navigation and object tasks, EDGI performs as well as or better than Diffuser with the full training set.On Kuka, both diffusion methods clearly outperform BCQ and CQL.Source: §4 · p. 8

Sample efficiency

  • 10% of Kuka training data and 0.01% of navigation training data yielded EDGI rewards as strong as with the full dataset; Diffuser was much less sample-efficient.Data augmentation partially closes Diffuser’s gap, but EDGI retains an edge.Source: Figure 3 · p. 9, §4 · p. 9

Effects of individual symmetries

  • Full-group EDGI outperforms variants equivariant to only SE(3) or only Sn, though both partial variants outperform Diffuser.The comparison studies the contribution of spatial and object-permutation symmetries to sample efficiency.Source: §4 · p. 9

Group generalization

  • Under restricted training orientations, EDGI remains robust to unseen goal positions and block configurations, while Diffuser performs substantially worse.Navigation training used goals at x = 0; manipulation training used red-block positions at x = 0.Source: §4 · p. 9

5 Related work

EDGI builds on diffusion-based planning and equivariant learning, adding a symmetry-aware denoising network for planning problems. The related work situates this design among equivariant RL and generative-model approaches.

Diffusion-based planning

  • EDGI’s key novelty over diffusion-based planning is a denoising network equivariant to SE(3) × Z × Sn.The original Diffuser is the baseline [Janner et al., 2022]; concurrent work used geometric algebra representations and a transformer [Brehmer et al., 2023].Source: §5 · p. 9

Equivariant deep learning

  • Equivariant deep learning spans geometric transformations, permutations, spherical geometry, irreducible representations, and message-passing methods for 3D symmetries.The cited approaches include Cohen and Welling [2016a], DeepSet, and geometric message passing [Satorras et al., 2021].Source: §5 · p. 9

Equivariance in RL

  • Prior RL work studies symmetries in MDP state-action spaces, while model-based approaches also report data-efficiency benefits from environmental symmetries.The passage cites van der Pol et al. [2020] and Deac et al. [2023], among others.Source: §5 · pp. 9–10

Equivariant generative models

  • Equivariant generative models include continuous normalizing flows and diffusion methods; EDGI extends diffusion to SE(3) × Z × Sn and multiple data representations.The passage describes equivariant vector fields, score matching, and diffusion approaches as related lines of work.Source: §5 · p. 10

6 Discussion

EDGI uses a product-group-equivariant diffusion model as a general blueprint for multi-representation data, while task-specific inputs permit soft symmetry breaking. The authors note that training and planning remain expensive.

Discussion

  • The paper frames spatial, temporal, and permutation structure as a way planning algorithms can improve sample efficiency and generalization, two weaknesses of RL algorithms.Source: §6 · p. 10
  • EDGI models an invariant base density while task-specific conditioning softly breaks symmetry to solve non-invariant tasks.The paper presents its multi-representation product-group network as a general blueprint for equivariant architectures.Source: §6 · p. 10
  • Training and planning are currently expensive; the authors identify more efficient denoising layers and accelerated continuous-time diffusion sampling as possible routes to improvement.These are proposed directions for addressing the cost, not demonstrated solutions.Source: §6 · p. 10

A Group theory

The appendix introduces group-theoretic concepts used in the paper, including Lie groups, matrix Lie groups, and SE(3), then reviews ways to represent 3D rotations.

Group theory

  • A Lie group combines a group structure with a smooth-manifold structure; a topological group has continuous composition and inverse maps.The appendix develops the topological-group definition before discussing Lie groups.Source: A Group theory · p. 14
  • A matrix Lie group is a closed subgroup of GLn(R); SE(3) represents 3D rotations and translations, excluding reflections.SE(3) is described as a subgroup of E(3) and as isomorphic to SO(3) ⋉ R3.Source: A Group theory · p. 14
  • Rotations can be represented by SO(3) matrices or rotation vectors, whose length gives the angle and whose direction gives the axis.The rotation-vector description uses a vector associated with the Lie algebra of SO(3).Source: A Group theory · p. 14
  • Unit rotation quaternions encode a rotation through scalar and vector parts, and rotate a point by q · p′ · q∗.The scalar part corresponds to cos(ω/2), while the vector part corresponds to n sin(ω/2).Source: A Group theory · pp. 14–15

B Architecture details

EDGI uses an SE(3) × Z × Sn-equivariant diffusion architecture built around a four-level U-net. Its blocks combine temporal, object, normalization, geometric, and conditioning operations, with equivariance argued for object permutations and spatial rotations.

Overall architecture

  • The equivariant U-net has four levels, with two residual standard blocks at each level before downsampling or upsampling.Downsampling and upsampling occur on the downward and upward passes, respectively.Source: B Architecture details · p. 15

Residual standard block

  • Each residual standard block combines equivariant processing, a diffusion-time and conditioning context block, residual addition, and linear attention over time.The equivariant block includes temporal, object, normalization, and geometric layers.Source: B Architecture details · p. 15

Temporal and normalization layers

  • Temporal layers use bias-free one-dimensional convolutions with kernel size 5, while normalization rescales each batch element to unit variance without mean-centering.The passage states that mean-centering would break SO(3) equivariance.Source: B Architecture details · p. 15

Geometric layers

  • Geometric layers form scalar invariants from pairwise inner products of 16 vector channels, then use two MLPs to produce scalar outputs and vector-map coefficients.A residual connection adds scalar and vector inputs to the outputs.Source: B Architecture details · p. 15

Linear attention over time

  • Linear attention over time is added residually at each U-net level to match the architecture used by Janner et al. [2022].The passage states that the attention mechanism is made equivariant by computing its weights accordingly.Source: B Architecture details · p. 15

Context and sampling layers

  • Context embeddings use a Mish nonlinearity and linear layer before zero-padding, while U-net paths temporally downsample or upsample and adjust channel counts.The downward path doubles channels; the upward path reduces them.Source: B Architecture details · p. 15

Equivariance

  • The geometric layer preserves object-permutation equivariance by leaving the object dimension untouched and preserves rotational equivariance because inner products are rotation-invariant.Translations are handled through canonicalization relative to the center of mass or robot base.Source: B Architecture details · p. 16

C Navigation experiments

The navigation environment places a spherical agent on a plane with a goal and ten spherical obstacles, with initial positions and velocities randomized in a rotation-invariant way. Expert offline trajectories are generated using TD3, and rewards combine goal distance, obstacle penalties, and control cost.

Environment

  • The environment has a spherical agent, a goal, and n = 10 spherical obstacles; agent position and velocity, obstacle positions, and goal position are randomly initialized in a rotation-invariant way.Dynamics are simulated with PyBullet [Coumans and Bai, 2016–2019].Source: C Navigation experiments · p. 16

Offline dataset

  • The offline dataset contains 105 expert trajectories generated after training a TD3 agent for 107 steps with default hyperparameters.The TD3 implementation is by Raffin et al. [2021], and the method is cited as [Fujimoto et al., 2018].Source: C Navigation experiments · p. 16

State and actions

  • The state includes agent position and velocity, goal position, and obstacle positions, while the two-dimensional action specifies force on the agent.Source: C Navigation experiments · p. 16

Rewards

  • Each step rewards negative Euclidean distance to the goal, adds a −0.1 obstacle-contact penalty and a −103 force control cost, then normalizes rewards so random policy scores 0 and expert policy 100.Source: C Navigation experiments · p. 16

D Kuka experiments

The Kuka experiments use three manipulation tasks and compare EDGI with a reimplemented Diffuser baseline. EDGI represents object poses with an equivariant Euclidean embedding, while the Diffuser baseline performs worse when given that reparameterization.

Tasks and baseline

  • The manipulation experiments cover unconditional stacking, conditional stacking, and block rearrangement, with Diffuser reimplemented with code fixes including proper environment resets.The environments and tasks are from Janner et al. [2022].Source: D Kuka experiments · p. 16

Robot state representation

  • The robot base orientation is represented by an xy-plane vector, and gravity is included as another vector so rotations can represent passive coordinate changes or active scene rotations.The paper treats this as a valid symmetry because the laws of physics are invariant to the transformation.Source: D Kuka experiments · p. 17

Object pose representation

  • EDGI represents each object’s rotation by the first two columns of its rotation matrix, which form an equivariant embedding into two unit orthogonal 3-vectors, and combines them with translation.A Gram-Schmidt procedure defines an equivariant map back to SO(3) almost everywhere.Source: D Kuka experiments · p. 17

Baseline result

  • The Diffuser baseline performs worse when tested on EDGI’s state reparameterization than with its original 39-dimensional Kuka-state parameterization.Source: D Kuka experiments · pp. 16–17

Hyperparameters

  • For each model and dataset, both linear and cosine noise schedules are tested, and the better result is reported.The cosine schedule follows Janner et al. [2022].Source: D Kuka experiments · p. 17
Loading 2303.12410v2…