In-context imitation learning enables robots to adapt to new tasks from a small number of demonstrations without additional training. However, existing approaches typically condition only on state-action trajectories and lack explicit representations of task intent. This limitation hinders performance in complex and ambiguous task settings where the same actions may be consistent with different objectives. To address this, we present In-Context Imitation Learning with Visual Reasoning (ICLR), a novel framework that augments demonstration prompts with structured visual reasoning traces representing anticipated future robot trajectories in image space. ICLR also jointly learns to generate reasoning traces and low-level actions within a unified autoregressive transformer, enabling the model to mimic not only action prediction but also the reasoning process that leads to those actions. We extensively evaluate ICLR in both simulation and real-world manipulation tasks and demonstrate consistent improvements in success rates and generalization to unseen tasks and novel object configurations compared to other in-context imitation learning methods. These results suggest that incorporating embodied visual reasoning represents a promising direction for enhancing the robustness and generalization of robotic in-context learning systems.
Figures & tables
Fig. 1: General framework overview. Our method augments prompt demos with keypoint-based visual reasoning traces in the image space, shown above with the overlaid polyline in the middle column. During inference, the model also performs visual reasoning before predicting the subsequent low-level robot action. The task’s language description is included for clarity.
Fig. 2: Method overview. (A) To generate the visual reasoning trace at a given time step, we uniformly sample five third-view images from that time step to the end of the trajectory and use Molmo2 to predict the gripper’s pixel location in each image. (B) Multi-view camera observations and proprioceptive states are encoded by a state encoder to produce state tokens fs . Visual reasoning traces are embedded by a reasoning encoder to produce reasoning tokens fr , and actions are embedded by an action encoder to produce action tokens fa . (C) These modality-specific tokens are interleaved and fed into a causal transformer, which autoregressively predicts the next reasoning trace followed by the corresponding action. During training, teacher forcing is applied over reasoning and action tokens. In inference, the model first generates a reasoning trace and then produces the action in a closed-loop manner.
LIBERO-90
Method
LIBERO-Object
Kitchen
Living
Study
Avg.
ICRT [ 5 ]
44.44
0.89
18.93
0.83
16.27
TO Dropout
62.22
17.56
16.67
25.00
30.36
TO
54.44
12.11
12.80
29.33
27.17
Ours Dropout
70.89
60.22
38.93
46.17
54.05
Ours
70.00
20.00
11.20
32.17
33.34
TABLE I: Simulation success rates (%). “Dropout” models are models that learn to generate visual reasoning traces in training but skip the reasoning steps for the target trajectory in inference. Results are reported on 3 unseen tasks of LIBERO-Object and 15 unseen tasks of LIBERO-90 (6 of kitchen scenes, 5 of living room scenes, and 4 of study scenes). In each column, the highest score is bolded , and the second-best performance is underlined .
Poking
Pick-and-Place
Method
Hippo
Dumpling
Jaguar
Monkey
Lion
Potato
Avg.
Dumpling to Red box
Zebra to Blue bowl
Tomato to Grey bowl
Monkey to Red box
Lion to Blue bowl
Potato to Grey bowl
Avg.
ICRT [ 5 ]
20
40
80
60
40
50
48.33
15
35
30
15
25
5
22.50
TO Dropout
0
60
40
40
60
20
40.00
55
50
30
5
40
10
31.67
TO
30
60
80
30
40
70
51.67
65
45
45
55
20
15
40.83
Ours Dropout
70
60
60
60
70
70
65.00
55
55
50
45
30
45
46.67
Ours
60
80
80
70
70
70
71.67
65
70
65
50
45
65
60.00
TABLE II: Real-world task success rates (%). Every model has 10 rollouts for each task. See Section V-C for our detailed evaluation methodology.
Fig. 3: Real robot setting. We use a Franka Research 3 robot arm equipped with a UMI gripper. Visual observations are captured by two RealSense cameras. Teleoperation for data collection and test-time prompt demonstration recording is performed using a GELLO system. Testing objects appearing in training episodes are shown in the bottom-left box, while completely unseen testing objects are shown in the bottom-right box.
Fig. 4: Qualitative results. Rollout examples of our complete ICLR model in simulation (first two rows) and real-world settings (two bottom rows). All presented visual traces are predicted by our model.
Fig. 5: Three types of prompt demonstrations. The task of picking up the tomato and putting it in the grey bowl is selected.
1 Demo
Method
0 Distr.
1 Distr.
Distr. Receptacle
2 Demos
3 Demos
ICRT [ 5 ]
35
30
40
35
45
TO Dropout
45
25
20
25
45
TO
40
40
50
50
30
Ours Dropout
35
35
40
55
40
Ours
55
65
55
70
65
TABLE III: Results of different prompt types (%). “Distr.” stands for distractor. Success rates are calculated over 10 trials for each experiment on the task of picking up the tomato and putting it in the grey bowl.
Fig. 6: Failure analysis. We report the proportion of visual trace errors ( red ) versus other failure types ( blue ), computed over 20 failed rollouts for each of the pick-and-place (tomato to grey bowl) and poking (hippo) tasks.
Fig. 7: Results of different reasoning intervals. The experiment is conducted on the task of putting the hedgehog into the red box, with 10 rollouts for each model variant.