PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Makoto Sato, Yusuke Iwasawa, Yujin Tang, So Kuroki

arXiv:2603.08269v2cs.ROcs.AI

TL;DR

One-shot robot trajectory generation can be fragile under environmental variation, motivating an approach that can refine trajectories during inference. SAIL combines trajectory-level MCTS with retrieved demonstrations and VLM-based trajectory and step-level scoring; increasing test-time compute improved success across six manipulation tasks, reaching 95% on HandOverBanana.

  • Problem

    One-shot trajectory generation can fail under new initial conditions, leaving in-context imitation dependent on a single prediction that cannot adjust during inference.

    Source: §I · p. 1

  • Method

    SAIL searches over complete trajectories with MCTS, using retrieved successful demonstrations, VLM scoring, and trajectory-aligned step feedback to guide refinement.

    Source: §I · p. 2

  • Results

    Success increased consistently with test-time compute across six manipulation tasks, reaching 95% on HandOverBanana and transferring with high success to a real block-moving task.

    Source: §I · p. 2, B. Scaling with Test-Time Compute · p. 5

  • Takeaways & Limitations

    The experiments suggest that scaling search at the trajectory level is a viable path toward more robust and generalizable robotic agents.

    Source: §I · p. 2, §VI · p. 7

Abstract

from arXiv · show

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.

I. INTRODUCTION

SAIL frames in-context robot imitation as trajectory-level iterative refinement that can scale with test-time compute, addressing the fragility of one-shot generation under environmental variation. It combines trajectory search, retrieval of successful demonstrations, and step-level feedback; experiments report improved success as compute increases.

I. INTRODUCTION

  • One-shot trajectory generation can fail under new initial conditions, and its success is capped by the model’s initial guess without an inference-time adjustment loop.Small errors in state estimation or object localization can cascade during execution.Source: §I · p. 1
  • SAIL treats complete trajectories as search objects in MCTS, retrieving visually similar successful demonstrations and using trajectory-aligned step feedback to refine proposals.The framework is presented as a way to scale imitation beyond one-shot prediction.Source: §I · p. 2
  • Success rates increased consistently with test-time compute across six manipulation tasks, with 95% success on HandOverBanana and high-success transfer to a real block-moving task.The evaluation covered simulation and a real robot using a trial-matched digital twin.Source: §I · p. 2

A. Trajectory Generation Using LLMs and VLMs

  • Prior trajectory-generation work used structured representations, waypoint constraints, or trajectory editing, while systematic test-time scaling across initial conditions remained largely unexplored.These approaches generally focused on improving a single prediction or lightly iterating for one environment instance.Source: A. Trajectory Generation Using LLMs… · p. 2

B. Planning and Test-Time Scaling

Prior planning and scoring work uses external evaluation, search, retrieval, and temporally structured VLM scores, but typically does not search continuous robot trajectories. The problem formulation considers generating a successful trajectory for an unseen seed from demonstrations available for a subset of seeds.

B. Planning and Test-Time Scaling

  • Prior test-time scaling methods focus mainly on task-level decisions or planning structures, whereas this work searches and evaluates continuous robot trajectories.Related approaches include external scoring, MCTS expansion, and iterative replanning.Source: B. Planning and Test-Time Scaling · p. 2

C. Scoring Using VLMs

  • VLM video scoring has progressed from binary success judgments to temporal progress and value estimates, and SAIL uses both trajectory-level scores and timestep-level feedback.The scoring signal supports search as well as trajectory refinement.Source: C. Scoring Using VLMs · pp. 2–3

A. Problem Formulation

  • The task is to generate a successful trajectory for an unseen environmental seed using initial demonstrations from a small subset of seeds.A seed indexes an initial condition, and the trajectory is represented as a sequence of end-effector states including position, orientation, and gripper status.Source: A. Problem Formulation · p. 3

B. Monte Carlo Tree Search

SAIL formulates trajectory refinement as independent MCTS trees per evaluation seed, with complete policy-VLM trajectories as nodes. Search uses prior-weighted selection, expands selected leaves with refined proposals, evaluates and backs up rewards, and returns the best trajectory within budget.

B. Monte Carlo Tree Search

  • Each seed has an independent MCTS tree whose nodes are complete trajectory proposals generated by the policy VLM from the initial state and node-specific demonstrations.Edges represent refinement operations that modify a previous trajectory.Source: B. Monte Carlo Tree Search · p. 3
  • Selection uses prior-weighted UCB to balance exploration and exploitation, incorporating a node-specific prior proportional to the node’s score.The loop comprises selection, expansion, evaluation, and backup.Source: B. Monte Carlo Tree Search · p. 3
  • At a selected leaf, SAIL samples B refined child trajectories, executes and scores them, then backs up rewards; search stops when budget is exhausted or success is found and returns the best trajectory encountered.Source: B. Monte Carlo Tree Search · p. 3

IV. METHODS

SAIL iteratively proposes, executes, scores, and refines complete trajectories in a simulator. Its components reuse successful trajectories, evaluate candidates with a VLM, and provide waypoint-aligned feedback for later refinements.

IV. METHODS

  • The policy VLM receives encoded initial state, reference demonstrations, and prior feedback, then proposes a complete trajectory that an inverse-kinematics controller executes waypoint by waypoint.Source: §IV · p. 3
  • A shared archive stores successful rollouts across seeds and retrieves visually similar trajectories as in-context demonstrations for candidate generation.Source: Figure 2 · p. 4, §IV · p. 4
  • A scoring VLM evaluates executed proposals with a scalar node value and aligns per-frame progress scores to waypoints to provide step-level refinement feedback.The overview also describes decomposing the task into ordered subtasks from one demonstration.Source: §IV · pp. 3–4, Figure 2 · p. 4

A. Archive Retrieval for In-Context Demonstrations

SAIL maintains an archive of successful trajectories across seeds so scene-compatible examples can be retrieved during search. For each target scene, it selects visually similar archive entries as in-context references.

A. Archive Retrieval for In-Context Demonstrations

B. Trajectory Scoring for Node Evaluation

SAIL uses a VLM to score simulated rollout videos, avoiding labor-intensive task-specific reward design. It converts progress through ordered subtasks into a scalar node reward for search.

B. Trajectory Scoring for Node Evaluation

C. Step-Level Feedback for Trajectory Refinement

SAIL supplies step-level feedback by aligning scorer progress with trajectory waypoints. The policy VLM is instructed to preserve high-scoring segments and modify low-scoring ones during refinement.

C. Step-Level Feedback for Trajectory Refinement

V. EXPERIMENTS

The experiments evaluate SAIL in simulation and the real world. They examine whether more MCTS nodes improve success, how retrieval and step-level feedback contribute to refinement, and whether simulation-refined trajectories transfer to real-world execution.

V. EXPERIMENTS

  • The experiments ask whether increasing the test-time compute budget, measured in MCTS nodes, improves task success rates.Source: §V · p. 5
  • The experiments examine how automated archive retrieval and step-level feedback contribute to iterative trajectory refinement.Source: §V · p. 5
  • The experiments test whether trajectories refined in simulation can transfer successfully to the real world.Source: §V · p. 5

A. Tasks and Experimental Setup

Across six simulated manipulation tasks, the experiments test compute scaling and compare search, retrieval, and feedback choices. Success rises with additional MCTS nodes, while similarity-based retrieval and step-level scores outperform the evaluated alternatives.

A. Tasks and Experimental Setup

  • Simulation evaluation uses 20 random-seed instances per task, with one or a few successful demonstrations from distinct seeds; real-world trials randomize block and bowl positions.The six simulation tasks use ground-truth verification, and the BlockIntoBowl real-world setup uses six trials.Source: A. Tasks and Experimental Setup · p. 5

B. Scaling with Test-Time Compute

  • Average success rises from 25% with a single rollout to 73% with 45 MCTS nodes.BowlOnRack reaches 100% by 6 nodes and remains there at larger budgets.Source: B. Scaling with Test-Time Compute · p. 5, Table I · p. 6
  • At 15 nodes, MCTS averages 65% success, versus 51% for breadth-first search and 37% for depth-first search.Breadth-first search performs better on LaptopClose, where stochastic exploration handles the screen-edge geometry.Source: B. Scaling with Test-Time Compute · p. 6

C. Effects of Retrieval Strategy and Feedback Modality

  • With K = 1, similarity-based retrieval averages 65% success, versus 45% for fixed demonstrations and 50% for random retrieval.The comparisons use a fixed search budget of 15 nodes.Source: C. Effects of Retrieval Strategy and… · p. 6
  • Increasing fixed demonstrations from K = 1 to K = 3 raises average success only from 45% to 49%, while random retrieval rises from 50% to 53%.Similarity-based retrieval with one demonstration reaches 65% on average; on DrawerOpen it reaches 40%, versus 10% for three fixed demonstrations.Source: C. Effects of Retrieval Strategy and… · p. 6
  • Step-level feedback averages 65% success, compared with 45–48% for dense unscored histories; sparse feedback reaches 15% on DrawerOpen and 30% on MarkerRemoveLid, versus 40% on both with step-level feedback.The step-wise annotations use the same scoring function as rollout evaluation, aligning refinement feedback with the objective.Source: C. Effects of Retrieval Strategy and… · p. 6

D. Real-World Validation

The real-world validation reconstructs the BlockIntoBowl workspace as a trial-specific digital twin, searches for a trajectory in simulation, then transfers it to the robot. The reconstruction uses RGB-D observations and object-pose estimates to match the simulated scene to the physical setup.

D. Real-World Validation

  • The BlockIntoBowl task is depicted as grasping the blue block, moving it to the bowl, and releasing it; the figure pairs simulation on the left with the real setup on the right.The physical workspace is reconstructed for each trial before simulation search and robot execution.Source: D. Real-World Validation · p. 6, Figure 3 · p. 7

1) Real2Sim Pipeline:

  • Real2Sim uses synchronized RGB-D observations to estimate object poses and place corresponding assets in a scene-matched simulator.GroundingDINO and SAM2 provide object masks; aligned depth points, calibrated camera parameters, and an ArUco marker support pose estimation and coordinate mapping.Source: 1) Real2Sim Pipeline · pp. 6–7

2) Sim2Real: Trajectory and Policy Deployment:

In six real-world BlockIntoBowl trials, both MCTS-based trajectory refinement and a policy distilled from its rollouts succeed in five. The distilled policy substantially reduces average execution time, while the paper identifies conventional physics simulators for data collection as a scope boundary.

2) Sim2Real: Trajectory and Policy Deployment:

  • Both MCTS-based refinement and the distilled policy succeed in 5 of 6 real-world trials; average execution time is 644.72 seconds for MCTS and 72.306 seconds for the policy.The policy is trained by behavior cloning on 120 successful MCTS-generated trajectories and deployed by selecting a successful rollout in the trial-specific Real2Sim scene.Source: Table III · p. 7, 2) Sim2Real: Trajectory and Policy… · p. 7
  • MCTS searches the trial-specific simulation scene with up to 15 nodes, stopping early when it finds a successful trajectory to execute on the robot.The authors attribute the single failed trial to pose-estimation errors and unmodeled real-world contact dynamics.Source: 2) Sim2Real: Trajectory and Policy… · p. 7

VI. CONCLUSIONS

  • The conclusion reports that higher compute budgets raise average success across six manipulation tasks; data collection in this work relies on conventional physics simulators.The authors identify integration with Gaussian-Splatting digital twins as a possible advancement for addressing remaining visual-fidelity and contact-dynamics gaps.Source: §VI · p. 7
Loading 2603.08269v2…