Prediction with Action: Visual Policy Learning via Joint Denoising Process

TL;DR

PAD addresses the coupling between predictive vision and robotic control by unifying image forecasting and action generation under a single diffusion-transformer framework. It jointly denoises multi-modal latent representations conditioned on current observations and language instructions to produce future frames and actions. The method enables co-training on robotic demonstrations and large-scale internet video data and supports additional modalities like depth. Empirically, PAD achieves a 26.3% relative improvement on Metaworld and a 28.0% relative gain in unseen real-world tasks, demonstrating strong data efficiency and generalization.

Abstract

Diffusion models have demonstrated remarkable capabilities in image generation tasks, including image editing and video creation, representing a good understanding of the physical world. On the other line, diffusion models have also shown promise in robotic control tasks by denoising actions, known as diffusion policy. Although the diffusion generative model and diffusion policy exhibit distinct capabilities--image prediction and robotic action, respectively--they technically follow a similar denoising process. In robotic tasks, the ability to predict future images and generate actions is highly correlated since they share the same underlying dynamics of the physical world. Building on this insight, we introduce PAD, a novel visual policy learning framework that unifies image Prediction and robot Action within a joint Denoising process. Specifically, PAD utilizes Diffusion Transformers (DiT) to seamlessly integrate images and robot states, enabling the simultaneous prediction of future images and robot actions. Additionally, PAD supports co-training on both robotic demonstrations and large-scale video datasets and can be easily extended to other robotic modalities, such as depth images. PAD outperforms previous methods, achieving a significant 26.3% relative improvement on the full Metaworld benchmark, by utilizing a single text-conditioned visual policy within a data-efficient imitation learning setting. Furthermore, PAD demonstrates superior generalization to unseen tasks in real-world robot manipulation settings with 28.0% success rate increase compared to the strongest baseline. Project page at https://sites.google.com/view/pad-paper

Figures (11)

• Figure 2: Diffusion models have achieved impressive success in visual generation tasks (a) and visual-motor control tasks (b). Image prediction and robot action are actually highly correlated since they share the same underlying physical dynamics. The PAD framework predicts the future and generates actions in a joint denoising process.
• Figure 3: Visualization of the PAD framework. Current observations in different modalities are first encoded into latent and concatenated with white noise channel-wise. These noised latent are then tokenized into tokens and perform a joint denoising process to predict the images and robot actions simultaneously. PAD can flexibly accommodate extra or missing modal inputs through a masked-attention mechanism
• Figure 4: We learn a single vision-language conditioned policy to solve all tasks in each domain with limited demonstrations, co-training with the bridge video data. In simulated MetaWorld, we learn a policy to tackle all 50 tasks. In real-world panda manipulations, we split objects into seen objects and unseen new objects to test the generalization ability of our policy.
• Figure 5: Generalization test under 3 levels of difficulties. The yellow bounding box suggests the target position. Our proposed PAD shows the strongest generalization abilities in unseen tasks.
• Figure 6: Comparisons on predicted images between PAD and GR-1. PAD generates more precise images than GR-1 which may potentially lead to more accurate control actions. Zoom in for better comparisons.