SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Johann Brehmer*
Qualcomm AI Research†
jbrehmer@qti.Qualcomm.com
Joey Bose*†
Qualcomm AI Research
Mila, McGill University
Pim de Haan
Taco Cohen
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  $\mathrm{SE}(3)$ , the discrete-time translation group  $\mathbb{Z}$ , and the object permutation group  $\mathrm{S}_n$ . EDGI follows the Diffuser framework by Janner et al. [2022] in treating both learning a world model and planning in it as a conditional generative modeling problem, training a diffusion model on an offline trajectory dataset. We introduce a new  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ -equivariant diffusion model that supports multiple representations. We integrate this model in a planning loop, where conditioning and classifier guidance let us softly break the symmetry for specific tasks as needed. On object manipulation and navigation tasks, EDGI is substantially more sample efficient and generalizes better across the symmetry group than non-equivariant models.
Our world is awash with symmetries. The laws of physics are the same everywhere in space and time—they are symmetric under translations and rotations of spatial coordinates as well as under time shifts. $^{4}$  In addition, whenever multiple identical or equivalent objects are labeled with numbers, the system is symmetric with respect to a permutation of the labels. Embodied agents are exposed to this structure, and many common robotic tasks exhibit spatial, temporal, or permutation symmetries. The gaits of a quadruped are independent of whether it is moving East or North, and a robotic gripper would interact with multiple identical objects independently of their labeling. However, most reinforcement learning (RL) and planning algorithms do not take this rich structure into account. While they have achieved remarkable success on well-defined problems after sufficient training, they are often sample-inefficient [Holland et al., 2018] and lack robustness to changes in the environment.
To improve the sample efficiency and robustness of RL algorithms, we believe it is paramount
*Equal contribution, order determined through a game of table tennis
$\dagger$ Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc.
$\ddagger$ Work done during an internship at 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.
to develop them with an awareness of their symmetries. Such algorithms should satisfy two key desiderata. First, policy and world models should be equivariant with respect to the relevant symmetry group. Often, for embodied agents this will be a subgroup of the product group of the spatial symmetry group SE(3), the group of discrete time shifts  $\mathbb{Z}$  and one or multiple object permutation groups  $\mathrm{S}_n$ . Second, it should be possible to softly break (parts of) the symmetry group to solve concrete tasks. For example, a robotic gripper might be tasked with moving an object to a specific point in space, which breaks the symmetry group SE(3). First works on equivariant RL have demonstrated the potential benefits of this approach [van der Pol et al., 2020, Walters et al., 2020, Mondal et al., 2021, Muglich et al., 2022, Wang and Wal
ters, 2022, Wang et al., 2022, Cetin et al., 2022, Rezaei-Shoshtari et al., 2022, Deac et al., 2023]. However, these works generally only consider small finite symmetry groups such as  $C_n$  and do not usually allow for soft symmetry breaking at test time based on a specific task.
In this paper, we introduce the Equivariant Diffuser for Generating Interactions (EDGI), an equivariant algorithm for model-based reinforcement learning and planning. EDGI consists of a base component that is equivariant with respect to the full product group  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$  and supports the multiple different representations of this group we expect to encounter in embodied environments. Moreover, EDGI allows for a flexible soft breaking of the symmetry at test time depending on the task.
Our work builds on the Diffuser method by Janner et al. [2022], who approach both the learning of a dynamics model and planning within it as a generative modeling problem. The key idea in Diffuser is to train a diffusion model on an offline dataset of state-action trajectories. To plan with this model, one samples from it conditionally on the current state, using classifier guidance to maximize the reward.
Our main contribution is a new diffusion model that is equivariant with respect to the product group  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$  of spatial, temporal, and permutation symmetries and supports data consisting of multiple representations. We introduce a new way of embedding multiple input representations into a single internal representation, as well as novel temporal, object, and permutation layers that act on the individual symmetries. When integrated into a planning algorithm, our approach allows for a soft breaking of the symmetry group through test-time task specifications both through conditioning and classifier guidance.
We demonstrate EDGI empirically in 3D navigation and robotic object manipulation environments. We find that EDGI greatly improves the performance in the low-data regime—matching the performance of the best non-equivariant baseline when using an order of magnitude less training data. In addition, EDGI is significantly more robust to symmetry transformations of the environment, generalizing well to unseen configurations.
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  $\mathrm{SE}(3)$ , discrete-time translations  $\mathbb{Z}$ , and object permutations, which are defined using the symmetric group of  $n$  elements  $\mathbb{S}_n$ . We now recall the main definition of equivariance and invariance of functions.
Definition 1. A function  $f: \mathcal{X} \to \mathcal{Y}$  is called  $G$ -equivariant if  $g \cdot f(x) = f(g \cdot x)$  for all  $g \in G$  and  $x \in \mathcal{X}$ . Here  $\mathcal{X}$  and  $\mathcal{Y}$  are input and output spaces that carry a  $G$  action denoted by  $\cdot$ . The function  $f$  is called  $G$ -invariant if the group action in  $\mathcal{Y}$  is trivial,  $g \cdot y = y$ .
We will focus on  $\mathcal{X} = \mathbb{R}^n$ ,  $\mathcal{Y} = \mathbb{R}^m$ , and linear group actions or representations, which are group
Figure 1: Schematic of EDGI in a navigation task, where the agent (red square) plans the next actions (red dots) to reach the goal (green star) without touching obstacles (grey circles). Top: planning as conditional sampling from a diffusion model. Bottom: effect of a group action. Equivariance requires the diagram to commute.
homomorphisms  $\rho : G \to \mathrm{GL}(\mathbb{R}^k)$ . Examples include rotation and permutation matrices. For a more comprehensive introduction to group and representation theory, we direct the interested reader to Appendix A or Esteves [2020] and Bronstein et al. [2021].
For generative modeling, we seek to model  $G$ -invariant densities. As proven in [Köhler et al., 2020, Bose and Kobyzev, 2021, Papamakarios et al., 2021], given a  $G$ -invariant prior density it is sufficient to construct a  $G$ -equivariant map to reach the desired  $G$ -invariant target density. In Sec. 3, we design  $G$ -equivariant diffusion architectures to model a distribution of trajectories that are known to be symmetric with respect to the product group  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ .
Diffusion models. Diffusion models [Sohl-Dickstein et al., 2015] are latent variable models that generate data by iteratively inverting a diffusion process. This diffusion process starts from a clean data sample  $x \sim q(x_0)$  and progressively injects noise for  $i \in [T]$  steps until the distribution is pure noise. The reverse, generative process takes a sample from a noise distribution and denoises it by progressively adding back structure, until we return to a sample that resembles being drawn from the empirical data distribution  $p(x)$ .
In diffusion models, it is customary to choose a parameter-free diffusion process (e.g. Gaussian noise with fixed variance). Specifically, we may define  $q(x_{t}|x_{t - 1})$  as the forward diffusion distribution modeled as a Gaussian centered around the sample at timestep  $x_{t - 1}$ :  $q(x_{t}|x_{t - 1}) = \mathcal{N}(x_t;\sqrt{1 - \beta_t} x_{t - 1},\beta_tI)$ , where  $\beta_{t}$  is a known variance schedule. The reverse generative process is learnable and can be parametrized using another distribution  $p_{\theta}(x_{t - 1}|x_t) = \mathcal{N}(x_{t - 1};\mu_\theta (x_t,t),\sigma_t^2 I)$  and the constraint that the terminal marginal at time  $T$  is a standard Gaussian—i.e.  $p(x_{T}) = \mathcal{N}(0,I)$ . The generative process can be learned by maximizing a variational lower bound on the marginal likelihood. In practice, instead of predicting the mean of the noisy data, it is convenient to predict the noise level  $\epsilon_{t}$  directly [Ho et al., 2020]. Furthermore, to perform low-temperature sampling in diffusion models it is possible to leverage a pre-trained classifier to guide the generation process [Dhariwal and Nichol, 2021]. To do so we can modify the diffusion score by including the gradient of the log-likelihood of the classifier  $\bar{\epsilon}_{\theta}(x_{t},t,y) = \epsilon_{\theta}(x_{t},t) - \lambda \sigma_{t}\nabla_{x_{t}}\log p(y|x_{t})$ , where  $\lambda$  is the guidance weight and  $y$  is the label.
Trajectory optimization with diffusion. We are interested in modeling systems that are governed by discrete-time dynamics of a state  $s_{h+1} = f(s_h, a_h)$ , given the state  $s_h$  and action  $a_h$  taken at timestep  $h$ . The goal in trajectory optimization is then to find a sequence of actions  $\mathbf{a}_{0:H}^*$  that maximizes an objective (reward)  $\mathcal{J}$  which factorizes over per-timestep rewards  $r(s_h, a_h)$ . Formally, this corresponds to the optimization problem  $\mathbf{a}_{0:H}^* = \arg \max_{a_0:H} \mathcal{J}(s_0, a_0:H) = \arg \max_{a_0:H} \sum_{h=0}^{H} r(s_h, a_h)$ , where  $H$  is the planning horizon and  $\tau = (s_0, a_0, \dots, s_H, a_H)$  denotes the trajectory.
A practical method to solve this optimization problem is to unify the problem of learning a model of the state transition dynamics and the problem of planning with this model into a single generative modeling problem. Janner et al. [2022] propose to train a diffusion model on offline trajectory data consisting of state-action pairs, learning a density  $p_{\theta}(\tau)$ . Planning can then be phrased as a conditional sampling problem: finding the distribution  $\tilde{p}_{\theta}(\tau) \propto p_{\theta}(\tau)c(\tau)$  over trajectories  $\tau$  where  $c(\tau)$  encodes constraints on the trajectories and specifies the task for instance as a reward function. Diffusion models allow conditioning in a way similar to inpainting in generative image modeling, and test-time reward maximization in analogy to classifier-based guidance.
3 Equivariant diffuser for generating interactions (EDGI)
We now describe our EDGI method. We begin by discussing the symmetry group  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$  and common representations in robotic problems. In Sec. 3.2 we introduce our key novelty, an  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant diffusion model for state-action trajectories  $\tau$ . 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  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ , which is a product of three distinct groups: 1. the group of spatial translations and rotations  $\mathrm{SE}(3)$ , 2. the discrete time translation symmetry  $\mathbb{Z}$ , and 3. the permutation group over  $n$  objects  $\mathrm{S}_n$ . It is important to note, however, that this symmetry group may be partially broken in an environment. For instance, the direction of gravity
Figure 2: Architecture of our  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant denoising network. Input trajectories (top left), which consist of features in different representations of the symmetry group, are first transformed into a single internal representation (green block). 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.
usually breaks the spatial symmetry group  $\mathrm{SE}(3)$  to the smaller group  $\mathrm{SE}(2)$ , and distinguishable objects in a scene may break permutation invariance. We follow the philosophy of modeling invariance with respect to the larger group and including any symmetry-breaking effects as inputs to the networks.
We require that spatial positions are always expressed relative to a reference point, for example, the robot base or center of mass. This guarantees equivariance with respect to spatial translations: to achieve  $\mathrm{SE}(3)$  equivariance, we only need to design an SO(3)-equivariant architecture.
Data representations. We consider 3D environments that contain an embodied agent as well as  $n$  other objects. We parameterize their degrees of freedom with two SO(3) representations, namely the scalar representation  $\rho_0$  and the vector representation  $\rho_1$ . Other representations that may occur in a dataset can often be expressed in  $\rho_0$  and  $\rho_1$ . For example, object poses may be expressed as 3D rotations between a global frame and an object frame. We can convert such rotations into two  $\rho_1$  representations through the inclusion map  $\iota: \mathrm{SO}(3) \hookrightarrow \mathbb{R}^{2 \times 3}$  that chooses the first two columns of the rotation matrix. Such an embedding avoids challenging manifold optimization over SO(3), while we can still uncover the rotation matrix by orthonormalizing using a Gram-Schmidt process to recover the final column vector. Thus, any SE(3) pose can be transformed to these two representations.
We assume that all trajectories transform under the regular representation of the time translation group  $\mathbb{Z}$  (similar to how images transform under spatial translations). Under  $\mathrm{S}_n$ , object properties permute, while robot properties or global properties of the state remain invariant. Each feature is thus either in the trivial or the standard representation of  $\mathrm{S}_n$ .
Overall, we thus expect that data in environments experienced by our embodied agent to be categorized into four representations of the symmetry group  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ : scalar object properties, vector object properties, scalar robotic degrees of freedom (or other global properties of the system), and vector robotic degrees of freedom (again including other global properties of the system).
3.2 Equivariant diffusion model
Our main contribution is a novel  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant diffusion model which leads to an invariant distribution over trajectories. Specifically, given an invariant base density with respect to our
chosen symmetry group—an isotropic Gaussian satisfies this property for  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$  and an denoising model  $f$  that is equivariant with respect to the same group, we arrive at a diffusion model that is  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -invariant [Köhler et al., 2020, Papamakarios et al., 2021]. Under mild assumptions, such an equivariant map that pushes forward the base density always exists [Bose and Kobyzev, 2021].
We design a novel equivariant architecture for the denoising model  $f$ . Implemented as a neural network, it maps noisy input trajectories  $\tau$  and a diffusion time step  $i$  to an estimate  $\hat{\epsilon}$  of the noise vector that generated the input. Our architecture does this in three steps. First, the input trajectory consisting of various representations is transformed into an internal representation of the symmetry group. Second, in this representation the data are processed with an equivariant network. Finally, the outputs are transformed from the internal representation into the original data representations present in the trajectory. We illustrate the architecture of our EDGI model in Fig. 2.
Step 1: Representation mixer. The input noisy trajectory consists of features in different representations of the symmetry group (see above). While it is possible to mirror these input representations for the hidden states of the neural network, the design of equivariant architectures is substantially simplified if all inputs and outputs transform under a single representation. Hence, we decouple the data representation from the representation used internally for the computation—in a similar fashion to graph neural networks that decouple the data and computation graphs.
Internal representation. We define a single internal representation that for each trajectory time step  $t \in [H]$ , each object  $o \in [n]$ , each channel  $c \in [n_c]$  consists of one<sup>5</sup> SO(3) scalar  $s_{toc}$  and one SO(3) vector  $v_{toc}$ . We write  $w_{toc} = (s_{toc}, v_{toc}) \in \mathbb{R}^4$ . Under spatial rotations  $g \in \mathrm{SO}(3)$ , these features thus transform as the direct sum of the scalar and vector representations  $\rho_0 \oplus \rho_1$ :
$$
w _ {t o c} = \binom {s _ {t o c}} {v _ {t o c}} \rightarrow w _ {t o c} ^ {\prime} = \binom {\rho_ {0} (g) s _ {t o c}} {\rho_ {1} (g) v _ {t o c}}. \tag {1}
$$
These internal features transform in the regular representation under time shift and in the standard representation under permutations  $\mathbb{P}$  as  $w_{to c} \rightarrow w_{t o^{\prime}c} = \sum_{o} \mathbb{P}_{o^{\prime}o} w_{to c}$ . 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  $\tau$ , which consists of features in different representations of  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$ , into the internal representation. On the one hand, we pair up  $\mathrm{SO}(3)$  scalars and  $\mathrm{SO}(3)$  vectors into  $\rho_0 \oplus \rho_1$  features. On the other hand, we distribute global features - those a priori unassigned to one of the  $n$  objects in the scene - over the  $n$  objects.
Concretely, for each object  $o \in [n]$ , each trajectory step  $t \in [H]$ , and each channel  $c = [n_c]$ , we define the input in the internal representation as  $w_{toc} \in \mathbb{R}^4$  as follows:
$$
w _ {t o c} = \left(\frac {\sum_ {c ^ {\prime}} \mathbf {W} _ {o c c ^ {\prime}} ^ {1} s _ {t o c ^ {\prime}}}{\sum_ {c ^ {\prime}} \mathbf {W} _ {o c c ^ {\prime}} ^ {2} v _ {t o c ^ {\prime}}}\right) + \left(\frac {\sum_ {c ^ {\prime}} \mathbf {W} _ {o c c ^ {\prime}} ^ {3} s _ {t \emptyset c ^ {\prime}}}{\sum_ {c ^ {\prime}} \mathbf {W} _ {o c c ^ {\prime}} ^ {4} v _ {t \emptyset c ^ {\prime}}}\right). \tag {2}
$$
The matrices  $\mathbf{W}^{1,2,3,4}$  are learnable and of dimension  $n\times n_c\times n_s^{\mathrm{object}}$ ,  $n\times n_c\times n_v^{\mathrm{object}}$ ,  $n\times n_c\times n_s^{\mathrm{global}}$ , or  $n\times n_c\times n_v^{\mathrm{global}}$ , respectively. Here  $n_s^{\mathrm{object}}$  is the number of SO(3) scalar quantities associated with each object in the trajectory,  $n_v^{\mathrm{object}}$  is the number of SO(3) vector quantities associated with each object,  $n_s^{\mathrm{global}}$  is the number of scalar quantities associated with the robot or global properties of the system, and  $n_v^{\mathrm{global}}$  is the number of vectors of that nature. The number of input channels  $n_c$  is a hyperparameter. We initialize the matrices  $\mathbf{W}^i$  such that Eq. (2) corresponds to a concatenation of all object-specific and global features along the channel axis at the beginning of training.
Step 2:  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant U-net. We then process the data with a  $\mathrm{SO}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant denoising network. Its key components are three alternating types of layers. Each type acts on the representation dimension of one of the three symmetry groups while leaving the other two invariant—i.e. they do not mix internal representation types of the other two layers:
- Temporal layers: Time-translation-equivariant convolutions along the temporal direction (i.e. along trajectory steps), organized in a U-Net architecture.
- Object layers: Permutation-equivariant self-attention layers over the object dimension.
Pairing up just one scalar and one vector is a 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.
- Geometric layers: SO(3)-equivariant interaction between the scalar and vector features.
In addition, we use residual connections, a new type of normalization layer that does not break equivariance, and context blocks that process conditioning information and embed it in the internal representation (see Appendix B for more details). These layers are combined into an equivariant block consisting of one instance of each layer, and the equivariant blocks are arranged in a U-net, as depicted in Fig. 2. Between the levels of the U-net, we downsample (upsample) along the trajectory time dimension by factors of two, increasing (decreasing) the number of channels correspondingly.
Temporal layers. Temporal layers consist of 1D convolutions along the trajectory time dimension. To preserve  $\mathrm{SO}(3)$  equivariance, these convolutions do not add any bias and there is no mixing of features associated with different objects nor the four geometric features of the internal SO(3) representation.
Object layers. Object layers enable features associated with different objects to interact via an equivariant multi-head self-attention layer. Given inputs  $w_{tot}$ , the object layer computes
$$
\mathbf {K} _ {t o c} = \sum_ {c ^ {\prime}} \mathbf {W} _ {c c ^ {\prime}} ^ {K} w _ {t o c}, \mathbf {Q} _ {t o c} = \sum_ {c ^ {\prime}} \mathbf {W} _ {c c ^ {\prime}} ^ {Q} w _ {t o c}, \mathbf {V} _ {t o c} = \sum_ {c ^ {\prime}} \mathbf {W} _ {c c ^ {\prime}} ^ {V} w _ {t o c},
$$
$$
w _ {t o c} ^ {\prime} \propto \sum_ {o ^ {\prime}} \operatorname {s o f t m a x} _ {o ^ {\prime}} \left(\frac {\mathbf {Q} _ {t o} \cdot \mathbf {K} _ {t o ^ {\prime}}}{\sqrt {d}}\right) \mathbf {V} _ {t o ^ {\prime} c}, \tag {3}
$$
with learnable weight matrices  $\mathbf{W}^{K,V,Q}$  and  $d$  the dimensionality of the key vector. There is no mixing between features associated with different time steps, nor between the four geometric features of the internal SO(3) representation. Object layers are SO(3)-equivariant, as the attention weights compute invariant SO(3) norms.