Learning Physics Priors for Deep Reinforcement Learing

1 INTRODUCTION . Recent advances in deep reinforcement learning ( RL ) have largely relied on model-free approaches , demonstrating strong performance on a variety of domains ( Silver et al. , 2016 ; Mnih et al. , 2013 ; Kempka et al. , 2016 ; Zhang et al. , 2018b ) . Unfortunately , model-free techniques do not have very good sample efficiency ( Sutton , 1990 ) and are difficult to adapt to new tasks or domains ( Nichol et al. , 2018 ) . This is mainly because a single value function represents both the agent ’ s policy and its knowledge of environment dynamics . On the other hand , decoupling the dynamics model from the policy , possible with model-based RL , allows for better generalization and transfer ( Zhang et al. , 2018a ) . However , estimating an accurate dynamics model of the environment while simultaneously using it to learn a policy is challenging and expensive , often leading to sub-optimal policies and slower learning . One way to facilitate this process would be to initialize dynamics models with generic priors for more efficient and stable model-based reinforcement learning . t t+1 new environment . Humans , on the other hand , have reliable priors that allow us to learn the dynamics of new environments very quickly ( Dubey et al. , 2018 ) . In this work , we demonstrate that learning a prior of general physics ( ex . concepts of velocity , mass , acceleration ) allows for better and more efficient estimation of the dynamics of new environments , thereby resulting in better control policies . In order to learn a prior for physical dynamics , we utilize unsupervised learning over raw videos containing moving objects . In particular , we first train a dynamics model to predict the next frame given the previous k frames , over a wide variety of scenarios with moving objects . This training allows the parameters of the dynamics model to capture the general laws of physics which are useful in predicting entity movements . We can then use this pre-trained dynamics model as a prior in a model-based RL algorithm for a control task . Specifically , we initialize the dynamics model with the pre-trained parameters and finetune them using transitions from the specific task . We utilize this dynamics model in order to predict future frames up to a finite horizon , which are then used as an additional context input into a policy network , similar to the approach of Weber et al . ( 2017 ) . We show that this results in both faster learning and more optimal policies . We characterize our dynamics model as an image pixel prediction model . Learning a good future frame model is challenging ( Mathieu et al. , 2015 ) , mainly because of two reasons - a ) the highdimensionality of the output space with arbitrary number of moving objects and interactions , and b ) the partial observability in environments . Existing techniques that use a combination of recurrent and convolutional networks ( Oh et al. , 2015 ) suffer from the problem of error compounding and poor generalization as they encode a image into a single vector . In contrast , we propose SpatialNet , a neural network that consists of a convolutional encoder , a spatial memory block , and a convolutional decoder that captures the local nature of dynamics . The spatial memory module operates by performing convolution operations over a temporal 3-dimensional state representation that keeps spatial information intact , similar to the model of Xingjian et al . ( 2015 ) . This allows the network to capture localized physics of objects such as directional movements and collisions in a more fine-grained manner , resulting in lower prediction error , better generalization and input size invariance . We evaluate our approach and compare it with existing techniques on two different RL scenarios . First , we consider PhysWorld , a suite of randomized physics-focused games , where learning object movement is crucial to a successful policy . Second , we also evaluate on a stochastic variant of the popular ALE framework consisting of Atari games Bellemare et al . ( 2013 ) . In both scenarios , we first demonstrate the value of learning a physics prior for model dynamics — our agent achieves up to 130 % higher performance on a shooting game , PhysShooter and 49.8 % higher on the Atari game of Asteroids , compared to the most competitive baseline . Further , we also show that the dynamics model fine-tuned on these tasks can transfer to new tasks effectively . For instance , our model achieves a relative score improvement of 37.4 % on transfer from PhysForage to PhysShooter , than compared to a relative score improvement of 5.4 % from a policy-transfer baseline . 2 RELATED WORK . Video prediction and reinforcement learning Our frame prediction model is closest in spirit to the ConvolutionalLSTM which has been applied to several domains including weather forecasting ( Xingjian et al. , 2015 ) , gesture recognition ( Zhu et al. , 2017 ) , and forecasting passenger demand ( Ke et al. , 2017 ) . Similar architectures that incorporate differentiable memory modules have been proposed ( Patraucean et al. , 2015 ) , with applications to deep RL ( Parisotto and Salakhutdinov , 2017 ) . We use a simpler architecture , which we show generalizes better at capturing dynamics , that captures entity movements more directly and demonstrate its use in learning useful physics priors for environment dynamics . Several recent methods have combined policy learning with future frame prediction in various ways . Action-conditioned frame prediction ( Oh et al. , 2015 ; Finn et al. , 2016 ; Weber et al. , 2017 ) has been used to simulate trajectories for policy learning . Predicted frames have also been used to incentivize exploration in agents , via hashing ( Yin et al. , 2017 ) or using the prediction error to provide intrinsic rewards ( Pathak et al. , 2017 ) . The main departure of our work from these papers are that we use a frame prediction model that is not conditioned on actions , which allows us to learn the model from videos independent of the task , and then use it as a prior for learning the specific dynamics of different tasks . Physics models for reinforcement learning Incorporating physics priors into learning dynamics models of environments has been the subject of some recent work . Cutler et al . ( 2014 ) ; Cutler and How ( 2015 ) learn Bayesian nonparametric priors for physical systems and use it to guide policy search . Scholz et al . ( 2014 ) model latent physical parameters like mass and friction and use an external physics simulator to predict future states for planning . Kansky et al . ( 2017 ) learn a generative physics simulator and demonstrate its usefulness in transferring policies across task variations , but require objects to be detected from raw images . Xie et al . ( 2016 ) develop a model that incorporates prior knowledge about the dynamics of rigidbody systems in order to perform model-predictive control . Nguyen-Tuong and Peters ( 2010 ) use parametric physics models as priors for learning environment dynamics . While all these approaches demonstrate the importance of having relevant priors to sample efficient model learning , they all require some manual parameterization . In contrast , we learn physics priors in the form of a predictive neural network , directly from observing raw videos . Decoupling dynamics from policy Our work also falls into the general direction of decoupling the agent ’ s knowledge of the environment dynamics from its task-oriented policy . Successor representations Dayan ( 1993 ) decompose the agent ’ s value function into a feature-based state representation and a reward projection operator , resulting in better exploration of the state space Kulkarni et al . ( 2016 ) ; Barreto et al . ( 2017 ) ; Machado et al . ( 2017b ) . While these state abstractions help with exploration , such representations do not explicitly capture dynamics models of the environment . Zhang et al . ( 2018a ) recently proposed an approach to learn two separate models for dynamics and rewards and use it to perform online planning . However , they learn a dynamics model using task-specific transitions , while we learn a prior from task-independent videos and demonstrate its usefulness in learning different environment dynamics . 3 FRAMEWORK . Our goal is to demonstrate that agents can learn useful physics priors from raw videos and employ them effectively to learn dynamics of new environments . To this end , our approach has two phases : 1 . Pre-training : we first design a suitable neural network architecture ( SpatialNet ) to predict pixels in the next frame given the previous k frames of a video . We train this network on a newly created video dataset that captures various physical phenomena between entities . 2 . Reinforcement learning : We use the pre-trained dynamics predictor from the previous phase to initialize the dynamics model for a reinforcement learning agent . This dynamics model is used by the agent to predict a few frames into the future and use them as additional context for its policy . The dynamics model is simultaneously fine-tuned using trajectories observed in the task environment . We first describe SpatialNet , and then demonstrate how we use it for reinforcement learning . 3.1 SPATIALNET . Predicting the physical behavior of an entity requires a model that can perform two crucial operations – 1 ) isolation of dynamics of an entity , and 2 ) accurately model localized spaces and interactions around the entity . LSTM-based recurrent networks used in prior work ( Oh et al. , 2015 ) are ill-suited for this task since they encode the entire scene into a single latent vector , thereby losing the localized spatio-temporal correlations that are important for making accurate physical predictions . On the other hand , the ConvLSTM ( Xingjian et al. , 2015 ) architecture has localized spatio-temporal correlations , but is not able to accurately maintain dynamics of entities due to LSTM state updates and limited separation of stationary and non-stationary objects . In order to satisfy both the above desiderata , SpatialNet uses a spatial memory that explicitly encodes dynamics that are updated with object movement through convolutions . This allows us to implicitly capture and maintain localized physics , such as entity velocities and collisions between entities , in our frame prediction model and allows significantly lower long term prediction error . Architecture SpatialNet is conceptually simple and consists of three modules . The first module is a standard convolutional encoder E that converts an input image x into a 3D representational map z . The second module is a spatial memory block , σ , that converts z and the hidden state h from the previous timestep into an output representation o and new hidden state h′ . Finally , we have a convolutional decoder D that predicts the next frame x′ from o . Both the encoder and decoder modules ( E and D ) use two convolutional layers each with residual connections . We implement the spatial memory block σ as a 2D convolution operation . The module takes in a previous hidden state ht and input zt at timestep t , both of shape k× n× n where k is the number of channels and n× n is the dimensionality of the grid . We then perform the following operations : it = f ( Ce ⊕ [ ht ; zt ] ) ; ut = f ( Cu ⊕ [ it ; ht ] ) ; ht = f ( Cdyn ⊕ ut ) ; ot = f ( Cd ⊕ [ zt ; ht+1 ] ) where ⊕ denotes a convolution operation , [ ; ] denotes concatenation , Ce , Cu , Cdyn , Cd are convolutional kernels and f is a non-linearity ( we use ELU ( Clevert et al. , 2015 ) ) . The module first encodes a combination of zt and ht into a proposal state ut , using two convolutions Ce , Cu . Cdyn acts like a dynamics simulator and generates a new hidden state ht+1 , which captures the localized predictions for the next state around each entity . Finally , Cd uses ht+1 and zt to produce ot , encoding information about the entire frame to be rendered by subsequent decoding . Intuitively , this architecture biases the module towards storing relevant physics information about each entity in a block of pixels at the entity ’ s corresponding location . This information is sequentially updated through the convolutions , while static information such as background texture are passed directly through the input encoding zt . We note that our spatial memory is not action-conditional , which allows us to learn from task-independent videos , as well as generalize better to new environ- ments . Given training videos D = { ( x ( i ) 1 , x ( i ) 2 , . . . , x ( i ) Ti } N i=1 , we learn the parameters of the model using a standard MSE-based loss function , L ( θ ) = ∑ i ∑ j ‖x̂ij − xij‖2 .

Summary: A method for learning physics priors is proposed for faster learning and better transfer learning. The key idea in learning physics priors using spatial net, which is similar to a convolutional LSTM model for making predictions. Authors propose to improve the sample efficiency of Deep RL algorithms, by augmenting PPO’s state input with 3 future frames predicted by the physics prior model.