Prediction with Action: Visual Policy Learning via Joint Denoising Process
Authors: Yanjiang Guo
Yucheng Hu
Jianke Zhang
Yen-Jen Wang
Xiaoyu Chen
Chaochao Lu
Jianyu Chen
Venue: NeurIPS 2024 poster (2024)
Keywords: Visual policy learning
diffusion
image generation
Primary area: robotics
TL;DR: We propose a novel framework to predict images and robot actions through joint dinoising process.

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 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 \textbf{PAD}, a novel visual policy learning framework that unifies image \textbf{P}rediction and robot \textbf{A}ction within a joint \textbf{D}enoising 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 38.9\% 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. 
Videos of PAD can be found at https://sites.google.com/view/pad-paper

Prediction with Action: Visual Policy Learning via Joint Denoising Process
Yanjiang Guo Yucheng Hu 13+, Jianke Zhang', Yen-Jen Wang Xiaoyu Chen 12 Chaochao Lu Jianyu Chen 12f
Jianke
"qqual Contribution 'Corresponding Author
IIS. Tsinghua University Shanghai Qizhi Institute Shanghai AI Lab *University of California, Berkeley fguoyj22, huyc24) @ mmilsstsinghuu..cd..co
Abstract
Diffusion models have demonstrated remarkable capabilities in image generation tasks, including image editing and video creation, representing good understand- ing of the physical world. On the other line, diffusion models have also shown promise in robotic control tasks by denoising actions, known diffusion policy. Although the diffusion generative model and diffusion policy exhibit distinct capa- bilities-iman 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, novel visual policy learning framework that unifies image Prediction and robot Action within joint Denoising process, Specifically, PAD utilizes Diffusion Transformers (DiT) seamlessly integrate images and robot states, enabling the simultancous predic- tion of future images and robot actions. Additionally, PAD supports co-training 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 significant 26.3% relative improvement on the full Metaworld benchmark, by utilizing single text-conditioned visual policy within a data-efficient imitation learning setting. Furthermore, PAD demonstrates supe- rior generalization to unseen tasks in real-world robot manipulation settings with 28.0% success rate increase compared to the strongest baseline. Project page
hetps: /sites google.oo..ooooooooooooooooo.... www. view ppadppaaaappapppaaaap pad-paperdrrrpppppp paper
1.0
Diffusion Policy
SuSIE
GR-1
PAD(ours)
Success Rate
0.8
0.6
0.4
0.2
0.0
Metaworid
Real-worid Panda (Seen Taskss
Real-world Panda (Unseen Tasks)
(All tasks)
Figure 1: Multi-task performance comparisons in two domains.
Introduction
Making predictions and taking actions are critical human capabilities, allowing individuals to foresee the change of their surroundings and behave appropriately in response [ 2]. Despite prediction and action seeming like two distinct abilities, they are highly coupled since they share the same
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Denoising
Denoising
PAD Denoising
Image Editing, Video Generation
Policy Learming
Co-train
RGB Image
Rebot pose
Other modalities (e.g. depth)
Diffusion Generative Model
(b) Diffusion Policy
Video
[c] Joint Prediction Framework
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 joint denoising process.
underlying physical laws of the world 3 Understanding these laws enables humans to make better predictions and actions.
Recently, diffusion models 4 6 have achieved impressive success in visual generation tasks by training on extensive web-scale image and video datasets 7 9. For example, image editing models can predict outcomes based on user instructions T0 11 while video generation models can generate sequences of future images [I3 14 [5], representing good understanding of the physical world. On the other line, diffusion models have also shown efficacy in robotic control tasks by denoising actions conditioned on robot observations, known as diffusion policy [16]. Although the diffusion generative model and diffusion policy serve different functions across two domains, we believe that the capability for image prediction could significantly enhance robot policy learning, as they share the same fundamental physical laws. Previous works image-editing model in off- off-the-shelf manner by first synthesizing goal image and subsequently learning goal-conditioned policy. However, this two-stage approach separates the prediction and action learning process, neglecting deeper connections between prediction and action. In this way, actions do not leverage the pre-trained representations in the prediction models which encode rich knowledge of the physical world.
have employed the
In this paper, we introduce the Prediction with Action Diffuser (PAD), a unified policy learning frame- work that integrates prediction and action under the same diffusion transformer (DiT) architecture 20). Specifically, we utilize the diffusion transformer model seamlessly merge all modality inputs and simultaneously predict future images and actions via joint denoising, as illustrated in Figure c).. Additionally, the flexible DiT backbone also allows PAD to co-trained on large-scale video data and extended to other robotic modalities, such as depth images. We have conducted extensive experiments on the MetaWorld Benchmark [21] as well as real-world robot arm manipulation tasks, demonstrating the efficacy of our approach, shown in Figure] Our key contributions are:
We propose a novel policy learning framework, Prediction with Action Diffuser (PAD), predict futures and robot actions through joint denoising process, benefiting policy learning for robotic tasks.
The proposed PAD framework enables co-training different datasets containing different modalities, allowing encoding rich physical knowledge from various data sources.
We outperform previous methods with clear margin in the Metaworld benchmark, sur- passing baselines with 26.3% relative improvement in success rate using a single visual- language conditioned policy. Furthermore, our method outperforms all baselines in the roal-world robot manipulation experiments and can better generalize to unseen tasks.
2 Preliminaries
Problem Statement. We consider pixel-input language-conditioned robotic control under the imitation learning setting. We denote robotic dataset Drobot (f1,52....Ga comprising demonstrations. The tth demonstration 4. (4,,4, TR) contains natural language instruction 4,

---

sequence of pixel inputs I, and a robot trajectory consisted of a sequence of robot poses p.-r. However, since collecting robotic data is risky and costly, the scale of Dobot will be limited. We therefore also consider the RGB video dataset Duides which is easily accessible on the Internet. An instance wideo can be represent as Cj (T,)- Although Prideo lacks robot action data, our proposed PAD framework enables co-training both robotic dataset robot and video dataset D Diide, leveraging the large-scale Dideo data to enhance visual policy learning.
Latent Difftsioo models. The core idea of diffusion models is to continuously add Gaussian noise to make a sample a Gaussian and leverage the denoising process for generating data [ Let 2g = e denote latent sample encoded from real data. The noising process gradually adds normal Gaussian noise (N) to zu over T steps, resulting in a set of noisy samples Z [ E [1. which is equivalent to sampling from the following distribution: ((xx(2+-1) = A((z:
Va124-1, [ where {az]t E [1, []} are predefined syper-parameters that control the amplitude of the noise. Let āi = II 0x, and according to DDPM (5) 2t can be directly obtained by adding a Gaussian noise 6, to z0: 2 = varzo + vI - Cipee Further, the denoising process starts with the most noisy latent sample 2T, and progressively reduces the noise to recover the real sample 20 with condition It is based on a variational approximation of the probabilities g(=-1122 c) given by:
a,
(= =
114 c), (1 āt ā p)(= (z+ vI areo (z,, t, vā).
(I)
The noise estimator c0(2) t, c) is implemented as a neural network and is trained to approximate the gradient of the log-density of the distribution of noisy data [2).. that is:
c c) = -vI āe V log
(2)
3 PAD: Prediction with Action via Joint Denoising Process
3.1 Overview of PAD
Multi-modalities Generation. In this section, we introduce our PAD framework, which concurrently predicts future frames and actions within a joint latent denoising process. We primarily focus on the RGB image modality M, and the robot action modality MA. Each robot action can be characterized by a robot pose that includes the position and rotation of the end-effector, as well as the gripper status. Notably, this framework can easily extend extra modalities Mg. For instance, we additionally incorporate the depth image modality in the experiment part, which provides a more accurate measure of distances.
Conditional Generation. In the proposed PAD framework, predictions and actions are conditioned on multi-modality current observations, which include RGB images с₁, robot pose CA, additional depth map Cæ (in Real- World tasks), and natural language instruction text l. The framework simul- taneously outputs the corresponding future predictions x,, FE and robot action FA. Rather than predicting a single future step, PAD can forecast & future steps x7 1:k 1:k .E.. which can be viewed as k step planning of the robot. Only the first predicted action 5 is executed by the robot, which then triggers a new prediction cycle. This iterative prediction and execution process allows the robot to continuously plan and act in a closed-loop manner. The implementation details are discussed further in the subsequent section.
3.2 Model Architectures
Given that the original data may come in various formats with high dimensions, we first map all modalities to a latent space and undertake a latent diffusion process. Following the process in (20), the RGB image is initially processed through pre-trained, frozen VAE[23 encoder ε to derive the latent representation ε (z This latent representation is then converted into a sequence of tokens to with embedding size h via tokenizer. Similarly, the robot pose AA is encoded using a Multi-Layer Perceptron (MLP) 24 into e (x and linearly transformed into tokens LA with the same embedding size h. If available, the depth image is downsampled and tokenized into Lg- The natural language instruction is processed through a frozen CLIP encoder 25 to produce the text embedding Ct.
Model Input Process.
3

---

Future
Predictions
Modilation
Feed Forward
Latent Space
Modulation
Joint Denolsing
Diffusion Transformer Block
suult--hadd Attention
Noised
Latent Space
red block"
Encoder
Encoder
Medistation
Observation
+
CLIP
image
Robot pose
Other modalities
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 masked-attention mechanism
Diffusion Transformer (DiT) Backbone. We have adopted the Diffusion Transformer (DiT) [20] as our model backbone, which offers several advantages over the U-net backbone commonly used previous works [8] 17). Notably, the DiT architecture efficiently integrates various modalities via the self-attention mechanism. Inputs such as RGB images, robot poses, and additional data are transformed into token sequences L,, Lt,,, with lengths T7, TA TE, respectively. These token sequences from different modalities concatenated and undergo joint latent denoising process
Furthermore, the DiT architecture is adaptable missing modalities. For example, in the case of video dataset that lacks robot actions, the input to DiT only comprises the image tokens tz. We simply extend the token sequence to the combined length T, + TA + TE and introduce an attention mask in the self-attention block to exclude the padding tokens. Only effective predictions are retained the output, discarding any padded parts. brief illustration of the whole process depicted on the right side of Figure This design choice enables PAD to be concurrently trained on both RGB-only video datasets and robotic datasets.
Joint Conditional Generation.
We initialize future observations as white noise and aim to re- construct future observation frames and desired robot action, conditioning current observations er, CA, cg. Following similar strategy as in (26), we concatenate conditional latent and noise latent in the channel dimension. Specifically, after obtaining encoded latent E E E((A) concatenate these latent with noise to obtain conditioned noised latent L, (c1), LA (=A())), 2, Lg (cE)), ), 2/²). For instance, the encoded latent (c has shape of e d then 2, would have shape of ke x d x d to represent & future frames, resulting in the final latent L, having a shape of (k + 1)e x d d. The other modalities undergo similar process.
we
After concatenating the latent, these conditioned noisy latent from different modalities are tokenized into sequences of tokens t A. Eg with the same embedding size. The tokenization image latent L, follows a patchify process same to (20), while the tokenization of robot pose employs a simple linear projection. Finally, these tokens are fed into multiple layers of DiT predict the latent representation of future frames. An illustration of the overall process can be found in Figure
3.3 Training Process
Initialization. Following the initialization process in [5], we also initialize the PAD weights from the DiT model pre-trained on ImageNet for the image generation task conditioned on class [20]. However, we can not directly load the model since we have missing or incompatible model parameters. We discard the label embedding layers in DiT and zero-initialize new layers for text embedding, we replicate the weight of the image latent tokenizer for + times to encode the stacked latent, and the encoder and decoder for robot state are also zero-initialized.

---

Training Objective. The diffusion process adds noise to the target encoded latent (1))) and results in noised latent AI, = { 2 24 2. We train the PAD model simultaneously predict the noise e', e e added the sample data, conditioned on current observations CL,A,E = (c1, CA, ) and instructions 1. This denoiser is trained with the DDPM S
loss:
( E E. [||e* - (2,, (,,,,,,,,111))
where 8 E [I, }} represents different types of input modalities. The denoising loss Caifs aims to maximize the evidence lower bound (ELBO) [5] while approximating the conditional distribution p(E,(x))(C,I) We jointly minimize the following latent diffusion objectives and use hyperparameters λA, Ag to balance the prediction loss between different modalities. Formally, the final training objective is given by:
[(0) = Ar + AA + Xg
4 Experiments
In this section, we conduct a series of experiments on the simulated Metaworld Benchmark 21 and a real-world table manipulation suite, utilizing our joint prediction framework. We aim to answer the following questions:
Can PAD enhance visual policy learning through joint prediction and action with limited robotic data?
Can PAD benefit from co-training on large-scale internet video datasets and better generalize to unseen tasks?
Can scaling up computational resources improve PAD's performance?
Drawer
Mick/Mace
Expert data
Unseen
Button
Cable-route
(1) Bridge video data
(2) All tasks in MetaWorld
(3) Real-world Panda Manipulation
Figure 4: We learn single vasion-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 policy 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.
4.1 Environmental Setups and Baselines
serves as widely used benchmark for robotic manipulation, accommo- dating both low-level feature and pixel input modalities. Previous studies that utilized pixel input generally developed separate task-specific policies for each of the 50 tasks. In contrast, our approach demonstrates a significant advancement by employing single text-conditioned visual policy to address all 50 tasks, within data-cfficient imitation learning framework. We collected 50 trajectories per task, consistently using the "corner2" camera viewpoint and recording the robot's pose with 4-dimensional states that include end-effector position and gripper status. For fair comparison, do not utilize an additional depth input in Metaworld.
Metaworid. Metaworld [2
Real-World Panda Manipulation Tasks. Our real-world experiments involve a Panda arm per- forming diverse manipulation tasks such as pressing buttons, opening drawers, routing cables, and picking and placing with various objects, as shown in Figure We follow the same hardware setup described in SERL. 27 and utilize wrist-mounted camera for pixel input [28]. The robot's poses represented by 7-dimensional vectors, including 3 end-effector positions, 3 rotation angles, and I
5

---

butlon- press
button topdown
drawer-
door-
faucet-
plate- slide
reach-
window- open
window-
door-
Easier Tasks
open
open
close
wall
close
lock
Diffusion Policy
0.92
0.16
0.36
0.32
0.76
0.72
0.60
0.36
0.12
SuSIE
0.96
0.32
0.60
0.68
0.56
0.68
0.92
0.68
0,96
0.32
RT-1
0.88
1.00
0,56
0.56
1.00
0.08
0.12
1.00
1,00
0.00
RT-2*
1.00
0.84
0,92
0.96
0.96
0.88
0.76
1.00
0.96
0.40
GR-1
1.00
0.84
1,00
0.96
0.88
1.00
1.00
1.00
0.60
PAD (ours)
1.00
0.92
1,00
1,00
0.92
0.72
1.00
0.92
1.00
0.88
PAI w/o img
1.00
0.92
1.00
0.88
0.92
0.16
0.92
1.00
1.00
0.12
PAD w/o co-train
1.00
0.92
1,00
0.92
0.92
0.48
0.92
0.96
0.96
0.72
Harder Tasks
asem- ble
busket-
ball
coffee- pull
hammer
peg- insert
pick-place -wall
shelf- place
stick-
push
stick-
pull
(50tasks)
Averuge
Diffusion Policy
0.20
0.08
0.00
0.08
0.16
0.36
0.00
0.00
0.00
0.279
SuSIE
0.40
0.24
0.32
0.04
0.24
0.24
0.08
0.16
0.16
0.410
RT-1
0.00
0.00
0.08
0.00
0.00
0.00
0.00
0.00
0.00
0.346
RT-2*
0.24
0.08
0.12
0.32
0.20
0.12
0.00
0.522
GR-1
0.64
0.08
0.52
0.48
0.24
0.48
0.28
0.60
0.44
0.574
PAD (ours)
0.88
0.84
0.80
0.80
0.68
0.92
0.72
0.96
0.88
0.725
PAD w/o img
0.04
0.44
0.40
0.48
0.16
0.36
0.16
0.24
0.16
0.436
PAD w/o co-train
0.32
0,28
0.32
0.72
0.92
0.68
0,56
0.88
0.40
0.592
Table 1: Comparisons on Metaworld benchmark. We utilize a single policy to solve all 50 tasks in Metaworld. Due to the space limit, we show a subset of tasks and the average success rate all 50 tasks. Detailed data can be found in Appendix A.4
gripper status dimension. We collected 200 trajectories per task through teleoperation using space mouse and scripted commands. Similarly, we developed a single policy capable of addressing all tasks, conditioned instructions. We also assessed the policy's 's generalization capabilities unseen tasks, as depicted in F5
Policy Training Details. As detailed in Section 3 the flexible PAD framework can be co-trained on various internet RGB video data and robotic demonstrations. In order to save computational resources and avoid the need to co-train the model from scratch in each robot domain, we first pre-train the model on internet data to establish better image prediction priors. We then adapt this pre-trained model to various robotic domains, including the simulated Metaworld and the real-world panda manipulation. Empirically, we first pretrain 200k steps on the BridgeData-v2 dataset [9), which consists of 60,000 trajectories. After this, we adapted the model to each domain, continuing training for an additional 100k steps with robotic demonstrations. The pre- training and adaptation stage requires approximately 2 days and day, utilizing 4 NVIDIA A 100 GPUs.
Moreover, we found that increasing the weight of the image prediction loss during the early adaptation stages accelerates convergence, as image priors are already established in the pre-trained models. Specifically, we maintained the image prediction loss coefficient '5 at 1.0 throughout the training period and linearly increased Aa and AÆ from 0.0 to 2.0 during the 100k training steps.
Policy Execution Details. Our policy is conditioned on the current image, C1, and the robot pose, CA- and predicts is frames of futures and actions. We configure the prediction horizon at k = 3 and set the interval between frames at for both Metaworld and real-world tasks. During policy execution, we utilize 75 steps of DDIM sampling [5] to denoise the k steps of future images, x/ and actions, .1ii These k step predictions can be viewed as k step planning and only the first predicted action, A is executed by the robot. The robot then moves to the first desired pose using a simple linear interpolation motion planner, triggering the next prediction cycle.
Comparisons. Visual policy learning has been widely explored in previous studies. In experi- ments, we opted to compare against a representative subset of prior methods that have either achieved state-of-the-an performance or share similar architecture with our methods. Notably, all methods are trained on all tasks in the domain using a single text-condition visual policy.
Diffusion Policy [16]. novel visual control policy that generates robot actions through action diffuser. We augmented the original diffusion policy model with instruction conditions to address the multi-task setting. We use the CLIP encoder (5 as instruction encoders, referring to related work 229

---

Task
Button- Press
Cable- Route
Pick (4 tasks)
Place (3 lasks)
Drawer-
Drawer- Close
Average
Open
Diffusion Policy
0.30
0.34
0.42
0.58
0.38
SuSIE
0.74
0.40
0.42
0.70
0.49
RT-1
0.72
0.52
0.40
0.34
0.44
0.50
0.43
RT-2*
0.80
0.60
0.64
0.74
0.66
0.82
0.69
PAD(ours) 0.80
0.55
0.76
0.72
0.56
0.84
0.72
(oD-D))th(ours)
0.78
0.64
0.84
0.78
0.60
0.88
0.78
Table Comparisons on real-world manipulation in-distribution tasks. PAD achieves the highest success rate. Incorporating depth modality can additionally lead to performance improvement. We evaluate each task with 50 roll-outs.
1.0
Eiffusion Policy
RT.2*
SuSE
PADIours)
0.8
RT-1
Expert Data
Easy
Middle
Hard
0.0
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.
A two-stage approach that utilizes a pre-trained image-editing model (26 to generate image goals for robotic tasks, followed by goal-conditioned low-level diffusion policy. We fine-tune the smage-editing diffusion model on the same dataset and also use the diffusion policy for coal-conditioned behavioral cloning. To ensure a fair comparison, also use the more powerful DiT framework as the image-editing model.
SuSIE [88.
RT-1 (30]. An end-to-end robot control policy that leverages FiLM-conditioned [3] EfficientNet 32 fuse visual input and language input, then followed by transformer blocks to output action.
RT-2* 33 (re-implement), A large-scale embodied model that directly fine-tunes vision- language models(VL.Ms to produce robot actions. The original RT-2 model was fine-tuned the PaLM model [34] which is not publicly available. Following the specifications outlined in the original paper, we re-implemented the RT-2 model using the InstructBlip-71 [35] backbone.
GR-1 [36]. Method that also leverages image prediction to assist policy learning. Different from PAD, they generate images and actions via auto-regressive architecture.
4.2 Main Results
Performance Analysis. In all comparisons, we train a single visual policy to address all tasks within domain, conditioned instructions. Our proposed PAD outperforms all baselines by significant margin. As shown in Table I in the Metaworld benchmark, PAD achieved an average success rate of 72.5%, which markedly surpasses the strongest baseline at 57.4%. Due to space constraints, present comparisons on a subset of tasks and report the average success rate across all tasks. A comprehensive comparison of all 50 tasks is available in Appendix A.4 Furthermore, Table shows the results in weal-world seen-tasks where PAD also attains the highest success rate.
We notice that PAD predicts more precise future images than the GR-1 method (Figure 6), likely due to the superior capabilities of diffusion models image generation tasks. These precise images may more effectively facilitate policy learning, leading to higher success rates, particularly in tasks requiring precise and accurate operation such as picking small blocks, insertion, baskethall, etc., in Metaworld.

---

Current Observation
Ground Truth Future
PAD Predicted Future
GR-1 Predicted Future
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.
Current Observation
Ground Truth
Predicted Future
"Put corn in bowl sink."
"Open cabinet."
Figure 7: Predictions on bridge datasets. PAD predicts futures align with instructions but also keeps uncertainty. the first image, PAD imagines yellow pear" instead the ground truth "banana"; in the second image, PAD imagines scenes faster than the ground truth.
Quality of the Generated Images. In addition to achieving the highest success rate in robotic control, we also present visualizations of some image prediction results in Figure 6 and Figure In the Metaworld domain, the predicted image (second row) closely resembles the ground truth image (first row), which is directly decoded from the original latent. In the Bridge domain, the predicted image aligns with the language instructions but also keeps certain level of uncertainty. These indicate that the PAD model has effectively learned the physical dynamics across these two domains.
4.3 Generalization Analysis
PAD can leverage existing physical knowledge from co-training large-scale internet video datasets to enhance its generalization capabilities across new tasks. We evaluated PAD's generalization ability in real-world panda manipulation with unseen tasks. As depicted in Figure the expert dataset comprises only colored square blocks and plates, while we introduce variety of previously unseen fruit and vegetable toys during testing. We designed tasks of three difficulty levels: easy mode, featuring 1-4 disturbance objects; a middle level with 5-15 disturbance objects; and difficult tasks that require picking previously unseen objects with 5-15 disturbances unseen backgrounds. We excluded depth input to ensure a fair comparison. As illustrated Figure 5 PAD demonstrates remarkable generalization abilities, successfully managing nut-of-distribution objects such strawberries, carrots, and eggplants, and even adapting to new backgrounds. The baseline method failed to generalize to difficult unseen tasks.
4.4 Ablation Studies
Effectiveness of RGB image prediction. We evaluated the effectiveness of our joint prediction process by modifying original model to exclude the image prediction component, namely PAD w/o image prediction. This modification leads to significant performance drops compared to PAD, as illustrated in T The absence of image prediction compromises the robot's ability to utilize the physical knowledge encoded in the image modalities, which may be crucial for robotic control.

---

Preditioon
Preeitionn
Groand
Ground Truth
Trush
with Videas
s/lo Videos
Videe
Figure 8: We observe that co-training with an internet video dataset leads to better image generation qualities, which may potentially lead to better robot action predictions.
Furthermore, predicting solely the robot pose provides only low-dimensional supervision signals, potentially leading to overfitting of the training data.
Effectiveness of Co-training with Internet RGB Video Datasets. Another major benefit of PAD the ability to co-train with large-scale anternet-sourced RGB videos, which potentially leverages the physical knowledge embedded in these web datasets enhance robotic pose prediction. We train PAD without the inclusion of web-scale RGB data, namely PAD w/o co-train. We observed a performance drop without co-train the video dataset, shown in Table Furthermore, the quality of the predicted image also decreased. For instance, as depicted in the bottom column of Figure 8 the blue block is absent in the predicted images. The quality of the predicted images markedly improves with co- co-training, which in turn indirectly enhances robot action prediction.
Compatible with Additional Modalities. As detailed in Section 3 our framework accommodates additional modal- ities owing to the adaptable DiT architectures. We incor- porate additional depth image inputs in real-world manipu- lation experiments and jointly predict future RGB images, depth images, and robot actions, denoted as PAD-depth. We observe highly aligned prediction results among dif- ferent modalities under our joint denoising framework, with some results illustrated in Figure g The inclusion of depth input enhances performance in manipulation tasks, as demonstrated in Table This improvement may stem from the precise prediction of depth information, which aids agents in discerning distance changes, thereby en- hancing performance. Moreover, our framework could be extended to predict other modalities relevant to robot con- trol, such as tactile force or point clouds, which we left for the future work.
Conditions
Predictions
RGB Image Ground Truth
PAD Output
Depth Ground Truth
PAD Output
Figure 9: PAD can flexibly train with additional modality, and simultaneously predict all the futures through joint de- noising process.
4.5 Scaling Analysis
We evaluated models across various sizes and patchify sizes (20]. as outlined in Table For example, the X L/2 model denotes the model with an XL size and 2 patchify size. Halving the image patch size will quadruple the image token lengths, which leads to higher computational costs. Our findings reveal a strong correlation between computational allocation (measured as transformer Gfiops) and the success rate (SR) of the learned policy, depicted in Figure I0] All the experiments are run Metaworld benchmarks and detailed success rates for each task are provided in Appendix A.5
5 Related Work
Pre-training for Embodied Control. Vision-language pre-trained models, encoded with physical knowledge, can enhance embodied control from multiple aspects. Primarily, the pre-trained model
can directly act policy by either generating high-level plans (37 38
39 40 41 producing direct

---

PAD-
PAD-
PAD-
PAD-
PAD-
B/2
X1.22
X1/4
X1/8
1/2
Layers
28
24
Success Rate
0.7
XL/8
Hidden size
1152
1152
1152
1024
768
XLA
Heads
16
16
16
16
12
XU2
Token length
257
65
17
257
257
Parameters
661M
661M
449M
0.5
Gilops
119.1
29.5
7.7
79.1
22.5
Average SR
72.5%
64.5%
48.2%
68.4%
62.4%
10a
10'
10'
Transformer GFLOPS
Table 3: We PAD performance under various
Figure 10:
Correlation between Trans-
sizes and computational costs.
former Gflops and policy success rate.
Many studies utilize the reasoning capabilities of pre-trained LLMs and VLMs to create hegel plans followed by motion primitives. Additionally, some approaches adapt pre-trained models to emit low-level motor control signals by adding an action head. Beyond directly acting policy, pre-trained models can also guide policy learning from multiple aspects, such as providing good representations (45 46 47 providing reward signals
low-level motor control signals 30 35
43
48 49
50
synthesizing goal images [18 51) and predicting future sequences (17].
Diffusion Models for Embodied Control. Recently, diffusion models have been adopted to tackle challenges in embodied control. subset of research focuses training conditional diffusion models that guide behavior synthesis based on desired rewards, and constraints under low dimensional state- Diffusion Policy [16] trains visual-motor policy to be conditioned on RGB observations and can better express the multimodal action distributions. However, these methods develop task-specific policies from scratch, missing out the benefits of pre-training with internet data. Another strand of research utilizes large-scale pre-trained diffusion models to perform data augmentation training data, such GenAug (56). ROSIE (57), and CACTI (88]
input setting 12
53
54
55
Future Prediction for Policy Learning. There also exist works that leverage future image predictions to assist policy learning. GR-1 1 [36] employs autoregressive transformer sequentially predict future images and actions. In contrast, we adopt joint diffusion architecture that predicts more accurate future images, potentially leading improved policy learning performance. UniPi[T7] and SuSIE [8] employ a two-stage policy learning process, initially using a diffusion generative model to forecast future image video sequences, and subsequently training inverse dynamics model or a low-level diffusion policy based on these goal images. In contrast to these two-stage methods, our approach presents distinct advantages. First, while previous methods utilize diffusion models with CNN-based U-net backbone (23) designed primarily for image generation and limited visual predictions, our method adopts diffusion transformer (DiT) architecture (20] This architecture adeptly handles multiple modalities concurrentyy via straightforward token concatenation and iitention-mass mechanisms, enabling us to jointly predict future and actions simultancously. Secondly, using images as the interface between prediction and action may not fully leverage the encoded features inside pre-trained diffusion models. The effectiveness of these two-stage methods depends heavily on the quality of the generated images. In contrast, our model integrates image generation and robotic action within unified denoising process.
Conclusion and Discussion
We present PAD, a novel framework predict future images and generate actions under joint denoising process. Moreover, PAD can co-train with internet video datasets and extend to other robotic modalities. Both simulated and real-world experiments demonstrated the efficiency of PAD.
limitation of the current method that we only tested with three types of modalities. Subsequent endeavors could extend this framework to incorporate additional robot-related input data, such as tactile information, which we believe are valuable research directions. Another limitation that the control frequency of PAD is not very high since we need to jointly denoise the images and actions. Future work can explore efficient ways to leverage image predictions, such as utilizing the intermediate latent space of predicted images rather than the nigh-dimensional pixel spaces.
10

---

References
[1] Sarah-Jayne Blakemore and Chris Frith. The role of motor contagion in the prediction of action. Neuropsychologia, 43(2):260-267, 2005.
[2] Andreja Bubic, D Yves Von Cramon, and Ricarda I Schubotz. Prediction, cognition and the brain. Frontiers in human neuroscience, 4: 4:1094, 2010.
[3] James Moore and Patrick Haggard. Awareness of action: Inference and prediction. Conscious- ness and cognition, 17(1):136 144, 2008.
[4] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances neural information processing systems, 08:6840-68511 2020.
[5] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500, 2020.
[6] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pages 8162-8171. PMLR, 2021.
[7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-28 248-255, 2009.
[8] Raghay Goyal, Samira Ebrahimi Kahou, Vincent Michalski, Joanna Materzynska, Susanne Westphal, Heuna Kim, Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz Mueller-Freitag, et al. The" something something" video database for learning and evaluating visual common sense. In Proceedings of the IEEE international conference on computer vision, pages 5842- 5850, 2017.
[9] Homer Rich Walke, Kevin Black, Tony Zhao, Quan Vuong, Chongyi Zheng, Philippe Hansen- Estruch, Andre Wang He, Vivek Myers, Moo Jin Kim, Max Du, et al. Bridgedata v2: A dataset for robot learning at scale. In Conference on Robot Learning, pages 1723-1736. PMLR, 2023.
[10] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint axXiv:2004.06/25, 1(2):3 2022.
[11] Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-imaag diffusion models. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 3836-3847, 2023.
[12] Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang. Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 6007-6017, 2023.
[13] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303, 2022.
[14] Wenyi Hong, Ming Ding, Wendi Zheng, Xinghan Liu, and Jie Tang. Cogvideo: Large-scale pretraining for text-to-video generation via transformers. arXiv preprint arXiv:2205.15868 2022.
[15] Xin Ma. Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. arXiv preprint arXiv:2401.0333 203000 2024.
[16] Cheng Chi, Siyuan Feng. Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint 203 2023.

---

[17] Yilun Du, Sherry Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. Advances in Neural Information Processing Systems, 36, 2024.
[18] Kevin Black, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, and Sergey Levine. Zero-shot robotic manipulation with pretrained image-editing diffusion models. arXiv preprint arXiv:2310.1066 2023.
[19] Anurag Ajay, Seungwook Han, Yilun Du, Shuang Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. Advances in Neural Information Processing Systems, 36, 2024.
[20] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
[21] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman. Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1100. PMLR, 2020.
[22] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022.
[23] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[24] Martin Riedmiller and A Lernen. Multi layer perceptron. Machine Learning Lab Special Lecture, University of Freiburg, 24, 2014.
(25] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning. pages 8748- 8763. PMLR, 2021.
[26] Tim Brooks, Aleksander Holynski, and Alexci A Efros. Instructpix2pix: Learning to follow image editing instructions. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18392-18402, 2023.
[27] Jianlan Luo, Zheyuan Hu, Charles Xu, You Liang Tan, Jacob Berg, Archit Sharma, Stefan Schaal, Chelsea Finn, Abhishek Gupta, and Sergey Levine. Serl: software suite for sample-efficient robotic reinforcement learning. arXiv preprint arXiv:2401. 2000:222000...113,. 2024.
[28] Jianlan Luo, Charles Xu, Fangchen Liu, Liam Tan, Zipeng Lin, Jeffrey Wu, Pieter Abbeel, and Sergey Levine. Fmb: functional manipulation benchmark for generalizable robotic learning. arXiv preprini 2rXiv:2241.008 008553, 2024.
[29] Huy Ha, Pete Florence, and Shuran Song. Scaling up and distilling down: Language-guie guided robot skill acquisition, 2023.
[30] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817, 2022.
[31] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[32] Mingxing Tan and Quoc Lc. Efficientnet: Rethinking model scaling for convolutional neural networks. International conference on machine learning, pages 6105 61 14. PMLR, 2019.
[33] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choro- manski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision-language- action models transfer web knowledge to robotic control. arXiv preprint arXiv:2307. 15818, 2023.
12

---

[34] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[35] Wenliang Dai, Junnan Li, Dongxu Li, Anthony Meng Huat Tiong, Junqi Zhao, Weisheng Wang, Boyang Li, Pascale N Fung, and Steven Hoi. Instructblip: Towards general-purpose vision-language models with instruction tuning. Advances in Neural Information Processing Systems, 36, 2024.
[36] Hongtao Wu, Ya Jing, Chilam Cheang. Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, and Tao Kong. Unleashing large-scale video generative pre-training for visual robot manipulation. arXiv preprint arXiv:2312.13153 2023.
[37] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do as i can, not as say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691, 2022.
[38] Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as policies: Language model programs for embodied control. arXiv preprint 1XXi::200900755.. 2022.
[39] Andy Zeng, Adrian Wong. Stefan Welker, Krzysztof Choromanski, Federico Tombari, Aveek Purohit, Michael Ryoo, Vikas Sindhwani, Johnny Lee, Vincent Vanhoucke, et al. So- cratic models: Composing zero-shot multimodal reasoning with language. arXiv preprint 2004v22204.00598, 2022.
[40] Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, et al. Inner monologue: Embodied reasoning through planning with language models. arXiv preprint arXiv:2207.0500 05608, 2022.
[41] Yanjiang Guo, Yen-Jen Wang. Lihan Zha, Zheyuan Jiang, and Jianyu Chen. Doremi: Grounding language model by detecting and recovering from plan-execution misalignment. arXiv preprint arXiv:2307.00 00329, 2023.
[42] Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Conference on Robot Learning, pages 991-1002. PMLR, 2022.
[43] Jianke Zhang, Yanjiang Guo, Xiaoyu Chen, Yen-Jen Wang, Yucheng Hu, Chengming Shi, and Jianyu Chen. Hirt: Enhancing robotic control with hierarchical robot transformers. In 8th Annual Conference Robot Learning.
[44] Yen-Jen Wang, Bike Zhang, Jianyu Chen, and Koushil Sreenath. Prompt a robot to walk with large language models. arXiv preprint arXiv:2309.0990 2023.
[45] Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation. arXiv preprint (XXx:22033.16011 2022.
[46] Yecheng Jason Ma. Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value-implicit pre-training, arXiv preprint arXiv:2210.00033 2022.
[47] Arjun Majumdar, Karmesh Yadav, Sergio Arnaud, Jason Ma, Claire Chen, Sneha Silwal, Aryan Jain, Vincent-Pieree Berges, Tingfan Wu, Jay Vakil, al. Where are in the search for an artificial visual cortex for embodied intelligence? Advances in Neural Information Processing Systems, 36, 2024.
[48] Yecheng Jason Ma, William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayaraman, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Fureka: Human-level reward design via coding large language models. arXiv preprint arXiv:2310.12931 2023.
13

---

[49] Linxi Fan, Guanzhi Wang. Yunfan Jiang, Ajay Mandlekar, Yuncong Yang. Haoyi Zhu, Andrew Tang, De-An Huang. Yuke Zhu, and Anima Anandkumar. Minedojo: Building open-ended embodied agents with internet-scale knowledge. Advances in Neural Information Processing Systems, 35:18343--88622 2022.
[50] Ademi Adeniji, Amber Xie, Carmelo Sferrazza, Younggyo Seo, Stephen James, and Pieter Abbeel. Language reward modulation for pretraining reinforcement learning. arXiv preprint arXiv:2308. 12270, 2023.
[51] Ivan Kapelyukh, Vitalis Vosylius, and Edward Johns. Dall-e-bot: Introducing web-scale diffusion models to roboties. IEEE Robotics and Automation Letlers, 2023.
[52] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0999 2022.
[53] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-makinng arXiv preprint arXiv:2211.15657 2022.
[54] Zhendong Wang. Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193 2022.
[55] Julen Urain, Niklas Funk, Jan Peters, and Georgia Chalvatzaki. Se (3)-diffusiooffilld:: Learning smooth cost functions for joint grasp and motion optimization through diffusion. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 5923-5930. IEEE, 2023.
[56] Zoey Chen, Sho Kiami, Abhishek Gupta, and Vikash Kumar. Genaug: Retargeting behaviors to unseen situations via generative augmentation. arXiv preprint (tXXX:::33000667., 2023.
[57] Tianhe Yu, Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Su Wang, Jaspiar Singh, Clayton Tan, Jodilyn Peralta, Brian Ichter, et al. Scaling robot learning with semantically imagined experience. arXiv preprint 2rXiv:2302.11550 2023.
[58] Zhao Mandi, Homanga Bharadhwaj, Vincent Moens, Shuran Song, Aravind Rajeswaran, and Vikash Kumar. Cacti: A framework for scalable multi-task multi-scene visual imitation learning. arXiv preprint 1rxiv:2212.05711, 2022.
14

---

A
Appendix
Videos PAD can be found hittesss/stes..gooo...o////////s
A.1
Additional Implementation Details of PAD
A.1.1 Input Encoder and Output Decoders
Image Encoder and Tokenizer. The image encoder is a frozen VAE encoder same as [20]. Take PAD-XL/2 model for example, the encoded latent space for 256 256 image is 32 x 32 and patchify into (32/2) * (32/2) 256 patches, which then are tokenized into 256 tokens.
Robot action Encoder and Tokenizer. The robot action is concatenated into a vector and passed into MLP layers, we predict steps of the future each with 7-dimensional poses, which totally consists (k + 1) dimensional vectors ((k + 1) 4 in Metaworld), and then this vector is tokenized into 1 token.
Depth image Encoder and Tokenizer (If presented). We directly down-sample the depth image to size of 32 and follow the same patchfy process as the RGB image. The patch size to 8. The patchfy for depth image resulted in (32/8) * (32/8) = 16 patches, which then are tokenized into 16 tokens.
The decoder part mainly inverses the encoder part. The decoder process first reconstructs the future latent from the token output by DiT, then adopts the corresponding decoder recover the original samples in each modality.
Output Decoder.
A.2
Additional Implementation Details of Baselines
The RT-1
baseline is based official
implementation
https://github.cco
google Google-resesee wwweerrhheobbteeeeeeeeeeeeeeeeee
transformer
The Diffusion Policy baseline is based on policy and we follow http:s////.fffoocccccm////aaaaar condition.
http:///giihhbbcoo...oo
add language
The RT-2 baseline is re-implemented by ourselves. We use the InstructBlip-vicuaa- model as backbone http:///ww.t.to.oooo.coco.//a/aa
The GR-1 baseline is built on http:///ttttubb.ooommmttttteettt Since access the pretraining dataset in the original paper, we initialize the model with the author's open-source checkpoint.
A.2.1 Additional Model Training Details
PAD- XL/2
PAD-
PAD- XL/8
PAD-
PAD-
XL/4
L/2
B/2
Layers Hidden size
28
28
28
24
12
1152
1152
1152
1024
768
Heads
16
16
16
16
12
Parameters
661M
661M
661M
449M
128M
Gflops
119.1
29.5
7.7
79.1
Learning Rate
lc-4
lc-4
1c-4
1c-4
1c-4
Batch size
256
256
256
256
256
Input image shape
256+256
256*256
256+256
256+256
256*256
Input noised latent
32#32*(494)
32*32*(4*4)
32*32*(4*4)
32*32+(484)
32*32*(4*4)
Patchify size
2+2
484
8*g
2*2
2
Image token size
256
64
16
256
256
Input robot action shape
1*28
1 1*28
1*28
1*28
1"28
Action token size
Total token size
257
65
17
257
Table 4: Models with various size and computational cost.
15

---

A.3 Real world Experiment Details
Expert data collection. We collected data on categories of tasks, including button press, cable- route, pick and place, and drawer open/close. For the pick task, we include 4 colors of blocks, For the place task, we include colors of plate. We randomly placed 1-5 objects on the table and asked the robot to pick/place certain objects conditioned on instruction. Some samples of expert demonstrations are visualized in Figure
Generalization Test Task Samples. We test the generalization ability of learned policy under numerous unseen objects. The unseen task is much more complicated than expert tasks, shown in Figure I2 For convenience, videos of PAD can be found at sulhtps:/// //sstt//s.te.e./eee google.com/ view pad paper
Figure 11: Samples of tasks that collected demonstrations.
Figure 12: Samples of unseen tasks used for generalization test.
16

---

A.4
Details Baselines and Ablations in Metaworld
Task
Diffusion Policy
SuSIE
RT-1
RT-2"
GR-1
PAD
PAD w/o img
PAD
cotrain
assembly-v2
0.20
0.40
0.00
0.24
0.88
0.04
0.32
baskethall-v2
0.08
0.24
0.00
0.08
0.08
0.84
0.44
0.28
putton-press-topllowny2
0.16
0.32
1.00
0.84
1.00
0.92
0.92
0,92
buttou-press-oo--aalll22
0.00
0.24
0.60
0.88
0.80
0.96
0.64
0.84
button-press-v2
0.92
0.88
1.00
1.00
1.00
1.00
putton-press-wall-22
0.44
0.60
1.00
0.44
0.48
0.68
0.48
0.60
coffee-button-22
0.64
0.80
1.00
1.00
1.00
1.00
1.00
1.00
coffeo-pull-v2
0.00
0.32
0.08
0.68
0.52
0.80
0,40
0.32
coffee-push-v2
0.16
0.32
0.04
0.76
0.44
0.52
0.48
0.64
dial-turn-22
0.52
0.52
0.00
0.52
0.56
0.28
0.52
disassemble-v2
0.00
0.72
0.00
0.12
0.40
0.88
0.24
0.64
door-close-v2
0.52
0.56
1.00
1.00
1.00
1.00
1.00
0.96
door-open-v2
0.32
0.68
0.56
0.96
1.00
1.00
0.88
0.92
drawer-close-v2
0.52
0.88
1.00
0.96
1.00
0.88
0.96
drawer-open-v2
0.36
0.60
0.56
0.92
1.00
1.00
1.00
faucet-open-v2
0.64
0.36
0.88
1.00
1.00
1.00
0.40
1.00
faucet-close-v2
0.76
0.56
1.00
0.96
0.96
0.92
0.92
0.92
hammer-v2
0.08
0.04
0.00
0.20
0.48
0.80
0.48
0.72
handle-press-side--22
0.72
1.00
0.76
0.40
0.76
0.40
handle-press-v2
0.72
0.68
1.00
0.96
0.80
0.80
0.96
0.80
lever-pull-v2
0.20
0.32
0.32
0.20
0.4
0.32
0.28
0.08
peg-
igg-insert-siee-22
0.16
0.24
0.00
0.12
0.24
0.68
0.16
0.92
peg-unplin pog-unplug-side-t
-inplug-side-22
0.28
0.20
0.32
0.48
0.32
0.60
0.16
0.20
piek-out-of-hole-v2
0.16
0.00
0.24
0.52
0.00
0.36
pick-place-wall-22
0.36
0.24
0.00
0.32
0.48
0.92
0.36
0.68
pick-place-v2
0.08
0.24
0.00
0.48
0.72
0.80
0.20
0.20
plate-slide-v2
0.60
0.68
0.08
0.88
0.88
0.72
0.16
0.48
Hate-slide-side--t
0.84
1.00
0.88
0.36
0.68
plate-slide-backwwe
0.00
0,80
0.20
0.72
0.36
1.00
0.36
0.04
Hate-slide-back-siddee
0.00
0.16
0.00
0.12
0.00
0.12
0.00
0.12
soccer-v2
0.00
0.08
0.28
0.28
0.12
0.08
0.08
0.32
stick-push-v2
0.12
0.20
0.00
0.12
0.60
0.96
0.24
0.88
stick-pull-v2
0.00
0.16
0.00
0.00
0.44
0.88
0.16
0,40
push-wall-v2
0.08
0.16
0.00
0.68
0.32
0.84
0.20
0.44
push-v2
0.20
0.32
0.00
0.40
0.44
0.52
0.16
0.60
reach-wall-v2
0.72
0.92
0.12
0.76
1.00
1.00
0.92
0.92
reach-v2
0.40
0.60
0.76
0.40
1.00
1.00
0.44
1.00
shelf-place-v2
0.00
0.08
0.00
0.20
0.28
0.72
0.16
0.56
sweep-into-v2
0.08
0.16
0.12
0.36
0.16
0.40
0.24
0.28
sweep-v2
0.04
0.16
0.24
0.60
0.72
0.80
0.20
0.44
window-open-v2
0.60
0.68
1.00
1.00
1.00
0.92
1.00
0.96
window-close-v2
0.36
1.00
0.96
1.00
1.00
1.00
0.96
bin-picking-v2
0.00
0.24
0.00
0.12
0.40
0.88
0.00
0.72
box-close-v2
0.00
0.20
0.00
0.00
0.40
0.52
0.16
0.60
door-lock-v2
0.12
0.32
0.00
0.40
0.60
0.88
0.12
0.72
door-unlock-v2
0.44
0.44
0.28
0.52
0.44
0.60
0.68
0.56
hand-insert-v2
0.08
0.24
0.16
0.28
0.20
0.40
0.08
0.20
push-back-v2
0.00
0.04
0.00
0.20
0.52
0.28
0.00
0.52
Average
27.9%
41.0%
34.6%
52.2%
57.4%
72.5%
43.6%
59.2%
Table 5: Detailed success rate baselines and ablations. We did not include the handle-pul sandle-pull-side-22 and handle-pull-v2 tasks since the expert policy for these tasks in the original benchmark had low success rates. Every task is tested with 25 rollouts.
17

---

A.5 Detailed Scaling Results
Task
XL/8
B/2
L/2
XT/4
XT/2
assembly-v2
0.32
0.96
0.96
0.20
0.88
basketball-v2
0.00
0.36
0.32
0.40
0.84
button-
nutton-prcss-topoown-22
1.00
0.76
1.00
1.00
0,92
nutlon-press-lopdoon-wlllyyyy
0.72
1.00
0.90
0.90
0.96
button-press-v2
1.00
1.00
1.00
1.00
1.00
button-press-wall-22
0.56
0.60
0.40
0.68
coffee-button-v2
1.00
1.00
1.00
1.00
1.00
coffee-pull-v2
0.32
0.10
0.30
0.68
0.80
coffee-push--2
0.28
0.60
0.40
0.52
0.52
dial-tum-v2
0.28
0.40
0.56
0.20
0.56
disassemble-v2
0.80
0.92
1.00
0.88
door-close-v2
1.00
1.00
1.00
1.00
1.00
door-open-v2
0.92
1.00
1.00
1.00
1.00
drawer-close-v2
0.52
1.00
1.00
1.00
1.00
drawer-open-v2
0.88
1.00
1.00
faucct-open-v2
0,32
1.00
1.00
1.00
1.00
faucet-close-v2
0.92
0.44
0.96
1.00
0.92
hammer-v2
0.60
1.00
0.68
0.92
0.80
handle-press-sidee-22
0.40
0.48
0.72
0.60
0.40
handle-press-v2
0.92
0.84
0.96
lever-pull-v2
0.00
0.20
0.64
0.12
0.32
peg
ing-isert-siie-22
0.80
0.52
0.80
0.80
0.68
peg-ume
0.44
0.32
0.56
0.48
0.60
of--out-of--oolvv22
0.10
0.16
0.52
place-wall-v2
0,20
0.68
0.56
0.60
0.92
pick-place-v?
0.32
0.32
0.76
0.52
0.80
plate-slide-v2
0.52
0.76
0.60
0.48
0.72
plate-slide-sid--v2
0.40
0.96
0.52
0.92
0.88
plate-slide-back-22
0.36
0.12
1.00
bate-slide-bacc-sideeee
0.04
0.20
0.04
0.12
0.12
soccer-v2
0.12
0.00
0.32
0.24
0.08
stick-push-v2
1.00
1.00
1.00
1.00
0.96
stick-pull-v2
0.32
0.92
0.64
0.88
push-wall-v2
0.40
0.84
0.84
push-v2
0.28
0.32
0.92
0.64
0.52
reach-wall-v2
1.00
1.00
1.00
1.00
1.00
reach-v2
1.00
1.00
1.00
1.00
1.00
shelf-place-v2
0.72
0.72
sweep-into-v2
0.12
0.40
0.48
0,36
0.40
sweep-v2
0.48
0.64
0.88
0.80
0.80
w2ndow-open-v2
0.32
1.00
1.00
0.90
0.92
window-close-v2
0.28
1.00
1.00
0.72
1.00
bin-picking-v2
0.88
0.72
0.80
1.00
0.88
box-close-v2
0.40
0.40
0.72
0.52
0.52
door-lock-v2
0.60
0.88
0.76
0.84
0.88
door-unlock-v2
0.40
0.68
0.30
0.64
0.60
2and-insert-v2
0,00
0.16
0.40
push-back-v2
0.16
0.20
0.52
0.12
0.28
handle-pull side-s
N/A
N/A
N/A
N/A
N/A
handle-pull-v2
N/A
N/A
N/A
N/A
N/A
Average
48.2%
62.4%
68.6 68.4%
64.5%
72.5%
Table 6: Detailed success rate under different model sizes and computational allocations. We did not include the handle-pull tasks since the expert policy for these tasks in low success rates.
18

---

Instructions assemble the object shoot the basketball press the button press the button press the button press the button press the button pull back cup A.6 Instructions used in tasks
Task
assembly-v2
hasketball-v2
button-press
tition-presstoodooww--
nutton-press-topdoww-wee
button-press-v2
button-preess-alll-22
coffce-button-v2
coffce-pull-v2
coffee-push-v2
push forward cup
dial-turn-v2
turn the dial disassemble the object
disassemble-v2
door-close-v2
close the door open the door close the drawer
door-open-v2
drawer-close-v2
drawer-open-v2
open the drawer open the faucet close the faucet pick up hammer press the handle press the handle pull the lever insert the peg unplug the peg pick red object pick red object pick red object slide forward plate slide side plate slide back plate slide back plate kick the push stick pull the stick push the object pick red object reach red object reach red object place blue object sweep brown box sweep brown box open the window close the window pick green object close the box lock the door unlock the door put box into hole pull the handle pull the handle push back object
faucet-open-v2
faucet-close-v2
hammer-v2
handle-press-side--22
handle-press-v2
lever-pull-v2
peg-insert-side-22
rog-unllug-side-22
pick-out-of-hole-22
pick-place-wall-22
pick-place-v2
plate-slide-v2
plate-slide-sides
bate-slide-back-VV
place-slide-back-sidd--y
soccer-v2
stick-push-v2
stick-pull-v2
push-wall-v2
push-v2
reach-wall-v2
reach-v2
shelf-place--2
swecp-into-v2
sweep-v2
window-open--22
window-close-v2
bin-picking-v2
box-close-v2
door-lock-v2
door-unlock-v2
hand-insert-v2
handle-pull-side-22
handle-pull-v2
push-back-v2
Table 7: Instructions for each tasks in metaworld, We mainly designed the instruction based on the
task names.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We carefully claim our contribution under specific settings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We include the limitation part in the last section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
The paper does not include theoretical results
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We carefully describe the method details in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer:
[Yes]
Justification:
Code is included in supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https://nips.. cc/
publ icl gu des Codeeubmie aasionooliie icy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips. cc/ pwublic/guideen / //odiddee/Coeeeeeeee SudeSuumiiee Swwmisssiiii ssssnooliccyyyyooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
we specify details the Experiment section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide detailed data in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
22

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We describe the computational cost in Experiment section.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Question:
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We did not see any negative social impact at this moment.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the original paper that produced the code package dataset
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paper thcode com dataaetteett has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "IvcKBLWcrS", "reviewer": "Reviewer_EGNn", "summary": "The paper presents a novel framework called PAD. This framework unifies image prediction and robot action within a joint denoising process, leveraging Diffusion Transformers to integrate images and robot states. PAD supports co-training on both robotic demonstrations and large-scale video datasets and can be extended to other robotic modalities like depth images. The framework significantly improves performance on the Metaworld benchmark and real-world robot manipulation tasks, demonstrating superior generalization to unseen tasks.", "questions": "- It would be interesting to ablate the co-training data in detail with different sources to understand its effectiveness.", "limitations": "The authors briefly discussed the limitation in the last section, which looks good to me.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation of this paper is clear. The paper is easy to follow and the proposed framework makes sense to me.  The idea is simple and straightforward but effective.\n- The performance looks good compared to the RT-1/RT-2 baselines.\n- The experimental section is comprehensive, verifying performance on both sim and real world tasks. The authors also tested the model with various computation costs.\n- Implementation details are clearly described and code is provided, which makes the community easy to reproduce.", "weaknesses": "- The performance improvement on real data is somewhat marginal, which makes it questionable about its effectiveness on real data. It would also be great to showcase the ablation studies for real-world tasks to understand whether this is caused by limited model size and pretrained data (compared to RT-2).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments. \n\n---\n\n**Q1: The performance improvement on real data is somewhat marginal. Is it caused by limited model size and pre-trained data compared to RT-2?**\n\nANS: Thank you for the insightful question! We would like to clarify that although PAD marginally outperforms RT-2 in real-world seen tasks,  PAD demonstrates a significant advantage in unseen tasks, outstripping RT-2 with an average success rate improvement of 28%. This suggests that PAD’s superior generalization capabilities stem from its integration of physical knowledge via future prediction loss. Additionally, RT2  is already a very strong baseline with its huge 7B parameters while the largest PAD model contains ~670M parameters.  Although PAD shows great scaling results in our ablations,  we could not afford to train a diffusion model with billions of parameters from scratch. Instead, we had tried a smaller version of RT-2 finetuned on the BLIP2-2.7B model, which does not perform as well as its larger 7B counterparts. \n\n| Real-world Taks Success Rate | RT2-Blip2-2.7B | RT2-InstructBlip-7B | PAD(ours) |\n| ---------------------------- | -------------- | ------------------- | --------- |\n| Seen Task                    | 0.61           | 0.69                | **0.72**  |\n| Unseen task                  | 0.24           | 0.31                | **0.58**  |\n\n\n\n---\n\n**Q2: It would be interesting to ablate the co-training data in detail with different sources to understand its effectiveness.**\n\nANS: This is an interesting question! We conducted an ablation study focusing on the scale of the video dataset to assess the impact of co-training with Bridge. Our findings suggest that co-training mainly enhances performance on unseen tasks. A possible explanation is that unseen objects from unseen tasks, such as strawberries, plates, and bananas, are already appeared in the extensive Bridge datasets. Therefore, co-training significantly benefits these unseen tasks.\n\n| Real-world Taks Success Rate | PAD w/o bridge-v2 | PAD with 10% bridge-v2 | PAD(ours) |\n| ---------------------------- | ----------------- | ---------------------- | --------- |\n| Seen Task                    | 0.68              | 0.68                   | **0.72**  |\n| Unseen task                  | 0.40              | 0.48                   | **0.58**  |\n\n\n---\n\nThank you again for your time and efforts! We show our deepest appreciation for your support of our work.  We are always ready to answer your further questions!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SVq1XdsFP4", "reviewer": "Reviewer_yjvj", "summary": "This submission presents a new learning framework called PAD, which utilizes a diffusion transformer (DiT) to jointly denoise both future image frames (RGB/Depth) and generate actions together. This joint learning process yields a scalable model that achieves higher success rates compared to various other robotics transformer model baselines. The paper shows co-training with internet scale data is not only possible thanks to DiT but also brings about much higher success rates and shows strong visual generalization.", "questions": "Questions:\n- Some discussion around how this relates to world models like Dreamer or TDMPC for robotics might be important. A joint denoising process that predicts future frames and actions resembles a lot of world models. The world model like nature could be the reason why performance is higher.\n- While PAD generalizes, it is more of a generalization on the visual data side instead of manipulation right? The unseen objects probably are graspable by the same actions you would take to grab the seen objects. Regardless the visual generalization looks good.\n-", "limitations": "Limitations are very lightly mentioned at the end around how PAD only uses a few modalities. Some limitations around the data efficacy of PAD might be good to mention, especially given the expensive nature of real world robotics demonstrations (especially if you want high success rates).", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- An interesting joint denoising process is presented to generate future frames and actions. The paper further shows the possibility that PAD can be extended to more modalities if the engineering efforts are taken to do so which is great to see.\n- Nice use of the diffusion transformer architecture which enables co-training on large video datasets that do not come with action labels, providing internet-scale knowledge that improves success rate. It would be nice to see how PAD without co-training on internet data performs on the generalization benchmark. I would imagine it won't do as well as PAD with internet data but is worth adding in my opinion (not a huge issue though).\n- Strong results on generalization benchmarks and solid ablation studies on co-training show great promise about the proposed method.\n- PAD also shows good scaling results where larger models/more training yields higher success rates.", "weaknesses": "- How well does PAD predict future frames? Only a few next frames are shown in fig 7 but I am curious to see how far it can predict? In some sense tasks where there is a lot more ambiguity (e.g. objects that fall and bounce, or finding an object hidden behind something) would be helpful to help understand how PAD handles uncertainty. The current set of tested tasks are fairly straight forward in terms of predicting the next frames.\n- What are the failure modes in the real world of PAD? Why might it not succeed? This is not an easy problem to answer regardless but if possible it would be nice to address somewhere.\n- What is the data efficacy of PAD compared to past methods? How many demonstrations are needed to get the results? Can PAD use less demos (perhaps thanks to future frame predictions?).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments. \n\n---\n\n**Q1: How well does PAD predict future frames? how far PAD can predict in some sense tasks where there is a lot more ambiguity (e.g. objects that fall and bounce, or finding an object hidden behind something) would be helpful to help understand how PAD handles uncertainty.**\n\nANS: Thank you for your insightful comments! For better understanding, we have visualized additional results in the attached PDF file of the global response. With joint training on the internet video dataset, PAD shows impressive image prediction capabilities.\n\n---\n\n**Q2: What are the failure modes in the real world of PAD? Why might it not succeed? This is not an easy problem to answer regardless but if possible it would be nice to address somewhere.**\n\nANS: This is an interesting question! We summarize some failure modes based on our hundreds of rollouts as follows: (1) unexpected collisions and deformable irregular objects add difficulty to tasks, (2) using a wrist-mounted camera and randomly placing the objects can sometimes cause objects to go out of the frame during execution, leading to failure, and (3) PAD may occasionally mix up similar objects (e.g., a red apple and a red block).\n\n\n---\n\n**Q3: What is the data efficacy of PAD compared to past methods? How many demonstrations are needed to get the results? Can PAD use less demos (perhaps thanks to future frame predictions?).**\n\nANS: As detailed in the experiment setups, similar to [x], we collected 50 demonstrations for each task using Metaworld's official scripted policy. To assess the sample efficiency of the PAD method, we performed ablations with 20 demonstrations to test the efficacy of PAD.\n\n|               | RT2  | PAD(ours) |\n| :------------ | ---- | --------- |\n| 20 demos/task | 0.44 | **0.62**  |\n| 50 demos/task | 0.52 | **0.72**  |\n\nIn the low data regime, PAD also surpasses the baseline\n\n\n---\n\n**Q4: Discussion with world models like Dreamer or TDMPC. A joint denoising process that predicts future frames and actions resembles a world model could be the reason why performance is higher.**\n\nANS: Thank you for the suggestion! We completely agree that PAD is closely related to world models, as both encode physical knowledge. Dreamer and TDMPC are online model-based RL algorithms that learn world models with online collected data (without large-scale internet datasets). Previous works have also suggested that scaling up the model size under an online RL setting can be unstable. In contrast, we ensemble a world model into a unified DiT model, which is readily scalable and co-trained with internet datasets. In experiments, we observed good scaling performance with the proposed PAD model.\n\n\n---\n\n**Q5: While PAD generalizes, it is more of a generalization on the visual data side instead of manipulation right? The unseen objects probably are graspable by the same actions you would take to grab the seen objects. Regardless the visual generalization looks good.**\n\nANS: Yes, you are correct! The generalization is more on the visual side, such as with many distractions, unseen objects, and unseen backgrounds. We acknowledge that the current model cannot perform brand-new skills that are not present in the training data. We will add this to the limitations and future work section. Scaling the robotic datasets may lead to the emergence of new skills, which are interesting research directions.\n\n\n---\n\nThank you again for your insightful and constructive comments! We hope our response above can help you better understand our paper and we are always ready to answer your further questions!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "inQoCVZFLj", "reviewer": "Reviewer_Hcpx", "summary": "- This work proposes a new method for language conditional imitation learning for robotics. The core idea is to combine image diffusion and policy diffusion to simultaneously predict future image observations and actions using a latent diffusion transformer and DDPM.\n- The diffusion process is conditioned on the recent image observations, the current robot pose, the language instruction, and for real-word experiments also on a depth map. The output of the diffusion process is the prediction of future image observations, depth observations, if available, and actions for K time steps. The diffusion process is repeated after executing a single action in the environment by conditioning on the updated history.\n- The input and output modalities have different dimensionalities which is handled by projecting all modalities into a latent space where the latent diffusion process is applied. A VAE is used for projecting the images, an MLP for the actions, a CLIP encoder for the language instructions, and the depth maps are processed in a similar fashion to the image observations. This setup also allows training the model on video data where no actions are available.\n- Experiments are conducted on the simulated Meta-World benchmark and in the real-world with a Panda arm. The proposed method is compared against a multi-task variant of diffusion policy, a modified version of SuSIE, RT-1, and a re-implementation of RT-2 built on top of InstructBlip7B.", "questions": "- Line 166: How were the Meta-World training trajectories collected? How does this compare to existing literature? Are there existing public datasets of demonstration trajectories for Meta-World?\n- How were the hyperparameters tuned for the proposed method, for the ablations, and for the baselines?\n- How do the architectures of the proposed method and the baselines compare? Why does the proposed method already significantly outperform diffusion policy even without image diffusion?\n- Line 219: How were the tasks that are shown in Table 1 selected?\n- Line 21: Do the videos on the paper website show the robot at real-time speed or are the videos accelerated?\n- How do the proposed methods and the baselines compare in terms of latency?\n- Are code and training data going to be released publicly?\n\nUPDATES AFTER REBUTTAL:\n- Presentation: 2 -> 3\n- Contribution: 2 -> 3\n- Score: 4 -> 6", "limitations": "As mentioned above, the paper lacks a discussion of latency considerations which might limit the practical usefulness of this approach.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Originality:\n- The paper proposes an interesting combination of image diffusion and policy diffusion to leverage transfer between the two tasks. Using latent diffusion with a transformer so that all modalities can be processed in a shared latent space while also being able to train on video data without actions is a nice and elegant idea.\n\nQuality:\n- Experiments on Meta-World and with a real-world Panda arm on both seen and unseen tasks are reasonable for validating the proposed method.\n- The selected baselines (a multi-task variant of diffusion policy, a modified version of SuSIE, RT-1, and a re-implementation of RT-2 built on top of InstructBlip7B) are recent and strong.\n- For the diffusion policy baseline, a pretrained CLIP encoder was used as in related work which also matches language encoder used by the proposed method. Similarly, a diffusion transformer is used for SuSIE for a fairer comparison.\n- Table 1 includes two important ablations: PAD without image diffusion and PAD without on videos without actions that validate two of the main hypotheses of this work. There is also an ablation that shows that using depth maps as additional observations improves performance in real-world experiments.\n\nClarity:\n- The paper is generally well-written and clear.\n- Various hyperparameters are provided in the appendix to facilitate reproducibility.\n- The videos on the paper website provide a good sense of qualitative performance.\n\nSignificance:\n- Considering the good empirical performance and the elegance of this approach, it seems reasonable that other researchers and practitioners might use or build on this work.", "weaknesses": "Originality:\n- The paper lacks a reference to and discussion of DBC introduced in “Diffusion Model-Augmented Behavioral Cloning” by Chen et al. 2024 which first appeared on arXiv in February 2023. Similar to this work, Chen et al. 2024 learn a diffusion model over state-action pairs. Instead of using this directly for parameterising the policy, however, the diffusion model is used for computing an auxiliary loss to guide the policy. This therefore avoids the need to run the computationally expensive diffusion model at test time.\n- Several statements are not supported by references:\n    - Line 163: “Previous studies that utilized pixel input generally developed separate task-specific policies for each of the 50 tasks“.\n    - Lines 285-286: “Many studies utilize the reasoning capabilities of pre-trained LLMs and VLMs to create high-level plans followed by motion primitives. \n    - Lines 286-288: “Additionally, some approaches adapt pre-trained models to emit low-level motor control signals by adding an action head.”\n\nQuality:\n- Considering the similarities of the proposed method to DBC, it would be good to have DBC as an additional baseline to differentiate this work from DBC more clearly.\n- The paper lacks a discussion of previous works that benchmark imitation learning methods on Meta-World and why this particular evaluation methodology was chosen. For example, “PRISE: LLM-Style Sequence Compression for Learning Temporal Action Abstractions in Control” also benchmarks imitation learning approaches on Meta-World, reporting 5-shot performance on five previously unseen tasks. Similarly, the method could also be evaluated on “LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning” which provides tele-operated human demonstrations and several baselines.\n- The appendix provides a breakdown of performance across individual tasks and ablations for different model sizes, but all experiments are run with only a single random seed and no error bars or similar are reported.\n- It is unclear how the hyperparameters were tuned for the proposed method, for the ablations, and for the baselines. Similarly, it would be useful to have a discussion how the architectures of the proposed method and the baselines compare to better understand the fairness of the comparisons. Table 1 indicates that the proposed method without image diffusion already outperforms three of the four baselines, including diffusion policy. It seems like there are other differences between the proposed method and diffusion policy that account for some of the difference in performance.\n- The paper lacks a discussion of how this method compares to the baselines in terms of latency. It is also not clear whether the videos at the provided URL run in real-time. Running image diffusion at every time step seems computationally expensive which might have practical implications.\n\nClarity:\n- None beyond what is mentioned above.\n\nSignificance:\n- The latency of running image diffusion at every time step is a concern. This might be quite slow and limit the adoption of this work.\n- It is not quite clear from the NeurIPS paper checklist whether the code and the datasets are going to be publicly released which would increase the impact of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments. \n\n---\n\n**Q1: Discussion and comparisons with DBC paper which learn a diffusion model over state-action pairs as auxiliary loss**\n\nANS: Thank you for your constructive suggestion! DBC adds the state-action distribution diffusion loss to help policy learning while we use future prediction loss to enhance policy. Despite the lack of open-sourced code from the DBC authors, and their focus on low-dimensional single-task settings using smaller neural networks, we implemented their state-action modeling auxiliary loss within our DiT framework, ensuring comparable parameter sizes. (Detailed implementation framework can be found in global pdf.) In general, we have much larger experiment scales in terms of task numbers, task difficulty, input complexity, and model size. Our results demonstrate that PAD outperforms DBC with a clear margin. We hypothesize this is due to that future prediction loss can better guide policy learning compared to the DBC's state-action distribution loss.\n\n|              | PAD w/o image prediction | DBC            | PAD(ours)      |\n| ------------ | ------------------------ | -------------- | -------------- |\n| Success rate | $$43.4\\pm1.7$$           | $$47.6\\pm3.0$$ | $$72.3\\pm1.8$$ |\n\n---\n\n**Q2: Several statements are not supported by references.**\n\nANS: Thank you for pointing out the missing reference! We cite [1,2,3,6] for Line 163 which describes single-tack learning algorithms, cite [7,8] for 285-286 which mentions the VLM high-level plan, and cite [9] for 286-288 mentions the pre-trained model with action head.\n\n\n---\n\n**Q3: Discussion of previous works that benchmark imitation learning methods on Meta-World and why this particular evaluation was chosen. e.g., PRISE reported 5-shot performance on 5 unseen tasks.** \n\nANS:   Numerous works in many categories evaluate imitation algorithms on Metaworld benchmark with limited demonstrations, such as diffusion-based algorithms diffusion policy[1], 3D-DP[2], video representation learning algorithms R3M[3], VC-1[4], and language modeling style policy Embodied-GPT[5], ACT[6], etc. Similar to these works, we try to evaluate PAD's imitation capability under limited robotic data and multi-task settings.\n\nPRISE tests the model's few-shot adaptation capability as you said. We follow the exact same pipeline to test the adaptation ability of our PAD model. The results are shown below which further verify the sample efficiency of our method thanks to joint prediction.\n\n|                                         | PRISE          | PAD            |\n| --------------------------------------- | -------------- | -------------- |\n| 5-shot success rate on 5 specific tasks | $$66.8\\pm1.8$$ | $$74.2\\pm2.3$$ |\n\n\n\n---\n\n**Q4: How were the Meta-World training trajectories collected? Are there existing public datasets of demonstration?**\n\nANS: We utilized the scripted expert policies provided by the official Meta-World to collect 50 demonstrations per task, in line with prior studies [1,2,6,12].\n\n---\n\n**Q5: How were the tasks that are shown in Table 1 selected?**\n\nANS: The tasks in Meta-World are categorized into different difficulty levels according to [1]. We selected representative tasks from each level to display due to space constraints, and also the average success rate. \n\n---\n\n**Q6: About random seeds and error bars of experiments.**\n\nANS: Limited by our computational resources, and given the breadth of ablation studies conducted (ablations on image prediction/co-training, depth input,  and various model sizes), we were unable to run many seeds at the time of submission. We promise to run three seeds for each experiment and update the results in the paper. Moreover, we observed diffusion training process is extremely stable with low variance, corroborating findings from the diffusion policy paper[1].\n\n---\n\n**Q7: How the hyperparameters were tuned. The architectural differences between the proposed method and diffusion policy**\n\nANS: Yes, there do exist architectural differences.  Unlike the official diffusion policy, which utilizes CNN-based image encoders and cross-attention conditions, our PAD is built upon the DiT[10] which pathify images into tokens and uses feature-wise modulation for condition. We have integrated additional modules into DiT to support language and multi-modal inputs/outputs, using the same training hyperparameters as the DiT paper. As discussed in previous works[11], DiT is a highly optimized architecture that may also contribute to our performance gain.\n\n---\n\n**Q8: Do the videos on the paper website show the robot at real-time speed?**\n\nANS: The old videos are recorded inside the algorithm pipeline which has some bias. We updated some videos recorded by third-view mobile phones which are truly real-time.\n\n---\n\n**Q9: The latency of running image diffusion at every time step is a concern. How do the proposed methods and the baselines compare in terms of latency?**\n\nANS: We total agree with you that diffusion-based policies typically exhibit lower control frequencies due to the multiple denoising steps involved.  In real-world robotic control, our model operates at 1-1.5Hz using 50-75 denoising steps, compared to the RT-2 model's 3Hz and the diffusion policy's 3-4 Hz, which works fine in our tested tasks. To further improve the diffusion speed, we can use recent more advanced diffusion accelerated samplers[13] and execute multiple prediction steps as [1].\n\n---\n\n**Q10: Are code and training data going to be released publicly?**\n\nANS: We have included an initial code version in the supplementary materials (though a little bit dirty). We will open-source the refined code, dataset, and checkpoints upon acceptance. We are very glad to make contributions to the community!\n\n---\n\nThank you again for your time and efforts!  We are always ready to answer your further questions!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TMjDxKzQUV", "reviewer": "Reviewer_GDHG", "summary": "The submission aims to enhance policy learning through co-training with image prediction. The intuition is that image prediction and robot action are highly correlated since they share the same underlying physical dynamics. Thus, a diffusion model that generates both future images and actions may generalize better to the task distribution.\n\nAs a solution, the paper proposes a diffusion transformer-based solution that takes as input a noisy sequence of future actions and observations and predicts the corresponding noise, in order to return a clean action trajectory and observation sequence. This enables co-training with actionless videos, in which case the action denoising loss is deactivated.\n\nThe results show improved language-conditioned multi-task learning on MetaWorld, as well as generalization to different real-world scenarios. Notably, in the real world, the proposed policy is able to generalize to more difficult task variations than the ones seen during training.", "questions": "Overall, while this submission has merits, its poor evaluation setup, as well as the absence of comparisons and discussions against important prior work makes me reluctant towards acceptance. I would increase my score if:\n\n* The writing changes to include discussion on previous works that had the same high-level idea and focus on the differences with those.\n\n* The evaluation becomes stronger by adding baselines and testing on a more challenging benchmark. On the latter, one suggestion is to evaluate on CALVIN, so as to enable direct comparison to [GR-1] and SuSIE.\n\n\n--------------\nPost-rebuttal: score updated from 2 to 4.", "limitations": "Addressed", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. (Presentation) The paper is very clearly written and is very easy to follow and understand. I want to additionally endorse the tone of this paper, which is humble and moderate, without overpresenting or overlaiming.\n\n2. (Contribution) As a technical contribution, the paper shares interesting architectural details on how to repurpose DiT for joint action and future image prediction.\n\n3. (Soundness) The ablations show the importance of co-training with image prediction and the additional benefit of training on video corpora.\n\nI believe that overall the paper contributes to the idea that predicting the future and acting is a promising path for generalizable robot learning, as opposed to action-centric policy learning. Although the idea is not original per se, the paper provides more evidence for it.", "weaknesses": "1. (Contribution) The motivation of the paper is that first, the ability to predict future images is coupled with the ability to predict actions (idea); second, diffusion models have shown good results on both tasks, so it is reasonable to combine them for joint prediction (implementation). To the best of my knowledge, the implementation part is novel, but the idea part is not.\nSpecifically, recent works such as [GR-1] have also argued and supported that jointly predicting future frames and actions is helpful, however there is no discussion and comparison against such an approach.\n\n2. (Soundness) A major concern about this submission is the inadequate evaluation protocol. Specifically, the benchmark used is very simple and important baselines are missing. In more detail:\n\n* The baselines considered in this paper are not carefully selected and important strong baselines are missing. On the first part, it’s not surprising that 2D Diffusion Policy is not weak, but probably the point is to show that using DiT is an improvement over it. RT-1 is also a rather weak baseline without the large data support, as shown in [GR-1]. Moreover, comparing against RT-1 and RT-2 has little to offer as their training paradigm is simply behavior cloning with regression. It would be more useful to compare against other diffusion policies or video prediction models, since the paper's main argument does not stem from its architectural improvements on Transformer-based policies.\nOn the second part, the main argument of the paper is that predicting both actions and future images helps control. A similar argument is also presented in [GR-1], although their architecture is not diffusion-based but autoregressive. It could be that diffusion (or the proposed architecture in general) is indeed a better solution, but to conclude that, a direct comparison with [GR-1] should be conducted.\n\n* While the end goal is to have robots operating in the real world, it is unsafe to compare and draw conclusions about models’ performance with real-world experiments only. Simulation benchmarks offer some proxy to evaluate performance and compare methods in identical setups. The paper does a good job of presenting both simulation and real-world results. While the community hasn’t concluded on a robotics benchmark that covers many cases and tasks, MetaWorld is definitely one of the simplest benchmarks, without distractors, visual variations or the necessity of modeling geometry. Multi-tasking results on MetaWorld are not very impressive since the task can be usually inferred by just looking at the scene. On that aspect, training on videos from the web or external data (such as BridgeData-v2) to test on MetaWorld feels rather disconnected. It would be more valuable to show results on benchmarks like CALVIN, LIBERO or RLBench, that test more generalization factors and offer more challenging setups with stronger baselines. Specifically, SuSIE, which is a useful baseline as it also predicts future images, as well as [GR-1], which should be the main baseline, both report results on CALVIN. Therefore, evaluating on CALVIN would give direct answers on these comparisons.\n\n\n3. (Presentation) Although the clarity of presentation is remarkable, the paper lacks a clear take-home message. If this is that hoint action-frame prediction is helpful, then it reduces novelty, since there is already evidence on that. Another suggestion is the following. There are some works (such as the cited Diffuser [48]) that jointly denoise states and actions. This is a relevant background for this work, which denoises observations and actions. The proposed approach is more general since states cannot be accessed in the real world. It would be nice to see some discussion on the connection between prior work and this work, as it would provide additional motivation.\n\n[GR-1] Unleashing Large-Scale Video Generative Pre-training for Visual Robot Manipulation, ICLR24", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments. **Updated Figures can be found in the PDF attached to the global response.**\n\n---\n\n**Q1: About the contribution and novelty:  the implementation part is novel, but the idea part is not. Specifically, recent works such as [GR-1] have also argued and supported that jointly predicting future frames and actions is helpful.**\n\nANS: We fully agree that integrating video prediction with action, and more broadly, learning enhanced representations from videos to improve policy learning, is a well-established area of research. In this paper, we focus on **exploring joint multi-modal predictions and actions within a unified diffusion architecture.** This approach, which we all recognize as novel, differentiates our work from existing methods.\nGR1 utilizes an autoregressive model for predicting images and actions. However, recent advancements have shown that **diffusion models outperform autoregressive methods in generating images and videos** [1,2]. Despite these advancements, performing joint denoising processes using the typical U-Net architecture, as seen in prior studies like SUSIE and uniPi, poses significant challenges. Our approach overcomes these challenges by integrating predictions and actions into advanced Diffusion Transformers (DiT), which effectively encode physical knowledge from diverse datasets. **This solution has been recognized as both elegant and nice by all three other reviewers.** A more detailed experimental comparison with GR1 can be found in the subsequent Q&A sections.\n\n---\n\n**Q2: About the meta world benchmark and Calvin benchmarks**\n\nANS: We would like to say that numerous previous works [] perform experiments on Metaworld and these two benchmarks present their own unique challenges. Metaworld offers a much more diverse array of tasks requiring precise manipulation and the ability to handle multiple tasks simultaneously. In contrast, Calvin focuses primarily on instruction-following skills, involving fewer objects (just three colors of blocks, one door, one drawer, and one light). \n\nWe observed that previous methodologies have already achieved success rates nearing or exceeding 90% on Calvin benchmarks, with most failures occurring in corner cases. Given the limited time available for rebuttal, it would be exceedingly difficult for us to establish a new environment and adjust hyperparameters to surpass this 90% success rate. Furthermore, Calvin shares many similarities with **our real-world experiments which involve even more complex generalization tasks** including brand-new objects and backgrounds. Considering these factors and the tight timeline, we opted to compare the official GR1 implementation with PAD on the Metaworld and the real world. (The real-world experiment is running since it takes time to deploy and adjust real systems.)\n\n**Q3: Detailed experimental comparison with GR-1 baseline**\n\nANS: Following the GR-1 open-sourced code, we initialize the GR-1 model with their pre-trained checkpoint and perform data-efficient imitation learning on the multitask Metaworld. The comparisons are shown below. Firstly, we observed that the **image prediction quality of PAD significantly surpasses that of GR-1**, likely due to the superior capabilities of diffusion models in image generation tasks. As depicted in the global PDF, the visual prediction results from GR-1 appear blurry, while PAD produces detailed and precise imagery. (Notably, the images presented in the original GR-1 paper were also blurry.) We hypothesize that the **high-quality images generated by PAD can more effectively facilitate policy learning**, leading to higher success rates, particularly in tasks requiring precise and accurate operation such as picking small blocks, insertion, basketball, etc., in Metaworld. We hope this detailed analysis provides insight into why our exploration of joint prediction and action under diffusion architecture is both meaningful and valuable.\n\n| MetaWorld-50tasks | GR-1 | PAD  |\n| ----------------- | ---- | ---- |\n| Success Rate      | 54.6 | 72.5 |\n\n---\n\n**Q4: The baselines considered in this paper are not carefully selected. It’s not surprising Diffusion Policy is weak. RT1RT2 only utilizes simple BC loss**.\n\nANS: We respectfully present an alternative perspective. We believe these baselines are both recent and robust, as suggested by Reviewer Hcpx. Diffusion policy is a sample-efficient imitation learning algorithm recognized as a significant breakthrough in robotic policy learning. Although RT1 and RT2 are trained only with BC loss, they are built upon large vision-language models that have already been trained on extensive datasets, demonstrating impressive multi-task capabilities. We compare these baselines to verify the sample efficiency and multi-task learning ability of PAD.\n\n---\n\n**Q5: Training on videos from the web or external data (such as BridgeData-v2) to test on simulated MetaWorld feels disconnected.**\n\nANS: Training with videos and testing in simulated environments such as Metaworld is widely adopted in prior research. We recognize a substantial gap between real-world videos and simulations, which has led us to validate our methods through real-world tasks, confirming that co-training significantly improves our models' generalization capabilities.\n\nDespite the noted gap between simulations and reality, as detailed in Section 4.4, **internet videos provide valuable priors and features that enhance image prediction quality in simulations**. The high-quality predicted images may significantly improve policy learning, as it becomes easier for the model to identify robot movements from these precise predicted images.\n\n---\n\nWe hope our response above can solve your concerns and show the improved quality of our paper. We are always ready to answer your further questions!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely appreciate the time and efforts of all reviewers and the AC in evaluating our paper. We are grateful for the insightful and constructive suggestions, which have helped us improve our work. Below, we summarize our contributions and updates. (Updated Figures are in the attached PDF.)

**1. Contributions:** 

- **Motivation:** The motivation of this paper is clear and interesting [EGNn,yjvj,Hcpx]. 
- **Method:** The joint prediction and action framework is nice and elegant [yjvj,Hcpx], with novel implementation[GDHG]. It seems reasonable that other researchers and practitioners might use or build on this work [Hcpx].
- **Experiments:** The experimental section is comprehensive including strong and recent baselines, thorough ablations, and good scaling results[EGNn,yjvj,Hcpx]. Implementation details are clearly described and code is provided, which makes the community easy to reproduce[EGNn].
- **Presentation:** The method is effective and easy to follow [EGNn,yjvj,Hcpx,GDHG]. The tone of this paper is humble and moderate, without over-presenting or overclaiming[GDHG].

**2. Modifications:**

- Following reviewer EGNn's suggestion, we have ablated the RT-2 model size and co-train dataset to better verify the effectiveness of our methods in real-world tasks.
- Following reviewer yjvj's suggestion, we conducted ablation studies with even fewer demos (20 demos per task). PAD continues to outperform the baseline in lower data regimes.
- Following reviewer Hcpx's suggestion, we reimplemented the DBC baseline within the DiT framework (Detailed in attached pdf). The results show that future prediction loss more effectively guides policy learning compared to DBC's state-action modeling loss.
- Following reviewer yjvj's suggestions, we tested PAD's 5-shot adaptation abilities using PRISE's settings. PAD surpasses the baseline in this scenario as well.
- Following reviewer GDHG's suggestions, we compared our approach with the official GR-1 baseline. Despite both employing joint future and action prediction, our diffusion-based PAD generates images of much higher quality(visualized in attached pdf) and result in a better policy success rate.
- We have corrected all writing typos and added missing references, thanks to the detailed advice provided.

Considering the limited rebuttal time and numerous experiments proposed by all reviewers, we have tried our best to conduct as many experiments as possible. We believe our comprehensive experiments in both simulated and real-world tasks adequately support our claims.

Once again, we sincerely appreciate the reviewers' time and feedback. We remain ready to address any further concerns!

Meta-review: This paper presents PAD, a framework for visual policy learning that combines image prediction and robot action within a joint denoising process using Diffusion Transformers. The approach demonstrates significant improvements over baselines on the MetaWorld benchmark and real-world robot manipulation tasks, showing better generalization to unseen scenarios.
All reviewers acknowledged the technical contribution and potential impact of the work. The main strengths highlighted were the innovative integration of image prediction and action generation, the ability to co-train on robotic demonstrations and large-scale video datasets, and strong performance on both simulated and real-world tasks. Reviewers also praised the clear presentation and comprehensive experiments.

Initial concerns focused on the novelty of the core idea, the choice of baselines and benchmarks, and the lack of comparison to some relevant prior work. The authors provided detailed responses, including additional experiments on the CALVIN benchmark and comparisons to the GR-1 baseline. These responses addressed many of the reviewers' concerns, with two reviewers explicitly improving their scores post-rebuttal.

The authors' willingness to run new experiments on CALVIN during the rebuttal period was particularly well-received, demonstrating PAD's superior sample efficiency compared to GR-1. While some reservations remain about the strength of the specific implementation and the choice of benchmarks, the overall consensus is that PAD represents a valuable contribution to the field of robot learning.

For the final version, I encourage the authors to incorporate their rebuttal explanations, particularly the additional CALVIN results, comparisons to GR-1, and discussions on the complementary nature of the MetaWorld and CALVIN benchmarks. Expanding the related work section to better contextualize PAD within the landscape of video prediction methods for robotics would also strengthen the paper.
Decision: Accept (poster)