Freeze the Policy, Infer the Goal: Cross-Domain Imitation with World Models
Abstract
The paradigm of pretraining foundation models and subsequently finetuning them for downstream tasks has emerged as the prevailing approach in the decision-making community. Although imitation learning offers a powerful mechanism to adapt pretrained models using expert demonstrations, agents frequently struggle when data is collected under fundamentally different morphologies or viewpoints. This challenge, known as cross-domain imitation learning, typically requires complex algorithmic design or kinematic retargeting. In this paper, we propose Goal Inference Imitation Learning (GIIL), a framework that reframes cross-domain imitation as an efficient adaptation problem. First, we pretrain a goal-conditioned multi-task policy entirely on unlabelled, offline embodiment data using world models and self-supervised distance-based rewards. During the imitation phase, rather than updating the policy network, we freeze the policy and adapt the agent by inferring a continuous goal vector that best matches the expert demonstration. This inference is guided by an Optimal Transport reward evaluated within the world model's imagination. By decoupling the acquisition of embodiment-specific motor skills from the high-level intentions of the demonstrator, GIIL enables zero-shot transfer and lightweight finetuning without catastrophic forgetting. Empirical results on the DeepMind Control Suite demonstrate that GIIL outperforms baseline approaches across cross-embodiment and cross-viewpoint settings.