$\mathrm{SE}(3)$ Frame Equivariance in Dynamics Modeling and Reinforcement Learning

Symmetry groups and equivariance. A symmetry group is defined as a set $G$ together with a binary composition map satisfying the axioms of associativity, identity, and inverse. A (left) group action of $G$ on a set $\mathcal{X}$ is defined as the mapping $(g,x)\mapsto g\cdot x$ which is compatible with composition. Given a function $f:\mathcal{X}\to\mathcal{Y}$ and $G$ acting on $\mathcal{X}$ and $\mathcal{Y}$, then $f$ is $G$-equivariant if it commutes with group actions: $g\cdot f(x)=f(g\cdot x),\forall g\in G,\forall x\in\mathcal{X}$. In the special case the action on $\mathcal{Y}$ is trivial $g\cdot y=y$, then $f(x)=f(g\cdot x)$ holds, and we say $f$ is $G$-invariant.

We consider a dynamical system that models a robot in a 3D physical world. One example is locomotion, where a robot learns to move on a ground that provides support, with gravity directed downwards. The robot is represented by connected links (bodies) and joints (actuators). Suppose the system is modeled by a discrete-time Markov decision process (MDP) as $s^{\prime}=f(s,a)$. As done in Interaction Network [battaglia_interaction_2016] or Graph Network [sanchez-gonzalez_graph_2018], we could represent the robot state $s$ at each time step as a geometric graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$.

This geometric graph lives in a 3D geometric space, allowing for rotation and translation transformations. The graph nodes $v_{i}\in\mathcal{V}$ are links with features of positions, orientations, linear velocity, and angular velocity. The edges $e_{ij}\in\mathcal{E}$ describe the joints between links, where actuation is provided as edge feature. As an example, we can rotate the half cheetah $45^{\circ}$ along the vertical axis, where all nodes and edges are rotated correspondingly, as shown in the figure.

The frame symmetry in 3D is the focus of our interest. It implies that even if we change the reference frame by moving in different directions or locations, the robot does not need to relearn how to walk or run. The frame symmetry represents the proper transformations of the entire 3D world, represented as ${\mathrm{SE}(3)}$, the group of 3D continuous translations and rotations.

However, the robot is subject to external forces that may break symmetry, such as gravity, contact, and actuation forces. Despite these forces, the ${\mathrm{SE}(3)}$ symmetry can still be maintained if they are included as inputs to the system along with the robot state, during rotation or translation. The network learns from data that has only rotation symmetry along the gravity axis. We include these factors as global features to the geometric graph, which is transformed along with the graph. We explain how to transform the graph and its features in the next section.

In this section, we outlines the symmetry under consideration in the continuous control tasks and how we explore that. [ravindran2004algebraic, ravindran_symmetries_nodate, zinkevich_symmetry_2001] explore symmetry in MDPs with no function approximation like neural networks. [van2020mdp, mondal_group_2020] initiate the exploration of symmetry in model-free (deep) RL by using equivariant policy networks. [zhao_integrating_2022] study symmetry in value-based planning on 2D grid. We extend it and focus on MDPs with continuous state and action spaces and sampling-based control/planning algorithms.

To use frame equivariance in control, there are three key steps [zhao_integrating_2022]: (1) specify the symmetry group $G$ in the system, (2) learn a $G$-equivariant dynamics model $P(s^{\prime}\mid s,a)$, and (3) incorporate symmetry into continuous control. In the second step, using equivariant model induces equivariant Bellman operator: $\mathcal{T}[V]=\max_{a}R(s,a)+\gamma\sum_{s^{\prime}}P(s^{\prime}|s,a)V(s^{\prime})$. Albeit the previous two steps are analogous to prior work, the third step is not trivial because of the nature of sampling-based approaches. We explain it in the next subsection.

Since we use graph-based representation of the system, we make use of equivariant message passing networks for implementing the dynamics network. Specifically, we use steerable ${\mathrm{E}(3)}$-equivariant message passing network [brandstetter_geometric_2022]. Compared to using invariant scalar features, steerable networks allow to use higher-order equivariant steerable features, enabling better expressivity. We briefly explain how we use steerable message passing network and specify ${\mathrm{SO}(3)}$-representations for state and action features. We omit the details on hidden layers and more.

In prior work [zhao_integrating_2022], they developed Symmetric Value Iteration Network (SymVIN) for path planning on 2D grid ${\mathbb{Z}^{2}}$. The network iteratively apply learned Bellman operators to reach a fixed point $V^{\star}:{\mathbb{Z}^{2}}\to\mathbb{R}$. Since every step is equivariant, the entire value iteration process $\texttt{VI}(M)$ is equivariant:

where $M$ is the occupancy map and goal input ${\mathbb{Z}^{2}}\to\{0,1\}^{2}$. Since Bellman operators are a sequence of (linear) equivariant operations on steerable signals, the whole process can be implemented using steerable convolution networks [cohen_steerable_2016, weiler_general_2021].

Intuitively, SymVIN relate maps under four discrete rotations and two reflections, so if a map is rotated, it does not need to relearn the optimal plan or actions. Analogously, we aim to develop a control method that considers the frame equivariance of 3D space. One practical setup is locomotion on some terrain. The symmetry of locomotion task is that a robot does not to learn how to walk when facing different directions and locations, while instead share information between all directions and locations of ${\mathrm{SE}(3)}$.

Symmetry in MPC. Our goal is to exploit ${\mathrm{SE}(3)}$ frame symmetry of the underlying $\mathbb{R}^{3}$ space in the planning algorithm. For sampling-based planning, symmetries enable two types of benefits.

1. 1.

   Forward search. When sampling a trajectory $(s_{1},a_{1},s_{2},a_{2},\ldots)$ in the ground MDP, we equivariantly know the outcome of all trajectories under symmetries: $\{(g\cdot s_{1},g\cdot a_{1},g\cdot s_{2},g\cdot a_{2},\ldots)\mid g\in G\}$. This saves computation when the group is “large”. However, while this is important for discrete case, it is negligible for continuous state space.

2. 2.

   Value backup. Symmetries allow us to reuse knowledge between equivalent state, as $V^{\star}(s)=V^{\star}(g\cdot s)$ and $\pi(g\cdot s)=g\cdot\pi(s)$. Intuitively, facing different directions do not change how a robot walks.

In continuous control, the second type of consideration is crucial. Since the optimal policy and value functions are equivariant (or invariant), we constrain the function approximation to be only the set of equivariant functions.

Planning with MPC. We use model predictive control as a sampling-based planning approach for continuous actions. We use MPPI (Model Predictive Path Integral) control method [williams_model_2015, williams_aggressive_2016, williams_model_2017, williams_information_2017], as also done in [hansen_temporal_2022]. We sample $N$ trajectories with horizon $H$ using the learned dynamics model with actions from a learned policy, and estimate the expectation of total return $G_{\tau}$. Importantly, if we rotate the entire trajectory $\tau$ by $g\cdot\tau$, we know the return is invariant $G_{\tau}=G_{g\cdot\tau}$ since it is a scalar:

We can then update the action selection policy by using top-$k$ trajectories, which is a Gaussian distribution with learned mean and variance. To capture the symmetry here, we use equivariant policy network and invariant value network in estimating values and optimal actions.

Benefits of symmetry in continuous control. There are some benefits of explicitly considering symmetry in continuous control. The possibility of hitting orbits is negligible, so there is no need for orbit-search on symmetric states in forward search in continuous control. Additionally, the model predictive control algorithm implicitly plans in a smaller continuous MDP $\mathcal{M}/G$ [ravindran2004algebraic]. Furthermore, from equivariant network literature [elesedy_provably_2021], the generalization gap for learned equivariant policy and value networks are smaller, which allows them to generalize better.