Reflex: Reinforcement Learning with Reflection Symmetry Exploitation in State-Based Continuous Control

Contributions

• presents the theory behind symmetries in RL, group invariant MDPs and show that Bellman operator, value function and policy are all invariant under symmetry
• presents the 2 types of symmetries, axial and bilateral
• suggest a framework to exploit symmetry in the MDP for on-policy (PPO) and off-policy (SAC)

Group-Invariant MDPs

In short, the state-action pair and the transformed/equivalent state-action pair result in the same reward and in the transformed next state.

Under group-invariant MDPs, we have:

• equivariance of the Bellman operator
  $\mathcal{T}(V \circ g) = (\mathcal{T}V) \circ g$

• invariance of optimal value functions
  $V^{*}(s)=V^{*}(gs)$
  $Q^{*}(s,a)=Q^{*}(gs, ga)$

  • equivariance of optimal policies
    $\pi^{*}(a|gs) = \pi^{*}(g^{-1}a|s)$
    $\mu^{*}(gs) = g\mu^{*}(s)$

Reflection Symmetries

Axial Reflection

Transforms the elements of a state-action vector element-wise into an equivalent state-action vector based on the table below.

Bilateral Reflection

Case where the axis of reflection coincides with the agent's intrinsic symmetry axis. For example, the right-left sides of a humanoid. Mathematically it can be formulated as
$g_b \, . \, (s,a) = \sigma (g_as, g_aa)$
where $g_a$ denotes the axial reflection based on the element type, and $\sigma$ swaps the right and left components
$\sigma (s_{left}, s_{right}, a_{left}, a_{right}) \rightarrow (s_{right}, s_{left}, a_{right}, a_{left})$

Reflex PPO

For PPO, they introduce a reflection symmetry regularization that explicitly enforces equivariance in the learned policy:
$\mathcal{L}_{\pi}^{sym}(\theta) = \mathbb{E}_{s_t \sim \mathcal{T}, \; g \sim G} || \pi_{\theta}(gs_t) - g(\pi_{\theta}(s_t)) ||_2^2$
Intuitively, this regularization promotes information sharing between symmetric components, allowing the policy to leverage mirrored experiences without additional environment interaction.

They additionally regularize the critic by enforcing value consistency:
$\mathcal{L}_{V}^{sym}(\phi) = \mathbb{E}_{s_t \sim \mathcal{T}, \; g \sim G} \Big(V_{\phi}(gs_t) - V_{\phi}(s_t) \Big)^2$
The overall training objective is then the combination of the original PPO loss and the symmetry regularizations above:
$\mathcal{L}_{total}=\mathcal{L}_{\pi} + \mathcal{L}_{V} + w_t (\mathcal{L}_{\pi}^{sym}+\mathcal{L}_{V}^{sym})$
where $w_t$ controls the strength of the symmetry regularization compared to the original PPO loss. This parameter decays linearly over time:
$w_t=w_0 \Big(\frac{1-t}{T} \Big)$

Reflex SAC

For SAC, the idea is to incorporate symmetry in the critic only.

First a symmetric target is computed, which enforces consistency of value estimates across mirrored state-action pairs:

• sample a mini-batch for training
• calculate its critic target value $y=r(s,a) + \gamma V(s')$
• apply the symmetry transformation to obtain the mirrored next state $\tilde{s}'=gs'$
• calculate the reflected target $\tilde{y}=r(s,a) + \gamma V(\tilde{s}')$
• calculate the final target as the average of the original and reflected targets $\bar{y} = \frac{1}{2}(y + \tilde{y})$

Then a symmetry-regularized critic loss is defined:
$\mathcal{L}_Q = \mathbb{E}_{(s,a) \sim \mathcal{B}} \Big[ \big(Q_{\phi}(s,a)-\bar{y} \big)^2 + \big(Q_{\phi}(gs,ga)-\bar{y} \big)^2 \Big]$