ImitDiff: Transferring Foundation-Model Priors for Distraction Robust Visuomotor Policy

TL;DR

ImDiff addresses the challenge of distraction-robust visuomotor imitation by transferring vision-language foundation-model priors into pixel-level semantic masks that guide a dual-resolution perception pipeline. A consistency-driven diffusion transformer head then maps semantically grounded visual features to real-time actions, achieving large speedups without sacrificing performance. The approach combines an open-vocabulary detect-track-segment pipeline, a dual-encoder fusion scheme, and PF-ODE-based trajectory conditioning to deliver robust zero-shot generalization to novel objects and distractions. Extensive simulation and real-world experiments demonstrate superiority over state-of-the-art vision-language and visuomotor baselines, with significant gains in complex scenes and under strong visual distractors, along with practical inference-time improvements.

Abstract

Visuomotor imitation learning policies enable robots to efficiently acquire manipulation skills from visual demonstrations. However, as scene complexity and visual distractions increase, policies that perform well in simple settings often experience substantial performance degradation. To address this challenge, we propose ImitDiff, a diffusion-based imitation learning policy guided by fine-grained semantics within a dual-resolution workflow. Leveraging pretrained priors of vision-language foundation models, our method transforms high-level instructions into pixel-level visual semantic masks. These masks guide a dual-resolution perception pipeline that captures both global context (e.g., overall layout) from low-resolution observation and fine-grained local features (e.g., geometric details) from high-resolution observation, enabling the policy to focus on task-relevant regions. Additionally, we introduce a consistency-driven diffusion transformer action head that bridges visual semantic conditions and real-time action generation. Extensive experiments demonstrate that ImitDiff outperforms state-of-the-art vision-language manipulation frameworks, as well as visuomotor imitation learning policies, particularly under increased scene complexity and visual distractions. Notably, ImitDiff exhibits strong generalization in zero-shot settings involving novel objects and visual distractions. Furthermore, our consistency-driven action head achieves an order-of-magnitude improvement in inference speed while maintaining competitive success rates.

Figures (4)

• Figure 1: ImitDiff leverages the priors of vision-language foundation models to transform high-level instructions into pixel-level semantic masks, which finely constrain dual-resolution visual features within the dual-resolution workflow. Based on this, the consistency-driven diffusion transformer (DiT) action head generates executable trajectories in real time under conditional supervision.
• Figure 2: Overview of ImitDiff. Our framework comprises three components: a) Transferring Foundation-Model Priors. Given a user instruction and the initial observation, a VLM identifies task-relevant objects through a chain-of-thought process. An open-vocabulary detect-track-segment pipeline then produces visual semantic masks, which are injected into a shared latent space via a semantic-injection encoder, thereby transferring foundation-model priors to the dual-resolution visual features. b) Dual-Resolution Workflow. For each camera view, both high and low resolution observations are obtained and encoded by a dual-encoder system. Low-resolution features query candidate regions within high-resolution features at the patch level via attention, maximizing multi-scale information extraction while retaining efficiency. c) Consistency-Driven DiT Action Head. An EDM action decoder is first trained as a teacher model conditioned on visual observations and proprioceptive inputs. A CTM student decoder is then distilled along the same PF-ODE trajectory, achieving substantially faster inference while preserving task success rates.
• Figure 3: Tasks in simulation and real-world experiments. Real-world trials include foreground distractors at increasing strengths.
• Figure 4: Objects visualization for the training and test sets.