SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {S E} (3) = \left\{\left( \begin{array}{c c} R & x \\ 0 & 1 \end{array} \right): R \in \operatorname {S O} (3), x \in \left(\mathbb {R} ^ {3}, +\right) \right\} \tag {5}
$$
Represented by this  $4 \times 4$  matrix and with the group operation defined by matrix multiplication, this group can be seen as a subgroup of the general linear group  $\mathrm{GL}_4(\mathbb{R})$ .
SO(3) as a matrix Lie group. The group of 3D rotations is SO(3). It is a compact matrix Lie group where each element  $R \in \mathrm{SO}(3)$  is a rotation. To represent  $R$  there are various possible choices. The most familiar of them is perhaps  $R \in \mathbb{R}^{3 \times 3}$  is as a  $3 \times 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  $\mathfrak{g}$ . Each element of the Lie algebra can be distinctively associated with a vector  $\omega \in \mathbb{R}^3$ . Given any  $\mathbf{v} \in \mathbb{R}^3$ , it's related by  $\mathfrak{g}\mathbf{v} = \boldsymbol{\omega} \times \mathbf{v}$ , with the symbol  $\times$  representing the cross product. The length of this vector, symbolized as  $\omega = ||\omega||$ , represents the rotation angle, while the unit direction, expressed as  $e_{\omega} = \frac{\omega}{||\omega||}$ , indicates the rotation axis.
Yet another parametrization of rotations can be achieved using quaternions. A quaternion,  $q$ , is often represented as  $q = w + xi + yj + zk$  where  $w, x, y$ , and  $z$  are real numbers and  $i, j$ , and  $k$  are mutually orthogonal imaginary units vectors. The unit vectors obey the following rule:  $i^2 = j^2 = k^2 =ijk = -1$ . A subset of quaternions, known as rotation quaternions, can be used to represent 3D rotations. These specific quaternions have magnitude 1 and can also be thought of as an ordered
The field can also be complex  $\mathbb{C}$  to allow invertible complex matrices  $\mathrm{GL}_n(\mathbb{C})$ .
pair  $(w, v)$  where  $w$  is a scalar and  $v = (x, y, z)$  is a 3D vector. The scalar part  $w$  can be thought of as  $\cos(\omega/2)$  and the vector part  $v$  as  $n \sin(\omega/2)$ , where  $\omega$  is the angle of rotation and  $n$  is the unit vector along the axis of rotation.
Finally, to rotate a point  $p$  in 3D space using a quaternion  $q$ , the point can be represented as a pure quaternion  $p' = 0 + xi + yj + zk$ . The rotated point  $p''$  is then given by:  $p'' = q \cdot p' \cdot q^*$ . Where  $q^*$  is the conjugate of  $q$  —i.e. if  $q = w + xi + yj + zk$ , then  $q^* = w - xi - yj - zk$ .
B Architecture details
On a high level, EDGI follows Diffuser [Janner et al., 2022]. In the following, we will describe the key difference: our  $\mathrm{SE}(3)\times \mathbb{Z}\times \mathrm{S}_n$ -equivariant architecture for the diffusion model.
Overall architecture. We illustrate the architecture in After converting the input data in our internal representation (see Sec. 3.2), the data is processed with an equivariant  $U$ -net with four levels. At each level, we process the hidden state with two residual standard blocks, before downsampling (in the downward pass) or upsampling (in the upward pass).
Residual standard block. The main processing unit of our architecture processes the current hidden state with an equivariant block consisting of a temporal layer, an object layer, a normalization layer, and a geometric layer. In parallel, the context information (an embedding of diffusion time and a conditioning mask) is processed with a context block. The hidden state is added to the output of the context block and processes with another equivariant block. Finally, we process the data with a linear attention layer over time. This whole pipeline consists of an equivariant block, a context block, and another equivariant block is residual (the inputs are added to the outputs).
Temporal layers. Temporal layers consist of one-dimensional convolutions without bias along the time dimension. We use a kernel size of 5.
Normalization layers. We use a simple equivariant normalization layer that for each batch element rescales the entire tensor  $w_{toc}$  to unit variance. This is essentially an equivariant version of LayerNorm. The difference is that our normalization layer does not shift the inputs to zero means, as that would break equivariance with respect to SO(3).
Geometric layers. In the geometric layers, the input state is split into scalar and vector components. The vector components are linearly transformed to reduce the number of channels to 16. We then construct all SO(3) invariants from these 16 vectors by taking pairwise inner products and concatenating them with the scalar inputs. This set of scalars is processed with two MLPs, each consisting of two hidden layers and ReLU nonlinearities. The MLPs output the scalar outputs and coefficients for a linear map between the vector inputs and the vector outputs, respectively. Finally, there is a residual connection that adds the scalar and vector inputs to the outputs.
Linear attention over time. To match the architecture used by Janner et al. [2022] as closely as possible, we follow their choice of adding another residual linear attention over time at the end of each level in the U-net. We make the linear attention mechanism equivariant by computing the attention weights as
Context blocks. The embeddings of diffusion time and conditioning information are processed with a Mish nonlinearity and a linear layer, like in Janner et al. [2022]. Finally, we embed them in our internal representation by zero-padding the resulting tensor.
Upsampling and downsampling. During the downsampling path, there is a final temporal layer that implements temporal downsampling and increases the number of channels by a factor of two. Conversely, during the upsampling path, we use a temporal layer for temporal upsampling and a reduction of the number of channels.
Equivalence. 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 \in \mathbb{R}^{n \times H \times c \times 4}$  be data in our internal representation, such that the entries  $w_{tot}$  decompose into SO(3) scalars  $s_{tot}$  and SO(3) vectors  $v_{tot}$ . Let  $S(w_{to})$  be the set of all scalars and all pairwise  $SO(3)$  inner products between the vectors  $v_{to}$ , as defined in Eq. (4). The outputs of the geometric layer are
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)_{tot} = (\phi(S(w_{to})), \sum_{c'} \psi(S(w_{to})) v_{tot'})$ .
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.
Next, consider the behavior under spatial transformations. Like most (S)E(3)-equivariant architectures, we deal with translations through canonicalization, defining all coordinates with respect to the center of mass or the robot base, as applicable. This means we only have to analyze the behavior under rotations.
Let  $R \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)_{tot} = (\phi(S(R \cdot w_{to})), \sum_{c'} \psi(S(w_{to}))R \cdot v_{tot'}) = (\phi(S(w_{to})), R \cdot \sum_{c'} \psi(S(w_{to}))v_{tot'}) = (R \cdot f(w))_{tot}$ . Hence the geometric layer is equivariant with respect to SE(3).
C Navigation experiments
We introduce a new navigation environment. The scene consists of a spherical agent navigating a plane populated with a goal state and  $n = 10$  spherical obstacles. At the beginning of every episode, the agent position, agent velocity, obstacle positions, and goal position are initialized randomly (in a rotation-invariant way). We simulate the environment dynamics with PyBullet [Coumans and Bai, 2016-2019].
Offline dataset. To obtain expert trajectories, we train a TD3 [Fujimoto et al., 2018] agent in the implementation by Raffin et al. [2021] for  $10^{7}$  steps with default hyperparameters on this environment. We generate  $10^{5}$  trajectories for our offline dataset.
State. The state contains the agent position, agent velocity, goal position, and obstacle positions.
Actions. The action space is two-dimensional and specifies a force acting on the agent.
Rewards. At each time step, the agent receives a reward equal to the negative Euclidean distance to the goal state. In addition, a penalty of  $-0.1$  is added to the reward if the agent touches any of the obstacles. Finally, there is an additional control cost equal to  $-10^{3}$  times the force acting on the agent. We affinely normalize the rewards such that a normalized reward of 0 corresponds to that achieved by a random policy and a normalized reward of 100 corresponds to the expert policy.
D Kuka experiments
We use the object manipulation environments and tasks from Janner et al. [2022], please see that work for details on the environment. In our experiments, we consider three tasks: unconditional stacking, conditional stacking, and block rearrangement. Figure 4 visualizes the unconditional block stacking task. For a fair comparison, we re-implement the Diffuser algorithm while making bug fixes in the codebase of Janner et al. [2022], which mainly included properly resetting the environment.
State. We experiment with two parameterizations of the Kuka environment state. For the Diffuser baseline, we use the original 39-dimensional parameterization from Janner et al. [2022].
For our EDGI, we need to parameterize the system in terms of  $\mathrm{SE}(3) \times \mathbb{Z} \times \mathrm{S}_n$  representations. We, therefore, describe the robot and block orientations with  $\mathrm{SO}(3)$  vectors as follows. Originally, the
robot state is specified through a collection of joint angles. One of these encodes the rotation of the base along the vertical  $z$ -axis. We choose to represent this angle as a  $\rho_{1}$  vector in the  $xy$ -plane. In addition, we add the gravity direction (the  $z$ -axis itself) as another  $\rho_{1}$  vector, which is also the normal direction of the table on which the objects rest. Combined, these vectors define the pose of the base of the robot arm. Rotating gravity direction, and the robot and object pose by SO(3) can be interpreted as a passive coordinate transformation, or as an active rotation of the entire scene, including gravity. As the laws of physics are invariant to this transformation, this is a valid symmetry of the problem.
The  $n$  objects can be translated and rotated. Their pose is thus given by a translation  $t \in \mathbb{R}^3$  and rotation in  $r \in \mathrm{SO}(3)$  relative to a reference pose. The translation transforms by a global rotation  $g \in \mathrm{SO}(3)$  as a vector via representation  $\rho_1$ . The rotational pose transforms by left multiplication  $r \mapsto gr$ . The SO(3) pose is not a Euclidean space, but a non-trivial manifold. Even though diffusion on manifolds is possible De Bortoli et al. [2022], Huang et al. [2022], we simplify the problem by embedding the pose in a Euclidean space. This is done by picking the first two columns of the pose rotation matrix  $r \in \mathrm{SO}(3)$ . These columns each transform again as a vector with representation  $\rho_1$ . This forms an equivariant embedding  $\iota: \mathrm{SO}(3) \hookrightarrow \mathbb{R}^{2 \times 3}$ , whose image is two orthogonal 3-vectors of unit norm. Via the Gram-Schmidt procedure, we can define an equivariant map  $\pi: \mathbb{R}^{2 \times 3} \to \mathrm{SO}(3)$  (defined almost everywhere), that is a left inverse to the embedding:  $\pi \circ \iota = \mathrm{id}_{\mathrm{SO}(3)}$ . Combining with the translation, the roto-translational pose of each object is thus embedded as three  $\rho_1$  vectors.
We also tested the performance of the baseline Diffuser method on this reparameterization of the state but found worse results.
Hyperparameters. We also follow the choices of Janner et al. [2022], except that we experiment with a linear noise schedule as an alternative to the cosine schedule they use. For each model and each dataset, we train the diffusion model with both noise schedules and report the better of the two results.