Learning to Control Visual Abstractions for Structured Exploration in Deep Reinforcement Learning

1 INTRODUCTION . Exploration in environments with sparse feedback is a key challenge for deep reinforcement learning ( DRL ) research . In DRL , agents typically explore with local exploration strategies like epsilon greedy or entropy based schemes . We are interested in learning structured exploration algorithms , grounded in spatio-temporal visual abstractions given raw pixels . In human perception and its developmental trajectory , spatio-temporal pixel groupings is one of the first visual abstractions to emerge , which is also used for intrinsically motivated goal-driven behaviors ( Spelke & Kinzler , 2007 ) . Inspired by this insight , we develop a new agent architecture and loss functions to autonomously learn visual abstractions and ground temporally extended behaviors in them . Our approach and key contributions can be broken down into two parts : ( 1 ) an information theoretic loss function and a neural network architecture to learn visual groupings ( abstractions ) given raw pixels and actions , ( 2 ) a hierarchical action-value function agent which explores in the space of options grounded in the learned visual abstractions , instead of low level actions . In the first step , we pass images through an encoder which outputs spatial discrete vector-quantized ( VQ ‘ ) grids , with 1 of E discrete components . We train this encoder to maximize the mutual information between VQ layers at different time steps , in order to obtain a temporally consistent representation , that preserves controllability and appearance information . We extract segmentation masks from the VQ layers for the second step , referred to as visual entities . We compute affine geometric measurements for each entity , namely centroid and area of the corresponding segment . We use off-policy learning to train action-value function to minimize or maximize these measurements , referred collectively as the options bank . Controlling these measurements enable higher levels of behaviors such as approaching an object ( maximizing area ) , avoiding objects ( minimize area or minimize/maximize centroid coordinates ) , moving an object away towards the left ( minimize centroid x coordinate ) , controlling the avatars position on the screen etc . Finally , given a task reward , we use off-policy learning to train a meta action-value function that takes actions at fixed intervals and selects either one of the policies in the options bank or low-level actions . So effectively , this hierarchical action-value function setup solves a semi markov decision process as in ( Sutton et al. , 1999 ; Kulkarni et al. , 2016 ) . We demonstrate that our approach can scale to two different domains – navigation in a 3D environment and challenging Atari games – given raw pixels . Although much work remains in improving the visual and temporal abstraction discovery models , our results indicate that it is possible to learn bottom-up structured exploration schemes with simple spatial inductive biases and loss functions . 2 RELATED WORK . Learning visual abstractions has a long history in computer vision with some of the earlier successes relying on clustering either for inference ( Shi & Malik , 2000 ) or learning ( Ren & Malik , 2003 ) . More recently some of these intuitions were adapted to neural networks as in ( Xia & Kulis , 2017 ) . Instance segmentation algorithms have been developed to output spatio-temporal groupings of pixels from raw videos ( Romera-Paredes & Torr , 2016 ) . However , most of the existing deep learning based approaches require supervised data . Structured deep generative models ( van den Oord et al. , 2017 ) is another approach to learning disentangled representations from raw videos . Very recently segmentation has been cast as a mutual information maximization problem in ( Ji et al. , 2018 ) where the mutual information is computed between the original image segmentation and the output of its transformation ( additive color changes in HSV space , horizontal flips and spatial shifts ) . That approach uses the discrete mutual information estimate which makes it only applicable to enforce pixel-label constraints . For continuous variables recent papers have proposed very promising techniques . Most relevant to our work are ( van den Oord et al. , 2018 ) and ( Belghazi et al. , 2018 ) . In reinforcement learning research , semi-MDPs and the options framework ( Sutton et al. , 1999 ) have been proposed as a model for temporal abstractions of behaviors . Our work is most similar to hierarchical-DQN ( Kulkarni et al. , 2016 ) . However , this approach required hand-crafted instance segmentation and the agent architecture is not distributed to learn about many intrinsic rewards learners at the same time . Object-Oriented-MDPs ( Diuk et al. , 2008 ) uses object oriented representations for structured exploration but requires prebuilt symbolic representations . A recent paper also demonstrates the importance of object based exploration when humans learn to play video games ( Dubey et al. , 2018 ) . HRA ( Van Seijen et al. , 2017 ) is an agent that used prebuilt object representations to obtain state of the art policies on Pacman using object based structured exploration . Another interesting line of work is ( Gregor et al. , 2016 ) which formalizes a notion of empowerment which ends up as a mutual information between options in the same MDP . Count-based exploration algorithms have yielded impressive results on hard exploration Atari games ( Ostrovski et al. , 2017 ) . The Horde architecture ( Sutton et al. , 2011 ) proposed learning many value functions , termed as Generalized Value Functions ( GVFs ) , using off-policy learning . This work was later extended with neural networks by Schaul et al . ( 2015 ) . Our approach automatically constructs GVFs or a UVFA using the abstract entity based representations . Our work is also related to pixel control ( Jaderberg et al. , 2016 ) as an auxiliary task . However , we learn to control and compose abstract discrete representations . Notation . Unless explicitly stated otherwise , we will use calligraphic upper-case letters for sets and MDPs , upper-case letters for random variables , bold capitals for constants and lower-case letters for realizations/measurements . 3 MODEL . Consider a Markov Decision ProcessM = ( S , A , P , r ) ( MDP ) represented by states s ∈ S , actions a ∈ A , transition distribution function P : S × A × S → [ 0 , 1 ] and an extrinsic reward function defined as r : S × A → R. In a discrete MDP , an agent observes a state st at time t , produces and action at then the agent can observe st+1 ∼ P ( S′|S = st , A = at ) and a reward rt+1 = r ( st , at ) . The agent ’ s objective is to maximize the expected sum of rewards over time . In this work we are focusing on visual inputs thus we assume S ⊂ RH×W , where H , W are the height and width of the image . This is a very important special case in current applications but many of the intuitions and machinery we develop should carry over to different domains . Our agent also learns an separate abstract representation we call visual entities1 and denote with v ∈ V , where V ⊂ { 1 . . .E } H×W i.e . it assigns an entity id to each location . These are meant to capture useful information about the state s and form the basis for computing the intrinsic rewards re , m ( see Sec . 3.1 ) . Here e ∈ { 1 . . .E } denotes a discrete entity id , where E is the maximum number of possible entities , and m ∈ { 1 . . .M } denotes a geometric feature of e ’ s derived segmentation mask , from M possible measurements . In this work the measurements are fixed to what we consider a sufficient set that captures the essential information for natural 3D navigation and Atari game play , and settled on centroid cartesian coordinates and entity mask area , which we can both minimize and maximize , thus in all our experiments M = 6 . Temporal changes in these measurements constitute our intrinsic reward functions2 which induce E ×M additional MDPs Oe , m = ( S , A , P , re , m ) . These intrisic rewards will induce behaviors which should hopefully provide structured exploration in the original MDP i.e . picking a random Qe , m is more likely to lead to higher reward than epsilon greedy exploration . Our agent architecture tries to leverage this additional structure . The top level MDP M is represented by Qmeta which outputs action ametat at time t ( switched every T time steps ) , where the action space is discrete 1 of ( E ×M ) + 1 possible actions . In our implementation this is modelled by composite actions E + 1 and M i.e . ametat = ( et , mt ) with et ∈ { 1 . . .E + 1 } and mt ∈ { 1 . . .M } . We also learn ( E ×M ) + 1 separate Q functions : ( E ×M ) that each solve one of the Oe , m and one for the original MDP Qtask . These Q functions are defined over the environment action set and differ only in the reward function ( see Figure 2 for schematic representation ) . 3.1 VISUAL ABSTRACTIONS . Our agent relies on an abstraction model that assigns each pixel in the image to one of E separate abstractions or entities . To obtain this representation , the image is passed through a convolutional network ( CNN ) encoder to output a spatial grid of the same resolution as the original image . Then a vector quantized layer V ( van den Oord et al. , 2017 ) assigns the planar activations to 1 of E entities . From the agents perspective this means all the pixels that are being grouped together become indistinguishable thus providing the visual abstraction we desire . Let us define f : RH×W → { 1 . . .E } H×W the function that takes the observation st at time t and computes an abstract representation corresponding to it vt = f ( st ) . The key question is how to train f . One way is to make it representative of the current state . In most of the current literature , this is 1We use visual entities and visual abstractions interchangeably . Both terms capture the most important qualities of the representation : 1 ) emergent nature and 2 ) discreteness . 2If the mask is empty then the centroid is ( 0 , 0 ) and area is zero . achieved by training a function g : { 1 . . .E } H×W → RH×W such that g ( f ( st ) ) ' st . This ensures that the representation preserves all the information about the state which is a sufficient condition for the purpose . It is not hard to see however that it potentially wastes model capacity on unimportant factors of variation . In this work we aim instead for a necessary condition for constraining the representation . To this end we train f to be an injective function i.e . such that a decoding function g can distinguish between states by looking at the representation ∀x 6= y =⇒ g ( f ( x ) ) 6= g ( f ( y ) ) . Our approach is to formulate classification losses to promote distinguishability between different states at different levels of the representation . This was shown , in works like ( Barber & Agakov , 2004 ; Belghazi et al. , 2018 ; van den Oord et al. , 2018 ) , to be equivalent to maximizing a lower-bound mutual information between random variables of the representation3 . By choosing the appropriate random variables and sampling strategies from the joint distribution and marginals we can specify the right representation invariances as follows ( see Figure 1 for an overview ) : Preserving global information . The main term driving the representation learning is a global information term . To estimate it , the VQ layer output at time t , vt , is further processed by another CNN encoder to output a frame level embedding vector Gt . We train a non-parametric classifier to distinguish pairs of frames from the same trajectory from pairs of frames from different trajectories . For that we form pairs ( gt , gt+∆ ) , where ∆ is sampled randomly , from pairs ( gt , g′ ) , where g′ is sampled randomly from a different trajectory . Training this classifier lets us indirectly maximize a lower-bound on I ( Gt , Gt+∆ ) . This forces V to preserve enough information that distinguishes this particular trajectory from other ones which tends to remove all irrelevant information like textures and unchanging “ background ” elements and preserve useful moving elements . Note that as long as predictions are stable across time there is no pressure exerted by this cost to simplify the representation . 3We write our losses as I because of this fact but this is an approximation . Preserving controllable information . Secondly , we want the controllable information to be preserved in the abstract representation . That is , we want to know what aspects of the input were changed as a result of the action . We achieve this by training to predict which action was taken in a particular transition based on our representation . For that we add another loss , denoted by I ( ( Gt , Gt+1 ) , At ) , that maximizes a lower-bound on mutual information between the a pair of consecutive frames and the action that was taken in the transition . Preserving local appearance information . Finally , for hard exploration Atari games , where inputs change very little e.g . Montezuma ’ s Revenge much of the initial experience is from the first room , we add an appearance term4 . This term is meant to align the abstract representation with appearance changes i.e . promote abstractions have consistent colors . For that we feed the input image into a shallow CNN encoder which outputs an embedding C of the local color and texture structure of the image . We would like for the abstract representation to follow the appearance changes thus we maximize a lower bound on mutual information I ( Vt , Ct ) in the same way as before . The positive pairs spatially aligned V and C embeddings , and the negative pairs are obtained by sampling appearance embeddings from other spatial locations and pairing them with non-sampled V . This promotes V representations that represent well appearance under the representational constraints of the VQ representation . To learn the agents ’ abstract representation we minimize a weighted sum of these classification losses Labs ( θabs ) = ∑ t −αgEg0∼P ( G′|gt ) gi∼P ( G ) log qg ( Zg = 0| ( g0 , . . . , gK ) , gt ) ( 1 ) −αc log qc ( at| ( gt , gt+1 ) ) − αaEcit∼P ( Ct ) vt∼P ( Vt ) log qa ( Zc = 0|cit , vt ) ( 2 ) where we denote by lower case letter the samples from the corresponding uppercase random variables . P ( G ) is the time independent marginal distribution over Gt ’ s . We model qg as qg ( Zg = 0| ( g0 , . . . gK ) , gt ) = expφ ( g0 , gt ) ∑K i=0 expφ ( gi , gt ) ( 3 ) with φ ( . , . ) being the cosine similarity over the embeddings andK is a hyper parameter denoting the number of samples . We can similarly derive the third term . This is possible because the embeddings have the same dimensionality . The second term is a cross entropy based action classifier .

Summary: The appproach introduces visual abstractions that are used for reinforcement learning. The abstractions are learned using a lower bound on the mutual information and options are created to generate different measurements for each abstraction. The algorithm hence learns to "control" each abstraction as well as to select the options to achieve the overall task. The algorithm is tested on a 3D navigation task and a few Atari tasks which are known for difficult exploration.