Paper: Reasoning with Latent Diffusion in Offline Reinforcement Learning (ICLR-2024)

Reviewer comment (Evaluation):
The proposed algorithm is interesting, but might have difficulty with the execution and easy-ness of the approach. Primarily because the Q-function is now learnt over Q(s,z) latents, and it would be helpful to do some qualitative evaluation of the Q-values over the latents. This is because the learnt Q would now heavily depend on the quality of the latents, which is coming from a beta-vae, and prior works on representation learning has showed that VAEs are not necessarily good for learning good quality latents.

Author response:
Thanks for the insightful question. The $\beta$-VAE latents are trained to distill important information from the trajectory sequence. The usefulness of the latents is evidenced by the fact that the policy conditioned on the latents achieves significantly better performance than baselines that use regularized latents, such as OPAL [1].

Q learning is only performed with the high-level latents, and the low-level policy is kept fixed after initial VAE training. If the latent representation does not contain information that is useful, then the low-level policy decoder would collapse to behavior cloning of the dataset, and so Q learning with the high-level policy will not affect it at all and our performance would be similar to BC baseline. ***We have a added TD error plot during Q learning in Appendix K Figure 12 for AntMaze, Franka, Adroit, and Locomotion tasks.*** This plot shows that the td error steadily decreases, indicating stable convergence of the Q-learning.

Reviewer comment (Evaluation):
This draws to my other worry of how good the latents are when learnt from Beta-VAE, and if there is a qualitative analysis about the structure that Beta-VAE latents recover from these tasks. Recent works have shown that with representation learning objectives, the latent structure can be well recovered - this work only shows how the latents are distributed across the horizonz, but it is not clear if these latents can itself contain useful information specific to the task?

Author response:
Many of your questions seem to center around whether using a VAE leads to good representations. Vanilla VAEs are not necessarily the best forms of representation learning, many previous skill-learning papers use these for simplicity as well. Some of these papers were referenced by reviewer Agic (OPAL[1], SPiRL[2], PLAS[3]).

We use $\beta$-VAE with a low value of $\beta$, which is a common choice for training latent diffusion models for other modalities, such as high-resolution images. The low value of $\beta$ ensures that the VAE behaves like an auto-encoder that encodes compressed latents with high mutual information with the uncompressed high-dimensional inputs. As a result, we learn more information-rich latents than these prior works by lowering the KL regularization and sampling from the resulting non-Gaussian latent space using the LDM prior. As such, we do not use the VAE itself as a generative model but just as a way to learn the skill-dependent low-level policy and get an information-rich latent space to be sampled by the LDM. The ability for this LDM to capture the diverse multimodal skill support in the behavior dataset is what allows us to use BCQ, as discussed in sections 5.1-5.3.

Reviewer comment (Experiments):
Maybe not so related, but I wonder if an entropy-regularized objective can also be used in these settings to induce diverse skills learnt from the diffusion prior? The paper uses BCQ in offline, which is drawn from TD3 - so I agree that the algorithm is fundamentally using a deterministic policy approach; but for the discrete action tasks - maybe entropy regularized objectives to learn diverse skills may be introduced?

Author response:
This could be an exciting future research direction that is beyond the scope of our current work. It is unclear to us what form of entropy regularized objective would be directly applicable to the current approach, but it could be interesting to force diversity in learnt skills. Would this make more sense in online RL where skill discovery is important than offline RL, where we just require copying skills from the behavior support?

Reviewer comment (Experiments):
Paper claims that the algorithm can work for both continuous and discrete action spaces - but there aren’t enough results showing for discrete actions? Why is that so?

Author response:
We have not run experiments on discrete tasks because D4RL benchmark does not contain any such environment. We expect latent diffusion will work on discrete tasks as demonstrated in prior work [4].

Reviewer comment (Evaluation):
How can we evaluate the quality of the latents? This is of big importance for this work, since the Q function is now learnt over the latents; I would like to see quantitivate results showing the effectiveness of the Q-function, maybe TD errors too, and qualitative results showing how good are the latents. Figure 2 only shows that latents are diverse and more skills are introduced across the horizons; but I worry about the use of Beta-VAE in the algorithm, and think better representation objectives should instead be used in the algorithm pipeline. Results justifying this would strengthen the paper.

Author response:
Thanks for the insightful question. The $\beta$-VAE latents are trained to distill important information from the trajectory sequence. The usefulness of the latents is evidenced by the fact that the policy conditioned on the latents achieves significantly better performance than baselines that use regularized latents, such as OPAL [1].

Q learning is only performed with the high-level latents, and the low-level policy is kept fixed after initial VAE training. If the latent representation does not contain information that is useful, then the low-level policy decoder would collapse to behavior cloning of the dataset, and so Q learning with the high-level policy will not affect it at all and our performance would be similar to BC baseline. ***We have a added TD error plot during Q learning in Appendix K Figure 12 for AntMaze, Franka, Adroit, and Locomotion tasks.*** This plot shows that the td error steadily decreases, indicating stable convergence of the Q-learning.

Many of your questions seem to center around whether using a VAE leads to good representations. Vanilla VAEs are not necessarily the best forms of representation learning, many previous skill-learning papers use these for simplicity as well. Some of these papers were referenced by reviewer Agic (OPAL[1], SPiRL[2], PLAS[3]).

We use $\beta$-VAE with a low value of $\beta$, which is a common choice for training latent diffusion models for other modalities, such as high-resolution images. The low value of $\beta$ ensures that the VAE behaves like an auto-encoder that encodes compressed latents with high mutual information with the uncompressed high-dimensional inputs. As a result, we learn more information-rich latents than these prior works by lowering the KL regularization and sampling from the resulting non-Gaussian latent space using the LDM prior. As such, we do not use the VAE itself as a generative model but just as a way to learn the skill-dependent low-level policy and get an information-rich latent space to be sampled by the LDM. The ability for this LDM to capture the diverse multimodal skill support in the behavior dataset is what allows us to use BCQ, as discussed in sections 5.1-5.3.

Reviewer comment (Evaluation):
The phrases “improves credit assignment”, “faster reward propagation” describing the proposed work should be avoided, or backed by empirical evidence. I don’t see how either of these quantities can be measured. I understand how the proposed method will intuitively achieve both of these -- however, they are just intuitions, not proven.

Author response:
>  ___The phrases “improves credit assignment”, “faster reward propagation” describing the proposed work should be avoided, or backed by empirical evidence.___

We agree these are more intuitive claims rather than something which we quantitatively measured. Qualitatively, improved performance in long horizon tasks indicates better credit assignment. We can change the wording in the abstract to reflect this for the final camera ready version, since abstract changes are not permitted during the review phase.

Reviewer comment (Evaluation):
The paper describes the decoder or policy as a “powerful decoder” that can handle “high-frequency details”. This is true for LDMs used for computer vision dealing with high resolution images. However, I don’t see how this translates well into RL. Training the VAE encoder-decoder model in a separate first stage seems like it will produce a rigid, fixed policy (i.e. decoder). The paper itself acknowledges that the “limited decoder capacity” is probably what is leading to the saturation of performance in Figure 4. How can the decoder be “powerful” and “limited” at the same time? Is there a way to empirically measure the strength of a decoder?

Author response:
> ___The paper describes the decoder or policy as a “powerful decoder” that can handle “high-frequency details”___

We intended to refer to our decoder as “powerful” in comparison to other raw trajectory diffusion methods. Due to the separation of the low-level policy from the high-level latent diffusion prior, we can have more powerful decoder networks, which would be infeasible to train with end-to-end diffusion methods. We use an autoregressive decoder, which outputs the action dimensions one at a time sequentially. It is possible for us to use even larger decoders, but we chose not to do so because it would require extra computational resources and be time inefficient for a marginal gain in performance, and it can also result in overfitting because D4RL datasets are not very large or diverse.

Reviewer comment (Experiments):
In the latter category of environments that do not require trajectory stitching, more environments/tasks seem to have been silently omitted. The “-expert” suffix versions of Adroit pen, hammer and relocate tasks have been omitted. I would like to see an explanation of why certain environments in Table 1 were omitted, as otherwise it appears as cherry picking. Note that these results were not present in the Appendix either.

Author response:
>  ___The “-expert” suffix versions of Adroit pen, hammer and relocate tasks have been omitted.___

***We have added the Adroit ‘expert’ tasks as well as ‘kitchen-complete’ to Table 1.*** We did not cherry-pick D4RL tasks. We only had a limited compute budget, so the reasoning behind not including these previously was that the datasets for these tasks only consist of expert demonstrations and are better for testing behavior cloning ability rather than offline RL. Further, existing state-of-the-art benchmarks IQL[1] and CQL[2] also did not evaluate Adroit expert tasks in their papers. However, we agree that more task evaluations can highlight the robustness of the algorithm better, and we have added the experiments to Table 1. We found our method performed very well in these tasks compared to baselines, possibly because of the improved behavior cloning capabilities with diffusion models. We summarize the results here:
|                       | BC    | BCQ   | CQL   | IQL  | DT   | Diffuser | DD | LDCQ       | LDGC       |
|-----------------------|-------|-------|-------|------|------|----------|----|------------|------------|
| maze2d-umaze          | 3.8   | 12.8  | 5.7   | 47.4 | 27.3 | 113.5    | -  | 134.2+-4.0 | 141.0+-2.7 |
| maze2d-medium         | 30.3  | 8.3   | 5.0   | 34.9 | 32.1 | 121.5    | -  | 125.3+-2.5 | 139.9+-4.2 |
| antmaze-umaze-diverse | 45.6  | 55.0  | 84.0  | 62.2 | 54.0 | -        | -  | 81.4+-3.5  | 85.6+-2.4  |
| kitchen-complete      | 65.0  | 52.0  | 43.8  | 62.5 | -    | -        | -  | 62.5+-2.1  | -          |
| pen-expert            | 85.1  | 114.9 | 107.0 | -    | -    | -        | -  | 121.2+-3.6 | -          |
| hammer-expert         | 125.6 | 107.2 | 86.7  | -    | -    | -        | -  | 45.8+-10.5 | -          |
| door-expert           | 34.9  | 99.0  | 101.5 | -    | -    | -        | -  | 105.1+-0.3 | -          |
| relocate-expert       | 101.3 | 41.6  | 95.0  | -    | -    | -        | -  | 104.7+-1.4 |            |

Reviewer comment (Evaluation):
Some qualitative results such as visualization of trajectories at test time would have been helpful for supporting the trajectory stitching property.

Author response:
> ___Some qualitative results such as visualization of trajectories at test time___

Unlike Diffuser, we do not predict the full trajectory at once. Since we do Q learning, greedily predicting the next best H-length sequence of actions. As such, it is difficult to show the stitching ability that happens during diffusion since we only diffuse smaller sub-trajectories at a time. The complete trajectory for tasks like antmaze look similar for all algorithms, with the main failure mode mainly just toppling of the ant after hitting walls. Let us know if you had something else in mind, and we can try to create these visualizations.

Reviewer comment (Evaluation):
Not all reported performance values are state-of-the-art, this claim should be qualified.

Author response:
>  ___Not all reported performance values are state-of-the-art___

We have not claimed SOTA in the locomotion tasks, and in Appendix F we had provided potential reasons why our method seems to underperform in these tasks. To summarize, we suspect the high-frequency periodicity of the walking gaits does not require temporal abstraction, and our lack of addition of a perturbation function like in BCQ for simplicity of tuning means we do not trade off conservatism from the behavior support. As shown in Table 1, we do have very strong results in all non-Locomotion tasks compared to baselines.

Reviewer comment (Writing):
Batch-constrained Q (BCQ) is mentioned several times and the acronym is used once before it is explained in Section 4.2. I recommend briefly explaining BCQ as a part of “Section 3.2 Offline Reinforcement Learning”.

Author response:
Clarity:

>  ___Batch-constrained Q (BCQ) is mentioned several times and the acronym is used once before it is explained___
 
We now reference BCQ when we introduce the idea of batch-constrained action selection in the related work section.

Reviewer comment (Reproducibility):
Algorithms 1 and 2 explain the second stage of training and the policy execution phase respectively. However, there is no algorithm box for the first stage of training (VAE). The policy \pi also has a subscript \theta, but \theta is not used anywhere else in the paper. For completeness, I would recommend adding an Algorithm to describe the first stage of training where the policy parameters \theta are learned.

Author response:
>  ___I would recommend adding an Algorithm to describe the first stage of training___

***We have added the beta-VAE and diffusion prior training to Algorithm 2 in Appendix B, referenced in section 4.1.***

Reviewer comment (Experiments):
My major concern is that the results on offline RL benchmarks may be insufficient to show the advantage of planning with diffusion models on latent action space. While I appreciate that the authors have covered the most popular state-of-the-art baselines in Table 1, I think it is necessary to compare LDCQ with some literature that similarly performs planning on the learned action representation space learned by VAE [1, 2, 3, 4] or Flow [5, 6]. Also, I am curious about the ablation of LDCQ on diffusion models (Section 5.2 has shown some support, however, it compares the diffusion prior with the VAE conditional prior rather than a learned latent policy).

Author response:
We thank the reviewer for sharing their concerns. We aim to address all concerns in the following response.

> ___compare LDCQ with some literature that similarly performs planning on the learned action representation space learned by VAE or Flow___

A1: Thanks for raising this interesting point. We have added a comparison with OPAL and PLAS among VAE methods, and Flow to Control among offline RL methods that use normalizing flows. For OPAL scores we used the code provided to us by the authors. SPiRL is an online RL method, and if used offline would resemble OPAL. All other remaining references suggested by the reviewer also do not have open-source codebases, making them hard to reproduce. We will add fair comparisons with the remaining references in the camera-ready version. 

LDCQ (ours) vastly outperforms OPAL and PLAS (VAE methods) across tasks. While Flow to Control reports better scores, we found it difficult to tune our implementation of their method. We have contacted the authors for their code, and have currently compared with the reported scores in their paper. ***We have only made comparisons in the tasks these papers quoted in their respective papers. Please refer to Appendix D, Table. 5 of the revised paper on OpenReview*** (and our global rebuttal). We list here the scores in antmaze and kitchen. The other scores are in the paper:
|                        | OPAL | PLAS | Flow to Control | LDCQ      |
|------------------------|------|------|-----------------|-----------|
| antmaze-medium-diverse | 57.5 | 0.0  | 83.7            | 68.9+-0.7 |
| antmaze-large-diverse  | 52   | 0.0  | 52.8            | 57.7+-1.8 |
| kitchen-partial        | 55.5 | 43.9 | 74.9            | 67.9+-0.8 |
| kitchen-mixed          | 50.2 | 40.8 | 69.2            | 62.3+-0.5 |

Reviewer comment (Theory):
The paper has discussed some limitations of prior works depending on VAE representations. However, it is unclear to me why introducing latent diffusion can solve these limitations, as LDCQ also requires a learned VAE embedding. Is it because the diffusion model is more expressive so you can use less KL penalty or there are some other reasons? And if we learn the latent action space with a flow model, which is a lossless representation, does latent diffusion planning still have some advantages?

Author response:
> ___Is it because the diffusion model is more expressive so you can use less KL penalty or there are some other reasons?___

A2: That is correct. The reason this works with a $\beta=0.05$ is because the diffusion model is expressive enough to capture the more unstructured but information-rich latent space produced with small $\beta$. This was not possible with the previous methods which learnt VAE embeddings and used VAE prior. See Figure. 3 for a PCA visualization of the diffusion latent space and the comparison with VAE prior.

 > ___And if we learn the latent action space with a flow model, which is a lossless representation, does latent diffusion planning still have some advantages?___

A3: We believe latent space planning compresses information across longer timescales into low-dimensional embeddings, which helps with computational/time efficiency of diffusion. 

However, if the latent space is lossless, and is an invertible function of the original state-action timeseries, then it is probably better to directly diffuse the raw time-series sequence. Moreover, there exists an asymmetry in the encoding and decoding stages, as we encode state-action sequences, but we only decode action sequences. While this can still be implemented via conditional NFs, it is simpler and more stable to directly diffuse the action sequence from states. 

In fact, diffusion models themselves are a particular type of continuous normalizing flow (CNF) through the probability flow ODE [1, 2]. Diffusion models are deep autoregressive normalizing flows that also conveniently have a fixed, tractable, simulation-free noising process, which allows us to train each reverse step independently. They are also more flexible with network architectures while normalizing flows require special architectural considerations for a diagonal Jacobian. As such, even if NFs might be sufficient for smaller tasks such as in D4RL and comparable in performance with latent diffusion, we claim that as we scale to large offline datasets, we expect diffusion models to be the choice of generative models that the community will turn to. Diffusion models have been shown to scale very well with large scale datasets in other generative domains. Our paper demonstrates a simple way to leverage these latent diffusion models for offline skill learning.

[1] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021

[2] Lipman, Yaron, et al. "Flow Matching for Generative Modeling." The Eleventh International Conference on Learning Representations. 2022.

Please let us know if there are any remaining questions. We look forward to continuing the discussion.