OSVI-WM: One-Shot Visual Imitation for Unseen Tasks using World-Model-Guided Trajectory Generation

TL;DR

OSVI-WM tackles one-shot visual imitation for unseen tasks by learning a world-dynamics predictor that, given a single expert demonstration and the agent's initial observation, forecasts a latent trajectory $\{r_2,\dots, r_M\}$ in a shared latent space. This trajectory is decoded into physical waypoints for a waypoint controller, enabling task execution without environment-specific fine-tuning. The method trains in-domain with joint supervision on predicted latents and waypoints using a stop-gradient WM loss and differentiable soft-DTW, and it supports re-planning at test time. Experiments on simulated and real-world robotic benchmarks show significant improvements over prior methods, including strong generalization to unseen tasks and robustness to embodiment mismatch.

Abstract

Visual imitation learning enables robotic agents to acquire skills by observing expert demonstration videos. In the one-shot setting, the agent generates a policy after observing a single expert demonstration without additional fine-tuning. Existing approaches typically train and evaluate on the same set of tasks, varying only object configurations, and struggle to generalize to unseen tasks with different semantic or structural requirements. While some recent methods attempt to address this, they exhibit low success rates on hard test tasks that, despite being visually similar to some training tasks, differ in context and require distinct responses. Additionally, most existing methods lack an explicit model of environment dynamics, limiting their ability to reason about future states. To address these limitations, we propose a novel framework for one-shot visual imitation learning via world-model-guided trajectory generation. Given an expert demonstration video and the agent's initial observation, our method leverages a learned world model to predict a sequence of latent states and actions. This latent trajectory is then decoded into physical waypoints that guide the agent's execution. Our method is evaluated on two simulated benchmarks and three real-world robotic platforms, where it consistently outperforms prior approaches, with over 30% improvement in some cases.

Figures (20)

• Figure 1: OSVI-WM infers the task from the expert demonstration and, along with the agent's observation "foresees" future latent states using a world-model-guided trajectory generation module. The predicted trajectory is decoded into physical waypoints for control.
• Figure 2: OSVI-WM: The expert demonstration frames $E_1,\dots,E_N$ and the agent’s current observation $R_1$ are encoded into a latent space using a ResNet encoder. A world-model-guided trajectory generation module predicts future latent states $r_2,\dots,r_M$, which are decoded into physical waypoints for control. During training, supervision is applied on the predicted waypoints and latent states.
• Figure 3: WM-Guided Trajectory Generation: Starting with the agent’s initial observation and the expert demonstration, future states are recursively predicted using action and world models.
• Figure 4: Simulation Environments: Test tasks are different from the ones used for training. Additionally, Pick-and-Place uses different embodiments for expert and agent.
• Figure 5: Real-World Environments: (a) Pick-and-place setup with expert (gray) and agent (white) Franka arms mounted at different locations. (b) Similar setup with a human expert and Franka agent. (c) Box push setup with human expert and Franka agent.