Prediction with Action: Visual Policy Learning via Joint Denoising Process

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.

---

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.

ANS: 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. GR1 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.

---

Q2: About the meta world benchmark and Calvin benchmarks

ANS: 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).

We 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.)

Q3: Detailed experimental comparison with GR-1 baseline

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

---

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.

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

---

Q5: Training on videos from the web or external data (such as BridgeData-v2) to test on simulated MetaWorld feels disconnected.

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

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

---

We 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!

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)

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!

We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments.

---

Q1: Discussion and comparisons with DBC paper which learn a diffusion model over state-action pairs as auxiliary loss

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

---

Q2: Several statements are not supported by references.

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

---

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.

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

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

---

Q4: How were the Meta-World training trajectories collected? Are there existing public datasets of demonstration?

ANS: 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].

---

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]. We selected representative tasks from each level to display due to space constraints, and also the average success rate.

---

Q6: About random seeds and error bars of experiments.

ANS: 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].

---

Q7: How the hyperparameters were tuned. The architectural differences between the proposed method and diffusion policy

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

---

Q8: Do the videos on the paper website show the robot at real-time speed?

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

---

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?

ANS: 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].

---

Q10: Are code and training data going to be released publicly?

ANS: 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!

---

Thank you again for your time and efforts! 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

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

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

We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments.

---

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.

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

---

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.

ANS: 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).

---

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?).

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

In the low data regime, PAD also surpasses the baseline

---

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.

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

---

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.

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

---

Thank 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!

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

We sincerely appreciate your time and efforts in reviewing our paper! Based on your review, we added a detailed discussion and additional experiments.

---

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?

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

---

Q2: It would be interesting to ablate the co-training data in detail with different sources to understand its effectiveness.

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

---

Thank 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!

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