Model-based Offline Reinforcement Learning with Lower Expectile Q-Learning

[Q6] Similarly in Section 5.4, Does a reduced value of $\tau$ in the expectile regression (more conservative estimates) address this issue?

As shown in Table 16 of Appendix C, reducing the value of $\tau$ helps decrease the standard deviation in Hopper tasks. However, in Walker2D tasks, fluctuations persist, resulting in a large standard deviation even with a lower $\tau$. This occurs because all the world models consistently predict overly optimistic transitions, similar to the failure cases observed in AntMaze.

[Q7] Placement of Section 6.1

Thank you for the suggestion! Replacing the true terminal function with a learned terminal function leads to a significant performance drop ($461.8 \rightarrow 232.6$), especially in diverse datasets ($233.2 \rightarrow 66.8$), where termination signals are scarce because the dataset is collected by navigating to randomly selected goals.

We added the experimental results and their setup in Table 22 of Appendix D. We moved the limitation to Section 6 (between the experiments and conclusion sections) and explained these results in L520-522 of Section 6.

References

[1] Hafner et al. "Mastering diverse domains through world models.", Arxiv 2023.

[Q1-1] On line 256, the authors state "because of the expectile term ..., we cannot directly compute the gradient... thus we optimize the unnormalized version of the expectile". Please could the authors expand on this line? What is meant by the 'unnormalized' version of the expectile

Thank you for pointing this out. From the definition, $\mathbb{E}^{\tau}[X]$, expectile of a distribution $X$, is a minimizer of $L(y) = \mathbb{E}\_{x \sim X}[|\tau - 1(y>x)| (y-x)^2]$. Since the gradient of $L(\mathbb{E}^{\tau}[X])$ is zero (because $\mathbb{E}^{\tau}[X]$ is the minimum), we have the formula $L'(\mathbb{E}^{\tau}[X]) = \mathbb{E}\_{x \sim X}[|\tau - 1(\mathbb{E}^{\tau}[X]>x)| \cdot 2 \cdot (\mathbb{E}^{\tau}[X]-x)|] = 0$. By rearranging the equation, we can have $\mathbb{E}^{\tau}[X] = \frac{\mathbb{E}\_{x \sim X}[|\tau - 1(\mathbb{E}^{\tau}[X]>x)| \cdot x]}{\mathbb{E}\_{x \sim X}[|\tau - 1(\mathbb{E}^{\tau}[X]>x)|]}$. We remove the normalizer (denominator) and refer $\mathbb{E}\_{x \sim X}[|\tau - 1(\mathbb{E}^{\tau}[X]>x)| \cdot x]$ as “unnormalized version of the expectile”.

We added this explanation on the derivation of surrogate loss and clarified the meaning of the terms in L257-258 of Section 4.3.

[Q1-2] Also, on line 258 the authors say they approximate the expectation with the learned Q-estimator in order to derive the surrogate loss. But then on line 266 they say the surrogate loss provides a better approximation than the Q-estimator, which naively seems contradictory (without going through Appendix B). Could the authors clarify this please?

To clarify, L266-267 means that minimizing our surrogate loss (consists of $Q\_\phi(\mathbf{s}\_t, \mathbf{a}\_t)$ and $Q^\lambda\_t(\mathcal{T})$) leads to a better policy than directly minimizing $-Q\_\phi(\mathbf{s}\_t, \mathbf{a}\_t)$. This is not contradictory.

Specifically speaking, the learned Q-estimator provides a high-bias, low-variance estimate, while the surrogate loss derived from $\lambda$-returns serves as a low-bias, high-variance estimator. When learned Q-values are inaccurate, using $\lambda$-returns offers a better approximation of the true Q-values by relying on a multi-step return computed (although it also uses the inaccurate learned Q-values instead of true Q-values), reducing bias at the cost of increased variance. We note that if $\tau = 0.5$, the policy optimization is equivalent to optimizing the $\lambda$-returns instead of Q-values.

[Q2] In Equation 7, can you confirm that $\beta$ and $\beta_{\text{EMA}}$ are independent hyperparameters? Assuming so, it would be less confusing to use a different symbol for $\beta_{\text{EMA}}$.

Yes, those are independent hyperparameters. To make it clearer, we changed the notation $\beta_{\text{EMA}}$ to $\omega_{\text{EMA}}$. Thank you for your suggestion!

[Q3] ​​Could you explain the main difference between LEQ and IQL-TD-MPC and include this in the paper?

Please refer to [W4-1, Q3].

[Q4] In Section 5, LEQ is compared with the baselines in terms of performance. It would also be helpful to have some comparison in terms of computational cost/run-time, as this is also an important aspect of these algorithms (and LEQ appears quite strong in this regard from the appendix). Could this be included?

Thank you for the suggestion. Unfortunately, we did not have enough space for the main paper to include a comparison table for the computational cost in Section 5. Instead, we compared the run-time with MOBILE and CBOP in L310-311 of Section 5.2, and referred to Table 8 of Appendix A. Notably, LEQ is $6$x faster than recent model-based offline RL algorithms like CBOP and MOBILE.

[Q5-1] I found this confusing as my understanding is that LEQ uses a single world model - is this the case?

Following MOBILE, LEQ uses an ensemble of 5 world models by selecting 5 out of 7 models with the best validation score. Please refer to L732-735 of Appendix A for further details.

[Q5-2] In Section 5.3, the authors discuss failed trajectories that attempt to pass through walls, and that 'this could be addressed by ... improving the number of ensembles for the world models.' … Do the authors have any hypothesis or evidence that using an ensemble would address this issue?

Our hypothesis for the AntMaze-medium maze is that if 1) at least one of the world models can predict that the wall is impenetrable and 2) we apply sufficiently low $\tau$, LEQ will correctly evaluate the returns (since it will value the most conservative model prediction). Thus, we believe that better world models (which can predict that the wall is blocked) or increasing the number of world models (so that at least one of the world models can predict that the wall is impenetrable) can address the issue.

Thank you for your constructive feedback. We address your questions and concerns in detail below and our new manuscript reflects your suggestions. Please refer to the updated PDF for new results and revisions.

[W1, W2-1] None of the components of the model are particularly novel, and as a result, I wouldn't expect LEQ to be able to solve any problem that is not already solvable by existing offline RL algorithms. Other insights such as the benefit of using λ-returns are generally well-established.

First, we propose a novel methodology to get conservative return estimates via expectile regression. Please note that application of expectile regression in LEQ fundamentally differs from IQL; LEQ utilizes lower expectile regression to get conservative return estimates from potentially inaccurate model rollouts, while IQL employs upper expectile regression to approximate the max operation in $V(s) = \max_{a} Q(s, a)$. Moreover, we introduce a novel surrogate loss for policy optimization, designed to maximize the expectile of sampled returns, previously unexplored in the field. Although prior works have demonstrated the effectiveness of $\lambda$-returns, our extension of policy gradients for $\lambda$-returns to expectile is novel and crucial for offline MBRL. We believe this innovation will be beneficial for future offline MBRL research. Finally, LEQ is able to solve AntMaze tasks, which is nearly unsolvable for previous model-based methods. Specifically, most of the model-based offline RL shows zero score in AntMaze tasks; IQL-TD-MPC is the only model-based algorithm showing non-zero scores in AntMaze tasks, but still much worse than LEQ in AntMaze as well as HalfCheetah.

[W2-2, W3] The improvements in hyperparameters are likely to be environment specific and therefore not very generalizable. While standard practice in offline RL, the use of differing hyperparameters per environment (specifically τ, as provided in Table 10), also limits the significance and generality of the work.

We first would like to emphasize that LEQ with $\tau=0.1$ consistently shows the strong performance across all benchmarks, as shown in Table 15-17 in Appendix C.

Moreover, LEQ shares the same hyperparameters across all state-based environments, except a single hyperparameter $\tau$ for penalization. In contrast, prior works such as MOBILE and MOPO not only tune the penalization coefficient but also search for the optimal rollout length, with different hyperparameter ranges for different domains.

[W4-1, Q3] Could you explain the main difference between LEQ and IQL-TD-MPC and include this in the paper?

While both LEQ and IQL-TD-MPC are model-based offline RL methods that utilize expectile regression, these two approaches are fundamentally different. Specifically, IQL-TD-MPC does not penalize inaccurate model rollouts as LEQ. Instead, it mitigates model inaccuracies by removing random actions from MPPI and constraining the policy search to actions near the initial AWR-learned policy. This effectively limits the policy to predict actions similar to those present in the dataset, similar to IQL.

Accordingly, IQL-TD-MPC employs upper expectile regression to approximate the max operation in $V(s) = \max_{a} Q(s, a)$, for calculating the advantage for AWR, while LEQ utilizes lower expectile regression to get conservative return estimates from potentially inaccurate model rollouts.

We have clarified the differences between LEQ and IQL-TD-MPC in L121-125 of Section 2 and included their comparisons in Table 3 using HalfCheetah tasks, where official scores for IQL-TD-MPC are available. Notably, LEQ significantly outperforms IQL-TD-MPC in HalfCheetah tasks (240.0 vs 151.4).

[W4-2] The citation for IQL-TD-MPC is also wrong in Section 5.3

Thank you for pointing out our mistake. We corrected the citation for IQL-TD-MPC in L363 of Section 5.3 in the revised paper.

Thank you for your constructive feedback. We address your questions and concerns in detail below and our new manuscript reflects your suggestions. Please refer to the updated PDF for new results and revisions.

[W1] The article does not provide the theoretical analysis of how the use of low expectile regression impacts the final Q-values of the model data. Additionally, it does not offer theoretical guarantees for policy improvement , as seen in algorithms like COMBO and MOBILE.

Our extensive ablation studies in Table 5 demonstrate the impact of lower-expectile regression and the surrogate loss on policy improvement. Specifically, replacing lower-expectile Q-learning with MOBILE penalization leads to a substantial performance drop in AntMaze tasks ($461.8 \rightarrow 338.9$). Similarly, using $Q(s,a)$ ($461.8 \rightarrow 373.5$) or AWR ($461.8 \rightarrow 114.2$) in place of our surrogate loss for policy optimization results in notable performance degradation. These empirical findings support our claim that lower-expectile Q-learning effectively mitigates inaccuracies in model rollouts and that our surrogate loss better improves the policy over prior approaches.

We agree that theoretical analysis can definitely enhance our paper. However, we believe our comprehensive empirical results provide valuable insights into offline RL and can serve as a foundation for future research. We believe that these empirical findings can sufficiently support our claims, despite the lack of theoretical support.

[W2] The method presented in this paper directly scales the loss values of the overestimated model data by a factor of $\tau$, without considering the differences among the model data. For example, the scaling coefficient $\tau$ should vary depending on the degree of error in the model data.

To clarify, our method already penalizes states with high model errors more strongly than states with low model errors. This is because high uncertainty in model predictions results in increased variance in the target returns, leading to high variance in the target returns. For example, if the target returns follow the normal distribution, lower-expectile is equivalent to lower-confidence bound (LCB), which penalizes the mean estimate with its standard deviation.

We agree that developing a mechanism that further adapts the conservatism $\tau$ based on state-specific information is an exciting direction for future research. We do not have a good idea about how to implement it yet, and we note that similar limitations also exist in prior approaches, such as MOBILE, which applies a single LCB coefficient uniformly across all states, regardless of their individual characteristics.

[Q1] Does EMA regularization affect final performance?

Table 18 in Appendix D shows that replacing EMA regularization with a target network leads to a decrease in performance ($461.8 \rightarrow 379.1$), with an increase in Bellman error. Although we believe that further hyperparameter tuning can improve the performance of the one with a target network,we opted for EMA regularization since it works well without tuning.

[Q2] Why does the policy update use only the generated model data and not the precise offline data? Most model-based offline RL algorithms currently use accurate offline data to update policy.

While training a policy using both offline data and model-generated data (similar to the critic training in LEQ) could improve the policy, we opted for the simpler Dreamer-style policy update, which we found to be sufficient in practice.

[Q3] Error about the definition of $L_2^{\tau}$ in Eq. (1)

Thank you for pointing this out. We fixed the definition of $L_2^{\tau}$ in Eq. (1) and the explanation in L153-157 of Section 3 accordingly.

Thank you very much for the author's reply. The author conducted a lot of experiments to verify the effectiveness of the algorithm. I think this article is acceptable. However, due to the defects in the theoretical analysis, I still maintain my score.

Thank you for your constructive feedback. We address your questions and concerns in detail below and our new manuscript reflects your suggestions. Please refer to the updated PDF for new results and revisions.

[Q1] I do not find any reference of Figure 2 in the paper. Could you please insert the figure into the proper place to better explain your method?

Thank you for your suggestion. We added the references to Figure 2 in L120 and L213 in the revised version.

[Q2] How do you use the ensemble of world models for sampling? For example, do you randomly select one model for each step or for each rollout?

We select one model for each step, following the design choice of MOBILE [1]. We found that randomly choosing a model every step can make imaginary rollouts more robust to model biases, leading to better performance.

We clarified this in L739-L741 of Appendix A in the revised paper.

[Q3] Could you please explain why to include TT and TAP, and provide what we can learn from comparing LEQ against them?

Sequence modeling approaches (e.g. TT, TAP) serve as strong baselines on long-horizon tasks, such as AntMaze. We compare LEQ with these approaches and show that LEQ outperforms them in the most challenging task, AntMaze-Ultra. Moreover, while sequence modeling methods struggle on dense-reward tasks, LEQ maintains near-SOTA performance, demonstrating its broad applicability.

We clarified this in L367-368 of Section 5.3 and L413-414 of Section 5.4 in the revised paper.

[Q4] Error in the proof (last inequality)

We greatly appreciate that you pointed out our mistake in the proof in Appendix B. We could revise the proof that still ensures our original claim--$|Y - \hat{Y}\_\textrm{new}| \leq |Y-\hat{Y}|$, which means the proposed estimator in our surrogate loss is better than the learned estimator (Q). The revised theorem now have a looser bound $|Y - Y\_\textrm{new}| \leq (\frac{1-2\tau}{1-\tau}) * |Y-\hat{Y}|$ (previously $|Y - Y_\textrm{new}| \leq (\frac{1-2\tau}{1-\tau}) * p(Y \leq X \leq \hat{Y}) * |Y-\hat{Y}|$); but, this is still sufficient to support our original claim.

[Q5] What is the difference between MBPO and Dyna training schemes, as you classify in Table 8? Previous works usually regard MBPO as a Dyna-style algorithm.

Thank you for pointing out our mistakes again. We agree that describing our method’s training scheme as Dyna is not clear as MBPO is also a Dyna-style algorithm. We originally used Dyna to represent a Dyna-style training scheme but not MBPO. LEQ follows a Dreamer-style training scheme [2, 3], which uses the model for both expanding the dataset and calculating the target values ($\lambda$-returns), while MBPO only uses the model for dataset expansion.

To clearly distinguish LEQ from CBOP and MOBILE, we changed the “training scheme” of LEQ in Table 8 from Dyna to Dreamer.

References

[1] Sun et al. "Model-Bellman inconsistency for model-based offline reinforcement learning.”, ICML 2023.

[2] Hafner et al. "Mastering Atari with Discrete World Models.", ICLR 2021.

[3] Hafner et al. "Mastering diverse domains through world models.", Arxiv 2023.

I am very sorry for my late reply. Thanks for your answers which have cleared most of my confusion. However, after careful checking, I still have quesiton about the revised proof to Theorem 1. The deduction from Line 954 to Line 956 seems using $E[W^\tau(Y>X)] = \tau + (1-2\tau) p(Y \le X)$. This might still be wrong, as

In fact, here may be a counterexample to Theorem 1. Suppose $X$ follows the uniform distribution over $[0,1]$, then $Y = E^\tau[X] = \tau$. Let $\hat Y=0$ be an estimate of $Y$, resulting in $W^\tau(\hat Y > X)=\tau$ for any $X$. Now we can directly compute $\hat Y_{new}$ as

If $0 < \tau < 0.25$, we have $|\hat Y_{new} - Y| = |0.5 - \tau| > 0.25 > |0 - \tau| = |\hat Y - Y|$, which violates Theorem 1.

I think the motivation and the experiment results of LEQ are quite convincing. So if I misunderstood your proof or you could give a new proof, I will follow up as soon as possible.

Thank you for your constructive feedback. We address your questions and concerns in detail below and our new manuscript reflects your suggestions. Please refer to the updated PDF for new results and revisions.

[W1] Why did the authors choose to update the policy based on deterministic policy gradients? What considerations influenced this choice, and are there experimental differences between using a stochastic versus deterministic policy?

Thank you for pointing this out. We used a deterministic policy since lower-expectile regression of LEQ inadvertently penalizes the stochasticity of the policy while penalizing the uncertainty of model rollouts.

Empirically, LEQ with stochastic policies shows worse performance compared to LEQ with deterministic policies in AntMaze ($461.8 \rightarrow 380.4$). Although the performance of LEQ with stochastic policies could be further improved by increasing the stochasticity bonus (e.g. entropy), we suggest using a deterministic policy in LEQ for its simplicity.

We included the corresponding experimental results and their setup in Table 21, Appendix D.

[W2] How does BC pretraining impact the results, particularly in tasks like AntMaze?

This is a good point! Without BC pretraining, the performance of LEQ decreases, $461.8 \rightarrow 322.2$ in total. This is more apparent especially in medium mazes ($96.4 \rightarrow 4.6$) and ultra mazes ($81.6 \rightarrow 12.4$) due to early instability of training, matching the observation in CBOP [1].

In addition, we also conducted additional experiments on MOBILE$^*$ with BC training. However, the result shows that MOBILE$^*$ with BC pretraining degrades the performance ($285.3 \rightarrow 232.2$). We believe that careful initialization of the networks can mitigate this issue.

We added the corresponding experimental results in Table 19, Appendix D.

Reference

[1] Jeong et al. "Conservative Bayesian Model-Based Value Expansion for Offline Policy Optimization.", ICLR 2023.