Non-Markovian Predictive Coding For Planning In Latent Space

1 INTRODUCTION . Learning to control from high dimensional observations has been made possible due to the advancements in reinforcement learning ( RL ) and deep learning . These advancements have enabled notable successes such as solving video games ( Mnih et al. , 2015 ; Lample & Chaplot , 2017 ) and continuous control problems ( Lillicrap et al. , 2016 ) from pixels . However , it is well known that performing RL directly in the high-dimensional observation space is sample-inefficient and may require a large amount of training data ( Lake et al. , 2017 ) . This is a critical problem , especially for real-world applications . Recent model-based RL works ( Kaiser et al. , 2020 ; Ha & Schmidhuber , 2018 ; Hafner et al. , 2019 ; Zhang et al. , 2019 ; Hafner et al. , 2020 ) proposed to tackle this problem by learning a world model in the latent space , and then applying RL algorithms in the latent world model . The existing MBRL methods that learn a latent world model typically do so via reconstruction-based objectives , which are likely to encode task-irrelevant information , such as of the background . In this work , we take inspiration from the success of contrastive learning and propose Non-Markovian Predictive Coding ( NMPC ) , a novel information-theoretic approach for planning from pixels . In contrast to reconstruction , NMPC formulates a mutual information ( MI ) objective to learn the latent space for control . This objective circumvents the need to reconstruct and prioritizes the encoding of task-relevant components of the environment , thus make NMPC more robust when dealing with complicated observations . Our primary contributions are as follows : • We propose Non-Markovian Predictive Coding ( NMPC ) , a novel information-theoretic approach to learn latent world models for planning from high-dimensional observations and theoretically analyze its ability to prioritize the encoding of task-relevant information . • We show experimentally that NMPC outperforms the state-of-the-art model when deal- ing with complex environments dominated by task-irrelevant information , while remaining competitive on standard DeepMind control ( DMControl ) tasks . Additionally , we conduct detailed ablation analyses to study the empirical importance of the components in NMPC . 2 BACKGROUND . The motivation and design of our model are largely based on two previous works ( Shu et al. , 2020 ; Hafner et al. , 2020 ) . In this section , we briefly go over the key concepts in each work . Shu et al . ( 2020 ) proposed PC3 , an information-theoretic approach that uses contrastive predictive coding ( CPC ) to learn a latent space amenable to locally-linear control . Specifically , they present the theory of predictive suboptimality to motivate a CPC objective between the latent states of two consecutive time steps , instead of CPC between the frame and its corresponding state . Moreover , they use the latent dynamics F as the variational device in the lower bound ` cpc ( E , F ) = E 1 K ∑ i ln F ( E ( o ( i ) t+1 ) | E ( o ( i ) t ) , a ( i ) t ) 1 K ∑ j F ( E ( o ( i ) t+1 ) | E ( o ( j ) t ) , a ( j ) t ) ( 1 ) This particular design of the critic has two benefits , where it circumvents the instantiation of an auxiliary critic , and also takes advantage of the property that an optimal critic is the true latent dynamics . However , the author also shows that this objective does not ensure the learning of a latent dynamics F that is consistent with the true latent dynamics , therefore introduces the consistency loss to ensure the latent dynamics model F indeed approximates the true latent dynamics . ` cons ( E , F ) = Ep ( ot+1 , ot , at ) lnF ( E ( ot+1 ) | E ( ot ) , at ) ( 2 ) Since PC3 only tackles the problem from an optimal control perspective , it is not readily applicable to RL problems . Indeed , PC3 requires a depiction of the goal image in order to perform control , and also the ability to teleport to random locations of the state space to collect data , which are impractical in many problems . On the other hand , Dreamer ( Hafner et al. , 2020 ) achieves stateof-the-art performance on many RL tasks , but learns the latent space using a reconstruction-based objective . While the authors included a demonstration of a contrastive approach that yielded inferior performance to their reconstruction-based approach , their contrastive model applied CPC between the frame and its corresponding state , as opposed to between latent states across time steps . In this paper , we present Non-Markovian Predictive Coding ( NMPC ) , a novel latent world model that leverages the concepts in PC3 and apply them to Dream paradigm and RL setting . Motivated by PC3 , we formulate a mutual information objective between historical and future latent states to learn the latent space . We additionally take advantage of the recurrent model in Dreamer to model more complicated dynamics than what was considered in PC3 . The use of recurrent dynamics also allows us to extend Eqs . ( 1 ) and ( 2 ) to the non-Markovian setting , which was not considered in PC3 . 3 NON-MARKOVIAN PREDICTIVE CODING FOR PLANNING FROM PIXELS . To plan in an unknown environment , we need to model the environment dynamics from experience . We do so by iteratively collecting new data and using those data to train the world model . In this section , we focus on presenting the proposed latent world model , its components and objective functions , and provide practical considerations when implementing the method . 3.1 NON-MARKOVIAN PREDICTIVE CODING . We aim to learn a latent dynamics model for planning . To do that , we define an encoder E to embed high-dimensional observations into a latent space , a latent dynamics F to model the world in this space , and a reward function , as follows Encoder : E ( ot ) = st Latent dynamics : F ( st | s < t , a < t ) = p ( st | s < t , a < t ) Reward function : R ( rt | st ) = p ( rt | st ) ( 3 ) in which t is the discrete time step , { ot , at , rt } Tt=1 are data sequences with image observations ot , continuous action vectors at , scalar rewards rt , and st denotes the latent state at time t. To handle potentially non-Markovian environment dynamics , we model the transition dynamics using a recurrent neural network with a deterministic state ht = RNN ( ht−1 , st−1 , at−1 ) , which summarizes information about the past , followed by the stochastic state model p ( st|s < t , a < t ) = p ( st|ht ) . In practice , we use a deterministic encoder , and Gaussian distribution for dynamics and reward functions . The graphical model is presented in Figure 1 . We now introduce the three key components of our model : recurrent contrastive predictive coding , non-markovian consistency , and reward prediction . Recurrent contrastive predictive coding Instead of performing pixel prediction to learn E and F as in Hafner et al . ( 2019 ; 2020 ) , we take a different approach , where we maximize the mutual information ( MI ) between the past latent codes and actions against the future latent code I ( E ( Ot ) ; E ( O < t ) , A < t ) . This objective prioritizes the encoding of predictable components from the environment , which potentially helps avoid encoding nuance information when dealing with complex image observations . A similar MI objective was previously used in Shu et al . ( 2020 ) . However , Shu et al . ( 2020 ) only considered the Markovian setting and thus focused on the mutual information over consecutive time steps . In contrast , we consider a non-Markovian setting and wish to maximize the mutual information between the latent code at any time step t and the entire historical trajectory latent codes and actions prior to t. Given a trajectory over T steps , our non-Markovian MI objective thus sums over every possible choice of t ∈ { 2 , . . . , T } , T∑ t=2 I ( E ( Ot ) ; E ( O < t ) , A < t ) . ( 4 ) To estimate this quantity , we employ contrastive predictive coding ( CPC ) proposed by Oord et al . ( 2018 ) . We perform CPC by introducing a critic function f to construct the lower bound at a particular time step t , I ( E ( Ot ) ; E ( O < t ) , A < t ) ≥ E 1 K ∑ i ln exp f ( E ( o ( i ) t ) , E ( o ( i ) < t ) , a ( i ) < t ) 1 K ∑ j exp f ( E ( o ( i ) t ) , E ( o ( j ) < t ) , a ( j ) < t ) = : ` ( t ) r-cpc , ( 5 ) where the expectation is over K i.i.d . samples of ( ot , o < t , a < t ) . Note that Eq . ( 5 ) uses past ( E ( o ( j ) < t ) , a ( j ) < t ) from unrelated trajectories as an efficient source of negative samples for the con- trastive prediction of the future latent code E ( o ( i ) t ) . Following Shu et al . ( 2020 ) , we choose to tie f to our recurrent latent dynamics model F , exp f ( st , s < t , a < t ) = F ( st|s < t , a < t ) . ( 6 ) There are two favorable properties of this particular design . First , it is parameter-efficient since we can circumvent the instantiation of a separate critic f . Moreover , it takes advantage of the fact that an optimal critic is the true latent dynamics induced by the encoder E ( Poole et al. , 2019 ; Ma & Collins , 2018 ) . We denote our overall CPC objective as ` r-cpc ( E , F ) = ∑T t=2 ` ( t ) r-cpc ( E , F ) . Non-Markovian consistency Although the true dynamics is an optimal critic for the CPC bound , maximizing this objective only does not ensure the learning of a latent dynamics model F that is consistent with the true latent dynamics , due to the non-uniqueness of the optimal critic ( Shu et al. , 2020 ) . Since an accurate dynamics is crucial for planning in the latent space , we additionally introduce a consistency objective , which encourages the latent dynamics model to maintain a good prediction of the future latent code given the past latent codes and actions . Similar to the recurrent CPC objective , we optimize for consistency at every time step in the trajectory , ` cons ( E , F ) = T∑ t=2 Ep ( ot , o < t , a < t ) lnF ( E ( ot ) |E ( o < t ) , a < t ) . ( 7 ) Reward prediction Finally , we train the reward function by maximizing the likelihood of the true reward value conditioned on the latent state st = E ( ot ) , ` reward ( E , R ) = T∑ t=1 Ep ( ot ) lnR ( rt|E ( ot ) ) . ( 8 )

Summary: This paper proposes an information-theoretic framework for learning a world model that encodes task-relevant information of the world. It shows that the learned encoder and dynamics model can be used to train the policy and fitting the value function to agent to perform comparibly well to Dreamer on standard tasks and outperform them when there are distractions in the scene. The paper also provides a theoretical anylisis of the task-relevant information in the encoding.