Partially Equivariant Reinforcement Learning in Symmetry-Breaking Environments

AI summary

Overview

• Research area: Reinforcement learning (RL) with group-symmetry inductive biases, spanning discrete control (Grid-World) and continuous control (MuJoCo locomotion, Fetch Reach, and a UR5e manipulator).
• Technical level: Advanced. The paper combines MDP theory (Bellman operators, contraction bounds, group-invariant MDPs) with deep RL implementation (DQN, SAC, equivariant network design).
• Scope: The paper introduces a Partially Group-Invariant MDP (PI-MDP) framework and two algorithms (PE-DQN and PE-SAC) that apply equivariance only where symmetry actually holds in the environment, and evaluates them against vanilla, strictly equivariant, and approximately equivariant baselines. (arXiv:2512.00915v2 [cs.LG], Machine Learning category, CC BY 4.0; project page: https://pranaboy72.github.io/perl_page/)

What This Paper Is About

Most equivariant RL methods assume a fully group-invariant MDP, where rewards and transitions are unchanged under the symmetry group. Real environments rarely satisfy this: dynamics, actuation limits, obstacles, contacts, joint limits, and reward design break symmetry, often only in localized regions. The paper shows that even local symmetry violations create one-step Bellman backup errors that propagate across the whole state–action space and are amplified by a factor of 1/(1−γ), producing global value errors; the goal is to keep equivariance's sample-efficiency benefits while avoiding that global error.

Key Contributions

1. An analysis of how local symmetry violations become global value error. The authors derive a one-step Bellman error bound between the true MDP and its group-averaged approximation, then lift it to a global bound on the optimal value functions.
2. The Partially Group-Invariant MDP (PI-MDP) framework. A gating function λ(s,a) ∈ [0,1] interpolates between the group-invariant MDP and the true MDP per state–action pair, with a proof that the resulting operator is a γ-contraction with a unique fixed point and a bound on its distance from the true optimum.
3. Practical algorithms: PE-DQN and PE-SAC. These implement the framework inside standard DQN and SAC, using a learned gate that detects symmetry-breaking from disagreement between an equivariant and an unconstrained one-step predictor, plus gated critic and policy parameterizations.
4. Experiments across discrete and continuous control. Grid-World, MuJoCo locomotion, Fetch Reach, and UR5e Reach results show the methods retain equivariance's sample efficiency in symmetric regions and remain robust as symmetry-breaking increases, outperforming strict and approximate-equivariant baselines.

Main Findings

• Local violations produce global errors: The paper proves (Lemma 1) that the one-step Bellman error at a state–action pair is bounded by ϵ_R(s,a) + 2γ‖V_Q‖_∞ ϵ_P(s,a), where ϵ_R is the absolute reward difference and ϵ_P is the total-variation distance between next-state kernels of the true and group-averaged MDPs. Proposition 1 lifts this to ‖Q*N − Q*E‖∞ ≤ 1/(1−γ)‖δ‖∞, showing amplification by (1−γ)^(−1).
• Gating bounds the damage: Corollary 1 gives ‖Q*H − Q*N‖∞ ≤ 1/(1−γ)‖(1−λ)δ‖∞. The right-hand side is zero when λ(s,a) = 1 everywhere symmetry breaks, or when the group-invariant MDP coincides with the true MDP at that pair.
• Grid-World, increasing obstacles: With no symmetry-breaking factors, PE-DQN converges quickly to λ ≈ 0 and behaves like a purely equivariant agent, matching strictly equivariant DQN. As obstacles are added, strictly equivariant DQN degrades much more rapidly than other baselines, approximately equivariant methods gain only minor improvements over vanilla DQN, and PE-DQN maintains strong performance across all obstacle counts. Results are averaged over 100K steps with five random seeds.
• Reward-level symmetry-breaking: Using 10- and 30-obstacle layouts where half the obstacles are passable but return a negative reward (transitions unchanged, reward altered), PE-DQN achieves the highest returns. RPP-DQN is the strongest baseline but trails in both sample efficiency and final performance.
• Complex stochastic dynamics: In a 40-obstacle layout with stochastic transitions, random slips partly mask transition-level symmetry-breaking, letting strictly equivariant and approximately equivariant DQN recover reasonable performance; PE-DQN still attains the best returns and learning speed.
• Continuous control: PE-SAC consistently outperforms all baselines across MuJoCo locomotion, Fetch Reach, and UR5e Reach. Locomotion results are averaged over 1M training steps with eight random seeds; Fetch Reach at 30K steps and UR5e Reach at 500K steps with five random seeds. The paper notes that RPP was re-run from official code and that discrepancies with reported numbers arise because RPP reports "max over steps" rather than average performance. (Specific return values are not reported in the available excerpt; the discussion of Fig. 5 is truncated.)
• Hard gating is the practical choice: Although the theory holds for any measurable λ: S × A → [0,1], the algorithms use binary gating (λ ∈ {0,1}) for simplicity and empirically more stable training.
• State gate versus sampled-max gate: The sampled-max variant, which approximates max_a λ_ω(s,a) by maximizing over K sampled actions, performs similarly to the learned state gate λ_ζ(s) for K ∈ {4,8}, but the learned state gate is slightly more robust when symmetry-breaking is sparse and easy to miss with few sampled actions.
• Architecture detail: Separate trunks for critics, policy, and one-step predictors gave stable training and clean separation between equivariant and non-equivariant components; trunk-sharing variants did not consistently improve performance and sometimes harmed stability.

Methodology in Plain English

The authors start from two MDPs: the true environment and a "group-invariant" version built by averaging the true rewards and dynamics over the symmetry group (using the normalized Haar measure, which is the uniform measure for finite groups). The difference between them is measured pointwise by a reward discrepancy and a transition-kernel discrepancy.

To decide where the environment actually respects symmetry, they train two one-step predictors: one constrained to respect the group symmetries, one unconstrained. Where symmetry genuinely holds, both predictors should agree; where it breaks, the equivariant predictor can only represent the group-averaged surrogate while the unconstrained one can fit the true kernel, so their predictions diverge. The disagreement score is tracked with running statistics, and high-disagreement samples are treated as the upper tail of the distribution and labeled as symmetry-breaking. A gating network is trained with binary cross-entropy on those pseudo-labels and is not updated by RL gradients (stop-gradient).

The critic is then a gated mixture of an equivariant Q-network and an unconstrained one, with binary gating making it a hard switch. The policy uses a state-only gate and a product-of-experts blend of the equivariant and unconstrained policies, which arises from SAC policy improvement; the state gate is aligned with the state–action critic gate through an expectile loss with τ → 1, approximating a max over actions so that any broken action at a state activates the unconstrained policy. Training uses standard DQN and SAC objectives with these gated parameterizations, so no new optimizer or algorithm scaffold is required.

Why This Matters

• Impact on research: The paper reframes symmetry-breaking from "relax the architecture globally" to "decide locally where equivariance is valid," and gives a quantitative propagation bound that explains why strictly equivariant backups fail. It connects approximate-equivariance work in representation learning to the Bellman backup itself.
• Real-world applications:
  • Robotic manipulation arms (the paper's UR5e Reach task with SE(3) pose goals), where contacts, floor collisions, and kinematic singularities break rotational symmetry.
  • Mobile robot navigation and path planning around obstacles, as studied in the Grid-World tasks with increasing obstacle counts.
  • Legged locomotion control subject to external forces and reward perturbations, as in the MuJoCo locomotion benchmarks.
  • Industrial automation and assembly, where actuation limits and fixture geometry make the environment only locally symmetric.
• Industry relevance: Sample efficiency directly translates to fewer physical robot trials and lower data-collection cost. A method that keeps symmetric-region efficiency while degrading gracefully as real-world asymmetry grows is more deployable than either vanilla RL or strict equivariance, and being built on standard DQN and SAC makes it easier to integrate into existing training pipelines.

Future Directions

• Continuous rather than binary gating: The theory covers any measurable λ: S × A → [0,1], but the implementation uses hard gating; exploring soft gating with the same stability properties is a natural next step.
• Better state-only policy gates: The authors note that a full state–action gate in the policy would be theoretically appealing but intractable because the normalization constant of the product-of-experts would depend on the action. Improving this approximation remains open.
• Detecting sparse or partially masked symmetry-breaking: Random stochastic slips partly hid transition-level violations in the 40-obstacle Grid-World, and the authors observe that the learned state gate is more robust than sampled-max when violations are sparse and easy to miss.
• Extending beyond the studied groups and settings: The experiments cover C4 (Grid-World), SO(3) (reach tasks), and SE(3) (UR5e pose goals) with state-based observations; whether the disagreement-based gate scales to larger groups, higher-dimensional dynamics, and image-based inputs is not addressed.

Target Audience

Researchers and graduate students working on equivariant or symmetry-aware reinforcement learning and robot learning; practitioners building sample-efficient RL systems for manipulation or locomotion who need robustness when ideal symmetry assumptions fail; and readers interested in the theory linking MDP homomorphisms, Bellman contraction bounds, and approximate equivariance. Comfort with Bellman operators, group representations, and off-policy deep RL (DQN, SAC) is assumed.

Authors’ abstract

Group symmetries provide a powerful inductive bias for reinforcement learning (RL), enabling efficient generalization across symmetric states and actions via group-invariant Markov Decision Processes (MDPs). However, real-world environments almost never realize fully group-invariant MDPs; dynamics, actuation limits, and reward design usually break symmetries, often only locally. Under group-invariant Bellman backups for such cases, local symmetry-breaking introduces errors that propagate across the entire state-action space, resulting in global value estimation errors. To address this, we introduce Partially group-Invariant MDP (PI-MDP), which selectively applies group-invariant or standard Bellman backups depending on where symmetry holds. This framework mitigates error propagation from locally broken symmetries while maintaining the benefits of equivariance, thereby enhancing sample efficiency and generalizability. Building on this framework, we present practical RL algorithms -- Partially Equivariant (PE)-DQN for discrete control and PE-SAC for continuous control -- that combine the benefits of equivariance with robustness to symmetry-breaking. Experiments across Grid-World, locomotion, and manipulation benchmarks demonstrate that PE-DQN and PE-SAC significantly outperform baselines, highlighting the importance of selective symmetry exploitation for robust and sample-efficient RL. Project page: https://pranaboy72.github.io/perl_page/