SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Geometric layers. Geometric layers enable mixing between the scalar and vector quantities that are combined in the internal representation, but do not mix between different objects or across the time dimension. We construct an expressive equivariant map between scalar and vector inputs and outputs following Villar et al. [2021]: We first separate the inputs into  $\mathrm{SO}(3)$  scalar and vector components,  $w_{toc} = (s_{toc},v_{toc})^T$ . We then construct a complete set of  $\mathrm{SO}(3)$  invariants by combining the scalars and pairwise inner products between the vectors,
$$
S _ {t o} = \left\{s _ {t o c} \right\} _ {c} \cup \left\{v _ {t o c} \cdot v _ {t o c ^ {\prime}} \right\} _ {c, c ^ {\prime}}. \tag {4}
$$
These are then used as inputs to two MLPs  $\phi$  and  $\psi$ , and finally we get output scalars and vectors,  $w_{tot}^{\prime} = (\phi(S_{to})_{c}, \sum_{c'} \psi(S_{to})_{cc'} v_{tot'})$ .
Assuming full expressivity of the MLPs  $\phi$  and  $\psi$ , this approach can approximate any equivariant map between SO(3) scalars and vectors [Villar et al., 2021, Proposition 4]. In this straightforward form, however, it can become prohibitively expensive, as the number of SO(3) invariants  $S_{to}$  scales quadratically with the number of channels. In practice, we first linearly map the input vectors into a smaller number of vectors, apply this transformation, and increase the number of channels again with another linear map.
Step 3: Representation unmixer. The equivariant network outputs internal representations  $w_{tot}$  that are transformed back to data representations using linear maps, in analogy to Eq. (2). Global properties, e.g. robotic degrees of freedom, are aggregated from the object-specific internal representations by taking the elementwise mean across the objects. We find it beneficial to apply an additional geometric layer to these aggregated global features before separating them into the original representations.
Training. We train EDGI on offline trajectories without any reward information. We optimize for the simplified variational lower bound  $\mathcal{L} = \mathbb{E}_{\tau, i, \epsilon}[\| \epsilon - f(\tau + \epsilon; i) \|^2]$  [Ho et al., 2020]. where  $\tau$  are training trajectories,  $i \sim \mathrm{Uniform}(0, T)$  is the diffusion time step, and  $\epsilon$  is Gaussian noise with variance depending on a prescribed noise schedule.
3.3 Invariant sampling and symmetry breaking
We now discuss sampling from EDGI (and, more generally, from equivariant diffusion models). While unconditional samples follow an invariant density, conditional samples may either be invariant or break the symmetry of the diffusion model.
Invariant sampling. It is well-known that unconditional sampling from an equivariant denoising model defines an invariant density [Köhler et al., 2020, Bose and Kobyzev, 2021, Papamakarios et al., 2021]. We repeat this result without proof:
Proposition 1. Consider a group  $G$  that acts on  $\mathbb{R}^n$  with representation  $\rho$ . Let  $\pi$  be a  $G$ -invariant distribution over  $\mathbb{R}^n$  and  $\epsilon_t: \mathbb{R}^n \to \mathbb{R}^n$  be a  $G$ -equivariant noise-conditional denoising network. Then the distribution defined by the denoising diffusion process of sampling from  $\pi$  and iteratively applying  $\epsilon_t$  is  $G$ -invariant.
We now extend this result to sampling with classifier-based guidance [Dhariwal and Nichol, 2021], a technique for low-temperature sampling based on a classifier  $\log p(y|x_{t})$  with class labels  $y$ , or more generally any guide  $h(x)$ . When the guide is  $G$ -invariant, guided sampling retains  $G$  invariance:
Proposition 2. Consider a group  $G$  that acts on  $\mathbb{R}^n$  with representation  $\rho$ . Let  $\pi$  be a  $G$ -invariant density over  $\mathbb{R}^n$  and  $\epsilon_t: \mathbb{R}^n \to \mathbb{R}^n$  be a  $G$ -equivariant noise-conditional denoising network. Let the guide  $h: \mathbb{R}^n \to \mathbb{R}$  be a smooth  $G$ -invariant function. Further, assume  $\rho(g)$  is orthogonal  $\forall g \in G$ . Define the modified diffusion score  $\bar{\epsilon}_{\theta}(x_t, t, c) = \epsilon_{\theta}(x_t, t) - \lambda \sigma_t \nabla_{x_t} h(x_t)$  for some guidance weight  $\lambda \in \mathbb{R}$ . Then the distribution defined by the denoising diffusion process of sampling from  $\pi$  and iteratively applying  $\bar{\epsilon}$  is  $G$ -invariant.
Proof. The function  $h$  has a gradient  $\nabla_{x}h(x)$  that is  $G$ -equivariant [Papamakarios et al., 2021, Lemma 2]. Thus,
$$
\begin{array}{l} \bar {\epsilon} _ {\theta} (\rho (g) \left(x _ {t}\right), t) = \epsilon_ {\theta} (\rho (g) \left(x _ {t}\right), t) - \lambda \sigma_ {t} \nabla_ {x _ {t}} h (\rho (g) \left(x _ {t}\right)) \\ = \rho (g) \epsilon_ {\theta} \left(x _ {t}, t\right) - \rho (g) \lambda \sigma_ {t} \nabla_ {x _ {t}} h \left(x _ {t}\right)) \\ = \rho (g) \left(\epsilon_ {\theta} \left(x _ {t}, t\right) - \lambda \sigma_ {t} \nabla_ {x _ {t}} h \left(x _ {t}\right)\right). \\ \end{array}
$$
The modified diffusion score  $\bar{\epsilon}$  is therefore  $G$ -equivariant. Applying Prop. 1, we find that classifier-based guidance samples from a  $G$ -invariant distribution.
Proposition 2 applies to  $G = \mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$  as employed in EDGI, as each group within the product admits an orthogonal matrix representation. Both unconditional sampling from EDGI, as well as sampling guided by a  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ -invariant classifier (or reward model), is thus  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ -invariant.
Symmetry breaking. However, if the classifier  $h(x)$  or  $\log p(y|x)$  is not  $G$ -invariant, classifier-guided samples are in general also not  $G$ -invariant. For example, consider  $h(x) = -\| x - x_0\|^2$  for some  $x_0 \in \mathbb{R}^n$ , which will bias samples to  $x_0$ . Similarly, conditional sampling on components of  $x$  (similar to inpainting) clearly leads to non-invariant samples. As we will argue below, these properties are essential for robotic planning. $^6$
3.4 Planning with equivariant diffusion
A diffusion model trained on offline trajectory data jointly learns a world model and a policy. Following Janner et al. [2022], we use it to solve planning problems by choosing a sequence of actions to maximize the expected task rewards.
To do this, we use three features of diffusion models. The first is the ability to sample from them by drawing noisy trajectory data from the base distribution and iteratively denoising them with the learned network yielding trajectories similar to those in the training set. For such sampled trajectories to be useful for planning, they need to begin in the current state of the environment. We achieve this by conditioning the sampling process such that the initial state of the generated trajectories matches the current state, in analogy to inpainting. Finally, we can guide this sampling procedure toward solving concrete tasks specified at test time using classifier-based guidance where a regression model is trained offline to map trajectories to task rewards.
Task-specific symmetry breaking. By construction, our equivariant diffusion model learns a  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ -invariant density over trajectories. As shown in the previous section, both unconditional samples (and samples guided by an invariant classifier) reflect this symmetry property—it will be equally likely to sample a trajectory and its rotated or permuted counterpart. However, concrete tasks will often break this invariance, for instance by requiring that a robot or object is brought into a particular location or specifying an ordering over objects to be manipulated in a scene.
<sup>6</sup>Task-specific symmetry breaking would be more challenging to implement in classifier-free guidance. That would require training a diffusion model that jointly models equivariant unconditional and non-equivariant task-conditional distributions, which will in general be difficult.
Environment	Standard setting	SO(3) generalization
BCQ	CQL	Diffuser	EDGI (ours)	Diffuser	EDGI (ours)
Navigation	-	-	94.9±3.9	95.1±3.4	5.6±4.4	83.3±3.5
Unconditional	0.0	24.4	59.7±2.6	68.7±2.5	38.7±2.3	69.0±2.7
Conditional	0.0	0.0	46.0±3.4	52.0±3.6	16.7±2.0	35.9±3.5
Rearrangement	0.0	0.0	49.2±3.3	47.2±3.9	17.8±2.3	45.0±3.6
Average	0.0	8.1	51.6±1.8	56.0±2.0	24.4±1.3	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 our reproduction using their codebase. Left: Models trained on the standard datasets. Right: SO(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-based approach with classifier guidance allows us to elegantly break the symmetry at test time as required. Such a soft symmetry breaking both occurs through conditioning on the current state, by conditioning on a goal state, and through a non-invariant reward model used for guidance during sampling.
We demonstrate the effectiveness of incorporating symmetries as a powerful inductive bias in the Diffuser algorithm with experiments in two environments. The first environment is a 3D navigation task, in which an agent needs to navigate a number of obstacles to reach a goal state. Rewards are awarded based on the distance to the goal at each step, with penalties for collisions with obstacles. The position of the obstacles and the goal state are different in each episode and part of the observation. For simplicity, the actions directly control the acceleration of the agent; both the agent and the obstacles are spherical. Please see Fig. 1 for a schematic representation of this task and Appendix C for more details and the reward structure for this task.
In our remaining experiments, the agent controls a simulated Kuka robotic arm interacting with four blocks on a table. We use a benchmark environment introduced by Janner et al. [2022], which specifies three different tasks: an unconditional block stacking task, a conditional block stacking task where the stacking order is specified, and a rearrangement problem, in which the stacking order has to be changed in a particular way. For both environments, we train on offline trajectory datasets of roughly  $10^{5}$  (navigation) or  $10^{4}$  (manipulation) trajectories. We describe the setup in detail in Appendix D.
Algorithms. We train our EDGI on the offline dataset and use conditional sampling to plan the next actions. For the conditional and rearrangement tasks in the Kuka environment, we use classifier guidance following Janner et al. [2022].
As our main baseline, we compare our results to the (non-equivariant) Diffuser model [Janner et al., 2022]. In addition to a straightforward model, we consider a version trained with SO(3) data augmentation. We also compare two model-based RL baselines reported by Janner et al. [2022], BCQ [Fujimoto et al., 2019] and CQL [Kumar et al., 2020]. To study the benefits of the symmetry groups in isolation, we construct two EDGI variations: one is equivariant with respect to SE(3), but not  $S_{n}$ ; while the other is equivariant with respect to  $S_{n}$ , but not SE(3). Both are equivariant to temporal translations, just like EDGI and the baseline Diffuser.
Task performance. We report the results on both navigation and object tasks in Tab. 1. For each environment, we evaluate 100 episodes and report the average reward and standard error for each method. In both environments and across all tasks, EDGI performs as well as or better than the Diffuser baseline when using the full training set. In the navigation task, achieving a good performance for the baseline required substantially increasing the model's capacity compared to the hyperparameters used in Janner et al. [2022]. On the Kuka environment, both diffusion-based methods clearly outperform the BCQ and CQL baselines.
Sample efficiency. We study EDGI's sample efficiency by training models on subsets of the training
Figure 3: Average reward as a function of training dataset size for EDGI and Diffuser. Left: navigation environment. Right: Kuka object manipulation, averaged over the three tasks.
data. The results in Fig. 3 show that EDGI achieves just as strong rewards in the Kuka environment when training with only on  $10\%$  of the training data, and on the navigation task even when training on only  $0.01\%$  if the training data. The Diffuser baseline is much less sample-efficient. Training the Diffuser model with data augmentation partially closes the gap, but EDGI still maintains an edge. Our results provide evidence for the benefits of the inductive bias of equivariant models and matches similar observations in other works for using symmetries in an RL context [van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021, Rezaei-Shoshtari et al., 2022, Deac et al., 2023].
Effects of individual symmetries. In the left panel of Fig. 3, we also show results for EDGI variations that are only equivariant with respect to  $\mathrm{SE}(3)$ , but not  $\mathrm{S}_n$ , or vice versa. Both partially equivariant methods perform better than the Diffuser baseline, but not as well as the EDGI model equivariant to the full product group  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ . This confirms that the more of the symmetry of a problem we take into account in designing an architecture, the bigger the benefits in sample efficiency can be.
Group generalization. Finally, we demonstrate that equivariance improves generalization across the SO(3) symmetry group. On both environments, we train EDGI and Diffuser models on restricted offline datasets in which all trajectories are oriented in a particular way. In particular, in the navigation environment, we only use training data that navigates towards a goal location with  $x = 0$ . In the robotic manipulation tasks, we only use training trajectories where the red block is in a position with  $x = 0$  at the beginning of the episode. We test all agents on the original environment, where they encounter goal positions and block configurations unseen during training. We show results for these experiments in Tab. 1. The original Diffuser performs substantially worse, showing its limited capabilities to generalize to the new setting. In contrast, the performance of EDGI is robust to this domain shift, confirming that equivariance helps in generalizing across the symmetry group.
Diffusion-based planning. The closest work to ours is the original Diffuser paper [Janner et al., 2022], which we used as a baseline. Concurrent to our work, Diffuser was extended by Ajay et al. [2022], who used a separate inverse dynamics model and classifier-free guidance. The key novelty of our work is that we make this approach aware of the symmetry structure of planning problems through a new  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant denoising network.
Equivariant deep learning. Baking in symmetries into deep learning architectures was first studied in the work of Cohen and Welling [2016a] for geometric transformations, and the DeepSet architecture for permutations [Zaheer et al., 2017]. Followup work to group convolutional networks focused on both spherical geometry [Cohen et al., 2018] and building kernels using irreducible group representations [Cohen and Welling, 2016b, Weiler and Cesa, 2019, Cesa et al., 2021]. For symmetries of the 3D space—i.e. subgroups of  $\mathrm{E}(3)$ —a dominant paradigm is to use the message passing framework [Gilmer et al., 2017] along with geometric quantities like positions, velocities, and relative angles [Satorras et al., 2021, Schütt et al., 2021, Batatia et al., 2022].
Equivariance in RL. The role of symmetries has also been explored in reinforcement learning problems with a body of work focusing on symmetries of the joint state-action space of an MDP
[van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021, Muglich et al., 2022, Wang and Walters, 2022, Wang et al., 2022, Cetin et al., 2022, Rezaei-Shoshtari et al., 2022]. More recently, model-based approaches—like our proposed EDGI—have also benefited from increased data efficiency through the use of symmetries of the environment [Deac et al., 2023]. Concurrently to this work, Brehmer et al. [2023] also experiment with an equivariant Diffuser variation, but their denoising network is based on geometric algebra representations and a transformer architecture.
Equivariant generative models. Early efforts in learning invariant densities using generative models utilized the continuous normalizing flow (CNF) framework. A variety of works imbued symmetries by designing equivariant vector fields [Köhler et al., 2020, Rezende and Mohamed, 2015, Bose and Kobyzev, 2021]. As flow-based models enjoy exact density estimation, their application is a natural fit for applications in theoretical physics [Boyda et al., 2020, Kanwar et al., 2020] and modeling equivariant densities on manifolds [Katsman et al., 2021]. Other promising approaches to CNFs include equivariant score matching [De Bortoli et al., 2022] and diffusion models [Hoogeboom et al., 2022, Xu et al., 2022, Igashov et al., 2022]. Our proposed EDGI model extends the latter category to the product group  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$  and increases flexibility with respect to the data representations.
Embodied agents often solve tasks that are structured through the spatial, temporal, or permutation symmetries of our 3D world. Taking this structure into account in the design of planning algorithms can improve sample efficiency and generalization—notorious weaknesses of RL algorithms.
We introduced EDGI, an equivariant planning algorithm that operates as conditional sampling in a generative model. The main innovation is a new diffusion model that is equivariant with respect to the symmetry group  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$  of spatial, temporal, and object permutation symmetries. Beyond this concrete architecture, our work presents a general blueprint for the construction of networks that are equivariant with respect to a product group and support multiple representations in the data. Integrating this equivariant diffusion model into a planning algorithm allows us to model an invariant base density, but still solve non-invariant tasks through task-specific soft symmetry breaking. We demonstrated the performance, sample efficiency, and robustness of EDGI on object manipulation and navigation tasks.
While our work shows encouraging results, training and planning are currently expensive. Progress on this issue can come both from more efficient layers in the architecture of the denoising model as well as from switching to recent continuous-time diffusion methods with accelerated sampling.
We would like to thank Gabriele Cesa, Daniel Dijkman, and Pietro Mazzaglia for helpful discussions.
J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015. (Cited on page 3)
E. van der Pol, T. Kipf, F. A. Oliehoek, and others. Plannable approximations to MDP homomorphisms: Equivariance under actions. arXiv preprint arXiv, 2020. (Cited on pages 2, 9, and 10)
S. Villar, D. W. Hogg, K. Storey-Fisher, W. Yao, and B. Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics. Advances in Neural Information Processing Systems, 34:28848-28863, 2021. (Cited on page 6)
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 R. Walters. So (2) equivariant reinforcement learning. In International Conference on Learning Representations, 2022. (Cited on pages 2 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)-equivariant steerable cnns. Advances in Neural Information Processing Systems, 32, 2019. (Cited on page 9)
M. Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022. (Cited on page 10)
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)
A Group theory
In this appendix, we provide a basic introduction to Lie groups with an emphasis on the Lie group SE(3) which is used in the main paper.
A Lie group is both a group and a smooth manifold. A manifold is a topological space and thus a Lie group is a topological group that is formally defined below.
Definition 2. A topological group is a topological space  $G$  equipped with continuous maps that allow group composition  $\circ : G \times G \to G$  and group inverses  $(\circ)^{-1} : G \to G$ .
Now let  $g \in G$  be an element of the topological group and consider the map  $\bar{g}: G \to G$  given by  $\bar{g}(g) = g \circ g'$ . If  $g \circ g_1 = g \circ g_2$  then  $g_1 = g_2$  and thus  $\bar{g}$  is injective and for all  $g \in G$ , and there exists an element—i.e.  $g^{-1} \circ g$  such that  $\bar{g}(g^{-1} \circ g) = g \circ g^{-1} \circ g' = g'$ . Thus,  $\bar{g}$  is a bijection.
In a topological space  $\bar{g}$  and  $\bar{g}^{-1}$  are continuous. Thus,  $\bar{g}$  is a homomorphism. This means if  $U\subset G$  is open 1.)  $g\circ U$  is open 2.)  $U\circ g$  is open and inversion is also a homomorphism so  $U$  is open if and only if  $U^{-1}$  is open. A topological space  $\mathcal{X}$  is said to be homogenous if  $\forall x,y\in \mathcal{X}$  there exists a map  $h_{x,y}:\mathcal{X}\to \mathcal{X}$  that is a homomorphism such that  $h_{x,y}(x) = y$ . Note that a topological group is always homogenous. In words, a homogenous space for a topological group means any group element is reachable by a suitable group homomorphism.
Matrix Lie groups. A group of special interest is the group of  $n \times n$  invertible matrices with entries in  $\mathbb{R}$ ,  $\mathrm{GL}_n(\mathbb{R})$ . This is a topological group equipped with Euclidean topology. The subgroup  $\mathrm{SL}_n(\mathbb{R}) \subset \mathrm{GL}_n(\mathbb{R})$  is a closed subgroup with the property that  $\det = 1$ . As an example  $\operatorname{SE}(3) \subset \operatorname{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  $\mathrm{GL}_n(\mathbb{R})$ . This means  $\mathrm{SL}_n(\mathbb{R})$  is a matrix Lie group as are  $\mathrm{O}(n)$  and  $\mathrm{SO}(n)$ . Moreover,  $\mathrm{O}(n)$  and  $\mathrm{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  $\mathrm{SE}(3)\cong \mathrm{SO}(3)\ltimes \mathbb{R}^3$  and thus can be written as,