Latent Diffusion Planning for Imitation Learning

Dear Reviewer,

Thank you for your detailed feedback on our project. We are happy that you found the paper to be well-written and clear. We address your questions and comments below. Please let us know whether there are any other concerns you have that prevent you from increasing your score.

Q1: Choice of labeled / unlabeled / suboptimal trajectories are needed for optimal performance

For different robotics tasks, we tried to keep most of the parameters the same: 100 action-free trajectories for robomimic tasks and 500 suboptimal trajectories for all tasks. However, we chose:

# of demonstrations: Robomimic, by default, includes 200 proficient demonstrations per task, and ALOHA includes 50 demonstrations per task. To showcase the effectiveness of our method when the number of demonstrations is low, we chose to take half the number of demonstrations-- 100 for Robomimic tasks, and 25 for ALOHA tasks. However, Lift is a very simple task, and even with around 10 demonstrations, behavior cloning policies can achieve very high accuracy. Thus, specifically due to the simplicity of the Lift task, we chose 3 demonstrations to adhere to the low demonstration setting.

# of action-free trajectories: we chose to use the remaining half of the demonstrations as action-free trajectories. This meant 100 demonstrations for Robomimic tasks, and 25 demonstrations for ALOHA. Since Lift only used 3 demonstrations, we considered using only 3 action-free trajectories. However, because we wanted to try and remain consistent with the number of trajectories, we chose to use 100 action-free demonstrations.

Q2: Does using action-free and suboptimal trajectories at the same time further improve the results? Yes, using action-free and suboptimal trajectories further improves results and strongly exceeds our baselines. We have included updated results. Here is a how LDP + Action-Free + Suboptimal compares against just LDP without the additional trajectories:

In addition, we have reran UniPi with identical hyperparameters to LDP, per Reviewer PyaC’s suggestion.

Q3: Training the inverse dynamics on suboptimal trajectories...

We believe this does help, because LDP + Subopt > LDP, and LDP + Action-Free + Subopt > LDP + Action-Free. In both of these cases, the difference is the suboptimal data for the IDM, which improves policy rollouts. One metric we can use to quantitatively measure this is action MSE for a validation dataset, but we found that this is noisy and not extremely correlated with the policy performance, which is true for many behavior cloning methods. Thus, we don’t find action MSE metrics for the IDM to be particularly insightful, and instead refer to policy performance.

Q4: Are the action-free data always from expert policies? To my understanding, suboptimal action-free trajectories cannot be used because they would decrease the performance of the planner.

Yes, action-free data assumes that optimal behavior.

Dear Reviewer,

Thank you for your detailed feedback on our project. We provide the requested experiments and address the comments and questions in detail below. Please let us know whether there are any other concerns you have that prevent you from increasing your score.

Q1: Improvement in Experiments

We agree the experimental results in our submitted draft show that there is around a 10% improvement for LDP. However, we have now run updated experiments where LDP leverages both action-free and suboptimal data. While BC-like methods work well in settings with a lot of optimal demonstration, LDP effectively leverages additional data sources, which we believe is a fundamental advantage. The results strongly outperform baselines now, by around 30% or more. See Reviewer LWMgQ2

Q2: Novelty

The novelty of this paper is to propose a simple and scalable algorithm to address learning from heterogeneous data sources like suboptimal and action-free data. Our main contribution is in showing good performance with a model-based method, which will enable future work to build on it. Further, we provide a comprehensive evaluation of alternative methods for learning from heterogeneous data, and establish a novel finding that our planning-based is particularly suited for this setting.

In terms of comparison to UniPi, we agree that the method is fairly similar. However, we find that LDP strongly outperforms UniPi due to its method, with the main difference being the much lower-dimensional latent space. Due to this change in method, there is a large improvement in performance. We attach updated results below and highlighted in red in our PDF.

In terms of how to shape the space, we add additional experiments using pretrained DINOv2 embeddings, which investigates how pretrained embeddings compare to LDP’s latent space. In these experiments, we directly swap out VAE embeddings with frozen DINOv2 embeddings.

However, we found directly planning over the DINOv2 embeddings does not lead to good learned behaviors (0% success rate even for easy tasks), which we hypothesize is due to the challenges of planning over a large latent space (384 dimensional). Thus, as an alternative, we fix a random projection matrix that reduces the 384 dimensional feature space to 16 dimensions, matching LDP’s latent space. We find that LDP strongly outperforms using pretrained latent embeddings. See Reviewer PyaC Q4

Q3: Data Size in Experiments

We agree that using 1,000 or 10,000 trajectories would be interesting. However, most robotics benchmarks and tasks typically use 50-300 demonstrations (Diffusion Policy, Robomimic, Action Chunking with Transformers, etc.). For example. Robomimic, a popular simulated imitation learning benchmark, includes datasets with tasks that have 200 demonstrations. Simulated ALOHA tasks use only 50 demonstrations. It is typically difficult to find single-task imitation learning datasets with more than a few hundred trajectories, due to the difficulty of collecting large-scale expert demonstrations.

Larger datasets are often multi-task or language-conditioned, and they are often used to train large multi-task or language-conditioned policies. (Open X-Embodiment, DROID, Octo, OpenVLA). These datasets are expensive to train on, especially for video models. We do not explore these datasets, due to the lack of good pretrained models, and because video models are much more expensive to train. This is why we are starting with small datasets, but this is an important direction for future work.

Q4: Pretrained Video Models

We agree that LDP cannot use pretrained video models. However, for our UniPi baseline, we compare finetuning a pretrained model vs. training from scratch. To provide a comparison, we use a model with weights pretrained on THOR, and from scratch. In these results, we find that initializing from a pretrained model does not actually improve performance. See Reviewer 9mX6 Q1

The power of pretrained video models may lie in its ability to generalize and extrapolate to new tasks and scenes, often through language conditioning, which we leave to future work. Generalization is important for robot learning, but we find that for single-task imitation learning, which typically use smaller datasets of 50-300 demonstrations, using pretrained video models is not essential.

Q5: Outdated Beta-VAE and DDPM

We agree that there are advances to beta-VAE (VQ-VAE, VQ-GAN, etc.) and DDPM (DDIM, Consistency Models, etc.). However, we choose both of these implementations due to their simplicity. For our planner, fast inference is not crucial to our contribution, and hence, we don’t use faster samplers like DDIM. However, we agree that improvements to both beta-VAE and DDPM can lead to additional improvements and scalability, which we leave to future work.

Q6: Citations

Thank you for those references. We will include them in the draft.

Dear Reviewer,

Thank you for your detailed feedback on our project. We are happy that you found the paper well-presented with clear experimental evaluations. Please let us know whether there are any other concerns you have that prevent you from increasing your score.

Q1: Would the authors be able to explain why these vastly different amounts of trajectories were chosen?

For different robotics tasks, we tried to keep most of the parameters the same: 100 action-free trajectories for robomimic tasks and 500 suboptimal trajectories for all tasks.

# of demonstrations: Robomimic includes 200 proficient demonstrations per task, and ALOHA includes 50 demonstrations per task. To showcase the effectiveness of our method, we always take half the number of demonstrations-- 100 for Robomimic tasks, and 25 for ALOHA tasks. This is a consistent protocol we use to select the number of demonstrations.

# of action-free trajectories: we chose to use the remaining half of the demonstrations as action-free trajectories Lift is a very simple task, so we reduce the number of demonstrations to 3 increase the complexity.

Q2: Ablating the VAE Component

To address “directly planning in high-dimensional image-space,” we include the UniPi baseline, which learns a video prediction planner. UniPi-CL, in particular, mirrors LDP’s planner, in that UniPi-CL forecasts over dense states instead of subgoals. Let us know if that addresses your concern!

Q3.1 DP vs. LDP on ALOHA Cube

Specifically for ALOHA Cube, we found that a larger planner was imperative for reasonable performance. We did not find this necessary for DP, and thus, we kept the smaller DP architecture. We have updated results where DP is trained with the same architecture size as all LDP variants (down_dims = [512, 1024, 2048]). We use batch_size=16 for DP, since the end-to-end training of the encoder requires much more GPU memory than LDP, which uses frozen embeddings. We couldn't run DP with a larger batch size or down_dims.

Q3.2 LDP vs. UniPI: IDM size for Can, and IDM train time

We observed a bigger IDM is helpful for the Can task, likely because the scene is slightly more visually complex. As requested, we have now rerun Can experiments with UniPi to ensure the same IDM architecture.

We train the IDM model for longer in LDP because we use a smaller batch size. However, we agree that for consistency, rerun UniPi experiments with the same batch size and # gradient steps for LDP and UniPi IDM. In addition, we have retrained UniPi GCBC models with the same hyperparameters as LDP. See Reviewer LWMg Q2.

Q3.3 LDP 50% optimal and 50% suboptimal batches

Each batch consists of 50% optimal and 50% suboptimal trajectories.

Q3.4 IDM one expert demonstration

We meant to write, “The IDM is trained only ON expert demonstrations.” The IDM is only trained on action-labelled data, which in this case, are from the expert demonstrations.

Q3.5 LDP Hierarchical IDM

(Appendix A.2) Hierarchical LDP’s IDM is a Conditional U-Net with down-dims [256, 512], which has 1.67e7 parameters. (Appendix A.1) Non-Hierarchical LDP uses an MLP ResNet based off of IDQL, which has 1.79e6 parameters. The Hierarchical LDP IDM has more parameters, while still underperforming our method.

Q4: Pretrained Encoders

LDP uses a task-specific VAE, because existing image encoders or VAEs (e.g. DINOv2, R3M, Stable Diffusion VAE) produce high-dimensional embeddings. Our VAE produces a much more compact latent space (16 dimension), enabling much faster training and inference, as well as interpretability through decoding image latents.

To address this approach, we swap out VAE embeddings with frozen DINOv2 embeddings.

However, we found directly planning over the DINOv2 embeddings does not lead to good learned behaviors (0% success), which we hypothesize is due to the challenges of planning over a large latent space (384 dim). Thus, as an alternative, we fix a random projection matrix, reducing the 384 dim feature space to 16 dim, matching LDP. We choose this, because this is a straightforward way of using frozen embeddings. It may be possible to plan over the large embeddings space with a much larger and complex model, or it may be possible to learn alternative ways to project DINOv2 embeddings to a lower-dimensional space, but this may lead to fundamental changes in our method, so we don’t explore more complicated approaches to use pretrained embeddings.

Dear Reviewer,

Thank you for your detailed feedback on our project. We are happy that you found the paper easy to follow and relevant to learning from demonstrations. Please let us know whether there are any other concerns you have that prevent you from increasing your score.

Q1: Pretrained UniPi + Finetuning In-Domain

The comparison with UniPi is limited as the original UniPi model, code and pretraining data are not released, per Appendix A.1. To provide a comparison, we use AVDC-THOR, which is the best available reproduction of UniPi. We include results on Robomimic tasks. The experiments compare how finetuning a pretrained UniPi model, instead of training a UniPi model from scratch, affects overall performance. We finetune for an additional 50k steps at learning rate 1e-4.

We observe the pretrained AVDC/UniPi model does not substantially improve performance, likely because the pretrained model quality is limited. For UniPi-OL, there appears to be slightly improvement when pretraining from scratch, possibly because without closed-loop planning, the pretrained initialization may be helpful.

Q2: Action-Free Demonstrations

Thank you for this suggestion. We were unable to set up the UMI hardware within the short rebuttal time, but we agree that a more scalable and realistic way of obtaining action-free expert data is through data collection tools such as UMI (Universal Manipulation Interface).

Q3: Collecting Demonstrations in the Real World

The demonstrations were collected from one teleoperator. The average time for one demonstration was around 25 seconds, including resetting the cube. In total, for 82 trajectories, it would take around 34 minutes to collect all demos.

For suboptimal trajectories, the teleoperator merely supervises the policy to ensure safety, and does not need to actually teleoperate the robot at all. Each suboptimal trajectory takes around 75 seconds, including resets. For 85 trajectories, this would take around 1 hour 45 minutes.

The total time is longer than LANE from Zhao et. al., which collects limited demonstrations (requires teleoperator) and performs RL (does not require active teleoperation). It’s possible that our method may work with fewer demonstrations (less teleoperation time) or with less suboptimal data, but we did not sweep over these numbers or optimize for these metrics.

Q4: Additional Comments

“The tasks are rather simple with clean visuals.” -- We agree that the tasks are not extremely complex nor do they have complex backgrounds or visuals. However, we chose commonly used robotics benchmarks for this project and focused on showing the effect of action-free and suboptimal data.

“The main weakness is that the algorithm still seems to require many demonstrations (and more suboptimal data) for rather simple tasks.” -- We agree that this paper doesn’t operate in a few-shot demonstration setting. However, the standard number of demonstrations provided by robomimic is 200 per task, and ALOHA is 50 per task. The robomimic benchmark is saturated with 200 demos, and we chose to restrict the number of demonstrations to evaluate our method in a lower demonstration regime. In order to learn from even fewer demonstrations for these tasks, it would likely be necessary to learn from prior data (e.g. from retrieval), or reinforcement learning, which we do not focus on.

Q5: Improved Experiments

Please see Reviewer LWMg Q2. We include updated experiments where LDP leverages both action-free and suboptimal data. While BC-like methods work well in settings with a lot of optimal demonstration, LDP effectively leverages additional data sources, which we believe is a fundamental advantage. The results strongly outperform baselines now, by around 30% or more.

In addition, we have reran UniPi with identical hyperparameters to LDP, per Reviewer PyaC’s suggestion.