SO(2)-Equivariant Reinforcement Learning

Overview

The paper proposes incorporating rotational symmetry directly into the actor and critic networks of reinforcement learning agents. The main idea is that if rotating a visual state and its corresponding action does not change the reward or transition dynamics, then the agent should not need to learn each orientation independently. Equivariant networks enforce this relation by construction, which can improve sample efficiency compared with learning the same symmetry from data augmentation.

A group-invariant MDP is an MDP in which a group $G$ acts on the state and action spaces while preserving the reward and transition function. For every $g\in G$,

$$R(s,a)=R(gs,ga)$$

and

$$T(s,a,s')=T(gs,ga,gs').$$

The optimal Q-function is therefore invariant when the state and action are transformed together,

$$Q^*(s,a)=Q^*(gs,ga),$$

while the optimal policy is equivariant,

$$\pi^*(gs)=g\pi^*(s).$$

The paper studies a robotic manipulation task from top-down images. The continuous rotation transformation $G=SO(2)$ is approximated using a finite cyclic rotation group $C_n$. The state $\mathcal{F}_s$ is a top-down depth image centered on the gripper. The action is

$$a=(a_{\lambda},a_{xy},a_z,a_{\theta}),$$

where the planar displacement $a_{xy}$ rotates with the image, while the gripper command $a_{\lambda}$, vertical displacement $a_z$, and change in gripper orientation $a_{\theta}$ remain invariant. Thus,

$$A_{\mathrm{equiv}}=A_{xy}, \qquad A_{\mathrm{inv}}=A_{\lambda}\times A_z\times A_{\theta},$$

and the group acts on an action as

$$ga= \left( \rho_{\mathrm{equiv}}(g)a_{\mathrm{equiv}}, a_{\mathrm{inv}} \right).$$

Equivariant DQN

For DQN, the equivariant and invariant action variables are discretized into $\mathcal{A}_{\mathrm{equiv}}$ and $\mathcal{A}_{\mathrm{inv}}$. Rather than producing an unstructured vector of Q-values, the network outputs a function

$$\mathcal{F}_a: \mathcal{A}_{\mathrm{equiv}} \rightarrow \mathbb{R}^{|\mathcal{A}_{\mathrm{inv}}|}.$$

The spatial location of $\mathcal{F}_a$ represents an equivariant action such as an $xy$ displacement, while its channels represent the invariant action variables. The Q-network $q$ is constructed from equivariant convolutional layers and satisfies

$$q(g\mathcal{F}_s) = gq(\mathcal{F}_s) = g\mathcal{F}_a.$$

The transformed output is defined by

$$[g\mathcal{F}_a](a_{\mathrm{equiv}}) = \mathcal{F}_a \left( \rho_{\mathrm{equiv}}(g)^{-1}a_{\mathrm{equiv}} \right).$$

Therefore, rotating the input image rotates the spatial Q-map but leaves its invariant-action channels unchanged. This is equivalent to the desired joint invariance

$$Q(gs,ga)=Q(s,a).$$

the value assigned to an action before rotating the scene is assigned to the correspondingly rotated action afterward. The DQN training objective itself is unchanged; the symmetry is enforced by the network architecture.

Equivariant SAC

SAC learns a continuous policy and a scalar critic. The actor outputs the mean action together with the standard deviation variables of its Gaussian action distribution,

$$\bar{a} = \left( a_{\mathrm{equiv}}, a_{\mathrm{inv}}, a_{\sigma} \right).$$

The paper assumes that the standard deviation is invariant; so, only the equivariant action variables rotate:

$$g\bar{a} = \left( \rho_{\mathrm{equiv}}(g)a_{\mathrm{equiv}}, a_{\mathrm{inv}}, a_{\sigma} \right).$$

The actor is an equivariant network constrained to satisfy

$$\pi(g\mathcal{F}_s) = g\pi(\mathcal{F}_s).$$

Consequently, rotating the observation rotates the predicted planar action while preserving the invariant action components.

The critic must instead be invariant to jointly rotating its state and action. An equivariant encoder $e$ first maps the image to a group-structured state representation $\bar{s}$:

$$e(g\mathcal{F}_s) = ge(\mathcal{F}_s) = g\bar{s}.$$

The representation is concatenated with the action to form

$$w=(\bar{s},a), \qquad gw=(g\bar{s},ga).$$

The Q-network then maps $w$ to a scalar and is constrained by

$$q(gw)=q(w).$$

Combining the encoder and Q-network gives

$$q\left(e(g\mathcal{F}_s),ga\right) = q\left(e(\mathcal{F}_s),a\right).$$

Equivariant SACfD

Some sparse-reward tasks are difficult to solve without guided exploration. Equivariant SACfD extends Equivariant SAC by first adding expert demonstrations from a hand-coded planner to the replay buffer. It also adds a behavior-cloning term to the actor loss for transitions that come from an expert demonstration:

$$L_{\mathrm{actor}} = L_{\mathrm{SAC}} + \mathbb{1}_e \left[ \frac{1}{2} \left\lVert a-a_e \right\rVert^2 \right], \qquad a\sim\pi(s),$$

where $a_e$ is the expert action and $\mathbb{1}_e=1$ for an expert transition and $0$ otherwise. The additional term makes the sampled policy action imitate the expert action without changing the ordinary SAC loss for non-expert transitions.