Dynamic Rank Adjustment in Diffusion Policies for Efficient and Flexible Training

TL;DR

The paper introduces DRIFT, a framework for dynamically adjusting the complexity of diffusion policies during training, leading to more efficient and flexible robotic motion learning.

Contribution

It proposes a novel method using SVD for rank adjustment in diffusion policies, enabling seamless transition between offline and online training phases.

Findings

Improved sample efficiency in imitation learning.

Faster training times with minimal performance loss.

Effective dynamic complexity management in diffusion policies.

Abstract

Diffusion policies trained via offline behavioral cloning have recently gained traction in robotic motion generation. While effective, these policies typically require a large number of trainable parameters. This model size affords powerful representations but also incurs high computational cost during training. Ideally, it would be beneficial to dynamically adjust the trainable portion as needed, balancing representational power with computational efficiency. For example, while overparameterization enables diffusion policies to capture complex robotic behaviors via offline behavioral cloning, the increased computational demand makes online interactive imitation learning impractical due to longer training time. To address this challenge, we present a framework, called DRIFT, that uses the Singular Value Decomposition to enable dynamic rank adjustment during diffusion policy training. We implement and demonstrate the benefits of this framework in DRIFT-DAgger, an imitation learning algorithm that can seamlessly slide between an offline bootstrapping phase and an online interactive phase. We perform extensive experiments to better understand the proposed framework, and demonstrate that DRIFT-DAgger achieves improved sample efficiency and faster training with minimal impact on model performance. The project website is available at: https://apollo-lab-yale.github.io/25-RSS-DRIFT-website/.