Uncertainty-aware Distributional Offline Reinforcement Learning

Does the policy learn a fixed risk-averse distortion? If so, can the network be modifed to accept different risk level in test time?

We appreciate the reviewer's concern. In the current version, the proposed method needs to be optimized based on certain risk distortions (such as $CVaR_\alpha$, Wang etc). Our model can be modified to accept different $\alpha$ in CVaR, given that the $\alpha$ is a hyperparameter of the proposed method. Specifically, we can change the $\alpha$ of CVaR in Equation 6 (if we are using CVaR as the distortion). For your reference, the $CVaR_\alpha$ is defined as,

$$CVaR_\alpha (F_Z^{-1}(s,a;\tau)) = \frac{1}{\alpha} \int^{\alpha}_0 F_Z^{-1}(s,a;\tau)d \tau$$

where $a$ is the output of the deterministic policy $\pi$ with parameter $\phi$. Moreover, it would also be possible to adjust the risk level by modifying the testing environments. For example, we can adjust the stochastic reward setting in risk-sensitive D4RL to make it more risky.However, considering the alterations made to the test environments, it's possible that the effectiveness of the proposed method might be impacted. We view this as a potential avenue for promising future research. For your reference, we have also provided some experimental results about different $\alpha$ of CVaR.

We sincerely appreciate the reviewers' thorough evaluation of our paper. We acknowledge the concerns that have been raised regarding clarity, technical details, and experimental aspects. We have conducted a thorough revision to address these concerns. We have provided a detailed point-by-point response to your comments.

The paper uses a diffusion model, for the latter performs better than VAE, under the condition of a dataset with different modalities(in Section 4.2). This part seems to lack theoretical analysis for ”denoising a noisy mixed policy. How does the diffusion model decide which part of the policy is noisy? Or it's just the diffusion model has a better representational ability? Some theoretical analysis or a numerical one would be better.

We appreciate the reviewer's concerns. Our model operates under the assumption that the pre-collected trajectories form a mixture of different behavior policies, some of which may be sub-optimal. Regrettably, from a denoising perspective, to our knowledge in the RL domain, we are not aware of a close-up of existing works that try to control the segment that needs to be denoised. We believe this a valuable point, we will explore the potential relevant works in other domains such as CV to investigate the feasibility. Given that the diffusion model is trained together with the RL agent, we think the diffusion process can identify noisy segments based on the rewards obtained during ongoing training. We believe that diffusion offers superior representational capabilities compared to methods like VAE or cVAE within our target domain and our experiments also substantiate this claim. For instance, in a baseline comparison, ORAAC, utilizing VAE for behavior policy representation, displays lower performance across all tasks compared to our proposed method. Additionally, recent studies (Pearce et al., 2023; Wang et al., 2023) corroborate the assertion that diffusion possesses stronger representational abilities than VAEs. Regarding the numerical experiments, we are uncertain about the specific expectations of the reviewers. It would greatly assist us if the reviewers could provide more detailed specifications. Thanks again for your thoughtful review, we are committed to advancing the point that control the diffusion model to denoise certain segments in the RL domain.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam De- vlin. Imitating human behaviour with diffusion models. ICLR, 2023

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. ICLR, 2023

In eq.(7) the policy does not take the risk distortion parameter as an input, does this imply that the policy's risk level can't be tuned? If it can be tuned, studying this would be more innovative than studying $\lambda$

We appreciate the reviewer's concern. The purpose of eq.(7) is used to show how the actor's policy $\pi_\phi(s)$ is composed. The risk distortion parameter is an input and will be used to train the $\pi_\phi(s)$ via eq.(6). We hope this justification can address your concern.

Minor issues: In the last paragraph of Section 1, DR4L should be D4RL.

We appreciate the reviewer's comment. We have revised the typos and checked the whole paper to ensure all the typos are revised.

Please clarify how diffusion model learns from the behavior policy instead of cloning it.

We appreciate the reviewer's concern. As the diffusion process is kind of a denoising process, it indeed will perform like behavior cloning in the context of RL. However, it deviates from traditional behavior cloning in its approach. Instead of replicating the demonstrations directly, the diffusion method identifies segments of the data as noise and excludes them. This denoising process results in a learned representation that differs from the initial pre-collected data, having refined segments. This process leans more toward imitation rather than strict behavior cloning.

What are $\alpha_i, \overline{\alpha}_i$in eq.(9) (10) (11)?

We appreciate the reviewer's comment. The variables $\alpha_i$ and $\overline{\alpha}_i$are directly linked to the forward process variances $\beta_i$. Specifically, $\alpha_i = 1-\beta_i$, $\overline{\alpha}_i$ is defined as the sum of $\alpha_i$ from 0 to $t$ . We chose to omit this explanation as it mirrors the conventional diffusion model. Our intention was to conserve space by refraining from detailing how the MDP integrates into the diffusion process.

We sincerely appreciate the reviewers' thorough evaluation of our paper. We acknowledge the concerns that have been raised regarding clarity, technical details, and experimental aspects. We have conducted a thorough revision to address these concerns. We have provided a detailed point-by-point response to your comments.

the paper is not clearly written, for example, see first 4 lines of section 4.1.

We genuinely appreciate the reviewer's concerns. In order to improve the clarity, we have revised the first paragraph of section 4.1.

Novelty is limited since risk-averse distributional offline RL has been explored by other work like [1]. This work only extends prior work by using diffusion model to model the behavior policy instead of VAEs.

We appreciate the reviewer’s concern. We would like to make a justification for the novelty. As the reviewer pointed out, our work is an extension of the existing work [1] (i.e., ORAAC). The biggest limitation of the ORAAC would be the expressive capability of the behavior policy model. Moreover, replacing the VAE with the diffusion model is not an easy task. The diffusion model we used here is not just a standard DDPM. We define the overall diffusion as a conditional model with states as the condition and actions as the outputs. Different from the ORAAC which trains the VAE separately from the RL agent, UDAC trains the diffusion model and the RL agent simultaneously. By using such a way, UDAC is able to understand and adapt to the environment dynamic (i.e., aleatoric uncertainty) efficiently. Different from vanilla DDPM, the diffusion model used in UDAC only requires a few steps (less than 3) to learn a superior behavior policy model.

The authors claim that using their framework accurately models the behavior policy but does not provide any theoretical justification for the claim.

We appreciate reviewer’s concern regarding the theoretical justification for the accurate behavior policy modeling. In this study, we mainly focus on validating the quality of the learned behavior representation through comprehensive experiments. We can find that the proposed method outperforms than existing baselines in both risk-sensitive environments and risky environments. To the best of our knowledge, it is challenging to theoretically prove that the proposed behavior policy model is more accurate than the existing works, since the recent works (Pearce et al., 2023; Wang et al., 2023) about the theoretical justification of diffusion in RL for representation learning are still unexplored. In this work, we have provided some theoretical analysis of the UDAC from the convergence perspective and policy improvement perspective (please kindly refer to, Appendix A, Proposition 1 and Proposition 2). Given our empirical results along with theoretical analysis, we can conclude that the proposed behavior policy modeling method can learn the behavior policy more accurately or expressively.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. ICLR, 2023

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. ICLR, 2023

The proposed algorithm and baselines are compared at a single CVaR confidence level $\alpha=0.1$. It would be good to evaluate them on different confidence levels.

We appreciate the reviewer’s comment. Given that all the existing baselines are evaluated under the single CvaR confidence level 0.1, we choose to use 0.1 solely. To address your concern, we have run the different levels of CvaR on the proposed model. We choose the alpha from the range 0.05, 0.1, and 0.15 which represent the CvaR at the 5%, 10% and 15% quantile levels. Given the timely constraint of the rebuttal phase, we are not able to provide a comprehensive comparison of all of the baselines under different quantile levels. As a result, we have reported the most two relevant baselines for your reference (i.e., ORAAC and CODAC). Results are attached to the second part of the response.

While reporting the performance of the proposed algorithm and baselines, the authors do not provide the errors in the mean and cvar value of returns.

We appreciate the reviewer’s comment. We would appreciate it if the reviewer could give more details about this concern. As we reported in Table 1 and Table 2, all the reported values are in the form `mean(error)’. Moreover, we have added the errors in Table 4 in the revised version. We hope this can address your concern.

Additional Comments: The first line of 4.1 is grammatically incorrect.

We appreciate the reviewer’s comment. In order to improve the clarity, we have revised the first paragraph of section 4.1

noise policies (section 4.2) $\rightarrow$ noisy policies

We appreciate the reviewer’s comment. We have revised this typo.

We are pleased to know that our response has successfully addressed the majority of your concerns. We agree that the provided theoretical analysis may not specifically related to the behavior modeling. Given that the behavior modeling is kind of a representation learning, we decide to use the experiments to valid the feasibility. Considering the breadth of comprehensive experiments we have conducted, which demonstrate the superior performance of our proposed method over baseline models, we believe that our work can offer some valuable insights to the community.

Regarding your additional concern, please refers to the following response.

How does this framework guarantee robustness against epistemic uncertainty without explicitly optimizing any robust measure of expected returns?

We appreciate your questions. We acknowledge the epistemic uncertainty arising from limited data[1]. When constructing a model with this dataset for future state or reward prediction, there is a risk of bias if we solely rely on immediate returns as our optimization objective. In this work, we seek distributional reinforcement learning, which optimizes the distribution of states and rewards rather than individual outcomes. This approach enhances the model's ability to learn about unseen states, effectively mitigating the impact of epistemic uncertainty, and thereby ensuring a certain degree of robustness.

[1] Clements, W. R., Van Delft, B., Robaglia, B. M., Slaoui, R. B., & Toth, S. (2019). Estimating risk and uncertainty in deep reinforcement learning. arXiv preprint arXiv:1905.09638.

Thank you for the clarification and for responding to my concerns. The theory in the appendix is not specific to the behavior policy modelling and it just shows that the distributional bellman operator is a contraction mapping and proves that the distributional RL algorithm will improve the policy over time. I am not very convinced that simple replacing one generative model (VAE) with another (Diffusion model) is a novel enough contribution for this conference.

Can you please also answer the following question:

a) How does this framework guarantee robustness against epistemic uncertainty without explicitly optimizing any robust measure of expected returns?

figure 2 needs correction, y axis is not visible, also there are no error bands - how do we know what is the stability across reruns?

We are greatly appreciating the reviewer's comments. In response, we have made revisions to Figure 2, introducing error bars to illustrate the stability. We believe that these updates in the new figure adequately address your concerns.

Is the diffusion generative model absolutely necessary in order to achieve this performance? Would you switch it out for a CVAE in order to compare?

We genuinely appreciate the reviewer's concerns. Regarding the suggestion to replace the CVAE with the proposed diffusion generative model, in Section 4.2 on Page 4, we have provided a detailed discussion explaining why we opted for the diffusion method over the cVAE model. Moreover, we have also mentioned that the limitation of the cVAE in the introduction section paragraph 2 on Page 1. The key idea behind this is that the behavior policies on the offline dataset are normally collected from different agents and may have different modalities. It implies that some behavior policies may be sub-optimal. And recent works like Pearce et al., 2023 and Wang et al., 2023 also demonstrate that cVAE may not have a strong expressive ability than diffusion in offline RL. In order to further support our claim, we have provided the experiments here about replacing the diffusion model with cVAE. We hope this can address your concern.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. ICLR, 2023

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. ICLR, 2023

We sincerely appreciate the reviewers' thorough evaluation of our paper. We acknowledge the concerns that have been raised regarding clarity, citations, and experimental aspects. We have conducted a thorough revision to address these concerns. We have provided a detailed point-by-point response to your comments.

Elimination of Manual Behavior Policy Specification. ii) Accurate Modeling of the behavior policy is the same statement? You use diffusion generative models to model the behavior policy, that's it?}

We thank for reviewer’s comment. We would like to provide a further clarification about these two statements. The first one refers to the previous work (Ma et al. 2021) which requires us to define the behavior policy manually. To achieve this, we can use cVAE or diffusion model (as what we do in this work) to avoid the manual definition. The second statement refers to the previous works (Urpi et al., 2021 and Lyu et al., 2022) which employ the VAE/cVAE to represent the behavior policy. The reason why we do not omit the first statement is because the performance of the Ma et al. 2021 is quite good in some tasks (we would kindly ask the reviewer to refer the table 1). Hence, we decided to keep this statement as one of the main streams in the current progress of the distributional RL. Moreover, replacing the VAE with the diffusion model is not an easy task. The diffusion model we used here is not just a standard DDPM. We define the overall diffusion as a conditional model with states as the condition and actions as the outputs. Different from the ORAAC which trains the VAE separately from the RL agent, UDAC trains the diffusion model and the RL agent simultaneously. By using such a way, UDAC is able to understand and adapt to the environment dynamic (i.e., aleatoric uncertainty) efficiently. We observe significant performance improvement through our simultaneous learning framework.

Nuria Armengol Urpi, Sebastian Curi, and Andreas Krause. Risk-averse offline reinforcement learning. ICLR, 2021

Jiafei Lyu, Xiaoteng Ma, Xiu Li, and Zongqing Lu. Mildly conservative q-learning for offline reinforcement learning. NeurIPS, 2022.

Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning. NeurIPS, 2021.

p2.par2 - ”We conducted extensive experiments in three different environments covering both risk-sensitive” - it’s 3 different benchmarks, you use multiple environments of D4RL.

We appreciate the reviewer's comment. We apologize for the potential confusion of benchmarks and environments. In order to address this, we have revised the sentence to replace such ambiguous wordings to improve readability. We use three different benchmarks here: we conduct the experiments on risk-neutral D4RL (i.e., the traditional D4RL), risk-sensitive D4RL (i.e., the modified version of D4RL which includes the risk in the environment) and the risky robot navigation. We strictly/exactly follow the evaluation protocol in previous works (Ma et al., 2021, Urpi et al. 2021), and evaluate the algorithms on the same benchmarks.

Q: p3. par2 - "For simplify" - to simplify

We appreciate the reviewer's comment. We have made the revision accordingly.

In Section 4 and earlier you talk about "risk-distortion operators", but you do not define the concept in the background section, yet it seems essential for your method.

We appreciate the reviewer's comment. We have added a footnote to provide the background knowledge of the risk-distortion operators. It is a class of distortion operators that also consider the risk. And we have added the relevant citation (Wang 2000) to improve the clarity.

Shaun S Wang. A class of distortion operators for pricing financial and insurance risks. Journal of risk and insurance, pp. 15–36, 2000

In Section 4.1 you talk about implicit quantile functions, yet this concept has not been defined (what is implicit about them?) and no reference is given, I suppose you need [1].

We appreciate the reviewer's comment. We have added the relevant citation [1] to the revised version,

There is no citation for the quantile Huber loss, you should provide a citation for it and refer to the equations the first time it is mentioned (ideally the definition would preceed the first usage of the function).

We appreciate the reviewer's comment. We have added the citation of huber loss [Huber 1992]. Peter J Huber. Robust estimation of a location parameter. In Breakthroughs in statistics: Methodology and distribution, pp. 492–518. Springer, 1992.

p3, last par - you say that you use both uniform sampling and a quantile proposal network, then you say that you don't use a quantile proposal network and leave it for future work? If you don't use it, don't mention it here or mention it as a comment (an alternative is to use a quantile proposal network such as in (xyz citation)

We value the reviewer's feedback. The intent behind the paragraph was to highlight two distinct methods for controlling the quantile level. We've revised this section to delineate the second approach as a comment, along with referencing relevant citations (Yang et al. 2019), clarifying that it stands as a potential alternative for future exploration.

Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie-Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. NeurIPS 2019.

In eq. 5, I suppose that you have N fixed quantile levels , hence the indexing. However, this was not mentioned previously, yet it is part of the methodology (also along the lines of earlier work?). Moreover, the loss is a function o and there is no expectation defined over s, a, s', a', though the quantile td error depends on the former. Also, "J" is commonly used as return? Also normalizing factor should be $N^2$?

We appreciate the reviewer's consideration. Indeed, N represents the fixed quantile levels, although we regret not explicitly mentioning this earlier in the text, save for its inclusion in Algorithm 1. To address this, we've inserted a brief explanation regarding N following Equation 5. Additionally, we've revised Equation 5 to disambiguate the presence of two Ns, now denoted as N and K separately, indicating potential variations in quantile levels. Furthermore, to alleviate confusion, we've included the previously defined expectations over s, a, s', a' in Equation 5. In response to the concern about 'J,' we have removed it from the text to prevent any potential confusion. We truly appreciate the reviewer raising the concern about the normalizing factor, we have made the adjustment accordingly based on the new N and K notations. We hope this revision can address your concerns.

eq. 6 - is it "distortion risk function" or "risk-distortion function"? I also don't understand this loss function, you have expectations wrt. marginal state distribution, however you have an action appearing in the implicit quantile function, where does the action come from? You are missing a policy here, please define this correctly.

We thank the reviewer for the concerns. Addressing the query about the distortion risk function, we acknowledge and apologize for the confusion—it should indeed be termed the risk-distortion function, and we've rectified this in the paper. As for the policy clarification, we regret the oversight in representing a deterministic policy (i.e., $\pi_{\phi}(s) = a)$. To alleviate this confusion, we've updated the notation, replacing 'a' with $\pi_{\phi}(s)$. This revision aims to accurately reflect the involvement of the policy within the implicit quantile function. We hope this can address your concerns.

I think that your description of diffusion generative modelling needs significant adjustments. What you call an model is the denoiser, the Gaussian reverse process transition kernel arises from a continuous treatment of diffusion and exactly because of the choice of noise schedule. Since you don't provide contributions on level of diffusion generative modelling, this shouldn't take so much space of the paper.

We sincerely value the reviewer's input. Section 4.2 primarily delineates why we opted for diffusion over cVAE and delves into the actor's components. The section discussing the diffusion model in the second part of Section 4.2 spans about half a page and includes four important formulas (8-11). These equations play a crucial role in facilitating the comprehension and replication of how RL is integrated into the diffusion process in our paper. Furthermore, the diffusion model employed in this study is not the standard DDPM; we have adapted the DDPM to create a conditional model, using states as conditions and actions as outputs. Consequently, our model can efficiently generate the behavior model in just a few steps (in our experiments, the number of steps is less than 3).

p6 - CQL is not a distributional RL algorithm?

We are thankful for the reviewer's comment. CQL, despite not being a distributional algorithm, falls within the realm of offline RL algorithms. We included CQL as a baseline due to its widespread use in the field of offline RL as well as the recent advance in distributional RL algorithm, acknowledging its common usage as a benchmark in this context.

suggestion for table 1 - use colors instead of bold and *

We truly value the concerns raised by the reviewer. In response, we've made adjustments by changing the formatting, switching bold text to red and asterisks (*) to blue, aiming to enhance the overall readability.

. Moreover, ORAAC, CODAC and CQL are distributional RL algorithms where ORAAC and CODAC belong to offline RL.

Thank you for providing the result table for CVAE, this answers my question partially. Is the likelihood of the data in the case of the CVAE lower than in the case of the diffusion model? You mean to say that the CVAE can't fit the distribution of the data?

We appreciate the reviewer's concern. We agree that it would be possible that the likelihood of the data in this case of the cVAE is lower than the diffusion model. Based on our existing experiments, which reflect that the cVAE may not be able to perform as well as diffusion. However, we believe that more extensive experiments and theoretical analysis should be carried out to further investigate this problem. We do appreciate the reviewer's point and will target this as future work.

We can not certainly say that the cVAE can't fit the distribution of the data, the possible justification here could be that: the cVAE may not have a strong expressive as what the diffusion model has.

Dear reviewer LtLP,

Thank you again for your valuable comments on our paper. We hope that our recent response can address your additional concerns. We are mindful that you may be currently facing a demanding schedule and truly appreciate the time and effort you invested in reviewing our response. As the ICLR discussion deadline approaches (9 hours remaining), we kindly request your feedback on our submitted response. Your comments are crucial to us, and we understand that this may be the final opportunity for us to address any additional concerns you might have.

Thank you once again for your time and engagement.

Best regards, The Authors