Neural ODE and SDE Models for Adaptation and Planning in Model-Based Reinforcement Learning

Chao Han; Stefanos Ioannou; Luca Manneschi; T. J. Hayward; Michael Mangan; Aditya Gilra; Eleni Vasilaki

cs.LGMar 24, 2026

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

Organizations: School of Computer Science, The University of Sheffield, UK · Cancer Research UK National Biomarker Centre, The University of Manchester, UK · School of Chemical, Materials and Biological Engineering, The University of Sheffield, UK · Machine Learning group, Centrum Wiskunde & Informatica, Amsterdam, Netherlands · Institute for Ecological Economics, Vienna University of Economics and Business, Austria

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 the inherent stochasticity of transition dynamics, 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 with a GAN-trained stochastic component in latent space. Policies derived from this model provide a strong baseline, outperforming or matching general model-based and model-free approaches across stochastic continuous-control benchmarks. This work demonstrates 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

Figures & tables

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 (generation) stage, the initial latent

z_{0}

is sampled from the standard Gaussian prior distribution instead. Note that the

\tilde{o}_{i}

could denote either the real observation

o_{i}

or predicted one

\hat{o}_{i}

during training, depending on whether teacher forcing strategy is used, yet it is always set to be

\tilde{o}_{i}=\hat{o}_{i}

during inference. (b) The SDE generator and MLP critic of the latent SDE. In the SDE solver, we use the learnt drift function

f_{\bar{\theta}_{\mu}}

from the ODE decoder in (a) and train the diffusion function

g_{\theta_{\sigma}}

to capture the full stochasticity in the latent space. In both panels, the deterministic variables are represented as diamonds while the stochastic variables are depicted as circles.

Figure 2: Overview of the architecture for training and deploying a policy adapted from the source domain in the target domain. (a) Given an off-the-shelf source policy

\pi^{\text{src}}

and transition model

\mathcal{T}_{\theta}^{\text{src}}

, we use an inverse dynamics model

I_{\eta}

to generate the target action

a_{t}^{\text{tge}}

leading to the next target state

\hat{s}_{t+1}^{\text{tge}}

I_{\eta}

is trained via minimizing the mismatch between the predicted next state

\hat{s}_{t+1}^{\text{tge}}

and the desired state

\hat{s}_{t+1}^{\text{src}}

. (b) The trained inverse model, combined with the source policy and transition, is deployed in the real target environment as an adapted target policy.

Figure 3: Comparison between neural ODE and SDE on 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-ODE-based) and SDE (N-SDE-based), as well as in real transition (model-free), are evaluated in the real environment. N-SDE-based policy achieves similar performance to the model-free policy (used as the oracle here).

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 of returns over 4 runs. (a) Under deterministic source and target dynamics, for pole lengths between 1.8 and 3.2, both the N-ODE–adapted and ensemble-adapted policies consistently outperform the original non-adapted policy, and both substantially surpass the trained-from-scratch policy. The N-ODE–adapted and ensemble-adapted 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–adapted, 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-from-scratch policy performing worst.

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.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

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.

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.

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 ODE-based models.

Figure 9: (a) Scatter plots of 200 PCA-embedded 2D 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.

		Cartpole	Swimmer	Walker2d	Hopper
Training hparams	Max env. steps	250K	500K
	Model batch size	256	128
	Optimizer	Adam
Neural/Latent ODE/SDE	Planning horizon	N/A	10
	Search population	N/A	1000	1700
	Trajectory cut length	20	10

Appendix

Table 1: Hyperparameter summary of model-based and model-free baselines across environments. Model architectures are described by network type and number of hidden layers with their sizes. For MBPO in Hopper, the model horizon follows a capped linear function of the epoch

i

f(i)=\min(\max(1+0.175\times(i-20),1),15)

, as in ( Janner et al., 2019 ) .

Explore similar work

Flowing Through States: Neural ODE Regularization for Reinforcement Learning
Aug 6, 2026Mohamed Ghanem, Bernd FinkbeinerLatent Dynamics Modeling Representation Learning for RL
CISPA Helmholtz Center for Information Security · Technical University of Munich
Integrating Lagrangian Neural Networks into the Dyna Framework for Reinforcement Learning
Mar 9, 2026Shreya Das, Kundan Kumar, Muhammad Iqbal +4Model-Based RL
Department of Electrical Engineering and Automation, Aalto University, Finland · VTT Technical Research Centre of Finland · Department of Computer Science, University of Helsinki, Finland +1
Variational Smoothing and Inference for SDEs from Sparse Data with Dynamic Neural Flows
May 7, 2026Yu Wang, Arnab GangulyStochastic Differential Equations Bayesian Inference
Department of Mathematics Louisiana State University

Aug 6, 2026J/K move · Enter open · S save

Flowing Through States: Neural ODE Regularization for Reinforcement Learning

Mohamed Ghanem, Bernd Finkbeiner

CISPA Helmholtz Center for Information Security · Technical University of Munich

Latent Dynamics Modeling Representation Learning for RL State Representation Learning Neural ODEs

Open paper AskPDF arXiv

Neural networks applied to sequential decision-making tasks typically rely on latent representations of environment states. While environment dynamics dictate how semantic states evolve, the corresponding latent transitions are usually left implicit, creating a potential misalignment between the two. We propose to model latent dynamics explicitly by drawing an analogy between Markov decision process (MDP) trajectories and ordinary differential equation (ODE) flows: in both cases, the current state fully determines its successors. Building on this view, we introduce a neural ODE-based regularization method that enforces latent embeddings to follow consistent ODE flows, thereby aligning representation learning with environment dynamics. Although broadly applicable to deep learning agents, we demonstrate its effectiveness in reinforcement learning by integrating it into Actor-Critic algorithms. Our approach yields major performance gains across various standard Atari benchmarks for A2C and gridworld environments for PPO.

Abstract

Explore similar work

Flowing Through States: Neural ODE Regularization for Reinforcement Learning

Integrating Lagrangian Neural Networks into the Dyna Framework for Reinforcement Learning

Variational Smoothing and Inference for SDEs from Sparse Data with Dynamic Neural Flows