Hard Attention Control By Mutual Information Maximization

1 INTRODUCTION . Reinforcement learning ( RL ) algorithms have successfully employed neural networks over the past few years , surpassing human level performance in many tasks ( Mnih et al. , 2015 ; Silver et al. , 2017 ; Berner et al. , 2019 ; Schulman et al. , 2017 ) . But a key difference in the way tasks are performed by humans versus RL algorithms is that humans have the ability to focus on parts of the state at a time , using attention to limit the amount of information gathered at every step . We actively control our attention to build an internal representation of our surroundings over multiple fixations ( Fourtassi et al. , 2017 ; Barrouillet et al. , 2004 ; Yarbus , 2013 ; Itti , 2005 ) . We also use memory and internal world models to predict motions of dynamic objects in the scene when they are not under direct observation ( Bosco et al. , 2012 ) . By limiting the amount of input information in these two ways , i.e . directing attention only where needed and internally modeling the rest of the environment , we are able to be more efficient in terms of data that needs to be collected from the environment and processed at each time step . By contrast , modern reinforcement learning methods often operate on the entire state . Observing the entire state simultaneously may be difficult in realistic environments . Consider an embodied agent that must actuate its camera to gather visual information about its surroundings learning how to cross a busy street . At every moment , there are different objects in the environment competing for its attention . The agent needs to learn to look left and right to store locations , heading , and speed of nearby vehicles , and perhaps other pedestrians . It must learn to create a live map of its surroundings , frequently checking back on moving objects to update their dynamics . In other words , the agent must learn to selectively move its attention so as to maximize the amount of information it collects from the environment at a time , while internally modeling motions of the other , more predictable parts of the state . Its internal representation , built using successive glimpses , must be sufficient to learn how to complete tasks in this partially observable environment . We consider the problem of acting in an environment that is only partially observable through a controllable , fixed-size , hard attention window ( see figure 1 ) . Only the part of the state that is under the attention window is available to the agent as its observation . The rest must be inferred from previous observations and experience . We assume that the location of the attention window at every time step is under control of the agent and its size compared to the full environment state is known to the agent . We distinguish this task from that of learning soft attention ( Vaswani et al. , 2017 ) , where the full state is attended to , weighted by a vector , and then fed into subsequent layers . Our system must learn to 1 ) decide where to place the attention in order to gather more information about its surroundings , 2 ) record the observation made into an internal memory and model the motion within parts of the state that were unobserved , and 3 ) use this internal representation to learn how to solve its task within the environment . Our approach for controlling attention uses RL to maximize an information theoretic objective closely related to the notion of surprise or novelty ( Schmidhuber , 1991 ) . It is unsupervised in terms of environment rewards , i.e . it can be trained on offline data ( states and actions ) without knowing the task or related rewards . We discuss this in more detail in section 5 . Memory also plays a crucial role in allowing agents to solve tasks in partially observable environments . We pair our attention control mechanism with a memory architecture inspired largely by Du & Narasimhan ( 2019 ) ’ s SpatialNet , but modified to work in partially observable domains . This is described in section 4 . Empirically , we show in section 6.1 that our system is able to reconstruct the full state image including dynamic objects at all time steps given only partial observations . Further , we show in section 6.2 that the internal representation built by our attention control mechanism and memory architecture is sufficient for the agent to learn to solve tasks in this challenging partially observable environment . 2 RELATED WORKS . Using hard attention for image classification or object recognition is well studied in computer vision ( Alexe et al. , 2012 ; Butko & Movellan , 2009 ; Larochelle & Hinton , 2010 ; Paletta et al. , 2005 ; Zoran et al. , 2020 ; Welleck et al. , 2017 ) . Attention allows for processing only the salient or interesting parts of the image ( Itti et al. , 1998 ) . Similarly , attention control has been applied to tracking objects within a video ( Denil et al. , 2012 ; Kamkar et al. , 2020 ; Yu et al. , 2020 ) . Surprisingly , not a lot of recent work exists on the topic of hard attention control in reinforcement learning domains , where a sequential decision making task has to be solved by using partial observations from under the attention . Mnih et al . ( 2014 ) proposed a framework for hard attention control in the classification setting and a simple reinforcement learning task . Their approach consists of using environment rewards to train an attention control RL agent . Our approach differs mainly in that we train the attention control using our novel information theoretic objective as reward . Mnih et al . ( 2014 ) ’ s approach leads to a task specific policy for attention control , whereas our approach is unsupervised in terms of the task and can be applied generally to downstream tasks in the environment . Our approach also differs in that we use a memory architecture that is more suited to partially observable tasks with 2D images as input , compared to a RNN used by Mnih et al . ( 2014 ) . There has been much prior work on memory and world models for reinforcement learning ( Ha & Schmidhuber , 2018 ; Graves et al. , 2016 ; Hausknecht & Stone , 2015 ; Khan et al. , 2017 ) . The work closest to our own is Du & Narasimhan ( 2019 ) ’ s SpatialNet , which attempts to learn a task-agnostic world model for multi-task settings . Our memory architecture is largely inspired by SpatialNet and adapted to work in the partially observable setting . We also use their PhysEnv environment to evaluate our approach . Closely related work is Neural Map ( Parisotto & Salakhutdinov , 2017 ) , which uses a structured 2D memory map to store information about the environment . Their approach also applies to partially observable RL tasks , but the attention is fixed to the agent . In contrast , we consider the problem of learning to control an attention window that can move independently of the agent location . Recently , Freeman et al . ( 2019 ) showed that world models can be learnt by simply limiting the agent ’ s ability to observe the environment . They apply observational dropout , where the output of the agent ’ s world model , rather than the environment state , is occasionally provided to the policy . We consider the related scenario where only a part of the true environment state is provided to the agent at each time step and the rest must be modeled using previous observations . Finally , mutual information has been used to train self-supervised RL agents . This line of work originates in curiosity driven and intrinsically motivated RL ( Schmidhuber , 1991 ; Pathak et al. , 2017 ; Bellemare et al. , 2016 ) . Typically , some notion of predictive error or novelty about an aspect of the environment is optimized in lieu of environment rewards . Multiple papers have successfully used different formulations of mutual information to learn how to efficiently explore the environment without extrinsic rewards ( Mohamed & Rezende , 2015 ; Houthooft et al. , 2016 ; Achiam et al. , 2018 ; Gregor et al. , 2016 ; Eysenbach et al. , 2018 ; Sharma et al. , 2019 ) . We apply the idea of using mutual information to the problem of curiosity driven attention control . 3 PRELIMINARIES . In this work , we are concerned with tasks solved using reinforcement learning ( RL ) . RL is the study of agents acting in an environment to maximize some notion of long term utility . It is almost always formalized using the language of Markov decision problems ( MDPs ) . States , actions , rewards and transitions form the components of an MDP , often represented as the tuple 〈S , A , R , T 〉 . Maximizing the expected sum of ( discounted ) rewards over rollouts in the environment is usually set as the objective of learning . Once a problem is formulated as an MDP , such that the components 〈S , A , R , T 〉 are well defined , one can apply a range of model-free RL algorithms to attempt to solve it ( Schulman et al. , 2017 ; Mnih et al. , 2016 ; Wu et al. , 2017 ) . A solution to an MDP is typically sought as an optimal policy , π∗ : S → A , which is a function that maps every state to an action that maximizes the long term expected rewards . In this section , we will attempt to formalize the components of the partially observable reinforcement learning problem under study . 3.1 INTERNAL REPRESENTATION First , we give a brief description of how the agent stores observations in an internal representation of its surroundings ( figure 2 ) . An map ( µ ) tracks the full environment state ( s ) , which is never directly observed . The map is empty at the start of the episode and gets sequentially written into as observations are made . The agent also creates a reconstruction of the full state ( τ ) at every time step based on its map . This will become useful later for training . For now , it suffices that µt is the map before an observation is made and µwt is the one after the observation is made . µt+1 is the map after the dynamics of the system for the next time step are taken into account . 3.2 THE TWO AGENTS . We formulate the solutions to controlling the attention at every time step and completing the environment task as two separate reinforcement learning agents . The attention location is controlled by the glimpse agent ( the eye in figure 2 ) , and actions within the environment are taken by the regular agent . These two agents have their own separate MDPs with their 〈S , A , R〉 tuples defined below . The glimpse agent ’ s state at every timestep is µt . Its set of actions is all possible attention locations within the full state ( Width × Height actions ) . Its reward is based on the information theoretic objective discussed in section 5 . Thus , the glimpse agent is provided the map before an observation is made and it must decide where the observation should be made from in order to optimize its reward . The environment agent acting in the environment receives as input µwt , i.e . the internal representation after the observation has been recorded . Its actions are the normal set of actions in the environment and its reward is the normal environment reward . We emphasize that neither agent has access to the full state at any time . They must both act based on the internal representation alone . They also can not make multiple observations from the same environment state . Once an observation is made , an action must be selected that will change the environment state . In the next section , we will describe in detail how µwt , µt+1 and τt are formed and how the internal map is trained through a sequence of partial observations .

Summary: This work presents a method for learning a hard attention controller using an information maximization approach. As the authors point out, such a method could be very useful for reasoning in terms of high-dimensional observations, like vision. In brief, the method learning to choose the next attention position to be the most informative by maximizing the uncertainty of the next observation. Uncertainty is quantified using a spatial memory model that is trained to reconstruct and predict the scene. The authors validate this approach by showing that the resulting attention mechanism can be used for two simple downstream tasks. The resulting agent outperforms others trained using baseline attention mechanisms: a hard attention mechanism that is trained on task reward ("environment"; similar to Mnih et al 2014), as well as models that attend to random positions or to the agent's location.