SAIL: Test-Time Scaling for In-Context Imitation Learning with VLM

TL;DR

SAIL introduces a test-time iterative refinement framework for robot imitation learning, leveraging Monte Carlo Tree Search and vision-language models to improve success rates in manipulation tasks.

Contribution

SAIL presents a novel test-time scaling approach using MCTS and trajectory refinement guided by vision-language models for improved robot imitation.

Findings

Success rates up to 95% on complex tasks

Test-time compute scaling improves performance

Effective in both simulation and real-world settings

Abstract

In-context imitation learning allows robots to acquire skills from demonstrations, yet one-shot trajectory generation remains fragile under environmental variation. We propose SAIL, a framework that reframes robot imitation as an iterative refinement problem capable of scaling with test-time compute. SAIL utilizes Monte Carlo Tree Search, where each node is a complete trajectory and edges correspond to trajectory refinements. The process is guided by three core components: an automated archive of successful trajectories for contextually relevant retrieval, a vision language model-based scoring mechanism for trajectory evaluation, and a step-level feedback that provides trajectory-aligned scores for iterative refinement. Experiments across six diverse manipulation tasks in simulation and real-world validation clearly demonstrate that increasing test-time compute consistently improves success rates, achieving up to 95% on complex tasks. Our results suggest that trajectory-level test-time scaling is a robust path toward more generalizable robotic agents.