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: The paper proposes a variant of a Hamiltonian neural network for learning dynamics models for model-based RL. The proposed network architecture uses encoders/decoders to map the system state x = [q, p] from the observations s and vice versa. In the latent state, a standard HNN is used to compute dx/dt to increment the time step.   In the experiments, the paper applies this dynamics model to model-based rl and wants to show that this specific dynamics model increases the sample efficiency and leads to greater performance. It is especially noteworthy that the domains include locomotion tasks that include many contacts.