Prediction with Action: Visual Policy Learning via Joint Denoising Process

1 Introduction

• Prediction and action are linked because both depend on the same underlying physical laws, motivating their joint learning.Prior two-stage approaches synthesize a goal image before learning a policy, leaving prediction and action separated.Source: §1 · pp. 1–2
• PAD jointly predicts future images and robot actions through a shared diffusion-transformer denoising process.The framework integrates prediction and action rather than using a two-stage goal-image and policy pipeline.Source: §1 · p. 2
• PAD can co-train on robotic demonstrations and large-scale video datasets despite their differing available modalities.The DiT-based framework is also extendable to additional robotic modalities such as depth images.Source: §1 · p. 2
• On the Metaworld benchmark, the method surpasses previous methods and baselines with a 26.3% relative improvement in success rate, using a single visual-language-conditioned policy.Source: §1 · p. 2

2 Preliminaries

• Robotic demonstrations contain language instructions, pixel-input sequences, and robot trajectories, but their scale is limited because collection is risky and costly.The formulation denotes the robotic dataset as D_robot and each demonstration as a language-conditioned trajectory.Source: §2 · pp. 2–3
• PAD co-trains robotic demonstrations with RGB video data that lacks robot actions, using video to enhance visual policy learning.The framework handles missing action modalities while leveraging the larger accessible video dataset.Source: §2 · p. 3
• Latent diffusion progressively adds Gaussian noise to encoded data over T steps and then denoises from a highly noisy latent to recover the real sample under condition c.The noise amplitude is controlled by predefined α_t hyperparameters.Source: §2 · p. 3
• A neural noise estimator ε_θ(z_t,t,c) is trained to approximate the gradient of the log-density of noisy data.This estimator supports the denoising process used to generate conditioned samples.Source: §2 · p. 3

3.1 Overview of PAD

• PAD concurrently predicts future frames and robot actions from multimodal observations through a joint latent denoising process.The framework primarily models RGB images and robot poses, while also supporting depth images as an additional modality.Source: §3.1 · p. 3, §3.2 · p. 4
• PAD forecasts k future steps but executes only the first predicted action before repeating prediction, enabling closed-loop planning and acting.Source: §3.1 · p. 3

3.2 Model Architectures

• DiT concatenates token sequences from different modalities and jointly denoises them using self-attention.Image latents are encoded with a frozen VAE, robot poses with an MLP, depth images are tokenized when available, and text uses a frozen CLIP encoder.Source: §3.2 · pp. 3–4
• Masked self-attention lets PAD handle missing modalities, allowing concurrent training on RGB-only video datasets and robotic datasets.Padding tokens are excluded, and only effective predictions are retained.Source: §3.2 · p. 4, Figure 3 · p. 4
• Conditioned noisy latents concatenate encoded current observations with noise before tokenization and DiT processing to predict future latent representations.The stacked noise represents multiple future frames, while the same process is applied across modalities.Source: §3.2 · p. 4

3.3 Training Process

• PAD initializes from an ImageNet-pretrained class-conditioned DiT while adding or adapting modality-specific components.New text, robot-state, and stacked-latent components are initialized or replicated to match the joint architecture.Source: §3.3 · p. 4
• The training objective jointly predicts diffusion noise for image, action, and optional depth latents conditioned on current observations and instructions.Modality-specific losses are combined with λI, λA, and λE to balance prediction objectives.Source: §3.3 · p. 5

4.1 Environmental Setups and Baselines

• A single text-conditioned visual policy addresses all 50 MetaWorld tasks and real-world manipulation tasks using limited demonstrations and video co-training.Real-world evaluation separates seen and unseen objects to assess generalization.Source: §4.1 · p. 5, Figure 4 · p. 5

4.2 Main Results

• 72.5% average success rate on all 50 MetaWorld tasks surpasses the strongest baseline at 57.4%.The comparison uses one visual policy across the full benchmark.Source: Table 1 · p. 6, §4.2 · p. 7
• PAD predicts more precise future images than GR-1, potentially supporting higher success on manipulation tasks requiring precise operation.The comparison is visualized in Figure 6, including tasks such as small-block picking and insertion.Source: §4.2 · p. 7, Figure 6 · p. 8

4.3 Generalization Analysis

• PAD generalizes to unseen objects and backgrounds across three difficulty levels, while the baseline fails on difficult unseen tasks.Tests include strawberries, carrots, eggplants, disturbances, and unseen backgrounds without depth input.Source: Figure 5 · p. 7, §4.3 · p. 8

4.4 Ablation Studies

• Removing image prediction causes significant performance drops, indicating that image prediction contributes useful physical knowledge beyond low-dimensional robot-pose supervision.Predicting only robot pose may provide limited supervision and encourage training-data overfitting.Source: §4.4 · pp. 8–9
• Removing internet-video co-training lowers both policy performance and predicted-image quality, including missing objects in generated images.Co-training improves image quality and may indirectly enhance robot action prediction.Source: §4.4 · p. 9, Figure 8 · p. 9
• Adding depth enables joint prediction of RGB images, depth images, and actions, and improves manipulation performance.The framework is also described as extensible to modalities such as tactile force or point clouds.Source: Table 2 · p. 7, §4.4 · p. 9, Figure 9 · p. 9

4.5 Scaling Analysis

• Success rate correlates strongly with transformer Gflops, while halving image patch size quadruples image-token length and increases computational cost.Scaling experiments vary model size and patchify size on MetaWorld.Source: §4.5 · p. 9

Pre-training for Embodied Control

• Pre-trained vision-language models support embodied control through planning, low-level motor signals, representations, rewards, goal images, and future-sequence prediction.Source: §5 · pp. 9–10

Scaling Analysis

• Table 3 evaluates PAD across model sizes and computational costs, while Figure 10 relates Transformer Gflops to policy success rate.Source: Table 3 · p. 10, Figure 10 · p. 10

Diffusion Models for Embodied Control

• Diffusion policies model multimodal actions from RGB observations, but task-specific training from scratch misses pre-training with internet data.Source: §5 · p. 10

Future Prediction for Policy Learning

• Future-prediction methods use autoregressive or two-stage pipelines, whereas PAD jointly predicts images and actions through a multimodal DiT with token concatenation and attention masks.PAD integrates image generation and robotic action within one denoising process instead of relying on generated images as an interface between separate stages.Source: §5 · p. 10

Conclusion

• PAD jointly predicts future images and robot actions, co-trains with internet video datasets, and extends to other robotic modalities.Both simulated and real-world experiments demonstrated the efficiency of PAD.Source: §6 · p. 10

Limitations

• PAD was tested with only three modality types, and jointly denoising images and actions limits its control frequency.Future work may use intermediate image latents instead of high-dimensional pixel spaces.Source: §6 · p. 10

Appendix

• PAD videos are available on the project webpage.Source: A Appendix · p. 15

Image Encoder and Tokenizer

• The frozen VAE encodes a 256 × 256 image into a 32 × 32 × 4 latent, which is patchified into 256 tokens for PAD-XL/2.Source: §A.1.1 · p. 15

Robot Action Encoder and Tokenizer

• Robot actions are concatenated and projected through MLP layers into tokens representing future 7-dimensional poses.The encoded action vector contains (k + 1) ∗7 dimensions, or (k + 1) ∗4 in Metaworld.Source: §A.1.1 · p. 15

Depth Image Encoder and Tokenizer

• Optional depth images are down-sampled to 32 ∗32 ∗1 and patchified with size 8 into 16 tokens.Source: §A.1.1 · p. 15

Output Decoder

• The output decoder reconstructs future modality latents from DiT tokens and recovers the original samples with the corresponding modality decoders.Source: §A.1.1 · p. 15

• RT-1 uses the official Robotics Transformer implementation as its experimental baseline.Source: §A.2 · p. 15
• Diffusion Policy follows the official implementation and adds language conditioning using the ScalingUp codebase.Source: §A.2 · p. 15
• RT-2 is reimplemented with InstructBlip-vicuna-7b as its backbone.Source: §A.2 · p. 15
• GR-1 initializes from the authors’ open-source checkpoint, using independently written training code based on the inference model.Source: §A.2 · p. 15

A.2.1 Additional Model Training Details

• Model-size and computational-cost variants are reported in Table 4.Source: Table 4 · p. 15

A.3 Real world Experiment Details

• Real-world demonstrations cover six task categories with varied block, plate, and object configurations.The collection includes button press, cable routing, pick and place, and drawer opening or closing, with 1–5 objects randomly placed.Source: §A.3 · p. 16, Figure 11 · p. 16
• Generalization is tested on numerous unseen objects and more complicated tasks than those used for expert demonstrations.Source: §A.3 · p. 16, Figure 12 · p. 16

A.4 Details Baselines and Ablations in Metaworld

• Detailed Metaworld success rates for baselines and ablations are reported with 25 rollouts per task, excluding two low-success handle-pull tasks.Source: Table 5 · p. 17

A.5 Detailed Scaling Results

• Scaling results compare success rates across model sizes and computational allocations, again excluding the low-success handle-pull tasks.Source: Table 6 · p. 18

A.6 Instructions used in tasks

• Metaworld task instructions are primarily designed from the corresponding task names.Source: Table 7 · p. 19, §A.6 · p. 19