Title: EDGI: Equivariant Diffusion for Planning with Embodied Agents
Venue: NeurIPS 2023
Accepted: yes
Keywords: Planning

Diffusion models

Equivariance

Equivariant generative models
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 ℤ, and the object permutation group Sₙ. EDGI follows the Diffuser framework by 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) × ℤ × Sₙ-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.

EDGI: Equivariant Diffusion for Planning with Embodied Agents

Abstract

Embodied agents operate in a structured world, often solving tasks with spatial, temporal, and permutation symmetries. Most algorithms for planning and modelbased 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 S n . EDGI follows the Diffuser framework by 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) × Z × S n -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

Our world is awash with symmetries. The laws of physics are the same everywhere in space and time-they are symmetric under translations and rotations of spatial coordinates as well as under time shifts. 4 In addition, whenever multiple identical or equivalent objects are labeled with numbers, the system is symmetric with respect to a permutation of the labels. Embodied agents are exposed to this structure, and many common robotic tasks exhibit spatial, temporal, or permutation symmetries. The gaits of a quadruped are independent of whether it is moving East or North, and a robotic gripper would interact with multiple identical objects independently of their labeling. However, most reinforcement learning (RL) and planning algorithms do not take this rich structure into account. While they have achieved remarkable success on well-defined problems after sufficient training, they are often sample-inefficient [Holland et al., 2018] and lack robustness to changes in the environment.

To improve the sample efficiency and robustness of RL algorithms, we believe it is paramount

to develop them with an awareness of their symmetries. Such algorithms should satisfy two key desiderata. First, policy and world models should be equivariant with respect to the relevant symmetry group. Often, for embodied agents this will be a subgroup of the product group of the spatial symmetry group SE(3) , the group of discrete time shifts Z , and one or multiple object permutation groups S n . Second, it should be possible to softly break (parts of) the symmetry group to solve concrete tasks. For example, a robotic gripper might be tasked with moving an object to a specific point in space, which breaks the symmetry group SE(3) . First works on equivariant RL have demonstrated the potential benefits of this approach [van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021, Muglich et al., 2022, Wang and Wal-

g

ters, 2022, Wang et al., 2022, Cetin et al., 2022, Rezaei-Shoshtari et al., 2022, Deac et al., 2023]. However, these works generally only consider small finite symmetry groups such as C n and do not usually allow for soft symmetry breaking at test time based on a specific task.

In this paper, we introduce the Equivariant Diffuser for Generating Interactions (EDGI), an equivariant algorithm for model-based reinforcement learning and planning. EDGI consists of a base component that is equivariant with respect to the full product group SE(3) × Z × S n and supports the multiple different representations of this group we expect to encounter in embodied environments. Moreover, EDGI allows for a flexible soft breaking of the symmetry at test time depending on the task.

Our work builds on the Diffuser method by Janner et al. [2022], who approach both the learning of a dynamics model and planning within it as a generative modeling problem. The key idea in Diffuser is to train a diffusion model on an offline dataset of state-action trajectories. To plan with this model, one samples from it conditionally on the current state, using classifier guidance to maximize the reward.

Our main contribution is a new diffusion model that is equivariant with respect to the product group SE(3) × Z × S n of spatial, temporal, and permutation symmetries and supports data consisting of multiple representations. We introduce a new way of embedding multiple input representations into a single internal representation, as well as novel temporal, object, and permutation layers that act on the individual symmetries. When integrated into a planning algorithm, our approach allows for a soft breaking of the symmetry group through test-time task specifications both through conditioning and classifier guidance.

We demonstrate EDGI empirically in 3D navigation and robotic object manipulation environments. We find that EDGI greatly improves the performance in the low-data regime--matching the performance of the best non-equivariant baseline when using an order of magnitude less training data. In addition, EDGI is significantly more robust to symmetry transformations of the environment, generalizing well to unseen configurations.

2 Background

Equivariant deep learning . Equivariant networks directly encode the symmetries described by a group G in their architecture. For the purposes of this paper, we are interested in the symmetries of 3D space, which include translations and rotations and are described by the special Euclidean group SE(3) , discrete-time translations Z , and object permutations, which are defined using the symmetric group of n elements S n . Wenow recall the main definition of equivariance and invariance of functions. Definition 1. A function f : X → Y is called G -equivariant if g · f ( x ) = f ( g · x ) for all g ∈ G and x ∈ X . Here X and Y are input and output spaces that carry a G action denoted by · . The function f is called G -invariant if the group action in Y is trivial, g · y = y .

We will focus on X = R n , Y = R m , and linear group actions or representations, which are group

homomorphisms ρ : G → GL( R k ) . Examples include rotation and permutation matrices. For a more comprehensive introduction to group and representation theory, we direct the interested reader to Appendix A or Esteves [2020] and Bronstein et al. [2021].

For generative modeling, we seek to model G -invariant densities. As proven in [Köhler et al., 2020, Bose and Kobyzev, 2021, Papamakarios et al., 2021], given a G -invariant prior density it is sufficient to construct a G -equivariant map to reach the desired G -invariant target density. In Sec. 3, we design G -equivariant diffusion architectures to model a distribution of trajectories that are known to be symmetric with respect to the product group SE(3) × Z × S n .

Diffusion models . Diffusion models [Sohl-Dickstein et al., 2015] are latent variable models that generate data by iteratively inverting a diffusion process. This diffusion process starts from a clean data sample x ∼ q ( x 0 ) and progressively injects noise for i ∈ [ T ] steps until the distribution is pure noise. The reverse, generative process takes a sample from a noise distribution and denoises it by progressively adding back structure, until we return to a sample that resembles being drawn from the empirical data distribution p ( x ) .

In diffusion models, it is customary to choose a parameter-free diffusion process (e. g. Gaussian noise with fixed variance). Specifically, we may define q ( x t | x t -1 ) as the forward diffusion distribution modeled as a Gaussian centered around the sample at timestep x t -1 : q ( x t | x t -1 ) = N ( x t ; √ 1 -β t x t -1 , β t I ) , where β t is a known variance schedule. The reverse generative process is learnable and can be parametrized using another distribution p θ ( x t -1 | x t ) = N ( x t -1 ; µ θ ( x t , t ) , σ 2 t I ) , and the constraint that the terminal marginal at time T is a standard Gaussian-i.e. p ( x T ) = N (0 , I ) . The generative process can be learned by maximizing a variational lower bound on the marginal likelihood. In practice, instead of predicting the mean of the noisy data, it is convenient to predict the noise level ϵ t directly [Ho et al., 2020]. Furthermore, to perform low-temperature sampling in diffusion models it is possible to leverage a pre-trained classifier to guide the generation process [Dhariwal and Nichol, 2021]. To do so we can modify the diffusion score by including the gradient of the log-likelihood of the classifier ¯ ϵ θ ( x t , t, y ) = ϵ θ ( x t , t ) -λσ t ∇ x t log p ( y | x t ) , where λ is the guidance weight and y is the label.

Trajectory optimization with diffusion . We are interested in modeling systems that are governed by discrete-time dynamics of a state s h +1 = f ( s h , a h ) , given the state s h and action a h taken at timestep h . The goal in trajectory optimization is then to find a sequence of actions a ∗ 0: H that maximizes an objective (reward) J which factorizes over per-timestep rewards r ( s h , a h ) . Formally, this corresponds to the optimization problem a ∗ 0: H = arg max a 0: H J ( s 0 , a 0: H ) = arg max a 0: H ∑ H h =0 r ( s h , a h ) , where H is the planning horizon and τ = ( s 0 , a 0 , . . . , s H , a H ) denotes the trajectory.

A practical method to solve this optimization problem is to unify the problem of learning a model of the state transition dynamics and the problem of planning with this model into a single generative modeling problem. Janner et al. [2022] propose to train a diffusion model on offline trajectory data consisting of state-action pairs, learning a density p θ ( τ ) . Planning can then be phrased as a conditional sampling problem: finding the distribution ˜ p θ ( τ ) ∝ p θ ( τ ) c ( τ ) over trajectories τ where c ( τ ) encodes constraints on the trajectories and specifies the task for instance as a reward function. Diffusion models allow conditioning in a way similar to inpainting in generative image modeling, and test-time reward maximization in analogy to classifier-based guidance.

3 Equivariant diffuser for generating interactions (EDGI)

We now describe our EDGI method. We begin by discussing the symmetry group SE(3) × Z × S n and common representations in robotic problems. In Sec. 3.2 we introduce our key novelty, an SE(3) × Z × S n -equivariant diffusion model for state-action trajectories τ . We show how we can sample from this model invariantly and break the symmetry in Sec. 3.3. Finally, we discuss how a diffusion model trained on offline trajectory data can be used for planning in Sec. 3.4.

3.1 Symmetry and representations

Symmetry group . We consider the symmetry group SE(3) × Z × S n , which is a product of three distinct groups: 1. the group of spatial translations and rotations SE(3) , 2. the discrete time translation symmetry Z , and 3. the permutation group over n objects S n . It is important to note, however, that this symmetry group may be partially broken in an environment. For instance, the direction of gravity

Internal

representation:

Equivariant blocks:

Representation mixer / unmixer:

Context block:

=

Mish

Linear

Embedding in

internal rep

Time

Objects

Channels

(Geometry suppressed)

=

Attention

over

objects

Temporal

conv.

Geometric

layer

Norm.

…

Input

trajectory

Diffusion time,

conditioning

Output

trajectory

usually breaks the spatial symmetry group SE(3) to the smaller group SE(2) , and distinguishable objects in a scene may break permutation invariance. We follow the philosophy of modeling invariance with respect to the larger group and including any symmetry-breaking effects as inputs to the networks.

We require that spatial positions are always expressed relative to a reference point, for example, the robot base or center of mass. This guarantees equivariance with respect to spatial translations: to achieve SE(3) equivariance, we only need to design an SO(3) -equivariant architecture.

Data representations . We consider 3D environments that contain an embodied agent as well as n other objects. We parameterize their degrees of freedom with two SO(3) representations, namely the scalar representation ρ 0 and the vector representation ρ 1 . Other representations that may occur in a dataset can often be expressed in ρ 0 and ρ 1 . For example, object poses may be expressed as 3D rotations between a global frame and an object frame. We can convert such rotations into two ρ 1 representations through the inclusion map ι : SO(3) ↪ → R 2 × 3 that chooses the first two columns of the rotation matrix. Such an embedding avoids challenging manifold optimization over SO(3) , while we can still uncover the rotation matrix by orthonormalizing using a Gram-Schmidt process to recover the final column vector. Thus, any SE(3) pose can be transformed to these two representations.

We assume that all trajectories transform under the regular representation of the time translation group Z (similar to how images transform under spatial translations). Under S n , object properties permute, while robot properties or global properties of the state remain invariant. Each feature is thus either in the trivial or the standard representation of S n .

Overall, we thus expect that data in environments experienced by our embodied agent to be categorized into four representations of the symmetry group SE(3) × Z × S n : scalar object properties, vector object properties, scalar robotic degrees of freedom (or other global properties of the system), and vector robotic degrees of freedom (again including other global properties of the system).

3.2 Equivariant diffusion model

Our main contribution is a novel SE(3) × Z × S n -equivariant diffusion model which leads to an invariant distribution over trajectories. Specifically, given an invariant base density with respect to our

chosen symmetry group-an isotropic Gaussian satisfies this property for SE(3) × Z × S n -and an denoising model f that is equivariant with respect to the same group, we arrive at a diffusion model that is SE(3) × Z × S n -invariant [Köhler et al., 2020, Papamakarios et al., 2021]. Under mild assumptions, such an equivariant map that pushes forward the base density always exists [Bose and Kobyzev, 2021].

We design a novel equivariant architecture for the denoising model f . Implemented as a neural network, it maps noisy input trajectories τ and a diffusion time step i to an estimate ˆ ϵ of the noise vector that generated the input. Our architecture does this in three steps. First, the input trajectory consisting of various representations is transformed into an internal representation of the symmetry group. Second, in this representation the data are processed with an equivariant network. Finally, the outputs are transformed from the internal representation into the original data representations present in the trajectory. We illustrate the architecture of our EDGI model in Fig. 2.

Step 1: Representation mixer . The input noisy trajectory consists of features in different representations of the symmetry group (see above). While it is possible to mirror these input representations for the hidden states of the neural network, the design of equivariant architectures is substantially simplified if all inputs and outputs transform under a single representation. Hence, we decouple the data representation from the representation used internally for the computation-in a similar fashion to graph neural networks that decouple the data and computation graphs.

Internal representation . We define a single internal representation that for each trajectory time step t ∈ [ H ] , each object o ∈ [ n ] , each channel c ∈ [ n c ] consists of one 5 SO(3) scalar s toc and one SO(3) vector v toc . We write w toc = ( s toc , v toc ) ∈ R 4 . Under spatial rotations g ∈ SO(3) , these features thus transform as the direct sum of the scalar and vector representations ρ 0 ⊕ ρ 1 :

These internal features transform in the regular representation under time shift and in the standard representation under permutations P as w toc → w to ′ c = ∑ o P o ′ o w toc . There are thus no global (not object-specific) properties in our internal representations.

Transforming input representations into internal representations . The first layer in our network transforms the input τ , which consists of features in different representations of SE(3) × Z × S n , into the internal representation. On the one hand, we pair up SO(3) scalars and SO(3) vectors into ρ 0 ⊕ ρ 1 features. On the other hand, we distribute global features - those a priori unassigned to one of the n objects in the scene - over the n objects.

Concretely, for each object o ∈ [ n ] , each trajectory step t ∈ [ H ] , and each channel c = [ n c ] , we define the input in the internal representation as w toc ∈ R 4 as follows:

The matrices W 1 , 2 , 3 , 4 are learnable and of dimension n × n c × n object s , n × n c × n object v , n × n c × n global s , or n × n c × n global v , respectively. Here n object s is the number of SO(3) scalar quantities associated with each object in the trajectory, n object v is the number of SO(3) vector quantities associated with each object, n global s is the number of scalar quantities associated with the robot or global properties of the system, and n global v is the number of vectors of that nature. The number of input channels n c is a hyperparameter. We initialize the matrices W i such that Eq. (2) corresponds to a concatenation of all object-specific and global features along the channel axis at the beginning of training.

Step 2: SE(3) × Z × S n -equivariant U-net . We then process the data with a SO(3) × Z × S n -equivariant denoising network. Its key components are three alternating types of layers. Each type acts on the representation dimension of one of the three symmetry groups while leaving the other two invariant-i. e. they do not mix internal representation types of the other two layers:

Temporal layers : Time-translation-equivariant convolutions along the temporal direction (i. e. along trajectory steps), organized in a U-Net architecture.

Object layers : Permutation-equivariant self-attention layers over the object dimension.

Geometric layers : SO(3) -equivariant interaction between the scalar and vector features.

In addition, we use residual connections, a new type of normalization layer that does not break equivariance, and context blocks that process conditioning information and embed it in the internal representation (see Appendix B for more details). These layers are combined into an equivariant block consisting of one instance of each layer, and the equivariant blocks are arranged in a U-net, as depicted in Fig. 2. Between the levels of the U-net, we downsample (upsample) along the trajectory time dimension by factors of two, increasing (decreasing) the number of channels correspondingly.

Temporal layers . Temporal layers consist of 1 Dconvolutions along the trajectory time dimension. To preserve SO(3) equivariance, these convolutions do not add any bias and there is no mixing of features associated with different objects nor the four geometric features of the internal SO(3) representation.

Object layers . Object layers enable features associated with different objects to interact via an equivariant multi-head self-attention layer. Given inputs w toc , the object layer computes

with learnable weight matrices W K,V,Q and d the dimensionality of the key vector. There is no mixing between features associated with different time steps, nor between the four geometric features of the internal SO(3) representation. Object layers are SO(3) -equivariant, as the attention weights compute invariant SO(3) norms.

Geometric layers . Geometric layers enable mixing between the scalar and vector quantities that are combined in the internal representation, but do not mix between different objects or across the time dimension. We construct an expressive equivariant map between scalar and vector inputs and outputs following Villar et al. [2021]: We first separate the inputs into SO(3) scalar and vector components, w toc = ( s toc , v toc ) T . We then construct a complete set of SO(3) invariants by combining the scalars and pairwise inner products between the vectors,

These are then used as inputs to two MLPs ϕ and ψ , and finally we get output scalars and vectors, w ′ toc = ( ϕ ( S to ) c , ∑ c ′ ψ ( S to ) cc ′ v toc ′ ) .

Assuming full expressivity of the MLPs ϕ and ψ , this approach can approximate any equivariant map between SO(3) scalars and vectors [Villar et al., 2021, Proposition 4]. In this straightforward form, however, it can become prohibitively expensive, as the number of SO(3) invariants S to scales quadratically with the number of channels. In practice, we first linearly map the input vectors into a smaller number of vectors, apply this transformation, and increase the number of channels again with another linear map.

Step 3: Representation unmixer . The equivariant network outputs internal representations w toc that are transformed back to data representations using linear maps, in analogy to Eq. (2). Global properties, e. g. robotic degrees of freedom, are aggregated from the object-specific internal representations by taking the elementwise mean across the objects. We find it beneficial to apply an additional geometric layer to these aggregated global features before separating them into the original representations.

Training . We train EDGI on offline trajectories without any reward information. We optimize for the simplified variational lower bound L = E τ,i,ϵ [ ∥ ϵ -f ( τ + ϵ ; i ) ∥ 2 ] [Ho et al., 2020]. where τ are training trajectories, i ∼ Uniform(0 , T ) is the diffusion time step, and ϵ is Gaussian noise with variance depending on a prescribed noise schedule.

3.3 Invariant sampling and symmetry breaking

We now discuss sampling from EDGI (and, more generally, from equivariant diffusion models). While unconditional samples follow an invariant density, conditional samples may either be invariant or break the symmetry of the diffusion model.

Invariant sampling . It is well-known that unconditional sampling from an equivariant denoising model defines an invariant density [Köhler et al., 2020, Bose and Kobyzev, 2021, Papamakarios et al., 2021]. We repeat this result without proof:

Proposition 1. Consider a group G that acts on R n with representation ρ . Let π be a G -invariant distribution over R n and ϵ t : R n → R n be a G -equivariant noise-conditional denoising network. Then the distribution defined by the denoising diffusion process of sampling from π and iteratively applying ϵ t is G -invariant.

We now extend this result to sampling with classifier-based guidance [Dhariwal and Nichol, 2021], a technique for low-temperature sampling based on a classifier log p ( y | x t ) with class labels y , or more generally any guide h ( x ) . When the guide is G -invariant, guided sampling retains G invariance:

Proposition 2. Consider a group G that acts on R n with representation ρ . Let π be a G -invariant density over R n and ϵ t : R n → R n be a G -equivariant noise-conditional denoising network. Let the guide h : R n → R be a smooth G -invariant function. Further, assume ρ ( g ) is orthogonal ∀ g ∈ G . Define the modified diffusion score ¯ ϵ θ ( x t , t, c ) = ϵ θ ( x t , t ) -λσ t ∇ x t h ( x t ) for some guidance weight λ ∈ R . Then the distribution defined by the denoising diffusion process of sampling from π and iteratively applying ¯ ϵ is G -invariant.

Proof. The function h has a gradient ∇ x h ( x ) that is G -equivariant [Papamakarios et al., 2021, Lemma 2]. Thus,

The modified diffusion score ¯ ϵ is therefore G -equivariant. Applying Prop. 1, we find that classifierbased guidance samples from a G -invariant distribution.

Proposition 2 applies to G = SE(3) × Z × S n as employed in EDGI, as each group within the product admits an orthogonal matrix representation. Both unconditional sampling from EDGI, as well as sampling guided by a SE(3) × Z × S n -invariant classifier (or reward model), is thus SE(3) × Z × S n -invariant.

Symmetry breaking . However, if the classifier h ( x ) or log p ( y | x ) is not G -invariant, classifierguided samples are in general also not G -invariant. For example, consider h ( x ) = -∥ x -x 0 ∥ 2 for some x 0 ∈ R n , which will bias samples to x 0 . Similarly, conditional sampling on components of x (similar to inpainting) clearly leads to non-invariant samples. As we will argue below, these properties are essential for robotic planning. 6

3.4 Planning with equivariant diffusion

A diffusion model trained on offline trajectory data jointly learns a world model and a policy. Following Janner et al. [2022], we use it to solve planning problems by choosing a sequence of actions to maximize the expected task rewards.

To do this, we use three features of diffusion models. The first is the ability to sample from them by drawing noisy trajectory data from the base distribution and iteratively denoising them with the learned network yielding trajectories similar to those in the training set. For such sampled trajectories to be useful for planning, they need to begin in the current state of the environment. We achieve this by conditioning the sampling process such that the initial state of the generated trajectories matches the current state, in analogy to inpainting. Finally, we can guide this sampling procedure toward solving concrete tasks specified at test time using classifier-based guidance where a regression model is trained offline to map trajectories to task rewards.

Task-specific symmetry breaking . By construction, our equivariant diffusion model learns a SE(3) × Z × S n -invariant density over trajectories. As shown in the previous section, both unconditional samples (and samples guided by an invariant classifier) reflect this symmetry property-it will be equally likely to sample a trajectory and its rotated or permuted counterpart. However, concrete tasks will often break this invariance, for instance by requiring that a robot or object is brought into a particular location or specifying an ordering over objects to be manipulated in a scene.

As discussed in the previous section, our diffusion-based approach with classifier guidance allows us to elegantly break the symmetry at test time as required. Such a soft symmetry breaking both occurs through conditioning on the current state, by conditioning on a goal state, and through a non-invariant reward model used for guidance during sampling.

4 Experiments

We demonstrate the effectiveness of incorporating symmetries as a powerful inductive bias in the Diffuser algorithm with experiments in two environments. The first environment is a 3 D navigation task, in which an agent needs to navigate a number of obstacles to reach a goal state. Rewards are awarded based on the distance to the goal at each step, with penalties for collisions with obstacles. The position of the obstacles and the goal state are different in each episode and part of the observation. For simplicity, the actions directly control the acceleration of the agent; both the agent and the obstacles are spherical. Please see Fig. 1 for a schematic representation of this task and Appendix C for more details and the reward structure for this task.

In our remaining experiments, the agent controls a simulated Kuka robotic arm interacting with four blocks on a table. We use a benchmark environment introduced by Janner et al. [2022], which specifies three different tasks: an unconditional block stacking task, a conditional block stacking task where the stacking order is specified, and a rearrangement problem, in which the stacking order has to be changed in a particular way. For both environments, we train on offline trajectory datasets of roughly 10 5 (navigation) or 10 4 (manipulation) trajectories. We describe the setup in detail in Appendix D.

Algorithms . We train our EDGI on the offline dataset and use conditional sampling to plan the next actions. For the conditional and rearrangement tasks in the Kuka environment, we use classifier guidance following Janner et al. [2022].

As our main baseline, we compare our results to the (non-equivariant) Diffuser model [Janner et al., 2022]. In addition to a straightforward model, we consider a version trained with SO(3) data augmentation. We also compare two model-based RL baselines reported by Janner et al. [2022], BCQ [Fujimoto et al., 2019] and CQL [Kumar et al., 2020]. To study the benefits of the symmetry groups in isolation, we construct two EDGI variations: one is equivariant with respect to SE(3) , but not S n ; while the other is equivariant with respect to S n , but not SE(3) . Both are equivariant to temporal translations, just like EDGI and the baseline Diffuser.

Task performance . We report the results on both navigation and object tasks in Tab. 1. For each environment, we evaluate 100 episodes and report the average reward and standard error for each method. In both environments and across all tasks, EDGI performs as well as or better than the Diffuser baseline when using the full training set. In the navigation task, achieving a good performance for the baseline required substantially increasing the model's capacity compared to the hyperparameters used in Janner et al. [2022]. On the Kuka environment, both diffusion-based methods clearly outperform the BCQ and CQL baselines.

Sample efficiency . We study EDGI's sample efficiency by training models on subsets of the training

10

4

10

3

10

2

10

1

10 0

Relative dataset size

20

0

20

40

60

80

100

Reward

EDGI (ours)

SE(3)-equivariant only

Sn

-equivariant only

Diffuser

Diff. + augmentation

10

3

10

2

10

1

10 0

Relative dataset size

0

10

20

30

40

50

60

Reward (average over 3 tasks)

EDGI (ours)

Diffuser

data. The results in Fig. 3 show that EDGI achieves just as strong rewards in the Kuka environment when training with only on 10% of the training data, and on the navigation task even when training on only 0 . 01% if the training data. The Diffuser baseline is much less sample-efficient. Training the Diffuser model with data augmentation partially closes the gap, but EDGI still maintains an edge. Our results provide evidence for the benefits of the inductive bias of equivariant models and matches similar observations in other works for using symmetries in an RL context [van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021, Rezaei-Shoshtari et al., 2022, Deac et al., 2023].

Effects of individual symmetries . In the left panel of Fig. 3, we also show results for EDGI variations that are only equivariant with respect to SE(3) , but not S n , or vice versa. Both partially equivariant methods perform better than the Diffuser baseline, but not as well as the EDGI model equivariant to the full product group SE(3) × Z × S n . This confirms that the more of the symmetry of a problem we take into account in designing an architecture, the bigger the benefits in sample efficiency can be.

Group generalization . Finally, we demonstrate that equivariance improves generalization across the SO(3) symmetry group. On both environments, we train EDGI and Diffuser models on restricted offline datasets in which all trajectories are oriented in a particular way. In particular, in the navigation environment, we only use training data that navigates towards a goal location with x = 0 . In the robotic manipulation tasks, we only use training trajectories where the red block is in a position with x = 0 at the beginning of the episode. We test all agents on the original environment, where they encounter goal positions and block configurations unseen during training. We show results for these experiments in Tab. 1. The original Diffuser performs substantially worse, showing its limited capabilities to generalize to the new setting. In contrast, the performance of EDGI is robust to this domain shift, confirming that equivariance helps in generalizing across the symmetry group.

5 Related work

Diffusion-based planning . The closest work to ours is the original Diffuser paper [Janner et al., 2022], which we used as a baseline. Concurrent to our work, Diffuser was extended by Ajay et al. [2022], who used a separate inverse dynamics model and classifier-free guidance. The key novelty of our work is that we make this approach aware of the symmetry structure of planning problems through a new SE(3) × Z × S n -equivariant denoising network.

Equivariant deep learning . Baking in symmetries into deep learning architectures was first studied in the work of Cohen and Welling [2016a] for geometric transformations, and the DeepSet architecture for permutations [Zaheer et al., 2017]. Followup work to group convolutional networks focused on both spherical geometry [Cohen et al., 2018] and building kernels using irreducible group representations [Cohen and Welling, 2016b, Weiler and Cesa, 2019, Cesa et al., 2021]. For symmetries of the 3D space-i. e. subgroups of E (3) -a dominant paradigm is to use the message passing framework [Gilmer et al., 2017] along with geometric quantities like positions, velocities, and relative angles [Satorras et al., 2021, Schütt et al., 2021, Batatia et al., 2022].

Equivariance in RL . The role of symmetries has also been explored in reinforcement learning problems with a body of work focusing on symmetries of the joint state-action space of an MDP

[van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021, Muglich et al., 2022, Wang and Walters, 2022, Wang et al., 2022, Cetin et al., 2022, Rezaei-Shoshtari et al., 2022]. More recently, model-based approaches-like our proposed EDGI-have also benefited from increased data efficiency through the use of symmetries of the environment [Deac et al., 2023]. Concurrently to this work, Brehmer et al. [2023] also experiment with an equivariant Diffuser variation, but their denoising network is based on geometric algebra representations and a transformer architecture.

Equivariant generative models . Early efforts in learning invariant densities using generative models utilized the continuous normalizing flow (CNF) framework. A variety of works imbued symmetries by designing equivariant vector fields [Köhler et al., 2020, Rezende and Mohamed, 2015, Bose and Kobyzev, 2021]. As flow-based models enjoy exact density estimation, their application is a natural fit for applications in theoretical physics [Boyda et al., 2020, Kanwar et al., 2020] and modeling equivariant densities on manifolds [Katsman et al., 2021]. Other promising approaches to CNFs include equivariant score matching [De Bortoli et al., 2022] and diffusion models [Hoogeboom et al., 2022, Xu et al., 2022, Igashov et al., 2022]. Our proposed EDGI model extends the latter category to the product group SE(3) × Z × S n and increases flexibility with respect to the data representations.

6 Discussion

Embodied agents often solve tasks that are structured through the spatial, temporal, or permutation symmetries of our 3D world. Taking this structure into account in the design of planning algorithms can improve sample efficiency and generalization-notorious weaknesses of RL algorithms.

We introduced EDGI, an equivariant planning algorithm that operates as conditional sampling in a generative model. The main innovation is a new diffusion model that is equivariant with respect to the symmetry group SE(3) × Z × S n of spatial, temporal, and object permutation symmetries. Beyond this concrete architecture, our work presents a general blueprint for the construction of networks that are equivariant with respect to a product group and support multiple representations in the data. Integrating this equivariant diffusion model into a planning algorithm allows us to model an invariant base density, but still solve non-invariant tasks through task-specific soft symmetry breaking. We demonstrated the performance, sample efficiency, and robustness of EDGI on object manipulation and navigation tasks.

While our work shows encouraging results, training and planning are currently expensive. Progress on this issue can come both from more efficient layers in the architecture of the denoising model as well as from switching to recent continuous-time diffusion methods with accelerated sampling.

References

Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and P. Agrawal. Is conditional generative modeling all you need for decision-making? International Conference on Learning Representations , 2022. (Cited on page 9)

Batatia, D. P. Kovács, G. N. Simm, C. Ortner, and G. Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. arXiv preprint arXiv:2206.07697 , 2022. (Cited on page 9)

J. Bose and I. Kobyzev. Equivariant finite normalizing flows. arXiv preprint arXiv:2110.08649 , 2021. (Cited on pages 3, 5, 6, and 10)

Boyda, G. Kanwar, S. Racanière, D. J. Rezende, M. S. Albergo, K. Cranmer, D. C. Hackett, and P. E. Shanahan. Sampling using su ( n ) gauge equivariant flows. arXiv preprint arXiv:2008.05456 , 2020. (Cited on page 10)

Brehmer, P. de Haan, S. Behrends, and T. Cohen. Geometric algebra transformers. In Advances in Neural Information Processing Systems , volume 37, 2023. URL https://arxiv.org/abs/ 2305.18415 . (Cited on page 10)

M. Bronstein, J. Bruna, T. Cohen, and P. Veliˇ ckovi´ c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 , 2021. (Cited on page 3)

Cesa, L. Lang, and M. Weiler. A program to build e (n)-equivariant steerable cnns. In International Conference on Learning Representations , 2021. (Cited on page 9)

Cetin, B. Chamberlain, M. Bronstein, and J. J. Hunt. Hyperbolic deep reinforcement learning. arXiv [cs.LG] , Oct. 2022. (Cited on pages 2 and 10)

Cohen and M. Welling. Group equivariant convolutional networks. In International conference on machine learning , pages 2990-2999. PMLR, 2016a. (Cited on page 9)

S. Cohen and M. Welling. Steerable cnns. arXiv preprint arXiv:1612.08498 , 2016b. (Cited on page 9)

S. Cohen, M. Geiger, J. Köhler, and M. Welling. Spherical cnns. arXiv preprint arXiv:1801.10130 , 2018. (Cited on page 9)

Coumans and Y. Bai. PyBullet, a Python module for physics simulation for games, robotics and machine learning. http://pybullet.org , 2016-2019. (Cited on page 16)

De Bortoli, E. Mathieu, M. J. Hutchinson, J. Thornton, Y. W. Teh, and A. Doucet. Riemannian Score-Based generative modelling. Oct. 2022. URL https://openreview.net/pdf?id= oDRQGo8I7P . (Cited on pages 10 and 17)

Deac, T. Weber, and G. Papamakarios. Equivariant MuZero. Feb. 2023. (Cited on pages 2, 9, and 10)

Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021. (Cited on pages 3 and 7)

Esteves. Theoretical aspects of group equivariant neural networks. arXiv preprint arXiv:2004.05154 , 2020. (Cited on page 3)

Fujimoto, H. Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning , pages 1587-1596. PMLR, 2018. (Cited on page 16)

Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019. (Cited on page 8)

Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In International conference on machine learning , pages 1263-1272. PMLR, 2017. (Cited on page 9)

C. Hall. Lie groups, Lie algebras, and representations . Springer, 2013. (Cited on page 14)

Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020. (Cited on pages 3 and 6)

Z. Holland, E. J. Talvitie, and M. Bowling. The effect of planning shape on dyna-style planning in high-dimensional state spaces. arXiv preprint arXiv:1806.01825 , 2018. (Cited on page 1)

Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning , pages 8867-8887. PMLR, 2022. (Cited on page 10)

C.-W. Huang, M. Aghajohari, A. J. Bose, P. Panangaden, and A. Courville. Riemannian diffusion models. Aug. 2022. URL http://arxiv.org/abs/2208.07949 . (Cited on page 17)

Igashov, H. Stärk, C. Vignac, V. G. Satorras, P. Frossard, M. Welling, M. Bronstein, and B. Correia. Equivariant 3d-conditional diffusion models for molecular linker design. arXiv preprint arXiv:2210.05274 , 2022. (Cited on page 10)

Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022. (Cited on pages 1, 2, 3, 7, 8, 9, 15, 16, and 17)

Kanwar, M. S. Albergo, D. Boyda, K. Cranmer, D. C. Hackett, S. Racaniere, D. J. Rezende, and P. E. Shanahan. Equivariant flow-based sampling for lattice gauge theory. Physical Review Letters , 125(12):121601, 2020. (Cited on page 10)

Katsman, A. Lou, D. Lim, Q. Jiang, S.-N. Lim, and C. De Sa. Equivariant manifold flows. In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models , 2021. (Cited on page 10)

Köhler, L. Klein, and F. Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning , pages 5361-5370. PMLR, 2020. (Cited on pages 3, 5, 6, and 10)

Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020. (Cited on page 8)

K. Mondal, V. Jain, K. Siddiqi, and S. Ravanbakhsh. EqR: Equivariant representations for DataEfficient reinforcement learning. Nov. 2021. (Cited on pages 2, 9, and 10)

Muglich, C. S. de Witt, E. van der Pol, S. Whiteson, and J. Foerster. Equivariant networks for Zero-Shot coordination. Oct. 2022. (Cited on pages 2 and 10)

Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research , 22(1): 2617-2680, 2021. (Cited on pages 3, 5, 6, and 7)

Raffin, A. Hill, A. Gleave, A. Kanervisto, M. Ernestus, and N. Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22(268): 1-8, 2021. URL http://jmlr.org/papers/v22/20-1364.html . (Cited on page 16)

Rezaei-Shoshtari, R. Zhao, P. Panangaden, D. Meger, and D. Precup. Continuous mdp homomorphisms and homomorphic policy gradient. arXiv preprint arXiv:2209.07364 , 2022. (Cited on pages 2, 9, and 10)

J. Rezende and S. Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd international conference on Machine learning . ACM, 2015. (Cited on page 10)

G. Satorras, E. Hoogeboom, and M. Welling. E (n) equivariant graph neural networks. In International conference on machine learning , pages 9323-9332. PMLR, 2021. (Cited on page 9)

Schütt, O. Unke, and M. Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning , pages 93779388. PMLR, 2021. (Cited on page 9)

Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 22562265. PMLR, 2015. (Cited on page 3)

van der Pol, T. Kipf, F. A. Oliehoek, and others. Plannable approximations to MDP homomorphisms: Equivariance under actions. arXiv preprint arXiv , 2020. (Cited on pages 2, 9, and 10)

Villar, D. W. Hogg, K. Storey-Fisher, W. Yao, and B. Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics. Advances in Neural Information Processing Systems , 34:28848-28863, 2021. (Cited on page 6)

Walters, J. Li, and R. Yu. Trajectory prediction using equivariant continuous convolution. Oct. 2020. (Cited on pages 2, 9, and 10)

Wang and R. Walters. So (2) equivariant reinforcement learning. In International Conference on Learning Representations , 2022. (Cited on pages 2 and 10)

Wang, M. Jia, X. Zhu, R. Walters, and R. Platt. On-robot learning with equivariant models. In Conference on robot learning , 2022. (Cited on pages 2 and 10)

Weiler and G. Cesa. General e (2)-equivariant steerable cnns. Advances in Neural Information Processing Systems , 32, 2019. (Cited on page 9)

Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923 , 2022. (Cited on page 10)

Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. J. Smola. Deep sets. Advances in neural information processing systems , 30, 2017. (Cited on page 9)

A Group theory

In this appendix, we provide a basic introduction to Lie groups with an emphasis on the Lie group SE(3) which is used in the main paper.

A Lie group is both a group and a smooth manifold. A manifold is a topological space and thus a Lie group is a topological group that is formally defined below.

Definition 2. Atopological group is a topological space G equipped with continuous maps that allow group composition ◦ : G × G → G and group inverses ( ◦ ) -1 : G → G .

Now let g ∈ G be an element of the topological group and consider the map ¯ g : G → G given by ¯ g ( g ) = g ◦ g ′ . If g ◦ g 1 = g ◦ g 2 then g 1 = g 2 and thus ¯ g is injective and for all g ∈ G , and there exists an element-i.e. g -1 ◦ g such that ¯ g ( g -1 ◦ g ) = g ◦ g -1 ◦ g ′ = g ′ . Thus, ¯ g is a bijection.

In a topological space ¯ g and ¯ g -1 are continuous. Thus, ¯ g is a homomorphism. This means if U ⊂ G is open 1.) g ◦ U is open 2.) U ◦ g is open and inversion is also a homomorphism so U is open if and only if U -1 is open. A topological space X is said to be homogenous if ∀ x, y ∈ X there exists a map h x,y : X → X that is a homomorphism such that h x,y ( x ) = y . Note that a topological group is always homogenous. In words, a homogenous space for a topological group means any group element is reachable by a suitable group homomorphism.

Matrix Lie groups . A group of special interest is the group of n × n invertible matrices with entries in R , GL n ( R ) . This is a topological group equipped with Euclidean topology. The subgroup SL n ( R ) ⊂ GL n ( R ) is a closed subgroup with the property that det = 1 . As an example SE(3) ⊂ E (3) satisfies this definition and is the group of 3D Euclidean symmetries that include rotations and translations, but not reflections.

A matrix Lie group is a closed subgroup of GL n ( R ) . 7 This means SL n ( R ) is a matrix Lie group as are O ( n ) and SO ( n ) . Moreover, O ( n ) and SO ( n ) are compact. Finally, it is important to note that for matrix Lie groups the exponential and logarithmic maps correspond to the matrix exponential and matrix logarithm. Both of these are infinite power series and have precise connections to the representation theory of matrix Lie groups; namely, the Lie algebra associated with the Lie groups. While we pause further discussion on representation theory here the interested reader is encouraged to read Hall [2013].

SE(3) as a matrix Lie group . The group of Euclidean symmetries in 3D is known as the E (3) . A closed subgroup of this group is the special Euclidean group, SE(3) , and corresponds to rotations and translations in 3D . Commonly, elements of SE(3) can be used to represent rigid body transformations in 3 dimensions. This is because SE(3) ∼ = SO(3) ⋉R 3 and thus can be written as,

Represented by this 4 × 4 matrix and with the group operation defined by matrix multiplication, this group can be seen as a subgroup of the general linear group GL 4 ( R ) .

SO(3) as a matrix Lie group . The group of 3D rotations is SO(3) . It is a compact matrix Lie group where each element R ∈ SO(3) is a rotation. To represent R there are various possible choices. The most familiar of them is perhaps R ∈ R 3 × 3 is as a 3 × 3 rotation matrix.

An alternative parameterization is the rotation vector . Such a parametrization exploits the Lie algebra of SO(3) which are skew-symmetric matrices g . Each element of the Lie algebra can be distinctively associated with a vector ω ∈ R 3 . Given any v ∈ R 3 , it's related by g v = ω × v , with the symbol × representing the cross product. The length of this vector, symbolized as ω = || ω || , represents the rotation angle, while the unit direction, expressed as e ω = ω || ω || , indicates the rotation axis.

Yet another parametrization of rotations can be achieved using quaternions . A quaternion, q , is often represented as q = w + xi + yj + zk where w,x,y, and z are real numbers and i, j, and k are mutually orthogonal imaginary units vectors. The unit vectors obey the following rule: i 2 = j 2 = k 2 = ijk = -1 . A subset of quaternions, known as rotation quaternions, can be used to represent 3D rotations. These specific quaternions have magnitude 1 and can also be thought of as an ordered

pair ( w,v ) where w is a scalar and v = ( x, y, z ) is a 3D vector. The scalar part w can be thought of as cos( ω/ 2) and the vector part v as n sin( ω/ 2) , where ω is the angle of rotation and n is the unit vector along the axis of rotation.

Finally, to rotate a point p in 3D space using a quaternion q , the point can be represented as a pure quaternion p ′ = 0 + xi + yj + zk . The rotated point p ′′ is then given by: p ′′ = q · p ′ · q ∗ . Where q ∗ is the conjugate of q -i.e. if q = w + xi + yj + zk , then q ∗ = w -xi -yj -zk .

B Architecture details

On a high level, EDGI follows Diffuser [Janner et al., 2022]. In the following, we will describe the key difference: our SE(3) × Z × S n -equivariant architecture for the diffusion model.

Overall architecture . We illustrate the architecture in After converting the input data in our internal representation (see Sec. 3.2), the data is processed with an equivariant U -net with four levels. At each level, we process the hidden state with two residual standard blocks, before downsampling (in the downward pass) or upsampling (in the upward pass).

Residual standard block . The main processing unit of our architecture processes the current hidden state with an equivariant block consisting of a temporal layer, an object layer, a normalization layer, and a geometric layer. In parallel, the context information (an embedding of diffusion time and a conditioning mask) is processed with a context block. The hidden state is added to the output of the context block and processes with another equivariant block. Finally, we process the data with a linear attention layer over time. This whole pipeline consists of an equivariant block, a context block, and another equivariant block is residual (the inputs are added to the outputs).

Temporal layers . Temporal layers consist of one-dimensional convolutions without bias along the time dimension. We use a kernel size of 5.

Normalization layers . We use a simple equivariant normalization layer that for each batch element rescales the entire tensor w toc to unit variance. This is essentially an equivariant version of LayerNorm. The difference is that our normalization layer does not shift the inputs to zero means, as that would break equivariance with respect to SO(3) .

Geometric layers . In the geometric layers, the input state is split into scalar and vector components. The vector components are linearly transformed to reduce the number of channels to 16. We then construct all SO(3) invariants from these 16 vectors by taking pairwise inner products and concatenating them with the scalar inputs. This set of scalars is processed with two MLPs, each consisting of two hidden layers and ReLU nonlinearities. The MLPs output the scalar outputs and coefficients for a linear map between the vector inputs and the vector outputs, respectively. Finally, there is a residual connection that adds the scalar and vector inputs to the outputs.

Linear attention over time . To match the architecture used by Janner et al. [2022] as closely as possible, we follow their choice of adding another residual linear attention over time at the end of each level in the U-net. We make the linear attention mechanism equivariant by computing the attention weights as

Context blocks . The embeddings of diffusion time and conditioning information are processed with a Mish nonlinearity and a linear layer, like in Janner et al. [2022]. Finally, we embed them in our internal representation by zero-padding the resulting tensor.

Upsampling and downsampling . During the downsampling path, there is a final temporal layer that implements temporal downsampling and increases the number of channels by a factor of two. Conversely, during the upsampling path, we use a temporal layer for temporal upsampling and a reduction of the number of channels.

Equivariance . We now demonstrate the equivariance of the EDGI architecture explicitly. For concreteness, we focus on the geometric layers, as they are the most novel, and on both SE(3) transformations and permutations. Similar arguments can be made for the other layers and for equivariance with respect to temporal translations.

Let w ∈ R n × H × c × 4 be data in our internal representation, such that the entries w toc decompose into SO(3) scalars s toc and SO(3) vectors v toc . Let S ( w to ) be the set of all scalars and all pairwise SO (3) inner products between the vectors v to , as defined in Eq. (4). The outputs of the geometric layer are

First, consider what happens under permutations of the objects, w toc → w tπ ( o ) c for a permutation π ∈ S n . Wehave f ( π · w ) toc = ( ϕ ( S ( w tπ ( o ) )) , ∑ c ′ ψ ( S ( w tπ ( o ) )) v tπ ( o ) c ′ ) = f ( w ) tπ ( o ) c = ( π · f ( w )) toc . Thus, because this layer 'leaves the object dimension untouched', it is equivariant with respect to object permutations.

Next, consider the behavior under spatial transformations. Like most (S)E(3)-equivariant architectures, we deal with translations through canonicalization, defining all coordinates with respect to the center of mass or the robot base, as applicable. This means we only have to analyze the behavior under rotations.

Let R ∈ SO(3) , such that R · w = R · ( s, v ) = ( s, R · v ) . By definition, orthogonal matrices leave the inner product invariant, thus S ( R · w ) = S ( w ) . The geometric layer applied to rotated inputs then gives f ( R · w ) toc = ( ϕ ( S ( R · w to )) , ∑ c ′ ψ ( S ( w to )) R · v toc ′ ) = ( ϕ ( S ( w to )) , R · ∑ c ′ ψ ( S ( w to )) v toc ′ ) = ( R · f ( w )) toc . Hence the geometric layer is equivariant with respect to SE(3).

C Navigation experiments

We introduce a new navigation environment. The scene consists of a spherical agent navigating a plane populated with a goal state and n = 10 spherical obstacles. At the beginning of every episode, the agent position, agent velocity, obstacle positions, and goal position are initialized randomly (in a rotation-invariant way). We simulate the environment dynamics with PyBullet [Coumans and Bai, 2016-2019].

Offline dataset . To obtain expert trajectories, we train a TD3 [Fujimoto et al., 2018] agent in the implementation by Raffin et al. [2021] for 10 7 steps with default hyperparameters on this environment. We generate 10 5 trajectories for our offline dataset.

State . The state contains the agent position, agent velocity, goal position, and obstacle positions.

Actions . The action space is two-dimensional and specifies a force acting on the agent.

Rewards . At each time step, the agent receives a reward equal to the negative Euclidean distance to the goal state. In addition, a penalty of -0 . 1 is added to the reward if the agent touches any of the obstacles. Finally, there is an additional control cost equal to -10 3 times the force acting on the agent. We affinely normalize the rewards such that a normalized reward of 0 corresponds to that achieved by a random policy and a normalized reward of 100 corresponds to the expert policy.

D Kuka experiments

We use the object manipulation environments and tasks from Janner et al. [2022], please see that work for details on the environment. In our experiments, we consider three tasks: unconditional stacking, conditional stacking, and block rearrangement. Figure 4 visualizes the unconditional block stacking task. For a fair comparison, we re-implement the Diffuser algorithm while making bug fixes in the codebase of Janner et al. [2022], which mainly included properly resetting the environment.

State . We experiment with two parameterizations of the Kuka environment state. For the Diffuser baseline, we use the original 39-dimensional parameterization from Janner et al. [2022].

For our EDGI, we need to parameterize the system in terms of SE(3) × Z × S n representations. We, therefore, describe the robot and block orientations with SO(3) vectors as follows. Originally, the

robot state is specified through a collection of joint angles. One of these encodes the rotation of the base along the vertical z -axis. We choose to represent this angle as a ρ 1 vector in the xy -plane. In addition, we add the gravity direction (the z -axis itself) as another ρ 1 vector, which is also the normal direction of the table on which the objects rest. Combined, these vectors define the pose of the base of the robot arm. Rotating gravity direction, and the robot and object pose by SO(3) can be interpreted as a passive coordinate transformation, or as an active rotation of the entire scene, including gravity. As the laws of physics are invariant to this transformation, this is a valid symmetry of the problem.

The n objects can be translated and rotated. Their pose is thus given by a translation t ∈ R 3 and rotation in r ∈ SO(3) relative to a reference pose. The translation transforms by a global rotation g ∈ SO(3) as a vector via representation ρ 1 . The rotational pose transforms by left multiplication r ↦→ gr . The SO(3) pose is not a Euclidean space, but a non-trivial manifold. Even though diffusion on manifolds is possible De Bortoli et al. [2022], Huang et al. [2022], we simplify the problem by embedding the pose in a Euclidean space. This is done by picking the first two columns of the pose rotation matrix r ∈ SO(3) . These columns each transform again as a vector with representation ρ 1 . This forms an equivariant embedding ι : SO(3) ↪ → R 2 × 3 , whose image is two orthogonal 3-vectors of unit norm. Via the Gram-Schmidt procedure, we can define an equivariant map π : R 2 × 3 → SO(3) (defined almost everywhere), that is a left inverse to the embedding: π ◦ ι = id SO(3) . Combining with the translation, the roto-translational pose of each object is thus embedded as three ρ 1 vectors.

We also tested the performance of the baseline Diffuser method on this reparameterization of the state but found worse results.

Hyperparameters . We also follow the choices of Janner et al. [2022], except that we experiment with a linear noise schedule as an alternative to the cosine schedule they use. For each model and each dataset, we train the diffusion model with both noise schedules and report the better of the two results.

Peer reviews:
comments: SUMMARY:
The paper proposes a new $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$-equivariant diffusion model based on the symmetriesThe empirical results demonstrate that the proposed EDGI (Equivariant Diffusion for Generating Interactions) model exhibits enhanced efficiency and superior generalization capabilities, even when applied to unseen tasks.  Most RL methods have the issuses of sample-inefficient and lack robustness to the changes of the environment. This paper introduces spatial, temporal, and permutation symmetries into the diffusion model.

STRENGTHS:
* The idea is simple and effective. The paper introduces the equivariance into diffusion models, resulting in improvements in generalization performance for unseen tasks. 
* The paper designs a novel equivariant U-net architecture that incorporates temporal, object, geometric layers for symmeetries and internal representations for different symmetries.

WEAKNESSES:
* Equavirance has been explored in both RL and model-based approaches. What's the difference with the prior works with equivariance? Need comparison with other baselines that also use equivariance.
* It lacks analysis and ablation studies about the equivariance. For example, there are three types of equivariance. Which has the most improvement to the performance?

QUESTIONS:
* Needs more analysis and comparison with previous methods that also use equivariance at both method level and experiment level.
* Which has the most improvement to the performance among equivariance?
* Can you explain more clearly how to solve soft symmetry breaking and maybe provide some experimental results?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-17

comments: SUMMARY:
This paper proposes an enhancement to a planning/model-based RL method leveraging diffusion models. Specifically, the diffusion model is structured to be equivariant to the known symmetries of reasoning about objects in 3D space, namely translation symmetry, time shift symmetry, and permutation of object labels in the scene. The paper proposes a modeling approach which improves performance on navigation and block stacking benchmarks. The improvement is modest over the baseline Diffuser framework, but taking symmetries into account dramatically improves performance in low data regime and in regimes where the evaluation is performed in a setting that is enforced to be symmetric to the training setting.

STRENGTHS:
Very well argumented approach.
Clear writing.
Sound experimental results.

WEAKNESSES:
The paper up to section 3.1 is very repetitive and could be made more concise, leaving more space to introduce some of the modeling details that were left to the appendix.
The ROI of the approach (ratio of improvement over the additional complexity introduced by the model) is limited.

QUESTIONS:
'training and planning are currently expensive' - do you have more specific details about the overhead, e.g. compared to the Diffuser approach?

LIMITATIONS:
Limitations of the approach have been adequately discussed, modulo my question above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces the Equivariant Diffuser for Generating Interactions (EDGI), an $SE(3)\\times \\mathbb{Z} \\times S_n$-equivariant diffusion model for model-based reinforcement learning. The proposed method maintains equivariance with spatial symmetry as depicted by $SE(3)$, the discrete time translation symmetry signified by $\\mathbb{Z}$, and the object permutation symmetry symbolized by $S_n$. The paper further theoretically analysis on the conditions under which the samples from an equivariant diffusion model will be group invariant or symmetry-breaking. Finally, experimental evaluations were conducted in both manipulation and navigation tasks, demonstrating that the proposed method surpasses the performance of non-equivariant baselines.\n\nSTRENGTHS:\n1. The method considers a large variety of symmetries, all of which are common in many robotic tasks.\n2. The concept of using a sequence of three equivariant layers to handle the three distinct symmetries is novel and intriguing.\n3. The generalization experiment demonstrates convincing results.\n\nWEAKNESSES:\nThe experiment section could be more comprehensive.   \nFirst, the paper does not provide an ablation study to justify the three symmetries considered. Will removing one or two of the symmetries harms the performance? Which of the three symmetries contributes the most to the success of the architecture?   \nSecond, a data augmentation baseline can also be considered. Though it is widely demonstrated that the equivariant network architecture normally performs better than the learned equivariance through data augmentation, it is still a valuable experiment to validate the proposed network architecture. Moreover, will data augmentation + equivariant network yield even better performance?\n\nQUESTIONS:\n1. Although the proposed sequential approach for handling three different types of symmetries is conceptually sound, a theoretical understanding would be beneficial. Specifically, can the author provide a theoretical analysis that the operation on one of the three equivariant layers will not influence the other two equivariant properties?\n2. In some of the experiments (e.g., Navigation in Table 1), EDGI does not significantly outperforms the baseline Diffuser. This is different from what normally is observed in the equivariant learning literature when comparing an equivariant approach vs. a non-equivariant approach. It would be helpful if the authors could provide some analysis on this.\n3. Why is the hidden layer in the form of $\\rho_0 \\oplus \\rho_1$? Will adding higher frequency signal in the hidden layer (i.e., $\\oplus_0^k \\rho_k$ where $k>1$) improve the performance?\n4. Some figures of the experimental domains would be helpful (at least in the appendix) to better understand the environments.\n\nLIMITATIONS:\nThe paper addresses its limitation, but the discussion could be expanded. For instance, the proposed method seems highly constrained on the input data type, which could make it challenging to extend the proposed method to visual inputs.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces the Equivariant Diffuser for Generating Interactions (EDGI), a novel algorithm for model-based reinforcement learning (MBRL) and planning. It addresses the challenge of structured environments with spatial, temporal, and permutation symmetries, which are often overlooked by existing planning and MBRL algorithms. EDGI leverages the concept of equivariant diffusion to maintain symmetry under the product of SE(3), Z, and Sn symmetry groups. The algorithm achieves improved sample efficiency and generalization by incorporating a new SE(3) × Z × Sn-equivariant diffusion model that supports multiple representations.\n\nSTRENGTHS:\n1）Novel Approach: The idea of equivariant diffusion is innovative and introduces a fresh perspective on addressing symmetries in planning and MBRL.\n2） Multiple Representations: The EDGI algorithm supports multiple representations, which enhances its flexibility and applicability to a wider range of tasks.\n\nWEAKNESSES:\nLack of Clarity in Conceptual Explanation: The introduction of equivariant symmetries could have been more accessible, with clearer explanations of mathematical notations, making it easier for readers to comprehend the concept.\nInsufficient Emphasis on Sample Efficiency Improvement: The paper could provide a more explicit explanation of how equivariant symmetries contribute to improved sample efficiency. For instance, the relationship between symmetry breaking and the ability to transfer equivalent trajectories needs further clarification. Does the symmetry breaking approach allow for easy transfer of trajectories between equivalent states, as indicated by Figure 1?\nLimited Discussion on Network Design: The paper lacks a thorough discussion on how the network architecture is designed and whether it guarantees the preservation of equivariant symmetries. It would be beneficial to elaborate on the relationship between the network structure and the preservation of equivariant symmetries.\n\nQUESTIONS:\n1) Can you provide a more intuitive explanation of equivariant symmetries and their role in achieving sample efficiency and generalization?\nHow does the proposed equivariant diffusion model differ from traditional diffusion models, and how does it support multiple representations?\n2) Could you elaborate on the specific network design choices and how they ensure the preservation of equivariant symmetries?\n3) In practice, what are the computational costs associated with training and using the equivariant diffusion model, particularly for tasks with larger symmetry groups or high-dimensional state and action spaces?\n4) Are there any inherent trade-offs between achieving equivariant symmetries and other performance metrics, such as computational efficiency or convergence speed? How does EDGI address these trade-offs, if any?\n\nLIMITATIONS:\n1) Complexity of Symmetry Breaking: Although the paper mentions soft symmetry breaking through conditioning and classifier guidance, it does not delve into the challenges and limitations associated with breaking symmetries in complex environments. Further exploration of the limitations and potential difficulties in achieving effective symmetry breaking would provide a more realistic perspective.\n\n2) Scalability: The scalability of the proposed equivariant diffusion model is not thoroughly discussed. It remains unclear how the algorithm's performance scales with increasing problem complexity or the size of the symmetry group. A deeper investigation into the computational requirements and scalability of the approach would be valuable.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}