Dynamic Rank Adjustment in Diffusion Policies for Efficient and Flexible Training

2025-02-06

AI Summary

To address parameter redundancy, high training overhead, and poor adaptability to online interactive imitation learning in diffusion-based robotic policy learning, this paper proposes DRIFT-DAgger—a novel framework featuring the first SVD-based dynamic rank adjustment mechanism. This mechanism enables on-the-fly scaling of trainable parameters during training, facilitating seamless transitions between offline behavioral cloning and online DAgger-style interactive learning. By dynamically balancing representational capacity and computational efficiency, DRIFT-DAgger achieves substantial improvements across multiple robotic manipulation tasks: average training speedup of 2.1× and 37% higher sample efficiency, with negligible performance degradation (<0.8%). Its core contribution lies in pioneering the integration of low-rank dynamic updates into diffusion policy training—establishing a new paradigm for efficient, adaptive embodied policy learning.

Technology Category

• Intelligent Robots: Learning & Optimization for ROB
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• Dynamic rank adjustment in diffusion policies
• Balancing computational efficiency and representational power
• Enabling efficient online interactive imitation learning

Innovation

• Dynamic Rank Adjustment via SVD
• DRIFT-DAgger for imitation learning
• Balancing computational efficiency and representational power