Deterministic Uncertainty Propagation for Improved Model-Based Offline Reinforcement Learning
Authors: Abdullah Akgül
Manuel Haussmann
Melih Kandemir
Venue: NeurIPS 2024 poster (2024)
Keywords: offline reinforcement learning
offline model-based reinforcement learning
uncertainty propagation
moment matching
Primary area: reinforcement_learning
TL;DR: Prior model-based offline reinforcement learning suffers from high variance due to sampling-based estimation, but MOMBO addresses this by deterministically propagating uncertainties through the value function, providing novel suboptimality guarantees

Abstract: Current approaches to model-based offline reinforcement learning often incorporate uncertainty-based reward penalization to address the distributional shift problem. These approaches, commonly known as pessimistic value iteration, use Monte Carlo sampling to estimate the Bellman target to perform temporal difference-based policy evaluation. We find out that the randomness caused by this sampling step significantly delays convergence. We present a theoretical result demonstrating the strong dependency of suboptimality on the number of Monte Carlo samples taken per Bellman target calculation. Our main contribution is a deterministic approximation to the Bellman target that uses progressive moment matching, a method developed originally for deterministic variational inference. The resulting algorithm, which we call Moment Matching Offline Model-Based Policy Optimization (MOMBO), propagates the uncertainty of the next state through a nonlinear Q-network in a deterministic fashion by approximating the distributions of hidden layer activations by a normal distribution. We show that it is possible to provide tighter guarantees for the suboptimality of MOMBO than the existing Monte Carlo sampling approaches. We also observe MOMBO to converge faster than these approaches in a large set of benchmark tasks.

Deterministic Uncertainty Propagation for Improved Model-Based Offline Reinforcement Learning
Abdullah Akgül
Manuel Haußmann
Melih Kandemir
Department of Mathematics and Computer Science
University of Southern Denmark
Odense, Denmark
aakuul, haussmann, kandemi r : iimada... adu. ck
Abstract
Current approaches to model-based offline reinforcement learning often incorporate uncertainty-baat baasdd reward penalization to address the distributional shift problem. These approaches, commonly known as pessimistic value iteration, use Monte Carlo sampling estimate the Bellman target to perform temporal difference-based policy evaluation. We find out that the randomness caused by this sampling step significantly delays convergence. We present theoretical result demonstrating the strong dependency of suboptimality on the number of Monte Carlo samples taken per Bellman target calculation. Our main contribution is a deterministic approximation to the Bellman target that uses progressive moment matching, a method developed originally for deterministic variational inference. The resulting algorithm, which we call Moment Matching Offline Model- -ased Policy Optimiza- tion (MOMBO), propagates the uncertainty of the next state through a nonlinear Q-network in deterministic fashion by approximating the distributions of hidden layer activations by a normal distribution. We show that it is possible to provide tighter guarantees for the suboptimality of MOMBO than the existing Monte Carlo sampling approaches. We also observe MOMBO to converge faster than these approaches in large set of benchmark tasks.
Introduction
Offline reinforcement learning (Lange et al., 2012; Levine et al., 2020) aims to solve control task using an offline dataset without interacting with the target environment. Such an approach is essential in cases where environment interactions are expensive risky (Shortreed et al., 201 1; Singh al., 2020; Nie et al., 2021; Micheli al., 2022). Directly adapting traditional online off-policy reinforcement learning methods to offline settings often results in poor performance due to the adverse effects the distributional shift caused by the policy updates (Fujimoto et al., 2019; Kumar et al., 2019). The main reason for the incompatibility is the rapid accumulation of overestimated action-values during policy improvement. Pessimistic Value Iteration (PEVI) (Jin et al., 2021) offers theoretically justified solution to this problem that penalizes the estimated rewards of synthetic state transitions proportionally to the uncertainty of the predicted next state. The framework encompasses many of the-art offline reinforcement learning algorithms special cases (Yu et al., 2020; Sun et al., 2023).
Model-based offline reinforcement learning approaches first fit probabilistic model on the real state transitions and then supplement the real data with synthetic samples generated from this model throughout policy search (Janner et al., 2019). One line of work addresses distributional shift by constraining policy learning (Kumar et al., 2019; Fujimoto and Gu, 2021) where the policy is trained to mimic the behavior policy and is penalized based the discrepancy between its actions and those
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

of the behavior policy, similarly to behavioral cloning. A second strategy introduces conservatism to training by (i) perturbing the training objective with a high-entropy behavior policy (Kumar et al., 2020; Yu et al., 2021), (ii) penalizing state-action pairs proportionally to their estimated variance (Yu et al., 2020; An et al., 2021; Bai et al., 2022; Sun et al., 2023), or (iii) adversarially training the environment model to minimize the value function (Rigter et al., 2022) to prevent overestimation in policy evaluation for out-of-domain state-action pairs.
In the offline reinforcement learning literature, uncertainty- driven approaches exist for both model-free (An et al., 2021;
hal rcheetah -redd ux-expert--------ee expert-v2
Merrent Masching
Bai et al., 2022) and model-based settings (Yu et al., 2020;
sample
Sun al., 2023), all aimed learning a pessimistic value function (Jin et al., 2021) by penalizing it with an uncer- tainty estimator. The impact of uncertainty quantification has been investigated in both online (Abbas et al., 2020) and offline (Lu et al., 2021) scenarios, partieularly in model- based approaches, which is the focus of our work.
10 samples
camples
0.4
1000 sample
10000 samples
Dondit
0.8
0.2
Despite the strong theoretical evidence suggesting that the quality of the Bellman target uncertainties has significant influence on model quality (O' O'Donoghue et al., 2018; Luis et al., 2023; Jin et al., 2021), the existing implementations rely on Monte Carlo samples and heuristics-driven uncer- tainty quantifiers being used as reward penalizers (Yu et al., 2020; Sun et al., 2023) which results in high degree of randomness that manifests itself as significant delays in model convergence. Figure demonstrates why this is the case by toy example. The uncertainty around the esti- mated action-value of the next state does not shrink even after taking 10000 samples on the next state and passing them through Q-network. Our moment matching- based approach predicts similar mean with significantly smaller variance at the cost of only two Monte Carlo samples.
0.1-
0.0
642.5
647.5
600.0
652.5
655.0
Q(f.ce)
Figure 1: Monte Carlo Sampling. Moment match- ing offers sharp estimates of the action- value of the next state the cost of only two forward passes through a critic net- work. A similar sharpness cannot be reached even with 10000 Monte Carlo samples, which is 5000 times more costly. See Appendix C.I.1 for details.
Moment Matching versus
We identify the randomness introduced by Bellman target approximation via Monte Carlo sampling as the primary limiting factor for the performance of the PEVI approaches used in offline reinforcement learning. We have three main contributions:
(i) We present a new algorithm that employs progressive moment matching, an idea originally developed for deterministic variational inference of Bayesian neural networks (Wu et al., 2019a), for the first time to propagate deterministically environment model uncertainties through Q-function estimates. We refer to this algorithm as Moment Matching Offline MMddel-Base Policy Optimization (MOMBO).
(ii) We develop novel suboptimality guarantees for both MOMBO and existing sampling- based approaches highlighting that MOMBO is provably more efficient.
(iii) We present comprehensive experiment results showing that MOMBO significantly acceler- ates training convergence while maintaining asymptotic performance.
2 Preliminaries
We define a Markov Decision Process (MDP) as a tuple M é (S, A, r, P. P1, H) where S represents the state space and A denotes the action space. We have
Reinforcement learning.
S x A + [0, Rmax] as a bounded deterministic reward function and P:S x Ax x A(S) + j0, as the probabilistic state transition kernel, where A(S) denotes the prob- ability simplex defined over the state space. The MDP has an initial state distribution P1 E A(S), a discount factor E (0, 1). and an episode length H. We use deterministic policies π. : S - A and deterministic rewards in analytical demonstrations for simplicity and without loss of generality. Our results extend straightforwardly to probabilistic policies and reward functions. The randomness on the next state may result from the system stochasticity or the estimation error of probabilistic model trained on the data collected from deterministic environment. We define the action- value of a policy

---

for state-action pair (s,a) as
H Σ b=
- En
y h ((s), = s, a/ = a
This function maps the range (0, Humaz/1
Offline reinforcement learning algorithms perform policy search using an offline collected dataset a, s')} from the target environment via a behavior policy kg. The goal is to find a policy D = a,,
that minimizes the suboptimality, defined as SubOpt (TT; s) - (s)) - Q (s)) for an initial state s and optimal policy π². For brevity, we omit the dependency on s in SubOpt (-). Model-based reinforcement learning approaches often train state transition probabilities and reward functions on the offline dataset by maximum likelihood estimation with an assumed density, modeling these as an ensemble heteroscedastic neural networks (Lakshminarayanan et al., 2017). Mainstream offline reinforcement learning methods adopt the Dyna-style (Sutton, 1990; Janner et al., 2019), which suggests training an off-the-shelf model-free reinforcement learning algorithm on synthetic data generated from the learned environment model D using initial states from D. has been observed that mixing minibatches collected from synthetic and real datasets improves performance in both online (Janner et al., 2019) and offline (Yu et al., 2020, 2021) settings.
Scope and problem statement. We focus model-based offline reinforcement learning due to its superior performance over model-free methods (Yu et al., 2020; Kidambi et al., 2020; Yu et al., 2021: Rigter et al., 2022; Sun al., 2023). The primary challenge in offline reinforcement learning is distributional shift caused by a limited coverage of the state-action space. When the policy search algorithm probes unobserved state-action pairs during training, significant value approximation errors arise. In standard supervised learning, such errors diminish as new observations are collected. However, in reinforcement learning, overestimation errors are exploited by the policy improvement step and accumulate, resulting in a phenomenon known as the overestimation bias (Thrun and Schwartz, 1993). Techniques developed overcome this problem in the online setting such as min-clipping (Fujimoto et al., 2018) are insufficient for offline reinforcement learning. Algorithms that reward penalization proportional are commonly referred to as Pessimistic Value Iteration (Yu et al., 2020; Sun et al., 2023). These algorithms are grounded in a theory that provides performance improvement guarantees by bounding the variance of next state predictions (Jin et al., 2021; Uchara and Sun, 2022). We demonstrate the theoretical and practical limitations of PEVI-based approaches and address them with new solution.
the estimated uncertainty of unobserved next state
Bellman operators. Both our analysis and the algorithmic solution build on a problem that arises from the inaccurate estimation of the Bellman targets in critic training. The enror stems from the fact that during training the agent has access only to the sample Bellman operator,
B a, s') r(s,a) + 7Q(s', (s5)),
where s' is Monte Carlo sample of the next state. However, the actual quantity of interest is the exact Bellman operator that is equivalent the expectation of the sample Bellman operator,
B=Q(s, a) A EgrP(1a)) [..,((,,,
3 Pessimistic value iteration algorithms
A pessimistic value iteration algorithm (Jin et al., 2021), denoted as ApEVI Q, P). performs Dyna-style model-based learning using the following pessimistic sample Bellman target for temporal difference updates during critic training:
B Q a, 8') (, ((. -
where @ (s, a) admits learned state transition kernel P. and value function Q map state-action pair penalty score. Notably, this penalty term does not depend on an observed previously sampled as it quantifies the uncertainty around s' using P. The rest follows as running an off- the-shelf model-free online reinforcement learning algorithm, for instance Soft Actor-Critie (SAC)
3

---

(Haarnoja et al., 2018), on a replay buffer that comprises blend of real observations and synthetic data generated from P using the recent policy of intermediate training stage. We study the following two PEVI variants in this paper due to their representative value:
Model-based Offline Policy Optimization (MOPO) (Yu al., 2020) directly penalizes the uncertainty of state as inferred by the learned environment model r (s, a) - varp[s']. MOPO belongs to the family of methods where penalties are based on the uncertainty the next state.
ii) MOdel-Beilee Badel-Bellmann Inconsistency penalized offLinE Policy Optimization (MOBILE) (Sun et al., 2023) propagates the uncertainty the next state to the Bellman target through Monte Carlo sampling and uses it as a penalty:
(s, 1 BB
where var. 3 denotes the empirical variance of the quantity in brackets with respect to N samples drawn i.i.d. from P. MOBILE represents the family of methods that penalize rewards based on the uncertainty of the Bellman target.
Both approaches approximate the mean of the Bellman target by evaluating the sample Bellman operator B,, (s, a, s') with single s' available either from real environment interaction within the dataset or a single sample taken from P. The following theorem establishes the critical role the penalty term plays in the model performance.
Theorem (Suboptimality of PEVI (Jin et al., 2021)). For any derived with Apsvr [B P) that satisfies
B Qu,Q(ss - aB, s')l (s, V(s,a) E S.A.
with probability least 1 - for some error tolerance E (0, 1) where ~ P(-|8, a), the following inequality holds for any initial state 81 E S:
SubOpt(T) < E. (1)
When deep neural networks are used as value function approximators, calculating their variance becomes analytically intractable even for normally distributed inputs. Consequently, reward penal- ties 12 are typically derived from Monte Carlo samples, which are prone to high estimator variance (see Figure 1). Our key finding is that using vighavariance estimates for reward penalization introduces three major practical issues in the training process of offline reinforcement learning algorithms:
(i) The information content of the distorted gradient signal shrinks, causing delayed convergence the asymptotic performance.
(ii) The first two moments of the Bellman target poorly approximated for feasible sample counts.
(iii) Larger 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 unnecessarily restricts model capacity.
3.1 Theoretical analysis of sampling-basee based PEVI
We motivate the benefit of our deterministic uncertainty propagation scheme by demonstrating the prohibitive relationship of the sample Bellman operator when used in the PEVI context. The analysis of the distribution around the Bellman operator be characterized as follows. We are interested in the distribution of a deterministic map = f(x) for a normally distributed input æ ~ A (H, o²). We analyze this complex object via surrogate of it. For sample set SN { j obtained i.i.d. from (1, d²). let μN and of be its empirical mean and variance. Now consider the following
two random variables yn v EM fxx)) and v'y (('') for x' ~ ~N(x,, Note that yı

---

and gí are equal in distribution. Furthermore, both UN and v'v converge to the true as N tends to infinity. We perform our analysis on gyy where analytical guarantees are easier to build. Furthermore, v'y is tighter approximation of both y1 and v/. We construct the following suboptimality proof for the PEVI algorithmic family that approximates the uncertainty around the Bellman operator in the way v'y is built. In this theorem and elsewhere, A stands for the weights of the 1-th layer of a Multi-Layer Perceptron (MLP) which has 1-Lipschitz activation functions and ||-|| denotes L1-norm. See Appendix B.2 for the proof.
Theorem 2 (Suboptimality of sampling- baannd PEVI). For any policy Twe learned by Augae (B P) using N Monte Carlo samples to approximate the Bellman target with respect to an action- value network defined an L-layer MLP with 1-Lipschitz activation functions, the following inequality holds for any error tolerance E (0,1) with probability at least -8
L. SubOpt(TTc) S 2H IIIAII
8log()/4) [2/2] (1 -7)²
The bound in Theorem 2 is prohibitively loose since / - 7)2 is a large number in practice. For instance, the maximum realizable step reward in the Hal @elffheetaah.. environment of the MuJoCo physics engine (Todorov et al., 2012) is at least around 11.7 according to Hui et al. (2023). Choosing the usual 7 0.99, we obtain R / - (1 1.37 106, Furthermore, as the bound depends on the number of samples for the next state N, it becomes looser as N + 1. The bound is not defined for N 1, but the loosening trend is clear.
4 MOMBO: Moment matching offline model-based policy optimization
Our key contribution is the finding that quantifying the uncertainty around the Bellman target brings significant benefits the model-base based offline reinforcement learning setting. We obtain deterministic approximation using a moment matching technique originally developed for Bayesian inference. This method propagates the estimated uncertainty of the next state s' obtained from learned environment model P through an action-value network by alternating between an affine transformation of normally distributed input to another normally distributed linear activation and projecting the output of nonlinear activation to a normal distribution by analytically calculating its first two moments. The resulting normal distributed output is then used to build a lower confidence bound on the Bellman target, which an equivalent interpretation of reward penalization. Such a deterministic approximation is both a more accurate uncertainty quantifier and a more robust quantity to be used during training in a deep reinforcement learning setting. Its contribution to robust and sample-efficient training has been observed in other domains (Wang and Manning, 2013; Wu et al., 2019a). Prior work attempted to propagate full covariances through deep neural nets (see, c.g. Wu al., 2019a; Look et al., 2023; Wright et al., 2024) at a prohibitive computational cost (quadratic in the number of neurons) that does not bring a commensurate empirical benefit. Therefore, we choose to propagate only means and variances.
Assuming the input of a fully-connected layer to X N(Xlu,,")) the pre-activation Y for a neuron associated with weights 0 is given by Y X i.e., Y N 0 H, (02)To2), where we absorb the bias into 0 for simplicity and the square on 0 is applied element-wise. For a ReLU activation function r(z) 2 max(0, x). mean ğ and variance 62 of Y = max(0, X) are analytically tractable (Frey and Hinton, 1999). We approximate the output with a normal distribu- tion ~ (X (X))),,)) and summarize several properties regarding ju and ø2 below. See Appendix B. for proofs of all results in this subsection.
Lemma (Moment matching). For X N(X and = max X), we have (i) ū 2 E (a( + oΦ(a), 82 var [Y] (n² + +)))((()) + μoδ((a) ft2 where a = 11/a a and (6.), (b.)) are the probability density function (pdf) and cumulative distribution function (cdf) of the standard normal distribution, respectively. Additionally, we have that
(ii) 2 P and (ii) 32 0²
Althhough we build our algorithm with the ReLU activation function, it is also applicable to other activation functions whose first and second moments are tractable or can be approximated with sufficient accuracy. including all commonly used activation functions.
5

---

Algorithm outlines the moment matching process. This process involves two passes through the linear layer: one for the first moment and one the second, along with additional operations that take negligible computation time. In contrast, sampling- based methods require N forward passes, where N is the number of samples. Thus, for any N > sampling-han methods introduce computational overhead. MOBILE (Sun et al., 2023), e.g., uses N = 10. See Figure and Figure 3 for illustrations on the effect of varying N.
Algorithm Deterministic uncertainty propagation through moment matching 6, X)
function
8: Weights of the layer, b: bias of the layer
X = XN [X |ax, (Xxxxxx
> input a normal distribution > transform mean and variance output the transformed distribution
(UY, o + (0 ux + b, x
return Y = = (Yyu,
end function
function MOMENNTAATCHNNG/THOOUGRRRELL
X V(X|xx,03)
> input a normal distribution
< ux/ox
+ ux ((a) + +xx((a
b compute the first two moments of ReLU(X) > 6/4 are the normal pdf/edf & output a normal distribution with these moments
oy + (H'x + ox) D(a) (xxxx(a) n²
return y= ===WYlm -N(Mpx,
end function
We propagate the distribution of the next state predicted by the environment model through the value function network using moment matching. We define a moment matching Bellman target distribution as:
i a, s') E r(s. + 7Qmas x(s')
for P(-1s,a) a) and QMM (s', a') normal distribution with mean (s'a aa) and variance (s' a') obtained as the outcome of progressive application of Algorithm through the layers
of critic network Q. The sign d denotes equality in distribution, i.e., the expressions on the two sides share the same probability law. We perform pessimistic value iteration using lower confidence bound on , (( 1,s0 a, s') as the Bellman target
B BF e a(s,a) + = (700m((d,
for some radius B > 0.
4.1 Theoretical analysis of moment matching- uncertainty propagation
The following lemma provides a bound on the Wasserssteiin distance W; between the true distribution PY of random variable after being transformed by ReLU activation function and its moment matched approximation Px. See Appendix B.3 for proofs of results presented this subsection. Lemma (Moment matching bound). For the following three random variables
X ~ PX,
~xpx
y= ma
X X~N(817,79)
with ū = E Y and o2 var [Y] the following inequality holds
Wi(pr.P3) [ (x)d + Wi(px,p3)
where (-) of x. pfxx = N (2))) it can be further simplified to
Wi(py. s 76
7
+ 1p
-
Applying this to moment matching L-layer MLP yields the following deterministic bound.

---

Lemma 3 (Moment matching MLP bound). Let f(X) be an L-layer MLP with ReLU activation r(z) max(0, 2). For I.....L 1, the sampling -aanndd--bdded forward-pact pass is
Yo Xs: Yi - (1=r(((((-)))
YL fi(Yi-1)
where fi(-) is the 1-th layer and X, sample of NXX xx). Its moment matching pendant is
No N(Xollx,
X1 N [1 18 ((f((((-1)) var [r((((i-1i))
The following inequality holds for pr = PX2 = NXX var
Wi(py, -
Σ
(G(Xi-1) + CL-1) III
1=2
j=1
with G(Xi) = σ16
11
1
OISAA - Mudl A
where V and (-) are applied elementwise.
Relying on Lemma 6 again, this result provides an upper bound on the suboptimality of our moment matching approach.
Theorem 3 (Suboptimalityy of moment matching-ber -based PEVI algorithms). For any policy was derived with ApeevT [ Q P) learned by a penalization algorithm that uses moment matching approximate the Bellman target with respect to action-value network defined as an L-layer MLP with 1-Lipschitz activation functions, the following inequality holds
I. SubOpt(TTa) S 2H Σ (G(X--1++CL) II
5=1
The bound in Theorem 3 is much tighter than Theorem 2 in practice R R ((1--)) 7) is very large while the Lipschitz continuities || of the two-layer MLPs used in the experiments are in the order of low hundreds according to empirical investigations (Khromov and Singh, 2024). Another favorable property of Theorem 3 is that its statement is exact, as opposed to the probabilistic statement made in Theorem 2. The provable efficiency of MOMBO could be of independent interest for safety-critical use cases of offline reinforcement learning where the availability of analytical performance guarantees a fundamental requirement.
4.2 Implementation details of MOMBO
We adopt the model learning scheme from Model-Based Policy Optimization (MBPO) (Janner al., 2019) and use SAC (Haarnoja al., 2018) the policy search algorithm. These approaches represent the best practices in many recent model-based offline reinforcement learning methods (Yu 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 heteroscedastic neural ensemble of transition kernels modeled as Gaussian distributions over the next state and reward. We denote each ensemble element as P,, (s,'', r|s,, a) for E (I...., News }. After evaluating the performance of each model on a validation set, we select the Netuee best persotrerfoming ensemble elements for further processing. We use these elite models to generate k-step rollouts with the current policy and to create the synthetic dataset D. which then combine with the real observations D. We retain the mean and variance of the predictive distribution for the next state to propagate it through the action-value function while assigning zeto variance to the real tuples.
The lower confidence bound given in Equation (1) builds on our MDP definition, which assumes a deterministic reward function and deterministic policy for illustrative purposes. In our implemen- tation, the reward model also follows heteroscedastic ensemble of normal distributions. We also incorporate the uncertainty around the predicted reward of a synthetic sample into the Bellman target calculation. Furthermore, our policy function a squashed Gaussian distribution trained by SAC in the maximum entropy reinforcement learning setting. For further details, refer to the Appendix C.2.
7

---

Table 1: Performance evaluation the D4RL dataset. Normalized reward at 3M gradient steps and Area Under the Learning Curve (AULC) (mean-tud) scores are averaged across four repetitions. The highest means are highlighted in bold and underlined if they fall within one standard deviation of the best score. The average normalized score is the average across all tasks. The average ranking is based on the rank of the mean.
NORMALIZED REWARD (1)
AULC (t)
Datuset Type
Environment
MOPO
MOBILE
MOMBO
MOPO
MOBILE
MCMBO
halfcheszah
37.211.
41.2-11
43 43.6111
36.311.0
39.511.
41.411.
random
hopper
31.720.1
31.3=0.1
25.4
28.6+1.4
23.6+3.7
17.3=113
walker2d
8.2-6.6
22.1=0.9
21.5+0.1
5.4-a.2
18.0#a.4
19.2=0.5
Average on random
25.7
31.5
30.2
23.4
27.1
26.0
ha fcheett
72.4+4.2
75.8=0.8
76.1+0.8
70.9=2.0
72.1+1.0
73.0=0.9
am
hopper
62.8
103,6:1.0
104.2:0.5
37.0+:1.33
82.2+7.3
95. 2.55
wa walker2d
85.412.9
88.3
86,41112
77.611.3
79.0
8..0
Average on medium
73.6
89.3
88.9
61.8
77.8
84.3
halfcheszah
72.1:3.8
71.9=3.2
72.014.3
68.414.7
67 2..9.
3.39
hopper
92.7:2207
105. 11.3
104.8
81.724.6
78.724.0
87.3
walker2d
85.915.3
90.5
89.615.8
65.31
79.914.3
80.8
Average on medi um replay
83.4
89.2
88.8
72.4
75.5
78.9
fcheetah
13.6
100.9=1.5
103.3:
77.1
94.5
95.
hopper
74.9144.2
112.5:0.2
112.6:0.
55.6
82.717.3
84.3::477
walker2d
108. +...3
114,5+2.2
113.9+0.9
88.3+6.3
94.3+0.9
98.9=3.3
Average medi expext
88.9
109.3
109.9 .
73.6
90.5
92.8
Average Score
67.6
79.8
79.5
57.5
67.7
70.5
Average Ranking
2.7
1.7
1.7
2.7
2.2
1.2
High sundard deviation failure reporition, which mittgoted increaring Median memle
Experiments
We compare MOMBO against MOPO and MOBILE, the two representative PEVT variants, across twelve tasks from the D4RL. dataset (Fu et al., 2020), which consists of data collected from three MuJoCo environments (hal sralfcheetal) hopper, wa walker2d) with behavior policies exhibiting four degrees of expertise (random, medium, medi um replay, and medi um tm-uumpppt expert). We use the datasets collected with 'v2' versions of the MuJoCo environments to be commensurate with the state of the art. We evaluate model performance using two scores:
(i) Normalized Reward: Total episode reward collected evaluation mode after offline training ends, normalized by the performances of random and expert policies.
(ii) Area Under Learning Curve (AULC): Average normalized reward computed at the interme- diate steps of training.
AULC indicates how fast a model converges to its final performance. A higher AULC reflects more efficient learning other things being equal. We report the main results in Table 1 and provide the learning curves of the hal fcheetah environment in Figure 2 for visual investigation. We present the learning curves for the remaining tasks in Figure 4. We obtain the results for the baseline models from the log files provided by the OfflineRL library (Sun, 2023). We implement MOMBO into the MOBILE (Sun al., 2023) code base shared by its authors and use their experiment configurations wherever applicable. See Appendix C for details. The source code of our algorithm is available at https: gi thub. com/ adinlab ab/ ww..///// /OMBOO The OfflineRL library does not contain the log files for the random datasets for MOPO at the time of writing. We replicate these experiments using the MOBILE code based on the preseribed configurations.
Theorem indicates that the performance of PEVI algorithm depends on how tightly its reward penalizer r (s, a) upper bounds the Bellman operator error [B,Q(s, a)
We use this theoretical result to compare the quality of the reward penalizers of the three models based average values of the following two performance scores calculated across data points observed during ten evaluation episodes:

---

Ball wat.ehaaraaaaee
43-
33
2.5
1.5 50
0.5
2.5
2-3
Steps
Stope
Grodier Stepe
Figure 2: Evaluation results on hal fcheetah for four settings. The dashed, dotted, and solid curves represent the mean of the normalized rewards across ten evaluation episodes and four random seeds. The shaded area indicates one standard deviation from the mean.
Table Uncertainty quantification on the D4RL dataset. Accuracy and tightness (mean scores are averaged across four repetitions. The highest means are highlighted in bold and are underlined if they fall within one standard deviation of the best score.
ACCURACY (1)
TIGHTNESS (t)
Datasct Type
Environment
MOBILE
MOMBO
MOPO
MOBILE
MOMBO
hal talicheetah
0.02
0.25=0002
0,24=0oz
14.58=0.63
6,88=0.61
-6.21
random
hopper 0.1420.02
0.8520.03
0.61
-1.88
-0.44102
00.31
walker2d
00.01
0.55g0004
0.74
-17
+10.10"
-0.27=0.08
-0.14:0.02
hal Echeetah
0.06#0.01
0.25
0.34
15.71=0.4s
10.03
9.066.0.2
medium
hopper
0.01
0.4140.01
0.55
4.162122
3.14=0.0s
1.3:10.21
walker2d
0.02.0.0
0.1620.01
0.
5.0220.52
-4.03
hal Fcheetah
0.09
0.04
0.16
-11.67=091
-1.8522041
-10.4
pedium-reglas
hopper 0.02=0.01
0.03
0.17=0.01
-5.35
-3.4
-
-3.17
walker2d
0.0820.01
0.1410.01
0.36
4.4740.43
4.5620.33
3.78
hal fcheetah
0.1310002
0.3620.03
0.
21.25_287
13.22a0.47
-11.8810.5
medium redium-expert
hopper 0.04
0.4320.02
0.51
9.7777
3.55.00
-3.38
Nal kel2d
0.05.20.1
0.47:0.02
0.45=0.02
-9.7720.22
-5.5220.36
-5.29.0011
(i) Accuracy: 1 (s 2 [, B (s, a))) for the indicator function 1, i.e., how often the reward penalizer is valid unccertainty quantifier assumed by Theorem 1.
(ii) Tightness: (s, a) [B, ((, a) (s i.., how sharp a - uuncertaint quantifier the reward penalizer is.
Table 2 shows that MOMBO provides more precise uncertainty estimates compared the sampling- based approaches. It also indicates that MOMBO provides tighter estimates of the Bellman operator error, meaning that the sampling-ber for further details.
approaches use larger confidence radii. See Appendix C.1.2
The D4RL dataset is a heavily studied benchmark database where many hyperparameters, such as penalty coefficients, are tuned by trial and error. We argue that this is not feasible in a realistic offline reinforcement learning setting where the central assumption that policy search needs be performed without real environment interactions. Furthermore, it is more realistic to assume that collecting data from expert policies is more expensive than random exploration. One would then need to perform offline reinforcement learning on a dataset that comprises observations obtained at different expertise levels. The mixed offline reinforcement learning dataset (Hong al., 2023) satisfies both desiderata, as the baseline models are not engineered for its setting and its tasks comprise data from two policies: an expert or medium policy and random policy, presented in varying proportions. We compare MOMBO against MOBILE and MOPO on this dataset for a fixed and shared penalty coefficient for all models. We find that MOMBO outperforms both baselines in final performance and learning speed in this more realistic setup. See Appendix C.1.3 and Table 3 for further details and results.
Our key experimental findings can be summarized as follows:
(i) MOMBO converges faster and trains more stably. learns more rapidly and reaches its final performance earlier as visible from the AULC scores. Its moment matching approach

---

provides more informative gradient signals and better estimates of the first two moments of the Bellman target compared to sampling-bos based approaches, which suffer from high estimator variance caused by Monte Carlo sampling.
(ii) MOMBO delivers a competitive final training performance. It ranks highest across all tasks and outperforms other model-based offline reinforcement learning approaches, including COMBO (Yu et al., 2021) with an average normalized reward of 66.8. TT (Janner et al., 2021) with 62.3, and RAMBO (Rigter et al., 2022) with 67. 37.7. all evaluated on the same twelve tasks in the D4RL dataset. We took these scores from Sun et al. (2023).
(iii) MOMBO delivers superior final training performance when expert data is limited. MOMBO outperforms the baselines (at 1% in the mixed dataset) in both AULC and normalized reward. For normalized reward, MOMBO achieves 37.0, compared to 32.4 for MOBILE and 27.9 for MOPO, averaged across all tasks 1%. In terms of AULC, MOMBO attains 29.5, outperforming the 28.0 of MOBILE and the 23.5 of MOPO. See Table for details.
(iv) MOMBO provides more precise estimates of Bellman target uncertainty. It outperforms both baselines in accuracy for 9 out of 12 tasks and in tightness all 12 tasks in the D4RL dataset.
6
Conclusion
The main objective of MOMBO is to accurately quantify the uncertainty surrounding Bellman target estimate caused by the error in predicting the next state. We achieve this by deterministically propa- gating uncertainties through the value function with moment matching and introducing pessimism by constructing lower confidence bound 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-basee reinforcement learning algorithms spaces. Our algorithmic contributions can also be instrumental in both model-basee basell online and model-free offline reinforcement learning setups (An al., 2021; Bai et al., 2022).
continuous state-action
Limitations. The accuracy of the learned environment models sets bottleneck on the performance of MOMBO. The choice of the confidence radius is another decisive factor in model performance. model-based offline reinforcement
MOMBO shares these weaknesses with the other state-of-the-ee learning methods (Yu al., 2020; Lu et al., 2021; Sun et al., 2023) and does not contribute their mitigation. Furthermore, our theoretical analyses rely on the assumption of normally distributed Q-values, which may be improved with heavy-tailed assumed densities. Considering the bounds, a limitation is that claim a tighter bound compared However, the most important factor in that bound, n?2 / (1- arises from Hoeffding 's inequality, i.e., 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 either analytically tractable can be approximated with sufficient accuracy, which includes all commonly used activation functions. Extensions to transformations such as, c.g.. BatchNorm (loffe and Szegedy, 2015), other activation functions remove the analytical tractability for moment matching. However, these are usually not used in our specific applications.
bascline bound that we derived ourselves.
Acknowledgments and Disclosure of F Funding
AA and MH thank the Carlsberg Foundation for supporting their research under the grant number CF21-0250. We are grateful to Birgit Debrabant for her valuable inputs in the development of some theoretical aspects. We also thank all reviewers and the area chair for improving the quality of our work by their thorough reviews and active discussion during the rebuttal phase.
10

---

References
Abbas, Z., Sokota, S., Talvitie, E., and White, M. (2020). Selective Dyna-style planning under limited model capacity. In International Conference Machine Learning.
An, G., Moon, S., Kim, J.-H., and Song, H. O. (2021). Uncertainty based offline reinforcement learning with diversified Qenneeble. Advances Neural Information Processing Systems.
Bai, C., Wang, L., Yang, Z., Deng, Z., Garg, A., Liu, P., and Wang, Z. (2022). Pessimistic bootstrapping for
uncertainty- offline reinforcement leamning In International Conference on Learning Representations.
Casella, G. and Berger, R. (2024). Statistical inference. CRC Press.
Cetin, E., Tirinzoni, A., Pirotta, M., Lazaric, A., Ollivier, Y., and Touati, A. (2024). Simple ingredients for offline reinforcement learning. In International Conference Machine Learning.
Cheng, C.-A,, Xie, T., Jiang, N., and Agarwal, A. (2022). Adversarially trained actor critic for offline reinforce- ment learning. In International Conference on Machine Learning-
Chhachhi, S. and Teng, F. (2023). On the I-wasserstein distance between location scale distributions and the effect of differential privacy. arXiv preprint arXiv:2.004. 14869.
Feinberg, V., Wan, A., Stoica, I., Jordan, M. L., Gonzalez, J. E., and Levine, S. (2018). Model- based value ostimation for efficient model-free reinforeement learning. arXiv preprint wwx.ivv18030011110
Frey, B. J. and Hinton, G. F. (1999). Variational learning in nonlincar Gaussian belief networks. Neural Computation.
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. (2020). D4RL: Datasets for deep data-driven reinforcement learning.
Fujimoto, S. and Gu, S. S. (2021). A minimalist approach to offline reinforcement learning. In Advances in Neurud Information Processing Systems.
Fujimoto, S., Hoof, H., and Meger, D. (2018). Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning.
Fujimoto, S., Meger, D., and Precup, D. (2019). Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning.
Gast, J. and Roth, S. (2018). Lightweight probabilistic deep networks. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition.
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:
Haußmann, M., Hamprecht, F. A., and Kandemir, M. (2019). Deep active learning with adaptive acquisition. In International Joint Conference Artificial Intelligence.
Hoeffding, W. (1963). Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association.
Hong, Z., Agrawal, P., Combes, R. T. d., and Laroche, R. (2023). Harnessing mixed offline reinforcement learning datasets via trajectory weighting. In International Conference on Learning Representations.
Hui, D. Y.-T., Courville, A. C., and Bacon, P.L. (2023). Double Gumbel Q-learning. In Advances in Neural Information Processing Systems.
loffe, and Szegedy, C. (2015). Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning.
Janner, M., Fu, J., Zhang, M., and Levine, S. (2019). When to trust your model: Model-based policy optimization.
In Advances in Neural Information Processing Systems.
Janner, M., Li, Q.. and Levine, S. (2021). Offline reinforcement learning as one big sequence modeling problem.
In Advances in Newral Information Processing Systems.
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, Y., Yang. Z.. and Wang. Z. (2021). Is pessimism provably efficient for offline RL? In International Conference on Machine Learning.
Khromov, G. and Singh, S. P. (2024). Some fundamental aspects about Lipschitz continuity of neural networks. In International Conference on Learning Representations.
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims. T. (2020). MOReL: Mooeel-baee offline reinforce-
ment learning. In Advances Neural Information Processing Systems.
Kingma, D. P. and Ba, J. L. (2015). Adam: A method for stochastic optimization. In International Conference on Learning Representations.
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, A., Fu, J., Soh, M., Tucker, G., and Levine, S. (2019). Stabilizing off off-po Q-learnng via bootstrapping reduction. In Advances in Neural Information Processing Systems.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. (2020). Conservative Q-leaming for offline reinforcement learning. In Advances in Neural Information Processing Systems.
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, S., Gabel, T., and Riedmiller, M. (2012). Batch reinforcement learning. In Reinforcement learning: State-of-the-art Springer.
Levine, S., Kumar, A., Tucker, G., and Fu, J. (2020). Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0166
Look, A., Kandemir, M., Rakitsch, B.. and Peters, (2023). A deterministic approximation neural SDEs. IEEE Trunsactions on Pattern Analysis and Machine Intelligence.
Lu, C., Ball, P. J., Parker-Holder, J., Osborne, M. A., and Roberts, J. (2021). Revisiting design choices in offline model-based reinforeement learning. In International Conference on Learning Representations.
Luis, C. E., Bottero, A. G., Vinogradška, J., Berkenkamp, F., and Peters, (2023). Model-based uncertainty in value functions. In International Conference on Artificial Intelligence and Statistics.
Micheli, V., Alonso, E., and Fleuret, F. (2022). Transformers are sample-efficient world models. In International Conference on Learning Representations.
Nie, X., Brunskill, E., and Wager, S. (2021). Leaming when-to-treat policies. Journal of the American Statistical Association.
O'Donoghue, B., Osband, I., Munos, R., and Mnih, (2018). The uncertainty Bellman equation and exploration. In International Conference on Machine Learning.
Peng, X. B., Kumar, A., Zhang, G., Levine, (2019). wevintage-eeighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 1910.00177.
Pitcan, Y. (2017). A note on concentration inequalities for U-statistics. arXiv preprint 17206017
Rigter, M., Lacerda, B., and Hawes, N. (2022). RAMBO-RL: Robust adversarial model-based offline reinforce- ment learning. In Advances in Neural Information Processing Systems.
Shortreed, S. M., Laber, E., Lizotte, D. J., Stroup. T. S., Pineau, J., and Murphy, S. A. (2011). Informing sequential clinical macision through reinforcement learning: empirical study. Machine learning.
Siegel, N. Y., Springenberg. J. T., Berkenkamp. F., Abdolmaleki, A., Neunert, M., Lampe, 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, 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, Y. (2023). Offlinerl kit: An elegant pytorch offline reinforcement leamning library. attpss//gittee omtyyhhacuun124//
12

---

Sun. Y.. Zhang. J. Jia, C.. Lin, H., Ye, J.. and Yu. Y. (2023). Model-Bellman inconsistency for model-basee offline reinforcement learning. In International Conference Machine Learning.
Sutton, R. S. (1990). Integrated architectures for learning, planning. and reacting based approximating dynamic programming. In Machine Learning Proceedings 1990. Morgan Kaufmann.
Thrun, S. and Schwartz. A. (1993). Issues in using function approximation for reinforcement learning. In Proceedings of the 1993 Connectionist Models Summer School.
Todorov, E.,, Frez, T., and Tassa, Y. (2012). MuloCa: A physies engine for model-based control. In IEFF/RS) International Conference on Intelligent Robots and Systems.
Uehara, M. and Sun, W. (2022). Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference on Learning Representations.
Villani, C. al. (2009). Optimal transport: old and new. Springer.
Wang, S. and Manning, C. (2013). Fast dropout training. In International Conference on Machine Learning. Wasserman, L. (2004). All of Statistics. Springer New York.
Wright, O., Nakahira, Y., and Moura, J. (2024). analytic solution to covariance propagation neural networks. In International Conference on Artificial Intelligence and Statistics.
Wu, A., Nowozin, S., Meeds, E., Turner, E., Hernandez-Lobato, M., and Gaunt, A. L. (2019a). Determin- istic variational inference for robust Bayesian neural networks. In International Conference on Learning Representations.
Wu, Y., Tucker, G., and Nachum, O. (2019b). Behavior regularized offline reinforcement learning. arXiv preprint arXiv: /911./136.111
Xie, T., Cheng, C.-A,, Jiang. N., Mineiro, P. and Agarwal, A. (2021). reinforcement learning. In Advances in Neural Information Processing Systems.
pessimism for offline
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, T.. Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma. T. (2020). MOPO: Model-based offline policy optimization. In Advances in Neural Information Processing Syxtems.
13

---

APPENDIX
A
Related work
Offline reinforcement learning. The goal of offline reinforcement learning is to derive an optimal policy from fixed datasets collected through interactions with behavior policy in the environment. This approach is particularly relevant in scenarios where direct interaction with the environment is costly or poses potential risks. The applications of offline reinforcement learning span various domains, including roboties (Singh et al., 2020; Micheli al., 2022) and healthcare (Shortreed et al., 2011; I; Nie et al., 2021). Notably, applying off-policy reinforcement learning algorithms directly to offline reinforcement learning settings often fails due to issues such as overestimation bias and distributional shift. Research in offline reinforcement learning has evolved along two main avenues: model-free and model-based approaches.
Model-based offline reinforcement learning. Dyna-style model-based reinforcement learning (Sutton, 1990; Janner et al., 2019) algorithms employ a environment model to simulate the true transition model, generating a synthetic dataset that enhances sample efficiency and serves as a surrogate environment for interaction. Model-based offline reinforcement learning methods vary in their approaches to applying conservatism using the environment model. For instance, MOPO (Yu et al., 2020) and MORel (Kidambi et al., 2020) learn a pessimistic value function by penalizing the MDP generated by the environment model, with rewards being penalized based on different uncertainty measures of the environment model. COMBO (Yu et al., 2021) is a Dyna-style model- based approach derived from conservative Q-learning (Kumar et al., 2020). RAMBO (Rigter et al., 2022) introduces conservatism by adversarially training the environment model to minimize the value function, thus ensuring accurate predictions. MOBILE (Sun et al., 2023) penalizes the value function based the uncertainty of the Bellman operator through sampling. CBOP (Jeong et al., 2023) introduces conservatism by applying a lower confidence bound to the value function with an adaptive weighting of h-step returns in the model-based value expansion framework (Feinberg et al., 2018) 20.8).
Uncertainty-driven offline reinforcement learning approaches apply penalties in various ways. As model-free methods, EDAC (An et al., 2021) applies penalties based ensemble similar- ities, while PBRL (Bai et al., 2022) penalizes based on the disagreement among bootstrapped Q-functions. Uncertainty-driven model-based offline reinforcement learning algorithms adhere the meta-algorithm known as pessimistic value iteration (Jin et al., 2021). This meta-algorithm introduces pessimism through ξ-uncertainty quantifier to minimize the Bellman approximation error. MOPO (Yu et al., 2020) and MOBILE (Sun et al., 2023) can be seen as practical implementations of this meta-algorithm. MOPO leverages prediction uncertainty from the environment model in various ways, as explored by Lu et al. (2021), while MOBILE uses uncertainty in the value function for the synthetic dataset through sampling. MOMBO is also based on PEVI, and our g-uncertainty quantifier is derived from the uncertainty of the Bellman target value, which does not rely on sampling. Instead, directly propagates uncertainty from the environment models to the value function, thanks to moment matching.
Model-free offline reinforcement learning 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 behavior 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 out-of of-iistiibutiin actions to avoid overestimation. Similarly, EDAC (An et al., 2021), which employs ensemble value networks, applies a similar penalization strategy based on uncertainty measures over Q-values.
Uncertainty propagation via various moment matching- approaches has been well- researched, primarily in the area of Bayesian deep learning (Frey and Hinton, 1999; Wang and
14

---

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).
B Theoretical analysis
This section contains proofs for all the theoretical statements made in the main text.
B.1
Prerequisites
We first restate and extend the following result moment matching with the ReLU activation function. Our focus on the ReLU activation throughout this work is due to the fact that is currently the most commonly used activation function in deep reinforcement learning. One could derive similar results for any piecewise linear activation function and several others as well.
Lemma 1 (Moment matching). For X ~ NN(Xu, and Y max X), have
(i) T T ( + aδ(a),
22 var [M] (n² + + + pań(a) - pt
where - u//o, and ((-). (-)) are the probability density function (pdf) and cumulative distribution function (cdf) of the standard normal distribution, respectively. Additionally, we have that
(ii)
22#
μ and
(iii)
o².
Proof of Lemma 1. See Frey and Hinton (1999) for the derivation of (i), i... and 0².
Statement (ii) holds as
j F. Y ( xp(x)dd 2 [" xp(x))dx + 1/² xp(x))d = μ.
Note that
x'p(x)dx S 10 d'p(z)de + +/ E E [
Y =
Therefore, we get (iii) by
σ² var
Y =E y nr E x2 1 02
concluding the proof.
Although we do not require the following lemma in our final results, a helpful result is the follow- ing inequality between the edfs of two normally distributed random variables X and its matched counterpart X.
Lemma (An inequality between normal cdfs). For two normally distributed random vari- ables X (X o²) and X NN (X,0000 6²) with jıo pō, the following inequality holds:
Px(u) S Fx(u) for S O.
Given ji and 7 as in Lemma 1, the assumptions of this lemma hold and it is applicable to our moment matching propagation.
Proof of Lemma 4. Assume a < 0. We have that
Φ ("" +
(x( (u)
As is a monotonically increasing function, this in turn is equivalent to
"
+ uo - Na S uō μō
# u(u - ō) < júo μō
which holds as u(a -ö) < 0 via Lemmal for 0, and juo μ by assumption.
15

---

We provide the following definition, as we derive our bounds using Wasserstein distances.
Definition 1 (Wasserstein distance). Let (M, d) denote a metric space, and let p [1] The W-ssssssstein distance between two probability measures pt and p2 on M, with finite p-moments, is defined as
(P
inf TEx(P1.2))
(Teaas) 1)1) 4/
where -(: represents the set of all couplings of two measures, and p2.
For 1, the 1-Wasserstein distance, Wi, can be simplified (see, e.g. Villani et al., 2009) to
Wi(P1, / [F(()) - -F2(x)
(2)
where Pl, P2 are two probability measures on R, and Fi(-) and F2(-) are their respective cdfs. Given this definition, we prove the following two bounds on W1.
Lemma (Wasserstein inequalities). The following inequalities hold for the --assersteiin metric
(i) Wi (et pt)) S [[A][WI(P (V) U = AV + b. (i) Wp (2) & (().b..)).... for Z = g(U) with g(-) locally K-Lipschitz with K S 1 where Par is the density of x, E Rd_ V V Rd. A Rd.xd. and E Rd..
Proof of Lemma 5. Given the Cantorovich-Rubinee duality (Villani et al., 2009) for Wi we have
1
WH( = sup Ezzn [f(x)] Ey [f())] lip())
K
for any two probability measures μ1, and where lip(f) < K specifies the family of K K-Lipschitz functions.
We use this to first prove (ii). For Z = g(U) we have that
Wi (ob,,
sup
Ez-p' [f())] - E Ez~r3 [f))]
lip(S)S1
sup Eurpp [f((()))] Em [(((())),
change of variables
lip(S)11
< sup Eu-pt [b[u]] - Euppy [M(4]] = Wi Wibb, Pt) lip(h)
where the inequality holds as fo g has Lipschitz constant L < 1, i.e. we end up with supremum over smaller class of Lipschitz functions which we revert the inequality. This gives the desired inequality.
To prove (i) we use that for ||A||| < 1, the transformation g(V) & is I-Lipschitz and the result follows via (ii). For ||A|| > we use that g(-) is K K-Lipschitz and the result follows by adapting the proof accordingly.
Finally, the following lemma allows relate the 1-Wasserstein distance of two transformed distributions to an absolute difference between the respective transforming functions.
Lemma 6 (Wasserstein suboptimality). Consider two probability distributions Pa, P., defined
measarable space (S, (S)) and two transformed random variables = f(x) ~ Py (with x ~ Pa).
and g(u) P- (with for functions f.g R. If
(PG P²)S0
then FigeP [f(=)] En [g(w)] sC
for any constant C20 0.
Proof of Lemma 6. Assume that Wi(P, P.) < C for some C 0. The Rubbiissiinnnuinteeenn duality (Villani et al., 2009), with K gives us
WP,,,,PP) sup Ey-p, [h(v)] Ez-P, [h(2)] < C. 1ip(h)
16

---

As the identity function À + x: 1-Lipschitz, we can lower bound the supremum further with
Ey [g] P. [2]
sup Ev-,, [h(y)] Px~P. [h(=)] - C.
lip(k)
Using what is colloquially known the law of the unconscious statistician (Casella and Berger, 2024). we have that
Ey [y] = Fig [f(x)]
and analogously for P. [z]. Therefore
F. EymP, [ Eg [2] P. [f(x)] Eu [g(4)]
which gives us
Eg [f(x)] EurP. [a(u)] SC
As Wi is symmetric, we can follow the same argument and additionally have
EuxP, [g(u)] Ex-P, [f(x)] S C.
Using that for E It with < C and -x < C. it follows that [x|| C, we can combine these two inequalities and get
ExmP, [f(x)] EunP, [g(u]] S C
as desired.
B.2 Proof of theorem? 2
Proving Theorem 2 requires the following upper W1 bound for an MLP with 1-Lipschitz activation function.
Lemma 7 (W1 bound on an MLP). Consider L-layer MLP f(-)
= AT XI-1 + br, X r(=i(Y))
1= 1=
where AG bt are the weight matrix and bias vector for layer I, respectively, and r(-) denotes an elementwise 1-Lipschitz activation function. We write Y - Yr f(Xo). where Xo the input. Assuming two measures X1 ~ p'x and X2 pix, have that
Wi 13)
A// P2)
Proof of Lemma 7. We have that
W wi(pt,p3) < || P2
via Lemma 5 (i)
= WI (PhcYa-.): Ph(Ye-.)) S ||Ar (P via Lemma (ii)
1=
The W1 distance between measures the output distribution is upper bounded by the Wi distance between the corresponding input measures with a factor given by the product of the norms of the weight matrices As. This result allows us provide probabilistic bound the W, distance between a normal py and its empirical, sampling-ben based,, estimate pr
The following lemma allows to extend this result to probabilistic sampling-basee based upper bound. Lemma 8 (Sampling-based MLP bound). For an L-layer MLP f(-)with 1-Lipschitz activation function, let Y Yı f(X) where X ~ NN xxxx0x) the following bound holds:
8 12
P W
Wilpy,
P S II II
21 1
[1/2
1=1
where 8 (0, 1) and P is the empirical estimate given N i.i.d. samples.
17

---

Proof of Lemma 8. With Hoeffding's inequality (Wasserman, 2004), we have that
P PIAN 14) 8) 2exp
-2N2
Rmax
2 0 - )
2 exp
2 2=2N(11 - 7)²
R
for the empirical mean ju of Y. As variances are U-statistics with order 2 and ker- h 1 y)² using Hoeffding (1963); Pitcan (2017), we have for the empirical covariance of Y, N. that
52 [ exp
121N/2[11 7)2
P(O 222 S exp
2R,
c2 [2[N/2] - -)2
P( ṭ² < - cxp
2R max
where |1:1|oo denotes L -orm. Applying the inequality ēN V0² +0 VE, we get
E² 1211/21(11 2)
P(ôN 2 exp
2R
=2 [V/2][11 7)2
P(ōN 0 2 2) exp
2R
As the same inequality holds for o - ōN 2 c. union bound give us that
8² N/22(1 -)
P(lo dnl 2 x) s 2xxp
2/2
Using analytical upper bound (Chhachhi and Teng, 2023) the 1-Wasserstein distance between two normal distributions.
Wi W3) < | u22 la 02l
have that
23) < (Ip - Pul 0 ĀN 2 2s)
< P(IĀN - 14| 2 or Ôx - al 2 e)
< P(JĀN - -l 2 5) + P(lx al E),
union bound
< 2exp(-22cc2 + (x.5 c=*)
8,
where we use the shorthand notation c (1 7) R2 R
Solving for 2c gives us
(lg((%))
[N/2]c
The bound is then given
()4)
P wi (.. S
2 2-8
W//2cc
which gives us with Lemma that
L <
8 (log(%/)).t) R22 mmax
P Wi(py,pp)
II 1=1
Arll
2 21-
[1/2](1 7)2
Finally, applying this result to the relation Lemma allows us to prove the desired non-asymptotic bound on the performance of the camlinggbbseed model.
18

---

Theorem 2 (Suboptimality of sampling- basedd PEVI algorithms). For any policy Tue learned by A (BF P) using N Monte Carlo samples approximate the Bellman target with respect to action-value network defined as an L-layer MLP with L-Lipschitz activation functions, the following inequality holds for any error tolerance å E (0, 1) with probability at least
L SubOpt(TTc) & 2H TI
8 Blog($/4) max
(1-21(1-))
Then, combining the bound from Lemma 8 with the result from Lemma get that with probability 1-6
Proof of Theorem 2.
Define C = Il
8 (2g( Rags
[1
Ex~N(xx03))
BaQ(X) - Ex-N(ix.33)
B [x(())
sc
B,Q(s,a) - Ex 03)
[xQ(x)
sc
B=Q(s,a) Q,s,
SC
Note that X = (s, a,,,,)) s'), i.e., the state, action, and next state tuple, is distributed as described the main text. Therefore C' f-uncertainty quantifier, with { 6. Applying Theorem 1 finalizes the proof.
B.3 Proof' of theorem 3
In order to prove Theorem i.c., an upper bound on our proposed moment matching approach, require the following two lemmata.
The first lemma allows to upper bound the 1- W-Wasserstein distance between input distribution and its output distribution after transformation by the ReLU activation function.
Lemma (Moment matching bound). For the following three random variables
X ~ PX,
== ==m
~~N(X80,0)
with j = E [Y] and T2 = var [Y] the following inequality holds
Wi(pr,p3) 1
Fx(u)du Wi(px,px)
where (x is edf of X. Ifxx N(X)u,0)) it can be further simplified
Wi(py. s T6
7
+ Ip -
Proof of Lemma 2. For a generic X px. Y max(0, X) and X ~ N(ū, we have with (u) Fx (w)² and (2) that
Wi(pr.pg)
[Fy()) -F((a) JR
[ Ex(u)du+ + / 00 [x((u) Fx
Fx(u)du (
[x(()) (x) du + 1 [Fx(u)- Fx(u)) du
TTh indicator function 1 is detined ls 1 if the statement S is true and otherwise.
19

---

(u(uddu+ + Wi (ix,x
If px = ( (1. o), we can upper bound the W1 distance via Chhachhi and Teng (2023)
Wi(W(HI, of, oB)) S | u2l + loi - 02l
and evaluate the integral
[ (u)du [00 + 11
du σΦ 14
7
where we used that
/ D(a+ bx)d - 5
((a + (() (a + bx) + o(a + bx))
where - signifies equality up to an additive constant.
Together this gives us that
71 σ
Wilpy, S oō
+ 1u
N + 10
Lemma 3 (Moment matching MLP bound). Let f(X) be an L-layer MLP with ReLU activation r(z) = ma x). For = 1..... L - 1, the sampling campling-based based forward-pa pass is
Yo Xs: Yi = r(J((Vi-)))
Y,. = (Y((1,))
where fi(-) is the 1-th layer and X, sample of ((Xlux,03) Its moment matching pendant is
No ox) 8 N X E (((((())))]] var (
The following inequality holds for Po = Pxx = 1
varl (Xx
1 Wi(py,Pv pX) Σ
(G(Xi-1) + Cl-1) II
t=2
ju!
with G(Xi) = ōzΦ
1
C s A Tán +
| A 5
where V: and (-)2 are applied elementwise.
Proof of Lemma 3. Using Lemma we have that the W1 distance between the deterministic forward pass py; and the matching-baser forward pass Px, is given as
where G(X) A L (x(ud
Wi Wi S G(X)) + Wi(PnCn-D)P.))..
S G(X)) + WiGPACN .)
+ Wi (P
via the triangle inequality
< <G(x) + (Px + Ct.
Ci Wilp (x)))) Px,)
$ G(X) + |AN ((X
+ A|_ [| (pY (() + +Ci-1) + Ct.
That is, for the entire MLP have
Wilpy, pr) = Wi(py. ((x.)):Px))))))) || -14
< A (G(XL-1) + | ,:PRe..) + C2-1)
20

---

Σ (((x)-)) + Ci-1)
I IT Wi(en IIIIAp
where, finally,
16
Px( S oΦ
+ +A11x - -aal+ V A703 01
G(Xi) and C given
G(Xi) ō.
11
Cr S |A - -7sl+ | - - -
where V and (.)2 are applied elementwise.
Theorem 3 (Suboptimality of moment matching-based PEVI algorithms). For any policy whs Q, P) learned by penalization algorithm that uses moment matching to approximate the Bellman target with respect an action-value network defined as an L-layer MLP with 1-Lipschitz activation functions, the following inequality holds
derived with Avevi (
L I1
SubOpt(T S 2H
Σ ( 1=2
j=l
Proof of Theorem 3. The proof analogous to the one of Theorem 2. Note that X = (s, (,,0,8)) 8)) i.e. the state-action pair, is distributed as described in the main text.
Define C= = E1 (G(( + Cc-1) -1 By combining this theorem with the results from
Theorem 1, we have that
Ex-N(xxx0))
[xQ(()) Ex~V(ix
I [(x)
sc
B=Q(s,a) - Ex [+Q(())
cC
BaQ(s, a) -B. BB,
sC
holds deterministically. Therefore, C'is a f-uncertainty quantifier for & = 8 = 0. Applying Theorem finalizes the proof.
C Further details on experiments
C.1 Experiment procedures
C.1.1 Moment matching versus Monte Carlo sampling experiment
We sample an initial state from the environments. Next, we use the expert policies x², treated as optimal policies, provided in the D4RL datasets (Fu al., 2020) to sample action ai from the expert policy. Using the learned environment model P we predict the next state samples 82 ~ P(-|s1, a1). For each sample, we evaluate the next action-value Q(22, n (s2)) using the previously learned critics the value function. We then evaluate the next action-value using moment matching QMM(s2, n (82)) and visualize the distributions. For this experiment, we use the learned critics from seed We repeat this process for all tasks in the D4RL dataset and present the results in Figure 3.
"https://rall eecs bee.keeeeee.......eee......... http://...ee//aaaatett/e//aa///p
21

---

ba wechetabbrandieenn
hopper randowme2
sampla
33-
ammpo
0.55
aumpp
serple
10 serple
13 arrpls
10-
100 sample
110 sumpla
220 sumpla
1000 sample
1000 samplo
3023 samplo
10000 compless
2.5
12600
30000 comppie
Dentyy
Don2
80-
8..0
0.15-
13-
8.00
4050
400
095.5
907
167
9980
QU4
Q(6,2)
hoppar-radione
2il vulbac2i-nadiuur-
Marest Matcinn
1.5-
Motert Matchng
0.00
Museo Matshin
1.25
samples
smmpls
sample
sarpie
1.33
surples
23 smmpa
1.50-
comple
corplen
cirppee
samples
10000 sumpla
20000 sumplis
Dosity
8:00-
1..2
0.75
0.25
0.00
0.03
20
0.00
Q(".e))
Q(Viv))
QU(.4))
halr coeet wwweetah--aee www.hah-euue-eeeeee ue-d:ue---rrpp..... www.eerrpllg-aaee...... Ing--12
Sopper--ne wwd.- Lay--22
colkerr' meddeeddeee ux-oueie--ggg-g
0.12
Morret Manchin
2.5
Noment Matching
Moment Matching
1.75
corplee
cirple
1.50
cumples
23
cxmple
sample
1000 somples
sur plos
3002 samplees
10000 curplex
10000 evrple
DATS
10030 complee
1.3-
1.00-
0.00
14:00
0.02
0125
060
470
679
671
250
133
2.93
100
140
110
1.5
300
7.0
0.0
QU v/)
QU4
balf www.cah---aaae export-v22
bogger wod-ur---d------eo export-v2
vallerr
expext-v22
Matching comple
0.14
Momani Matching
0.5
cumplex
cmppec
100 comples
sarple
23
10 eurple
100 coeples
33 ariple
100 camples
1L4
10000
10000
0.10
20000
cample
1.5
Denaky
1133
DxTS-
0.2
0/31
113
112
642.5
060.0
647.5
600DD
592.5
010
340
0/00
Q(4.2))
Figure 3: Moment Matching versus Monte Carlo Sampling. A comparison of moment matching and Monte Carlo sampling methods for estimating the next value for all tasks in the D4RL dataset.
22

---

Table 3: Performance evaluation the mixed dataset. Normalized reward at 2M gradient steps and Area Under the Learning Curve (AULC) (mean-tud) scores are averaged across four repetitions. The highest means are highlighted in bold and underlined if they fall within standard deviation of the best score. The average normalized score is the average across all tasks, and the average ranking based the rank of the mean.
Dataset
NORMALIZED REWARD (1)
AULC (*)
Ratio (%)
Dutaset Type
Environment
MOBILE
MO3ILE
MOMBO
ha fchectah
42.1::12.0
51.312.7
55.4:21
41.4
46.0s2.1
48.5
randon-medis
hopper
48. 4
22.9
30.4
35.5230
21.8a2.
27.3228
walker2d
0.140.0
15.9
22. 22.0=0.4
0.1=0.0
13.1
18. 18.7
aahecheetah
33.5
51.8-20
41.0=16.9
31.2
38.2:s.0
32.3
random-sxport
hopper
42.8
30.7
51.5
32.3:
32.7.7.
32.0
walker2d
-.1a
21.61000
21.8
0.3
16.6.0.9
18.28110
Average on
27.9
32.4
37.0
23.5
28.0
29.5
fcheeaht
1227
62.4=1.4
62.2+33.
54.3a11
56.4
58 1=:11
random medium
hopper
32 02.0
16.8+11.8
60 0227.
29.5+3.7
38.3+a.
walker2d
5.248.4
10.6
21.7:0.1
4.7:45
10.3
18.5:007
fchestah
33.8 2
54.0122.0
39.1a197
33.612.
49.7:10.6
30.3
random eendom-expert
hopper
43 6:30
43.0 23,0
17.1a126
29.2
1447
39 62333
walker2d
1.6+2.0
13.8+7.4
24.0=2.6
2.0=1.2
13.3
18 44:18
Average on
29.7
35.0
halfchestah
66. 2:07
65.4+0.7
59.5+1.
59.9
2:112
random-nedium
hopper
37.3:90
45.0+220
58 18.4
33.0+2.1
58.4
60.1:
10
walker2d
17.8
24.5
10.9
16.8
halfcheetah
29. 29.5
74.5-8.2
58.1+16.
32.9
64.
62 Aac.
random-expert
hopper
20.0s5.7
80.3 3:214
77.8+20.5
19.2+2.1
40.3+6.1
45. 5:
walker2d
14.9
21.3:4.0
14.0:3.4
17.2::10
Average 10
33.6
48.5
50.2
28.7
41.3
43.9
halfcheetah
73. 4
72,0+0.8
73.5=1.0
68.4+9.2
67.5 07.3
69. 3:0.8
random random-mediing
hopper
67.2 2.345
106.5
0.0
34.1
80.9a5.
9
walker2d
61 1330.6
36.1#25..
14.7
58. 38.9=16.5
29. 29.62:1.2
37.7
halfcheetah
12.1
72.9 22.9:22.1
33.4
30.7
47. 47.8
44.1
randem-axpert
hopper
24.7
70.3+:9.8
61.33:33 a
22.51:18
62.5
walker2d
16.0a11.4
15.7
16.4=9.7
1.7=1.1
9.8
11.8
Average on
48.0
62.2
55.8
37.1
49.6
51.9
Average Score
38.5
Average Ranking
2.5
2.0
1.5
2.6
1.9
1.5
C.1.2 Uncertainty quantifier experiment
We generate 10 episode rollouts using the real environment and previously trained policies in evaluation mode. From these rollouts, we select state, action, reward, and next state tuples at every 1hth step, including the final step. For each selected tuple, we caleulate the tightness scores using the learned environment models. We repeat this process for each of the 4 seeds across every task and report the mean and standard deviation of the accuracies and tightness scores. We estimate the exact Bellman target by taking samples from the policies, and the number of samples 1000. Table shows the results of this experiment.
accuracy and
C.1.3 Mixed offline reinforcement learning dataset experiment
The behavior policy of the mixed datasets consists of combination of random and medi um or expert policies, mixed a specified demonstration ratio. Throughout the experiments, we use the configurations of Cetin et al. (2024). See Appendix C.2 for details regarding the hyperparameters. Results in Table 3 show that MOMBO converges faster with minimal expert input, indicating that it effectively leverages limited information while maintaining competitive performance. Compared to other PEVI approaches, MOMBO demonstrates faster learning, greater robustness, and provides strong asymptotic performance, particularly under the constraints of limited expert knowledge and fixed training budgets, where hyperparameter tuning not feasible. See Figures for visualizations of the learning curves.
23

---

Discount factor (?) Table 4: Common hyperparameters and sources used in the experimental pipeline.
D4RL
Mixed
Sources
Dataset
Fu al. (2020)
Hong al. (2023) Hong et al. (2023)
Configuration Implementation Expert policy
Sun et al. (2023)
Sun et al. (2023)
Fu et al. (2020)
Pretrained environment models
Sun (2023)
Policy learning
Seeds
[0. 2]
Learning rate actor Learning rate critic Hidden dimensions actor
0.0001
0.0003
[256, 256 [256, 256
[1024, 1024, 1024, 1024, 1024
Hidden dimensions critic
[256, 256, 256
0.99
Soft update (T) Number of critics Batch size Learning rate scheduler Epoch Number of steps per epoch
0.005
256
Cosine Annealing for Actor
3000
2000
1000 1000 steps
Rollout frequency
Rollout length (A)
See Table See Table 5
5
Penalty coefficient (3)
2.0
Rollout batch size
50000
Model retain epochs
5
10
Length of fake buffer [D]
5 50000
10 50000
Real ratio (p)
0.05 (except ha fcheetah -ediuu-t---diii----ea- cectun-experttee 0.5)
Environment model learning
Learning rate Early stop Hidden dimensions
0.001
2000 200, 200, 200
1.2 weight decay Number of ensembles N Nene
[2.5 x 10-5,5 10 7.5 x 10 7.5 10-5, 10-1)
Number of elites Natite
C.2 Hyperparameters and experimental setup
In this section, we provide all the necessary details to reproduce MOMBO. We evaluate MOMBO with four repetitions using the following seeds: [0, 1, 2, We run the experiments using the official repository of MOBILE (Sun et al., 2023),* as well our own model implementation, available https: / /iggttttuubb thub. com/ a a /OMBO. We list the hyperparameters related to the experimental pipeline in Table 4.
C.2.1 Environment model training
Following prior works (Yu et al., 2020, 2021; Sun al., 2023), we model the transition model neural network ensemble that predicts the next state and reward as Gaussian distribution. We formulate this Pa(s', ,r) = - (s,a), a). [s a)), where μ4 and Eg are neural networks that model the parameters of the Gaussian distribution. These neural networks are parameterized by 0 and learn diagonal covariance Ea(a, 8).
'nttps: con/ 7126611
37c0a95cd

---

We use Nens = 7 ensemble elements and select Nome = 5 elite elements from the ensemble based on a validation dataset containing 1000 transitions. Each model in the ensemble consists of 4-layer feedforward neural network with 200 hidden units, and we apply 7.2 weight decay with the following weights for each layer, starting from the first layer: [2.5x 10 x 10 7.5x 10-6, 5, 7.5 x 10-5, 1 10-2]. We train the environment model using maximum likelihood estimation with learning rate of 0.001 and a batch size of 256. We apply early stopping with 5 steps using the validation dataset. The only exception is for wa walker2d -wweedddddddddd medium um dataset, which has a fixed number of 30 learning episodes. We use Adam optimizer (Kingma and Ba, 2015) for the environment model.
We use the pre-trained environment models provided in Sun (2023) to minimize differences and reduce training time. For the mixed dataset experiments, we train environment models using four seeds (0, 1, 3) and use them for each baseline.
C.2.2 Policy training
Architecture and optimization details. We train MOMBO for 3000 episodes on the D4RI. dataset and 2000 episodes on the mixed dataset, performing updates to the policy and Q-function 1000 times per episode with a batch size of 256. We set the learning rate for the critics to 0.0003, while the learning rate for the actor is 0.0001. We use the Adam optimizer (Kingma and Ba, 2015) for both the actor and critics. We employ a cosine annealing learning rate scheduler for the actor. For the D4RL dataset, both the actor and critics architectures consist of 2-layer feedforward neural networks with 256 hidden units. For the mixed dataset, the actor uses a 5-layer feedforward neural network with 1024 hidden units, and the critics consist of 3-layer feedforward neural network with 256 hidden units. Unlike other baselines, our critic network is capable of deterministically propagating uncertainties via moment matching. The eritic ensemble consists of two networks. For policy optimization, we mostly follow the 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 0.99 and the soft update parameter to 0.005. The a parameter is learned during training with the learning rate of 0.0001, except for hopper-met -e um and hopper seppee--mddie i um replay, where a = 0.2. We set the target entropy to (A), where dim is the number of dimensions. We train all networks using a random mixture of the real dataset D and synthetically generated rollouts D with real and synthetic data ratios of p and 1 p, respectively. We set p 0.05, except for hal fcheetal wwwaaaiiiineeee iidiumm--- uanllum-ppppp experr,, where p = 0.5.
Synthetic dataset generation for policy optimization. We follow the same procedure for synthetic dataset generation as it is common in the literature. During this phase, we sample a batch of initial states from the real dataset D and sample the corresponding actions from the current policy. For each state-action pair in the batch, the environment models predict a Gaussian distribution over the next state and reward. We choose one of the elite elements from the ensemble uniformly at random and take sample from the predicted distribution for the next state and reward. We store these rollout transitions in synthetic dataset D. where we also keep their variances in the buffer. We generate new rollouts at the beginning of each episode and append them to D. We set the batch size for rollouts to 50000, and store only the rollouts from the last 5 and 10 episodes in DD for the D4RL dataset and the mixed dataset, respectively. For the mixed dataset, the rollout length (k) is 5 for all tasks. Table presents the rollout length for each task for the D4RL dataset.
Penalty coefficients. We adopt all configurations from MOBILE with the exception of 33 due to the differences in the scale of the uncertainty estimators. For the D4RL dataset, we provide our choices for 3 in Table 5. For the mixed dataset, we select penalty coefficient / as 2.0 for all tasks in order to apply a penalty corresponding to 95% confidence level. We provide all the other shared hyperparameters in Table 4.
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 1 GPU, AMD EPYC 7F72 CPU at 3.2 GHz, and 256 GB of memory; and 3) GeForce RTX 4090 GPU, Intel(R) Core(TM) 17-147000 CPU at 5.6 GHz, and 96 GB of memory. We measure the computation time for 1000 gradient steps to be approximately seconds for the D4RI. dataset and 11 seconds for the mixed dataset on the third device. Assuming the environment models
25

---

Table 5: Hyperparameters of MOMBO the D4RL dataset.
Dataset Type
Environment
Rollout Length &t
Penalty Coefficient 3
halfcheetah
5
0.5
random
hopper
4.5
walker2d
5
2.5
halfcheetah
5
0.75
medi um
hopper
5
3.0
walker2d
5
0.75
hal fcheetah
medium-replay
hopper
5
0.15
walker2d
0.5
halfcheetah
5
1.0
expeut-
hopper
5
1.5
walker2d
1
1.5
are provided and excluding evaluation time, the total time required reproduce all MOMBO repeti- tions is approximately 330 hours for D4RL and 600 hours for the mixed dataset. The computational cost for other experiments is negligible.
C.3 Visualizations of learning curves
Figures to 6 illustrate the learning curves of PEVI-based approaches, including ours, across gradient steps. In the figures, the thick (dashed/dooted/sollid) curve represents the mean normalized rewards across ten evaluation episodes and four random seeds, with the shaded area indicating one standard deviation from the mean. The legend provides the mean and standard deviation of the normalized reward and AULC scores in this order. We show the results for hal fcheetah in the left panel, hopper in the middle panel, and walker2d in the right panel. The horizontal axis represents gradient steps and the vertical axis shows normalized episode rewards.
26

---

balfcheetah randon-v2
Sopper randor v2
valier2d randen-F
46
25
Nomalzzd Epiade Roward
40-
20-
15-
12
KOPO (S7.2 2 16
36.3 a 1.00
10-
KOPO (GL. 1 1.1
28.6 1.4)
KOPO (8.2 5.6 5.4
1550
KONGO (43.5 + 41.4
5-
KDPBO (25, 10.2
17.3
1.3)
NENGO (21.5 D.1 192 0.5)
6.5
1.0
1.5
2.0
25
5.0
0.0
0.5
1.5
20
25
50
0.5
1.0
1.5
20
25
30
Grocient Steps
10
Grodient Steps
10
Grodient Steps
X10
hallchhetta-maadun-e
val.kor2d-zod:---
50-
120
100
300-
70-
Nomakaed Egicads
40-
40-
60
(72.4 4.
(62
HOBILE (15.8
NOBILS (10206
HORILE
HONTO (76.1 08 750 0.9)
NENTOO (104 0.5 06.9 2.5)
NEMTO (95.1
84.0
0.0
0.5
LO
L.
20
25
300
100
0.5
1.01
1.5
2.0
2.5
3,0
0.0
0.6
Lo
1.5
2.0
25
30
Gradient Steps
X10
Gradient Steps
XII
Gradient Steps
X10
bald whwethh--eeeeeeeecee weti---eeiu--- replay-v2
hopper-mediae meed---ddduun replay-pee
walker 2d medium www.unrrr---aae.
90
100
Noroaltaad Epikade
40-
20
ROSTLE (71
67.9 2.8)
20
NERTLE (105.1
787
MORTLE: (10.5 5
NOHBO (72.0 43 68.7
3.00
KONBO (IDLS 16 37.3
04
NCNGO (80.5 38 80.8
5.6)
0.0
10
15
20
2.5
30
0.0
0.5
1.5
2.0
25
30
0.0
0.5
1.0
1.5
2.0
25
50
Grocient Steps
X10P
Grodient Sceps
100
Gradient Steps
100
S..ffhhettah-madine expert-n2
hoppar-medt Un expaart--2
110
1.0-
100
Normakzed Epieztle Recard
90
100-
100
50
80
60-
60-
50
(83.6
10
ROBILE (100L
94.5
NOBILE (112.5 2: 02 3277
HOBILE (114.
94.3
NOARIC (1033 0.8
06.2
HOMBO (1 03
84.3
HONIO (1158
98.9 9 53)
0.0
0.5
1.0
L.S.
2.5
3.0
0.0
1.5
20
2.5
50
0.0
0.5
100
1.5
200
3.0
Gradient Steps
100
Grodient Steps
X10
Gradient Steps
x10
Figure Learning curves for the D4RL dataset.
27

---

halffhhostale wan.------------o-eeoo-o 000000000000-22
agppe--radoo--eciiee 0.01-12
Normatzel Rewaad
30-
-0-
NEPO (42.1 12.0 6.0))
HOFO (66.9
25.7
KaPO (0.1 = 00 0.00
30-
(SL.3
(15.9
MONBO (50.4 18.5
273
2.80
MONSO (220 0.4 0..7
0.5
1.5
2.0
1.0
Stops
Steps
x10
x10
ha.1f Shoetaab-randose expert-O........ 01-v2
hopper randon expeer-eepeet-OO 01-v2
walkeord 2oreddd--dddr randea--epert--oooooo
102
40-
60-
30
40-
HOFO (33.5 4
31.2 1.5)
20
MIFO (42.8 2 82.9
1222
2.4)
54
MOPO (-0.0 0.1 0.3
NOBILE (SL. a
28.2 30)
WOBILE (30.7 180
327
7.0
MOBILE (21.8 0.0. 16.8
0.0
1.0 Gradient Steps
Gradient Steps
Gradient Steps
hoppar-randox-sett t00-0.t00000
05-
70-
100
60-
Nemmalize Epioeds Reward
25-
90
64)
15-
30
10-
30-
(5.2
MOBELE (G2 4
0.
KOBELE (48.8 # 11.5
39.3 3.3)
MOBILE (LIVE 6.0
40.3
MOHBO 162.2 3.5
58.1 L.1]
NOKBO (6000 a 27.3
30.1
8:2)
NOMBO (21.7 0.1 183 2.7)
-5-
0.0
0.5
1.0
1.5
0.0
1.0 Gradient Stops
15
2.0
0.0
0.5
10
15
Gradient Steps
x10
x
Grocient Steps
x10
balf waw.ctta-rarddo-expprr.e. 066622
hoppar- wwweorr-rnndd-----o aoport-- 06-v2
Na walkerrd-raamd www.rradddddaaaoooo wwwenndo---rrrrreeeeeeee expert-
06-02
50-
90-
15-
40-
41-
10-
20
(SL.
ROBILS 10.7
MOBILE
MOBILE
133
KINBO (30.1 a 19.7 3303 9.11
0-
NOFBO [17.1 4: 12.6
30.6
3.3)
NOMBO (24.0 26 18.4 188)
0.0
2.0
L5
2.0
0.0
0.5
1.0
13
2.0
0.0
11.
1.0
15
2.0
Gradient Steps
X10
Gradient Steps
Grocient Steps
x10
Figure 5: Learning curves for the mixed dataset with trained policy demonstration ratios of 01 and 0. 05.
28

---

11-2
hopper-randor
So-keerd-raraee -ndd pediuu--
70
102
Normalied Epiodd Rowward
00)
60
50-
50-
40-
-0
40-
30-
70-
NOFO (66.2 0.7
50.0 10)
HOPC (37.3 9.0
20.
MOPO (4L. = S. 215 4.4)
10
MOBILE (65.4 0.7 389 (.99 0.4)
HOBILE (45.0 250 38.4
1888
HOBILE (1LL 6.9 10.9 3.5)
1o Gradient Steps
1.5
1.0 Gradiem Steps
Grodient Steps
tcheetahh expar
hopper-randox-exparrr
2d-randon expert--
100-
50-
MMM
201
102
13-
63-
10-
20-
NOFO (20.5 = D8 829 3.5)
20-
HDPE 12.0 2.1
NEPO (4.2 29
MOBILE (74.3 = 822 6L7 2.3
POBILE (80.3 214
00.3 6.1
0
ROBILE (1L. 7.9 3:44
0.2
1.5
Gradient Steps
Gradient Steps
Gradiem Steps
hoppar-randon-rodf un-
wwwerrrddrradd----a--- 5-522
Normalees Episcde Revord
90-
80-
20-
30-
1.4 22)
KOPC 34.5
(64.4
58.9 15.5)
10-
MOBTLE (12.0 8
MORTLE (106.5 8: 20
80.9
0-
PUBILS (36
25.0) 29.6 3.2)
NCMBO (0.3 08)
(108.0 89.9
3377 9.1)
0.0
1.5
1.0 Gradient Steps
0.0
Gradient Steps
Gradient Steps
-chaetab-randsn-orprrrt- 5-r22
hopper-raado--expers
1200
100-
WMW
40-
60-
00-
-0-
20-
(42.1 30.7
20-
(24.7
MOFO +
MORTLLE (72.9 20.1 47.3 0.7)
MORTLE (70.3 #: 30.8 m 9.3)
KOBTUK (15.7 8.5 99. 4.8)
NOKBO (38.4 44.1 96)
(88 37.9 66.4
(16.4
1/0 Gradient Stops
Gradient Steps
Gradient Steps
Figure 6: Learning curves for the mixed dataset with trained policy demonstration ratios of 0. and 0. 5.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We summarize our contribution in the abstract and in greater detail in the last paragraph of the introduction. These are theoretical (discussed in Section 3 and Section 4) and empirical (discussed in Section 5).
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations of the approach as part of our concluding discussion in Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
30

---

Answer: [Yes]
Justification: We provide all assumptions in the respective theoretical lemmata and theorems. Appendix includes all statements together with their respective proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We highlight all experimental details and hyperparameters ensure repro- ducibility in Appendix C. The experimental benchmarks come from standard libraries and for the public repository we base part of our evaluation on we reference the specific com- mit we worked with. Additionally, we make an implementation of our model available at https: /ttps://gihhan /gttttu..b..//.g// com/ adir.ab//iii /MOBBO.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
31

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We release version of our model. The data is generated from standard public libraries for this specific task which are referenced in Appendix C.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https: /tapps:://ii..../ cc/
pub 1 icc iuiddeee quides/A /et//Cooeuuaaa ssionnooooooos seiinPooicyy olicy) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips- cc/ public/ gttp//uiieeseeeeee CodeSubmis ss onPol ooiccy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide the necessary details about the experimental pipeline in Ap- pendix C, along with the source code needed to reproduce our results.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report error bars for each experiment together with details on what the respective numbers mean. We report mean + one standard deviation.
Guidelines:
The answer NA means that the paper does not include experiments.
32

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the computational resources and provide an approximate estimate the computation time for all the experiments in Appendix C, in the paragraph titled "Experiment compute resources".
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't it make it into the paper).
paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics https:/ neurips......e.. ce http/p/lc//thhcssueeeeeeee
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: We read the Code of Ethics carefully and ensured that our work complies with it. As the research is foundational do not anticipate any ethical problems.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
33

---

Justification:
The paper focuses on foundational research. As such there is no immediate short term social impact of the work be expected, neither positive nor negative.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: As research is foundational research there is no safeguard needed at the current level of research.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We do not rely on any fixed data sets. The public libraries and the repository that we rely on are referrenced and explicitly acknowledged.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
34

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, papersw teriithhhoot Choode com da aset has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
There are no new assests released with this work.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No crowdsourcing research with human subjects was performed part of this project.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No crowdsoureing or research with human subjects was performed as part of this project. As such no IRB approval was necessary.
35

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
36


Reviews:
[{"review_id": "o2FKCIo4qG", "reviewer": "Reviewer_3Ppu", "summary": "This work addresses the issue of sampling for uncertainty propagation that is the standard practice in offline RL and identifies the high variance of sampling-based estimates as an obstacle to better performance of uncertainty-aware offline RL methods. As an alternative, the authors propose propagating uncertainty with moment matching, which is a deterministic procedure. The method involves propagating the uncertainty of state transitions into the value function where the value function is parameterized with a ReLU network. The method, coined MOMBO, is further theoretically supported with tight bounds on value error approximation, and experiments on the standard Mujoco benchmarks indicate that MOMBO is on par or exceed existing approaches in offline RL.", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Uncertainty propagation in RL is a relevant problem to which there is not yet a definitive solution, mainly due to the interplay between uncertainty in state transition and uncertainty in value estimation. This work addresses this issue and proposes a method which is well motivated and theoretically supported. While the method itself does not touch upon the method of uncertainty quantification, rather only the propagation of uncertainty, I believe further investigation into the effect of various UQ methods when used with this propagation could be interesting.\n- Empirical results are fairly strong\n- Overall the paper is well written", "weaknesses": "While the authors conjecture that sampling based approaches to uncertainty propagation has inherent flaws, it would be nice to see more evidence, e.g. empirical evidence that digs into this phenomenon. \nAs such, the proposed method does not seem like a fundamentally novel approach, rather an improvement in estimation of a given framework. Still, the suggestion of moment matching as a scalable and efficient method of propagating uncertainty seems useful to the community and this paper would fit well alongside existing work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback. \n\n> While the authors conjecture that sampling based approaches to uncertainty propagation has inherent flaws, it would be nice to see more evidence, e.g. empirical evidence that digs into this phenomenon\n\nWe include an additional evaluation in our general answer above, to demonstrate both our improvements in accuracy as well as tightness.", "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": "rXEiuL65iE", "reviewer": "Reviewer_czcW", "summary": "** I am unfamiliar with the methods/related works in this paper.**\n\nThis paper proposed a model-based method for offline RL, based on moment matching. Improved numerical results are presented.", "questions": "The authors mention that the uncertainty-based approach can be overly pessimistic, but some recent paper [1,2] are showing that the uncertainty-based methods achieve the minimax optimal sample complexity, under the tabular setting. Can you explain the reason of such a statement? \n[1]Li, Gen, et al. \"Settling the sample complexity of model-based offline reinforcement learning.\" The Annals of Statistics 52.1 (2024): 233-260.\n[2]Wang, Yue, Jinjun Xiong, and Shaofeng Zou. \"Achieving the Asymptotically Optimal Sample Complexity of Offline Reinforcement Learning: A DRO-Based Approach.\"", "limitations": "n/a", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The method of moment matching is quite novel, which aims to improve the accuracy of the first two moment estimation.", "weaknesses": "I am quite confused about Sec 4. Why the Gaussian RVs are related? \n\nI feel the presentation of the paper can be improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough reading despite the unfamiliarity. \n\n> I feel the presentation of the paper can be improved.\n\nGiven the overall feedback provided via the reviews, we updated the section title and presentation for Section 3, and the _discussion and results_ paragraph in Section 5 (see, e.g., answer to BmgH). Please let us know if further parts of the paper could benefit from presentation improvements in your opinion.\n\n> Why the Gaussian RVs are related? \n\nAssuming that the overall input follows a normal distribution is a common one in the literature coming from the environmental model.\nAssuming in turn that the pre-activation output of a neural layer is also Gaussian can be justified via a central limit theorem argumentation. This is because the input to a specific neuron is a sum of random variables, i.e., the central limit theorem holds and converges to a Gaussian as the width increases. Additionally, if the computation of the first two moments is analytically tractable, the Gaussian distribution has the property of being the maximum entropy distribution, i.e., it imposes the least additional number of assumptions.\nThe fact that Algorithm 1 ignores covariance terms is common practice in moment-matching approaches (see the references in the paper), as the computational complexity and memory cost scales quadratically in the number of neurons otherwise.\n\n> The authors mention that the uncertainty-based approach can be overly pessimistic, but some recent papers [1,2] are showing that the uncertainty-based methods achieve the minimax optimal sample complexity, under the tabular setting. \n\nPessimism is required, however, current uncertainty-based approaches, e.g., MOPO and MOBILE, penalize too much. We do not contradict the two references, but show that our uncertainty estimator is tighter than our baselines'. See the additional experiments we provide in the main rebuttal answer. \nNote that the the term \"samples\" in those references refers to the size of the underlying data set, see, e.g., Section 2.2 in Li et al. (2024).\n\n_____\nLi et al., _Settling the sample complexity of model-based offline reinforcement learning_ (Annals of Statistics, 2024)", "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": "gK42dw7QPb", "reviewer": "Reviewer_BmgH", "summary": "This paper proposes a new uncertainty estimation method for model-based offline reinforcement learning, which uses moment matching to deterministically propagate uncertainties through the Q-function, rather than relying on sampling-based uncertainty estimation. The resulting model, Moment Matching Offline Model-Based Policy Optimization (MOMBO), significantly accelerates training compared to its sampling-based counterparts while maintaining performance.", "questions": "To summarize the main points/questions raised in the weaknesses section:\n\n- Could the authors evaluate their method on more challenging tasks to show that it can achieve better asymptotic performance compared to MOBILE?\n- Could the authors provide evidence to demonstrate that MOMBO can provide a more correlated/tighter estimation of the ideal uncertainty, i.e., the Bellman error?\n\nI‘m willing to raise my score if all these concerns can be addressed by the authors.", "limitations": "The authors have thoroughly discussed the limitations of their method in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written, allowing readers to follow the main arguments.\n- It provides a detailed theoretical analysis to demonstrate why moment matching is better than Monte Carlo sampling.", "weaknesses": "- The novelty of this method seems relatively weak. It appears to be a minor modification of the uncertainty quantification method used in MOBILE [1]. Additionally, as noted by the authors in the limitations section, this method is also limited to certain activation functions. I want to clarify that I do not think minor modifications lack value or novelty. Instead, the novelty of a method should be evaluated along with its empirical performance. However, the experiments in this paper do not show a strong improvement over the baseline (MOBILE).\n- The advantage claimed by the authors is sample efficiency. Firstly, using “sample efficiency” here is incorrect. This criterion is used to evaluate an online RL method. In offline RL, we train an algorithm on a fixed dataset, so “convergence speed” would be more appropriate. Furthermore, this criterion is not critical in an offline setting, where we are more concerned with the asymptotic performance.\n- According to PEVI [2], the key factor for an uncertainty-based offline RL method is how tightly it can estimate the Bellman error, which directly determines the performance of the derived policy. The authors do not provide direct evidence (theoretical or empirical) to show that MOMBO can provide a tighter estimation of the Bellman error.\n\n[1] Sun et al. \"Model-Bellman Inconsistency for Model-based Offline Reinforcement Learning\" (ICML 2023)\n\n[2] Jin et al. \"Is Pessimism Provably Efficient for Offline RL?\" (ICML 2021)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review. \n\n> The novelty of this method seems relatively weak. It appears to be a minor modification of the uncertainty quantification method used in MOBILE.\n\nMOPO, MOBILE, and MOMBO are three algorithms derived from the meta-algorithm known as Pessimistic Value Iteration (Jin et al., 2021). They all attempt to quantify uncertainty to learn a pessimistic value function. Therefore, a high degree of similarity among them is expected. However, we demonstrate that our approach surpasses the SOTA (MOBILE) both analytically and empirically.\n\n>  as noted by the authors in the limitations section, this method is also limited to certain activation functions.\n\nLet us clarify that while this limitation exists, it is less restrictive than it sounds. The derivations required to compute the first two moments of a ReLU activation directly generalize to all other piecewise linear activation functions, e.g., Leaky-Relu or PReLU. \nWhile requiring a more expensive forward pass, the first two moments are also tractable for more modern activation functions such as the ELU and its SELU counterpart. \nTo summarize, while the first two moments cannot be computed for all activation functions, they can be for the most relevant and most popular ones covering all mainstream architectures used in critic training in modern RL.\n\n\n\n> The advantage claimed by the authors is sample efficiency. Firstly, using “sample efficiency” here is incorrect. This criterion is used to evaluate an online RL method. In offline RL, we train an algorithm on a fixed dataset, so “convergence speed” would be more appropriate. Furthermore, this criterion is not critical in an offline setting, where we are more concerned with the asymptotic performance.\n\nWe agree, that _sample efficiency_ was poorly worded. We adapt the discussion paragraph in the final version of the paper accordingly to make this clear. (Table 2 refers to the new results provided in this rebuttal.) \n\n```\nDiscussion and results. To summarize, our experimental findings are:\n(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).\n(ii) MOMBO has a faster convergence speed. It learns faster and reaches its final performance earlier. This is reflected in its AUC score, where MOMBO outperforms the baselines.\n(iii) MOMBO provides improved uncertainty estimates for the Bellman target. Its moment-matching approach provides it with tighter and more accurate estimates, as shown in Table 2.\n(iv) MOMBO is more robust. It has a lower standard deviation than 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.\n```\n\n> However, the experiments in this paper do not show a strong improvement over the baseline (MOBILE).\n\nUpon three million gradient steps, both MOBILE and MOMBO perform very similar, with a tendency for minor improvements in MOMBO. However, the AUC results, being calculated over ten evaluation results throughout the training process, show a clear distinction in the speed of convergence with MOMBO outperforming MOBILE in all but one environment. Convergence speed is a crucial metric in real applications where it is not feasible to engineer the training duration based on evaluation results, which is the common practice for standard offline RL benchmarks.\nAdditionally, this means that for a fixed amount of computational time/training data, we can train more powerful models than MOBILE, which can be expected to then translate into a larger normalized reward as well.\nFinally, the experiments performed for the rebuttal (see general answer) show improvements in accuracy and tightness indicating a better training signal for our proposed approach.\n\n\n> Could the authors evaluate their method on more challenging tasks to show that it can achieve better asymptotic performance compared to MOBILE?\n\n\nWe provide additional performance metrics in our general answer demonstrating the improvements our method provides as well as preliminary results on the adroit domain.\n\n\n> Could the authors provide evidence to demonstrate that MOMBO can provide a more correlated/tighter estimation of the ideal uncertainty, i.e., the Bellman error?\n\nWe have included an additional experiment in our main response that demonstrates improved performance in terms of accuracy and tightness upon both MOBILE and MOPO.\n\nPlease let us know if this answer clarifies your concerns.\n\n\n\n\n-----\nJin et al., _Is Pessimism Provably Efficient for Offline RL?_ (ICML 2021)\n\nYu et al., _COMBO: Conservative offline model-based policy optimization_ (NeurIPS 2021)\n\nJanner et al., _Offline reinforcement learning as one big sequence modeling problem_ (NeurIPS 2021)\n\nRigter et al., _RAMBO-RL: Robust adversarial model-based offline reinforcement learning_ (NeurIPS 2022)\n\nSun et al., _Model-bellman inconsistency for model-based offline reinforcement learning_ (ICML 2023)", "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": "ZGN4y4nmk4", "reviewer": "Reviewer_igqb", "summary": "To address the high variance issue of model-based offline RL methods, which rely on sampling-based uncertainty estimation, this manuscript proposes MOMBO, a model-based policy optimization method based on moment matching. MOMBO learns a Q-function using moment matching and deterministically propagates uncertainties through the Q-function. Theoretically, using moment matching has tighter upper bounds on value approximation than using Monte Carlo sampling. The authors evaluate MOMBO’s performance across various environments.", "questions": "See weaknesses.", "limitations": "None.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. MOMBO is a more stable and sample-efficient approach.\n\n2. There is a tight upper bound on value approximation errors using moment matching compared to Monte Carlo sampling.", "weaknesses": "1.There is a clear necessity for the authors to enhance their writing proficiency. \n\n2.The manuscript lists many contributions, which obscures its central thesis. The primary contribution is learning a Q-function using moment matching. However, this approach needs to strike me as sufficiently innovative.\n\n3.How can the alignment of moments from one distribution with those from another ensure the approximation or equivalence of the two?\n\n4.OOD data seems more penalized by high variance, whereas IND data does not generally have high variance. Moreover, there is a need for diversification to improve the variance of Ensemble RL (cf. EDAC, NeurIPS 2021). The moment matching proposed by the authors is similar to transforming uncertainty/variance. Furthermore, the experimental results do not prove the method's validity due to the lack of comparison with the latest SOTA method.\n\n5.Section 3, entitled \"Limitations of SOTA model-based offline RL approaches,\" does not quite capture the issue's essence. Ideally, it should be reframed to reflect a specific shortcoming within a methodology, technique, or theoretical framework. Moreover, the term \"SOTA\" is inherently time-bound, making it an imprudent choice for a section title, as it may quickly become outdated.\n\n6.The authors follow \\mu - \\beta \\cdot \\sigma, where \\beta is the weight parameter seems to differ little from most offline RL methods based on uncertainty estimation.\n\n7.The core of the proposed method is based on moment matching; however, the citations to moment matching in the manuscript are too old. Is there any recent work?\n\n8.Only some tests on mujoco environments are insufficient. Isn't it common to perform experimental validation on environments like Antmaze, Adroit, etc.? How do the authors explain this? I suspect that the proposed method will not work in these neglected environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this thorough review. We will address your concerns below. \n\n> 1.There is a clear necessity for the authors to enhance their writing proficiency. \n\nCould you be specific as to where the writing proficiency is suboptimal?\n\n> 2.The manuscript lists many contributions, which obscures its central thesis.\n\nOur contribution is three-fold. (i) We propose a novel moment matching and propagation approach for uncertainty-based model-based offline RL. \n(ii) We highlight theoretical differences between it and a sampling-based approach. \n(iii) We provide empirical evidence for the usefulness of our proposal.\nWe will restructure the relevant paragraph to make this more clear.\n\n> 3.How can the alignment of moments from one distribution with those from another ensure the approximation or equivalence of the two?\n\nThe Wasserstein 1-distance (Def 1) which we consider in this paper measures closeness between two distributions in terms of their cdf not their moments. However, e.g., Wasserstein 2-distance between two univariate Gaussian distributions can be shown to be the sum of the squared difference between their mean and variance terms, i.e., alignment of the moments minimizes their distance based on the $W_2$ metric. Does that answer your question?\n\n> The moment matching proposed by the authors is similar to transforming uncertainty/variance.\n\nIndeed, as we propose a sampling-free approach, this necessitates a way to transform uncertainty as it passes throughout the network. Can the reviewer clarify why this analytical transformation is a weakness and not a feature?\n\n> Furthermore, the experimental results do not prove the method's validity due to the lack of comparison with the latest SOTA method.\n\nEDAC (An et al., 2021) and its improvements, such as SPQR (Lee et al., 2023), focus on model-free learning, whereas we consider model-based algorithms. To our knowledge, MOBILE is the state-of-the-art method for model-based offline learning. Most recent follow-up work such as Luo et al. (2024) and Song et al. (2024) take MOBILE as the base algorithm and complements it with additional features which can be applied also to our method.\nIn our general answer, we have included further metrics that highlight the difference between MOBILE's performance and ours.\n\n\n\n> Section 3, entitled \"Limitations of SOTA model-based offline RL approaches,\" does not quite capture the issue's essence. Ideally, it should be reframed to reflect a specific shortcoming within a methodology, technique, or theoretical framework.\n\nThank you for the proposal. We will change the section title to _Limitations of sampling-based uncertainty propagation_ to highlight the specific limitation we are addressing. \n\n> The authors follow \\mu - \\beta \\cdot \\sigma, where \\beta is the weight parameter seems to differ little from most offline RL methods based on uncertainty estimation.\n\nWe agree, using a lower confidence bound is a standard practice throughout offline RL. That is not a contribution of ours. What is relevant is how this bound is constructed, i.e., in our case $\\hat \\sigma$. \nSee also Theorem 2 in Jin et al. (2021) regarding which derives the desired suboptimality bound on a lower bound estimator. \n\n> The core of the proposed method is based on moment matching; however, the citations to moment matching in the manuscript are too old. Is there any recent work?\n\nThere are still minor variations explored in the literature, but these tend to be either only theoretical or too costly, such that there have been no further fundamental new developments. E.g., moment-matching with complete covariance matrices is explored in the current work by Wright et al. (2024). But the cost scales quadratically in the layer width, i.e., becomes intractable for common network sizes. As the authors in that work acknowledge (see their Section 4.3), using the moment matching developed in our references performs just as well or even better. Similarly, Look et al. (2023) explored a way to propagate these covariances in neural stochastic differential equations using Stein's lemma. We will clarify the current state of the literature at the beginning of Section 4.\n\n\n> 8.Only some tests on mujoco environments are insufficient. Isn't it common to perform experimental validation on environments like Antmaze, Adroit, etc.? How do the authors explain this? I suspect that the proposed method will not work in these neglected environments.\n\nWe focused on mujoco as a popular and extensive benchmark. See our general answer above for more metrics and a preliminary set of adroit results.\n\n\n\n_____\nAn et al., _Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble_ (NeurIPS 2021)\n\nJin et al., _Is Pessimism Provably Efficient for Offline RL?_ (ICML 2021)\n\nLee et al., _SPQR: Controlling Q-ensemble Independence with Spiked Random Model for Reinforcement Learning_ (NeurIPS 2023)\n\nLook et al., _A deterministic approximation to neural SDEs_ (PAMI 2023)\n\nLuo et al., _Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning_ (ICLR 2024)\n\nSong et al., _Compositional Conservatism: A Transductive Approach in Offline Reinforcement Learning_ (ICLR 2024)\n\nWright et al., An Analytic Solution to Covariance Propagation in Neural Networks (AISTATS 2024)", "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:
We thank all reviewers for their feedback. 

We address all the reviewer comments in our individual responses. To summarize the main changes during the rebuttal phase:
- We conducted an experiment to compare the quality of uncertainty quantifiers among the baselines and our MOMBO, which supports our thesis. _(see below)_
- We provide additional results on a new benchmark (Adroit). _(see below)_
- We reorganized Section 3, _discussion and results_ paragraph of Section 5, and contributions paragraph. _(see, e.g., answer to BmgH)_
## Experiment on Uncertainty Quantifiers
Per the reviewers' request, we conducted an experiment to provide empirical evidence that MOMBO offers a tighter estimation of the ideal uncertainty. To achieve this, we followed the definition of the $\xi$-uncertainty quantifier (see Definition 4.1 in Jin et al. (2021)):
$$| \hat{\mathcal{T}}\hat{Q}(s, a) - \mathcal{T}\hat{Q}(s, a) | \leq \beta \cdot u(s, a)$$
for all $(s, a) \in \mathcal{S} \times \mathcal{A}$.
We aimed to measure the quality of uncertainty penalizers and determine their tightness in model-based offline reinforcement learning algorithms, specifically comparing MOPO, MOBILE, and our algorithm, MOMBO. To achieve this, we report two scores:
1. **Accuracy**: For a given state-action pair, we check whether the equation above holds or not. This measures the quality of the uncertainty quantifier.
2. **Tightness score**: For a given state-action pair we calculate
$$\text{Tightness score}= \beta \cdot u(s, a) - | \hat{\mathcal{T}}\hat{Q}(s, a) - \mathcal{T}\hat{Q}(s, a) |.$$
This score measures how tightly the uncertainty quantifiers can estimate the possible errors.

Here, $\beta \cdot u(s, a)$ represents the penalty applied by the algorithm for a specific state-action pair. $\hat{\mathcal{T}}\hat{Q}(s, a)$ is the learned estimated Bellman target, calculated using learned dynamics, while $\mathcal{T}\hat{Q}(s, a) = r(s, a) + \gamma E_{s' \sim T} [E_{a' \sim \pi}[Q(s', a')]]$ is the true Bellman target. The true Bellman target requires calculating the real expectation for the next state and next action. Since the environments we used are deterministic, we used the actual next state from the environment dynamics and the actual reward. For the expectation of the next action, we used a large number of samples to approximate it accurately.
### Experiment Procedure:
We load the learned policies and dynamics and generate 10 episode rollouts using the real environment in evaluation mode. From these rollouts, we select state, action, reward, and next state tuples at every 10th step, including the final step. For each of these tuples, we calculate the mean accuracy and tightness scores. This process is repeated for each of the 4 seeds across every task, and we report the mean and standard deviation of the accuracies and tightness scores.
### Experiment Results
#### Accuracy
| | MOPO | MOBILE | MOMBO |
|---:|:---:|:---:|:---:|
| halfcheetah-random | 0.02 ± 0.01 | **0.25 ± 0.02** | 0.24 ± 0.07 |
| hopper-random | 0.04 ± 0.01 | **0.85 ± 0.03** | 0.62 ± 0.07 |
| walker2d-random | 0.0 ± 0.0 | 0.55 ± 0.04 | **0.74 ± 0.06** |
| halfcheetah-medium | 0.05 ± 0.01 | 0.25 ± 0.0 | **0.34 ± 0.04** |
| hopper-medium | 0.04 ± 0.01 | 0.41 ± 0.01 | **0.55 ± 0.03** |
| walker2d-medium | 0.02 ± 0.0 | 0.16 ± 0.01 | **0.38 ± 0.02** |
| halfcheetah-medium-replay | 0.08 ± 0.01 | 0.04 ± 0.0 | **0.16 ± 0.0** |
| hopper-medium-replay | 0.02 ± 0.01 | 0.03 ± 0.01 | **0.17 ± 0.01** |
| walker2d-medium-replay | 0.08 ± 0.01 | 0.14 ± 0.01 | **0.36 ± 0.02** |
| halfcheetah-medium-expert | 0.11 ± 0.01 | 0.36 ± 0.03 | **0.44 ± 0.03** |
| hopper-medium-expert | 0.04 ± 0.02 | 0.43 ± 0.02 | **0.51 ± 0.04** |
| walker2d-medium-expert | 0.05 ± 0.01 | **0.47 ± 0.02** | 0.45 ± 0.02 |
#### Tightness Score
| | MOPO | MOBILE | MOMBO |
|---:|:---:|:---:|:---:|
| halfcheetah-random | -14.92 ± 1.14 | -6.88 ± 0.61 | **-6.21 ± 0.28** |
| hopper-random | -1.81 ± 0.16 | -0.64 ± 0.2 | **-0.31 ± 0.65** |
| walker2d-random | -98521227.61 ± 91264317.73 | -0.27 ± 0.08 | **-0.14 ± 0.02** |
| halfcheetah-medium | -15.76 ± 0.66 | -10.03 ± 0.53 | **-9.06 ± 0.26** |
| hopper-medium | -4.08 ± 1.16 | -3.14 ± 0.08 | **-1.3 ± 0.21** |
| walker2d-medium | -8.98 ± 0.63 | -5.02 ± 0.52 | **-4.03 ± 0.24** |
| halfcheetah-medium-replay | -12.16 ± 1.18 | -11.85 ± 0.41 | **-10.4 ± 0.66** |
| hopper-medium-replay | -5.45 ± 0.59 | -3.4 ± 0.04 | **-3.17 ± 0.04** |
| walker2d-medium-replay | -4.45 ± 0.43 | -4.56 ± 0.13 | **-3.78 ± 0.39** |
| halfcheetah-medium-expert | -21.45 ± 2.79 | -13.22 ± 0.47 | **-11.88 ± 0.5** |
| hopper-medium-expert | -7.12 ± 2.6 | -3.5 ± 0.03 | **-3.38 ± 0.02** |
| walker2d-medium-expert | -9.72 ± 0.24 | -5.52 ± 0.36 | **-5.29 ± 0.14** |
### Discussion on Results
To summarize the experiment results: MOMBO provides improved uncertainty estimates for the Bellman target on almost all data sets (9/12 in terms of accuracy, 12/12 in terms of tightness). This improved estimation leads to faster convergence due to clearer gradient signals and results in higher final performance.
## Adroit
Per the reviewer igqb's suggestion, we conducted an experiment on the Adroit domain of the D4RL dataset, specifically reporting on 'pen-human-v1' and 'pen-cloned-v1'. We skipped the other tasks because MOBILE does not provide their hyperparameters, and conducting an extensive hyperparameter search requires too much time, which we do not have for the rebuttal period. We report the MOPO and MOBILE results from the MOBILE paper.
Even without detailed hyperparameter tuning, we are competitive with MOBILE. (MOBILE provides hyperparameters in the appendix, however, these do not agree with those reported in their repository. We rely on the latter.)
| | MOPO | MOBILE | MOMBO |
|---:|:---:|:---:|:---:|
| pen-human | 10.7 | 30.1 ± 14.6 | **32.0 ± 10.7** |
| pen-cloned | 54.6 | **69.0 ± 9.3** | 63.5 ± 19.5 |
_____
Jin et al., _Is Pessimism Provably Efficient for Offline RL?_ (ICML,  2021)

Meta-review: There is some variation in the overall assessments of the reviews, with most reviewers at least leaning toward acceptance and one reviewer firmly on the reject side. Because of this, in my meta-review I will do my best to synthesize key points from the discussion and corroborate with my own read of the paper.

There is general agreement amongst the reviewers that the application of moment-matching to improve uncertainty estimates is interesting and potentially impactful. They also generally appreciated the theoretical findings that support the hypothesis that moment-matching may provide benefits over Monte Carlo estimation.

The quality of the presentation was somewhat in dispute, with some reviewers expressing that the paper was clear and easy to follow and others expressing concerns about clarity. In my own reading, I did not find clarity of presentation to be a significant barrier to understanding.

Some of the reviewers expressed uncertainty about whether the empirical results sufficiently supported the hypotheses that Monte Carlo estimates were overly conservative and that tighter uncertainty estimates would result in practical benefits. Many of these questions were resolved in the discussion period. 

For instance, the authors clarified that, while MOMBO and MOBILE perform similarly after many training steps (and that this was to be expected), the AUC comparison indicates that MOMBO generally outperformed MOBILE for smaller numbers of training steps. I tend to agree with the reviewers that this point was not especially clear in the paper and that this issue was exacerbated by the incorrect use of the term "sample efficiency," rather than the real claim of faster improvement/convergence. The improvements proposed by the authors to address this confusion will improve the paper.

There is still a lingering question of how much faster improvement actually matters in the offline setting, where the central premise is that computational resources are far cheaper than data collection. The authors make the point that it is not necessarily easy to tell when to stop training and that therefore faster improvement implies better performance if training is stopped too early. However, this is a somewhat conjectural claim. There is a long history of research on stopping conditions for neural network training, albeit largely in the supervised learning setting. It's not clear that improving convergence rate is the best or only way to address this issue; that seems like a claim that requires more study. All that said, the empirical findings do at least corroborate the stated intuition and theoretical claims and faster convergence is pretty clearly desirable, even if it's not clear exactly how critical it is.

There were also questions about whether the performance results could be clearly connected to the central hypothesis that MC estimates were overly conservative. The additional experimental results offered by the authors seem to compellingly demonstrate that, indeed, the moment-matching estimates are tighter than the MC estimates. These results seem fairly straightforward to incorporate as they could reside in an appendix with a brief summary of the findings in the main text.

Overall, it seems to me that the paper presents a novel, conceptually sound, well-analyzed, and potentially impactful approach to uncertainty estimation in MBRL and is thus acceptable under the NeurIPS guidelines. It does seem that there is room for improvement in the area of significance, with more work to be done to determine the extent of the practical importance of the theoretical and algorithmic insight. For that reason, I think it's fair for this paper to be considered somewhat borderline. If the paper is ultimately not accepted, I recommend that the authors attend to that particular issue as they continue this line of work.
Decision: Accept (poster)