ContextFlow: In-Context Flow Matching for Robot Manipulation

Conference: ECCV 2026
Paper: ECCV Official
Project Page: Project Page
Area: Image Generation
Keywords: Robot Manipulation, In-Context Imitation Learning, Flow Matching, Multimodal Context Compressor, Compounding Error Mitigation

TL;DR

Addressing the quantization errors of action discretization and compounding errors of step-by-step next-token prediction in autoregressive in-context imitation learning, ContextFlow introduces conditional flow matching for continuous action chunk prediction combined with perceiver-style multimodal context compressors, outperforming autoregressive baselines by 35% on unseen LIBERO configurations and matching fine-tuned VLA models without test-time parameter updates.

Background & Motivation

In-context learning empowers models to adapt to novel tasks at inference time given only a few demonstrations without updating neural network parameters. When extended to robotics as in-context imitation learning, conditioning policies on demonstration trajectories consisting of observations, proprioceptive states, and actions holds immense promise for flexible generalist manipulation, eliminating the costly task-specific fine-tuning required by conventional imitation learning. However, robotic manipulation demands producing long-horizon, continuous control actions under tight closed-loop visual feedback. Under unseen task configurations—such as novel target objects or shifted spatial layouts—policies encounter significant distribution shifts.

Most prior in-context imitation learning methods, such as ICRT and RoboPrompt, adopt the autoregressive paradigm inspired by large language models. Autoregressive policy modeling suffers from two fundamental bottlenecks in robotics: first, predicting actions token by token condition on previous outputs exacerbates early prediction errors, causing severe compounding error drift over long horizons especially under distribution shift; second, action tokenizers that discretize continuous actions introduce irreversible reconstruction and quantization errors, impairing high-frequency, precision-critical tasks.

While continuous flow-matching models (like \(\pi_0\)) predict action chunks in parallel and naturally avoid discretization and compounding errors, formulating in-context imitation learning within a flow-matching framework remains underexplored, particularly because processing long-horizon multimodal demonstration sequences incurs prohibitive attention complexity. Core idea: ContextFlow presents a conditional flow matching framework for in-context imitation learning that decouples context understanding and flow action generation via a dual-expert architecture, leveraging Perceiver-style multimodal context compressors to distill long demonstration sequences into compact latent representations and predicting continuous action chunks via block-wise causal cross-attention.

Method

Overall Architecture

ContextFlow targets predicting a chunk of \(H\) future continuous actions \(a_t = a_{t:t+H-1} \sim \pi(a_t \mid o_t, d)\) given the current observation \(o_t = (I_t, q_t)\) (camera image \(I_t\) and proprioceptive state \(q_t\)), task instruction \(T\), and a multimodal demonstration trajectory \(d = [T, X, Q, A]\). The architecture follows a mixture-of-experts formulation consisting of two dedicated networks: a Context Expert that processes task instructions, current observations, and compressed demonstration latents, and an Action Expert that iteratively denoises continuous action chunks via flow matching.

Demonstration sequences spanning hundreds of steps across visual frames, proprioceptive states, and control actions are first subsampled temporally and then passed through modality-specific Perceiver-style context compressors. The resulting compact latent tokens are concatenated with current visual features from a SigLIP ViT and language instruction tokens, forming the input to the Context Expert. The Action Expert receives current proprioception and Gaussian noise action tokens, attending to the Context Expert's output through a block-wise causal attention mechanism across \(K\) flow-matching numerical integration steps to generate a clean continuous action chunk in parallel.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    InDemo["Multimodal Demonstration Sequence<br/>Images X + Proprio Q + Actions A"] --> Comp["Multimodal Context Compressor<br/>Perceiver Cross/Self-Attention Latents"]
    InCur["Current Observation & Instruction<br/>Camera Image It + Instruction T"] --> CtxExp["Context Expert Representation<br/>Aggregates Demonstration & Scene Context"]
    Comp --> CtxExp
    InAct["Current Proprio State qt +<br/>Gaussian Noise Action Chunk a0"] --> ActExp["Action Expert Flow Matching<br/>Block-Wise Causal Cross-Attention"]
    CtxExp -->|Cached Causal Cross-Attention| ActExp
    ActExp --> OutAct["Continuous Action Chunk Output<br/>Parallel Future H-step Trajectory"]

Key Designs

1. Conditional Flow Matching Action Generation: Eliminating Discretization and Compounding Errors

To avoid the compounding drift and quantization loss inherent in autoregressive action decoders, ContextFlow models action chunk generation via conditional continuous-time flow matching. The model defines a linear probability path connecting standard Gaussian noise \(p_0(a_t^0) = \mathcal{N}(0, I)\) at time \(s=0\) to the target action distribution \(p_1(a_t^1) = \pi(a_t \mid o_t, d)\) at \(s=1\). Intermediate states are sampled via linear interpolation \(a_t^s = (1 - s)\epsilon + s a_t\) with \(\epsilon \sim \mathcal{N}(0, I)\), producing a tractable closed-form conditional velocity vector field \(u(a_t^s \mid a_t) = a_t - \epsilon\) pointing straight from noise to clean action. The Action Expert parameterizes a learned vector field \(u_\theta(a_t^s, o_t, d, s)\) conditioned on noisy actions, observation, compressed context, and sinusoidal time embeddings \(s\). During inference, continuous actions are integrated from \(s=0\) to \(s=1\) in just \(K=10\) Euler forward steps:

\[a_t^{s + \delta} = a_t^s + \delta u_\theta(a_t^s, o_t, d, s), \quad \delta = 1/K\]

Because the entire future action horizon \(H\) is generated in parallel in continuous space, the policy suppresses the accumulation of rollout drift and enables robust execution under unseen task configurations.

2. Perceiver Multimodal Context Compressors: Distilling Long-Horizon Multimodal Demonstrations

Demonstration trajectories often span hundreds of timesteps and multi-view video streams, creating an intractable quadratic attention bottleneck if passed directly into large transformers. ContextFlow develops dedicated context compressors \(C_{\text{img}}\), \(C_{\text{state}}\), and \(C_{\text{action}}\) for visual, proprioceptive, and action modalities. For images, after temporal subsampling to \(N_1=8\) frames, features are extracted across camera views. A set of \(N_q=32\) learnable query tokens initializes a compact latent representation, alternating across stacked cross-attention and self-attention layers (4 layers for images, 2 layers each for proprioceptive states and actions). Cross-attention allows queries to absorb spatial-temporal features, while self-attention mixes information across the fixed-size latent tokens. The outputs from all three compressors are concatenated into a unified latent context:

\[T_{\text{context}} = [Z_{\text{img}}; Z_{\text{state}}; Z_{\text{action}}] \in \mathbb{R}^{N_{\text{context}} \times D}\]

Empirical attention visualizations reveal that the learnable queries automatically learn to shift focus over task phases—concentrating on the gripper before grasping, attending to the manipulated object post-grasp, and focusing on the target receptacle during placement.

3. Block-Wise Causal Cross-Attention: Efficient Inference via Context Representation Caching

In the dual-expert framework, interaction between context representations and action noise tokens is essential for policy conditioning. If dense bidirectional attention were used across the entire sequence, every step of the \(K\)-step flow-matching ODE solver would require recomputing attention across both the Context Expert and Action Expert, resulting in severe inference latency. ContextFlow imposes a block-wise causal attention constraint: action tokens can attend to context expert tokens, but context tokens are strictly prohibited from attending back to action tokens. This unidirectional dependency ensures that Context Expert activations and Key-Value caches can be computed once per observation timestep and reused across all \(K\) flow-matching integration steps, drastically reducing inference latency while preserving expressive conditioning.

Loss & Training

The network is trained using the Conditional Flow Matching (CFM) objective. Time \(s\) is sampled from a Beta distribution \(s \sim \text{Beta}(\alpha, \beta)\) with \(\alpha=1.5, \beta=1\), action chunk \(a_t \sim \pi(a_t \mid o_t, d)\), and Gaussian noise \(\epsilon \sim \mathcal{N}(0, I)\):

\[\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{s, a_t, \epsilon} \left[ \left\| u_\theta(a_t^s, o_t, d, s) - (a_t - \epsilon) \right\|^2 \right]\]

The vision encoder (SigLIP) and Action Expert are initialized from pretrained \(\pi_0\) weights, with the Action Expert fine-tuned using LoRA (rank 32, \(\alpha=32\)). The 300M Context Expert and multimodal context compressors are randomly initialized and fully trained. Training runs for 20,000 iterations with batch size 32, optimized via AdamW with a cosine learning rate schedule (peak \(2.5 \times 10^{-5}\), minimum \(2.5 \times 10^{-6}\), 1,000-step linear warmup). Inference uses \(K=10\) forward Euler steps.

Key Experimental Results

Main Results

ContextFlow is evaluated in simulation on the LIBERO benchmark (LIBERO-Spatial and LIBERO-Object suites) under unseen task configurations (evaluated over 50 trials per task), and on a physical ALOHA dual-arm robot across single-arm pick-and-place and fine-grained bimanual manipulation suites (10 trials per task).

Benchmark / Task Suite	Evaluation Metric	ContextFlow (Ours)	ContextAR (AR Baseline+)	ICRT (Standard AR)	\(\pi_0\) FT (1-shot FT)
LIBERO-Spatial (Unseen)	Success Rate (%)	64.0	53.0	27.0	63.0
LIBERO-Object (Unseen)	Success Rate (%)	83.0	54.0	50.0	82.0
LIBERO Average (Unseen)	Success Rate (%)	73.5	53.5	38.5	72.5
Real ALOHA: Pick & Place (Orange)	Success Rate (10 trials)	5/10 (50%)	-	0/10 (0%)	-
Real ALOHA: Uncap Pen (Left Hand)	Success Rate (10 trials)	4/10 (40%)	-	0/10 (0%)	-
Real ALOHA: Put Egg in Box (Red Egg)	Success Rate (10 trials)	6/10 (60%)	-	0/10 (0%)	-

In LIBERO simulation, ContextFlow surpasses the strongest autoregressive baseline ContextAR by 20.0 percentage points and outperforms ICRT by 35.0 percentage points. Remarkably, without any test-time gradient updates, ContextFlow matches or slightly exceeds the 1-shot fine-tuned VLA model \(\pi_0\text{ FT}\) (73.5% vs. 72.5%). In real-world ALOHA tasks requiring closed-loop coordination, ICRT failed in all bimanual trials (0/10), whereas ContextFlow achieved 40% to 60% success.

Ablation Study

1. Multimodal Context Compressor Ablation

Model Configuration	Image Compressor \(C_{\text{img}}\)	State & Action Compressor \(C_{\text{state}} \& C_{\text{action}}\)	Spatial Success (%)	Object Success (%)	Average Success Rate (%)
ContextFlow-Plain	✗	✗	55.0	73.0	64.0
Image-Only Compressor	✓	✗	51.0	85.0	68.0
ContextFlow (Full Model)	✓	✓	64.0	83.0	73.5

2. Action Head Comparison (20k iterations unless noted)

Action Head Architecture	Iterations	Spatial Success (%)	Object Success (%)	Average Success Rate (%)	Cumulative MSE (\(m^2\))
ContextAR (Autoregressive Discrete)	20k	53.0	54.0	53.5	4.7
MLP \(L_1\) Regression (Continuous)	20k	26.0	90.0	58.0	-
Diffusion Head	100k	30.0	37.0	33.5	-
ContextFlow-Plain (Flow Matching w/o Comp.)	20k	55.0	73.0	64.0	3.0
ContextFlow (Flow Matching + Compressors)	20k	64.0	83.0	73.5	2.2

Key Findings

• Compounding Error Reduction: Continuous flow matching with action chunking fundamentally suppresses compounding errors. Trajectory tracking analysis shows that ContextAR accumulates a cumulative MSE of \(4.7\,m^2\), whereas ContextFlow reduces this to \(2.2\,m^2\). Autoregressive errors compound steeply past step 30, while flow matching maintains smooth trajectory tracking aligned with expert demonstrations.
• Criticality of State and Action Modalities: Removing proprioceptive states and actions from the demonstration while keeping text and images drops unseen average success from 73.5% to 29.0% (a 44.5 percentage point collapse), demonstrating that visual prompts alone lack the kinesthetic fidelity required for precision imitation.
• Compressor Efficiency and Feature Regularization: Compressing visual and state-action sequences with Perceiver queries not only bounds computational complexity but also filters spatio-temporal redundancy, yielding a 9.5 percentage point boost over uncompressed ContextFlow-Plain.

Highlights & Insights

• Flow Matching as a Remedy for Compounding Errors: Replaces tokenized autoregressive prediction with continuous parallel flow matching, resolving quantization drift and establishing continuous generative modeling for in-context imitation learning.
• Decoupled Dual-Expert Architecture with Block-Wise Causal Caching: Context and action experts are decoupled, allowing KV caching of context tokens during numerical integration and achieving high execution throughput without losing conditioning fidelity.
• Perceiver-Style Compression for Multi-View Trajectories: Compresses long-horizon multimodal demonstrations into a fixed set of latent tokens, where cross-attention automatically transitions from grasping tools to manipulating objects and target locations.

Limitations & Future Work

• Generalization Limited to Related Primitives: Demonstrated in-context adaptation is primarily restricted to unseen objects and novel spatial configurations within known task primitives, rather than generalizing to fundamentally distinct manipulation categories.
• Dependence on Clean Demonstrations: The framework relies on high-quality human demonstrations; handling sub-optimal, noisy, or truncated trajectories without reinforcement-guided self-correction remains an open challenge.
• Future Directions: Integrating preference learning (e.g., flow-based DPO) and scaling pretraining over diverse cross-embodiment robot datasets could unlock open-world few-shot generalization.

Related Work & Insights

• vs ICRT: ICRT collapses visual and proprioceptive latents into a single token and predicts discretized actions autoregressively via LLaMA, leading to catastrophic compounding drift; ContextFlow predicts continuous action chunks via flow matching, outperforming ICRT by 35% on unseen configurations.
• vs \(\pi_0\) / \(\pi_0\text{-FT}\): Standard \(\pi_0\) relies on text prompts and requires gradient fine-tuning to adapt to unseen configurations; ContextFlow incorporates multimodal demonstration prompting through a context expert, matching 1-shot fine-tuned \(\pi_0\text{-FT}\) without any parameter updates.
• vs OpenVLA-OFT: While OpenVLA-OFT adds continuous action regression, it remains constrained by language-only instructions; ContextFlow shows that multimodal demonstration prompts are superior for spatial and geometric generalization.

Rating

• Novelty: ⭐⭐⭐⭐☆ Pioneers conditional flow matching with Perceiver-style context compression for robotic in-context imitation learning.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Rigorous validation across LIBERO simulation suites, error accumulation tracking, modality ablations, and real-world ALOHA bimanual manipulation.
• Writing Quality: ⭐⭐⭐⭐⭐ Well-structured narrative with sharp empirical insights into autoregressive failure modes versus continuous flow dynamics.
• Value: ⭐⭐⭐⭐⭐ Offers a high-impact, principled generative foundation for zero-shot and few-shot embodied robotic control.