Neural Stochastic Differential Equations for Uncertainty-Aware Offline RL
Authors: Cevahir Koprulu
Franck Djeumou
ufuk topcu
Venue: ICLR 2025 Poster (2025)
Keywords: neural stochastic differential equations
offline reinforcement learning
physics-informed machine learning
Primary area: reinforcement learning
TL;DR: We develop an uncertainty-aware, offline model-based reinforcement learning approach with neural stochastic differential equations that outperforms the state-of-the-art in continuous control benchmarks, particularly in low-quality datasets.

Abstract: Offline model-based reinforcement learning (RL) offers a principled approach to using a learned dynamics model as a simulator to optimize a control policy. 
Despite the near-optimal performance of existing approaches on benchmarks with high-quality datasets, most struggle on datasets with low state-action space coverage or suboptimal demonstrations.
We develop a novel offline model-based RL approach that particularly shines in low-quality data regimes while maintaining competitive performance on high-quality datasets.
Neural Stochastic Differential Equations for Uncertainty-aware, Offline RL (NUNO) learns a dynamics model as neural stochastic differential equations (SDE), 
where its drift term can leverage prior physics knowledge as inductive bias.
In parallel, its diffusion term provides distance-aware estimates of model uncertainty by matching the dynamics' underlying stochasticity near the training data regime while providing high but bounded estimates beyond it.
To address the so-called model exploitation problem in offline model-based RL, NUNO builds on existing studies by penalizing and adaptively truncating neural SDE's rollouts according to uncertainty estimates.
Our empirical results in D4RL and NeoRL MuJoCo benchmarks evidence that NUNO outperforms state-of-the-art methods in low-quality datasets by up to 93% while matching or surpassing their performance by up to 55% in some high-quality counterparts.

Published as a conference paper at ICLR 2025
NEURAL STOCHASTIC DIFFERENTIAL EQUATIONS FOR UNCERTAINTY- AWARE, OFFLINE RL
Cevahir Koprulu Franck Djeumou Ufuk Topeu'
'Te University of Texas at Austin Rensselaer Polytechnic Institute
ABSTRACT
Offline model-based reinforcement learning (RL) offers a principled approach to using learned dynamics model simulator optimize control policy. Despite the near-optimal performance of existing approaches on benchmarks with high- quality datasets, most struggle on datasets with low state-action space coverage suboptimal demonstrations. We develop novel offline model-based RL approach that particularly shines in low-quality data regimes while maintaining competitive performance on high- high-quality datasets. Neural Stochastic Disferrntial Equations for UNcertainty-awar, Offline RI. (NUNO) learns a dynamics model as neural stochas- tic differential equations (SDE), where its drift term leverage prior physics knowledge as inductive bias. In parallel. its diffusion term provides distance-awarc estimates of model uncertainty by matching the dynamics' underlying stochasticity near the training data regime while providing high but bounded estimates beyond it. To address the so-called model exploitation problem in offline model-based RL, NUNO builds on existing studies by penalizing and adaptively truncating neural SDE's rollouts according uncertainty estimates. Our theoretical results show that penalization via a distance-aware uncertainty estimator incentivizes the pol- icy to stay close the offline data. Our empirical results in D4RL and NeoRL MuloCo benchmarks evidence that NUNO outperforms state-of-the thet-artt methods in low- low-quality datasets by up to 93% while matching or surpassing their performance by up to 55% in some high-quality counterparts.
INTRODUCTION
Offline reinforcement learning (RL) concerns the problem of learning control policies from offline datasets of interactions (Lange et 2012 Levine al., 2020). This paradigm captures safety-critical
real-world settings such as healthcare (Tseng et al. 2017; Wang al. 2018), robotics (Levine et al.) 2018 Rafailov et al. 202T) and autonomous driving Yuu et al. 2020a), where logged data is abundant, simulators are computationally expensive, online learning causes hazardous behavior. Although off-policy RL algorithms can, in principle, address settings with a priori available data, they fail the offline setting due to the distribution shift between the dataset and learned policies (Fujimoto et al.. 2019; Kumar et al., 2019). To resolve distribution shift, model-free offline RL methods introduce
conservatism via constraining learned policies available data (Jaques al. 2019 Wu al. 2019
Fujimoto & Gu (2021) or penalizing fitto---iittibutiinn actions (Kumar et al. 2020
Bai et al. 2022).
However, such approaches struggle with sub-optimal behavior policies YYu et al. (020b).
Offline model-based RL trains a control policy via synthetic data generated by a learned dynamics model (Kidambi et al. 2020 Yu et al. 202T). Compared to offline model-free RL, employing the learned model improves generalization beyond the training data regime |Ritter et al. (002). However, naive application of model-based RL causes a phenomenon called model exploitation: Learned control policies exploit the parts of the state-action space where the model is inaccurate and overestimates the return Janner et al. (2019); Yu et al. (2020b): Kurutach et al. (2018). Model exploitation can result in learning policies that perform worse than data data-logging policies.
Prior works in offline model based RL address model exploitation by enforcing conservatism for learning policies Janner al. 2019 Yu et al. 2020b Kidambi et al. 2020 Yu et al. 2021
"Corr. to: Cevahir Koprulu (ccyahir.kopmluuGutecaaa Franck Djeumou (djeumf2@rpi.edu)
Code

---

Published as a conference paper at ICLR 2025
NUNO: Noural SDES for UNeertaiat--awire Offlnne RL.
fo(s.a) F(s, 99,
Training
dynamies
Desotspors into vdlndle posiiion
Frioe knowledge
Data coverage informe uncertiiity
Inform structure
Physies-Infann infooreed driitt term Si..
Distance-aware uncertainty estimate
(s. a
(5,
Policy +
J(
O
Uncertainty Penalization
fi(s,a)
+ (ns
ds
fa
Tg + hol a))
dW
Origial
No
Truncate?
Uncertaintty truncation threshold based on CVait
Rollout
Figure 1: NUNC learns dynamics model neural stochastic differential equations, where its drift term can leverage prior physics knowledge as inductive bias, and its diffusion term provides distance- aware estimates of uncertainty. NUNO addresses model exploitation inherent in offline model-based RL by penalizing and adaptively truncating neural SDE's rollouts according to uncertainty estimates.
2022) A standard methodology of imposing conservatism is to dynamics models (Rigter et al.
penalize the agent with respect the predicted uncertainty of the learned model on taken transition Yu et al. (2020b); Kidambi et al. (2020); Yang et al. (2021); Zhang et al. (2023b). Given true admissible error estimator for learned dynamies, these approaches provide theoretical guarantees for lower bounds on the expected cumulative reward in the groundtruth environment (Yu et al., (020b). In practice, the standard architecture for learning dynamics models is deep probabilistic ensembles. The eгror estimators rely heuristics such maximum alcatoric uncertainty, i.e., the maximum standard deviation of learned models in the ensemble, the maximum pairwise difference between predictions of ensemble members, variance of the log-likelihood of members. (Lu al. 2021).
Inspired by Djeumou et al. (2023) that shows neural stochastic differential equa-
80
NUNO
Human normalized score
NUNO"
TATU+ MOPO
60
tions improve uncertainty estimates and prediction accuracy over probabilistic ensembles, we develop an offline model-based RL approach that leverages them: Neural Stochastic Differential Equations for UNcertainty-aware Offline RL (NJNO; see Figure IT. NUNO learns dynamics model as neural stochastic differential equations (SDE) and introduces conservatism through its uncertainty estimates. Neural SDEs consist of two main terms: drift and diffusion. A priori available physics knowledge imposes inductive biases on the drift term as a differentiable composition of separately parameterized known and unknown functions. At the same time, the diffusion term provides aleatoric and distance-awar estimate of the model uncertainty: It emulates the stochasticity of groundtruth dynamics around the training data regime while corresponding to conservative estimates uncertainty beyond the dataset. Building on Yu et al. (2020b); Zhang et al. (2023b), NUNO addresses model exploitation by penalizing control policies and truncating training rollouts according to distance-aware uncertainty estimates of neural SDEs. NUNO provides consistently high-pe perforerfo framework. especially in randomly collected datasets (see Figure 2). by exploiting neural SDEs' capability of accurate predictions over long horizons with separate and theoretically motivated quantification of aleatoric and epistemic uncertainty.
MOPO
40
20
hallcheetab-randomwz2
hopper-random-v2
Figure Comparison in random datasets of D4RL MuJoCo bench- mark Fu et al. (2020). MOPO and TATU+MOPO penalize and trun- cate, rollouts based uncertainty estimates from Gaussian ensembles, whereas NUNO achieves SOTA results in all environments via distance- aware uncertainty estimates of learned neural SDEs (see Fig. 1 NUNO predicts rewards, whereas NUNO uses the groundtruth reward function.
2

---

Published as a conference paper at ICLR 2025
Contribution. Our contribution is fhree-fold:
(1) We develop an uncertainty-awant offline model-based RL approach, NUNO, that (i) learns dynamics model as neural SDE, where the drift term leverages minimal prior physics knowledge as inductive bias, and the diffusion term provides distance-aware estimates of model uncertainty, and (ii) addresses model exploitation by penalizing and adaptively truncating synthetic rollouts based on a distance-aware uncertainty estimator.
We theoretically show that penalization via a distance-aware uncertainty estimator enables conservatism by encouraging the policy to stay close to the convex hull of the offline data. (3) In control benchmarks D4RL (Fu et al. 2020) and NeoRL (Qin et al.) 2022), NUNO imposes structure on the drift term by exploiting the fact that MuloCo environments are governed by rigid body dynamies, and decomposing state into position and velocity components.
(4) Our empirical results evidence that NUNO outperforms bate-of-the-ar methods in low- quality datasets ("random-v2' in D4RL and 'Low' in NcoRL) by up to 93% while either matching or surpassing their performance by up to 55% in high-quality counterparts.
2 RELATED WORK
Our work focuses on the intersection of offline RL and physics-informet learning of dynamics models. Appendix A investigates existing works on offline model-free RL.
The objective is to learn a dynamics model from a static dataset of environment interactions in a supervised manner and subsequently generate synthetic data to train a control policy. To tackle model exploitation, most impose conservatism by constraining the learning
Offline model-based RL:
policy to the behavior policy (Matsushima et af., 2020 Swazinna et al. 2021 Cang et al. 202 1
Bhardwaj et al. 2023). learning conservative value functions (Yu et al. 2021 Rigter et al., 2022).
learning pessimistic policies via biased sampling from a belief distribution over dynamics (Guo
сt al. (222) or via uncertainty penalization (Yu et al. 2020b Kidambi et al. 2020 Yang et al. 2021 1
Ratailov et al. |02 Zhang et al. 2023b Sun et al.] 2023). To address compounding estimation errors,
Jeong al. (2022) propose methodology that combines model-based and model-free value estimates
for policy evaluation based on their epistemic uncertainties. A recent line of works casts offline model-based RL sequence modeling problem and learns dynamics model transformer (Chen et al., 202T] Janner et al. [021 Yamagata et al. (023) without enforcing conservatism. NUNO inherits its principle of uncertainty penalization and rollout truncation from MOPO (Yu et al. 2020b) and TATU+MOPO (Zhang et al. 2023b). respectively. Both train a deep probabilistic ensemble as a dynamics model and penalize the reward based on their uncertainty estimator. TATU+MOPO extends MOPO by truncating synthetic trajectories if the accumulated uncertainty exceeds threshold based on single-step estimates on the training data. NUNO builds on them by training a neural SDE as a dynamics model to improve uncertainty estimation and prediction accuracy (Deuumou et al., 2023).
Neural differential equations for physics-informed learning: Neural ordinary differential equations (ODEs) specify structure that parameterizes differential equation via neural networks using priori known physics knowledge. Many existing works utilize neural ODE-based thysics-informer
architectures to learn dynamics models for control tasks Liu & Wang (202 Shi al. (2019); Plaza
et al. 2002)) Furieri et al. (2022); Wong et al. (2022) : Menda et al. (2019): Gupta et al. (2020g Duong
& Atanasov (2021) Lutter al. (2019). Although not for control tasks, some inform the structure of
neural ODEs via Hamiltonian Greydanus et al. 20119; Chen et al. (2019); ; Zhu et al. (2020); Zhong
al. 2020 : Eidnes et al. (2023), Port-Hamiltonian Desai al. (2021) ) Neary & Topcu| (2023).
or Lagrangian Roehrl et al. 2200)) Finzi et al. (2020); Cranmer et al. (000); Allen-
al. (2020); Zhong et al. (222bb a) formulation of dynamics. Neural bDE-d structures are
commonly determiistic and, hence do not provide a notion of uncertainty. In contrast, neural SDEs models, and previous works investigate their use for learning dynamics of stochastic systems (Jia & Benson| 2019 Yang et al. 2023). estimating the uncertainty in parameters of neural networks Kong et al. (2020); et al. (2000) : Kidger (2022); Xu et al. (2022). and generative modeling Kidger et al. (2021). ). However, these approaches do not deal with epistemic uncertainty in a way tailored to offline RL. Djeumou et al. (202)) propose using neural SDEs to leverage a priori physics knowledge and capture epistemic uncertainty to control dynamical systems and offline model-based RL. In comparison to our work, Djeumou et al. (2023) does not address the model exploitation problem and does not a loss function that properly models aleatoric uncertainty.
allow uncertainty-awar
3

---

Published as a conference paper at ICLR 2025
3 PRELIMINARIES
3.1 MARKOV DECISION PROCESSES
We formalize the environments of interest in this work Markov decision processes (MDP), specified by a tuple M = A, R, Y Pn). where S and A are state and action spaces, T : S x A -+ A(S) the transition distribution, R: state distribution, i.e., A(S). A policy : S - ((A) in an MDP M outputs a probability simplex over the action space A given state $ E S. The objective of RL is to learn an optimal policy T², which maximizes the expected discounted return in M, i.e., x² E arg max, (T), where (M) = .o. 1_ 7 Rss,, at)] is the expected discounted return, a ~ I ((:)) is the policy's action, and 81+1 ~ (8s a.) is the new state at time I. by starting in So Po-
A R. 7 E (0, 1) is the discount factor, and po is the initial
3.2 OFFLINE MODEL-BASED RL
The offline RL problem assumes access dataset D {ri}e ; of interactions == {(s. a, r, (t) with the environment M. Multiple behavior policies The optimal or suboptimal, can contribute to D. The objective is to learn a policy æ that minimizes the sub-optimality gap, namely, MM ( Offline model-base RL methods approach this problem by first learning a dynamics model 7 from the dataset D. Then, they utilize the learned dynamies model optimize the policy. Depending on the access, one can learn a reward function R or an initial state distribution Po from the dataset D.
A naive way to optimize the policy is to interact with the learned MDP, c.g., as online RL algorithms. However, such an approach can cause model exploitation, i.c., the estimated return in the learned environment is greater than the true return: 72,2 (TT) T) (T) > 0 YYu et al. 2020b). Due to the finite coverage of the dataset D, the policy æ learns exploit regions of the state-action space where the epistemic uncertainty of the learned model / and the estimated return 7, (T) are high.
A common strategy to undertake model exploitation is to penalize the agent in correlation to the estimated model uncertainty, as in Model-based Offline Policy Optimization (MOPO) (Yu et al.) 2020b), which defines a pessimistic reward function: it((s, a) R(s, a) - Apenuss, a), where u(s, a) the estimation of the model uncertainty at the state-action pair (s, a) and Xpen is the regularization coefficient for the uncertainty penalty. Utilizing the pessimistic reward function, MOPO constructs a pessimistic learned MDP M (S. ... f. R, Y, Po) and modifies the policy optimization objective max, (T)) Zhang et al. (2023b) proposes Trajectory Truncation with Uncertainty, TATU, which truncates model rollouts If the accumulated uncertainty exceeds a predetermined threshold. The theoretical results follow similar line of argument in MOPO and construct a pessimistic MDP, then provide suboptimality bounds for policies learned in pessimistic MDPs. TATU's uncertainty truncation threshold depends on uncertainty estimates over the single-step transitions from the datasets, which is limitation considering that TATU trains policies with longer rollouts.
3.3 NEURAL STOCHASTIC DIFFERENTIAL EQUATIONS AS DYNAMICS MODELS
Stochastic differential equations (SDEs) offer a principled approach to modeling uncertain, real-world, and time-varying stochastic processes. Their continuous time nature and ability to encode prior physics knowledge (world models) as inductive bias make them suitable for modeling dynamical systems from data. A neural SDE is an SDE parameterized by neural networks as follows
ds = fo a) di + E (s, a) * dW,
where fo : S A A-+R" R" and Σ. : S A R x ah are the drift and diffusion terms parameterized by 0 and di W is the nuw-dimensional Wiener process, and * expresses that the SDE is either in Ito Ito et al. 195T) or Stratonovich (Stratonvich (1966) form. The reader unfamiliar with these forms
should feel free ignore the distinction (Van Kampen, 1981 Massaroli et al., 2021 Kidger 2022).
which becomes an arbitrary modeling choice when fe and E are learned.
Given dataset D of interactions of behavior policy ma in MDP M, we seek the unknown functions fe and Σg of a neural SDE that best fit the sequences of states and actions in the dataset. Specifically, we build on the framework proposed in Djeumou et al. (2023) and extend it to train neural SDEs in such a way that the diffusion term E, captures aleatoric uncertainty as well as epistemic uncertainty in the form -ware estimates of model uncertainty.
4

---

Published as a conference paper at ICLR 2025
4 NUNO
In this section, we discuss the details of NUNO's design. First, we provide insight into the use of a distance-aware uncertainty estimator and discuss our parametric estimator and its corresponding training algorithm. Then, we introduce our physics- pessccs-based neural SDE approach for modeling the dynamics of the MDP while capturing aleatoric and epistemic uncertainty. Finally, we show how our uncertainty estimator can efficiently enforce conservatism when training the RL policy. In the remainder of the paper, we assume access to the dataset D of realized acteonstions trajectories. We also assume access to the time steps At between consecutive states 5, and St+1.
4.1 DISTANCE-AWARE UNCERTAINTY ESTIMATOR
By investigating particle-based estimators of the cross entropy between the learned model's transition distribution and the unknown transition distribution, Zhang et al.] (2023a) provides a theoretical framework for characterizing model uncertainty u(s, a) as function of the distance, in the appropriate
space, between the query point (s, a) and its k-th nearest neighbor (KNN) in the dataset D. We build on this idea and propose a parametric distance-aware uncertainty estimator пe S A + IR that captures such distance to the closest k-th neighbor in the dataset without the need for KNN search. Besides bypassing intractable KNN search, our parametric estimator can be trained alongside the neural SDE model (see Section 4.2) such that the model can capture both aleatoric and epistemic uncertainty in the dynamics. The estimator smooth and differentiable and thus blends well with the requirements for numerical integration of the neural SDE model.
simple choice for 77 for which we can provide theoretical guarantees is given by
Nio argmin, B(s.a)DD [E(',, "a) =UUnifffm(((S (n a') - (1 - as))a
The optimal solution Hi of equation 2 is convex function with respect to (s, a) and is an upper bound of the distance to the state-acties centroid of the training dataset. Additionally, we have that the negative gradient v, at any point (s, a) points inside the convex hull D.
Lemma 1.
We provide the proof of Lemma] T Appendix The first property above illustrates that Ti is a suitable choice for a distance- aware uncertainty estimator. In contrast, the second property enables conservatism by suggesting that any reward penalization with 17e will encourage the policy to stay within the convex hull of the training dataset. However, the estimator Vo approximates only the distance to the centroid of the entire dataset, which may not be sufficient to accurately capture the uncertainty in the model's predictions if the geometry of the dataset has multiple clusters.
To address this limitation, we enforce additional constraints to encourage nh to cluster the dataset properly. Informally speaking, model the term T1(()) with neural networks such that when evaluated near points in the training dataset D, such term provides low values with almost-zero gradients. In contrast, it provides high but bounded values when evaluated far from the training data. Specifically, a strong property of our approach is that by sampling only locally around the training dataset,
train the parameters of Te enforce the desired distance-based properties globally. In particular, we translate the distance-aware requirement into several mathematical properties that 7) must satisfy, and propose a loss function that encourages the neural network to learn these properties.
(a) Increasing nH along state-action paths that move away from the training data. As the query point (s, a) moves away from the training data, the distance-aware term Πe should monotonically increase accordingly. Let I' be any path along which the distance from the current point to the nearest training datapoint always increases. Then, along 1, the entries of no should monotonically increase. We enforce this property via local strong convexity constraints near the training dataset. Specifically, for every state action (st, a:) DD and a fixed radius r > O, we enforce strong convexity of пe within a ball B, (s,, a.) := {s..a) [|(s. a) - (st, at)l S r) with a convexity constant μ4 >> 0. More specifically, we want to enforce that for any (s, a). (s', a') E (s,, the convexity constraint (s, s', a') 2 0 holds, where the constraint is defined as
(s, a. S.". a') a'),. = no(s', a') (s a) - V (o.a)) a) ((s, a) (s',a') - (4 a) (s', a
We parametrize a function Ha S x A + R, using a neural network to predict the strong convexity constants 14 = P (St, ax) for each (st. a:) E D instead of manually tuning them. Thus, we define
5

---

Published as a conference paper at ICLR 2025
the following loss functions to enforce the desired properties sample (st. au) E DD as follows:
2 0. if as, a') otherwise
Σ
and Ey
Σ
Hb (St, as)
(s .a) a.)
(ar,a:)E
where (((ss,,)) is Gaussian distribution with mean (st, ar) and standard deviation r, and Ey a regularization loss term that encourages high values of pa. Intuitively, such regularization ensures that the distance-aware term ne reaches its maximum value close possible to the boundaries of the training dataset, enabling dataset clustering.
(b) Zero-gre Zero-gradient and distance-aware estimate near training data. We enforce that the distance- aware term To has almost zero gradients near the training dataset such that, with the local convexity constraints, points in the dataset become local minima of ne and the negative gradient of no near cluster is directed towards the cluster. This constraint can be enforced at a sample (st, az) E 'D as
Lgadd
V (o az)l (o az)?
where the last term encourages Te to be zero when evaluated on the training data. Appendix B.2 provides insights about the distance-aware uncertainty estimator, well as toy 2-D dataset examples to demonstrate how Th efficiently clusters the training dataset to capture datapoints distance.
4.2 THYSICS-INSPIRES NEURAL SDES
We aim to learn a neural SDE's drift and diffusion terms that best fit the sequences of states and actions in the dataset D. Specifically, we first consider the following black-box neural SDE
ds = fo a) dt (s, a) + (s, a))) dW,
(5)
where we simplify the diffusion term Eg from equation I to be diagonal matrix composed of two complementary terms. The first term de : A R an unconstrained neural network that captures the aleatoric uncertainty of the dynamics, while we design the second term hg to estimate heterogeneous epistemic uncertainty in the model's ' predictions. Here h& : bounded, monotonic, and learnable transformation that ensures the diffusion term is positive and monotonically increasing in the proposed distance-aware term 7/ S x - R. In the following, use Σe to refer to 0s + (16) when the distinction is unnecessary.
Rd is
Monotonicity and boundedness of hag To ensure globally monotonic and bounded diffusion values as function of fi. we adopt a simple design choice for hg: A scaled sigmoid function to transform 77 into a heterogeneous diffusion term. Specifically, we define h& (1+) = sigmoid Wne+ + b), where sigmoid (x) (1 + exp(-x)))- W E Th and b E R are learnable parameters of the neural network. We constrain W to be greater than 1. Besides, the term Wmmax is hyperparameter that controls the desired maximum value of the diffusion term outside the training data regime. We emphasize that this design choice works well in our experiments, but others possible.
Training the neural SDE. In contrast to the standard approaches such as probabilistic ensembles where the model fits single-step transition, the proposed neural SDE is designed and trained to fit sequences of states and the uncertainty in the model's predictions. Given a sequence of states and actions (st, ag,..., .11 with H being the prediction horizon, we aim to minimize the negative log-likelln big-liieldhood (NLL.) of the sequence under the neural SDE model. However, estimating the NLL of neural SDE-generated sequences challenging due to the intractability of computing the kernel density of the underlying stochastic process. To address this issue, we adopt numerical integration schemes to approximate the sequence's 's NLL through Monte Carlo sampling. Specifically, assuming approximate Gaussian transitions between discrete time steps of the stochastic process, e.g., when employing the Suler-Maruyama sampler, we can approximate the NLL as
Ldata Egg
I+H-1
||s
28,0 5K41 (x)+)) + log(det(Ex)))
(6)
where (Eo) (s ax), and 8 $/41.... H44 the sample states obtained by any differential SDE numerical integration scheme. Note that the accuracy the NLL approximation depends on the quality of the numerical integration scheme. the stepsize discretize the SDE between two
0.0 0.0

---

Published as a conference paper at ICLR 2025
consecutive states, and the number of samples used to estimate the expectation. In practice, though, we can fit accurate neural SDE models to the data even with galer-Maruyama and a single sample.
The problem of learning the SDE model parameters with distance-aware uncertainty estimates can be formulated as the following optimization problem: (d Cdata Xschae + Agrad Egrad Xy
minimize 8.6
Incorporating prior physics knowledge. We can incorporate prior physics knowledge into the neural SIDE model by designing the drift term So to encode structural knowledge from first principles or domain expertise. To this end, we represent the drift term fe as the composition of a known function - derived from priori knowledge and collection of unknown functions that must be learned from data. That is, we write fo a) == F(s, 9o, (-).....oo.) (od(-)) where F is known differentiable function and go, (-)..... .oo..).god()) are unknown terms within the underlying model. The imputs to these functions could themselves be arbitrary functions of the states and control inputs. Additionally, known constraints on go; can be enforced during training using the augmented Lagrangian method.
We exploit the fact that rigid body dynamics govern our benchmark environments to constrain the structure on the drift term. We typically decompose the state as $ = [Spus: Syel], where Spas and Syel are the position and velocity components, respectively, and we define the drift term
f (s. a) = Svel. (s, a) Go((vel) a + Ha(S)Seel, fe = [ fo
f0
(8)
where Go and Ho are learnable neural networks, and f& is a residual term capturing unmodeled dynamics. We penalize the residual term to ensure minimal deviation from the drift term. Note that this formulation integrates minimal prior knowledge into the neural SDE model, and such knowledge does not affect modeling performance in the large dataset regime seen in our experiments. For discussion on general prior physics knowledge, we refer the reader to Djeumou et al. (0023).
Incorporating reward learning. We can incorporate reward learning into the neural SDE model by augmenting the state representation with variable representing cumulative rewards. Specifically, we define the new state as s = [Spos. Svel, re]. where тe is the cumulative reward up to the current time step. We then augment the neural SDE model with dr. fg (s. a) dt, where frew is a learnable neural network that captures the reward dynamies. We can then train the combined neural SDE model to minimize the NLL of the sequence of states, actions, and rewards under the model.
4.3 DISTANCE-ANI AWARE REGULARIZED OFFLINE RL
We now discuss incorporating the distance-aware uncertainty estimate To into the offline RI. frame- work to enforce conservatism in the learned policy. Specifically, we build on the work by Zhang ctt al. 2003bb and use our distance-aware uncertainty estimate to penalize and truncate the transitions generated by the learned neural SDE model during the RL policy training.
Reward penalty, Following MOPO penalization criteria, we use the distance-aware uncertainty to define the pessimistic reward as R(s, a) = R(s, a) - 7pp a).
Trajectory truncation. During the RL agent training, use the current policy and the neural SDE model to generate synthetic trajectories for policy improvement. To figure out whether the synthetic trajectory is reliable, we set truncating threshold
the accumulated distance-aware estimate No over the sequence. Specifically, we compute T Ei= [s az), and we compare its value with the threshold e. If the accumulated quantity exceeds the threshold, we truncate the trajectory and do not use it for policy optimization. The choice of the threshold is a crucial hyperparameter that varies accross environments or tasks while enforcing the level of conservatism in the policy training. To account for different task and environment complexities, we propose automatically setting the threshold based on the entire training dataset. We propose to use user-defined Conditional Value-at atee-at-Ree -isk (CVaR) as the threshold to compute the hyperparameter e via performing statistics on the entire dataset over all possible sequences of horizon h.
5 EXPERIMENTAL RESULTS
We empirically evaluate NUNO against state-of-the the-thh-rtt art (SOTA) offline model-based and model free approaches in continuous control benchmarks, namely MuJoCo datasets in D4RL. |u et al. 2020))
7

---

Published as a conference paper at ICLR 2025
Table Average human-normalized scores of NUNO and other model-based and model-free offline RL approaches on D4RL MuJoCo v2 datasets. Due to limited space, we use abbreviations of task and dataset names: he halfcheetah, hp hopper, wk walker2d; r = random, m medium, mr medium-replay, me medium-expert For NUNC, we provide the mean and standard deviation (following +) best scores among independent runs. Bold scores indicate the best for each task.
Task
NUNO (Ours)
NUNO" (Ouns)
MOBILE
MOPO
MОPO
COMBO
MORET
RAMBO
EDAC
ho-r
52. 2
52. 3...5
39.313.0
33.3
35.9
38.8
38.9
39.5
28.4
hp-r
73 7..2
53.7++13.9
31.9+0.8
31.9
16.7
17.9
38.1
25.4
25.3
27.720.9
28.1:1.2
17.9+6.6
10.4
0.0
he-m
68 68.840.4
64.720.5
74.6#1.2
61.9
73.1
54.2
60.7
77.9
65.9
hp-m
104.6:20.2
104.4:0.5
00..6
104.3
38.3
97.2
84.0
87.0
101.6
wk-m
85.1400.9
92.61:13
87.7
77.9
11.2
81.9
72.8
81.9
92.5
he-mr
66.510.2
64.6220.3
71.71
67.2
69.2
55.1
44.5
68.7
61.3
hp-mr
107.8:1.2
106. 6 1,9
103.9:1.0 1.0
104.4
32.7
89.5
81.8
99.5
101.0
wk-mr
97.0+11.4
101.1:3.9
8..9
75.3
73.7
56.0
40.8
89.2
87.1
97.0+0.5
95.8+1.2
222.
95.4
hp-me
112.2:0.3
111.9=0.5
112.6:0.2
107.0
60.6
111.1
105.6
88.2
110.7
wk me 113.2=0.5
112.620.6
115.2=0.7
107.9
77.4
103.3
107.5
56.7
114.7
Average 83.8
82.4
80.0
71.3
49.4
66.8
64.3
67.7
76.0
and NeoRL Qin et al. (2022). Through our empirical evaluation, we answer the following questions: 1) How does NUNO perform in terms of human normalized score? 2) Can NUNO'S uncertainty estimator, i.e., distance-aware estimate of neural SDE, effectively quantify uncertainty? 3) How does NUNO address the model exploitation phenomenon in contrast MOPU+MOPO and MOPO?
5.1 How DOES NUNO PERFORM IN STANDARD CONTROL BENCHMARKS?
5.1.1 D4RL
We run experiments on 12 D4RL tasks, combining three MuloCo environments (halfcheetah, hop- per, and walker2d) and four datasets (random, medium, medium-replay, and medium-expert) per environment. We compare NUNO against a model-free method called EDAC An et al.) (202T), that penalizes Q-values based on the estimated uncertainty of a Q-function ensemble; and model-based methods: МOPO Yu et al. (2020b) and TATU+MOPO Zhang et al. (2023b). from which NUNO inherits uncertainty penalization and truncation, respectively, COMBO Kumar et al. (202)), which equally penalizes samples that are out-of of-distribution according to model uncertainty, MOBILE Sun al. (023), which penalizes the Bellman estimation based the inconsistency of Bellman estimations by an ensemble of learned dynamics models, RAMBO that adversarially learns a policy and dynamics model, and finally, MOREL Kidambi et al. (2020), which penalizes a transition when refer to TATU+MOPO MOPOT
estimated uncertainty exceeds threshold. In Table
Table demonstrates the mean and standard deviation of maximum human-normalizz scores that NUNO and NUNOR, which predicts the reward, reach in D4RL MuJoCo tasks (v2) during five independent runs of one million gradient steps. In the random task involving datasets collected by randomly initialized policies, NUNO and NUNO outperform all approaches by a significant margin. NUNO achieves this by building onto uncertainty penalization and truncation principles proposed by MOPO and TATU+MOPO. NUNC's advantage comes from leveraging minimal prior physics knowledge and exploiting the diffusion term's 's capability of estimating alcatoric and epistemic uncertainty. Given higher-quality datasets, namely, netter-perffrming data logging policies, NUNO either reaches SOTA results or closely follows existing ones. Specifically. NUNO reaches of the-art results in
and hopper-mediium-epper v2. Table [ also demonsirates that NUNO achieves the second-best results in medium tasks of halfcheetah and hopper, as well as medium-expern task of hopper. Overall, NUNO and NUNO yield the highest average human normaalied scores in the D4RL MuJoCo benchmark. Figure visualize the progression of human normalized score for NUNO and NUNO See Appendix E.2 for more details.
malker-medium--eplay
8

---

Published as a conference paper at ICLR 2025
Table Average human-normalized scores of NUNO and other model-based and model-free offline RL approaches on NeoRL MuJoCo datasets. Due to limited space, we use abbreviations of dataset names: L low, M medium, H = high. For NUNO, we provide the mean and standard deviation (following +) of best scores among independent runs. Bold scores indicate the best for each task.
Task
(Ours)
NUNO" (Ours)
MOBILE
MOPO
BC
CQL
TD3+BC
EDAC
he-L
02.5
58. 2...5
54.743.0
40.1
29.1
38. 2
30.0
31.3
bp-L. 26.943.8
26.4+6.8
17.4+3.9
6.2
15.1
16.0
15.8
18.3
wk-1.
52.5:2.4
49.411.9
37.622.0
11.6
28.5
44.7
43.0
40.2
he-M
73.4=0.6
78.8
77.8+1.4
62.3
49.0
54.6
52.3
54.9
hp-M
103.3:2.2
92.3=1.7
51.11113.3
51.3
64.5
70.3
44.9
wk-M
65.8+-0.4
49.4+11.9
62.241.6
39.9
48.7
37.3
58.5
57.6
he-H
85 2.10..6
84.910.4
83 6
65.9
71.3
77.4
75.3
81.4
hp-H
103.0:38
97.915.5
87.842260
11.5
43.1
76.6
75.3
52.5
wk-H
72. 1.6
74.5111.
74.913.4
18.0
72.0
75.3
69.6
75.5
Average
70.6
68
60.7
28.5
45.4
56.1
54.5
50.7
[tuman aormelizdd acare
0.4
0.4
75-
Ussertainty
0.3
Uscirtaiity
0.3
progross
0.2
0.2
25-
NUNO
NUNO
0.1
0
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number of epochs
Number of epocts
Number of spochas
(a) Human normalized score
(b) Uncertainty-NUNO
(c) Uncertainty-NUNO"
(a) We report the progression of human normalized score in evaluation episodes during training. (b-c) We demonstrate how the uncertainty estimates of neural SDEs in NUNO and NUNO evolve when evaluated with trained policies' actions in one-step rollouts from states in the dataset. 'random' and 'data' refer to the uncertainty estimates of the learned model given actions from random policy and the dataset, respectively.
Figure 3: Training progression in hopper-medinee
5.1.2 NEORL
We further evaluate NUNO in NeoRL al. (2022), benchmark developed to reflect real-world characteristics by logging data via conservative policies. We investigate nine datasets involving three environments (HaliCheetah--v,, Hopper-v3, Walker2d-v3) and three types of datasets (low, medium, high) per environment with 1000 trajectories each. We compare NUNO against MOBILE, MОPO, EDAC, CQL which penalizes OOD samples' Q-values equally, behavior cloning (BC), which imitates data-logging policies, and TD3+BC Fujimoto & Gu (2021). which extends TD3 Fujimoto et al. (0018) by regularizing the policy optimization objective via behavioral cloning term.
Table 2 reports the mean and standard deviation of maximum reaches in NeoRL MuloCo tasks during four independent runs of one million gradient steps. NUNO achieves the highest scores in the low tasks of NeoRL by outperforming existing SOTA results in hopper and walker2d by a significant margin, as in the random tasks D4RL. In addition, NUNO or NUNO8 reach the highest scores in medium and high tasks of all MuJoCo environments. Overall, NUNO and N collect the highest average human normalized scores across nine tasks in NeoRL.
scores that NUNO
5.2 CAN NUNO'S UNCERTAINTY ESTIMATOR EFFECTIVELY QUANTIFY UNCERTAINTY?
and 3c demonstrate the evolution of uncertainty estimates of trained neural SDEs for NUNO and NUNO during training. Neural SDE's uncertainty estimator assigns the largest values to random actions and the smallest to the dataset, evidencing that the uncertainty estimators correctly identify out-of-distribution and in-distribution actions, respectively. As trained policies progress, see Figure Bal the model uncertainty for learned policies' actions approaches the uncertainty of in-distribution samples because learned policies avoid ofitt--diitributiin actions through penalization pprovidde
Figures 3D
and truncation based on distance-ave aware uncertainty estimates. Appendix[D] provides an ablation study on the choice of the uncertainty estimator for penalization and truncation policy learning.

---

Published as a conference paper at ICLR 2025
Pooman normalond scurs
Human mormell alland scrves
Fair
Pate
NUNOF
NUNOP
MOPS
(a) Random
(b) Medium-replay
Figure 4: Model exploitation: We evaluate NUNO, NUNOR, MOPU and MOPO in rollouts from their learned dynamics models in (a) random and (b) medium-reple tasks, and report the average score per step with (pessimistic, Pess) and without (groundtruth, GT) uncertainty penalization.
Predietton Brvor
109.
walkerdd experss22
walkorrad modtram replaying
wallberidd medram vi
Prediction Eeror
10
wallearii modians explayend
Mallerrdd random
asde mooombll
1o>
wrallcrrdd modians cogeett-i
to+
I0"
Io:
109
10a
Horizoe
Hortcon
(a) D4RL Walker2d: In-distribution
(b) D4RL Walker2d: Out dfittr-iiittiutttin
Figure 5: Model analysis: We illustrate the evolution of model prediction error in different datasets
for D4RL Walker2d. (a) In-distribution: Evaluation of the datasets in which the models are trained. (b) Out-of-distribution: Evaluation of models, trained via random, in trajectories from other datasets.
5.3 HOw DOES NUNO ADDRESS THE MODEL EXPLOITATION PIENOMENNO?
We assess how NUNO addresses the model exploitation phenomenon based on two aspects: (1) con- servativeness of the reward function of pessimistic learned MDPs, and (2) prediction accuracy of learned dynamics models. Figures and 5 evidence that NONO enables less conservativeness and better accuracy over longer horizons. Figure addresses the first aspect in two sets of D4RL tasks: random and medium-replay. Based on the gap between the groundtruth score and the pessimistic score, we observe that NUNO and NUNO construct pessimistic learned MDPs that are less conser- vative than their counterparts MOPO and TATU+MOPO, which use Gaussian ensembles. The only exception is sopper-medium-replayy which may be why TATU+MOPO and MOPO perform slightly better, as reported in Table Model accuracy results in Figure show that neural SDEs are significantly more accurate than Gaussian ensemble over longer horizons.
CONCLUSION
We develop a novel uncertainty- offline model-based RL algorithm, NUNO, that learns a single dynamics model, in contrast to probabilistic ensembles in most existing work, neural SDE and addresses model exploitation phenomenon by penalizing and adaptively truncating model rollouts based on its uncertainty estimates. NUNO achieves this by imposing minimal prior physics knowledge into the drift term of neural SDE as inductive bias and learning distance-aware uncertainty estimates via its diffusion term, which matches the dynamics' underlying stochasticity around the training data regime while providing high but bounded estimates beyond it. Through our empirical evaluations of NUNO in these benchmarks, we demonstrate that NUNO outperforms state-of of-the--rt methods, particularly in low-quality datasets with low state-action space coverage suboptimal demonstrations ("random-v2) in D4RL. and 'low' in NeoRL.) by up 93%. tasks involving higher quality datasets, NUNO matches exceeds the state-of-the-a art performances in environments by up to 55%.
Limitations and future work. Although we can extend our formulation to address partially observed Markov decision processes, experiments utilize full knowledge of the system state in MuJoCo environments from both benchmarks. In the future, we aim to extend our uncertainty-awar approach to address different settings, e.g., environments with image observations. Additionally, future work can investigate formally proving properties our distance-aware uncertainty estimator and extend formulation for non-Fuclidean state-action spaces by adjusting the distance metrics accordingly.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
For the theoretical analysis of this work, we state all assumptions made in Section 4 and Appendix B For the hyperparameters and detailed settings of the experiments, please refer to Appendix Lastly, put the core code of our approach in the supplementary details. The code includes dataloaders, execution code, and links to download the datasets and models used.
ACKNOWLEDGMENTS
This work is supported by the Air Force Office of Scientific Research (AFOSR) under the grant FA9550-19-1-0005 and the National Science Foundation (NSF) under grant number 2214939.
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In ICML, pp. 114. PMLR, 2020.
Christine Allen-BAl Mlen-Blanchette, Sushant Veer, Anirudha Majumdar, and Naomi Ehrich Leonard. Lagnetvip: A lagrangian neural network for video prediction. arXiv preprint arXiv:2010.12932. 2020.
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-bau offline reinforcement learning with diversified g-ensemble. NeurIPS, 34:7436-7447, 2021.
Chenjia Bai, Lingxiao Wang. Zhuoran Yang. Zhi-Hong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In ICLR, 2022. httt://s//ppeneer..eeeeeeeeroooe
Mohak Bhardwaj. Tengyang Xie, Byron Boots, Nan Jiang, and Ching-An Cheng. Adversarial model for offline reinforcement learning, arXiv preprint arXiv:2302.11044 2023.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
/tttp://gitee ccw./ooolee/iiiocoooooooo
Catherine Cang, Aravind Rajeswaran, Pieter Abbeel, and Michael Laskin. Behavioral priors and dynamics models: Improving performance and domain transfer in offline rl. arXiv preprint arXiv:2106.09119, 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch, Decision transformer: Reinforcement learning via sequence modeling. NeurIPS, 34: 14:15084-150979 2021.
Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neural networks. arXiv preprint 19Xxv:1909.13:: 2019.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.044 04630, 2020.
Shaan A Desai, Marios Mattheakis, David Sondak, Pavlos Protopapas, and Stephen J Roberts. Port- hamiltonian neural networks for learning explicit time-dependent dynamical systems. Physical Review E, 104(3):034312, 2021.
Franck Djeumou, Cyrus Neary, and Ufuk Topcu. How to learn and generalize from three minutes of data: Physics-constrained and uncertainty-awant neural stochastic differential equations. CoRL, 2023.
Thai Duong and Nikolay Atanasov. Hamiltonian baaste neural ode networks on the se (3) manifold for dynamics learning and control. arXiv preprint arXiv:2/06, /2782, 2021.
11

---

Published as a conference paper at ICLR 2025
Sølve Eidnes, Alexander J Stasik, Camilla Sterud, Eivind Bøhn, and Signe -erner-Srrensn.. Pseudo- hamiltonian neural networks with state-dependent external forces. Physica D: Nonlinear Phenom- ena, 1446:13363,, 2023.
Mare Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying hamiltonian and lagrangian neural networks via explicit constraints. NeurIPS, 33:13880-1388, 2020.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep
data-driven reinforcement learning. arXiv preprini arXiv:2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. NeurIPS, 34:20132-20145,, 2021.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In ICML, pp. 1587-1596. PMLR, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML, pp. 2052-2062. PMLR, 2019.
Luca Furieri, Clara Lucfa Galimberti, Muhammad Zakwan, and Giancarlo Ferrari-Trecate. Distributed neural network control with dependability guarantees: a compositional port-hamiltonian approach. In L4DC. pp. 57 1-583, PMLR, 2022.
Carles Gelada and Mare Bellemare, Off-policy deep reinforcement learning by bootstrapping the covariate shift. In AAAI, volume 33, pp. 3647-3655, 2019.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. NeurIPS, 32, 2019.
Kaiyang Guo, Shao Yunfeng, and Yanhui Geng. Model-based offline reinforcement learning with dynamics belief. NeurlPS, 35:449-461 2022.
Jayesh K Gupta, Kunal Menda, Zachary Manchester, and Mykel Kochenderfer. Structured mechanical models for robot learning and control. In L4DC, pp. 328-33 337. PMLR, 2020.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICML, pp. 1861-1870. PMLR, 2018.
Kiyosi Ito, Kiyosi Itô, Kiyosi Itô, Japon Mathématicien, Kiyosi Itô, and Japan Mathematician. On stochastic differential equations, volume 4. American Mathematical Society New York, 1951.
Michael Janner. Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. NeurIPS, 32, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. NeurIPS, 34: 14:1127-1 2021.
Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Way off off-policy batch deep reinforcement learning of implicit human preferences in dialog. arXiv preprint urXiv:/901) 004.000456 2019.
Jihwan Jeong, Xiaoyu Wang, Michael Gimelfarb, Hyunwoo Kim, Scott Sanner, et al. Conservative bayesian model-based value expansion for offline policy optimization. In ICLR, 2022.
Junteng Jia and Austin Benson. Neural jump stochastic differential equations. NeurIPS, 32, 2019.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. NeurIPS, 33:21810-21822 2020.
Patrick Kidger. On neural differential equations. arXiv preprint arXiv:2202.02453 2022.
Patrick Kidger, James Foster, Xuechen Li, and Terry Lyons. Neural sdes as dimnite-dimessional gans. In ICML, pp. -453-5463. PMLR, 2021.
12

---

Published as a conference paper at ICLR 2025
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Lingkai Kong. Jimeng Sun, and Chao Zhang. Sde-net: Equipping deep neural networks with uncertainty estimates. arXiv preprint urXiv:2008. 10546, 2020.
Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In ICML, pp. 5774-7788 PMLR, 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. NeurfPS, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. NeurlPS, 33:1 17911191, 2020.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv:1802. 10592, 2018.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-ar pp. 45-73. Springer, 2012.
Sergey Levine, Peter Pastor, Alex Krizhevsky, Julian Iharz, and Deirdre Quillen. Learning hand- eye coordination for robotic grasping with deep learning and large-scale data collection. The International journal of robotics research, 37(4-5):421-436 2018.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement Iearning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.00 00.400005./64, 2020.
Xuechen Li, Kam Leonard Wong, Ricky TQ Chen, and David K Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In Symposium Advances in Approximate Bayesian Inference, pp. 1-28. PMLR, 2020.
Jeremiah I. Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance aware- ness. NeurIPS, 13:7498-7512, 2020.
Xin- Sin-Yang/aa Liu and Jian-Xun Wang. Physics-informed dyna-style model-based deep reinforcement learning for dynamic control. Proceedings of the Royal Society A, 477(2255):202166) 2021.
Cong Lu, Philip Ball, Jack Parker-Holder, Michael A Osborne, and Stephen J Roberts. Revisiting design choices in offline model-based reinforcement learning. arXiv preprint arX:v:2110.04135, 2021.
Michael Lutter, Christian Ritter, and Jan Peters. Deep lagrangian networks: Using physics as model prior for deep learning. arXiv preprint arXiv: (900:0407999999 2019.
Stefano Massaroli, Michael Poli, Stefano Peluchetti, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Learning stochastic optimal policies via gradient descent. IEEE Control Systems Letters, 6:1094-1099, 2021.
Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment- efficient reinforcement learning via model-based offline optimization. arXiv preprint 203Y::2006.03647, 2020.
Kunal Menda. Jayesh K Gupta, Zachary Manchester, and Mykel Kochenderfer. Structured me- chanical models for efficient reinforcement learning. In Workshop on Structure and Priors in Reinforcement Learning, ICLR, volume 11, 2019.
Ofir Nachum, Yinlam Chow, Bo Dai. and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. NeurIPS, 32, 2019a.
Ofir Nachum, Bo Dai, llya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Algaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.024 0200444 2019b.
13

---

Published as a conference paper at ICLR 2025
Cyrus Neary and Ufuk Topcu. Compositional learning of dynamical system models using port- hamiltonian neural networks. In L4DC, pp. 679-691. PMLR, 2023.
Santiago Sanchez-Esealonila Plaza, Rodolfo Reyes- Bacz, and Bayu Jayawardhana. Total energy shaping with neural interconnection and damping assignment-passivii based control. In L4DC, pp. 520-531. PMLR, 2022.
Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy temporal-diffcen -ifferenne learning with function approximation. In ICML., pp. 417-424, 2001.
Rong-Jun Qin, Xingyuan Zhang, Songyi Gao, Xiong-Hui Chen, Zewen Li, Weinan Zhang, and Yang Yu. Neorl: A near real-world benchmark for offline reinforcement learning. NeurIPS, 35:
24753-24765, 2022.
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In L4DC, pp. 1154-1168. PMLR, 2021.
Paria Rashidinejad, Hanlin Zhu, Kunhe Yang, Stuart Russell, and Jiantao Jiao. Optimal conservative offline RL with general function approximation via augmented lagrangian. In ICLR, 2023. URL
httts:://ppenreval
net. mett/orrumid=-svvtitte
Mare Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. NeurlPS, 1582--66776,7 2022.
Manuel A Roehrl, Thomas A Runkler, Veronika Brandtstetter, Michel Tokic, and Stefan Obermayer. Modeling system dynamics with physics- informed neural networks based on lagrangian mechanics. PFne--apersnnnl.ine 53(2):915592200, 2020.
Guanya Shi, Xichen Shi, Michael O'Connell, Rose Yu, Kamyar Azizzadenesheli, Animashree Anandkumar, Yisong Yue, and Joon-Jo Chung. Neural lander: Stable drone landing control using learned dynamics. In ICRA, pp. 9784-9790. IEEE, 2019.
RL196814 Stratonovich. A new representation for stochastic integrals and equations. SIAM Journal on Control, 4(2):362-371, 1966.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-bellman inconsistency for model-based offline reinforcement learning. In ICMI., pp. 33177-33194, PMLR, 2023.
Richard S Sutton. A Rupam Mahmood, and Martha White. An emphatic approach to the problem of off-po temporal-difference learning. JMLR, (711):263326311 2016.
Phillip Swazinna, Steffen Udluft, and Thomas Runkler. Overcoming model bias for robust offline deep reinforcement learning. Engineering Applications of Artificial Intelligence, 104:104366, 2021.
Huan-Hsin Tseng, Yi Luo, Sunan Cui, Jen-Tzung Chien, Randall K Ten Haken, and Issam El Naqa. Deep reinforcement learning for automated radiation adaptation in lung cancer. Medical physics, 44(12):6690-6705, 2017.
Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. Uncertainty estimation using single deep deterministic neural network. In ICML, pp. 9690-9700. PMLR, 2020.
Nicolaas G Van Kampen. Itō versus stratonovich. Journal of Statistical Physics, 14:175-187, 1981.
Lu Wang, Wei Zhang, Xiaofeng He, and Hongyuan Zha. Supervised reinforcement learning with recurrent neural network for dynamic treatment recommendation. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 2447-244445 2018.
Josiah Wong, Viktor Makoviychuk, Anima Anandkumar, and Yuke Zhu. Oscar: Data-driven operational space control for adaptive and robust robot manipulation. In ICRA, pp. 10519-10526. IEEE, 2022.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning.
arXiv preprint arXiv: / 1361, 2019.
14

---

Published as a conference paper at ICLR 2025
Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua M Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. In ICML, 2021. URL http:t//ppcc/cee/nngg
http://wlll.... mlr. wwwppppsssse.... wttees///13//mu2l.ktttt html
Winnie Xu, Ricky TQ Chen, Xuechen Li, and David Duvenaud. Infinitely deep bayesian neural networks with stochastic differential equations. In AISTATS, pp. 721-738. PMLR, 2022.
Taku Yamagata, Ahmed Khalil, and Raul Mantos-Rodriguez Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline In ICML, pp. 38989-39007. PMLR, 2023.
Luxuan Yang. Ting Gao, Yubin Lu, Jinqiao Duan, and Tao Liu. Neural network stochastic differential equation models with applications to financial data forecasting. Applied Mathematical Modelling, 115:279-2999 29927992999220 2023.
Yijun Yang. Jing Jiang, Tianyi Zhou, Jie Ma, and Yuhui Shi. Pareto policy pool for model-based offline reinforcement learning. In International Conference on Learning Representations, 2021.
Fisher Yu, Haofeng Chen, Xin Wang. Wenqi Xian, Yingying Chen, Fangchen Liu, Vashisht Madhavan, and Trevor Darrell. Bdd100k: diverse driving dataset for heterogeneous multitask learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 2636-2645 2020a.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. NeurlPS, 33:14129-14142, 2020b.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. NearIPS, 34:28954-28961 2021.
Hongchang Zhang, Jianzhun Shao, Shuncheng He, Yuhang Jiang, and Xiangyang Ji. Darl: distance- aware uncertainty estimation for offline reinforcement learning. In AAAL, volume 37, pp. 11210- 11218, 2023a.
Junjie Zhang, Jiafei Lyu, Xiaoteng Ma, Jiangpeng Yan, Jun Yang, Le Wan, and Xiu Li. Uncertainty- driven trajectory truncation for data augmentation in offline reinforcement learning. In ECAI 2023, pp.
IOS Press, 2023b.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Dissipative symoden: En coding hamiltonian dynamics with dissipation and control into deep learning. arXiv preprint arXiv:2002.08880, 2020.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Benchmarking energy-conserving neural networks for learning dynamics from data. In L4DC, pp. 1218-1229. PMLR, 2021a.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Extending lagrangian and hamilto- nian neural networks with differentiable contact models. NeurIPS. 34:21910-21922, 2021b.
Aiqing Zhu, Pengzhan Jin, and Yifa Tang. Deep hamiltonian networks based on symplectic integrators. arXiv preprint 2004 13830, 2020.
15

---

Published as a conference paper at ICLR 2025
A RELATED WORK EXTENSION
Model-free RL algorithms in the offline setting aim to learn an optimal policy within the available data coverage without learning dynamics model. Existing methods focus on learning policies that stay close to the data logging policy and avoid out-of-distribution actions by constraining the policy explicitly to the data logging policy Wu et al. (2019); Fujimoto et al. [ & Gu (2021), importance sampling Precup et al. : Sutton et al. (2016);
Offliine model-free RL:
Fujimoto
Gelada & Bellemare (2019); Nachum et al. (2019a); Rashidinejad al. (023). learning conservative
value functions Nachum et al. (2019b) Kumar et al. (2020): Kostrikov et al. 2021). and uncertainty
quantification Kumar al. (2019); Agarwal al. (2020); An al. (2021 Wu et al. (202T); Bai
et al. (2022). Although they refrain from the computational expense of learning a dynamics model,
model-free approaches commonly struggle when the data logging policy is sub-optimal because optimal actions become out-of-distribttion.
B DISTANCE-AWART UNCERTAINTY ESTIMATOR
B. PROOF OF LEMMA
This section provides results supporting the claim of Lemma To this end, we analyze the extrema of the optimization problem
argmin, (ss,a)< [Es ( Unittorrss (n(s,, a') (l a) (s', a')1
(9)
Under the assumption that S x A is compact, reformulate the optimization problem according to Fubini's theorem as
min, E(sr, ((stta))~ualfor((xx/ [B(a)) in(s', a') - (| a) (s', a'),,a))11 =
(10)
Let (s, and z' = , We can rewrite the objective function
J(((z)) SxA S Eu- || z'111² dz.
(11)
The extrema of the objective function are solutions of
aJ(n(=')) =0
(12)
0n(22)
On(2²) (BznD(I( (n(22) -121 2'113) =
(13)
EEND ((((=)) - 1|2- 2/11] = 0.
(14)
Thus, by expanding the expectation, we have
Ex-D (M( (2')] - EEND |||= - 2 =0
(15)
(>²) E EEzpDll [||= =/1
(16)
We can then conclude that the optimal solution is a convex function since a linear combination of convex functions.
Additionally, we have through Jensen's 's inequality that
((=)) = E.zDDll |||= - ='1l] 2 [|E=~D[z] 2'11 = || - z'1l = 0,
(17)
where 2p = E.>DD[R [z] is the state-action centroid of the dataset. Thus, the first property of Lemma I is proven.
Finally, let's prove that the negative of the gradient points inside the convex hull. By lincarity of the gradient, we have
Vn((z' = EZ-D ||
(18)
This implies that for any point z' that lies outside of the training dataset, the negative of the gradient z'=a that points inside the convex hull of the dataset. -
is non-negative combination of vectors
This concludes the proof of Lemmal]
16

---

Published as a conference paper at ICLR 2025
B.2 ILLUSTRATION OF THE DISTANCEAAARR TERM ON TOY 2-D EXAMPLES
Distance-aware estimator as model epistemic uncertainty estimator. Most offline model-based reinforcement learning approaches employ either Monte Carlo (MC) dropout or model ensemble for epistemic uncertainty estimation. Although such approaches have demonstrated incredible results, Liu al. (2020); Van Amersfoort et al. (2020) show that MC Dropout or model ensembles are unaware of the distance between unseen samples and training datasets, even in simple toy examples. Besides, these uncertainty estimators are parametric models targeted for reconstruction or regression objectives solely based on in in-distribution data rather than directly tasked for uncertainty estimation. Therefore, they might discard relevant information, such as the distances between different samples or distances to out-of-sample data. We also refer to Figure from Zhang et al. (023a), where the authors demonstrate how these methods could not yield accurate distance- based uncertainty estimates.
Instead, using distance- uncertainty estimator preserves the data's mutual relations while providing the ability to detect out-of-sample data. Besides, we can theoretically relate the problem of estimating the cross entropy between learned model dynamics and ground truth dynamics to calculating data point distances to a k-nearest neighbor clustering of the training dataset. Such cross-ents cross-entropy is erucial for enforcing pessimism when training offiin RL policies and for providing tight performance bounds. One of our goals is to design parametric distance-aware uncertainty estimator that can efficiently cluster the dataset without performing k-nearest neighbor clustering and that can directly be embedded into the neural SDE formulation.
Our term The accurately provides distance-based uncertainty estimates. We seek demonstrate that the loss functions in equation 3 and equation 4 are sound and, upon convergence, provide distance-basee uncertainty estimate term T7 that can efficiently cluster the training dataset. To this end, we generate three datasets of a two-dimensional system as illustrated in Figure O Our approach can cluster the training dataset all examples while providing clear delimitation between in-sample and out-of-sample data points.
Dateset
10
Detaset
1.0
Dataset
1.0
03
06
0.
0.8
0.4
0.4
0.
Figure 6: Visualization of the distance-aware uncertainty estimate The on three generated datasets. The red points represent the state-action samples in the dataset. Yellow indicates high uncertainty, while dark blue represents low uncertainty. X and y-axis denote the states of the system, respectively.
C EXPERIMENTAL DETAILS
C.1 BENCHMARKS
We empirically evaluate NUNO in two continuous control benchmarks: D4RL and NeoRL. We utilize three MuJoCo environments from both benchmarks: halfcheetah, hopper, and walker2d. D4RL, each environment comes with four types of datasets: (1) random-v2, where randomly initialized policy collects the samples; (2) medium-v2, where an early-stopped policy trained via SAC Haarnoja et al. (0018) for one million steps is the data-logging policy; (3) medium-replay-v2. where the datasets comprises of the samples from the buffer of the early-stopped policy used for medium-v22 from medium-level policy and the other
(4) medium-eexpert
where half of the samples from an expert one. Our experiments use the v2 version of D4RL datasets. We report the results of MOBILE, RAMBO, and EDAC from their original papers, the experiments were v2 datasets.
17

---

Published as a conference paper at ICLR 2025
For the rest, we provide the scores reported in TATU+MOPO paper (hhag et al. 2023b), codebase is based on theirs. We note that MOBILE and EDAC train four independent runs for three million gradient steps, and RAMBO reports five runs of two million gradient steps, in contrast to five runs of one million steps for the rest.
In comparison, the NeoRL (Near real-world offline RL) benchmark consists of datasets collected by policies with validated performance. More specifically, NeoRL trains a policy via SAC until convergence and uses several checkpoints from the training to collect data. These checkpoints correspond to policies with three levels of sub-optimality: 25%, 50%, and 75% of expert returns, which NeoRL calls low, medium, and high. The datasets we investigate consist of 1000 trajectories. We report the results of BC, CQL, and MOPO from the paper proposing NeoRL (Qin et al. (2022). For the rest, we provide the scores reported in the MOBILE paper (Sun et al. 2023). We exclude TATU+MOPO, MATU+MOPO, COMBO, and MOREL because NeoRL paper does not report any results for them, and also it would be extremely time-consuming to carry out hyperparameter search for each approach. Similar to the experiments with D4RL, MOBILE results come from four independent training runs three million gradient steps.
C.2 POLICY OPTIMIZATION
Our implementation heavily relies on the codebase of Zhang et al. (202b)), which proposes uncertainty truncation, e.g., PATU+MOPE and TATU+COMBO. We use the default parameters of SAC described in Zhang et al. (2023b). We train RL agents on a cluster with NVIDIA RTX A5000 GPUs and Intel(R) X Gold 6226R CPU @ 2.90GHz. Given memory of approximately 6GB, training run of million gradient steps take around 12 hours.
C.3 HYPERPARAMETERS OF NUNC
NUNO has four hyperparameters: real ratio 3, rollout length h. CVaR coefficient a to set a truncation threshold, and uncertainty penalization threshold Xpen- The real ratio parameter 3 refers to the ratio of samples from the real dataset in mini-batch used to update the SAC policy, We set , to 0.05, as MATU+MOPO, for all tasks in our experiments. For the rest of the parameters, we run search over the following set of values: h E {5, 10, 15, 20). a E {0.9, 0.95, 0.98, 0.99, 1.0). and Apen E {0.001, 0.1, Our hyperparameter search procedure starts by tuning for rollout length with a = 0.9 and Xpen = 0.001. Using the best performing, namely, the highest huwan- normalized score yielding rollout length, we tune for a. Finally, we run search for Xpen Table reports the best-performing values for each task in our experiments. We use the same values for NUNOR.
C.4 NEURAL SDE TRAINING
We implement all the numerical experiments using the python library JAX Bradbury et al. (2018), in order to take advantage of its automatic differentiation and just-in-time compilation features. We use Python 3.8.5 for the experiments and train all our models laptop computer with Intel i9-9900 3.1 GHz CPU with 32 GB of RAM and a GeForce RTX 2060, TU106.
For training the neural SDE, we use randomly sampled sequences of horizon 2 for all the environments. We take the timestep of the ground truth environment and use it as the time step to integrate the neural SDE models. We use Euler-Maut Maller-Maruyama as the integration scheme in all our experiments and generate one particle during each integration step to compute the expectation defined Ldata- For the regularization loss term C,., we define (4 (St, ax) = e (N, ensuring that the output is positive, where No is a neural network parametrized by d.
For the neural SDE architecture, we parameterize 77 as a neural network with two hidden layers of size 64 with swish activation functions. We parameterize the uncertainty term 0g a neural network with two hidden layers of size 256 with tanh activation functions. The reward's 's drift term freward is parameterized as a neural network with three hidden layers of size 64 with swish activation functions while the other drift terms are parameterized with three hidden layers of size 256 and swish activation functions. Finally, the strong convexity neural network is parameterized with two hidden layers of size 32 with swish activation functions.
18

---

needed to Published as a conference paper at ICLR 2025
Table 3: Hyperparameters of NUNO in D4RL and NeoRL MuJoCo tasks
Task
h
Apeo
be-random
20
0.001
bp-random
10
0.001
wk-random
0.99
0.001
he-medium
0.99
hp-medium
10
0.99
wk-medium
10
0.98
he medium replay
5
0.9
hp-medium-replay
10
0.90
wk-mmedium-replay
10
0.95
ne-mediunn-xpert
hp medium expert
10
0.99
wk-mmeiim expert
0.98
be-low
10
bp-low
0.001
wk-low
5
0.99
0.001
ne--medium
10
0.99
hp-medium
5
0.9
0.1
wk--mediom
0.99
0.001
he-high
5
0.99
hp high
0.9
0.1
wk-high
5
0.99
The distance-awar diffusion term trained with a ball radius of 1.1 in all environments and a strong convexity coefficient of 1. We use Àgrad 10-4 and 20 samples to obtain the enforce the strong convexity constraint.
Training optimizer hyperparameters. We use the Adam optimizer (Kingma & Ba 2014) for all optimization problems. We use the default hyperparameters for the optimizer, except for the learning rate, which we linearly decay from 0.01 0.001 over the first 5000 gradient steps. We use early stopping criteria for all our experiments. We use a batch size of 128 for the neural SDE training.
D ABLATION STUDY FOR UNCERTAINTY ESTIMATORS IN POLICY TRAINING
NUNO incorporates the distance-aware uncertainty estimate ne into the offline RL. framework to enforce conservatism in the learned policy. Specifically, NUNO uses n7 to penalize and truncate the transitions generated by the learned neural SDE model during the RL policy training. In an ablation study, we investigate whether the choice of the uncertainty estimator impacts the learned policy. We compare NUNO, which uses the distance-aware uncertainty estimate Ma. corresponding to epistemic uncertainty, against NUNCa1 which utilizes the alcatoric uncertainty estimate dg (5). We evaluate NUNO and NUNOa1 in two types of datasets of the D4RL benchmark: random and medium-experr. Random datasets have low-quality trajectories, the data-logging policies sub-optimal. At the same time, random datasets have high coverage, as the trajectories showcase random behavior. In comparison, medium-exper datasets have high-quality trajectories yet low coverage as the data- logging policies are not random, and they act expert-like. D4RL benchmarks do not consist of very noisy datasets. Hence, expect have low aleatoric uncertainty in both datasets. However, data coverage determines epistemic uncertainty. We expect low epistemic uncertainty in random datasets and high medium-expert expert ones.
Table shows the best average Humann-normallo scores NUNO and NUNOal achieve, whereas Figure demonstrates their performance progression. random datasets, NUNO and NUNO both achieve SOTA results, with NUNC performing better in halfcheetah and walker2d. In contrast, NUNOa1 performs significantly worse in medium-exper datasets, except in hopper, where NUNO and NUNO"1 achieve similar scores. These results align with expectations based on the coverage properties of random and medium-expert datasets. A critical remark is that the reward penalty coefficient Apen is set to low value, Apen 0.001, in random datasets. Hence, NUNO and NUNO" practically do not penalize the agent, except when the uncertainty is estimated to be very high. In comparison. in medium-expert datasets, the reward penalty coefficient is Xpen 1, hence they frequently penalize
19

---

Published as a conference paper at ICLR 2025
Table 4: Average human-normalized scores NUNO and NUNO#1 in D4RL benchmarks. Due to limited space, we use abbreviations of task and dataset names: he halfcheetah, hp hopper, wk walker2d; = random, me = medium-expert. We report the mean and standard deviation (following =) of best scores among independent runs. Bold scores indicate the best for each task.
Task & Data
he-r
hp-r
wk-r
he-me
hp-me
wk-me
NUNO
52. 23.4 50.6+2.8
73.2.0.8
27.7=0.9
97.0
112.2*0.3
113.220.5
NUNO
71.7
18.3+1.7
10.5=a.4
112.6*0.9
48.3
Human 3ormalized scure
80
Human Sormallzd scure
Hotman aommslzzd score
30
150
2.5
20
2.0
15
20
20
NUNO
NUNO
NUNO
1.0
0.2
0.6
1.0
0.8
1.
opochs
opochs
opochs
x10
(a) hopper-random-v2
(b) hallcheetah-randomee
(c) walker2d-random--22
Elumen murmalzzed score
score
100
aleeelleedd scorts
100
Blumen normalized
75-
60
50
40
NUNG"
55
25
NUNO NUNO"
20
NUNO
NUNO"
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number of upochs
x10
Number epochs
X10
Number upochs
x10
(d) hopper-medduum-expert222
(e) mallcheetth-mmediu---xer-r--
(f) walker22-meeduum--xxprty
Figure Ablation study: The impact of the choice of uncertainty estimator on policy learning.
the agent. As the epistemic uncertainty is expected be high in this setting, the distance-aware uncertainty estimate is superior to the aleatoric uncertainty estimate.
E DETAILED RESULTS
This section provides detailed results on model accuracy and training progression.
E.1 MODEL ANALYSIS
Figure provides model analysis results for D4RL Hopper and HalfCheetah (see Section 5 for results in D4RL Walker2d). In-distribution evaluation demonstrates how learned dynamics models perform over varying prediction horizons in trajectories from datasets with which the models are trained. In D4RL Hopper, probabilistic ensembles yield significantly higher prediction errors than neural SDEs as the horizon lengths increase. D4RL HalfChectah, the same pattern occurs, except in seffcheetaa-mediim--eeee
where all models provide low prediction error. Dutto--distriiuution evaluation assesses how learned dynamics models trained with low-quality datasets, i.e., random, perform trajectories collected by behavior policies that better than a random policy. All models perform well in trajectories from the random task, which is in-distribution. However, in D4RL Hopper, ensembles yield high prediction error in medium-ree replayy and medium-exper as the horizon length increases. In D4RL HalfCheetah, we observe the same results, except in medium-replay.
E.2 TRAINING PROGRESSION
Figures 10 and demonstrate the training progression D4RL domains, Hopper, HalfCheetah and Walker2d, respectively. The first columns illustrate the progression of human normalized score in evaluation episodes ran after every epoch during training. In most tasks, NUNO and NUNO achieve similar human normalized scores at the end of the training, with some exceptions such hopper-
20

---

Published as a conference paper at ICLR 2025
104
Doa mids poy
Sea hale nate
Prodiction Livert
1013
rode
10"
2013
prediction
20
10%
104
10"
Horlaon
Hotiton
(a) D4RL Hopper: In distribution
(b) D4RL Hopper: Out-of distribubtt
taifthwitthh modhim
balth randim
Prediction Breee
talffiteiiih modium
balttibes modthm
telffieotthh modiuam expering
onsemble
20"
battihectih modlamm implay
satfcbectib modiam oapers-11
10%
200
104
20
26
20
(c) D4RL HalfCheetah:
In-distribution
(d) D4RL HalfCheetah: Out-of-distribution
Figure 8: Model analysis: We illustrate the evolution of model prediction error in different datasets for D4RL Hopper and HalfCheetah. (a) In-distribution: Evaluation of the datasets in which the models trained. (b) -itto---istiibtiion: Evaluation of models, trained via random, in trajectories from other datasets.
aifichettah-mediuu--ppphaa and walker2d-medium The second and third columns show the progression of the uncertainty estimates of neural SDEs trained in NUNO and NUNO*, as well as those models' uncertainty estimates for random actions and actions from offline datasets. In random datasets, progress, random, and data curves close each other, as these datasets consist of trajectories from random policy. In the rest of the tasks, neural SDEs can distinguish in-distribution actions (data) from diitriibittiitioonn actions (random). Furthermore, as the trained policy progresses, the corresponding uncertainty estimates of neural SDEs approach the data curve. This is expected as neural SDEs generate synthetic trajectories close to offline data, and the policies' replay buffers are initially augmented with the offline dataset.
random-v2, suffiheetta-meddim-ee
Figure demonstrates the progression of human normalized score in NeoRL tasks. Like D4RL, in most tasks, NUNO and NUNO reach similar scores. In Medium- 1000 NUNC outperforms NUNOR, The opposite occurs in HalfCheetah--3-Loov and HalfCheetah-v2 -3- 13-Medium-100
1000 and Walker2d-v3-
21

---

Published as a conference paper at ICLR 2025
Humaz normalzed score
0.25
0.20
data
runiomi
Uncertalaty
0.20
0.15
0.15
NUNO
0.10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number of epochs
x10
Number of egochs
x10>
Number of egochs
x10
(a) Random: Score
(b) Random: NUNO uncertainty
(c) Random: NUNO uncertainty
Zuman norma aiizaalddddd score
restom dats
rundom
Unceralicy
9.4
progrexs
Unzertaiaty
0.4
data
progees
60
0.3
0.3
0.2
NUNO
0.1
0.0
0.2
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number of epochs
K10
Number of egochs
x10>
Number of egochs
x10
(d) Medium: Score
(e) Medium: NUNO uncertainty
(f) Medium: NUNO uncertainty
Humao nomallzzd soor
100
0.4
data
random
data
random
Uge Uaxertaiuty
0.3
progree
0.3
proje66
0.2
0.2
20
NUNO
NUNO
0.1
0.1
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
0.0
0.6
Number of epochs
x10
Number of egochs
x10
Number opochs
(g) MR: Score
MR: NUNG uncertainty
(i) MR: NUNO" uncertainty
Human normalind 50000
1.00
random
78
Uncertalaty
0.3
0.3
progress
0.2
0.2
28
NUNO
NUNO
0.1
Number epochs
Number egochs
opochs
x10
(j) ME: Score
(k) ME: NUNO uncertainty
()) ME: NUNO uncertainty
In the first column, we report the progression of in evaluation episodes during training. the second and third columns, demonstrate how the uncertainty estimates NSDES in NUNO and NUNO evolve when evaluated with trained policies' actions in one-step rollouts from states in the dataset. 'random' and 'data' refer to the uncertainty estimates of the learned model given actions from random policy and the dataset, respectively. Each row corresponds to progression in different task: random, medium, medium-replay, and medium-expert.
Figure 9: Training progression in D4RL Hopper:
human normalized
22

---

Published as a conference paper at ICLR 2025
Humaa pormalzed score
rundom
6.-
data
0.3
progress
0.3
Uncertalen
0.2
20
1.0
NUNO
0.1
NUNO
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
03.6
0.8
1.0
Number of epochs
x10
Number of egochs
x10>
Number egochs
x10
(a) Random: Score
(b) Random: NUNO uncertainty
(c) Random: NUNO uncertainty
2human norma slieed 50ors
1.0
60-
Uncortalacy
40-
0.6
0.8
2.
0.4
0.4
NUNO
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
0.8
Number of eguchs
K10
Number of egochs
x10>
Number of egoths
x10
(d) Medium: Score
(e) Medium: NUNO uncertainty
(f) Medium: NUNO uncertainty
Homan namallad 5006
random
caniom
data
Uso copcrriiity
progrexx
Uscurtiisty
progresa
0.4
0.2
2.0
NUNO
0.0
NUNO
0.2
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
Number of epochs
x10
Numbar egoctss
x10
Number opochs
x10
(g) MR: Score
MR: NUNG uncertainty
(i) MR: NUNO" uncertainty
Htmmas normallzed score
cardom data
progrees
0.8
Uncoralaity
80
0.6
0.6
dats
randdom
NUNO
NUNO
0.2
0.2
1.0
Number epochs
Number egochs
opochs
(j) ME: Score
(k) ME: NUNO uncertainty
()) ME: NUNO uncertainty
Figure 10: Training progression in D4RL HalfCheetah: In the first column, we report the progression of human normalized score in evaluation episodes during training. In the second and third columns, we demonstrate how the uncertainty estimates of NSDEs in NUNO and NUNO evolve when evaluated with trained policies' actions in one-step rollouts from states in the dataset. 'random' and 'data' refer to the uncertainty estimates of the learned model given actions from random policy and the dataset, respectively. Each row corresponds to progression in different task: random, medium, medium-replay, and medium-expert.
23

---

Published as a conference paper at ICLR 2025
Human 20mallooo momalltedd score
30
0.30
0.25
rasiom
2.5
Usoortainn aartaainy
0.25
0.20
data
progress
0.20
0.15
0.15
10
NUNO
0.85
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0 x10>
0.0
0.2
0.4
0.6
0.8
1.0
Number of epochs
x10
Number of egochs
Number of egochs
x10
(a) Random: Score
(b) Random: NUNO uncertainty
(c) Random: NUNO uncertainty
2iuman normallzed 50ors
.... rentom dats
rundom
data
Unceralsty
progrexs
Uncertaiaty
progees
0.4
0.3
0.3
2.0
0.2
NUNO
0.1
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
0.8
Number of eguchs
K10
Number of egochs
x10>
Number of egoths
x10
(d) Medium: Score
(e) Medium: NUNO uncertainty
(f) Medium: NUNO uncertainty
Humao nomallzzd soore
100
rundon dele
00
Uncortaiitt
Uacertalozy
60
0.3
20
NUNO
NUNO
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
Number of epochs
x10
Number of egochs
x10
Number opochs
(g) MR: Score
MR: NUNG uncertainty
(i) MR: NUNO" uncertainty
Human normalind sco06
0.6
Utsveriaiity
progress
0.5
progress
75
0.4
28
NUNO
0.2
NUNO
Number epochs
Number egochs
opochs
x10
(j) ME: Score
(k) ME: NUNO uncertainty
()) ME: NUNO uncertainty
Figure 11: Training progression in D4RL Walker2d: In the first column, we report the progression of
human normalized score in evaluation episodes during training. the second and third columns, demonstrate how the uncertainty estimates NSDES in NUNO and NUNO evolve when evaluated with trained policies' actions in one-step rollouts from states in the dataset. 'random' and 'data' refer to the uncertainty estimates of the learned model given actions from random policy and the dataset, respectively. Each row corresponds to progression in different task: random, medium, medium-replay, and medium-expert.
24

---

Published as a conference paper at ICLR 2025
Euuman 302malzz0 score
Human bormallze scoze
30
100
1.00
80
80
80
1.0
NUND
NUNO
20
NUNO
NUND
NUNO
1.0
1
opochs
epochs
epochs
(a) 33- 1000
(b) Hopper- v3 wwewe---MMMae Mee--MM-dimmo 1000
(c) H3pper-v3-Iigge High--1000 1000
31uman pormalzed score
Human normallzed scure
80
100
50
60
40
NUNO
20
NUNO NUNO
Human
NUNO
NUNO
NUNO
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number eguchs
K10>
Number
epochs
x10>
Number of epochs
x10
(d) HalsCheetah--y--Low-
(e) ShallCheetaa-v3-MMediie 1000
(i) HallCheetah--33-HHHe -ightthhhh--------O00-
Hummn noraalized scora
Humann aormaltzed ltzed score
80-
20
40
2.0
20
Human
3.0
NUND
NUND
NUNO
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number of epochs
x10>
Number of epochs
x10>
Number of opochs
x10
gg) Walker2d-v3-Loo-11n0 1000
(h) Walker2d 1000
Walker2ddvv--ighh-1000
Figure 12: Training progression in NeoRL tasks: Each subfigure reports the progression of human normalized score evaluation episodes during training.
25


Reviews:
[{"review_id": "w658wX1Cla", "reviewer": "Reviewer_wnnp", "summary": "The paper presents NUNO, a method for offline reinforcement learning that uses neural SDEs to model dynamics with uncertainty. The SDE structure consists of a drift term that incorporates prior physics knowledge (as an inductive bias) and a diffusion term that captures uncertainty. Further, the uncertainty term is designed to account for both aleatoric and epistemic uncertainties explicitly to obtain accurate uncertainty estimates. By incorporating physics-based inductive biases and adaptive rollout truncation, NUNO mitigates model exploitation, where policies may overfit unreliable parts of the model.", "questions": "same as weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Strengths**\nI like the paper overall. Here are some general comments. \n- Impressive Results on Noisy Datasets: The paper achieves particularly strong performance on challenging, low-quality datasets, highlighting the robustness of NUNO’s uncertainty modeling in noisy conditions.\n- Clear Presentation: The paper is well-written, with a logical structure that effectively communicates complex methods and results, making it accessible even with its technical depth.\n- Novel Modeling Approaches: The decomposition of uncertainty into aleatoric and epistemic components, and learning them separately, stands out as an interesting and valuable contribution. This approach advances the understanding of uncertainty in offline RL and reinforces the model’s reliability.", "weaknesses": "**Weaknesses**\n- Lack of Ablations for Design Decisions: A detailed analysis of each design decision, particularly the separate uncertainty terms, , would clarify their individual contributions to the model’s performance.\n- Limited to Low-Dimensional Environments: The experiments are confined to low-dimensional MuJoCo environments, without tests on high-dimensional settings. Such environments may reveal limitations of the approach, where ensemble-based uncertainty methods could potentially perform better.\n- Limited Scope on Partially Observable/Image-Based Domains: Although the paper notes the absence of experiments in partially observable or image-based environments, this remains a notable gap. These domains may challenge NUNO’s performance, and ensemble-based approaches might excel here due to their structure. Expanding experiments to these settings could provide a more comprehensive evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C1IldNuywj", "reviewer": "Reviewer_jFi7", "summary": "In this work, the authors propose a method for offline model-based reinforcement learning that leverages uncertainty to enhance performance in settings with low-quality data or limited samples. The approach involves learning a model of the environment from samples using Neural Differential Equations (NDE), incorporating uncertainty and physics-informed learning. To account for uncertainty in the NDE, the authors introduce a parametric, distance-aware uncertainty estimator. Unlike kernel k-th nearest neighbor, this estimator avoids expensive search requirements while maintaining desirable properties such as monotonicity and boundedness. The estimator is used to limit the differential equation’s rollouts when distance is sufficiently high, restricting state-action paths to high-quality model rollouts and penalizing high-uncertainty paths, thereby encouraging the agent to favor paths where the NDE model is well-behaved. Simulations on D4RL and NeoRL, varying the quality of the policy used for offline training, demonstrate that the proposed method outperforms baselines in many cases. These baselines consist of other uncertainty-aware model-based methods that also penalize and truncate paths based on uncertainty.", "questions": "- Q1: How does incorporating prior physics knowledge differ from simply parameterizing the state space in a convenient way? Choosing a state representation is a standard design choice in most RL problems. How does this incorporation of prior knowledge apply to applications outside of MuJoCo? The parameterization in Equation (8) appears very specific to this case. Could you incorporate other types of physics knowledge, such as gravity, joint resistance, or environmental viscosity? One could argue that these environmental properties also represent physics knowledge.\n\n- Q2: What are the implications of using dataset distance as an uncertainty estimator? The dynamical SDE described in Equation (5) can separate epistemic and aleatoric uncertainty in the dynamical model, where the epistemic uncertainty term depends on the distance-aware uncertainty estimation. It is possible that points further from the modes of the training set are influenced by noise, leading to high distance, and thus high epistemic uncertainty (since h in Equation 5 is monotonic). Conversely, a high-distance point could be \"stable\" (not noisy), with the dynamical model learning it well, so true epistemic uncertainty should be low for this point but is high according to the distance estimator.\n\n- Q3: Related to the previous question, as the pessimistic reward is discounted by the distance-aware estimator, why isn’t aleatoric uncertainty considered here (as in Equation 5)? Highly noisy trajectories from the dynamical model cannot be learned by the agent, but they are not penalized, only those that the model hasn’t learned well enough based on distance are. This also relates to the previous question.\n\n- Q4: How critical is this particular distance-aware uncertainty estimator? Why not use other uncertainty surrogates, such as prediction error from the dynamical model or the frequency of updates in a region (number of visits)? These other types of uncertainty estimation have been used in exploration-exploitation trade-offs in RL agents. This, and other ablation studies could be explored in a toy model.\n\n- Q5: Typo in Equation 8: I believe the first variable should be s_{pos}? Additionally, the uncertainty penalization in Figure 1 doesn’t match the equation, and some notation in the figure is inconsistent with that used in the main text.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- (S1): The problem of offline training with few samples and low-quality policies is a relevant one with significant applications in real-world settings.\n\n- (S2): Extensive simulations, including hyperparameter searches and comparisons to baselines, demonstrate promising results.\n\n- (S3): The manuscript is well-explained and relatively easy to follow, and the authors provide a thorough literature review that clarifies the position of their proposed method within the existing body of work.\n\n- (S4): The proposed Distance-Aware uncertainty estimator is particularly interesting and may be useful for other problems. The authors ensure that their estimator is well-behaved through mathematical analysis, and they impose desirable properties on it through careful considerations in the optimization process.", "weaknesses": "- (W1): Some of the main contributions listed by the authors are not representative of the actual content presented in the paper. This could be improved by shifting the focus to other key aspects of the proposed method (see questions).\n\n- (W2): The concept of uncertainty is underdeveloped in the manuscript, despite being a vital component of the proposed method. Although the authors clearly define distance-aware uncertainty, it is unclear what this concept represents in a broader context, given that the notion of uncertainty varies across applications and methods. This could be improved by examining the inner workings of their proposed uncertainty estimator in the simulated environments.\n\n- (W3): While the results on the selected environments across different data regimes are promising (though other methods perform similarly to NUNO), there is a lack of exploration into the inner workings of the method. Obviously, simulations in these environments are computationally expensive, setting up a simpler problem with a known ground truth could enable multiple regime tests and ablation studies of the proposed method. This is especially important given the method's multiple components, and providing a thorough explanation of why and how the method works would significantly enhance the paper’s contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zYjasXZdie", "reviewer": "Reviewer_KRec", "summary": "This paper introduces an offline model-based reinforcement learning method that can work well in low-quality data regimes while maintaining competitive performance on high-quality datasets. NUNO, the proposed  Neural Stochastic Differential Equations for\nUNcertainty-aware Offline RL,  learns a dynamics model as neural stochastic differential equations (SDE), where the drift term integrates   physics knowledge to provide inductive bias, and the diffusion term estimates uncertainty to control the effects of uncertain predictions. Additionally, the paper addresses model exploitation by penalizing and truncating predictive rollouts, achieving strong results on several control benchmarks.", "questions": "In Equation 3, on the definition of \\mathcal{L}_{\\mu}, how is it guaranteed that the denominator $\\mu$ is not zero in experiments?\n\nIn Equation 6, what is the intuition of the term log(det(\\Sigma)) in the context of \\mathcal{L}_{\\text{data}}?\n\nCould you provide examples of \"unmodeled dynamics\" as mentioned in Lines 342–343?\n\nIf the unmodeled dynamics are significant, would it still be beneficial to penalize the residual term?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes  a Neural Stochastic Differential Equation based method to address the important challenge of underrepresented state-action pairs in offline datasets and achieves strong experimental results. It has a solid theoretic foundation, and it is innovative and interesting. \n\nThe proposed method is well thought out, and the design for the uncertainty estimator has good local theoretical properties. \n\nThe paper is mostly well-written and easy to follow.", "weaknesses": "The proposed method relies on decomposing the state into position and velocity, and it imposes strong assumptions on the underlying environment dynamics (e.g., the dynamics must satisfy rigid body dynamics). I suggest that the authors also elaborate on how to apply the  physics-informed approach to other setting (not just position/velocity based). \nWhile incorporating such prior physics knowledge is beneficial in some cases, clearly it is  not possible in all RL applications. \n\n Regarding the design of uncertainty estimator in 4.1, it would help to include simulation experiments under simplified settings of MDP and datasets to explicitly quantify the effects of the design choices.  This would shed light on the design of uncertainty estimator and hence the SDE model.\n\nThe computational overhead can be very high. Specifically, what is the computational overhead in training due to sampling around each data point for the uncertainty estimator? How does this affect the scalability of the method in high-dimensional datasets?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)