Value function estimation using conditional diffusion models for control

We thank the reviewer for their thorough feedback on our paper, and their interesting questions. We provide our responses below:

Policy conditioning

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. Policy conditioning has been extensively studied by works treating RL problems as reward-conditioned sequence modelling [3,4], and we heavily borrowed architectural design choices from this literature, notably the use of transformers (specifically, Perceiver IO due to its ability to accept arbitrary input-output modalities) as opposed to MLPs, which primarily used in model-free RL.

Comparison to other diffusion-based baselines

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.

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 hyper-parameters for all runs.

Missing references

We thank the reviewer for their thoroughness, and added the suggested missing references to Section 5 of the revised paper.

References:

[1] Kakade, Sham, and John Langford. "Approximately optimal approximate reinforcement learning." Proceedings of the Nineteenth International Conference on Machine Learning. 2002.

[2] Brandfonbrener, David, et al. "Offline rl without off-policy evaluation." Advances in neural information processing systems 34 (2021): 4933-4946.

[3] Chen, Lili, et al. "Decision transformer: Reinforcement learning via sequence modeling." Advances in neural information processing systems 34 (2021): 15084-15097.

[4] Janner, Michael, Qiyang Li, and Sergey Levine. "Offline reinforcement learning as one big sequence modeling problem." Advances in neural information processing systems34 (2021): 1273-1286.

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

[7] 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).

We thank the reviewer for their thorough review of our work, and provide an itemized response below:

Method clarifications and averaging step

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.

Typos and correctness of equations

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.

Computational costs of DVF and architecture

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.

Comparison to other diffusion-based approaches

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.

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.

We thank the reviewer for their questions, as well as suggestions of relevant baselines. We provide a detailed response below:

Knowledge of behavior policy $\mu$

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.

Additional diffusion baselines

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.

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).

We thank the reviewer for their constructive feedback. We provide a point-by-point response below.

General presentation of the paper

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.

Missing $\Delta t$

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

Wording of the conditioning operator

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.

Equation 8 dependence on $\theta$

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

Computational cost of maximizing $Q(s,a,\phi(\pi))$ directly

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.

Baselines

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

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 [4] in our setting, specifically using a single set of hyperparameters for all runs, which the authors do not do in their code.

I would like to thank the authors for the detailed response - I think the additional baselines and the improvements in presentations are good, I have decided to increase to score to 8.