Paper: Value function estimation using conditional diffusion models for control (ICLR-2024)

Reviewer comment (Theory):
The methods section is extremely difficult to follow due to many typos, notation mistakes, and a mathematical error. Specifically, equation 12 is wrong — there needs to be a term with the dynamics as well.

Author response:
We appreciate the reviewer’s comments on catching bad phrasing, e.g. “performance of deep reinforcement learning scales with model capacity”, where it should read “performance of offline reinforcement and imitation learning”. Recent works in offline RL [1] and imitation learning [2,3] do show that offline settings (essentially supervised learning as pointed out by the reviewer) do exhibit such trends. Next, Equation 12 indeed has a typo, where the gradient from the dynamics is ignored. We added this term back, per reviewer’s suggestion, and we want to point out that the dependence of dynamics in $a_t$ is implicitly encoded via policy conditioning in $\rho$. Our claim still stands, in that the gradient of $Q$ depends on the gradient of the reward function, which is less computationally cumbersome to compute than if we directly took the gradient of $\epsilon_\theta$ with respect to the action.

Reviewer comment (Experiments):
The authors fail to compare against any diffusion-based baselines, which could mislead readers into believing that the proposed DVF method is a state-of-the-art method for offline RL. For example, Table 1 in the Diffusion-QL paper shows that DVF (non-pooled, which is the fair comparison) performs worse than Diffusion-QL in every task.

Author response:
We agree that our experimental section could benefit from additional baselines that are also based on generative models such as [6,7]. We present the results below, as well as in Table 1 of the revised manuscript.

| **Task**                | **BC** | **CQL** | **IQL** | **BRAC-p** | **BEAR** | **Diffusion BC** | **Diffusion QL** | **DVF**       | **DVF (pooled)** |
|:-----------------------:|:------:|:-------:|:-------:|:----------:|:--------:|:----------------:|:----------------:|:-------------:|:----------------:|
| **kitchen-complete-v0** | 65.0   | 43.8    | 62.5    | 0.0        | 0.0      | 53.8 $\pm$ 10    | 67.2 $\pm$ 8     | 74.8 $\pm$ 13 | 79.1 $\pm$ 12    |
| **kitchen-partial-v0**  | 38.0   | 49.8    | 46.3    | 0.0        | 0.0      | 42.5 $\pm$ 3     | 48.1 $\pm$ 4     | 50.2 $\pm$ 7  | 72.2 $\pm$ 9     |
| **kitchen-mixed-v0**    | 51.5   | 51.0    | 51.0    | 0.0        | 0.0      | 48.3 $\pm$ 2     | 53.5 $\pm$ 6     | 51.1 $\pm$ 7  | 67.4 $\pm$ 6     |


The new baselines based on diffusion roughly match the performance of DVF, while being trained on x3 more compute. Note that we had to re-run the author’s code for [7] in our setting, specifically using a single set of hyperparameters for all runs. 
The main differences of DVF with existing baselines are in the distribution that they model: Diffuser [8] generates action plans using diffusion, Diffusion BC [6] imitates data contained in the training set, and Diffusion QL [9] replaces the SAC [9] policy with a diffusion model.

Reviewer comment (Theory):
Equation 12 is wrong because it omits taking the gradient of the expected value with respect to the action (the dynamics uses \( a_t \)).

Author response:
We appreciate the reviewer’s comments on catching bad phrasing, e.g. “performance of deep reinforcement learning scales with model capacity”, where it should read “performance of offline reinforcement and imitation learning”. Recent works in offline RL [1] and imitation learning [2,3] do show that offline settings (essentially supervised learning as pointed out by the reviewer) do exhibit such trends. Next, Equation 12 indeed has a typo, where the gradient from the dynamics is ignored. We added this term back, per reviewer’s suggestion, and we want to point out that the dependence of dynamics in $a_t$ is implicitly encoded via policy conditioning in $\rho$. Our claim still stands, in that the gradient of $Q$ depends on the gradient of the reward function, which is less computationally cumbersome to compute than if we directly took the gradient of $\epsilon_\theta$ with respect to the action.

Reviewer comment (Writing):
In the abstract, the statement "A fairly reliable trend in deep reinforcement learning is that performance scales with the number of parameters, provided a complimentary scaling in amount of training data. As the appetite for large models increases, it is imperative to address, sooner than later, the potential problem of running out of high-quality demonstrations." is not entirely correct and suggests possible confusion — perhaps the authors meant a reliable trend in "supervised learning"? Furthermore, it is unclear why demonstrations are mentioned when the first sentence discusses reinforcement learning.

Author response:
We appreciate the reviewer’s comments on catching bad phrasing, e.g. “performance of deep reinforcement learning scales with model capacity”, where it should read “performance of offline reinforcement and imitation learning”. Recent works in offline RL [1] and imitation learning [2,3] do show that offline settings (essentially supervised learning as pointed out by the reviewer) do exhibit such trends.

Reviewer comment (Writing):
In Figure 5, the phrase "Returns are normalized by average per-task data performance" is unclear — its precise meaning should be explained.

Author response:
We agree with the reviewer that some additional details can be added to the paper. Specifically, the averaging of $\Delta t$ in the value function estimation is mentioned in Algorithm 1 line 7, which we didn’t mention separately in the main text due to lack of space. The networks architecture and choice is discussed in Appendix Section 7.1, to which we added the architecture of the reward and policy networks as well (4-layer MLPs with mish activation functions, layernorm and residual connections). Figure 5 shows the performance of each method rescaled to the 0-1 interval for each run using the minimum and maximum statistics of the dataset’s logging policy (available on the dataset’s website https://github.com/takuseno/d4rl-pybullet). The mean and standard deviation are then computed on top of these normalized scores.

Reviewer comment (Novelty):
The paper should provide concrete discussion of DVF differences and tradeoffs relative to other diffusion-based offline RL methods.

Author response:
We agree that our experimental section could benefit from additional baselines that are also based on generative models such as [6,7]. We present the results below, as well as in Table 1 of the revised manuscript.

| **Task**                | **BC** | **CQL** | **IQL** | **BRAC-p** | **BEAR** | **Diffusion BC** | **Diffusion QL** | **DVF**       | **DVF (pooled)** |
|:-----------------------:|:------:|:-------:|:-------:|:----------:|:--------:|:----------------:|:----------------:|:-------------:|:----------------:|
| **kitchen-complete-v0** | 65.0   | 43.8    | 62.5    | 0.0        | 0.0      | 53.8 $\pm$ 10    | 67.2 $\pm$ 8     | 74.8 $\pm$ 13 | 79.1 $\pm$ 12    |
| **kitchen-partial-v0**  | 38.0   | 49.8    | 46.3    | 0.0        | 0.0      | 42.5 $\pm$ 3     | 48.1 $\pm$ 4     | 50.2 $\pm$ 7  | 72.2 $\pm$ 9     |
| **kitchen-mixed-v0**    | 51.5   | 51.0    | 51.0    | 0.0        | 0.0      | 48.3 $\pm$ 2     | 53.5 $\pm$ 6     | 51.1 $\pm$ 7  | 67.4 $\pm$ 6     |


The new baselines based on diffusion roughly match the performance of DVF, while being trained on x3 more compute. Note that we had to re-run the author’s code for [7] in our setting, specifically using a single set of hyperparameters for all runs. 
The main differences of DVF with existing baselines are in the distribution that they model: Diffuser [8] generates action plans using diffusion, Diffusion BC [6] imitates data contained in the training set, and Diffusion QL [9] replaces the SAC [9] policy with a diffusion model.

Reviewer comment (Novelty):
A rationale should be given for why Perceiver I/O was used instead of a standard diffusion U-Net architecture.

Author response:
The computational costs of DVF primarily depend on the choice of the noise network’s architecture. We investigated the trade-offs between a classical U-net architecture [4], and a cross-attentional model using Perceiver IO [5]. While the U-net is the go-to choice for diffusion in image space, Perceiver worked best in our setting, where policy conditioning requires operating over (potentially) long context windows composed of vectors that vary according to different scales (which is not the case for pixels). We would like to specify that Perceiver IO and U-net are architectural design choices, that are not critical to the overall contribution. The computational cost also depends on the context length which, in our experiments, were fairly short, due to the determinism of the MuJoCo environments. Specifically, we only needed a context of length 4 for all experiments. This makes DVF computationally heavier than expected TD-based approaches, but on-par with diffusion based approaches. One advantage DVF has over other methods is that by decoupling the diffusion model training from the policy training, we can pool together datasets with missing actions or rewards, e.g. train the diffusion model on video data (and using only states as context), and then combine it with a small amount of action-reward labeled data to obtain a more confident statistical estimator of the value function.

Reviewer comment (Reproducibility):
The methods section should clearly state which networks are being trained, their objectives, and their inputs and outputs, as this is currently difficult to parse.

Author response:
We agree with the reviewer that some additional details can be added to the paper. Specifically, the averaging of $\Delta t$ in the value function estimation is mentioned in Algorithm 1 line 7, which we didn’t mention separately in the main text due to lack of space. The networks architecture and choice is discussed in Appendix Section 7.1, to which we added the architecture of the reward and policy networks as well (4-layer MLPs with mish activation functions, layernorm and residual connections). Figure 5 shows the performance of each method rescaled to the 0-1 interval for each run using the minimum and maximum statistics of the dataset’s logging policy (available on the dataset’s website https://github.com/takuseno/d4rl-pybullet). The mean and standard deviation are then computed on top of these normalized scores.

Reviewer comment (Theory):
The assumption that the behavior policy $\mu$ is known is not usual in offline RL. It seems like this paper only utilized the dataset $\mathcal{D}$. So can this assumption be removed without affecting the result?

Author response:
You are correct, in that, generally speaking, the behavior policy that generated the dataset is not known in offline RL. The experiments in Figures 3 and 4, which use the discrete policy representation (i.e. $\phi(\pi_i)=i$ is saved during dataset generation), are provided only for purposes of illustrating the diffusion model’s capabilities to capture long-horizon dependencies and disambiguate data coming from different sources. In practice, we can treat the history of all available data from $\mu$ (or $\mu_i$ if data was collected from a mixture of policies) up to timestep $t$ as the continuous policy representation (or context) of state $s_t$, mimicking the belief state in POMDPs. Such an approach embeds the policy using its trajectories, which is expensive but fairly universal as to forego additional assumptions about the data.

Reviewer comment (Experiments):
There have been many works that apply diffusion models on offline RL [1,2,3,etc.]. Could you please include more baselines that use diffusion models for more convincing experiments? Such works are also worth discussing in related works or other parts of the paper.  

[1] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning. https://arxiv.org/abs/2208.06193  

[2] Is Conditional Generative Modeling all you need for Decision-Making? https://arxiv.org/abs/2211.15657  

[3] IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies. https://arxiv.org/abs/2304.10573

Author response:
We agree that our experimental section could benefit from additional baselines that are also based on generative models such as [1,2]. We present the results below, as well as in Table 1 of the revised manuscript.

| **Task**                | **BC** | **CQL** | **IQL** | **BRAC-p** | **BEAR** | **Diffusion BC** | **Diffusion QL** | **DVF**       | **DVF (pooled)** |
|:-----------------------:|:------:|:-------:|:-------:|:----------:|:--------:|:----------------:|:----------------:|:-------------:|:----------------:|
| **kitchen-complete-v0** | 65.0   | 43.8    | 62.5    | 0.0        | 0.0      | 53.8 $\pm$ 10    | 67.2 $\pm$ 8     | 74.8 $\pm$ 13 | 79.1 $\pm$ 12    |
| **kitchen-partial-v0**  | 38.0   | 49.8    | 46.3    | 0.0        | 0.0      | 42.5 $\pm$ 3     | 48.1 $\pm$ 4     | 50.2 $\pm$ 7  | 72.2 $\pm$ 9     |
| **kitchen-mixed-v0**    | 51.5   | 51.0    | 51.0    | 0.0        | 0.0      | 48.3 $\pm$ 2     | 53.5 $\pm$ 6     | 51.1 $\pm$ 7  | 67.4 $\pm$ 6     |


The new baselines based on diffusion roughly match the performance of DVF, while being trained on x3 more compute. Note that we had to re-run the author’s code for [2] in our setting, specifically using a single set of hyperparameters for all runs, which the authors do not do in their code.

## References:

[1] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." arXiv preprint arXiv:2303.04137 (2023).

[2] Wang, Zhendong, Jonathan J. Hunt, and Mingyuan Zhou. "Diffusion policies as an expressive policy class for offline reinforcement learning." arXiv preprint arXiv:2208.06193(2022).

Reviewer comment (Writing):
The presentation of the work can be improved as the manuscript is a bit hard to understand in its current form. It would be good to dissect and analyze each sentence with a bit more care when rewriting. The authors' work could be presented with much more clarity if written more clearly.

Author response:
We agree with the reviewer that some parts of the paper warrant clarifications. We made edits to the rebuttal PDF (marked in red), to make the manuscript easier to follow, and correct typos pointed out by the reviewers.

Reviewer comment (Writing):
In Section 2, $\Delta t$ suddenly appears without being defined, and the readers are left to figure out what it is.

Author response:
Thanks for catching this, we added the definition back in the footnote of Section 2 page 2.

Reviewer comment (Writing):
The wording of explicit conditioning in Section 2 is a bit strange and requires some domain expertise to understand what the authors mean. The occupation measure is always conditioned on a policy; the choice of whether to implicitly or explicitly do it seems like a choice of implementation. Perhaps it would be better to say something like "rather than statistically estimating the occupation measure through Monte Carlo sampling, we choose to directly learn a map that can infer the occupation measure given some features of the policy."

Author response:
We modified the text of Section 3 page 3 to reflect your suggestion, with which we agree. Since our work primarily tackles the challenges arising in offline settings, it is impossible to collect on-policy samples during policy optimization. We’ve experimented with two potential solutions to this issue. The first strategy is to simply condition the occupancy measure (i.e. noise network in diffusion) on data available in the training set. This approach maximizes the value of the behavior policy, and yields a single step of policy improvement which, in most realistic MDPs, is enough to find a policy with higher returns [1,2]. The second strategy is re-generate the policy embedding $\phi(\pi)$ for the learned policy, which can be done with continuous embeddings by using $\rho$ and $\pi$ to predict the future states in an autoregressive manner (see end of Section 3.4). Since the context for deterministic environments like MuJoCo is rather small, this doesn’t pose a significant overhead.

Reviewer comment (Writing):
In Equation 8, $l_{diffusion}$ should explicitly be noted as a function of $\theta$.

Author response:
Thank you for the observation, we fixed it in the rebuttal PDF Equation 8.

Reviewer comment (Writing):
In Section 3, the authors state that maximizing $Q(s,a,\phi(\pi))$ directly is expensive, but the readers lack the context to understand this at that point in the manuscript. The parameters being maximized and the representation of a diffusion model have not yet been described, and $Q(s,a,\phi(\pi))$ has never been defined anywhere.

Author response:
The reason for the computational overhead of the direct maximization is due to the fact that the soft policy update (e.g. using SAC’s actor loss[3]) requires to compute $\nabla_{a}Q(s,a,\phi(\pi))$, which in turns needs $\nabla_{a}\epsilon_\theta(s,a,\phi(\pi))$. When $\epsilon_\theta$ is a large network, especially using cross-attention, this computation can be quite memory intensive and FLOP heavy. Instead, we can use a trick similar to [3], where we represent $Q^\pi(s_t,a_t) = r(s_t,a_t)+\gamma E_{s_{t+1}}[V^\pi(s_{t+1})]$, making the gradient only depend on the reward network, which is much smaller. Note that this problem doesn’t arise when the policy is implicitly encoded in the Q-value, e.g. with discrete action spaces. We acknowledge that this point was understated in the paper (Section 3.2), and added a clarifying discussion at the beginning of Section 3.

Reviewer comment (Experiments):
The computational aspect of the approach has not been discussed. It is unclear whether DVF is cheaper or more expensive to train compared to other baselines.

Author response:
The computational costs of DVF primarily depend on the choice of the noise network’s architecture. We investigated the trade-offs between a classical U-net architecture [5], and a cross-attentional model using Perceiver IO [6]. While the U-net is the go-to choice for diffusion in image space, Perceiver worked best in our setting, where policy conditioning requires operating over (potentially) long context windows composed of vectors that vary according to different scales (which is not the case for pixels). Perceiver IO also gives a way to scale complexity using the number of latents in addition to the number of layers. The computational cost also depends on the context length which, in our experiments, were fairly short, due to the determinism of the MuJoCo environments. Specifically, we only needed a context of length 4 for all experiments. This makes DVF computationally heavier than expected TD-based approaches, but on-par with diffusion based approaches. One advantage DVF has over other methods is that by decoupling the diffusion model training from the policy training, we can pool together datasets (see last column of Table 1) with missing actions or rewards, e.g. train the diffusion model on video data (and using only states as context), and then combine it with a small amount of action-reward labeled data to obtain a more confident statistical estimator of the value function.

Reviewer comment (Novelty):
It is unclear whether there are other interesting use cases of having a generative model for the occupation measure besides estimating the value function.

Author response:
Thanks for the question - there are indeed a few additional applications, with the most prominent one potentially being using the generative model to estimate state coverage and uncertainty. We took a first step towards this idea in Section 4.4, where we analyzed how well various offline RL methods fare under reward-shaping for exploration. In practice, we ran all methods in Table 2 using some form of reward shaping, and features learned by DVF’s diffusion model were useful for exploration from offline data (i.e. no updates on online samples).