PaperlayerView exact revision on arXiv

Source-linked AI summary

$\mathrm{SO}(2)$-Equivariant Reinforcement Learning

Dian Wang, Robin Walters, Robert Platt

arXiv:2203.04439v1cs.RO

TL;DR

Visual reinforcement learning can require costly environmental interactions, and data augmentation relies on assumptions about transformed dynamics. This paper derives symmetry properties of optimal solutions and builds equivariant DQN and SAC methods for robotic manipulation. The authors report that these methods substantially outperform competitive baselines, with Equivariant SAC particularly strong.

  • Problem

    Improving sample efficiency matters in robotics, where collecting experience can require interaction with a physical environment, and prior equivariant reinforcement-learning evaluations were limited to toy settings.

    Source: §1 · p. 1

  • Method

    The paper defines group-invariant MDPs, derives invariant optimal Q-functions and equivariant optimal policies, and uses equivariant architectures in DQN and SAC.

    Source: §1 · p. 1, §4.1 · p. 4, §5.1 · p. 5

  • Results

    The robotic-manipulation experiments show that the proposed methods substantially surpass competitive baselines, with Equivariant DQN faster and higher-reward across three environments and Equivariant SAC outperforming other methods significantly.

    Source: §6.1 · p. 7, §6.2 · p. 8, §7 · p. 9

  • Takeaways & Limitations

    The authors suggest Equivariant SAC's performance could make reinforcement learning feasible for a larger class of robotics problems.

    Source: §1 · p. 2

    The policy formulation assumes action variance is invariant to the reference frame, a simplifying assumption that constrains the action-policy setup.

    Source: §5.2 · p. 5

Abstract

from arXiv · show

Equivariant neural networks enforce symmetry within the structure of their convolutional layers, resulting in a substantial improvement in sample efficiency when learning an equivariant or invariant function. Such models are applicable to robotic manipulation learning which can often be formulated as a rotationally symmetric problem. This paper studies equivariant model architectures in the context of $Q$-learning and actor-critic reinforcement learning. We identify equivariant and invariant characteristics of the optimal $Q$-function and the optimal policy and propose equivariant DQN and SAC algorithms that leverage this structure. We present experiments that demonstrate that our equivariant versions of DQN and SAC can be significantly more sample efficient than competing algorithms on an important class of robotic manipulation problems.

1 INTRODUCTION

The paper addresses sample efficiency for visual reinforcement learning, especially where environmental interaction is costly. It proposes equivariant architectures for robotic manipulation and reports that they outperform competitive data-augmentation methods.

1 INTRODUCTION

  • Improving sample efficiency matters in robotics because gathering experience can require interacting with a physical environment.Visual-state data augmentation commonly applies transformations such as image translations or rotations, implicitly assuming invariant transition and reward dynamics.Source: §1 · p. 1
  • Equivariant architectures encode desired symmetry directly in the model, whereas data augmentation modifies training data and requires the model to learn equivariance alongside the task.The paper argues that augmentation can require more training time and model capacity, while equivariant networks guarantee equivariance rather than only approximating it.Source: §1 · p. 1
  • The paper defines group-invariant MDPs and introduces equivariant DQN, SAC, and learning-from-demonstration variants for robotic manipulation.Its contribution also includes analyzing the invariance of optimal values and the equivariance of optimal policies.Source: §1 · p. 1
  • The proposed methods outperform recent competitive data-augmentation approaches, with Equivariant SAC showing a particularly dramatic advantage.The authors suggest this result could make reinforcement learning feasible for a larger class of robotics problems.Source: §1 · p. 2

2 RELATED WORK

Prior work applies equivariant networks, data augmentation, and contrastive learning to improve reinforcement learning or robotic manipulation. This paper extends equivariant learning to robotic visual-motor control, including continuous-action methods and learning from demonstration.

Equivariant Learning

  • Earlier reinforcement-learning uses of equivariant networks focused on toy game settings or a small set of tasks, while robotic-manipulation work learned equivariant Q-functions only for spatial actions.This paper adds Equivariant SAC, equivariant learning from demonstration, and a new Equivariant DQN variation for visual-motor control.Source: §2 · p. 2

Data Augmentation

  • Random-crop and random-translate augmentation can improve reinforcement-learning performance, but the paper characterizes augmentation as often less sample efficient than equivariant networks.In robotic manipulation, augmentation has also been used to generate additional samples.Source: §2 · p. 2

Contrastive Learning

  • Contrastive learning with image augmentation has improved sample efficiency, but it learns invariant feature encoders rather than equivariant functions.The passage cites applications of this approach in reinforcement learning and robotic manipulation.Source: §2 · p. 2

Close-Loop Robotic Control

  • Prior closed-loop robotic-control work includes grasping and more complex policies, and this paper applies equivariant learning to that action-space setting.Closed-loop actions control end-effector displacement rather than a target pose.Source: §2 · p. 2

3 BACKGROUND

The background formalizes planar rotations as group actions on feature maps and their channels. Equivariant layers preserve the prescribed transformation relationship, enabling networks to be built by composing such layers.

SO(2) and Cn

  • SO(2) is the group of continuous planar rotations, and C_n is its cyclic subgroup used to represent discrete rotations.The paper reasons about rotations using both the continuous group and this discrete subgroup.Source: §3 · p. 2

Cn actions

  • A group action specifies how group elements transform objects; under the trivial representation, scalar pixel values remain unchanged as image pixels rotate.Figure 1 illustrates this unchanged channel mapping for a one-channel feature map.Source: §3 · p. 2, Figure 1 · p. 3
  • The standard representation rotates vectors, while the regular representation cyclically permutes channel coordinates.The standard representation describes vector fields; Figure 1 illustrates the vector rotation and channel permutation.Source: §3 · p. 3

Feature maps as functions

  • Feature maps can be written as functions from spatial coordinates to channel values, and the group acts on those maps according to their feature semantics.A vector is treated as a feature map with one spatial location.Source: §3 · p. 3

Equivariant convolutional layer

  • An equivariant layer must transform its output consistently with the transformed input, according to the representations assigned to each feature map.The constraint applies to feature maps or vectors acted on by the group.Source: §3 · p. 3
  • A convolutional kernel must satisfy the equivariance constraint, and stacking such layers with equivariant nonlinearities yields an equivariant network.The cited kernel construction is given for the group C_n.Source: §3 · p. 3

4 PROBLEM STATEMENT

The paper defines group-invariant MDPs and shows that their optimal Q-functions are invariant while their optimal policies are equivariant. It specializes this framework to visual robotic manipulation, with rotated image coordinates and factored actions.

4.1 GROUP-INVARIANT MDPS

  • A G-invariant MDP has reward and transition functions unchanged when the group acts jointly on states and actions.The transition condition also transforms the next state by the same group element.Source: §4.1 · p. 4
  • For a group-invariant MDP, the optimal Q-function is invariant and the optimal policy is equivariant under the group action.The paper notes that this result follows from the connection to MDP homomorphisms, whose abstract solutions can be lifted to the original MDP.Source: §4.1 · p. 4

4.2 SO(2)-INVARIANT MDPS IN VISUAL STATE SPACES

  • For visual state spaces, the paper focuses on SO(2)-invariant MDPs and approximates SO(2) using its discrete subgroup C_n.The state is an image whose pixels rotate while pixel feature values remain unchanged.Source: §4.2 · p. 4
  • The robotic-manipulation setup uses a gripper-centered depth image and separates actions into rotationally equivariant xy displacement and invariant components.The invariant components are gripper aperture, height change, and orientation change; the setup assumes a C_n-invariant reward to obtain a C_n-invariant MDP.Source: §4.2 · p. 4
  • The paper illustrates the manipulation scene and its visual state space as the setting for this formulation.The caption identifies the two panels as the manipulation scene and the visual state space.Source: Figure 2 · p. 4

5 APPROACH

The approach builds rotational structure into DQN and SAC networks: DQN represents an invariant Q-function with an equivariant network, while SAC uses an equivariant actor and invariant critic. SAC from Demonstration adds expert demonstrations and an actor-loss term compatible with equivariance.

5.1 EQUIVARIANT DQN

  • The DQN Q-network uses equivariant layers to represent an invariant Q-function, with equivariant action variables mapped to Q-values for invariant action variables.For the robotic manipulation example, the equivariant action variables are xy positions and the invariant variables are gripper aperture, height, and rotation.Source: §5.1 · p. 5

5.2 EQUIVARIANT SAC

  • SAC models the optimal policy with an equivariant actor and the optimal Q-function with an invariant critic.The actor’s action standard deviation is assumed invariant; the critic uses a nonlinear maxpool over the group space to avoid additional undesired symmetries from a linear mapping.Source: §5.2 · pp. 5–6

5.3 EQUIVARIANT SACFD

  • The SACfD actor loss adds an L2 term for expert transitions, and the term is compatible with equivariance because sampled and expert actions transform equivalently.Training also pre-populates the replay buffer with expert demonstrations generated by a hand-coded planner.Source: §5.3 · p. 6

6 EXPERIMENTS

The experiments evaluate equivariant DQN and SAC on robotic manipulation tasks formulated as SO(2)-invariant MDPs. Across the reported settings, equivariant methods outperform baselines, including when learning from demonstrations and when training uses only one orientation.

6.1 EQUIVARIANT DQN

  • Equivariant DQN learns faster and converges to higher discounted reward in Block Pulling, Object Picking, and Drawer Opening.The tasks use the group C4, with expert demonstrations preloaded into the replay buffer.Source: §6.1 · pp. 6–7

6.2 EQUIVARIANT SAC

  • Equivariant SAC significantly outperforms other methods; Object Picking and Drawer Opening appear infeasible for baselines, while FERM also solves Block Pulling.The comparison uses SO(2) buffer augmentation and 20 expert demonstration episodes before training.Source: §6.2 · p. 8

6.3 EQUIVARIANT SACFD

  • Equivariant SACfD performs best on all four challenging tasks, and only the equivariant method solves the three hardest tasks.These experiments use SACfD and prioritized replay because the tasks are too difficult without guided exploration.Source: §6.3 · p. 8

6.4 COMPARING WITH LEARNING EQUIVARIANCE USING AUGMENTATION

  • Equivariant SACfD outperforms RAD and DrQ variants that apply SO(2) augmentation to both state and action.This comparison disables buffer augmentation for the online transitions in those augmentation baselines.Source: §6.4 · p. 9

6.5 GENERALIZATION EXPERIMENT

  • Equivariant SACfD generalizes from a single training orientation to random evaluation orientations, while none of the baselines do.SO(2) buffer augmentation is disabled in this experiment.Source: §6.5 · p. 9

7 DISCUSSION

The paper defines group-invariant MDPs, characterizes the optimal solutions, and reports that its proposed methods substantially surpass competitive baselines in robotic manipulation. Its applicability is bounded by the required invariance assumptions.

7 DISCUSSION

  • The paper defines group-invariant MDPs and identifies invariance and equivariance properties of their optimal solutions.It proposes Equivariant SAC and a new Equivariant DQN variation for continuous and discrete action spaces, respectively.Source: §7 · p. 9
  • The methods substantially surpass competitive baselines in the evaluated robotic manipulation domains.Source: §7 · p. 9
  • The MDP definition requires invariant rewards and transitions, limiting applicability in domains such as some ATARI games.Non-top-down observations or non-equivariant structures such as the robot arm also do not directly satisfy the invariant assumptions.Source: §7 · p. 9

A PROOF OF PROPOSITION 4.1

The proof establishes that the optimal Q-function is invariant under the group action, while the section also states the corresponding optimal-policy equivariance relation and develops consequences for equivariant maps and robotic manipulation settings.

A PROOF OF PROPOSITION 4.1

  • The transformed optimal Q-function satisfies the same Bellman equation, so uniqueness gives Q*(s, a) = Q*(gs, ga).The proof first re-indexes states under the group action and uses reward and transition invariance.Source: A PROOF OF PROPOSITION 4.1 · p. 13
  • The policy proof uses Q-function invariance to substitute Q*(gs, ā) with Q*(s, g^-1ā) when comparing policies at related states.The stated target relation is π*(gs) = gπ*(s).Source: A PROOF OF PROPOSITION 4.1 · p. 13

B EQUIVARIANCE OVERCONSTRAIN

  • A linear Cn-equivariant map from two regular representations to a trivial representation factors through their trivial components, yielding sums of each input’s coordinates.The section also states that these maps are Cn × Cn-equivariant.Source: B EQUIVARIANCE OVERCONSTRAIN · p. 14

C ENVIRONMENT DETAILS

  • The manipulation environments use a 0.4m × 0.4m × 0.24m workspace, random object positions and orientations, and sparse rewards; tasks include pulling, picking, opening, stacking, and building.The gripper starts from a fixed configuration; Corner Picking uses compliant joints that let the gripper apply force on the block.Source: C ENVIRONMENT DETAILS · p. 14

D.1 EQUIVARIANT DQN ARCHITECTURE

  • The Equivariant DQN uses a seven-layer Steerable CNN over C4, mapping its input to an 18-channel 3 × 3 output whose channels encode invariant actions and spatial positions encode xy actions.The models are implemented with E2CNN and PyTorch; a rotation-invariant binary channel indicates whether the gripper is holding an object.Source: §D.1 · p. 14, D NETWORK ARCHITECTURE · p. 14, C ENVIRONMENT DETAILS · p. 14

D.2 EQUIVARIANT SAC ARCHITECTURE

The Equivariant SAC uses separate steerable actor and critic networks over C8, with representation choices for action components and a pooling operation intended to prevent equivariant overconstraint.

D.2 EQUIVARIANT SAC ARCHITECTURE

  • The C8 actor maps the state to a mixed action representation, while the critic combines state and action features and produces two Q estimates.The actor has eight layers and the critic nine; max pooling transforms regular representations into trivial ones to prevent equivariant overconstraint.Source: §D.2 · pp. 14–15
  • The architecture diagram distinguishes regular-representation layers, marked R, from trivial-representation layers, marked T.The figure depicts Equivariant DQN and Equivariant SAC; ReLU nonlinearities are omitted.Source: Figure 12 · p. 15

E BASELINE DETAILS

  • RAD Crop, CURL, and FERM use random crops, while DrQ Shift uses random shifts; FERM also pretrains its contrastive encoder on expert data.The crop baselines resize 142 × 142 images to 128 × 128, and DrQ uses two augmentations for both target and loss calculations.Source: E BASELINE DETAILS · p. 15

F TRAINING DETAILS

The training setup specifies shared simulator and target-update settings, alongside separate optimization and replay configurations for DQN and SAC-family experiments.

F TRAINING DETAILS

  • Both DQN and SAC use soft target updates with τ = 10^-2, and training uses five parallel environments.The PyBullet visual state is 128 × 128, except crop-based baselines that begin at 142 × 142; the workspace is 0.4m × 0.4m × 0.24m.Source: F TRAINING DETAILS · p. 15
  • DQN uses Adam with a 10^-4 learning rate, Huber TD loss, γ = 0.95, batch size 32, and a 100,000-transition replay buffer.Source: F TRAINING DETAILS · p. 15
  • The baseline CNNs have parameter counts similar to the equivariant models: DQN has 3.9M versus 2.6M, and SAC has 2.6M versus 2.3M.Source: Figure 13 · p. 16
  • SAC and SACfD use Adam at 10^-3, γ = 0.99, batch size 64, and a 100,000-transition buffer; SACfD adds prioritized replay and an expert-transition priority bonus.The prioritized replay exponents are 0.6 and β0 = 0.4, with expert-transition bonus ϵd = 1.Source: F TRAINING DETAILS · p. 16

G ADDITIONAL EXPERIMENTAL RESULTS FOR EQUIVARIANT SACFD

Additional experiments report Equivariant SACfD outperforming all baselines in Block Pulling and Object Picking, and outperforming SO(2)-augmented RAD and DrQ in those environments.

G ADDITIONAL EXPERIMENTAL RESULTS FOR EQUIVARIANT SACFD

H.1 USING EQUIVARIANT NETWORK ONLY IN ACTOR OR CRITIC

Across these ablations, using equivariant networks in both actor and critic performs best, while larger tested symmetry groups generally perform better. The performance gain is less significant in non-symmetric tasks, and rotational buffer augmentation does not make the baselines outperform the proposal.

H.3 EQUIVARIANT SACFD IN NON-SYMMETRIC ENVIRONMENTS

  • With fixed initial orientations in both training and evaluation, equivariant SACfD's performance gain is less significant in non-symmetric tasks.The plots report discounted reward, evaluated every 500 training steps and averaged over four runs; shading denotes standard error.Source: Figure 17 · p. 18, §H.3 · p. 18

H.4 ROTATIONAL AUGMENTATION + BUFFER AUGMENTATION

  • None of the rotational-augmentation baselines with buffer augmentation outperforms the proposal in any task.The buffer adds four randomly SO(2)-rotated transitions per new transition; it hurts RAD and DrQ compared with the results without buffer augmentation, which the passage attributes to redundant augmentation.Source: Figure 18 · p. 18, §H.4 · p. 19

H.1 USING EQUIVARIANT NETWORK ONLY IN ACTOR OR CRITIC

  • Using equivariant networks in both actor and critic always demonstrates the best performance, and actor-only use generally helps more than critic-only use in 5 of 6 experiments.The four variations combine equivariant or conventional CNN networks in each of the actor and critic.Source: §H.1 · pp. 17–18

H.2 DIFFERENT SYMMETRY GROUPS

  • C8 generally outperforms C4, followed by C2, in the SACfD comparison of symmetry groups.Source: §H.2 · p. 18

Related papers

Loading 2203.04439v1…