Efficient Training of Generalizable Visuomotor Policies via Control-Aware Augmentation

Improving generalization is a key challenge in Embodied AI, where obtaining large-scale datasets from diverse scenarios is costly. Visuomotor policies trained with weak augmentations provide only marginal improvements when applied to new environments. Strong augmentations, such as random overlay, can disrupt task-relevant information and degrade performance. To overcome these challenges, we introduce EAGLE-an Efficient trAining framework for GeneraLizablE visuomotor policies. EAGLE enhances generalization by applying augmentation only to control-related regions using a self-supervised, control-aware mask. It also boosts training efficiency and stability by transferring knowledge from an expert to a student policy, enabling deployment in new environments without further fine-tuning. Experiments on the DMControl Generalization Benchmark (DMC-GB) demonstrate the effectiveness of our approach. Project website at https://vrl-eagle.github.io/