Model-Predictive Policy Learning with Uncertainty Regularization for Driving in Dense Traffic

1 INTRODUCTION . In recent years , model-free reinforcement learning methods using deep neural network controllers have proven effective on a wide range of tasks , from playing video or text-based games ( Mnih et al. , 2015 ; 2016 ; Narasimhan et al. , 2015 ) to learning algorithms ( Zaremba et al. , 2015 ) and complex locomotion tasks ( Lillicrap et al. , 2015 ; Zhang et al. , 2015 ) . However , these methods often require a large number of interactions with the environment in order to learn . While this is not a problem if the environment is simulated , it can limit the application of these methods in realistic environments where interactions with the environment are slow , expensive or potentially dangerous . Building a simulator where the agent can safely try out policies without facing real consequences can mitigate this problem , but requires human engineering effort which increases with the complexity of the environment being modeled . Model-based reinforcement learning approaches try to learn a model of the environment dynamics , and then use this model to plan actions or train a parameterized policy . A common setting is where an agent alternates between collecting experience by executing actions using its current policy or dynamics model , and then using these experiences to improve its dynamics model . This approach has been shown empirically to significantly reduce the required number of environment interactions needed to obtain an effective policy or planner ( Atkeson & Santamaria , 1997 ; Deisenroth & Rasmussen , 2011 ; Nagabandi et al. , 2017 ; Chua et al. , 2018 ) . Despite these improvements in sample complexity , there exist settings where even a single poor action executed by an agent in a real environment can have consequences which ∗Equal contribution . are not acceptable . At the same time , with data collection becoming increasingly inexpensive , there are many settings where observational data of an environment is abundant . This suggests a need for algorithms which can learn policies primarily from observational data , which can then perform well in a real environment . Autonomous driving is an example of such a setting : on one hand , trajectories of human drivers can be easily collected using traffic cameras ( Halkias & Colyar , 2006 ) , resulting in an abundance of observational data ; on the other hand , learning through interaction with the real environment is not a viable solution . However , learning policies from purely observational data is challenging because the data may only cover a small region of the space over which it is defined . If the observational data consists of stateaction pairs produced by an expert , one option is to use imitation learning ( Pomerleau , 1991 ) . However , this is well-known to suffer from a mismatch between the states seen at training and execution time ( Ross & Bagnell , 2010 ) . Another option is to learn a dynamics model from observational data , and then use it to train a policy ( Nguyen & Widrow , 1989 ) . However , the dynamics model may make arbitrary predictions outside the domain it was trained on , which may wrongly be associated with low cost ( or high reward ) as shown in Figure 1 . The policy network may then exploit these errors in the dynamics model and produce actions which lead to wrongly optimistic states . In the interactive setting , this problem is naturally self-correcting , since states where the model predictions are wrongly optimistic will be more likely to be experienced , and thus will correct the dynamics model . However , the problem persists if the dataset of environment interactions which the model is trained on is fixed . In this work , we propose to train a policy while explicitly penalizing the mismatch between the distribution of trajectories it induces and the one reflected in the training data . We use a learned dynamics model which is unrolled for multiple time steps , and train a policy network to minimize a differentiable cost over this rolled-out trajectory . This cost contains two terms : a policy cost which represents the objective the policy seeks to optimize , and an uncertainty cost which represents its divergence from the states it is trained on . We measure this second cost by using the uncertainty of the dynamics model about its own predictions , calculated using dropout . We apply our approach in the context of learning policies to drive an autonomous car in dense traffic , using a large-scale dataset of real-world driving trajectories which we also adapt into an environment for testing learned policies 1 . We show that model-based control using this additional uncertainty regularizer substantially outperforms unregularized control , and enables learning good driving policies using only observational data with no environment interaction or additional labeling by an expert . We also show how to effectively leverage an action-conditional stochastic forward model using a modified posterior distribution , which encourages the model to maintain sensitivity to input actions . 2 MODEL-PREDICTIVE POLICY LEARNING WITH UNCERTAINTY REGULARIZATION . We assume we are given a dataset of observational data which consists of state-action pairs D = { ( st , at ) } t. We first describe our general approach , which consists of two steps : learning an actionconditional dynamics model using the collected observational data , and then using this model to train a fast , feedforward policy network which minimizes both a policy cost and an uncertainty cost . 2.1 ACTION-CONDITIONAL FORWARD MODEL . In this work , we consider recent approaches for stochastic prediction based on Variational Autoencoders ( Kingma & Welling , 2013 ; Babaeizadeh et al. , 2017 ; Denton & Fergus , 2018 ) . The stochastic model fθ ( s1 : t , at , zt ) takes as input a sequence of observed or previously predicted states s1 : t , an action at , and a latent variable zt which represents the information about the next state st+1 which 1Dataset and environment can be found at https : //github.com/Atcold/pytorch-PPUU is not a deterministic function of the input . During training , latent variables are sampled from a distribution whose parameters are output by a posterior network qφ ( s1 : t , st+1 ) conditioned on the past inputs and true targets . This network is trained jointly with the rest of the model using the reparameterization trick , and a term is included in the loss to minimize the KL divergence between the posterior distribution and a fixed prior p ( z ) , which in our case is an isotropic Gaussian . The per-sample loss used for training the stochastic model is given by : L ( θ , φ ; s1 : t , st+1 , at ) = ‖st+1 − fθ ( s1 : t , at , zt ) ‖22 + βDKL ( qφ ( z|s1 : t , st+1 ) ‖p ( z ) ) ( 1 ) After training , different future predictions for a given sequence of frames can be generated by sampling different latent variables from the prior distribution . Recent models for stochastic video prediction ( Babaeizadeh et al. , 2017 ; Denton & Fergus , 2018 ) do not use their model for planning or training a policy network , and parameterize the posterior distribution over latent variables using a diagonal Gaussian . In our case , we are training an actionconditional video prediction model which we will later use to train a policy . This leads to an additional requirement : it is important for the prediction model to accurately respond to input actions , and not use the latent variables to encode factors of variation in the outputs which are due to the actions . To this end we propose to use a mixture of two Gaussians , with one component fixed to the prior , as our posterior distribution : ( µφ , σφ ) = qφ ( s1 : t , st+1 ) u ∼ B ( pu ) zt ∼ ( 1− u ) · N ( µφ , σφ ) + u · p ( z ) This can be seen as applying a form of global dropout to the latent variables at training time 2 , and forces the prediction model to extract as much information as possible from the input states and actions by making the latent variable independent of the output with some probability . In our experiments we will refer to this parameterization as z-dropout . 2.2 TRAINING A POLICY NETWORK WITH UNCERTAINTY MINIMIZATION . Once the forward model is trained , we use it to train a parameterized policy network πψ , which we assume to be stochastic . We first sample an initial state sequence s1 : t from the training set , unroll the forward model over T time steps , and backpropagate gradients of a differentiable objective function with respect to the parameters of the policy network ( shown in Figure 2 ) . During this process the weights of the forward model are fixed , and only the weights of the policy network are optimized . This objective function contains two terms : a policy cost C , which reflects the underlying objective the policy is trying to learn , and an uncertainty cost U , which reflects how close the predicted state induced by the policy network is to the manifold which the data D is drawn from . Training the policy using a stochastic forward model involves solving the following problem , where latent variables are sampled from the prior and input into the forward model at every time step : argmin ψ [ T∑ i=1 C ( ŝt+i ) + λU ( ŝt+i ) ] , such that : zt+i ∼ p ( z ) ât+i ∼ πψ ( ŝt+i−1 ) ŝt+i = f ( ŝt+i−1 , ât+i , zt+i ) The uncertainty cost U is applied to states predicted by the forward model , and could reflect any measure of their likelihood under the distribution the training data is drawn from . We propose here a general form based on the uncertainty of the dynamics model , which is calculated using dropout . Intuitively , if the dynamics model is given a state-action pair from the same distribution asD ( which 2If the variances of both Gaussians are zero , this becomes equivalent to applying dropout to the latent code where all units are either set to zero or unchanged with probability pu . it was trained on ) , it will have low uncertainty about its prediction . If it is given a state-action pair which is outside this distribution , it will have high uncertainty . Dropout ( Hinton et al. , 2012 ; Srivastava et al. , 2014 ) is a regularization technique which consists of randomly setting hidden units in a neural network to zero with some probability . The work of ( Gal & Ghahramani , 2016 ) showed that estimates of the neural network ’ s uncertainty for a given input can be obtained by calculating the covariance of its outputs taken over multiple dropout masks . We note that this uncertainty estimate is the composition of differentiable functions : each of the models induced by applying a different dropout mask is differentiable , as is the covariance operator . Furthermore , we can summarize the covariance matrix by taking its trace ( which is equal to the sum of its eigenvalues , or equivalently the sum of the variances of the outputs across each dimension ) , which is also a differentiable operation . This provides a scalar estimate of uncertainty which is differentiable with respect to the input . More precisely , let fθ1 , ... , fθK denote our prediction model with K different dropout masks applied to its hidden units ( this can also be viewed as changing its weights ) . We define our scalar measure of uncertainty U as follows : U ( ŝt+1 ) = tr [ Cov [ { fθk ( s1 : t , at , zt ) } Kk=1 ] ] = d∑ j=1 Var ( { fθk ( s1 : t , at , zt ) j } Kk=1 ) where d is the dimensionality of the output . Minimizing this quantity with respect to actions encourages the policy network to produce actions which , when plugged into the forward model , will produce predictions which the forward model is confident about 3 . A simple way to define U given an initial sequence of states s1 : t from D would be to set U ( ŝt+k ) = ‖ŝt+k − st+k‖2 , which would encourage the policy network to output actions which lead to a similar trajectory as the one observed in the dataset . This leads to a set of states which the model is presumably confident about , but may not be a trajectory which also satisfies the policy cost C unless the dataset D consists of expert trajectories . If this is the case , setting C ( ŝt+i ) = U ( ŝt+i ) = 1 2‖ŝt+k − st+k‖2 gives a model-based imitation learning objective which simultaneously optimizes the policy cost and the uncertainty cost . The problem then becomes : 3In practice , we apply an additional step to normalize this quantity across different modalities and rollout lengths , which is detailed in Appendix D. argmin ψ [ T∑ i=1 ‖ŝt+i − st+i‖2 ] , such that : zt+i ∼ qφ ( z|s1 : t , st+1 ) ât+i ∼ πψ ( ŝt+i−1 ) ŝt+i = f ( ŝt+i−1 , ât+i , zt+i ) We call the first approach MPUR , for Model-Predictive Policy with Uncertainty Regularization , and the second MPER , for Model-Predictive Policy with Expert Regularization . A key feature of both approaches is that we optimize the objective over T time steps , which is made possible by our learned dynamics model . This means that the actions will receive gradients from multiple time steps ahead , which will penalize actions which lead to large divergences from the training manifold further into the future , even if they only cause a small divergence at the next time step .

Summary: The paper addresses the difficulty of covariate shift in model-based reinforcement learning. Here, the distribution over trajectories during is significantly different for the behaviour or data-collecting policy and the target or optimised policy. As a mean to address this, the authors propose to add an uncertainty term to the cost, which is realised by the trace of the covariance of the outputs of a MC dropout forward model. The method is applied to driving in dense traffic, where even single wrong actions can be catastrophic.