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, theyrequire the neural network to learn translational and rotational invariance from the augmented data.
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 creates an inductive bias that can signiﬁcantly improve the sample efﬁciency of the model, the number of environmental steps needed tolearn a policy.
Moreover, it enables us to generalize in a very precise way: everything learned withrespect to one camera viewpoint is automatically also represented in other camera perspectives viaselectively tied parameters in the model architecture.
We focus our work on Q learning in spatialaction spaces, where the agent’s action space spans SE(2) or SE(3) .
We make the following contributions.
First, we identify the conditions under which the optimal Q function is SE(2) equivariant.
Second, we propose neural network model architectures that encode SE(2) equivariance in the Qfunction.
Third, since most policy learning problems are only equivariant in some of the state variables, we propose partially equivariant model architectures that can accommodate this.
Finally, wecompare equivariant models against non-equivariant counterparts in the context of several robotic manipulation problems.
The results show that equivariant models are more sample efﬁcient thannon-equivariant models, often by 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 used togenerate additional samples [1, 10, 11].
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.
Thissection identiﬁes conditions under which the Q function is SE(2) -invariant.
The next section describes how these invariance properties translate into equivariance properties in the neural network.
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 and action.
To do that, we deﬁne rotation and translation of stateand action as follows.
Let g 2 SE(2) be an arbitrary rotation and translation in the plane and let [Figure] s =
( I, s rbt )
2 S world ⇥ S rbt be a state.
g operates on s by rotating and translating the image I , butleaving s rbt unchanged:
gs = ( gI, s rbt ) , where gI denotes the image I translated and rotated byg .
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 97 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]
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 u ⇥ m generally takes
“ﬂat” image as input:h inequiv : R
[Equation]
n
[Equation]
[Equation]
u
[Equation]
{+w . 
+}
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 convolutional 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 equivariant 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]
Experimental Comparison With Baselines:
Weevaluate against the following baselines: 1) Conventional FCN: FCN with 1-channel input and6 -channel output where each output channel corresponds to a Q map for one rotation in the actionspace (similar to Satish et al.
[19] but without thez dimension).
2) RAD [7] FCN: same architectureas 1), while at each training step, we augmenteach transition in the minibatch with a rotationrandomly sampled from C 12 . 3) DrQ [8] FCN:same architecture as 1), while at each trainingstep, the Q targets and Q outputs are calculatedby averaging over multiple augmented versions of the sampled transitions.
Random rotationssampled from C 12 are used for the augmentation.
Rot FCN: FCN with 1-channel input and1-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 dynamic kernelgenerated by a separate FCN with an input of an image crop at the pick location.
Baseline 2)and 3) are data augmentation methods that aim to learn the symmetry encoded in our equivariantnetwork using rotational data augmentation sampled from the same symmetry group ( C 12 ) as usedby our equivariant model.
All baselines have the same action space as the proposal.
More detailon 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 TD loss term withpenalties on non-expert actions [29].
(Transporter Network is a behavior cloning method.)
Table 1shows the number of demonstration steps.
Those expert transitions are augmented by 9 random [Table] [Table] SE(2) transformations.
See Appendix F.2 for more parameter detail.
Fig 4 shows the results.
Ourequivariant 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 to a similarlevel as RAD and DrQ. This is because the domain itself is already partially rotationally equivariantbecause 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 settingof the SE(2) action space.
See Fig 5 for an illustration.
As before, actions ˆ a sp = ( x, ✓ ) are elements of the space ˆSE(2) , the ﬁnite approximation of SE(2) as in Section 4.1.
However, the Qfunction is now computed using two separate functions, the positionfunction 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 that takes an n -channelimage I as input and produces a 1 -channel Q map 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 stateis encoded using the image patch P =
CROP ( I, x ) 2 R n ⇥ h 0 ⇥ w 0 cropped from I and centeredat x .
We model Q 2 using the network q 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 values during 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 ) .
Notethat Q maps produced by q 1 and q 2 are of size u + hw , signiﬁcantly smaller than the Q map in the FCN approach which is size uhw .
Essentially the ASR method takes advantage of the fact that theoptimal ✓ depends only on the local patch P given an optimal position x .
[Figure]
Equivariant architecture for ASR in SE(2) :
We decompose the SE(2) equivariance property of Eq.
1into two equivariance properties for q 1 and 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 q 1 is similar to that of Eq.
2 except that the output of q 1 has only one channel which is invariant to rotations (since it is a maximum over all rotations).
This means we can rewritethe q 1 equivariance property as q 1 ( gI ) =
gq 1 ( I ) ,where g on the RHS of this equation translates androtates the output Q map.
In practice, we obtainedthe best performance for q 1 by enforcing equivariance to the Dihedral group D 4 , which is generated by 90 degree rotation and reﬂections over thecoordinate axis.
For q 2 , we used an equivariant feature map that outputs a single u -dimensionalvector of Q values corresponding to the ﬁnite cyclic group C u used.
(We use C 12 /C 2 and C 32 /C 2in our experiments below).
We handle the partial equivariance using the same strategies as earlier.
Appendix D describes the model details.
[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 minibatch isaugmented with a random rotation. 3): DrQ [8] ASR: same architecture as (1) but each Q target and Q estimate are calculated by averaging over several augmented versions of the sampled transition.
The augmentation in (2) and (3) is by random rotations sampled from C 32 , the same group usedin the equivariant model.
4) the Transporter Network [1].
See Appendix E.2 for the networkarchitecture for the baselines.
The results in Fig 7 show that equivariant ASR outperforms the othermethods on all tasks, followed by DrQ, RAD, and Transporter Net, followed by Conventional ASR.
[Table]
Robot Experiment: We evaluate the trained equivariant ASR [Table] models 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 robot arranges the bottles outside of the tray.
In the House Building task, therobot succeeds in all 20 episodes.
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 positionis offset slightly from the the rest of the stack and there is no room to place the last box.
The sameproblem happens in another successful episode, where the ﬁngers squeeze the boxes and make roomfor 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 be extended into SE(3) by adding networks similar to q 2that encode Q values for additional dimensions ofthe action space [29].
Speciﬁcally, we add threenetworks to the SE(2) -equivariant architecture described in Section 4.3, q 3 , q 4 , and q 5 encoding Qvalues for Z (height above the plane), and angles (rotation in XZ plane) and (rotation in YZ plane)dimensions of SE(3) .
Each of these networks takesas input a stack of 3 orthographic projections of apoint cloud along the coordinate planes.
The pointcloud is re-centered and rotated to encode the partial SE(3) actions.
(see [29] for details).
Unfortunately, we cannot easily make these additional networks equivariant using the same methods we have proposed so far because they encodevariation outside of SE(2) .
Instead, we create anencoding that is approximately equivariant by explicitly transforming the input to these networks ina way that corresponds to a set of candidate robotpositions or orientations (called a “deictic” encoding in [22]).
We will describe this idea using q 2 as an example.
Deﬁne q 0
( P ) [-
 Rto-]2 R to output thesingle Q value of the rotation encoded in the input image P .
Then q 2 can be deﬁned as a vectorvalued 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 approximatelyequivariant because everything learned by q 0 2 is automatically replicated in all orientations.
Wedesign 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- 351 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) .