[Paper Review] Equivariant $Q$ Learning in Spatial Action Spaces

Dian Wang, Robin Walters+2 more|arXiv (Cornell University)|Oct 28, 2021

Robot Manipulation and Learning39 references4 citations

TL;DR

This paper proposes equivariant Q-learning for spatial action spaces by leveraging SE(2) symmetry in visual states and actions, using steerable convolutional networks to enforce equivariance. The method achieves significantly higher sample efficiency than standard deep Q-learning, with state-of-the-art performance on robotic manipulation tasks including real-world deployment.

ABSTRACT

Recently, a variety of new equivariant neural network model architectures have been proposed that generalize better over rotational and reflectional symmetries than standard models. These models are relevant to robotics because many robotics problems can be expressed in a rotationally symmetric way. This paper focuses on equivariance over a visual state space and a spatial action space -- the setting where the robot action space includes a subset of $ m{SE}(2)$. In this situation, we know a priori that rotations and translations in the state image should result in the same rotations and translations in the spatial action dimensions of the optimal policy. Therefore, we can use equivariant model architectures to make $Q$ learning more sample efficient. This paper identifies when the optimal $Q$ function is equivariant and proposes $Q$ network architectures for this setting. We show experimentally that this approach outperforms standard methods in a set of challenging manipulation problems.

Motivation & Objective

• To identify conditions under which the optimal Q-function is equivariant under SE(2) transformations (rotations and translations).
• To design neural network architectures that enforce SE(2) equivariance in the Q-function for spatial action spaces.
• To develop partially equivariant models that accommodate partial symmetry in state variables common in robotics.
• To empirically evaluate the proposed method against non-equivariant baselines on challenging visuo-motor control tasks.
• To demonstrate generalization and sample efficiency improvements in both simulation and real-world robotic deployment.

Proposed method

• Formalize the conditions under which the optimal Q-function is equivariant under finite subgroups of SE(2), ensuring that transformations in the visual state induce corresponding transformations in the action space.
• Design Q-network architectures using steerable convolutional layers that explicitly encode SE(2) equivariance, reducing free parameters and improving inductive bias.
• Introduce a modular architecture where different components (q1, q2, q3–q5) can be independently equivariant, deictic, or conventional, enabling flexible design for partial symmetry.
• Implement and compare multiple configurations: Equi+Equi+Equi, Equi+Deic+Deic, and baselines like Conv+Conv+Conv, using ablation studies to isolate performance gains.
• Apply the method to dense pixel-action space learning with fully convolutional networks (FCN), extending standard approaches with equivariance in both state and action heads.
• Use data augmentation (e.g., random rotations, translations) and contrastive learning as baselines for comparison, while emphasizing that the proposed method hardcodes symmetry into the architecture.

Experimental results

Research questions

• RQ1 Under what conditions is the optimal Q-function equivariant under SE(2) transformations in visual state and spatial action spaces?
• RQ2 How can neural network architectures be designed to enforce SE(2) equivariance in Q-learning for robotic manipulation?
• RQ3 What is the impact of partial equivariance—where only some state variables are symmetric—on policy learning performance?
• RQ4 How does equivariant Q-learning compare to non-equivariant and data-augmented baselines in terms of sample efficiency and final performance?
• RQ5 Can equivariant Q-learning generalize effectively from simulation to real-world robotic control tasks?

Key findings

• The optimal Q-function is SE(2)-equivariant when the environment's dynamics and reward are invariant under rotations and translations in the visual state space.
• Equivariant Q-networks outperform standard non-equivariant models in sample efficiency, achieving higher success rates with fewer environment interactions across all tested tasks.
• In the House Building and Bin Packing tasks, the Equi+Equi+Equi architecture achieved 100% success rate in simulation, significantly outperforming the Conv+Conv+Conv baseline.
• In real-world Bottle Arrangement and Box Palletizing tasks, the equivariant model with deictic encoding in q3–q5 achieved 90% and 85% success rates respectively, demonstrating strong sim-to-real transfer.
• The runtime cost of equivariant inference is higher (0.72s per step for Equivariant FCN vs. 0.08s for conventional FCN), but the performance gain justifies the overhead in sample efficiency.
• The ablation study shows that replacing conventional layers with equivariant or deictic components improves performance, with the best results achieved when both q1 and q2 are equivariant, and q3–q5 use deictic encoding.