Prediction with Action: Visual Policy Learning via Joint Denoising Process

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.

Strengths

1. (Presentation) The paper is very clearly written and is very easy to follow and understand. I want to additionally endorse the tone of this paper, which is humble and moderate, without overpresenting or overlaiming. 2. (Contribution) As a technical contribution, the paper shares interesting architectural details on how to repurpose DiT for joint action and future image prediction. 3. (Soundness) The ablations show the importance of co-training with image prediction and the additional benefit of training on video corpora. I believe that overall the paper contributes to the idea that predicting the future and acting is a promising path for generalizable robot learning, as opposed to action-centric policy learning. Although the idea is not original per se, the paper provides more evidence for it.

Weaknesses

1. (Contribution) The motivation of the paper is that first, the ability to predict future images is coupled with the ability to predict actions (idea); second, diffusion models have shown good results on both tasks, so it is reasonable to combine them for joint prediction (implementation). To the best of my knowledge, the implementation part is novel, but the idea part is not. 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

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.

Rating

4

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

Addressed

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.

Strengths

Originality: - The paper proposes an interesting combination of image diffusion and policy diffusion to leverage transfer between the two tasks. Using latent diffusion with a transformer so that all modalities can be processed in a shared latent space while also being able to train on video data without actions is a nice and elegant idea. Quality: - Experiments on Meta-World and with a real-world Panda arm on both seen and unseen tasks are reasonable for validating the proposed method. - The selected baselines (a multi-task variant of diffusion policy, a modified version of SuSIE, RT-1, and a re-implementation of RT-2 built on top of InstructBlip7B) are recent and strong. - For the diffusion policy baseline, a pretrained CLIP encoder was used as in related work which also matches language encoder used by the proposed method. Similarly, a diffusion transformer is used for SuSIE for a fairer comparison. - Table 1 includes two important ablations: PAD without image diffusion and PAD without on videos without actions that validate two of the main hypotheses of this work. There is also an ablation that shows that using depth maps as additional observations improves performance in real-world experiments. Clarity: - The paper is generally well-written and clear. - Various hyperparameters are provided in the appendix to facilitate reproducibility. - The videos on the paper website provide a good sense of qualitative performance. Significance: - Considering the good empirical performance and the elegance of this approach, it seems reasonable that other researchers and practitioners might use or build on this work.

Weaknesses

Originality: - 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

- 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

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

As mentioned above, the paper lacks a discussion of latency considerations which might limit the practical usefulness of this approach.

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.

Strengths

- An interesting joint denoising process is presented to generate future frames and actions. The paper further shows the possibility that PAD can be extended to more modalities if the engineering efforts are taken to do so which is great to see. - Nice use of the diffusion transformer architecture which enables co-training on large video datasets that do not come with action labels, providing internet-scale knowledge that improves success rate. It would be nice to see how PAD without co-training on internet data performs on the generalization benchmark. I would imagine it won't do as well as PAD with internet data but is worth adding in my opinion (not a huge issue though). - Strong results on generalization benchmarks and solid ablation studies on co-training show great promise about the proposed method. - PAD also shows good scaling results where larger models/more training yields higher success rates.

Weaknesses

- How well does PAD predict future frames? Only a few next frames are shown in fig 7 but I am curious to see how far it can predict? In some sense tasks where there is a lot more ambiguity (e.g. objects that fall and bounce, or finding an object hidden behind something) would be helpful to help understand how PAD handles uncertainty. The current set of tested tasks are fairly straight forward in terms of predicting the next frames. - 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: - 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. -

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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

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.

Strengths

- The motivation of this paper is clear. The paper is easy to follow and the proposed framework makes sense to me. The idea is simple and straightforward but effective. - The performance looks good compared to the RT-1/RT-2 baselines. - The experimental section is comprehensive, verifying performance on both sim and real world tasks. The authors also tested the model with various computation costs. - Implementation details are clearly described and code is provided, which makes the community easy to reproduce.

Weaknesses

- The performance improvement on real data is somewhat marginal, which makes it questionable about its effectiveness on real data. It would also be great to showcase the ablation studies for real-world tasks to understand whether this is caused by limited model size and pretrained data (compared to RT-2).

Questions

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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors briefly discussed the limitation in the last section, which looks good to me.

References for rebuttal

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

To reviewer GDHG : Please respond to rebuttal

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

To reviewer yjvj : Please respond to rebuttal

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

To reviewer EGNn : Please respond to rebuttal

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

Welcome any discussion if you want at the last moment!

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