EMI: Exploration with Mutual Information Maximizing State and Action Embeddings

1 INTRODUCTION . The central task in reinforcement learning is to learn policies that would maximize the total reward received from interacting with the unknown environment . Although recent methods have demonstrated to solve a range of complex tasks ( Mnih et al. , 2015 ; Schulman et al. , 2015 ; 2017 ) , the success of these methods , however , hinges on whether the agent constantly receives the intermediate reward feedback or not . In case of challenging environments with sparse reward signals , these methods struggle to obtain meaningful policies unless the agent luckily stumbles into the rewarding or predefined goal states . To this end , prior works on exploration generally utilize some kind of intrinsic motivation mechanism to provide a measure of surprise . These measures can be based on density estimation via generative models ( Bellemare et al. , 2016 ; Fu et al. , 2017 ; Oh et al. , 2015 ) , predictive forward models ( Stadie et al. , 2015 ; Houthooft et al. , 2016 ) , or more ad-hoc measures that aim to approximate surprise ( Pathak et al. , 2017 ) . Methods based on predictive forward models and generative models must model the distribution over state observations , which can make them difficult to scale to complex , high-dimensional observation spaces , while models that eschew direct forward predictive or density estimation rely on heuristic measures of surprise that may not transfer effectively to a wide range of tasks . Our aim in this work is to devise a method for exploration that does not require a direct generation of high-dimensional state observations , while still retaining the benefits of being able to measure surprise based on the forward prediction . If exploration is performed by seeking out states that maximize surprise , the problem , in essence , is in measuring surprise , which requires a representation where functionally similar states are close together , and functionally distinct states are far apart . In this paper , we propose to learn compact representations for both the states ( φ ) and actions ( ψ ) simultaneously satisfying the following criteria : First , given the representations of state and the corresponding next state , the uncertainty of the representation of the corresponding action should be minimal . Second , given the representations of the state and the corresponding action , the uncertainty of the representation of the corresponding next state should also be minimal . Third , the action embedding representation ( ψ ) should seamlessly support both the continuous and discrete actions . Finally , we impose the linear dynamics model in the representation space which can also explain the rare irreducible error under the dynamics model . Given the representation , we guide the exploration by measuring surprise based on forward prediction and relative increase in diversity in the embedding representation space . Figure 1 illustrates an example visualization of our learned state embedding representations ( φ ) and sample trajectories in the representation space in Montezuma ’ s Revenge . We present two main technical contributions that make this into a practical exploration method . First , we describe how compact state and action representations can be constructed via Donsker & Varadhan ( 1983 ) estimation of mutual information without relying on generative decoding of full observations . Second , we show that imposing linear topology on the learned embedding representation space ( such that the transitions are linear ) , thereby offloading most of the modeling burden onto the embedding function itself , provides an essential informative measure of surprise when visiting novel states . For the experiments , we show that we can use our representations on a range of complex image-based tasks and robotic locomotion tasks with continuous actions . We report significantly improved results compared to recent intrinsic motivation based exploration methods ( Fu et al. , 2017 ; Pathak et al. , 2017 ) on several challenging Atari tasks and robotic locomotion tasks with sparse rewards . 2 RELATED WORKS . Our work is related to the following strands of active research : Unsupervised representation learning via mutual information estimation Recent literature on unsupervised representation learning generally focus on extracting latent representation maximizing approximate lower bound on the mutual information between the code and the data . In the context of generative adversarial networks ( Goodfellow et al. , 2014 ) , Chen et al . ( 2016 ) ; Belghazi et al . ( 2018 ) aims at maximizing the approximation of mutual information between the latent code and the raw data . Belghazi et al . ( 2018 ) estimates the mutual information with neural network via Donsker & Varadhan ( 1983 ) estimation to learn better generative model . Hjelm et al . ( 2018 ) builds on the idea and trains a decoder-free encoding representation maximizing the mutual information between the input image and the representation . Furthermore , the method uses f -divergence ( Nowozin et al. , 2016 ) estimation of Jensen-Shannon divergence rather than the KL divergence to estimate the mutual information for better numerical stability . Oord et al . ( 2018 ) estimates mutual information via autoregressive model and makes predictions on local patches in an image . Thomas et al . ( 2017 ) aims to learn the representations that maximize the causal relationship between the distributed policies and the representation of changes in the state . Exploration with intrinsic motivation Prior works on exploration mostly employ intrinsic motivation to estimate the measure of novelty or surprisal to guide the exploration . Bellemare et al . ( 2016 ) utilize density estimation via CTS ( Bellemare et al. , 2014 ) generative model and derive pseudocounts as the intrinsic motivation . Fu et al . ( 2017 ) avoids building explicit density models by training K-exemplar models that distinguish a state from all other observed states . Some methods train predictive forward models ( Stadie et al. , 2015 ; Houthooft et al. , 2016 ; Oh et al. , 2015 ) and estimate the prediction error as the intrinsic motivation . Oh et al . ( 2015 ) employs generative decoding of the full observation via recursive autoencoders and thus can be challenging to scale for high dimensional observations . VIME ( Houthooft et al . ( 2016 ) ) approximates the environment dynamics , uses the information gain of the learned dynamics model as intrinsic rewards , and showed encouraging results on robotic locomotion problems . However , the method needs to update the dynamics model per each observation and is unlikely to be scalable for complex tasks with high dimensional states such as Atari games . Other approaches utilize more ad-hoc measures ( Pathak et al. , 2017 ; Tang et al. , 2017 ) that aim to approximate surprise . ICM ( Pathak et al . ( 2017 ) ) transforms the high dimensional states to feature space and imposes cross entropy and Euclidean loss so the action and the feature of the next state are predictable . However , ICM does not utilize mutual information like VIME to directly measure the uncertainty and is limited to discrete actions . Our method ( EMI ) is also reminiscent of ( Kohonen & Somervuo , 1998 ) in a sense that we seek to construct a decoder-free latent space from the high dimensional observation data with a topology in the latent space . In contrast to the prior works on exploration , we seek to construct the representation under linear topology and does not require decoding the full observation but seek to encode the essential predictive signal that can be used for guiding the exploration . 3 PRELIMINARIES . We consider a Markov decision process defined by the tuple ( S , A , P , r , γ ) , where S is the set of states , A is the set of actions , P : S × A × S → R+ is the environment transition distribution , r : S → R is the reward function , and γ ∈ ( 0 , 1 ) is the discount factor . Let π denote a stochastic policy over actions given states . Denote P0 : S → R+ as the distribution of initial state s0 . The discounted sum of expected rewards under the policy π is defined by η ( π ) = Eτ [ ∑ t=0 γtr ( st ) ] , ( 1 ) where τ = ( s0 , a0 , . . . , aT−1 , sT ) denotes the trajectory , s0 ∼ P0 ( s0 ) , at ∼ π ( at | st ) , and st+1 ∼ P ( st+1 | st , at ) . The objective in policy based reinforcement learning is to search over the space of parameterized policies ( i.e . neural network ) πθ ( a | s ) in order to maximize η ( πθ ) . Also , denote PπSAS′ as the joint probability distribution of singleton experience tuples ( s , a , s′ ) starting from s0 ∼ P0 ( s0 ) and following the policy π . Furthermore , define PπA = ∫ S×S′ dP π SAS′ as the marginal distribution of actions , PπSS′ = ∫ A dP π SAS′ as the marginal distribution of states and the corresponding next states , PπS′ = ∫ S×A dP π SAS′ as the marginal distribution of the next states , and PπSA = ∫ S′ dP π SAS′ as the marginal distribution of states and the actions following the policy π . 4 METHODS . Our goal is to construct the embedding representation of the observation and action ( discrete or continuous ) for complex dynamical systems that does not rely on generative decoding of the full observation , but still provides a useful predictive signal that can be used for exploration . This requires a representation where functionally similar states are close together , and functionally distinct states are far apart . We approach this objective from maximizing mutual information under several criteria . 4.1 MUTUAL INFORMATION MAXIMIZING STATE AND ACTION EMBEDDING REPRESENTATIONS . We first introduce the embedding function of states φα : S → Rd and actions ψβ : A → Rd with parameters α and β ( i.e . neural networks ) respectively . We seek to learn the embedding function of states ( φα ) and actions ( ψβ ) satisfying the following two criteria : 1 . Given the embedding representation of states and the actions [ φα ( s ) ; ψβ ( a ) ] , the uncertainty of the embedding representation of the corresponding next states φα ( s′ ) should be minimal and vice versa . 2 . Given the embedding representation of states and the corresponding next states [ φα ( s ) ; φα ( s ′ ) ] , the uncertainty of the embedding representation of the corresponding actions ψβ ( a ) should also be minimal and vice versa . Intuitively , the first criterion translates to maximizing the mutual information between [ φα ( s ) ; ψβ ( a ) ] , and φα ( s′ ) which we define as IS ( α , β ) in Equation ( 2 ) . And the second criterion translates to maximizing the mutual information between [ φα ( s ) ; φα ( s′ ) ] and ψβ ( a ) defined as IA ( α , β ) in Equation ( 3 ) . maximize α , β IS ( α , β ) : = I ( [ φα ( s ) ; ψβ ( a ) ] ; φα ( s′ ) ) = DKL ( PπSAS′ ‖ PπSA ⊗ PπS′ ) ( 2 ) maximize α , β IA ( α , β ) : = I ( [ φα ( s ) ; φα ( s′ ) ] ; ψβ ( a ) ) = DKL ( PπSAS′ ‖ PπSS′ ⊗ PπA ) ( 3 ) Mutual information is not bounded from above and maximizing mutual information is notoriously difficult to compute in high dimensional settings . Motivated by Hjelm et al . ( 2018 ) ; Belghazi et al . ( 2018 ) , we compute Donsker & Varadhan ( 1983 ) lower bound of mutual information . Concretely , Donsker-Varadhan representation is a tight estimator for the mutual information of two random variables X and Z , derived as in Equation ( 4 ) . I ( X ; Z ) = DKL ( PXZ ‖ PX ⊗ PZ ) ≥ sup ω∈Ω EPXZTω ( x , z ) − logEPX⊗PZ expTω ( x , z ) , ( 4 ) where Tω : X × Z → R is a differentiable transform with parameter ω . Furthermore , for better numerical stability , we utilize a different measure between the joint and marginals than the KL- divergence . In particular , we employ Jensen-Shannon divergence ( JSD ) ( Hjelm et al. , 2018 ) which is bounded both from below and above by 0 and log ( 4 ) 1 . Theorem 1 . I ( JSD ) ( X ; Z ) ≥ supω∈Ω EPXZ [ −sp ( −Tω ( x , z ) ) ] − EPX⊗PZ [ sp ( Tω ( x , z ) ) ] + log ( 4 ) Proof . I ( JSD ) ( X ; Z ) = DJSD ( PXZ ‖ PX ⊗ PZ ) ≥ sup ω∈Ω EPXZ [ Sω ( x , z ) ] − EPX⊗PZ [ JSD∗ ( Sω ( x , z ) ) ] = sup ω∈Ω EPXZ [ −sp ( −Tω ( x , z ) ) ] − EPX⊗PZ [ sp ( Tω ( x , z ) ) ] + log ( 4 ) , ( 5 ) where the inequality in the second line holds from the definition of f -divergence ( Nowozin et al. , 2016 ) . In the third line , we substituted Sω ( x , z ) = log ( 2 ) − log ( 1 + exp ( −Tω ( x , z ) ) ) and Fenchel conjugate of Jensen-Shannon divergence , JSD∗ ( t ) = − log ( 2− exp ( t ) ) . From Theorem 1 , we have , maximize α , β I ( JSD ) S ( α , β ) ≥ maximize α , β sup ωS∈ΩS EPπ SAS′ [ −sp ( −TωS ( φα ( s ) , ψβ ( a ) , φα ( s′ ) ) ) ] ( 6 ) − EPπSA⊗PπS′ [ sp ( TωS ( φα ( s ) , ψβ ( a ) , φα ( s̃ ′ ) ) ) ] + log 4 , maximize α , β I ( JSD ) A ( α , β ) ≥ maximize α , β sup ωA∈ΩA EPπ SAS′ [ −sp ( −TωA ( φα ( s ) , ψβ ( a ) , φα ( s′ ) ) ) ] ( 7 ) − EPπ SS′⊗P π A [ sp ( TωA ( φα ( s ) , ψβ ( ã ) , φα ( s ′ ) ) ) ] + log 4 , where sp ( z ) = log ( 1 + exp z ) . The expectations in Equation ( 6 ) and Equation ( 7 ) are approximated using the empirical samples trajectories τ . Note , the samples s̃′ ∼ PπS′ and ã ∼ PπA from the marginals are obtained by dropping ( s , a ) and ( s , s′ ) in samples ( s , a , s̃′ ) and ( s , ã , s′ ) from PπSAS′ . Figure 2 illustrates the computational architecture for estimating the lower bounds on IS and IA .

Summary: This paper introduces actions as a co-predictor of next-states and the predicted (from current and next state) in the context of (model-based) RL. In addition they incorporate the idea of using a JSD-based objective do prediction (as the Deep InfoMax paper), which is novel to RL. The enforce a linear structure between current / next states and actions with an additional sparse nonlinear term computed from both current states and actions. From this, they are able to quantify the amount of novelty in the representation space as a measure of exploration, which can be used as an intrinsic reward.