VILP: Imitation Learning with Latent Video Planning

TL;DR

VILP addresses the challenge of leveraging video data to train robot policies by integrating a latent video diffusion model into imitation learning. It introduces latent video planning, where multiview observations condition a latent-space diffusion model to generate time-aligned future frames that are mapped to action sequences via a low-level policy, enabling real-time receding horizon planning. The approach demonstrates faster inference, lower training costs, and robust performance with limited task-specific action data, including successful real-world demonstrations with real-time closed-loop control. These results highlight the potential of video-to-action pipelines and cross-domain video pretraining to scale robot learning beyond action-labeled datasets.

Abstract

In the era of generative AI, integrating video generation models into robotics opens new possibilities for the general-purpose robot agent. This paper introduces imitation learning with latent video planning (VILP). We propose a latent video diffusion model to generate predictive robot videos that adhere to temporal consistency to a good degree. Our method is able to generate highly time-aligned videos from multiple views, which is crucial for robot policy learning. Our video generation model is highly time-efficient. For example, it can generate videos from two distinct perspectives, each consisting of six frames with a resolution of 96x160 pixels, at a rate of 5 Hz. In the experiments, we demonstrate that VILP outperforms the existing video generation robot policy across several metrics: training costs, inference speed, temporal consistency of generated videos, and the performance of the policy. We also compared our method with other imitation learning methods. Our findings indicate that VILP can rely less on extensive high-quality task-specific robot action data while still maintaining robust performance. In addition, VILP possesses robust capabilities in representing multi-modal action distributions. Our paper provides a practical example of how to effectively integrate video generation models into robot policies, potentially offering insights for related fields and directions. For more details, please refer to our open-source repository https://github.com/ZhengtongXu/VILP.

Figures (8)

• Figure 1: VILP is capable of generating predictive robot videos that adhere to temporal consistency to a good degree. Within the VILP framework, these videos can be further mapped to robot actions. Through extensive experiments, we demonstrate the effectiveness and real-time performance of VILP. The video plans depicted in this figure are some frames extracted from videos produced by VILP.
• Figure 2: Illustration of our proposed video planning pipeline. For the architecture of $\epsilon_\theta$, we employ a UNet model with 3D convolution layers. We adopt the cross-attention conditioning mechanism proposed in rombach2022high and extend it to visual conditioning for 3D video data. $k$ represents the step index of the denoising process, where $k = 1, \ldots, K$. Variables with $\hat{}$ indicate the estimated/predicted variables produced by the model.
• Figure 3: Overview of the low-level policy that maps the predicted video to predicted action sequence.
• Figure 4: Tasks for experiments. Move-the-Stack xu2024leto, Push-T chi2023diffusionpolicy, Towers-of-Hanoi zeng2021transporter are for video planning experiments (no policy involved). Nut-Assembly mandlekar2022matters, Can-PickPlace mandlekar2022matters, Sim Push-T chi2022iterative, and Arrange-Blocks chi2022iterative are for the evaluation of the whole process which consists of video planning and low-level policy. Move-the-Stack and Push-T are in real environment while the rest tasks are in simulation.
• Figure 5: Example frames extracted from the generated videos. For more synthesised video visualizations, please see our supplementary video.