# Deterministic Uncertainty Propagation for Improved Model-Based Offline
Reinforcement Learning

Abdullah Akgül
Department of Mathematics and Computer Science
University of Southern Denmark
Odense, Denmark
<EMAIL_ADDRESS>
&Manuel Haußmann
Department of Mathematics and Computer Science
University of Southern Denmark
Odense, Denmark
<EMAIL_ADDRESS>
&Melih Kandemir
Department of Mathematics and Computer Science
University of Southern Denmark
Odense, Denmark
<EMAIL_ADDRESS>

###### Abstract

Current approaches to model-based offline Reinforcement Learning (RL) often
incorporate uncertainty-based reward penalization to address the
distributional shift problem. While these approaches have achieved some
success, we argue that this penalization introduces excessive conservatism,
potentially resulting in suboptimal policies through underestimation. We
identify as an important cause of over-penalization the lack of a reliable
uncertainty estimator capable of propagating uncertainties in the Bellman
operator. The common approach to calculating the penalty term relies on
sampling-based uncertainty estimation, resulting in high variance. To address
this challenge, we propose a novel method termed Moment Matching Offline
Model-Based Policy Optimization (MOMBO). MOMBO learns a Q-function using
moment matching, which allows us to deterministically propagate uncertainties
through the Q-function. We evaluate MOMBO’s performance across various
environments and demonstrate empirically that MOMBO is a more stable and
sample-efficient approach.

##  1 Introduction

Offline reinforcement learning (RL) (Lange et al., , 2012; Levine et al., ,
2020) is a paradigm for extracting the optimal policy from a dataset
containing transitions derived from interactions with a behavioral policy in a
given environment. This approach proves valuable in various domains where
direct interaction with the environment is impractical or risky (Shortreed et
al., , 2011; Singh et al., , 2020; Nie et al., , 2021; Micheli et al., ,
2022). Applying traditional online off-policy RL methods directly to offline
RL settings often leads to poor performance due to the distributional shift
problem (Fujimoto et al., , 2019; Kumar et al., , 2019) caused by the
accumulation of extrapolation errors during policy evaluation into a positive
bias exploited by the policy search step.

It is commonplace to perform offline RL via both model-free and model-based
approaches. Model-based offline RL approaches fit a probabilistic model on the
real state transitions as a first step and then supplement real data with
synthetic samples generated from this model on the fly through the course of
policy search (Janner et al., , 2019). One line of work addresses
distributional shift by constraining the policy learning (Kumar et al., ,
2019; Fujimoto and Gu, , 2021), where the policy is trained to mimic the
behavioral policy and penalized based on the discrepancy between its actions
and those of the behavioral policy, akin to behavioral cloning. A second
strategy is to introduce conservatism to training by _(i)_ perturbing the
training objective via a high-entropy behavior policy (Kumar et al., , 2020;
Yu et al., , 2021), _(ii)_ penalizing state-action pairs proportionally to
their estimated variance (An et al., , 2021; Yu et al., , 2020; Bai et al., ,
2022; Sun et al., , 2023), or _(iii)_ adversarially training the dynamics
model to minimize the value function (Rigter et al., , 2022) to prevent
overestimation in policy evaluation for out-of-domain (OOD) state-action
pairs.

In the offline RL literature, uncertainty-driven approaches exist for both
model-free (An et al., , 2021; Bai et al., , 2022) and model-based settings
(Yu et al., , 2020; Sun et al., , 2023), all aimed at learning a pessimistic
value function (Jin et al., , 2021) by penalizing it with an uncertainty
estimator. The impact of uncertainty quantification has been investigated in
both online (Abbas et al., , 2020) and offline (Lu et al., , 2021) scenarios,
particularly in model-based approaches, which is the setting our work focuses
on. Despite model-based research demonstrating the significance of uncertainty
estimators and their influence on policy performance, current model-based
offline RL methods still lack a reliable uncertainty estimator. These methods
struggle to propagate uncertainty from the predicted next state and reward by
the learned dynamics model to the Bellman targets, even though policy
performance is directly linked to uncertainty in the Bellman error (O’Donoghue
et al., , 2018; Luis et al., , 2023). State-of-the-art approaches (Sun et al.,
, 2023) approximate the variance of Bellman targets by sampling from the
learned environment model, which incurs high estimator variance and causes
both over-penalization and instability in training.

##### Our contribution.

We identify sampling variance as a common limiting factor for the performance
of uncertainty-based model-based offline RL approaches. We characterize the
reasons behind this bottleneck by a theoretical analysis of the concentration
properties of the estimators in commonplace use. We address this bottleneck by
employing progressive moment matching, an idea that has, e.g., useful
applications in deterministic variational inference of Bayesian deep learning
(Wu et al., 2019a, ), for the first time to deterministically propagate
environment model uncertainties through Q-function estimates. The resulting
model, _Moment Matching Offline Model-Based Policy Optimization (MOMBO)_ ,
significantly accelerates training compared to its sampling-based counterparts
while maintaining performance. Furthermore, it has favorable theoretical
guarantees.

##  2 Preliminaries

##### Reinforcement Learning.

We define an episodic Markov Decision Process (MDP) as a tuple
${\mathcal{M}=(\mathcal{S},\mathcal{A},r,\mathrm{T},\gamma,H)}$, where
$\mathcal{S}$ represents the state space and $\mathcal{A}$ denotes the action
space. The function
${r(s,a):\mathcal{S}\times\mathcal{A}\rightarrow[0,R_{\max}]}$ is a bounded
reward function, and
$\mathrm{T}(s^{\prime}|s,a):\mathcal{S}\times\mathcal{A}\rightarrow\Delta(\mathcal{S})$
represents the state transition probability density where $\Delta$ denotes the
probability simplex defined on the state space. The parameter $\gamma\in(0,1)$
is the discount factor and $H$ is the episode length. We define return for a
policy $\pi$ as
$J_{t}(\pi)\triangleq\mathds{E}_{\pi}[\sum_{i=t}^{H}\gamma^{i-t}(s_{i},a_{i})]$.
The objective of RL is to choose a policy that maximizes the cumulative
episode reward for the initial state $J_{0}(\pi)$ within an environment
(Sutton and Barto, , 2018). Actor-critic algorithms achieve this by
alternating between _policy evaluation_ and _policy improvement_ phases.
During the policy evaluation phase, they rely on a deep neural network
$Q_{\psi}^{\pi}(s_{t},a_{t}):\mathcal{S}\times\mathcal{A}\rightarrow\mathds{R}_{+}=J_{t}(\pi)$,
to approximate the _Q-function_. The networks parameters $\psi$ are inferred
by minimizing the _Bellman error_ ,

$\psi^{*}=\operatorname*{arg\,min}_{\psi\in\Psi}\mathds{E}_{(s,a,r,s^{\prime})}[(Q_{\psi}^{\pi}(s,a)-\mathcal{T}^{\pi}Q_{\psi}(s,a))^{2}],$

where $\mathcal{T}^{\pi}Q$ is the _Bellman operator_ ,

$\mathcal{T}^{\pi}Q(s_{t},a_{t})\triangleq
r(s_{t},a_{t})+\gamma\mathds{E}_{\vbox{\Let@\restore@math@cr\default@tag\halign{\hfil$\m@th\scriptstyle#$&$\m@th\scriptstyle{}#$\hfil\cr
s_{t+1}&\sim\mathrm{T}(s_{t},a_{t})\\\
a_{t+1}&\sim\pi(a_{t+1}|s_{t+1})\crcr}}}\left[Q^{\pi}(s_{t+1},a_{t+1})\right].$

In the policy improvement phase, actor-critic algorithms learn a policy using
a deep neural network parameterized by $\xi$ through maximizing the Q-value:
$\xi^{*}=\operatorname*{arg\,min}_{\xi\in\Xi}\mathds{E}_{s}[\big{.}-Q_{\psi^{*}}^{\pi_{\xi}}(s,\pi_{\xi}(a|s))]$.

##### Model-based offline RL algorithms.

Offline RL algorithms perform policy search using only a dataset
${\mathcal{D}=\left\\{(s_{i},a_{i},r_{i},s_{i}^{\prime})\right\\}_{i=1}^{\mid\mathcal{D}\mid}}$
collected from the target environment beforehand by a behavioral policy
$\pi_{\beta}$. Unlike online RL, the policy search algorithm does not interact
with the environment during training. In model-based approaches for offline
RL, a dynamics model ${\widehat{\mathrm{T}}}$ is commonly trained through
maximum likelihood estimation (MLE) using the provided dataset
${\mathcal{D}}$, i.e.,
${\min_{\widehat{\mathrm{T}}}\mathds{E}_{s,a,s^{\prime}\sim\mathcal{D}}[-\log\widehat{\mathrm{T}}(s^{\prime}|s,a)}]$.
Additionally, a reward model $\widehat{r}(s,a)$ can also be jointly trained
using $\mathcal{D}$. Together, these models form an estimated MDP,
${\widehat{\mathcal{M}}=\\{\mathcal{S},\mathcal{A},\widehat{r},\widehat{\mathrm{T}},\gamma,H\\}}$.
Subsequently, an optimal policy can be learned from ${\widehat{\mathcal{M}}}$
through planning or RL algorithms. Mainstream model-based offline RL methods
adopt the Dyna approach (Sutton, , 1990; Janner et al., , 2019), which
suggests training an off-the-shelf model-free RL algorithm on synthetic data
generated from the learned environment model $\widehat{\mathcal{D}}$ using
initial states from $\mathcal{D}$. It has been observed that shuffling
minibatches collected from synthetic and real data sets improves performance
in both online (Janner et al., , 2019) and offline (Yu et al., , 2020)
settings. The state-of-the-art approach models state transition probabilities
and reward distributions as an ensemble (Lakshminarayanan et al., , 2017) of
$N_{\text{ens}}$ heteroscedastic neural networks assuming a normal density.

##  3 Limitations of SOTA model-based offline RL approaches

Offline RL methods face the primary challenge of _distributional shift_
because the static dataset lacks coverage of a large portion of the state-
action space. When the policy search algorithm probes unobserved state-action
pairs in the course of training, significant value approximation errors occur.
While in standard supervised learning scenarios these errors shrink as the
model observes more data, in RL, overestimation errors are picked up during
the policy improvement step and stay within the system causing the well-known
_overestimation bias_ (Thrun and Schwartz, , 1993). While this issue can be
mitigated in online setups by min-clipping (Fujimoto et al., , 2018),
additional measures are required in offline RL setups, such as adding
penalties to rewards proportional to the uncertainty around the next state
estimates. Algorithms belonging to this family, known commonly as pessimistic
policy iteration (Yu et al., , 2020; Sun et al., , 2023), build on a theory
that provides performance improvement guarantees for an upper bound on the
next state variance (Jin et al., , 2021; Uehara and Sun, , 2022). The state-
of-the-art practice exploits the fact that the theoretical guarantees are
tighter when applied on the uncertainty of the value estimate of the next
state, rather than the next state itself (Sun et al., , 2023). When deep
neural networks are used as value function approximators, calculating their
variance even for normal distributed input becomes analytically intractable.
Hence, reward penalties are obtained by Monte Carlo samples, which are prone
to high estimator variance.

It is our key finding that using high-variance estimates for reward
penalization causes three significant problems in the training process of
offline RL algorithms:

  * _(i)_

The information content of the distorted gradient signal shrinks, causing
delayed convergence to the asymptotic performance.

  * _(ii)_

The first two moments of the Bellman target are poorly approximated for
reasonable sample counts.

  * _(iii)_

Excessively large confidence radii need to be used to counter the instability
caused by _(i)_ and the high overestimation risk caused by the inaccuracy
described in _(ii)_ , which restricts model capacity at unnecessary levels.

We propose a novel approach centered on deterministic uncertainty propagation
using moment-matching, originally developed for Bayesian inference, combined
with a lower confidence bound. Through moment matching, we can effectively
propagate uncertainty from the predictions of the dynamics to the Q-function
without encountering the high variance caused by sampling errors. By employing
a robust uncertainty estimator, we establish a lower confidence bound on the
target Q-value, enabling a more efficient, sampling-free penalization with
improved training times.

The common practice in model-based offline RL algorithms is to generate
imaginary rollouts at every episode start. The generated data are shuffled
with the real observations and used to update the policy. During this
generation process, samples from the predictive distributions of the next
state and reward are used. These predicted samples are later used again to
calculate the Bellman error when updating the Q-function. However, estimators
calculated from these samples are subject to sampling errors, which may
amplify overestimation biases. To address these sampling errors and to tackle
distributional shift, model-based offline RL methods often adopt conservative
policies and penalize the state-action values based on the target predictions.
The target Q-value ($\mathcal{T}^{\pi}Q(s,a)$) for the model-based offline RL
algorithm is

$\displaystyle\mathcal{T}^{\pi}Q(s,a)=\mathds{E}_{\vbox{\Let@\restore@math@cr\default@tag\halign{\hfil$\m@th\scriptstyle#$&$\m@th\scriptstyle{}#$\hfil\cr
r&\sim\widehat{r}(s,a)\\\
s^{\prime}&\sim\widehat{\mathrm{T}}(s^{\prime}|s,a)\crcr}}}\left[\big{.}r+\gamma\mathds{E}_{a^{\prime}\sim
s^{\prime}}\left[Q^{\pi}(s^{\prime},a^{\prime})\right]\right],$ (3)

where $\widehat{\mathrm{T}}$ and $\widehat{r}$ are learned transition and
reward models respectively. Despite various options for conservatism in the
model-based offline RL literature, such as reward penalization by prediction
uncertainty of the dynamics model (Yu et al., , 2020; Lu et al., , 2021), or
by using uncertainty on the Q-value of the next state by taking samples (Sun
et al., , 2023), there is no analytical solution. Furthermore, these
uncertainty estimators are not able to accurately reflect the real uncertainty
on the target Q-value because they cannot propagate the uncertainty through
the Q-function. Instead, they have to approximate it by taking samples, which
incurs high estimator variance, leading to over-penalization and limiting the
learning speed. In the following section, we provide a theoretical analysis of
sampling-based uncertainty estimation and present a tail bound for one of the
most successful state-of-the-art methods, namely MOBILE (Sun et al., , 2023).

We assume that the estimation error surrounding the Bellman target estimate of
a state-action pair is normally distributed
$y\sim\mathcal{N}(\mu,\sigma^{2})$. The lower bound of a confidence set
surrounding its mean is then given as $y=\mu-\beta\sigma$ for some radius
$\beta>0$. We refer to this quantity as a $\beta-$Lower Confidence Bound
($\beta-$LCB). For a sample $S=\\{y_{1},\ldots,y_{N}\\}$ of $N$ independently
identically distributed (i.i.d.) draws from this distribution, denote the
empirical estimate of the $\beta-$LCB as
$\widehat{y}_{s}\triangleq\widehat{\mu}_{N}-\beta\widehat{\sigma}_{N}$ for the
empirical mean and variance estimates

$\widehat{\mu}_{N}\triangleq\frac{1}{N}\sum_{n=1}^{N}y_{i},\quad\widehat{\sigma}^{2}_{N}\triangleq\frac{\sum_{i=1}^{N}(y_{i}^{2}-\widehat{\mu}_{N})}{N-1}.$
(4)

Offline RL algorithms that apply pessimism to their target Q-value estimate
use a $\beta$-LCB estimate. E.g., the target Q-value estimate of MOBILE (Sun
et al., , 2023) can be expressed as
${\widehat{y}_{s}\triangleq\widehat{\mu}_{N}-\beta\widehat{\sigma}_{N}}$,
which is a $\beta$-LCB estimate, where

$\widehat{\mu}_{N}=r+\gamma(1-d)Q^{\pi}(s^{\prime},a^{\prime}),\quad\widehat{\sigma}_{N}=\sqrt{\mathrm{var}\left[Q^{\pi}(s^{\prime},a^{\prime})\right]}.$
(5)

We discuss additional work related to our proposal in Appendix A.

##  4 MOMBO: Moment matching offline model-based policy optimization

##### Deterministic uncertainty propagation via moment matching.

To compute target Q-values for the synthetic dataset one needs to perform
mathematical operations on uncertain inputs. Typically, the impact of
uncertain inputs on the output is approximated by sampling and evaluation.
However, this method introduces sampling errors and computational overhead. We
propose to instead propagate uncertainty through Q-function operations using
moment matching. Moment matching-based uncertainty propagation offers better
accuracy, computational efficiency, and reduces estimator variance (Wu et al.,
2019a, ; Wang and Manning, , 2013).

Assuming the input of a fully-connected layer to be
$X\sim\mathcal{N}(X|\mu,\sigma^{2})$, the pre-activation $Y$ for a neuron
associated with weights $\theta$, is given as $Y=\theta^{\top}X$, i.e.,
$Y\sim\mathcal{N}(Y|\theta^{\top}\mu,(\theta^{2})^{\top}\sigma^{2})$, where we
absorbed the bias into $\theta$ for simplicity and the square on $\theta$ is
applied element-wise. For a ReLU activation function
$r(x)\triangleq\max(0,x)$, mean $\widetilde{\mu}$ and variance
$\widetilde{\sigma}^{2}$ of $Y=\max(0,X)$ are analytically tractable (Frey and
Hinton, , 1999). We approximate the output with a normal distribution,
${\widetilde{X}\sim\mathcal{N}(\widetilde{X}|\widetilde{\mu},\widetilde{\sigma}^{2})}$,
and summarize several properties about $\widetilde{\mu}$ and
$\widetilde{\sigma}^{2}$ below.

######  Lemma 1 (_Moment matching_).

For $X\sim\mathcal{N}(X|\mu,\sigma^{2})$ and $Y=\max(0,X)$, we have

$\displaystyle\textit{(i)}\hskip
10.00002pt\begin{split}\widetilde{\mu}\triangleq\mathds{E}\left[Y\right]&=\mu\Phi(\alpha)+\sigma\phi(\alpha),\\\
\widetilde{\sigma}^{2}\triangleq\mathrm{var}\left[Y\right]&=(\mu^{2}+\sigma^{2})\Phi(\alpha)+\mu\sigma\phi(\alpha)-\widetilde{\mu}^{2},\end{split}$

where $\alpha=\mu/\sigma$, and $\phi(\cdot)$, $\Phi(\cdot)$ are the
probability density function (pdf) and cumulative distribution function (cdf)
of the standard normal distribution, respectively. Additionally, we have that

$\textit{(ii)}\hskip
10.00002pt\widetilde{\mu}\geq\mu\qquad\text{and}\qquad\textit{(iii)}\hskip
10.00002pt\widetilde{\sigma}^{2}\leq\sigma^{2}.$

We can show the following inequality between the cdfs,
$F_{U}(u)\triangleq\mathds{P}(U\leq u)$, of $X$ and $\widetilde{X}$.

######  Lemma 2 (_An inequality between normal cdfs_).

For two normally distributed random variables
${X\sim\mathcal{N}(X|\mu,\sigma^{2})}$ and
$\widetilde{X}\sim\mathcal{N}(\widetilde{X}|\widetilde{\mu},\widetilde{\sigma}^{2})$
with $\widetilde{\mu}\sigma\geq\mu\widetilde{\sigma}$, the following
inequality holds:

$F_{\widetilde{X}}(u)\leq F_{X}(u),\qquad\text{for $u\leq 0$}.$

Given $\widetilde{\mu}$ and $\widetilde{\sigma}^{2}$ as in Lemma 1, the
assumptions of this lemma hold and it is applicable to our moment-matching
propagation. See Appendix B for proofs on these two lemmata.

As the input to our Q-function, which propagates uncertainties
deterministically through moment matching, we provide a distribution of the
next state, with known mean and variance as predictions of the dynamics model,
along with an action sampled from the policy for the mean of the next state
and a zero variance. I.e., we treat the action as if it followed a delta
distribution and simply propagate its mean throughout the neural net. The
moment matching procedure is outlined in Algorithm 1.

Algorithm 1 Deterministic uncertainty propagation through moment matching

function MomentMatchingThroughLinear($\theta,b,X$)

$\theta$: Weights of the layer, $b$: bias of the layer

$X=\mathcal{N}(X|\mu_{X},\sigma_{X}^{2})$ $\triangleright$ input a normal
distribution

$(\mu_{Y},\sigma_{Y}^{2})\leftarrow(\theta^{\top}\mu_{X}+b,{\theta^{2}}{}^{\top}\sigma_{X}^{2})$
$\triangleright$ transform mean and variance

return $Y=\mathcal{N}(X|\mu_{Y},\sigma_{Y}^{2})$ $\triangleright$ output the
transformed distribution

end function

function MomentMatchingThroughReLU($X$)

$X=\mathcal{N}(X|\mu_{X},\sigma_{X}^{2})$ $\triangleright$ input a normal
distribution

$\alpha\leftarrow\nicefrac{{\mu_{X}}}{{\sigma_{X}}}$

$\mu_{Y}\leftarrow\mu_{X}\Phi(\alpha)+\sigma_{X}\phi(\alpha)$ $\triangleright$
compute the first two moments of $\text{ReLU}(X)$

$\sigma_{Y}^{2}\leftarrow\left(\mu_{X}^{2}+\sigma_{X}^{2}\right)\Phi(\alpha)+\mu_{X}\sigma_{X}\phi(\alpha)-\mu_{Y}^{2}$$\triangleright$
$\phi$/$\Phi$ are the normal pdf/cdf

return $Y=\mathcal{N}(X|\mu_{Y},\sigma_{Y}^{2})$ $\triangleright$ output a
normal distribution with these moments

end function

##### Lower confidence bound on target Q-value.

We construct a lower confidence bound on the target Q-value (Equation 3) by
representing it as a normally distributed random variable. We use ensemble
dynamics models, reward models, and a value function capable of propagating an
input distribution to the output in order to construct this LCB. Leveraging
moment matching on the value function and ensemble dynamics models, the target
Q-value can be expressed as a summation of normal distributions without
requiring samples. We formulate the $\beta$-LCB for MOMBO as

$\displaystyle\widehat{\mu}=\dfrac{1}{N}\sum_{n=1}^{N}\mu_{r}^{n}+\gamma\mu_{Q^{\prime}}^{n},\qquad\widehat{\sigma}=\sqrt{\dfrac{1}{N}\sum_{n=1}^{N}\left((\mu_{r}^{n}+\gamma\mu_{Q^{\prime}}^{n})-\widehat{\mu}\right)^{2}},$

where $N$ is the number of learned transition models,

$\displaystyle\widehat{r}(s,a)$
$\displaystyle=\mathcal{N}\big{(}\mu_{r}^{n},(\sigma_{r}^{n})^{2}\big{)},\qquad\widehat{\mathrm{T}}(s^{\prime}|s,a)=\mathcal{N}\big{(}\mu_{s^{\prime}}^{n},(\sigma_{s^{\prime}}^{n})^{2}\big{)},$
$\displaystyle Q^{\prime}$
$\displaystyle=\mathcal{N}\big{(}\mu_{Q^{\prime}}^{n},(\sigma_{Q^{\prime}}^{n})^{2}\big{)}=f_{L}\big{(}\mathcal{N}(\mu_{s^{\prime}}^{n},(\sigma_{s^{\prime}}^{n})^{2}),a^{\prime
n}\big{)},$

and $a^{\prime n}\sim\pi(a^{\prime n}|\mu_{s^{\prime}}^{n})$. Our pessimistic
target Q-value estimate of MOMBO is
$\mathcal{T}^{\pi}Q(s,a)=\widehat{\mu}-\beta\widehat{\sigma}$.

###  4.1 Practical implementation of MOMBO

We adopt the model learning scheme from model-based policy optimization (MBPO)
(Janner et al., , 2019) and use Soft Actor-Critic (SAC) (Haarnoja et al., ,
2018) as the policy search algorithm, as these are the best practices used in
many recent model-based offline RL methods (Yu et al., , 2020, 2021; Sun et
al., , 2023). However, we note that most of our findings are more broadly
applicable.

Following MBPO, we train an independent neural ensemble of dynamics models.
These models output Gaussian distributions over the next state and reward via
MLE. We denote each member of the ensemble is as
$\widehat{\mathrm{T}}^{i}(s^{\prime},r|s,a)=\mathcal{N}\big{(}\mu^{i}(s,a),(\sigma^{i})^{2}(s,a)\big{)}$
for ${i\in\\{1,\ldots,N_{\text{ens}}\\}}$. After evaluating their performance
on a validation set, we select the $N_{\text{elite}}$ best-performing ensemble
elements for further processing. These elite models are used for generating
$k$-step rollouts with the current policy to create the synthetic dataset
$\widehat{\mathcal{D}}$, which we then combine with the offline dataset
$\mathcal{D}$. Additionally, we construct LCB on the target Q-value. For
clarity, we distinguish the dynamics model and reward model separately as
$\widehat{\mathrm{T}}^{i}(s^{\prime}|s,a)=\mathcal{N}(\mu_{s^{\prime}}^{i},(\sigma_{s^{\prime}}^{i})^{2})$
and $\widehat{r}^{i}(r|s,a)=\mathcal{N}(\mu_{r}^{i},(\sigma_{r}^{i})^{2})$ for
$i\in\\{1,\ldots,N_{\text{elite}}\\}$.

For policy learning, we modify SAC and switch its Q-function approximator,
which is a neural network, with our proposed deterministic moment propagating
neural network. Additionally, we add a pessimistic target Q-value calculation.
The critic networks $Q_{\psi_{k}}^{\pi_{\xi}}(\cdot,\cdot)$ for the policy
$\pi_{\xi}$ are trained with the following loss

$\mathcal{L}(Q_{\psi_{k}}^{\pi_{\xi}})=\mathds{E}_{(s,a,r,s^{\prime})\sim\mathcal{D}\cup\widehat{\mathcal{D}}}\left[\left(\mu_{Q_{\psi_{k}}^{\pi_{\xi}}}-\mathcal{T}^{\pi_{\xi}}Q(s,a)\right)^{2}\right],\quad\text{
for }k=1,2,$ (6)

where we use the mean of the Q-value distribution
$Q_{\psi_{k}}^{\pi_{\xi}}(s,a)=\mathcal{N}(\mu_{Q_{\psi_{k}}^{\pi_{\xi}}},\sigma_{Q_{\psi_{k}}^{\pi_{\xi}}}^{2})$.
The target is

$\displaystyle\mathcal{T}^{\pi_{\xi}}Q(s,a)=\begin{cases}r+\gamma\mathds{E}_{a^{\prime}\sim\pi_{\xi}(a^{\prime}|s^{\prime})}\left[\min_{k=1,2}Q_{\bar{\psi}_{k}}^{\pi_{\xi}}(s^{\prime},a^{\prime})\right],&(s,a,r,s^{\prime})\in\mathcal{D}\\\
\widehat{\mu}-\beta\widehat{\sigma},&(s,a,r,s^{\prime})\in\widehat{\mathcal{D}}\\\
\end{cases},$

where $\bar{\psi}_{k}$ indicates that it is not updated during a stochastic
gradient update. Additionally,

$\displaystyle\widehat{\mu}$
$\displaystyle=\dfrac{1}{N_{\text{elite}}}\sum_{i=1}^{N_{\text{elite}}}\mu_{r}^{i}+\gamma\mathds{E}_{a^{\prime}\sim\pi_{\xi}(a^{\prime}|s^{\prime})}\left[\min_{k=1,2}\mu_{Q_{\bar{\psi}_{k}}^{\pi_{\xi},i}}\right],$
$\displaystyle\text{and}\qquad\widehat{\sigma}$
$\displaystyle=\sqrt{\dfrac{1}{N_{\text{elite}}}\sum_{i=1}^{N_{\text{elite}}}\Big{(}\Big{(}\mu_{r}^{i}+\gamma\mathds{E}_{a^{\prime}\sim\pi_{\xi}(a^{\prime}|s^{\prime})}\Big{[}\min_{k=1,2}\mu_{Q_{\bar{\psi}_{k}}^{\pi_{\xi},i}}\Big{]}\Big{)}-\widehat{\mu}\Big{)}^{2}}.$

Finally, we choose the policy $\pi_{\xi}$ by optimizing

$\mathcal{L}(\pi_{\xi})=-\mathds{E}_{s\sim\mathcal{D}\cup\widehat{\mathcal{D}}}\left[\min_{k=1,2}Q_{\psi_{k}}^{\pi_{\xi}}(s,a)-\alpha\log\pi_{\xi}(a|s)\right],$
(7)

where $\alpha$ is either a fixed or learnable hyperparameter.

###  4.2 Theoretical analysis

The key outcome of this analysis is that it is easier to obtain tight upper
bounds on value approximation errors using moment matching (Theorem 2)
compared to Monte Carlo sampling (Theorem 1). We provide proofs for all
statements in this section in Appendix B.

We evaluate the approximation error in terms of the Wasserstein distance,
which is defined as follows.

######  Definition 1 (_Wasserstein distance_).

Let $(M,d)$ denote a metric space, and let $p\in[1,\infty]$. The _Wasserstein
$p$-distance_ between two probability measures $\rho_{1}$ and $\rho_{2}$ on
$M$, with finite $p$-moments, is

$W_{p}(\rho_{1},\rho_{2})\triangleq\inf_{\gamma\in\Gamma(\rho_{1},\rho_{2})}\left(\mathds{E}_{(x,y)\sim\gamma}\left[d(x,y)^{p}\right]\right)^{\nicefrac{{1}}{{p}}},$
(8)

where $\Gamma(\cdot,\cdot)$ represents the set of all couplings of the two
measures, $\rho_{1}$, and $\rho_{2}$.

For $p=1$, the Wasserstein 1-distance, $W_{1}$, can be simplified (see, e.g.,
Villani et al., , 2009) to

$W_{1}(\rho_{1},\rho_{2})=\int_{\mathds{R}}\big{|}F_{1}(x)-F_{2}(x)\big{|}d(x).$
(9)

where $\rho_{1},\rho_{2}$ are two probability measures on $\mathds{R}$ and
$F_{1}(\cdot)$, $F_{2}(\cdot)$, their respective cdfs.

Given this definition, we can prove the following two bounds on $W_{1}$.

######  Lemma 3 (_Wasserstein inequalities_).

The following inequalities hold for the 1-Wasserstein metric

  * (i)

$W_{1}(\rho_{U}^{1},\rho_{U}^{2})\leq||A||_{1}W_{1}(\rho_{V}^{1},\rho_{V}^{2})$,
for $U=AV+b$, where $||\cdot||\triangleq||\cdot||_{1}$ is the $1$-norm.

  * (ii)

$W_{1}(\rho_{Z}^{1},\rho_{Z}^{2})\leq W_{1}(\rho_{U}^{1},\rho_{U}^{2})$, for
$Z=f(U)$ with $f(\cdot)$ locally $K$-Lipschitz with $K\leq 1$,

where $\rho_{x}$ is the density of $x$, $U\in\mathds{R}^{d_{u}}$,
$V\in\mathds{R}^{d_{v}}$, $A\in\mathds{R}^{d_{u}\times d_{v}}$, and
$b\in\mathds{R}^{d_{u}}$.

That is, for a fully connected linear layer with ReLU activation function, a
$1$-Lipschitz mapping, we have an analytical upper bound. We generalize this
to an $L$ layer multi-layer perceptron (MLP).

######  Lemma 4 (_$W_{1}$ bound on an MLP_).

Consider a MLP $f(\cdot)$ with $L$ layers

$Y_{l}=A_{l}^{\top}X_{l-1}+b_{l},\qquad X_{l}=r(Y_{l})\qquad l=1,\ldots,L,$

where $A_{l}$, $b_{l}$ are the weight matrix and bias vector for layer $l$,
and $r(\cdot)$ denotes an elementwise $1$-Lipschitz activation function. We
write $Y\triangleq Y_{L}=f(X_{0})$, where $X_{0}$ is the input. Assuming two
measures $X_{1}\sim\rho_{X}^{1}$ and $X_{2}\sim\rho_{X}^{2}$, we have that

$W_{1}(\rho_{Y}^{1},\rho_{Y}^{2})\leq\prod_{l=1}^{L}||A_{l}||W_{1}(\rho_{X}^{1},\rho_{X}^{2}).$

I.e., the $W_{1}$ distance between two measures on the output distribution is
upper bounded by the $W_{1}$ distance between the corresponding input measures
with a factor given by the product of the norms of the weight matrices
$A_{l}$. This result allows us to provide a probabilistic bound on the $W_{1}$
distance between a normal $\rho_{Y}$ and its empirical, sampling-based,
estimate $\hat{\rho}_{Y}$.

######  Theorem 1 (_Sampling-based MLP bound_).

For an $L$ layer MLP $f(\cdot)$ as defined in Lemma 4, with $Y\triangleq
Y_{L}=f(X)$ where ${X\sim\mathcal{N}(X|\mu_{X},\sigma^{2}_{X})}$, the
following bound holds:

$\mathds{P}\left(W_{1}(\rho_{Y},\rho_{\hat{Y}}^{N})\leq\prod_{l=1}^{L}||A_{l}||\sqrt{-\frac{8\log(\nicefrac{{\delta}}{{4}})R_{\max}^{2}}{\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}}\right)\geq
1-\delta.$

where $\delta\in(0,1)$ and $\hat{\rho}_{Y}^{N}$ is the empirical estimate
given $N$ i.i.d. samples.

While Lemma 3 focused on generic densities and the relation between their
$W_{1}$ distances, we can prove a concrete, analytical bound for our specific
moment-matching pipeline. In this case, we bound the $W_{1}$ distance between
the true density $\rho_{Y}$ and its moment-matching approximation
$\rho_{\widetilde{Y}}$.

######  Lemma 5 (_Moment-matching bound_).

For the following three random variables

$X\sim\rho_{X},\qquad
Y=\max(0,X),\qquad\widetilde{X}\sim\mathcal{N}(\widetilde{X}|\widetilde{\mu},\widetilde{\sigma}^{2}),$

with $\widetilde{\mu}=\mathds{E}\left[Y\right]$ and
$\widetilde{\sigma}^{2}=\mathrm{var}\left[Y\right]$ the following inequality
holds

$W_{1}(\rho_{Y},\rho_{\widetilde{X}})\leq\int_{-\infty}^{0}F_{\widetilde{X}}(u)du+W_{1}(\rho_{X},\rho_{\widetilde{X}}).$

If $\rho_{X}=\mathcal{N}(X|\mu,\sigma)$, it can be further simplified to

$W_{1}(\rho_{Y},\rho_{\widetilde{X}})\leq\widetilde{\sigma}\phi\left(\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)-\widetilde{\mu}\Phi\left(-\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)+|\mu-\widetilde{\mu}|+|\sigma-\widetilde{\sigma}|.$

where $\Phi(\cdot)$ is the cdf of a standard normal distribution.

Applying this to moment-matching MLP with $L$ layers, we get the following
deterministic bound.

######  Theorem 2 (_Moment-matching MLP bound_).

Let $f(X)$, be an MLP with $L$ fully connected layers, and ReLU activation
${r(x)=\max(0,x)}$. For $l=1,\ldots,L-1$, the sampling-based forward-pass is

$Y_{0}=X_{s},\qquad Y_{l}=r(f_{l}(Y_{l-1})),\qquad Y_{L}=f_{L}(Y_{L-1}),$

where $f_{l}(\cdot)$ is the $l$-th layer and $X_{s}$ a sample of
$\mathcal{N}(X|\mu_{X},\sigma^{2}_{X})$. Its moment matching pendant is

$\widetilde{X}_{0}\sim\mathcal{N}(\widetilde{X}_{0}|\mu_{X},\sigma_{X}^{2}),\qquad\widetilde{X}_{l}\sim\mathcal{N}\left(\widetilde{X}_{l}\Big{|}\mathds{E}\left[r(f_{l}(\widetilde{X}_{l-1}))\right],\mathrm{var}\left[r(f_{l}(\widetilde{X}_{l-1}))\right]\right).$

The following inequality holds for
$\widetilde{\rho}_{Y}=\rho_{\widetilde{X}_{L}}=\mathcal{N}(\widetilde{X}_{L}|\mathds{E}[f(\widetilde{X}_{L-1})],\mathrm{var}[f(\widetilde{X}_{L-1})])$.

$W_{1}(\rho_{Y},\widetilde{\rho}_{Y})\leq\sum_{l=2}^{L}\big{(}G(\widetilde{X}_{l-1})+C_{l-1}\big{)}\prod_{j=l}^{L}||A_{j}||,$
$\text{with}\hskip
10.00002ptG(\widetilde{X}_{l})=\widetilde{\sigma}_{l}\phi\left(\frac{\widetilde{\mu}_{l}}{\widetilde{\sigma}_{l}}\right)-\widetilde{\mu}_{l}\Phi\left(-\frac{\widetilde{\mu}_{l}}{\widetilde{\sigma}_{l}}\right)\leq
1,\qquad
C_{l}\leq\left|A_{l}\widetilde{\mu}_{l-1}-\widetilde{\mu}_{l}\right|+\left|\sqrt{A_{l}^{2}\widetilde{\sigma}^{2}_{l-1}}-\widetilde{\sigma}_{l}\right|,$

where $\sqrt{\cdot}$ and $(\cdot)^{2}$ are applied elementwise.

Comparing Theorem 2 and Theorem 1, the former is tighter than the latter, as
${R_{\max}^{2}/(1-\gamma)^{2}}$ is in practice a lot larger than the
corresponding terms in Theorem 2. For instance, the maximum realizable step
reward of the HalfCheetah-v4 environment of the MuJoCo physics engine (Todorov
et al., , 2012) is around $11.7$ (Hui et al., , 2023). If the discount factor
is chosen as $\gamma=0.99$ following common practice, one would have, e.g.,
$R_{\max}^{2}/(1-\gamma)^{2}=$1.37\text{\times}{10}^{6}$$. Contrarily, the
Lipschitz continuities of MLPs with two layers, used in the experiments, are
much smaller and in the low hundreds (Khromov and Singh, , 2024). Remarkably,
Theorem 2 the bound makes an exact statement on the maximum possible value
estimation error on a moment-matching MLP, as opposed to the high-probability
bound presented in Theorem 1 for the sampling-based MLP. Hence, a sampling-
based MLP would need a significantly large number of samples to match the
tightness of the moment-matching MLP bound and only with high probability. We
therefore expect our method to be more robust than sampling-based approaches,
which is confirmed by our experimental results.

##  5 Experiments

Table 1: Normalized reward at 3M gradient steps and Area Under the normalized
reward Curve (AUC) (mean
$\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>scores are averaged across four repetitions for the MuJoCo domain of the D4RL
offline RL dataset. The highest normalized reward and AUC scores are
highlighted in bold and are underlined if they fall within one standard
deviation of the best score. The average normalized score is the average
across all tasks, and the average ranking is calculated based on the rank of
the mean.

Dataset Type Environment Normalized Reward ($\uparrow$) AUC ($\uparrow$) MOPO
MOBILE MOMBO MOPO MOBILE MOMBO random halfcheetah
$37.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.6$
$41.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.1$ $\bf
43.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.1$
$36.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.0$
$39.5\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.2$ $\bf
41.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.0$ hopper $\bf
31.7\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.1$
$31.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.1$
$25.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>10.2$∗ $\bf
28.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.4$
$23.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.7$
$17.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.3$ walker2d
$8.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>5.6$ $\bf
22.1\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.9$
$\underline{21.5\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.1}$
$5.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.2$
$18.0\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.4$ $\bf
19.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.5$ medium halfcheetah
$72.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.2$
$\underline{75.8\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.8}$ $\bf
76.1\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.8$
$70.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.0$
$\underline{72.1\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.0}$ $\bf
73.0\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.9$ hopper
$62.8\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>38.1$
$103.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.0$ $\bf
104.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.5$
$37.0\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>15.3$
$82.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>7.3$ $\bf
95.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.5$ walker2d
$85.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.9$ $\bf
88.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.5$
$\underline{86.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.2}$
$77.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.3$
$79.0\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.3$ $\bf
84.0\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.1$ medium-replay halfcheetah $\bf
72.1\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.8$
$\underline{71.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.2}$
$\underline{72.0\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.3}$
$\underline{68.4\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.7}$
$\underline{67.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.8}$ $\bf
68.7\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.9$ hopper
$92.7\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>20.7$ $\bf
105.1\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.3$
$\underline{104.8\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.0}$
$81.7\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.6$
$78.7\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.0$ $\bf
87.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.0$ walker2d
$85.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>5.3$ $\bf
90.5\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.7$
$\underline{89.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.8}$
$65.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>12.7$
$\underline{79.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.3}$ $\bf
80.8\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>5.6$ medium-expert halfcheetah
$83.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>12.5$
$100.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.5$ $\bf
103.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.8$
$77.1\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.0$
$\underline{94.5\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>1.8}$ $\bf
95.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.7$ hopper
$74.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>44.2$
$\underline{112.5\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.2}$ $\bf
112.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.3$
$55.6\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>17.3$
$\underline{82.7\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>7.3}$ $\bf
84.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.7$ walker2d
$108.2\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>4.3$ $\bf
114.5\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>2.2$
$\underline{113.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.9}$
$88.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>6.3$
$94.3\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>0.9$ $\bf
98.9\color[rgb]{0.2,0.2,0.2}\definecolor[named]{pgfstrokecolor}{rgb}{0.2,0.2,0.2}\pgfsys@color@gray@stroke{0.2}\pgfsys@color<EMAIL_ADDRESS>3.3$ Average Normalized Reward $67.9$ $79.8$ $79.4$ $57.7$ $67.7$ $70.5$
Average Ranking $2.7$ $1.7$ $1.7$ $2.7$ $2.2$ $1.2$

_∗ High standard deviation due to failure in one repetition, which can be
mitigated by increasing $\beta$. Median result: $31.3$_

We assess the performance of our model on twelve data sets of the D4RL offline
RL benchmark (Fu et al., , 2020). The D4RL dataset consists of data collected
from three MuJoCo environments (halfcheetah, hopper, walker2d) with behavior
policies having four degrees of expertise (random, medium, medium-replay, and
medium-expert). We use the ’v2’ versions of all data sets. To run the
experiments, we use the official repository of MOBILE (Sun et al., , 2023), in
addition to our own model implementation, which is available at anonymous.
Throughout the experiments, we follow MOBILE’s configurations, with additional
hyperparameter optimization. See Appendix C for details regarding the
hyperparameters and experimental configurations.

We compare our model with the two main uncertainty-driven, model-based offline
RL algorithms, namely MOPO (Yu et al., , 2020), and MOBILE (Sun et al., ,
2023). These methods are considered the best representatives of model-based
offline RL. They employ reward penalization based on state uncertainty (MOPO),
and value uncertainty (MOBILE) respectively. MOBILE additionally serves as a
representative of the state of the art in model-based offline RL in general.
We obtain performance scores for the baseline models from log files provided
by the OfflineRL library (Sun, , 2023). As the library does not provide log
files for MOPO with random expertise, we reran them ourselves using the
provided implementation and suggested hyperparameters.

Our results are summarized in Table 1. The performance of all three models is
evaluated based on two metrics: the _normalized reward_ and the learning
efficiency, which is quantified using the _Area Under the Curve (AUC)_. The
AUC score is calculated as the mean of the average normalized scores over ten
evaluation episodes throughout the training process. The normalized score
represents the performance of the algorithms in the final episode of training.
To further illustrate these results, we provide a visualization of the
evaluation curves for the halfcheetah environment in Figure 1. See Appendix D
for larger versions of these four plots as well as for additional evaluation
curve visualizations on the remaining environments.

##### Discussion and results.

To summarize, our experimental findings are: _(i) MOMBO matches the state of
the art_ MOBILE in normalized reward. Additionally, it outperforms other
model-based offline RL approaches like COMBO (Yu et al., , 2021) with $66.8$,
TT (Janner et al., , 2021) with $62.3$, and RAMBO (Rigter et al., , 2022) with
$67.7$ as their respective average normalized reward scores across the same
twelve tasks. These numbers are cited from Sun et al.,  (2023). _(ii) MOMBO is
more sample efficient._ It learns faster and reaches its final performance
earlier. This is reflected in its AUC score, where MOMBO outperforms the
baselines. _(iii) MOMBO is more robust._ It has a lower standard deviation of
the normalized reward in six out of twelve tasks. This indicates better
stability compared to MOBILE, which has a lower standard deviation in only
three tasks. Note that for hopper-random, our model failed to learn in one
repetition, leading to a high standard deviation and low average normalized
reward. In conclusion, our MOMBO achieves state-of-the-art performance,
exhibiting robustness and fast learning capabilities, aligning with the
theory.

Figure 1: Evaluation results on halfcheetah for four settings. The dashed,
dotted, and solid curves visualize the mean of the normalized rewards across
ten evaluation episodes and four random seeds. The shaded area covers up to
one standard deviation distance to the mean.

##  6 Conclusion

We introduced MOMBO, a novel method for model-based offline RL. The main
objective of our method is to accurately quantify the uncertainty surrounding
a Bellman target estimate caused by the error in the prediction of the next
state. We achieve this by deterministically propagating uncertainties through
the value function via moment matching and introducing pessimism by
constructing LCB on the target Q-values. We analyze our model theoretically
and evaluate its performance in various environments. Our findings may lay the
groundwork for further theoretical analysis of model-based RL algorithms in
continuous state-action spaces. Our algorithmic contributions can also be
instrumental in both model-based online and model-free offline RL setups (An
et al., , 2021; Bai et al., , 2022).

##### Limitations.

The accuracy of the learned dynamics models sets a bottleneck on the
performance of MOMBO. The choice of the confidence set radius $\beta$ is also
an overly decisive factor in model performance. MOMBO shares these two
weaknesses with the other state-of-the-art model-based offline RL methods (Yu
et al., , 2020; Lu et al., , 2021; Sun et al., , 2023), however, does not
contribute to their mitigation. Furthermore, our theoretical analyses are
based on the assumption of normal distributed Q-values, which may be improved
with heavy-tailed assumed densities. Considering the bounds, a limitation is
that we claim a tighter bound compared to a baseline bound that we derived
ourselves. However, the most important factor in that bound,
$R_{\max}^{2}/(1-\gamma)^{2}$, arises via Hoeffding’s inequality, i.e., a
classical statement. As such we assume our bound to be rigorous and not
inadvertently biased. Finally, our moment matching method is limited to
activation functions for which the first two moments are either analytically
tractable or can be approximated with sufficient accuracy. Extensions to
transformations such as, e.g., BatchNorm (Ioffe and Szegedy, , 2015), or other
activation functions remove the analytical tractability of the current
proposal. However, these are usually not used in our specific applications.

## References

  * Abbas et al., (2020) Abbas, Z., Sokota, S., Talvitie, E., and White, M. (2020).  Selective dyna-style planning under limited model capacity.  In International Conference on Machine Learning.
  * An et al., (2021) An, G., Moon, S., Kim, J., and Song, H. (2021).  Uncertainty-based offline reinforcement learning with diversified q-ensemble.  In Advances in Neural Information Processing Systems.
  * Bai et al., (2022) Bai, C., Wang, L., Yang, Z., Deng, D., Garg, A., Liu, P., and Wang, Z. (2022).  Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning.  In International Conference on Learning Representations.
  * Cheng et al., (2022) Cheng, C., Xie, T., Jiang, N., and Agarwal, A. (2022).  Adversarially trained actor critic for offline reinforcement learning.  In International Conference on Machine Learning.
  * Chhachhi and Teng, (2023) Chhachhi, S. and Teng, F. (2023).  On the 1-wasserstein distance between location-scale distributions and the effect of differential privacy.  arXiv preprint arXiv:2304.14869.
  * Feinberg et al., (2018) Feinberg, V., Wan, A., Stoica, I., Jordan, M. I., Gonzalez, J., and Levine, S. (2018).  Model-based value estimation for efficient model-free reinforcement learning.  arXiv preprint arXiv:1803.00101.
  * Frey and Hinton, (1999) Frey, B. and Hinton, G. (1999).  Variational learning in nonlinear gaussian belief networks.  Neural Computation.
  * Fu et al., (2020) Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. (2020).  D4rl: Datasets for deep data-driven reinforcement learning.
  * Fujimoto and Gu, (2021) Fujimoto, S. and Gu, S. (2021).  A minimalist approach to offline reinforcement learning.  In Advances in Neural Information Processing Systems.
  * Fujimoto et al., (2018) Fujimoto, S., Hoof, H., and Meger, D. (2018).  Addressing function approximation error in actor-critic methods.  In International Conference on Machine Learning.
  * Fujimoto et al., (2019) Fujimoto, S., Meger, D., and Precup, D. (2019).  Off-policy deep reinforcement learning without exploration.  In International Conference on Machine Learning.
  * Gast and Roth, (2018) Gast, J. and Roth, S. (2018).  Lightweight probabilistic deep networks.  In Conference on Computer Vision and Pattern Recognition.
  * Haarnoja et al., (2018) Haarnoja, T., Zhou, A., Hartikainen, K., Tucker, G., Ha, S., Tan, J., Kumar, V., Zhu, H., Gupta, A., Abbeel, P., et al. (2018).  Soft actor-critic algorithms and applications.  arXiv preprint arXiv:1812.05905.
  * Haußmann et al., (2019) Haußmann, M., Hamprecht, F., and Kandemir, M. (2019).  Deep active learning with adaptive acquisition.  In Joint Conference on Artificial Intelligence.
  * Hoeffding, (1963) Hoeffding, W. (1963).  Probability inequalities for sums of bounded random variables.  Journal of the American Statistical Association.
  * Hui et al., (2023) Hui, D., Courville, A., and Bacon, P. (2023).  Double gumbel q-learning.  Advances in Neural Information Processing Systems.
  * Ioffe and Szegedy, (2015) Ioffe, S. and Szegedy, C. (2015).  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In International conference on machine learning.
  * Janner et al., (2019) Janner, M., Fu, J., Zhang, M., and Levine, D. (2019).  When to trust your model: Model-based policy optimization.  In Advances in Neural Information Processing Systems.
  * Janner et al., (2021) Janner, M., Li, Q., and Levine, S. (2021).  Offline reinforcement learning as one big sequence modeling problem.  In Advances in Neural Information Processing Systems.
  * Jeong et al., (2023) Jeong, J., Wang, X., Gimelfarb, M., Kim, H., Abdulhai, B., and Sanner, S. (2023).  Conservative bayesian model-based value expansion for offline policy optimization.  In International Conference on Learning Representations.
  * Jin et al., (2021) Jin, Y., Yang, Z., and Wang, Z. (2021).  Is pessimism provably efficient for offline rl?  In International Conference on Machine Learning.
  * Khromov and Singh, (2024) Khromov, G. and Singh, S. (2024).  Some fundamental aspects about lipschitz continuity of neural networks.  In International Conference on Learning Representations.
  * Kidambi et al., (2020) Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. (2020).  MOReL : Model-based offline reinforcement learning.  In Advances in Neural Information Processing Systems.
  * Kingma and Ba, (2015) Kingma, D. and Ba, J. (2015).  Adam: A method for stochastic optimization.  In International Conference on Learning Representations.
  * Kostrikov et al., (2021) Kostrikov, I., Fergus, R., Tompson, J., and Nachum, O. (2021).  Offline reinforcement learning with fisher divergence critic regularization.  In International Conference on Machine Learning.
  * Kumar et al., (2019) Kumar, A., Fu, J., Soh, M., Tucker, G., and Levine, S. (2019).  Stabilizing off-policy q-learning via bootstrapping error reduction.  In Advances in Neural Information Processing Systems.
  * Kumar et al., (2020) Kumar, A., Zhou, A., Tucker, G., and Levine, S. (2020).  Conservative q-learning for offline reinforcement learning.  In Advances in Neural Information Processing Systems.
  * Lakshminarayanan et al., (2017) Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017).  Simple and scalable predictive uncertainty estimation using deep ensembles.  In Advances in Neural Information Processing Systems.
  * Lange et al., (2012) Lange, S., Gabel, T., and Riedmiller, M. (2012).  Batch reinforcement learning.  In Reinforcement learning: State-of-the-art. Springer.
  * Levine et al., (2020) Levine, S., Kumar, A., Tucker, G., and Fu, J. (2020).  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  arXiv preprint arXiv:2005.01643.
  * Lu et al., (2021) Lu, C., Ball, P., Parker-Holder, J., Osborne, M., and Roberts, S. (2021).  Revisiting design choices in offline model-based reinforcement learning.  In International Conference on Learning Representations.
  * Luis et al., (2023) Luis, C., Bottero, A., Vinogradska, J., Berkenkamp, F., and Peters, J. (2023).  Model-based uncertainty in value functions.  In International Conference on Artificial Intelligence and Statistics.
  * Micheli et al., (2022) Micheli, V., Alonso, E., and Fleuret, F. (2022).  Transformers are sample-efficient world models.  In International Conference on Learning Representations.
  * Nie et al., (2021) Nie, X., Brunskill, E., and Wager, S. (2021).  Learning when-to-treat policies.  Journal of the American Statistical Association.
  * O’Donoghue et al., (2018) O’Donoghue, B., Osband, I., Munos, R., and Mnih, V. (2018).  The uncertainty bellman equation and exploration.  In International Conference on Machine Learning.
  * Peng et al., (2019) Peng, X., Kumar, A., Zhang, G., and Levine, S. (2019).  Advantage-weighted regression: Simple and scalable off-policy reinforcement learning.  arXiv preprint arXiv:1910.00177.
  * Pitcan, (2017) Pitcan, Y. (2017).  A note on concentration inequalities for u-statistics.  arXiv preprint arXiv:1712.06160.
  * Rigter et al., (2022) Rigter, M., Lacerda, B., and Hawes, N. (2022).  RAMBO-RL: Robust adversarial model-based offline reinforcement learning.  In Advances in Neural Information Processing Systems.
  * Shortreed et al., (2011) Shortreed, S. M., Laber, E., Lizotte, D., Stroup, T., Pineau, J., and Murphy, S. (2011).  Informing sequential clinical decision-making through reinforcement learning: an empirical study.  Machine learning.
  * Siegel et al., (2020) Siegel, S., Jost, T., Berkenkamp, F., Abdolmaleki, A., Neunert, M., Lampe, T., Hafner, R., Heess, N., and Riedmiller, M. (2020).  Keep doing what worked: Behavior modelling priors for offline reinforcement learning.  In International Conference on Learning Representations.
  * Singh et al., (2020) Singh, A., Yu, A., Yang, J., Zhang, J., Kumar, A., and Levine, S. (2020).  COG: Connecting new skills to past experience with offline reinforcement learning.  In Conference on Robot Learning.
  * Sun, (2023) Sun, Y. (2023).  Offlinerl-kit: An elegant pytorch offline reinforcement learning library.  https://github.com/yihaosun1124/OfflineRL-Kit.
  * Sun et al., (2023) Sun, Y., Zhang, J., Jia, C., Lin, H., Ye, J., and Yu, Y. (2023).  Model-bellman inconsistency for model-based offline reinforcement learning.  In International Conference on Machine Learning.
  * Sutton, (1990) Sutton, R. (1990).  Integrated architectures for learning, planning, and reacting based on approximating dynamic programming.  In Machine Learning Proceedings 1990. Morgan Kaufmann.
  * Sutton and Barto, (2018) Sutton, R. and Barto, A. (2018).  Reinforcement learning: An introduction.  MIT press.
  * Thrun and Schwartz, (1993) Thrun, S. and Schwartz, A. (1993).  Issues in using function approximation for reinforcement learning.  In Proceedings of the Fourth Connectionist Models Summer School.
  * Todorov et al., (2012) Todorov, E., Erez, T., and Tassa, Y. (2012).  Mujoco: A physics engine for model-based control.  In IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE.
  * Uehara and Sun, (2022) Uehara, M. and Sun, W. (2022).  Pessimistic model-based offline reinforcement learning under partial coverage.  In International Conference on Learning Representations.
  * Villani et al., (2009) Villani, C. et al. (2009).  Optimal transport: old and new.  Springer.
  * Wang and Manning, (2013) Wang, S. and Manning, C. (2013).  Fast dropout training.  In International Conference on Machine Learning.
  * Wasserman, (2004) Wasserman, L. (2004).  All of Statistics.  Springer New York.
  * (52) Wu, A., Nowozin, S., Meeds, E., Turner, R., Hernandez-Lobato, J., and Gaunt, A. (2019a).  Deterministic variational inference for robust bayesian neural networks.  In International Conference on Learning Representations.
  * (53) Wu, Y., Tucker, G., and Nachum, O. (2019b).  Behavior regularized offline reinforcement learning.  arXiv preprint arXiv:1911.11361.
  * Xie et al., (2021) Xie, T., Cheng, C., Jiang, N., Mineiro, P., and Agarwal, A. (2021).  Bellman-consistent pessimism for offline reinforcement learning.  Advances in Neural Information Processing Systems.
  * Yu et al., (2021) Yu, T., Kumar, A., Rafailov, R., Rajeswaran, A., Levine, S., and Finn, C. (2021).  COMBO: Conservative offline model-based policy optimization.  In Advances in Neural Information Processing Systems.
  * Yu et al., (2020) Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J., Levine, S., Finn, C., and Ma, T. (2020).  MOPO: Model-based offline policy optimization.  In Advances in Neural Information Processing Systems.

Appendix

##  Appendix A Related work

##### Offline RL.

The goal of offline RL is to derive an optimal policy from fixed datasets
collected through interactions with a behavioral policy in the environment.
This approach is mostly relevant in scenarios where direct interaction with
the environment is costly or poses potential risks. The applications of
Offline RL span various domains, including robotics (Singh et al., , 2020;
Micheli et al., , 2022) and healthcare (Shortreed et al., , 2011; Nie et al.,
, 2021). Notably, applying online off-policy RL algorithms directly to offline
RL settings often fails due to issues such as overestimation bias and
distributional shift. Research in Offline RL has evolved along two main
avenues: model-free and model-based approaches.

##### Model-based offline RL.

Dyna-style model-based RL (Sutton, , 1990; Janner et al., , 2019) algorithms
employ a dynamics model to simulate the true transition model, which generates
a synthetic dataset to enhance sample efficiency and serves as a surrogate
environment for interaction. Model-based offline RL methods vary in their
approach to applying conservatism using the dynamics model. For instance, MOPO
(Yu et al., , 2020) and MORel (Kidambi et al., , 2020) simply learn a
pessimistic value function by penalizing the MDP generated by the dynamics
model, with rewards being penalized based on different uncertainty measures of
the dynamics model. COMBO (Yu et al., , 2021) is a Dyna-style model-based
approach derived from CQL (Kumar et al., , 2020). RAMBO (Rigter et al., ,
2022) introduces conservatism by adversarially training the dynamics model to
minimize the value function, thus ensuring accurate predictions. MOBILE (Sun
et al., , 2023) penalizes the value function based on the uncertainty of the
Bellman operator through sampling. CBOP (Jeong et al., , 2023) introduces
conservatism by applying LCB to the value function with adaptive weighting of
$h$-step returns in the model-based value expansion framework (Feinberg et
al., , 2018).

##### Uncertainty-driven offline RL.

Uncertainty-driven offline RL approaches adhere to the meta-algorithm known as
pessimistic value iteration (PEVI) (Jin et al., , 2021). This meta-algorithm
introduces pessimism through a $\xi$-uncertainty quantifier to minimize the
Bellman error. In the realm of offline RL literature, practical
implementations of this meta-algorithm exist for both model-free methods, such
as EDAC (An et al., , 2021) and PBRL (Bai et al., , 2022), and model-based
approaches like MOPO (Yu et al., , 2020) and MOBILE (Sun et al., , 2023). PBRL
constructs its $\xi$-uncertainty quantifier based on the disagreement among
bootstrapped Q-functions, while EDAC incorporates uncertainty from ensemble
value functions. Model-based versions of these methods construct their
$\xi$-uncertainty quantifiers using information from the dynamics model. For
instance, MOPO leverages prediction uncertainty from the dynamics model in
various ways, as explored by Lu et al.,  (2021), while MOBILE uses uncertainty
in the value function for the synthetic dataset by sampling. Our model, MOMBO,
also implements PEVI, and our $\xi$-uncertainty quantifier is based on the
uncertainty of the target Q-value, which does not rely on sampling. Instead,
it directly propagates uncertainty from the dynamics models to the value
function thanks to moment matching.

##### Model-free offline RL.

Model-free offline RL algorithms can be broadly categorized into two groups:
policy regularization and value regularization. Policy regularization methods
aim to constrain the learned policy to prevent deviations from the behavioral
policy (Fujimoto et al., , 2019; Kumar et al., , 2019; Wu et al., 2019b, ;
Peng et al., , 2019; Siegel et al., , 2020). For instance, during training,
TD3-BC (Fujimoto and Gu, , 2021) incorporates a behavior cloning term to
regulate its policy. On the other hand, value regularization methods
(Kostrikov et al., , 2021; Xie et al., , 2021; Cheng et al., , 2022) introduce
conservatism during the value optimization phase using various approaches. For
example, CQL (Kumar et al., , 2020) penalizes Q-values associated with OOD
actions to avoid overestimation. Similarly, EDAC (An et al., , 2021),
employing ensemble value networks, applies a similar penalization strategy
based on uncertainty measures over Q-values.

##### Uncertainty propagation.

Uncertainty propagation via varying moment matching-based approaches has been
well-researched primarily in the area of Bayesian deep learning (Frey and
Hinton, , 1999; Wang and Manning, , 2013; Wu et al., 2019a, ), with
applications, e.g., in active learning (Haußmann et al., , 2019) and computer
vision (Gast and Roth, , 2018).

##  Appendix B Proofs

This section contains proofs for all theoretical statements made in the main
text.

##### Lemma 1

_(Moment matching)_. _For $X\sim\mathcal{N}(X|\mu,\sigma^{2})$ and
$Y=\max(0,X)$, we have_

$\displaystyle\textit{(i)}\hskip
10.00002pt\begin{split}\widetilde{\mu}\triangleq\mathds{E}\left[Y\right]&=\mu\Phi(\alpha)+\sigma\phi(\alpha),\\\
\widetilde{\sigma}^{2}\triangleq\mathrm{var}\left[Y\right]&=(\mu^{2}+\sigma^{2})\Phi(\alpha)+\mu\sigma\phi(\alpha)-\widetilde{\mu}^{2},\end{split}$

_where $\alpha=\mu/\sigma$, and $\phi(\cdot)$, $\Phi(\cdot)$ are the
probability density function (pdf) and cumulative distribution function (cdf)
of the standard normal distribution, respectively. Additionally, we have that_

$\textit{(ii)}\hskip
10.00002pt\widetilde{\mu}\geq\mu\qquad\text{and}\qquad\textit{(iii)}\hskip
10.00002pt\widetilde{\sigma}^{2}\leq\sigma^{2}.$

###### Proof of Lemma 1..

See Frey and Hinton,  (1999) for a derivation of _(i)_ , i.e.,
$\widetilde{\mu}$ and $\widetilde{\sigma}^{2}$.

Statement _(ii)_ holds as

$\widetilde{\mu}=\mathds{E}\left[Y\right]=\int_{0}^{\infty}xp(x)dy\geq\int_{-\infty}^{0}xp(x)dx+\int_{0}^{\infty}xp(x)dx=\mu.$

Note that

$\displaystyle\mathds{E}\left[Y^{2}\right]=\int_{0}^{\infty}x^{2}p(x)dx\leq\int_{-\infty}^{0}x^{2}p(x)dx+\int_{0}^{\infty}x^{2}p(x)dx=\mathds{E}\left[X^{2}\right].$

Therefore, we get _(iii)_ by

$\widetilde{\sigma}^{2}=\mathrm{var}\left[Y\right]=\mathds{E}\left[Y^{2}\right]-\widetilde{\mu}^{2}\leq\mathds{E}\left[X^{2}\right]-\mu^{2}=\sigma^{2},$

concluding the proof. ∎

##### Lemma 2

_(An inequality between normal cdfs)_ _For two normally distributed random
variables $X\sim\mathcal{N}(X|\mu,\sigma^{2})$ and
$\widetilde{X}\sim\mathcal{N}(\widetilde{X}|\widetilde{\mu},\widetilde{\sigma}^{2})$
with $\widetilde{\mu}\sigma\geq\mu\widetilde{\sigma}$, the following
inequality holds:_

$F_{\widetilde{X}}(u)\leq F_{X}(u),\qquad\text{for $u\leq 0$}.$

###### Proof of Lemma 2..

Assume $u\leq 0$. We have that

$F_{\widetilde{X}}(u)\leq
F_{X}(u)\quad\Leftrightarrow\quad\Phi\Big{(}\frac{u-\widetilde{\mu}}{\widetilde{\sigma}}\Big{)}\leq\Phi\Big{(}\frac{u-\mu}{\sigma}\Big{)},$

where
$\Phi(x)=\frac{1}{\sqrt{2\pi}}\int_{-\infty}^{x}\exp(-t^{2}/2)\operatorname{d}\\!{t}$
is the cumulative distribution function (cdf) of the standard normal
distribution. As it monotonically increases, this in turn is equivalent to

$\displaystyle\frac{u-\widetilde{\mu}}{\widetilde{\sigma}}\leq\frac{u-\mu}{\sigma}\quad$
$\displaystyle\Leftrightarrow u\sigma-\widetilde{\mu}\sigma\leq
u\widetilde{\sigma}-\mu\widetilde{\sigma}$ $\displaystyle\Leftrightarrow
u(\sigma-\widetilde{\sigma})\leq\widetilde{\mu}\sigma-\mu\widetilde{\sigma}$

which holds as $u(\sigma-\tilde{\sigma})<0$ via Lemma 1 for $u<0$, and
$\widetilde{\mu}\sigma\geq\mu\widetilde{\sigma}$ by assumption. ∎

##### Lemma 3

_(Wasserstein inequalities)._ _The following inequalities hold for the
1-Wasserstein metric_

  * (i)

$W_{1}(\rho_{U}^{1},\rho_{U}^{2})\leq||A||_{1}W_{1}(\rho_{V}^{1},\rho_{V}^{2})$_,
for_ $U=AV+b$, where $||\cdot||\triangleq||\cdot||_{1}$ is the $1$-norm.

  * (ii)

$W_{1}(\rho_{Z}^{1},\rho_{Z}^{2})\leq W_{1}(\rho_{U}^{1},\rho_{U}^{2})$_, for_
$Z=f(U)$ with $f(\cdot)$ locally $K$-Lipschitz with $K\leq 1$,

_where $\rho_{x}$ is the density of $x$, $U\in\mathds{R}^{d_{u}}$,
$V\in\mathds{R}^{d_{v}}$, $A\in\mathds{R}^{d_{u}\times d_{v}}$, and
$b\in\mathds{R}^{d_{u}}$._

###### Proof of Lemma 3..

Given the Kantorovich-Rubinstein duality (Villani et al., , 2009) for $W_{1}$
we have

$W_{1}(\mu,\nu)=\frac{1}{K}\sup_{||f||_{L}\leq
K}\mathds{E}_{x\sim\mu}\left[f(x)\right]-\mathds{E}_{y\sim\nu}\left[f(y)\right],$

for any two probability measures $\mu$, $\nu$, with $K>0$ and $||\cdot||_{L}$
is the $L$-Lipschitz norm.

This gives us for $K=1$ and $Y=g(X)$ that

$\displaystyle W_{1}(\mu_{Y},\nu_{Y})$ $\displaystyle=\sup_{||f||_{L}\leq
1}\mathds{E}_{y\sim\mu_{Y}}\left[f(y)\right]-\mathds{E}_{y\sim\nu_{Y}}\left[f(y)\right]$
$\displaystyle=\sup_{||f||_{L}\leq
1}\mathds{E}_{x\sim\mu_{X}}\left[f(g(x))\right]-\mathds{E}_{x\sim\nu_{X}}\left[f(g(x))\right],$
change of variables $\displaystyle\leq\sup_{||h||_{L}\leq
1}\mathds{E}_{x\sim\mu_{X}}\left[h(x)\right]-\mathds{E}_{x\sim\nu_{X}}\left[h(x)\right]=W_{1}(\mu_{X},\nu_{X})$

where the inequality holds as $f\circ g$ has a Lipschitz constant $L\leq 1$,
i.e., we end up with a supremum over a smaller class of Lipschitz functions
which we revert in the inequality. This gives us _(ii)_. We get _(i)_
analogously for $|A|\leq 1$ and for $|A|>1$ via $K=|A|$. ∎

##### Lemma 4

_( $W_{1}$ Bound on an MLP)._ _Consider a multi-layer perceptron (MLP)
$f(\cdot)$ with $L$ layers_

$Y_{l}=A_{l}^{\top}X_{l-1}+b_{l},\qquad X_{l}=r(Y_{l})\qquad l=1,\ldots,L,$

_where $A_{l}$, $b_{l}$ are the weight matrix and bias vector for layer $l$,
and $r(\cdot)$ denotes an elementwise $1$-Lipschitz activation function. We
write $Y\triangleq Y_{L}=f(X_{0})$, where $X_{0}$ is the input. Assuming two
measures $X_{1}\sim\rho_{X}^{1}$ and $X_{2}\sim\rho_{X}^{2}$, we have that_

$W_{1}(\rho_{Y}^{1},\rho_{Y}^{2})\leq\prod_{l=1}^{L}||A_{l}||W_{1}(\rho_{X}^{1},\rho_{X}^{2}).$

###### Proof of Lemma 4.

We have that

$\displaystyle W_{1}(\rho_{Y}^{1},\rho_{Y}^{2})$
$\displaystyle\leq|A_{L}|W_{1}(\rho_{X_{L-1}}^{1},\rho_{X_{L-1}}^{2}),$ via
Lemma 3 (i)
$\displaystyle=|A_{L}|W_{1}\big{(}\rho_{h(Y_{L-1})}^{1},\rho_{h(Y_{L-1})}^{2}\big{)}\leq|A_{L}|W_{1}\big{(}\rho_{Y_{L-1}}^{1},\rho_{Y_{L-1}}^{2}\big{)},$
via Lemma 3 (ii)
$\displaystyle\leq\cdots\leq\prod_{l=1}^{L}||A_{l}||W_{1}(\rho_{X}^{1},\rho_{X}^{2}).$

∎

##### Lemma 5

_(Moment-matching bound)_. _For the following three random variables_

$X\sim\rho_{X},\qquad
Y=\max(0,X),\qquad\widetilde{X}\sim\mathcal{N}(\widetilde{X}|\widetilde{\mu},\widetilde{\sigma}^{2}),$

_with $\widetilde{\mu}=\mathds{E}\left[Y\right]$ and
$\widetilde{\sigma}^{2}=\mathrm{var}\left[Y\right]$ the following inequality
holds_

$W_{1}(\rho_{Y},\rho_{\widetilde{X}})\leq\int_{-\infty}^{0}F_{\widetilde{X}}(u)du+W_{1}(\rho_{X},\rho_{\widetilde{X}}).$

_If $\rho_{X}=\mathcal{N}(X|\mu,\sigma)$, it can be further simplified to_

$W_{1}(\rho_{Y},\rho_{\widetilde{X}})\leq\widetilde{\sigma}\phi\left(\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)-\widetilde{\mu}\Phi\left(-\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)+|\mu-\widetilde{\mu}|+|\sigma-\widetilde{\sigma}|.$

_where $\Phi(\cdot)$ is the cdf of a standard normal distribution._

###### Proof of Lemma 5..

For a generic $X\sim\rho_{X}$, $Y=\max(0,X)$ and
$\widetilde{X}\sim\mathcal{N}(\widetilde{\mu},\widetilde{\sigma}^{2})$, we
have with $F_{Y}(u)=\mathds{1}_{u\geq 0}F_{X}(u)$111The indicator function
$\mathds{1}$ is defined as $\mathds{1}_{S}=1$ if the statement $S$ is true and
$0$ otherwise. and (9) that

$\displaystyle W_{1}(\rho_{Y},\rho_{\widetilde{X}})$
$\displaystyle=\int_{\mathds{R}}|F_{Y}(u)-F_{\widetilde{X}}(u)|du$
$\displaystyle=\int_{-\infty}^{0}F_{\widetilde{X}}(u)du+\int_{0}^{\infty}|F_{X}(u)-F_{\widetilde{X}}|du$
$\displaystyle\leq\int_{-\infty}^{0}F_{\widetilde{X}}(u)du+\int_{0}^{\infty}|F_{X}(u)-F_{\widetilde{X}}(u)|du+\int_{-\infty}^{0}|F_{X}(u)-F_{\widetilde{X}}(u)|du$
$\displaystyle=\int_{-\infty}^{0}F_{\widetilde{X}}(u)du+W_{1}(\rho_{X},\rho_{\widetilde{X}}).$

If $\rho_{X}=\mathcal{N}(X|\mu,\sigma)$, we can upper bound the $W_{1}$
distance as (Chhachhi and Teng, , 2023)

$W_{1}(\mathcal{N}(\mu_{1},\sigma_{1}^{2}),\mathcal{N}(\mu_{2},\sigma_{2}^{2}))\leq|\mu_{1}-\mu_{2}|+|\sigma_{1}-\sigma_{2}|,$

and evaluate the integral as

$\int_{-\infty}^{0}F_{\widetilde{X}}(u)du=\int_{-\infty}^{0}\Phi\left(\frac{u-\widetilde{\mu}}{\widetilde{\sigma}}\right)du=\widetilde{\sigma}\phi\left(\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)-\widetilde{\mu}\Phi\left(-\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right),$

where we used that

$\int\Phi(a+bx)dx=\frac{1}{b}\big{(}(a+bx)\Phi(a+bx)+\phi(a+bx)\big{)}+C,$

where $C$ is a constant. Together this gives us that

$W_{1}(\rho_{Y},\rho_{\tilde{x}})\leq\sigma\phi\left(\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)-\widetilde{\mu}\Phi\left(-\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)+|\mu-\widetilde{\mu}|+|\sigma-\widetilde{\sigma}|.$

∎

##### Theorem 1

_(Sampling-based MLP bound)_. _For an $L$ layer MLP $f(\cdot)$ as defined in
Lemma 4, with $Y\triangleq Y_{L}=f(X)$ where
${X\sim\mathcal{N}(X|\mu_{X},\sigma^{2}_{X})}$, the following bound holds:_

$\mathds{P}\left(W_{1}(\rho_{Y},\rho_{\hat{Y}}^{N})\leq\prod_{l=1}^{L}||A_{l}||\sqrt{-\frac{8\log(\nicefrac{{\delta}}{{4}})R_{\max}^{2}}{\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}}\right)\geq
1-\delta.$

_where $\delta\in(0,1)$ and $\hat{\rho}_{Y}^{N}$ is the empirical estimate
given $N$ i.i.d. samples._

###### Proof of Theorem 1..

With Hoeffding’s inequality (Wasserman, , 2004), we have that

$\mathds{P}\left(|\hat{\mu}_{N}-\mu|\geq\varepsilon\right)\leq
2\exp\left(-\frac{2\varepsilon^{2}}{\sum_{i=1}^{n}(\frac{R_{\max}}{1-\gamma}-0)^{2}}\right)=2\exp\left(-\frac{2\varepsilon^{2}N(1-\gamma)^{2}}{R_{\max}^{2}}\right),$

for the empirical mean $\hat{\mu}_{N}$ of $Y$. As variances are U-statistics
with order $m=2$ and kernel ${h=\frac{1}{2}(x-y)^{2}}$ using Hoeffding,
(1963); Pitcan,  (2017), we have for the empirical covariance of $Y$,
$\hat{\sigma}_{N}^{2}$, that

$\displaystyle\mathds{P}(\hat{\sigma}_{N}^{2}-\sigma^{2}\geq\varepsilon)$
$\displaystyle\leq\exp\left(-\frac{\varepsilon^{2}\lfloor\nicefrac{{N}}{{m}}\rfloor}{2||h||_{\infty}^{2}}\right)=\exp\left(-\frac{\varepsilon^{2}\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}{2R_{\max}^{2}}\right)$
$\displaystyle\Leftrightarrow\quad\mathds{P}(\hat{\sigma}_{N}^{2}-\sigma^{2}\leq\varepsilon)$
$\displaystyle\geq
1-\exp\left(-\frac{\varepsilon^{2}\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}{2R_{\max}^{2}}\right)$

Applying the inequality
$\hat{\sigma}_{N}\leq\sqrt{\sigma^{2}+\varepsilon}\leq\sigma+\sqrt{\varepsilon}$,
we get

$\displaystyle\mathds{P}(\hat{\sigma}_{N}\leq\sigma+\varepsilon)$
$\displaystyle\geq
1-\exp\left(-\frac{\varepsilon^{2}\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}{2R_{\max}^{2}}\right)$
$\displaystyle\Leftrightarrow\quad\mathds{P}(\hat{\sigma}_{N}-\sigma\geq\varepsilon)$
$\displaystyle\leq\exp\left(-\frac{\varepsilon^{2}\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}{2R_{\max}^{2}}\right).$

As the same inequality holds for $\sigma-\hat{\sigma}_{N}\geq\varepsilon$, a
union bound give us that

$\displaystyle\mathds{P}(|\sigma-\hat{\sigma}_{n}|\geq\varepsilon)$
$\displaystyle\leq
2\exp\left(-\frac{\varepsilon^{2}\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}{2R_{\max}^{2}}\right).$

Using an analytical upper bound (Chhachhi and Teng, , 2023) on the
Wasserstein-1 distance between two normal distributions,

$W_{1}(\mathcal{N}(\mu_{1},\sigma_{1}^{2}),\mathcal{N}(\mu_{2},\sigma_{2}^{2}))\leq|\mu_{1}-\mu_{2}|+|\sigma_{1}-\sigma_{2}|,$

we have that

$\displaystyle\mathds{P}\big{(}$ $\displaystyle
W_{1}(\mathcal{N}(\mu,\sigma^{2}),\mathcal{N}(\hat{\mu}_{N},\hat{\sigma}_{N}))\geq
2\varepsilon\big{)}\leq\mathds{P}\left(|\mu-\hat{\mu}_{N}|+|\sigma-\hat{\sigma}_{N}|\geq
2\varepsilon\right)$
$\displaystyle\leq\mathds{P}(|\hat{\mu}_{N}-\mu|\geq\varepsilon\text{ or
}|\hat{\sigma}_{N}-\sigma|\geq\varepsilon)\leq\mathds{P}(|\hat{\mu}_{N}-\mu|\geq\varepsilon)+\mathds{P}(|\hat{\sigma}_{N}-\sigma|\geq\varepsilon),$
union bound $\displaystyle\leq
2\exp(-2Nc\varepsilon^{2})+2\exp(-0.5\lfloor\nicefrac{{N}}{{2}}\rfloor
c\varepsilon^{2})\leq 4\exp(-0.5\lfloor\nicefrac{{N}}{{2}}\rfloor
c\varepsilon^{2})\triangleq\delta,$
$\displaystyle{\color[rgb]{0.4,0.4,0.4}\definecolor[named]{pgfstrokecolor}{rgb}{0.4,0.4,0.4}\pgfsys@color@gray@stroke{0.4}\pgfsys@color@gray@fill{0.4}c\triangleq(1-\gamma)^{2}/R_{\max}^{2}}.$

Solving for $\bar{\varepsilon}=2\varepsilon$ gives us

$\bar{\varepsilon}=\sqrt{-\frac{8\log(\nicefrac{{\delta}}{{4}})}{\lfloor\nicefrac{{N}}{{2}}\rfloor
c}}.$

The bound is then given as

$\displaystyle\mathds{P}\left(W_{1}(\mathcal{N}(\mu,\sigma^{2}),\mathcal{N}(\hat{\mu}_{N},\hat{\sigma}_{N}^{2})\leq\sqrt{-\frac{8\log(\nicefrac{{\delta}}{{4}})}{\lfloor\nicefrac{{N}}{{2}}\rfloor
c}}\right)\geq 1-\delta,$

which gives us with Lemma 4, that

$\displaystyle\mathds{P}\left(W_{1}(\rho_{Y},\rho_{\hat{Y}}^{N})\leq\prod_{l=1}^{L}||A_{l}||\sqrt{-\frac{8\log(\nicefrac{{\delta}}{{4}})R_{\max}^{2}}{\lfloor\nicefrac{{N}}{{2}}\rfloor(1-\gamma)^{2}}}\right)\geq
1-\delta.$

∎

##### Theorem 2

_(Moment-matching MLP bound)_. _Let $f(X)$, be an MLP with $L$ fully connected
layers, and ReLU activation ${r(x)=\max(0,x)}$. For $l=1,\ldots,L-1$, the
sampling-based forward-pass is_

$Y_{0}=X_{s},\qquad Y_{l}=r(f_{l}(Y_{l-1})),\qquad Y_{L}=f_{L}(Y_{L-1}),$

_where $f_{l}(\cdot)$ is the $l$-th layer and $X_{s}$ a sample of
$\mathcal{N}(X|\mu_{X},\sigma^{2}_{X})$. Its moment matching pendant is_

$\widetilde{X}_{0}\sim\mathcal{N}(\widetilde{X}_{0}|\mu_{X},\sigma_{X}^{2}),\qquad\widetilde{X}_{l}\sim\mathcal{N}\left(\widetilde{X}_{l}\Big{|}\mathds{E}\left[r(f_{l}(\widetilde{X}_{l-1}))\right],\mathrm{var}\left[r(f_{l}(\widetilde{X}_{l-1}))\right]\right).$

_The following inequality holds for
$\widetilde{\rho}_{Y}=\rho_{\widetilde{X}_{L}}=\mathcal{N}(\widetilde{X}_{L}|\mathds{E}[f(\widetilde{X}_{L-1})],\mathrm{var}[f(\widetilde{X}_{L-1})])$._

$W_{1}(\rho_{Y},\widetilde{\rho}_{Y})\leq\sum_{l=2}^{L}\left(G(\widetilde{X}_{l-1})+C_{l-1}\right)\prod_{j=l}^{L}||A_{j}||,$

_with_

$G(\widetilde{X}_{l})=\widetilde{\sigma}_{l}\phi\left(\frac{\widetilde{\mu}_{l}}{\widetilde{\sigma}_{l}}\right)-\widetilde{\mu}_{l}\Phi\left(-\frac{\widetilde{\mu}_{l}}{\widetilde{\sigma}_{l}}\right)\leq
1,\qquad
C_{l}\leq\left|A_{l}\widetilde{\mu}_{l-1}-\widetilde{\mu}_{l}\right|+\left|\sqrt{A_{l}^{2}\widetilde{\sigma}^{2}_{l-1}}-\widetilde{\sigma}_{l}\right|,$

_where $\sqrt{\cdot}$ and $(\cdot)^{2}$ are applied elementwise._

###### Proof of Theorem 2..

Using Lemma 5 we have that the $W_{1}$ distance between the deterministic
forward pass $\rho_{Y_{l}}$ and the matching-based forward pass
$\rho_{\widetilde{X}_{l}}$ is given as

$\displaystyle W_{1}(\rho_{Y_{l}},\rho_{\widetilde{X}_{l}})$
$\displaystyle\leq
G(\widetilde{X}_{l})+W_{1}(\rho_{f_{l}(Y_{l-1})},\rho_{\widetilde{X}_{l}}),$
where $G(X)\triangleq\int_{-\infty}^{0}F_{X}(u)du$ $\displaystyle\leq
G(\widetilde{X}_{l})+W_{1}(\rho_{f_{l}(Y_{l-1})},\rho_{f_{l}(\widetilde{X}_{l-1})})$
$\displaystyle\qquad+W_{1}(\rho_{f_{l}(\widetilde{X}_{l-1})},\rho_{\widetilde{X}_{l}}),$
via the triangle inequality $\displaystyle\leq
G(\widetilde{X}_{l})+||A_{l}||W_{1}(\rho_{Y_{l-1}},\rho_{\widetilde{X}_{l-1}})+C_{l}$
$\displaystyle{\color[rgb]{0.4,0.4,0.4}\definecolor[named]{pgfstrokecolor}{rgb}{0.4,0.4,0.4}\pgfsys@color@gray@stroke{0.4}\pgfsys@color<EMAIL_ADDRESS>W_{1}(\rho_{f_{l}(\widetilde{X}_{l-1})},\rho_{\widetilde{X}_{l}})}$
$\displaystyle\leq
G(\widetilde{X}_{l})+||A_{l}||\Big{(}G(\widetilde{X}_{l-1})$
$\displaystyle\qquad+||A_{l-1}||W_{1}(\rho_{Y_{l-2}},\rho_{\widetilde{X}_{l-2}})+C_{l-1}\Big{)}+C_{l}$

I.e., for the whole MLP we have

$\displaystyle W_{1}(\rho_{Y},\widetilde{\rho}_{Y})$
$\displaystyle=W_{1}(\rho_{Y_{L}},\rho_{\widetilde{X}_{L}})=||A_{L}||W_{1}(\rho_{Y_{L-1}},\rho_{\widetilde{X}_{L-1}})$
$\displaystyle\leq||A_{L}||\left(G(\widetilde{X}_{L-1})+||A_{L-1}||W_{1}(\rho_{Y_{L-2}},\rho_{\widetilde{X}_{L-2}})+C_{L-1}\right)$
$\displaystyle\leq\sum_{l=3}^{L}\left(G(\widetilde{X}_{l-1})+C_{l-1}\right)\prod_{j=l}^{L}||A_{j}||+W_{1}(\rho_{Y_{1}},\rho_{\widetilde{X}_{1}})\prod_{l=2}^{L}||A_{l}||,$

where, finally,

$W_{1}(\rho_{Y_{1}},\rho_{\widetilde{X}_{1}})\leq\sigma\phi\left(\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)-\widetilde{\mu}\Phi\left(-\frac{\widetilde{\mu}}{\widetilde{\sigma}}\right)+\left|A_{1}\mu_{X}-\widetilde{\mu}_{1}\right|+\left|\sqrt{A_{1}^{2}\sigma^{2}_{X}}-\widetilde{\sigma}_{1}\right|.$

$G(\widetilde{X}_{l})$ and $C_{l}$ are given as

$\displaystyle G(\widetilde{X}_{l})$
$\displaystyle=\widetilde{\sigma}_{l}\phi\left(\frac{\widetilde{\mu}_{l}}{\widetilde{\sigma}_{l}}\right)-\widetilde{\mu}_{l}\Phi\left(-\frac{\widetilde{\mu}_{l}}{\widetilde{\sigma}_{l}}\right),$
$\displaystyle C_{l}$
$\displaystyle\leq\left|A_{l}\widetilde{\mu}_{l-1}-\widetilde{\mu}_{l}\right|+\left|\sqrt{A^{2}\widetilde{\sigma}_{l-1}^{2}}-\widetilde{\sigma}_{l}\right|,$

where $\sqrt{\cdot}$ and $(\cdot)^{2}$ are applied elementwise. ∎

##  Appendix C Experimental details

###  C.1 Hyperparameters and experimental setup

In this section, we provide all the necessary details to reproduce our method.
We evaluate our model with four repetitions using the following seeds:
$[0,1,2,3]$. We run the experiments using the official repository of MOBILE
(Sun et al., ,
2023),222https://github.com/yihaosun1124/mobile/tree/4882dce878f0792a337c0a95c27f4abf7e926101
as well as our own model implementation, available at anonymous.

####  C.1.1 Dynamics model training

Following prior works (Yu et al., , 2020, 2021; Sun et al., , 2023), we model
the transition model $\mathrm{T}$ as an ensemble of neural networks that
predict the next state and reward as a Gaussian distribution. We formulate
this as
$\widehat{\mathrm{T}}_{\theta}(s^{\prime},r)=\mathcal{N}(\mu_{\theta}(s,a),\Sigma_{\theta}(s,a))$,
where $\mu_{\theta}$ and $\Sigma_{\theta}$ are neural networks that model the
parameters of the Gaussian distribution. These neural networks are
parameterized with $\theta$. Note that we learn a diagonal covariance
$\Sigma_{\theta}(a,s)$.

We use $N_{\text{ens}}=7$ ensemble elements and select $N_{\text{elite}}=5$
elite elements from the ensemble based on a validation dataset containing
$1000$ transitions. Each model in the ensemble is a $4$-layer feedforward
neural network with $200$ hidden units, and $L2$ weight decay is applied with
the following weights for each layer, starting from the first layer:
$[$2.5\text{\times}{10}^{-5}$,$5\text{\times}{10}^{-5}$,$7.5\text{\times}{10}^{-5}$,$7.5\text{\times}{10}^{-5}$,$1\text{\times}{10}^{-4}$]$.
The dynamics model is trained using Maximum Likelihood Estimation with a
learning rate of $0.001$, a batch size of $256$, and early stopping is applied
for $5$ steps using the validation dataset. The only exception is for
walker2d-medium dataset, which has a fixed number of $30$ learning episodes.
We use Adam optimizer (Kingma and Ba, , 2015) for the dynamics model.

Note that, to minimize differences and reduce training time, we use the pre-
trained dynamics models provided in Sun,  (2023), specifically those from
MOBILE.

####  C.1.2 Policy training

##### Architecture and optimization details.

Our model is trained over $3000$ episodes, with updates to the policy and
Q-function performed $1000$ times per episode, using a batch size of $256$.
The learning rate for the critic is set to $0.0003$, while for the actor, it
is $0.0001$. We use Adam optimizer (Kingma and Ba, , 2015) for both actor and
critics. We employ a cosine annealing learning rate scheduler for the actor.
Both the actor and critic architectures consist of 2-layer feedforward neural
networks with 256 hidden units. Unlike other baselines, our critic network is
capable of deterministically propagating uncertainties via moment matching.
The critic ensemble consists of two networks.

For policy optimization, we mostly follow the Soft Actor-Critic (SAC)
(Haarnoja et al., , 2018), with a difference in the policy evaluation phase.
Using MOBILE’s configuration, we apply the deterministic Bellman backup
operator instead of the soft Bellman backup operator. We set the discount
factor to $\gamma=0.99$ and the soft update parameter to $\tau=0.005$. The
$\alpha$ parameter is learned during training with the learning rate of
$0.0001$, except for hopper-medium and hopper-medium-replay, where
$\alpha=0.2$. The target entropy is set to $-\dim(\mathcal{A})$, where dim is
the number of dimensions.

Both the policy and actor are trained using a random mixture of the real
dataset $\mathcal{D}$ and synthetically generated rollouts
$\widehat{\mathcal{D}}$ with probabilities $p$ and $1-p$, respectively. The
real ratio $p=0.05$, except for halfcheetah-medium-expert, where $p=0.5$.

##### Synthetic dataset generation for policy optimization.

We follow the same procedure for synthetic dataset generation as is common in
the literature. During this phase, a batch of initial states is sampled from
the real dataset $\mathcal{D}$, and the corresponding actions are taken by the
current policy. For each state-action pair in the batch, the dynamics models
predict a Gaussian distribution over the next state and reward. One of the
elite elements from the ensemble is chosen at random, and a sample is drawn
from the predictive distribution for the next state and reward. These rollout
transitions are stored in a synthetic dataset $\widehat{\mathcal{D}}$. New
rollouts are generated at the beginning of each episode and appended to
$\widehat{\mathcal{D}}$. The batch size for rollouts is set to $50000$, and
only the rollouts from the last $5$ episodes are stored in
$\widehat{\mathcal{D}}$. The rollout lengths ($k$) for each task are shown in
Table 2.

##### Hyperparameters:

We adopted all configurations from MOBILE with the exception of $\beta$ due to
differences in the scale of the uncertainty estimators. Our choices for
$\beta$ and rollout length $k$ for our model are provided in Table 2.

Table 2: Hyperparameters for MOMBO.

Dataset Type Environment Rollout Length $k$ Penalty Coefficient $\beta$ random
halfcheetah $5$ $0.5$ hopper $5$ $4.5$ walker2d $5$ $2.5$ medium halfcheetah
$5$ $0.75$ hopper $5$ $3.0$ walker2d $5$ $0.75$ medium-replay halfcheetah $5$
$0.2$ hopper $5$ $0.15$ walker2d $1$ $0.5$ medium-expert halfcheetah $5$ $1.0$
hopper $5$ $1.5$ walker2d $1$ $1.5$

##### Time complexity.

To propagate uncertainties through a neural network, we propagate the first
two moments through the network, as shown in Algorithm 1. This procedure
involves passing through the linear layer twice, once for the first moment and
once for the second moment, along with some additional operations whose
computation time is negligible. In contrast, sampling-based methods require
$N$ forward passes where $N$ is the number of samples. Therefore, any $N\geq
2$ creates a computational overhead for sampling-based methods; for instance,
$N=10$ for MOBILE.

##### Experiment compute resources.

We perform our experiments on three computational resources: 1) Tesla V100
GPU, Intel(R) Xeon(R) Gold 6230 CPU at 2.10 GHz, and 46 GB of memory; 2)
NVIDIA Tesla A100 GPU, AMD EPYC 7F72 CPU at 3.2 GHz, and 256 GB of memory; and
3) GeForce RTX 4090 GPU, Intel(R) Core(TM) i7-14700K CPU at 5.6 GHz, and 96 GB
of memory. We measure the computation time for 1000 gradient steps as
approximately 8 seconds on the third device, which would require approximately
330 hours of computation for all the repetitions of MOMBO, assuming the
dynamics models are provided.

##  Appendix D Further results

Performance of the uncertainty-driven model-based offline RL approaches and
ours through gradient steps are provided in Figures 2 to 5.

Figure 2: Evaluation results of uncertainty-driven model-based offline RL
algorithms on random datasets: left panel halfcheetah, middle panel hopper,
and right panel walker2d. The horizontal axis depicts the gradient steps, and
the vertical axis shows the normalized episode reward. The thick
(dashed/dotted/solid) curve represents the mean of the normalized rewards
across ten evaluation episodes and four random seeds. The shaded area
indicates one standard deviation from the mean. We also report the mean and
standard deviation of the normalized reward and AUC score in the legend in
this order. Figure 3: Evaluation results of uncertainty-driven model-based
offline RL algorithms on medium datasets: left panel halfcheetah, middle panel
hopper, and right panel walker2d. The horizontal axis depicts the gradient
steps, and the vertical axis shows the normalized episode reward. The thick
(dashed/dotted/solid) curve represents the mean of the normalized rewards
across ten evaluation episodes and four random seeds. The shaded area
indicates one standard deviation from the mean. We also report the mean and
standard deviation of the normalized reward and AUC score in the legend in
this order. Figure 4: Evaluation results of uncertainty-driven model-based
offline RL algorithms on medium-replay datasets: left panel halfcheetah,
middle panel hopper, and right panel walker2d. The horizontal axis depicts the
gradient steps, and the vertical axis shows the normalized episode reward. The
thick (dashed/dotted/solid) curve represents the mean of the normalized
rewards across ten evaluation episodes and four random seeds. The shaded area
indicates one standard deviation from the mean. We also report the mean and
standard deviation of the normalized reward and AUC score in the legend in
this order. Figure 5: Evaluation results of uncertainty-driven model-based
offline RL algorithms on medium-expert datasets: left panel halfcheetah,
middle panel hopper, and right panel walker2d. The horizontal axis depicts the
gradient steps, and the vertical axis shows the normalized episode reward. The
thick (dashed/dotted/solid) curve represents the mean of the normalized
rewards across ten evaluation episodes and four random seeds. The shaded area
indicates one standard deviation from the mean. We also report the mean and
standard deviation of the normalized reward and AUC score in the legend in
this order.