Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning

We are deeply grateful for the reviewer's detailed and accurate summary. We also appreciate the time and effort the reviewer has devoted. As for the weakness, we prepared the following responses, which are presented as follows.

... However, improvements in the presentation would greatly enhance the clarity and comprehensibility of the manuscript.

We apologize for any unclear expressions or improper organization in the article. We'll adjust our presentation in the updated version.

In particular, several equations within the paper omit the definitions of variables, which can lead to confusion—for example, $a_0$ in Line 119. In addition, some variables overlap while they are independent. For instance, in Line 155, the variable of the integration overlaps with the $a$ in the numerator.

We apologize for omitting some definitions of variables and will include them in the updated version. In Line 119, $a_0$ represents the diffused action where the footprint stands for diffusion timestep. We'll also address the overlap of independent variables in the revised version.

We appreciate the reviewer's time and effort dedicated to evaluating our paper, as well as the constructive feedback provided. In response to the concerns and questions raised, we have prepared detailed answers, which are outlined separately below.

Missing comparison in Table 1 of more optimal Gaussian policy methods, e.g. EDAC [1]

We list the results of EDAC and Diffusion-DICE in the following table. For EDAC's results, we copy the results on MuJoCo locomotion tasks from its original paper and copy the results on AntMaze navigation tasks from CORL[1]. It's evident from the results that on MuJoCo locomotion tasks, Diffusion-DICE is comparable to EDAC. While on AntMaze navigation tasks, EDAC totally fails[3]. We suggest that this is because such ensemble-based uncertainty estimation heavily depends on the dataset distribution and on AntMaze datasets, it's not reliable.

Selection of D4RL datasets is limited, e.g. what about expert/random datasets? Further environments like Adroit would also be interesting

For the experiments of MuJoCo locomotion tasks, we only choose "medium", "medium-replay", "medium-expert" datasets because "random" dataset hardly exists in real-world tasks, and "expert" data is typically used for imitation learning settings. Furthermore, these datasets are also rarely used in other diffusion-based offline RL methods. To demonstrate Diffusion-DICE's superiority against other methods, we evaluate Diffusion-DICE on 2 tasks from kitchen and 2 tasks from adroit, following the same experimental setting in Appendix D. Note that we only choose 4 tasks in total due to the limited rebuttal preiod.

The results and the chosen hyperparameters are listed in the tables above. We compare Diffusion-DICE with other Offline RL baselines (either diffusion-based or not). The results are either taken from their original papers (if available) or from LD[2] (if not). The results on these more complex environments consistently show Diffusion-DICE's superiority.

Discussion of hyperparameter choice in the Appendix is important and should be included in the experimental section.

We apologize that due to the page limit, the discussion of hyperparameter choice is placed in the appendix. In the updated version, we'll include this discussion in the experimental section.

A comparison of the inference speed of Diffusion-DICE and prior baselines would be valuable.

In the following table, we compare the inference time (seconds/100 actions) of Diffusion-DICE and other baselines. It's worth noting that because we mainly focus on diffusion-based methods, these baselines also contain only diffusion-based algorithms. The results are based on a single RTX 4090 GPU, under antmaze-large-diverse-v2 environment.

Line 70: “M” -> “M=”

We'll fix this typo in the updated version.

Line 73: LP abbreviation not explained

We're sorry that due to the page limit, the explanation of LP abbreviation is omitted. In fact, we refer to LP as the expected return's linear programming form (linear with $d^\pi$). By the definition of $d^\pi(s, a)$ in Line 74, $d^\pi(s, a)$ represents the discounted sum of probability that the agent takes action $a$ on state $s$ over all steps $t$. Then it's obvious that $E_{(s, a) \sim d^\pi} [r(s, a)]$ equals to the discounted sum of rewards., i.e. $E[\sum_{t=0}^\infty \gamma^t \cdot r(s_t, a_t) ]$. Due to the bijection between $\pi$ and $d^\pi$, maximizing expected return over $\pi$ is equivalant to maximizing $E_{(s, a) \sim d^\pi} [r(s, a)]$ with respect to $d^\pi$. The latter exactly possesses a linear programming form.

There is concurrent related work [2] which also performs an analogous transformation between the behavior distribution to an online policy with diffusion models for synthetic data generation.

Thanks for pointing out. We'll add this to the discussion in the updated version.

[1]: CORL: Research-oriented Deep Offline Reinforcement Learning Library. Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, Sergey Kolesnikov. NeurIPS, 2024.

[2]: Efficient Planning with Latent Diffusion. Wenhao Li. ICLR, 2024.

[3]: Revisiting the Minimalist Approach to Offline Reinforcement Learning. Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, Sergey Kolesnikov. NeurIPS, 2023.

We appreciate the reviewer's time and effort dedicated to evaluating our paper, as well as the constructive feedback provided. In response to the concerns and questions raised, we have prepared detailed answers, which are outlined separately below.

... The way authors presented the guide-then-select is confusing. I’d suggest authors provide a bit more background and carefully define the “guidance term”, and how it relates to the RL before using it in both abstract and introduction.

We apologize for the lack of explanation of 'guidance term' in the abstract and introduction due to the page limit. The 'guidance term' refers to the log-expectation term defined in Eq. (6), which 'guides' the diffused action towards high-value regions. In the updated version, we will provide more background information on this term and its relation to RL in the introduction.

Achieving in-sample learning of offline diffusion RL is not new. Efficient Diffusion Policy (EDP) has introduced an IQL-based variant which naturally allows training Q-values using only in-sample data, without querying out-of-distribution actions during policy evaluation. I’d suggest the authors carefully check the claims and avoid over claiming.

We acknowledge that EDP is also a diffusion-based algorithm that avoids querying OOD actions' value, and we will add it to the discussion in the updated version. However, the major difference between EDP and Diffusion-DICE is that EDP has no guarantee of the form of optimal policy. According to EDP's original paper, it discusses two types of approaches: "direct policy optimization" and "likelihood-based policy optimization". It's obvious that the former can not guarantee the form of optimal policy. The latter replaces $\log \pi_\theta(a|s)$ with its lower bound in the optimization objective, which consequently loses guarantee for the form of optimal policy. On the other hand, Diffusion-DICE directly calculates the score function of the optimal policy induced from DICE's objective. As both terms in Eq. (6) can be estimated unbiasedly, the policy distribution induced by Diffusion-DICE can match the exact optimal policy distribution.

The D4RL experiments are only conducted on the locomotion tasks and the antmaze tasks. The commonly tested kitchen and adroit tasks are missing, which weakens the claim of the paper.

To validate that Diffusion-DICE also demonstrates superior performance on other more complex tasks, we evaluate Diffusion-DICE in kitchen and adroit environments. Due to limited rebuttal period, we choose 2 tasks from kitchen and 2 from adroit. We compare Diffusion-DICE with other Offline RL baselines (either diffusion-based or not). The results are either copied from their original papers (if available) or LD[1] (if not). The results and the chosen hyperparameters are as follows:

Note that we follow the same experimental settings in Appendix D. The results further substantiate our claim that Diffusion-DICE achieves optimal policy transformation while keeping minimal error exploited, and thus exhibits SOTA performance even on more complex tasks.

Could you provide a bit more discussion about the differences of using guidance for inference with the Diffusion-QL style inference, which directly guides the sampling process towards actions with high returns? It would be interesting to understand the pros and cons of these two different paradigms

The major difference comes from the way score function is modeled. Diffusion-QL represents algorithms that directly model optimal policy's score function with one neural network. Diffusion-DICE represents algorithms that indirectly model it as a "transformed" score function of the behavior policy, with possibly more than one neural network.

For Diffusion-QL, as the guidance towards high-value actions has already been encoded in the score network, simply running reverse diffusion process with the learned score network could bring high-value actions, which makes it easier to implement and faster to inference. However, as the marginal probability $\log \pi_\theta(a|s)$ for diffusion model is hard to compute, the policy improvement of such algorithms almost relies on surrogate objectives (See Eq. (3) in Diffusion-QL[2] and Eq. (10), Eq. (12) in EDP[1]). Consequently, the exact distribution of diffused action after inference is unknown, and the form of optimal policy is not guaranteed.

On the other hand, using guidance for inference allows for the decoupled and exact learning of both the behavior policy's score function and the guidance term. This provides guarantee for the action distribution after inference. Moreover, due to the decoupling between behavior policy's score function and guidance term, it's possible to combine different guidance flexibly during inference, without training the desired score function from scratch. This is especially useful for aligning large diffusion models in the future. However, given a limited amount of data, using guidance for inference may introduce auxiliary models, which increase the computational burdens.

[1]: Efficient Planning with Latent Diffusion. Wenhao Li. ICLR, 2024.

[2]: Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning. Zhendong Wang, Jonathan J Hunt, Mingyuan Zhou. ICLR, 2023.