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

1 INTRODUCTION . A key challenge in reinforcement learning is to improve sample efficiency – that is to reduce the amount of environmental interactions that an agent must take in order to learn a good policy . This is particularly important in robotics applications where gaining experience potentially means interacting with a physical environment . One way of improving sample efficiency is to create “ artificial ” experiences through data augmentation . This is typically done in 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-invariant feature representations ( Laskin et al. , 2020b ; Zhan et al. , 2020 ) . Recent work in geometric deep learning suggests that it may be possible to learn transformationinvariant policies and value functions in a different way , using equivariant neural networks ( Cohen & Welling , 2016a ; b ) . The key idea is to structure the model architecture such that it is constrained only to represent functions with the desired invariance properties . In principle , this approach aim at exactly the same thing as the data augmentation approaches described above – both methods seek to 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 training 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 augmentation results only in approximate equivariance whereas equivariant neural networks guarantee it and often have stronger generalization as well ( Wang et al. , 2020b ) . While equivariant architectures have recently been applied to reinforcement learning ( van der Pol et 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 important class of MDPs that we call group equivariant MDPs . Second , we introduce a new variation of the Equivariant DQN ( Mondal et al. , 2020 ) , and we further introduce equivariant variations of SAC ( Haarnoja et al. , 2018 ) , and learning from demonstration ( LfD ) . Finally , we show that our methods convincingly outperform recent competitive data augmentation approaches ( Laskin et al. , 2020a ; Kostrikov et al. , 2020 ; Laskin et al. , 2020b ; Zhan et al. , 2020 ) . Our Equivariant SAC method , in particular , outperforms these baselines so dramatically ( Figure 7 ) that it could make reinforcement learning feasible for a much larger class of robotics problems than is currently the case . 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 GConvolution ( Cohen & Welling , 2016a ) . The extension work proposes an alternative architecture , Steerable CNN ( Cohen & Welling , 2016b ) . Weiler & Cesa ( 2019 ) proposes a general framework for implementing E ( 2 ) -Steerable CNNs . In the context of reinforcement learning , Mondal et al . ( 2020 ) investigates the use of Steerable CNNs in the context of two game environments . van der Pol et al . ( 2020b ) 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 et al . ( 2021 ) learns equivariant Q-functions but is limited in the spatial action space . In contrast to prior work , this paper proposes an Equivariant SAC algorithm , an equivariant LfD algorithm , and a 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 augmentation . 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 samples ( Kalashnikov et al. , 2018 ; Lin et al. , 2020 ; Zeng et al. , 2020 ) 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 ) to improve feature extraction . Laskin et al . ( 2020b ) show significant sample-efficiency 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 an invariant feature encoder and is not capable of learning equivariant functions . Close-Loop Robotic Control : There are two typical action space definitions when learning policies that control the end-effector of a robot arm : the spatial action space that controls the target pose of the end-effector ( Zeng et al. , 2018b ; a ; Satish et al. , 2019 ; Wang et al. , 2020a ) , or the closeloop action space that controls the displacement of the end-effector . The close-loop action space is widely used for learning grasping policies ( Kalashnikov et al. , 2018 ; Quillen et al. , 2018 ; Breyer et al. , 2019 ; James et al. , 2019 ) . Recently , some works also attempt to 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 in the close-loop action space by using equivariant learning to improve the sample efficiency . 3 BACKGROUND . SO ( 2 ) and Cn : We will reason about rotation in terms of the group SO ( 2 ) and its cyclic subgroup Cn ≤ SO ( 2 ) . SO ( 2 ) is the group of continuous planar rotations { Rotθ : 0 ≤ θ < 2π } . Cn is the discrete subgroup Cn = { Rotθ : θ ∈ { 2πin |0 ≤ i < n } } of rotations by multiples 2π n . Cn actions : A groupGmay be equipped with an action on a setX by specifying a map · : G×X → X satisfying g1 · ( g2 · x ) = ( g1g2 ) · x and 1 · x = x for all g1 , g2 ∈ G , x ∈ X . Note that closure , gx ∈ X , and invertibility , g−1gx = x , follow immediately from the definition . We are interested in actions of Cn 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 ρ0 . Let g ∈ Cn and x ∈ R. Then ρ0 ( g ) x = x . For example , the trivial representation describes how pixel color/depth values change when an image is rotated , i.e . they do not change ( Figure 1 left ) . 2 . R2 through the standard representation ρ1 . Let g ∈ Cn and v ∈ R2 . Then ρ1 ( g ) v = ( cos g − sin g sin g cos g ) v. This describes how elements of a vector field change when rotated ( Figure 1 middle ) . 3 . Rn through the regular representation ρreg . Let g ∈ Cn and ( x1 , x2 , . . . , xn ) ∈ Rn . Then ρreg ( g ) ( x1 , x2 , . . . , xn ) = ( xn , x1 , x2 , . . . , xn−1 ) cyclically permutes the coordinates of Rn ( Figure 1 right ) . Feature maps as functions : In deep learning , images and feature maps are typically expressed as tensors . However , it will be convenient here to sometimes express these as functions . Specifically , we may write an h× w one-channel image F ∈ R1×h×w as a function F : R2 → R where F ( x , y ) describes the intensity at pixel x , y . Similarly , an m-channel tensor F ∈ Rm×h×w may be written as F : R2 → Rm . We refer to the domain of this function as its “ spatial dimensions ” . Cn actions on vectors and feature maps : Cn acts on vectors and feature maps differently depending upon their semantics . We formalize these different ways of acting as follows . Let F : R2 → Rm be an m-channel feature map and let V ∈ Rm×1×1 = Rm be a vector represented as a special case of a feature map with 1× 1 spatial dimensions . Then g is defined to act on F by ( gF ) ( x , y ) = ρj ( g ) F ( ρ1 ( g ) −1 ( x , y ) ) . ( 1 ) For a vector V ( considered to be at ( x , y ) = ( 0 , 0 ) ) , this becomes : gV = ρj ( g ) V. ( 2 ) In the above , ρ1 ( g ) rotates pixel location and ρj ( g ) transforms the pixel feature vector using the trivial representation ( ρj = ρ0 ) , the standard representation ( ρj = ρ1 ) , the regular representation ( ρj = ρreg ) , or some combination thereof . Equivariant convolutional layer : A Cn-equivariant layer is a function h whose output is constrained to transform in a defined way when the input feature map is transformed by a group action . Consider an equivariant layer h with an input Fin : R2 → R|ρin| and an output Fout : R2 → R|ρout| , where ρin and ρout denote the group representations associated with Fin and Fout , respectively . When the input is transformed , this layer is constrained to output a transformed version of the same output feature map : h ( gFin ) = g ( h ( Fin ) ) = gFout . ( 3 ) where g ∈ Cn acts on Fin or Fout through Equation 1 or Equation 2 , i.e. , this constraint equation can be applied to arbitrary feature maps F or vectors V . A linear convolutional layer h satisfies Equation 3 with respect to the group Cn if the convolutional kernel K : R2 → R|ρout|×|ρin| has the following form ( Cohen et al. , 2018 ) : K ( ρ1 ( g ) v ) = ρ −1 out ( g ) K ( v ) ρin ( g ) . ( 4 ) Since the composition of equivariant maps is equivariant , a fully convolutional equivariant network can be constructed by stacking equivariant convolutional layers that satisfy the constraint of Equation 3 and together with equivariant non-linearities ( Weiler & Cesa , 2019 ) .

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