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
encoder encoder encoder
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:13:50

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

Main paper: 
 introduction: 
 Making predictions and taking actions are critical human capabilities, allowing individuals to foresee the change of their surroundings and behave appropriately in response [1, 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 (NeurIPS 2024). underlying physical laws of the world [3] . Understanding these laws enables humans to make better predictions and actions.Recently, diffusion models [4, 5, 6] have achieved impressive success in visual generation tasks by training on extensive web-scale image and video datasets [7, 8, 9] . For example, image editing models can predict outcomes based on user instructions [10, 11, 12] , while video generation models can generate sequences of future images [13, 14, 15] , representing a 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 [17, 18, 19] have employed the image-editing model in an off-the-shelf manner by first synthesizing a goal image and subsequently learning a 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.In this paper, we introduce the Prediction with Action Diffuser (PAD), a unified policy learning framework that integrates prediction and action under the same diffusion transformer (DiT) architecture [20] . Specifically, we utilize the diffusion transformer model to seamlessly merge all modality inputs and simultaneously predict future images and actions via joint denoising, as illustrated in Figure 2(c ). Additionally, the flexible DiT backbone also allows PAD to be 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, as shown in Figure 1 . Our key contributions are:• We propose a novel policy learning framework, Prediction with Action Diffuser (PAD), to predict futures and robot actions through a joint denoising process, benefiting policy learning for robotic tasks.• The proposed PAD framework enables co-training of different datasets containing different modalities, allowing encoding rich physical knowledge from various data sources.• We outperform previous methods with a clear margin in the Metaworld benchmark, surpassing baselines with a 26.3% relative improvement in success rate using a single visuallanguage conditioned policy. Furthermore, our method outperforms all baselines in the real-world robot manipulation experiments and can better generalize to unseen tasks.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 be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
preliminaries: 
 Problem Statement. We consider pixel-input language-conditioned robotic control under the imitation learning setting. We denote a robotic dataset D robot = {ζ 1 , ζ 2 , ...ζ n } comprising n demonstrations. The i th demonstration ζ i = (I i , l i , τ i ) contains a natural language instruction l i , a sequence of pixel inputs I i , and a robot trajectory τ i consisted of a sequence of robot poses p 1:T i . However, since collecting robotic data is risky and costly, the scale of D robot will be limited. We therefore also consider the RGB video dataset D video which is easily accessible on the Internet. An instance in D video can be represent as ζ j = (I j ). Although D video lacks robot action data, our proposed PAD framework enables co-training on both robotic dataset D robot and video dataset D video , leveraging the large-scale D video data to enhance visual policy learning.Latent Diffusion 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 [4] . Let z 0 = ε(x 0 ) denote a latent sample encoded from real data. The noising process gradually adds normal Gaussian noise (N ) to z 0 over T steps, resulting in a set of noisy samples Z = {z t |t ∈ [1, T ]}, which is equivalent to sampling from the following distribution:q(z t |z t-1 ) = N (z t ; √ α t z t-1 , (1 -α t )I),where {α t |t ∈ [1, T ]} are predefined hyper-parameters that control the amplitude of the noise. Let ᾱt = t i=1 α i , and according to DDPM [5] , z t can be directly obtained by adding a Gaussian noise ϵ t to z 0 :z t = √ ᾱt z 0 + √ 1 -ᾱt ϵ t .Further, the denoising process starts with the most noisy latent sample z T , and progressively reduces the noise to recover the real sample z 0 with condition c. It is based on a variational approximation of the probabilities q(z t-1 |z t , c) given by:EQUATIONThe noise estimator ϵ θ (z t , 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 [22] ., that is:EQUATION3 PAD: Prediction with Action via Joint Denoising ProcessQuestion: 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 a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
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 I and the robot action modality M A . 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 to extra modalities M E . 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 c I , robot pose c A , an additional depth map c E (in Real-World tasks), and natural language instruction text l. The framework simultaneously outputs the corresponding future predictions x I , x E and robot action x A . Rather than predicting a single future step, PAD can forecast k future steps x 1:k I , x 1:k A , x 1:k E , which can be viewed as k step planning of the robot. Only the first predicted action x 1A 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.Model Input Process. 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 x I is initially processed through a pre-trained, frozen VAE [23] encoder ε I to derive the latent representation ε I (x I ). This latent representation is then converted into a sequence of tokens t I with embedding size h via tokenizer. Similarly, the robot pose x A is encoded using a Multi-Layer Perceptron (MLP) [24] into ε A (x A ) and linearly transformed into tokens t A with the same embedding size h. If available, the depth image is downsampled and tokenized into t E . The natural language instruction is processed through a frozen CLIP encoder [25] to produce the text embedding c l . 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 in previous works [18, 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 t I , t A , t E with lengths T I , T A , T E , respectively. These token sequences from different modalities are concatenated and undergo a joint latent denoising process Furthermore, the DiT architecture is adaptable to missing modalities. For example, in the case of a video dataset that lacks robot actions, the input to DiT only comprises the image tokens t I . We simply extend the token sequence to the combined length T I + T A + T E and introduce an attention mask in the self-attention block to exclude the padding tokens. Only effective predictions are retained in the output, discarding any padded parts. A brief illustration of the whole process is depicted on the right side of Figure 3 . 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 reconstruct future observation frames and desired robot action, conditioning on current observations c I , c A , c E . Following a similar strategy as in [26] , we concatenate conditional latent and noise latent in the channel dimension. Specifically, after obtaining encoded latent ε I (c I ), ε A (c A ), ε E (c E ), we concatenate these latent with noise to obtain conditioned noised latentL I = [ε I (c I ), z I t ], L A = [ε A (c A ), z A t ], L E = [ε E (c E ), z E t ].For instance, if the encoded latent ε I (c I ) has a shape of c × d × d, then z I t would have a shape of kc × d × d to represent k future frames, resulting in the final latent L I having a shape of (k + 1)c × d × d. The other modalities undergo a similar process.After concatenating the latent, these conditioned noisy latent from different modalities are tokenized into sequences of tokens t I , t A , t E with the same embedding size. The tokenization of image latent L I 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 to predict the latent representation of future frames. An illustration of the overall process can be found in Figure 3 .Initialization. Following the initialization process in [15] , 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 k + 1 times to encode the stacked latent, and the encoder and decoder for robot state are also zero-initialized.The diffusion process adds noise to the target encoded latent {ε I (x I ), ε A (x A ), ε E (x E )} and results in noised latent Z I,A,E = {z I t , z A t , z E t }. We train the PAD model to simultaneously predict the noise ϵ I , ϵ A , ϵ E added to the sample data, conditioned on current observations C I,A,E = {c I , c A , c E } and instructions l. This denoiser is trained with the DDPM [5] loss:EQUATIONwhere δ ∈ {I, A, E} represents different types of input modalities. The denoising loss L δ dif f aims to maximize the evidence lower bound (ELBO) [5] while approximating the conditional distribution p(ε δ (x δ )|C, l). We jointly minimize the following latent diffusion objectives and use hyperparameters λ I , λ A , λ E to balance the prediction loss between different modalities. Formally, the final training objective is given by:EQUATION)4 ExperimentsIn 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? We learn a single vision-language conditioned policy to solve all tasks in each domain with limited demonstrations, co-training with the bridge video data. In simulated MetaWorld, we learn a policy to tackle all 50 tasks. In real-world panda manipulations, we split objects into seen objects and unseen new objects to test the generalization ability of our policy.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA]• 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 a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not 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. 
environmental setups and baselines: 
 Metaworld. Metaworld [21] serves as a widely used benchmark for robotic manipulation, accommodating 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 a single text-conditioned visual policy to address all 50 tasks, within a data-efficient 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 a fair comparison, we do not utilize an additional depth input in Metaworld.Real-World Panda Manipulation Tasks. Our real-world experiments involve a Panda arm performing diverse manipulation tasks such as pressing buttons, opening drawers, routing cables, and picking and placing with various objects, as shown in Figure 4 . We follow the same hardware setup described in SERL [27] and utilize a wrist-mounted camera for pixel input [28] . The robot's poses are represented by 7-dimensional vectors, including 3 end-effector positions, 3 rotation angles, and 1 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 on all 50 tasks. Detailed data can be found in Appendix A.4.gripper status dimension. We collected 200 trajectories per task through teleoperation using a space mouse and scripted commands. Similarly, we developed a single policy capable of addressing all tasks, conditioned on instructions. We also assessed the policy's generalization capabilities on unseen tasks, as depicted in Figure 5 .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 1 day, utilizing 4 NVIDIA A100 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 λ I at 1.0 throughout the training period and linearly increased λ A and λ E from 0.0 to 2.0 during the 100k training steps.Policy Execution Details. Our policy is conditioned on the current image, c I , and the robot pose, c A , and predicts k frames of futures and actions. We configure the prediction horizon at k = 3 and set the interval between frames at i = 4 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 1:K I , and actions, x 1:k A . These k step predictions can be viewed as k step planning and only the first predicted action, x 1 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 our experiments, we opted to compare against a representative subset of prior methods that have either achieved state-of-the-art performance or share a 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] . A novel visual control policy that generates robot actions through an action diffuser. We augmented the original diffusion policy model with instruction conditions to address the multi-task setting. We use the CLIP encoder [25] as instruction encoders, referring to related work [29] . • SuSIE [18] . A two-stage approach that utilizes a pre-trained image-editing model [26] to generate image goals for robotic tasks, followed by a goal-conditioned low-level diffusion policy. We fine-tune the image-editing diffusion model on the same dataset and also use the diffusion policy for goal-conditioned behavioral cloning. To ensure a fair comparison, we also use the more powerful DiT framework as the image-editing model.• RT-1 [30] . An end-to-end robot control policy that leverages FiLM-conditioned [31] EfficientNet [32] to 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 visionlanguage models(VLMs) to produce robot actions. The original RT-2 model was fine-tuned on 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-7B [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.Performance Analysis. In all comparisons, we train a single visual policy to address all tasks within a domain, conditioned on instructions. Our proposed PAD outperforms all baselines by a significant margin. As shown in Table 1 , 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, we present comparisons on a subset of tasks and report the average success rate across all 50 tasks. A comprehensive comparison of all 50 tasks is available in Appendix A.4. Furthermore, Table 2 shows the results in real-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 in 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, basketball, etc., in Metaworld.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.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. In the first image, PAD imagines "a yellow pear" instead of 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 7 . 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 a certain level of uncertainty. These indicate that the PAD model has effectively learned the physical dynamics across these two domains.PAD can leverage existing physical knowledge from co-training on 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 4 , the expert dataset comprises only colored square blocks and plates, while we introduce a 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 or unseen backgrounds. We excluded depth input to ensure a fair comparison. As illustrated in Figure 5 , PAD demonstrates remarkable generalization abilities, successfully managing out-of-distribution objects such as strawberries, carrots, and eggplants, and even adapting to new backgrounds. The baseline method failed to generalize to difficult unseen tasks.Effectiveness of RGB image prediction. We evaluated the effectiveness of our joint prediction process by modifying the original model to exclude the image prediction component, namely in PAD w/o image prediction. This modification leads to significant performance drops compared to PAD, as illustrated in Table 1 . 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. Effectiveness of Co-training with Internet RGB Video Datasets. Another major benefit of PAD is the ability to co-train with large-scale internet-sourced RGB videos, which potentially leverages the physical knowledge embedded in these web datasets to 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 on the video dataset, as shown in Table 1 . 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-training, which in turn indirectly enhances robot action prediction. Compatible with Additional Modalities. As detailed in Section 3, our framework accommodates additional modalities owing to the adaptable DiT architectures. We incorporate additional depth image inputs in real-world manipulation experiments and jointly predict future RGB images, depth images, and robot actions, denoted as PAD-depth. We observe highly aligned prediction results among different modalities under our joint denoising framework, with some results illustrated in Figure 9 . The inclusion of depth input enhances performance in manipulation tasks, as demonstrated in Table 2 . This improvement may stem from the precise prediction of depth information, which aids agents in discerning distance changes, thereby enhancing performance. Moreover, our framework could be extended to predict other modalities relevant to robot control, such as tactile force or point clouds, which we left for the future work.We evaluated models across various sizes and patchify sizes [20] , as outlined in Table 3 . For example, the XL/2 model denotes the model with an XL size and a 2 × 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 Gflops) and the success rate (SR) of the learned policy, as depicted in Figure 10 . All the experiments are run in Metaworld benchmarks and detailed success rates for each task are provided in Appendix A.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 as policy by either generating high-level plans [37, 38, 39, 40, 41] or producing direct low-level motor control signals [30, 33, 42, 43, 44] . Many studies utilize the reasoning capabilities of pre-trained LLMs and VLMs to create high-level 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 as policy, pre-trained models can also guide policy learning from multiple aspects, such as providing good representations [45, 46, 47] , providing reward signals [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. A subset of research focuses on training conditional diffusion models that guide behavior synthesis based on desired rewards, and constraints under low dimensional stateinput setting [52, 53, 54, 55] . Diffusion Policy [16] trains a 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 on the benefits of pre-training with internet data. Another strand of research utilizes large-scale pre-trained diffusion models to perform data augmentation on training data, such as GenAug [56] , ROSIE [57] , and CACTI [58] .Learning. There also exist works that leverage future image predictions to assist policy learning. GR-1 [36] employs an autoregressive transformer to sequentially predict future images and actions. In contrast, we adopt a joint diffusion architecture that predicts more accurate future images, potentially leading to improved policy learning performance. UniPi [17] and SuSIE [18] employ a two-stage policy learning process, initially using a diffusion generative model to forecast future image or video sequences, and subsequently training an 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 a CNN-based U-net backbone [23] , designed primarily for image generation and limited to visual predictions, our method adopts a diffusion transformer (DiT) architecture [20] . This architecture adeptly handles multiple modalities concurrently via straightforward token concatenation and attention-mask mechanisms, enabling us to jointly predict future and actions simultaneously. 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 a unified denoising process. 
conclusion and discussion: 
 We present PAD, a novel framework to predict future images and generate actions under a 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.A limitation of the current method is 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 is 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 high-dimensional pixel spaces. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We describe the computational cost in Experiment section.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
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 or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: We cite the original paper that produced the code package or dataset Guidelines:• The answer NA means that the paper does not use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS 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. 
 
 Appendix: Videos of PAD can be found at https://sites.google.com/view/pad-paper.A. Output Decoder. 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 to recover the original samples in each modality.The RT-1 baseline is based on official implementation https://github.com/ google-research/robotics_transformer.The Diffusion Policy baseline is based on https://github.com/real-stanford/diffusion_ policy, and we follow https://github.com/real-stanford/scalingup to add language condition.The RT-2 baseline is re-implemented by ourselves. We use the InstructBlip-vicuna-7b model as backbone https://huggingface.co/Salesforce/instructblip-vicuna-7b.The GR-1 baseline is built on https://github.com/bytedance/GR-1. Since we can not access the pretraining dataset in the original paper, we initialize the model with the author's open-source checkpoint. Table 4 : Models with various size and computational cost.Expert data collection. We collected data on 6 categories of tasks, including button press, cableroute, pick and place, and drawer open/close. For the pick task, we include 4 colors of blocks. For the place task, we include 3 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 11 .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, as shown in Figure 12 . For convenience, videos of PAD can be found at https://sites.google.com/ view/pad-paper. Table 5 : Detailed success rate of baselines and ablations. We did not include the handle-pull-side-v2 and handle-pull-v2 tasks since the expert policy for these two tasks in the original benchmark had low success rates. Every task is tested with 25 rollouts.A. 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 two tasks are in low success rates.NeurIPS Paper ChecklistJustification: The research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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).
questions: questions: - It would be interesting to ablate the co-training data in detail with different sources to understand its effectiveness.
limitations: limitations: The authors briefly discussed the limitation in the last section, which looks good to me.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- 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?).
questions: questions: Questions:
- 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.
- 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.
-
limitations: 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).
paper_summary: paper_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.
- 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.
- 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.
- 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.
strength_weakness: strength_weakness: Originality:
- 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.
- Several statements are not supported by references:
    - Line 163: “Previous studies that utilized pixel input generally developed separate task-specific policies for each of the 50 tasks“.
    - Lines 285-286: “Many studies utilize the reasoning capabilities of pre-trained LLMs and VLMs to create high-level plans followed by motion primitives. 
    - Lines 286-288: “Additionally, some approaches adapt pre-trained models to emit low-level motor control signals by adding an action head.”

Quality:
- 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.
- 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.
- 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.
- 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.
- 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.

Clarity:
- None beyond what is mentioned above.

Significance:
- 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.
- 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.
questions: 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?
- How were the hyperparameters tuned for the proposed method, for the ablations, and for the baselines?
- 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?
- Line 219: How were the tasks that are shown in Table 1 selected?
- Line 21: Do the videos on the paper website show the robot at real-time speed or are the videos accelerated?
- How do the proposed methods and the baselines compare in terms of latency?
- Are code and training data going to be released publicly?

UPDATES AFTER REBUTTAL:
- Presentation: 2 -> 3
- Contribution: 2 -> 3
- Score: 4 -> 6
limitations: limitations: As mentioned above, the paper lacks a discussion of latency considerations which might limit the practical usefulness of this approach.
paper_summary: paper_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.

As 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.

The 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.
strength_weakness: strength_weakness: 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.
Specifically, 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.

2. (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:

* 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.
On 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.

* 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.


3. (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.

[GR-1] Unleashing Large-Scale Video Generative Pre-training for Visual Robot Manipulation, ICLR24
questions: 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:

* The writing changes to include discussion on previous works that had the same high-level idea and focus on the differences with those.

* 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.


--------------
Post-rebuttal: score updated from 2 to 4.
limitations: limitations: Addressed

Comments: title: Welcome any discussion if you want at the last moment!
Thank you for your support!
Response
Thank you for your endorsement of our paper!
Dear Reviewer Hcpx: Thank you for your time and effort!
Thank you for the comprehensive feedback!
Thank you for your answers
About the reference
Welcome to see our new experiments on CALVIN benchmark!
Response to Reviewer Hcpx's further questions
To reviewer EGNn : Please respond to rebuttal
To reviewer yjvj : Please respond to rebuttal
To reviewer GDHG : Please respond to rebuttal
References for rebuttal
comment: Dear Reviewers and Area Chairs,

As the discussion period draws to a close, with less than two hours remaining, we remain available to engage in any further discussions you may wish to have. We deeply appreciate your insights and feedback and are grateful that all reviewers have acknowledged our contributions and novelty to varying extents. We sincerely thank you for the time and effort you have dedicated to reviewing our paper.

Best regards,

Authors
Dear Reviewer yjvj,

We sincerely appreciate your support and engagement with our work!

Regarding the visualizations, Metaworld's requirement for precise movements means it cannot adequately demonstrate how PAD handles uncertainty. As we also co-train on bridge datasets, we have included visualizations of our model's predictions on bridge at the top of our website (as noted in the abstract), **which better illustrate how PAD manages uncertainty in long-horizon prediction**. Visualization include the following scenarios: (1) After opening the cabinet door, a blue bowl and banana are imagined inside the cabinet. (2) A carrot, previously obscured, appears in the sink, though the ground truth is a pepper. PAD infers the identity of the object based on its exposed portion. 

Regarding the novelty, we acknowledge that the high-level idea of robot learning involving prediction and action—more broadly, deriving better representations from video datasets—is a well-explored area. Our contribution and novelty lie in proposing a unified diffusion-based architecture that jointly predicts futures and actions. This approach is inspired by the impressive recent performance of DiT in image/video generation [1,2]. Our findings suggest that DiT can also handle multi-modal predictions in robotic tasks with appropriate technical design.

Thank you once again for your time and effort in reviewing our work! We are committed to continually refining our research to meet the highest standards.

Best regards,
The Authors

---

[1] Esser P, Kulal S, Blattmann A, et al. Scaling rectified flow transformers for high-resolution image synthesis[C]//ICML 2024.

[2] Ma X, Wang Y, Jia G, et al. Latte: Latent diffusion transformer for video generation[J]. arXiv preprint arXiv:2401.03048, 2024.
Thanks for the additional visualizations. I don't think MetaWorld is hard enough or has any ambiguous like tasks really so it's hard to tell how PAD handles situations where objects may unpredictably move around or occlusions.

I currently do not have any other concerns and will bump up my score to 8 as I don't really share the concerns as some of the other reviewers upon reviewing their comments. I do agree with other reviewers that this idea is not really novel, it's a mix of established techniques and in my opinion is very much an "expected idea". But given it performs much better than related work already, is well worth at minimum a score of 8. I can see other papers citing this paper because they really use it as a baseline which is impactful enough for me.

I leave my confidence at 4 since I am experienced with robot learning (small/large scale) and simulation, although not experienced with diffusion models specifically.
Dear Reviewer EGNn:

We would like to express our profound gratitude for your endorsement of our paper. We will keep polish the paper to meet highest standard. Once again, thank you for your time and effort in reviewing our paper!

Best, Authors
Dear Reviewer Hcpx,

Thank you for your interest in our paper and for your detailed feedback over the past few days!

We are extremely grateful for your patient and thoughtful insights, which have significantly improved our manuscript. We will release the code upon acceptance, enabling you to examine the specifics more thoroughly. As it approach the end of the rebuttal phase, if you find the content of our paper and our responses to your questions satisfactory, we would sincerely appreciate any consideration for a score improvement.

Once again, we deeply appreciate your efforts in reviewing our work and for the insightful questions that have been invaluable in enhancing our research. We still remain open to addressing any further queries you may have!

Best regards,

The Authors
Dear Reviewer GDHG:

Thank you for your comprehensive feedback! We are delighted to hear from you.

First of all, we are very glad that we reach a consensus on the technical contribution of the paper. As you mentioned, there is a growing integration of video prediction methods in robotic learning. We have bet and explored a unified diffusion architecture, as we all recognized, the first of its kind in this context. Thank you for your suggestions which have helped refine and clarify our contributions!

As for the CALVIN benchmarks, we totally agree with you that CALVIN can better **evaluate the instruction following ability** and easier to compare with previous works, which had encourage us to conducted additional experiments on it. Thank you for your constructive suggestions!

Concerning the MetaWorld benchmarks, we would like to offer a perspective that this benchmark is particularly effective at **evaluating the dexterity of the learned policy** due to its requirement for precise actions with numerous small objects. Sometimes robot learning policy can still failed even given a single task to accomplish (as you said, without distractions) due to a lack of precision in execution. We believe that **PAD offers unique advantages in enhancing the dexterity** (or precision) of the policy, as it predicts precise future images and actions through joint denoising. The protocol for this benchmark includes an official scripted policy provided by the authors of MetaWorld to collect demonstrations for each task, which is a method commonly used in previous research.

In summary, while the Calvin benchmarks prioritize language comprehension and do not focus heavily on dexterity and precision, as the tasks involve only large blocks and objects such as drawers or doors, MetaWorld requires precise execution of dexterous tasks without a strong emphasis on following instructions. We believe that **both language comprehension and dexterity** are crucial for developing intelligent robots. We hope this detailed analysis can explain why our experiments on Metaworld is also meaningful, which assess different aspect of the learned policy. 

Thank you once again for your detailed feedback and the improved score. We will add all the new experiments and discussions to the paper to meet the highest standards. As the rebuttal phase draws to a close, we still remain open to addressing any further queries you may have!

Best regards,
Authors
Thanks for the authors' response. I have also read other reviewers' comments and decide to keep my score as 7
Thank you for your effort on the rebuttal.

Regarding my concerns on contribution: I also recognized the technical contribution of this paper. My major concern was the absence of related alternatives, specifically GR1, from the discussion. Of course working on an existing problem and advancing it is a valid contribution, however it should be placed in the context of related work. The rebuttal discussed GR1 and I believe it placed itself better in the space of relevant literature.

Regarding MetaWorld: while arguing on which benchmark is more challenging is of little value, my take on MetaWorld in my original review was that the task can be usually inferred by just looking at the scene, so that makes multi-tasking results less impressive. This is because of the lack of distractions in the scene, mostly the relevant objects per task are there. CALVIN scenes do not offer a wide variety of tasks to be deployed, but the scene is the same and language understanding is necessary for task execution. At the same time, I thought that evaluation on CALVIN would further allow for easier comparisons against stronger related baselines: SuSIE, which is very related, GR1 and RT1's results are already there for free. Since SuSIE was adapted to MetaWorld, I thought it could be as easy to take its data loader and adapt it to support PAD, thus enabling all those comparisons for free. Lastly, CALVIN offers an established train/test setup, where baselines have shown the benefits of using video data (GR1), future prediction (SuSIE), foundation models (RoboFlamingo). In contrary, MetaWorld papers often define their own setup in terms of tasks and demos used and I'm not aware of a standardized and widely adopted protocol there.

I'm eventually glad to see that experiments on CALVIN were ran. I agree that the performance on ABCD->D is already quite high. The SOTA on this split is actually much better than GR1, but showing SOTA results on CALVIN was never my expectation for this submission, the comparison to previous works was. It's good to see that PAD is more sample-efficient. It would be more meaningful to also show results on ABC->D, which is more challenging and could better showcase the advantages of using joint video training.

Regarding the baselines: I still believe that GR1 is the most related baseline, then also SuSIE. It's ok to include RT-1 and 2 as well as training paradigms or transformer-based architectures. For diffusion policies, if CALVIN was adopted from the beginning it would allow comparison with 3D Diffuser Actor, which is a much stronger baseline. But it's ok that this is not included as long as other video prediction approaches are included in a fair comparison, which is better satisfied after seeing the CALVIN results.

Overall, I appreciate the effort for the rebuttal, especially after seeing the CALVIN results. As more and more video prediction methods for robotics show up, the questions that make the difference are what is the right representation and which implementation details are important. This paper bets on pixels and offers a valid technical contribution. I respect that, even if I'm not fully convinced about the strength of the specific implementation. It's better to let multiple voices be heard, so I'm removing my strong objections against this paper, even if I'm not giving it a full acceptance score for now.

Score updated from 2->4.
Dear Reviewer Hcpx:

Sorry for the wrong reference. The correct reference should be [14]. On page 18 of the experiment details of [14], there is a table provided categorizes tasks into easy, medium, hard, and very hard. In our table, the "easier tasks" row includes some of the easy tasks, while the "harder tasks" row encompasses medium, hard, and very hard tasks. This categorization is also employed in the 3D Diffusion Policy paper [15], specifically in Section IV.A, line 17.

We have check all other citations for accuracy.

Thank you for your time, and we greatly appreciate your detailed feedback! 

Best regards,

Authors

[14] Seo, Y., Hafner, D., Liu, H., Liu, F., James, S., Lee, K., & Abbeel, P. Masked world models for visual control. In Conference on Robot Learning (pp. 1332-1344). PMLR.

[15] Ze Y, Zhang G, Zhang K, et al. 3d diffusion policy[J]. arXiv preprint arXiv:2403.03954, 2024. RSS2024
> "Q5: How were the tasks that are shown in Table 1 selected? ANS: The tasks in Meta-World are categorized into different difficulty levels according to [1]."

I think this might be the wrong reference as [1] does not evaluate on Meta-World?
Dear Reviewer GDHG:

After we submit the initial rebuttal message, we continue to conduct experiments in the Calvin benchmark for better comparisons as you suggested. We are pleased to share some promising results. Following the settings in the GR-1 paper, we conducted imitation learning on the whole training datasets and 10% training datasets respectively.  When using the full dataset, our results were comparable to those reported for GR-1, given that GR-1's original success rate was already high (approaching 100%) and we had limited time for parameter tuning. However, **in the reduced 10% data regimen, PAD clearly outperforms GR-1**, indicating that PAD is a more sample-efficient learning algorithm with joint prediction and action under a unified diffusion architecture. The detailed comparison of task completion numbers in a row is shown below: (some baseline data is sourced from GR-1)

| **Whole datasets** | 1     | 2     | 3     | 4     | 5     | Avg. Len.     |
| ------------------ | ----- | ----- | ----- | ----- | ----- | ------------- |
| RT-1               | 0.844 | 0.617 | 0.438 | 0.323 | 0.227 | 2.45          |
| HULC               | 0.889 | 0.733 | 0.587 | 0.475 | 0.383 | 3.06          |
| MT-R3M             | 0.752 | 0.527 | 0.375 | 0.258 | 0.163 | 2.08          |
| GR-1               | 0.949 | 0.896 | 0.844 | 0.789 | 0.731 | **4.21**      |
| PAD(ours)          | 0.965 | 0.904 | 0.842 | 0.768 | 0.707 | **4.19**      |
| **10% datasets**   | **1** | **2** | **3** | **4** | **5** | **Avg. Len.** |
| RT-1               | 0.249 | 0.069 | 0.015 | 0.006 | 0.000 | 0.34          |
| HULC               | 0.668 | 0.295 | 0.103 | 0.032 | 0.013 | 1.11          |
| MT-R3M             | 0.408 | 0.146 | 0.043 | 0.014 | 0.002 | 0.61          |
| GR-1               | 0.778 | 0.533 | 0.332 | 0.218 | 0.139 | 2.00          |
| PAD(ours)          | 0.813 | 0.606 | 0.412 | 0.294 | 0.189 | **2.31**      |

Thank you once again for your time and effort in reviewing our paper. We hope that these additional experiments address your remaining concerns and demonstrate the improved quality of our work. We are always ready to answer your further questions!
Dear Reviewer Hcpx:

We are very delighted to hear from you! Here are our replies:

Q1: Did you sweep the lambda hyperparameter for weighting the two loss terms in DBC or did you use the value used in the original paper?

Ans: We used the value of lambda from the original paper. As shown in Figure 6 of the DBC paper, the authors performed a hyperparameter sweep and found that a lambda value around 1 yields the best performance. Consequently, we adopted lambda = 1 for our experiments.

Q7:  For the baselines, did you tune the hyperparameters on the benchmarks  that are being used here or did you use the same hyperparameters as in the original papers? For the ablations, did you re-tune other  hyperparameters or are they kept the same?

Ans: For diffusion policy, we use their config for simulation envs( https://diffusion-policy.cs.columbia.edu/data/experiments/image/lift_ph/diffusion_policy_transformer/config.yaml). We tried to tune the predicti_lens and transformer layer parameters in their config but find their original configs performs the best. For  RT-1, and Susie, we utilized the parameters from their open-source implementations. As for RT-2, since the code was not publicly available, we carefully implemented it based on the descriptions provided in the paper.

Regarding our ablation study, we maintained the same hyperparameters and only removed the specific components being ablated.

---

We hope these answers address your questions. Thank you again for your valuable time! We are always ready to answer any of your new questions!

Best regards,

Authors
Hi reviewer EGNn,

Thank you for your initial review. Please kindly respond to the rebuttal posted by the authors.
Does the rebuttal answer your questions/concerns? If not, why?

Best,
AC
Hi reviewer yjvj,

Thank you for your initial review. Please kindly respond to the rebuttal posted by the authors.
Does the rebuttal answer your questions/concerns? If not, why?

Best,
AC
Hi reviewer GDHG,

Thank you for your initial review. Please kindly respond to the rebuttal posted by the authors.
Does the rebuttal answer your questions/concerns? If not, why?

Best,
AC
Thank you for the rebuttal.

Q1: Did you sweep the lambda hyperparameter for weighting the two loss terms in DBC or did you use the value used in the original paper?

Q7: For the baselines, did you tune the hyperparameters on the benchmarks that are being used here or did you use the same hyperparameters as in the original papers? For the ablations, did you re-tune other hyperparameters or are they kept the same?
[1] Chi C, Feng S, Du Y, et al. Diffusion policy: Visuomotor policy learning via action diffusion[J]. RSS2023

[2] Ze Y, Zhang G, Zhang K, et al. 3d diffusion policy[J]. RSS2024

[3] Nair S, Rajeswaran A, Kumar V, et al. R3m: A universal visual representation for robot manipulation[J]. CoRL2022

[4] Majumdar A, Yadav K, Arnaud S, et al. Where are we in the search for an artificial visual cortex for embodied intelligence?[J]. Neurips2023

[5] Mu Y, Zhang Q, Hu M, et al. Embodiedgpt: Vision-language pre-training via embodied chain of thought[J]. Neurips2023

[6] Zhao T Z, Kumar V, Levine S, et al. Learning fine-grained bimanual manipulation with low-cost hardware[J]. RSS2023

[7] Ahn M, Brohan A, Brown N, et al. Do as i can, not as i say: Grounding language in robotic affordances[J]. CoRL 2033

[8] Driess D, Xia F, Sajjadi M S M, et al. Palm-e: An embodied multimodal language model[J]. arXiv preprint arXiv:2303.03378, 2023.

[9] Chen W, Mees O, Kumar A, et al. Vision-language models provide promptable representations for reinforcement learning[J]. arXiv preprint arXiv:2402.02651, 2024.

[10] Peebles W, Xie S. Scalable diffusion models with transformers[C]//ICCV 2023: 4195-4205.
[11] Esser P, Kulal S, Blattmann A, et al. Scaling rectified flow transformers for high-resolution image synthesis[C]//ICML 2024.
[12] Wang H C, Chen S F, Hsu M H, et al. Diffusion model-augmented behavioral cloning[J]. ICML 2024.

[13] Song Y, Dhariwal P, Chen M, et al. Consistency models[J]. arXiv preprint arXiv:2303.01469, 2023.