Deterministic Uncertainty Propagation for Improved Model-Based Offline Reinforcement Learning

Authors: abdullah  akgül
manuel  haußmann
melih  kandemir
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:02:52

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.

Main paper: 
 introduction: 
 Offline reinforcement learning (Lange et al., 2012; Levine et al., 2020) aims to solve a control task using an offline dataset without interacting with the target environment. Such an approach is essential in cases where environment interactions are expensive or risky (Shortreed et al., 2011; Singh et al., 2020; Nie et al., 2021; Micheli et al., 2022) . Directly adapting traditional online off-policy reinforcement learning methods to offline settings often results in poor performance due to the adverse effects of the distributional shift caused by the policy updates (Fujimoto et al., 2019; Kumar et al., 2019) . The chief reason for the incompatibility is the rapid accumulation of overestimated actionvalues during policy improvement. Pessimistic Value Iteration (PEVI) (Jin et al., 2021) offers a 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 state-of-the-art offline reinforcement learning algorithms as special cases (Yu et al., 2020; Sun et al., 2023) .Model-based offline reinforcement learning approaches first fit a probabilistic model on the real state transitions and then supplement the real data with synthetic samples generated from this model throughout the course of 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 on the discrepancy between its actions and those 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. Moment matching offers sharp estimates of the action-value of the next state at the cost of only two forward passes through a critic network. A similar sharpness cannot be reached even with 10000 Monte Carlo samples, which is 5000 times more costly. See Appendix C.1.1 for details.In the offline reinforcement learning literature, uncertaintydriven approaches exist for both model-free (An et al., 2021; Bai et al., 2022) and model-based settings (Yu et al., 2020; Sun et al., 2023) , all aimed at learning a pessimistic value function (Jin et al., 2021) by penalizing it with an uncertainty estimator. The impact of uncertainty quantification has been investigated in both online (Abbas et al., 2020) and offline (Lu et al., 2021) scenarios, particularly in model-based approaches, which is the focus of our work.Despite the strong theoretical evidence suggesting that the quality of the Bellman target uncertainties has significant influence on model quality (O'Donoghue et al., 2018; Luis et al., 2023; Jin et al., 2021) , the existing implementations rely on Monte Carlo samples and heuristics-driven uncertainty quantifiers being used as reward penalizers (Yu et al., 2020; Sun et al., 2023) which results in a high degree of randomness that manifests itself as significant delays in model convergence. Figure 1 demonstrates why this is the case by a toy example. The uncertainty around the estimated action-value of the next state does not shrink even after taking 10000 samples on the next state and passing them through a Q-network. Our moment matching-based approach predicts a similar mean with significantly smaller variance at the cost of only two Monte Carlo samples.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 that has been developed originally for deterministic variational inference of Bayesian neural networks (Wu et al., 2019a) , for the first time to deterministically propagate environment model uncertainties through Q-function estimates. We refer to this algorithm as Moment Matching Offline Model-Based 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 accelerates training convergence while maintaining asymptotic performance. 
preliminaries: 
 Reinforcement learning.We define a Markov Decision Process (MDP) as a tuple M ≜ (S, A, r, P, P 1 , γ, H) where S represents the state space and A denotes the action space. We have r : S × A → [0, R max ] as a bounded deterministic reward function and P : S × A × ∆(S) → [0, 1] as the probabilistic state transition kernel, where ∆(S) denotes the probability simplex defined over the state space. The MDP has an initial state distribution P 1 ∈ ∆(S), a discount factor γ ∈ (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 a probabilistic model trained on the data collected from a deterministic environment. We denote the action-value of a policy π for a state-action pair (s, a) as:Q π (s, a) ≜ E π H i=h γ i-h r(s i , a i )|s h = s, a h = a .This function maps to the range [0, Rmax /1-γ].Offline reinforcement learning algorithms perform policy search using an offline collected dataset D = {(s, a, r, s ′ )} from the target environment via a behavior policy π β .The goal is to find a policy π that minimizes the suboptimality, defined as SubOpt(π; s) ≜ Q π * (s, π * (s)) -Q π (s, π(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 of heteroscedastic neural networks (Lakshminarayanan et al., 2017) . Mainstream model-based 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. It 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 (Yu et al., , 2021) ) settings.Scope and problem statement. We focus on 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 et al., 2023) . A primary challenge in offline reinforcement learning is the distributional shift caused by the 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 to overcome this problem in the online setting such as min-clipping (Fujimoto et al., 2018) are insufficient for offline reinforcement learning. Algorithms that use reward penalization proportional to the estimated uncertainty of an unobserved next state 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; Uehara and Sun, 2022) . We demonstrate the theoretical and practical limitations of PEVI-based approaches and address them with a new solution.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 error stems from the fact that during training the agent has access only to the sample Bellman operator:B π Q(s, a, s ′ ) ≜ r(s, a) + γQ(s ′ , π(s ′ ))where s ′ is a Monte Carlo sample of the next state. However, the actual quantity of interest is the exact Bellman operator that is equivalent to the expectation of the sample Bellman operator:B π Q(s, a) ≜ E s ′ ∼P(•|s,a) B π Q(s, a, s ′ ) . 
pessimistic value iteration algorithms: 
 A pessimistic value iteration algorithm (Jin et al., 2021) , denoted as A PEVI ( B Γ π Q, P), performs Dynastyle model-based learning using the following pessimistic sample Bellman target for temporal difference updates during critic training:B Γ π Q(s, a, s ′ ) ≜ B π Q(s, a, s ′ ) -Γ Q P (s, a)where Γ Q P (s, a) admits a learned state transition kernel P and a value function Q to map a state-action pair to a penalty score. Notably, this penalty term does not depend on an observed or previously sampled s ′ as it quantifies the uncertainty around s ′ using P. The rest follows as running an offthe-shelf model-free online reinforcement learning algorithm, for instance Soft Actor-Critic (SAC) (Haarnoja et al., 2018) , on a replay buffer that comprises a blend of real observations and synthetic data generated from P using the recent policy of an intermediate training stage. We study the following two PEVI variants in this paper due to their representative value: i) Model-based Offline Policy Optimization (MOPO) (Yu et al., 2020) (Sun et al., 2023) propagates the uncertainty of the next state to the Bellman target through Monte Carlo sampling and uses it as a penalty:Γ Q P (s, a) ≜ var N s ′ ∼ P [ B π Q(s, a, s ′ )]where var N s ′ ∼ P 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 π Q(s, a, s ′ ) with a 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 1 (Suboptimality of PEVI (Jin et al., 2021) ). For any π derived withA PEVI ( B Γ π Q, P) that satisfies |B π Q(s, a) -B π Q(s, a, s ′ )| ≤ Γ Q P (s, a), ∀(s, a) ∈ S × Awith probability at least 1-δ for some error tolerance δ ∈ (0, 1) where s ′ ∼ P(•|s, a), the following inequality holds for any initial state s 1 ∈ S:SubOpt(π) ≤ 2 H i=1 E π * Γ Q P (s i , a i ) s 1 .When deep neural networks are used as value function approximators, calculating their variance becomes analytically intractable even for normally distributed inputs. Consequently, reward penalties Γ Q P are typically derived from Monte Carlo samples, which are prone to high estimator variance (see Figure 1 ). Our key finding is that using high-variance 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 to the asymptotic performance. (ii) The first two moments of the Bellman target are 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.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 can be characterized as follows. We are interested in the distribution of a deterministic map y = f (x) for a normally distributed input x ∼ N (µ, σ 2 ).We analyze this complex object via a surrogate of it. For a sample setS N = {x i } N i=1 obtained i.i.d.from N (µ, σ 2 ), let µ N and σ 2 N be its empirical mean and variance. Now consider the following two random variables yN = 1 N N i=1 f (x i ) and y ′ N = f (x ′ ) for x ′ ∼ N (µ N , σ 2 N ).Note that y 1 and y ′ 1 are equal in distribution. Furthermore, both y N and y ′ N converge to the true y as N tends to infinity. We perform our analysis on y ′ N where analytical guarantees are easier to build. Furthermore, y ′ N is a tighter approximation of both y 1 and y ′ 1 . We construct the following suboptimality proof for the PEVI algorithmic family that approximates the uncertainty around the Bellman operator in the way y ′ N is built. In this theorem and elsewhere, A l stands for the weights of the l-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-based PEVI). For any policy π MC learned by A PEVI ( B Γ π Q, P) using N Monte Carlo samples to approximate the Bellman target with respect to an action-value network defined as an L-layer MLP with 1-Lipschitz activation functions, the following inequality holds for any error tolerance δ ∈ (0, 1) with probability at least 1 -δSubOpt(π MC ) ≤ 2H L l=1 A l - 8 log( δ /4)R 2 max N /2 (1 -γ) 2 .The bound in Theorem 2 is prohibitively loose since R 2 max /(1 -γ) 2 is a large number in practice. For instance, the maximum realizable step reward in the HalfCheetah-v4 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 γ = 0.99, we obtain R 2 max /(1 -γ) 2 = 1.37 × 10 6 . 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. 
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 to the model-based offline reinforcement learning setting. We obtain a 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 a learned environment model P through an action-value network by alternating between an affine transformation of a normally distributed input to another normally distributed linear activation and projecting the output of a 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 is 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 sampleefficient 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, e.g. Wu et 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 be X ∼ N (X|µ, σ 2 ), the pre-activation Y for a neuron associated with weights θ is given by Y= θ ⊤ X, i.e., Y ∼ N (Y |θ ⊤ µ, (θ 2 ) ⊤ σ 2 ),where we absorb the bias into θ for simplicity and the square on θ is applied element-wise. For a ReLU activation function r(x) ≜ max(0, x)1 , mean µ and variance σ 2 of Y = max(0, X) are analytically tractable (Frey and Hinton, 1999) . We approximate the output with a normal distribution X ∼ N ( X| µ, σ 2 ) and summarize several properties regarding µ and σ 2 below. See Appendix B.1 for proofs of all results in this subsection. Lemma 1 (Moment matching). For X ∼ N (X|µ, σ 2 ) and Y = max(0, X), we have(i) µ ≜ E [Y ] = µΦ(α) + σϕ(α), σ 2 ≜ var [Y ] = (µ 2 + σ 2 )Φ(α) + µσϕ(α) -µ 2where α = µ/σ, and ϕ(•), Φ(•) are the probability density function (pdf) and cumulative distribution function (cdf) of the standard normal distribution, respectively. Additionally, we have that (ii) µ ≥ µ and (iii) σ 2 ≤ σ 2 .Algorithm 1 outlines the moment matching process. This process involves two passes through the linear layer: one for the first moment and one for the second, along with a few 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 > 2, sampling-based methods introduce computational overhead. For example, MOBILE (Sun et al., 2023) uses N = 10. See Figure 1 and Figure 3 for illustrations on the effect of varying N .Algorithm 1 Deterministic uncertainty propagation through moment matching function MOMENTMATCHINGTHROUGHLINEAR(θ, b, X) θ: Weights of the layer, b: bias of the layerX = N (X|µ X , σ 2 X ) ▷ input a normal distribution (µ Y , σ 2 Y ) ← (θ ⊤ µ X + b, θ 2⊤ σ 2 X ) ▷ transform mean and variance return Y = N (Y |µ Y , σ 2 Y ) ▷ output the transformed distribution end function function MOMENTMATCHINGTHROUGHRELU(X) X = N (X|µ X , σ 2 X ) ▷ input a normal distribution α ← µ X /σ X µ Y ← µ X Φ(α) + σ X ϕ(α) ▷ compute the first two moments of ReLU(X) σ 2 Y ← µ 2 X + σ 2 X Φ(α) + µ X σ X ϕ(α) -µ 2 Y ▷ ϕ/Φ are the normal pdf/cdf return Y = N (Y |µ Y , σ 2 Y )▷ output a normal distribution with these moments end functionWe propagate the distribution of the next state predicted by the environment model through the action-value function network using moment matching. We define a moment matching Bellman target distribution as:B π Q(s, a, s ′ ) d = r(s, a) + γQ M M (s ′ , π(s ′ ))for s ′ ∼ P(•|s, a) and Q M M (s ′ , a ′ ) a normal distribution with mean µ M M (s ′ , a ′ ) and variance σ 2 M M (s ′ , a ′ ) obtained as the outcome of a progressive application of Algorithm 1 through the layers of a 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 a lower confidence bound on B π Q(s, a, s ′ ) as the Bellman targetB Γ π Q(s, a, s ′ ) ≜ r(s, a) + γµ M M (s ′ , π(s ′ )) -βγσ M M (s ′ , π(s ′ ))(1) for some radius β > 0.The following lemma provides a bound on the 1-Wasserstein distance W 1 between the true distribution ρ Y of a random variable after being transformed by a ReLU activation function and its moment matched approximation ρ X . See Appendix B.3 for the proofs of all results presented in this subsection. Lemma 2 (Moment matching bound). For the following three random variablesX ∼ ρ X , Y = max(0, X), X ∼ N ( X| µ, σ 2 ) with µ = E [Y ] and σ 2 = var [Y ] the following inequality holds W 1 (ρ Y , ρ X ) ≤ 0 -∞ F X (u)du + W 1 (ρ X , ρ X )whereF X (•) is cdf of X . If ρ X = N (X|µ, σ 2 ), it can be further simplified to W 1 (ρ Y , ρ X ) ≤ σϕ µ σ -µΦ - µ σ + |µ -µ| + |σ -σ|.Applying this to a 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(x) = max(0, x). For l = 1, . . . , L -1, the sampling-based forward-pass isY 0 = X s , Y l = r(f l (Y l-1 )), Y L = f L (Y L-1 )where f l (•) is the l-th layer and X s a sample of N (X|µ X , σ 2 X ). Its moment matching pendant isX 0 ∼ N ( X 0 |µ X , σ 2 X ), X l ∼ N X l E r(f l ( X l-1 )) , var r(f l ( X l-1 )) .The following inequality holds forρ Y = ρ X L = N ( X L |E[f ( X L-1 )], var[f ( X L-1 )]). W 1 (ρ Y , ρ Y ) ≤ L l=2 G( X l-1 ) + C l-1 L j=l A j , with G( X l ) = σ l ϕ µ l σ l -µ l Φ - µ l σ l ≤ 1, C l ≤ |A l µ l-1 -µ l | + A 2 l σ 2 l-1 -σ lwhere √ • and (•) 2 are applied elementwise.Relying on Lemma 6 again, this result provides an upper bound on the suboptimality of our moment matching approach.Theorem 3 (Suboptimality of moment matching-based PEVI algorithms). For any policy π MM derived with A PEVI ( B Γ π Q, P) learned by a penalization algorithm that uses moment matching to approximate the Bellman target with respect to an action-value network defined as an L-layer MLP with 1-Lipschitz activation functions, the following inequality holdsSubOpt(π MM ) ≤ 2H L l=2 G( X l-1 ) + C l-1 L j=l A j .The bound in Theorem 3 is much tighter than Theorem 2 in practice as R 2 max /(1 -γ) 2 is very large while the Lipschitz continuities A j 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 is a fundamental requirement.We adopt the model learning scheme from Model-Based Policy Optimization (MBPO) (Janner et al., 2019) and use SAC (Haarnoja et al., 2018) as the policy search algorithm. These approaches represent the best practices in many recent model-based offline reinforcement learning methods (Yu et al., 2020 (Yu et al., , 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 n (s ′ , r|s, a) for n ∈ {1, . . . , N ens }. After evaluating the performance of each model on a validation set, we select the N elite best-performing 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 we 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 zero 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 a deterministic policy for illustrative purposes. In our implementation, the reward model also follows a 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 is a squashed Gaussian distribution trained by SAC in the maximum entropy reinforcement learning setting. For further details, refer to the implementation.  
experiments: 
 We compare MOMBO against MOPO and MOBILE, the two representative PEVI variants, across twelve tasks from the D4RL dataset (Fu et al., 2020) , which consists of data collected from three MuJoCo environments (halfcheetah, hopper, walker2d) with behavior policies exhibiting four degrees of expertise (random, medium, medium-replay, and medium-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 in 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 intermediate 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 halfcheetah 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 et 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://github.com/adinlab/MOMBO. 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 prescribed configurations.Theorem 1 indicates that the performance of a PEVI algorithm depends on how tightly its reward penalizer Γ Q P (s, a) upper bounds the Bellman operator error|B π Q(s, a) -B π Q(s, a)|.We use this theoretical result to compare the quality of the reward penalizers of the three models based on average values of the following two performance scores calculated across data points observed during 10 evaluation episodes: (i) Accuracy: 1(Γ Q P (s, a) ≥ |B π Q(s, a) -B π Q(s, a)|)for the indicator function 1, i.e., how often the reward penalizer is a valid ξ-uncertainty quantifier as assumed by Theorem 1.(ii) Tightness:Γ Q P (s, a) -|B π Q(s, a) -B π Q(s, a)|, i.e., how sharp a ξ-uncertainty quantifier the reward penalizer is.Table 2 shows that MOMBO provides more precise uncertainty estimates compared to the samplingbased approaches. It also indicates that MOMBO provides tighter estimates of the Bellman operator error, meaning that the sampling-based approaches use larger confidence radii. See Appendix C.1.2 for details.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 is that policy search needs to 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 et 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 a 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 (Rigter et al., 2022) with 67.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 MO-BILE and 27.9 for MOPO, averaged across all tasks at 1%. In terms of AULC, MOMBO attains 29.5, outperforming the 28.0 of MOBILE and the 23.5 of MOPO. See Table 3 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.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We release a 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 The main objective of MOMBO is to accurately quantify the uncertainty surrounding a Bellman target estimate caused by the error in predicting the next state. We achieve this by deterministically propagating uncertainties through the value function with moment matching and introducing pessimism by constructing a 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-based reinforcement learning algorithms in continuous state-action spaces. Our algorithmic contributions can also be instrumental in both model-based online and model-free offline reinforcement learning setups (An et al., 2021; Bai et al., 2022) .Limitations. The accuracy of the learned environment models sets a bottleneck on the performance of MOMBO. The choice of the confidence radius β is another decisive factor in model performance. MOMBO shares these weaknesses with the other state-of-the-art model-based offline reinforcement learning methods (Yu et al., 2020; Lu et al., 2021; Sun et al., 2023) and does not contribute to 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 we claim a tighter bound compared to a baseline bound that we derived ourselves. However, the most important factor in that bound, R 2 max /(1 -γ) 2 , arises from Hoeffding's inequality, i.e., a classical statement. As such, we assume our bound to be rigorous and not inadvertently biased. Finally, our moment matching method is limited to activation functions for which the first two moments are either analytically tractable or can be approximated with sufficient accuracy, which includes all commonly used activation functions. Extensions to transformations such as, e.g., BatchNorm (Ioffe and Szegedy, 2015) , or other activation functions remove the analytical tractability for moment matching. However, these are usually not used in our specific applications.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We provide the necessary details about the experimental pipeline in Appendix C, along with the source code needed to reproduce our results.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of 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, in appendix, or as supplemental material. 
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.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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] Justification: The paper focuses on foundational research. As such there is no immediate short term social impact of the work to 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 no societal 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 not tied to particular applications, let alone deployments. However, if there is a 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, it is not needed to point out that a 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 as 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 or 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). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: As our research is foundational research there is no safeguard needed at the current level of research. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: 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? 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 or dataset.Guidelines:• The answer NA means that the paper does not involve crowdsourcing 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 we 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 as the institution conducting the review. 
 
 Appendix: Offline reinforcement learning. The goal of offline reinforcement learning is to derive an optimal policy from fixed datasets collected through interactions with a 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 robotics (Singh et al., 2020; Micheli et al., 2022) and healthcare (Shortreed et al., 2011; 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 on 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) .Uncertainty-driven offline reinforcement learning approaches apply penalties in various ways. As model-free methods, EDAC (An et al., 2021) applies penalties based on ensemble similarities, while PBRL (Bai et al., 2022) penalizes based on the disagreement among bootstrapped Qfunctions. Uncertainty-driven model-based offline reinforcement learning algorithms adhere to the meta-algorithm known as pessimistic value iteration (Jin et al., 2021) . This meta-algorithm introduces pessimism through a ξ-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 ξ-uncertainty quantifier is derived from the uncertainty of the Bellman target value, which does not rely on sampling. Instead, it 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-distribution 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-based approaches has been wellresearched, primarily in the area of Bayesian deep learning (Frey and Hinton, 1999; Wang and Manning, 2013; Wu et al., 2019a) , with applications, e.g., in active learning (Haußmann et al., 2019) and computer vision (Gast and Roth, 2018) .This section contains proofs for all the theoretical statements made in the main text.We first restate and extend the following result on moment matching with the ReLU activation function. Our focus on the ReLU activation throughout this work is due to the fact that it 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 ∼ N (X|µ, σ 2 ) and Y = max(0, X), we have(i) µ ≜ E [Y ] = µΦ(α) + σϕ(α), σ 2 ≜ var [Y ] = (µ 2 + σ 2 )Φ(α) + µσϕ(α) -µ 2where α = µ/σ, and ϕ(•), Φ(•) are the probability density function (pdf) and cumulative distribution function (cdf) of the standard normal distribution, respectively. Additionally, we have that(ii) µ ≥ µ and (iii) σ 2 ≤ σ 2 .Proof of Lemma 1. See Frey and Hinton (1999) for the derivation of (i), i.e., µ and σ 2 .Statement (ii) holds asµ = E [Y ] = ∞ 0 xp(x)dy ≥ 0 -∞ xp(x)dx + ∞ 0 xp(x)dx = µ.Note thatE Y 2 = ∞ 0 x 2 p(x)dx ≤ 0 -∞ x 2 p(x)dx + ∞ 0 x 2 p(x)dx = E X 2 .Therefore, we get (iii) byσ 2 = var [Y ] = E Y 2 -µ 2 ≤ E X 2 -µ 2 = σ 2concluding the proof.Although we do not require the following lemma in our final results, a helpful result is the following inequality between the cdfs of two normally distributed random variables X and its matched counterpart X. Lemma 4 (An inequality between normal cdfs). For two normally distributed random variables X ∼ N (X|µ, σ 2 ) and X ∼ N ( X| µ, σ 2 ) with µσ ≥ µ σ, the following inequality holds:F X (u) ≤ F X (u) for u ≤ 0.Given µ and σ 2 as in Lemma 1, the assumptions of this lemma hold and it is applicable to our moment matching propagation.Proof of Lemma 4. Assume u ≤ 0. We have thatF X (u) ≤ F X (u) ⇔ Φ u -µ σ ≤ Φ u -µ σ .As it is a monotonically increasing function, this in turn is equivalent tou -µ σ ≤ u -µ σ ⇔ uσ -µσ ≤ u σ -µ σ ⇔ u(σ -σ) ≤ µσ -µ σwhich holds as u(σ -σ) < 0 via Lemma 1 for u < 0, and µσ ≥ µ σ by assumption.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 p-Wasserstein distance between two probability measures ρ 1 and ρ 2 on M , with finite p-moments, is defined asW p (ρ 1 , ρ 2 ) ≜ inf γ∈χ(ρ1,ρ2) E (x,y)∼γ [d(x, y) p ] 1 /pwhere χ(•, •) represents the set of all couplings of two measures, ρ 1 and ρ 2 .For p = 1, the 1-Wasserstein distance, W 1 , can be simplified (see, e.g., Villani et al., 2009) toW 1 (ρ 1 , ρ 2 ) = R F 1 (x) -F 2 (x) dx (2)where ρ 1 , ρ 2 are two probability measures on R, and F 1 (•) and F 2 (•) are their respective cdfs.Given this definition, we prove the following two bounds on W 1 . Lemma 5 (Wasserstein inequalities). The following inequalities hold for the 1-Wasserstein metric(i) W 1 (ρ 1 U , ρ 2 U ) ≤ A W 1 (ρ 1 V , ρ 2 V ), for U = AV + b. (ii) W 1 (ρ 1 Z , ρ 2 Z ) ≤ W 1 (ρ 1 U , ρ 2 U ), for Z = g(U ) with g(•) locally K-Lipschitz with K ≤ 1 where ρ x is the density of x, U ∈ R du , V ∈ R dv , A ∈ R du×dv , and b ∈ R du .Proof of Lemma 5. Given the Kantorovich-Rubinstein duality (Villani et al., 2009) for W 1 we haveW 1 (µ, ν) = 1 K sup lip(f )≤K E x∼µ [f (x)] -E y∼ν [f (y)]for any two probability measures µ, ν, and where lip(f ) < K specifies the family of K-Lipschitz functions.We use this to first prove (ii). For Z = g(U ) we have thatW 1 (ρ 1 Z , ρ 2 Z ) = sup lip(f )≤1 E z∼ρ 1 Z [f (z)] -E z∼ρ 2 Z [f (z)] = sup lip(f )≤1 E u∼ρ 1 U [f (g(u))] -E u∼ρ 2 U [f (g(u))] , change of variables ≤ sup lip(h)≤1 E u∼ρ 1 U [h(u)] -E u∼ρ 2 U [h(u)] = W 1 (ρ 1 U , ρ 2 U )where the inequality holds as f • g has a Lipschitz constant L ≤ 1, i.e., we end up with a supremum over a smaller class of Lipschitz functions which we revert in the inequality. This gives us the desired inequality.To prove (i) we use that for A ≤ 1, the transformation g(V ) ≜ 1 is 1-Lipschitz and the result follows via (ii). For A > 1 we use that g(•) is K-Lipschitz and the result follows by adapting the proof accordingly.Finally, the following lemma allows us to relate the 1-Wasserstein distance of two transformed distributions to an absolute difference between the respective transforming functions. Lemma 6 (Wasserstein to suboptimality). Consider two probability distributions P x , P u defined on a measurable space (S, σ(S)) and two transformed random variables y = f (x) ∼ P y (with x ∼ P x ), and z = g(u) ∼ P z (with u ∼ P u ) for functions f, g :S → R. If W 1 (P y , P z ) ≤ C then |E x∼Px [f (x)] -E u∼Pu [g(u)] | ≤ C for any constant C ≥ 0.Proof of Lemma 6. Assume that W 1 (P y , P z ) ≤ C for some C ≥ 0. The Kantorovich-Rubinstein duality (Villani et al., 2009) , with K = 1 gives usW 1 (P y , P z ) = sup lip(h)≤1 E y∼Py [h(y)] -E z∼Pz [h(z)] ≤ C.As the identity function λ : x → x is 1-Lipschitz, we can lower bound the supremum further withE y∼Py [y] -E z∼Pz [z] ≤ sup lip(h)≤1 E y∼Py [h(y)] -E z∼Pz [h(z)] ≤ C.Using what is colloquially known as the law of the unconscious statistician (Casella and Berger, 2024) , we have thatE y∼Py [y] = E x∼Px [f (x)]and analogously for E z∼Pz [z] . ThereforeE y∼Py [y] -E z∼Pz [z] = E x∼Px [f (x)] -E u∼Pu [g(u)] which gives us E x∼Px [f (x)] -E u∼Pu [g(u)] ≤ C.As W 1 is symmetric, we can follow the same argument and additionally haveE u∼Pu [g(u)] -E x∼Px [f (x)] ≤ C.Using that for x ∈ R with x < C and -x < C it follows that |x| < C, we can combine these two inequalities and getE x∼Px [f (x)] -E u∼Pu [g(u)] ≤ C as desired.Proving Theorem 2 requires the following upper W 1 bound for an MLP with a 1-Lipschitz activation function.Y l = A ⊤ l X l-1 + b l , X l = r(Y l ), l = 1, . . . , Lwhere A l , b l are the weight matrix and bias vector for layer l, respectively, and r(•) denotes an elementwise 1-Lipschitz activation function. We write Y ≜ Y L = f (X 0 ), where X 0 is the input. Assuming two measures X 1 ∼ ρ 1 X and X 2 ∼ ρ 2 X , we have thatW 1 (ρ 1 Y , ρ 2 Y ) ≤ L l=1 A l W 1 (ρ 1 X , ρ 2 X ).Proof of Lemma 7. We have thatW 1 (ρ 1 Y , ρ 2 Y ) ≤ A L W 1 (ρ 1 X L-1 , ρ 2 X L-1 ), via Lemma 5 (i) = A L W 1 ρ 1 h(Y L-1 ) , ρ 2 h(Y L-1 ) ≤ A L W 1 ρ 1 Y L-1 , ρ 2 Y L-1 , via Lemma 5 (ii) ≤ • • • ≤ L l=1 A l W 1 (ρ 1 X , ρ 2 X ).The W 1 distance between two measures on the output distribution is upper bounded by the W 1 distance between the corresponding input measures with a factor given by the product of the norms of the weight matrices A l . This result allows us to provide a probabilistic bound on the W 1 distance between a normal ρ Y and its empirical, sampling-based, estimate ρY .The following lemma allows us to extend this result to a probabilistic sampling-based upper bound. Lemma 8 (Sampling-based MLP bound). For an L-layer MLP f (•)with 1-Lipschitz activation function, let Y ≜ Y L = f (X) where X ∼ N (X|µ X , σ 2 X ), the following bound holds:P W 1 (ρ Y , ρ N Ŷ ) ≤ L l=1 A l - 8 log( δ /4)R 2 max N /2 (1 -γ) 2 ≥ 1 -δwhere δ ∈ (0, 1) and ρ N Ŷ is the empirical estimate given N i.i.d. samples.Proof of Lemma 8. With Hoeffding's inequality (Wasserman, 2004) , we have thatP (|μ N -µ| ≥ ε) ≤ 2 exp -2N 2 ε 2 N i=1 R max 1 -γ -0 2 = 2 exp - 2ε 2 N (1 -γ) 2 R 2 maxfor the empirical mean μN of Y . As variances are U-statistics with order m = 2 and kernel h = 1 2 (x -y) 2 using Hoeffding (1963); Pitcan (2017), we have for the empirical covariance of Y , σ2N , thatP(σ 2 N -σ 2 ≥ ε) ≤ exp - ε 2 N /m 2 h 2 ∞ = exp - ε 2 N /2 (1 -γ) 2 2R 2 max ⇔ P(σ 2 N -σ 2 ≤ ε) ≥ 1 -exp - ε 2 N /2 (1 -γ) 2 2R 2 max where • ∞ denotes L ∞ -norm. Applying the inequality σN ≤ √ σ 2 + ε ≤ σ + √ ε, we get P(σ N ≤ σ + ε) ≥ 1 -exp - ε 2 N /2 (1 -γ) 2 2R 2 max ⇔ P(σ N -σ ≥ ε) ≤ exp - ε 2 N /2 (1 -γ) 2 2R 2 max .As the same inequality holds for σ -σN ≥ ε, a union bound give us thatP(|σ -σn | ≥ ε) ≤ 2 exp - ε 2 N /2 (1 -γ) 2 2R 2 max .Using an analytical upper bound (Chhachhi and Teng, 2023) on the 1-Wasserstein distance between two normal distributions,W 1 (N (µ 1 , σ 2 1 ), N (µ 2 , σ 2 2 )) ≤ |µ 1 -µ 2 | + |σ 1 -σ 2 | we have that P W 1 (N (µ,σ 2 ), N (μ N , σN )) ≥ 2ε ≤ P (|µ -μN | + |σ -σN | ≥ 2ε) ≤ P(|μ N -µ| ≥ ε or |σ N -σ| ≥ ε) ≤ P(|μ N -µ| ≥ ε) + P(|σ N -σ| ≥ ε), union bound ≤ 2 exp(-2N cε 2 ) + 2 exp(-0.5 N /2 cε 2 ) ≤ 4 exp(-0.5 N /2 cε 2 ) ≜ δ,where we use the shorthand notation c = (1 -γ) 2 /R 2 max . Solving for ε = 2ε gives usε = - 8 log( δ /4) N /2 c .The bound is then given asP W 1 N (µ, σ 2 ), N (μ N , σ2 N ) ≤ - 8 log( δ /4) N /2 c ≥ 1 -δwhich gives us with Lemma 7, thatP W 1 (ρ Y , ρ N Ŷ ) ≤ L l=1 A l - 8 log( δ /4)R 2 max N /2 (1 -γ) 2 ≥ 1 -δ.Finally, applying this result to the relation in Lemma 6 allows us to prove the desired non-asymptotic bound on the performance of the sampling-based model.Theorem 2 (Suboptimality of sampling-based PEVI algorithms). For any policy π MC learned by A PEVI ( B Γ π Q, P) using N Monte Carlo samples to approximate the Bellman target with respect to an action-value network defined as an L-layer MLP with 1-Lipschitz activation functions, the following inequality holds for any error tolerance δ ∈ (0, 1) with probability at least 1 -δSubOpt(π MC ) ≤ 2H L l=1 A l - 8 log( δ /4)R 2 max N /2 (1 -γ) 2 . Proof of Theorem 2. Define C = L l=1 A l - 8 log( δ /4)R 2 max ⌊ N /2⌋(1-γ) 2 .Then, combining the bound from Lemma 8 with the result from Lemma 6, we get that with probability 1 -δE X∼N (µ X ,σ 2 X ) B π Q(X) -E X∼N ( µ X , σ 2 X ) B π Q(X) ≤ C B π Q(s, a) -E X∼N ( µ X , σ 2 X ) B π Q(X) ≤ C B π Q(s, a) -B π Q(s, a, s ′ ) ≤ C.Note that X = (s, a, s ′ ), i.e., the state, action, and next state tuple, is distributed as described in the main text. Therefore C is a ξ-uncertainty quantifier, with ξ = δ. Applying Theorem 1 finalizes the proof.In order to prove Theorem 3, i.e., an upper bound on our proposed moment matching approach, we require the following two lemmata.The first lemma allows us to upper bound the 1-Wasserstein distance between an input distribution and its output distribution after transformation by the ReLU activation function.Lemma 2 (Moment matching bound). For the following three random variablesX ∼ ρ X , Y = max(0, X), X ∼ N ( X| µ, σ 2 ) with µ = E [Y ] and σ 2 = var [Y ] the following inequality holds W 1 (ρ Y , ρ X ) ≤ 0 -∞ F X (u)du + W 1 (ρ X , ρ X )whereF X (•) is cdf of X . If ρ X = N (X|µ, σ 2 ), it can be further simplified to W 1 (ρ Y , ρ X ) ≤ σϕ µ σ -µΦ - µ σ + |µ -µ| + |σ -σ|.Proof of Lemma 2. For a generic X ∼ ρ X , Y = max(0, X) and X ∼ N ( µ, σ2 ), we have withF Y (u) = 1 u≥0 F X (u) 2 and (2) that W 1 (ρ Y , ρ X ) = R F Y (u) -F X (u) du = 0 -∞ F X (u)du + ∞ 0 F X (u) -F X (u) du ≤ 0 -∞ F X (u)du + ∞ 0 F X (u) -F X (u) du + 0 -∞ F X (u) -F X (u) du = 0 -∞ F X (u)du + W 1 (ρ X , ρ X ).If ρ X = N (X|µ, σ), we can upper bound the W 1 distance via Chhachhi and Teng ( 2023)W 1 (N (µ 1 , σ 2 1 ), N (µ 2 , σ 2 2 )) ≤ |µ 1 -µ 2 | + |σ 1 -σ 2 | and evaluate the integral 0 -∞ F X (u)du = 0 -∞ Φ u -µ σ du = σϕ µ σ -µΦ - µ σwhere we used thatΦ(a + bx)dx c = 1 b (a + bx)Φ(a + bx) + ϕ(a + bx)where c = signifies equality up to an additive constant. Together this gives us thatW 1 (ρ Y , ρ x) ≤ σϕ µ σ -µΦ - µ σ + |µ -µ| + |σ -σ|.Lemma 3 (Moment matching MLP bound). Let f (X) be an L-layer MLP with ReLU activation r(x) = max(0, x). For l = 1, . . . , L -1, the sampling-based forward-pass isY 0 = X s , Y l = r(f l (Y l-1 )), Y L = f L (Y L-1 )where f l (•) is the l-th layer and X s a sample of N (X|µ X , σ 2 X ). Its moment matching pendant isX 0 ∼ N ( X 0 |µ X , σ 2 X ), X l ∼ N X l E r(f l ( X l-1 )) , var r(f l ( X l-1 )) .Y = ρ X L = N ( X L |E[f ( X L-1 )], var[f ( X L-1 )]). W 1 (ρ Y , ρ Y ) ≤ L l=2 G( X l-1 ) + C l-1 L j=l A j with G( X l ) = σ l ϕ µ l σ l -µ l Φ - µ l σ l ≤ 1, C l ≤ |A l µ l-1 -µ l | + A 2 l σ 2 l-1 -σ lwhere √ • and (•) 2 are applied elementwise.Proof of Lemma 3. Using Lemma 2 we have that the W 1 distance between the deterministic forward pass ρ Y l and the matching-based forward pass ρ X l is given asW 1 (ρ Y l , ρ X l ) ≤ G( X l ) + W 1 (ρ f l (Y l-1 ) , ρ X l ), where G(X) ≜ 0 -∞ F X (u)du ≤ G( X l ) + W 1 (ρ f l (Y l-1 ) , ρ f l ( X l-1 ) ) + W 1 (ρ f l ( X l-1 ) , ρ X l ), via the triangle inequality ≤ G( X l ) + A l W 1 (ρ Y l-1 , ρ X l-1 ) + C l , C l ≜ W 1 (ρ f l ( X l-1 ) , ρ X l ) ≤ G( X l ) + A l G( X l-1 ) + A l-1 W 1 (ρ Y l-2 , ρ X l-2 ) + C l-1 + C l .That is, for the entire MLP we haveW 1 (ρ Y , ρ Y ) = W 1 (ρ Y L , ρ X L ) = A L W 1 (ρ Y L-1 , ρ X L-1 ) ≤ A L G( X L-1 ) + A L-1 W 1 (ρ Y L-2 , ρ X L-2 ) + C L-1 ≤ L l=3 G( X l-1 ) + C l-1 L j=l A j + W 1 (ρ Y1 , ρ X1 ) L l=2 A lwhere, finally,W 1 (ρ Y1 , ρ X1 ) ≤ σϕ µ σ -µΦ - µ σ + |A 1 µ X -µ 1 | + A 2 1 σ 2 X -σ 1 .G( X l ) and C l are given asG( X l ) = σ l ϕ µ l σ l -µ l Φ - µ l σ l , C l ≤ |A l µ l-1 -µ l | + A 2 σ 2 l-1 -σ lwhere √ • and (•) 2 are applied elementwise.Theorem 3 (Suboptimality of moment matching-based PEVI algorithms). For any policy π MM derived with A PEVI ( B Γ π Q, P) learned by a penalization algorithm that uses moment matching to approximate the Bellman target with respect to an action-value network defined as an L-layer MLP with 1-Lipschitz activation functions, the following inequality holdsSubOpt(π MM ) ≤ 2H L l=2 G( X l-1 ) + C l-1 L j=l A j .Proof of Theorem 3. The proof is analogous to the one of Theorem 2. Note that X = (s, a, s ′ ), i.e., the state-action pair, is distributed as described in the main text.Define C = l=2 G( Xl-1 ) + C l-1 L j=l A j .By combining this theorem with the results from Theorem 1, we have thatE X∼N (µ X ,σ 2 X ) B π Q(X) -E X∼N (μ X ,σ 2 X ) B π Q(X) ≤ C B π Q(s, a) -E X∼N (μ X ,σ 2 X ) B π Q(X) ≤ C B π Q(s, a) -B π Q(s, a, s ′ ) ≤ Cholds deterministically. Therefore, C is a ξ-uncertainty quantifier for ξ = δ = 0. Applying Theorem 1 finalizes the proof.We sample an initial state s 1 from the environments. Next, we use the expert policies π * , treated as optimal policies, provided3 in the D4RL datasets (Fu et al., 2020) to sample action a 1 from the expert policy. Using the learned environment model P, we predict the next state samples s 2 ∼ P(•|s 1 , a 1 ). For each sample, we evaluate the next action-value Q(s 2 , π * (s 2 )) using the previously learned critics as the value function. We then evaluate the next action-value using moment matching Q M M (s 2 , π * (s 2 )) and visualize the distributions. For this experiment, we use the learned critics from seed 0. We repeat this process for all tasks in the D4RL dataset and present the results in Figure 3 . 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 10 th step, including the final step. For each selected tuple, we calculate the mean accuracy and 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 is 1000. Table 2 shows the results of this experiment.The behavior policy of the Mixed datasets consists of a combination of random and medium or expert policies, mixed at 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 is not feasible. See Figures 5 to 6 for visualizations of the learning curves. We use N ens = 7 ensemble elements and select N elite = 5 elite elements from the ensemble based on a validation dataset containing 1000 transitions. Each model in the ensemble consists of a 4-layer feedforward neural network with 200 hidden units, and we apply L2 weight decay with the following weights for each layer, starting from the first layer: [2.5×10 -5 , 5×10 -5 , 7.5×10 -5 , 7.5×10 -5 , 1× 10 -4 ]. We train the environment model using maximum likelihood estimation with a 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 walker2d-medium 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, 2, 3) and use them for each baseline.Architecture and optimization details. We train MOMBO for 3000 episodes on the D4RL 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 a 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 critic 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 α parameter is learned during training with the learning rate of 0.0001, except for hopper-medium and hopper-medium-replay, where α = 0.2. We set the target entropy to -dim(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 halfcheetah-medium-expert, 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 a sample from the predicted distribution for the next state and reward. We store these rollout transitions in a 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 D for the D4RL dataset and the Mixed dataset, respectively. For the Mixed dataset, the rollout length (k) is 5 for all tasks. Table 5 presents the rollout length for each task for the D4RL dataset.Penalty coefficients. We adopt all configurations from MOBILE with the exception of β due to the differences in the scale of the uncertainty estimators. For the D4RL dataset, we provide our choices for β 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 A100 GPU, AMD EPYC 7F72 CPU at 3.2 GHz, and 256 GB of memory; and 3) GeForce RTX 4090 GPU, Intel(R) Core(TM) i7-14700K CPU at 5.6 GHz, and 96 GB of memory. We measure the computation time for 1000 gradient steps to be approximately 8 seconds for the D4RL dataset and 11 seconds for the Mixed dataset on the third device. Assuming the environ- Figures 4 to 6 illustrate the learning curves of PEVI-based approaches, including ours, across gradient steps. In the figures, the thick (dashed/dotted/solid) 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 halfcheetah 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. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Although 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.The indicator function 1 is defined as 1S = 1 if the statement S is true and 0 otherwise.https://rail.eecs.berkeley.edu/datasets/offline_rl/ope_policies/onnx/https://github.com/yihaosun1124/mobile/tree/4882dce878f0792a337c0a95c2 7f4abf7e926101 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
As 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.
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: ** I am unfamiliar with the methods/related works in this paper.**

This paper proposed a model-based method for offline RL, based on moment matching. Improved numerical results are presented.
strength_weakness: strength_weakness: I am quite confused about Sec 4. Why the Gaussian RVs are related? 

I feel the presentation of the paper can be improved.
questions: 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? 
[1]Li, Gen, et al. "Settling the sample complexity of model-based offline reinforcement learning." The Annals of Statistics 52.1 (2024): 233-260.
[2]Wang, Yue, Jinjun Xiong, and Shaofeng Zou. "Achieving the Asymptotically Optimal Sample Complexity of Offline Reinforcement Learning: A DRO-Based Approach."
limitations: limitations: n/a
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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).
- 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.
- 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.

[1] Sun et al. "Model-Bellman Inconsistency for Model-based Offline Reinforcement Learning" (ICML 2023)

[2] Jin et al. "Is Pessimism Provably Efficient for Offline RL?" (ICML 2021)
questions: questions: To summarize the main points/questions raised in the weaknesses section:

- Could the authors evaluate their method on more challenging tasks to show that it can achieve better asymptotic performance compared to MOBILE?
- 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?

I‘m willing to raise my score if all these concerns can be addressed by the authors.
limitations: limitations: The authors have thoroughly discussed the limitations of their method in the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.There is a clear necessity for the authors to enhance their writing proficiency. 

2.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.

3.How can the alignment of moments from one distribution with those from another ensure the approximation or equivalence of the two?

4.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.

5.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.

6.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.

7.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?

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.
questions: questions: See weaknesses.
limitations: limitations: None.

Comments: title: Further Clarifications
Further Clarifications
Clarifications about Pessimism and SOTA
Clarifications about Experiments and Results
Clarifications about Contribution Statement and Motivation
Discussion Period
comment: 3. “Based on my understanding, a lower standard deviation typically indicates a faster rate of convergence.“

Can you specify the problem with treating the area under learning curve scores we reported in Table 1 as a quantitative measure of convergence rate and drawing conclusions accordingly? We do understand the merit of looking carefully into individual figures, however, as scientists, we believe more in numbers than bare eye analysis. From the fact that in 11 of the 12 tasks MOMBO delivers higher AUC and similar cross-experiment variance compared to MOBILE, it converges faster.

MOMBO and MOBILE perform similarly only when both models are trained for 3 million steps. This training time is typically chosen by performing online evaluations. However, in a realistic offline RL setup, we do not have access to the real environment on which we can check whether our model is sufficiently trained. Furthermore, a model that converges faster would fit better to data when its capacity is increased.

The significance of a low standard deviation is improved training robustness. This comes as a side benefit of eliminating a factor of randomness from the training process and an explanation of why our method reaches faster convergence, which is our main result. MOBILE and MOPO perform comparably only when both are trained long enough. This outcome serves as a consistency check and is not the main point of our contribution.

We hope we were able to answer your questions. Thanks again for the truly enjoyable scientific discussion we had together. We fully trust your fair scientific judgment and respect any score you give in your final evaluation.
1. We promise once again to implement all of your suggestions, which comprise: 
    1. Shortening the sentences that are hard to read
    2. Unifying the allocation of terminology, sticking to their usage in the literature
    3. Correcting wrongly constructed sentences
    4. Rephrasing sentences written in passive voice in active voice in the most understandable way
    5. Correcting improperly used writing conventions throughout the text.
    
We do think they are extremely important to improve the quality of our work. Their implementation does require a thorough and careful pass over the text, which is a tedious task and we are ambitious to fulfill it following your guidance. We kindly express our opinion that all these corrections are typically handled during the preparation of a camera-ready versions of conference papers. We encounter corresponding period in journal publication processes as **accept with minor revision**.

2. Here are our answers to your questions:

We would like to first make sure that we agree on the terminology. We study the **variance** of *the Monte Carlo estimators used in the penalty term calculation of model-based offline RL algorithms*. Using an estimator is inevitable while calculating a penalty term based on the Bellman operator. In mathematical terms, there is uncertainty on the next state $s’$ stemming from the error in the learned environment dynamics. This uncertainty needs to be propagated through the state-action value function $Q(s’,a’)$ which appears in the Bellman target calculation. This operation is analytically intractable if $Q$ is a neural network. 

> Is high variance the only issue with ineffective Monte Carlo sampling? 

We are sure it is not the only issue. Our claim is that it is one of the many important issues as it plays a significant role in training efficiency. We demonstrate why it is the case by developing concentration inequalities for the Monte Carlo estimator and our moment matching based alternative.

> Given that the authors have only analyzed the variance issue, it naturally raises the following concern:” Does the variance problem only arise with Monte Carlo sampling?

As we pose the problem in the first place as the estimator variance of a Monte Carlo estimate, the answer is yes. Reducing estimator variance is major issue in a wide range of problems in statistics. It is handled using methods such as introducing control variates or Rao-Blackwellization. See for instance: Lemieux (2017).


The state-of-the-art approach to approximate the uncertainty of the state-action value of the next state is to use Monte Carlo integration, which is represented in our paper with the MOBILE algorithm (Sun et al., 2023). Sampling-free approaches do exist, but only for penalty terms that are computed directly on the estimated error of the model dynamics. We represent this alternative in our paper by the MOPO algorithm (Yu et al., 2021). However, as proven in Theorem 4.3 of the MOBILE paper, this option over-penalizes the reward and causes suboptimal results. The tightness experiment we reported in our global response demonstrates that MOPO applies much higher penalty than MOBILE as this theorem suggests, and our MOMBO applies even less penalty, hence brings about much faster training. 

> Compared to other means of addressing the variance issue, what is the motivation for using a sampling-free approach in this paper

There does not exist an alternative way to address this issue yet. Our key novelty is that we are the first to detect its existence and prominence for offline reinforcement learning and provide a solution to mitigate it.


-----
Lemieux, C., _Control Variates_ (Wiley StatsRef: Statistics Reference OnlineL 1-8 2017)

Yu et al., _COMBO: Conservative offline model-based policy optimization_ (NeurIPS 2021)

Sun et al., _Model-bellman inconsistency for model-based offline reinforcement learning_ (ICML 2023)
1. Could you please clarify why only minor writing changes are considered? If there is no substantial evidence to counter the multiple shortcomings I have highlighted, then my concerns remain unresolved.

2. The authors have yet to adequately address my concerns regarding the motivation behind the research. I kindly request that you respond directly to the three questions I posed in the Clarification of Research Motivation.

3. Based on my understanding, a lower standard deviation typically indicates a faster rate of convergence. However, I am concerned that from Figures 2 through 5, the proposed method demonstrates a fast convergence rate in only a few tasks in 4/12  (Figure 1: HalfCheetah, Figure 2: HalfCheetah, Walker, and Figure 3: Hopper). Given that the convergence rate is not consistently guaranteed, fluctuations are expected, and performance does not exceed that of MOBILE, what is the significance of a low standard deviation? From the current graphs, it appears that MOMBO, despite having a high AUC, does not outperform MOBILE when considering the same number of training steps.
Thank you for the additional experiments. I maintain my original rating, thank you.
Thanks for your detailed response! I’m hence keeping my positive rating.
Thank you indeed for your response. 

> Thank you for the response. The experiments do show some improvement, but I am still not convienced that the previous methods are too pessimistic.

We believe the tightness scores we reported in the global response provide direct and strong evidence in favor of our claim that MOMBO applies less penalty than both MOPO and MOBILE. The experiments demonstrate that MOPO applies significantly higher penalty than MOMBO and MOBILE, which is an expected consequence of Theorem 3.6 of Sun et al. (2023). To remind, tightness score measures the difference between the penalty applied by a model and the actual error the Bellman estimator makes. See our global response for details. If there still remains an open issue regarding the relative pessimism of the models in comparison, we are happy to discuss further.

> Considering the different counstruction of [Li et al. 2024] and the MOPO baseline, I am wondering if the MOPO is the SOTA of current uncertainty-based approaches.

We actually consider MOBILE as the SOTA of the model-based offline RL models, not MOPO, given its publication date (ICML, 2023) and the lack of a widely known method that improves its performance using commensurate resources. That said, MOPO is still a reference model representing the option of reward penalization based on environment dynamics. MOBILE improves on MOPO by deriving a penalty term based on a sampling-based estimate of the Bellman target. Our contribution is to identify and characterize the problems this sampling-based approach causes and present a novel solution that addresses these problems both theoretically (see our Theorems 1 and 2, and comparison paragraph between lines 247-258.) and empirically (see our AUC results in Table 1 that demonstrate faster convergence).



-----
Sun et al., _Model-bellman inconsistency for model-based offline reinforcement learning_ (ICML 2023)
Thank you for the response. The experiments do show some improvement, but I am still not convienced that the previous methods are too pessimistic. 

Considering the different counstruction of [Li et al. 2024] and the MOPO baseline, I am wondering if the MOPO is the SOTA of current uncertainty-based approaches.
Hi reviewer igqb,

Thanks so much for your lively engagement with the authors so far. We are about to lose access to them in the discussion period (August 13, 11:59pm AoE). I'm hoping that you might be able to weigh in one more time regarding your concerns about the comparison between MOBILE and MOMBO. The authors have made several points in their most recent comment. They argue that the key improvements are learning with fewer training steps and improved uncertainty estimates (see the PDF from the rebuttal); they have expanded upon the halfcheetah results; and they have clarified how the Adroit environments were selected.

Could you take a look and let us know the extent to which your concerns have been alleviated by these responses?

Thanks!
> 4. Further Experimental Validation and Concerns:

We would like to reiterate our reply to Reviewer BmgH about the interpretation of our key experimental findings. 

In a realistic offline RL scenario, it is not possible to measure the evaluation-time performance, hence not trivial to decide how long a model should be trained. Hence, we argue that the steepness of the learning curve is a crucial performance criterion, and the area under the learning curve is an appropriate means to quantify its steepness. As visible in the right-most column of Table 1, MOMBO reaches a significantly steeper learning curve than MOBILE, which is what we would expect from an offline RL algorithm that calculates its reward penalty with less approximation error. The normalized reward results, in turn, serve as a sanity check. They demonstrate that our scheme maintains the original MOBILE performance when trained long enough. Notably, we expect an algorithm with a higher AUC score to deliver better results for a fixed amount of computational time or training data.

We would also like to bring the uncertainty quantification results we reported in the global response to your attention. These results provide direct evidence about why moment matching would address the sampling error more effectively.
 
> Figure 1 of the manuscript, it appears that only the halfcheetah-random task shows a faster convergence, with no significant advantage in the convergence rate for the other three tasks when compared to the high performance achieved by MOBILE. In the halfcheetah-medium-replay task, although it starts converging faster, it experiences significant fluctuations later, reflecting instability in the training process.

We believe there is a misunderstanding here about the implications of the results we presented. It is true that MOMBO fluctuates after 1 million steps during halfcheetah-medium-replay training a little more than MOBILE. However, this fluctuation is still within the one standard deviation distance of MOBILE's mean performance during the far majority of the training process. The eventual AUC scores of the two models are nearly identical even with similar standard deviations across experiment repetitions: 68.7 ± 3.9 for MOMBO and 67.9 ± 2.8 for MOBILE. We would like to remind that halfcheetah-medium-replay is among the few environments where MOMBO fluctuates more than MOBILE. In the majority of the tasks, MOMBO delivers clearly higher AUC than MOBILE, providing strong evidence of its improved training stability and supporting the central claim of our work. 

> Additionally, the standard deviation in the appendix's Figure 4 for both the halfcheetah-medium-replay and walker2d-medium-replay tasks is also significant.

As reported in Table 1, MOMBO delivers smaller standard deviations across experiment repetitions than other models in six of the twelve environments with respect to final reward, while MOBILE does the same for only four. However, one can see MOMBO's advantage over MOBILE more directly from the AUC scores, where it outperforms MOBILE in all environments but one. The legends of Figures 2, 3, and 4 convey the same information as Table 1 for individual environments. You can see in the final and AUC scores presented in the Figure 4 legends that MOMBO's variance across experiment repetitions is comparable to MOBILE, as expected. 

> Adroit experiments

We reported new experiments only on the two Adroit environments for which the official GitHub source of the MOBILE paper provides sufficient details regarding the configuration. We did not select the tasks based on other criteria.
Thanks for your meticulous review and valuable suggestions. We believe there remains a few points calling for further clarification, but we also fully respect your final decision.

> 1. Improvement Suggestions for Writing:

We agree with your points and appreciate your rigor in pointing them out greatly. We promise to handle them in the revised version of our work. However, we kindly express our view that their scope stays within the limits of a minor revision.

> 2. Clarification and Enhancement of Contribution Statement:
> 3. Clarification of Research Motivation:

We agree about the need for further improvement here. We believe presenting our core contribution in Section 1, `Our contribution` paragraph as follows would address both points 2 and 3, as well as your concern about the central focus of our work:

Estimating tight $\xi$-uncertainty quantifiers and using them for reward penalization is a provable recipe for improved offline reinforcement learning (see Theorem 4.2 of Jin et al. (2021)). Reward penalties calculated on the Bellman target are provably tighter $\xi$-uncertainty quantifiers than those calculated on the environment model (see Theorem 3.6 of Sun et al. (2023)). However, the calculation of this quantity necessitates propagating the uncertainty of the estimated next state through a non-linear value function, which is typically modeled by a deep neural net. The commonplace approach is to approximate this quantity by Monte Carlo sampling. Our key novelties include:

i) detecting the problem of approximation errors on the reward penalty for the first time in the literature
ii) characterizing the problem theoretically by deriving novel concentration inequalities on their commonplace Monte Carlo estimator,
iii) proposing an original solution to this problem by using a moment matching approach originally developed for Bayesian neural net inference for the first time in the context of reward penalty calculation,
iv) quantifying the approximation error incurred by the Monte Carlo estimate on a large set of benchmarks, showing a significant reduction of this error by our proposed moment matching approach, and demonstrating the impact of the error on convergence speed (improved AUC score).


-----
Jin et al., _Is Pessimism Provably Efficient for Offline RL?_ (ICML 2021)

Sun et al., _Model-bellman inconsistency for model-based offline reinforcement learning_ (ICML 2023)
4. Further Experimental Validation and Concerns:

I appreciate the additional experiments provided by the authors. However, even with the selected tasks, I fail to discern the advantages of MOMBO. The performance of the proposed method does not surpass MOBILE; the only advantage seems to be a slightly smaller standard deviation. However, the high standard deviation in the pen-cloned scenario contradicts the motivation of this paper. Moreover, the authors have corrected the expression from sample efficiency to convergence rate. However, in Figure 1 of the manuscript, it appears that only the halfcheetah-random task shows a faster convergence, with no significant advantage in the convergence rate for the other three tasks when compared to the high performance achieved by MOBILE. In the halfcheetah-medium-replay task, although it starts converging faster, it experiences significant fluctuations later, reflecting instability in the training process. Additionally, the standard deviation in the appendix's Figure 4 for both the halfcheetah-medium-replay and walker2d-medium-replay tasks is also significant.

In summary, I believe there is room for improvement in this paper's experimental validation, including, but not limited to, a broader comparison of experimental tasks, verification of method effectiveness (for example, modifying linear or ReLU layers to construct ablation studies), and the reliability of statistical results. These concerns also prevent me from confidently assigning a high score.
Thanks for the author's response, which has resolved some of my initial concerns. I have thoroughly read the comments from other reviewers and the author's replies, including the general response. Below, I provide explanations for the points raised by the author and list the concerns that are still unclear to me:


1. Improvement Suggestions for Writing:

Here are the areas I believe would benefit from refinement in writing:

- Sentences are too long and hard to read: Line 20-23 "Applying traditional …"， line 24-27 "Model-based …"， line 171-174 "As the input …"

- Conflation of expressions: distributional shift & distribution shift, behavioral policy & behavior policy. Refer to literatures MOBILE and COMBO.

- Sentence Construction: On lines 89-104, the frequent use of "when" and consecutive "while" at the beginning of sentences, coupled with a near absence of connecting words, raises concerns about the flow and coherence of the writing.

- Passive Voice: Upon reviewing lines 121-125, it has been observed that the text relies heavily on the passive voice. While the passive voice can be appropriate in certain contexts, an overreliance on it can obscure the clarity and directness of the writing.

- Writing Conventions: The placement of "E.g.," (Line 146) and "I.e.," (Lines 174, 226) at the beginning of sentences is not in accordance with standard academic writing practices. While these abbreviations are commonly used to provide examples or clarify meanings, their position at the sentence's onset can disrupt the flow and formality expected in scholarly articles.

Improvement of Pronoun Usage: The manuscript exhibits multiple instances of ambiguous or incorrect pronoun usage, which can lead to confusion for the reader (as seen on line 135; a comprehensive list is not provided here for the sake of brevity).

- Refinement of Expressions for Scholarly Rigor: The manuscript frequently employs expressions that may not meet the standards of precision expected in academic writing, such as "via" (the specific instances are not exhaustively listed here).

In summary, the numerous writing issues are a significant shortcoming of this paper. The abovementioned concerns may only encompass some areas that require attention; thus, the authors must invest further effort into meticulously reviewing the entire manuscript.


2. Clarification and Enhancement of Contribution Statement:

The conciseness of the stated contributions needs to be improved. The author's response appears somewhat superficial and has not yet captured the essence of the contributions. If I understand correctly, this paper's core contribution lies in employing a moment-matching method to address the variance issues arising from Monte Carlo sampling. However, the authors have not adequately reflected these critical elements, especially considering using the term "bottleneck" as a substitute.

3. Clarification of Research Motivation:

The motivation behind the study is somewhat ambiguous. If I have understood correctly, the paper's logic is as follows: Monte Carlo sampling leads to variance issues, which in turn prompts the proposal of a sampling-free method. It seems the authors aim to improve upon Monte Carlo sampling. However, I have concerns:

- Is high variance the only issue with ineffective Monte Carlo sampling? Given that the authors have only analyzed the variance issue, it naturally raises the following concern:
- Does the variance problem only arise with Monte Carlo sampling?
- Compared to other means of addressing the variance issue, what is the motivation for using a sampling-free approach in this paper?

In summary, the current manuscript's analysis of motivation is incomplete. I have not grasped the authors' central focus, which is why solving the variance problem necessitates using a sampling-free approach.
Thank you for your reply, my concerns have been mostly resolved. I have already improved my score.
Hi all! Just your friendly area chair checking in.

First, thanks to the reviewers for your work so far on this paper. The discussion period has begun, which includes both reviewers and authors. This is our opportunity to clear up misunderstandings, get answers to questions, and generally gather enough information to make an informed decision grounded in the acceptance criteria.

It seems there is general agreement that the MOMBO approach is novel and interesting. The main critiques raised seem to concern whether there is clear evidence of meaningful empirical advantages of MOMBO over existing methods, and some clarity issues in the manuscript. Does that seem like a fair high-level summary?

Reviewers: please carefully read the other reviews and the authors' responses and let us know what follow-up questions you have and to what extent your evaluation might change as a result of the additional context. Please especially raise any points of disagreement with other reviewers or the authors, as these are opportunities for us to clarify misunderstandings and detect misconceptions.