PaperlayerView exact revision on arXiv

Source-linked AI summary

ContextFlow: In-Context Flow Matching for Robot Manipulation

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

arXiv:2609.06852v1cs.RO

TL;DR

Autoregressive in-context imitation methods discretize actions and can accumulate prediction errors, while flow matching for in-context imitation remains underexplored. ContextFlow uses conditional flow matching and multimodal context compressors; it outperforms autoregressive baselines and matches fine-tuned VLA performance on unseen configurations. Its demonstrated generalization is mainly within task configurations involving unseen objects or spatial arrangements.

  • Problem

    Robotic in-context imitation must predict continuous actions from few demonstrations, but autoregressive methods discretize actions and compound errors, and flow matching in this setting remains underexplored.

    Source: §1 · pp. 2–3

  • Method

    ContextFlow conditions continuous-action flow matching on observations and demonstrations, using multimodal compressors to distill demonstration streams into compact latent tokens.

    Source: §1 · p. 3, §3.1 · p. 5, §3.2 · p. 7

  • Results

    ContextFlow outperforms autoregressive in-context imitation baselines and matches task-specific fine-tuned VLA performance without fine-tuning on unseen configurations.

    Source: §1 · p. 4

  • Takeaways & Limitations

    The authors report that flow matching appears beneficial for generalizing to unseen configurations, while multimodal compressors further improve performance.

    Source: §7 · p. 15

    The formulation assumes each task goal is described by a textual instruction and provides in-context demonstrations.

    Source: §3.1 · p. 5

Abstract

from arXiv · show

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 $\pi_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/.

1 Introduction

ContextFlow brings conditional flow matching to in-context robot imitation, predicting continuous action chunks from observations and demonstrations. It addresses error accumulation and action discretization in autoregressive approaches, and evaluates the method in simulation and on real robots.

1 Introduction

  • Autoregressive policies predict actions sequentially, allowing early errors to compound under task-distribution shift; action discretization also introduces reconstruction error.These limitations are particularly pronounced when evaluation configurations differ from training configurations.Source: §1 · pp. 2–3
  • ContextFlow uses conditional flow matching to predict continuous action chunks in parallel from current observations and in-context demonstrations.This avoids action discretization and is intended to mitigate compounding errors on unseen configurations.Source: §1 · p. 3
  • Perceiver-style multimodal context compressors distill visual, proprioceptive, and action sequences into fixed-size latent tokens, reducing encoding and cross-attention costs.The compressed tokens are designed to preserve task-relevant information from long demonstrations.Source: §1 · p. 3
  • The authors curate an ALOHA dataset of 1,318 multimodal trajectories across 25 task configurations, covering single-arm skills and fine-grained bimanual tasks.They evaluate ContextFlow in LIBERO simulation and on ALOHA real robots.Source: §1 · pp. 3–4
  • ContextFlow outperforms autoregressive in-context imitation baselines and matches task-specific fine-tuned VLA models without fine-tuning on unseen configurations.The introduction also reports lower errors on unseen configurations.Source: §1 · pp. 3–4

2 Related Work

Prior robotic in-context imitation work includes autoregressive approaches, tokenized prompts, full-state methods, and 3D-based methods. Related VLA systems use language prompts or predict high-level actions, while ContextFlow uses multimodal demonstrations for low-level action prediction.

2 Related Work

  • Robotic few-shot imitation includes tokenized external-LLM methods [42], full-state methods, and 3D-based approaches [37] [44].ContextFlow instead uses RGB observations and directly outputs executable actions.Source: §2 · p. 4
  • ICRT [11] and LipVQ-VAE [38] use autoregressive action tokens, whereas ContextFlow models continuous action chunks with conditional flow matching.RICL [34] extends autoregressive VLA models with in-context adaptability.Source: §2 · p. 4
  • Many VLA models use textual instructions, while VIMA predicts 2–5 high-level actions per episode rather than hundreds of low-level action steps.ContextFlow uses multimodal in-context demonstrations as its task prompt.Source: §2 · p. 4
  • ContextFlow’s architecture routes compressed demonstration and observation tokens to a context expert, while an action expert uses proprioception and action noise to predict actions.The action expert performs flow-matching steps.Source: Figure 2 · p. 5

3 Method

ContextFlow models future action chunks conditioned on the current observation and demonstrations using continuous flow matching. Its architecture compresses multimodal demonstrations into latent context tokens that interact with action tokens during prediction.

3.1 Problem Formulation

  • Flow matching models the conditional distribution of future action chunks directly in continuous action space, avoiding action discretization.The model transports samples from a Gaussian prior toward the target distribution conditioned on the observation and demonstration.Source: §3.1 · pp. 5–6
  • At inference, the learned velocity field is integrated from Gaussian noise to generate the policy’s next action chunk.The implementation uses a forward Euler solver with K = 10 steps.Source: §3.1 · p. 6

3.2 Model Architecture

  • The context and action experts exchange information through attention as the action expert predicts a clean action chunk over K flow-matching steps.The action expert receives proprioception and noisy action tokens; context tokens are cached during inference because attention is causal from context to action.Source: §3.2 · pp. 6, 8
  • Perceiver-style compressors reduce long multimodal demonstrations to fixed-size latent representations using alternating cross-attention and self-attention.Separate image, state, and action compressors produce representations that are combined into a unified context.Source: §3.2 · pp. 6–7

4 Experimental Setup

The evaluation measures generalization to unseen configurations in LIBERO simulation and ALOHA real-robot task suites. The setup also specifies implementation choices and compares ContextFlow with autoregressive and fine-tuned VLA baselines.

4.1 Simulation Benchmark Setup

  • LIBERO evaluation uses two unseen configurations per suite, held out from eight seen tasks, and reports average success over 50 trials per task.The evaluated suites are LIBERO-Spatial and LIBERO-Object; LIBERO-Goal and LIBERO-10 demonstrations are additionally used in training.Source: §4.1 · p. 8

4.2 Real-World Experiment Setup

  • ALOHA experiments cover three task suites with 21 configurations and 1,064 demonstrations, including long-horizon bimanual tasks requiring closed-loop feedback.Four additional bimanual tasks with 254 demonstrations are used for training only.Source: §4.2 · pp. 8–9

4.3 Implementation Details

  • The implementation samples 8 image frames, 128 proprioceptive states, and 128 actions for demonstrations, with 32 learnable query tokens per modality.The image compressor has 4 layers, while the state and action compressors each have 2 layers.Source: §4.3 · p. 9

4.4 Baselines

  • Baselines include ICRT, an improved ContextAR autoregressive model, large-scale VLA models, and π0 fine-tuned with one unseen-task demonstration.ContextAR shares ContextFlow’s vision backbone and is initialized from π0; the VLA models receive unseen-task text instructions.Source: §4.4 · pp. 9–10, Table 1 · p. 10

5 Results

Section 5 compares flow-matching and autoregressive baselines, then tests context modalities, action heads, training breadth, and real-robot configurations.

5.1 Comparison with Baselines

  • In the controlled comparison, ContextFlow-Plain outperforms ContextAR by 10.5 percentage points in success rate.ContextFlow-Plain and ContextAR share a vision backbone and π0 initialization; ContextAR also beats OpenVLA-OFT by 13.5 percentage points.Source: §5.1 · p. 11
  • ContextFlow beats ContextFlow-Plain by 9.5 and ContextAR by 20 percentage points in average success rate, while matching π0 (FT) at 73.5% vs. 72.5%.Adding context compressors improves performance over the uncompressed flow-matching model.Source: §5.2 · p. 11, §5.1 · p. 11

5.2 Ablations

  • Removing proprioceptive states and actions drops unseen-task performance by 44.5 points on average; removing text drops it by 9.0 points.Using only the image compressor improves average success rate by 4 points; jointly compressing visual and state/action context brings the largest gains in spatial reasoning and average success rate.Source: §5.2 · pp. 11–12
  • Increasing demonstration frames from 4 to 8 improves average success rate by 3.0 points, and 128 is the reported optimal number of states/actions in context.The frame comparison keeps the number of states/actions constant.Source: §5.2 · p. 12
  • Among the tested continuous parallel action heads, flow matching performs best at 73.5% success on LIBERO.MLP L1 regression reaches 58%; diffusion reaches 16.5% after 20k iterations and 33.5% after 100k iterations.Source: §5.2 · p. 12
  • Adding LIBERO-90 raises the average on the evaluated LIBERO-Goal and LIBERO-10 tasks from 36.8 to 43.0 points.These were previously unsolved tasks, with two unseen tasks evaluated from each suite.Source: §5.2 · p. 12

5.3 Results on Real-World Experiments

  • On real robots, ContextFlow performs reasonably on single-arm pick-and-place and fine-grained bimanual configurations, while ICRT succeeds in only 2/10 pear pick-and-place trials.The reported results suggest ContextFlow can leverage demonstrations on unseen real-world configurations; ICRT struggles with fine-grained bimanual tasks.Source: §5.3 · pp. 12–13

6 Analysis

The analysis examines error accumulation, rollout failures, and the image compressor’s attention across task phases.

Compounding Error

  • Flow matching reduces cumulative MSE by 1.7, and compression reduces it by a further 0.8; ContextAR’s error grows noticeably beyond the first 30 steps in the evaluated task.The MSE decrease begins around step 60, near the end of the first task stage.Source: §6 · pp. 13–14

Failure Mode

  • ICRT often misses the target during pick-up, whereas ContextFlow failures are more evenly distributed across stages, with more at placement.The visualized failed trajectories include detours during pick-up and placement.Source: §6 · p. 14, Figure 5 · p. 14, Figure 4 · p. 14

Attention in Image Compressor

  • Image-compressor attention shifts from the gripper before grasping, to the object after grasping, and then to the target plate during placing.The visualized attention averages weights over 32 queries and shows selected demonstration frames 3, 6, and 8.Source: §6 · p. 15, Figure 6 · p. 15

7 Conclusion

The conclusion presents ContextFlow as conditional flow matching for in-context imitation, with continuous parallel action prediction and multimodal context compression. Its reported generalization is mainly to unseen objects and spatial arrangements within tasks sharing similar primitives, not fundamentally different task types.

Conclusion

  • ContextFlow is an in-context imitation framework based on conditional flow matching.It predicts continuous actions in parallel from the current observation and in-context demonstrations.Source: §7 · p. 15
  • The empirical findings suggest flow matching benefits generalization to unseen configurations, while multimodal context compressors reduce demonstration redundancy and further improve performance.The conclusion associates continuous parallel prediction with alleviating compounding error.Source: §7 · p. 15
  • The demonstrated generalization is mainly limited to unseen objects and spatial arrangements within tasks sharing similar primitives, rather than fundamentally different task types.The conclusion identifies greater training-data diversity and quantity as a possible route to broader generalization.Source: §7 · p. 15

Supplementary Material

The supplementary material lists additional implementation and dataset details, experimental results, behavioral analyses, and a discussion.

Supplementary Material

S1 Additional Implementation Details

The supplement details ContextFlow’s multimodal context compression and model setup, baseline training procedures, and real-world evaluation data. It also documents implementation choices that shape the comparisons.

S1.1 Implementation Details of ContextFlow

  • ContextFlow compresses demonstration image, state, and action sequences into 32-token representations, then combines context tokens with current-observation tokens and text for the context expert.LIBERO uses 128 demonstration context tokens and ALOHA uses 160; current-observation image tokens remain uncompressed.Source: S1.1 Implementation Details of… · p. 19
  • The action horizon is 50; the context expert is trained from scratch, while the action expert and SigLIP encoder are initialized from π0 and fine-tuned.The action expert uses LoRA; the SigLIP encoder is fully fine-tuned.Source: S1.1 Implementation Details of… · p. 19

S1.2 Implementation Details of Baselines

  • ICRT uses a fine-tuned Llama2-Base model, trained for five epochs; ten epochs performed worse.The supplement also evaluates a LoRA-fine-tuned Llama2-7B variant, which performs worse than Llama2-Base in this setting.Source: S1.2 Implementation Details of Baselines · pp. 19–20
  • π0 is fine-tuned for 30K iterations on seen configurations, while its one-shot unseen-task baseline receives 100 additional iterations on one demonstration.The one-shot comparison requires task-specific weight updates; ContextFlow does not require such updates on unseen tasks.Source: S1.2 Implementation Details of Baselines · p. 20
  • OpenVLA-OFT training is reduced to 20,000 iterations and four GPUs, with sub-optimal seen-task performance attributed as likely related to the lower compute budget and modified settings.Source: S1.2 Implementation Details of Baselines · p. 20
  • ContextAR projects each discrete in-context state and action into a continuous token to preserve the pretrained π0-FAST input format without substantial architectural changes.Source: S1.2 Implementation Details of Baselines · p. 20

S1.3 Details of Real-World Experiments

S2 Supplementary Experimental Results

Across LIBERO and real-robot evaluations, ContextFlow performs comparably to π0 on seen tasks and generalizes more strongly than the baselines to unseen configurations. The supplement also examines rollout robustness and computational trade-offs.

S2.1 Results on unseen and seen tasks across all task suites

S2.2 Comparison on Real-world experiments

  • On real robots, ContextFlow generalizes more strongly to unseen configurations than π0 and ICRT; π0 fails on uncap-red-pen, while ContextFlow retains a reasonable success rate.π0 performs well when the unseen put-egg-in-box configuration differs only by egg color, but remains below ContextFlow on single-arm pick-and-place tasks with larger configuration gaps.Source: S2.2 Comparison on Real-world… · p. 23

S2.3 Visualization of ContextFlow Rollout

S2.4 Computation cost for training and inference

  • Flow-matching models are usually faster than autoregressive models because they predict action chunks in parallel, while adding image tokens to ContextFlow-Plain degrades performance and increases computation.The supplement reports training time in A100 GPU hours and inference throughput in Hz.Source: S2.4 Computation cost for training and… · pp. 24–25

S3 Supplementary Analysis

The supplementary analyses relate grasp-height adaptation to performance on an unseen LIBERO object, catalogue failure modes, and visualize changes in image-compressor attention across task phases.

S3.1 Model Behavior Analysis on LIBERO-Object

  • Fig. S3 illustrates rollouts with a banana at a different location and with a cut-in-half appearance relative to the in-context demonstration.Source: Figure S3 · p. 26
  • On the unseen milk task, ContextFlow reaches 76% success versus ICRT’s 4.0% and shifts grasp heights upward toward the object, unlike ICRT’s concentration near the seen-task average.The analysis finds that grasp-height variation within a spatial cluster can substantially influence performance.Source: S3.1 Model Behavior Analysis on… · pp. 25–26

S3.2 Additional Failure Mode Analysis

  • The three described failure modes are wrong-object selection amid similar nearby objects, grasp failure after identifying the object, and failure to place a grasped object at the target.The authors note that grasp failure may relate to inaccurate action prediction.Source: S3.2 Additional Failure Mode Analysis · pp. 26–27

S3.3 Attention in Image Compressor

S4 Discussion

The discussion reports stronger accuracy and generalization for flow-matching policies in in-context imitation, especially on unseen configurations. It also notes a training-efficiency trade-off with autoregressive heads and proposes exploring hybrid AR–flow training.

S4 Discussion

  • Flow-matching policies achieve higher accuracy and more robust generalization than autoregressive models in the reported in-context imitation evaluations, particularly on unseen configurations.The discussion attributes reduced long-horizon error accumulation to parallel continuous action-chunk prediction rather than token-by-token dependence.Source: S4 Discussion · p. 27
  • For the unseen milk object, ContextFlow aligns grasp height upward while ICRT stays near the seen-task average, accompanying success rates of 76% and 4.0%, respectively.Source: Figure S4 · p. 28
  • Autoregressive heads can train more efficiently through simpler next-token objectives and faster convergence, motivating exploration of staged or joint AR–flow training.The proposed staged direction is to train an in-context imitation model autoregressively, then fine-tune a flow-matching head.Source: S4 Discussion · p. 28

Related papers

Loading 2609.06852v1…