Zero-Shot Policy Transfer with Disentangled Attention

1 INTRODUCTION . RL agents learn to maximize rewards within a task by taking actions based on observations . The advent of deep learning has enabled RL agents to learn in high dimensional feature spaces and complex action spaces ( Krizhevsky et al. , 2012 ; Mnih et al. , 2016 ) . Deep RL methods have beat human performance in a variety of tasks , such as Go ( Silver & Hassabis , 2016 ) . However , deep RL has two crippling drawbacks : high sample complexity and specificity to training task . There are many domains where data collection is expensive and time consuming , such as healthcare , autonomous vehicles , and robotics ( Gottesman et al. , 2019 ; Barrett et al. , 2010 ) . Thus , agents are often trained in simulation and must transfer the resulting knowledge to reality . While this solves the issue of sample complexity , reality and simulated domains are sufficiently different that it is infeasible to naively train a deep RL agent in a simulation and transfer . This is known as the reality gap ( Sadeghi & Levine , 2016 ) . Jumping the reality gap is difficult for two orthogonal reasons . The first is that dynamics of a simulation are an approximation to the dynamics of the real world . Prior work has shown success in transfer between domains with different dynamics ( Killian et al. , 2017 ; Yao et al. , 2018 ; Doshi-Velez & Konidaris , 2016 ) . In this paper , we address the second difficulty : the difference in visual observations of states . Due to limitations in current photorealistic rendering , simulation and the real world are effectively two different visual domains ( Sadeghi & Levine , 2016 ) . We present a method of robust transfer between visual RL domains , using attention and a β variational autoencoder to automatically learn a state representation sufficient to solve both source and target domains . By learning disetangled and relevant state representation , our approach does not require target domain samples when training . The state representation enables the RL agent to attend to only the relevant state information and ignore all other , potentially distracting information . 2 RELATED WORK . Domain randomization is currently the most popular transfer method between different visual domains in RL ( Tobin et al. , 2017 ; OpenAI et al. , 2018 ; Sadeghi & Levine , 2016 ) . By training on many source domains , an RL agent implicitly learns to ignore factors of variation present . Sample complexity scales with the number of source domains and randomized factors . Modes of variation are maually selected . When faced with different variation , domain randomization will fail . Furthermore , there is evidence that domain randomization can destablize the training of some RL algorithms , such as A3C and DDPG ( Matas et al. , 2018 ; Zhang et al. , 2019 ) . Though domain randomization has shown success in transfer , it has three downsides : high sample complexity , manual selection of irrelevant features , and ustable training ( Tobin et al. , 2017 ; OpenAI et al. , 2018 ) . Some visual domain adaptation work is based on image to image translation ( Ajakan et al. , 2014 ; Tzeng et al. , 2017 ; Hoffman et al. , 2017 ; Zhu et al. , 2017 ; Pan et al. , 2017 ; Isola et al. , 2017 ; Liu et al. , 2017 ) . These methods learn a mapping from source domain inputs to target domain inputs , utilizing adversarial methods , such as Generative adversarial Networks trained on data from both the source and target domain . During inference , the target domain input is translated to the source domain before being processed . While these methods have shown promising results , they have two downsides . First , translation incurs additional overhead at inference time . This overhead is especially large since the inputs are high-dimensional images . Second , these methods require samples from the target domain . In applications such as robotics , sampling the target domain is expensive and difficult to obtain . Further , the target domain may not be known at training time . Some recent work focuses on modifying the inputs from the source domain so that they are similar to the target domain , effectively matching their distributions ( Shrivastava et al. , 2017 ) . This approach uses a GAN to refine images from a ( simulated ) source domain to match images from the target domain : the real world . While this directly addresses domain shift and does not add any additional overhead at inference time in the target domain , it still assumes that samples from the target domain are present during training . It enables transfer between two differing visual domains , but it does not enable adaptation to previously unseen domains . There is a body of work that takes this idea further . Rather than translating from the target domain to the source domain , it learns to map all domains to a canonical domain in which the RL agent operates ( James et al. , 2019 ) . By learning to translate from randomized domains to a canonical domain , effectively training a translation network using domain randomization . While this approach has shown simulation to reality transfer , it icurrs overhead through the use of the mapping network ; Rather than learning to map to a canoical domain , it would be more efficient to map to a disentagled representation of the canoical environment . Other work learns to map image inputs to a latent space for use in RL . Specifically , it utilizes a β Variational Autoencoder ( β-VAE ) to learn a compressed , disentagled state representation ( Higgins et al. , 2017 ) . While this work focuses on domain adaptation in the RL setting , it utilizes a generalpurpose generative model . Thus , the state representation is not tailored to the task posed ; it preserves information needed to reconstruct the image that may not be needed to solve the RL task given . Additionally , unlike image translation , this method does not attempt to match the distribution of inputs from source and target domains . Since the β-VAE is trained to reconstruct input states , the distributions of compressed and input states will match . Accordingly , the distribution shift between source and target domains will be preserved , even when using compressed states . To circumvent this issue , the RL agent is trained on multiple source domains , as in domain randomization . 3 BACKGROUND . We formalize the concept of transfer between related MDPs . We denote the source domain as DS and the target domain as DT . Both domains are MDPs , defined as D = ( S , A , T , R ) , where S are states , A are actions , T is the transition function , and R is the reward function . We assume that both DS and DT share a discount factor γ . In general transfer , states and actions may be quite different , while transition and reward functions are similar . We denote a family of related MDPs as M . Within M , transitions , rewards , and actions are the same across MDPs , and states are different . These states are generated from a set of latent state factors Sz , each passed through a different MDPspecific renderer GU for some MDP U ∈M . In the case of visual cartpole , Sz contains the position and velocity of the cart , the angle and position of the pole , and the colors of the cart , pole , axle , and track . The observed states are then rendered by the renderer . Specifically , SU = GU ( Sz ) and SV = GV ( S z ) , where GU and GV are any two MDPs within M . In this paper , we target transfers where TS ≈ TT , RS = RT , AS = AT , but SS = GS ( Sz ) 6= ST = GT ( S z ) . Namely , when target dynamics are similar to source dynamics , rewards , actions , and true states are identical , but observed states SS and ST differ due to differing renderers GS and GT . 4 SADALA . 4.1 FEATURE EXTRACTION WITH β-VARIATIONAL AUTOENCODERS Deep RL agents learn a mapping from states s ∈ S to actions a ∈ A , either by directly maximizing reward or through a Q function , using a deep neural network . The deep network implicitly learns a mapping function F : SU → Ŝz , for some MDP U ∈ M and a policy function π : Ŝz → A that maps from latent state factors to actions . These functions are learned as a result of the feed-forward nature of deep neural networks . The policy π is a function of the last layer of the network . The output of the last layer of the neural network is a highly processed feature representation of the input that is optimized for use in generating π . While this implicit mapping is sufficient for a deep RL agent to solve a single MDP from M , it is not general enough to transfer to other MDPs within M . By optimizing for π F overfits to the specific SU from the source domain U . To circumvent this overfitting , DARLA proposed learning a feature extractor F ′ that maps from any observed state SX to the latent state factors Sz ( Higgins et al. , 2017 ) . The RL agent is the trained on the ( approximate ) extracted latent state factors Ŝz , separating the “ vision ” and “ action ” of the RL agent . DARLA uses a β variational autoencoder to learn the mapping F ′ ( Higgins et al. , 2017 ) . The extracted latent state factors Ŝz are ( mostly ) disentangled in practice , leading to better transfer results than a baseline RL agent . The factors each correlate with an independent factor in the true state Sz of the task . However , it preserves all latent state variables z ∈ Sz , incorrectly assuming that they are all relevant to solving the MDP . When attempting to locate an object such as an enemy in a video game , an object for a robot , or a road for an autonomous vehicle , the position and orientation of the object are the most important factors . Visual distractions such as the color of the enemy , opacity of the object , or texture of the road are entirely irrelevant to its localization . Dependence on these distracting factors causes the RL agent to fail to robustly transfer to target visual domains . 4.2 ADDING AN ATTENTION MECHANISM . SADALA is similar to DARLA , but has one key difference—it explicitly learns to ignore irrelevant extracted latent state variables . SADALA learns the same mapping function F ′ : SU → Ŝz as DARLA , also utilizing a β variational autoencoder . It then learns a policy π : Ŝz → A , but utilizes a soft attention mechanism to ignore irrelevant features . Using this mechanism , it implicitly learns two functions : A : Ŝz → Sa that maps from extracted latent state factors to latent state factors relevant to solving the MDP , and πA : S a → A that maps from relevant latent state factors to actions . Soft attention mechanisms learn an importance weight wi ∈ [ 0 , 1 ] for each feature zi in part of a neural network . These weights are then multiplied element-wise with the corresponding feature . Thus , the attention layer can be defined as follows : f ( Z ) =W Z , where W is the vector of importance weights wi and is element-wise multiplication . W is defined as W = σ ( h ( Z ) ) where h is a differentiable function parameterized by dense layers of a neural network . Thus , the values wi ∈ W are constrained between 0 and 1 , inclusive . This allows the RL agent to ignore a feature zi by setting wi to 0 . It also allows the agent to down-weight the importance of a feature by setting wi to a fractional value or pay full attention by setting wi to 1 .

Summary: The paper proposes adding an attention mechanism to the DARLA beta-VAE approach to transfer learning. The beta-VAE, soft attention and policy are trained on appropriate source tasks and evaluated zero-shot on target tasks, using two more difficult continuous control domains with RGB observations. Results indicate some improvements to compared to the immediate relevant baseline which may be statistically significant, but it is not clear whether over 10% in practice.