Retrieve, Don't Retrain: Extending Vision Language Action Models to New Tasks at Test Time

2026-06-14

AI Summary

This work addresses the high cost of scaling conventional vision-language-action models to new tasks, which typically requires extensive teleoperated demonstrations and per-task fine-tuning. The authors propose a retrieval-augmented frozen policy: during training, the model is trained once on paired demonstrations from both target embodied agents and low-cost surrogates (e.g., human hand videos) and then fully frozen; at deployment, new tasks are incorporated simply by adding their trajectories to a retrieval pool, enabling action generation through retrieved examples without any retraining. This approach shifts the burden of task adaptation from parameter updates to data indexing, substantially improving cross-embodiment generalization. The method demonstrates strong performance on PushT and RoboTwin 2.0 benchmarks and successfully transfers to real robots, achieving notable results when integrated with video generative world models such as the Cosmos Policy.

Abstract

Extending a vision-language-action (VLA) policy to a new task typically requires task-specific teleoperated demonstrations and per-task fine-tuning, making adaptation costly in both data collection and compute. In this paper, we show that this target-side per-task adaptation cost can be replaced by retrieval. Our retrieval-augmented policy is trained once on paired demonstrations from the target embodiment (query) and a cheaper embodiment (pool, e.g., human-hand video), then frozen. New tasks are added at deployment by appending pool-side demonstrations to a retrieval pool. The frozen policy conditions on retrieved trajectories at every control step, so new tasks are absorbed by indexing data rather than updating parameters. Fine-tuning is needed only to take on a new, unseen embodiment, not for each new task. We show that retrieval improves policies beyond a specific backbone, including standard VLA policies, but its effect is especially pronounced in Cosmos Policy, a video-generation-based world-action model (WAM). In this setting, retrieval supplies coarse task progression, while the WAM's future-image objective provides an additional visual consistency signal that strengthens the retrieval-conditioned actions. On PushT, we study how retrieval provides a reusable high-level motion prior for cross-embodiment generalization to unseen goal angles, while on RoboTwin 2.0 our method outperforms cross-embodiment baselines on unseen tasks, and we additionally demonstrate the method on a real robot.

Problem

• vision-language-action
• task adaptation
• retrieval
• embodiment generalization
• policy extension

Innovation

• retrieval-augmented policy
• vision-language-action models
• cross-embodiment generalization
• world-action model
• test-time task extension