PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jeongeun Park, Juhan Park, Taekyung Kim, Sungjoon Choi, Dongyoon Han, Sangdoo Yun

arXiv:2606.15631v1cs.ROcs.AI

TL;DR

Adapting vision-language-action policies to new tasks typically requires costly task-specific demonstrations and fine-tuning. RECAP replaces repeated retraining with retrieval over a reusable pool, and outperforms cross-embodiment baselines on unseen tasks while generalizing across benchmarks and a physical robot.

  • Problem

    New tasks for a new embodiment still require teleoperated demonstrations and per-task fine-tuning, causing adaptation costs to grow with each task.

  • Method

    RECAP trains once to bridge source and target embodiments, then freezes the policy and adds new tasks by retrieving demonstrations from an expandable pool.

  • Results

    RECAP consistently outperformed cross-embodiment baselines on unseen tasks, reaching 31.5% versus 26.0% on RoboTwin 2.0.

  • Takeaways & Limitations

    Indexing cheaper pool-embodiment data at deployment can substitute for collecting new target-robot demonstrations and support broader cross-embodiment generalization.

  • Takeaways & Limitations

    RECAP assumes shared end-effector action spaces and trajectory data, while substantially mismatched execution speeds or temporal scales can reduce residual reliability.

Abstract

from arXiv · show

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.

1 Introduction

RECAP replaces repeated target-embodiment teleoperation and per-task fine-tuning with retrieval from a reusable source-embodiment demonstration pool. A frozen retrieval-conditioned policy uses retrieved trajectories for high-level task progression while learning embodiment-specific corrections through a world-action model.

  • Motivation: Target-embodiment demonstrations are roughly 18× slower to acquire than equivalent human-hand demonstrations, while per-task Cosmos Policy fine-tuning costs roughly 24 GPU-hours/task.These data and compute costs make conventional adaptation expensive as new tasks are added.
  • RECAP: RECAP trains once to bridge source and target embodiments, then freezes the policy while new tasks are added by appending source-embodiment demonstrations to retrieval memory.Adaptation shifts from repeated optimization to retrieval over a reusable pool.
  • Method: In the retrieval-conditioned residual policy, retrieved trajectories provide coarse high-level motion and task progression, while the policy learns embodiment-specific dynamics for the target robot.The action latents are parameterized as a residual over retrieved trajectories.
  • Method: RECAP uses a world-action model, specifically Cosmos Policy, whose future-image prediction objective reinforces consistency between retrieved trajectories and predicted evolution.This visual consistency signal strengthens the retrieval-conditioned action prediction.

2 Related Work

Prior work includes world-action models that integrate actions into pretrained video generation, as well as retrieval-based imitation and cross-embodiment transfer. However, existing retrieval methods may still require task-specific training, motivating adaptation through retrieval without retraining.

  • World action models: World-action models fold actions into a pretrained video model, whose internet-scale backbone encodes semantics and physical dynamics while learning control.Examples include DreamZero, Cosmos Policy, mimic-video, and Fast-WAM.
  • Retrieval and Cross-embodiment policy transfer: Retrieval-based imitation adapts policies using relevant demonstrations instead of training from scratch for every task.This line of work includes retrieval-based imitation and within-embodiment methods such as FlowRetrieval and STRAP.
  • Retrieval and Cross-embodiment policy transfer: Cheaper human demonstrations reduce data costs, but existing approaches still fine-tune or train a policy for each new task.Hong et al. uses one human-hand demonstration to retrieve matching robot sub-trajectories and fine-tune efficiently for fast adaptation.

3 Problem Formulation

The problem is cross-embodiment imitation: a target query embodiment is controlled at deployment using demonstrations from a cheaper, never-deployed pool embodiment. At test time, new pool demonstrations are retrieved to predict query actions without collecting additional query data or updating model parameters.

  • Setting: The query embodiment is deployed, while the cheaper pool embodiment supplies demonstrations and is never deployed.The embodiments differ in geometry, contacts, dynamics, and data-collection cost.
  • Train and test access: Training uses paired query- and pool-side demonstrations, whereas test tasks are outside the fixed training distribution and provide only pool demonstrations.No additional query data is collected at test time.
  • Retrieval-conditioned action prediction: At each control step, the method retrieves the nearest pool state in feature space and feeds its H-step action chunk with the current query state to predict a query action chunk.The distance function is specified in Section 4.2, and the same retrieval rule is used during training and deployment.
  • Retrieval-conditioned action prediction: The deployment retrieval pool can be extended at any time without touching θ, enabling new demonstrations to be added without parameter updates.The model is rolled out on the target embodiment while observing the current query state at every control step.

4 Proposed Method

RECAP adapts frozen policies to new tasks at test time by retrieving relevant pool-embodiment demonstrations instead of retraining. Retrieved trajectories provide a shared high-level plan, while the policy learns embodiment-specific corrections and can absorb newly indexed demonstrations during deployment.

  • Test-Time Adaptation: RECAP replaces per-task retraining with retrieval, using pool trajectories to provide a shared high-level task plan while the policy predicts embodiment-specific corrections.The target and pool embodiments are assumed to agree on task requirements but differ in execution.
  • Test-Time Adaptation: New tasks are added by extending the external pool while keeping the policy frozen, allowing pool coverage to grow without updating model weights.At inference, newly added pool demonstrations replace the active pool and are reindexed.
  • World Action Model Conditioning: The current observation retrieves a matching state-action chunk, which is prepended with the query frame as temporal conditioning for jointly denoising future actions and observations.The language instruction enters through cross-attention, and this extends standard image-to-video conditioning without architectural modification.
  • World Action Model Conditioning: A single flow-matching loss jointly supervises action and future-image latents, aligning predicted actions with the next predicted state for more grounded outputs.The future-image objective supplies a visual consistency signal for learning action behavior.
  • Residual Action Parameterization: Residual action parameterization makes action latents encode how query actions differ from retrieved pool actions, focusing learning on embodiment-specific effects such as contact and gripper closure.State prediction provides a dense visual signal for learning this correction.
  • Retrieval: At every control step, retrieval selects a pool index whose surrounding chunk best matches the current query context, and retrieval reruns as the indexed pool grows.Candidate selection uses a composite descriptor involving the initial frame, language goal, object positions, and proprioception.

5 Experiments

Experiments evaluate RECAP across controlled PushT, multi-task RoboTwin 2.0, and real-robot settings, showing that frozen retrieval-conditioned policies absorb unseen tasks through added pool demonstrations. Retrieval improves cross-embodiment generalization, especially with a WAM backbone, while enabling deployment-time task expansion without retraining.

  • Experimental settings: RECAP is evaluated on PushT, RoboTwin 2.0 dual-arm manipulation, and a real robot using pool demonstrations from discs, UR5s, and human hands.The target policy is trained with paired demonstrations and then evaluated on unseen tasks or goal angles after test-time pool additions.
  • PushT: 6.0% to 34.9%: PushT unseen-angle success rises monotonically as the retrieval pool expands without parameter updates.Training uses paired triangle-disc demonstrations at ±45°, while evaluation covers nine goal angles, seven unseen during training.
  • PushT: Retrieval benefits both π0.5 and Cosmos backbones, with a larger effect for Cosmos because its future-image objective reinforces visual consistency with the retrieved trajectory.The WAM objective provides an additional signal for adapting coarse retrieved plans to embodiment-specific dynamics.
  • PushT: PushT angles achieve much of their final success before matching pool angles are added, indicating interpolation over neighboring demonstrations rather than nearest-example memorization.The retrieved trajectory supplies a coarse motion prior that the policy adapts to target embodiment dynamics.
  • RoboTwin 2.0: 43.5% seen and 31.5% unseen: RECAP leads RoboTwin 2.0, versus 32.5% and 26.0% for the strongest baseline.Aloha-Agilex is the target embodiment and UR5 is the retrieval pool; evaluation covers five seen and five unseen tasks.
  • RoboTwin 2.0: 9.0% to 31.5%: RoboTwin unseen-task success increases monotonically as the frozen policy receives an expanded retrieval pool.Once all five held-out tasks are retrievable, the frozen policy matches its supervised unseen-task average.
  • Real robot: On two real-robot held-out tasks, retrieval raises success from 10% and 0% without retrieval to 80% and 30% for placing the bottle and closing the cabinet.The policy is fine-tuned only on open-cabinet and receives 10 human-hand demonstrations per held-out task at test time.

6 Discussions

The paper extends a world-action-model policy to new tasks without retraining by adding pool-embodiment demonstrations at deployment. It improves cross-embodiment generalization on PushT, RoboTwin, and a physical robot, while exposing action-space, trajectory, and temporal-mismatch limitations.

  • Summary: The frozen policy absorbs new tasks by conditioning on a retrieval pool and predicting an embodiment-specific residual over retrieved trajectories.Deployment adaptation requires adding cheap pool-embodiment demonstrations rather than retraining the policy.
  • Summary: Across cross-embodiment PushT, RoboTwin, and a physical robot, retrieval improves generalization to unseen angles and tasks over cross-embodiment baselines.The policy is trained once and then extended through deployment-time retrieval.
  • Limitations and Future Work: Target and pool embodiments must share an end-effector action space because the residual refines retrieved chunks in a common low-level representation.Structurally different action spaces would require an embodiment-agnostic interface or learned action translator.
  • Limitations and Future Work: The pool must contain trajectories rather than video alone, limiting direct use of video-only sources such as raw human video.The limitation motivates future methods that can use broader video sources.
  • Limitations and Future Work: The residual becomes less reliable when retrieved motions differ substantially in execution speed or temporal scale, especially for larger chunks where errors accumulate.Robustness to significant temporal or dynamical mismatch remains an important direction for future work, including scaling beyond curated pools to in-the-wild video.

Appendix

The appendix collects additional results, analyses, and implementation details supporting the main paper. It expands PushT comparisons, action-parameterization ablations, and decoder analyses beyond the main text.

  • The appendix provides additional results, analyses, and implementation details that support the main paper.
  • Section A: Section A reports the full PushT baseline comparison across all nine goal angles and the complete 2 × 2 action-parameterization × next-state ablation.
  • Section B: Section B expands the decoder analysis with the full attention probe, no-retrieval comparison, and causal masking interventions summarized only in the main text.

A Additional PushT Results

Additional PushT results compare retrieval against three cross-embodiment baselines and dissect how action parameterization and next-state supervision affect unseen-angle success. Residual action parameterization benefits most from next-state prediction, while the same auxiliary loss is harmful without retrieval.

  • PushT generalization: Table 2 evaluates PushT success across goal angles, with training covering only ±45° and seven additional angles unseen during training.Success rate is measured in percent, with higher values better.
  • Cross-embodiment baselines: The comparison includes Cosmos Policy without pool data, Retrieval Only using the nearest pool action sequence, and Co-train (all) jointly training on target and pool trajectories.All three baselines share the retrieval method’s backbone but differ in how they incorporate pool data.
  • Action and next-state ablation: Residual action parameterization outperforms absolute action parameterization under both next-state-objective settings.The full ablation crosses absolute versus residual actions with the auxiliary next-state objective and includes a same-backbone non-retrieval baseline.
  • Action and next-state ablation: Next-state prediction adds +1.1 on unseen under absolute parameterization and +7.4 under residual, but reduces non-retrieval performance by −2.3.These results indicate that next-state supervision benefits specifically from the interaction between the WAM and retrieval.

B PushT Mechanism Analysis

On PushT, retrieval induces a two-stage decoder routing mechanism: L10 reads the retrieved task trajectory, while L15 commits using the policy’s own predicted endpoint. Both stages are absent without retrieval and are causally necessary for behavior.

  • Two-stage routing: The action slot routes retrieval in two stages: L10 intake reads the retrieved task region, and L15 commit uses the policy’s predicted end-of-chunk pose.The two attention peaks occur specifically at decoder layers 10 and 15.
  • Two-stage routing: At L10, attention reaches 5.1× uniform on the retrieved chunk’s initial T-block and 5.4× on its endpoint pose.The earlier frame anchors the block’s current location, while the later frame anchors the retrieved plan’s destination.
  • Two-stage routing: At L15, attention shifts from retrieved frames to the policy’s own predicted endpoint, reaching 3.9× uniform attention.This commit step makes retrieval more than additional context: it induces action from the policy’s own prediction.
  • No-retrieval comparison: Without retrieval, the same backbone keeps ROI ratios near 1.0, flattening the ours peaks of approximately 3.8 at L15 and 2.0–2.5 at L10.The matched comparison uses the five unseen angles common to both runs.
  • Causal necessity: Masking L10 retrieval attention on successful rollouts drops success by 20–80 percentage points, establishing causal necessity for retrieval intake.The ablation independently blocks action-to-retrieval cross-attention at L10 and L15 across seven unseen angles.

C PushT Failure Case Analysis

PushT failures separate into a universal L15 commit weakening and a goal-angle-dependent L10 intake distortion. These attention signatures correspond to distinct behavioral clusters and motivate different fixes for commitment versus adaptive retrieval control.

  • Attention-axis failures: L15 commit weakens on failures at all seven unseen angles, with success-minus-failure gaps from +0.14 at +30° and +60° to +1.01 at 0°.This is the universal, condition-independent failure signature.
  • Attention-axis failures: L10 intake shows condition-specific under-anchoring near seen angles, with failure gaps of −0.15 to −0.47 versus matched successes.Its failure direction depends on the goal angle’s distance from the seen training band.
  • Behavioral failure clusters: Failures split behaviorally into over-anchored rollouts with coverage ≥0.10 that move but never finish, and non-engaging rollouts with coverage <0.10 that remain largely inactive.The over-anchored cluster mirrors L10 over-anchoring and weakened L15 commit; non-engaging failures reflect collapse of the policy’s own prediction.
  • Implications for future work: Because L15 weakens uniformly, future work can directly supervise action-to-predicted-end attention to strengthen commitment under unseen conditions.The proposed intervention targets the universal direction of the L15 failure.
  • Implications for future work: Because L10 fails in opposite directions across angles, it requires adaptive control such as an attention temperature or gate.The proposed intervention addresses L10’s condition-specific sign and magnitude.

D RoboTwin Setup Details

RoboTwin fine-tuning uses 246 paired episodes across five Aloha-Agilex target tasks and UR5 retrievals, then evaluates on five held-out crosstask tasks. Test-time evaluation progressively expands the retrieval pool while keeping the policy frozen, after excluding catastrophically misaligned query episodes.

  • Train and test tasks: The policy is fine-tuned on five paired Aloha-Agilex–UR5 tasks totaling 246 training episodes.The tasks are Place Cans Plasticbox, Move Can Pot, Open Microwave, Grab Roller, and Pick Dual Bottles.
  • Train and test tasks: Crosstask evaluation uses five held-out tasks unseen during fine-tuning: Move Pillbottle Pad, Lift Pot, Click Bell, Hand-over Mic, and Place Bread Skillet.
  • Excluded query episodes: Four query episodes are excluded because cross-embodiment retrieval is catastrophically mis-aligned and would inject systematically wrong supervision.The exclusions are Move Can Pot episodes 32 and 45, Open Microwave episode 34, and Grab Roller episode 38.
  • Progressive retrieval pool: The frozen-policy pool-progression evaluation uses five strictly nested retrieval-pool levels, expanding from 11 to 35 UR5 tasks.Each level starts with the five training tasks and appends six further UR5 tasks in a fixed order; one held-out evaluation task becomes retrievable at each level.

E Retrieval Rule Details · F Hyperparameters

The retrieval rule uses an initial-scene trajectory filter followed by state-based subframe matching, with training-time action information removed at inference. This frozen, reindexed pipeline supports deployment-time pool extension without retraining, while hyperparameters vary by benchmark.

  • E Retrieval Rule Details: Stage 1 retains the top-K pool trajectories whose initial scenes best match the query using language, object positions, and proprioception.The filter improves scalability and precision by excluding mismatched instructions, layouts, and initial arm poses before step-wise retrieval.
  • E Retrieval Rule Details: Stage 2 retrieves the subframe minimizing a weighted match over object pose, DINO visual features, proprioception, and training-only action chunks.The action term aligns the retrieved demonstration with what it is about to execute, giving the WAM a tightly aligned training signal.
  • E Retrieval Rule Details: Squared L2 is used for geometric features, cosine distance for DINO embeddings, and a small visual weight keeps DINO as a sanity check rather than the dominant metric.On PushT, geometric terms dominate and DINO primarily breaks ties between geometrically comparable subframes.
  • E Retrieval Rule Details: At deployment, the unknown future query action is omitted from the subframe cost while the Stage 1 trajectory filter remains unchanged.Training with the richer state-and-action match transfers the train-to-inference mismatch into the action latents.
  • E Retrieval Rule Details: New pool-embodiment demonstrations are reindexed for both retrieval stages, adding unseen tasks by indexing data rather than updating model parameters.Because both stages use the current index, the pool can grow during a session and expand coverage without restarting.
  • E Retrieval Rule Details: The inference loop fixes the trajectory candidates once, re-runs subframe retrieval every control step, predicts a residual action chunk, executes its first K actions, and replans as the scene evolves.The policy therefore does not commit to one retrieved trajectory for the entire episode.
  • F Hyperparameters: Table 10 consolidates the backbone, RoboTwin 2.0 retrieval, and training settings; PushT instead uses the T-block pose with dominant geometric terms and low visual weight.RoboTwin settings include benchmark-specific feature weights, proprioception encoding, and trajectory prefilter details.
Loading 2606.15631v1…