Model-based Reinforcement Learning with a Hamiltonian Canonical ODE Network

Model-based reinforcement learning usually suffers from a high sample complexity in training the world model , especially for the environments with complex dynamics . To make the training for general physical environments more efficient , we introduce Hamiltonian canonical ordinary differential equations into the learning process , which inspires a novel model of neural ordinary differential auto-encoder ( NODA ) . NODA can model the physical world by nature and is flexible to impose Hamiltonian mechanics ( e.g. , the dimension of the physical equations ) which can further accelerate training of the environment models . It can consequentially empower an RL agent with the robust extrapolation using a small amount of samples as well as the guarantee on the physical plausibility . Theoretically , we prove that NODA has uniform bounds for multi-step transition errors and value errors under certain conditions . Extensive experiments show that NODA can learn the environment dynamics effectively with a high sample efficiency , making it possible to facilitate reinforcement learning agents at the early stage . 1 INTRODUCTION . Reinforcement learning has obtained substantial progress in both theoretical foundations ( Asadi et al. , 2018 ; Jiang , 2018 ) and empirical applications ( Mnih et al. , 2013 ; 2015 ; Peters & Schaal , 2006 ; Johannink et al. , 2019 ) . In particular , model-free reinforcement learning ( MFRL ) can complete complex tasks such as Atari games ( Schrittwieser et al. , 2020 ) and robot control ( Roveda et al. , 2020 ) . However , the MFRL algorithms often need a large amount of interactions with the environment ( Langlois et al. , 2019 ) in order to train an agent , which impedes their further applications . Model-based reinforcement learning ( MBRL ) methods can alleviate this issue by resorting to a model to characterize the environmental dynamics and conduct planning ( van Hasselt et al. , 2019 ; Moerland et al. , 2020a ) . In general , MBRL can quench the thirst of massive amounts of real data that may be costly to acquire , by using rollouts from the model ( Langlois et al. , 2019 ; Deisenroth & Rasmussen , 2011 ) . It has witnessed numerous works on approximating the model with various strategies , such as leastsquares temporal difference ( Boyan , 1999 ) , guided policy search ( GPS ) ( Levine & Abbeel , 2014 ) , dynamic Bayesian networks ( DBN ) ( Hester & Stone , 2012 ) , and deep neural networks ( Fujimoto et al. , 2018 ) . However , the sample efficiency of MBRL can still be limited due to the high sample complexity of learning a world model when the environment is complex . Traditional methods such as the Gaussian Processes based method ( Deisenroth & Rasmussen , 2011 ) can perform well on some problems with high sample efficiency , but they are not easy to scale to high-dimensional problems ( Plaat et al. , 2020 ) . High-capacity models scale well , but they often have low sample efficiency ( Plaat et al. , 2020 ) . The trade-off between scalability and sample complexity remains as a critical issue for model-based RL . To address the aforementioned issue , we propose to introduce physical knowledge to reduce the sample complexity for learning high-dimensional dynamics in physical environments . We focus on reinforcement learning in an environment whose dynamics can be formulated by Hamiltonian canonical equations ( Goldstein et al. , 2002 ) . Up till now , Hamiltonian dynamics have been successfully applied in numerous areas of physics from robotics to industrial automation . Specifically , we formulate the environments dynamics as ordinary differential equations ( ODEs ) , and then use a novel network architecture called Neural Ordinary Differential Auto-encoder ( NODA ) as our world model , which is naturally induced by physical equations . In particular , NODA consists of two parts — an auto-encoder and an ODE network . We use the auto-encoder to get the underlying physical variables , and use the ODE network to learn the dynamics over physical variables . By using NODA , we can enjoy its ability of modeling the physical world as well as its flexibility of combining physical knowledge ( if any ) , such as the dimension of physical variables . Theoretically , we provide uniform bounds for both multi-step transition errors and value errors for NODA by extending the former study of Lipschitz models ( Asadi et al. , 2018 ) to continuous action spaces . It is noted that NODA can be combined with both MFRL methods like SAC ( Haarnoja et al. , 2018 ) and MBRL methods like Dreamer ( Hafner et al. , 2019 ) by facilitating the learning of the world model . Extensive experiments show that we can learn NODA well using a small number of data with an appropriate structure encoded , which can boost the sample efficiency by using imaginary trajectories over the environment models ( Todorov et al. , 2012 ; Schulman et al. , 2015 ) . 2 BACKGROUND . We start by presenting the background knowledge of reinforcement learning , and then explain the relationship between MBRL and Hamiltonian mechanics . 2.1 MODEL-BASED REINFORCEMENT LEARNING . We consider the Markov decision process ( MDP ) model for reinforcement learning . Formally , an MDP is formulated as a tuple 〈S , A , T , R , γ〉 , where S is the state space , A is the action space , T : S ×A → P ( S ) is the transition function , R : S ×A → R is the reward function , and γ ∈ [ 0 , 1 ) is a discount factor . We denote the set P ( · ) as all probability measures on the space in the bracket . Our goal is to find a policy π that can choose an action to maximize the accumulated reward . Here we focus on the challenging tasks with continuous state and action spaces ( i.e. , there are infinite states and actions ) . MBRL aims to learn a policy by integrating planning with the aid of a known or learned model ( Moerland et al. , 2020b ) , and an essential part of MBRL is to learn a transition function that characterizes the environment . The transition function above is defined over a given state , but we can generalize it to represent the transition from a state distribution z ∈ P ( S ) . By calling the generalized transition function recursively , we can get the n-step generalized transition function , which is defined as : Definition 1 ( Transition Functions ) In a metric state space ( S , dS ) and an action spaceA , we can define the generalized transition function of TG ( over state distribution z ) , and the n-step generalized transition function of TnG ( for fixed sequence of actions ) ( Asadi et al. , 2018 ) as TG ( s ′ | s , a ) = ∫ T ( s′ | s , a ) z ( s ) ds TnG ( · | z ) = TG ( · | · · ·TG ( · | z , a0 ) · · · , an−1 ) ︸ ︷︷ ︸ n recursive calls . ( 1 ) Here the generalized transition gives the distribution of outcome under a certain state distribution . For MBRL , it is nontrivial to learn the transition function ( i.e. , the dynamics for a physical environment ) because S can be high-dimensional . Several attempts have been made to learn such dynamics , while they have various limitations . For instance , probabilistic inference for learning control ( PILCO ) ( Deisenroth & Rasmussen , 2011 ) uses Gaussian processes to model the transition function , but the inference part does not scale well to high dimensions ( Langlois et al. , 2019 ) . Stochastic ensemble value expansion ( STEVE ) ( Buckman et al. , 2018 ) and adaptation augmented model-based policy optimization ( AMPO ) ( Shen et al. , 2020 ) use scalable machine learning models , but there is no further discussion on how to learn such a model efficiently . Actually , high-dimensional state and action spaces usually require much more data samples ( Plaat et al. , 2020 ) . Monte Carlo tree search ( MCTS ) ( Silver et al. , 2017 ) can introduce human knowledge of the transition function to enhance learning , but it is restricted to cases where transition functions are totally known . 2.2 HAMILTONIAN MECHANICS . Methods of analytical mechanics have been introduced to predict the evolution of dynamic systems such as pendulums . For example , Lagrangian neural networks ( Lutter et al. , 2019 ; Cranmer et al. , 2020 ) and Hamiltonian neural networks ( Greydanus et al. , 2019 ) can be used to simulate dynamic systems . These papers focus on how to model the Lagrangian or the Hamiltonian . One challenge for such methods is that the equations for the Lagrangian or the Hamiltonian are second-order differential equations , which are difficult to model in a general way . Besides , numerical solutions of second-order differential equations are prone to error accumulation . One natural idea is to reformulate second-order differential equations into first-order ones . Then , we can use an existing neural network ( Chen et al. , 2018 ) to model these equations in the ODE form . In this paper , we concentrate on the Hamiltonian case , for which the first-order representation corresponds to Hamiltonian canonical equations ( Junkins & Schaub , 2009 ) .1 Specifically , in Hamiltonian mechanics , we use pairs of generalized coordinate and generalized momentum ( qk , pk ) to completely describe a dynamic system ( k ∈ { 0 , 1 , · · · , K } ) , where K is the dimension of generalized coordinates.It is noted thatK can be intuitively interpreted as the degree of freedom . We denote pk and qk as canonical states , and they are minimal independent parameters which can describe the state of the system . We define Hamiltonian H : R2K+1 → R as a function of these variables and time t. Then the evolution of the system satisfies Hamiltonian canonical equations ( Junkins & Schaub , 2009 ) : dqk dt = ∂H ∂pk , dpk dt = − ∂H ∂qk +Qk ( t ) , ( 2 ) where k ∈ { 0 , 1 , · · · , K } , andQk ( t ) are the generalized forces which describe the effects of external forces . One advantage of using these equations is that they can describe general dynamic systems . Moreover , it is possible to incorporate prior knowledge into Hamiltonian canonical equations , e.g. , by assuming a specific form of energy ( Sprangers et al. , 2014 ) . In our case , prior knowledge can be available . For example , we may know the dimension K of the generalized coordinate and generalized momentum , which only requires the knowledge of the ’ degree of freedom ’ for a given system but makes a difference in training . Another example is about transfer learning : If we already learn the underlying dynamics , we can combine the learnt dynamics with other modules to transfer the learnt knowledge to a variety of different tasks . We empirically examine the these advantages in the experimental section . 3 METHODOLOGY . We now formally present NODA , which consists of an auto-encoder and an ODE network . NODA aims to serve as a simulator of the real environment ( i.e. , a dynamic system ) by learning the transition function and the reward function . Here we assume that the transition is deterministic , otherwise an SDE network ( Li et al. , 2020 ) can be used instead . Then we can use NODA to assist reinforcement learning by generating imaginary trajectories , as outlined in Algorithm 1 . Besides , we also discuss the prior knowledge that can be incorporated with NODA .

Summary: This paper proposes a model of neural ODE auto-encoder, NODA, which incorporates Hamiltonian mechanics to learn a world model. The paper shows theoretical results of transition errors and value errors. NODA is tested on a range of RL tasks.