$\mathrm{SO}(2)$-Equivariant Reinforcement Learning
Authors: Dian Wang
Robin Walters
Robert Platt
Venue: ICLR 2022 Spotlight (2022)
Keywords: Reinforcement Learning
Equivariance
Robotic Manipulation

Abstract: 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.

Published as a conference paper at ICLR 2022
-QUUA-TQUI REINFORCEMENT LEARNING
Dian Wang, Robin Walters, and Robert Platt Khoury College of Computer Sciences
Northeastern University Boston, MA 02115, USA
{wang. dian, r. w.wattess, r. e.platt/enorttesssern..e
ABSTRACT
Equivariant neural networks enforce symmetry within the structure of their con- volutional layers, resulting a substantial improvement in sample efficiency when learning equivariant invariant function. Such models applicable to robotic manipulation learning which can often be formulated as a rotationally symmetric problem. This paper studies equivariant model architectures in the con- text 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.
INTRODUCTION
A key challenge in reinforcement learning is to improve sample efficiency that is reduce the amount of environmental interactions that an agent must take in order to learn a good policy. This particularly important in robotics applications where gaining experience potentially means inter- acting with physical environment. One way of improving sample efficiency is to create "artificial" experiences through data augmentation. This is typically done visual state spaces where an affine transformation (e.g.. translation or rotation of the image) is applied to the states experienced during a transition (Laskin et al., 2020a; Kostrikov et al. 2020). These approaches implicitly assume that the transition and reward dynamics of the environment are invariant to affine transformations of the visual state. In fact, some approaches explicitly use a contrastive loss term to induce the agent to learn translation-invariint feature representations (Laskin et al. 2020b Zhan et al. 2020).
Recent work in geometric deep learning suggests that it may be possible to learn transformation- invariant policies and value functions in different way, using equivariant neural networks (Cohen & Welling (0166a]b) (). The key idea is to structure the model architecture such that it constrained only represent functions with the desired invariance properties. In principle, this approach aim exactly the same thing as the data augmentation approaches described above both methods seek improve sample efficiency by introducing an inductive bias. However, the equivariance approach achieves this more directly by modifying the model architecture rather than by modifying the train- ing data. Since with data augmentation, the model must learn equivariance in addition to the task itself, more training time and greater model capacity are often required. Even then, data augmenta- tion results only in approximate equivariance whereas equivariant neural networks guarantee it and often have stronger generalization as well (Wang al.) 2020b). While equivariant architectures have recently been applied reinforcement learning (van der Pol al.) 2020a b. Mondal et al., 2020), this has been done only in toy settings (grid worlds, etc.) where the model is equivariant over small finite groups, and the advantages of this approach over standard methods is less clear.
This paper explores the application of equivariant methods to more realistic problems in robotics such as object manipulation. We make several contributions. First, we define and analyze an im- variation
portant class of MDPs that call group -ivariii MDPs. Second, we introduce of the Equivariant DQN (Mondal et al.) 2020). and we further introduce equivariant variations of SAC (Haarnoja et al. 2018), and learning from demonstration (Lff). Finally, we show that

---

Published as a conference paper at ICLR 2022
methods convincingly outperform recent competitive data augmentation approaches (Lskin et al. 2020a Kostrikov et al. 2020; Laskin et al. 2020b; Zhan et al., (020). Our Equivariant SAC method, in particular, outperforms these baselines dramatically (Figure Z that it could make reinforcement learning feasible for a much larger class of robotics problems than is currently the case. Supplemen- tary video and code available at httpt::: /p ntw wtt..github. io/equi r Baapuuiiir_ag__///_
2 RELATED WORK
Equivariant Learning: Encoding symmetries in the structure of neural networks can improve both generalization and sample efficiency, The idea of equivariant learning is first introduced in G- Convolution (Cohen & Welling 20) The extension work proposes an alternative architecture, Steerable CNN (Cohen & Welling 2016b). Weiler & Cesa (2019) proposes general framework for implementing E(2)-Steerable CNNs. the context of reinforcement learning, MMndal et al. (2020) investigates the use of Steerable CNNs in the context of two game environments. vvan der Pol et al. 22020b) proposes MDP homomorphic networks to encode rotational and reflectional equivariance of an MDP but only evaluates their method in a small set of tasks. In robotic manipulation, Wang al. (20TT learns equivariant Q-functions but is limited in the spatial action space. In contrast to prior work, this paper proposes Equivariant SAC algorithm, an equivariant L.f algorithm, and novel variation of Equivariant DQN (Mondal et al. 2020) focusing on visual motor control problems.
Data Augmentation: Another popular method for improving sample efficiency is data augmenta- tion. Recent works demonstrate that the use of simple data augmentation methods like random crop or random translate can significantly improve the performance of reinforcement learning (Laskin et al. 2020a, Kostrikov et al. 2020). Data augmentation is often used for generating additional sam- ples
Kalashnikov al. [018 Lin al. 2020 Zeng al., (020) in robotic manipulation. However, data augmentation methods are often less sample efficient than equivariant networks because the latter injects an inductive bias to the network architecture.
Contrastive Learning: Data augmentation is also applied with contrastive learning (Oord et al. 2018 8) to improve feature extraction. [askin et al.) (2020b) show significant sample-effciency improvement by adding an auxiliary contrastive learning term using random crop augmentation. Zhan et al. (2020)
use a similar method in the context of robotic manipulation. However, contrastive learning is limited to learning invariant feature encoder and is not capable of learning equivariant functions.
There are two typical action space definitions when learning policies that control the end-effector of robot arm: the spatial action space that controls the target pose of the end-effector ZZng et al. 2018b a Satish et al. 2019 Wang et al. 2020a), or the close-loop action space that controls the displacement of the end-effectur effector. The close- loop action space is widely
Close-Loop Robotic Control:
used for learning grasping policies (Kalashnikov et al. 2018 Quillen et al. [018 Breyer et al.) 2019
James al. 2019). Recently, some works also learn more complex policies than grasping ( (Viereck et al. 2020 Kilinc et al.) 2019: Cabi et al.) 2020; Zhan et al. (2020). This work extends prior works the close-loop action space by using equivariant learning to improve the sample efficiency.
3 BACKGROUND
SO(2) and C₁: We will reason about rotation in terms of the group SO(2) and its cyclic subgroup C, S (0))). SO(2) is the group of continuous planar rotations (Rot : < < 2 Cn is the discrete subgroup C, = {Rots: : 0 E { 2 00 < < n}} of rotations by multiples 2x. n
On actions: A group G may be equipped with an action on set X by specifying map G X satisfying 91 (92 x) = (g)g2)) and x = x for all 91 92 E G, x E X. Note that closure, gz E X,, and invertibility, g gr x, follow immediately from the definition. We are interested in actions of C, which formalize how vectors or feature maps transform under rotation. The group Cn acts in three ways that concern us (for a more comprehensive background, see Bronstein et al. 2021
1. R through the trivial representation Pu- Let g Cn and R. Then Po(g)æ x. For example, the trivial representation describes how pixel color/depth values change when an image is rotated, L.e. they do not change T
2

---

Published as a conference paper at ICLR 2022
2
Figure 1: Illustration of how an element g E C2, acts on the feature map by rotating the pixels and transforming the channel space through Po. Pt. and Preg. Left: Cn acts on the channel space of 1-channel feature map by identical mapping. Middle: C, acts on the channel space of vector field by rotating the vector at each pixel through P1. Right: C, acts on the channel space of a 4-channel feature map by permuting the order of the channels by Preg.
2. R through the standard representation Pt. Let E Cn and R Then pi(g)o -sin sing cosg )u. This describes how elements of vector field change when rotated (Figure mid- dle).
R" through the regular representation Preg- Let g = pm E C. = [1, r and (r), 22,........ Fnn) R", Then Preg (g) (In (In-m=1:--. Ln, U1, I22 ................................ In In-m) eyclically per- mutes the coordinates of IR" (Figure right).
Feature maps functions: In deep learning, images and feature maps typically expressed as lensors. However, it will be convenient here to sometimes express these as functions. Specifically, may write h one-channel image F x describes the intensity at pixel x, y. Similarly, an m-channel tensor F E mxxxxx may be written as F: R + Rm We refer the domain of this function its "spatial dimensions".
function F R - R where F(x, y)
On actions on vectors and feature maps: C₁ acts on vectors and feature maps differently depending upon their semantics. We formalize these different ways of acting as follows. Let F: R - IRm be an m-channel feature map and V mmxx1x = R" vector represented as a special case of feature map with 1 x spatial dimensions. Then y is defined to act on J. by
(gF) (p))))
For vector V (considered to be at (x, = (0, 0)), this becomes:
gV Pg(gV
In the above, P1 (g) rotates pixel location and pو(g) transforms the pixel feature vector using the trivial representation (p, = Po). the standard representation (p, = P1). the regular representation (pj Preg), some combination thereof.
Equivariant convolutional layer: A C, -n-equivariant layer is a function h whose output constrained to transform in defined way when the input feature map is transformed by group action. Consider an equivariant layer h with an input Fin R Ipl and output F.nt. R2 + IR Powel where Pin and Pout denote the group representations associated with Fin and Fout, respectively. When the input is transformed, this layer is constrained to output transformed version of the same output feature map:
h(gFm) g(h(Fim) = gFout.
(3)
where g E Cn acts on Fin or Fout through Equation can be applied to arbitrary feature maps F or vectors
I or Equation i.e., this constraint equation
A linear convolutional layer l satisfies Equation 3 with respect to the group C₁ if the convolutional kernel K: IR? + Poual x Paal has the following form (Cohen et al. 2018):
(g)((g)v) = Poul (g) (L))
Since the composition of equivariant maps is equivariant, fully convolutional equivariant network can be constructed by stacking equivariant convolutional layers that satisfy the constraint Equa- tion and together with equivariant non-rineasssies (Weiler & Cesal 2019)
3

---

Published as a conference paper at ICLR 2022
4 PROBLEM STATEMENT
4.1 GROUP-INVARIANT MDPs
In group-invariant MDP, the transition and reward functions are invariant to group elements y E G acting on the state and action space. For state E S, action a E A, and g E G, let gs E S denote the action of g on and ga E A denote the action of g on a.
Definition 4.1 (G-invariant MDP). A G-invariant MDP MG = (S, an MDP M (S, A, T, R) that satisfies the following conditions:
1. Reward Invariance: The reward function is invariant to the action of the group element E G, R(s, a) R(gs.ga).
2. Transition Invariance: The transition function is invariant to the action of the group element
g E G, T(s, a(s,,,,')) = Tigs, ga,gs') gs').
A key feature of a G-invariant MDP is that its optimal solution is also G-invariant (proof in Ap- pendix A):
Proposition 4.1. Let MG be a group-invariant MDP. Then its optimal Q-function group invuri- ant, (s, Q" (gs,ga), and its optimal policy is group-equivarian, 't" (gs) gx' (s), for any g E
It should be noted that the G-invariant MDP of Definition 4.1 is in fact special case of an MDP homomorphism (Ravindran & Barto 200T 2004). broad class of MDP abstractions. MDP ho- momorphisms are important because optimal solutions to the abstract problem can be "lifted" to produce optimal solutions the original MDP (Ravindran & Barto 2004). As such, Proposition 4.1 follows directly from those results.
4.2 INVARIANN MDPs VISUAL STATE SPACES
In the remainder of this paper, focus exclusively important class of -nvvariant MDPs where the state is encoded as an image. We approximate SO(2) by its subgroup Cn.
State space: State is expressed as m-channel image, F, TR2 R The group opera- g On acts on this image as defined in Equation [ where we set Pj po: gF. (p( y)), i.e., by rotating the pixels but leaving the pixel feature vector unchanged.
Action space: We assume we are given a factored action space Anv x Aequiv = AS R embedded in k-dimensional Euclidean space where Ainv C Rawe and Aequiv C k We require the variables Ain to invariant with the rotation operator and the variables in Aequiv to rotate with the representation Pequiv P1- Therefore, the rotation operator g Cn acts ga = (Pequiv (g) lequiv, ajav) where anv E Any and Gequiv E Aequiv.
A via
Application to robotic manipulation: We express the state as a depth image centered on the gripper position where depth is defined relative the gripper. The orientation of this image is relative to the base reference frame not the gripper frame. We require the fingers of the gripper and objects grasped by the gripper to be visible in the im- age. Figure shows an illustration. The action a tuple, (ax, azy az, ag) E A C. RR,, where a AĄ denotes the commanded gripper aperture, axy E Any denotes the commanded change in gripper zy position, a E A_ denotes the commanded change in gripper height, and ag E Aa denotes the commanded change in gripper orientation. Here, the ry action is equivariant with E Cn. Aequiv = Azy. and the rest of the action variables are invariant, Any = AA x A X Ag. Notice that the transition dynamics Newtonian physics of the interaction are invariant to the choice reference frame. If we constrain the reward function be C, -n-invariant as well, then the resulting MDP is -.niinariaan.
(a)
(b)
Figure 2: The manipulation scene (a)
and the visual state space (b).
C, -invariaant (i.e. T(s, a, s') T(gs, ga, gs')) because the
4

---

Published as a conference paper at ICLR 2022
5 APPROACH
5.1 EQUIVARIANT DQN
In DQN, we assume we have a discrete action space, and we learn the parameters of Q-network that maps from the state onto ac- tion values. Given a G-invariant MDP, Proposition tells us that the optimal Q-function is G-invariant. Therefore, we encode the Q-function using an equivariant neural network that is constrained to represent only G-invariant Q-functions. First, in order to use DQN, we need to discretize the action space. Let Aequiv C Aequiv and Anw C Aiw be discrete subsets of the full equivariant and invariant action spaces, respectively. Next, we define a function Fa Aequiv + R- from the equivariant action variables Aequiv to the Q values of the invariant action variables in Auv- For example. in the robotic manipulation domain described Sec- tion 4.2 we have Aequiv Any and Amny A, x A₂ x Ag and Poquiv Pt. and we define Aequiv and Ainv accordingly. We now encode the Q network q stack of equivariant layers that each encode the equivariant constraint of Equation 3] Since the compo- sition of equivariant layers is equivariant, q satisfies:
Jou
Dew
y.
4
Figure 3: Illustration of Q- map equivariance. The output Q-map rotates with the input image.
((gF,) = g(a(Fa)) = gJa:
(5)
where we have substituted Fin. F3 and Fout = Fa. the above, the rotation operator g E (o(1))a(l1(())1((
On is applied using Equation as (May) equivariance constraint for the robotic manipulation example with Aequtvy - | 9. When the state (represented
Figure 3 illustrates this
image the left) is rotated by 90 degrees, the values associated with the action variables in Ary are also rotated similarly. The detailed network architecture is shown in Appendix [.I]] Our architecture is different from that in Mondal et al. (2020) in that we associate the action of on Aequiv and Any with the group action on the spatial dimension and the channel dimension of feature map F., which is more efficient than learning such mapping using FC layers.
5.2 EQUIVARIANT SAC
In SAC, we assume the action space is continuous. We learn the parameters for two networks: policy net- work (the actor) and an action-value network Q (the eritic) (Haarnoj et al. 20. The critic Q : S x A + R approximates Q values in the typical way. However, the actor II: dard deviation of action for a given state. Here, we define A, * to be the domain of the standard deviation vari- ables over the k-dimensional action space defined in Sec- tion 4.2 Since Proposition 4.1 tells us that the optimal Q is invariant and the optimal policy is equivariant, we must model Q as an invariant network and II as an equivariant network.
Gegir
Fire fnsture
Pape formu
B
Nott
Equiverient Inveriant Grite Natwork
A x A, estimates both the mean and stan-
i
F
Policy network: First, consider the equivariant constraint of the policy network. As before, the state is encoded by the function F.. However, we must now express the action vector over A = A x A.. Factoring A into its equivariant and invariant components, we have A = Aequiv X Any x Ag. In order to identify the equivariance relation for Á, we must define how the group operator g E acts on 0g E A,, Here, make the simplifying assumption that ag is invariant to the group operator. This choice makes sense in robotics domains where we would expect the variance of our policy to be invariant to the choice of reference frame. As result, have that the group element g G acts on à E Á via:
Figure 4: Illustration of the equivari- ant actor network (top) and the invariant critic network (bottom).
ga= g(uequiv, Ginv; ax) (Pequiv (gquui Ginv; 0.))
5

---

Published as a conference paper at ICLR 2022
We can now define the actor network to be mapping F. a (Figure 4 top) that satisfies the following equivariance constraint (Equation B)
ग (g.F.) 9( gā.
(7)
The critic network takes both state and action input and maps onto real value. We define two equivariant networks: a state encoder and a Q network q. The equivariant state encoder, e, maps the input state Fs onto a regular representation E (R where each of n group elements is associated with an a-vector. Since 8 has a regular representation, we have g8 Preg (g) Writing the equivariance constraint of Equation for e, we have that must satisfy e(gFF.) - The output state representation & is concatenated with the action a E A, producing w = (5. (). The action of the group operator now gw (gs,go, where ga = (Pequiv (equuy ainv). Finally, the network maps from onto R, real-valued estimate of the Q value for w. Based on proposition 4.1 this network must be invariant to the group action: g(gw) = g(w). All together, the critic satishes the following invariance equation:
Critic network:
gs.
((((gF)), ga) {(e(FF),
This network is illustrated at the bottom of Figure For a robotic manipulation domain in Sec- tion 4.2 we have Aequiv = -Ay and Any = A x A. x Ag and Pequiv P1. The detailed network
architecture in Appendix D.2
Preventing the critic from becoming overconstrained: In the model architecture above, the hidden layer of g is represented using a vector in the regular representation and the output of g is encoded using the trivial representation. However, Schur's ' Lemma (see e.g. Dummit & Foote (1991) implies there only exists a one-dimensional space of linear mappings from a regular representation to a trivial representation (i.e., v is a regular representation). This implies that a linear mapping ʃ : IR" x Rn + IR from two regular representations to trivial representation that satisfies f(gu, will also satisfy f(g1v, w) = f(u, w) and ((vggw)) = f(u, w) for all g1. 92 E G. (See details in Appendix B In principle, this could overconstrain the last layer of g to encode additional undesired symmetries. To avoid this problem we use non-linear equivariant mapping, maxpool, over the group space to transform the regular representation to the trivial representation.
E. Us where is a trivial representation,
is a constant, and
f(u, w) for all g G
5.3 EQUIVARIANT SACED
Many of the problems want to address cannot be solved without guiding the agent's exploration somehow. order to evaluate our algorithms in this context, we introduce the following simple strategy for learning from demonstration with SAC. First, prior training, we pre-populate the replay buffer with a set of expert demonstrations generated using a hand-coded planner. Second, we introduce the following L2 term into the SAC actor's 's loss function:
Lactor LSAC + le 2/2 ((a ~ (()))()>)))))))>>>>))) - ax)²
where LSAC is the actor's loss term in standard SAC, 1. 1 if the sampled transition is an expert demonstration and 0 otherwise, a ~ (s) is an action sampled from the output Gaussian distribution of (s)) and de is the expert action. Since both the sampled action (s)) and the expert action a. transform equivalently, Lactor is compatible with the equivariance we introduce in Section 5.2
We refer this method as SACfD (SAC from Demonstration).
6 EXPERIMENTS
We evaluate Equivariant DQN and Equivariant SAC in the manipulation tasks shown in Figure 5 These tasks can be formulated as -On))-invariant MDPs. All environments have sparse rewards (+1
when reaching the goal and 0 otherwise). See environment details in Appendix
6.1 EQUIVARIANT DQN
We evaluate Equivariant DQN in the Block Pulling, Object Picking, and Drawer Opening tasks
for the group C4. The discrete action space is A {OPEN, CLOSE);; Any ()r,,)

---

Published as a conference paper at ICLR 2022
Block Pulling
Object Picking
(c) Drawer Opening
(d) Block Stacking
(e) House Building
(l) Corner Picking
Figure 5: The experimental environments implemented in PyBullet (Coumans & Bai 2016). The left image in each sub figure shows initial state of the environment; the right image shows the goal state. The poses of the objects are randomly initialized.
OWN
CIN DER
CIN
6-d Dan DO'
Erd SWP DON
Erd ShP DON
1OO0 mece
LOND mece
LONO CeCT MUPTER OT FHER
(a) Block Pulling
(b) Object Picking
(c) Drawer Opening
Figure 6: Comparison of Equivariant DQN (blue) with baselines. The plots show the evaluation performance of the greedy policy in terms of the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four runs. Shading denotes standard error.
(-0.02m, Om, 0.02m }}: A, = f- (-0.02m Om, 0.02m): Ao = f- is.
]- Note that the def- C1 satisfies the closure requirement of the action space in a way that
inition of Azy and Vany Ary. Vg E C4. Pi (g)a.y E Any. We compare Equivariant DQN (Equi DQN) against the following baselines: 1) CNN DQN: DQN with conventional CNN instead of equivariant network, where the conventional CNN has similar amount of trainable parameters (3.9M) as the equivari- ant network (2.6M). 2) RAD Crop DQN (Laskin al.) (020a):: same network architecture CNN DQN. At each training step, each transition the minibatch is applied with a random-crop data aug- mentation. 3) DrQ Shift DQN (Kostrikov et al. 2020): same network architecture CNN DQN. At each training step, both the Q-targets and the TD losses are calculated by averaging over two random-shift augmented transitions. 4): CURL DQN (Laskin et al., (020bb:: similar architecture as CNN DQN with an extra contrastive loss term that learns an invariant encoder from random crop augmentations. See the baselines detail in Appendix E At the beginning of each training process, we pre-populate the replay buffer with 100 episodes of expert demonstrations.
Figure compares the learning curves of the various methods. Equivariant DQN learns faster and converges at a higher discounted reward in all three cnvironments.
6.2 EQUIVARIANT SAC
In this experiment, we evaluate the performance of Equivariant SAC (Equi SAC) for the group
C'.. The continuous action space is:A, = =00, 1]; Azy = ()o,8)]=,, E [-0.05m, 0.05m]] J:
A, |-0.05m, 0.05m]; Ag -- T. We compare against the following baselines: 1) CNN
SAC: SAC with conventional CNN rather than equivariant networks, where the conventional CNN has similar amount of trainable parameters (2.6M) the equivariant network (2.3M). 2) RAD
Crop SAC (Laskin et al.) 2020a): same model architecture as CNN SAC with random crop data augmentation when sampling transitions. 3) DrQ Shift SAC (Kostrikov et al. 2020): same model
7

---

Published as a conference paper at ICLR 2022
Ciill una YERM
chu sit: tro 2h.::SG
cit kit
Dep Se
SOO8 1260 1780 2016
200 Mi0
1200 17M0 2015
200 VGa
30-0 1240 16-38 11M6 20000
(a) Block Pulling
(b) Object Picking
(c) Drawer Opening
Figure 7: Comparison of Equivariant SAC (blue) with baselines. The plots show the evaluation performance of the greedy policy in terms of the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four runs. Shading denotes standard error.
OlN SACE
OlN SACE
ChN SADE
(ND) NIT NA13
(ND vi NALRO
Do pVn 5400
JUM LoW 2009 198 oed
209 DM0 .00 LOW: 2009 de00
200 190 olt J.W L00: olic 0960
000 200 JAOW KOMK ORCCP CHRIC
Drawer Opening
(b) Block Stacking
(c) House Building
(d) Corner Picking
Figure 8: Comparison of Equivariant SACID (blue) with baselines. The plots show the evaluation performance of the greedy policy in terms of the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four runs. Shading denotes standard error.
architecture as CNN SAC with random shift data augmentation when calculating the Qttarget and the loss. 4) FERM (Zhan et al. 2020): combination of SAC, contrastive learning, and random crop augmentation (baseline details Appendix E All methods use SO(2) data augmentation buffer, where every time new transition is added, we generate 4 more augmented transitions by applying random continuous rotations to both the image and the action (this data augmentation in the buffer in addition to the data augmentation that is performed in the RAD DrQ. and FERM baselines). Prior to each training run, we pre-load the replay buffer with 20 episodes expert demonstration.
shows the comparison among the various methods. Notice that Equivariant SAC outper- forms the other methods significantly. Without the equivariant approach, Object Picking and Drawer Opening appear to be infeasible for the baseline methods. In Block Pulling, FERM is the only other method able to solve the task.
Figure 7
6.3 EQUIVARIANT SACED
We want explore our equivariant methods the context more challenging tasks such as those in the bottom row of Figure However, since these tasks are too difficult to solve without some kind of guided exploration, we augment the Equivariant SAC as well as all the baselines two ways: 1) use SACID as described in Section [.3 2) use Prioritized Experience Replay (Schaul et al. 2015) rather than standard replay buffer. As in Section 6.2 we use the SO(2) data augmentation in the buffer that generates 4 extra (agg)eaugeented transitions whenever new transition is added. Figure g shows the results. First, note that our Equivariant SACID does best on all four tasks, followed by FERM, and other baselines. Second, notice that only the equivariant method solve the last three (most challenging tasks). This suggests that equivariant models are important not only for unstructured reinforcement learning, but also for learning from demonstration. Additional results for Block Pulling and Object Picking environments are shown in Appendix G
6.4 COMPARING WITH LEARNING EQUIVARIANCE USING AUGMENTATION
In the previous experiments, we compare against the data augmentation baselines using the same data augmentation operators that the authors proposed (random crop in RAD (Laskin al. 2020a)
8

---

Published as a conference paper at ICLR 2022
Paa 0r0-00/1 5420
NAE3 DrO SKO0
Prom NACRS
DVO SOGI
5A0-5001 1077
50.1
200 1a0 30-60 10000 de-da 2alo
2009 130
34-88 1mi0 1640 7960
2009 19E 200 OCON 1ANP 1960 1360
100 200 1a0e 1760 1807 1383
(a) Drawer Opening
(b) Block Stacking
(c) House Building
(d) Corner Picking
Figure 9: Comparison of Equivariant SACID (blue) with baselines. The plots show the evaluation performance of the greedy policy in terms of the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four runs. Shading denotes standard error.
and random shift in DrQ (Kostrikov et al. 2020)). However, those two methods can also be modified to learn SO( equivariance using S0(2) data augmentation. Here, we explore this idea as an alternative to our equivariant model. Specifically, instead of augmenting on the state as in Laskin iet al. (2020a) and Kostrikov et al. (2020) using only translation, we apply the SO(2) augmentation in both the state and the action. Since the RAD and DrQ baselines in this section running SO(2) augmentations themselves, we disable the SO(2) buffer augmentation for the online transitions in those baselines. (See the result of RAD and DrQ with the SO(2) data augmentation buffer in Appendix H..4.. We compare the resulting version of RAD (RAD SO(2) SACID) and DrQ (DrQ (O(2) SACAD) with Equivariant SACfI) in Figure Our method outperforms both RAD and DrQ equipped with SO(2) data augmentation. Additional results for Block Pulling and Object Picking are shown in Appendix G
already
6.5 GENERALIZATION EXPERIMENT
This experiment evaluates the ability for the equivariant model to generalize over the equiv- ariance group. We use similar experimental setting as Section 6.3 However, now the training environment is always initialized with a fixed orientation rather than a random ori- entation. For example, in Block Pulling, the two blocks are initialized with fixed relative orientation; in Drawer Opening, the drawer is initialized with a fixed orientation. In the eval- uation environment, however, these objects are initialized with random orientations. To suc- ceed, the agent needs to generalize over varied orientations while being trained with a fixed orienta- tion. To prevent the agent from generalizing via augmentation, we disable the SO(2) augmentation the buffer. As shown in Figure T0, Equivariant SACID generalizes better than the baselines. Even though the equivariant network is presented with only one orientation during training, it successfully generalizes over random orientation whereas none of the basclines can.
MALTI
TAD Cri
3IO SE
Eau SAUE
RAD Onea
tanter of tu saay
Tambe of 1al Akay (b) Drawer Opening
(a) Block Pulling
Figure 10: Comparison of Equivariant SACID with baselines. Results are averaged over four runs.
DISCUSSION
This paper defines a class of group-invariant MDPs and identifies the invariance and equivariance characteristics of their optimal solutions. This paper further proposes Equivariant SAC and new variation of Equivariant DQN for continuous action space and discrete action space, respectively. We show experimentally in the robotic manipulation domains that our proposal substantially surpasses the performance of competitive baselines. A key limitation of this work is that our definition of G-invariant MDPs requires the MDP to have an invariant reward function and invariant transition function. Though such restrictions are often applicable in robotics, they limit the potential of the proposed methods in other domains like some ATARI games. Furthermore, if the observation is from a non-top-down perspective, or there are non-equivariant structures in the observation (c.g., the robot arm), the invariant assumptions of G-invariant MDP will not be directly satisfied.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
This work supported in part by NSF 1724257, NSF 1724191, NSF 1763878, NSF 1750649, and NASA 80NSSC19K 1474. R. Walters is supported by the Roux Institute and the Harold Alfond Foundation and NSF grants 2107256 and 2134178.
REFERENCES
Michel Breyer, Fadri Furrer, Tonci Novkovic, R. Siegwart, and Nieto. Comparing task simplifica- tions to learn closed-loop object picking using deep reinforcement learning. IEEE Robotics and Automation Letters, 4:1549-1556, 2019.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličkovié. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint Mrxiv:22144 13478, 2021.
Serkan Cabi, Sergio G'omez Colmenarejo, Alexander Novikov, Ksenia Konyushkova, Scott E. Reed, Rae Jeong, Konrad Zolna, Yusuf Aytar, D. Budden, Mel Vecerík, Oleg O. Sushkov, David Barker, Jonathan Scholz, Misha Denil, N. D. Freitas, and Ziyun Wang. Scaling data-driven robotics with reward sketching and batch reinforcement learning, arXiv: Robotics, 2020.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confer- ence on machine learning, pp. 2990-2999. PMLR, 201
Taco Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant cmns homoge- neous spaces. arXiv preprint arXiv: 1.020/7, 2018.
Taco S Cohen and Max Welling. Steerable cmns. arXiv preprint arXiv:1612.08443 2016b.
Erwin Coumans and Yunfei Bai. Pybullet, python module for physics simulation for games, robotics and machine learning. Gitflub repository, 2016.
David S Dummit and Richard M Foote. Abstract algebra, volume 1999. Prentice Hall Englewood Cliffs, NJ, 1991.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018
Brian C Hall. Lie groups, Lie algebras, and representations: an elementary introduction, volume 10. Springer, 2003.
Peter J. Huber. Robust estimation of a location parameter. The Annals of Mathematicaal Statistics, 35 (1):73-101, 1964. ISSN 00034851,
Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, A. Irpan, J. Ibarz, Sergey Levine, R. Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to-sim: Data-efficient robotic grasping via randomized- -o-podooto-ccnccaa adaptation networks. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12619-12629, 2019.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, al. Qt-opt: Scalable deep reinforcement learning for vision-bar based robotic manipulation. arXiv preprint arXiv:1806. 10293, 2018.
Ozsel Kilinc, Yangfan Hu, and G. Montana. Reinforcement learning for robotic manipulation using simulated locomotion demonstrations. ArXiv, abs/1910.072994, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 141.::1412.6980 2014.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004 / 13649, 2020.
10

---

Published as a conference paper at ICLR 2022
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel. and Aravind Srinivas. Re- inforcement learning with augmented data. arXiv preprint arXiv:2004. 2020a.
Michael Laskin, Aravind Srinivas, and Picter Abbeel. Curl: tions for reinforcement learning. In International Conference on Machine Learning, pp. 5639- 5650. PMIR, 2020b.
Contrastive unsupervised representa-
Yijiong Lin, Jiancong Huang, Matthieu Zimmer, Yisheng Guan, Juan Rojas, and Paul Weng. In- variant transform experience replay: Data augmentation for deep reinforcement learning. IEEE Robotics and Automation Letters, 5(4):6615 6622, 2020.
Arnab Kumar Mondal, Pratheeksha Nair, and Kaleem Siddiqi. Group equivariant deep reinforce- ment learning. arXiv preprint arXiv:2007.03437, 2020.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. Xrxiv preprint arXi:: 2018.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. In NIPS Autodiff Workshop, 2017.
Deirdre Quillen, Eric Jang, Ofir Nachum, Chelsea Finn, J. Ibarz, and Sergey Levine. Deep inforcement learning for vision-based robotic grasping: A simulated comparative evaluation of off-policy methods. 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 6284-6291, 2018.
Balaraman Ravindran and Andrew Barto. Symmetrics and model minimization in markov deci- sion processes, 2001.
Balaraman Ravindran and Andrew G Barto. Approximate homomorphisms: A framework for non- exact minimization in markov decision processes. 2004.
Vishal Satish, Jeffrey Mahler, and Ken Goldberg. On-policy dataset synthesis for learning robot grasping policies using fully convolutional deep networks. IEEE Robotics and Automation Let- ters, 4(2): 2019.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprint arXiv: 15t.01551..0992,, 2015.
Elise van der Pol, Thomas Kipf, Frans A Olichock, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. arXiv preprint arXiv:2002.11963 2020a.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Olichock. and Max Welling. Mdp homo- Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020b.
morphic networks:
UIrich Viereck, Kate Saenko, and Robert W. Platt. Learning visual servo policies via planner cloning. ArXiv, abs/2005. 1810, 2020.
Dian Wang, Colin Kohler, and Robert Platt. Policy learning in se (3) action spaces, arXiv preprint arXiv:2010.02798, 2020a.
Dian Wang, Robin Walters, Xupeng Zhu, and Robert Platt. Equivariant $q$ learning in spatial action spaces. In 5th Annual Conference Robot Learning, 2021. URL htttss:/
net I==
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arXiv preprint arXiv:2002.03061 2020b.
Maurice Weiler and Gabriele Cesa. General e(2) steerable cnns. 1r1iv:1911.08255 2019.
arXiv preprint
11

---

Published as a conference paper at ICLR 2022
Andy Zeng, Shuran Song, Stefan Welker, Johnny Lee, Alberto Rodriguez, and Thomas Funkhouser. Learning synergies between pushing and grasping with self seff-supervised deep reinforcement learn- ing. In 2018 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp- IREE, 2018a 8a.
4238-422 4245.
Andy Zeng, Shuran Song, Kuan-Ting Yu, Elliot Donlon, Francois R Hogan, Maria Bauza, Daolin Ma, Orion Taylor, Melody Liu, Eudald Romo, et al. Robotic pick-and placeppcce of novel objects in clutter with aulti-affordance grasping and cross domain image matching. In 2018 IEEE interna- tional conference on robotics and automation (ICRA), pp. 3750-3757. IEEE, 2018b.
Andy Zeng, Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong, Ivan Krasin, Dan Duong, Vikas Sindhwani, et al. Transporter networks: Rear- ranging the visual world for robotic manipulation. arXiv preprint arXiv:2010.14406, 2020.
Albert Zhan, Philip Zhao, Lerrel Pinto, Pieter Abbeel, and Michael Laskin. A framework for effi- cient robotic manipulation. arXiv preprin! arXiv:2012.07975. 2020.
12

---

Published as a conference paper at ICLR 2022
A PROOF OF PROPOSITION 4.1
The proof in this section follows Wang et al. (021). Note that the definition of group action :: G x X + X implies that elements g G act by bijections on X since the action of g gives a two- sided inverse for the action of That is, g permutes the elements of X.
Proof of Proposition 4.1 For g E G. we will first show that the optimal Q-function is G-invariant, ie., (s, Q" [gs. ga), then show that the optimal policy is G-equivariant, i... n² (gs)
gr²(().
(1) (s, (gs, ga): The Bellman optimality equations for Q" a) and Q" ga) are, respectively:
/
(²( = R(s, a) + sup d'EA WES
T(s, (os.a,s))QGG0 (s',)))
(10)
and
Q' gg) R(gs, ga) + sup a'EA dESS
T(gs,ga (
(11)
Since g E G merely permutes the elements of S, we can re-index the integral using = gs':
Q" (gs. ga) R(ggs,g) +y sup a'EnA NeEs
T(gs, ga, 5 5')Q"" (5', a')
(12)
-Rggs ga) ty sup a'EA S'ES
T(gs, ga. gs'(0'''00000 (gs'' ga').
(13)
Using the Reward Invariance and the Transition Invariance in Definition 4.1 this can be written:
/ T(s, a, s )²' (gs'
(r = R(s, a) + sup d'EA J
(14)
Now, define a new function Q such that Vs S A, Q a) Qggs ga) and substitute into Eq. 14 resulting in:
(s( R(s, + sup
T(s a, s' (s',
(15)
REAS
Notice that Eq. 15 and Eq. 10 the same Bellman equation. Since solutions the Bellman equation are unique, we have that Vs, a E S x A, Q"(s" a) = Q* as) = Q (gs,ga).
(2) nt (gs) gT² (s): The optimal policy for n" (s()) and ' (gs) be written in terms of the optimal funnctio Q", as:
(s) = arg max (*(ssa) acA
(16)
and
(gs) arg max (gs,
(17)
Using the invariant property of Q* we can substitute Q" (gs, a) with Q" (s.gg ''a) in Equation 17
x (gs) = arg max Q" (s, ()
(18)
aeA
Let ga, Equation 18 can be written
(gs) g aarg max (* g ]9)]
(19)
ac
Cancelling g and and substituting Equation 16 we have.
t (gs) gT (s).
(20)
n
13

---

Published as a conference paper at ICLR 2022
B EQUIVARIANCE OVERCONSTRAIN
Proposition B.1. Let f: Vreg @ Vreg + Viriu be a linear -n--quiivriant function. Then f(v, w)
Σ, U/ Di W/g
Proof. By Weyl decomposibility (Hall, 2003). Vrey decomposes into irreducible representations for Cn each with multiplicity determined by dimension. Among these the trivial representation with multiplicity 1. By Schur's lemma (Dummit & Foote, 1991), the mapping Vreg @ Vreg + Virin must factor through the trivial representation embedded in Vreg. The projection onto the trivial representation is given v + a E, v4. The result follows by linearity.
As a corollary, find that Cj -quqqivvara maps Vreg Vreg - Virie are actually Cn x Ca- equivariant. Let (2) E Cn x Cn, then applying the Proposition u. = =0 (gv)+ b (gw), Ei Us E; wy [V, w)
C ENVIRONMENT DETAILS
In all environments, the environment reset is conduced by randomly initializing the objects with random positions and orientations in- side the workspace. The arm is always initialized at the same con- figuration. The workspace has size of 0.4m X 0.4m 0.24m. All environments have a sparse reward, i.e., the agent acquires a +1 reward when reaching the goal state, and 0 otherwise. In the Py- Bullet simulator, the robot joints have enough compliance to allow the gripper apply force the block in the Corner Picking task.
We augment the state image with an additional binary channel (i.e., either all pixels are 1 or all pixels are 0) indicating if the gripper is holding an object. Note that this additional channel is invariant to rotations (because all pixels have the same value) so won't break the proposed equivariant properties.
Figure 11: The object set for
Object Picking environment
The Block Pulling requires the robot to pull one block to make contact with the other block. The Object Picking requires the robot the pick up object randomly sampled from set of 11 objects (Figure IT]. The Drawer Opening requires the robot to pull open a drawer. The Block Stacking requires the robot to stack one block on top of another. The House Building requires the robot to stack triangle roof on top of a block. The Corner Picking requires the robot to slide the block from the corner and then pick it up.
D NETWORK ARCHITECTURE
Our equivariant models are implemented using the E2CNN (Weiler & Cesa 2019) library with PyTorch (Paszke et al. 2017).
D. EQUIVARIANT DQN ARCHITECTURE
In the Equivariant DQN, we use a 7-layer Steerable CNN defined in the group Ca (Figure 12a). The
input F, encoded as 2-channel Po feature map, and the output is 18-channel 3 x po feature map where the channel encodes the invariant actions Ainv and the spatial dimension encodes Aay
D.2 EQUIVARIANT SAC ARCHITECTURE
In the Equivariant SAC, there are two separate networks, both are Steerable CNN defined in the group Cg. The actor л (Figure 12b top) is an 8-layer network that takes in a 2-channel Po feature map (F.)) and outputs mixed representation type 1 x 1 feature map (ā) consisting of P1 feature for ary and 8 Po features for ain and a,, The critic (Figure [26 bottom) 9-layer network that takes in both F, as 2-channel Po feature map and a as a mixed representation feature map
14

---

Published as a conference paper at ICLR 2022
lsve Nudoo gag
LBk 'UlDo LMs
p200 SouL
1ood stu
poot you
poco sou
HZE
(a) Equivariant DQN Network Architecture
LSn Nudto SRo
LiP Nutto SX
ldkg Nudoo CAS
Hale Wedds gAS
8000 vou
podt VELL
teo XuGO Ch
poss draun
uize NuCO
"ADCO
(b) Equivariant SAC Network Architecture
Figure 12: The architecture of the Equivariant DQN (a) and the Equivariant SAC (b). ReLU nonlin- earity is omitted in the figure. A convolutional layer with a suffix of R indicates a regular represen- tation layer (e.g.. 16R is 16-channel regular representation layer); a convolution layer with suffix of T indicates trivial representation layer (e.g., IT is I-channel trivial representation layer).
consisting of feature for Gny and for anv The upper path encodes F, into 64-channel regular representation feature map with Two separate Q-value paths take in the concatenated feature map and generate two Q-estimates the form of representations into trivial representations to prevent the equivariant overconstraint (Section 5.2 Note that there are two Q outputs based on the requirement of the SAC algorithm.
spatial dimensions, then concatenates it with a.
Po feature. The non-linear maxpool layer is used for transforming regular
E BASELINE DETAILS
Figure shows the baseline network architectures for DQN and SAC. The RAD (Laskin et al. 2020a) Crop baselines, CURL (Laskin et (0020b) baselines, and FERM (Zhan et al. (0200 base- lines use random crop for data augmentation. The random crop crops 142 142 state image to the size of 128 x 128. The contrastive encoder of CURL baselines has size of 128 as in Laskin al. (2200b), and that for the FERM baselines has size of 50 as in Zhan al. (2220). The FERM baseline's ' contrastive encoder is pretrained for 1.6k steps using the expert data as Zhan et al. 2222)). The DrQ (Kostrikov et al., 2020) Shift baselines use random shift of +:4 pixels for data aug- mentation in the original work. In all DrQ baselines, the number of augmentations for calculating the target K and the number of augmentations for calculating the loss M are both 2 as Kostrikov al. 2020).
F TRAINING DETAILS
We implement our experimental environments in the PyBullet simulator (Coumans & Bai 2016 The workspace's size 0.4m 10..4..0.24mmm The pixel size of the visual state / x 128 (except for the RAD Crop baselines, CURL baselines, and FERM basclines, where I's size is 142 142 and will be cropped to 128 x 128). I's FOV 0.6m x 0.6m. During training, we use 5 parallel environments. We implement all training in PyTorrh (Paszke et al. 2017). Both DON and SAC use soft target update with f = 2
In the DQN experiments, of 10 We use Huber loss (Huber 1964) for calculating the TD loss. We use a discount factor 7 0.95. The batch size 32. The buffer has capacity of 100,000 transitions.
the Adam optimizer (Kingma & Ba 2014 with learning rate
15

---

Published as a conference paper at ICLR 2022
Bd: Nuto gag
Bsz buoo pxg
bis Nute pxg
1000 maW
1ood seu
Fodt 2010
1o0d 3004
B
Bascline DQN Network Architecture
pMc
SPS 04
pood
2B
19
821
962
962
821
gi pog
"uuos
pood
pood
100d
621
2e
19
B62
021
(b) Baseline SAC Network Architecture
Figure 13: The architecture of baseline conventional CNN DQN (a) and the baseline conven- tional CNN SAC (b). The baseline CNN architectures have similar amount of trainable parameters as the equivariant architectures. Specifically, Equivariant DQN has 2.6M parameters, and baseline DQN has 3.9M parameters; Equivariant SAC has 2.3M parameters, and baseline SAC has 2.6M parameters. ReLU nonlinearity omitted in the figure.
SACE Cra
Che 5AGP RAD
TENN SHOT
500
M021 skm
MX07
(a) Block Pulling
(b) Object Picking
(c) Block Pulling
(d) Object Picking
Figure 14: (a)-(b): additional results for Section 6.3
(c)-dd): additional results for Section 6.4
The plots show the evaluation performance the greedy policy in terms of the discounted reward. The evaluation performed every 500 training steps. Results averaged over four Shading denotes standard error.
In the SAC (and SACID) experiments, entropy temperature a is initialized at 10-2 The target entropy -5. The discount factor = 0.99. The batch size is 64. The buffer has capacity of 100.000 transitions. SACID uses the prioritized replay buffer (Schaul et al. (015)) with prioritized replay exponent of 0.6 and prioritized importance sampling exponent Bg 0.4 in /Schaul et al. (2015). The expert transitions are given a priority bonus of < 1.
the Adam optimizer with learning of 10-3 The
G ADDITIONAL EXPERIMENTAL RESULTS FOR EQUIVARIANT SACFD
Figure 14 (a)-(b) shows the results for the experiment of Section 6.3 in Block Pulling and Object Picking environments. The Equivariant SACID outperforms all baselimes those two environments.
Figure 14 (d) shows the results for the experiment of Section 6.4 in block Pulling and Object Picking environments. Similarly as the results in Figure [9] our Equivariant SACID outperforms both RAD and DrQ equipped with SO(2) dat augmentation.
16

---

Published as a conference paper at ICLR 2022
OM Cin Crdl
CR kater
Sou Adar 6e. CTBc
Eoul Adar Siga CrBc
Adter Cr
Adber Cell 12M0 29-00 palla
Use 2008 1200 1760
VG0 30:00 12M00
2015
Mid
(a) Block Pulling
(b) Object Picking
Drawer Opening
Rea CNN Cits hotor CNN CISS:
Eaull Aste Eaull Cros ow adter Cols
Coal Adar Eaal Cric Adwr Cril
TON0 1710 2010
12000 ITMD
(d) Block Stacking
(c) House Building
(f) Corner Picking
Figure 15: Ablation of using equivariant network solely in actor network or critic network. The plots show the evaluation performance terms discounted reward during training. The evaluation performed every 500 training steps. Results are averaged over four runs. Shading denotes standard error.
D00 200
1000
1900
JOOM LO00
200 20m
20/8 (000 10W
(a) Block Pulling
(b) Object Picking
(c) Drawer Opening
0
Sumber citrati vicis
Wumbe of
(d) Block Stacking
(e) House Building
( Corner Picking
Figure 16: Ablation of using different symmetry groups (Cs, C4. or C2). in Equivariant SACID.
The plots show the evaluation performance of the greedy policy in terms of the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four runs. Shading denotes standard error.
H ABLATION STUDIES
H.1 USING EQUIVARIANT NETWORK ONLY ACTOR CRITIC
In this experiment, we investigate the effectiveness of the equivariant network in SACID by only applying the actor network or the critic network. We evaluate four variations: 1) Equi Actor Equi Critic that uses equivariant network in both the actor and the critic; 2) Equi Actor + CNN Critic that uses equivariant network solely in the actor and uses conventional CNN in the critic; 3) CNN Actor Equi Critic that uses conventional CNN in the actor and equivariant network in the Critic; 4) CNN Actor + CNN Critic that uses the conventional CNN in both the actor and the critic. Other experimental setup mirrors Section 6.3 As is shown in Figure 15 applying the equivariant network
17

---

Published as a conference paper at ICLR 2022
OVR 5007
010 5007
130 34-06 EMi0 1alia TIME 7860
2003 FINE
UCO8 12i0 1ali 1360
(a) Block Pulling
(b) Drawer Opening
Figure 17: Ablation of using equivariant architecture in non-symmetric tasks. The plots show the evaluation performance of the greedy policy in terms the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four runs. Shading denotes standard emro.
50.21 SACI 50L2 SACI
SMAT 2022 54C0
5503 30.23 54O0
(a) Block Pulling
(b) Object Picking
(c) Drawer Opening
sarm
sorth
(d) Block Stacking
(e) House Building
(l) Corner Picking
Figure 18: Ablation of comparing against rotational augmentation baselines applied with rotational buffer augmentation. The plots show the evaluation performance of the greedy policy in terms of the discounted reward. The evaluation is performed every 500 training steps. Results are averaged over four rums. Shading denotes standard eтror.
in the actor generally helps more than applying the equivariant network in the critic (in 5 out of 6 experiments), and using the equivariant network in both the actor and the critic always demonstrates the best performance.
H.2 DIFFERENT SYMMETRY GROUPS
This experiment compares the equivariant networks defined in three different symmetry groups: Cs- C4, and C2. We this experiment in SACID with the same setup in Section 6.3 As is shown in Figure [6 the network defined in Css generally outperforms the network defined in C4. followed by the network defined in C2.
H.3 EQUIVARIANT SACED IN CON-SYMMETIIC ENVIRONMENTS
This experiments evaluates the performance of Equivariant SACID in non-symmetric tasks where the initial orientation of the environments are fixed rather than random. (Similarly as in Section 6.5 but both the training and the evaluation environments have the fix orientation.) Specifically, in Block Pulling, the two blocks in the training environment is initialized with a fixed relative orientation; in Drawer Opening, the drawer is initialized with fixed orientation. As is shown in Figure 17 when the environments do not contain SO(2) symmetries, the performance gain using equivariant network is less significant.
18

---

Published as a conference paper at ICLR 2022
H.4 ROTATIONAL AUGMENTATION + BUFFER AUGMENTATION
Section 6.4 compares our Equivariant SACID with rotational data augmentation baselines. This experiment shows the performance of those baselines (and extra CNN SACID baseline that conventional CNN) equipped with the data augmentation buffer. As is mentioned Section 6.2 the data augmentation baseline creates 4 extra augmented transitions using random SO(2) rotation every time a new transition is added. Figure 18 shows the result, where none of the baselines outperform our proposal in any tasks. Compared with Figure 9 the data augmentation buffer hurts RAD and DrQ because of the redundancy of the same data augmentation.
19


Reviews:
[{"review_id": "vRaffxad5AL", "reviewer": "Reviewer_pbdx", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found the paper to be clearly written and generally well-executed. The main strength of the paper is the experimental results which are very supportive in a range of different tasks. Its main weakness is lack of novelty, on the theoretical side. This is because, as noted by the paper itself, equivariant networks have been used for learning equivariant policies and invariant Q functions in previous work. However, given the importance of this domain, I believe the focus on this group and proposed methodology still has interesting contributions (e.g., in learning from demonstrations).\n\nQuestions/comments:\n \nProposition 4.1 is not new, and proper references should give credit to early works on symmetric MDPs (e.g., “Symmetries and Model Minimization in Markov Decision Processes”)\n\nOn “preventing critic from becoming overconstrained”: while I understand the proposed argument, it doesn’t seem valid to me, since you have the option of having many equivariant layers before the scaled mean pooling operation. While using max pooling instead of min pooling is fine, the argument seems tangential. Any comments?\n\nSince the results are surprisingly good compared to the model that replaces C8/C4 equivariant layer with ordinary convolution, could you please also provide the results for C2 equivariant layer? Such a layer would only have a factor of 2 improvement over the CNN (in data-efficiency) and therefore it would help make sense of the results.\n", "summary_of_the_paper": "The paper uses rotation equivariant CNNs for model-free reinforcement learning. More specifically, it is argued that for robotic manipulation tasks the corresponding MDP is invariant under translation and rotation, and therefore one could more efficiently learn equivariant/invariant policy/value functions. This idea is applied to both DQN and SAC for control with finite and continuous action spaces respectively. Experimental results on several tasks demonstrate substantial improvement over the competing methods.\n", "summary_of_the_review": "see above.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "pVKffQMyJ_y", "reviewer": "Reviewer_nGwx", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The proposed method and propositions are backed by solid theoretical proofs. \n\n2. The claims made about the proposed method are validated through extensive experimentations and comparison against strong baselines.\n\n3. The related works are well-covered. This paper is novel in its identification of equivariant properties of optimal Q-functions and optimal policies and its application to problems in robotics such as object manipulation (as opposed to previous applications to toy settings). It shows the improved performance of an equivariant DQN on these tasks and introduces an equivariant SAC model. It is closely related to Wang et al. (2021) which focuses on the translational and rotational invariance of Q-functions whereas this paper focuses on group equivariance. \n\n4. Overall presentation of the paper is good\n\nWeaknesses:\n\n1. Is the claim of proposing an equivariant version of DQN as a novel contribution correct? The authors reference Mondal et al. (2020) which claims to introduce equivariant DQN. Is what the authors proposing different in nature from Mondal et al. (2020)? The two works seem to use different symmetry groups (this paper talks about SO(2) and Mondal et al. talks about E(2)) but this and any/all other differences need to be better highlighted in the paper. \n\n2. The improvement in the reward shows the ability of the Equi DQN and Equi SAC to faster learn the policies but can it be ensured this is because of the inductive bias of the model architecture as claimed? One suggestion is to include a generalization experiment to test the robustness to equivariance where vanilla DQN fails. Additionally, the equivariant models are said to be applicable to rotationally symmetric problems. How would they perform in other problem settings without the symmetry assumption/ in the presence of other types of symmetry? Would they still beat the vanilla models? \n\n3. In the experiment section, the baseline CNN DQN only says “DQN with conventional CNN instead of equivariant network”. Is there a reason why the Equi DQN and baseline DQN have different numbers of convolutional layers and is this a fair comparison? What is the significance of the extra two fully-connected layers? Does this mean the Equi DQN needs more parameters to learn compared to the baseline DQN. If so, then there should be a baseline DQN that has the same number of parameters.  \n\n4. The paper has minor grammatical errors/omitted words. A careful proof-reading should be helpful. \n\n", "summary_of_the_paper": "This paper defines and theoretically characterizes a class of group-equivariant MDPs and studies its invariance and equivariance properties. It introduces two new models, one for discrete action spaces (equivariant DQN) and one for continuous action spaces (equivariant SAC). It compares the performance of the proposed methods against strong competitive baselines for multiple robotic manipulation tasks and establishes their superior performance. ", "summary_of_the_review": "The paper has an interesting premise about the advantage of encoding symmetry as an inductive bias in the data. It is novel in terms of introducing equivariant SAC but the claim of introducing equivariant DQN is ambiguous as it has already been introduced in previous work (Mondal et al. 2020).  Their experiments show improved performance of equi DQN and equi SAC over other methods for rotationally symmetric problems in robotic manipulation (which is different from previous works).  Although their results look very promising,  their claims may be better supported by additional experimentation  such as evaluating the performance of the proposed models in non-symmetric environments. Additionally, in order to ensure fair comparison, their equi DQN needs to be compared with a baseline DQN model with the same number of parameters. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "hxEkuaQuQZQ", "reviewer": "Reviewer_PtaK", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Although I liked the paper, I would like to kindly ask the authors the following questions: \n\n1. The authors have provided a definition of G-invariant MDPs and demonstrated their effectiveness in visual tasks. This definition rather introduces restrictions on the class of MDPs considered. My question is apart from visual tasks like those solved in this paper, can the authors provide examples and intuitions into when those assumptions hold? Is it the case that G-invariant MDPs are exclusive to image state representations? \n\n2. Similar question as the above goes for the action spaces. What assumptions are we really making on the reward and transition functions? Do these relate to Lipshitzness at all? Can we handle collision-avoidance scenarios for example? \n\n3. In the experiments, the authors have considered relatively low-dimensional action spaces. Can we please see a demonstration in higher dimensional action spaces? Would the method apply to ATARI domains as well? Or is it that the type of discontinuities present in ATARI prohibit the application of the proposed method? It would be great to have a section discussing the limitations in terms of MPDs that could be considered when using the current approach. \n\n4. Although a minor point, I am a bit confused about the novelty of the current approach beyond applying equivariant architectures to reinforcement learning. Can the authors please help me understand the major contributions of this work? ", "summary_of_the_paper": "In this paper, the authors present an application of equivariant neural networks to reinforcement learning. They demonstrate new DQN and SAC architectures that make use of equivariant representations in order to improve sample efficiency. Results on manipulation tasks demonstrate sample complexity reduction compared to other techniques. ", "summary_of_the_review": "Please see above.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SjrGmwt6VQJ", "reviewer": "Reviewer_27S2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n+ Very clearly motivates sample efficiency + generalization and articulates inductive bias from the perspective of adding equivariance to the model vs doing data augmentation.\n+ Results do show better data efficiency compared to baselines.\n\nWeakness/Comments\n- Results don't get into generalization but the claims/motivation in intro and related work leave the impression that this axis would be investigated,\n- The exposition of the approach is a bit hard to parse and sprinkling in some intuition and ground things to the application would be helpful.\n- What is the input? Is the kinds of images shown in Fig 5 or the depth map in Fig 2b? Some of the assumptions for this depth map don't make sense from a real application perspective: (i) if a depth camera is mounted near the robot end-effector it won't be possible to get the image in base frame, if it is overhead mounted then there will be occlusions so the object may not be always visible; (ii) depending on the task the gripper will often partially/fully occlude the object -- how much is this a problem? (iii) both Fig 5 and 2b are very idealized and in practice many other visual distractors would be present.\n- The claim in Sec 7 about \"transfers easily to real-world environments\" is not well supported and given the points above (and possibly other) this is in fact going to be quite challenging. A more nuanced coverage of limitations would be helpful to constraint where this method works and where it doesn't, rather that surface level limitations of learning in simulation and delta state based action.", "summary_of_the_paper": "This paper develops equivariant architecture for RL, specifically DQN and SAC. The theoretical exposition is general and a specific SO2 instantiation is shown to outperform baselines on image based RL tasks.\n", "summary_of_the_review": "Bringing inductive bias to mostly unstructured policies is useful for RL. The results are promising, but some claims are overblown or not well supported.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RhnW1wbC8lF", "reviewer": "Reviewer_cinR", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n\n- Sound theory of group-invariant MDP. The theory also characterizes both the invariance and equivariance cases.\n- Applications on realistic robotic manipulation\n- Showcasing on different RL (SAC and DQN) and robot learning (LfD) algorithms.\n\nWeakness:\n- The description sometimes is hard to follow\n\n\nIn general, the paper pursues an interesting research problem. it's well written. The proposed idea address the main well-understood challenge in robot learning. I would have only following comments.\n\nRegarding to the invariant-MDP, the proposed approach only deals with visual state spaces, so is it possible to extend to include proprioceptive information in the state space?\n\nIs the group element $G$ pre-defined and does it contain only a single rotation operator $g$ as defined in Experiment? If so would it be too simplistic? More ablation studies regarding this choice would be more helpful.\n", "summary_of_the_paper": "This paper explores to integrate equivariance deep learning to robotic applications. The authors propose two main contributions: 1) define and theoretically characterize an important class of group-equivariant MDPs, 2) and integrate equivariant variations to DQN, SAC, and LfD. The paper provides many different sets of experiments. The results are promising which shows benefits of the proposed approach.\n", "summary_of_the_review": "The paper proposes interesting ideas that might be useful for robot learning tasks. Experiment results are positive. There are some questions as raised in the above section.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper investigates the use of equivariant neural network architectures for model-free reinforcement learning in the context of visuomotor robot manipulation tasks, exploiting rotational symmetries in an effort to improve sample efficiency. The paper first provides a formal definition and theoretical evaluation of a class of MDPs for which the reward and transition are invariant to group elements ("group-invariant MDPs"). It goes on to describe equivariant versions DQN, SAC, and learning from demonstration (LfD). Experiments on a set of different manipulation tasks reveal that the proposed architectures outperform contemporary baselines in terms of sample complexity and generalizability, while ablations demonstrate the contribution of the different model components.

The idea of structuring a neural architecture to exploit symmetry present in a domain as a means of improving sample complexity is compelling and principled. The contributions of the paper are two-fold. First, while the idea of exploiting symmetry in the context of deep RL is not new, the paper describes a variation of equivariant DQN that is effective for visual control domains (visuomotor control) that are more challenging and realistic than those considered previously. Second, the paper proposes novel equivariant versions of SAC and LfD and validates their effectiveness through extensive experiments. Following a detailed author response to the initial reviews together with the inclusion of additional experiments and other updates to the paper, the reviewers largely agree on the significance of these contributions and value of the paper as a whole.
Decision: Accept (Spotlight)