Although highly effective in vision and language domains, applying in-context learning to robotics remains challenging. Existing autoregressive in-context imitation methods discretize continuous actions and exacerbate the accumulation of early prediction errors through next-token prediction, limiting their generalization on unseen task configurations. Meanwhile, flow-matching policies have been explored for continuous robot control and can help mitigate compounding errors; however, in-context imitation learning within a flow-matching framework remains underexplored. To address these limitations, we introduce ContextFlow, a conditional flow-matching model that learns continuous action distributions for in-context imitation learning. ContextFlow conditions flow-based action prediction on demonstrations and observations, enabling robust generation from noisy action distributions. To better encode multimodal in-context demonstrations, we adapt perceiver-style multimodal context compressors that distill visual, proprioceptive, and action sequences into compact, task-relevant latent representations. On LIBERO, ContextFlow outperforms ICRT by 35 percentage points in average success rate on unseen task configurations, while matching the performance of the task-specific fine-tuned VLA model π0 without any fine-tuning on unseen tasks. On real robots, it generalizes to unseen configurations of both single-arm and bimanual tasks, achieving 40% success on a new pen-uncapping configuration. Project Page: https://dingjiansw101.github.io/contextflow-page/.
Figures & tables
Figure 1: Autoregressive (a) vs. Flow Matching (b) for In-Context Imitation Learning on action decoding and compounding errors. In-context imitation learning refers to performing a task conditioned on in-context demonstrations without updating model parameters. (a) In-context autoregressive models generate discretized actions step by step, which leads to compounding errors, especially under task distribution shift. (b) Our proposed in-context flow matching framework instead predicts a chunk of continuous actions, alleviating error accumulation and improving generalization. (c) Comparison of compounding errors under seen task configuration between autoregressive and flow-matching approaches. (d) Comparison of compounding errors under unseen task configuration.
Figure 2: ContextFlow Model Architecture. We follow the mixture-of-experts design and employ two experts: the context expert and the action expert . The current observation, text prompt, and in-context demonstrations are first encoded and passed through multimodal context compressors , and the resulting compressed tokens are routed to the context expert, while the current proprioceptive state and action noise tokens are routed to the action expert. The context expert encodes the current observation and contextual information, while the action expert performs flow matching steps to predict the actions.
Figure 3: Examples of our real-world in-context experimental setup using the ALOHA robot platform. We evaluate three task suites. Each suite comprises multiple task configurations that differ in manipulated objects or the required left–right hand motion.
Model
LIBERO-Spatial
LIBERO-Object
Avg.
S 1 ( on )
S 2 ( next to )
Avg.
O 1 ( in )
O 2 ( in )
Avg.
OpenVLA-OFT [ 17 ]
94.0
32.0
63.0
0.0
34.0
17.0
40.0
π0 [ 4 ]
52.0
0.0
26.0
46.0
80.0
63.0
44.5
π0 [ 4 ] FT
94.0
32.0
63.0
80.0
84.0
82.0
72.5
ICRT (AR baseline) [ 11 ]
54.0
0.0
27.0
4.0
96.0
50.0
38.5
ContextAR (AR baseline + )
94.0
12.0
53.0
16.0
92.0
54.0
53.5
Table 1: Comparison of models on LIBERO simulation. Models are trained on seen tasks and evaluated on unseen task configurations. The two unseen task configurations on LIBERO-Spatial are “pick up the black bowl on the cookie box and place it on the plate” and “pick up the black bowl next to the plate and place it on the plate.” The two unseen task configurations on LIBERO-Object task suite are “pick up the milk and place it in the basket” and “pick up the tomato sauce and place it in the basket.” FT means fine-tuning on unseen task configurations with 1 demonstration. ContextAR is our improved in-context autoregressive baseline for fair comparison with ContextFlow.
Table 3 : Ablation study on different modalities in the demonstrations. * indicates the default model. State Image Text Spatial Object Avg. ✗ ✓ ✓ 30.0 28.0 29.0 ✓ ✗ ✓ 42.0 48.0 45.0 ✓ ✓ ✗ 59.0 70.0 64.5 ∗ ✓ ✓ ✓ 64.0 83.0 73.5
Table 4 : Ablation study on the number of images and states in in-context demonstrations. # Images # States Spatial Object Avg. 4 128 62.0 79.0 70.5 8 64 53.0 87.0 70.0 ∗ 8 128 64.0 83.0 73.5 8 256 52.0 81.0 66.5
Table 7: Results with Expanded Training Data and Evaluation Tasks
Model
Single-arm Pick-and-place tasks
Fine-grained bimanual tasks
Left hand
Right hand
Uncap (Left Hand)
Put in box
Pear
Orange
Banana
Kiwi
Red Pen
Red Egg
ICRT [ 11 ]
2/10
0/10
0/10
0/10
0/10
0/10
ContextFlow
2/10
5/10
4/10
4/10
4/10
6/10
Table 8: Performance on unseen task configurations on a real-world robot. See Sec. 4.2 for task instruction templates. Elements in different task configurations are listed in the table.
Figure 8
Figure 6: Context compressor selects task-relevant, compact information. Visualization of image-compressor attention on selected frames (3/6/8), showing that the attended image regions vary across time. In-context demonstration frames and their corresponding attention maps are shown interleaved.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Pick & Place : “pick up ⟨ object ⟩ and place it in the basket with ⟨ left/right ⟩ hand”
Seen (14) ⟨ left ⟩ : apple, corn, gray milk, carrot, chips ⟨ right rangle : pear, orange juice, gray milk, cucumber, corn, red apple, chips, blue milk, carrot
Unseen (4) ⟨ left ⟩ : pear, orange juice ⟨ right ⟩ : kiwi, banana
Pen Uncap : “pick up ⟨ pen ⟩ with ⟨ left/right ⟩ hand, grasp the cap with another hand and uncap it”
Seen (6) ⟨ left & right ⟩ : gray pen, blue pen v2 ⟨ right ⟩ : red pen ⟨ left ⟩ : blue pen
Unseen (1) ⟨ left ⟩ : red pen
Put Egg in Box : “pick up the ⟨ object ⟩ with right hand, place it in the box, and close the box”
Appendix
Table S1 : Seen and unseen task configurations in the real-world aloha experiments.
Task Suite
Task configurations
Episodes
Avg Length
Pick & Place
apple / ⟨ left ⟩
25
200
corn / ⟨ left ⟩
25
200
gray milk / ⟨ left ⟩
25
200
carrot / ⟨ left ⟩
25
200
chips / ⟨ left ⟩
25
200
pear / ⟨ right ⟩
50
250
Appendix
Table S2 : Training data statistics for seen task configurations.
Model
Seen Task Configurations
Unseen Task Configurations
Spatial
Object
Goal
10
Avg.
Spatial
Object
Goal
10
Avg.
π0 [ 4 ]
96.0
97.0
94.0
87.0
93.5
26.0
63.0
0.0
0.0
22.3
ICRT [ 11 ]
85.0
99.0
90.0
65.0
84.8
27.0
50.0
1.0
0.0
19.5
ContextFlow (Ours)
91.0
97.0
90.0
82.0
90.0
64.0
83.0
0.0
0.0
36.8
Appendix
Table S3 : Success rates on seen and unseen task configurations across all task suites in LIBERO.
Task Suites
Task Indexes
Language Descriptions
Spatial
35
Pick up the black bowl on the cookie box and place it on the plate
36
Pick up the black bowl next to the plate and place it on the plate
Object
25
Pick up the milk and place it in the basket
28
Pick up the tomato sauce and place it in the basket
Goal
10
Put the bowl on the plate
17
Put the bowl on the stove
Appendix
Table S4 : Unseen task configurations in LIBERO.
Figure S1 : Rollout of ContextFlow in real-world experiments. We show the in-context demonstrations and the corresponding ContextFlow rollouts from high camera.
Figure S2 : Rollout of ContextFlow in LIBERO simulation. We show the in-context demonstrations and the corresponding ContextFlow rollouts from both agent and wrist views. “Pick up black bowl next to the plate and place it on the plate” belongs to LIBERO Spatial task suite and “Pick up the milk and place it in the basket” belongs to LIBERO Object task suite.
Figure S3 : Generalization of ContextFlow under distribution shifts from the in-context demonstration. The first row shows the in-context demonstration. The second row shows the rollout where the banana is placed at a different location and orientation. The third row shows the rollout where the banana is cut in half, introducing a novel object appearance.
Model
Single-armed Pick-and-place tasks
Bimanual tasks
Avg.
Left hand
Right hand
Uncap
Put in box
Pear
Orange
Banana
Kiwi
Red Pen
Red Egg
ICRT [ 11 ]
2/10
0/10
0/10
0/10
0/10
0/10
3.3
π0 [ 4 ]
2/10
2/10
3/10
2/10
0/10
7/10
26.7
ContextFlow (Ours)
2/10
5/10
4/10
4/10
4/10
6/10
41.7
Appendix
Table S5 : Success rates on unseen task configurations on a real-world robot. See Sec. S1.3 for task instruction templates. Elements in different task configurations are listed in the table. Note that the pen uncap task here uses the left hand to pick up the pen (i.e., “pick up the red pen with left hand, grasp the cap with the other hand and uncap it”.)
Model
Training (Hours) ↓
Inference (Hz) ↑
OpenVLA-OFT [ 17 ]
30
77.8
π0 [ 4 ]
21
569.0
ICRT [ 11 ]
16
44.3
ContextAR
65
3.7
ContextFlow-Plain
21
310.2
ContextFlow
40
130.6
Appendix
Table S6 : Training and inference computational costs of different models. We report the training time in A100 GPU hours and the inference throughput (Hz) measured on a single A100 GPU.
Figure S4 : Grasp position analysis of the LIBERO-object task suite. (a) XY distribution of grasp positions across all demonstrations, revealing two distinct spatial clusters. (b) Z-coordinate distributions for task configurations belonging to the cluster that contains the unseen tomato sauce task, whose object height is close to the training average. (c) Z-coordinate distributions for task configurations belonging to the cluster that contains the unseen milk task, whose object height is notably higher than the training average. In (b) and (c), green triangles mark the mean Z-coordinate across multiple demonstrations, yellow vertical lines denote the median, and dots indicate outliers. Tasks configurations in the lower green region correspond to seen expert demonstrations (ED); the middle blue region corresponds to unseen task configurations in LIBERO-Object; and the upper red region corresponds to ICRT and ContextFlow inference rollouts (RO). We observe that when the height of an unseen object differs from the training average, ICRT rollouts tend to follow the average grasp position and fail to match the unseen object’s height, whereas our proposed ContextFlow generalizes well and aligns its grasp height with the unseen objects.
Figure S5 : Visualization of three types of failure modes. The task for the visualization is “pick up the black bowl next to the plate and place it on the plate.”
Figure S6 : Visualization of attention in image compressor on unseen task configurations. (a) Frames in In-context Demonstration (b) Image compressor attention. The left shows the task “pick up the tomato sauce and place it in the basket”, while the right shows the task “pick up the black bowl next to the plate and place it on the plate”.
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.
Toan Nguyen, Weiduo Yuan, Songlin Wei +3
University of Southern California · Autodesk Research
Learning long-horizon robotic manipulation requires jointly achieving expressive behavior modeling, real-time inference, and stable execution, which remains challenging for existing generative policies. Diffusion-based approaches offer strong modeling capacity but incur high inference latency, while flow matching enables fast, near-single-step generation yet often suffers from unstable execution when operating directly in the raw action space. We propose Continuous Latent Action Flow Policy (CoLA-Flow Policy), a trajectory-level imitation learning framework that performs flow matching in a continuous latent action space. By encoding action sequences into temporally coherent latent trajectories and learning an explicit latent-space flow, CoLA-Flow Policy decouples global motion structure from low-level control noise, enabling smooth and reliable long-horizon execution. The framework further integrates geometry-aware point cloud conditioning and execution-time multimodal modulation, using visual cues as a representative modality to enhance real-world robustness. Experiments in simulation and on real robots show that CoLA-Flow Policy achieves near-single-step inference, improves trajectory smoothness by up to 93.7% and task success by up to 25 percentage points over raw action-space flow baselines, while remaining significantly faster than diffusion-based policies.
Wu Songwei, Jiang Zhiduo, Sun Wandong +4
Harbin Institute of Technology, Harbin 150001, China · Honor Device Co., Ltd., Shenzhen, China
Flow matching has recently become a new standard for behavior cloning in robotic manipulation. However, state-of-the-art flow matching policies suffer from a systematic structural mismatch: they rely on a globally fixed isotropic source distribution despite the strongly fragmented and heteroscedastic structure of robotic action spaces. This agnostic initialization forces the model to learn highly entangled vector fields, bottlenecking training efficiency and limiting overall policy performance. To address this limitation, we introduce Latent Action Guided Flow Matching (LAFM), a novel framework that replaces the monolithic Gaussian with an adaptive library of learned prior distributions. By grounding these distributions using a latent action model, LAFM maps current observations to discrete motion primitives, selecting a specialized base distribution that provides an informed, structurally aligned initialization for the denoising process. This dynamic adaptivity naturally accommodates heteroscedasticity in human demonstrations and makes transport trajectories shorter and less entangled. Empirically, LAFM substantially outperforms standard flow matching formulations, increasing task success rates by 23.4% in real-world robotic deployments and by 10.4% on the LIBERO-90 benchmark. Furthermore, we demonstrate that LAFM achieves state-of-the-art results, surpassing massively pre-trained vision-language-action models while utilizing significantly smaller architectures.
Bruno Machado, Alexandre Chapin, Emmanuel Dellandrea +1