Contingency-Aware Exploration in Reinforcement Learning

1 INTRODUCTION . The success of reinforcement learning ( RL ) algorithms in complex environments hinges on the way they balance exploration and exploitation . There has been a surge of recent interest in developing effective exploration strategies for problems with high-dimensional state spaces and sparse rewards ( Schmidhuber , 1991 ; Oudeyer & Kaplan , 2009 ; Houthooft et al. , 2016 ; Bellemare et al. , 2016 ; Osband et al. , 2016 ; Pathak et al. , 2017 ; Plappert et al. , 2018 ; Zheng et al. , 2018 ) . Deep neural networks have seen great success as expressive function approximators within RL and as powerful representation learning methods for many domains . In addition , there have been recent studies on using neural network representations for exploration ( Tang et al. , 2017 ; Martin et al. , 2017 ; Pathak et al. , 2017 ) . For example , count-based exploration with neural density estimation ( Bellemare et al. , 2016 ; Tang et al. , 2017 ; Ostrovski et al. , 2017 ) presents one of the state-of-the-art techniques on the most challenging Atari games with sparse rewards . Despite the success of recent exploration methods , it is still an open question on how to construct an optimal representation for exploration . For example , the concept of visual similarity is used for learning density models as a basis for calculating pseudo-counts ( Bellemare et al. , 2016 ; Ostrovski et al. , 2017 ) . However , as Tang et al . ( 2017 ) noted , the ideal way to represent states should be based on what is relevant to solving the MDP , rather than only relying on visual similarity . In addition , there remains another question on whether the representations used for recent exploration works are easily interpretable . To address these questions , we investigate whether we can learn a complementary , more intuitive , and interpretable high-level abstraction that can be very effective in exploration by using the ideas of contingency awareness and controllable dynamics . ∗Equal contributions , listed in alphabetical order . †Now at DeepMind . 1Examples of the learned policy and the contingent regions are available at https : //coex-rl.github.io/ . The key idea that we focus on in this work is the notion of contingency awareness ( Watson , 1966 ; Bellemare et al. , 2012 ) — the agent ’ s understanding of the environmental dynamics and recognizing that some aspects of the dynamics are under the agent ’ s control . Intuitively speaking , this can represent the segmentation mask of the agent operating in the 2D or 3D environments ( yet one can think of more abstract and general state spaces ) . In this study , we investigate the concept of contingency awareness based on self-localization , i.e. , the awareness of where the agent is located in the abstract state space . We are interested in discovering parts of the world that are directly dependent on the agent ’ s immediate action , which often reveal the agent ’ s approximate location . For further motivation on the problem , we note that contingency awareness is a very important concept in neuroscience and psychology . In other words , being self-aware of one ’ s location is an important property within many observed intelligent organisms and systems . For example , recent breakthroughs in neuroscience , such as the Nobel Prize winning work on the grid cells ( Moser et al. , 2015 ; Banino et al. , 2018 ) , show that organisms that perform very well in spatially-challenging tasks are self-aware of their location . This allows rats to navigate , remember paths to previously visited places and important sub-goals , and find shortcuts . In addition , the notion of contingency awareness has been shown as an important factor in developmental psychology ( Watson , 1966 ; Baeyens et al. , 1990 ) . We can think of self-localization ( and more broadly self-awareness ) as a principled and fundamental direction towards intelligent agents . Based on these discussions , we hypothesize that contingency awareness can be a powerful mechanism for tackling exploration problems in reinforcement learning . We consider an instantiation of this hypothesis evaluated on the Arcade Learning Element ( ALE ) . For example , in the context of 2D Atari games , contingency-awareness roughly corresponds to understanding the notion of controllable entities ( e.g. , the player ’ s avatar ) , which Bellemare et al . ( 2012 ) refer to as contingent regions . More concretely , as shown in Figure 1 , in the game FREEWAY , only the chicken sprite is under the agent ’ s control and not the multiple moving cars ; therefore the chicken ’ s location should be an informative element for exploration ( Bellemare et al. , 2012 ; Pathak et al. , 2017 ) . In this study , we also investigate whether contingency awareness can be learned without any external annotations or supervision . For this , we provide an instantiation of an algorithm for automatically learning such information and using it for improving exploration on a 2D ALE environment ( Bellemare et al. , 2013 ) . Concretely , we employ an attentive dynamics model ( ADM ) to predict the agent ’ s action chosen between consecutive states . It allows us to approximate the agent ’ s position in 2D environments , but unlike other approaches such as ( Bellemare et al. , 2012 ) , it does not require any additional supervision to do so . The ADM learns in an online and self-supervised fashion with pure observations as the agent ’ s policy is updated and does not require hand-crafted features , an environment simulator , or supervision labels for training . In experimental evaluation , our methods significantly improve the performance of A2C on hardexploration Atari games in comparison with competitive methods such as density-based exploration ( Bellemare et al. , 2016 ; Ostrovski et al. , 2017 ) and SimHash ( Tang et al. , 2017 ) . We report very strong results on sparse-reward Atari games , including the state-of-the-art performance on the notoriously difficult MONTEZUMA ’ S REVENGE , when combining our proposed exploration strategy with PPO ( Schulman et al. , 2017 ) , without using expert demonstrations , explicit high-level information ( e.g. , RAM states ) , or resetting the environment to an arbitrary state . We summarize our contributions as follows : • We demonstrate the importance of learning contingency awareness for efficient exploration in challenging sparse-reward RL problems . • We develop a novel instance of attentive dynamics model using contingency and controllable dynamics to provide robust localization abilities across the most challenging Atari environments . • We achieve a strong performance on difficult sparse-reward Atari games , including the state-ofthe-art score on the notoriously challenging MONTEZUMA ’ S REVENGE . Overall , we believe that our experiments confirm the hypothesis that contingency awareness is an extremely powerful concept for tackling exploration problems in reinforcement learning , which opens up interesting research questions for further investigations . 2 RELATED WORK . Self-Localization . The discovery of grid cells ( Moser et al. , 2015 ) motivates working on agents that are self-aware of their location . Banino et al . ( 2018 ) emphasize the importance of self-localization and train a neural network which learns a similar mechanism to grid cells to perform tasks related to spatial navigation . The presence of grid cells is correlated with high performance . Although grid cells seem tailored to 2D or 3D problems that animals encounter in their life , it is speculated that their use can be extended to more abstract spaces . A set of potential approaches to self-localization ranges from ideas specific to a given environment , e.g. , SLAM ( Durrant-Whyte & Bailey , 2006 ) , to methods with potential generalizability ( Mirowski et al. , 2017 ; Jaderberg et al. , 2017 ; Mirowski et al. , 2018 ) . Self-supervised Dynamics Model and Controllable Dynamics . Several works have used forward and/or inverse dynamics models of the environment ( Oh et al. , 2015 ; Agrawal et al. , 2016 ; Shelhamer et al. , 2017 ) . Pathak et al . ( 2017 ) employ a similar dynamics model to learn feature representations of states that captures controllable aspects of the environment . This dense representation is used to design a curiosity-driven intrinsic reward . The idea of learning representations on relevant aspects of the environment by learning auxiliary tasks is also explored in ( Jaderberg et al. , 2017 ; Bengio et al. , 2017 ; Sawada , 2018 ) . Our presented approach is different as we focus on explicitly discovering controllable aspects using an attention mechanism , resulting in better interpretability . Exploration and Intrinsic Motivation . The idea of providing an exploration bonus reward depending on the state-action visit-count was proposed by Strehl & Littman ( 2008 ) ( MBIE-EB ) , originally under a tabular setting . Later it has been combined with different techniques to deal with high-dimensional state spaces . Bellemare et al . ( 2016 ) use a Context-Tree Switching ( CTS ) density model to derive a state pseudo-count , whereas Ostrovski et al . ( 2017 ) use PixelCNN as a state density estimator . Martin et al . ( 2017 ) also construct a visitation density model over a compressed feature space rather than the raw observation space . Alternatively , Tang et al . ( 2017 ) propose a localitysensitive hashing ( LSH ) method to cluster states and maintain a state-visitation counter based on a form of similarity between frames . We train an agent with a similar count-based exploration bonus , but the way of maintaining state counter seems relatively simpler in that key feature information ( i.e. , controllable region ) is explicitly extracted from the observation and directly used for counting states . Another popular family of exploration strategies in RL uses intrinsic motivation ( Schmidhuber , 1991 ; Singh et al. , 2004 ; Oudeyer & Kaplan , 2009 ; Barto , 2013 ) . These methods encourage the agent to look for something surprising in the environment which motivates its search for novel states , such as surprise ( Achiam & Sastry , 2017 ) , curiosity ( Pathak et al. , 2017 ; Burda et al. , 2018 ) , and diversity ( Eysenbach et al. , 2018 ) , or via feature control ( Jaderberg et al. , 2017 ; Dilokthanakul et al. , 2017 ) . 3 APPROACH . 3.1 DISCOVERING CONTINGENCY VIA ATTENTIVE DYNAMICS MODEL . To discover the region of the observation that is controllable by the agent , we develop an instance of attentive dynamics model ( ADM ) based on inverse dynamics finv . The model takes two consecutive input frames ( observations ) st−1 , st ∈ S as input and aims to predict the action ( at−1 ∈ A ) taken by the agent to transition from st−1 to st : ât−1 = finv ( st−1 , st ) . ( 1 ) Our key intuition is that the inverse dynamics model should attend to the most relevant part of the observation , which is controllable by the agent , to be able to classify the actions . We determine whether each region in a H ×W grid is controllable , or in other words , useful for predicting the agent ’ s action , by using a spatial attention mechanism ( Bahdanau et al. , 2015 ; Xu et al. , 2015 ) . An overview of the model is shown in Figure 1 . Model . To perform action classification , we first compute a convolutional feature map φst = φ ( st ) ∈ RH×W×K based on the observation st using a convolutional neural network φ . We estimate a set of logit ( score ) vectors , denoted et ( i , j ) ∈ R|A| , for action classification from each grid cell ( i , j ) of the convolutional feature map . The local convolution features and feature differences for consecutive frames are fed into a shared multi-layer perceptron ( MLP ) to derive the logits as : et ( i , j ) = MLP ( [ φst ( i , j ) − φst−1 ( i , j ) ; φst ( i , j ) ] ) ∈ R|A| . ( 2 ) We then compute an attention mask αt ∈ RH×W corresponding to frame t , which indicates the controllable parts of the observation st . Such attention masks are computed via a separate MLP from the features of each region ( i , j ) , and then converted into a probability distribution using softmax or sparsemax operators ( Martins & Astudillo , 2016 ) : αt = sparsemax ( α̃t ) where α̃t ( i , j ) = MLP ( φst ( i , j ) ) , ( 3 ) so that ∑ i , j αt ( i , j ) = 1 . The sparsemax operator is similar to softmax but yields a sparse attention , leading to more stable performance . Finally , the logits et ( i , j ) from all regions are linearly combined using the attention probabilities αt : p ( ât−1 | st−1 , st ) = softmax ( ∑ i , j αt ( i , j ) · et ( i , j ) ) ∈ R|A| . ( 4 ) Training . The model can be optimized with the standard cross-entropy loss Laction ( a∗t−1 , ât−1 ) with respect to the ground-truth action a∗t−1 ∈ A that the agent actually has taken . Based on this formulation , the attention probability αt ( i , j ) should be high only on regions ( i , j ) that are predictive of the agent ’ s actions . Our formulation enables learning to localize controllable entities in a selfsupervised way without any additional supervisory signal , unlike some prior work ( e.g. , ( Bellemare et al. , 2012 ) ) that adopts simulators to collect extra supervisory labels . Optimizing the parameters of ADM on on-policy data is challenging for several reasons . First , the ground-truth action may be unpredictable for given pairs of frames , leading to noisy labels . For example , actions taken in uncontrollable situations do not have any effect ( e.g. , when the agent is in the middle of jumping in MONTEZUMA ’ S REVENGE ) . Second , since we train the ADM online along with the policy , the training examples are not independently and identically distributed , and the data distribution can shift dramatically over time . Third , the action distribution from the agent ’ s policy can run into a low entropy2 , being biased towards certain actions . These issues may prevent the ADM from generalization to novel observations , which hurts exploration . Generally , we prefer models that quickly adapt to the policy and learn to localize the controllable regions in a robust manner . To mitigate the aforementioned issues , we adopt a few additional objective functions . We encourage the attention distribution to attain a high entropy by including an attention entropy regularization loss , i.e. , Lent = −H ( αt ) . This term penalizes over-confident attention masks , making the attention closer to uniform whenever action prediction is not possible . We also train the logits corresponding to each grid cell independently using a separate cross-entropy loss : p ( âi , jt−1 | et ( i , j ) ) = softmax ( et ( i , j ) ) . These additional cross-entropy losses , denoted Li , jcell , allow the model to learn from unseen observations even when attention fails to perform well at first . The entire training objective becomes : LADM = Laction + ∑ i , j L i , j cell + λentLent ( 5 ) where λent is a mixing hyperparameter .

Summary: This paper investigates the problem of extracting a meaningful state representation to help with exploration in RL, when confronted to a sparse reward task. The core idea consists in identifying controllable (learned) features of the state, which in an Atari game for instance typically corresponds to the position of the player-controlled character / vehicle on the screen. Once this position is known (as x, y coordinates on a custom low-resolution grid), one can use existing count-based exploration mechanisms to encourage the agent to visit new positions (NB: in addition to the x, y coordinates, extra information is also used to disambiguate the states for counting purpose, namely the current score and the state’s cluster index obtained with a basic clustering scheme). To find the position, the algorithm trains one inverse dynamics model per x, y cell on the grid: each model tries to predict the action taken by the agent given two consecutive states, both represented by their feature map (at coordinate x, y) learned by a convolutional network applied to the pixel representation. The outputs of these inverse dynamics models are combined through an attention mechanism to output the final prediction for the action: the intuition is that the attention model will learn to focus on the grid cell with best predictive power (for a given state), which should correspond to where the controllable parts of the state are. Experiments on several Atari games (including Montezuma’s Revenge) indeed show that this mechanism is able to track the true agent’s coordinates (obtained from the RAM state) reasonably well. Using these coordinates for count-based exploration (in A2C) also yields significantly better results compared to vanilla A2C, and beats several previously proposed related techniques for exploration in sparse reward settings.