Deterministic Uncertainty Propagation for Improved Model-Based Offline Reinforcement Learning

Thank you for this thorough review. We will address your concerns below.

1.There is a clear necessity for the authors to enhance their writing proficiency.

Could you be specific as to where the writing proficiency is suboptimal?

2.The manuscript lists many contributions, which obscures its central thesis.

Our contribution is three-fold. (i) We propose a novel moment matching and propagation approach for uncertainty-based model-based offline RL. (ii) We highlight theoretical differences between it and a sampling-based approach. (iii) We provide empirical evidence for the usefulness of our proposal. We will restructure the relevant paragraph to make this more clear.

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

The Wasserstein 1-distance (Def 1) which we consider in this paper measures closeness between two distributions in terms of their cdf not their moments. However, e.g., Wasserstein 2-distance between two univariate Gaussian distributions can be shown to be the sum of the squared difference between their mean and variance terms, i.e., alignment of the moments minimizes their distance based on the $W_2$ metric. Does that answer your question?

The moment matching proposed by the authors is similar to transforming uncertainty/variance.

Indeed, as we propose a sampling-free approach, this necessitates a way to transform uncertainty as it passes throughout the network. Can the reviewer clarify why this analytical transformation is a weakness and not a feature?

Furthermore, the experimental results do not prove the method's validity due to the lack of comparison with the latest SOTA method.

EDAC (An et al., 2021) and its improvements, such as SPQR (Lee et al., 2023), focus on model-free learning, whereas we consider model-based algorithms. To our knowledge, MOBILE is the state-of-the-art method for model-based offline learning. Most recent follow-up work such as Luo et al. (2024) and Song et al. (2024) take MOBILE as the base algorithm and complements it with additional features which can be applied also to our method. In our general answer, we have included further metrics that highlight the difference between MOBILE's performance and ours.

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.

Thank you for the proposal. We will change the section title to Limitations of sampling-based uncertainty propagation to highlight the specific limitation we are addressing.

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.

We agree, using a lower confidence bound is a standard practice throughout offline RL. That is not a contribution of ours. What is relevant is how this bound is constructed, i.e., in our case $\hat \sigma$. See also Theorem 2 in Jin et al. (2021) regarding which derives the desired suboptimality bound on a lower bound estimator.

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?

There are still minor variations explored in the literature, but these tend to be either only theoretical or too costly, such that there have been no further fundamental new developments. E.g., moment-matching with complete covariance matrices is explored in the current work by Wright et al. (2024). But the cost scales quadratically in the layer width, i.e., becomes intractable for common network sizes. As the authors in that work acknowledge (see their Section 4.3), using the moment matching developed in our references performs just as well or even better. Similarly, Look et al. (2023) explored a way to propagate these covariances in neural stochastic differential equations using Stein's lemma. We will clarify the current state of the literature at the beginning of Section 4.

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.

We focused on mujoco as a popular and extensive benchmark. See our general answer above for more metrics and a preliminary set of adroit results.

---

An et al., Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble (NeurIPS 2021)

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

Lee et al., SPQR: Controlling Q-ensemble Independence with Spiked Random Model for Reinforcement Learning (NeurIPS 2023)

Look et al., A deterministic approximation to neural SDEs (PAMI 2023)

Luo et al., Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning (ICLR 2024)

Song et al., Compositional Conservatism: A Transductive Approach in Offline Reinforcement Learning (ICLR 2024)

Wright et al., An Analytic Solution to Covariance Propagation in Neural Networks (AISTATS 2024)

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)

Thank you for the detailed review.

The novelty of this method seems relatively weak. It appears to be a minor modification of the uncertainty quantification method used in MOBILE.

MOPO, MOBILE, and MOMBO are three algorithms derived from the meta-algorithm known as Pessimistic Value Iteration (Jin et al., 2021). They all attempt to quantify uncertainty to learn a pessimistic value function. Therefore, a high degree of similarity among them is expected. However, we demonstrate that our approach surpasses the SOTA (MOBILE) both analytically and empirically.

as noted by the authors in the limitations section, this method is also limited to certain activation functions.

Let us clarify that while this limitation exists, it is less restrictive than it sounds. The derivations required to compute the first two moments of a ReLU activation directly generalize to all other piecewise linear activation functions, e.g., Leaky-Relu or PReLU. While requiring a more expensive forward pass, the first two moments are also tractable for more modern activation functions such as the ELU and its SELU counterpart. To summarize, while the first two moments cannot be computed for all activation functions, they can be for the most relevant and most popular ones covering all mainstream architectures used in critic training in modern RL.

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.

We agree, that sample efficiency was poorly worded. We adapt the discussion paragraph in the final version of the paper accordingly to make this clear. (Table 2 refers to the new results provided in this rebuttal.)

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

However, the experiments in this paper do not show a strong improvement over the baseline (MOBILE).

Upon three million gradient steps, both MOBILE and MOMBO perform very similar, with a tendency for minor improvements in MOMBO. However, the AUC results, being calculated over ten evaluation results throughout the training process, show a clear distinction in the speed of convergence with MOMBO outperforming MOBILE in all but one environment. Convergence speed is a crucial metric in real applications where it is not feasible to engineer the training duration based on evaluation results, which is the common practice for standard offline RL benchmarks. Additionally, this means that for a fixed amount of computational time/training data, we can train more powerful models than MOBILE, which can be expected to then translate into a larger normalized reward as well. Finally, the experiments performed for the rebuttal (see general answer) show improvements in accuracy and tightness indicating a better training signal for our proposed approach.

Could the authors evaluate their method on more challenging tasks to show that it can achieve better asymptotic performance compared to MOBILE?

We provide additional performance metrics in our general answer demonstrating the improvements our method provides as well as preliminary results on the adroit domain.

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?

We have included an additional experiment in our main response that demonstrates improved performance in terms of accuracy and tightness upon both MOBILE and MOPO.

Please let us know if this answer clarifies your concerns.

---

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

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

Janner et al., Offline reinforcement learning as one big sequence modeling problem (NeurIPS 2021)

Rigter et al., RAMBO-RL: Robust adversarial model-based offline reinforcement learning (NeurIPS 2022)

Sun et al., Model-bellman inconsistency for model-based offline reinforcement learning (ICML 2023)

Thank you for your thorough reading despite the unfamiliarity.

I feel the presentation of the paper can be improved.

Given the overall feedback provided via the reviews, we updated the section title and presentation for Section 3, and the discussion and results paragraph in Section 5 (see, e.g., answer to BmgH). Please let us know if further parts of the paper could benefit from presentation improvements in your opinion.

Why the Gaussian RVs are related?

Assuming that the overall input follows a normal distribution is a common one in the literature coming from the environmental model. Assuming in turn that the pre-activation output of a neural layer is also Gaussian can be justified via a central limit theorem argumentation. This is because the input to a specific neuron is a sum of random variables, i.e., the central limit theorem holds and converges to a Gaussian as the width increases. Additionally, if the computation of the first two moments is analytically tractable, the Gaussian distribution has the property of being the maximum entropy distribution, i.e., it imposes the least additional number of assumptions. The fact that Algorithm 1 ignores covariance terms is common practice in moment-matching approaches (see the references in the paper), as the computational complexity and memory cost scales quadratically in the number of neurons otherwise.

The authors mention that the uncertainty-based approach can be overly pessimistic, but some recent papers [1,2] are showing that the uncertainty-based methods achieve the minimax optimal sample complexity, under the tabular setting.

Pessimism is required, however, current uncertainty-based approaches, e.g., MOPO and MOBILE, penalize too much. We do not contradict the two references, but show that our uncertainty estimator is tighter than our baselines'. See the additional experiments we provide in the main rebuttal answer. Note that the the term "samples" in those references refers to the size of the underlying data set, see, e.g., Section 2.2 in Li et al. (2024).

---

Li et al., Settling the sample complexity of model-based offline reinforcement learning (Annals of Statistics, 2024)

Thank you for your positive feedback.

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

We include an additional evaluation in our general answer above, to demonstrate both our improvements in accuracy as well as tightness.