Approximate Equivariance in Reinforcement Learning

Overview

Exactly equivariant agents assume that transforming a state and action by a symmetry leaves the reward and transition dynamics unchanged. This assumption can improve sample efficiency because the agent shares what it learns between symmetric situations. However, exact equivariance can be too restrictive especially in real-world scenarios when the environment contains symmetry-breaking factors, such as asymmetric gravity, differences in actuator response, etc...
The paper proposes using approximate equivariance as a middle ground between exactly equivariant and unconstrained RL agents. The network retains a strong symmetry bias but can learn deviations from it when the data indicates that transformed situations are not identical.

Approximate Equivariance in MDPs

For a group $G$ acting on the input space $\mathcal{X}$ and output space $\mathcal{Y}$, the equivariance error of a function $f:\mathcal{X}\rightarrow\mathcal{Y}$ for a transformation $g\in G$ and input $x\in\mathcal{X}$ is

$$ee(f,g,x)=\left\lVert f(g(x))-g (f(x))\right\rVert.$$

This measures the difference between transforming the input before applying $f$ and applying $f$ before transforming its output. An exactly equivariant function has zero equivariance error for every input and every group transformation.

The $\varepsilon$-stabilizer contains the transformations whose equivariance error is bounded by $\varepsilon$ for all inputs:

$$\operatorname{Stab}_{\varepsilon}(f,G) =\left\{g\in G\mid ee(f,g,x)\leq\varepsilon,\ \forall x\in\mathcal{X}\right\}.$$

The function is approximately $G$-equivariant when

$$\operatorname{Stab}_{\varepsilon}(f,G)=G.$$

Therefore, every transformation in the group must satisfy the equivariance relation up to a bounded error. This differs from partial equivariance, where only a subset of the group satisfies the bound.

The paper extends this definition to an infinite-horizon discounted MDP $\mathcal{M}=(\mathcal{S},\mathcal{A},P,R,\gamma)$. A group $G$ acts on both states and actions, mapping $(s,a)$ to $(gs,ga)$. The MDP is $(G,\epsilon_R,\epsilon_P)$-invariant if its reward and transition functions satisfy

$$\left|R(gs,ga)-R(s,a)\right|\leq\epsilon_R$$

and

$$d_{\mathscr{F}}\left(P(gs'\mid gs,ga),P(s'\mid s,a)\right)\leq\epsilon_P$$

for every $g\in G$. The first condition says that symmetric state-action pairs have approximately equal rewards. The second says that their next-state distributions are approximately transformed versions of one another. The distance $d_{\mathscr{F}}$ is an integral probability metric; i.e. a measure of distance between probability distributions.

The main theoretical result shows that approximate symmetry of the MDP produces approximate invariance of its optimal value functions. If rewards are bounded and $0\leq\gamma<1$, then

$$\left|Q^*(s,a)-Q^*(gs,ga)\right|\leq\alpha$$

and

$$\left|V^*(s)-V^*(gs)\right|\leq\alpha,$$

where

$$\alpha=\frac{\epsilon_R+\gamma\rho_{\mathscr{F}}(V^*)\epsilon_P}{1-\gamma}.$$

Approximately Equivariant Actor-Critic

The main change from exactly equivariant RL is to replace the exactly equivariant group or steerable convolution layers in the encoder, actor, and critic with relaxed group convolutions.
An exact group convolution is defined as

$$(f\star\psi)(g)=\sum_{h\in G}\psi(g^{-1}h)f(h).$$

Because the kernel depends only on the relative group element $g^{-1}h$, the same weights are shared across symmetry-related inputs, which enforces exact equivariance. A relaxed group convolution instead uses $L$ exactly equivariant kernels and transformation-dependent relaxed weights:

$$(f\widetilde{\star}\psi)(g) =\sum_{h\in G}f(h)\sum_{l=1}^{L}w^l(h)\psi^l(g^{-1}h).$$

Each kernel $\psi^l$ retains the equivariant structure, while the weights $w^l(h)$ are allowed to depend on the specific group element $h$. This breaks strict weight sharing and allows the layer to learn different responses for symmetry-related inputs. If the relaxed weights remain equal across group elements, the layer behaves like an exactly equivariant layer; differences between them represent learned symmetry breaking.

The agent is organized as follows:

• the encoder $E$ is approximately equivariant and maps the observation to a latent state $z$
• the latent state transforms according to copies of the regular representation of $G$, which permutes its group-indexed channels when the input is transformed
• the policy $\pi$ is approximately equivariant, so transforming the state approximately produces the corresponding transformed action
• the action transformation is specified separately for each environment because different action components can respond differently to rotations or reflections
• the critic is approximately invariant and produces a scalar $Q(s,a)$ that should remain approximately unchanged when both the state and action are transformed