# COLLAGE: Adaptive Fusion-based Retrieval for Augmented Policy Learning

arXiv: https://arxiv.org/abs/2508.01131 · HTML: https://mlsanity.com/paper/2508.01131/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 42 of its 54 references that are in the map.
- First posted on arXiv 2 Aug 2025.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://mlsanity.com/paper/1706.02413/) (2017)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [π0: A Vision-Language-Action Flow Model for General Robot Control](https://mlsanity.com/paper/2410.24164/) (2024)
- [Voyager: An Open-Ended Embodied Agent with Large Language Models](https://mlsanity.com/paper/2305.16291/) (2023)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning](https://mlsanity.com/paper/2306.03310/) (2023)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)
- [Mobile ALOHA: Learning Bimanual Mobile Manipulation with Low-Cost Whole-Body Teleoperation](https://mlsanity.com/paper/2401.02117/) (2024)
- [Attention Bottlenecks for Multimodal Fusion](https://mlsanity.com/paper/2107.00135/) (2021)
- [GMFlow: Learning Optical Flow via Global Matching](https://mlsanity.com/paper/2111.13680/) (2021)
- [Deep Imitation Learning for Complex Manipulation Tasks from Virtual Reality Teleoperation](https://mlsanity.com/paper/1710.04615/) (2017)
- [Characterizing and Avoiding Negative Transfer](https://mlsanity.com/paper/1811.09751/) (2018)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [ALOHA Unleashed: A Simple Recipe for Robot Dexterity](https://mlsanity.com/paper/2410.13126/) (2024)
- [Human-in-the-Loop Imitation Learning using Remote Teleoperation](https://mlsanity.com/paper/2012.06733/) (2020)
- [BAKU: An Efficient Transformer for Multi-Task Policy Learning](https://mlsanity.com/paper/2406.07539/) (2024)
- [Learning and Retrieval from Prior Data for Skill-based Imitation Learning](https://mlsanity.com/paper/2210.11435/) (2022)
- [MUTEX: Learning Unified Policies from Multimodal Task Specifications](https://mlsanity.com/paper/2309.14320/) (2023)
- [MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts](https://mlsanity.com/paper/2407.21770/) (2024)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [RAM: Retrieval-Based Affordance Transfer for Generalizable Zero-Shot Robotic Manipulation](https://mlsanity.com/paper/2407.04689/) (2024)
- [Distilling and Retrieving Generalizable Knowledge for Robot Manipulation via Language Corrections](https://mlsanity.com/paper/2311.10678/) (2023)
- [DINOBot: Robot Manipulation via Retrieval and Alignment with Vision Foundation Models](https://mlsanity.com/paper/2402.13181/) (2024)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)

## Built on by (most cited first, 5 of 5)

- [VLA-Pro: Cross-Task Procedural Memory Transfer for Vision-Language-Action Models](https://mlsanity.com/paper/2605.29562/) (2026)
- [Gentle Object Retraction in Dense Clutter Using Multimodal Force Sensing and Imitation Learning](https://mlsanity.com/paper/2508.19476/) (2025)
- [Domain Arithmetic: One-Shot VLA Adaptation under Environmental Shifts](https://mlsanity.com/paper/2607.00666/) (2026)
- [Reuse Before You Retrieve: Diagnosing Headroom and Complementarity for Test-Time Augmentation of Embodied Multimodal Policies](https://mlsanity.com/paper/2608.17484/) (2026)
- [Hierarchical Skill Retrieval for Data-Efficient Adaptation of Vision-Language-Action Models](https://mlsanity.com/paper/2608.24042/) (2026)

## Abstract (arXiv)

In this work, we study the problem of data retrieval for few-shot imitation learning: selecting data from a large dataset to train a performant policy for a specific task, given only a few target demonstrations. Prior methods retrieve data using a single-feature distance heuristic, assuming that the best demonstrations are those that most closely resemble the target examples in visual, semantic, or motion space. However, this approach captures only a subset of the relevant information and can introduce detrimental demonstrations, e.g., retrieving data from unrelated tasks due to similar scene layouts, or selecting similar motions from tasks with divergent goals. We present COLLAGE, a method for COLLective data AGgrEgation in few-shot imitation learning that uses an adaptive late fusion mechanism to guide the selection of relevant demonstrations based on a task-specific combination of multiple cues. COLLAGE follows a simple, flexible, and efficient recipe: it assigns weights to subsets of the dataset that are pre-selected using a single feature (e.g., appearance, shape, or language similarity), based on how well a policy trained on each subset predicts actions in the target demonstrations. These weights are then used to perform importance sampling during policy training, sampling data more densely or sparsely according to estimated relevance. COLLAGE is general and feature-agnostic, allowing it to combine any number of subsets selected by any retrieval heuristic, and to identify which subsets provide the greatest benefit for the target task. In extensive experiments, COLLAGE outperforms state-of-the-art retrieval and multi-task learning approaches by 5.1% in simulation across 10 tasks, and by 16.6% in the real world across 6 tasks, where we perform retrieval from the large-scale DROID dataset. More information at https://robin-lab.cs.utexas.edu/COLLAGE .