ContextFlow: In-Context Flow Matching for Robot Manipulation

Jian Ding; Xianjie Dai; Roei Herzig; Nussair Hroub; Jinjie Mai; Dengxin Dai; Bernard Ghanem; Mohamed Elhoseiny

cs.ROSep 6, 2026

ContextFlow: In-Context Flow Matching for Robot Manipulation

Authors: Jian Ding, Xianjie Dai, Roei Herzig, Nussair Hroub, Jinjie Mai, Dengxin Dai, Bernard Ghanem, Mohamed Elhoseiny

Organizations: KAUST, Thuwal, Saudi Arabia · University of California, Berkeley, USA · Zurich, Switzerland

Abstract

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.
Model	S 1 ( on )	S 2 ( next to )	Avg.	O 1 ( in )	O 2 ( in )	Avg.	Avg.
OpenVLA-OFT [ 17 ]	94.0	32.0	63.0	0.0	34.0	17.0	40.0
$\pi_{0}$ [ 4 ]	52.0	0.0	26.0	46.0	80.0	63.0	44.5
$\pi_{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 2 : Ablation study on context compressors. Model $\mathcal{C}_{\text{img}}$ $\mathcal{C}_{\text{state}}$ & $\mathcal{C}_{\text{action}}$ Spatial Object Avg. ContextFlow-Plain ✗ ✗ 55.0 73.0 64.0 ✓ ✗ 51.0 85.0 68.0 ContextFlow ✓ ✓ 64.0 83.0 73.5
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 5 : Ablation of action heads. Action Head Iterations Spatial Object Avg. MLP L1 20k 26.0 90.0 58.0 Diffusion 100k 30.0 37.0 33.5 Flow Matching 20k 64.0 83.0 73.5	Table 6 : Cumulative MSE. Model MSE ( $\mathrm{m}^{2}$ ) ContextAR (baseline + ) 4.7 ContextFlow-Plain 3.0 ContextFlow 2.2

Table 2: Ablation study on context compressors.

Training Data	Spatial	Object	Goal	10	Avg
4 suites	64.0	83.0	0.0	0.0	36.8
4 suites + LIBERO-90	55.0	85.0	29.0	3.0	43.0

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 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

\langle

object

\rangle

and place it in the basket with

\langle

left/right

\rangle

hand”

Seen (14)

\langle

left

\rangle

: apple, corn, gray milk, carrot, chips

\langle

right

\ rangle

: pear, orange juice, gray milk, cucumber, corn, red apple, chips, blue milk, carrot

Unseen (4)

\langle

left

\rangle

: pear, orange juice

\langle

right

\rangle

: kiwi, banana

Pen Uncap : “pick up

\langle

pen

\rangle

with

\langle

left/right

\rangle

hand, grasp the cap with another hand and uncap it”

Seen (6)

\langle

left & right

\rangle

: gray pen, blue pen v2

\langle

right

\rangle

: red pen

\langle

left

\rangle

: blue pen

Unseen (1)

\langle

left

\rangle

: red pen

Put Egg in Box : “pick up the

\langle

object

\rangle

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 / $\langle$ left $\rangle$	25	200
	corn / $\langle$ left $\rangle$	25	200
	gray milk / $\langle$ left $\rangle$	25	200
	carrot / $\langle$ left $\rangle$	25	200
	chips / $\langle$ left $\rangle$	25	200
	pear / $\langle$ right $\rangle$	50	250

Appendix

Table S2 : Training data statistics for seen task configurations.

Model	Seen Task Configurations					Unseen Task Configurations
Model	Spatial	Object	Goal	10	Avg.	Spatial	Object	Goal	10	Avg.
$\pi_{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
Spatial	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
Object	28	Pick up the tomato sauce and place it in the basket
Goal	10	Put the bowl on the plate
Goal	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
$\pi_{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) $\downarrow$	Inference (Hz) $\uparrow$
OpenVLA-OFT [ 17 ]	30	77.8
$\pi_{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”.

Explore similar work

ICLR: In-Context Imitation Learning with Visual Reasoning
Mar 8, 2026Toan Nguyen, Weiduo Yuan, Songlin Wei +3Robot Imitation Learning ICL for Robotics
University of Southern California · Autodesk Research
CoLA-Flow Policy: Temporally Coherent Imitation Learning via Continuous Latent Action Flow Matching for Robotic Manipulation
Jan 30, 2026Wu Songwei, Jiang Zhiduo, Sun Wandong +4Flow Matching Latent Action Learning
Harbin Institute of Technology, Harbin 150001, China · Honor Device Co., Ltd., Shenzhen, China
Flowing With Purpose: Latent Action Guided Flow Matching Policies For Robotic Manipulation
Jun 22, 2026Bruno Machado, Alexandre Chapin, Emmanuel Dellandrea +1Flow Matching Robotic Manipulation
École Centrale de Lyon · LIRIS, UMR5205

Mar 8, 2026J/K move · Enter open · S save

ICLR: In-Context Imitation Learning with Visual Reasoning

Toan Nguyen, Weiduo Yuan, Songlin Wei, Hui Li+2

University of Southern California · Autodesk Research

Robot Imitation Learning ICL for Robotics Robotic Manipulation Generalization in Robotic Manipulation Embodied Reasoning

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

ICLR: In-Context Imitation Learning with Visual Reasoning

CoLA-Flow Policy: Temporally Coherent Imitation Learning via Continuous Latent Action Flow Matching for Robotic Manipulation

Flowing With Purpose: Latent Action Guided Flow Matching Policies For Robotic Manipulation