Neural ODE and SDE Models for Adaptation and Planning in Model-Based Reinforcement Learning
Authors: Chao Han
Stefanos Ioannou
Luca Manneschi
T.J. Hayward
Michael Mangan
Aditya Gilra
Eleni Vasilaki
Venue: Accepted by TMLR ()

Abstract: We investigate neural ordinary and stochastic differential equations (neural ODEs and SDEs) to model stochastic dynamics in fully and partially observed environments within a model-based reinforcement learning (RL) framework. Through a sequence of simulations, we show that neural SDEs more effectively capture transition dynamics’ inherent stochasticity, enabling high-performing policies with improved sample efficiency in challenging scenarios. We leverage neural ODEs and SDEs for efficient policy adaptation to changes in environment dynamics via inverse models, requiring only limited interactions with the new environment. To address partial observability, we introduce a latent SDE model that combines an ODE and a GAN-trained stochastic component in latent space. Policies derived from this model offer a strong baseline, outperforming or matching general model-based and model-free approaches across stochastic continuous-control benchmarks. This work illustrates the applicability of action-conditional latent SDEs for RL planning in environments with stochastic transitions. Our code is available at: https://github.com/ChaoHan-UoS/NeuralRL.

Published in Transactions on Machine Learning Research (10/2025)
Neural ODE and SDE Models for Adaptation and Planning in Model-Based Reinforcement Learning
Chao Han
School of Computer Science, The University of Sheffield, UK
Present address: Cancer Research UK National Biomarker Centre, The University of Manchester, UK
Stefanos Toannou'
si3860com.ac.
School of Computer Science, The Universily of Sheffield, UK
Luca Manneschi
Lmannestii@sheffiilld.cc uk
School of Computer Science, The University of Sheffield, UK
T.J. Hayward
t.haguand@shcffiildd.co uk
School of Chemical, Materials and Biological Engineering, The University of Sheffield, UK
Michael Mangan
m.manggan@shffiilldd.ccc
School of Computer Science, The University of Sheffield, UK
Aditya Gilra
adityo.gilmm@uun.ccc
Machine Learning group, Centrum Wiskunde & Informatica, Amsterdom, Netherlands
School of Computer Science, The University of Sheffield, UK
Present address:
Institute for Bcological Economics, Vienna University of Economics and Business, Austria
Eleni Vasilaki
e. Lwasslak@@shffielld.cccu
School of Computer Science, The University of Sheffield, UK
Reviewed on OpenReview: https://oopntei openrevee.. forum? id= T60rPlyPV4
'Equal Contribution.
Abstract
We investigate neural ordinary and stochastic differential equations (neural ODEs and SDEs) to model stochastic dynamics in fully and partially observed environments within a model- based reinforcement learning (RL) framework. Through sequence of simulations, we show that neural SDEs more effectively capture transition dynamics' inherent stochasticity, en- abling begformrforming policies with improved sample efficiency in challenging scenarios. We leverage neural ODEs and SDEs for efficient policy adaptation changes in environment dy- namics via inverse models, requiring only limited interactions with the new environment. To address partial observability. we introduce a latent SDE model that combines an ODE and GAN-trained stochastic component in latent space. Policics derived from this model offer a strong baseline, outperforming or matching general model-based and model-free approaches across stochastic continuous-control benchmarks. This work illustrates the applicability of action-conditional latent SDEs for RL planning in environments with stochastic transitions. Our code is available at: https://gii
1
Introduction
In recent years, the family of neural differential equations (neural DEs) (Chen et al. 2018 Rubanova al. 2019 Li et al. 2020: Kidger al. 2020 2021 have emerged powerful framework for modelling dynamical

---

Published in Transactions on Machine Learning Research (10/2025)
systems. The general idea of these models is to use neural networks to parameterise the derivatives of the system's is dynamics, while the state evolution is computed by a numerical differential equation solver. Such a way of decoupling the modelling of dynamics from the discretisation scheme leads to models with increased fidelity for capturing complex, contimuous-timg transitions. As a result, ncural DEs are especially well-suited for learning and representing transition dynamics in reinforcement learning (RI.).
Recent work has demonstrated that integrating neural ordinary differential equations (neural ODEs) (Chen et al. 2018) and stochastic differential equations (neural SDEs) Li et al. (2020); Kidger et al. (2021) into RL frameworks can improve the modelling of continuous- time dynamics. In partially observable Markov decision processes (POMDPs) and model-free RL, recurrent neural ODEs have shown robustness to irregularly sam- pled observations, due to their capacity to model non-uniform time series (Zhao et al. 2023). Latent neural ODEs have also been employed as data-driven dynamics models in model-based RL, providing higher sam- ple efficiency compared to model-free basclines (Du et al. 2020 Additionally, incorporating neural ODEs with control-theoretic constraints has supported the development of safe and stable RL in contimous--ime domains (Zhao et al. 2025
Neural SDEs extend the neural ODE framework by jointly modelling the deterministic and stochastic com- ponents of system dynamics, allowing for uncertainty-awaro policy learning. Recent studies show that offline model-based RL with neural SDEs, particularly when incorporating physics priors, can outperform state-of- the-art algorithms on low-quality datasets (Koprulu et al. 2025). , Plasiic- neural SDEs, which explicitly represent model uncertainty, have also been shown to support generalisation beyond the training distribution (Djeumou et al. 2023). While both these neural SDE frameworks could, in principle, be applied to POMDPs, practical demonstrations to date have focused on environments with full observability.
Despile recent advances, the application of latent neural SDEs to partially observed, stochastic RL environ- ments remains underexplored. Unified frameworks that combine controlled latent neural ODEs and neural SDEs in both fully and partially observed environments are lacking, and the practical benefits of modelling stochasticity in latent dynamics for planning have not been systematically assessed. Additionally, there has been limited exploration of neural ODEs and neural SDEs in settings that require adaptation to changes in the environment.
In this work, we first demonstrate sample-efficient policy adaptation to changes in environment configuration by employing an inverse dynamics approach (Christiano et al. 2016) based on neural ODF and SDE transition models. We then introduce a unified latent ODE/SDE framework for model-based RL, motivated by the need to handle partial observability in complex, stochastic environments. The framework uses a two-phase procedure: the mean latent dynamics are learned using a latent ODE, and stochasticity is handled via a GAN-trained latent SDE. In experiments on stochastic continuous control tasks, we show that model-basod RL with latent neural SDEs improves sample efficiency compared to both ODE-based models and model- free SAC baseline in the most challenging fully partially observed environment included in our study.
To our knowledge, there has been no prior empirical validation of action-conditioned latent SDEs for planning in partially observed stochastic control tasks. Our results highlight the potential benefit of latent neural SDEs for flexible, noise planning in stochastic domains.
2 Background
Markov decision process (MDP) and partially observable MDP (POMDP). An MDP (Cassandra) ct al. 1994) defined by tuple AS A.T,R), RR), where and are sets of states and actions respectively, T x A + A(S) is the probabilistic transition function (dynamics), with 1( a.) representing the probability of transitioning into $111 from st under a,, R : S x A x S + R is the deterministic reward function. The initial state follows certain distribution Po; the horizon is T E N+ (+00)) and the discount factor is E [0, 1], with 1 for finite horizon. When the state s, is not fully observable, MDP tuples are extended to POMDP (Cassandra et al., 1994) AS O,T, O, R) by adding a set of observations O and an emission function O: S:S- + A(O), which probabilistically maps a state 8t to an observation Oz.
Specifically, we consider POMDPs with two sources of uncertainty in the transition dynamics of observed states T(o+11 or, at): (i) aleatoric stochasticity in the underlying MDP transition T(s++1 8t; 0r) that
2

---

Published in Transactions on Machine Learning Research (10/2025)
produces the observations, and (fi) epistemic uncertainty from partial observability, which often be reduced by conditioning decouple it from aleatoric uncertainty in observation emission, we assume deterministic emission functions O,i., i.c., O2 = ((s)). Nonetheless, our approach readily generalizes to noisy observation emissions by replacing the deterministic observation model with a probabilistic neural network that outputs the statistics of the predicted observation distribution.
history of observations and actions. To focus on transition stochasticity and
In such POMDPs, introduce latent representation 2, that summarizes the observed history up to t {d0:--1, 00.c) as proxy of the state (Du et al.) 2020 Ni et al. 2024). The transition function of the latent 24 is (2 [2, Qz, 4t) with the observation or deterministically emitted from 2, via = o(2)). We are interested in learning the encoder mapping the history to the latent variable, the latent transition function and the decoder mapping the latent back to the observation, which together enable the prediction of the next observation the POMDP.
Neural Ordinary Differential Equation (Neural ODE). Neural differential equations are a family of differential equations whose rate functions approximated by learnable neural networks {Kidger 2022 A classic example is neural ODE (Chen et al. 2018)
ds; dt
- So(s,,t, where 8to = 80
(1)
where St denotes the state at time t as the solution of the ODE initial-value problem (IVP). The initial state at to sg. The rate function fo is mostly parameterized by a multi-layer perceptron (MLP) with learnable parameters 0. The general regime of neural ODE allows any off-the-shelf numerical integrator to solve for the above ODE IVP. More concretely:
8:1 ODEsolvo(fo(S,,t, to iit)), i= O,...,TT
(2)
where neural ODEsolve refers to the ODE numerical integrator used. 8:41 is a step forward of the neural ODEsolve from & with a pre-defined, fixed step size At.
A latent ODE (Rubanova et al. 2019) utilizes the neural ODE to generate latent trajectories in a variational
autoencoder (VAE) manner (Kingma 2013 which can be formulated
2g qs Oo:)) 2++1
ODEsolve(fo(2, to + iAt)), à = (8()), i= 0, , T-1
(3)
Here, an (RNN) encoder parameterized by Φ maps observations ourr to a distribution of the initial latent state A decoder parameterized by 0 reconstructs observations oo:r via a latent-variable ODE determined by sampled zn and an emission function (-)-.
Neural Stochastic Differential Equation (Neural SDE). A neural SDE consists of a parameterized deterministic drift term and a parameterized stochastic diffusion term:
ds; fo(st, t)dt + +go(s ()dw,, where
Sto
(4)
where (1e4 is a continuous-tin stochastic process, whose initial state St, is drawn from some probability distribution Jr. (wr)20 is the Brownian motion. The paper follows the idea of training a neural SDE as a (Wasserstein) generative adversarial net (GAN) (Goodfellow et al. [0014] Arjovsky et al. 2017). where the real and generated data samples are (interpolated) infiniteddimensionall paths (Kidger et al. 2021
3 Methodology
In this section, we exlend the vanilla autonomouss neural ODE/SDE-based models to their controlled variants by incorporating actions to model the MDP state transition dynamics T [ss+11|s; at) and POMDP latent transition T(2++112,, a,, oz). We also describe how train a policy for planning based on the learned transition model, as well as how adapt the model and policy to similar environment without retraining from scratch.
3

---

Published in Transactions on Machine Learning Research (10/2025)
3.1 Dynamics model learning
We use the termss neural ODE/SDE and latent ODE to overload the names of original autonomous models and here denote their controlled variants. In similar vein to latent neural ODEs, we also propose the latent SDE by employing the noural SDE in the encoded latent space, which, to the best of our knowledge, has not been proposed in the literature.
Neural ODE. Let the MDP transition dynamics of deterministic environment be defined as an au- tonomous neural ODE (L is not explicitly given an argument to fa):
8 ODEsolve(So(Ss, a:)), i= 0, -
5
We optimize the parameters 0 by minimizing the following mean squared error (MSE) between the predicted So:7 and observed state trajectories $o:7:
min (N:
Σ 88 8
(6)
where real state-action trajectories are sampled from a replay buffer Dm
Similar to Du et al. (2000) we modify the vanilla latent ODE in Eq. latent transition by taking actions a, and (predicted) observations o into account for the underlying evolution of RNN hidden states h4 in the encoder and ODE latent states 24 in the decoder. Specifically, the adjusted latent ODE consists of the following RNN encoder parameterized by d:
Latent ODE.
g to model the POMDP
h = R.NCCell (hi, ai, ōi), à 00))), 0
[Map 0z = MLP (hr). 20(2olou:: ao: = N(oo Mzo diag(o?
and neural ODE decoder parameterized by
& MLPP (z,, ai o.), 241 (DDEsolve(fo()) à; 00(22) po Zi, dg, o.) (0011:00(2441)) T), i= . T
(8)
where o0(-) denotes the parameterized emission function, 2; has the same dimension as 2; and represents the transformed latent state at which the derivative latent ODE fo is evaluated. o, could be either the real observation o or predicted observation ô; during training. while always setting o = ô, during inference. In this paper, models are typically trained using teacher forcing, i.e., o, 0,, this has been shown to be more sample-efficient and effective for training.
The above latent ODE is trained end-to-end by maximizing the evidence lower bound (ELBO) over obser- vation trajectories o
max F -o:T--)---
E Ezo (soloor -a:--1)
Σ log po
0,9
DKL (go(2olon.T. (1) p(2o))
(9)
Here, the joint generative distribution is decomposed as Po loorr, 40:r-1) = -IT po a; 6.) and p(z)) = 0, /) is the standard Gaussian prior.
We employ a GAN-based SDE to generate synthetic trajectorics that closely match those sampled from a stochastic MDP transition. More precisely, the adapted neural SDE trained as Wasserstein GAN (WGAN) consists of neural SDE generator parameterized by 0:
Neural SDE.
8:41 SDEsolve(fa a.), (3, ai), Aw;), 5; MLLP((8) i= 0....., 1
(10)
4

---

Published in Transactions on Machine Learning Research (10/2025)
and an MLP critic (discriminator) parameterized by th:
y MLP & (Forr,
(11)
Here, the projection from & to 8; allows for more flexible generated states. Different from Kidger et al. (2021) that uses neural controlled differential equation (CDE) a discriminator, we simply use an MLP that takes flattened state-action trajectories input. This change makes optimization much easier, but also limits the discriminator's capacity for processing varying-length trajectories.
Let Go (Worr, -1) + so:T denote the overall map of the generator from paths of noises and actions to that of synthetic observations, and Da (S0,T, 00:7-1) + y denote the overall map of the critic from paths of real/syuthetic observations and actions to a scalar score. The training dynamics of the WGAN are formulated follows:
min max E(sorr: D)-D, (So.. ao:r-1)] - E warrPu DD (Go(Wo:r; Go.r-1), corr-1)] 0
(12)
|DollaSi
where |/ooll.. i. < represent the Lipschitz continuity constraiit enforced the critic function Dø We use the gradient penalty (Gulrajani et al. 2017) to achieve the Lipschitz constraint.
a
RNN Encoder
ODE Decoder
ò
AT
Aaa aa
ODE So,ve(Ja,.)
ODE Solve(/o.)
Go
Fy
by
b
MLP Critic
óp
de
0
SDE
SDE
Solve(Ja )
Solvel )
ao
6n
Pr
SDE Generator
Figure 1: Computational graph of the latent SDE. (a) The RNN encoder and ODE decoder from the latent ODE. The encoder is only utilized during the model training. In the inference (gencration) stage, the initial latent #n is sampled from the standard Gaussian prior distribution instead. Note that the o; could denote either the real observation o, or predicted one ô, during training, depending on whether teacher forcing strategy is used, yet it always set to be o; = ô; during inference. (b) The SDE generator and MLP critic of the latent SDE. In the SDE solver, we use the learnt drift function fø from the ODE decoder in (a) and train the diffusion function ga, to capture the full stochasticity in the latent space. In both panels, the deterministic variables are represented diamonds while the stochastic variables are depicted as circles.
weakness of latent ODE that it can only model deterministic MDP transitions (or the mean dynamics of stochastic MDP transitions) of states underlying a POMDP. To learn the full stochastic MDP state transition, employed phased training framework. The latent. ODE (Eqs. and S) is first
Latent SDE.
5

---

Published in Transactions on Machine Learning Research (10/2025)
employed to capture the mean transition dynamics in the latent space. Specifically, the encoder-decoder pair of the latent ODE is described as follows:
20 us (2olooT,, 40:T-1), 2141 ODEsolve(fe, (zi, ai o:)), 6, 08, (zi), i= i (13)
where the recurrent encoder and the ODE decoder respectively parameterized by and 0ur which are trained in an end-to-end VAE-style setup (Eq. 9
We then freeze the learned ODE decoder and deploy it as the drift function of an SDE in the latent space. The diffusion function of the SDE and the emission function are further employed to model the variance of the latent transition dynamics and the latent-to-observation projection, respectively. Putting together, we have the following equations:
20 ~ p(zo), Z141 SDEssolve(a, (Zi, at, 0z), 90, (zi, ai oi), Aw;) o, = =0, (2i), i= T (14)
a -1)
where is sampled from the Gaussian prior p(=o) used in the VAE training. Sa, and go, are the drift and diffusion function of the SDE, where 02 denotes the detached optimal 0g from the first phase. The diftusion ge, and emission function 08, of the SDE generator are trained in pairs with the MLP critic parameterized by 2o in GAN-style sctup (Eq. 12
We term this model framework that combines latent ODE and neural SDE latent SDE, and depict its computation graph in Fig. !
3.2 Planning and policy learning with learned dynamics
When addressing sample-intensive tasks (such as Mujoco), model-free RL method is often trained off-policy using near-optimal transitions collected by a plamning strategy driven by learned, task-specific dynamics, thereby reducing sample complexity. In this paper, following Du et al. (2020). we adopt framework that combines model predictive control (MPC) (Nagabandi et al. 2017 Chua et al. 2018), which searches for the most promising exploratory action, with soft actor-critic (SAC) (Haarnoja et al. 2018 which learns the optimal policy (actor) and value function (critic).
At each time step t, MPC simulates K trajectories over a planning horizon H, where actions and next states are sampled from the actor and the transition model, respectively. To prevent the model rollouts from becoming shortsighted, the critic estimates the cumulative reward beyond the planning horizon. This estimate, together with the cumulative reward from the MPC rollouts, forms the return for the state at time t. Subsequently, the first action of the trajectory that achieves the highest return is selected to interact with the environment.
This approach interleaves training of the transition model (when needed), data collection, and optimization of both the policy and the value function, thereby enabling policy learning from a limited number of data samples. The complete algorithm for POMDPs, adapted from Du et al. (2020), is provided in A Algorithm with the MPC planning component highlighted in blue.
3.3 Model and Policy Adaptation
In this subsection, we introduce an efficient way to adapt the learned transition model and policy of source environment to a target environment, at the cost of minimal amount of data from the target environment. The source environment typically refers to simulated or controlled setting while the target environment denotes the real-world or modified simulated scenario, usually presenting novel dynamics or disturbances. Inspired by prior work by Christiano et al. (2016), we employ an adaptation architecture based on an inverse dynamics model, which allows us to leverage the high-level characteristies of the source policy while adapting to the specifics of the target domain.
In what follows, we use the superscripts "src" and "tge" to denote variables and models associated with the source and target environments, respectively. We assume that both environments characterized by MDP
6

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm 1: MPC with SAC for POMDPs.
Input
The empty replay buffer for model learning Dn and for policy learning Dp, the initial latent transition model T ae, ot), the reward function R(r, at, 0++1), the initial actor (0t,a)) [a)) the two critics Q (ot, ae) and their target networks Q (ot, az), i == 1,2.
Params: The planning horizon H, the search population K, the number of environment steps M and the number of epochs E.
Output: Learned Pa, Tw and min;-1,2 Q5.
Collect trajectories using a uniformly distributed policy and save them into pmi for i= to E do
Update To(=++1 zt; at, 0z) using data from Dm, detailed in Section 3.1 Observe the initial state on and initialize the latent state 20
for j = 1 to M do
6
for k = 1o to K do
(k) (k)
7
20 < 1-11
for h = 1 to H do
Select the action ah-1 ~ (-10% (k)
9
10
Transit to next latent 2h - To 1/-
(h)
(-%-
h1111 m11) and emit the observation
(k)
(2)
0f (k)) oz)h
Calculate the reward (()) h-1 R(o
(8) (k)
h-1:
12
end
Sclect the action a (-10 )
13
14
end
15
Solect the index of the sequence with the maximum return among the K sequences:
k- + arg max; DH ) h-1 + + min; Q (o "H
(k)
Seleet the first action of the best sequence 0j-1 f- d0 (k*).
16
17
Execute aj-i and observe the next state 0j;
18
Transit to the next latent 2j - To(-1 -.
19
Calculate the reward rj-1 + R(0j-1,, if Og is not the terminal state then
20
21
Store (0j-1,
in Dps
22
else
I, Store (oj-1, ( NULL) in D,, resct the environment; Observe the initial state on and initialize the latent state 20 continue;
24
25
28
end
27
Update Q 2% and using data from Dpi
28
Update target networks 07 07
29
end
30
Store the trajectories collected using the current actor in Dm,
end
transitions and share the same actuated degrees of froedom. As illustrated in Fig. 2 given a current target state st we first compute the corresponding source action ar = n policy nree Next, we estimate the subsequent. source state gare ( !ge ,aor) via the parameterized source transition model The conditioned on the current state 1ge and action arre, Finally, the current target action tge = (s agr gare is computed using the inverse dynamics model I,, which maps the current
" 120
(s
using the predefined source
target state st and the desired next state s to a target action that drives the next target state towards
the desired state 1441 close as possible.

---

Published in Transactions on Machine Learning Research (10/2025)
a
Target
Luv (%) Age
sate
Toge
Agu
I
Age
Target
Env
5gc
TsiC
Tore
111
π Ste
π
Source
Target Policy
Figure 2: Overview the architecture for training and deploying a policy adapted from the source domain in the target domain. (a) Given an off off-the-shelf source policy œa and transition model T we use an inverse dynamics model I to generate the target action ige leading to the next target state slge Jn is atge trained via minimizing the mismatch between the predicted next state $(11 and the desired state $t11 (b) The trained inverse model, combined with the source policy and transition, is deployed in the real target environment as an adapted target policy.
In order to minimize the mismatch between the next target state led by the target action and the anticipated next state, we need a differentiable transition model for the real transition dynamics the target domain. Specifically, we adapt the deterministic component of the source transition model parameterized by 0 that of the target transition model parameterized by the augmented parameter set 0+ I aug- which includes additional parameters Gang that capture the variation in the deterministic component of transition dynamics between the source and target environments. We train only the additional parameters 0aug of the augmented neural network that models the deterministic component of the target transition, using a small dataset collected under a random policy from the target environment and stored in the buffer Dm Similar to the loss for a deterministic Neural ODE defined in Eq. we minimize the following MSE loss between the predicted and real larget state trajectories:
Data collection/trainnic
Ago Age Ago -1) = E 1 - ) (cm (Oaug; D:7 -))-1)) L (Oang; 80:r
(15)
where
Ege 6m (Oangi $0.7 dD:T
tge
Σ
||5t
stge
8t 18
Σ
1
(8c- (86-1: a1-1))1)
Ige
tge
st t
(16)
Empirically only require small amount of data from the target domain to train the target transition model augmented from the source transition model, compared with training a target transition model from scratch. For example, in the cartpole adaptation task described in section 4.2 training a neural ODE from scratch requires 500k transitions, whereas fine-tuning only the final layer of a source-trained model achieves similar validation loss with just 2k target transitions. We freeze the target transition model once it sufficiently accurate, and then use for training the inverse dynamics model, which interleaved with data collection. Specifically, we repeat the following data collection/ tirooontttiinngggggg loop: during the data collection phase, given current target state, obtain the desired next state, and then use the learned-so-far inverse dynamics model to produce the target action, which interacts with the real target environment Icading to the next target state. We repeat the tgg tgo sare 81+1 " ... into the buffer
the preliminary source policy and transition model to
larget environment steps and save the sequence along the sleps {... s Din. Target states along the trajectory collected in such an on-the-fly fashion are near the optimal target trajectory, which will lead much faster convergence of the inverse dynamics model compared with training data collected by random policy. During the training phase, we optimize the following mean squared error (MSE) loss between the predicted next. target state and the desired next state, using mini-batch of
8

---

Published in Transactions on Machine Learning Research (10/2025)
real-desired transition pairs sampled from Dine:
Linv (m 8t (8+1) = 18 (ogo 177 1-D Knv (m (n: s,00 ((41)]
(17)
where
Ein Pmv((1; s (841)
||ste1
(s , In s44 8
(18)
It worth noting that the above MSE loss not only applies to the next states generated by deterministic transition dynamics but also to the means of the next states when the transition dynamics are stochastic. The averaged dynamics of the stochastic transition described by an SDE (i.e., the drift term of the SDE) be captured by the Neural ODE model. Therefore we are only interested to adapt the model and policy when the mean-field (drift) dynamics vary between the source and target domain as the diffusion dynamics capture the alcatoric uncertainty, which is irreducible (Chua et al. 2018 Han et al 2024).
Deployment. As illustrated in Fig. 23, can the learned inverse dynamics model a target policy, which is computed via adaptation from the source policy follows:
π 120 (g0) = (s 1go
(s) n (s
ige (F))))
(19)
4 Experiments
We empirically evaluate our ODE- and SDE-based models in stochastic environments of increasing com- plexity. In section 4.1 we show that the neural SDE captures stochastic transitions more accurately than the neural ODE, which only learns the deterministic drift corresponding to the averaged dynamics. In sec- tion 4.2 we demonstrate that neural ODE/SDE models enable policy adaptation to environmental changes via the inverse model with fewer data. Finally, in section 4.3 show that SDE-based policies generally achieve higher asymptotic rewards and faster convergence than other model- based and model-free baselines, under stochastic environments with full and partial observability.
Environments. For evaluation of our methods, modify the standard deterministic OpenAT Gym envi- ronments (Towers et al. 2024]: ] cartpole from the classic control task, swimmer, hopper and walker2d from the Mujoco locomotion task into stochastic versions by adding noises to their MDP transition dynamics. In each modified environment, both the action and observation spaces are continuous. When fully observable, the observation space includes the positions (or angles) and velocitios (or angular velocitios) of every degree of freedom.
Stochastic cartpole. The goal is to balance a pole on a moving cart by applying forces on the cart at each step. We convert the original discrete actions to continuous ones and apply independent and identically distributed (i.id.)) standard Gaussian force to the cart at each step, which can be formulated as an SDE of cart velocity.
Stochastic swimmer. observation space and actuators. We introduce i.idd Gaussian noise sampled from Ao0, 500?) the stiffnnss parameter of the first actuated joint of the Swimmer.
3-link swimmer is propelled forward in a fluid, with a 10-dimensional
Stochastic hopper. A single-legged robot is made to hop forward as far as possible, with a 12- dimensional observation space and actuators. We apply stochastic winds to the hopper at each step. The magnitude of the wind parameter is i.i.d. sampled from (0)
Stochastic walker2d. A bipedal robot, characterized by an 18-dimensional observation space and 6 actuators, is tasked with walking forward while subjected to stochastic winds. The wind's magnitude is sampled i.i.d. at each step from (0,5522.
9

---

Published in Transactions on Machine Learning Research (10/2025)
We further hide the position and velocity features from the observation space of the Mujoco to evaluate learning in POMDPs. It is worth noting that the POMDPs considered here satisfy our assumptions on the aleatoric and epistemic nature of the uncertainties in the POMDP transition dynamies, as described in [ Specifically, in the partially observable stochastic swimmer, we mask the positions and angle of the front tip (stochastic, swimmer (no position)), and its positional and angular velocities (stochastic swimmer (no velocity)) from the observation space (corresponding to the first three and the 5th to 8th dimensions of the observation respectively). In the partially observable stochastic hopper, we mask the positions of the torso (stochastic hopper (no position)), and its angular velocity (stochastic hopper (no velocity)), which correspond to the first two and the 8th dimensions of the observation, respectively. Similarly, for the partially observable stochastic walker2d, we hide the torso's positional information in the first two dimensions (stochastic walker2d (no position)) and its angular velocity in the 11th dimension of the observation space (stochastic walker2d (no velocity)).
Baselines. We compare ODE-based and SDE-based models in stochastic environments with respect to both dynamics learning and model-based policy optimization. In the simpler CartPole task, we evaluate neural ODE (N-ODE) and neural SDE (N-SDE) as proxies for the true stochastic transition dynamics, testing their effectiveness for policy learning and adaptation. Tn the more challenging MuJoCo tasks, we additionally consider latent ODE (L-ODE) and latent SDE (L-SDE) models for planning and policy opti- mization. Across all environments, we use Soft Actor-Crithe (SAC) (Haarnoja et al., 2018 as model-free baseline, training Gaussian policies but evaluating their deterministic mean. As the model-based counter- part, we adopt Model-Based Policy Optimization (MBPO) (Janner et al., 2019), which improves sample efficiency by training on short-horizon imaginary rollouts generated from an ensemble of dynamics models, where each model is a probabilistic neural network parameterizing a Gaussian distribution (Chua et al. 2018 Janner et al. 2019).
It is worth noting that, although for computational efficiency we set the integrator step size equal to the simulator time step, the TRE//DE-aased models can in principle achieve finer temporal resolution by using smaller integrator steps. This added error control in time-series modeling is not available to general model- based RL basclines such as MBPO (Janner et al. 2019), which operate only on discrete-time MDPs with a fixed step size. A detailed experimental setup is provided in the Appendix C
4.1 Neural ODE/SDE modelling transition dynamics
Here we demonstrate the capacity of the neural ODE and SDE in mimicking the stochastic MDP transitions of the stochastic cartpole environment. We use the learned transition models as a proxy for the real transition dynamics. A model-free agent can therefore be trained in the approximated transition dynamics without interacting with the real environment, To achieve the oracle performance given by the agent trained in the real environment, the transition model should be accurate enough to recover the distribution of the possible next states.
Neural ODEs model the mean of stochastic transitions while Neural SDEs capture the full stochastic dynamics. Fig. 3 depicts the evolution of marginal distributions of cart velocity across time predicted by the neural ODE and SDE, in comparison with the real marginal distribution. The ODE-based distributions are peakier at the mean than the real distribution and fail to recover values away from the mean. On the other hand, the SDE--based distributions mostly cover the real distributions. In addition, Fig. 3p shows the sample paths from the ODE and SDE-based distributions against the paths from the real distributions. Once again the ODE-based paths capture only the averaged lendency of the real paths, while SDE-based paths show better agreement with the real ones. For other features in the observation space of the stochastic cartpole, since their transition dynamics are deterministic, the neural SDE learns almost the same transition dynamics as the neural ODE (see Fig. and 7 for distribution and path matching respectively in the Appendix A
We evaluate the performance of agents trained on modeled transition dynamics in the real environment (Fig. B-) Performance of SDE-based policy is much closer to the near-optimal performance of the model-free oracle policy than the ODe-based policy, in terms of the resemblance of their distribution of returns to the
10

---

Published in Transactions on Machine Learning Research (10/2025)
3.0
N-00t-pred
Densty
0.6
0.2
Velocity
N-SDE-one
0.0
Denstry
Velocity
b
C
Cart Walocky
Cart
200
100
0.00
0.05
0.10
0.35
0.20
0.25
0,30
6.35
0.00
0.05
6.30
6.35
8.20 3.25
0.30
0.35
9-00E banM)
500 catee
Figure 3: Comparison between neural ODE and SDE stochastic cartpole. (a) Predicted histograms of cart velocity at 16%, 37%, 58%, 79%, 100% of the total timesteps by the neural ODE and SDE against the real histograms. (b) 50 sample paths from the distributions predicted by neural ODE and SDE against 50 paths from the real distributions. The neural SDE learns to match the real distributions and sample paths better than the neural ODE. (c) Policies trained in the approximate transition dynamics of neural ODE (N-ODF-basod) and SDE (N-SDE-basod), as well as in real transition (model-free), are evaluated in the real environment. N-SDE-based policy achieves similar performance the model-free policy (used as the oracle here).
oracle distribution. The poor performance ODF-based policy when deployed in the actual environment is due to the low fidelity of the model used for training the policy.
4.2 Policy adaptation via inverse dynamics model
In this section, we empirically test the policy adaptation framework described in section 3.3 in both deter- ministic and stochastic cartpole environments with increasing pole length. We compare diflerent transition models: the neural ODE, the neural SDE, and the ensemble of Gaussian neural networks. The results show that adapted source policies based these transition models outperform non-adapted source policies in the target environment, while also achieving greater sample efficiency than training new target policies from scratch under limited interaction with the target environment.
Deterministic transition. We first consider the setting where the cartpole dynamies are deterministic (Fig. 4.). We compare two adapted policies against the non-adapted source policy (red curve) as the pole length increascs. The adapted policies are based on neural ODEs (blue curve) and deterministic ensemble networks using the Gaussian mean (purple curve), both modeling the transition dynamics in the source and target domains under determinism. To evaluate sample efficiency, also include policy trained
11

---

Published in Transactions on Machine Learning Research (10/2025)
b
Deterministic cartpole
Stochastic cartpole
500
500
A
420
400
Eval. Rewards
Eval. Rewards
300
300
100
Pole Lengths
Pole Lengths
N ODE adapted
N N-SDE drift adapted
Ensemble adapted
Non adapted
Scratch-tralined
Figure 4: Evaluated performance of different policies in target cartpole environments with increasing pole lengths (source environment pole length = 1.0). Shaded regions denote the standard deviation returns over runs. (a) Under deterministic source and target dynamics, for pole lengths between 1.8 and 3.2, both the addDeedd and ensemblo-adapted policics consistently outperform the original non-adapted policy, and both substantially surpass the trained- fromsscraom-cc-ccccchccc policy. The N-ODE adaated and ensemble-adapled policies show nearly identical performance. (b) Under stochastic source and target dynamics, obtained by adding zero-mean Gaussian noise to cart velocity, for pole lengths between 2.6 and 3.8, the ODE-adapted, SDE-drift SDE-drift-ddaptedd and ensemble-adapted policies achieve similarly strong performance, with the ensemble- adapted policy showing occasional drops at certain pole lengths. These are followed by the non-adapted policy, with the trained-frome wol---foo---rrtce policy performing worst.
from scratch (brown curve) for 2k iterations, using the same number of target enviromment interactions as for training the augmented transition models. The results show that the two adapted policies perform similarly, consistently outperforming the non-adapted baseline, and remain significantly better than the scratch-trained policy when the gap between source and target environments moderate. Although the scratch-trained policy is likely undertrained given the limited iterations, this comparison highlights that, with the same amount. of target data, adapted policies perform substantially better, while from-scratch policy would require far more samples to reach comparable performance. Performance of both adapted and non-adapted policics declines as pole length increases, likely because the limited action force cannot counteract the larger gravitational torque induced by a longer pole.
Stochastic transition. In the stochastic cartpole setting, both source and target transitions share the same level of stochasticity, with pole length the only varying factor (Fig. 4p). We evaluate three adapted policies: ODE-adapted, SDE-drift-adapee use, respectively, augmented neural ODEs, the augmented drift function of neural SDEs, and the mean of each augmented Gaussian network in the ensemble to model the averaged target dynamics. The ODE- and SDE- drift-adapted policies achieve nearly identical performance, indicating that in stochastic environments, the neural ODE: effectively captures the same deterministic component as the SDE drift. The ensemble-adapted policy performs comparably, though with occasional drops likely due suboptimal ensemble members. As in the deterministic case, all adapted policies outperform the non-adapted and scratch-trained baselines, though their performance also declines pole length increases.
and ensemble-adapted (blue, green, and purple curves). These
Interestingly, the source policy trained in a stochastic environment (red curve in Fig. is more robust to environmental changes than the one trained in a deterministic environment. (red curve in Fig. 44). This ro- bustness likely arises because environmental stochasticity encourages broader exploration the state action space to counteract uncertainty, yielding a more generalizable policy. Consequently, adapted target policies benefit more from stronger source policies, as reflected by the larger performance gains over the adn-add baseline in Fig. compared to Fig. 4.
12

---

Published in Transactions on Machine Learning Research (10/2025)
4.3 Model-based policy learning
For more complex stochastic MuJoCo environments, adopt framework that interleaves exploration via MPC plamning with model-free policy optimization, described in Section 3.2 We compare the performance of policies derived from MPC planning using neural ODE/SDE models and their latent variants as transition models, under both full and partial observability. Additionally, we evaluate the family of neural differential equation-basee policy optimization methods against a model- baseline (SAC) and a model- baseline (MBPO).
Stochastic swimmer
Stochastic swimmer (no position)
Stochastic swimmer (no velocity)
60
Eval. Rewards
55
Eval. Rewards
40
Eval. Rewards
40.
40
30
20
10
Environment Steps
x10
Environment Steps
x10
Environment Stops
x10
Stochastic hopper
Stochastic hopper (no position)
Stochastic hopper (no velecity)
4000
Eval. Rewards
5000
Eval. Rewards
J000
Eval. Rewards
3000
Wor
2000
2000
2000
1.000
1000
1000
Environment Steps
x10
Environment Steps
X10*
Environment Steps
X10
Stochastic
Stochastic walker2d (no position)
Stochastic walker2d velacity)
5000
5000
Eval. Rewards
4000
Rewards
4000
Eval. Rewards
3000
2000
2000
2000
1000
1000
1000
0
Environment Steps
x10
Environment Steps
Environment Steps
x10!
N-ODE
L-ODE
N-SDE
L-SDE
SAC
MBPO
Figure 5: Learning curves of model-based and model-free policies in fully and partially observable stochastic Mujuco environments. Shaded regions indicate the standard deviation of evaluated returns over 5 runs, with evaluations conducted every 5k environment steps. The policy derived from the latent SDE model consistently achieves the best asymptotic performance across all environments. SDE-based policies also exhibit greater sample efficiency than the model-free baseline in more complex environments.
SDE-based policies outperform their ODE-based counterparts. Fig. illustrates the learning curves for all baselines across different MuJoCo environments. SDE-based policies generally converge more rapidly and achieve superior solutions compared to ODE-bascd policies. The SDE-based models explicitly represent both the mean and variance of the stochastic transition dynamics, enabling more accurate pre- diction of the next (partial) observation than ODE-based models (Fig. g in Appendix B). This, in turn, allows SDE-bascd policies to better account for transition uncortainty and avoid high-risk actions. In con- trast, based models capture only the mean dynamics. As result, planning with them may deviate from the true environment trajectories, leading to suboptimal action choices and limiting exploration of new
13

---

Published in Transactions on Machine Learning Research (10/2025)
state regions; policies trained on such transitions are therefore more prone to local minima during policy optimization.
SDE-based policies achieve higher sample efficiency than model-free methods in sufficiently complex environments. In the stochastic hopper and walker2d environments (middle and bottom rows of Fig. 5 the policies with SDE-based planning demonstrate their advantage over the model-free baseline by roquiring fower interactions with the real environment (environment steps) to converge. This advantage is not observed in the simpler Swimmer environments. In particular, the model using latent SDE is the best performing in the POMDP stochastic Hopper (no position). We further observe that the next partial obscrvation predicted by the SDE transition in the latent space (latent SDE) aligns better with the real partial observations than those predicted by the SDE transition directly in the observation space (neural SDF) (see Fig. in the Appendix B) which might explain the performance advantage.
MBPO performs no better than SDE- policies under stochastic transitions and fails under partial observability. While neural differential equation models have primarily shown advantages in modeling irregularly sampled time series (Rubanova et al. 2019 Kidger et al., 2020), here we also demonstrate their effectiveness in capturing stochastic dynamics within framework, surpassing the MBPO baseline. In MBPO, transition noise can degrade model-generated samples, and the resulting errors may be exploited by the policy when trained on augmented buffers containing such samples. The slower convergence we observe in hopper is likely due to the capped, gradually increasing model horizon (see Appendix Ch: in stochastic settings, this leads to accumulated model errors, causing the model- generated rollouts to deviate from real trajectories and thereby degrading policy performance. By contrast, in the deterministic hopper, Janner al. (2019) report that such a capped horizon accelerates learning without significantly deviating the trajectories. In stochastic walker2d,
model-based planning and policy learning
single-step model horizon as in Janner et al. (019)) which avoids error accumulation from long horizons and yields performance comparable to the L-SDE baseline (left panel, bottom row of Fig. 5). However, MBPO completely fails in partially obscrvable environments (e.g,, stochastic Hopper without position or velocity; middle and right panels, middle row of Fig. 5). This failure arises in POMDP settings with early termination (such as hopper and walker2d), where the transition model cannot reliably predict termination signals from partial observations. As a result, model uncertainty is exploited by the policy, which converges to local minima corresponding to unhealthy state regions that trigger termination. In contrast, our model-based RL framework does not suffer from this issue in POMDPs, since models are used only for planning, while the policy is trained directly environment transitions with real termination signals.
Finally, we note the computational overhead. SAC consistently requires the least training time, while ODE/SDE-based methods are slower SDE models take longer than ODE models, and latent variants are slower than their state-basod counterparts. MBPO is by far the most demanding. For instance, in stochastic Walker2d, N-ODE requires about 2x the training time of SAC; L-ODE and N-SDE about 2x that of N- ODE; L-SDE about 1.5x that of N-ODE; and MBPO roughly 6x longer than L-SDE. MBPO also consumes about 3x more CPU memory than the other baselines. All experiments were run a single NVIDIA A100 GPU. This overhead arises because ODE/SDE-barr based policies train an additional transition model compared to SAC, while MBPO is even more expensive due to the need to train and maintain an ensemble of models.
5 Discussion
We systematically evaluate action-conditional neural ODE and SDE models for policy adaptation and op- timization in stochastic continmous-control environments. Policy adaptation via inverse dynamics, using neural ODE or SDE transition models, is substantially more sample-efficient than training from scratch when environment configurations change. While the original inverse dynamics work by Christiano el al. (016)) did not specify how the target transition is estimated or the training loss for the inverse model, our contribution to make this explicit: we show how neural differentiallcquation transition models can be instantiated within the inverso-dnnmmies franework and contrastod with MBPO's ensemble transition models. Empirically, neural ODF/SDE models achieve adaptation performance comparable to the ensemble
14

---

Published in Transactions on Machine Learning Research (10/2025)
model in simple environment; more importantly, even in stochastic domain where the change stems from deterministic factor, a simple neural ODE matches a neural SDE with transformed drift.
For policy optimization, SDE-based models generally outperform other model-based and model-free ap- proaches, highlighting the effectiveness of diffusion in modeling transition stochasticity. In particular, latent- SDE-based policies consistently exhibit strong sample efficiency and high asymptotic returns that are difficult. to beat, suggesting the contribution of the latent space for robust policies under both full and partial ob- servability. Moreover, using action planning to collect near-optimal transitions for policy training in our ODE/SDE-based framework provides clear advantages over the model-generated data used in MBPO, which tends to let the policy exploit model errors when transition dynamics or observations are noisy. These find- ings underscore the potential of policy planning over latent SDE for general robotic tasks subject to noisy transition dynamics and partial observability.
However, there are also situations where ODE/SDE-basce policies may not offer advantages over standard model-based or model-free baselines. For instance, in continuous-control lasks with deterministic dy dynamics very weak noise, ODE/SDE-band based policies may not improve over model-based baselines such as MBPO. In stochastic walker2d, where the transition noise is relatively weak (N 5²) wind noise), MBPO with single- step rollout horizon achieves performance comparable to the latent-SDE bascline. The slower convergence of MBPO in the stochastic Hopper arises instead from its capped and increasing rollout horizon: while this strategy improves sample efficiency in deterministic settings (Janner et al., 2019), it causes the policy to exploit model error in the stochastic case. By contrast, when the noise amplitude is large, as in the stochastic swimmer with (i(0, 5002) joint stiffness noise, MBPO even with a single-step horizon underperforms the SDE-based models.
A second limitation may arise in tasks without early termination conditions. Our results suggest that the sample-eficiency advantages of planning with ODE/SDE-basee models are more pronounced relative to the model-free SAC in environments with termination conditions, where foresight helps agents avoid irreversible outcomes. Specifically, improved sample efficiency SDE-based policies over SAC is observed only tasks with termination conditions (stochastic hopper and walker2d), but not in tasks without them (stochastic swimmer). Nonetheless, alternative explanations for these differences are possible, and further investigation is roquired before drawing conclusion.
Finally, there are task regimes where ODEeSDE- models are inherently ill-suited. ODE/SDE continuous-tine integrators struggle with bybrid or discontinuous dynamies, such as mode switches or resets, which degrade both planning and policy learning. In strongly chaotic regimes, small modeling errors can rapidly amplify under ODE/SDE rollouts, making one-step ensembles or model-free approaches more robust. Similarly, systems with actuation delays or dead zones are better described by delay differential equations or discrete-time bistory-based models (e.g., RNNs (Cho et al., 2014p or Transformers (Zerveas et al. 2021 rather than by latent Markov ODE/SDE formulations.
Acknowledgments
C. Han, S. Ioannou, A. Gilra, and E. Vasilaki acknowledge support from the CHIST-ERA project Causal Explanations in Reinforcement Learning (CausalXRL) CHHST-ERA-I9-XAAIE XAL-002), funded by the Engineer- ing and Physical Sciences Research Council (EPSRC), United Kingdom (grant EP EEP/V55520/1)) C. Han, S. Ioannou, L. Manneschi, M. Mangan, and E. Vasilaki acknowledge support from the EPSRC project Active Learning and Selective Atlention for Robust, Transparent and Efficient AI (ActiveAl) (grant EP BP/S030964/1) C. Han, L. Manneschi, T. J. Hayward, and E. Vasilaki acknowledge support from the EPSRC project Mag- netic Architectures for Reservoir Computing Hardware (MARCH) (grant EP EP/V /V006339/1).
References
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223, PMLR, 2017.
Anthony R Cassandra, Leslie Pack Kaellling, and Michael L Littman. Acting optimally in partially observ- able stochastic domains. In Acai, volume 94, pp. 1023-1028, 1994.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Kyunghyun Cho, Bart Van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rmn encoder-decoder for statisti- cal machine translation. arXiv preprint arXiw:1406.1000 1078, 2014.
Paul Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojeiech Zaremba. Transfer from simulation real world through learning deep inverse dynamics model. arxiv preprint arxiv:1610. 1ad:11600035188 2016.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
Achille Djeumou, Lingfan Huang, Alex Buisson, et al. Plysios-constrained neural SDEs for efficient RL & control. In Conference on Robot Learning (CoRL), 2023.
Jianzhun Du, Joseph Futoma, and Finale Doshi-Velez. Model-based reinforcement learning for scmi-markov decision processes with neural odes. Advances in Neural Information Processing Systems, 18:19805-19881 2020.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Thomas Haarnoja, Aurick Zhou, Picter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum
entropy deep reinforcement learning with a stochastic actor. arXiv [cs.LG], January 2018. URL http:
//aaxiv.rrg/abs//// wttg/a/////..//.00 01290
Chao Han, Debabrota Basu, Michael Mangan, Fleni Vasilaki, and Aditya Gilra. Dynamical-vae-bsse hind- sight to learn the causal dynamics of factored-pomdps. arXiv preprint arXiv:2411.07882 2024.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model- based policy optimization. Advances in neural information processing systems, 32, 2019.
Patrick Kidger. On neural diffcrential equations. arXiv preprint arXiv:2202.022 02435, 2022.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time scries. Advances in neural information processing systems, 33:6696-6707, 2020.
Patrick Kidger, James Foster, Xucchen Li, Harald Oberhauser, and Terry Lyons Neural SDEs as infinite- dimensional GANs. arXiv Lcs.LG), February 2021. URL http:://arriv..o///bb////0///// 03657
Diederik P Kingma. Auto-cncoding variational bayes. arXiv preprint 1rXih:1812.611 2013.
Mchmet Koprulu, Aditya Dhariwal, Surya Ganguli, and Sanjeev Arora. Nuno: Ncural uncortainty-awaro neural SDE for oflline model-based RL. In Intervational Conference Learning Representations (ICLR), 2025.
Xuechen Li, Ting- Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3870 3882. PMLR, 2020.
Anusha Nagabandi, Gregory Kahn, Ronald Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-fres fine-tuning. arXiv LGs, August 2017. URL
16

---

Published in Transactions on Machine Learning Research (10/2025)
Tianwei Ni, Benjamin Eysenbach, Erfan Seyedsalchi, Michel Ma, Clement Gehring, Aditya Mahajan, and Pierre-Tuc Bacon. Bridging state and history representations: Understanding self-predictive rl. arXiv preprint arXiv:2401.0888 2024.
Yulia Rubanova, Ricky T
Chen, and D Duvenaud. Latent ODEs for irregularly-sampled time series. ArXiv, abs/ http//997.00307,, July 2019. ISSN 2331-8422. URL https://prrceediggss http://ceeeden....ooooooooooo//o now.pps..ccpppppr.
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Delou, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiu:2477.17082 2024.
George Zerveas, Srideepika Jayaraman, Dhaval Patel, Amuradha Bhamidipaty, and Carsten Eickhoff. A transformer-basee framework for multivariate time scries representation learning. In Proceedings of the 27th ACM SIGKDD conference knowledge discovery & dala mining, pp. 2114 2124, 2021.
Liqun Zhao, Keyan Miao, Hongpeng Cao, Konstantinos Gatsis, and Antonis Papachristodoullu. NLBAC: A neural 3DE-based algorithm for state-wise stable and safe reinforcement learning. Neurocomputing, 638: 130041, 2025.
Xinyu Zhao, Tianyu Ma, Ting Liu, Ziyan Wang, and Haoran Li. Ode-based recurrent reinforcement learning for partial observability. In NewrIPS, 2023.
17

---

Published in Transactions on Machine Learning Research (10/2025)
A Stochastic cartpole environment
a
T=0.06
Tw0.14
T-0.22
T=0.30
T-0.30
000 prol
1.00
0.75
0.50
Pastan
T-0,06
T-0.14
Tw0,22
T-0.30
T-0.38
3.6
Pole Angle
0.000
Reaa
Dentr
T-D 06
Tw0.14
T-0.22
T=0.39
T-0.30
prod
1.00
0.00
0.25
3.5
5.0
5.0
Carl Postion
T-0.00
T-014
Tw0.22
T-0 30
T-0 38
3.6
Priee Angle
T-o.14
Two.30
Tw0.90
0.2
-2
23
-1 3
-2
2
0
Pole Argular Vilicity
Figure 6: (a) Histograms of cart position, pole angle and pole angular velocity at 16%, 37%, 58%, 79%. 100% of the timesteps predicted by the neural ODE against the real histograms. (b) Similarly, histograms predicted by neural SDE against the real histograms.
18

---

Published in Transactions on Machine Learning Research (10/2025)
N-000 pred Real
p Real
M-032-pre
Raai
pole and monoen Weltoo
Cart Postion
0.15 Time
6.15
0.25
0.30
6,30
Time
Time
b
1: 505 prod
N 586 przd Rnal
N-53E prol
Real
Real
Cart Pook.on
0.10 0.15 0.20
0.35
0.25
0.30
0.25
6.30
Time
Timee
Tlmme
Figure 7: (a) 50 sample paths of cart position, pole angle and pole angular velocity predicted by neural ODE against 50 paths from the real distributions. (b) Similarly, paths predicted by neural SDE against paths from the real distributions. Since the transition dynamics of these features are deterministic, the neural ODE predicts almost the same distributions and their sample paths as the neural SDE does.
19

---

Published in Transactions on Machine Learning Research (10/2025)
B
Stochastic hopper environments
Stochastic hopper
Real
N-C0E-pred
-2.5
-5.0
+2.5
0.0
2.5
5.0
pred
-75
-5.
-25
0.0
2.5
5.0
75
b
Stochastic hopper
T-2
T-6
T-10
Beal
L-Dt pred
T=2 T-6
T-10
-1
Teal L-DDE
Figure 8: (a) Scatter plots of 200 PCA-embedded 2D representation of observations predicted by neural ODE and SDE, compared with that of the real observations across 20%, 60%, 100% of the timesteps in the stochastic hopper. (b) Similarly, scatter plots of 2D features predicted by latent ODE and SDE, against real ones. Predicted features by the SDE-based models demonstrate a better alignment with the real ones than those by the ODF-basee models.
20

---

Published in Transactions on Machine Learning Research (10/2025)
Stochastic hopper position)
T+2
T-6
T-10
Real
N-CDE-ored
Tw2
Tw6
T-10
Reil
L-GBE-ared
-3
Stochastic hopper (no velecity)
SDE ared
-50
-25
0.0
2.5
7.5
10.0
T-2 T-6
T-10
Roal
LSB6 pred
-2
0.0
16.0
Figure (a) Scatter plots of 200 PCA-embedded representation of partial observations predicted by neural SDE and latent SDE, compared with that of the real partial observations, across 20%, 60%, 100% of the timesteps in the stochastic hopper without position observations. (b) Similar scatter plots in the stochastic hopper without velocity observations. Features predicted by the latent SDE align more closely with the real ones than those by the neural SDE.
21

---

Published in Transactions on Machine Learning Research (10/2025)
C Experimental setup
Table 1: Hyperparameter summary of model-based and modcl-free basclines across environ- ments. Model architectures are described by network type and number of hidden layers with their sizes. For MBPO in Hopper, the model horizon follows capped linear function of the epoch [(i) min(max(1 + 0.175 (i 20), 1), 15), as in [Janner et al. 2019).
Cartpole
Swimmer
Walker2d
Hopper
Max env, sleps Model batch size Optimizer
250K
500K
Training hparams
256
128
Adam
Planning horizon
N/A
10
Search population
N/A
1000
1700
Trajectory cut length
20
10
ODE learning rate
Se-4
1e-3
SDE generator learning rate
Sc-4
Neural/ Latent ODE/SDE
SDE critic learning rate
86-5
4e-5
Encoder arch.
N/A
GRU [128]
Decoder arch.
N/A
Linear MLP
ODE (SDE drift func.) arch. SDE diffusion fune. arch.
MLP [100, 100, 100, 100]
MLP [100, 100, 100]
MLP [32, 32]
MLP [100, 100]
SDE critic arch. Solver
MLP [100, 100, 100, 100, 100
MLP [100, 64, 64
Euler
Discount factor Smoothing coef.
0.99
0.005
Learning rate
le-3
3e-4
SAC
Batch size
32
128
Temperature
0.3
0.2
0.25
Actor/Critic arch.
MLP [200, 200]
Eusemble size
7 1c-3
Model learning rate
МВРО
Model rollouts per env. step Model horizon
400
Capped linear func.
Madel arch.
MLP [200, 200, 200, 200]
Model training stopping criteria. For the Mujoco tasks, we apply early stopping to the ODE-basod models to avoid overfitting: training is stopped if the MSE reconstruction error on the validation set does not decrease for consecutive epochs. Following Du al. (2020), we use a linear decay schedule e = max(15-i, 3), where i is the epoch index in Algorithm 1. For the simpler carpole task, we train the ODE-based model for 50K iterations without early stopping, and terminate training of the Gaussian ensemble model when all members stop improving the holdout MSE, as in Janner et al. 20119 The trained ODE- model is then frozen and used as the drift component of the SDE-based models, while their additional components (the diffusion function of the SDE generator and the discriminator) are trained for 8K iterations.
GAN-based training. We mitigate the instability of GAN-based SDE training by phasing the learning of the drift and diffusion components. First, we pre-train the drift as an ODE, optimized with MSE in the state space or ELBO in the latent space. We then freeze this drift and train the diffusion with a critic in GAN- style setup. The ODEdb model is used the drift since it empirically captures the mcan dynamics of stochastic transitions, corresponding to the SDE drift function. Pretraining the drift substantially improves the stability of subsequent GAN-style training of the diffusion component. In addition, we found that ignoring actions in the critic input further stabilizes GAN training in the more complex Mujoco tasks.
22


Reviews:
[{"review_id": "uLn9B5NABX", "reviewer": "Reviewer_gsrz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n- I believe that physics constrained models are very relevant for model-based RL, and the work does a good job at introducing the use of neural SDEs in latent space, which seems like a powerful approach.\n- The proposed framework splitting the latent dynamics learning and the stochasticity modelling via a GAN seems novel and is intuitive.\n- The paper is clear, well written and provides a good cover of other ODE/SDE learning methods.\n\n## Weaknesses\n- The main weakness is the weak empirical evaluation. Given that the paper does not derive any novel theoretical results, I would expect a thorough study on more diverse environments, ablations between models, comparisons between ODE and SDE approaches and, perhaps more importantly, comparisons between ODE/SDE methods and other general model based and model free RL approaches. I think the paper could be a very good contribution if it would attempt to answer, more thoroughly, the question of \"when are ODE/SDE models necessary, and when are they not?\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present a study of neural ODE and SDE applications to model-based RL. They argue these are powerful tools to model dynamical systems since they restrict the model complexity to a physics informed class that should be able to learn better models for RL. They first demonstrate how these models result in better adaptation to changes in the environment, and then propose a framework to use neural ODE/SDE in latent space, and test these methods in a set of RL tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would like to see a deeper empirical evaluation. I would be very interested in seeing how well the methods proposed compare to other state of the art model-based RL baselines. Please note that I do not expect the proposed method to beat any benchmark at any task. I simply believe that a thorough study of when does using neural ODE/SDE methods make sense, versus when it is perhaps not necessary or not valuable, would greatly increase the quality and impact of the work. I understand the authors make an attempt to this with the environments considered, but in my opinion these are not really sufficient. \n\nI would expect neural ODE/SDE methods to be more sample efficient, but also more computationally expensive in some instances. There are problems where perhaps modelling the dynamics as a conditioned gaussian in latent space (like eg MBPO) would suffice, but I assume in other environments it would not. I also believe that the complexity constrained nature of physics based models should yield improvements across other robotics tasks, which are sometimes notoriously difficult for model-based RL. \n\nIn summary, I would like the work to provide a more thorough set of results providing intuition on the utility and necessity for ODE/SDE models in model-based RL.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EPIZdsstcq", "reviewer": "Reviewer_JiwX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n* This paper does a good job describing neural ODE and SDE, which may not be familiar to all RL readers. The authors also included important references such as Chen et al, 2018, Du et al, 2020 for readers to get familiar with the subject.\n* The experiments and adaptation strategy are described in (almost) sufficient detail to be reproduced. \n\n**Weakness**\n* Although not explicitly called out, I believe this paper consider discrete time MDP rather than continuous time MDP as in Du et al, 2020. Why is using neural ODE and SDE important in this setting, as neural ODE is usually argued to be only stronger in irregularly sampled time series?\n* In the same spirit as above, I believe a model-based RL baseline is missing from all experiments. A model-based RL baseline such as MBPO, PETS, or dreamer (for POMDP) can also be applied in the same inverse dynamics model based adaptation strategy.\n* Some parts of the methods and experiments section can be more detailed. I describe these below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes neural ODE and SDE based dynamics models for model based RL and adaptation in changing environments. Environments with both full and partial observability with removed/masked out feature dimension are considered. Adaptation is based on predicting desired next state in the original environment and training an inverse dynamics model to find an action in the target environment. Experiments are conducted in the stochastic and masked versions of cartpole, swimmer, and hopper. It was shown that in stochastic environments, neural SDE better captures environment stochasticity than neural ODE. The inverse dynamics model based adaptation strategy significantly outperforms training from scratch. The authors also conducted regular model based RL experiments, using a model free algorithm as baseline and ODE vs SDE and latent vs no-latent as ablation. It was shown that latent SDE performs best overall.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* The author should clarify that they consider discrete time MDP, if that's the case, and describe why experimenting neural ODE and SDE in this setting is important.\n* The authors should add a regular model-based RL baseline as a comparison for both planning and adaptation, or argue why such comparison is not needed.\n* \"Empirically we only require a small amount of data from the target domain to train the target transition\nmodel augmented from the source transition model, compared with training a target transition model from\nscratch.\" Do you have examples or empirical data to show this. For example, how many transitions do you need to train the target transition model from scratch?\n* \"We freeze the target transition model once it is sufficiently accurate\" How accurate is sufficiently accurate? What's your model training stopping criteria?\n* In Figure 4, when you train a policy from scratch for each pole length, how many training steps did you use? It's hard for readers to know if the from scratch policy was sufficiently trained. Generally cartpole is an easy environment and we would expect most RL algorithms to work no matter the pole length.\n* \"In contrast, ODE-based models capture only the mean dynamics, and policies planned on them may be more prone to local minima\" Could you expand on why ODE based models would be more prone to local minima? Is it the model optimization local minima or the policy optimization local minima?", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "maFKfOS410", "reviewer": "Reviewer_Pq2S", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n- The method is novel and practically relevant.\n- While the method consists of a relatively complex structure, the paper is reasonably well-written and well-explained.\n- The empirical evaluation is thorough in evaluating each of the proposed contributions.\n\n## Weaknesses\n- The evaluation focuses primarily on demonstrating that the approach works and that the claims are valid. It mostly omits a comparison with prior work, i.e., lacks a comparison on whether it is a significant contribution to the literature.\n- The model is rather complex.\n- None of the experiments are reproducible as there is zero discussion on architecture, training process, optimizer, hyperparameters, etc..\n- No public implementation is provided.\n\n### Minor\n- The GAN-based approach has the potential problem of being rather unstable, which is neither evaluated nor discussed.\n- The assumption of deterministic observation emission seems relatively unrealistic for real-world applications and also not a requirement for the proposed model. While causing another source of randomness, the approach seems like it should be stable against it.\n- There is no real ablation on the individual parts of the model.\n- Red/green color coding in Figure 3 could pose a problem for some color-blind people. Should be double-checked with somebody who has it or with an online tool.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors focus on model-based RL (MBRL) of a partially observable MDP with stochastic transitions.\nThey introduce an approach to model these dynamics in a latent-variable model by embedding\nordinary/stochastic differential equations (ODEs/SDEs) in its latent space. \nRelying on a GAN-based approach. they are able to effectively adapt learned policies to new environments.\nThe authors evaluate their contributions in a series of detailed experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "# Critical for acceptance recommendation\n- A proper discussion on the experimental setup in the appendix. \n\n# Nice to have\n- A public implementation.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)