# STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning

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

## Key facts

- Cited by 44 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning and post-training, one of 13 directions in Vision-language-action models.
- Builds on 32 of its 39 references that are in the map.
- First posted on arXiv 19 Dec 2024.

## Placed in

- Vision-language-action models / [Reinforcement learning and post-training](https://mlsanity.com/vla/reinforcement-learning-and-post-training/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (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)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)
- [GMFlow: Learning Optical Flow via Global Matching](https://mlsanity.com/paper/2111.13680/) (2021)
- [RT-H: Action Hierarchies Using Language](https://mlsanity.com/paper/2403.01823/) (2024)
- [The Surprising Effectiveness of Representation Learning for Visual Imitation](https://mlsanity.com/paper/2112.01511/) (2021)
- [Toward General-Purpose Robots via Foundation Models: A Survey and Meta-Analysis](https://mlsanity.com/paper/2312.08782/) (2023)
- [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)
- [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)
- [HRL4IN: Hierarchical Reinforcement Learning for Interactive Navigation with Mobile Manipulators](https://mlsanity.com/paper/1910.11432/) (2019)
- [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)
- [PoCo: Policy Composition from and for Heterogeneous Robot Learning](https://mlsanity.com/paper/2402.02511/) (2024)
- [R+X: Retrieval and Execution from Everyday Human Videos](https://mlsanity.com/paper/2407.12957/) (2024)
- [FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning](https://mlsanity.com/paper/2408.16944/) (2024)
- [Core Challenges in Embodied Vision-Language Planning](https://mlsanity.com/paper/2106.13948/) (2021)
- [Learning Options via Compression](https://mlsanity.com/paper/2212.04590/) (2022)
- [Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation](https://mlsanity.com/paper/2408.16228/) (2024)
- [Offline Imitation Learning Through Graph Search and Retrieval](https://mlsanity.com/paper/2407.15403/) (2024)
- [Zero-shot Imitation Policy via Search in Demonstration Dataset](https://mlsanity.com/paper/2401.16398/) (2024)

## Built on by (most cited first, 30 of 44)

- [Robometer: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons](https://mlsanity.com/paper/2603.02115/) (2026)
- [DataMIL: Selecting Data for Robot Imitation Learning with Datamodels](https://mlsanity.com/paper/2505.09603/) (2025)
- [Graph-based Agent Memory: Taxonomy, Techniques, and Applications](https://mlsanity.com/paper/2602.05665/) (2026)
- [A Survey on Imitation Learning for Contact-Rich Tasks in Robotics](https://mlsanity.com/paper/2506.13498/) (2025)
- [PEEK: Guiding and Minimal Image Representations for Zero-Shot Generalization of Robot Manipulation Policies](https://mlsanity.com/paper/2509.18282/) (2025)
- [Data Retrieval with Importance Weights for Few-Shot Imitation Learning](https://mlsanity.com/paper/2509.01657/) (2025)
- [SPECI: Skill Prompts based Hierarchical Continual Imitation Learning for Robot Manipulation](https://mlsanity.com/paper/2504.15561/) (2025)
- [Human2LocoMan: Learning Versatile Quadrupedal Manipulation with Human Pretraining](https://mlsanity.com/paper/2506.16475/) (2025)
- [Uni-Skill: Building Self-Evolving Skill Repository for Generalizable Robotic Manipulation](https://mlsanity.com/paper/2603.02623/) (2026)
- [Projected Coupled Diffusion for Test-Time Constrained Joint Generation](https://mlsanity.com/paper/2508.10531/) (2025)
- [RT-Cache: Training-Free Retrieval for Real-Time Manipulation](https://mlsanity.com/paper/2505.09040/) (2025)
- [COLLAGE: Adaptive Fusion-based Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2508.01131/) (2025)
- [Retrieve-then-Steer: Online Success Memory for Test-Time Adaptation of Generative VLAs](https://mlsanity.com/paper/2605.10094/) (2026)
- [Retrieve, Don't Retrain: Extending Vision Language Action Models to New Tasks at Test Time](https://mlsanity.com/paper/2606.15631/) (2026)
- [RealDrive: Retrieval-Augmented Driving with Diffusion Models](https://mlsanity.com/paper/2505.24808/) (2025)
- [Factorizing Diffusion Policies for Observation Modality Prioritization](https://mlsanity.com/paper/2509.16830/) (2025)
- [Policy Compatible Skill Incremental Learning via Lazy Learning Interface](https://mlsanity.com/paper/2509.20612/) (2025)
- [Dejavu: Towards Experience Feedback Learning for Embodied Intelligence](https://mlsanity.com/paper/2510.10181/) (2025)
- [Learning from the Best: Smoothness-Driven Metrics for Data Quality in Imitation Learning](https://mlsanity.com/paper/2604.23000/) (2026)
- [SimpleMemVLA: A Simple but Effective Native-Video Memory for Vision-Language-Action Models](https://mlsanity.com/paper/2609.05533/) (2026)
- [Retrieval-Augmented Robots via Retrieve-Reason-Act](https://mlsanity.com/paper/2603.02688/) (2026)
- [Tether: Autonomous Functional Play with Correspondence-Driven Trajectory Warping](https://mlsanity.com/paper/2603.03278/) (2026)
- [Grounding Robot Generalization in Training Data via Retrieval-Augmented VLMs](https://mlsanity.com/paper/2603.11426/) (2026)
- [Belief-Conditioned One-Step Diffusion: Real-Time Trajectory Planning with Just-Enough Sensing](https://mlsanity.com/paper/2508.12166/) (2025)
- [Quality over Quantity: Demonstration Curation via Influence Functions for Data-Centric Robot Learning](https://mlsanity.com/paper/2603.09056/) (2026)
- [ReGIL: Retrieval-Guided Imitation Learning from a Single Demonstration](https://mlsanity.com/paper/2606.09381/) (2026)
- [Difference-Aware Retrieval Policies for Imitation Learning](https://mlsanity.com/paper/2606.09758/) (2026)
- [ATHENA: Accelerated Multi-Task Heterogeneous Influence Functions for Robot Data Curation](https://mlsanity.com/paper/2606.16208/) (2026)
- [End-to-End LLM Flight Planning with RAG-based Memory and Multi-modal Coach Agent](https://mlsanity.com/paper/2607.06964/) (2026)
- [Retrieve in Time, Correct in Frequency](https://mlsanity.com/paper/2608.04527/) (2026)

## Abstract (arXiv)

Robot learning is witnessing a significant increase in the size, diversity, and complexity of pre-collected datasets, mirroring trends in domains such as natural language processing and computer vision. Many robot learning methods treat such datasets as multi-task expert data and learn a multi-task, generalist policy by training broadly across them. Notably, while these generalist policies can improve the average performance across many tasks, the performance of generalist policies on any one task is often suboptimal due to negative transfer between partitions of the data, compared to task-specific specialist policies. In this work, we argue for the paradigm of training policies during deployment given the scenarios they encounter: rather than deploying pre-trained policies to unseen problems in a zero-shot manner, we non-parametrically retrieve and train models directly on relevant data at test time. Furthermore, we show that many robotics tasks share considerable amounts of low-level behaviors and that retrieval at the "sub"-trajectory granularity enables significantly improved data utilization, generalization, and robustness in adapting policies to novel problems. In contrast, existing full-trajectory retrieval methods tend to underutilize the data and miss out on shared cross-task content. This work proposes STRAP, a technique for leveraging pre-trained vision foundation models and dynamic time warping to retrieve sub-sequences of trajectories from large training corpora in a robust fashion. STRAP outperforms both prior retrieval algorithms and multi-task learning methods in simulated and real experiments, showing the ability to scale to much larger offline datasets in the real world as well as the ability to learn robust control policies with just a handful of real-world demonstrations.