Thank you for your questions. We have merged answers to some of them as they are related, and have attempted to address all of them:

Experiment results are difficult to follow.... It seems the goal conditioned tasks are not the standard ones in GCRL literature, and it is not clear what the key takeaway is other than the constrained offline RL results + qualitative result

The takeaway of D4RL CARLA experiment on carla-lane-v0 was to show that LDCQ scales more graciously than prior offline RL algorithms to image-based environments (last row of Table 1). The takeaway of GCRL experiment was to make a fair comparison with the Diffuser baseline, which uses goal-conditioned inpainting strategy. Our goal-conditioned latent diffusion method (LDGC) outperforms Diffuser by 65% on maze2D-large (Table 1, last column).

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

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.

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

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.

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

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?

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?

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

We hope we have sufficiently answered your questions. Please let us know if you would like any further clarification, or if you think we have not addressed any particular point..

References follow in the next comment

Dear Reviewer,

We hope that you've had a chance to read our responses and clarification. As the end of the discussion period is approaching, we would greatly appreciate it if you could confirm that our updates have addressed your concerns.

Dear Reviewer,

We are rapidly approaching the end of the discussion phase. We hope we have answered your clarifying questions. We have made revisions to the paper to include the remaining D4RL tasks in Table 1. Our novel approach leveraging latent diffusion for TD learning using BCQ to train high level policies outperforms other offline RL methods in a majority of these tasks (maze2d, antmaze, franka, adroit and carla). We have also included further comparisons against other latent skill methods in Table 5, Appendix D. We would appreciate it greatly if the reviewer can acknowledge that we have addressed their concerns.

Dear Reviewer,

As the discussion phase is about to end, we would like to again kindly request if you could acknowledge the updates we have made to the paper and our response to your questions, and whether it has addressed your concerns.

We believe most of your concerns stem from the usage of the beta-VAE, the choice which we have justified in our earlier response. We have also added TD error plots during training as per your suggestion as a preliminary analysis to the stability of Q learning with beta-VAE latents. We found training to be stable and reproducible, and the strong results in a majority of the D4RL benchmark suggest that the latents are information rich and that this information is easily extracted by the high level Q function during Q-learning. We acknowledge that there are other representation learning methods which can yield high quality representations, but our work instead focuses on sampling these weakly regularized information rich representations with latent diffusion (similar to image latent diffusion generative models), and how these skill priors can be used to learn high value policies in offline RL.

We thank you for your feedback

READ PART 1 FIRST

Confidence intervals (error bars) are missing in Figure 4 and Adroit environment/tasks of Table 1.

We have added the missing error bars in Adroit, over 5 runs of the algorithm.

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

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.

Questions:

What would happen if z-i is re-sampled after every action is taken by the policy

We found that planning skills at a higher frequency did not seem to affect performance in any measurable way in antmaze tasks, and improved slightly but within the margin of error for Franka tasks. As we show in Appendix I, increasing diffusion steps seems to meaningfully improve performance, and high-frequency high-level planning with more diffusion steps would be very inefficient/slow.

We hope we have addressed your concerns. Please let us know if you would like us to clarify anything else.

Thanks for your detailed review. We have tried to address most of your points in our revised draft, and have provided some more context below:

Stretched claims:

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.

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.

Empirical evidence:

In Table 1, not all environments/tasks from D4RL are included. For example, the medium and umaze versions of Maze-2D are missing.

Thanks for pointing this out. We have added these experiments and antmaze-umaze-diverse to Table 1. Our proposed LDCQ methods outperforms other baselines in the smaller mazes as well. We agree that they can still be useful results alongside the harder antmaze tasks.

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:

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.

RESPONSE CONTINUED IN NEXT COMMENT

[1] Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement learning with implicit Q-learning. In International Conference on Learning Representations, 2022.

[2] Kumar, A., Zhou, A., Tucker, G., and Levine, S. Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020.

Dear Reviewer,

We hope that you've had a chance to read our responses and clarification. As the end of the discussion period is approaching, we would greatly appreciate it if you could confirm that our updates have addressed your concerns.

Dear Reviewer,

We are rapidly approaching the end of the discussion phase. We have addressed the comments in the review, and have made revisions to the paper to include the remaining D4RL tasks in Table 1. Our novel approach leveraging latent diffusion for TD learning using BCQ to train high level policies outperforms other offline RL methods in a majority of these tasks (maze2d, antmaze, franka, adroit and carla). We have also included further comparisons against other latent skill methods in Table 5, Appendix D. We would appreciate it greatly if the reviewer can acknowledge that we have addressed their concerns.

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:

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.

We are happy to answer these questions.

Does this result hold for trajectories during evaluation time that are not exactly the same as training data?

1. Yes, this result holds during evaluation time as well. We show this in Section 5.3 (see Fig. 4), where we compare the D4RL evaluation performance of the same generative models shown in Fig. 3, i.e., latent diffusion (LDCQ) and VAE (BCQ-H). Both LDCQ and BCQ-H follow the same batch-constrained Q-learning (BCQ) setup to learn latent policies, with the only difference being the choice of the generative model to approximate the prior. This result suggests that diffusion models are not just better than VAEs at modeling the prior distribution in the latent space (as shown in Fig. 3), but they also generate policies that evaluate to a higher performance on the offline RL tasks (as shown in Fig. 4/Table 5).

Does the trained VAE latent policy (as in OPAL) can capture the multi-modal distribution better?

2. This is a great question. In theory, both VAE based methods such as OPAL and diffusion models with sufficient capacity can approximate any target distribution. However, it is commonly observed that diffusion models are better at handling multi-modal distributions than VAEs, as suggested by prior works in image generative modeling[1,2] and imitation learning[3]. Our improved results (shown in Table 5.) compared to OPAL suggest this is true in the offline RL setting as well.

Do these answers fully address the reviewer's concerns? We look forward to continuing the discussion.

References:

[1] “Hierarchical Text-Conditional Image Generation with CLIP Latents”; Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, Mark Chen; https://arxiv.org/abs/2204.06125

[2] “High-Resolution Image Synthesis with Latent Diffusion Models”; Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer; https://arxiv.org/abs/2112.10752

[3] “Imitating Human Behaviour with Diffusion Models”; Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, Sam Devlin; https://arxiv.org/abs/2301.10677