Diffusion Transformer Policy

## Abstract

Abstract Recent large visual-language action models pretrained on diverse robot datasets have demonstrated the potential for generalizing to new environments with a few in-domain data. However, those approaches usually predict discretized or continuous actions by a small action head, which limits the ability in handling diverse action spaces.
In contrast, we model the continuous action with a large multi-modal diffusion transformer, dubbed as Diffusion Transformer Policy, in which we directly denoise action chunks by a large transformer model rather than a small action head.
By leveraging the scaling capability of transformers, the proposed approach can effectively model continuous end-effector actions across large diverse robot datasets, and achieve better generalization performance.
Extensive experiments demonstrate Diffusion Transformer Policy pretrained on diverse robot data can generalize to different embodiments, including simulation environments like Maniskill2 and Calvin, as well as the real-world Franka arm.
Specifically, without bells and whistles, the proposed approach achieves state-of-the-art performance with only a single third-view camera stream in the Calvin novel task setting (ABC → → \rightarrow → D), improving the average number of tasks completed in a row of 5 to 3.6, and the pretraining stage significantly facilitates the success sequence length on the Calvin by over 1.2. The code will be publicly available.

## 1 Introduction

Traditional robot learning paradigm usually relies on large-scale data collected for a specific robot and task, but collecting robot data for generalist tasks is time-consuming and expensive due to the limitations of robot hardware in the real world.
Nowadays, the foundational models OpenAI (; ; ); Rombach et al. () in Natural Language Process and Computer Vision, pretrained on broad, diverse, task-agnostic datasets, have demonstrated powerful ability in solving downstream tasks either zero-shot or with a few task-specific samples.
It is principally possible that a general robot policy exposed to large scale diverse robot datasets improves generalization and performance on downstream tasks Brohan et al. (; ). However, it is challenging to train a general robot policy on a large scale of cross-embodiment datasets with diverse sensors, action spaces, tasks, camera views, and environments.

Toward a unified robot policy, existing works directly map visual observation and language instructions to actions with large visual-language-action models for robot navigation Shah et al. (; ) or manipulation Brohan et al. (; ); Kim et al. (); Team et al. (), and demonstrate zero-shot or few-shot generation to new environments.
Robot Transformers Brohan et al. (; ); Padalkar et al. () present robot policy based on transformer architecture, and demonstrate robust generalization by training on the large scale of Open X-Embodiment Dataset Padalkar et al. (). Octo Team et al. () follows the autoregressive transformer architecture with a diffusion action head, while OpenVLA Kim et al. () discretizes the action space and leverage the pretrained visual-language model to build VLA model exposed to Open X-Embodiment Dataset Padalkar et al. ().
Though those Visual-Language-Action (VLA) models Team et al. (); Kim et al. () have shown the potential to learn robot policy from the large cross embodiment datasets Padalkar et al. (), the diversity of robot space among the cross embodiment datasets still limits the generalization of previous models with autoregressive transformers.

Recent diffusion policy Chi et al. (); Ze et al. (); Ke et al. () has shown its stable ability in robot policy learning for single task imitation learning with UNet or cross attention architecture, and diffusion transformer demonstrates its scalability in multi-modal image generation Peebles & Xie (). Specifically, Octo Team et al. () presents a generalist policy that denoises the action with a small MLP network conditioned on a single embedding of auto-regressive multi-modal transformer. However, the robot space of large-scale cross-embodiment datasets contains various cameras views and diverse action spaces, which poses a significant challenge for a small MLP to denoise the continuous action conditioned on a single action head embedding.

In this paper, we design a Diffusion Transformer architecture for generalist robot policy learning.
Similar to previous robot transformer models Brohan et al. (; ); Padalkar et al. (); Team et al. (); Kim et al. (), we leverage the transformer as our base module to retain the scalability on the large-scale cross-embodiment datasets.
Different from Brohan et al. (; ); Padalkar et al. (); Team et al. (); Kim et al. (), we present an in-context conditional diffusion transformer architecture to denoise the action chunks, rather than utilizing a small MLP to denoise action embedding to continuous actions as illustrated in Figure 1. The Diffusion Transformer Policy retains the scalability of transformer for diffusion, and thus more effectively generalizes action policy from the large diverse datasets.

In a nutshell, we present a Diffusion Transformer Policy, that incorporates a causal transformer as an in-context conditional diffusion backbone and denoise continuous action chunks with the transformers.
Extensive experiments demonstrate Diffusion Transformer Policy achieves considerably better performance on two large-scale Sim datasets, Maniskill2 and Calvin, compared to corresponding diffusion action head Team et al. () and discretized action head Brohan et al. () baselines.
Meanwhile, The proposed model trained on the Open X-Embodiment Dataset achieves better generalization performance compared to the baseline methods on the Real Franka platform.

Figure: Figure 1: Illustrations of different robot policy architectures. (a) is the common robot transformer architecture with discretization actions, e.g., Robot Transformer Brohan et al. (; ) and OpenVLA Kim et al. (). (b) is the transformer architecture with diffusion action head which denoises the continuous action with a small MLP with embedding from the causal transformer, e.g., Octo Team et al. (). (c) is the proposed Diffusion Transformer architecture that utilizes the large transformer to denoise actions in an in-context conditioning style.
Refer to caption: x1.png

## 2 Related Work

Diffusion Policy
Denoise diffusion techniques Ho et al. (); Rombach et al. () are pioneering image generation, and recent Diffusion Policy Chi et al. (); Ze et al. (); Ke et al. () has exhibited a powerful ability in modeling multimodal actions compared to previous robot policy strategies in both 2D Chi et al. () and 3D observations Ze et al. (); Ke et al. (). Current diffusion policy approaches usually follow an Unet structure or a shallow cross-attention network for a single manipulation task, leaving large-scale multimodal diffusion policy poorly investigated. For example, 3D diffusion Policy Ze et al. () presents a diffusion approach conditional on a 3D point cloud, while 3D diffuser actor Ke et al. () proposes a 3D diffusion strategy based on point cloud with cross-attention. Differently, we present a scalable in-context conditioning diffusion transformer architecture for generalist robot policy. Recent generalist policy Octo Team et al. () conditions the denoise process on the embedding from the Transformer model with a small MLP diffuser. By contrast, the diffuser in Diffusion Transformer Policy is a large Transformer architecture.

Generalist Robot Policies
The embodiment community has shown increasing interest in generalist robot policy with foundational multi-modal models for both robot navigation Shah et al. (; ); Yang et al. (); Sridhar et al. (); Huang et al. () and manipulation Bousmalis et al. (); Brohan et al. (); Shah et al. (); Shridhar et al. (); Brohan et al. (); Padalkar et al. (); Kim et al. (); Team et al. (). Recent approaches Brohan et al. (; ); Padalkar et al. (); Kim et al. (); Team et al. () aim to achieve generalist policy with scalable Visual-Language-Action models. We follow this paradigm to approach generalist and adaptive robot policy.  Brohan et al. (; ); Padalkar et al. (); Kim et al. () construct the action token by discretizing each dimension of the robot actions separately into 256 bins. However, this discretization strategy incurs internal deviation in robot execution. Unlike those methods, we present a Diffusion Transformer Generalist Policy, which denoises the continuous actions with a large Transformer model. The proposed approach retains the scalability of the Transformer and meanwhile facilitates the modeling of cross-embodiment action chunk representations. Meanwhile, the Diffusion Transformer Policy aligns robot action together with the language instructions and image observations as an in-context conditional style.

## 3 Method

We describe the proposed architecture of diffusion transformer policy in this section, a DiT-based generalist diffusion policy model that can be adapted to new environments and embodiment. The model, built from a diffusion transformer architecture, achieves better generalization on both novel camera views and environments from large amounts of diverse robot data.

### 3.1 Architecture

Instruction Tokenization. The language instructions are tokenized by a frozon CLIP Radford et al. () model.

Image observation Representations. The image observations first pass into the DINOv2 Oquab et al. () to obtain the image patch features. Note that DINOv2 Oquab et al. () is trained on the web data which is different from the robot data, we thus jointly optimize the DINOv2 parameters together with Transformers through an end-to-end way.

Q-Former. To reduce the computation cost, a Q-Former together with FiLM Perez et al. () conditioning is incorporated to select image features from the patch features of DIDOv2 Oquab et al. () by instruction context.

Action Tokenization. We use the end-effector action and represent each action with a 7D vector, including 3 dimensions for translation vector, 3 dimensions for rotation vector, and a dimension for gripper position. To align the dimension with image and language tokens, we simply pad the action vector with zeros to construct the action token. We only add the noise into the 7D action vector during denoise diffusion optimization.

Architecture. Our core design is the Diffusion Transformer structure Peebles & Xie () which denoises action token chunks, instead of each single action token, conditioned on image observation and instruction tokens by an in-context conditioning style with a causal transformer network, i.e., we simply concatenate image tokens, language tokens, and timestep token in the front of the sequence, equally treating the noisy action tokens from the image/instruction tokens as illustrated in Figure 2. This design retains the scaling properties of transformer networks. The model, conditioned on language instructions and image observations with the causal transformer structure, is supervised by the noise that we add to the continuous actions. In other words, we conduct the diffusion objective directly in the action chunk space with a large transformer model, differently from a diffusion action head with a few MLP layers Team et al. ().

The proposed Diffusion Transformer Policy is a general design that can be scaled to different datasets, and demonstrate excellent performance. Meanwhile, we can also add additional observation tokens and input into the transformer structure. Appendix A provides more details.

### 3.2 Training Objective

In our architecture, the denoising network $\epsilon_{\theta}({\bm{x}}^{t},c_{obs},c_{instruction},t)$ is the whole causal transformer, where $c_{obs}$ is the image observation, $c_{instruction}$ is the language instruction, and $t\in{1,2,...T}$ is step index in our experiments. During the training stage, we sample a Gaussian noise vector ${\bm{x}}^{t}\in\mathcal{N}(\mathbf{0},{\bm{I}})$ at timestep $t$, where $T$ is the number of denoising timesteps, and add it to action ${\bm{a}}$ as $\hat{{\bm{a}}}$ to construct the noised action token, finally predicting the noise vector $\hat{{\bm{x}}}$ based on the denoising network $\epsilon_{\theta}(\hat{{\bm{a}}},c_{obs},c_{instruction},t)$, where $t$ is randomly sampled during training. We optimize the network with MSE loss between ${\bm{x}}^{t}$ and $\hat{{\bm{x}}^{t}}$.

To generate an action, we apply $T$ steps of denoising with the optimized transformer architecture $\epsilon_{\theta}$ from a sampled gaussian noise vector ${\bm{x}}^{T}$ as follows,

$$ ${\bm{x}}^{t-1}=\alpha({\bm{x}}^{t}-\gamma\epsilon_{\theta}({\bm{x}}^{t},c_{obs },c_{instruction},t)+\mathcal{N}(\mathbf{0},\sigma^{2}{\bm{I}})).$ $$

where $\alpha$, $\gamma$, $\sigma$ is the noise scheduler Ho et al. (). In our experiments, $\epsilon_{\theta}$ is to predict the noise that adds to the action.

### 3.3 PreTraining Data

To evaluate the proposed Diffusion Transformer Policy, we choose Open X- Embodiment datasets Padalkar et al. () for pretraining the model. We mainly follow Team et al. (); Kim et al. () to choose the datasets and set the weights for each dataset. We normalize the actions similarity to Padalkar et al. () and filter out outlier actions in the dataset. Additional details are provided in Appendix B.

### 3.4 Pretraining Details

We devise the proposed Diffusion Transformer architecture and evaluate the pretraining approach in the large cross-embodiment datasets Padalkar et al. (). We use the DDPM Ho et al. () diffusion objective in the pretraining stage with $T=1000$ for the Open X-Embodiment dataset Padalkar et al. (), while we set $T=100$ with DDIM Song et al. () for zero-shot evaluation to accelerate the inference. According to the preliminary experiment from Maniskill2 Gu et al. (), we use 2 observation images and predict 32 action chunks. We filter out the action chunks that include outlier values. We train the network with AdamW Loshchilov () by 100,000 steps. We set the learning rate of the casual transformer and Q-Former as 0.0001, the learning rate of the pretrained image tokenizers as 0.00001, and the batch size as 8902. More pretraining details are provided in the Appendix.

Figure: Figure 2: Our model is a Transformer diffusion structure. The model first incorporates a pretrained CLIP network to obtain instruction tokens. Meanwhile, we use the DINO-V2 Oquab et al. () model to encode image observations, followed by a Q-Former to query observation features for each image observation. Next, we concatenate instruction tokens, image observation tokens, timestep token, and noised action tokens together to construct a token sequence as the input for transformer network to denoise the raw actions.
Refer to caption: x2.png

## 4 Experiments

We evaluate the proposed methods with two baselines in three environments. We leverage Maniskill2 to present the ability of Diffusion Transformer Policy on large scale novel view generalization. Meanwhile, we demonstrate the generalization of the pretrained Diffusion Transformer Policy on CALVIN benchmark. Lastly, we further show the generalization of Diffusion Transformer Policy on Real Franka Arm.

### 4.1 Baselines

Discretization Action Head We implement the RT-1 Brohan et al. () style baseline models with a similar structure as ours. We maintain the Instruction Tokenization and Image Tokenization strategy in the proposed method. Different from ours, we follow RT-1 Brohan et al. () to discretize each dimension of the action into 256 bins, and leverage the transformer network to predict the action bin indexes. Following Brohan et al. (; ), we use cross-entropy loss to optimize the network.

Diffusion Action Head. We also implement a diffusion action head strategy  Team et al. () in our architecture. Specifically, we utilize a three-layer MLP network as our denoising network condition on the output of each action token embedding by the same transformer architecture as ours.

### 4.2 Maniskill2

Maniskill2 Gu et al. () is the next generation of the SAPIEN Maniskill benchmark Mu et al. (), which is widely used to evaluate the generalized manipulation ability of the embodied models. It contains 20 different manipulation tasks families and over 4 million demonstration frames with different settings, including rigid/soft body, single/dual arm, etc. Maniskill2 also provides a fast and easy way to change the camera view and replay the trajectories. It is useful for the researchers working on generalized policy.

Setup. In our experiments, we select 5 tasks (PickCube-v0, StackCube-v0, PickSingleYCB-v0, PickClutterYCB-v0, PickSingleEGAD-v0) from Maniskill2, and then construct a camera pool with 300,000 random cameras, then sample 20 cameras from the camera pool to render a trajectory each time, and finally obtain about 40K trajectories totally. Moreover, we split the dataset into training set and validation set according to a ratio of 19:1. During the splitting, it is guaranteed that the single trajectory rendered under different camera views will appear in either training set or validation set in order to avoid data leaking. Specially, there are 74 different categories to pick and place in the task family PickSingleYCB-v0. In addition to the mentioned rules, we ensure each category can be found in both training set and validation set. After that, we sample 100 trajectories for each task family randomly from the validation set, constructing a close loop evaluation dataset with 500 trajectories in total. While training, considering the balance between different task families, we adjust the number of data pieces to the same by simply copying the trajectories from task family with fewer trajectories originally. The ability of the model is measured by the success rate of each task family while executing the close loop evaluation dataset.

Optimization Details. We optimize the network with AdamW Loshchilov () by 50,000 steps on Maniskill2 and we set the learning rate as 0.0001. The number of training timesteps $T$ is 100 in Maniskill2 and the global batch size is 1024.

Comparisons. Table 1 compares the proposed method with discretized action head Brohan et al. (; ) and diffusion action head Team et al. (). Here, the proposed method use the same backbone and transformer architecture as the baseline methods. The experiments demonstrate Diffusion Transformer Policy achieves better results compared to Discretization Action Head strategy Brohan et al. () under the large scale novel view maniskill2 benchmark. Specifically, we observe the proposed Diffusion Transformer Policy achieves clear better performance compared to the baselines. Meanwhile, Diffusion Transformer Policy demonstrates better performance in more complex tasks, e.g., Diffusion Transformer Policy improves diffusion action head Team et al. () by 20% in task PickSingleYCB and by 12% task PickClutterYCB. Those experiments show that Diffusion Transformer Policy achieves better scalability in the large scale diverse datasets, and meanwhile achieves better generation in camera view generalization.

**Table 1: Comparision with two baseline methods on Maniskill2 under success rate. SingleYCB indicates PickSingleYCB, ClutterYCB indicates PickClutterYCB, SingleEGAD indicates PickSingleEGAD. Disc ActionHead indicates Discretized Action Head strategy Brohan et al. (), while Diff ActionHead shows Diffusion Action Head Team et al. ().**
| Method | All | PickCube | StackCube | SingleYCB | ClutterYCB | SingleEGAD |
| --- | --- | --- | --- | --- | --- | --- |
| Disc ActionHead | 30.2% | 41.0% | 33.0% | 22.% | 1.0% | 54.0% |
| Diff ActionHead | 58.6% | 86.0% | 76.0% | 37.0% | 24.0% | 70.0% |
| DiT Policy(ours) | 65.8% | 79.0% | 80.0% | 62.0% | 36.0% | 72.0% |

### 4.3 CALVIN

CALVIN (Composing Actions from Language and Vision) Mees et al. () is an open-source simulated benchmark to learn long-horizon language-conditioned tasks. CALVIN Mees et al. () includes four different scenes tagged as ABCD and presents a novel scene evaluation benchmark, ABC$\rightarrow$D, i.e., trained on environments A, B, and C and evaluated on environment D. The goal of CALVIN is to solve up to 1000 unique sequence chains with 5 distinct subtasks. The benchmark requires successfully solving the task sequence with 5 continuous subtasks, and one of the important evaluation indicators is the success sequence length.

Setup. In this section, we utilize CALVIN (ABC$\rightarrow$D) to evaluate the novel task generalization of Diffusion Transformer Policy architecture. Specifically, we directly apply the proposed method on CALVIN with a single static RGB camera and predict the end-effector action, including 3 dimensions for translation, 3 dimensions for Euler angles rotation and 1 dimension for gripper pose. We evaluate Diffusion Transformer Policy and Diffusion Action Head Team et al. () on CALVIN, and leverage the pretrained model on Open X-Embodiment to initialize the model for CALVIN.

Optimization Details. While training Calvin, 2 history images are used as input. For each iteration, the model predicts 10 future frames supervised by MSE loss. An AdamW optimizer is used together with a decayed learning rate with half-cycle cosine scheduler after several steps of warming up. The learning rate is initialized as 1e-4. We use 4 NVIDIA A100 GPUs(80GB) to train the model for 15 epochs with a global batch size of 128.

Comparisons. Table 2 presents the comparisons with previous methods on Calvin and the proposed methods. Without whistles and bells, the proposed Diffusion Tranformer Policy achieves the state-of-the-art results. Particularly, we only use RGB camera stream for observation. The superior demonstrates the effectiveness of Diffusion Transformer Policy. Meanwhile, the pretraining on Open X-Embodiment Datasets significantly facilitates the performance by 1.23, which demonstrates the transferability of Diffusion Transformer Policy. By contrast, the performance of diffusion action head is worse than Diffusion Transformer Policy by 0.45, though we load similar pretraining weights for diffusion head architecture.
Diffusion Transformer Policy can scale across different environments, e.g., transferring the knowledge from the diverse real datasets to the CALVIN dataset.

**Table 2: The comparisons with state-of-the-art approaches on Calvin Benchmark under success rate and average success length. ’S’ indicates Static RGB, ’G’ indicate Gripper RGB. ’SD’ indicate Static RGB-D, ’GD’ indicates Gripper RGB-D, ’P’ indicates Proprio, i.e., the observation arm position, ’C’ indicates camera parameters.**
| Method | Input | No. Instructions in a Row (1000 chains) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 1 | 2 | 3 | 4 | 5 | Avg.Len. |  |  |
| SPIL  Zhou et al. () | S,G | 74.2% | 46.3% | 27.6% | 14.7% | 8.0% | 1.71 |
| RoboFlamingo Li et al. () | S,G | 82.4% | 61.9% | 46.6% | 33.1% | 23.5% | 2.47 |
| SuSIE Black et al. () | S | 87.0% | 69.0% | 49.0% | 38.0% | 26.0% | 2.69 |
| GR-1 Wu et al. () | S,G,P | 85.4% | 71.2% | 59.6% | 49.7% | 40.1% | 3.06 |
| 3D Diffuser Ke et al. () | SD,GD,P,C | 92.2% | 78.7% | 63.9% | 51.2% | 41.2% | 3.27 |
| diffusion head w/o pretrain | S | 75.5% | 44.8% | 25.0% | 15.0% | 7.5% | 1.68 |
| diffusion head | S | 94.3% | 77.5% | 62.0% | 48.3% | 34.0% | 3.16 |
| Ours w/o pretrain | S | 89.5% | 63.3% | 39.8% | 27.3% | 18.5% | 2.38 |
| Ours | S | 94.5% | 82.5% | 72.8% | 61.3% | 50.0% | 3.61 |

### 4.4 Real Franka Arm

We finally evaluate the proposed method on the Real Franka Arm. Specifically, we train the model on Open X-Embodiment Datasets Padalkar et al. (), and evaluate it in our Franka Arm environment under zero-shot generalization and few-shot finetuning generalization.

setup. We set up the franka on the table with a black background. Meanwhile, we use a single third-person RGB camera about 1.5 meters away from the Franka Arm. Please refer to Figure 3 for the visualized demonstration. Considering the environment of our Franka setup is different from the scenes in Open X-Embodiment Stone et al. (), we mainly evaluate the proposed method on out-of-the-box generation and few-shot generation. In our experiments, we evaluate each model with the same scene, and the object is placed in 9 similar positions in a 9-grid format in front of the franka arm. Meanwhile, we maintain a small variance in those positions placing the objects for evaluation.

To evaluate the ability of the proposed model on new environments with a few demonstrations. We set 5 tasks, including “Pick up the green block”, “Pick up the kiwifruit”, “Pick up the banana”, “Pick up the tiny green block”, “Pick up the pink block”. Meanwhile, we collect 50 trajectories for each task in the first three tasks, while leaving the remaining two tasks (“Pick up the tiny green block” and “Pick up the pink block”) for out-of-distribution evaluation. Figure 3 presents the scenes and tasks. The image in Figure 3 is the exact image for our model. Our real-world environment is challenging since the object is small compared to the whole scene.

Figure: Figure 3: Illustration of Franka environment and corresponding objects.
Refer to caption: x3.png

Finetuning details. In our experiments, we finetune the proposed method on the real Franka Arm with Lora Hu et al. () and AdamW  for 10,000 steps. We set the number of timesteps as 100 for DDPM Ho et al. (), and batch size as 512. Meanwhile, we finetune all the networks with one observation and one step prediction.

Zero-shot and 10-shot Generalization We directly take the models pretrained on Open X-Embodiment to evaluate zero-shot and 10-shot generalization in our environments. We compare the proposed method with the discretization action head baseline Brohan et al. () (RT-1) that are optimized on the same dataset mixtures, as well as the released OpenVla Kim et al. () and Octo-base Team et al. () models. Besides, we evaluate it on three simple tasks: “pick up the green block”, “pick up the banana”, and “pick up the kiwifruit”. Table 3 shows OpenVLA and Octo without in-domain finetuning fail to pick up the object, while the proposed method can pick the green block with a success rate of 10%. We think it is because the proposed model is able to scale better the mixture data from Open X-Embodiment. Meanwhile, we find we still achieve similar performance to OpenVLA Kim et al. () (a 7B model) when finetuning it with only 10 samples on task “pick up the banana”. We provide a visualized analysis in Appendix C.1, where we find the OpenVLA and Octo fail to pick due to wrong grasp pose, while the proposed model can grasp the object with a small success rate.

**Table 3: Comparision with Baseline, Octo and Openvla on Real Franka Arm (pick up tasks) under success rate. Here the baseline is the model with discretization actions, similar to RT-1 Brohan et al. (), that we implement and train on Open X-Embodiment. 10-shot finetuning indicates we finetune the model with only 10 samples for each task.**
| Method | Zero-shot | 10-shot |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Pick Block | Pick Banana | Pick Kiwifruit | Pick Block | Pick Banana | Pick Kiwifruit |  |
| Baseline | 0 | 0 | 0 | 0 | 0 | 0 |
| Octo | 0 | 0 | 0 | 0 | 0 | 0 |
| OpenVLA | 0 | 0 | 0 | 0 | 33.3% | 0 |
| Ours | 11.1 % | 14.8 % | 0 | 14.8 % | 33.3% | 0 |

In domain Finetuning Generalization Table 4 presents the performance of the proposed Diffusion Transformer Policy compared to baseline methods. We observe different objects demonstrate various performances according to their attributes. The banana is the easiest object to pick up because the banana is longer, while kiwifruit is fat compared to other objects and all models achieve poor performance. The proposed Diffusion Transformer Policy effectively improves the diffusion action head according to Table 4. We find the discretized action head baseline achieves poor performance. Meanwhile, the Diffusion Transformer Policy is still able to pick up the novel object (e.g., the pink object) with a low success rate, while the baseline methods totally fail.

**Table 4: Comparision on Real-Franka tasks with few-shot fine-tuning. We represent the values in success rate. Block is “Pick up the green block”, Banana is “Pick up the banana”, Kiwi is “Pick up the Kiwifruit”, Tiny is unseen object task “Pick up the green tiny block”, Oval is unseen object task “Pick up the pink oval block”.**
| Method | All | Block | Banana | Kiwi | Tiny | Oval |
| --- | --- | --- | --- | --- | --- | --- |
| Discretized Action Head Brohan et al. () | 19.3% | 29.6% | 51.9% | 14.8% | 0 | 0 |
| Diffusion Action Head  Team et al. () | 34.8 % | 40.7% | 85.2% | 25.9% | 22.2% | 0 |
| DiT Policy (ours) | 46.9% | 55.6% | 90.3% | 44.4% | 37.0% | 7.4% |

**Table 5: Ablation on Maniskill2 about the number of history observation images and the length of the trajectory. #obs indicates the number of history observation images. #traj shows the length of trajectory, i.e., the sum length of observation and action prediction chunks. SingleYCB indicates PickSingleYCB, ClutterYCB indicates PickClutterYCB, SingleEGAD indicates PickSingleEGAD.**
| #obs | #traj | All | PickCube | StackCube | SingleYCB | ClutterYCB | SingleEGAD |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 2 | 2 | 40.8% | 68.0% | 54.0% | 33.0% | 9.0% | 40.0% |
| 2 | 4 | 51.6% | 81.0% | 69.0% | 44.0% | 11.0% | 53.0% |
| 2 | 8 | 62.4% | 89.0 % | 78.0% | 54.0% | 25.0% | 66.0% |
| 2 | 16 | 65.6% | 83.0% | 80.0 % | 70.0 % | 25.0% | 70.0% |
| 2 | 32 | 65.8 % | 79.0% | 80.0 % | 62.0% | 36.0 % | 72.0% |
| 1 | 32 | 61.6% | 78.0% | 76.0% | 64.0% | 24.0% | 66.0% |
| 1 | 1 | 51.0% | 79.0% | 66.0% | 42.0% | 19.0% | 49.0% |
| 3 | 3 | 35.4% | 54.0% | 49.0% | 27.0% | 5.0% | 42.0% |

Figure: Figure 4: Visualized comparison between Diffusion Transformer Policy and Diffusion Action Head baseline on Maniskill2 (PickClutterYCB). The first raw is Diffusion Tranformer Policy, while the second raw is the baseline method with Diffusion Action Head.
Refer to caption: x4.png

### 4.5 Ablation Study

In this section, we ablate some of the important designs of the model architecture, including the length of horizon, the length of observation, execution steps for evaluation on Maniskill2.

Trajectory length. The length of action chunks has an important effect on the performance of different tasks. Table 5 demonstrates that the performance increases as the increase of trajectory length. Meanwhile, we notice the performance of more complex tasks, e.g., PickClutterYCB, significantly increases with the increase of trajectory length, while the easy task, e.g., PickCube, maintains high performance after the trajectory length is more than 4. Meanwhile, the long horizon optimization significantly facilitates the performance since long horizon optimization is able to provide the target object position and help the model understand the localization of the object. For example, task PickClutterYCB with multiple YCB objects, requires the model to understand which one is the corresponding object.

Observation length. In our experiments, we find the length of history observation images also significantly affects the performance. At first, the performance significantly drops when we increase the length of observation history to 3. It might because it is more difficult for the model to converge with more observations since the number of corresponding image tokens also increases. Secondly, we observe using two image observations is more helpful for the performance when the prediction horizon is long. For example, when the length of trajectory is 32, the experiment with two observations achieves better performance. We think two observations can provide the visualized difference between two positions, and the difference of continunous gripper position indicates the action. The visualized difference is beneficial for future action prediction. However, for short horizon, the model majorly learns the projection from current observation to the corresponding actions.

Execution steps.
Since the proposed model is able to predict multiple future actions, we can execute multiple steps in one inference. Here, we ablate the effect of execution steps using a model that has a trajectory length of 32 for quick evaluation in Table 6. The ablation study shows that the short execution steps are slightly better longer execution steps, i.e., the farther away from the current frame, the worse the prediction quality.

**Table 6: The effect of the number of execution steps on Maniskill2. #steps indicates the number of steps that we execute each prediction.**
| #steps | 1 | 2 | 4 | 8 | 16 |
| --- | --- | --- | --- | --- | --- |
| All | 61.6% | 60.8 % | 60.6 % | 60.0 % | 58.0 % |

### 4.6 Visualized Comparison

Maniskill2. The proposed Diffusion Transformer Policy is able to model better action sequences. We conduct visualized analysis between the proposed method and the diffusion action head baseline on Maniskill2 in Figure 4. We select two trajectories from PickClutterYCB task, which is the most challenging task in Maniskill2. Figure 4 presents the grasp position is significantly important for picking up successfully, and the main reason that the baseline fails to pick up is the wrong grasp position. Meanwhile, we observe the major challenge of task PickClutterYCB is the grasping position prediction, especially when the target object is near by other objects. Compared to the diffusion action head baseline, Diffusion Transformer Policy is able to predict better action chunks for correctly picking the object with a suitable end-effector pose.

Real Franka Arm. We also illustrate the comparison between the Diffusion Transformer Policy and diffusion action head baseline on real Franka Arm in Figure 5. We demonstrate the experimental results under few-shot finetuning setting. We find the proposed method achieves better action prediction when the Gripper is approaching the object and finally picks up the small green object successfully, while the baseline fails to pick up due to the inaccurate grasp position. Meanwhile, we observe that failures are usually caused by a tiny position bias and we can not even directly discriminate the position by eyes from the image. For those cases, we argue the diffusion transformer policy has learnt better grasp position during the pretraining stage, and thus reduce failure rate due to the wrong grasp pose, while it is difficult for the diffusion action head. We demonstrates more failure cases to analyze the challenges in the Franka Arm.

Figure: Figure 5: Visualized comparison between Diffusion Transformer Policy and Diffusion Action Head strategy on Real Franka Arm (Pick up the green block). The first raw is Diffusion Tranformer Policy, while the second raw is the baseline method with Diffusion Action Head.
Refer to caption: x5.png

## 5 Conclusion, Discussion and Future Work

In this paper, we present a Diffusion Transformer architecture for generalist robot learning, named as Diffusion Transformer Policy. Diffusion Transformer Policy directly utilizes the large transformers as a denoising network to denoise the continuous actions conditioned on language instruction and image observations. The proposed architecture retains the scale attribute of the transformer, thus is capable of generalizing to different datasets with a unified architecture. Extensive experiments on Maniskill2, CALVIN, real Frank Arm demonstrate the effectiveness of the proposed method. Particularly, the proposed approach achieves state-of-the-art performance in CALVIN (ABC$\rightarrow$D) with only a single observation.

A limitation of the Diffusion Transformer Policy is that it requires multiple denoising steps during inference, which will impede the inference speed in the real application. In this paper, we focus on the modeling of complex and diverse actions. We think it is possible to improve the finetuning strategy with a few denoising steps to accelerate the inference speed. Meanwhile, the failure cases indicate it is important for the model to plan a right trajectory and grasp pose for picking object.

## Appendix A Additional Training and Model Details

The structure of our Model is shown in Figure 2. In this section, we will talk about some extra details of our model. The language instruction is encoded by a pretrained clip model and freeze the encoder in the training loop. We then resize the input images into 224$*$224 and feed it into a pretrained ViT model. The selected ViT is the base version of DinoV2. All the parameters in DinoV2 are trained. After the above process, we use a Q-Former to decrease the size of image features. The Q-Former is from scratched with a depth of 4. In each block, we insert the text token as a FiLM Condition to get the image features containing language information. The query length of the image features will reduce to 32 when out of the Q-Former. Then, we concatenate the processed text features and image features, together with the action nosied by a DDPM scheduler with 100 time-step. The multimodal inputs then pass through a causal Transformer and predict the added noise, which will execute on the robot arm after a series of post-processing. The Transformer is a from scratched Llama2-type model with 12 self-attention blocks. The hidden size is set to 768. All the modules mentioned are trained except the text encoder of clip. In summary, we have 334M parameters in total and 221M trained. This is pioneering to get this performance with such a small-sized model.

## Appendix B Additional Details about Pretraining Data

We choose 15 large datasets from Open X-Embodiment Padalkar et al. () as illustrated in Table 7. We mainly follow Team et al. (); Kim et al. () to set the weights. Following Kim et al. (), we further resize the image to the size of 224.

**Table 7: The training dataset mixture**
| Fractal Brohan et al. () | DobbE Shafiullah et al. () | Droid Khazatsky et al. () |
| --- | --- | --- |
| 16.15 | 1.94 | 13.69 |
| Robo Set Kumar et al. () | Viola Zhu et al. () | Kuka Kalashnikov et al. () |
| 2.99 | 1.30 | 17.47 |
| BridgeV2 Walke et al. () | NYU Franka  Cui et al. () | Furniture Bench Heo et al. () |
| 21.86 | 1.14 | 6.73 |
| StanfordHydra Belkhale et al. () | DLR EDAN  Quere et al. () | BerkeleyFanuc  Zhu et al. () |
| 6.11 | 0.08 | 1.07 |
| Jaco Dass et al. () | LanguageTable Lynch et al. () | toto Zhou et al. () |
| 0.67 | 6.01 | 2.78 |

**Table 8: The Ablation of learning rate scheduler on Calvin Benchmark.**
| strategy | No. Instructions in a Row (1000 chains) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| w lr decay | 94.5% | 82.5% | 72.8% | 61.3% | 50.0% | 3.61 |
| w/o lr decay | 91.8% | 80.0% | 68.0% | 56.9% | 45.9% | 3.43 |

## Appendix C Analysis in Calvin

As illustrated in the main paper, we use a common learning rate scheduler to decay the learning rate in the experiments in Calvin, rather than a fixed learning rate of 0.0001 in our pre-training stage. We demonstrate that this can slightly improve the performance in Table 8.

## Appendix D Visualization Analysis

### D.1 Zero-shot generalization

Figure 6 presents the proposed method is able to grasp the object, while OpenVLA and Octo fails. We observe OpenVLA and Octo fail to rightly approach the right grasp position. This experiment demonstrates the proposed Diffusion Transformer structure achieves more robust policy learning compared to discretized actions or diffusion action head. The denoising transformer model has built a better mapping from the image observation to corresponding action chunks.

Figure: Figure 6: Visualization of zero-shot genearalizaton of different models. The first raw is Diffusion Transformer Policy(ours). The second raw of demonstration is OpenVLA Kim et al. (), the third raw is Octo Team et al. (). Our Model can complete the tasks successfully while both of the others fails.
Refer to caption: x6.png

### D.2 Failure cases

We visualize some failure cases and make some further research to explore why model fails on the cases. Figure 7 shows the visualizations of some of these cases. We realize that, models can localize the object and move forward to it in nearly all of the cases, but it is in trouble with predicting a correct pose for the end-effector to grasp the object. Most of the cases fails due to the indiscernible differences on the image or the leak of depth information sometimes. We will explore how to increase the fineness of the prediction for the grasp pose and add the depth information to the model in the future.

Figure: Figure 7: Failure examples on Real Robot. Models can not predict a totally correct pose for grasping. Many trajectories fail due to indiscernible differences or the lack of depth information.
Refer to caption: x7.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - What is the point of using the Q-former?
- What happens if the diffusion action head is larger? It is not exactly clear what the main benefit of using the full causal Transformer as the diffusing model if the two approaches have roughly the same parameter count.
soundness: 2
strengths: - Extensive experimental results provided in several simulation environments and real-world task
- Reasonable set of baselines for comparison
- Reports state-of-the-art performance on the Calvin benchmark
- Clearly states training hyperparameters and experimental design to replicate their experiments
summary: This paper presents a different architecture / training paradigm for learning generalist robot policies from large scale data that uses a Causal Transformer model that denoises actions using a diffusion model training objective. Unlike other robot VLAs, their approach treats the entire model to denoise actions, while prior work only denoises with a small MLP layer. The authors present results in simulation environments, Maniskill 2 and Calvin, as well as real-world results on the Franka arm after pretraining on the Open X-Embodiment dataset.
weaknesses: - Several typos and sentence style issues throughout the paper
- Motivation for the architecture modification is not clearly stated, not clear from reading the introduction.
- Figure 2 is easy to follow, but visually poorly constructed.
- Action tokenization procedure is unclear. Are the continuous actions being converted into discrete bins similar to some of the prior works? If it is not, then what is action tokenization is Section 3.1?
- Some architectural design choices are not motivated, i.e. what is the purpose of the Q-Former?
- If the issue is scalability, what happens if you increase the size of the MLP in the diffusion action head setup? 3-layer MLP seems small. Making the MLP head larger could provide a more fair comparison.
- Improvement in Table 1 seems to come from the SingleYCB case while the improvements in other tasks are modest and even worse for the PickCube task.
- The zero success rate performance seems questionable for OpenVLA and Octo in Table 3.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. For Maniskill2 experiment, are all policies OXE pre-trained or not? Please specify the experiment settings clearly.
2. MDT [3] also explores the use of multimodal-conditioned diffusion transformer model. Please clarify the difference between MDT and DiT policy.
3. In the real-world experiments, why is all the network fine-tuned with one observation and one step prediction (L375-376)? Especially for one-step prediction, since it has been proved that action chunking [8] (i.e., multi-step action prediction) is an effective technique for robotic manipulation policies [6, 7, 8, 9, 10], and the trajectory length ablation (L418-420) also supports the argument.
4. For the real world experiments, “it is a challenging setup with small object (L370-371)”. How small are the objects?
5. For the real world experiments, how many evaluation rollouts are conducted for zero-shot generalization and few-shot generalization, especially for the few-shot generalization experiments? What is the success metric, does the policy need to perform successful grasps at the first time, or multiple trials for the one rollout is allowed?
soundness: 2
strengths: 1. The paper utilizes DiT in an in-context conditioning style for continuous action chunk denoising, incorporating the strong scaling ability of Transformers for generalist robot policy learning. While several concurrent works also uses DiT to improve the manipulation policy, the idea has its novelty. The policy architecture is also well explained and illustrated. 
2. Thorough ablations on trajectory length, observation length and execution steps.
summary: This paper presents Diffusion Transformer Policy, a DiT-based generalist diffusion policy. Different from previous diffusion policies, DiT policy directly utilizes a causal transformer to perform diffusion denoising conditioned via in-context multimodal tokens. Specifically, language instruction and image observations are tokenized by CLIP and DINOv2, respectively. Then, the image tokens are compressed with Q-Former, conditioned by the language tokens via FiLM. After that, language tokens, compressed image tokens, diffusion timestep and the noisy actions are passed to a causal transformer for action denoising. Experiments in simulation environments (ManiSkill2, CALVIN) and real-world robot platform (Franka robot) demonstrate the performance of the proposed policy.
weaknesses: 1. **Real-World Experiments**.
    - **Task Setting**. The task settings are too easy with only “picking” operations. Despite that the authors claim that it is a challenging setup with small object (L370-371), I believe that only picking small objects would not be so difficult. I would suggest adding real-world manipulation experiments besides picking (and pick-and-place), *e.g.*, open drawer/door, pouring (with rotation actions), and long-horizon tasks, etc.
    - **Few-Shot Setting**. Collecting 50 trajectories for each task in the first three tasks (L368-369) results in totally 150 demonstrations, which is not a few-shot settings. Previous works [1, 2] has proved that 50 demonstrations in the real world are sufficient for the policy to learn complex behaviors far beyond picking, even w/o pre-training. Generally, using less than 10 demonstrations can be considered as “few-shot” [3, 4] for robot learning.
    - **Few-Shot Baselines**. While the authors employ Octo [5] and OpenVLA [6] as zero-shot generalization baselines, it would be consistent and convincing to have both methods also as few-shot generalization baselines.
    - **Few-Shot Results**. The current “few-shot fine-tuning” results in Table 4 seems a little low for the simple picking operations. Please provide additional object informations (Question 4 & 5) for task difficulty evaluations. The videos in the supplementary results lack unseen object evaluations (task-4 and task-5) and baseline performance.
    
2. **Improper Baseline Choices in CALVIN Experiment**. Developing an expressive generalist policy with a better generalization performance seems to be the key motivation of the paper. The generalization ability of the DiT policy is heavily boosted by large-scale OXE [7] pre-training (in CALVIN experiment, the policy outperform its w/o pre-training variant by 1.2 in success sequence length, L022-023), which is reasonable since large-scale pre-training exposes the policy to more robot data. Therefore, the baselines in CALVIN experiment should also access to a similar amount of robot data. Specifically,
    - **(a)** If the authors want to emphasize on the generalization ability of the DiT policy itself (not boosted by the large-scale pre-training), they should compare the DiT policy w/o pre-training  with some of the following policies: 3D diffusor actors [8], MDT [9], MT-ACT [10], Multitask-DP [11], etc., all trained with similar amount of data.
    - **(b)** If the authors want to emphasize that the DiT policy is suitable for large-scale pre-training and able to gain better generalization performance after large-scale pre-training, they should compare the DiT policy w/ pre-training with some of the following policies: RT-1/2-X [7], Octo [5], OpenVLA [6], MT-ACT [10] w/ OXE pre-training, Multitask-DP [11] w/ OXE pre-training, all fine-tuning with the same amount of data.
    
    Also, it is worth noticing that MDT [9] performs well in CALVIN benchmark even w/o large-scale robot data pre-training, though it has not been evaluated under the ABC→D settings in its paper. Evaluating MDT under this setting can serve as a strong baseline in CALVIN benchmark, even for situation (b). For most image-based policies [5, 6, 7, 9, 10, 11], the observation modalities can be easily aligned with the DiT policy (single-view RGB images) for fairness considerations.
    
3. **Writings**. There are several typos and unclear statement in the paper:
    - L137: QFormer is not properly cited for the first appearance.
    - L276: Table I, Disc ActionHead, SingleYCB “22.%”. Missing number.
    - L311: “…, and 1 dimension for gripper pose.” Is it “gripper position”?
    - L358: the first letter of “setup” should be capital.
    - L361: “… different from the scenes in Open X-Embodiment (Stone et al.)” Incorrect citation.

[1] Tony Z. Zhao, et al., “Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware”, RSS 2023.

[2] Cheng Chi, et al., “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion”, RSS 2023.

[3] Vivek Myers, et al., “Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation”, CoRL 2024.

[4] Kourosh Hakhamaneshi, et al., “Hierarchical Few-Shot Imitation with Skill Transition Models”, ICLR 2022.

[5] Octo Model Team et al., “Octo: An Open-Source Generalist Robot Policy”, RSS 2024.

[6] Moo Jin Kim, et al., “OpenVLA: An Open-Source Vision-Language-Action Model”, arXiv 2024.

[7] Open X-Embodiment Collaboration, et al., “Open X-Embodiment: Robotic Learning Datasets and RT-X Models”, ICRA 2024.

[8] Tsung-Wei Ke, “3D Diffuser Actor: Policy Diffusion with 3D Scene Representations”, CoRL 2024.

[9] Moritz Reuss et al., “Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals”, RSS 2024.

[10] Homanga Bharadhwaj, et al., “RoboAgent: Generalization and Efficiency in Robot Manipulation via Semantic Augmentations and Action Chunking”, ICRA 2024.

[11] Huy Ha, et al., “Scaling Up and Distilling Down: Language-Guided Robot Skill Acquisition”, CoRL 2023.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the weaknesses section.
soundness: 2
strengths: The writing and organization of this paper are good.

The simulation experiments on Calvin and Maniskill are solid.
summary: This paper presents Diffusion Transformer Policy,  which directly utilizes the larger transformers as a denoising network to denoise the continuous actions conditioned on language instruction and image observations.

The proposed model pre-trained on the OpenX-Embodiment Dataset achieves better generalization performance.

The authors validate the effectiveness of diffusion modeling and the larger transformer policy in two simulation environments.
weaknesses: The fact that conducting pre-training on the Open-X dataset is a significant advantage of this paper. However, the real-world evaluation is too simple, only the  picking skill. This is not sufficient to prove the effectiveness of the methods in this paper compared to OpenVLA and OCto. If the evaluation scenarios include additional skills, it would be better.

It lacks the parameter scaling experiments of the Causal Transformer part. For example, the impact of different amounts of parameters and pre-training data on the results.

Why is it necessary to perform action decoding in a causal transformer with shared parameters? What if we simply use a larger DiT to replace the diffusion action head on the basis of Octo (such as [1] ) ? One advantage of doing so is that there is no need to pad the action vector. This paper lacks these experiments. I want to know which of these two design paradigms is better.

From the perspective of contribution to the community, this paper does not propose a powerful foundation model; at least the real-world experiments do not support this claim. On the other hand, the technology presented in this paper does not seem to be innovative; the initial diffusion policy could also utilize a transformer version.

[1] RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Can the authors clarify how their transformer-based approach differs from and improves upon previous work like the 'Diffusion Policy' paper and the "transformer" papers?
Especially in terms of architectural innovation?
soundness: 3
strengths: The paper shows remarkable generalization capabilities across different simulated and real-world environments. Utilizing a large-scale Transformer model for action denoising demonstrates superior scalability compared to previous small-scale MLP-based denoising models and discretization actions. It shows significant advantages over existing methods, even when facing changes in camera perspectives or environmental variations.
summary: This paper identifies the challenges posed by large-scale robotic datasets such as the Open X-Embodiment Dataset, where variations in camera perspectives and changes across different robot embodiments create an enormous action space that traditional MLP action heads struggle to capture, leading to limited generalization capabilities. To address this, the authors introduce a causal transformer as the backbone of their conditional diffusion model. This model denoises action blocks using both language instructions and visual observations as input in an end-to-end fashion. This approach maintains scalability across large, cross-embodiment datasets while achieving superior generalization performance in new camera perspectives and environments.
weaknesses: I greatly respect the experimental design and the impressive generalization results presented in this paper. This work offers meaningful contributions on the engineering front. However, I have some questions regarding the academic novelty of the research.
 
Regarding the entire pipeline, employing the casual Transformer to accept both text tokens and image tokens as multimodal inputs is an approach that has been extensively used in prior research. Moreover, this paper's encoding of natural language and image features directly relies on existing models (e.g., CLIP and DINOv2) without showing any distinct insights or contributions from this study.
 
The core idea of this work is to use the Transformer as an end-effector to replace the conventional MLP head to capture complex action spaces and achieve higher generalization capabilities. However, the advantages of the Transformer architecture in capturing complex action spaces have already been extensively explored in the "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion" paper. How does introducing the Transformer to improve scalability in this work differ from previous approaches?
 
In summary, this work significantly contributes to engineering and its applications within the related field. However, its academic innovation, particularly regarding theoretical contributions and architectural use, is not particularly strong.