The Ingredients for Robotic Diffusion Transformers

TL;DR

This paper introduces a novel diffusion transformer architecture for robotics that improves task performance and scalability, reducing the need for extensive hyper-parameter tuning across diverse robotic tasks and embodiments.

Contribution

The paper identifies key design choices for diffusion transformer policies, proposes an improved architecture called extit{method}, and demonstrates superior performance on long-horizon robotic tasks.

Findings

Outperforms state-of-the-art in long-horizon dexterous tasks

Shows improved scaling with multi-modal, language-annotated data

Reduces hyper-parameter tuning for diverse robotic setups

Abstract

In recent years roboticists have achieved remarkable progress in solving increasingly general tasks on dexterous robotic hardware by leveraging high capacity Transformer network architectures and generative diffusion models. Unfortunately, combining these two orthogonal improvements has proven surprisingly difficult, since there is no clear and well-understood process for making important design choices. In this paper, we identify, study and improve key architectural design decisions for high-capacity diffusion transformer policies. The resulting models can efficiently solve diverse tasks on multiple robot embodiments, without the excruciating pain of per-setup hyper-parameter tuning. By combining the results of our investigation with our improved model components, we are able to present a novel architecture, named \method, that significantly outperforms the state of the art in solving long-horizon ($1500+$ time-steps) dexterous tasks on a bi-manual ALOHA robot. In addition, we find that our policies show improved scaling performance when trained on 10 hours of highly multi-modal, language annotated ALOHA demonstration data. We hope this work will open the door for future robot learning techniques that leverage the efficiency of generative diffusion modeling with the scalability of large scale transformer architectures. Code, robot dataset, and videos are available at: https://dit-policy.github.io