MTDP: A Modulated Transformer based Diffusion Policy Model

TL;DR

This work tackles the challenge of effectively integrating guiding conditions in diffusion-policy learning for robot manipulation, focusing on Transformer-based architectures. It introduces Modulated Attention, embedded in MTDP for Transformer backbones and MUDP for UNet backbones, to fully leverage conditioning signals during generation. The approach yields superior performance over DP-Transformer and DP-DIT across six tasks (notably Toolhang with +12%), and demonstrates that DDIM-based variants MTDP-I and MUDP-I substantially accelerate sampling with minimal loss in performance. The results suggest that Modulated Attention generalizes across architectures and can accelerate diffusion-based policy learning with potential broader applicability beyond robotics.

Abstract

Recent research on robot manipulation based on Behavior Cloning (BC) has made significant progress. By combining diffusion models with BC, diffusion policiy has been proposed, enabling robots to quickly learn manipulation tasks with high success rates. However, integrating diffusion policy with high-capacity Transformer presents challenges, traditional Transformer architectures struggle to effectively integrate guiding conditions, resulting in poor performance in manipulation tasks when using Transformer-based models. In this paper, we investigate key architectural designs of Transformers and improve the traditional Transformer architecture by proposing the Modulated Transformer Diffusion Policy (MTDP) model for diffusion policy. The core of this model is the Modulated Attention module we proposed, which more effectively integrates the guiding conditions with the main input, improving the generative model's output quality and, consequently, increasing the robot's task success rate. In six experimental tasks, MTDP outperformed existing Transformer model architectures, particularly in the Toolhang experiment, where the success rate increased by 12\%. To verify the generality of Modulated Attention, we applied it to the UNet architecture to construct Modulated UNet Diffusion Policy model (MUDP), which also achieved higher success rates than existing UNet architectures across all six experiments. The Diffusion Policy uses Denoising Diffusion Probabilistic Models (DDPM) as the diffusion model. Building on this, we also explored Denoising Diffusion Implicit Models (DDIM) as the diffusion model, constructing the MTDP-I and MUDP-I model, which nearly doubled the generation speed while maintaining performance.

Figures (5)

• Figure 1: The architecture of MTDP. Both the timestep and image features serve as guiding conditions and are input into the encoder. During the training phase, guiding conditions combined with noisy action features and fed into the decoder; while during the generation phase, combined with noise. Through iterative denoising, executable actions are obtained.
• Figure 2: We experimented with four different structures for Modulated Attention. (a) M-SelfAttention represents a model with both modulated self-attention and cross-attention; (b) M-CrossAttention represents a model with both self-attention and modulated cross-attention; (c) DIT-SelfAttention is DIT, represents a model with only modulated self-attention; (d) DIT-CrossAttention represents a model with only modulated cross-attention. We found that (a) performed the best.
• Figure 3: Tasks used for the experiment.
• Figure 4: Comparison of task execution performance between existing Transformer-based diffusion policy models and our proposed model.
• Figure 5: Average success rates of the MUDP-I model on the PushT and Square tasks at different sampling timesteps.