Decoupled Action Head: Confining Task Knowledge to Conditioning Layers
Jian Zhou, Sihao Lin, Shuai Fu, Qi WU
TL;DR
The paper tackles data scarcity in robotic Behavior Cloning by introducing a decoupled training paradigm that pretrains a general Action Head on observation-free trajectories and then freezes it while adapting task-specific conditioning. It demonstrates two-stage training (Stage 1: JP-conditioned action head pretraining; Stage 2: conditioning finetuning with a frozen backbone) and investigates backbone designs, notably DP-T-FiLM and a lightweight DP-MLP, showing substantial training efficiency gains and cross-task generalization. A key finding is that task knowledge concentrates in the conditioning modules, allowing a lightweight action generator backbone to perform competitively. The work offers practical guidance on efficient policy learning for manipulation, revealing that FiLM-based conditioning is crucial in decoupled setups and that lightweight backbones can deliver large speedups without sacrificing much performance.
Abstract
Behavior Cloning (BC) is a data-driven supervised learning approach that has gained increasing attention with the success of scaling laws in language and vision domains. Among its implementations in robotic manipulation, Diffusion Policy (DP), with its two variants DP-CNN (DP-C) and DP-Transformer (DP-T), is one of the most effective and widely adopted models, demonstrating the advantages of predicting continuous action sequences. However, both DP and other BC methods remain constrained by the scarcity of paired training data, and the internal mechanisms underlying DP's effectiveness remain insufficiently understood, leading to limited generalization and a lack of principled design in model development. In this work, we propose a decoupled training recipe that leverages nearly cost-free kinematics-generated trajectories as observation-free data to pretrain a general action head (action generator). The pretrained action head is then frozen and adapted to novel tasks through feature modulation. Our experiments demonstrate the feasibility of this approach in both in-distribution and out-of-distribution scenarios. As an additional benefit, decoupling improves training efficiency; for instance, DP-C achieves up to a 41% speedup. Furthermore, the confinement of task-specific knowledge to the conditioning components under decoupling, combined with the near-identical performance of DP-C in both normal and decoupled training, indicates that the action generation backbone plays a limited role in robotic manipulation. Motivated by this observation, we introduce DP-MLP, which replaces the 244M-parameter U-Net backbone of DP-C with only 4M parameters of simple MLP blocks, achieving a 83.9% faster training speed under normal training and 89.1% under decoupling.