Modeling the Long Term Future in Model-Based Reinforcement Learning

1 INTRODUCTION . Reinforcement Learning ( RL ) is an agent-oriented learning paradigm concerned with learning by interacting with an uncertain environment . Combined with deep neural networks as function approximators , deep reinforcement learning ( deep RL ) algorithms recently allowed us to tackle highly complex tasks . Despite recent success in a variety of challenging environment such as Atari games ( Bellemare et al. , 2013 ) and the game of Go ( Silver et al. , 2016 ) , it is still difficult to apply RL approaches in domains with high dimensional observation-action space and complex dynamics . Furthermore , most popular RL algorithms are model-free as they directly learn a value function ( Mnih et al. , 2015 ) or policy ( Schulman et al. , 2015 ; 2017 ) without trying to model or predict the environment ’ s dynamics . Model-free RL techniques often require large amounts of training data and can be expensive , dangerous or impossibly slow , especially for agents and robots acting in the real world . On the other hand , model-based RL ( Sutton , 1991 ; Deisenroth & Rasmussen , 2011 ; Chiappa et al. , 2017 ) provides an alternative approach by learning an explicit representation of the underlying environment dynamics . The principal component of model-based methods is to use an estimated model as an internal simulator for planning , hence limiting the need for interaction with the environment . Unfortunately , when the dynamics are complex , it is not trivial to learn models that are accurate enough to later ensure stable and fast learning of a good policy . The most widely used techniques for model learning are based on one-step prediction . Specifically , given an observation ot and an action at at time t , a model is trained to predict the conditional distribution over the immediate next observation ot+1 , i.e p ( ot+1 | ot , at ) . Although computationally easy , the one-step prediction error is an inadequate proxy for the downstream performance of model-based methods as it does not account for how the model behaves when com- ♦ Mila , Université de Montréal , ♥ Facebook AI Research , ♣ Polytechnique Montréal , ¶ CIFAR Senior Fellow , ‡ Work done at Facebook AI Research , * Georgia Institute of Technology Corresponding authors : rosemary.nan.ke @ gmail.com posed with itself . In fact , one-step modelling errors can compound after multiple steps and can degrade the policy learning . This is referred to as the compounding error phenomenon ( Talvitie , 2014 ; Asadi et al. , 2018 ; Weber et al. , 2017 ) . Other examples of models are autoregressive models such as recurrent neural networks ( Mikolov et al. , 2010 ) that factorize naturally as log pθ ( ot+1 , at+1 , ot+2 , at+2 , . . . | ot , at ) = ∑ t log pθ ( ot+1 , at+1 | o1 , a1 , . . . ot , at ) . Training autoregressive models using maximum likelihood results in ‘ teacher-forcing ’ that breaks the training over one-step decisions . Such sequential models are known to suffer from accumulating errors as observed in ( Lamb et al. , 2016 ; Bengio et al. , 2015 ) . Our key motivation is the following – a model of the environment should reason about ( i.e . be trained to predict ) long-term transition dynamics pθ ( ot+1 , at+1 , ot+2 , at+2 , . . . | ot , at ) and not just single step transitions pθ ( ot+1 | ot , at ) . That is , the model should predict what will happen in the long-term future , and not just the immediate future . We hypothesize ( and test ) that such a model would exhibit less cascading of errors and would learn better feature embeddings for improved performance . One way to capture long-term transition dynamics is to use latent variables recurrent networks . Ideally , latent variables could capture higher level structures in the data and help to reason about long-term transition dynamics . However , in practice it is difficult for latent variables to capture higher level representation in the presence of a strong autoregressive model as shown in Gulrajani et al . ( 2016 ) ; Goyal et al . ( 2017 ) ; Guu et al . ( 2018 ) . To overcome this difficulty , we leverage recent advances in variational inference . In particular , we make use of the recently proposed Z-forcing idea ( Goyal et al. , 2017 ) , which uses an auxiliary cost on the latent variable to predict the long-term future . Keeping in mind that more accurate long-term prediction is better for planning , we use two ways to inject future information into latent variables . Firstly , we augment the dynamics model with a backward recurrent network ( RNN ) such that the approximate posterior of latent variables depends on the summary of future information . Secondly , we force latent variables to predict a summary of the future using an auxiliary cost that acts as a regularizer . Unlike one-step prediction , our approach encourages the predicted future observations to remain grounded in the real observations . Injection of information about the future can also help in planning as it can be seen as injecting a plan for the future . In stochastic environment dynamics , unfolding the dynamics model may lead to unlikely trajectories due to errors compounding at each step during rollouts . In this work , we make the following key contributions : 1 . We demonstrate that having an auxiliary loss to predict the longer-term future helps in faster imitation learning . 2 . We demonstrate that incorporating the latent plan into dynamics model can be used for planning ( for example Model Predictive Control ) efficiently . We show the performance of the proposed method as compared to existing state of the art RL methods . 3 . We empirically observe that using the proposed auxiliary loss could help in finding sub-goals in the partially observable 2D environment . 2 PROPOSED MODEL . We consider an agent in the environment that observes at each time step t an observation ot . The execution of a given action at causes the environment to transition to a new unobserved state , return a reward and emit an observation at the next time step sampled from p ? ( ot+1|o1 : t , a1 : t ) where o1 : t and a1 : t are the observation and action sequences up to time step t. In many domains of interest , the underlying transition dynamics p ? are not known and the observations are very high-dimensional raw pixel observations . In the following , we will explain our novel proposed approach to learn an accurate environment model that could be used as an internal simulator for planning . We focus on the task of predicting a future observation-action sequence ( o1 : T , a1 : T ) given an initial observation o0 . We frame this problem as estimating the conditional probability distribution p ( o1 : T , a1 : T |o0 ) . The latter distribution is modeled by a recurrent neural network with stochastic latent variables z1 : T . We train the model using variational inference . We introduce an approximate posterior over latent variables . We maximize a regularized form of the Evidence Lower Bound ( ELBO ) . The regularization comes from an auxiliary task we assign to the latent variables . 2.1 GENERATIVE PROCESS . The graphical model in Fig . 1 illustrates the dependencies in our generative model . Observations and latent variables are coupled by using an autoregressive model , the Long Short Term Memory ( LSTM ) architecture ( Hochreiter & Schmidhuber , 1997 ) , which runs through the sequence : ht = f ( ot , ht−1 , zt ) ( 1 ) where f is a deterministic non-linear transition function and ht is the LSTM hidden state at time t. According the graphical model in Fig . 1 , the predictive distribution factorizes as follows : pθ ( o1 : T , a1 : T | o0 , h0 ) = ∫ T∏ t=1 pθ ( ot | at−1 , ht−1 , zt ) ︸ ︷︷ ︸ observation decoder pθ ( at−1 | ht−1 , zt ) ︸ ︷︷ ︸ action decoder pθ ( zt | ht−1 ) ︸ ︷︷ ︸ latent prior dz ( 2 ) where 1. pθ ( ot | at−1 , ht−1 , zt ) is the observation decoder distribution conditioned on the last action at−1 , the hidden state ht and the latent variable zt . 2. pθ ( at−1 | ht−1 , zt ) is the action decoder distribution conditioned on the the hidden states ht−1 and the latent variable zt . 3. pθ ( zt | ht−1 ) is the prior over latent variable zt condition on the hidden states ht−1 All these conditional distributions , listed above , are represented by simple distributions such as Gaussian distributions . Their means and standard variations are computed by multi-layered feedforward networks . Although each single distribution is unimodal , the marginalization over sequence of latent variables makes pθ ( o1 : T , a1 : T |o0 ) highly multimodal . Note that the prior distribution of the latent random variable at time step t depends on all the preceding inputs via the hidden state ht−1 . This temporal structure of the prior has been shown to improve the representational power ( Chung et al. , 2015 ; Fraccaro et al. , 2016 ; Goyal et al. , 2017 ) of the latent variable . 2.2 INFERENCE MODEL . In order to overcome the intractability of posterior inference of latent variables given observationaction sequence , we make use of amortized variational inference ideas ( Kingma & Welling , 2013 ) . We consider recognition or inference network , a neural network which approximates the intractable posterior . The true posterior of a given latent variable zt is p ( zt|ht−1 , at−1 : T , ot : T , zt+1 : T ) . For the sake of an efficient posterior approximation , we make the following design choices : 1 . We drop the dependence of the posterior on actions at−1 : T and future latent variables zt+1 : T . 2 . To take into account the dependence on ht−1 , we share parameters between the generative model and the recognition model by making the approximate posterior , a function of the hidden state ht−1 computed by the LSTM transition module f of the generative model . 3 . To take into account the dependence on future observations ot : T , we use an LSTM that processes observation sequence backward as bt = g ( ot , bt+1 ) , where g is a deterministic transition function and bt is the LSTM backward hidden state at time t. 4 . Finally , a feed-forward network takes as inputs ht−1 and bt and output the mean and the standard deviation of the approximate posterior qφ ( zt | ht−1 , bt ) . In principle , the posterior should depend on future actions . To take into account the dependence on future actions as well as future observations , we can use the LSTM that processes the observationaction sequence in backward manner . In pilot trials , we conducted experiments with and without the dependencies on actions for the backward LSTM and we did not notice a noticeable difference in terms of performance . Therefore , we chose to drop the dependence on actions in the backward LSTM to simplify the code . Now using the approximate posterior , the Evidence Lower Bound ( ELBO ) is derived as follows : log pθ ( o1 : T , a1 : T | o0 , h0 ) ≥ Eqφ ( z1 : T |o0 : T , a0 : T ) [ log pθ ( o1 : T , a1 : T , z1 : T | o0 , h0 ) log qφ ( z1 : T | o0 : T , a0 : T ) ] ( 3 ) = Eqφ ( z1 : T |o0 : T , a0 : T ) [ log pθ ( o1 : T , a1 : T | o0 , h0 , z1 : T ) ] ( 4 ) −KL ( qφ ( z1 : T | o0 : T , a0 : T ) ‖pθ ( z1 : T | o0 , h0 ) ) Leveraging temporal structure of the generative and inference network , the ELBO breaks down as : L ( o1 : T , a1 : T ; θ , φ ) = ∑ t Eqφ ( zt|ht−1 , bt ) [ log pθ ( ot | at−1 , ht−1 , zt ) + log pθ ( at−1 | ht−1 , zt ) ] ( 5 ) −KL ( qφ ( zt | ht−1 , bt ) ‖pθ ( zt | ht−1 ) )

Summary: The paper introduces an interesting approach to model learning for imitation and RL. Given the problem of maintaining multi-step predictions in the context of sequential decision making process, and deficiencies faced during planning with one-step models [1][2], it’s imperative to explore approaches that do multi-step predictions. This paper combines ideas from learning sequential latent models with making multi-step future predictions as an auxiliary loss to improve imitation learning performance, efficiency of planning and finding sub-goals in a partially observed domain.