Paper deep dive
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
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 3/26/2026, 1:45:32 AM
Summary
The paper introduces a unified framework for model-based reinforcement learning (RL) using latent neural ordinary differential equations (ODEs) and stochastic differential equations (SDEs). The authors propose a two-phase training approach: first, learning mean latent dynamics via a latent ODE, and second, modeling stochasticity using a GAN-trained latent SDE. This approach addresses partial observability and improves sample efficiency in stochastic continuous-control environments by enabling noise-aware planning.
Entities (6)
Relation Signals (3)
Latent SDE → addresses → POMDP
confidence 95% · To address partial observability, we introduce a latent SDE model
Latent SDE → combines → Neural ODE
confidence 95% · we introduce a latent SDE model that combines an ODE with a GAN-trained stochastic component
Neural SDE → improves → Sample Efficiency
confidence 90% · neural SDEs more effectively capture the inherent stochasticity of transition dynamics, enabling high-performing policies with improved sample efficiency
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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: this https URL
Tags
Links
- Source: https://arxiv.org/abs/2603.23245v1
- Canonical: https://arxiv.org/abs/2603.23245v1
Trouble viewing inline? Open PDF directly →
Full Text
67,958 characters extracted from source content.
Expand or collapse full text
Neural ODE and SDE Models for Adaptation and Planning in Model-Based Reinforcement Learning Chao Han† chao.han@cruk.manchester.ac.uk School of Computer Science, The University of Sheffield, UK Present address: Cancer Research UK National Biomarker Centre, The University of Manchester, UK Stefanos Ioannou† si386@cam.ac.uk School of Computer Science, The University of Sheffield, UK Luca Manneschi l.manneschi@sheffield.ac.uk School of Computer Science, The University of Sheffield, UK T.J. Hayward t.hayward@sheffield.ac.uk School of Chemical, Materials and Biological Engineering, The University of Sheffield, UK Michael Mangan m.mangan@sheffield.ac.uk School of Computer Science, The University of Sheffield, UK Aditya Gilra aditya.gilra@wu.ac.at Machine Learning group, Centrum Wiskunde & Informatica, Amsterdam, Netherlands School of Computer Science, The University of Sheffield, UK Present address: Institute for Ecological Economics, Vienna University of Economics and Business, Austria Eleni Vasilaki e.vasilaki@sheffield.ac.uk School of Computer Science, The University of Sheffield, UK 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. †Equal Contribution. 1 Introduction In recent years, the family of neural differential equations (neural DEs) (Chen et al., 2018; Rubanova et al., 2019; Li et al., 2020; Kidger et al., 2020; 2021) have emerged as a powerful framework for modelling dynamical systems. The general idea of these models is to use neural networks to parameterise the derivatives of the system’s 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, continuous-time transitions. As a result, neural DEs are especially well-suited for learning and representing transition dynamics in reinforcement learning (RL). 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 sampled 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 sample efficiency compared to model-free baselines (Du et al., 2020). Additionally, incorporating neural ODEs with control-theoretic constraints has supported the development of safe and stable RL in continuous-time domains (Zhao et al., 2025). Neural SDEs extend the neural ODE framework by jointly modelling the deterministic and stochastic components of system dynamics, allowing for uncertainty-aware 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). Physics-constrained 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. Despite recent advances, the application of latent neural SDEs to partially observed, stochastic RL environments 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 ODE 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-based RL with latent neural SDEs improves sample efficiency compared to both ODE-based models and a model-free SAC baseline in the most challenging fully or 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-aware planning in stochastic domains. 2 Background Markov decision process (MDP) and partially observable MDP (POMDP). An MDP (Cassandra et al., 1994) is defined by a tuple ⟨,,,R⟩ ,A,T,R , where S and A are sets of states and actions respectively, :×→Δ()T:S×A→ (S) is the probabilistic transition function (dynamics), with (st+1|st,at)T(s_t+1|s_t,a_t) representing the probability of transitioning into st+1s_t+1 from sts_t under ata_t, R:×→ℝR:S×A×S is the deterministic reward function. The initial state s0s_0 follows certain distribution ρ0 _0, the horizon is T∈ℕ+∪+∞T ^+∪\+∞\ and the discount factor is γ∈[0,1]γ∈[0,1], with 1 for finite horizon. When the state sts_t is not fully observable, MDP tuples are extended to POMDP (Cassandra et al., 1994) ⟨,,,,O,R⟩ ,A,O,T,O,R by adding a set of observations O and an emission function O:→Δ()O:S→ (O), which probabilistically maps a state sts_t to an observation oto_t. Specifically, we consider POMDPs with two sources of uncertainty in the transition dynamics of observed states (ot+1∣ot,at)T(o_t+1 o_t,a_t): (i) aleatoric stochasticity in the underlying MDP transition (st+1∣st,at)T(s_t+1 s_t,a_t) that produces the observations, and (i) epistemic uncertainty from partial observability, which can often be reduced by conditioning on a history of observations and actions. To focus on transition stochasticity and decouple it from aleatoric uncertainty in observation emission, we assume deterministic emission functions O, i.e., ot=o(st)o_t=o(s_t). 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. In such POMDPs, we introduce a latent representation ztz_t that summarizes the observed history up to t a0:t−1,o0:t\a_0:t-1,o_0:t\ as a proxy of the state sts_t (Du et al., 2020; Ni et al., 2024). The transition function of the latent ztz_t is (zt+1|zt,at,ot)T(z_t+1|z_t,a_t,o_t), with the observation oto_t deterministically emitted from ztz_t via ot=o(zt)o_t=o(z_t). 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 in the POMDP. Neural Ordinary Differential Equation (Neural ODE). Neural differential equations are a family of differential equations whose rate functions are approximated by learnable neural networks (Kidger, 2022). A classic example is a neural ODE (Chen et al., 2018): dstdt=fθ(st,t),wherest0=s0 ds_tdt=f_θ(s_t,t), s_t_0=s_0 (1) where sts_t denotes the state at time t as the solution of the ODE initial-value problem (IVP). The initial state at t0t_0 is s0s_0. The rate function fθf_θ is mostly parameterized by a multi-layer perceptron (MLP) with learnable parameters θ. The general regime of neural ODE allows any off-the-shelf numerical integrator to solve for the above ODE IVP. More concretely: s^i+1=ODEsolve(fθ(s^i,t0+iΔt)),i=0,…,T−1 s_i+1=ODEsolve(f_θ( s_i,t_0+i t)), i=0,…,T-1 (2) where neural ODEsolve refers to the ODE numerical integrator used. s^i+1 s_i+1 is a step forward of the neural ODEsolve from s^i s_i with a pre-defined, fixed step size Δt t. 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 as: z0∼qϕ(z0|o0:T),zi+1=ODEsolve(fθ(zi,t0+iΔt)),o^i=oθ(zi),i=0,…,T−1 gatheredz_0 q_φ(z_0|o_0:T), z_i+1=ODEsolve(f_θ(z_i,t_0+i t)), o_i=o_θ(z_i), i=0,…,T-1 gathered (3) Here, an (RNN) encoder parameterized by ϕφ maps observations o0:To_0:T to a distribution of the initial latent state z0z_0. A decoder parameterized by θ reconstructs observations o^0:T o_0:T via a latent-variable ODE determined by sampled z0z_0 and an emission function oθ(⋅)o_θ(·). Neural Stochastic Differential Equation (Neural SDE). A neural SDE consists of a parameterized deterministic drift term and a parameterized stochastic diffusion term: dst=fθ(st,t)dt+gθ(st,t)dwt,wherest0∼μ _t=f_θ(s_t,t)dt+g_θ(s_t,t)dw_t, s_t_0 μ (4) where stt≥t0\s_t\_t≥ t_0 is a continuous-time stochastic process, whose initial state st0s_t_0 is drawn from some probability distribution μ. wtt≥0\w_t\_t≥ 0 is the Brownian motion. The paper follows the idea of training a neural SDE as a (Wasserstein) generative adversarial net (GAN) (Goodfellow et al., 2014; Arjovsky et al., 2017), where the real and generated data samples are (interpolated) infinite-dimensional paths (Kidger et al., 2021). 3 Methodology In this section, we extend the vanilla autonomous neural ODE/SDE-based models to their controlled variants by incorporating actions to model the MDP state transition dynamics (st+1|st,at)T(s_t+1|s_t,a_t) and POMDP latent transition (zt+1|zt,at,ot)T(z_t+1|z_t,a_t,o_t). We also describe how to train a policy for planning based on the learned transition model, as well as how to adapt the model and policy to a similar environment without retraining from scratch. 3.1 Dynamics model learning We use the terms neural ODE/SDE and latent ODE to overload the names of original autonomous models and here denote their controlled variants. In a similar vein to latent neural ODEs, we also propose the latent SDE by employing the neural 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 a deterministic environment be defined as an autonomous neural ODE (t is not explicitly given as an argument to fθf_θ): s^i+1=ODEsolve(fθ(s^i,ai)),i=0,…,T−1 s_i+1=ODEsolve(f_θ( s_i,a_i)), i=0,…,T-1 (5) We optimize the parameters θ by minimizing the following mean squared error (MSE) between the predicted s^0:T s_0:T and observed state trajectories s0:Ts_0:T: minθ(s0:T,a0:T−1)∼m[∑i=0T‖s^i−si‖22] _θ\,E_(s_0:T,a_0:T-1) _m [ _i=0^T\| s_i-s_i\|_2^2 ] (6) where real state-action trajectories are sampled from a replay buffer mD_m. Latent ODE. Similar to Du et al. (2020), we modify the vanilla latent ODE in Eq. 3 to model the POMDP latent transition by taking actions aia_i and (predicted) observations o~i o_i into account for the underlying evolution of RNN hidden states hih_i in the encoder and ODE latent states ziz_i in the decoder. Specifically, the adjusted latent ODE consists of the following RNN encoder parameterized by ϕφ: hi+1=RNNCellϕ(hi,ai,o~i),o^i=oϕ(hi),i=0,…,T−1,[μz0,σz0]=MLPϕ(hT),qϕ(z0|o0:T,a0:T−1)=(z0;μz0,diag(σz02)) gatheredh_i+1=RNNCell_φ(h_i,a_i, o_i), o_i=o_φ(h_i), i=0,…,T-1,\\ [ _z_0, _z_0]=MLP_φ(h_T), q_φ(z_0|o_0:T,a_0:T-1)=N(z_0; _z_0,diag( _z_0^2)) gathered (7) and neural ODE decoder parameterized by θ: z0∼qϕ(z0|o0:T,a0:T−1),z~i=MLPθ(zi,ai,o~i),zi+1=ODEsolve(fθ(z~i)),o^i=oθ(zi),pθ(oi+1|zi,ai,o~i)=(oi+1;oθ(zi+1),I),i=0,…,T−1 gatheredz_0 q_φ(z_0|o_0:T,a_0:T-1), z_i=MLP_θ(z_i,a_i, o_i), z_i+1=ODEsolve(f_θ( z_i)),\\ o_i=o_θ(z_i), p_θ(o_i+1|z_i,a_i, o_i)=N(o_i+1;o_θ(z_i+1),I), i=0,…,T-1 gathered (8) where oθ(⋅)o_θ(·) denotes the parameterized emission function, z~i z_i has the same dimension as ziz_i and represents the transformed latent state at which the derivative of latent ODE fθf_θ is evaluated. o~i o_i could be either the real observation oio_i or predicted observation o^i o_i during training, while always setting o~i=o^i o_i= o_i during inference. In this paper, models are typically trained using teacher forcing, i.e., o~i=oi o_i=o_i, as 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 observation trajectories o0:To_0:T: maxϕ,θ(o0:T,a0:T−1)∼m[ _φ,θ\,E_(o_0:T,a_0:T-1) _m [ z0∼qϕ(z0|o0:T,a0:T−1)[∑i=0T−1logpθ(oi+1|zi,ai,o~i)] _z_0 q_φ(z_0|o_0:T,a_0:T-1) [ _i=0^T-1 p_θ(o_i+1|z_i,a_i, o_i) ] −DKL(qϕ(z0|o0:T,a0:T−1)∥p(z0))] -D_KL (q_φ(z_0|o_0:T,a_0:T-1) p(z_0) ) ] (9) Here, the joint generative distribution is decomposed as pθ(o0:T|z0:T,a0:T−1)=∏i=0Tpθ(oi+1|zi,ai,o~i)p_θ(o_0:T|z_0:T,a_0:T-1)= _i=0^Tp_θ(o_i+1|z_i,a_i, o_i) and p(z0)=(z0;0,I)p(z_0)=N(z_0;0,I) is the standard Gaussian prior. Neural SDE. We employ a GAN-based SDE to generate synthetic trajectories that closely match those sampled from a stochastic MDP transition. More precisely, the adapted neural SDE trained as Wasserstein GAN (WGAN) consists of a neural SDE generator parameterized by θ: s^i+1=SDEsolve(fθ(s~i,ai),gθ(s~i,ai),Δwi),s¯i=MLPθ(s^i)i=0,…,T−1 s_i+1=SDEsolve(f_θ( s_i,a_i),g_θ( s_i,a_i), w_i), s_i=MLP_θ( s_i) i=0,…,T-1 (10) and an MLP critic (discriminator) parameterized by ψ: y=MLPψ(s¯0:T,a0:T−1) y=MLP_ψ( s_0:T,a_0:T-1) (11) Here, the projection from s^i s_i to s¯i s_i allows for more flexible generated states. Different from Kidger et al. (2021) that uses a neural controlled differential equation (CDE) as a discriminator, we simply use an MLP that takes flattened state-action trajectories as input. This change makes optimization much easier, but also limits the discriminator’s capacity for processing varying-length trajectories. Let Gθ:(w0:T,a0:T−1)→s0:TG_θ:(w_0:T,a_0:T-1)→ s_0:T denote the overall map of the generator from paths of noises and actions to that of synthetic observations, and Dψ:(s0:T,a0:T−1)→yD_ψ:(s_0:T,a_0:T-1)→ y denote the overall map of the critic from paths of real/synthetic observations and actions to a scalar score. The training dynamics of the WGAN are formulated as follows: minθmaxψ‖Dψ‖L≤1(s0:T,a0:T−1)∼m[Dψ(s0:T,a0:T−1)]−w0:T∼pwa0:T−1∼[Dψ(Gθ(w0:T,a0:T−1),a0:T−1)] _θ\, _ subarraycψ\\ \|D_ψ\|_L≤ 1 subarray\,E_(s_0:T,a_0:T-1) _m[D_ψ(s_0:T,a_0:T-1)]-E_ subarraycw_0:T p_w\\ a_0:T-1 subarray[D_ψ(G_θ(w_0:T,a_0:T-1),a_0:T-1)] (12) where ‖Dψ‖L≤1\|D_ψ\|_L≤ 1 represent the Lipschitz continuity constraint enforced on the critic function DψD_ψ. We use the gradient penalty (Gulrajani et al., 2017) to achieve the Lipschitz constraint. 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 z0z_0 is sampled from the standard Gaussian prior distribution instead. Note that the o~i o_i could denote either the real observation oio_i or predicted one o^i o_i during training, depending on whether teacher forcing strategy is used, yet it is always set to be o~i=o^i o_i= 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θ¯μf_ θ_μ from the ODE decoder in (a) and train the diffusion function gθσg_ _σ 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. Latent SDE. A weakness of latent ODE is that it can only model deterministic MDP transitions (or the mean dynamics of a stochastic MDP transitions) of states underlying a POMDP. To learn the full stochastic MDP state transition, we employed a phased training framework. The latent ODE (Eqs. 7 and 8) is first employed to capture the mean transition dynamics in the latent space. Specifically, the encoder-decoder pair of the latent ODE is described as follows: z0∼qϕ(z0|o0:T,a0:T−1),zi+1=ODEsolve(fθμ(zi,ai,o~i)),o^i=oθμ(zi),i=0,…,T−1 gatheredz_0 q_φ(z_0|o_0:T,a_0:T-1), z_i+1=ODEsolve(f_ _μ(z_i,a_i, o_i)), o_i=o_ _μ(z_i), i=0,…,T-1 gathered (13) where the recurrent encoder and the ODE decoder are respectively parameterized by ϕφ and θμ _μ, which are trained in an end-to-end VAE-style setup (Eq. 3.1). 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: z0∼p(z0),zi+1=SDEsolve(fθ¯μ(zi,ai,o~i),gθσ(zi,ai,o~i),Δwi)o^i=oθσ(zi),i=0,…,T−1,y=MLPψ(o^0:T,a0:T−1) gatheredz_0 p(z_0), z_i+1=SDEsolve(f_ θ_μ(z_i,a_i, o_i),g_ _σ(z_i,a_i, o_i), w_i) o_i=o_ _σ(z_i), i=0,…,T-1,\\ y=MLP_ψ( o_0:T,a_0:T-1) gathered (14) where z0z_0 is sampled from the Gaussian prior p(z0)p(z_0) used in the VAE training. fθ¯μf_ θ_μ and gθσg_ _σ are the drift and diffusion function of the SDE, where θ¯μ θ_μ denotes the detached optimal θμθ_μ from the first phase. The diffusion gθσg_ _σ and emission function oθσo_ _σ of the SDE generator are trained in pairs with the MLP critic parameterized by ψ in GAN-style setup (Eq. 12). We term this model framework that combines latent ODE and neural SDE as latent SDE, and depict its computation graph in Fig. 1. 3.2 Planning and policy learning with learned dynamics When addressing sample-intensive tasks (such as Mujoco), a model-free RL method is often trained off-policy using near-optimal transitions collected by a planning strategy driven by learned, task-specific dynamics, thereby reducing sample complexity. In this paper, following Du et al. (2020), we adopt a 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). Input : The empty replay buffer for model learning mD_m and for policy learning pD_p, the initial latent transition model θ(zt+1|zt,at,ot)T_θ(z_t+1|z_t,a_t,o_t), the reward function R(ot,at,ot+1)R(o_t,a_t,o_t+1), the initial actor πω(at|ot) _ω(a_t|o_t), the two critics Qφiπ(ot,at)Q^π_ _i(o_t,a_t) and their target networks Qφi′π(ot,at),i=1,2Q^π_ _i(o_t,a_t),\ 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 PθP_θ, πω _ω and mini=1,2Qφiπ _i=1,2Q^π_ _i. Collect trajectories using a uniformly distributed policy and save them into mD_m; for i=1i=1 to E do Update θ(zt+1|zt,at,ot)T_θ(z_t+1|z_t,a_t,o_t) using data from mD_m, as detailed in Section 3.1; Observe the initial state o0o_0 and initialize the latent state z0z_0 ; for j=1j=1 to M do for k=1k=1 to K do o^0(k),z0(k)←oj−1,zj−1 o^(k)_0,z^(k)_0← o_j-1,z_j-1; for h=1h=1 to H do Select the action ah−1(k)∼πω(⋅|o^h−1(k))a^(k)_h-1 _ω(·| o^(k)_h-1); Transit to next latent zh(k)←θ(⋅|zh−1(k),ah−1(k),o^h−1(k))z^(k)_h _θ(·|z^(k)_h-1,a^(k)_h-1, o^(k)_h-1) and emit the observation o^h(k)←oθ(zh(k)) o^(k)_h← o_θ(z^(k)_h); Calculate the reward r^h−1(k)←R(o^h−1(k),ah−1(k),o^h(k)) r^(k)_h-1← R( o^(k)_h-1,a^(k)_h-1, o^(k)_h); end for Select the action aH(k)∼πω(⋅|o^H(k))a^(k)_H _ω(·| o^(k)_H); end for Select the index of the sequence with the maximum return among the K sequences: k∗←argmaxk∑h=1Hγh−1r^h−1(k)+γHmini=1,2Qφiπ(o^H(k),aH(k))k ← *arg\,max_k _h=1^Hγ^h-1 r^(k)_h-1+γ^H _i=1,2Q^π_ _i( o^(k)_H,a^(k)_H); Select the first action of the best sequence aj−1←a0(k∗)a_j-1← a^(k )_0; Execute aj−1a_j-1 and observe the next state ojo_j; Transit to the next latent zj←θ(⋅|zj−1,aj−1,oj−1)z_j _θ(·|z_j-1,a_j-1,o_j-1); Calculate the reward rj−1←R(oj−1,aj−1,oj)r_j-1← R(o_j-1,a_j-1,o_j); if ojo_j is not the terminal state then Store (oj−1,aj−1,rj−1,oj)(o_j-1,a_j-1,r_j-1,o_j) in pD_p; else Store (oj−1,aj−1,rj−1,NULL)(o_j-1,a_j-1,r_j-1,NULL) in pD_p, reset the environment; Observe the initial state o0o_0 and initialize the latent state z0z_0 ; continue; end if Update Qφ1πQ^π_ _1, Qφ2πQ^π_ _2 and πω _ω using data from pD_p; Update target networks Qφ1′πQ^π_ _1, Qφ2′πQ^π_ _2; end for Store the trajectories collected using the current actor in mD_m; end for Algorithm 1 MPC with SAC for POMDPs. 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 Algorithm 1, 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 a source environment to a target environment, at the cost of a minimal amount of data from the target environment. The source environment typically refers to a 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 characteristics 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 are characterized by MDP transitions and share the same actuated degrees of freedom. As illustrated in Fig. 2A, given a current target state sttges_t^tge, we first compute the corresponding source action atsrc=πsrc(sttge)a_t^src=π^src(s_t^tge) using the predefined source policy πsrcπ^src. Next, we estimate the subsequent source state s^t+1src=θsrc(sttge,atsrc) s_t+1^src=T_θ^src(s_t^tge,a_t^src) via the parameterized source transition model θsrcT_θ^src, conditioned on the current state sttges_t^tge and action atsrca_t^src. Finally, the current target action attge=Iη(sttge,s^t+1src)a_t^tge=I_η(s_t^tge, s_t+1^src) is computed using the inverse dynamics model IηI_η, which maps the current target state sttges_t^tge and the desired next state s^t+1src s_t+1^srcto a target action that drives the next target state towards the desired state s^t+1src s_t+1^src as close as possible. 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 πsrcπ^src and transition model θsrcT_θ^src, we use an inverse dynamics model IηI_η to generate the target action attgea_t^tge leading to the next target state s^t+1tge s_t+1^tge. IηI_η is trained via minimizing the mismatch between the predicted next state s^t+1tge s_t+1^tge and the desired state s^t+1src s_t+1^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. Data collection/training. 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 in the target domain. Specifically, we adapt the deterministic component of the source transition model parameterized by θ to that of the target transition model parameterized by the augmented parameter set θ+=θ∪θaugθ^+=θ∪ _aug, which includes additional parameters θaug _aug that capture the variation in the deterministic component of transition dynamics between the source and target environments. We train only the additional parameters θaug _aug 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 mD_m. Similar to the loss for a deterministic Neural ODE defined in Eq. 6, we minimize the following MSE loss between the predicted and real target state trajectories: ℒm(θaug;s0:Ttge,a0:T−1tge)=(s0:Ttge,a0:T−1tge)∼m[ℓm(θaug;s0:Ttge,a0:T−1tge)] _m( _aug;s_0:T^tge,a_0:T-1^tge)=E_(s_0:T^tge,a_0:T-1^tge) _m[ _m( _aug;s_0:T^tge,a_0:T-1^tge)] (15) where ℓm(θaug;s0:Ttge,a0:T−1tge)=∑t=1T‖s^ttge−sttge‖22=∑t=1T‖θ+tge(st−1tge,at−1tge)−sttge‖22 _m( _aug;s_0:T^tge,a_0:T-1^tge)= _t=1^T\| s^tge_t-s^tge_t\|_2^2= _t=1^T\|T_θ^+^tge(s_t-1^tge,a_t-1^tge)-s^tge_t\|_2^2 (16) Empirically we only require a 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 is sufficiently accurate, and then use it for training the inverse dynamics model, which is interleaved with data collection. Specifically, we repeat the following data collection/training loop: during the data collection phase, given a current target state, we use the preliminary source policy and transition model to 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 leading to the next target state. We repeat the target environment steps and save the sequence along the steps ⋯,sttge,s^t+1src,st+1tge,s^t+2src,⋯\·s,s_t^tge, s_t+1^src,s_t+1^tge, s_t+2^src,·s\ into the buffer invD_inv. Target states along the trajectory collected in such an on-the-fly fashion are near the optimal target trajectory, which will lead to much faster convergence of the inverse dynamics model compared with training data collected by a 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 a mini-batch of real-desired transition pairs sampled from invD_inv: ℒinv(η;sttge,s^t+1src)=(sttge,s^t+1src)∼inv[ℓinv(η;sttge,s^t+1src)] _inv(η;s_t^tge, s_t+1^src)=E_(s_t^tge, s_t+1^src) _inv[ _inv(η;s_t^tge, s_t+1^src)] (17) where ℓinv(η;sttge,s^t+1src)=‖s^t+1tge−s^t+1src‖22=‖θ+tge(sttge,Iη(sttge,s^t+1src))−s^t+1src‖22 _inv(η;s_t^tge, s_t+1^src)=\| s_t+1^tge- s_t+1^src\|_2^2=\|T_θ^+^tge(s_t^tge,I_η(s_t^tge, s_t+1^src))- s_t+1^src\|_2^2 (18) It is 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) can 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 aleatoric uncertainty, which is irreducible (Chua et al., 2018; Han et al., 2024). Deployment. As illustrated in Fig. 2B, we can use the learned inverse dynamics model as a target policy, which is computed via adaptation from the source policy as follows: πtge(sttge)=Iη(sttge,θsrc(sttge,πsrc(sttge))) π^tge(s_t^tge)=I_η(s_t^tge,T_θ^src(s_t^tge,π^src(s_t^tge))) (19) 4 Experiments We empirically evaluate our ODE- and SDE-based models in stochastic environments of increasing complexity. 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 section 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, we 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, we modify the standard deterministic OpenAI Gym environments (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 velocities (or angular velocities) 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.i.d.) standard Gaussian force to the cart at each step, which can be formulated as an SDE of cart velocity. • Stochastic swimmer. A 3-link swimmer is propelled forward in a fluid, with a 10-dimensional observation space and 2 actuators. We introduce i.i.d. Gaussian noise sampled from (0,5002)N(0,500^2) to the stiffness parameter of the first actuated joint of the Swimmer. • Stochastic hopper. A single-legged robot is made to hop forward as far as possible, with a 12-dimensional observation space and 3 actuators. We apply stochastic winds to the hopper at each step. The magnitude of the wind parameter is i.i.d. sampled from (0,52)N(0,5^2). • 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,52)N(0,5^2). 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 dynamics, as described in 2. 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. In the more challenging MuJoCo tasks, we additionally consider latent ODE (L-ODE) and latent SDE (L-SDE) models for planning and policy optimization. Across all environments, we use Soft Actor-Critic (SAC) (Haarnoja et al., 2018) as a model-free baseline, training Gaussian policies but evaluating their deterministic mean. As the model-based counterpart, 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 ODE/SDE-based 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 baselines 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. 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). Neural ODEs model the mean of stochastic transitions while Neural SDEs capture the full stochastic dynamics. Fig. 3a 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. 3b 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 tendency 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. 6 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. 3c). 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 oracle distribution. The poor performance of ODE-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 deterministic and stochastic cartpole environments with increasing pole length. We compare different transition models: the neural ODE, the neural SDE, and the ensemble of Gaussian neural networks. The results show that adapted source policies based on 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. 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. Deterministic transition. We first consider the setting where the cartpole dynamics are deterministic (Fig. 4a). We compare two adapted policies against the non-adapted source policy (red curve) as the pole length increases. 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, we also include a policy trained from scratch (brown curve) for 2k iterations, using the same number of target environment 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 is 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 a from-scratch policy would require far more samples to reach comparable performance. Performance of both adapted and non-adapted policies 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 as the only varying factor (Fig. 4b). We evaluate three adapted policies: ODE-adapted, SDE-drift–adapted, and ensemble-adapted (blue, green, and purple curves). These 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 to suboptimal ensemble members. As in the deterministic case, all adapted policies outperform the non-adapted and scratch-trained baselines, though their performance also declines as pole length increases. Interestingly, the source policy trained in a stochastic environment (red curve in Fig. 4b) is more robust to environmental changes than the one trained in a deterministic environment (red curve in Fig. 4a). This robustness likely arises because environmental stochasticity encourages broader exploration of 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 non-adapted baseline in Fig. 4b compared to Fig. 4a. 4.3 Model-based policy learning For more complex stochastic MuJoCo environments, we adopt a framework that interleaves exploration via MPC planning with model-free policy optimization, as 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–based policy optimization methods against a model-free baseline (SAC) and a model-based baseline (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. 5 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-based policies. The SDE-based models explicitly represent both the mean and variance of the stochastic transition dynamics, enabling more accurate prediction of the next (partial) observation than ODE-based models (Fig. 8 in Appendix B). This, in turn, allows SDE-based policies to better account for transition uncertainty and avoid high-risk actions. In contrast, ODE-based models capture only the mean dynamics. As a result, planning with them may deviate from the true environment trajectories, leading to suboptimal action choices and limiting exploration of new 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 requiring fewer 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 of stochastic Hopper (no position). We further observe that the next partial observation 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 SDE) (see Fig. 9 in the Appendix B), which might explain the performance advantage. MBPO performs no better than SDE-based 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 our model-based planning and policy learning 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 C): 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 et al. (2019) report that such a capped horizon accelerates learning without significantly deviating the trajectories. In stochastic walker2d, we use a single-step model horizon as in Janner et al. (2019), 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 observable 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 on 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-based counterparts. MBPO is by far the most demanding. For instance, in stochastic Walker2d, N-ODE requires about 2× the training time of SAC; L-ODE and N-SDE about 2× that of N-ODE; L-SDE about 1.5× that of N-ODE; and MBPO roughly 6× longer than L-SDE. MBPO also consumes about 3× more CPU memory than the other baselines. All experiments were run on a single NVIDIA A100 GPU. This overhead arises because ODE/SDE-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 optimization in stochastic continuous-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 et al. (2016) did not specify how the target transition is estimated or the training loss for the inverse model, our contribution is to make this explicit: we show how neural differential-equation transition models can be instantiated within the inverse-dynamics framework and contrasted with MBPO’s ensemble transition models. Empirically, neural ODE/SDE models achieve adaptation performance comparable to the ensemble model in a simple environment; more importantly, even in a stochastic domain where the change stems from a 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 approaches, 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 observability. 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 findings 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–based policies may not offer advantages over standard model-based or model-free baselines. For instance, in continuous-control tasks with deterministic dynamics or very weak noise, ODE/SDE-based policies may not improve over model-based baselines such as MBPO. In stochastic walker2d, where the transition noise is relatively weak ((0,52)N(0,5^2) wind noise), MBPO with a single-step rollout horizon achieves performance comparable to the latent-SDE baseline. 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 (0,5002)N(0,500^2) 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-efficiency advantages of planning with ODE/SDE-based 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 of SDE-based policies over SAC is observed only in 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 required before drawing a conclusion. Finally, there are task regimes where ODE/SDE-based models are inherently ill-suited. ODE/SDE continuous-time integrators struggle with hybrid or discontinuous dynamics, 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 history-based models (e.g., RNNs (Cho et al., 2014) 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) (CHIST-ERA-19-XAI-002), funded by the Engineering and Physical Sciences Research Council (EPSRC), United Kingdom (grant EP/V055720/1). C. Han, S. Ioannou, L. Manneschi, M. Mangan, and E. Vasilaki acknowledge support from the EPSRC project Active Learning and Selective Attention for Robust, Transparent and Efficient AI (ActiveAI) (grant EP/S030964/1). C. Han, L. Manneschi, T. J. Hayward, and E. Vasilaki acknowledge support from the EPSRC project Magnetic Architectures for Reservoir Computing Hardware (MARCH) (grant EP/V006339/1). References M. Arjovsky, S. Chintala, and L. Bottou (2017) Wasserstein generative adversarial networks. In International conference on machine learning, p. 214–223. Cited by: §2. A. R. Cassandra, L. P. Kaelbling, and M. L. Littman (1994) Acting optimally in partially observable stochastic domains. In Aaai, Vol. 94, p. 1023–1028. Cited by: §2. R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud (2018) Neural ordinary differential equations. Advances in neural information processing systems 31. Cited by: §1, §1, §2. K. Cho, B. Van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio (2014) Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078. Cited by: §5. P. Christiano, Z. Shah, I. Mordatch, J. Schneider, T. Blackwell, J. Tobin, P. Abbeel, and W. Zaremba (2016) Transfer from simulation to real world through learning deep inverse dynamics model. arXiv preprint arXiv:1610.03518. Cited by: §1, §3.3, §5. K. Chua, R. Calandra, R. McAllister, and S. Levine (2018) Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems 31. Cited by: §3.2, §3.3, §4. A. Djeumou, L. Huang, A. Buisson, et al. (2023) Physics-constrained neural SDEs for efficient RL & control. In Conference on Robot Learning (CoRL), Cited by: §1. J. Du, J. Futoma, and F. Doshi-Velez (2020) Model-based reinforcement learning for semi-markov decision processes with neural odes. Advances in Neural Information Processing Systems 33, p. 19805–19816. Cited by: Appendix C, §1, §2, §3.1, §3.2, §3.2. I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio (2014) Generative adversarial nets. Advances in neural information processing systems 27. Cited by: §2. I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville (2017) Improved training of wasserstein gans. Advances in neural information processing systems 30. Cited by: §3.1. T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine (2018) Soft actor-critic: off-policy maximum entropy deep reinforcement learning with a stochastic actor. arXiv [cs.LG]. External Links: Link, 1801.01290 Cited by: §3.2, §4. C. Han, D. Basu, M. Mangan, E. Vasilaki, and A. Gilra (2024) Dynamical-vae-based hindsight to learn the causal dynamics of factored-pomdps. arXiv preprint arXiv:2411.07832. Cited by: §3.3. M. Janner, J. Fu, M. Zhang, and S. Levine (2019) When to trust your model: model-based policy optimization. Advances in neural information processing systems 32. Cited by: Appendix C, Table 1, Table 1, §4, §4, §4.3. P. Kidger, J. Foster, X. Li, H. Oberhauser, and T. Lyons (2021) Neural SDEs as infinite-dimensional GANs. arXiv [cs.LG]. External Links: Link, 2102.03657 Cited by: §1, §1, §2, §3.1. P. Kidger, J. Morrill, J. Foster, and T. Lyons (2020) Neural controlled differential equations for irregular time series. Advances in neural information processing systems 33, p. 6696–6707. Cited by: §1, §4.3. P. Kidger (2022) On neural differential equations. arXiv preprint arXiv:2202.02435. Cited by: §2. D. P. Kingma (2013) Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114. Cited by: §2. M. Koprulu, A. Dhariwal, S. Ganguli, and S. Arora (2025) NUNO: neural uncertainty-aware neural SDE for offline model-based RL. In International Conference on Learning Representations (ICLR), Cited by: §1. X. Li, T. L. Wong, R. T. Chen, and D. Duvenaud (2020) Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, p. 3870–3882. Cited by: §1, §1. A. Nagabandi, G. Kahn, R. S. Fearing, and S. Levine (2017) Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. arXiv [cs.LG]. External Links: Link, 1708.02596 Cited by: §3.2. T. Ni, B. Eysenbach, E. Seyedsalehi, M. Ma, C. Gehring, A. Mahajan, and P. Bacon (2024) Bridging state and history representations: understanding self-predictive rl. arXiv preprint arXiv:2401.08898. Cited by: §2. Y. Rubanova, R. T. Q. Chen, and D. Duvenaud (2019) Latent ODEs for irregularly-sampled time series. ArXiv abs/1907.03907. External Links: Link, 1907.03907, ISSN 2331-8422 Cited by: §1, §2, §4.3. M. Towers, A. Kwiatkowski, J. Terry, J. U. Balis, G. De Cola, T. Deleu, M. Goulão, A. Kallinteris, M. Krimmel, A. KG, et al. (2024) Gymnasium: a standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032. Cited by: §4. G. Zerveas, S. Jayaraman, D. Patel, A. Bhamidipaty, and C. Eickhoff (2021) A transformer-based framework for multivariate time series representation learning. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & data mining, p. 2114–2124. Cited by: §5. L. Zhao, K. Miao, H. Cao, K. Gatsis, and A. Papachristodoulou (2025) NLBAC: a neural ODE-based algorithm for state-wise stable and safe reinforcement learning. Neurocomputing 638, p. 130041. Cited by: §1. X. Zhao, T. Ma, T. Liu, Z. Wang, and H. Li (2023) ODE-based recurrent reinforcement learning for partial observability. In NeurIPS, Cited by: §1. Appendix A Stochastic cartpole environment 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. Appendix B Stochastic hopper environments 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. Appendix C Experimental setup 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×(i−20),1),15)f(i)= ( (1+0.175×(i-20),1),15), as in (Janner et al., 2019). 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 ODE learning rate 8e-4 1e-3 SDE generator learning rate 8e-4 SDE critic learning rate 8e-5 4e-5 Encoder arch. N/A GRU [128] Decoder arch. N/A Linear MLP ODE (SDE drift func.) arch. MLP [100, 100, 100, 100] MLP [100, 100, 100] SDE diffusion func. arch. MLP [32, 32] MLP [100, 100] SDE critic arch. MLP [100, 100, 100, 100, 100] MLP [100, 64, 64] Solver Euler SAC Discount factor 0.99 Smoothing coef. 0.005 Learning rate 1e-3 3e-4 Batch size 32 128 Temperature 0.3 0.2 0.25 Actor/Critic arch. MLP [200, 200] MBPO Ensemble size 7 Model learning rate 1e-3 Model rollouts per env. step 400 Model horizon 1 Capped linear func. Model arch. MLP [200, 200, 200, 200] Model training stopping criteria. For the Mujoco tasks, we apply early stopping to the ODE-based models to avoid overfitting: training is stopped if the MSE reconstruction error on the validation set does not decrease for e consecutive epochs. Following Du et al. (2020), we use a linear decay schedule e=max(15−i,3)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 on the holdout MSE, as in Janner et al. (2019). The trained ODE-based 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 a GAN-style setup. The ODE-based model is used as the drift since it empirically captures the mean 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.