Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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).

heading: 2 BACKGROUND
text: We start by presenting the background knowledge of reinforcement learning, and then explain the relationship between MBRL and Hamiltonian mechanics.

heading: 2.1 MODEL-BASED REINFORCEMENT LEARNING
text: 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.

heading: 2.2 HAMILTONIAN MECHANICS
text: 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.

heading: 3 METHODOLOGY
text: 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.

heading: 3.1 MODELING HAMILTONIAN CANONICAL EQUATIONS
text: We focus on modeling Hamiltonian canonical equations, since after getting these equations, we know the continuous evolution of the system. For many RL environments such as MuJoCo (Todorov et al., 2012), the discretization of the continuous evolution is just the transition function. However, it may be non-trivial to get these equations for a real dynamic system2. So we propose to use neural ODE or ODE networks (Chen et al., 2018; 2021) to model these first-order differential equations.
1Hamiltonian canonical equations are widely used with clear physical meanings and more symmetric forms than the Lagrangian case.
2We provide an example to show the derivation of these equations in Appendix A, where we derive these equations under the setting of Pendulum-v0 task in Gym (Brockman et al., 2016).
Specifically, Neural ODE (Chen et al., 2018) provides an elegant framework to model differential equations. It is defined by dh(t)dt = f(h(t), t, θ). Here f is a neural network with parameter θ, and we use an integrator to transform the input h(0) to output h(τ), where τ is a given time horizon. Neural ODE can be viewed as a continuous version of ResNet (He et al., 2016). Here we assume that the agent affects the environment through forces. So the generalized forces Q(t) in Eq. (2) correspond to the effects of action a in reinforcement learning, since generalized forces can describe all physical perturbations.

heading: 3.2 NODA MODEL STRUCTURE
text: Hamiltonian canonical equations describe the evolution over canonical states, but the observed state s is not necessarily composed of canonical states q,p. Assuming that our state s ∈ S contains the complete information of the canonical states, we firstly construct an auto-encoder to map state s to canonical states q,p.3 Specifically, in the auto-encoder part, we use a function f : S → R2K to get the canonical states. Here we denote the concatenation of (q,p) as u, so u = f(s). We also need to restore the state from canonical states, so we further define a function g : R2K → S. After going through the ODE network (defined later), we can decode the evolved canonical states s′ to get the next state: snext = g(s′).
Since the canonical states are concise but complete to describe the system, the auto-encoder can serve as a dimension reducer to extract and refine state s when there is much redundancy in the state space. We further assume that the reward is related to the canonical states and the action. Assuming A ⊂ Rm, we can use another decoder g′ : R2K × A → R to get predicted reward r. More specifically, r = g′(u,a). Then we let the predicted canonical states go through an ODE network to evolve.
In the ODE network part, the encoded latent state u evolves through the ODE over a time interval [t0, t0 + τ ]. We can define the ODE network as a function ODE(h,u,a, t0, t0 + τ), which returns the value of the solution of neural ODE at time t0 + τ . The neural ODE is described as:
du dt = h(u,a, t), t ∈ [t0, t0 + τ ], (3)
where h is a neural network.
After combining the ODE network and the auto-encoder, we can get a full structure of NODA, which is shown in Figure 1. The inputs of NODA are the state s and the action a, and the outputs of NODA
3The auto-encoder is widely used in machine learning (Ballard, 1987; Ng et al., 2011; Kingma & Welling, 2013).
are the predicted state snext and the predicted reward r. Our model is formally characterized as:
snext = g(ODE(h,u,a, t0, t0 + τ)),
r = g′(u,a). (4)
To learn the unknown parameters in f, g and h, we define the total loss function for NODA as a convex combination of the state loss and the reward loss:
L = µ(||g(u)− s||22 + ||snext − s∗next||22) + (1− µ)||r − r∗||22, (5)
where µ ∈ (0, 1), and ∗ denotes the ground truth. By minimizing this objective, we jointly train the RL agent and our model using interactions with the environment, and use our model together with the real environment to provide training data for the RL agent after certain training steps. More details are provided in Appendix C.
Algorithm 1 NODA for Reinforcement Learning Input: Reinforcement learning agent s, environment env, NODA model m repeat
Collect data from interactions with env Use interactions with env to train m Use interactions with both env and m to train s
until Certain number of steps return Reinforcement learning agent s, NODA model m As a model of Hamiltonian canonical equations, NODA can describe general dynamic systems, and the form of canonical states can be very general. Moreover, NODA corresponds with real dynamics of the environment, which allows us to use our knowledge of the Hamiltonian. For example, we can determine 2K with little physical knowledge, which is the dimension of the canonical state space of a dynamic system, and we demonstrate the effects of such dimension in the experimental part. For stronger prior knowledge, we can replace the ODE network by the derived form of Hamiltonian canonical equations, or do transfer learning between similar dynamic systems such as transferring the auto-encoder or the ODE network. We also demonstrate the effects of transfer learning in the experimental part. This makes it possible for us to make trade-offs between human expert knowledge and the required number of training samples.

heading: 4 THEORETICAL ANALYSIS
text: In this section, we present error bounds for both multi-step transition and values for NODA.

heading: 4.1 LIPSCHITZ DYNAMIC SYSTEMS AND NODA
text: Former work (Asadi et al., 2018) studied using Lipschitz function groups to model the transition of the environment. However, there are error bounds only if the real environment is Lipschitz, and that work did not mention what kind of environments have a Lipschitz transition. Theorem 1 proves that under certain conditions, the transition functions of dynamic systems are Lipschitz, which is the basis of transition and state value error bounds.
Before stating the theorem, we firstly define Lipschitz continuity, which measures the maximum magnitude of enlargement of the perturbation of input at the output side. For reinforcement learning, Definition 2 is the condition of uniformly Lipschitz continuous in the former work (Asadi et al., 2018), and it plays a central role in our theorems.
Definition 2 (Lipschitz Models) In a metric state space (S, dS) and an action space A, we say that f is a Lipschitz model if the Lipschitz constant (Asadi et al., 2018)
KF := sup a∈A sup s1,s2∈S
dS (f (s1,a) , f (s2,a))
dS (s1, s2) <∞. (6)
Theorem 1 (Lipschitz Dynamic Systems) For a dynamic system with a C2 continuous Hamiltonian H : R2K+1 → R, if the state s is in a bounded closed set S ⊂ Rl, the evolution time equals τ , the generalized force Qk is C1 continuous with respect to states and bounded (for any dimension
k), and the transformation from states to canonical states f∗ : S → R2K is Lipschitz, then the canonical states are Lipschitz with respect to time, and the environment with respect to canonical states is Lipschitz. Additionally, if the transformation from canonical states to states g∗ : R2K → S is Lipschitz, then the environment is Lipschitz, which means (here s 6= s′)
sup a∈A sup s,s′∈S
dS (snext, s ′ next)
dS (s, s′) <∞. (7)
We provide a detailed proof in Appendix B, and here we only give a sketch.
Proof. By noting that the composition of Lipschitz functions is Lipschitz, we only need to analyse the ODE part. The conditions of H and Qk guarantees that the right hand side of the ODE is Lipschitz, which leads to the Lipschitz condition in this part.
Actually, the conditions in Theorem 1 are not hard to satisfy. The Hamiltonian H itself already has derivatives with respect to q, p and t. We just further assume that the Hamiltonian H’s derivatives with respect to q, p and t is differentiable continuous. Actually, for many dynamic systems such as spring-mass systems or three-body systems, the Hamiltonian is C∞ continuous.
As an imitation of Hamiltonian canonical equations, it is natural to validate if the NODA model is Lipschitz. We analyse this in Theorem 2 in which the Lipschitz model will be one of the conditions in error bounds.
Theorem 2 (Lipschitz NODA) For the NODA model, if the state s is in a bounded closed set S ⊂ Rl, f : S → R2K is Lipschitz, the evolution time equals τ , the action am is C1 continuous with respect to states and bounded (for any dimension m), function h is C1 continuous, then canonical states are Lipschitz with respect to time, and NODA with respect to canonical states is Lipschitz. Additionally, if the transformation from canonical states to states g : R2K → S is Lipschitz, then NODA is Lipschitz, which means (here s 6= s′)
sup a∈A sup s,s′∈S
dS (snext, s ′ next)
dS (s, s′) <∞. (8)
The proof is in Appendix B with the similar idea as Theorem 1.

heading: 4.2 UNIFORM ERROR BOUNDS FOR MULTI-STEP TRANSITION
text: Here we firstly define a metric called Wasserstein Metric in Definition 3. Wasserstein metric describes how to move one distribution to another with the least cost. It has been applied to generative adversarial networks (Arjovsky et al., 2017) and reinforcement learning (Asadi et al., 2018).
Definition 3 (Wasserstein Metric) In a metric space (X, d), the Wasserstein metric between two probability distributions z1 and z2 in P(X) is
W (z1, z2) := inf j∈Λ
∫∫ j (s1, s2) d (s1, s2) ds2ds1, (9)
where Λ is a set of all joint distributions j onX×X with marginal distributions z1 and z2 (Arjovsky et al., 2017).
Under the conditions of Theorem 1 and Theorem 2, the dynamic system and the NODA model are Lipschitz models. On that basis, we can give uniform error bounds for multi-step transition, which is shown in Theorem 3. The theorem tells us that the multi-step transitions of the NODA model and the environment do not differ much under certain conditions, which gives support for using the NODA model as the imaginary environment.
Theorem 3 (Transition Error Bounds) Under the conditions in Theorem 1 and Theorem 2, we already know that the transition function TG (s′ | s,a) of the environment and the transition function
T̂G (s ′ | s,a) of the NODA model are Lipschitz. We denote the Lipschitz constants of these transition functions as K1 and K2, respectively. Let K̄ = min{K1,K2}, then ∀n ≥ 1:
δ(n) := W ( T̂nG (· | µ), TnG (· | µ) ) ≤ ∆ n−1∑ i=0 (K̄)i, (10)
where ∆ is an upper bound of Wasserstein metric between T̂ (· | s,a) and T (· | s,a). This upper bound is tight for linear and deterministic transitions.
The original theorem (Asadi et al., 2018) gives a a bound for a fixed action sequence. However, here our definitions of Lipschitz environments and models are uniform for all actions. So, by using the original proof, we give a same error bound for all possible action sequences. Thus, we get a uniform error bound under the continuous action space. This concludes the proof.
Theorem 4 (Value Error Bounds) Under all the conditions in Theorem 3, if the reward function R(s) is uniformly Lipschitz for any action, which means we can define
KR := sup a∈A sup s1,s2∈S |R (s1,a)−R (s2,a)| dS (s1, s2) <∞. (11)
If we define state values as
VT (s) := ∞∑ n=0 γn ∫ TnG (s ′ | δs)R (s′) ds′, (12)
where δs means the probability that the state being s equals 1. Then ∀s ∈ S and K̄ (defined in Theorem 3)∈ [0, 1γ ], we have ∣∣VT (s)− VT̂ (s)∣∣ ≤ γKR∆(1− γ)(1− γK̄) . (13) We put the theorem of the uniform error bounds of state values here, and defer its proof and more detailed discussions of theorems to Appendix B.

heading: 5 EXPERIMENTS
text: In this section, we validate the efficiency of NODA’s learning and its ability of boosting both MFRL and MBRL methods.

heading: 5.1 EXPERIMENTAL SETUP
text: Baseline methods One natural baseline method for NODA is called AE, which means replacing the ODE network in NODA by an MLP. For physical simulation tasks, we choose AE and Hamiltonian neural network (HNN) (Greydanus et al., 2019) as baseline methods. For MFRL methods, we choose SAC and TD3 as our baselines. We modify the code from spinning up (Achiam, 2018) to use GPU in training. For MBRL methods, we compare our approach with state-of-the-art methods in the literature. Specifically, we focus on two methods on visual control tasks, Dreamer (Hafner et al., 2019) and BIRD (Zhu et al., 2020), as our MBRL baselines. We reproduce their results by their released codes respectively.
Implementation We implement NODA mainly by using pytorch (Paszke et al., 2019) and torchdiffeq (https://github.com/rtqichen/torchdiffeq) as the implementation of the ODE network (Chen et al., 2018; 2021). We integrate NODA with MFRL methods by using it to generate imaginary trajectories and compare with MFRL baselines. When comparing with the two MBRL baselines, we combine NODA with Dreamer to get our approach: NODA-Dreamer, and we implement this in TensorFlow as Dreamer does. We use tfdiffeq (https://github.com/titu1994/tfdiffeq) as the implementation of the ODE network, which runs entirely on Tensorflow Eager Execution. We refer to Appendix C for more implementation details.
Tasks For physical simulation tasks, we use the setting of pixel pendulum and real pendulum in the paper of HNN (Greydanus et al., 2019). The pixel pendulum task aims at predicting the next frame of the pendulum using former adjacent frames, and the real pendulum task does the same thing using physical values as input. Besides, we use NODA to learn the transition part in the Ant-v3 task (actions are sampled randomly) in Gym (Brockman et al., 2016) to study the parameter sensitivity and transferability of NODA. For MFRL methods, we choose 4 MuJoCo environments in Gym to compare the performance of TD3, SAC, AE-SAC and NODA-SAC. For MBRL methods, since the two baselines themselves are evaluated on DeepMind Control Suite (DMC) (Tassa et al., 2018), we also evaluate NODA-Dreamer on 4 environments in DMC, in order to make a fair comparison with them. More experimental details can be found in Appendix D.

heading: 5.2 SIMULATION EFFECTIVENESS
text: The testing loss curves of HNN, AE and NODA over two physical environments are shown in Figure 2 (a) and (b). As the training loss functions among models can be different, we choose to compare the well-defined testing loss. In each experiment, the number of NODA’s parameters equals to that of AE’s parameters, and it is no more than that of HNN’s parameters. In both tasks, NODA converges quickly, and achieves the best testing loss over only thousands of training iterations. It suggests that the inductive bias introduced by NODA that the system obeys Hamiltonian canonical equations accelerates training.
Figure 2 (c) shows the effect of the latent dimension on NODA’s testing error after 3,000 iterations. Here we use NODA to learn the transition part in the Ant-v3 task. By mechanics, we know that the dimension of Hamiltonian canonical equations (corresponding to the dimension of the latent space in NODA) is approximately 40 in this task. We can see that a latent dimension of 40 is sufficient for NODA to achieve a high performance. This suggests that if we know the approximate dimension, we can use the prior knowledge to get a sweet point.
Figure 2 (d) shows the testing loss of the original NODA and the transferred NODA in the Ant-v3 task. We train the latter model for 100 steps over the one step transition of Ant-v3, and the final task is to predict the two-step transition. The figure suggests that we can transfer NODA over similar environments to accelerate training.

heading: 5.3 SAMPLE EFFICIENCY
text: The episode returns achieved by MFRL methods and MBRL methds (with/without NODA) are shown in Table 1, Figure 3 and Figure 4. For MFRL methods, NODA-SAC achieves the best performance with good robustness (see the InvertedPendulum-v2 task in Figure 3). This means that NODA can enhance the performance of the state-of-the-art algorithm SAC by imaginary trajectories. Meanwhile, if we substitute the ODE part using an MLP (denoted as the AE-SAC model), the performance will fall. This suggests that introducing an ODE network, i.e., the prior knowledge of the existence of Hamiltonian indeed helps. But for environments that SAC do not perform well (such as the Ant-v3 task), the enhancement of NODA is not significant. For MBRL methods, NODA-Dreamer can also achieve a higher return than BIRD and Dreamer in many environments. In the task ’Walker Run’, during the first half of training, NODA-Dreamer lags far behind Dreamer, but it quickly catches up in the second half. This suggests that NODA has the power to accelerate training after some warm-up steps. Generally, NODA can enhance the performance for both MFRL and MBRL methods. More analyses and experiments can be found in Appendix D.

heading: 6 CONCLUSION
text: This paper proposes a novel simulator called NODA for reinforcement learning. Motivated by Hamiltonian canonical equations in physics, NODA has clear physical meanings. This allows us to incorporate prior knowledge or do transfer learning, which can improve sample efficiency. In the theoretical part, we fill the gap between former theorems about Lipschitz models and the physical world by proving that dynamic systems can be Lipschitz and extending former limited theorems to continuous action spaces. Besides, we give uniform transition error bounds and value error bounds for NODA. In the experimental part, we verify that NODA provides a more efficient modeling than HNN, and we can use prior knowledge or transfer learning to further boost its training. Not only is NODA itself sample-efficient, but NODA can improve the sample efficiency of both MFRL and MBRL methods such as SAC and Dreamer.

heading: 7 ETHICS STATEMENT
text: In this paper, we only use data sets that we are permitted to use, and we present our method and results in a transparent, honest and reproducible way. We fully acknowledge any contributions to this work. The possible harm that our method can bring originates from the modeling errors of NODA, which can lead the agent to go wrong. However, this can be mitigated by theoretical guarantees as well as human supervision. Generally, we believe that our work will make reinforcement learning more accessible by reducing the number of training samples that may be costly to acquire.

heading: 8 REPRODUCIBILITY STATEMENT
text: We give formal proofs for our theorems in Appendix B. Our code is provided in supplementary material. The NODA folder contains experiments about simulation effectiveness and experiments in Gym environments, and the NODA-Dreamer folder contains experiments in DMC environments.

heading: A HAMILTONIAN CANONICAL EQUATIONS: PENDULUM-V0 AS AN EXAMPLE
text: The settings of the Pendulum-v0 task can be found on GitHub4.
Let x be the axis horizontally to the right, y be the axis stright up. The pendulum’s position can be described by generalized position θ, where θ is the angle from the y axis to the pendulum.
For the free end, x = −l sin θ, y = l cos θ. For a force perpendicular to the pendulum (the angle from the y axis to u is θ+π/2) with its 2-norm equaling u, its generalized forceQ = ∂r∂xux+ ∂r ∂yuy = −l cos θ · (−u cos θ)− l sin θ · (−u sin θ) = ul, which is just the torque of the force. 4https://github.com/openai/gym/blob/master/gym/envs/classic_control/ pendulum.py
The kinetic energy T = 1/6ml2(dθ/dt)2, and the potential energy V = 1/2mgl cos θ, so we get the Lagrangian in Equation (14).
L ( θ, dθ
dt , t
) = T − V = 1
6 ml2
( dθ
dt
)2 − 1
2 mgl cos θ (14)
We denote θ as q, the canonical momentum for θ as p, then
p = ∂L ∂ (
dq dt
) = 1 3 ml2 dq dt (15)
Then we can write the Hamiltonian in Equation (16).
H(q, p, t) = dq dt p− L = 3p
2
2ml2 +
1 2 mgl cos q (16)
Then Hamiltonian canonical equations give out the dynamics, as is shown in Equation (17). dq dt = ∂H ∂p = 3p ml2 dp
dt = −∂H ∂q +Q(t) = 1 2 mgl sin q + u(t)
(17)
This is the general case, but the code implementation is a little different. In the code, u(t) is clipped to [−2, 2], θ is clipped to [−π, π), and dθ/dt is clipped to [−8, 8]. This exactly gives the bounds of the state space, since states are [cos θ, sin θ,dθ/dt]. The clips can be added to Hamiltonian canonical equations, so the code implementation can be viewed as a solver for Equation (18). dq dt = 3p ml2 dp dt =  min( 1 2 mgl sin q + clip(u(t),−2, 2), 0), if p > 8 3ml2 max( 1 2 mgl sin q + clip(u(t),−2, 2), 0), if p < − 8 3ml2 1
2 mgl sin q + clip(u(t),−2, 2), otherwise
(18)
The encoder f is a function from [cos θ, sin θ,dθ/dt] to [θ, 1/3ml2 · dθ/dt], which is invertible and Lipschitz. Similarly, the decoder g is Lipschitz. Besides, the right hand side of Hamiltonian canonical equations is continuous and bounded.

heading: B PROOFS
text: Theorem 5 (Lipschitz Dynamic Systems) For a dynamic system with a C2 continuous Hamiltonian H : R2K+1 → R, if the state s is in a bounded closed set S ⊂ Rl, the evolution time equals τ , the generalized force Qk is C1 continuous with respect to states and bounded (for any dimension k), and the transformation from states to canonical states f∗ : S → R2K is Lipschitz, then the canonical states are Lipschitz with respect to time, and the environment with respect to canonical states is Lipschitz. Additionally, if the transformation from canonical states to states g∗ : R2K → S is Lipschitz, then the environment is Lipschitz, which means (here s 6= s′)
sup a∈A sup s,s′∈S
dS (snext, s ′ next)
dS (s, s′) <∞. (19)
Proof. We know that u = f∗(s), where u is the concatenation of q and p. Because s is in a bounded closed set, and function f∗ is continuous (a Lipschitz function is continuous), each dimension of f∗ has a maximum value and a minimum value. So each dimension of f∗ is bounded, which means each dimension of q and p is bounded.
Now we are going to look into Hamiltonian canonical equations. BecauseH is C2 continuous, ∂H∂pk and − ∂H∂qk are C 1 continuous (for each k ∈ {1, · · · ,K}). Because qk and pk are bounded for all
the values of k, the concatenation of p, q and t lies in a bounded closed set. So ∂H∂pk and − ∂H ∂qk have a maximum value and a minimum value, which means they are bounded. Note that Qk is also bounded, then the right hand side of Hamiltonian canonical equations is bounded.
We denote the bound of ∣∣∣ ∂H∂pk ∣∣∣ as Mk, then we can give the Lipschitz constant for qk:
|qk(t2)− qk(t1)| = ∣∣∣∣∣ ∂H∂pk ∣∣∣∣ t=ξ (t2 − t1) ∣∣∣∣∣ ≤Mk|t2 − t1|. (20) Here ξ is a value between t1 and t2.
Similarly, we denote the bound of ∣∣∣− ∂H∂pk +Qk∣∣∣ as M ′k, then we can give the Lipschitz constant for pk:
|pk(t2)− pk(t1)| = ∣∣∣∣∣ ( − ∂H ∂pk +Qk ) ∣∣∣∣ t=ξ (t2 − t1) ∣∣∣∣∣ ≤M ′k|t2 − t1|. (21) From Equation (20) and Equation (21), we know that the canonical states are Lipschitz with respect to time. Let the right hand side of Hamiltonian canonical equations be I(u, t). Since dI(u, t)/du is continuous, and u and t lie in a bounded closed set, dI(u, t)/ du is bounded. Then for different initial canonical states ut0 and u
′ t′0 , we have (here the absolute value and the supremum for a vector/matrix is taken over each element of a vector/matrix, and the ≤ between two vectors means the relation between each element):∣∣∣∣∣d(ut0(t)− u ′ t′0 (t)) dt
∣∣∣∣∣ = |I(ut0(t), t)− I(u′t′0(t), t)| ≤ sup
u,t
∣∣∣∣dI(u, t)du ∣∣∣∣ |ut0(t)− u′t′0(t)|
≤ ∣∣∣∣∣∣∣∣sup
u,t ∣∣∣∣dI(u, t)du ∣∣∣∣∣∣∣∣∣∣∣∣ ∞ |ut0(t)− u′t′0(t)|.
(22)
So
d∞(ut0(t+ τ),u ′ t′0
(t+ τ)) ≤ exp { τ ∣∣∣∣∣∣∣∣sup u,t ∣∣∣∣dI(u, t)du ∣∣∣∣∣∣∣∣∣∣∣∣ ∞ } · d∞(ut0 ,u′t′0). (23)
Since dI(u,t)du is bounded, exp { τ ∣∣∣∣∣∣supu,t ∣∣∣dI(u,t)du ∣∣∣∣∣∣∣∣∣∞} <∞. This means the canonical states are Lipschitz with respect to initial canonical states. Because the composition of two Lipschitz functions is Lipschitz and f∗ is Lipschitz, we know that the environment is Lipschitz with respect to canonical states.
If we further assume that g∗ is Lipschitz, we know the environment is Lipschitz uniformly for all the action a with respect to states (the composition of two Lipschitz functions is Lipschitz), which is just Equation (19). This concludes the proof.
Theorem 6 (Lipschitz NODA) For the NODA model, if the state s is in a bounded closed set S ⊂ Rl, f : S → R2K is Lipschitz, the evolution time equals τ , the action am is C1 continuous with respect to states and bounded (for any dimension m), function h is C1 continuous, then canonical states are Lipschitz with respect to time, and NODA with respect to canonical states is Lipschitz. Additionally, if the transformation from canonical states to states g : R2K → S is Lipschitz, then NODA is Lipschitz, which means (here s 6= s′)
sup a∈A sup s,s′∈S
dS (snext, s ′ next)
dS (s, s′) <∞. (24)
Proof. The proof is similar to the proof of Theorem 5, and here we just give a sketch. We know u = f(s), where s is in a bounded closed set. Because function f is continuous, each dimension of q and p is bounded.
Because qk and pk are bounded for all the values of k, the concatenation of q, p, t and a lies in a bounded closed set. As a result, the output of the continuous function h is bounded. We denote the bound of the kth output of the right hand side of the ODE as Wk.
|uk(t2)− uk(t1)| ≤Wk|t2 − t1| (25)
Equation (25) tells us that u is Lipschitz with respect to time t.
After that, we can just follow the corresponding part (the part after Equation (21)) in Theorem 5’s proof to get final results. This concludes the proof.
With this bound of multi-step transition, we can give bounds for state values. These bounds tell us that under certain conditions, the optimal policy learned from the simulator and the optimal policy learned from the environment do not differ much.
Theorem 7 (Transition Error Bounds) Under the conditions in Theorem 5 and Theorem 6, we already know that the transition function TG (s′ | s,a) of the environment and the transition function T̂G (s
′ | s,a) of the NODA model is Lipschitz. We denote the Lipschitz constant of these transition functions as K1 and K2. Let K̄ = min{K1,K2}, then ∀n ≥ 1:
δ(n) := W ( T̂nG (· | µ), TnG (· | µ) ) ≤ ∆ n−1∑ i=0 (K̄)i. (26)
Here ∆ is defined as the upper bound of Wasserstein metric between T̂ (· | s,a) and T (· | s,a).
The original theorem (Asadi et al., 2018) gives a a bound for a fixed action sequence. However, here our definitions of Lipschitz environments and models are uniform for all actions. So, by using the original proof, we give a same error bound for all possible action sequences. Thus, we get a uniform error bound under the continuous action space. This concludes the proof. Now we are going to give bounds for state values.
Theorem 8 (Value Error Bounds) Under all the conditions in Theorem 7, if the reward function R(s) is (uniformly) Lipschitz, which means we can define
KR := sup a∈A sup s1,s2∈S |R (s1,a)−R (s2,a)| dS (s1, s2) <∞. (27)
If we define state values as
VT (s) := ∞∑ n=0 γn ∫ TnG (s ′ | δs)R (s′) ds′, (28)
where δs means the probability that the state is s equals 1. Then ∀s ∈ S and K̄ (defined in Theorem 7)∈ [0, 1γ ], we have ∣∣VT (s)− VT̂ (s)∣∣ ≤ γKR∆(1− γ)(1− γK̄) . (29) The original theorem (Asadi et al., 2018) is for action space A = {a}, which means there is only one action. Besides, the original theorem assumes that the reward function only depends on state s.
However, these strict limitations can be overcome by our uniformly Lipschitz models and requiring a uniformly Lipschitz reward function, as is shown in Equation (27). As long as these conditions are satisfied, we can just follow the path of the original proof. Specifically, for any action sequence, Equation (29) holds, so we get a uniform error bound. This concludes the proof.
Here we can find that a crucial condition in the proofs of Theorem 7 and Theorem 8 is that both the model and the environment is Lipschitz. In fact, Theorem 5 and Theorem 6 lay the foundations of uniform transition error bounds and value error bounds.

heading: C ALGORITHM DETAILS
text: Algorithm 2, 3 and 4 show how we combine SAC (Haarnoja et al., 2018) with NODA. Generally, we use interactions with the environment to train a NODA model, and use it to generate imaginary trajectories to facilitate the training of SAC. We combine AE with SAC in the same way. For efficiency, we run 50 interactions with the real environment, and then update the agent and our model 50 times, which is also implemented by spinning up (Achiam, 2018). We also reduce the number of imaginary training batches per epoch when interactions with the real environment is sufficient.
Combining NODA with MBRL methods is easier. For Dreamer (Hafner et al., 2019), it itself works in an auto-encoder’s fashion, that is, it encodes observations and actions into compact latent states and makes decisions and predictions in the latent state space using its world model, after which the observations are reconstructed via representation learning. Therefore, in order to combine NODA with Dreamer, we just need to replace the deterministic path(a GRU cell) of the recurrent state space model(RSSM) in Dreamer with an ODE network, and we get NODA-Dreamer. The input compact state to the original deterministic state model in Dreamer, which is believed to be a canonical state, evolves by going through the ODE network. We also similarly implement NODA-BIRD by replacing the deterministic path of the RSSM in BIRD (Zhu et al., 2020) with an ODE network, since BIRD and Dreamer have the same RSSM. All other algorithm details remain the same as in Dreamer and in BIRD respectively.

heading: D EXPERIMENTAL DETAILS
text: In the pixel pendulum task, we run all the models for 3,000 batches with a batch size of 200. we decode the output of h as the velocity prediction and set the next state to the evolution of the current state after a very short time (assuming the velocity does not change) as a weak supervision. More details can be found in our code. In the real pendulum task, we run all the models for 2,000 batches with a batch size of 200. The learning rate for both tasks is 1e-3.
When comparing different NODA models in the Ant-v3 task, we generate 20,000 steps for training and 20,000 steps for testing by a random policy. Each NODA model is trained by 3,000 batches. The batch size is 256, and the learning rate is still 1e-3.
For TD3 (Fujimoto et al., 2018), SAC (Haarnoja et al., 2018), AE-SAC and NODA-SAC, we use the code provided by spinning up (Achiam, 2018) (the PyTorch (Paszke et al., 2019) version). We modify it to use GPU, and use batch size=100 for TD3 and SAC (original setting), and batch size=256 for AE-SAC and NODA-SAC. Other parameters are the same as what are used in spinning up. We
Algorithm 2 NODA-SAC for Reinforcement Learning Input: SAC actor π, SAC critics q1, q2, environment env, NODA model m (with optimizer optm), warmup steps N1, model generate intervals N2, model planning steps N3, interaction bufferN4, batch sizeB1, model batch sizeB2, SAC training function T1, NODA training function T2, data generation function for NODA G done← True step← 0 D ← [ ] Dm ← [ ] repeat
if done then s← env.reset() end if if step ≤ N1 then a← env.sample(B1) s2, r, done← env.step(a) D.append({s, a, s2, r, done}) else a← π(s) s2, r, done← env.step(a) D.append({s, a, s2, r, done}) batch← D.sample(B1) π, q1, q2 ← T1(π, q1, q2, batch) m← T2(m, batch, optm) if step%N2 == 0 then Dm ← G(π, q1, q2,m,D,Dm, B2, N3) batch← D.sample(B1) π, q1, q2 ← T1(π, q1, q2, batch)
end if end if step← step+ 1
until step = N4 return SAC actor π, SAC critics q1, q2, NODA model m
Algorithm 3 Interaction function G in NODA-SAC Input: SAC actor π, SAC critics q1, q2, NODA model m, environment buffer D, model buffer Dm, batch size B, model planning steps N , learning rate lr (default=1e-5) batch← D.sample(B) s, a← batch.s, batch.a planning step← 0 repeat s2, r, done← m.step(s, a) Dm.append({s, a, s2, r, done}) s← s2 a← π(s) planning step← planning step+ 1
until planning step = N return model buffer Dm
implement NODA by PyTorch, and we use torchdiffeq5 as the implementation of the ODE network (Chen et al., 2018; 2021).
In addition to MuJoCo (Todorov et al., 2012; Schulman et al., 2015) environments, we also compare these algorithms over a simple physical environment, Pendulum-v0 in Gym (Brockman et al., 2016),
5https://github.com/rtqichen/torchdiffeq
Algorithm 4 Training function T2 for NODA Input: NODA model m, a batch of buffer batch, optimizer opt s, a, s2, r ← batch.s, batch.a, batch.s2, batch.r losssrecon ← (||m.decoder(m.encoder(s))− s||22).mean() m.set state(s) s′2, r
′ ← m.step(a) lossspred ← (||s′2 − s||22).mean() losss ← losssrecon + lossspred lossr ← (||r′ − r||22).mean() loss← 0.5 · losss + 0.5 · lossr opt.zero grad() loss.backward() opt.step() return NODA model m
and the result is shown in Figure 5. It shows that the return of NODA-SAC converges quickly, and is quite stable.
For NODA-Dreamer, all the experimental details remain the same as those in Dreamer, except that we run only 200,000 steps for every task we try, and we evaluate the performance of the model every 5,000 steps. What’s more, we use tfdiffeq6 as the implementation of the Neural ODE, which runs entirely on Tensorflow Eager Execution. We evaluate Dreamer, BIRD and NODA-Dreamer on 6 tasks in total in Figure 6. We also evaluate NODA-BIRD, which is mentioned in the previous section, on 3 tasks, as shown in Figure 7.
6https://github.com/titu1994/tfdiffeq

Review summaries:
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.

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.