Model-based Reinforcement Learning with a Hamiltonian Canonical ODE Network
Authors: Yao Feng
Yuhong Jiang
Hang Su
Dong Yan
Jun Zhu
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement learning
Hamiltonian canonical equation
ODE
World model
Sample efficiency

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.

Under review as a conference paper at ICLR 2022
MODEL-BASED REINFORCEMENT LEARNING WITH A HAMILTONIAN CANONICAL ODE NETWORK
Anonymous authors Paper under double-blind review
ABSTRACT
Model-based reinforcement learning usually suffers from a high sample complex- ity in training the world model, especially for the environments with complex dy- namics. 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 impose Hamiltonian mechanics (e.g., the dimension of the physical equations) which can further accelerate training of the environment models. It can consequentially em- power an RI. agent with the robust extrapolation using 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 envi- ronment dynamics effectively with a high sample efficiency, making it possible facilitate reinforcement learning agents at the early stage.
INTRODUCTION
Reinforcement learning has obtained substantial progress in both theoretical foundations (Asadi et al.)
2018; Diang 2018) and empirical applications (Mnih et al., 2013: 2015; Peters & Schaal Johannink et al. (019). In particular, model-free reinforcement learning (MFRL) can com- plete complex tasks such as Atari games (Schrittwieser et al.) 2020) and robot control (Rovedal 2020p. 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 applica- tions. Model-based reinforcement learning (MBRL) methods can alleviate this issue by resorting to model to characterize the environmental dynamics and conduct planning (van Hasselt et al. 2019
2006
jtt al.
Moerland et al. (220a)).
In general, MBRL 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] (011))
It has witnessed numerous works approximating the model with various strategies, such least- squares 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 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- nigh-dimensional problems (Plaat et al. 2020). High-capacity models scale well, but they often have low sample efficiency (Plaat al. 2020). The trade-off between scalability and sample complexity remains as 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 suc- cessfully applied in numerous areas of physics from robotics to industrial automation. Specifically, formulate the environments dynamics as ordinary differential equations (ODEs), and then use

---

Under review as a conference paper at ICLR 2022
novel network architecture called Neural Ordinary Differential Auto-encoder (NODA) our world model, which is naturally induced by physical equations.
auto-encoder and ODE network. We use the auto-encoder to get the underlying physical variables, and use the ODE network to learn the dy- namics over physical variables. By using NODA, 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) continuous action spaces. is noted that NODA can be combined with both MFRL methods like SAC (Haarnoja et al.] (018)) 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 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 al.)
In particular, NODA consists of two parts
2015).
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
tuple (S, A, T, R, 7), where S is the state space, A is the action space, T: S x A + P(S) is the transition function, R: S x A - IR is the reward function, and E [0, 1) discount factor. We denote the set P(-) all probability measures 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 learn a policy by integrating planning with the aid of known or learned model (Mo- erland et al. 2020b), and an essential part of MBRL is to learn transition function that characterizes the environment. The transition function above defined over a given state, but we can generalize it to represent the transition from state distribution P(S). By calling the generalized transition function recursively, we can get the nestep generalized transition function, which defined as:
Definition (Transition Functions) In a metric state space ds) and an action space A, we can define the generalized transition function of Tg (over state distribution 2). and the n-step general- ized transition function of TG (for fixed sequence of actions) (Asadi et al. 2018) as
Tg (8' s,
T (s'
(1)
Tp (- ()) =Tg(( 1......1. .g (: 12, ao) an-1)
recursive calls
Here the generalized transition gives the distribution of outcome under certain state distribution. For MBRL, it is nontrivial to learn the transition function (i.e., the dynamics for physical environ- ment) because S can be high-dimensional. Several attempts have been made to learn such dynam- ics, while they have various limitations. For instance, probabilistic inference for learning control (PILCO) (Deisenroth tion, 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- sigh-dimensional state and action spaces usually require much more data samples (Plaat et al. 2020). Monte Carlo tree search (MCTS) (Silver et al. (017) can introduce human knowledge of the transition function to enhance learning, but it is restricted to cases where transition functions are totally known.
Rasmussen] ( uses Gaussian processes to model the transition func-
2

---

Under review as a conference paper at ICLR 2022
2.2 HAMILTONIAN MECHANICS
Methods of analytical mechanics have been introduced to predict the evolution of dynamic systems such pendulums. For example, Lagrangian neural networks (Lutter al. 2019: Cranmer al., 2020) and Hamiltonian neural networks (Greeddauus 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 dif- ferential 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 Specifically. in Hamil- tonian mechanics, we use pairs of generalized coordinate and generalized momentum (gk. P&) to completely describe a dynamic system (k E {0, (..... K) }). where K is the dimension of general- ized coordinates.Ir It is noted that can be intuitively interpreted as the degree of freedom. We denote Pa and ds. canonical states, and they are minimal independent parameters which can describe the state of the system. We define Hamiltonian H : 2R2K+1 + R as a function of these variables and time Then the evolution of the system satisfies Hamiltonian canonical equations (Junkins & Schaub 2009)
dgk OH
df
Op.-
dpk
8H
+ Q=(1),
di
0y:
where & E {0, 1, .... KK, and (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. More- over, it is possible to incorporate prior knowledge into Hamiltonian canonical equations, e.g., by assuming 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 general- ized momentum, which only requires the knowledge of the "degree of freedom" for a given system but makes 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 variety of different tasks. We empirically examine the these advantages in the experimental section.
METHODOLOGY
We now formally present NODA, which consists of auto-encoder and 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 deterministic, otherwise an SDE network (Li ct al. 2020) can be used instead. Then we can use NODA to assist reinforcement learning by generating imaginary trajectories, as outlined in Algorithm [ Besides, we also discuss the prior knowledge that can be incorporated with NODA.
3.1 MODELING HAMILTONIAN CANONICAL EQUATIONS
We focus on modeling Hamiltonian canonical equations, since after getting these equations, we know the continuous evolution of the system. For many RL Rnviironmesnts such as MuJoCo (Todorov) cet al.) 2012), the discretization of the continuous evolution is just the transition function. However, it may be non-trivial to get these equations for real dynamic system? So we propose to use neural ODE or ODE networks (Chen et al. 2018 202T) to model these first-order differential equations.
Hamiltooiin canonical equations are widely used with clear physical meanings and more symmetric forms than the I Lagrangian case.
We provide example show the derivation of these equations in Appendix A, where we derive these equations under the setting of Pendulum-vO task in Gym (Brockman al. 2016).
3

---

Under review as a conference paper at ICLR 2022
Staess
Actima
Pacader
ODE Netwock
du alta]
Mu a, ().( (ooto
(0.p)
Decoder y
Plascca Dynanics
Decoder
dg
07t
Up
JN
Reward
db
Q211
Figure 1: NODA model structure. The state goes through the encoder f to get latent states q. P. The evolution of u (q.p) p) defined by the ODE network is decoded by decoder g to get the next state &neat. The decoder g predicts reward using q. p and action a.
Specifically. Neural ODE (Chen al.] 2018) provides an elegant framework to model differential equations. is defined by da(r)) de = ( 0). Here f is a neural network with parameter 0, and we use an integrator transform the input (0) to output /(T)). where 7 is given time horizon. Neural ODE can be viewed as a continuous version of ResNet (He et al. (016). Here we assume that the agent affects the environment through forces. So the generalized forces Q in Eq. (2) correspond to the effects of action a in reinforcement learning, since generalized forces can describe all physical perturbations.
3.2 NODA MODEL STRUCTURE
Hamiltonian canonical equations deseribe the evolution over canonical states, but the observed state is not necessarily composed of canonical states 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 Specifically, in the auto-encoder part, we use function f S -> R22 to the canonical states. Here we denote the concatenation of (q, p) as u. so u = fs). We also need to restore the state from canonical states, so we further define function g I22 + S. After going through the ODE network (defined later), we can decode the evolved canonical states s' get the g(s').
next state: Snext
Since the canonical states are concise but complete to describe the system, the auto-encoder can serve dimension reducer to extract and refine state 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 Rmm More specifically, r g' g'(u, a). Then we let the predicted canonical states go through an ODE network to evolve.
another decoder g' 222K A IR to get predicted reward
In the ODE network part, the encoded latent state u evolves through the ODE over a time interval (oo, + 7). We can define the ODE network function ODE(h, u, a, lo- 4o + +), which returns the value of the solution of neural ODE at time to + The neural ODE is described as:
du h(u, a, t), t E [to. to + 7],
(3)
dt
where h is neural network.
After combining the ODE network and the auto-encoder, can get full structure of NODA, which shown Figure The inputs of NODA are the state s and the action a, and the outputs of NODA
The auto-encoder is widely used in machine learning
(Ballard
1 1987
Ng et al. 2011
Kingma & Welling
2013
4

---

Under review as a conference paper at ICLR 2022
the predicted state next and the predicted reward Our model is formally characterized as:
B = (ODE u, a, to. to T)),
(4)
r r= g= (u, a).
To learn the unknown parameters in f.g and h, define the total loss function for NODA convex combination of the state loss and the reward loss:
L م ( 2|12 + || S 123) + -
(5)
where μ E (0, and denotes the ground truth. By minimizing this objective, we jointly train the RL agent and model using interactions with the environment, and use model together with the real environment to provide training data for the RI. agent after certain training steps. More details provided in Appendix C.
As model of Hamiltonian canonical equa- tions, NODA can describe general dynamic systems, and the form of canonical states can be very general. Moreover, NODA cor- responds with real dynamics of the environ- ment, which allows us to use our knowl- edge 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 sys- tem, and we demonstrate the effects of such dimension in the experimental part. For stronger prior knowledge, we 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 learn- ing the experimental part. This makes it possible for us make trade-off's between human expert knowledge and the required number of training samples.
Algorithm NODA for Reinforcement Learning
Input: Reinforcement learning agent 8, environ- ment env, NODA model m repeat
Collect data from interactions with env Use interactions with env to train m Use interactions with both env and train until Certain number of steps return Reinforcement learning agent 8, NODA model m
4 THEORETICAL ANALYSIS
In this section, we present error bounds for both multi-step transition and values for NODA.
4.1 LIPSCIITZ DYNAMIC SYSTEMS AND NODA
Former work (Asadi al. 2018) studied using Lipschitz function groups to model the transition of the environment. However, there are eтror bounds only if the real environment is Lipschitz, and that work did not mention what kind environments have a Lipschitz transition. Theorem proves that under certain conditions, the transition functions of dynamic systems are Lipschitz, which is the basis of transition and state value eтror bounds.
Before stating the theorem, 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 (Lipschitz Models) In a metric state space (S, ds) and an action space A, we say that ff is Lipschitz model the Lipschitz constant (Asadi al. 2018)
ds (ʃ (81, a) , (82, a)) <0.
Ky:= sup sup DEA...SEE
(6)
ds (81,
Theorem (Lipschitz Dynamic Systems) For dynamic system with C² continuous Hamilto- nian H. 12 - R. if the state is a bounded closed set S C R., the evolution time equals the generalized force Qs O continuous with respect states and bounded (for any dimension
5

---

Under review as a conference paper at ICLR 2022
k), and the transformation from states canonical states f S B2K 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" RR2K + S is Lipschitz, then the environment is Lipschitz, which means (here & + 8')
ds (Seea: Spor)) < OC.
sup sup a. ds (8,
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 analyse the ODE part. The conditions of H and Qs guarantees that the right hand side of the ODE is Lipschitz, which leads the Lipschitz condition in this part.
Actually, the conditions in Theorem are not hard to satisfy. The Hamiltonian H. itself already has derivatives with respect to q. p and f. We just further assume that the Hamiltonian H''s derivatives with respect to and is differentiable continuous. Actually, for many dynamic systems such as spring-mass mass systems or three-body systems, the Hamiltonian C continuous.
As an imitation of Hamiltonian canonical equations, it is natural to validate the NODA model is Lipschitz. We analyse this Theorem which the Lipschitz model will be one of the conditions in error bounds.
Theorem (Lipschitz NODA) For the NODA model, if the state s is in a bounded closed set S C R., f S + I22K is Lipschitz. the evolution time equals т, the action Om is C continuous with respect to states and bounded (for any dimension m), function states are Lipschitz with respect to time, and NODA with respect to canonical states is Lipschitz. Lipschitz, then
C continuous, then canonical
Additionally, if the transformation from canonical states to states g B2K NODA is Lipschitz, which means (here + s')
ds (Suexs Suen)
sup sup aEA.
00.
(8)
ds (s))88
The proof is in Appendix B with the similar idea Theorem
4.2 UNIFORM ERROR BOUNDS FOR MULTI-STEP TRANSITION
Here we firstly define a metric called Wasserstein Metric in Definition 3 Wasserstein metric de- scribes how 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 (Wasserstein Metric) In a metric space (X,d), the Wasserstein metric between two probability distributions 21 and 22 in P(X) is
// (8), 82) d (8), 82) digdss,
W (2) 22) 22 : inf jea
where A is set of all joint distributions j on x with marginal distributions žι and 22 (Arjovsky et al. 2017)
Under the conditions of Theorem and Theorem 2 the dynamic system and the NODA model Lipschitz models. On that basis, we can give uniform error bounds for multi-step transition, which shown in Theorem 3 The theorem tells 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 the imaginary environment.
Theorem (Transition Error Bounds) Under the conditions in Theorem and Theorem 2 we already know that the transition function Tg (s' 88 a) the environment and the transition function

---

Under review as a conference paper at ICLR 2022
Tg (s' a) of the NODA model are Lipschitz. We denote the Lipschitz constants these transition functions as K1 and Kg, respectively. Let min K1, K2), then vn 2 1:
8(n) = W ( (T)
A
Σ (RY, 1=0
(10)
where A is upper bound of Wasserstein metric between T 8) and T 8, This upper bound tight for linear and deterministic transitions.
The original theorem (Asadi al. 2018) gives 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 uniform error bound under the continuous action space. This concludes the proof.
Theorem 4 (Value Error Bounds) Under all the conditions in Theorem if the reward function R(+) is uniformly Lipschitz for any action, which means we can define
R a. - ((2,
Kn == sup sup a 81,
(11)
ds (81,
we define state values
50 Vi((s = Σ 7m
Tē (s' 62) R2))
(12)
where da means the probability that the state being s equals Then Vs E and K (defined in
Theorem E [0, 1), we have
RKRA (1- 71 RK)
(T(()) Vy(())
(13)
We put the theorem the uniform error bounds of state values here, and defer its proof and more detailed discussions of theorems to Appendix B.
5 EXPERIMENTS
In this section, we validate the efficiency of NODA's learning and its ability of boosting both MFRL and MBRL methods.
5.1 EXPERIMENTAL SETUP
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 Hamilto- nian neural network (HNN) (Greydanus et al.) 2019h bascline methods. For MFRL methods, choose SAC and TD3 as our baselines. We modify the code from spinning up (Achiam) 2018) to GPU in training. For MBRL methods, in the literature. Specifically, we focus on two methods on visual control tasks, Dreamer (Hafner al. 2019) and BIRD (Zhu et al. 2020), our MBRL baselines. We reproduce their results by their released codes respectively.
compare approach with state-of-the-ar methods
Implementation We implement NODA mainly by using pytorch (Paszke et al., 2019) and torchd- iffeq (https:/github..cmm) http:///whhbb.omm/prgeeetoocchce
the implementation of the ODE network (Chen) ct al.) 2018 202I). 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 tfdiffey (https:/giitubb.om//iuu htt../////iuu19///ffffffff. mentation of the ODE network, which runs entirely on Tensorflow Eager Execution. We refer to Appendix C for more implementation details.
the imple-

---

Under review as a conference paper at ICLR 2022
Poul Pendullun
Pendskem
1o
NeDi--rigital NODA- Trad--e
1.3
NODAA
10
NODA
0.5
10
16'
16
16'
216
430
930
252
510
750
1000
Step-
Suep
Sueze
(a)
(b)
(C)
(d)
Figure 2: (a, b): Results of NODA and other physical simulators in modeling dynamic systems. (c): Sensitivity NODA's latent dimension when learning the transition the Ant-v3 task. (d): Transferability of NODA in the Ant-v3 task.
the setting of pixel pendulum and real pendulum in the paper of HNN (Greedaanus 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) 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 make a fair comparison with them. More experimental details can be found in Appendix D.
Tasks For physical simulation tasks,
5.2 SIMULATION EFFECTIVENESS
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. 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 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 to predict the two-step transition. The figure suggests that we can transfer NODA over similar environments to accelerate training.
7000
201
34c
1921
120g
150
100
1806
2006
201h
A5-BAC
MIDA-AUC
Figure 3: Results of TD3, SAC, AE-SAC and NODA-SAC on MuJoCo environments in Gym. The models are evaluated every 4k steps for 10 episodes, and the means and standard deviations computed across 10 episodes across 4 seeds. The dashed lines are the best asymptotic performances of 1M steps.
8

---

Under review as a conference paper at ICLR 2022
Table Experiments in MuJoCo environments (200k steps, 4 seeds)
Algorithms
Environments
MFRL Algorithms
InvertedPendullm
HalfCheetah
Hopper
Ant
TD3
1000.0:0.0
7154.8+364.9
283.14
2218.84539.4
SAC
1000.0:0.0
9150..5642.99
476.9+297.6
492.3+209.4
AE-SAC
1000.0=0.0
8703.04738.4
255.6134.4
400.5+112.2
NODA-SAC
1000.0=0.0
9241.8+567.7
832.5+283.4
402.4+54.9
MBRL Algorithms
Hopper-Stand
Walker- Walk
Finger-Spin
Walker-Run
Dreamer
198. 03254.3
642.9458.5
366.2+252.5
229.1=11.4
BIRD
236.4=138.2
601.69
361.74246.8
216.0 +23,4
RODA-Dreamer
260.2+284.9
764.9#146.9
428.5+96.3
224.2+34.7
5.3
SAMPLE EFFICIENCY
Hopper Stand
Walker Walk
Finger Spin
W Walkee Run
Egpoode Retum
500
500
Epptode
0k
50k
100k
150k
200k
0k
5tk
100k
150k
200k
0k
50k
130k
150k
200k
0k
50k
100k
155k
20%
Eortronmten Stops
Eavirorment Stops
Eavirommene Stops
Saaiionmmre Steps
Dreamer
BIRD
NODA Dreamer
Figure 4: Results of Dreamer, BIRD and NODA -rrem--r on DMC. The models are evaluated every 5k steps for 10 episodes, and the means and standard deviations are computed across 10 episodes across 4 seeds. The dashed lines are the asymptotic performances 5M steps of Dreamer mentioned in (Hafner al. 2019).
The episode returns achieved by MFRL methods and MBRL methds (with/without NODA) shown in Table [ Figure and Figure 4 For MFRL methods, NODA-SAC achieves the best per- formance with good robustness (see the InvertedPenduluu-v22 task in Figure 3). This means that NODA can enhance the performance of the fate-of-thee-a art algorithm SAC by imaginary trajecto- ries. 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 knowl- edge of the existence of Hamiltonian indeed helps. But for environments that SAC 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-Dreamee lags far behind Dreamer, but quickly catches up the second half. This suggests that NODA has the power to accelerate training after warm-up up steps. Generally, NODA enhance the performance for both MFRL and MBRL methods. More analyses and experiments can be found in Appendix D.
6 CONCLUSION
This paper proposes novel simulator called NODA for reinforcement learning. Motivated by Hamiltonian canonical equations in physics, NODA has clear physical meanings. This allows us incorporate prior knowledge 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 prior knowledge transfer learning to further boost its training. Not only HNN, and sample-efficient, but NODA can improve the sample efficiency of both MFRL and NODA itself MBRL methods such as SAC and Dreamer.

---

Under review as a conference paper at ICLR 2022
7 ETHICS STATEMENT
In this paper, we only use data sets that we are permitted to use, and we present our method and results in 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 emrors 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.
REPRODUCIBILITY STATEMENT
We give formal proof' 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.
REFERENCES
Joshua Achiam. Spinning Up in Deep Reinforcement Learning. 2018.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning (ICML), pp. 214-223. PMLR, 2017.
Kavosh Asadi, Dipendra Misra, and Michael Littman. Lipschitz continuity in model-based rein- forcement learning. In International Conference on Machine Learning (ICML), pp. 264-273. PMLR, 2018.
Dana H Ballard. Modular learning in neural networks. In AAAI Conference on Artificial Intelligence, volume 647, pp. 279-284, 1987.
Justin A Boyan. Least-squares temporal difference learning. In International Conference Ma- chine Learning (ICMI.. pp. 49-56. Citeseer, 1999.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.015400 2016.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Advances in Neural 2018.
Information Processing Systems, pp.
Ricky T. Q. Chen, Brandon Amos, and Maximilian Nickel. Learning neural event functions for ordinary differential equations. International Conference on Learning Representations (ICLR), 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, pp. 6571-6583, 2018.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.046300 2020.
Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach policy search. In International Conference on Machine Learning (ICML), pp. 465-472, 2011.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference on Machine Learning (ICML), pp. 1587-1596.
PMLR, 2018.
Herbert Goldstein, Charles Poole, and John Safko. Classical mechanics, 2002.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, pp. 15379-15389, 2019.
10

---

Under review as a conference paper at ICLR 2022
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: off-policy max- imum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning (ICML), pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint (rXiv:1912.0166 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European Conference on Computer Vision, pp. 345.645. Springer, 2016.
Todd Hester and Peter Stone. Learning and using models. Reinforcement Learning, pp. 111-141. Springer, 2012.
Nan Jiang. Pac reinforcement learning with an imperfect model. In AAAI Conference Artificial Intelligence, pp. 3334-3341, 2018.
Tobias Johannink, Shikhar Bahl, Ashvin Nair, Jianlan Luo, Avinash Kumar, Matthias Loskyll, Juan Aparicio Ojea, Eugen Solowjow, and Sergey Levine. Residual reinforcement learning for robot control. In International Conference on Robotics and Automation (ICRA), pp. 6023-6029.
IEEE, 2019.
John Junkins and Hanspeter Schaub. Analytical mechanics of space systems. American Institute of Aeronautics and Astronautics, 2009.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv:1312.61 14, 2013.
arXiv preprint
Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforcement learning. arXiv preprint arXiv: 2019.
Sergey Levine and Picter Abbeel. Learning neural network policies with guided policy search under unknown dynamics. In Advances in Neural Information Processing Systems, pp. 1071-1079, 2014.
Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3880- PMLR, 2020.
Michael Lutter, Christian Ritter, and Jan Peters. Deep lagrangian networks: using physics as model prior for deep learning. In International Conference on Learning Representations (ICLR), 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 5 518(7540;:59-5333 2015.
Thomas M Mocrland, Joost Brockens, and Catholijn M Jonker. A framework for reinforcement learning and planning. arXiv preprint ar arXiv:2066.15009 2020a.
Thomas M Moerland, Joost Broekens, and Catholijn M Jonker. Model-based reinforcement learn- ing: survey, arXiv preprint arXiv:2006.1677 16712, 2020b.
Andrew Ng et al. Sparse autoencoder. CS294A Lecture notes, 72(2011):1-12 2011.
2011.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: an imperative style, high-performant deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer. F. d'Alché-Bue E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systems, pp. 8024-8035. Curran Associates, Inc., 2019.
11

---

Under review as a conference paper at ICLR 2022
Jan Peters and Stefan Schaal. Policy gradient methods for robotics. In International Conference Intelligent Robots and Systems, pp. 2219-2225. IEEE, 2006.
Aske Plaat, Walter Kosters, and Mike Preuss. Model-based deep reinforcement learning for high- dimensional problems, survey. arXiv preprint arXiv:2008. 05598, 2020.
Loris Roveda, Jeyhoon Maskani, Paolo Franceschi, Arash Abdi, Francesco Braghin, Lorenzo Moli- nari Tosatti, and Nicola Pedrocchi. Model-based reinforcement learning variable impedance con- trol for human-robot collaboration. Journal of Intelligent & Robotic Systems, pp. 1-17, 2020.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7833):604-609, 2020.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438, 2015.
Jian Shen, Han Zhao, Weinan Zhang, and Yong Yu. Model-based policy optimization with unsu- pervised model adaptation. H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 2823-2834. Curran Associates, Inc., 2020.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. Nature, 550(7676):354-359.. 2017
Olivier Sprangers, Robert Babuška, Subramanya Nageshrao, and Gabriel AD Lopes. Reinforce- ment learning for port-hamiltonian systems. IEEE Transactions on Cybernetics, 45(5):1017- 1027, 2014.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite, arXiv preprint 1rXiv:1801.000 000900, 2018.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
Hado P van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in rein- forcement learning? In Advances in Neural Information Processing Systems, pp. 14332-14333, 2019.
Guangxiang Zhu, Minghao Zhang, Honglak Lee, and Chongjie Zhang. Bridging imagination and reality for model-based deep reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.). Advances Neural Information Processing Systems, volume 33, pp. 8993-9006. Curran Associates, Inc., 2020.
A HAMILTONIAN CANONICAL EQUATIONS: PENDULUM-VO AS AN EXAMPLE
The settings of the Pendulum-vO task be found GitHut
Let æ be the axis horizontally to the right, y be the axis stright up. The pendulum's 's position can be described by generalized position where 0 is the angle from the y axis to the pendulum.
For the free end, x = -i 0, y = cos 0. For a force perpendicular to the pendulum (the angle dr
from the y axis is 0+ with its 2-norm equaling u. its generalized force Bi u, + By uy
dr
cl 0 (-u cos 0) - sin (-u sin 0) = ul, which is just the torque the force.
pendulum,
12

---

Under review as a conference paper at ICLR 2022
The kinetic energy T the Lagrangian in Equation (14
1/6ml (d0/ dt)?, and the potential energy V 1/2mgl cos A.
get
<(0.
ml d0
2
mgll cos
(14)
di
2
We denote 0 q, the canonical momentum for p, then
p OC 1
dq
ml
(15)
3
dt
Then we can write the Hamiltonian in Equation 16.
dq
3p²
H(q. p, t)
mgl cos q
(16)
2ml2
Then Hamiltonian canonical equations give out the dynamics, as is shown in Equation 17
dq 07
3p
dt 0p
ml?
(17)
dp
07l
+ Q()) = 2 mgl sing + u(t)
dt
3q
This is the general case, but the code implementation is a little different. In the code, u(t) is clipped [-2, 2). clipped to [- :T), and d0/ dt is clipped to [-8, 8). This exactly gives the bounds of the state space, since states are [cos 0, sin 0, d0/ dt). The clips can be added to Hamiltonian canonical equations, so the code implementation can be viewed as a solver for Equation (18).
dg
3p ml? min( mgl sin cl)p(u(()) -2,2), if
dt
8
3mnl2
(18)
dp
8
max( 2
mgl sin clip(((())))) 22, 2,2)),)) 0), if
dt
3ml2
mgl sin clip(u(()), 22, otherwise
The encoder f is function from [cas 0, sin 0, d0 dt] to [0, 1/3ml2 d0/ 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.
B PROOFS
Theorem (Lipschitz Dynamic Systems) For a dynamic system with a C2 continuous Hamiltonian H. 1g2K+1 - IR, if the state 8 is in bounded closed S C R', the evolution time equals T, the generalized force Qs is C² continuous with respect to states and bounded (for any dimension k), and the transformation from states canonical states f* : 222K is Lipsehitz, 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 stales g" .R2KK S is Lipschitz, then the environment is Lipschitz, which means (here # 8')
ds (Snent S'uet)
sup sup a A.
00.
(19)
ds (8,
Proof. We know that = (s), where the concatenation of q and p. Because is in a bounded closed set, and function /* is continuous (a Lipschitz function is continuous), each dimension of f* has maximum value and minimum value. So each dimension f is bounded, which means each dimension of q and p is bounded.
Now we are going to look into Hamiltonian canonical equations. Because H is C2 continuous, 8H
Tps
and
8H C continuous (for each & {I,... K}). Because qs and Pk: bounded for all Bek.
13

---

Under review as a conference paper at ICLR 2022
the values of k, the concatenation of p. 9 and lies in a bounded closed set. So UK and 8n have maximum value and minimum value, which means they bounded. Note that Qx. is also bounded, then the right hand side Hamiltonian canonical equations is bounded.
Spa. Bqu.
We denote the bound of
8t Bps
as Mg, then we can give the Lipschitz constant for qli
0H
(t2) qu (ti))
(t2 - (1)
S Malt2 - 4//-
(20)
Oyk le=
Here & is value between 1 and tz.
Similarly, we denote the bound of pk:
0H
Qu as Mg. then we can give the Lipschitz constant for
Opk
(tx) ptt(tt]
0H
Qx
(12 12) s Ma t//-
(21)
Opk
l==
From Equation (20) and Equation (21 know that the canonical states Lipschitz with respect to time. Let the right hand side of Hamiltonian canonical equations be [u, t). Since d/(u. t)/ du continuous, and u and t lie bounded closed set, du, t)/ du is bounded. Then for different initial canonical states Uta and u.g we have (here the absolute value and the supremum for a vec- tor/matrix is taken over each element of vector/matrix, and the s between two vectors means the relation between each element):
d(uro (t) u/y (t))
[(ueo (t), [ulg (t).
dl
dl dl(u,t) t)
sup u,f
Ueo (t) u/g (t)
(22)
du
cI,u,, t) du
11ee (t) - u/g
(t)l (t)
sup u.t
So
dd(u,t)
doo (ur (t++) TT, w/; (t + T)) S exp
sup
(oo ulz).
(23)
du
Since dl(u.)) bounded, cxp
df(u.t) du
cc. This means the canonical states are Lipschitz with respect 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.
du
supu,t
If we further assume that g" is Lipschitz, we know the environment is Lipsehitz uniformly for all the action a with respect to states (the composition of two Lipschitz functions is Lipschitz), which just Equation 19). This concludes the proof.
Theorem (Lipsehitz NODA) For the NODA model, if the state is bounded closed set S C R! f:S. S + I22K is Lipschitz, the evolution time equals T, the action 0m is C² continuous with respect to states and bounded (for any dimension m), function is C² 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 8 + s')
ds (8nexs )yea) < 00.
sup sup a s.''
ds (8,88)
(24)
Proof. The proof is similar the proof of Theorem 5 and here we just give sketch. We know = f(), where s is in a bounded closed set. Because function f is continuous, each dimension of q and p is bounded.
14

---

Under review as a conference paper at ICLR 2022
Because 9k and pk are bounded for all the values of k. the concatenation of q. p, t and lies in 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-
[x us (f1) & Walte t1
(25)
Equation 25 tells us that u Lipschitz with respect to time
After that, can just follow the corresponding part (the part after Equation (2T) in Theorem S proof to get final results. This concludes the proof.
With this bound of multi-step transition, that under certain conditions, the optimal policy learned from the simulator and the optimal policy learned from the environment do not differ much.
give bounds for state values. These bounds tell us
Theorem' (Transition Error Bounds) Under the conditions Theorem and Theorem we already know that the transition function Tg (s' 8, a) the environment and the transition function Tg (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 Vn 2 1:
n
5(n) = w (Te( (T0(11,) T5, S
AAE(Y (KY'
(26)
1=0
Here A is defined as the upper bound of Wasserstein metric between T 8, s, and T (- 8, a).
The original theorem (Asadi et al. 2018) gives bound for fixed action sequence. However, here our definitions Lipschitz environments and models are uniform for all actions. So, by using the original proof, give same error bound for all possible action sequences. Thus, get uniform error bound under the continuous action space. This concludes the proof. Now we are going to give bounds for state values.
Theorem (Value Error Bounds) Under all the conditions in Theorem if the reward function R(s) is (uniformly) Lipschitz, which means we can define
[ a) (82, a)l
KR = sup sup a $1.0ESS
< 00.
(27)
ds (811
we define state values
20 Vr(s) = Tnm T (s' (('))
(28)
n=0
where 8g means the probability that the state is s equals 1. Then vs E S and K (defined in
Theorem 7 [0, 2)), we have
RKRA (s) Vy(s)) (1- 71 - RK)
(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 uniformly Lipschitz models and requiring a uniformly Lipschitz reward function, as is shown Equation (27). As long as these conditions satisfied, 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 find that crucial condition in the proofs of Theorem and Theorem X is that both the model and the environment is Lipschitz. In fact. Theorem and Theorem 6 lay the foundations of uniform transition error bounds and value error bounds.
15

---

Under review as a conference paper at ICLR 2022
C ALGORITHM DETAILS
Algorithm and show how we combine SAC (Haarnoja et al. 2018) with NODA. Generally, we use interactions with the environment to train NODA model, and use 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 casier. For Dreamer (Hafner et al. 2019). itself works in auto-encoder's 's fashion, that is, it encodes observations and actions into compact latent states and makes decisions and predictions the latent state space using its world model, after which the observations 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 modcl(RSMM) in Dreamer with an ODE network, and get NODA-Dreamer. The input com- pact state to the original deterministic state model in Dreamer, which is believed to be canonical state, evolves by going through the ODE network. We also similarly implement NODA-BIL 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 in Dreamer and in BIRD respectively.
Pendulum-v0
Episode Return
-1000
TD3
SAC
-2000
AE
NODA
-3000
0k
50k
100k
150k
200k
Environment Steps
Figure 5: Results of TD3, SAC, AE and NODA on the Pendulum-vO environment in Gym. The models are evaluated every 4k steps for 10 episodes, and the means and standard deviations computed across 10 episodes across 4 seeds. The performance of NODA-SAC is stable, and NODA able to facilitate the training of SAC in an early stage.
D EXPERIMENTAL DETAILS
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, generate 20,000 steps for training and 20,000 steps for testing by a random policy. Each NODA model trained by 3,000 batches. The batch size is 256, and the learning rate is still 1e-3.
For TD3 FFjjimoto et al. 2018), SAC (Haarnoja et al. 2018). AE-SAC and NODA-SAC, we use the code provided by spinning up (Achiam (2018) (the Py Torch (Paszke al. 2019) version). We mod- ify 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 what are used in spinning up. We
16

---

Under review as a conference paper at ICLR 2022
Algorithm 2 NODA-SAC for Reinforcement Learning
Input: SAC actor SAC critics 42, environment env, NODA model m (with optimizer oplm), warmup steps Ni, model generate intervals N2, model planning steps N3, interaction buffer N4, batch size B1, model batch size B2, SAC training function Ti, NODA training function
T2. data generation function for NODA
done +- True
step t-
D -
Dm -[]
repeat
if done then
s st env.reset() end if
if step S N₁ then a < enw.sample(B),
82, r, done- +- env.step(a)
D.append{{s, a, 82, done))
else
a(s) 82, r, dome + env.step(al
D.append( a, 82, r, done))
batch + D.sample(() (B1))
t, 91 Q2 + Ti (T,, qı, 92, batch) + T2 batch, optm)
if step% == 0 then
Dm + G(T,, 91, 92, m. D, Dm, B2, N₂) batch D. (.a)
n, q, + Ti (T,, 41, 42, butch)
end if
end if
step + step + 1
until step N4
return SAC actor x, SAC critics Q1, 12. NODA model m
Algorithm 3 Interaction function G in NODA-SAC
Input: SAC actor SAC critics Qı, NODA model m, environment buffer D, model buffer Dms batch size B, model planning steps N, learning rate Ir (default=le-5) batch + D.sample(B)
8, a. +- batch.s, batch.a
plarming_sted <- 0
repeat
$2, r, done +- m.step(s. a)
mm appendl (s, a, $2, r, done))
82 a - (s))
planning_sted planninngstess until planning _step = N
return model buffer D.
implement NODA by PyTorch, and we use torchdiffed the implementation of the ODE net- work (Chen et al. 2018 2021).
In addition MuJoCo (ooooro et al. 2012 Schulman et al. 2015 environments, also compare
these algorithms over a simple physical environment, Pendulum-vO in Gym (Brockman et al. 2016),
https://github.coo htte::///ttttto.coo////cceeeeeet
17

---

Under review as a conference paper at ICLR 2022
Algorithm 4 Training function T2 for NODA
Input: NODA model m, batch of buffer batch, optimizer opt 8, O, 82, r +- batch.s, s, batch.a, a butch.s2, balch.r loss srecon (I|m.decoderl www.eeermmeeeneodeeee (s))nodder(s) - ( mean()
m.set_state(s)
$22, + m.stepp()
los8 + (11$2 -11B) mm)n)))
loss, -- 10s8 socoo + loss spres
loss, (1|r' ----r1 1011))rmemma meann)
loss t loss, + 0.5 loss,
opt.zero grad()
loss loss.backwoar(() 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 quite stable.
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 tfdiffed 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-BIRE which is mentioned in the previous section, on 3 tasks, as shown in Figure
For
Happer Stand
Walker Walk
Finger Spin
1000
Ppisode Retum
Episode Return
Iipisode Return
500
500
500-
250-
0k
50k
100%
150k
200k
0%
50k
100k
150k
200k
0%
50%
100%
150k
200k
Environment Steps
Eavironment Steps
Eaviroment Steps
Walker Run
Cheelah Run
Quadruped Run
750
Episode Retum
750
Egirede Retum
750
Epiide
500
500
500
250
250
250
0-
0-
50k
100k
150k
200k
0k
50k
100k
150k
200%
0k
100%
150k
200k
Envirenment Steps
Environment Steps
Environment Steps
Dreamer
BIRD
NODA-Dreames
Figure 6: Results of Dreamer, BIRD and NODA-Dreamed on tasks in DeepMind Control Suite. The models are evaluated every 5k steps for 10 episodes, and the means and standard deviations computed across 10 episodes across 4 seeds. The dashed lines are the asymptotic performances of 5M steps of Dreamer mentioned in (Hafner et al. 2019). In many tasks, NODA-Dreames has higher episode return starting early in the training process, which shows that NODA can assist in improving sample efficiency. In "Walker Run", during the first half of training, NODA- -Dremmee lags far behind Dreamer and BIRD, but it quickly catches up in the second half. This suggests that NODA has the power to accelerate training after some warm-up steps. In "Quadruped Run", 200,000 steps are not sufficient for any algorithm to learn meaningful information.
httpps://gitthbb..mm//ttttttu
18

---

Under review as a conference paper at ICLR 2022
Hopper Stand
Walker Walk
Walker Run
Tipioode Retur
Episede Itetur
750
500
500
500
Fyiende
250
0k
50k
100k
1.5kk
200k
50k
100k
150k
200k
0k
50k
100k
150k
200k
Environment Steps
Enviromment Steps
Eaviroument Steps
Dreamer
BIRD
NODA-Dreamer
NODA-BIRD
Figure 7: Results of Dreamer, BIRD, NODA-Dreamer and NODA-BIRD on 3 tasks in DeepMind Control Suite. The models evaluated every 5k steps for 10 episodes, and the means and standard deviations are computed across 10 episodes across 4 seeds. The dashed lines are the asymptotic performances of 5M steps Dreamer mentioned in (Hafner al., 2019). The results suggest that NODA is very promising, in the sense that it can be easily combined with different MBRI. methods, especially auto-encoder-type methods, by simply modifying their transition model.
19


Reviews:
[{"review_id": "v7GwgpaOLC", "reviewer": "Reviewer_grUX", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My main concern is about technical soundness.\n\nIn section 3, the authors mentions \n>For many RL environments such as MuJoCo (Todorovet al., 2012), the discretization of the continuous evolution is just the transition function.\n\nThis is incorrect. For tasks such as pendulum, cart-pole and double pendulum, the underlying dynamics is continuous and the discretization of that continuous dynamics is indeed the transition function. However, for many of the tasks that are tested in this paper (Ant-v3, HalfCheetah-v3, Hopper-v3, Walker), the underlying state evolution is not continuous, since these tasks involves collisions and contact, which results in instantaneous velocity change. In fact, the major contribution of Mujoco is its contact model. The authors seems to know Neural Event Functions (Chen et al, 2021), which specifically addresses this kind of discontinuities in the system. However, only bouncing balls examples are shown in Chen et al. 2021, as addressing contacts in deep learning for those systems in MuJoCo is difficult. \n\nI'm very surprised that the authors didn't comment on this gap between their motivation/models and the systems in the experiments. The ODE-based model assumes continuous dynamics and this is an incorrect inductive bias for systems such as HalfCheetah, which involves contacts. The q and p learned in these systems have no physical meaning at all. Moreover, the authors couldn't justify their choice of Hamiltonian mechanics. A well-known property of Hamiltonian mechanics is that the energy conserves along a trajectory (assuming no action/control is applied). For Ant-v3, HalfCheetah-v3, Hopper-v3, Walker, energy is obviously not conserved due to the inelastic collisions. \n\nThe whole section 4 is dedicated to theoretical analysis. These analysis seems valid to me. However, they are not applicable to most systems in the experiments, since state evolution in those systems are not Lipschitz continuous (again because of contacts). From this perspective, I don't see how the theoretical analysis help us gain any insight into the application. A majority of the theory are extended from Asadi et al, 2018. However, Asadi et al, 2018 uses pendulum and cartpole as experiments to demonstrate how the theory works. And both pendulum and cartpole are both continuous dynamics, which means the theory applies to them. The authors seems to not realize that the theory is not applicable to the systems the study and inappropriately use the theory to gain technical soundness. \n\n", "summary_of_the_paper": "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. ", "summary_of_the_review": "Despite the technical soundness, the proposed model seems to outperform existing models in the experiments. However, due to the concerns on technical soundness, I think this paper is not ready for ICLR yet. Thus I recommend rejection.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "nSK82lSmaXW", "reviewer": "Reviewer_9nPp", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "If the paper does what it claims to do, it would be amazing. In this case, the authors solved the two major problems of physics-inspired deep networks (e.g., HNN & DeLaN), which are the assumption of observing generalized coordinates & forces as well as not including contacts. However, the authors do not provide any experimental validation of this claim and even the architecture shows that the proposed model architecture should not be able to model the locomotion tasks. \n\nFor example, consider the network architecture in Figure 1. This model cannot describe the walker or hopper dynamics as these dynamics include contacts and are dissipative (assuming a(t) = \\forall t). However, the architecture in Figure 1 is conservative, which means that the walker and hopper would need to bounce ground plane (which makes no sense). Therefore, this model architecture is not suitable for the performed tasks and should not work. \n\nThe theoretical evaluation in Section 4 has no relevance. Unfortunately including random theorems, which don't provide any insights, has become an annoying habit in too many ML papers. Could the authors please elaborate on why these theorems are useful or have any relevance? Furthermore, I even question the underlying assumption of Lipschitz dynamics as the dynamics model is applied to locomotion tasks that have discrete switching points at contact. If the theorems have no relevance these should be removed from the paper. It would be beneficial if the paper discussed the benefits and disadvantages of the model instead.     \n\nThe experimental section also does not evaluate the proposed model architecture, it just wraps everything together in the model-based RL paradigm, which evaluates the proposed dynamics model architecture impossible. Therefore, one cannot make any claims about the dynamics models when comparing the learning curves. The advantage of the NODA model is also not significant. One would need to perform different experiments to isolate the impact of the model, which goes beyond comparing learning curves as these include many different aspects. The authors would need to provide a quantitative and especially QUALITATIVE evaluation of the learned dynamics model. Therefore, the authors need to provide video renderings of the model predictions. Furthermore, ablation studies are needed that look at when these models will break? How many steps can they predict? \n\nMinor details:\n- Eq 2 should be one line\n- The description of the Noda + MFRL (which btw. is not MFRL afterwards) is a bit unclear. From the main paper, I cannot understand how the imagined data is used? Is the policy only trained on imagined data like dreamer or more in the direction of MBPO where true and imagined data is mixed.\n- what is the real pendulum? Where does the data come from?\n- The submitted code does not contain the data, e.g., learned models & datasets\n", "summary_of_the_paper": "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. \n\nIn 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.", "summary_of_the_review": "In summary, the paper should be rejected as the contribution on physics-inspired networks is marginal and the claims of the paper are not backed-up by experiments. Just adding a model and comparing the learning curves is not sufficient to isolate the impact of the proposed model learning approach.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "iaXgGFWxkA", "reviewer": "Reviewer_cybw", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of modeling the Hamiltonian in a latent is appealing and the best way of doing this is still largely unsolved in the community. This paper suggests a way to do this that augments existing RL methods (Fig 3/Table 1). My main concern with the paper is that it's difficult to interpret where the main improvement is coming from. The method (Alg 1) uses the Hamiltonian to augment data that an otherwise model-free agent is trained on, but the paper omits significantly related works that report similar results, such as [MBPO](https://arxiv.org/abs/1906.08253). Even without this, the improved sample efficiency this method makes at the beginning of training is marginal (the first 200k timesteps in Table 1/Figure 4), and no results are reported to the asymptotic performance. I think this is an important comparison to make, as if the Hamiltonian is truly modeling the system better than the existing dynamics models (e.g., Dreamer), then I would expect the asymptotic performance to be better as well.\n\nSection 4.1 also talks about Lipschitz models and continuity. It's not clear to me that the MuJoCo environments are reasonably continuous, or are closed/conservative systems that can be modeled with a Hamiltonian, as there is contact and friction between the agents and the floor.\n\nFor one last minor comment: Definition 3 cites Arjovsky et al. (2017) for the definition of the Wasserstein metric, but this metric dates back much further in the optimal transport community. Villani is a more standard reference.", "summary_of_the_paper": "This paper considers Hamiltonian models for reinforcement learning, and suggests to use it in Alg 1 to augment the data a policy is trained on. Fig 2 shows that this approach can converge faster than Hamiltonian Neural Networks on pendulum tasks, and Fig 3/Table 1 compares to standard methods on data-limited MuJoCo tasks.", "summary_of_the_review": "Modeling the Hamiltonian is a promising direction but I do not feel a proper comparison to related approaches has been made here, such as to MBPO or to Dreamer's asymptotic performance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers raised concerns and the authors have not provided a response. All reviewers concur that this paper should be rejected at this time, and I agree.
Decision: Reject