STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning

2024-12-19

arXiv.org

AI Summary

In offline robotic reinforcement learning, pre-trained multi-task generalist policies often underperform task-specific expert policies due to negative transfer. To address this, we propose a dynamic adaptation paradigm that—during deployment—retrieves and fine-tunes relevant sub-trajectories on-demand, overcoming limitations of full-trajectory retrieval and static generalist policies. We introduce the first non-parametric, sub-trajectory-level retrieval method, leveraging cross-task sharing of low-level behaviors. By integrating vision foundation model representations with Dynamic Time Warping (DTW), our approach enables robust, variable-length, cross-task sub-trajectory matching and online policy enhancement. Extensive experiments in both simulation and real-robot settings demonstrate significant improvements over state-of-the-art retrieval-based and multi-task methods. Our framework supports scalable expansion to massive offline datasets and learns robust control policies from only a few real-world demonstrations.

Technology Category

• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Learning to Search
• Machine Learning: Online Learning & Bandits

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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.

Problem

• Retrieving sub-trajectories to improve policy adaptation
• Addressing negative transfer in multi-task robot learning
• Enhancing data utilization for novel problem generalization

Innovation

• Sub-trajectory retrieval for policy training
• Leveraging vision foundation models
• Dynamic time warping for robust retrieval