Paper: DMBP: Diffusion model-based predictor for robust offline reinforcement learning against state observation perturbations (ICLR-2024)

Reviewer comment (Novelty):
The literature review section lacks a discussion of another taxonomy: training-time and testing-time robustness. While this paper and many others focus on offline RL for testing-time robustness, there is another group of offline RL works that investigate training-time robustness [1] [2] [3]. These works involve corrupting the offline dataset and evaluating it in a clean environment. It would be beneficial to include the training-time works for a more comprehensive literature review.

Author response:
1. Thanks for pointing this out. In the Introduction, we added the review on training-time robustness and testing-time robustness as follows:

Robust RL can be categorized into two taxonomies: training-time and testing-time
robustness. Training-time robust RL involves perturbations during the training process, while
evaluating the agent in a clean environment (Zhang et al., 2022; Ye et al., 2023). Conversely, testing-time robust RL focuses on training the agent with unperturbed datasets or environments and then
testing its performance in the presence of disturbances (Yang et al., 2022; Panaganti et al., 2022). Our
work primarily aims at enhancing the testing-time robustness of existing offline RL algorithms.

Reviewer comment (Reproducibility):
It is not clear whether the states in this paper are normalized. Normalizing the observations would ensure fair observation corruption and may potentially impact the recovery ability of the diffusion model.

Author response:
2. We do not normalize the observations in our experiments.
In most experimental settings for testing the robustness of RL algorithms against uncertain observations, perturbed observations are selected within an $\ell_{\infty}$ ball, or the noise is added randomly with a specific scale. This approach ensures that the noise scale for each dimension of the state is the same.
However, if we choose to normalize the states perturbed by the aforementioned noises, the normalized states would have different noise scales for different dimensions, since the distribution of each state dimension is distinct.
Therefore, normalization of the perturbed states can complicate the denoising process for DMBP.
We have added a paragraph in Appendix B.2 to clarify this point.

Reviewer comment (Experiments):
Currently, it seems that the diffusion model is trained using complete observations even in the experiments with incomplete observations. I am curious to know if the diffusion model can handle training with incomplete observations as well.

Author response:
3. To our knowledge, DMBP does not possess the capability to handle training with incomplete observations. The primary objective of DMBP is to accurately recover the actual states for the agent to make decisions. However, it lacks the ability to reason about the unknown dimensions of states in offline datasets. Consequently, DMBP may require the baseline agent to be trained with unobserved dimensions in such scenarios.
Unfortunately, due to the time constraints, we were unable to validate this approach within the revision period. We firmly believe that this avenue holds promising value for further exploration and research.

Reviewer comment (Evaluation):
The authors are suggested to report the training time and the inference time of the proposed method.

Author response:
4. We have included a new section (Appendix D.5), which provides the computing costs associated with DMBP (including training and testing costs).

Reviewer comment (Writing):
Typo: page 2 'A diagram of the proposed approach is shown in the fight subplot of Figure 1' --> 'A diagram of the proposed approach is shown in the right subplot of Figure 1'.

Author response:
5. We apologize for the typos we made. We have carefully checked entire paper to correct all typos.

Reviewer comment (Theory):
The paper has the assumption that perturbation on state space follows a Gaussian distribution. However, in practical settings, such perturbations might be biased and skewed. For example, a water drop on a camera might lead to distortion of captured images. Therefore, it might be better if the authors could elaborate on the assumption of the perturbations and how they are produced.

Author response:
We acknowledge that the assumption of Gaussian distributed noise has certain limitations. It is commonly used as a standard setting for benchmark in many related research works [1, 2, 3]. Moreover, our experiments have incorporated different types of non-Gaussian distributed noises, including uniformly distributed random noise, maximum action-difference attack, and minimum Q-value attack. Under these adversarial noise distributions that have been commonly adopted in the literature, the proposed DMBP, along with the baseline algorithms, can achieve satisfactory denoising performance
(cf. Tables 2, 7, 8, 9, and Fig. 9).

Reviewer comment (Reproducibility):
Introducing a diffusion model for noise reduction could increase the complexity of the model, which may lead to longer training time and resource-intensive computations. The authors might want to include more details regarding training time and computational resources.

Author response:
In the revised manuscript, we have added a new section (Appendix D.5) on the computing costs associated with DMBP (including training and testing costs).

Reviewer comment (Experiments):
Scaling to very large state spaces or handling very noisy environments might pose challenges, especially when using diffusion models. It would be better for the authors to consider testing on tasks with larger state spaces such as the humanoid.

Author response:
We agree with the reviewer's point that larger state spaces may pose challenges for DMBP in accurately recovering actual states from noisy ones. Further investigation is required to confirm this. However, it is noted that most offline RL algorithms struggle to learn satisfactory policies when the state dimension is extremely large, such as in the humanoid environment with a total state dimension of 376. 
As a benchmark that is commonly used in most
offline DRL works, D4RL does not provide datasets for the humanoid environment. To our knowledge, the environment with the highest state space in the D4RL benchmark is Franka Kitchen, with a state dimension of 59. The experimental results in Table 9 and Fig. 14 demonstrate that the proposed DMBP performs well in this high-dimensional environment.

Reviewer comment (Theory):
Eq.(2) is wrong as the recovered state should not be a probability.

Author response:
We apologize for the mistake we made. In the revised manuscript, it has been fixed as follows:

The denoised state $\hat{{s}}\_{t}$ is sampled from the reverse denoising process, which can be expressed as a Markov chain:
    $\hat{{s}}\_{t} \sim p\_{\theta}(\tilde{{s}}^{0:k}\_{t} \mid {a}\_{t-1}, \tau^{\hat{{s}}}\_{t-1}) = f\_k(\tilde{{s}}\_{t}) \prod\limits\_{i=1}^{k} p\_{\theta}(\tilde{{s}}^{i-1}\_{t} \mid \tilde{{s}}^{i}\_{t}, {a}\_{t-1}, \tau^{\hat{{s}}}\_{t-1}),$

Reviewer comment (Experiments):
One main concern is that observation perturbations can be potentially addressed by traditional methods, especially when we assume observation perturbations does not affect reward and transition functions. The reviewer is curious of how offline RL methods with Kalman Filter perform compared to DMBP?

Author response:
We do agree with your point that traditional methods (e.g., the Kalman Filter for discrete control) have been proposed to address the issue of observation perturbations for discrete scenarios. However, the Kalman filter requires exact knowledge of the governing equation of the controlled system and the related parameters. 
In offline RL settings, on the other hand, we have no prior knowledge on the system's governing equations or parameters (otherwise, the agent can be trained online with a simulator). The only information available is the offline dataset. Modeling the environment for high action/state space problems with limited data can be prohibitively challenging. As discussed in the Introduction, modeling environment (called model-based) methods are mostly used for data augmentation in offline RL. To our knowledge, there are no relevant references on offline RL methods using the Kalman Filter for continuous control problems.

Reviewer comment (Evaluation):
Minor issue: all results in tables should be shown with standard deviations.

Author response:
Thanks for pointing this out! We have fixed the issue and supplemented the data in all tables with standard deviations.

Reviewer comment (Experiments):
In the experiments, results show DMBP significantly improve the robustness performance of existing offline RL methods. The review noticed that most offline RL methods used are value-based. How will DMBP perform when used with weighted imitation learning methods, e.g., IQL?

Author response:
We agree that imitation learning is a class of highly effective and well-studied offline RL algorithms. Adding an imitation learning algorithm to the baseline algorithms would be valuable to demonstrate the generality of DMBP. However, given the space constraints and limited time, we struggled to do it for now. It should be noted that a class of value-based offline RL algorithms, which utilize policy regularization methods, achieve a similar effect as imitation learning, i.e., selectively cloning the policy in datasets. The baseline algorithms adopted in our experiments, like BCQ, TD3+BC, and Diffusion QL, belong to this category.  Experimental results demonstrate that DMBP performs well for baseline algorithms in this category.
In the revised manuscript, we have added a new section (Appendix D.6) to discuss the observed experimental results, where we analyze the performance of DMBP in comparison to policy regularization category and conservative Q estimation category offline RL algorithms.
Below we present partial experimental results on the suggested algorithm IQL (Implicit Q Learning), which 
is trained on expert dataset in Mujoco domain.

 Noise scale | Noise type | HalfCheetah|                 |Hopper      |                    |Walker2d |                  |
| ------------  |------------  | ------------   | ------------| ------------| ------------| ------------| ------------ |
|                    |                   | Base           | DMBP       | Base        | DMBP       | Base        | DMBP        |
|      0      |        -     | 84.7$\pm$14.6  |             | 103.4$\pm$21.5 |           | 109.3$\pm$0.9  |             |
|      0.05      |        G-rand     | 15.2$\pm$9.9  | 53.0$\pm$17.6 | 28.5$\pm$27.4 | 58.7$\pm$15.1 | 105.5$\pm$8.4 | 105.6$\pm$11.6 |
|           |  U-rand     | 21.1$\pm$13.0 | 61.6$\pm$18.6 | 43.3$\pm$34.3 | 69.2$\pm$28.4 | 108.5$\pm$1.3 | 109.4$\pm$1.0  |
|               | MAD | 14.9$\pm$11.0 | 37.5$\pm$14.2 | 19.8$\pm$16.8 | 38.3$\pm$14.4 | 104.3$\pm$10.1 | 101.4$\pm$20.4 |
|           | MinQ | 25.5$\pm$15.8 | 46.8$\pm$14.2 | 26.0$\pm$24.6 | 31.6$\pm$18.8 | 106.5$\pm$1.5 | 103.4$\pm$11.4 |
|  0.10  | G-rand | 9.3$\pm$6.1 | 20.6$\pm$11.8 | 11.5$\pm$6.6 | 24.4$\pm$11.1 | 58.5$\pm$23.7 | 83.9$\pm$21.8 |
|          |  U-rand |  12.8$\pm$9.4 |  29.8$\pm$17.0 |  26.6$\pm$22.5 |  32.6$\pm$19.1 |  79.9$\pm$17.4 |  103.5$\pm$6.5 | 
|         |  MAD |  7.7$\pm$5.3 |  17.8$\pm$9.8 |  10.7$\pm$8.2 |  17.7$\pm$8.5 |  42.2$\pm$30.4 |  72.5$\pm$19.6 | 
|         |  MinQ |  10.0$\pm$7.1 |  26.6$\pm$14.3 |  8.5$\pm$8.6 |  15.6$\pm$9.5 |  77.5$\pm$35.7 |  98.6$\pm$8.8 |

Reviewer comment (Experiments):
Is DMBP sensitive to hyperparameters (i.e., a, b, c) of the variance schedule? Any guidance to choose these hyperparameters?

Author response:
There is no need to individually select the hyperparameters related to the variance schedule (i.e., $a, b, c$), and one can use the same set of predetermined values (as listed in Table 3) for all environments and datasets.
We redesigned our variance schedule with the specific aim of effectively handling information with small to medium scale noises. In other words, we intentionally set $\sqrt{\bar{\alpha}\_i}$  to be large and $\sqrt{1 - \bar{\alpha}\_{i}}$ to be small for smaller diffusion timesteps $i$, which ensures that ${s}\_{t}^i=\sqrt{\bar{\alpha}\_i} {s}\_{t} + \sqrt{1-\bar{\alpha}\_i} {\epsilon}$ maintains significant raw state information.

Reviewer comment (Theory):
It is intutitve that diffusion denoising can perform well for random noise. Can the authors provide insights why it also works well for adversarial attacks?

Author response:
In classical diffusion models, one typically assumes that the noise follows a Gaussian distribution and uses batch-generated Gaussian distributed noise for training. However, during the testing phase, the noised states are generated individually. In other words, for a single instance of $\tilde{{s}}_t = {s}_t + {\epsilon}_t$, we cannot determine the exact distribution of ${\epsilon}_t$ based on a single sample.
For DMBP, during the single-step denoising process in testing, when the noise scale is given, it is trained to reconstruct ${s}_t$ from the observed $\tilde{{s}}_t$ for all possible values of the noise ${\epsilon}_t$ within a given range, which means that DMBP aims to learn a mapping that can effectively denoise the observed state within a certain range, regardless of the specific distribution of the noise. This may explain why DMBP also works well for adversarial attacks.

Reviewer comment (Experiments):
When training different algorithms, do you use the same dataset size and batch size for all algorithms?

Author response:
When we train different baseline algorithms (including BCQ, CQL, TD3+BC, Diffusion QL, and RORL), we use the same dataset size and batch size.  
By utilizing the datasets in the D4RL benchmark, we ensure that the dataset size remains fixed, thereby ensuring fairness in the comparisons made.  
Furthermore, we adhere to the suggested hyperparameter settings for all baseline algorithms, including setting the batch size to 256.

Reviewer comment (Experiments):
In offline RL, especially the theoretical community, we are often concerned with the data coverage problem (full coverage vs partial coverage). I am interested in the performance of DMBP in this scenario. In particular, D4RL has millions of transitions in its dataset. Although one can argue that this is not that much for a continuous control task, we sometimes see more extreme cases in the tabular setting where the data is extremely limited, e.g., barely supporting the optimal policy's state-action visitation. I conjecture that if DMBP is paired with an algorithm like VI-LCB [1], which is proven to be statistically efficient, it can also facilitate offline RL with scarce data settings. Do you have any insight on this?  

[1] Li, G., Shi, L., Chen, Y., Chi, Y., and Wei, Y. (2022a). Settling the sample complexity of model-based offline reinforcement learning. arXiv preprint arXiv:2204.05275

Author response:
We agree that the data convergence issue of a RL algorithm significantly affects its stability and performance. In the D4RL benchmark, datasets  labelled with "medium-replay" in the Mujoco domain have only 1/5 of the data compared to other datasets. In addition, a significant portion of the state-action visitations are from suboptimal strategies, and many are from random policy. In our experiments, it is found that DMBP still performs well under such circumstances (cf. Tables 1 and 2). The baseline algorithm VI-LCB recommended by the reviewer has strong theoretical guarantees. However, due to the time constraints, we are currently unable to present the simulation results of DMBP assisting VI-LCB in a scarce data setting with perturbed state observations. We will validate this in our future work.