Equivariant Q Learning in Spatial Action Spaces
Anonymous Author(s)
AfﬁliationAddress email
Abstract:
Recently, a variety of new equivariant neural network model architectures have been proposed that generalize better over rotational and reﬂectionalsymmetries than standard models.
These models are relevant to robotics becausemany robotics problems can be expressed in a rotationally symmetric way.
Thispaper focuses on equivariance over a visual state space and a spatial action space– the setting where the robot action space includes a subset of SE(2) .
In this situation, we know a priori that rotations and translations in the state image shouldresult in the same rotations and translations in the spatial action dimensions of theoptimal policy.
Therefore, we can use equivariant model architectures to make Qlearning more sample efﬁcient.
This paper identiﬁes when the optimal Q function is equivariant and proposes Q network architectures for this setting.
We showexperimentally that this approach outperforms standard methods in a set of challenging manipulation problems.
Introduction
A key question in policy learning for robotics is how to leverage structure present in the robotand the world to improve learning.
This paper focuses on a fundamental type of structure presentin visuo-motor policy learning for most robotics problems: translational and rotational invariancewith respect to camera viewpoint.
Speciﬁcally, the reward and transition dynamics of most roboticsproblems can be expressed in a way that is invariant with respect to the camera viewpoint fromwhich
the agent observes the scene.
In spite of the above, most visuo-motor policy learning agentsdo not leverage this invariance in camera viewpoint.
The agent’s value function or policy typicallyconsiders different perspectives on the same scene to be different world states.
A popular wayto combat this problem is through visual data augmentation, i.e., to create additional samples orexperiences by randomly translating and rotating observed images [1] but keeping the same labels.
This can be used in conjunction with a contrastive term in the loss function which helps the systemlearn an invariant latent representation [2, 3].
While these methods can improve generalization, theytypically require more sample steps for learning the invariance from the augmented transition.
Our key idea in this paper is to model rotational and translation invariance in policy learning using neural network model architectures that are equivariant over ﬁnite subgroups of SE(2) .
Theseequivariant model architectures reduce the number of free parameters using steerable convolutionallayers [4].
Compared with traditional methods, this approach can signiﬁcantly improve sample efﬁciency, i.e., learning to achieve the best performance with fewer training steps.
Moreover, it enablesus to generalize in a very precise way: everything learned with respect to one camera viewpoint isautomatically also represented in other camera perspectives via selectively tied parameters in themodel architecture.
We focus our work on Q learning in spatial action spaces, where the agent’saction space spans SE(2) or SE(3) .
We make the following contributions.
First, we identify theconditions under which the optimal Q function is SE(2) equivariant.
Second, we propose neuralnetwork model architectures that encode SE(2) equivariance in the Q function.
Third, since mostpolicy learning problems are only equivariant in some of the state variables, we propose partiallyequivariant model architectures that can accommodate this.
Finally, we compare equivariant modelsagainst non equivariant in the context of several robotic manipulation problems.
The results show that equivariant models are more sample efﬁcient than non-equivariant models, oftenby a signiﬁcant margin.
Related Work
Data Augmentation: Data augmentation techniques have long been employed in computer vision toencode the invariance property of translation and reﬂection into neural networks [5, 6].
Recent workdemonstrates the use of data augmentation improves the data efﬁciency and the policy’s performancein reinforcement learning [7, 8, 9].
In the context of robotics, data augmentation is often usedto generate additional samples [1, 10, 11].
Data augmentation improves the data efﬁciency of asystem, but not the sample efﬁciency.
In contrast to learning the equivariance property using dataaugmentation, our work utilizes the equivariant network to hard code the symmetries in the structureof the network to achieve better sample efﬁciency.
Contrastive Learning: Another approach to learning a representation that is invariant to translationand rotation is to add a contrastive learning term to the loss function [2].
This idea has been appliedto reinforcement learning in general [3] and robotic manipulation in particular [12].
While thisapproach can help the agent learn an invariant encoding of the data, it does not necessarily improvethe sample efﬁciency of policy learning.
Equivariant Learning: Equivariant model architectures hard-code E(2) symmetries into the structureof the neural network and have been shown to be useful in computer vision [13, 4, 14].
In reinforcement learning, some recent work applies equivariant model architectures to structure-ﬁnding problems involving MDP homomorphisms [15, 16].
In addition, Mondal et al.
[17] recently applied an
E(2) -equivariant model to Q learning in an Atari game domain, but showed limited improvement.
To our knowledge, equivariant model architectures have not been explored in the context of roboticsapplications.
Spatial Action Representations: Several researchers have applied policy learning in spatial actionspaces to robotic manipulation.
A popular approach is to do Q learning with a dense pixel actionspace using a fully convolutional neural network (this is the FCN approach we describe and extendin Section 4.2) [18, 19, 20, 21].
Variations on this approach have been explored in [22, 23].
The
FCN approach has been adapted to a variety of different manipulation tasks with different actionprimitives [24, 25, 26, 27, 28, 1, 29, 30, 31].
In this paper, we extend the work above by proposingnew equivariant architectures for the spatial action space setting.
Problem Statement
We are interested in solving complex robotic manipulation problems such as the packing and construction problems shown in Fig 1.
We focus on problems expressed in a spatial action space.
Manipulation as an MDP in over a visual state space and a spatial action space: We assume that themanipulation problem is formulated as a Markov decision process (MDP): M =
( S, A, T, R,  ) .
We focus on MDPs in visual state spaces and spatial action spaces [29, 20, 31].
The state spaceis factored into the state of the objects in the world, expressed as an n -channel h ⇥ w image
I 2 S world =
R n ⇥ h ⇥ w , and the state of the robot (including objects held by the robot) s rbt 2 S rbt ,expressed arbitrarily.
The total state space is S = S world ⇥ S rbt .
The action space is expressedas a cross product of SE(2) (hence it is spatial) and a set of additional arbitrary action variables:
A = SE(2) ⇥
A arb .
The spatial component of action expresses where the robot hand is to move andthe additional action variables express how it should move or what it should do.
For example, in thepick/place domains shown in Fig 1, A arb = { PICK , PLACE } , giving the agent the ability to moveto a pose and close the ﬁngers (pick) or move and open the ﬁngers (place).
We will sometimes decompose the spatial component of action a sp 2 SE(2) into its translation and rotation components,a sp
=
( x, ✓ ) .
The goal of manipulation is to achieve a desired conﬁguration of objects in the world,as expressed by a reward function R : S ⇥ A !
R .
Translation and Rotation in SE(2) :
We are interested in learning policies that are invariant to translation and rotation of the state.
To do that, we deﬁne rotation and translation of state and action as follows.
Let g 2 SE(2) be an arbitrary rotation and translation in the plane and lets =
( I, s rbt ) 2 S world ⇥ S rbt be a state.
g operates on s by rotating and translating the image [Figure] I , but leaving s rbt unchanged: gs =
( gI, s rbt ) , where gI denotes the image I translated and rotatedby g .
For action a = ( a sp , a arb ) , g rotates and translates a sp but not a arb : ga =
( ga sp , a arb ) .
Noticethat both S and A are closed under g 2 SE(2) , i.e. that 8 g 2 SE(2) , a 2 A = ) ga 2 A and 96 s 2 S = )
gs 2 S .
Assumptions: We assume that the reward and transition dynamics of the system are invariant withrespect to translation and rotation of state and action as deﬁned above, and that the translation androtation operations on state and action are invertible.
Assumption 3.1 (Goal Invariance) .
The manipulation objective is to achieve a desired conﬁgurationof objects in the world without regard to the position and orientation of the scene.
That is, R ( s, a ) =
R ( gs, ga ) for all g 2 SE(2) .
Assumption 3.2 (Transition Invariance) .
The outcome of robot actions is invariant to translationsand rotations of both the scene and the action.
Speciﬁcally, T ( s, a, s 0 ) =
T ( gs, ga, gs 0 ) for allg 2 SE(2) .
Assumption 3.3 (Invertibility) .
Translations and rotations in state and action are invertible.
Thatis, 8 g 2 SE(2) , g  1 ( gs ) =
s and g  1 ( ga )
= a .
Assumptions 3.1 and 3.2 are satisﬁed in problem settings where the objective and the transitiondynamics can be expressed intrinsically to the world without reference to an external coordinateframe imposed by the system designer.
These assumptions are satisﬁed in many manipulation domains including all those shown in Fig 1.
In House Building, for example, the reward and transitiondynamics of the system are independent of the coordinate frame of the image or the action space.
Assumption 3.3 is needed to guarantee the Q function invariance described in the next section.
Approach
Assumptions 3.1, 3.2, and 3.3 imply that the optimal Q function is invariant to translations androtations in SE(2) .
Proposition 4.1.
Given an MDP M =
( S, A, T, R,  ) for which Assumptions 3.1, 3.2, and 3.3 aresatisﬁed, the optimal Q function is invariant to translation and rotation, i.e. Q ⇤ ( s, a ) =
Q ⇤ ( gs, ga ) ,for all g 2 SE(2) .
(Proof in Appendix A.)
Our key idea is to use the invariance property of Proposition 4.1 to structure Q learning (and makeit more sample efﬁcient) by deﬁning a neural network that is hard-wired to encode only invariant Qfunctions.
However, in order to accomplish this in the context of DQN, we must allow for the factthat state is an input to the neural network while action values are an output .
This neural networkis therefore a function q : S !
R A , where R A denotes the space of functions { A ! R } .
Theinvariance property of Proposition 4.1 now becomes an equivariance property,
[Equation]
[Equation]
) where q ( s )( a ) denotes the Q value of action a in state s .
We implement this constraint using equivariant convolutional layers as described below.
4.1 Equivariant Convolutions
Equivariance over a ﬁnite group: In order to implement the equivariance constraint, it is standard inthe literature to approximate SE(2) by a ﬁnite subgroup [4, 14].
Recall that the spatial componentof an action is a sp
= ( x, ✓ ) 2 SE(2) .
We constrain position to be a discrete pair of positive integersx 2 { 1 . . .
h } ⇥ { 1 . . .
w } ⇢
Z 2 , corresponding to a pixel in the input image I .
We constrainorientation to be a member of a ﬁnite cyclic group ✓ 2 C u , i.e. one of u discrete orientations.
For example, if u = 8 , then C 8 = { 0 , ⇡ 4 , 2 ⇡ 4 , 3 ⇡ 4 , 4 ⇡ 4 , 5 ⇡ 4 , 6 ⇡ 4 , 7 ⇡ 4 } .
Our ﬁnite approximation ofa sp 2 SE(2)
is ˆ a sp in the subgroup ˆSE(2) generated by translations Z 2 and rotations C u .
h
R u ⇥ m h w
of network
generally takes a “ﬂat” image as input:h inequiv : R [Equation] n [Equation]
Equivariance constraint: Let h
i ( I )( x ) denote the outputof convolutional layer h at channel i and pixel x given input I .
For an equivariant layer, h i ( I )( x ) 2 R C u describesfeature values for each element of C u .
For an element g 2ˆSE(2) , denote the rotational part by g ✓ 2 C u .
If we identify functions R C u with vectors R u , then the group actionof g ✓ 2 C u on R C u becomes left multiplication by a permutation matrix ⇢ ( g ✓ ) that performs a circular shift on the vector in R u .
Then the group action ofˆSE(2) on a feature map h ( I ) 2 R u ⇥ m ⇥ h ⇥ w can be expressed as g ( h i ( I ))
( x ) = ⇢ ( g ✓ ) h
i ( I )( g  1 x ) .
The mapping h is equivariant if and only if
[Equation]
for each i 2 { 1 . . .
m } .
This is illustrated in Fig 2.
We can calculate the output feature map in thelower right corner by transforming the input by g and then doing the convolution (left side of Eq.
2)or by doing the convolution ﬁrst and then taking the value of g  1 x and circular-shifting the outputvector (right side of Eq. 2).
In order to create a network that enforces the constraint of Eq. 1, we cansimply stack equivariant convolutions layers that each satisfy Eq. 2.
Kernel constraint: The equivariance constraint of Eq. 2 can be implemented by strategically tyingweights together in the convolutional kernel [4].
Since the standard convolutional kernel is alreadytranslation equivariant [13], we must only enforce rotational ( C u ) equivariance [33]:
[Equation]
where ⇢ in ( g ✓ ) and ⇢ out ( g ✓ ) are the permutation matrix of the group element g ✓ (note that for theﬁrst layer, K ( y ) will be a 1 ⇥ u matrix, and ⇢ in ( g ✓ ) will be 1).
More details are in Appendix B.
4.2 Equivariant Fully Convolutional Q Functions in SE(2)
A baseline approach to encoding the Q function over a spatial action space is to use a fully convo- lutional network (FCN) that stacks convolutional layers to produce an output Q map with the same resolution as the input image.
If we ignore the non-image state variables s rbt and the non-spatialaction variables a arb , then we have all the tools we need – we simply replace all convolutional layerswith equivariant convolutions and the Q network becomes fully equivariant.
Partial Equivariance: Unfortunately, in realistic robotics problems, Q function is generally not equiv- ariant with respect to all state and action variables. For example, the non-equivariant parts of state [Figure]
Green blocks are equivariant and white blocks are not. and action in Section 3 are s rbtand a arb .
We encode a arb by simply having a separate output headfor each.
However, to encode s rbt ,we need a mechanism for inserting the non-equivariant information into the neural network modelwithout “breaking” the equivariance property.
We explored twoapproaches: the lift expansion approach and the dynamic ﬁlter approach.
In the lift expansion, we tile the non-equivariant informationacross the equivariant dimensions of the feature map as additional channels (Fig 3a).
In the dynamicﬁlter approach [34], the non-equivariant data is passed through a separate pathway that outputs theweights of an equivariant kernel that is convolved into the main equivariant backbone.
We constrainthis ﬁlter to be equivariant by enforcing the kernel constraint of Eq. 3 (Fig 3b).
We empirically ﬁndthat both methods have similar performance (Appendix G.2).
In the remainder of this paper, we usethe dynamic ﬁlter approach because it is more memory efﬁcient.
Encoding Gripper Symmetry Using Quotient Groups: Another symmetry that we want to leverageis the bilateral symmetry of the gripper.
The outcome of a pick action performed using a twoﬁnger gripper in orientation ✓ is the same as for the gripper in orientation ✓ + k ⇡ for any integer k .
Similarly, it is often valid to assume that the outcome of place actions is invariant 3 .
We model thisinvariance using the quotient group C u /C 2 .
The C 2 = { 0 , ⇡ } action equates rotations which differby multiples of ⇡ in C u /C 2 .
The steerable layer deﬁned under the quotient group is applied withthe same constraint as in Eq. 3, except that the output space will be in C u /C 2 .
Experimental Domains: We evaluate the equivariant FCN approach in the Block Stacking and Bottle Arrangement tasks shown in Fig 1.
Both environments have sparse rewards (+1 at goal and 0 otherwise).
The world state is encoded by a 1-channel heightmap I 2 R 1 ⇥ h ⇥ w and robot state is encodedby an image patch H that describes the contents of the robotic hand.
The non-spatial action variablea arb 2 { PICK , PLACE } is selected by the gripper state, i.e., a arb = PLACE if the gripper is holding an object, and PICK otherwise.
The equivariant layers of the FCN are deﬁned over group C 12where the output is with respect to the quotient group C 12 /C 2 to encode the gripper symmetry.
See Appendix C and D for detail on the experimental domains and the FCN architecture respectively.
[Figure]
Shading denotes standard error.
Experimental Comparison With Baselines: Weevaluate against the following baselines: 1)
Conventional FCN: FCN with 1-channel inputand 6 -channel output where each output channel corresponds to a Q map for one rotation inthe action space (similar to Satish et al.
[19] butwithout the z dimension). 2) RAD [7] FCN:same architecture as 1), while at each trainingstep, we augment each transition in the minibatch with a rotation randomly sampled from C 12 . 3) DrQ [8] FCN: same architecture as 1),while at each training step, the Q targets and Qoutputs are calculated by averaging over multiple augmented versions of the sampled transitions. Random rotations sampled from C 12 are used for the augmentation.
Rot FCN: FCN with 1channel input and 1-channel output, the rotation is encoded by rotating the input and output for each✓ [21]. Transporter Network [1], an FCN-based architecture with the last layer being a dynamickernel generated by a separate FCN with an input of an image crop at the pick location.
Baseline2) and 3) are data augmentation methods that aim to learn the symmetry encoded in our equivariant network using rotational data augmentation.
All baselines have the same action space as theproposal.
More detail on the baselines is in Appendix E.1.
All methods except the Transporter
Network use SDQfD, an approach to imitation learning in spatial action spaces that combines a TDloss term with penalties on non-expert actions [29].
(Transporter Network is a behavior cloningmethod.)
Table 1 shows the number of demonstration steps.
Those expert transitions are augmentedby 9 random SE(2) transformations.
See Appendix F.2 for more parameter detail.
Fig 4 shows the [Table] [Table] results.
Our equivariant FCN outperforms all baselines in the block stacking task.
Notice that in the Bottle Arrangement task, the equivariant network learns faster than the baselines but converges toa similar level as RAD and DrQ. This is because the domain itself is already partially rotationallyequivariant because the bottles are cylindrical and therefore our network has less of an advantage.
4.3 Equivariant Augmented State Q Functions in SE(2)
The FCN approach does not scale well to challenging manipulation problems.
Therefore, we designan equivariant version of the augmented state representation (ASR) method of [29], which has beenshown to be faster and have better performance.
The ASR method transforms the original MDPwith a high dimensional action space into a new MDP with an augmented state space but a lowerdimensional action space.
Instead of encoding the value of all dimensions of action in a singleneural network, this model encodes the value of different factorized parts of the action space suchas position and orientation using separate neural networks conditioned on prior action choices.
ASR in SE(2) :
We explain the ASR method in the example setting of the SE(2) action space. See Fig 5 for an illustration.
Asbefore, actions ˆ a sp =
( x, ✓ ) are elements of the space ˆSE(2) .
However, the Q function is now computed using two separate functions,the position function Q 1 ( s, x ) = max ✓ Q ( s, ( x, ✓ )) and the orientation function Q 2 (( s, x ) , ✓ ) = Q ( s, ( x, ✓ )) .
Q 1 is encoded using a fully convolutional network q 1 : R n ⇥ h ⇥ w ! R 1 ⇥ h ⇥ w thattakes an n -channel image I as input and produces a 1 -channel Qmap that describes Q 1 ( s, x ) for all x .
We evaluate Q 2 on the“augmented state” ( s, x ) which contains the state s and the chosen x .
The augmented state is encoded using the image patch
P = CROP ( I, x ) 2 R n ⇥ h 0 ⇥ w 0 cropped from I and centered at x . We model Q 2 using the networkq 2 : R n ⇥ h 0 ⇥ w 0 !
R u that takes input P and outputs Q 2 (( s, x ) , ✓ ) for all u different orientation✓ .
These two networks are used together for both action selection and evaluation of target valuesduring learning.
We evaluate x ⇤ = arg max x q 1 ( I ) , calculate P = CROP ( I, x ⇤ ) , and then evaluate✓ ⇤ = arg max ✓ q 2 ( P ) and Q ⇤ = max ✓ q 2 ( P ) .
Note that Q maps produced by q 1 and q 2 are of sizeu + hw , signiﬁcantly smaller than the Q map in the FCN approach which is size uhw .
Essentiallythe ASR method takes advantage of the observation that effective Q learning in this environmentdoes not require computing exact Q values for all combinations of x and ✓ .
[Figure]
Equivariant architecture for ASR in SE(2) :
Wedecompose the SE(2) equivariance property of Eq. 1 into two equivariance properties for q 1and q 2 , respectively: q 1 ( gI )( x ) =
q 1 ( I )( g  1 x )where g 2 ˆSE(2) , and q 2 ( g ✓ P ) = ⇢ ( g ✓ ) q 2 ( P )where g ✓ 2 C u .
The equivariance property of 268 q 1 is similar to that of Eq. 2 except that theoutput of q 1 has only one channel which is invariant to rotations (since it is a maximum overall rotations).
This means we can rewrite theq 1 equivariance property as q 1 ( gI ) =
gq 1 ( I ) ,where g on the RHS of this equation translatesand rotates the output Q map.
In practice, weobtained the best performance for q 1 by enforcing equivariance to the Dihedral group D 4 , which isgenerated by 90 degree rotation and reﬂections over the coordinate axis.
For q 2 , we used an equivariant feature map that outputs a single u -dimensional vector of Q values corresponding to the ﬁnitecyclic group C u used.
(We use C 12 /C 2 and C 32 /C 2 in our experiments below).
We handle thepartial equivariance using the same strategies as earlier.
Appendix D describes the model details.
[Figure] [Figure] Experimental Comparison with Equivariant FCN: Fig 6 shows a comparison between equivariant ASR (this section) and equivariant FCN (Section 4.2) for the Block Stacking and Bottle Arrangementtasks.
The network q 2 is deﬁned using C 12 and its quotient group C 12 /C 2 to match Section 4.2.
The
ASR method surpasses the FCN method in both tasks. More Challenging Experimental Domains: The equivariant ASR method is able to solve more challenging manipulation tasks than equivariant FCN can. In particular, we could not run the FCN withas large a rotation space because it requires more GPU memory.
We evaluate on the following fouradditional domains: House Building, Covid Test, Box Palletizing (introduced in [1]), and Bin Packing (Fig 1(c-f)).
All domains except Bin Packing have sparse rewards.
In Bin Packing, the agentobtains a positive reward inversely proportional to the highest point in the pile after packing all objects.
See Appendix C for more details about the environments.
We now deﬁne q 2 using the group C 32 and its quotient group C 32 /C 2 , i.e., we now encode 16 orientations ranging from 0 to ⇡ .
Asin Section 4.2, we use the SDQfD loss term to incorporate expert demonstrations (except for Transporter Net which uses standard behavior cloning exclusively).
The number of expert transitionsprovided is shown in Table 1.
9 random SE(2) augmentations are applied to the expert transitions.
Experimental Comparison with Non-Equivariant Baselines: We compare equivariant ASR againstthe following non-equivariant baselines: 1) Conventional ASR: ASR in SE(2) with conventional
CNNs rather than equivariant layers.
RAD [7] ASR: same architecture as (1) but each minibatchis augmented with a rotation randomly sampled from C 32 . 3): DrQ [8] ASR: same architecture as(1) but each Q target and Q estimate are calculated by averaging over several augmented versionsof the sampled transition. The augmentation is by random rotations sampled from C 32 . the Transporter Network [1].
See Appendix E.2 for the network architecture for the baselines.
Theresults in Fig 7 show that equivariant ASR outperforms the other methods on all tasks, followed by
DrQ, RAD, and Transporter Net, followed by Conventional ASR.
[Table] Robot Experiment: We evaluate the trained equivariant ASRmodels for Bottle Arrangement, House Building, and Box Palletizing on a Universal Robots UR5 arm equipped with a Robotiq 2F-85 gripper.
The observation is provided by an Occipital Structure sensor mounted on top of the workspace.
Table 2shows the results.
In Bottle Arrangement, the robot shows a 90%success rate.
In one of the two failures, the arrangement is notcompact enough, leaving no enough space left for the last bottle.
In the other failure, the robotarranges the bottles outside of the tray.
In the House Building task, the robot succeeds in all 20episodes.
In the Box Palletizing task, the robot demonstrates a 95% success rate.
In the failure, the robot correctly stacks 16 of 18 boxes, but the 17th box’s placement position is offset slightly fromthe the rest of the stack and there is no room to place the last box.
The same problem happens inanother successful episode, where the ﬁngers squeeze the boxes and make room for the last box.
Fig 8 shows two example episodes in the robot experiment.
4.4 Equivariant Augmented State Q Functions in SE(3)
[Figure] [Figure]
A strength of the ASR method is that it can beextended into SE(3) by adding networks similar to q 2 that encode Q values for additional dimensions of the action space [29].
Speciﬁcally,we add three networks to the SE(2) -equivariantarchitecture described in Section 4.3, q 3 , q 4 ,and q 5 encoding Q values for Z (height abovethe plane), and angles  (rotation in XZ plane)and (rotation in YZ plane) dimensions of
SE(3) .
Each of these networks takes as inputa stack of 3 orthographic projections of a pointcloud along the coordinate planes.
The pointcloud is re-centered and rotated to encode thepartial SE(3) actions. (see [29] for details).
Unfortunately, we cannot easily make these additional networks equivariant using the samemethods we have proposed so far because theyencode variation outside of SE(2) .
Instead, wecreate an encoding that is approximately equivariant by explicitly transforming the input tothese networks in a way that corresponds to a set of candidate robot positions or orientations (calleda “deictic” encoding in [22]). {+We will describe this idea using q 2 as an example. Deﬁne q 0 
 ( P ) 
+}
Rto output the single Q value of the rotation encoded in the input image P .
Then q 2 can be deﬁned as avector-valued function: ˆ q 2 ( P ) = ( q 0 2 ( g  1 1 P ) , . . . , q 0 2 ( g  1 u P )) , where g 1 , . . , g u 2 C u .
ˆ q 2 is approximately equivariant because everything learned by q 0 2 is automatically replicated in all orientations.
We design deictic q 3 , q 4 , and q 5 similarly by selecting a ﬁnite subset of { g 1 , . .
. , g K } ⇢ SE(3)corresponding to the dimension of the action space encoded by each q i .
q i can then be deﬁned byevaluating a network q 0 i over input P transformed by ( g i )
Ki
=1 .
For q 3 , we evaluate over 36 translations g k ( z ) =
z + k (0 . 18 / 36) + 0 .
where 0  k  35 .
For q 4 and q 5 , we use rotationsg k 2 {  ⇡ 8 ,  ⇡ 12 ,  ⇡ 24 , 0 , ⇡ 24 , ⇡ 12 , ⇡ 8 } .
Note we use q 2 for explanation, while our model uses equiv- 349 ariant q 1 , q 2 and deictic q 3 - q 5 .
For an ablated version using deictic q 2 , see Appendix G.4.3 and G.4.4.
Comparison to Non-Equivariant Approaches: We evaluate ASR in SE(3) in the House Building and
Box Palletizing domains.
We modiﬁed those environments so that objects are presented randomlyon a bumpy surface with a maximum out of plane orientation of 15 degrees (Fig 9).
In order tosucceed, the agent must correctly perform pick and place actions with the needed height and out ofplane orientation.
We evaluated the Equivariant ASR in comparison with a baseline Conventional
ASR (same as [29]).
Both methods use SDQfD with 2000 expert demonstration steps.
The resultsare shown in Fig 10.
Our proposed approach outperforms the baseline by a signiﬁcant margin.
Discussion
In this paper, we show that equivariant neural network architectures can be used to improve Qlearning in spatial action spaces.
We propose multiple approaches and model architectures that canbe used to accomplish this and demonstrate improved sample efﬁciency and performance on severalrobotic manipulation applications both in simulation and on a physical system.
This work has severallimitations and directions for future research.
First, our methods apply directly only to problems inspatial action spaces.
While many robotics problems can be expressed this way, it would clearly beuseful to develop equivariant models for policy learning that can be used in other settings.
Second,although we extend our ASR approach from SE(2) to SE(3) in the last section of this paper, thissolution is not fully equivariant in SE(3) and it may be possible to do better by exploiting methodsthat are directly equivariant in SE(3) .