Paper: Uncertainty-aware Distributional Offline Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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

Reviewer comment (Novelty):
In eq.(7) the policy $\pi_\phi$ 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$.

Author response:
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.

Reviewer comment (Theory):
Please clarify how diffusion model learns from the behavior policy $\pi_b$ instead of cloning it.

Author response:
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.

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

Author response:
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.

Reviewer comment (Experiments):
Does the policy learn a fixed risk-averse distortion? If so, can the network be modified to accept different risk level in test time?

Author response:
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.

|              |           |                     |                    |                     |       Expert        |                    |                     |
|              |           | $CVaR_{0.05}$ | $CVaR_{0.1}$  | $CVaR_{0.15}$  | $CVaR_{0.05}$  | $CVaR_{0.1}$  | $CVaR_{0.15}$  |
| Half-Cheetah |   UDAC    |       200(14)       |      276(29)       |       303(35)       |       602(55)       |      732(104)      |      782(201)       |
|              |   ORAAC   |       188(30)       |      214(36)       |       277(43)       |      588(104)       |      595(191)      |       599(99)       |
|              |   CODAC   |       -33(12)       |      -41(17)       |       -40(10)       |      620(104)       |      687(152)      |      699(83)       |
|  Walker-2D   |   UDAC    |      1242(100)      |      1368(89)      |      1343(165)      |      1257(94)       |      1398(67)      |      1402(178)      |
|              |   ORAAC   |      610(130)       |      663(124)      |       732(99)       |      1086(59)       |      1172(71)      |      1153(91)       |
|              |   CODAC   |      1002(89)       |      1159(57)      |      1172(100)      |      1190(73)       |      1298(98)      |      1320(64)       |
|    Hopper    |   UDAC    |      1443(41)       |      1502(19)      |      1742(44)       |      1267(50)       |      1302(28)      |      1399(34)       |
|              |   ORAAC   |      1388(34)       |      1416(28)       |      1423(44)       |       853(40)       |      980(28)       |       992(30)       |
|              |   CODAC   |       955(28)       |      976(30)       |       988(25)       |       903(44)       |      990(19)       |      1001(32)       |

Reviewer comment (Writing):
The paper is not clearly written, for example, see first 4 lines of section 4.1.

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

Reviewer comment (Novelty):
Novelty is limited since risk-averse distributional offline RL have 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.

Author response:
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.

Reviewer comment (Theory):
The authors claim that using their framework accurately models the behavior policy but does not provide any theoretical justification for the claim.

Author response:
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

Reviewer comment (Experiments):
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.

Author response:
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.

|              |           |                     |                    |                     |                     |                    |                     |
|:------------:|:---------:|:-------------------:|:------------------:|:-------------------:|:-------------------:|:------------------:|:-------------------:|
|              | Algorithm |       Medium        |                    |                     |       Expert        |                    |                     |
|              |           | $CVaR_0.05$ | $CVaR_0.1$ |$CVaR_0.15$ | $CVaR_0.05$ |$CVaR_0.1$| $CVaR_0.15$ |
| Half-Cheetah |   UDAC    |       200(14)       |      276(29)       |       303(35)       |       602(55)       |      732(104)      |      782(201)       |
|              |   ORAAC   |       188(30)       |      214(36)       |       277(43)       |      588(104)       |      595(191)      |       599(99)       |
|              |   CODAC   |       -33(12)       |      -41(17)       |       -40(10)       |      620(104)       |      687(152)      |       699(83)       |
|  Walker-2D   |   UDAC    |      1242(100)      |      1368(89)      |      1343(165)      |      1257(94)       |      1398(67)      |      1402(178)      |
|              |   ORAAC   |      610(130)       |      663(124)      |       732(99)       |      1086(59)       |      1172(71)      |      1153(91)       |
|              |   CODAC   |      1002(89)       |      1159(57)      |      1172(100)      |      1190(73)       |      1298(98)      |      1320(64)       |
|    Hopper    |   UDAC    |      1443(41)       |      1502(19)      |      1742(44)       |      1267(50)       |      1302(28)      |      1399(34)       |
|              |   ORAAC   |      1388(34)       |      1416(28)      |      1423(44)       |       853(40)       |      980(28)       |       992(30)       |
|              |   CODAC   |       955(28)       |      976(30)       |       988(25)       |       903(44)       |      990(19)       |      1001(32)       |

Reviewer comment (Writing):
The first line of 4.1 is grammatically incorrect.

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

Reviewer comment (Theory):
How does this framework guarantee robustness against epistemic uncertainty without explicitly optimizing any robust measure of expected returns?

Author response:
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.

Reviewer comment (Evaluation):
Why choose confidence level $\alpha=0.1$ while evaluating the performance of the framework with CVaR metric?

Author response:
Thank you for your advice about the aleatoric uncertainty and epistemic uncertainty. We would like to sincerely apologize for the confusion between the aleatoric uncertainty and epistemic uncertainty. In the proposed method, we are targeting epistemic uncertainty from the risk-sensitive perspective (i.e., risk aversion). The proposed method utilizes the perturbation model to ask the network to learn a risk-sensitive policy that can reduce the epistemic uncertainty.