Task Definition
In-Context Imitation Learning: Learn the skill of an unseen task through an in-context demonstration immediately, without updating model weights.
In-Context Imitation Learning: Learn the skill of an unseen task through an in-context demonstration immediately, without updating model weights.
In-context learning is an emergent capability of large language models (LLMs), which are typically formulated as autoregressive models. However, when applied to robot manipulation, autoregressive approaches often require discretizing continuous actions and can compound early prediction errors through sequential next-token prediction. In this work, we explore an alternative formulation of in-context imitation learning based on flow matching, which directly predicts continuous actions in parallel.
ContextFlow models the conditional action distribution directly in continuous space. Given one in-context demonstration and the current observation, it initializes a noisy action chunk and refines the full chunk through ten conditional flow-matching updates. Unlike autoregressive policies, the model does not discretize actions or predict future controls token by token. Perceiver-style compressors reduce redundant visual, proprioceptive, and action sequences to compact latent tokens. The context expert fuses them with the task instruction and current images, and the action expert predicts a 50-step action horizon. Block-wise causal attention enables context caching during inference.
Models are trained on 32 seen tasks and evaluated on 4 unseen task configurations from LIBERO-Spatial and LIBERO-Object. FT means fine-tuning on unseen task configurations with one demonstration. ContextAR is our improved in-context autoregressive baseline for a fair comparison with ContextFlow.
| Model | LIBERO-Spatial | LIBERO-Object | Avg. | ||||
|---|---|---|---|---|---|---|---|
| S1 | S2 | Avg. | O1 | O2 | Avg. | ||
| OpenVLA-OFT | 94.0 | 32.0 | 63.0 | 0.0 | 34.0 | 17.0 | 40.0 |
| π0 | 52.0 | 0.0 | 26.0 | 46.0 | 80.0 | 63.0 | 44.5 |
| π0 FT | 94.0 | 32.0 | 63.0 | 80.0 | 84.0 | 82.0 | 72.5 |
| ICRT (AR baseline) | 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 |
| ContextFlow-Plain (Ours) | 90.0 | 20.0 | 55.0 | 92.0 | 54.0 | 73.0 | 64.0 |
| ContextFlow (Ours) | 86.0 | 42.0 | 64.0 | 76.0 | 90.0 | 83.0 | 73.5 |
The real-world dataset contains 25 seen task configurations (including both single-arm and bimanual tasks), and tested on 6 unseen task configurations.
| Model | Pear | Orange | Banana | Kiwi | Red pen | Red egg |
|---|---|---|---|---|---|---|
| ICRT | 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 |
Red egg manipulation. Put the red egg in the box and close the box.
Red pen manipulation. Pick up the red pen with the left hand and uncap it.
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.
