Attention Privileged Reinforcement Learning for Domain Transfer

1 INTRODUCTION . Deep Reinforcement Learning ( RL ) has recently provided significant successes in a range of areas , including video games ( Mnih et al. , 2015 ) , board games ( Silver et al. , 2017 ) , simulated continuous control tasks ( Lillicrap et al. , 2015 ) , and robotic manipulation ( Haarnoja et al. , 2018 ; Haarnoja , 2018 ; Riedmiller et al. , 2018 ; OpenAI et al. , 2018 ; Schwab et al. , 2019 ; Andrychowicz et al. , 2017 ) . However , application to physical systems has proven to be challenging in general , due to expensive and slow data generation as well as safety challenges when running untrained policies . A common approach to circumvent these issues is to transfer models trained in simulation to the real world ( Tobin et al. , 2017 ; Rusu et al. , 2016 ; Held et al. , 2017 ) . However , simulators only represent approximations of a physical system . Due to physical , visual , and behavioural discrepancies , naively transferring RL agents trained in simulation onto the real world can be challenging . To bridge the gap between simulation and the real world , we can either aim to align both domains ( Ganin et al. , 2016 ; Bousmalis et al. , 2016 ; Wulfmeier et al. , 2017 ) or ensure that the real system is covered by the distribution of simulated training data ( OpenAI et al. , 2018 ; Tobin et al. , 2017 ; Pinto et al. , 2018 ; Sadeghi & Levine , 2016 ; Viereck et al. , 2017 ) . However , training under a distribution of randomised visual attributes of the simulator , such as textures and lighting ( Sadeghi & Levine , 2016 ; Viereck et al. , 2017 ) , as well as physics ( OpenAI et al. , 2018 ) , can be substantially more difficult and slower due to the increased variability of the learning domain ( OpenAI et al. , 2018 ; Tobin et al. , 2017 ) . The more structured and informative the input representation is with respect to the task , the quicker the agent can be trained . A clear example of this effect can be found when an agent is trained with image inputs , versus training with access to the exact simulator states ( Tassa et al. , 2018 ; Pinto et al. , 2018 ) . However , visual perception is more general and access to more compressed representations can often be limited . When exact states are available during training but not deployment , we can make use of information asymmetric actor-critic methods ( Pinto et al. , 2018 ; Schwab et al. , 2019 ) to train the critic faster via access to the state while providing only images for the actor . 1Videos comparing the policy behaviours of APRiL to the asymmetric DDPG baseline can be found here By introducing Attention Privileged Reinforcement Learning ( APRiL ) , we aim to further leverage access to exact states . APRiL leverages states not only to train the critic , but indirectly also for an image-based actor . Extending asymmetric actor-critic methods , APRiL concurrently trains two actor-critic systems ( one symmetric , state-based agent , and another asymmetric agent with imagedependent actor ) . Both actors utilise an attention mechanism to filter input data and by having access to the simulation rendering system , we can optimise image and state based attention masks to align . By additionally sharing the replay buffer between both agents , we can accelerate the learning process of the image-based actor by training on better performing states that are more quickly discovered by the state-based actor due to its lower dimensional input that is invariant to visual randomisation . The key benefits of APRiL lie in its application to domain transfer . When training with domain randomisation for transfer , bootstrapping via asymmetric information has displayed crucial benefits ( Pinto et al. , 2018 ) . Visual randomisation substantially increases the complexity of the imagebased actor ’ s task . Under this setting , the attention network can support invariance with respect to the irrelevant , but highly varying , parts of the image . Furthermore , the convergence of the statespace actor remains unaffected by visual randomisation . We experimentally demonstrate considerable improvements regarding learning convergence and more robust transfer on a set of continuous action domains including : 2D navigation , 2D locomotion and 3D robotic manipulation . 2 PROBLEM SETUP . Before introducing Attention Privileged Reinforcement Learning ( APRiL ) , this section provides a background for the RL algorithms used . For a more in-depth introduction please refer to Lillicrap et al . ( 2015 ) and Pinto et al . ( 2018 ) . 2.1 REINFORCEMENT LEARNING . We describe an agent ’ s environment as a Partially Observable Markov Decision Process which is represented as the tuple ( S , O , A , P , r , γ , s0 ) , where S denotes a set of continuous states , A denotes a set of either discrete or continuous actions , P : S×A×S → { x ∈ R|0 ≤ x ≤ 1 } is the transition probability function , r : S × A → R is the reward function , γ is the discount factor , and s0 is the initial state distribution . O is a set of continuous observations corresponding to continuous states in S. At every time-step t , the agent takes action at = π ( ·|st ) according to its policy π : S → A . The policy is optimised as to maximize the expected return Rt = Es0 [ ∑∞ i=t γ i−tri|s0 ] . The agent ’ s Q-function is defined as Qπ ( st , at ) = E [ Rt|st , at ] . 2.2 ASYMMETRIC DEEP DETERMINISTIC POLICY GRADIENTS . Asymmetric Deep Deterministic Policy Gradients ( asymmetric DDPG ) ( Pinto et al. , 2018 ) represents a type of actor-critic algorithm designed specifically for efficient learning of a deterministic , observation-based policy in simulation for sim-to-real transfer . This is achieved by leveraging access to more compressed , informative environment states , available in simulation , to speed up and stabilise training of the critic . The algorithm maintains two neural networks : an observation-based actor or policy πθ : O → A ( with parameters θ ) used during training and test time , and a state-based Q-function ( also known as critic ) Qπφ : S ×A→ R ( with parameters φ ) which is only used during training . To enable exploration , the method ( like its symmetric version ( Silver et al. , 2014 ) ) relies on a noisy version of the policy ( called behavioural policy ) , e.g . πb ( o ) = π ( o ) + z where z ∼ N ( 0 , 1 ) ( see Appendix C for our particular instantiation ) . The transition tuples ( st , ot , at , rt , st+1 , ot+1 ) encountered during training are stored in a replay buffer ( Mnih et al. , 2015 ) . Training examples sampled from the replay buffer are used to optimize the critic and actor . By minimizing the Bellman error loss Lcritic = ( Q ( st , at ) − yt ) 2 , where yt = rt + γQ ( st+1 , π ( ot+1 ) ) , the critic is optimized to approximate the true Q values . The actor is optimized by minimizing the loss Lactor = −Es , o∼πb ( o ) [ Q ( s , π ( o ) ) ] . 3 ATTENTION PRIVILEGED REINFORCEMENT LEARNING ( APRIL ) . APRiL proposes to improve the performance and sample efficiency of an observation-based agent by using a quicker learning actor that has access to exact environment states , sharing replay buffers , and aligning attention mechanisms between both actors . While we focus in the following sections on extending asymmetric DDPG ( Pinto et al. , 2018 ) , these ideas are generally applicable to off-policy actor-critic methods ( Konda & Tsitsiklis , 2000 ) . APRiL is comprised of three modules as displayed in Figure 1 . The first two modules , As and Ao , are actor-critic algorithms with an attention network incorporated over the input to each actor . For the state-based module As we use standard symmetric DDPG , while the observation-based module Ao builds on asymmetric DDPG . Finally , the third partAT represents the alignment process between attention mechanisms of both actor-critic agents to more effectively transfer knowledge between the quicker and slower learners , As and Ao , respectively . As consists of three networks : Qπs , πs , hs ( respectively critic , actor , and attention ) with parameters { φs , θs , ψs } . Given input state st , the attention network outputs a soft gating mask ht of same dimensionality as the input , with values ranging between [ 0 , 1 ] . The input to the actor is an attentionfiltered version of the state , sat = hs ( st ) st. To encourage a sparse masking function , we found that training this attention module on both the traditional DDPG loss as well as an entropy loss helped : Lhs = −Es∼πb [ Qs ( s , πs ( sa ) ) − βH ( hs ( s ) ) ] , ( 1 ) where β is a hyperparameter to weight the additional entropy objective , and πb is the behaviour policy used to obtain experience ( in this case from a shared replay buffer ) . The actor and critic networks πs and Qs are trained with the symmetric DDPG actor and Bellman error losses respectively . Within AT , the state-attention obtained in As is converted to corresponding observation-attention T to act as a self-supervised target for the observation-based agent in Ao . This is achieved in a twostep process . First , state-attention hs ( s ) is converted into object-attention c , which specifies how task-relevant each object in the scene is . Second , object-attention is converted to observation-space attention by performing a weighted sum over object-specific segmentation maps : c =M · hs ( s ) ( 2 ) T = N−1∑ i=0 ci · zi ( 3 ) Here , M ∈ { 0 , 1 } N×ns ( where ns is the dimensionality of s ) is an environment-specific , predefined adjacency matrix that maps the dimensions of s to each corresponding object , and c ∈ [ 0 , 1 ] N is then an attention vector over the N objects in the environment . ci corresponds to the ith object attention value . zi ∈ { 0 , 1 } W×H is the binary segmentation map2 of the ith object segmenting the object with the rest of the scene , and has the same dimensions as the image observation . zi assigns values of 1 for pixels in the image occupied by the ith object , and 0 elsewhere . T ∈ [ 0 , 1 ] W×H is the converted state-attention to observation-space attention to act as a target to train the observationattention network ho on . The observation-based module Ao also consists of three networks : Qπo , πo , ho ( respectively critic , actor , and attention ) with parameters { φo , θo , ψo } . The structure of this module is the same as As except the actor and critic now have asymmetric inputs . The input to the actor is the attentionfiltered version of the observation , oat = ho ( ot ) ot3.The actor and critic networks πo and Qo are trained with the standard asymmetric DDPG actor and Bellman error losses respectively defined in Section 2.2 . The main difference between Ao and As is that the observation attention network ho is trained on both the actor loss and an object-weighted mean squared error loss : Lho = Eo , s∼πb 1 2 ∑ ij 1 wij ( ho ( o ) − T ) ij 2 − νQo ( s , πo ( oa ) ) , ( 4 ) where weights wij correspond to the fraction of the partial observation o that the object present in oi , j,1:3 occupies , and ν represents the relative weighting of both loss components . The weight terms , w , ensure that the attention network becomes invariant to the size of objects during training and does not simply fit to the most predominant object in the scene . Combining the self-supervised attention loss and the RL loss leverages efficient state-space learning unaffected by visual randomisation . During training , experiences are collected evenly from both state and observation based agents and stored in a shared replay buffer ( similar to Schwab et al . ( 2019 ) ) . This is to ensure that : 1 . Both statebased criticQs and observation-based criticQo observe states that would be visited by either of their respective policies . 2 . The attention modules hs and ho are trained on the same data distribution to better facilitate alignment . 3 . Efficient discovery of highly performing states from πs are used to speed up learning of πo . Algorithm 1 shows pseudocode for a single actor implementation of APRiL . In practice , in order to speed up data collection and gradient computation , we parallelise the agents and environments and ensure data collection from state- and image- based agents is even .

Summary: Building on top of the domain randomization principle (used to train policies robust to domain-variations) to learn policies which transfer well to new domains, the paper proposes an approach to improve and speed-up learning / training over randomized environments. The paper operates in a settings where the policy to be transferred only has access to observations -- images, etc -- and not the complete underlying state of a (simulated) environment. The underlying idea is to -- (1) maintain two sets of actor-critic networks - a symmetric pair where the actor has access to the underlying state and an asymmetric pair where the actor has access only to image observations; (2) evenly gather experiences from behavioral policies of both actors and store them in a shared replay buffer and (3) learn to align the attention placed by the policies over objects in the environment for the state and observation based actors. The idea is to leverage privileged information about the state (which is strictly more informative compared to observations) to learn robust observation based policies. Experimental results indicate the proposed approach improves generalization performance compared to several ablations of the same on both in-distribution and out-of-distribution environments.