EDGI: Equivariant Diffusion for Planning with Embodied Agents
Authors: Johann Brehmer
Joey Bose
Pim De Haan
Taco Cohen
Venue: NeurIPS 2023 poster (2023)
Keywords: Planning
Diffusion models
Equivariance
Equivariant generative models
TL;DR: We embed a new SE(3) × ℤ × Sₙ-equivariant architecture in a diffusion model, solving planning problems for embodied agents.

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 A Agents
Johann Brehmer' Qualcomm Al Research
Joey Bose"
Qualcomm Al Research Mila, McGill University
jbrehmerêqti. qua comm. com
Pim de Haan Qualcomm Al Research
Taco Cohen Qualcomm AI Research
Abstract
Embodied agents operate 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 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 dataset. We introduce a new SE(3) x x S,--qqiivariant 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 needed. On object manipulation and navigation tasks, EDGI is substantially more sample efficient and generalizes better across the symmetry group than non-eot equivvaavai models.
offline trajectory
Introduction
Our world is awash with symmetries. The laws physies are the same everywhere in space and time-they are symmetric under translations and rotations of spatial coordinates 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 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 quadruped are independent of whether it is moving East North, and 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 often sample-ieefficien [Holland 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
Equual contribution, order determined through game of table tennis "Qualcomm Al Research an initiative Qualcomm Technologies, Inc.
Work done during internship Qualcomm AI Research
"This is true in the approximately flat spacetime on Earth, as long as all velocities are much smaller than the speed of light. A machine learning researcher who finds herself close to a black hole may disagree.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

to develop them with an awareness of their symmetries. Such algorithms should sat- isfy two key desiderata. First, policy and world models should be equivariant with respect to the relevant symmetry group. Of- ten, for embodied agents this will be sub- group of the product group of the spatial symmetry group SE(3), the group of dis- crete time shifts Z. and one multiple object permutation groups Sn, Second, it should be possible to softly break (parts of) the symmetry group to solve concrete tasks. For example, robotic gripper might be tasked with moving an object to a spe- eific point in space, which breaks the sym- metry group SE(3). First works on equiv- ariant 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- ters, 2022. Wang et al., 2022, Cetin al., 2022, Rezaci- Shoshtari et al., 2022, Deac al.. 2023]. However, these works generally only consider small finite symmetry groups such as Cn and do not usually allow for soft symmetry breaking test time based on a specific task.
P
P
Schematic of EDGI in navigation task, where the agent (red square) plans the next actions (red dots) to reach the goal (green star) without touching ob- stacles (grey circles). Top: planning as conditional sam- pling from diffusion model. Bottom: effect of group action. Equivariance requires the diagram to commute.
Figure I:
In this paper, we introduce the Equivariant Diffuser for Generating Interactions (EDGI), an equivari- algorithm for model-based reinforcement learning and planning. EDGI consists of base compo- nent that is equivariant with respect to the full product group SE(3) x & x S,, and supports the multi- ple different representations of this group we expect encounter in embodied environments. More- over, EDGI allows for flexible soft breaking of the symmetry at test time depending on the task.
Our work builds the Diffuser method by Janner al. [2022], who approach both the learning of dynamics model and planning within it as generative modeling problem. The key idea in Diffuser is to train diffusion model 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 new diffusion model that is equivariant with respect to the product group SE(3) x S, of spatial, temporal, and permutation symmetries and supports data consisting of multiple representations. We introduce a new way of embedding multiple input representations into single internal representation, as well as novel temporal, object, and permutation layers that act on the individual symmetries. When integrated into planning algorithm, 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 perfor- mance of the best -quivaaraai baseline when using an order of magnitude less training data. In addition, EDGI is significantly more robust to symmetry transformations of the environment, general- izing 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,. We now recall the main definition of equivariance and invariance of functions. Definition 1. A function f X + y is called G-equivariant if g f(z) f(g x) for all g G' and I E X. Here X' and y are input and output spaces that carry a G action denoted by s. The function f is called G-invariant if the group action in y is trivial,
y.
We will focus on X' _R.. y R and linear group actions representations, which are group

---

homomorphisms p: - GL((R). Examples include rotation and permutation matrices. For 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, seek to model G-invariant densities. As proven in [Köhler et al., 2020, Bose and Kobyzev, 2021, Papamakarios et al., 2021], given G-invariant prior density it is sufficient construct a G-equivariant map reach the desired G-invariant target density. In Sec. 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,,
Diffusion models. Diffusion models [Sohl-Dickstein et al., 2015] are latent variable models that generate data by iteratively inverting diffusion process. This diffusion process starts from a clean data sample x ~ ((20) and progressively injects noise for i E [T) steps until the distribution is pure noise. The reverse, generative process takes sample from 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 g as the forward diffusion dis- tribution modeled as a Gaussian centered around the sample at timestep It-1: ( = N2;; vI- VI= -1 B,I). where 8, is known variance schedule. The reverse generative process is learnable and can be parametrized using another distribution pe (Z (2) = N (z; (T t), of I). and the constraint that the terminal marginal at time 7 is standard Gaussian-i.e. p(ær) = (O T). 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 prediet the noise level directly [Ho et al., 2020]. Furthermore, to perform low-temperature sampling in diffusion models possible leverage 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-likeliin likelihoeho of the classifier ( (.) eg(2,, () λo Vz, log p(y((r)) 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 8441 = [(A,, as) ). given the state 8y and action On, taken at timestep h. The goal in trajectory optimization then to find sequence of actions a that maximizes an objective (reward) J which factorizes over per- per-timestep rewards r r(s,, 0n)). ). Formally, this corresponds
to the optimization problem 26.1 = arg maN I [s0, (n.H) = arg maXag. Σn=o r(8n; ux), where H is the planning horizon and r (#o; ao: 81, an) denotes the trajectory.
A practical method to solve this optimization problem is to unify the problem of learning 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 density (T).). Planning can then be phrased as a conditional sampling problem: finding the distribution (T) x (o(T))(T) over trajectories 7 where c(T) encodes constraints on the trajectories and specifies the task for instance as a reward function. Diffusion models allow conditioning in 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) x Z x S, and common representations in robotic problems. In Sec. 3.2 we introduce our key novelty, an SE(3) x Z x S., S,,equiivariant 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) x Z x Sn, 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,, It is important to note, however, that this symmetry group may be partially broken in an environment. For instance, the direction of gravity
3

---

input trojectory
Cutout trajectory
Diffusion time.
conditioning
Equivariant
Representation
representation:
Charnie's
Context block:
Time
Objects
Temporal conk.
over
Norm
Geometric
(Geornetry suppremse)
oijects
layer
Architecture of our SE(3) Z S, equivariant denoising network. Input trajectories (top left), which consist of features in different representations of the symmetry group, are first transformed into single internal representation (green block). The data are then processed with equivariant blocks (blue), which consist of convolutional layers along the time dimension, attention over objects, normalization layers, and geometric layers, which mix scalar and vector components of the internal representations. These blocks are combined into a U-net architecture. For simplicity, we leave out many details, including residual connections, downsampling, and upsampling layers; see Appendix B.
Figure 2:
usually breaks the spatial symmetry group SE(3) to the smaller group SE(2), and distinguishable objects in scene may break permutation invariance. We follow the philosophy of modeling invariance with respect to the larger group and including any symmetry- symmetry-breaking effects as inputs to the networks.
We require that spatial positions are always expressed relative to 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 (333eeequivariant architecture.
Data representations. We consider 3D environments that contain embodied agent well other objects. We parameterize their degrees of freedom with two SO(3) representations, namely the scalar representation Po and the vector representation P1- Other representations that may occur in a dataset can often be expressed in and P1. For example, object poses may be expressed as 3D rotations between global frame and an object frame. We can convert such rotations into two P1 representations through the inclusion map 50(3) + 22x3 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 Gran-Schmidt process to recover the final column vector. Thus, any SE(3) pose can transformed 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 Sn, object properties permute, while robot properties or global properties of the state remain invariant. Each feature is thus either in the trivial the standard representation of S,,
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,y 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 novel SF(() x Z x Squ-eqvivar diffusion model which leads an in- variant distribution over trajectories. Specifically, given an invariant base density with respect to

---

chosen symmetry group an isotropic Gaussian satisfies this property for SE(3) x X Sy Snannd an de- noising model f that equivariant with respect the same group, we arrive at diffusion model that is SF(33 x ž x -ivvarvvri [Köhler et al., 2020, Papamakarios et al., 2021]. Under mild assumptions, such an equivariant map that pushes forward the hase density always exists [Bose and Kobyzev, 2021].
We design a novel equivariant architecture for the denoising model f. Implemented as neural network, it maps noisy input trajectories T and a diffusion time step 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 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 representa- tions 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 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 E [H]. each object O E [m]. each channel c E [mc] consists of one" (3)3) scalar Stoc and one (0(3) vector Stoe- We write Wtoc (Stoez Uroc) E R... Under spatial rotations g E SO(3), these features thus transform as the direct sum of the scalar and vector representations Po @ P1:
Whoc
SStcc Utoc/
whoc =
(Po(g)sMu Stoce P1 (g)
These internal features transform in the regular representation under time shift and in the standard representation under permutations IP Wtoee - Who'e Poo Wtoce- 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) x Z Sa, into the internal representation. On the one hand, we pair up SO(3) scalars and SO(3) vectors into Po @ #1 features. On the other hand, we distribute global features- - those a priori unassigned to one of the objects in the scene over the objects.
Concretely, for each object o E [n]. each trajectory step t E [H], and each channel c [ne]. we define the input in the internal representation as Wtoc E RK as follows:
Whoc (Eess
W1 W Wegerstood oce' Utoee
W3 W Stee oce Urtte
(2)
The matrices 11,2,3,4 are learnable and of dimension n X ni, X flg object Jt X ni, X fie object
JI n. x ng global
global respectively. Here object is the number of S0(3) scalar quantities associated or n x ne x n with each object in the trajectory, n., each object, global of the system, and global is the number of vectors of that nature. The number of input channels ne is a hyperparameter. We initialize the matrices W such that Eq. (2) corresponds to a concatenation of all spjecc-specific and global features along the channel axis the beginning of training.
obbject is the number of SO(3) vector quantities associated with
n.global is the number of scalar quantities associated with the robot or global properties
Step 2: SE(3) x Z x -qqiivarrr U-net. We then process the data with a SO(3) x Z x S,-- 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-is i. e. they do not mix internal representation types of the other two layers:
Temporal layers: time-translation-et twa.nsaationnequuivriannt equivariant convolutions along the temporal direction (i. e. along trajectory steps), organized in a U-Net architecture.
Object layers: Permutation-cquivarini
self-attention layers over the object dimension.
Pairing up just one scalar and one vector is design choice; for systems in which scalar or vectorial quantities play a larger role, it may be beneficial to use multiple copies of either representation here.
5

---

* Geometric layers: -quivvaraa 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 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 1D convolutions along the trajectory time dimension. To preserve SO 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-hean head self-attention layer. Given inputs Wroes the object layer computes
Kroc Σ
w ce' Wtoc Qtoe
Σ
oat Whoc
Vioe =
Σ W Wed Wrac
(Rio Vd
Qto Kid
Wfoc x
Σ softmax,
V
with learnable weight matrices WK.VD and d the dimensionality of the key vector. There is no mixing between features associated with different time steps. nor between the four geometrie features of the internal (3(3) representation. Object layers are compute invariant SO(3) norms.
as the attention weights
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 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, Wtoe (Stoc: Utoc)* We then construct complete set of SO(3) invariants by combining the scalars and pairwise inner products between the vectors,
Sto {Stoc)e {Vroe Vtoc'' Ja,e'
These are then used as inputs to two MLPs and vb, and finally we get output scalars and vectors, w/oc (o(Stw) Ee (SSo))ee ).oer).
Assuming full expressivity of the MLPs á and v, this approach can approximate any equivariant map between (3(3) scalars and vectors [Villar et al., 2021, Proposition 4]. In this straightforward form, however, can become prohibitively expensive, as the number of SO(3) invariants Sto scales quadratically with the number of channels. In practice, first linearly map the input vectors into 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 Whoe that transformed back to data representations using linear maps, in analogy to Eq. (2). Global proper- ties, 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 geomet- ric 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 [...,.lle f(T + [ [Ho al., 2020J. where are training trajectories, i ~ Uniform(O, (0, 7) 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 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 equivariant denoising model defines an invariant density [Köhler et al., 2020, Bose and Kobyzev, 2021, Papamakarios et al., 1.. We repeat this result without proof:

---

Proposition 1. Consider a group G that acts on R" with representation Let π be G-invariant distribution over R" and e : R + E be a G-equivariant noise-conditional denoising network. Then the distribution defined by the denoising diffusion process of sampling from π and iteratively applying e is G-invariant.
We now extend this result to sampling with dassifier-bascd guidance [Dhariwal and Nichol, 2021], technique for low-temperature sampling based on a classifier p(yp(z))) with class labels y. or more generally any guide h(x). When the guide is G-iivariant, guided sampling retains G' invariance:
Proposition 2. Consider a group G that acts on T with representation Let π be G-invariant density over IR'h and 4 IR. + R be a G-equivariant noise-conditional denoising network, Let the guide h R - R be smooth G-invariant function. Further, assume p(g) is orthogonal Vg E G. Define the modified diffusion score ( I, t) e c0(x;; t) Ao; Va, h(xe) for some guidance weight À E R. Then the distribution defined by the denoising diffusion process of sampling from and iteratively applying ε G-invariant.
Proof. The function h has gradient V Vah(z) that is G-equivariant [Papamakarios et al., 2021, Lemma 2]. Thus.
(o(p())(Z))t)) co(p(g)(z)).t) Ao ((,(((((( =(g))c( - ( p(g)(ce((Eet)) Ao, v,,,,),
The modified diffusion score is therefore G-equivariant. Applying Prop. 1, we find that classifier- based guidance samples from a G-invariant distribution.
Proposition applies to G SE(3) Z Sn as employed in EDGI, each group within the product admits an orthogonal matrix representation. Both unconditional sampling from EDGI, as well as sampling guided by a SE(3) x E S, iinvariant classifier (or reward model), is thus SE(3) x Z Sn--
invariant.
Symmetry breaking. However, if the classifier h(x) or log p(ylz) is nor G-invariant, classifier- guided samples are in general also not G-invariant. For example, consider h(x) = -||x - z for some 2g E IR", which will bias samples to zy. Similarly, conditional sampling on components of æ (similar to inpainting) clearly lcads to non non-in con-invariant samples. As we will argue below, these properties are essential for robotic planning. 6
3.4 Planning with equivariant diffusion
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 regression model trained offline to map trajectories to task rewards.
Task-specific symmetry breaking. By construction, our equivariant diffusion model learns SE(3) x x S,, -nvvirriannt 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 robot or object is brought into a particular location or specifying an ordering over objects be manipulated in a scene.
flask-specific symmetry breaking would be more challenging to implement in classifier-free guidance. That would require training diffusion model that jointly models equivariant unconditional and eeuievaiviin task- conditional distributions, which will in general be difficult.
7

---

Standard setting
SO(3) generalization
Environment
BCQ
CQL
Diffuser
EDGI (ours)
Diffuser
EDGI (ours)
Navigation
94.9 43.9
95.1
5.624.4
83.3
Unconditional
0.0
24.4
59.7
68.7
38.7
69.0
Conditional
0.0
0.0
46.0
23.0
16.7
35.9
Rearrangement
0.0
0.0
49.2
49.2
23.3
47.2.3.9
17.8
45.0
Average
0.0
8.1
51.6
56.0=2.0
24.4
50.0+1.9
Table 1: Performance on navigation tasks and block stacking problems with a Kuka robot. We report normalized cumulative rewards, showing the mean and standard errors over 100 episodes. Results consistent with the best results within the errors are bold. BCQ and CQL results are taken from Janner et al. [2022]; for Diffuser, we show reproduction using their codebase. Left: Models trained on the standard datasets. Right: S0(3) generalization experiments, with training data restricted to specific spatial orientations such that the agent encounters previously unseen states at test time.
As discussed in the previous section, our diffusion- approach with classifier guidance allows us elegantly break the symmetry test time required. Such a soft symmetry breaking both occurs through conditioning on the current state, by conditioning on goal state, and through a non-invariant reward model used for guidance during sampling.
4 Experiments
We demonstrate the effectiveness of incorporating symmetries as powerful inductive bias in the Diffuser algorithm with experiments in two environments. The first environment is 3D navigation task, in which an agent needs to navigate number of obstacles to reach 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 accelcration of the agent: both the agent and the obstacles are spherical. Please see Fig. 1 for 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 benehmark 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 (navigation) or 10a (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-equivarraat) 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- 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 SF(3), but not S.; while the other is equivariant with respect to Sn, but not SF))). 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. 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
8

---

Reward (average over 3 tasks)
60
100
EDGI (ours)
Diffuser
80
Reward
60
EDGI (ours)
30
40
equivvaraat only
20
equivaraaa only
20
Diffuser
0
10
Diff. augmentation
20
10
10
10
10
10
10 -3
10
Relative dataset size
Relative dataset size
Figure 3: Average reward function of training dataset size for EDGI and Diffuser. Left: navigation environment. Right: Kuka object maniplation, averaged over the three tasks.
data. The results in Fig. 3 show that EDGI achieves just 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 other works for using symmetries in an RL context [van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021,
et al., 2022, Deac et al., 2023].
Effects of individual symmetries. the left panel of Fig. 3, we also show results for EDGI variations that are only equivariant with respect to SE(3), but not Sn. 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) x Z x S.. This confirms that the more of the symmetry of a problem we take into account designing an architecture, the bigger the benefits 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 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 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 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 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 Su-equivariant denoising network.
Equivariant deep learning. Baking in symmetries into deep learning architectures was first studied in the work of Cohen and Welling 6a] for geometric transformations, and the DeepSet archi- tecture 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, 201 6b, Weiler and Cesa, 2019, Cesa et al., 2021]. I. For dominant paradigm is to the message passing framework [Gilmer et al., 2017] along with geometric quantities like positions, velocities, and relative angles [Satorras et al., 2021, Schüit al., 2021, Batatia al., 2022].
symmetries of the 3D space-i.e i.e. subgroups of E(3)
Equivariance RL. The role of symmetries has also been explored in reinforcement learning problems with body of work focusing 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 al., 2022, Rezaei-Shoshiie 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. 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 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) X Z x S,, 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 generalizaiionnnooorruu notorious weaknesses RL algorithms.
We introduced EDGI, an equivariant planning algorithm that operates as conditional sampling in generative model. The main innovation a new diffusion model that is equivariant with respect to the symmetry group SE(3) Z S,, 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 product group and support multiple representations in the data. Integrating this equivariant diffusion model into 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 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-timg diffusion methods with accelerated sampling.
Acknowledgements
We would like to thank Gabriele Cesa, Daniel Dijkman, and Pietro Mazzaglia for helpful discussions.
10

---

References
A. Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and P. Agrawal. Is conditional generative mod- eling all you need for decision-making" International Conference on Learning Representations, 2022. (Cited on page 9)
I. Batatia, D. P. Kovács, G. N. Simm, C. Ortner, and G. Csányi. Mace: Higher order equivariant mes- sage passing neural networks for fast and accurate force fields. arXiv preprint arXiv:2206.076677 2022. (Cited page 9)
A. Bose and I. Kobyzev, Equivariant finite normalizing flows, arXiv preprint arXiv:2110.08649, 2021. (Cited pages 3, 5, 6, and 10)
D. Boyda, G. Kanwar, S. Racanière, D. Rezende, M. S. Albergo, K. Cranmer, D. C. Hackett, and E. Shanahan. Sampling using su(n) gauge equivariant flows. arXiv preprint arXiv:2008.05493 2020. (Cited page 10)
Brehmer, P. Haan, S. Behrends, and Cohen. Geometric algebra transformers. In Advances in Neural Information Processing Systems, volume 37, 2023. URL 2305. 18415. (Cited on page 10)
org. /
M. M. Bronstein, J. Bruna, T. Cohen, and P. Veličković, Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13778 2021. (Cited on page 3)
G. Cesa, L. Lang, and M. Weiler. A program to build e (n)-eqiivariant steerable cnns. In International Conference Learning Representations, 2021. (Cited page 9)
E. Cetin, Chamberlain, M. Bronstein, and Hunt. Hyperbolic deep reinforcement learning. arXiv [cs.LG), Oct. 2022. (Cited pages and 10)
T. Cohen and M. Welling. Group equivariant convolutional networks. In International conference machine learning, pages 2990-2999, PMLR, 2016 (Cited on page 9)
T.S. Cohen and M. Welling. Steerable cnns. arXiv preprint arXiv: arXiv:1612.084p00 2016b. (Cited page
S. Cohen, M. Geiger, J. Köhler, and M. Welling. Spherical cnns. arXiv preprint arXiv:1801.1011 2018. (Cited on page 9)
E. Coumans and Y. Bai. PyBullet, Python module for physics simulation for games, robotics and machine learning. http: httpppyyblllleeeeeeeeeeeeeeeeeee org, 2016-2019, (Cited on page 16)
V. De Bortoli, E. Mathieu, M. J. Hutchinson. J. Thortoo.. Y. W. Tch, and A. Doucet. Riemannian Score-Based generative modelling. Oct. 2022. URL https: btctts::/cpeereeeee http://pppneeee..ceeeeeee. wtwpeenreviee..e//aaPP2 net/pdf?id- oDRQGo8 7P. (Cited pages 10 and 17)
A. Deac, T. Weber, and G. Papamakarios. Equivariant MuZero. Feb. 2023. (Cited on pages 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 pages 3 and 7)
C. Esteves. Theoretical aspects of group equivariant neural networks. arXiv:2004.05154, 2020. (Cited on page 3)
arXiv preprint
S. 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)
S. Fujimoto, D. Meger, and D. Precup. Off Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062. PMLR, 2019. (Cited on page 8)
J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In International conference machine learning, pages 1263-1272. PMLR, 2017. (Cited page 9)
B. C. Hall. Lie groups, Lie algebras, and representations. Springer, 2013. (Cited page 14)

---

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 116840-6551, 2020. (Cited on pages and 6)
G. 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:/806.001 1806.0182, 2018. (Cited on page 1)
E. 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: //axxiivvvv. www.ars.///....oooo. attaaabss2200 2208 07949. (Cited on page 17)
I. Igashov, H. Stärk, C. Vignac, V. G. Satorras, P. Frossard, M. Welling, M. Bronstein, and B. Cor- reia. Equivariant 3d-conditional diffusion models for molecular linker design. arXiv preprint arXiv:2210.05274, 2022. (Cited page 10)
M. Janner, Y. Du, B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint 20Xi::205500991,, 2022. (Cited on pages 3,
7, 8, 9, 15, 16, and 17)
G. 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)1121601 2020. (Cited page 10)
I. Katsman, A. Lou, 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)
J. 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)
A. 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)
A. K. Mondal, V. Jain, K. Siddiqi, and S. Ravanbakhsh. EqR: Equivariant representations for Data- Efficient reinforcement learning. Nov. 2021. (Cited pages 2, 9, and 10)
D. Muglich, C. S. de Witt, van der Pol, S. Whiteson, and J. Foerster. Equivariant networks for Ecro-Shot coordination. Oct. 2022. (Cited on pages 2 and 10)
G. Papamakarios, F. Nalisnick, D. 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 pages 3, 5, 6, and 7)
A. Raffin, A. Hill, A. Gleave, A. Kanervisto, M. Ernestus, and N. Dormann. 3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268): 1-8, 2021. URL http: / jmlr..... . org http//ppprr///// 220/222220-22 1364. html. (Cited on page 16)
S. Rezaei-Shoshtari R. Zhao, P. Panangaden, D. Meger, and D. Precup. Continuous mdp homo- morphisms and homomorphic policy gradient. arXiv preprint arXiv:2209.07334 2022. (Cited on pages 2, 9, and 10)
D. J. Rezende and Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd international conference on Machine learning. ACM, 2015. (Cited on page 10)
V. 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)
K. Schütt, O. Unke, and M. Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. International Conference on Machine Learning, pages 9377- 9388. PMLR. 1.21. (Cited on page
12

---

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, pages 2256- 2265. PMLR, 2015. (Cited on page
E. van der Pol, T. Kipf, F. A. Oliehock, and others. Plannable approximations to MDP homomor- phisms: 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: Equivari- ant machine learning, structured like classical physics. Advances in Neural Information Processing Systems, 34:28848-28861 2021. (Cited on page 6)
R. Walters, J. Li, and R. Yu. Trajectory prediction using equivariant continuous convolution. Oct. 2020. (Cited on pages 2, 9, and 10)
D. Wang and Walters. So (2) equivariant reinforcement learning. In International Conference on Learning Representations, 2022. (Cited on pages and 10)
D. 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)
M. Weiler and G. Cesa. General e (2)equivvariant steerable enns. Advances in Neural Information Processing Systems, 32, 2019. (Cited on page
M. Xu, L. Yu, Y. Song. C. Shi. S. Ermon, and Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022. (Cited on page 10)
M. 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)
13

---

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.
Lie group is both a group and smooth manifold. A manifold is a topological space and thus Lie group is topological group that is formally defined below.
Definition 2. A topological group is a topological space G equipped with continuous maps that allow group composition G G G and group inverses (o)-1 + G.
Now let G be an element of the topological group and consider the map G - G given by ((g) = o g'. If g 91 = ! o 92 then 9i = 92 and thus .j is injective and for all g E G, and there exists an element- -.e. g y such that " g) g g - o g Thus, 0 is a bijection.
In topological space 9 and g are continuous. Thus, 9 is homomorphism. This means if U C G is open 1.) g U is open 2.) U o g is open and inversion is also homomorphism so U is open if and only if U 1-11111 is open. A topological space X' is said be homogenous if Va, y X there exists a map /z.. X' + X' that is a homomorphism such that fig.g (z) = Note that topological group is always homogenous. In words, a homogenous space for a topological group means any group element is reachable by suitable group homomorphism.
Matrix Lie groups. A group of special interest is the group of n invertible matrices with entries in R. .L (R). This is a topological group equipped with Euclidean topology. The subgroup SL,, C GL,, (TR) is closed subgroup with the property that det 1. As example SF((3 c 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, (R.)7 This means SL, (R) is a matrix Lie group as (nn) and (O(n). Moreover, (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) x R and thus can be written as,
SE(3) - {
R SO(3), E (R. +)
(5)
Represented by this 4 x 4 matrix and with the group operation defined by matrix multiplication, this group can be seen as subgroup of the general linear group (LR)
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 E SO(3) is a rotation. To represent R there are various possible choices. The most familiar of them is perhaps R 3 is 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 vector w E T Given any E TR3, it's related by gv w x v, with the symbol x representing the cross product. The length of this vector, symbolized as = ||O|| represents the rotation angle, while the unit direction, expressed as ew = | indicates the rotation axis.
Yet another parametrization of rotations can achieved using quaternions. A quaternion, q, is often represented as y mutually orthogonal imaginary units vectors. The unit vectors obey the following rule: 12 = =² 2² = ijk - 1. A subset of quaternions, known rotation quaternions, can be used to represent 3D rotations. These specific quaternions have magnitude 1 and can also be thought of as an ordered
+ ai + yj + zk where w, y. and are real numbers and i, and k are
The field can also be complex C to allow invertible complex matrices GLn (C).
14

---

pair (w, v)) where w is scalar and v = (z, 2) is a 3D vector. The scalar part w can be thought of as c and the vector part as sin(w//2), where the angle of rotation and the unit vector along the axis of rotation.
Finally, rotate point p in 3D space using a quaternion q. the point can be represented as a pure quaternion p' 0 + vi + yj + zk. The rotated point p" is then given by: p" q.p' - Where g" the conjugate of q-i.e. if q + xi yj + zk, then q" xi yd zk.
p''"."
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) x Z x S,pequivarriant architecture for the diffusion model.
Overall architecture. We illustrate the architecture in After converting the input data 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) upsampling (in the upward pass).
Residual standard block. The main processing unit our architecture processes the current hidden state with an equivariant block consisting of temporal layer, an object layer, 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, context block, and another equivariant block is residual (the inputs are added to the outputs).
Temporal layers. Temporal layers consist of one dieddmmeniooall 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 Whoe unit variance. This is essentially 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 (3(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 linear map between the vector inputs and the vector outputs, respectively. Finally, there is 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 linear layer, like in Janner 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 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: E #xxx be data in our internal representation, such that the entries whoe decompose into SO(3) scalars Stoc and SO(3) vectors Vtoc- Let S S(wa.) be the set of scalars and all pairwise SO(3) inner products between the vectors Uto. as defined in Eq. (4). The outputs of the geometric layer
15

---

Figure 4: Rendering of the robotic manipulation environment. We show four frames from a single trajectory solving the unconditional block stacking task.
then f(w) (o(S(wto))) Ee (((
First, consider what happens under permutations of the objects, Whoe Wir(o)e for a permutation Sg. We have f(T-w w) toc = ((((w)x(o)) En. ((S(wex())) (w)(())))))s)) Uoz(()ee) f(w) = (T- (TTffw))tto Thus, because this layer "leaves the object dimension untouched", it equivariant with respect to object permutations.
Next, consider the behavior under spatial transformations. Like most (3)E(3) -qquivvraaa architectures, 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 E SO(3), such that w = R. (s. v) (s. R RR.00 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) toe = (R wra)), (S((ww))) 'ioo') (#(S(w:o)) Le ((S(we))))tee) = (R- ((w) Hence the geometric layer is equivariant with respect to SEE().
C Navigation experiments
We introduce new navigation environment. The scene consists of a spherical agent navigating plane populated with goal state and 10 spherical obstacles. At the beginning of every episode, agent position, agent velocity, obstacle positions, and goal position are initialized randomly (in 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 al. [2021] for 10" steps with default hyperparameters on this environment. We generate 10s 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 the goal state. addition, a penalty of -0.1 added the reward the agent touches any of the obstacles. Finally, there is an additional control cost equal to -10º times the force acting on the agent. We affinely normalize the rewards such that normalized reward of corresponds that achieved by random policy and normalized reward of 100 corresponds to the expert policy.
D Kuka experiments
We use the object manipulation environments and tasks from Janner 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 visualizes the unconditional block stacking task. For fair comparison, we re- 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, the original 39-dimensional parameterization from Janner al. [2022].
For our EDGI, we need to parameterize the system in terms of SE(3) S, representations. We, therefore, describe the robot and block orientations with SO(3) vectors as follows. Originally, the
16

---

robot state is specified through collection of joint angles. One of these encodes the rotation of the base along the vertical 2-axis, We choose to represent this angle pi vector in the zy-plane. In addition, we add the gravity direction (the z-axis itself) as another on vector, which is also the normal direction of the table on which the objects rest. Combined, these vectors define the pose of the hase of the robot arm. Rotating gravity direction, and the robot and object pose by SO(3) can be interpreted as passive coordinate transformation, or as an active rotation the entire scene, including gravity. As the laws of physics are invariant to this transformation, this is valid symmetry of the problem.
The n objects can be translated and rotated. Their pose thus given by a translation t e R and rotation in (3(3) relative to reference pose. The translation transforms by global rotation g E (3) as vector via representation P1. The rotational pose transforms by left multiplication gr. The (3)3) pose is not a Euclidean space, but non-trivial manifold. Even though diffusion on manifolds possible De Bortoli et al. [2022], Huang et [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 E SO(3). These columns each transform again as a vector with representation fi- This forms equivariant embedding S3)3) + 22xx whose image is two orthogonal 3-vectors of unit norm. Via the Trram-Schmide procedure, we can define an equivariant map π 2x2 + SO(3) (defined almost everywhere), that is left inverse to the embedding:
idso(3)- Combining with the translation, the noto-translations pose of each object is thus embedded as three P1 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 linear noise schedule 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.
17


Reviews:
[{"review_id": "5ElK8CFqUd", "reviewer": "Reviewer_FPx8", "summary": "The 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.", "questions": "1) 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?", "limitations": "1) 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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1）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.", "weaknesses": "Lack 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and nuanced comments. We are glad that the reviewer found our contribution to be novel and that it provides “a fresh perspective on symmetries in planning and MBRL”. We also appreciate the fact that the reviewer valued the “flexibility” provided by EDGI through multiple representations allowing it to extend to “a wider range of tasks.” We now address their key questions and concerns.\n\n**Improving accessibility on equivariance**\n\nWe thank the reviewer for their constructive criticism regarding our exposition of symmetries and equivariance. We agree with the reviewer that our coverage of the background material may not have been sufficiently clear in the initial manuscript. Toward this end, we will update the paper to include a more comprehensive coverage of the exact groups we consider in this work.\n\nIn particular, we will add more detail on the groups SO(3) and $S_n$ along with their representations and actions on our state and action spaces. We will further supplant this with a short primer on group and representation theory in the appendix along with detailed pointers to complete works for the interested reader in the main text. \n\n**Sample efficiency**\n\nWe thank the reviewer for highlighting an important aspect, ultimately a main selling point for equivariant architectures: equivariance improves the sample efficiency of training. Essentially, the equivariance constraint ensures that when the network encounters a single sample in training, it learns not only how to transform that sample, but also all data points that can be generated from the sample through symmetry transformations like rotations or permutations. Hence, less samples are necessary to reach a strong, robust performance.\n\nThis benefit of equivariance for sample efficiency has been demonstrated both theoretically and empirically in fields from molecular dynamics to robotics, see for instance [1] [2] [3] [4]. Sample efficiency is particularly important in the robotics case, as collecting training data can be expensive.\n\nThe reviewer also asks about the role of symmetry breaking in the transfer of trajectories between equivalent states. This is a subtle point (and we will improve its discussion in the final version of the paper). Symmetry breaking consists of two aspects: conditioning on the initial state, as well as on guidance from a reward model. The conditioning on the initial state, together with the equivariance of the denoising model, is what ensures that we can easily transfer learned behaviors to transformed conditions.\n\nThe guidance from the reward model, on the other hand, can be used to make the model behave *differently* on rotated situations, if this is desired. For instance, if a task consists of moving an object in a particular direction, then after the initial state is rotated, the agent’s behavior needs to adapt by more than just rotating. Reward guidance allows us to achieve that.\n\n**Network design choices and equivariance**\n\nWe value the reviewer's feedback and agree that the paper could be improved by a richer discussion on our network design. How our architecture ensures equivariance is indeed crucial. We address it in the global  response. In a nutshell, we can explicitly show that each layer in EDGI is equivariant to $\\mathrm{SE(3)} \\times S_n \\times \\mathbb{Z}$.\n\n**Computational complexity and scaling** \n\nWe agree that the computational costs of EDGI is an important aspects that deserve further discussion. We address it in the global response. Essentially, in its current implementation EDGI has substantial overhead over the baseline Diffuser, but we see plenty of room for optimization. What’s more, EDGI converges substantially faster, offsetting the computational overhead. Please see our global response for a more thorough discussion.\n\nAs for scalability, EDGI has favorable scaling properties in the three relevant directions – number of objects $n$, number of time steps $H$, and number of channels $c$. Thanks to its modular structure with alternating layers attending to the different data axes, it has a time complexity of $\\mathcal{O}(n^2 H c^2)$. When an efficient attention implementation is used, the memory complexity is even just linear in $n$. We will add a thorough discussion of this scaling behaviour to the final version of our paper.\n\n**Symmetry breaking**\n\nWe thank the reviewer for pointing out that our explanations of symmetry breaking and its complexities were not sufficiently clear. We will expand the description in the paper, but we believe that the essential idea follows quite naturally from the framework of treating RL as a diffusion problem. The overall behavior of the agent consists of three components – a trained, equivariant diffusion model; a task-specific reward guide; and an initial or current state. The first component is by constructions equivariant, while the latter two allow us to softly break the symmetry *if this is desired*, for instance because of a non-invariant task specification. All this requires is a non-equivariant reward model, for instance a simple MLP.\n\nWe would like to thank the reviewer again for their thorough review. We hope we have sufficiently addressed their comments, and we respectfully ask the reviewer to reconsider their impression of the paper and potentially improve the given score. We look forward to discussing any further questions. \n\n\n**References**\n\n[1] Bietti, Alberto, Luca Venturi, and Joan Bruna. \"On the sample complexity of learning under geometric stability.\" NeurIPS 2021. \n\n[2] Behboodi, Arash, Gabriele Cesa, and Taco S. Cohen. \"A pac-bayesian generalization bound for equivariant networks.\" NeurIPS 2022.\n\n[3] Jumper, John, et al. \"Highly accurate protein structure prediction with AlphaFold.\" Nature 2021.\n\n[4] Wang, Dian, et al. \"On-robot learning with equivariant models.\" arXiv:2203.04923."}, {"review_id": "VXkJKmXgRG", "reviewer": "Reviewer_LvgJ", "summary": "The 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.", "questions": "1. 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", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "The 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review and constructive criticism. We are glad that the reviewer thought the way that we incorporate symmetries in EDGI to be “novel and intriguing” and that the reviewer found our generalization experiments to demonstrate “convincing results”. We now address the main questions raised by the reviewer.\n\n\n**Will removing one or two of the symmetries harm the performance?**\n\nGood question. We answer it in detail in our global response to all reviewers. In summary, we implemented a new ablation model  that is equivariant with respect to $S_n$, but not SE(3), and a second ablation model that is equivariant with respect to SE(3), but not to $S_n$. As we show in Fig. 3 of the rebuttal response page, these methods perform much better than the Diffuser baseline, but are not quite as sample-efficient as EDGI with the full symmetry group. This provides more evidence that it is worthwhile to be equivariant with respect to each symmetry found in our robotics environments.\n\n\n**Data augmentation baseline**\n\nThanks for the suggestion. We trained an additional Diffuser model on the navigation environment, augmenting the data with random SO(3) rotations for each sample. Fig. 3 in the rebuttal result page shows the results. We find that data augmentation substantially improves the Diffuser model in terms of sample efficiency and generalization across the symmetry group, though EDGI still maintains a performance benefit in terms of sample efficiency.\n\n\n**Data augmentation for equivariant networks**\n\nThis is an interesting suggestion. However, as long as we only consider data augmentation through symmetry transformations (like rotating the scene or permuting the objects), data augmentations cannot help equivariant methods. Informally, the reason for this is that original and transformed data lie in the same orbit of the symmetry group. This means that both the original and transformed data are acted on by the same network weights in the same way, and lead to the same loss and network updates. Thus we expect no benefit from data augmentation for equivariant networks.\n\n\n**Theoretical analysis of equivariance**\n\nIn our global response, we demonstrate explicitly that our novel network layers leave the dimensions of the data it does not act on unaffected, and that it is equivariant with respect to the whole symmetry group.\n\n\n**Performance sometimes on par with baseline Diffuser**\n\nWe agree with the reviewer that in our navigation environment when training on a large dataset the performance of EDGI is roughly on par with the baseline. However, we still find two key advantages for our equivariant method: it is much more sample efficient than Diffuser, as shown in Fig. 3 of our paper, and it generalizes more robustly under the symmetry group, as shown in the right column of Table 1 of our paper.\n\n\n**Choice of representations**\n\nWe thank the reviewer for this great question. We chose to use  a direct sum of the $\\rho_0$ and $\\rho_1$ representations for three reasons: it is simple, computations are relatively cheap, and most real-world geometric quantities can be expressed in these representations as they are scalars and vectors. \n\nHowever, EDGI *can* be extended to include higher representations. Doing so would certainly increase the computational cost of the geometric layers. It is an intriguing question whether that will be offset by benefits in expressivity.  In some domains like molecular dynamics, including higher representations has been very advantageous, see e.g. [1] and the architectures referenced in there. At the same time, there are some results that show that latent scalar and vector representations are enough to represent any equivariant map between vectors [2]. Our EDGI architecture is built on the theoretical foundations of [2] and we leave investigations of higher order representations in the MBRL domain as natural directions for future work.\n\n\n**Illustrations of environments**\n\nThank you for the suggestion. In Fig. 1 of the rebuttal response page, we show a rendering of the robotic manipulation environment. For the navigation environment, please have a look at Fig. 1 of the paper. We will update our appendix to include these figures.\n\n\n**Constraints on the input data type**\n\nIndeed, we currently assume access to a representation of the data in terms of group representations. We will stress this in the final version of the paper. We consider learning from raw visual inputs and using symmetries in this setting as a natural direction for future work.\n\n\nWe thank the reviewer again for their time and efforts reviewing our work. We hope that our rebuttal was successful in addressing all the great points raised by the reviewer and allows the reviewer to consider a fresher evaluation given our rebuttal as context. Finally, please also have a look at our global response. In addition to the important baselines and ablations you suggested, we also show there that EDGI converges faster than the Diffuser baseline. \n\n**References**\n\n[1] Joshi, Chaitanya K., et al. \"On the expressive power of geometric graph neural networks.\" arXiv:2301.09308.\n\n[2] Villar, Soledad, et al. \"Scalars are universal: Equivariant machine learning, structured like classical physics.\" NeurIPS 2021.\n"}, {"review_id": "hRPNBUgFqK", "reviewer": "Reviewer_28Jy", "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.", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Very well argumented approach.\nClear writing.\nSound 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.\nThe ROI of the approach (ratio of improvement over the additional complexity introduced by the model) is limited.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive appraisal of our work! We are delighted to hear that the reviewer finds our approach to have solid motivations and arguments. We are also heartened to hear that the reviewer views our manuscript to be “clearly written” and to contain “sound experimental results”. We now address the main clarification points.\n\n\n**Paper is repetitive up to section 3.1**\n\nWe acknowledge the reviewer's concern regarding the lack of concision on our ideas up to section 3.1. We believe it is useful to reiterate the key design objectives in our architecture, especially considering that our work is at the intersection of MBRL, diffusion generative models, and equivariant machine learning. By emphasizing the role of equivariance, we aimed to motivate the key modeling decisions in EDGI. For instance, SE(3) equivariance directly leads to our geometric layer, while permutation symmetry leads to the permutation layer.  Nevertheless, we agree with the reviewer that these points could be more concise. We will update the manuscript to give much more detail on our specific symmetric groups, their corresponding group actions, and representations as well as a primer on group and representation theory in the appendix. We will also add more details on our architecture and explicitly show its equivariance properties, following the logic we sketch in the global  response.\n\n\n**Limited ROI**\n\nWe respect the reviewer's apprehension regarding the ROI with respect to the complexity added by EDGI. However, we would like to politely push back as the inclusion of symmetries into the Diffuser architecture is not an optional design decision but rather a key aspect of modeling since the state and action spaces of our environments are *guaranteed* to contain this rich geometric structure. Furthermore, EDGI follows a simple design philosophy and is plug and play in the sense that each layer operates on one specific symmetry group. Please see our global response for a detailed description on EDGI’s equivariant network design including new ablations that highlight the increased sample efficiency of being equivariant to the full product group.\n\n\n**Training and planning complexity**\n\nWe appreciate the reviewer's comment regarding the training and planning complexity of EDGI. We provide extensive details on this aspect in our global response as well as the 1-page document which contains training loss curves and inference costs for EDGI as well as the original Diffuser. In summary, we find that EDGI converged about 4x faster during training measured by the number of iterations until we observe a plateau in the loss. Moreover, as we demonstrate in Fig 3. our experiments in section EDGI also enjoys significant gains in sample efficiency over the vanilla Diffuser architecture. However, EDGI does incur a large overhead per iteration. We believe much of this additional overhead can be overcome with implementation optimizations in line with recent efforts [1][2], which we leave as exploration for future work.\n\nWe thank the reviewer again for their valuable feedback and great questions. We hope that our rebuttal addresses their questions and concerns and we kindly ask the reviewer to consider upgrading their score if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise.\n\n\n**References**\n\n[1] Liao, Yi-Lun, and Tess Smidt. \"Equiformer: Equivariant graph attention transformer for 3d atomistic graphs.\" arXiv:2206.11990.\n\n[2] Esteves, Carlos, Jean-Jacques Slotine, and Ameesh Makadia. \"Scaling Spherical CNNs.\" arXiv:2306.05420. "}, {"review_id": "FI0DxIUIUI", "reviewer": "Reviewer_kPhe", "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.", "questions": "* Needs more analysis and comparison with previous methods that also use equivariance at both method level and experiment level.\n* Which has the most improvement to the performance among equivariance?\n* Can you explain more clearly how to solve soft symmetry breaking and maybe provide some experimental results?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The idea is simple and effective. The paper introduces the equivariance into diffusion models, resulting in improvements in generalization performance for unseen tasks. \n* 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.\n* 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very appreciative of the reviewer’s time and effort in reviewing our manuscript. We are grateful to hear that the reviewer finds our idea to be “simple and effective” and that the reviewer thinks the addition of equivariance is indeed reflected in “improved generalization” on unseen tasks. We also pleased to hear that the reviewer found our EDGI architecture to be “novel” in the manner in which we incorporate temporal, object, and SE(3) symmetries.\nWe thank the reviewer again for their constructive criticism and now respond to the main clarification points below.\n\n\n**Comparison to other equivariant methods in RL**\n\nWe agree that more baselines would be beneficial. We would be grateful for any concrete suggestions of appropriate baselines to compare to.\n\nWe are not aware of any RL baselines that are equivariant with respect to similar product groups of all relevant symmetries. To the best of our knowledge, our work is the first equivariant RL method with this property, and as such we believe there is no direct comparison. Similarly, we are not aware of any equivariant methods that solve (MB)RL through generative modeling. \n\nWe will however extend our discussion of equivariant RL methods. These involve model-free methods that train SE(3)-equivariant DQN or SAC versions. In contrast, EDGI is a model-based approach that leverages recent advances in conditional generative modeling and is guaranteed to be equivariant to a full product group.\n\nOther recent work has focused on finite subgroups [3] which is an easier modeling problem. We further argue none of the prior work considers equivariance in an offline RL/ conditional generative modeling setting for robotics.\n\nAs the reviewer notes, there have been recent developments in equivariant MBRL [4][5]. In [4] the authors consider symmetry with respect to $C_4$, the group of 90-degree rotations, though their evaluation is limited to toy 2D environments. [5] on the other hand uses SE(2)-equivariant steerable convolutions, suitable for image data, but not 3D environments like we consider.\n\nFor now, we added one additional baseline: the Diffuser architecture trained with data augmentation. We discuss this experiment and show results in the overall response.\n\n\n**Equivariance ablation studies**\n\nWe thank the reviewer for this suggestion. We designed ablation models that are equivariant with respect to only parts of the product symmetry group, and experimented with them on the navigation task.\n\nOur results are shown in Fig. 3 of the rebuttal result page and discussed in the overall response. Essentially, we find that a model that is equivariant with respect to SE(3), but not $S_n$, and a model that is equivariant with respect to $S_n$, but not SE(3), both perform stronger than the Diffuser baseline, but do not quite reach the sample efficiency of EDGI. This provides evidence that both symmetries are important.\n\n\n**Soft symmetry breaking**\n\nWe thank the reviewer for pointing out that our explanations of symmetry breaking were not sufficiently clear. We will expand the description in the paper, but the essential idea is the following: Because EDGI builds on the Diffuser paradigm and treats (MB)RL as a diffusion problem, the final behaviour consists of three components – a trained, equivariant diffusion model; a task-specific reward guide; and an initial or current state. The first component is by constructions equivariant, while the latter two allow us to softly break the symmetry *if this is desired*, for instance because of a non-invariant task specification. In the EDGI framework, this form of soft symmetry breaking is very natural.\n\nWe already demonstrate this feature in our experiments. In the robotic manipulation environment, we consider both unconditional block stacking, which does not use a reward guide and maintains permutation equivariance, as well as two conditional stacking tasks, in which equivariance is broken by the task specification (e.g. “stack the blue block on top of the red block”). Specifically, each of the conditional and rearrangement tasks in Table 1 requires the Kuka arm to interact with blocks in the environment by stacking them in a particular order (see Fig 1 in our 1-page PDF for a visual of the environment) which requires breaking permutation symmetry due to the specific order of blocks in stack needed to get reward. Please see Fig. 1 in our 1-page PDF for a visual illustration of the task. EDGI achieves good rewards on all tasks, showing the strength of the combination of an equivariant base model and a non-equivariant task-specific reward guide. \n\n\nWe thank the reviewer for their valuable feedback and great questions. We hope that our rebuttal fully addresses all the important salient points raised by the reviewer and we kindly ask the reviewer to potentially upgrade their score if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise. \n\n\n**References**\n\n[1] ​​Wang, Dian, Robin Walters, and Robert Platt. \"$\\mathrm {SO}(2)$-Equivariant Reinforcement Learning.\" arXiv:2203.04439.\n\n[2] Mondal, Arnab Kumar, et al. \"Eqr: Equivariant representations for data-efficient reinforcement learning.\" ICML 2022.\n\n[3] Zhu, Xupeng, et al. \"On Robot Grasp Learning Using Equivariant Models.\" arXiv:2306.06489.\n\n[4] Deac, Andreea, Théophane Weber, and George Papamakarios. \"Equivariant MuZero.\" arXiv:2302.04798.\n\n[5] Zhao, Linfeng, et al. \"Integrating Symmetry into Differentiable Planning with Steerable Convolutions.\" ICLR 2022."}]

Author rebuttal:
We would like to thank all reviewers for their thorough reviews and valuable feedback. We are encouraged that they found our approach of equivariant diffusion for planning “innovative” (reviewer **FPx8**), “simple and effective” (**kPhe**), and “very well argumented” (**28Jy**). In particular, they appreciated that our method supported products of multiple symmetry groups (**LvgJ**) and multiple representations (**FPx8**). We are also happy to hear that they found the experimental results “convincing” (**LvgJ**) and the writing “clear” (**28Jy**).


**Showing equivariance (FPx8, LvgJ)**

We now show more thoroughly that our architecture is equivariant.This is best demonstrated explicitly, by proving for each network layer $f(w)$ and each symmetry group $G$ that $f(g \cdot w) = g \cdot f(w)$ for any data $w$ and group elements $g \in G$. Here $\cdot$ denotes the group action.

In the following we will do this for our geometric layers, as they are the most novel, for both SE(3) transformations and permutations. 

Let $w \in \mathbb{R}^{n \times H \times c \times 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 discussed in line 221 in the main paper.

The outputs of the geometric layer are then $f(w)\_{toc} = (\phi( S(w\_{to}) ), \sum_{c’} \psi( S(w\_{to}) ) v\_{toc’} )$.

First, consider what happens under permutations of the objects, $w_{toc} \to w_{t\pi(o)c}$ for a permutation $\pi \in S_n$. We have $f(\pi \cdot w)\_{toc} = (\phi( S(w\_{t\pi(o)}) ), \sum\_{c’} \psi( S(w\_{t\pi(o)}) ) v\_{t\pi(o)c’} ) = f(w)\_{t \pi(o) c} = (\pi \cdot f(w))\_{toc}$. Thus, because this layer “leaves the object dimension untouched”, it is equivariant with respect to object permutations.

An analogous argument can be made for temporal translations.

Finally, 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 \in \mathrm{SO(3)}$, such that $R \cdot w = R \cdot (s, v) = (s, R \cdot v)$. By definition, orthogonal matrices leave the inner product invariant, thus $S(R \cdot w) = S(w)$. The geometric layer applied to rotated inputs then gives $f(R \cdot w)\_{toc} = (\phi( S(R \cdot w\_{to}) ), \sum\_{c’} \psi( S(w\_{to}) ) R \cdot v\_{toc’} ) = (\phi( S(w\_{to}) ), R \cdot \sum\_{c’} \psi( S(w\_{to}) ) v_{toc’} ) = (R \cdot f(w))\_{toc}$. Hence the geometric layer is equivariant with respect to SE(3).


**Computational cost (FPx8, 28Jy)**

There are two aspects to the computational cost of EDGI.
First, EDGI has a more complex computational graph than the baseline Diffuser. In the current implementation, each step (forward plus backward pass) takes roughly 5.5 times as long as that of a baseline model.

We believe that this is not a fundamental property of our method, but just reflects the lack of optimization. We see lots of potential for further speed-ups by optimizing the memory layout of our representations,the contraction paths in tensor multiplications, using optimized attention implementions, or compiling the computational graph. We plan to investigate them in the future.

Second, thanks to its stronger inductive biases, EDGI requires substantially fewer optimizer steps to converge, roughly by a factor of 4. We illustrate this in Fig. 1 of the attached PDF. This at least partially makes up for the computational overhead.


**Additional baselines (kPhe, LvgJ)**

As suggested by reviewer **LvgJ**, we ran the navigation experiment with a version of the Diffuser model trained with data augmentation. We focused on the spatial symmetry group and transformed each sample with SO(3) rotations sampled uniformly from the Haar measure.

The results are shown in Fig. 3 on the attached PDF. We find that data augmentation substantially improves the sample efficiency of the Diffuser model. However, EDGI still maintains a performance benefit in the low-data regimes..


**Ablating the effect of different symmetry groups (LvgJ, kPhe)**

As suggested by the reviewers we investigated the importance of each of the three symmetry group towards final performance. We designed two ablation models: one equivariant with respect to SE(3), but not $S_n$; the other equivariant with respect to $S_n$, but not SE(3). Both models are equivariant to temporal translations, just like EDGI and the baseline Diffuser.

In Fig. 3 on the attached PDF we show how these models perform on the navigation task. Both of these partially equivariant models outperform the baseline Diffuser in terms of sample efficiency, and come closer to EDGI. To our surprise, the permutation-equivariant architecture performed slightly better, indicating large benefits of permutation equivariance. However, using EDGI’s full symmetry group is still the most sample efficient in low-data settings. 

We hope that with this global response and the individual responses to the reviewers, we were able to address all questions. We look forward to further discussion. 

Meta-review: This paper introduces the equivariant diffuser, a model respecting spatial symmetries to generated interactions in model based LR. It initially received 4 favorable reviews, ratings increased to (7,6,5,6) during the discussion phase. Reviewers generally appreciated the novelty and different perspective this paper brings to the community.

Some weaknesses were raised, linked to clarity / accessibility, limitations of the experiments. The reviewers provided a rebuttal, which answered most of the concerns.

The AC judges that the weaknesses are by far outweighted by the paper strengths, novelty, and recommends acceptance.
Decision: Accept (poster)