STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning
Authors: Marius Memmel
Jacob Berg
Bingqing Chen
Abhishek Gupta
Jonathan Francis
Venue: ICLR 2025 Poster (2025)
Keywords: dynamic time warping
few-shot imitation learning
retrieval
foundation models
Primary area: applications to robotics, autonomy, planning
TL;DR: Subsequence-DTW for sub-trajectory retrieval to augment few-shot policy learning

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.

Published as a conference paper at ICLR 2025
STRAP: ROBOT SUB-TRAJECTORY RETRIEVAL FOR AUGMENTED POLICY LEARNING
Marius Memmel'. Jacobb Berg" Bingging Abbishekk Gupta Gupthai,Jonthhan
Jonathan Francis
1 Paul G. Allen School of Computer Science & Engineering, University of Washington Roobo Learning Lab, Bosch Center for Artificial Intelligence Roobotics Institute, Carnegie Mellon University {memme Ima, jacob33, abhgupta fccs wwshington. {bi iinginggg................n chen, jon. f francis Hus www.bbo.cho.... com
ABSTRACT
Robot learning is witnessing significant increase in the size, diversity, and com- plexity of pre-collected datasets, mirroring trends in domains such as natural lan- guage processing and computer vision. Many robot learning methods treat such datasets as multi-task expert data and learn a multi-task, generalist policy by train- ing broadly across them. Notably, while these generalist policies can improve the average performance across many tasks, the performance of generalist poli- cies on any one task is often suboptimal due to negative transfer between parti- tions of the data, compared to task- task-specific specialist policies. In this work, we argue for the paradigm training policies during deployment given the scenar- ios they encounter: rather than deploying pre-trained policies to unseen prob- lems a zero-shot manner, we non-parametrically retrieve and train models di- rectly relevant data at test time. Furthermore, show that many robotics tasks share considerable amounts of low-level behaviors and that retrieval at the "rub"-trajectoy granularity enables significantly improved data utilization, gen- eralization, and robustness in adapting policies to novel problems. In contrast, existing full-trajectory retrieval methods tend underutilize the data and miss out on shared cross-task content. This work proposes STRAP, technique for leveraging pre- pre-trained vision foundation models and dynamic time warping to re- trieve sub- sequences of trajectories from large training corpora in robust fash- ion. 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 learn robust con- trol policies with just handful of real-world demonstrations. Project website at hww.bbu..itthuut github. wtt.uu...o//see/ io/strap/
INTRODUCTION
Offinn Dataant
Few Demonstrations
Robust Palicy
Robot learning techniques have shown the ability to shift the process of designing robot controllers from large manual or model-based process to a data-driven one (Francis et al.] 2022 Hu et al.) 2023). Especially, end-to-end imitation learning with, e.g. diffusion models (Chi et al.]
Augmented Policy Learning
Cynomic mipe warsing
2023 Wang et al.] (2024) and trans- formers (Haldar et al. 2024). have shown impressive success. While imitation learning can be effective for
Retrieval Dataset
Figure 1: STRAP retrieves sub-trajectories using subsequence dy- namic time warping for training robust policies during deployment.
cqual contribution equal advising

---

Published as a conference paper at ICLR 2025
performing particular tasks with targeted in-domain data collection, this process can be expensive and time-consuming in terms of human effort. This becomes a challenge as we deploy robots into dynamic environments such homes and offices, where new tasks and environments are commonplace and constant data collection is impractical.
Multi-task policy learning is often applied in such situations, where data across multiple tasks is used to train a large task- or prrructionnconditionnn model that has the potential to generalize to new problems. While multi-task learning has seen successes in certain settings (Reed et al.) [022] Brohan et al. 2022). the performance of a multi-task, generalist policy is often lower than task- specific, specialist policies. This can be attributed the model suffering from negative transfer and sacrificing per-task performance to improve the average performance across tasks. This challenge is exacerbated in unseen tasks or domains since zero-shot generalization challenging and collecting large amounts of in-domain fine-tuning data can be expensive. In this work, we consider training expert models during test time as a better way to use pre-collected datasets and enable few-shot imitation learning for new tasks.
In particular, we build on the paradigm of non-parametric data retrieval, where a small amount of in-domain data collected at test-time is used to retrieve subset of particularly "relevant" data from the training corpus. This retrieved data can then be used for robust and performant model training on new tasks. In this sense, the retrieved data can guide learned models towards desired behavior, however, the question becomes: How do we sub-select which data to retrieve from a large, pre- existing corpus?
Several prior techniques have studied the problem of non hormmaramitricc retrieval from the perspective of learning latent embeddings that encode states (Du et al. [002)) skills (Nasiriany et al. 2022), optical flow (Lin et al.] 2024), and learned affordances (Kuang et al. 2024). Most techniques are challeng- ing to apply out of the box for two primary reasons. Firstly, they require training Nommin--specife encoders to embed states, skills, or affordances: this makes it challenging to apply to demonstra- tions collected in the open world, where visual appearance can show wide variations. Secondly, they often retrieve entire trajectories, limiting the policies' ability to use data from other tasks that may share common components with the desired test-time behavior. These challenges limit both the broad applicability of these retrieval methods and the amount of cross-task data sharing. How can we design easy-to-use off-the- retrieval methods that maximally utilize the training data for test-time adaptation?
The key insight in this work is that retrieval methods do not need to measure the similarity between entire trajectories (or individual states), but rather between tubjerajecsories of the desired behavior at test-time and corresponding sub-trajectories of the training data. Notably, these sub-trajectories do not need to come from tasks that are similar in entirety the desired test-time tasks. Instead, sub-components of many related tasks can be shared to enable robust, test-time policy training. For example, as shown in Fig. [ for the multi-stage task of "pick up the mug, put it in the drawer, and close it", both "pick up the mug, put in on top of the drawer" and "close the bottom drawer, open the top drawer" contain sub-tasks that when retrieved provide useful training data. Our pro- posed method, Sub-sequence Trajectory Retrieval for Augmented Policy Learning (STRAP), uses small amount of in-domain trajectories collected at test-time to retrieve and train on these relevant sub-trajectories across a large multi-task training corpus. The resulting policies show considerable improvements in robustness and generalization over previous retricval methods, zero-shot multi-task policies, policies that are trained purely on test-time in-domain data.
We show how STRAP can be used with minimal effort across training and evaluation domains with non-trivial visual differences. Our method first compares sub-trajectory similarity using features from off-the-shelf foundation models, e.g., DINOv2 (Oquab et al.) 2023); these features capture strong notions of "object-ness", discarding spurious visual differences such as lighting, texture, and local changes in object appearance. Secondly, our method leverages time-invariant alignment techniques, such as dynamic time warping (Giorginol 2009), to compute the similarity between sub- trajectories of different lengths, removing requirements for retrieved trajectories to have a similar length and increasing the applicability of STRAP across tasks and domains. Lastly, we show how STRAP can be applied to arbitrary test corpora, with sub-trajectories being automatically extracted by our framework, thereby removing the requirement for manual segmentation of relevant sub- trajectories from the training corpus. We demonstrate how STRAP can be used out of the box to augment any few-shot imitation learning algorithm, providing significant gains in generalization
2

---

Published as a conference paper at ICLR 2025
at test-time, while avoiding expensive, test-time in-domain data collection. We instantiate STRAP with transformer-bas based imitation learning policies and show the benefits of few-shot sub-trajectory retrieval on the LIBERO (Liu et al. 2024) benchmark in simulation and real-world imitation learning problems.
2 RELATED WORK
Retrieval for Behavior Replay: A considerable body of work has explored retrieval-basee ap- proaches for robotic manipulation, where the retrieval of relevant past demonstrations aids in re- playing past experiences. The choices of embedding space hereby range from off-the-shelf mod- els (Di Palo & Johns 2024, Malato et al. 2024) like DINO (Caron et al. (02T), training encoders on the offline dataset (Pari et al. 2022) to abstract representation like object shapes (Sheikh et al.) 2023). Some works do not directly replay actions but add layer of abstraction following sub- goals Zhang et al.) (024b)) affordances (Kuang et al., (024) or keypoints (Papagiannis al., 2024). key assumption of these methods is that the offline data either exactly resembles expert demon- bridge the gap. These drawbacks limit the usage of large multi-task datasets collected in various domains.
strations collected in the test environment that intermediate representations
Retrieval for policy learning tries mitigate these is- sues by learning policies from the retrieved data. While retrieval has shown to benefit policy learning
Retrieval for Few-shot Imitation Learning:
from sub-optimal single-task data (Yin & Abbeel 2024). most work focuses retrieving from large
multi-task datasets like DROID KKazatsky et al. 2024 or OpenX (O'Neil et al. (023) containing expert demonstrations. BehaviorRetrieval (BR) (Du et al., 2023) and FlowRetrieval (FR) (Lin al.,
2024) train an encoder-decoder model state-action and optical flow respectively. Related to our work, SAILOR (Nasiriany et al.) (202)) imposes skill constraints on the embedding space, clustering similar skills together to later retrieve those. A significant downside of training custom represen- tations is that these methods do not scale well to the increasing size of available offline datasets and are unable to deal with significant visual and semantic differences. Moreover, techniques like BehaviorRetrieval and FlowRetrieval retrieve individual states, rather than uub-trajectories like our work, where sub-trajectory retrieval enables maximal data sharing between seemingly different tasks while capturing temporal information.
Several works propose to decompose demonstrations into reusable sub-trajeed uub-trajectories, e.g., based on and-effector-centric or full proprioceptive state-action tran- sitions (Li et al. 2020 Belkhale et al. 2024|| [Sankaar et al.) 2022 Myers al. 2024 Francis al.) 2022). Belkhale et al. (2024) propose to decompose demonstrations into end-effector-centre sub- tasks, e.g., "move forward" or "rotate left". The authors show that by decomposing and re-labeling the language instructions into shared vocabulary, knowledge from multi-task datasets can be better shared when training multi-task policies. |Myers et al. (2024) leverage VL.Ms to decompose demon- strations into sub-trajectories to better learn to imitate them. To our knowledge, we propose the first robot sub-trajectory retrieval mechanism, for partitioning large offline robotics datasets and for enabling cross-task positive transfer during policy learning.
Learning from Sub-trajectories:
PRELIMINARIES
3.1 DYNAMIC TIME WARPING
To match sequences of potentially variable length during retrieval, we build on an algorithm called dynamic time warping (DTW) (Miller 202T). DTW methods compute the similarity between two time series that may vary in time or speed, e.g., different video or audio sequences. This algorithm aligns the varying length sequences by warping the time axis of the series using a set of step sizes to minimize the distance between corresponding points while obeying boundary conditions.
DTW algorithms are given two sequences, X = {xI, 22..,........ )n and = {y1 9/,, }. where 7 n, and a corresponding cost matrix C(zi, 2;) that assigns the cost of assigning element x; of sequence X to correspond with element y, of sequence Y. The goal of DTW is to find mapping between X and Y that minimizes the total cumulative distance between the assigned elements of both sequences while obeying boundary and continuity conditions. Dynamic time warping methods solve this problem efficiently using dynamic programming methods.
3

---

Published as a conference paper at ICLR 2025
Prw-ehot Demmetradimaaaar
Snood
Paee
Segment Darpet into mab- trajectaniee
Sec. 4.2
Ret Ratriemnluugmeeddd Pollcy Leurning
Duayi
Xocodot
voblle pai chdl pilstle cap
Dudsset Pprie
Dynante Time Warping
Sec. 4.
Pesttoal
ea Bihbva, CliP
Dragee
Thask pww.ir
Eneeder
Poter
Match rogmenta Darree to sub trajectuziie in Pyrter to obtala Pravieal
Bec.
Figure 2: Overview of STRAP: 1) demonstrations Purget and offline datasets Pprior are encoded into a shared embedding space using a vision foundation model, 2) automatic slicing generates sub- trajectories which 3) S-DTW matches to corresponding -utttrajectories in Prrior creating Pecrieaal. 4) training policy on the union of rretiival and Durget results in better performance and robustness.
A cumulative distance matrix D is computed via dynamic programming as follows: D(0, 0) C(0, 0), D(n, 1) = Σ C(k, 1) for n E [1 : N] and D(1, m) = k= C(1, for E [1 M]. Then the following dynamic programming calculation is performed:
(i..j) = C(.,, ui) + min
1,5), Dii, 1), Dii 1,, 1 1)];
where Cl yj) is the distance between points H; and Vy- We assume this cost matrix is pre- provided, and we describe how we compute this from raw camera images in Sec. 4.3 The optimal alignment between the sequences is found by backtracking from D(n, m) to D(O, 0). This guaran- tees that the start is matched to the start and the end is matched to the end or that the pairs (0o. yo) and (Zn, Dm.) are the start and end of the path. This optimal paring path consists of the best possible alignment between X and Y such that the cumulative cost between all matched pairs is minimized. DTW as described, widely used time-series analysis, speech recognition, and other domains where temporal variations exist between sequences. In the context of our retrieval problem, DTW is used to go beyond retrieving exactly matched sequences to matching variable length subsequences, as we describe below.
Subsequence dynamic time warping (S-DTW) is an extension of the DT' algorithm for sce- narios where shorter query sequence must be matched to portion of longer reference quence. Given a query sequence X = {Ii, .2.............................. In} and a much longer reference sequence
Y {v1, 1/2 ...... .... ym J. the goal of S-DTW is find a subsequence of Y (of a potentially different length from X), denoted Yı.5 where i S J. that has the minimal DDTW distance to X.
The cumulative cost matrix D for S-DTW is computed similarly the traditional DTW described above but allows alignment to start and end at any point in R. D is initialized as
D(0,0) = C),,
" D(n, 1) = Σ k=1 D(1,m) C(1,m) for m E M
1111))((0) C(K,1) for e [1:
and then completed using dynamic programming following Eq. ( This ensures that the query can match any sub-sequence the reference. Once the cumulative cost matrix computed, the optimal alignment is found by backtracking from the minimal value in the last row of the matrix, ie., min( j)) for j E [I.....m.. This gives the subsequence of Y that best aligns with X, obeying only temporality while relaxing the boundary condition. As we will show, using S-DTW for data retrieval enables the maximal retrieval of data across tasks in a retrieval-augmented policy training setting, described in Sec.
4

---

Published as a conference paper at ICLR 2025
"Put the black bowl in the battom drawer of the cabinet and close it'
DD target
S-DTW Matching
D prior
*Put the black bowl on top of the cabinot'
"Close the bottom drawer of the cabinet and open the top drawer"
Figure 3: Sub-trajectory matching: S- S-DTW matches the sub-trajectories of Dusreet (top) to the relevant segments D Prrior- A feature of is that the start and end of the trajectories do not have to align, finding optimal matches for each pairing.
4 STRAP: SUB-SEQUENCE ROBOT TRAJECTORY RETRIEVAL FOR AUGMENTED POLICY TRAINING
4.1 PROBLEM SETTING: RETRIEVAL-AUGMENTED POLICY LEARNING
We consider few-shot learning setting where we're given target dataset Diarget ((sps dg; s4, si,ai,...,&$, H, "H,: ! T') containing expert trajectories of states (e.g., observations like camera views o and propriception x), actions a (such as robot controls), and task-speciffing language instructions I. This target dataset is collected in the test environment and task, but there is only a small set of N trajectories, which limits generalization for models trained purely on such small dataset. Since Derree is often insufficient solve the task alone, we posit that generaliza- tion can be accomplished by non con-parametrically retrieving data from an offline dataset Preirr- This offline dataset D prior (s6 ag. st. .......,81 Hy, d'1, PP) P)) can contain data from different envi- ronments, scenes, levels of expertise, tasks, embodiments. Notably, the of tasks in the offline dataset do not need to overlap with the set of tasks in the target dataset. We assume that the offline dataset shares matching embodiment with the target dataset and consists of expert-level trajectories, but may consist of diversity of scenes and tasks that vary widely from the target dataset Duaget--
Given Pprior and Dusget: the goal to learn optimal actions in the target environment when prompted with the current state and language instruction I. Assuming we can obtain measure of success (such as task completion), and broad set of initial conditions 8o ~ (ro)(so) in the test environment. The objective of policy learning is to determine the policy parameters
policy aa(ass 1) that can predict
maximize the expected success metric when evaluated on test conditions, under the policy Te and test-time environment dynamics. Since a limited corpus of data, Deareet- in the target domain, these policy parameters cannot be learned Durget- Instead, we will present approach where
only provided
by simply performing maximum likelihood a smaller, "relevant" subset of the offline dataset 'D. Dretrieval C Dpriur is retrieved non-parametrically and then mixed with the smaller in-domain dataset Diarget to construct larger, augmented training dataset, i.e., Durget U Dretrieval. which is most relevant to the desired test-time conditions ( This can then be used for training policies via imitation learning, Doing so avoids an expensive generalist training procedure and rather focuses the learned model to being tigh-performing specialist in particular setting of interest. The key questions becomes How can we define what subsel of the offline datasel Dprior is relevant to construcí 'D, Detrienaa?
we will describe in Sec.
To handle the unique nature robotic data, e.g., multi-modal and temporally dependent, we design STRAP for -uggeettd policy learning. Firstly, we need to define the unit of retrieval. Rather than retrieving individual state-action pairs or entire trajectories, STRAP crucially retrieves sub-trajectories. We also propose method to automatically segment trajectories in Drarget into such sub-trajectories (Sec. 4.2) p. Secondly, we need to define a suitable distance metric for a pair of sub- trajectories (Sec. 4.3 Then, we need computationally efficient algorithm to retrieve relevant sub- trajectories non- paraamatriccallyy from the training set (Sec. 4.4). Finally, we put everything together and train policies based retrieved data (Sec. 4.5
5

---

Published as a conference paper at ICLR 2025
4.2 SUB- TRAJECTRESS FOR DATA RETRIEVAL
To make the best use of the training dataset while capturing temporal task-specific dynamics, we ex- pand the notion of retrieval from being able to retrieve entire trajectories or single states to retrieving variable-length sub-trajectories. In doing so, retrieval can capture the temporal dynamics of the task, while still being able to share data between seemingly different tasks with potentially different task instruction labels. We define a sub-trajectory as consecutive subset of trajectory ty S T with the sub- -tajeerayetory aa = (s s sat1:..... .(.) including timestep to b of the whole trajectory T of length H,. Most Long-horizon et al. 2024 O'Neiil et al. 2023) naturally contain multiple such sub-trajectories. For instance, the task shown in Fig. can be decomposed into "put the bowl in the drawer" and "close the drawer". Note that we do not require these trajectories to explicitly have a specific semantic meaning, but semantically meaningful sub-trajectories often coincide with those most commonly encountered across tasks as we see in our experimental evaluation.
problems observed in robotics datasets (Liu et al. (024) Khazatsky
Given this definition of a sub-trajectory, our proposed retrieval technique only requires segmenting the target demonstrations into sub-trajectories Turga {ti :a t '......., ".....,,,H,,p..H... VT D Dirget but not the much larger offline training dataset Dprior- Instead, appropriate sub-sequences will be retrieved from this dataset using a DTW based retrieval algorithm (Sec. (.4.. This makes the pro- posed methodology far more practical since Preso is much larger than Parger- While this separation into sub-trajectories can be done manually during data collection, we propose an automatic tech- nique for sub-trajectory separation that yields promising empirical results. Building on techniques proposed by [ellkale al. (2024), we split the demonstrations into atomic chunks, i.e., lower-level motions, before retrieving similar uubetrajectories with our matching procedure (Sec. (.4). In partic- ular, we propose simple proprioception-bss. in the robot's end-effector motion indicating the transition between two chunks. For example, Pick&Place task can be split into picking and placing separated by a short pause when grasping the object. Let 24 be vector describing the end-effector position timestep t. We define "transition states" where the absolute velocity drops below a threshold: || < = We empirically find that this proprioception-driven segmentation can perform reasonable temporal segmentation of target trajectories into sub-components. This procedure can certainly be improved further via techniques
segmentation technique that optimizes for changes
in action recognition using iision-foundation models (Team al., 2023 Zhang al., 2024a). or
information-theorrtic segmentation methods (Jiang el al. 2022).
4.3 FOUNDATION MODEL-DRIVEL RELEVANCE METRICS FOR RETRIEVAL
Given the definition and automatic segmentation of sub-trajectories, we must define measure of similarity that allows for the retrieval of appropriate relevant sub-trajectory data from Dprior and at the same time is robust to variations in visual appearance, distractors, and irrelevant spurious features. While prior work has suggested objectives to train such similarity metrics through repre- sentation learning (Du et al., 2023 Lin et al., 2024 Kuang et al. 2024), these methods are often trained purely in-domain, making them particularly sensitive to aforementioned variations. While using more lossy similarity metrics based on optical flow (c.f. (Lin et al.) 2024)) language (Zhal iet al. 2024) can help with this fragility, it often fails to capture the necessary task-specific or seman- tie details. This suggests the need for a robust, domain- agnossic similarity metric that can easily be applied out-of-the-boo
In this work, we will adopt the insight that vision(-language foundation models (Oquab et al.) 2023 Radford et al., 2021) offer off-the-shelf solutions to this problem of measuring the seman- tie and visual similarities between sub-trajectories. capturing object- and task-centric affordances, while being robust to low-level variations in scene appearance. Trained on web-scale rorl-world image(-text) data, these models are typically robust to low-level perceptual variations, while pro- viding semantically rich representations that naturally capture a notion of object-ness and semantic correspondence. Denoting a vision foundation model as F(-), we can compute the pairwise distance of two camera views with an L2 norm in embedding space, i.e., [F...)) - F(o,)) 1l22. While aggre- gation methods such as temporal averaging could be used to go from embedding of a single image
For trajectories involving "stop-motion", this heuristic returns many short chunks as the end-effector idles, waiting for the gripper to close. To ensure a minimum length, we merge neighboring chunks until all are 2 20. OOter cost metrics such as (1-cosine similarity) could he used here as well.

---

Published as a conference paper at ICLR 2025
Algorithm STRAP (D (Derget, Pprices K, E, F)
Require: demos Durget, offline dataset Ppriors vision foundation model F, # retrieved chunks K,
chunking threshold e;
1* Pre- Pre-processing */
2: Turget + Sub.raajeementtti.oo (argott e);
> Heuristic demo chunking
3: Eprior f- [[F(o:))o E THIT D
price
Embed Prior
4: Eurget t- ([F(oo)) TT E largel
Embed chunked D Dunget
5: 1 Sub-trajectory Retrieval using
S-DTWW
6: for Suiget E Durger do
7: M
M+L: +l;
Initialize empty match storage
8: for Tprior E prior do
9:
D conpptteeosttaatiiileeeeee Eprior);
Eq.
10:
- Mira prior:;
Dynamic Programming
11: end for
end for
13: Deetrieal + #=eetrleveeoooKlacchheeMMMM.
Sec. 4.4
1* Policy Learning */
15: repeat
sample B Durget Desiieal update policy me with loss L(B: 8)
Eq.
17: until converged; return
to that of sub-trajectory, they lose out the actions and dynamics. We instead opt for sub- trajectory matching procedure based on the idea DTW (Gorgino 2009) and use the embeddings for finding maximally relevant sub- sub-trajectories. Given two sub-trajectories, t; and łj. compute pairwise cost matrix C IR [x] where its value as computed by:
C(i,)) ((i.))= F(0;)ll2
4.4 UB-TRAJECTORY RETRIEVAL WITH SUBSEQUENCE DYNAMIC TIME WARPING
Given the above-mentioned definitions of sub-trajectories and goundation-modll--ie similarity metrics, we instantiate an algorithm to find the K most relevant sub-trajectories Tmatch from the offline dataset Dprior for each sub-trajectory t segmented from Dunget- Sub-trajectories may have variable lengths and temporal positioning within trajectory caused by varying tasks, platforms, or demonstrators. We employ (S-DTW to match the target sub-trajectories Tragett to appropriate segment Tmarch in Dprior (Sec. 1j. S-DTW scales naturally with these challenges and allows for retrieval from diverse, multi-task datasets. On deployment, subsequence dynamic time warping accepts query sub-sequences from the target dataset, i.e., harges. and uses dynamic programming to compute matches that are maximally aligned with the query Tmatch (SDTW(f, Dprior), Vt E Turgert along with matching costs, D. To construct D Deeirival. we select the K matches with the lowest cost uniformly across the sub sub-traject uub-trajectories in Turget. i.e., the same number of matches for each query until K matches are retrieved. We note that the resulting set of matches can contain duplicates if the demonstrations share similar chunks, but argue that if a chunk occurs multiple times in the demonstrations, it important to the task and should be "up-weighted"" this case through duplicated retrieval. For each match, we also retrieve its corresponding language instruction. The training dataset then contains a union of the target dataset Piarget and the retrieved dataset DD, Dreuieval: Durget -Deuiivvll This significantly larger, retrieval-augmentt dataset then be used to leam policies via imitation learning, leading to robust, generalizable policies.
the training set in
4.5 PUTTING ALL TOGETHERR; STRAP
To start the retrieval process, we encode image observations in Diarget and D prior using vision foun- dation model, e.g., DINOv2 (Oquab et al. 2023) or CLIP (Radford et al. 2021). To best leverage the multi-task trajectories in D Prioor. split the demonstrations in Dearget into atomic chunks based on a low-level motion heuristic. Then we generate matches between chunked D Prarget and D and construct D renieval by selecting the top K matches uniformly across all chunks. Combining Dremival

---

Published as a conference paper at ICLR 2025
Stove-Mokn
Bowl-Dree
Soup&Cherse
Mug-Mug
Book-Caddy
Table
Sink
Stove
Figure 12:
Tasks in Durget LIBERO-10 (left) and real-world DROID-Kitchen (right). For Poiior
please refer to LIBERO-90 (Liu al. 2024 for the simulated and Fig. 28 for real-world tasks.
with Pureer forms our dataset for learning policy. In standard policy learning setting, noisy re- trieval data can lead to negative transfer, e.g., when observations similar to the target data are labeled with actions that achieve different task. Without conditioning, such contaminated samples hurt the policy's 's downstream performance. We propose to use a language-condition this inconsistency. With conditioning, the policy can distinguish between samples from different tasks, separating misleading from expert actions while benefiting from positive transfer from the additional training data and context of the language conditioning.
policy to deal with
We use behavior cloning (BC) learn visuomotor policy similar to Haldar et al. (2024); Nasiri- any et al. (2024). We choose a ransformer-basee (Vaswami 2017) architecture feeding in a his- tory of the last h observations St-het and predicting chunk of h future actions using Gaus- sian mixture model action head. We sample batches from the union of Darges and Dreunival. as in Derrget U Pretrieval- As proposed in Haldar et al. (2024) we compute the multi-step action loss and add an L2 regularization term over the model weights U, resulting in the following loss function:
()
=
B Σ
log( no (0z:4h (aihinlsssssss h.s 1)) + A|10112
(es (es ()) .1)
with policy and hyperparameter À controlling the regularization.
EXPERIMENTS AND RESULTS
5.1 EXPERIMENTAL SETUP
Task Definition: We demonstrate the efficacy of STRAP in simulation on the LIBERO bench- mark (Liu et al. 2024), in two real-world scenarios following the DROID (Khazatsky et al. 2024) hardware setup. Fig. 12 shows the target tasks and samples from the retrieval datasets. For more task details please refer to Appendix A.2.
LIBERO: We evaluate STRAP on 10 long-horizon tasks of the LIBERO benchmark (Liu et al., 2024) which include diverse objects, layouts, and backgrounds. We randomly sample 5 demon- strations from LIBERO-10 as D Drurgor and utilize the 4500 trajectories in LIBERO-90 as 'D Dprior The evaluation environments randomize the target object poses, providing an ideal test bed for robust- ness. We report the top five tasks and average LIBERO-10 performance in Tab. and provide the remaining ones in the appendix (Tab.
DROID-Kitchen Scaling STRAP to more realistic scenarios, we evaluate STRAP vegetable leck-and-paace in three real kitchen environments. We collect 150 multi-task demonstrations across scenes in D Pprior (Kitchen). with 50 per scene. Each demo includes two unique tasks, ran- domly sampled from three options, with varied object poses and appearances. Dearget contains three demos of single task per environment. To assess STRAP 's scalability, we sample 5000 demos from the DROID dataset and 50 from D Disget 's environment (Kitchen-DDe -ROIDD....D During evaluation, object poses vary within 20 x 20cm grid. Success rates measured as partial (pick, 0.5) and full (pick and place, 1.0), normalized to a - 100 scale.
Baselines and Ablations: We compare STRAP to the following baselines and ablations and refer the reader to Appendix A. for implementation details and Appendix A.3 for extensive ablations.
Behavior Cloning (BC) behavior cloning using a transformer based policy trained on D Durrgeti
Fine-tuning (FT) behavior cloning using transformer-basee policy pre-trained Dprior and
fine-tuned on murgett
Multi-task Policy (MT) transformer-hased multi-task policy trained D Dprior and Diarget'
8

---

BC Published as a conference paper at ICLR 2025
Table 1: Baselines: Performance of baselines, ablations and variations STRAP on the LIBERO-10 tasks
(Fig. [2). DINOV2 and CLIP features perform similarly, making STRAP flexible in the encoder choice. Addi- Bold indicates best and underline runner-up up results.
tional results in Tab.
Task
Stove-Moka
Bowl-Cabinet
Soup-Cheese
Mug-Mug
Book-Caddy
LIBERO- 10
77.3% + 1.4
71.3% = 5.7
27.3% 2.2
38.0% + 5.7
75.3% + 1.4
37.9% = 27.2
FT
56.0%
1.4
91.0% 0.7
38.0% 2.8
43.0% 0.7
100.0% + 0.0
51.7% 4:
MT
66.0% 14.1
45.0% 1 30.4
19.0% 1 13.4
31.0% 16.3
100.0%
+: 0.0
37.7% 1 32.6
BR
80.0% +: 1.6
72.0% = 7.7
26.0% 1: 5.3
8.6 29.3%
40.0%
16.0% +: 1.9
33.4% 1: 25.1
FR
76.0% + 6.6
51.7% 12.0
21.7% +: 8.6
52.0% + 5.9
33.1% = 23.0
D-S
70.7% + 7.9
65.3% = 2.0
18.0% + 3.4
16.0% + 0.9
57.3% # 2.9
28. # 26.4
D-T
78.1% + 2.7
75.3% = 2.7
37.3% + 6.6
63.3% +: 3.6
79.0% + 5.0
41.4% + 30.8
STRAP (CLIP, K=100)
56.0% + 4.1
90.7% 2.2
12.0% + 0.9
51.7% + 3.3
83.3% + 3.0
41.9% = 32.7
STRAP (DINOv2, K=100)
85.3% + 2.2
91.3% - 2.2
42.7% + 7.2
57.3% + 7.7
85.3% + 2.8
45.6% + 32.6
STRAP (DINOv2, best K)
94.0% + 1.4
96.0% 1: 0.0
42.7% + 7.2
69.0% d: 0.7
94.0% + 4.2
58.1 + 32.0
BR (BehaviorRetrieval) (Du al. 2023) prior work that trains a VAE on state-action pairs for retrieval and uses cosine similarity to retrieve single state-action pairs;
FR (FlowRetrieval) (Lin et al.] 2024) same setup BR but VAE is trained on pre-computed
optical flow from GMFlow (Xu et al. 2022);
D-S (DINOv2 state) same as BR and FR but uses off-the-shelf DINOv2 (Oquab et al. 2023)
features instead of training a VAE;
D-T (DINOv2 trajectory) retrieves trajeteco (rather than sub-trajectories) with S S-DTW W and
DINOv2 features;
5.2 EXPERIMENTAL EVALUATION
Our evaluation aims to address the following questions: (1) Does sub-trajectory retrieval improve performance in few-shot imitation learning? (2) How effective are the representations from vision- foundation models for retrieval? (3) What types of matches are identified by SS-D7W?
Does sub-trajectory retrieval improve performance in few-shot imitation learning? STRAP out- performs the retrieval baselines BR and FR on average by + +24.7% and +25.0% across all 10 tasks (Tab. D. These results demonstrate the policy's robustness to varying object poses. BC represents strong baseline on the LIBERO task as the benchmark's difficulty comes from pose variations during evaluation. By memorizing the demonstrations, BC achieves high success rates, outperforming BR and FR by +4.5% and +4.8% across all 10 tasks. In our real-world experiments. BC performs much worse due the increased randomization during evaluation. The policy replays the demonstrations in Dunget failing to adapt to new object poses.
Pre-training a policy on Dprior and fine- tuning it on Dunget (FT) emerges as the
Table 2: Real-world results: DROID-Kitchen
most competitive baseline underperform- ing STRAP by only -6.4%. Training a multi-task policy (MT) matches BC on average but shows improvements when Dprior contains demonstrations or envi- ronments overlapping with Dunget- Intro- ducing larger randomization to Dprior and the evaluation part of the real-world experiments hurts the performance of both methods. Specif- ically, we find checkpoint selection for FT difficult as it's easy for the policy to under- or overfit the demonstrations, leading degenerate behavior or exact replay. We found MT training challenging as the policy (Haldar et al.) (2024] Nasiriany et al. (2024) sometimes ignores the language instruction and solves task seen in D parior instead of the conditioned Durgen- We hypothesize that this behavior emerges due to the data imbalance of 'D prior and Dusget- Finally, augmenting Dpriio with 5000 tra- jectories from the DROID dataset amplifies these challenges leading to an even larger performance gap. In our real-world evaluations, we find STRAP to experience surprising generalization behavior to poses unseen in Durget- The policy further shows recovery behavior, completing the task even when the initial grasp fails and alters the object's pose. Since STRAP's 's policy training stage is inde- pendent of the size of D Dprior but the dataset size is determined by hyperparameter K, it can naturally deal with adding in larger datasets like DROID maintaining performance.
Kitchen Sink
Kitchen+DROID
Table
Stove
Table
Sink
Stove
BC
12.50
10.00
14.28
12.50
10.00
14.28
FT
20.00
27.27
30.43
28.00
8.60
22.72
MT
4.34
31.57
45.00
2.00
0.00
0.00
STRAP
36.36
61.36
57.12
56.81
63.04
45.45
9

---

Published as a conference paper at ICLR 2025
RehaviorRetrieval
Flow Retrivval
STRAP (ours)
Parcentuge (%)
close the bettom drawer cabneet open top drawer
put bowl top cabinet
torn on the stove and put the trytng pon on
close the top drawer of the cabmet
pick up the book and place in the left comportment of the caday
put the black bowd in the bottom drower of the cabinet and close
close the
drower culicet
black bowl in bottom drower the cabtnet
pick up the
place the right compariment of cuddy
others
Figure 13: Tasks distribution in Dretrieval for different retrieval methods with target task "put the black bowl the bottom drawer of the cabinet and close it".
How important are sub-trajectories for retrieval and how many should be retrieved? To in- vestigate the efficacy of sub-trajectories. we compare sub- sub-trajectory retrieval (S to retrieving full trajectories with (S-DTW (D-T) in Tab. T We find sub sub-trajectory retrieval im- prove performance by +4.1% across all TO tasks. We hypothesize that full trajectories can contain segments irrelevant to the task, effectively hurting performance and reducing the accuracy of the cumulative cost. Varying the number of retrieved segments K, we find that the optimal value for K is highly task-dependent with some tasks benefiting from retrieving less and some from retrieving more data. We hypothesize that K depends whether tasks leverage (positive transfer) or suf- fer (negative transfer) from multi-task training. We report the results for the best K in Tab. I and provide the full search in Tab.
How effective are the representations from vision-foundation models for retrieval? Next, we ablate the choice of foundation model representation in STRAP with fixed K = 100. We com- pare CLIP, a model trained through supervised learning on image-text pairs, with DINOv2, a self- supervised model trained on unlabeled images. We don't 't find any representation to significantly outperform the other with DINOv2 separated from CLIP by only +0.7% across all 10 tasks. To show the efficacy of tors from prior work (BR, FR) trained on ' prior with an off-the-shelf DINOv2 encoder model (D-S). Comparing them their natural configuration, Le., state-based retrieval using cosine similarity al- lows for side-by-side comparison of the representations. Tab. Tab..ssaat shows the choice of representation to depend on the task with no method outperforming the others on all tasks. Since D-S has no notion of dynamics and task semantics due to single-state retrieval, BR and FR outperform it by +5.0% and +4.7%, respectively. We highlight that vision foundation models don't 't have to be trained on D Dpnir and scale much better with increasing amounts of trajectory data and unseen tasks.
models for retrieval. we replace the in-domain feature extrac-
What types of matches are identified by S-DDWW To understand what data STRAP retrieves, we visualize the distribution over tasks function of Dreurieval proportion in Figure 13 The figure visualizes the top five tasks retrieved and accumulates the rest into the "others" category. categor.. It becomes clear that STRAP retrieves semantically relevant data each task shares least one sub-task with the target task. For example. "put the black bowl in the bottom drawer of the cabinet", "close the bottom drawer of the cabiner... ..." (Fig. 3. Furthermore, STRAP's retrieval sparse, only selecting data from 5/90 semantically relevant tasks and ignoring irrelevant ones. We observe that DINOv2 features are surprisingly agnostic to different environment textures, retrieving data from the same task but in a different environment (c.f. Fig. [3] "pul the black bowl in the botlom drawer of the cabinet and close it"). Furthermore, DINOv2 is robust to object poses retrieving sub-trajectories that "close the drawer" with the bowl either on the table or in the drawer (c.f. Fig. 31 "close the bottom drawer of the cabinet and open the top drawer"). Trained on optical flow, FR has no notion of visual appearance, failing to retrieve most of the semantically relevant data.
6 CONCLUSION
We introduce STRAP as an innovative approach for leveraging visual foundation models in few- shot robotics manipulation, eliminating the need to train on the entire retrieval dataset, and allowing scale with minimal compute overhead. By focusing on sub-trajectory retrieval using DDDWW STRAP improves data utilization and captures dynamics more effectively. Overall, our approach out- performs standard fine-tuning and multi-task approaches as well as state-of-the-ar retrieval methods by a large margin, showcasing robust performance in challenging real-world scenarios.
10

---

Published as a conference paper at ICLR 2025
Acknowledgements: The authors thank the anonymous reviewers for their valuable feedback, Soroush Nasiriany for answering our many questions on the robocasa and robosuite codebase, Li- Heng Lin and Yuchen Cui for helping understand policy learning with retrieved data, colleagues and fellow interns at Bosch for the fruitful discussions. The work has received funding from Ama- zon (CC OTH 00203375 2024 TR), ARL (W91 1NF2222222)))) and Robert Bosch LLC (most of this work was done during Marius Memmel's s 2024 summer internship project).
REFERENCES
Suneel Belkhale, Tianli Ding, Ted Xiao, Pierre Sermanet, Quon Vuong, Jonathan Tompson, Yevgen Chebotar, Debidatta Dwibedi, and Dorsa Sadigh. Rt-h: Action hicrarchies using language. arXiv preprint arXiv:2403.01823; 2024.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.068111 2022.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650-9660, 2021.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scatt scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 245. leee, 2009.
Norman Palo and Edward Johns. Dinobot: Robot manipulation via retrieval and alignment with vision foundation models. arXiv preprint arXiv:2402. 2024.
Maximilian Du, Suraj Nair, Dorsa Sadigh, and Chelsea Finn. Behavior retrieval: Few-shot imitation learning by querying unlabeled datasets. arXiv preprint arXiv:2304.08742, 2023.
Jonathan Francis, Nariaki Kitamura, Felix Labelle, Xiaopeng Lu, Ingrid Navarro, and Jean Oh, Core challenges in embodied vision- Dision-Language planning. Journal of Artificial Intelligence Research, 74: 459-515, 2022.
Toni Giorgino. Computing and visualizing dynamic time warping alignments in R: The dtw package. Journal of Statistical Sofwware 31(7):1-24 2009, doi: 10. (0.1187/jss.v031.... v07...07.
Siddhant Haldar, Zhuoran Peng, and Lerrel Pinto. Baku: An efficient transformer for multi-task policy learning. arXiv preprint arXiv:2406.07539 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. arxiv e-prints. arXiv preprint arXiv: 151.03355, 10, 2015.
Yafei Hu, Quanting Xie, Vidhi Jain, Jonathan Francis, Jay Patrikar, Nikhil Keetha, Seungchan Kim, Yaqi Xie, Tianyi Zhang, Zhibo Zhao, et al. Toward general-purpose robots via foundation models: A survey and meta-analysis. arXiv preprint arXiv:2312.0877 2023.
Yiding Jiang, Evan Liu, Benjamin Eysenbach, Zico Kolter, and Chelsea Finn. Learning options via compression. Advances in Neural Information Processing Systems, 15:21184-21 199, 2022.
Alexander Khazatsky, Karl Pertsch, Suraj Nair, Ashwin Balakrishna, Sudeep Dasari, Siddharth Karamcheti, Soroush Nasiriany, Mohan Kumar Srirama, Lawrence Yunliang Chen, Kirsty El- lis, et al. Droid: large-scale in-the-wild robot manipulation dataset. arXiv preprint 24X3::2403.12945 2024.
Yuxuan Kuang, Junjie Ye, Haoran Geng, Jiageng Mao, Congyue Deng, Leonidas Guibas, He Wang, and Yue Wang. Ram: Retrieval-based affordance transfer for generalizable zero-shot robotic manipulation. In 8th Annual Conference on Robot Learning, 2024.
11

---

Published as a conference paper at ICLR 2025
Chengshu Li, Fei Xia, Roberto Martin-Martin, and Silvio Savarese. Hrl4in: Hierarchical rein- forcement learning for interactive navigation with mobile manipulators. In Conference on Robot Learning. pp. 603-616. PMLR, 2020.
Li-Heng Lin, Yuchen Cui, Amber Xie, Tianyu Hua, and Dorsa Sadigh. Flowretrieval: Flow-guided data retrieval for few-shot imitation learning. In 8th Anmual Conference on Robot Learning, 2024.
Bo Liu, Yifeng Zhu, Chongkai Gao, Yihao Feng, Qiang Liu, Yuke Zhu, and Peter Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning. Advances in Neural Information Processing Systems, 36, 2024.
Federico Malato, Florian Leopold, Andrew Melnik, and Ville Hautamäki.. Zero-shot imitation policy via search in demonstration dataset. In ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7590-7594, IEEE, 2024.
Vivek Myers, Chunyuan Zheng, Oier Mees, Kuan Fang, and Sergey Levine. Policy adaptation via language optimization: Decomposing tasks for few-shot imitation. In 8th Annual Conference on Robot Learning, 2024.
Meinard Müller. Fundamentals of Music Processing: Using Python and Jupyter Notebooks. Springer Cham, 2 edition, 2021. ISBN 978-3-030-69807-2. URL 1007 7 1978-1981 3- 030- 6 59808-9
Soroush Nasiriany, Tian Gao, Ajay Mandlekar, and Yuke Zhu. Learning and retrieval from prior data for skill-based imitation learning. In Conference on Robot Learning, 2022.
Soroush Nasiriany, Abhiram Maddukuri, Lance Zhang, Adeet Parikh, Aaron Lo, Abhishek Joshi, Ajay Mandlekar, and Yuke Zhu. Robocasa: Large-scale simulation of everyday tasks for gener- alist robots. arXiv preprint arXiv:2406. 02523, 2024.
Abby O'Neill, Abdul Rehman, Abhinav Gupta, Abhiram Maddukuri, Abhishek Gupta, Ab- hishek Padalkar, Abraham Lee, Acorn Pooley, Agrim Gupta, Ajay Mandlekar, et al. Open x- embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2310.08864, 2023.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.077 000000,,711,, 2023.
Georgios Papagiannis, Norman Di Palo, Pietro Vitiello, and Edward Johns. R+ x: Retrieval and execution from everyday human videos. In RSS 2024 Workshop: Data Generation for Robotics, 2024.
Jyothish Pari, Nur Muhammad Mahi Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation. In 18th Robotics: Science and Systems, RSS 2022. MIT Press Journals, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 8748-8763 PMLR, 2021.
Scott E. Reed, Konrad Zoina, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando Freitas. A generalist agent. Trans. Mach. Learn. Res., 2022, 2022.
net www.oorrmmmddiiiooodooe ikkokk=
Tanmay Shankar, Yixin Lin, Aravind Rajeswaran, Vikash Kumar, Stuart Anderson, and Jean Oh. Translating robot skills: Learning unsupervised skill correspondences across robots. In Interna- fional Conference on Machine Learning, pp. 19626 19644. PMLR, 2022.
Jannik Sheikh, Andrew Melnik, Gora Chand Nandi, and Robert Haschke. Language-conditionn semantic search-basee policy for robotic manipulation tasks. In NewrIPS 2023 Foundation Models for Decision Making Workshop, 2023.
12

---

Published as a conference paper at ICLR 2025
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- lean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint 1rXiv:2312.11805, 2023.
A Vaswani. Attention all you need. Advances in Neural Information Processing Systems, 2017.
Lirui Wang, Jialiang Zhao, Yilun Du, Edward H. Adelson, and Russ Tedrake. Poco: Policy compo- sition from and for heterogeneous robot learning. CoRR, abs/2402.02511, 2024. doi: 10.48550/
ARXIV.240200211. URL httpsstss:///...ooo htt./oto..rggg///o 0e0e://1..8855..0 48550/85500/00 arxivv 2402.
Haofei Xu, Jing Zhang, Jianfei Cai, Hamid Rezatofighi, and Dacheng Tao. Gmflow: Learning optical flow via global matching. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 8121-8130, 2022.
Chao-Heng Yin and Pieter Abbeel. Offline imitation learning through graph search and retrieval. arXiv preprint arXiv:2407.15003, 2024.
Lihan Zha, Yuchen Cui, Li-Heng Lin, Minae Kwon, Montserrat Gonzalez Arenas, Andy Zeng. Fei Xia, and Dorsa Sadigh. Distilling and retrieving generalizable knowledge for robot manipulation via language corrections. 2024 IEFE International Conference on Robotics and Automation (ICRA), pp. 15172-15179. IEEE, 2024.
Jesse Zhang, Minho Heo, Zuxin Liu, Erdem Biyik, Joseph J Lim, Yao Liu, and Rasool Fakoor. Extract: Efficient policy learning by extracting transferable robot skills from offline data. arXiv preprint arXiv:2406. 171::2406611768. 2024a.
Yuying Zhang, Wenyan Yang, and Joni Pajarinen. Demobot: Deformable mobile manipulation with vision basse sub-goat retrieval. arXiv preprint 24X 15919, 2024b.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 SIMULATION EXPERIMENTS
Table 3: Baselines (sim): Performance of different methods on LIBERO-10 tasks in simulation. Bold indi- cates best and underline runner-up results.
Method
Mug-Mierowave
Moka-Moka
Soup-Sauce
Cream-Cheese-Butte
Mug-Pudding
BC
28.0% 0.9
0.0% + 0.0
17.3% 4.5
26.7% +: 4.3
18.0%
Fine-Tuning
38.0% 5.7
0.0% = 0.0
5.0% + 2.1
81.0% + 3.5
35.0% 3.5
Multi-Task
10.0% 7.1
0.0% + 0.0
24.0% + 17.0
73.0% + 13.4
9.0% = 2.1
BR (Du al.
2023
28.7% 3.9
0.0% + 0.0
13.3%
32.0% +: 4.3
26.0% +: 1.9
FR (Lin et al.
2024
27.3% 1.4
0.0% + 0.0
11.3% 3.0
41.3% + 5.5
14.7% + 1.1
D-S
30.0% + 3.4
0.0% + 0.0
4.7% + 0.5
16.0% + 5.7
6.0% + 0.9
D-T
34.7% 1. 2.0
0.0% 1: 0.0
4.7% 1:
27.3% 1. 4.5
14.0% 3.4
STRAP (CLIP, K=100)
30.0% 2.5
0.0% + 0.0
8.7% + 6.3
29.3%
10.5
24.0% 4.3
STRAP (DINOv2, K=100 100)
29.3% 1: 2.7
0.0% +: 0.0
16.7% 2.0
29.3%
11.3
18.7% +: 1.4
STRAP (DINOv2, best K)
32.0% 1: 5.7
0.0% =: 0.0
61.0% + 6.4 61.0% + 0.7
31.0% +
Hyperparameters: All results are reported over training and evaluation seeds (1234, 42, 4325). We fixed both the number of segments retrieved to 100, the camera viewpoint to the agent view image for retrieval, and the number of expert demonstrations to 5. We use the agent view (exocentric) observations for the retrieval and train policies on both agent view and in-hand observations. Our transformer policy trained for 300 epochs with batch size 32 and an epoch every 200 gradient steps.
Baseline implementation details: Following Lin et al. 2024). we retrieve single-state action pairs for the state-based retrieval baselines (BR, FR, D-S) and pad them by also retrieving the states from / - h to L. h -1 1 to make the samples compatible with our transformer-based policy. We refer the reader Appendix A.3 for extensive ablation.
Remaining results on LIBERO-10 Tab. 3 shows the results for the remaining LIBERO- 10 task not reported in the main sections. Both FR and BR outperform STRAP on the (ream-Cheese-Buttee task. We hypothesize that chunking heuristic generates sub-optimal tajetrajecssries (too long) causing them to contain multiple different semantic tasks, leading to worse matches in our retrieval datasets and eventually in decreasing downstream performance.
A.2 REAL-WORLD EXPERIMENTS
Setup & hyperparameters:
For retrieval, we average the embeddings per time-step across the left, right, and in-hand camera observations and fix the number of segments retrieved 100. We train policies on all three image observations for 200 epochs with batch size 32 and an epoch every 100 gradient steps. We initialize the ResNet- 18 (He al.) 2015) vision encoders of our policy with weights pre-trained on ImageNet (Deng et al.) (2009) We use the DROID-setup -etup (Khazatsky et al.) running at 15Hz collect demonstrations and deploy our policies on the real robot.
2024
A.2.1 FRANKA-PEN-IN-CUPP
Task: We evaluate STRAP's ability to retrieve from "unrelated" tasks in pen-in-cup scenario. The offline dataset prirrr consists of 100 single-task demonstrations across 10 tasks and Diarget contains 3 demonstrations for the pen-in-cup task. While some share sub-tasks, e.g., "put the pen next to the markers". "put the screwdriver into the cup", "put the green marker into the mug others are unrelated, e.g., "make hotdog". "push over the box". "pull the cable the right". Note that the target task is not part of D ! We demonstrate STRAP's 's ability to an unseen pose, and object and target appearance.
14

---

Published as a conference paper at ICLR 2025
Table Tranka-Pen-ine Cup: real-world results.
Pen-in-Cup
base
OOD
Pick
Place
Pick
Place
BC
100%
100%
0%
0%
STRAP
100%
90%
100
100 %
Figure 14: chess
Figure 15:
Figure 16:
Figure 17: knock_over.box
Figure 18: marker in._mu
cube_stacking
hotdog
Figure 19: medicine.pmp
Figure 20: dispense_sone soap
Figure 21:
Figure 22:
Figure 23: screwdriver
pull_cable, Lightlllleighh
ven_next__o_ppem
Figure 24: Franka-Pen-in- Pen-in-Cup: Environment setup for the tasks in D prior
Language Instructions D Ppzior
chess cube_stacking
"Move the king to the top right of the chess board"
"Stack the blue eube on top of the tower"
hotdog knock_over.bov marker in mug medicine -ppp dispense_son _soa pull_cable right
"Put the hotdog in the bun"
"Knock over the box"
"Put the marker in the mug"
"Pick up the medicine box on the right and put it next the other medicine boxes"
"Press the soap dispenser"
"Pull the cable to the right"
pen_next_to__ _epens screwdriver
"Put the pen next the markers"
"Pick up the screwdriver and put the cup"
Table 5: Franka-Pen-in--CA Task/language instructions for the real-world dataset D prior.
15

---

Published as a conference paper at ICLR 2025
A.2.2 DROID DROID-KITCHEN
Figure 25: table
Figure 26: sink
Figure 27: stove
Figure 28: DROID-Kifchen: Environment setup for the tasks in D peior- Task-relevant objects are marked by red eircles. From left to right, the objects table: carrot, chips bag, can; sink: pepper, soap dispenser, sponge; stove: chili, chicken, utensil. We randomize the object pose and type, e.g., color or style during data collection.
Language Instructions D prior
Language Instruction Dearget
table
"press the soap dispenser", "pick up the sponge and put
"pick up the pepper and put it the sink"
sink".
"pick up the pepper and put it in the sink"
sink
"pick the carrot and put it on the plate".
"pick up the can and move it next to the
pick up the chips bag and put the plate". "pick up the and
table"
next the table"
stove
"pick up chicken wing and put in pan",
"pick up the chili and put it the pan"
"pick up the utensil and put it in the pan".
"pick the chili and put it in the pan"
Table DROID- Task/language instructions for the real-world datasets Dprior and Durret- Each task in Dprior consists of two unique tasks randomly sampled from the available instructions. The chosen instructions are concatenated by the filler" " then". the table environment, for example, a task might be "pick up the sponge and put it in the sink, then press the soup dispenser".
A.3 ABLATIONS
Table Ablations Retrieval Method: explore different approaches for trajectory- retrieval. Be- sides the heuristic reported in the main paper, we experiment with a sliding window approach that segments a trajectory into sub- trajectories of equal length (here: 30). We use S-DTW for both sliding window sub- trajectories and full trajectories.
Method
Stove-Moka
Bowl-Cabenet
Mug-Microwave
Moka-Moka
Soup-Cream-Chees
Sub-traj Full ura
76.0% 4.71 78.67% 2.72
75.33% + 2.72 68.67% 1.44
26.0% + 1.89 34.67% 1.96
0.0% 0.00
37.33% + 6.62 28.67% 3.81
0.0%
Method
Soup-Sauce
Cream- Chees-Cheseaaae Butter
Mug-Mug
Mug-Pudding
Book-Cadee
Sub-traj Full traj
40.00% + 0.94
27.33% + 2.18
63.33% + 3.57 43.33%
30.00% + 3.40
79.0% 4.95
4.67% 1.09
27.33% + 4.46
14.0% + 3.4
68.0% 5.66
16

---

Published as a conference paper at ICLR 2025
Stove-Pot
Success Rate (%))
66
66-
.60
Soup-Cheess Cheese
Mug-Mug
08
46-
Success Rate
28
08-
76
Book-Caddy
Mug-Microwave
Success Rate (%6)
100
1000
Pots Oats-On-See Stove
Soup- Sauce
Success Rate (96)
08
10-
000
20
cheese Butter BueteeBButter
106
Success Rate (96)
48
18
Figure 29: Ablations K (Num Segments Retrieved): The figure shows success and stan- dard deviation for different values K.
17

---

Published as a conference paper at ICLR 2025
Table Ablations K (Num Segments Retrieved): Tuning K improves the success rates reported in Tab. on 8/10 tasks. The optimal value for K is highly task-dependent with some tasks benefiting from retrieving less (Stove-Pi) Bowl-Cabiner, Soup- Cheese, Mug- (lug-Microwave) and some from re- trieving more (Mug-Mug. Book- Caddy, Soup-Sauce, Cream cheese- Mug-Pudding) data. We hypothesize that depends on whether tasks leverage (positive transfer) or suffer (negative transfer) from multi-task training.
Stove- Moka
Bowl-Cahinet
Soup-Cheese
Mug-Mug
Book-Caddy
K 100
86.0% 4.3
94.5% + 0.83
35.0% + 2.5
61.0% + 5.12
89.5% 3.56
K 200
94.0% = 1.41
96.0% + 2.83
33.0% + 2.12
56.0% 4.24
89.0% + 4.95
K 500
88.0%
2.83
95.0% = 0.71
31.0% = 6.36
65.0% 2.12
88.0% 4.24
K 1000
86.0%
1.41
96.0% = 0.00
32.0% 2.83
69.0% + 0.71
94.0% 4.24
K 2000
72.0%
4.24
74.0% + 1.41
25.0% 4.95
42.0% 4.24
73.0% 2.12
K 4000
76.0% 1.41
73.0% + 6.36
29.0% 2.12
39.0% 0.71
78.0% +
Mig-Microwave
Pots-On-Stove
Soup-Sauce
Cream chcese-Butte
Mug-Pudding
K 100
27.5% + 3.11
0.0% 0.00
45.5% + 10.52
50.0% + 5.61
24.0% 2.12
K 200
32.0% + 5.66
0.0% 1.00
31.0% + 2.12
46.0% + 2.83
22.0% 8.49
K 500
27.0% 3.54
0.0% 0.00
37.0% 3.54
44.0% 9.9
31.0% - 2.12
K 1000
23.0% 0.71
0.0% 0.00
61.0% + 6.36
61.0% 0.71
27.0% 2.12
K 2000
27.0% 3.54
0.0% 0.0
12.0%
1.41
27.0% 3.54
14.0% 1.41
K 4000
32.0%
0.0% + 0.0
12.0% 2.83
30.0% 1.41
13.0% 3.54
Table 9: Ablations Retrieval Seeds:
We S STRAP different retrieval seeds subset of LIBERO-10
tasks using 10 expert demos as Dremievl. retricval seeds
We report results over all possible combinations training and
Method
Stove-Moka
Mug-Cabinet
Book-Caddy
BC Baseline
81.78% 2.6
83.11% + 2.69 88.67% 2.11
93.11 = 1.57
STRAP
86.89% + 1.51
98.0% + 1.04
A.4
ADDITIONAL BASELINES
Table 10: Diffusion Policy Baselines:
Performance on LIBERO-10 tasks using diffusion policies (DP) Chi without language conditioning for BehaviorRetrieval (BR) (Du al. 2023). FlowRetrieval
al.
2023
(FR)
Lin 2024..
These experiments replicate the training setup BR FR.
Method
Stove-Moka
Bowl-Cabinet
Soup-Cheese
Mug-Mug
Book-Caddy
BR
36.67% + 1.44 68.67% 1 2.37
68.0% + 2.49 56.0% + 4.32
34.0% + 2.49
55.33% 1.44 56.0% 1 3.4
42.0% = 1.63 35.33% 1 6.28
FR
18.0% 1 3.4
Method
Mug- Microwaee
Pots-On-Stove
Soup-Sauce
Cream cheese- Butter
Mlug-Pudding
BR
30.67% + 0.54
0.00% + 0.00 68.0% +:2 2.49
10.67% +: 1.96
24.0% + 0.94 35.33% +: 0.54
9.33% + 1.44 8.0% 1.89
FR
32.67% 4 3.31
6.0% 4: 0.00
A.5
COMPLEXITY AND SCALABILITY
A.5.1 ENCODING TRAJECTORIES WITH VISION FOUNDATION MODELS
The embeddings for D Pprior can be precomputed and reused for every new retrieval process. Embed- ding dataset with total timesteps T and number of camera views V scales linear with O(T V). We benchmark Huggingface's DINOv2 implementation on an NVIDIA L40S 46GB using batch size 32. Encoding a single image takes 2.83ms + 0.08 (average across 25 trials). The wall clock time for encoding the entire DROID dataset 18.9M timesteps, single-view) therefore sums up to only 26h.
18

---

5000 Published as a conference paper at ICLR 2025
A.5.2 SUBSEQUENCE DYNAMIC TIME WARPING
In contrast to the embedding process, retrieval must be run for every deployment scenario, S-DTW consists of two stages: computing the distance matrix D and finding the shortest path via dynamic programming. Computing the distance matrix has complexity O(n m E) with E the embedding dimension, the length of the sub-trajectory in Durget, and m the length of the trajectory in D prior- DTW and backtracking have complexities of On .m) and (n), respectively. These stages have to be run sequentially for each sub-trajectory (E Dtarget) and trajectory (E 'py Dprior) but don' don't depend on the other (sub- suub))trajectories Therefore, STRAP has a runtime complexity of ( * +M) with the number of sub- trajtrraccttoies in Dtarget and M the number of trajectories in Dprior. Our implementa- tion largely follows We use numb to compile python functions into optimized machine code and warm-start every method by running it three times. Following the statistics of DROID, we choose trajectory length of 250 (Drarget and Prrior) and single demonstration from Dtarget split into 5 sub-trajectories of length 50 each and embed each timestep into a 768-dimensional vector mim- icking DINOv2 embeddings. We benchmark (S-DDTW] and report the wall clock time (average 10 trials) in Fig. 30 For an offline dataset the size of DROID (76k), retrieval takes approximately 300sec. Note that computing the distance matrix can be expressed as matrix multiplications and can leverage GPU deployment and custom CUDA kernels for even greater speedup.
A.5.3 POLICY TRAINING
The training process also has to be repeated for every deployment scenario. We use robomimic and train policies for 200 epochs with batch size of 32 and 100 gradient steps per epoch NVIDIA L40S 46GB. Training a single policy takes 35min 4 (average over 10 trials).
A.5.4 INCREASING THE DATASET SIZE
Overall, STRAP scales linearly with new trajectories added to D pprsor. STRAP encodes the new tra- jectories using off off-the-shelf vision foundation model, eliminating the need to re-train the encoder like in previous approaches (BR. FR). Retrieving data with DDDT scales linearly with the size of Dprior. allowing for retrieval within 5min even from the largest available datasets like DROID. Finally, STRAP 's policy learning stage is independent of the size of Dprior and only depends on the amount of retrieved data K, making more scalable than common pre-training + fine-tuning or multi-task approaches that have to be re-trained when new trajectories are added to D Drrior.
Linear Regression
Average Runtime (sec)
DTW Runtime
30
25
15
1000 2000 0-
6000 7000 3000 4000 Trajectories in Dorior
8000 3000 19000
Figure 30: STRAP Runtime: We benchmark STRAP's retrieval step on varying sizes of D report the average wall clock time over 10 trials.
prior and
SubsecuenceeTT..atmll
httrs://humbb..yyytta.oor
http:://github.coo//M11
19

---

Published as a conference paper at ICLR 2025
A.6
QUALITATIVE RESULTS
20
10
Count (k)
6
10
5-
2
100
100
25
Matches Start Index
Matches End Index
Matches Length
Figure 31: Match distribution Perior for STRAP with target task: "put the black bowl in the bottom drawer of the cabinet and close it". [S-DTW finds the best matches regardless of start and end points trajectory length. This results distribution over start and end points well variety of trajectory lengths retrieved.
20


Reviews:
[{"review_id": "KPXaJbwvPx", "reviewer": "Reviewer_Ndo3", "summary": "The paper introduces STRAP, a novel method for trajectory retrieval to find similar sub-trajectory segments in large-scale datasets for efficient policy learning in few-shot settings. The method's key contribution lies in combining pretrained visual encoders with Dynamic Time Warping to encode sub-trajectories of variable lengths for improved retrieval. The proposed method is tested against several retrieval baselines and BC ones on LIBERO-10 simulation and real world pick and place tasks and achieves good performance.", "questions": "- Could STRAP be combined with fine-tuning of MT policies on the retrieved dataset to potentially achieve even better performance than domain-specific fine-tuning alone?\n\n- How does STRAP's performance compare against standard fine-tuning approaches when controlling for the total amount of data used?\n\n- What are the memory and computational requirements for deploying STRAP on very large trajectory datasets like Droid?\n\n- Can you add the average performance for LIBERO-10 results to the main table?\n\n- Can you provide a few more real world tasks with required baselines?\n\n- Real world retrieveal is conducted with the same robot embodiment and gripper. How does STRAP Perform when retrieving similar data from other robot datasets like BridgeV2 that does not share the same robot and scenes?\n\nSTRAP presents an novel and well designed approach to few-shot learning, that tackles several drawbacks of prior methods through sub-trajectory retrieval with dynamic time-warping. However, more comprehensive comparisons against fine-tuned baselines and clearer analysis of computational requirements would strengthen the paper's contributions. Thus, I recommend weak reject pending addressing the following concerns: (1) comparisons against fine-tuned baselines, (2) clearer analysis of computational requirements, and (3) better justification of parameter choices.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is well motivated and achieves strong results across multiple experiments, both in simulation and real-world settings\n\n- The use of Dynamic Time Warping for sub-trajectory matching is novel and well-suited for the problem domain\n\n- Comprehensive evaluation against recent retrieval baselines demonstrates the method's effectiveness\n\n- Thorough ablation studies on different pretrained encoders provide valuable insights into architecture choices\n\n- The paper is well written and includes several illustrative figures that enhance the text", "weaknesses": "- The baseline comparison against multi-task policy appears weak, as it only uses pretrained weights without fine-tuning. This seems like an artificially weak baseline since fine-tuning is standard practice for all MT-policies. \n\n- The paper's argument that retrieval is more efficient than expensive pretraining needs stronger empirical support, especially given that the robotics community regularly fine-tunes general policies for downstream tasks\n\n- The computational cost of STRAP's retrieval process on large-scale datasets like Droid is not adequately addressed, raising questions about real-world scalability. Some more clarity is necessary here\n\n- The choice of K for constructing D-retrieval lacks sufficient explanation and ablations. The paper should explore how different K values affect both retrieval quality and computational overhead and policy performance, as this parameter likely presents a trade-off between performance and efficiency. A discussion about the retrieved data quantity would provide valuable insights and strengthen then paper. \n\n- Small number of tested tasks in real world setting and missing baselines of MT policy and finetuned MT policy", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L6oNwkDdYm", "reviewer": "Reviewer_JdfR", "summary": "This paper focuses on the setting of generalizing a policy to an unseen task with few-shot demonstrations. Instead of deploying zero-shot, the paper proposes STRAP, training a model on task-relevant data augmented by retrieval. STRAP first retrieves sub-trajectories from a large pretraining dataset that are similar to the new task demonstrations, then combines them with the few-shot demos to train a policy. Results on sim and real environments show that STRAP outperforms other retrieval methods and pure behavioral cloning. Ablations show that STRAP is compatible with various vision encoders and justify each of its component.", "questions": "- Baseline: how does STRAP compare to the pretrain-then-finetune setup? (Pretrain on the prior dataset, then fine-tune on the few-shot target demonstrations?)\n- Baseline: how does STRAP compare to a multitask policy trained on all available data?\n- Generalization: to what extent does STRAP generalize across environments or embodiments?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Sub-trajectory retrieval for the few-shot demo behavioral cloning setting is a well-motivated and novel idea.\n- The method is clear and straightforward to implement.\n- Results show that matching with DTW on vision foundation model features are robust to variations and capture task semantics.\n- Real and simulated environments show that STRAP outperforms other retrieval methods and pure behavioral cloning.", "weaknesses": "- STRAP requires few-shot demos and model training at test time for a new task.\n- It would be good to see more sim and real environments for evaluations.\n- It would be more convincing to see a behavioral cloning baseline that uses all available data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GM4g10qfvA", "reviewer": "Reviewer_RToH", "summary": "In this paper, the authors propose a task-specific robot learning framework using pre-collected datasets. Unlike the many robot learning methods that train the generalist policy model with multi-task expert data, the proposed method (STRAP) train a task-specific policies, which can yields better performance on single task. When the few-shot target demo, in addition to prior dataset is given, STRAP filters the task-relevant data from prior data and use it with target demo to train the model. One of the key features of STRAP is that it retrieves the data measuring the similarity between sub-trajectories, rather than the whole trajectories. Also, it utilize subsequence dynamic time warping (S-DTW) to match between the data. As a result, the proposed method shows improved performance compared to the previous methods, generalist policy models, and specialist models that only use the target data.", "questions": "- Compared to the prior framework (Behavior Retrieval, Du et al., 2023), STRAP seems to have three main differences in retrieval system. (a) use non-parametric retrieve vs. VAE (b) use sub-trajectory wise retrieve vs. single state-action pairs and (c) use DTW. Among these, what gives the most / least performance gains?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- To deal with potentially variable length during retrieval, STRAP use dynamic time warping (DTW) to match the sub-sequences\n- STRAP shows improved performance compared to the prior framework (Behavior Retrieval, Du et al., 2023) which retrieves single state-action pairs using VAE.", "weaknesses": "- The idea of using only data relevant to the target task, rather than learning a generalist policy through multi-task data, is interesting. However, retrieving new data from prior dataset and training a policy each time a new scene is encountered is highly computationally costly.\n- Comparing the entire prior data with the target data one-to-one to measure similarity is not scalable with the dataset size. Moreover, since this retrieval process requires computationally intensive neural network operations, such as DINO, it raises questions about whether this process can be performed at test time. In particular, there is no mention of how to handle an increase in offline dataset size, nor are there any discussions about limitations in this regard.\n- There is no discussion about computational cost.\n- STRAP uses a top-k retrieval dataset. Increasing this k could bring in more data but might reduce relevance, whereas a smaller k would provide more refined data but with a smaller amount. However, there is a lack of analysis on how changing this k value affects performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dSumYEJ17W", "reviewer": "Reviewer_HWGB", "summary": "## Paper Review Summary\n\nThis work advocates for training policies dynamically during deployment, utilizing the encountered scenarios to improve model performance. Instead of relying on pre-trained policies to tackle new problems in a zero-shot fashion, the authors propose a non-parametric approach that retrieves relevant data and trains models directly at test time. The paper introduces SRTAP, a method built on pre-trained Vision-Language Models (VLM) and dynamic time wrapping, which combines sub-trajectories into a policy. The approach involves some training on a set similar in language to the test set, and has been demonstrated in both real-world and simulated environments.\n\n### Strengths:\n1. **Innovative Approach**: The authors present a compelling intuition, demonstrating robustness in solving multitask generalization challenges.\n2. **Efficient Data Usage**: The method shows improvement in the way data is leveraged for robotics tasks, particularly in sub-trajectory retrieval.\n3. **Thorough Experiments**: The experiments are detailed and show promising results for sub-trajectory retrieval and policy creation.\n\n### Weaknesses:\n1. **Project Incompleteness**: There is no accessible website or supplementary information, suggesting the project might still be unfinished.\n2. **Visual Readability**: The images in the paper are difficult to interpret, potentially detracting from the clarity of the results.\n3. **Writing Quality**: The paper's writing needs improvement, especially in terms of clarity and readability.\n4. **Generalization**: Some imitation learning methods, such as [Sparse Diffusion Policy](https://forrest-110.github.io/sparse_diffusion_policy/), [HPT](https://liruiw.github.io/hpt/), [RDT-Robotics](https://rdt-robotics.github.io/rdt-robotics), and [Humanoid Manipulation](https://humanoid-manipulation.github.io/), appear to show more generalization in similar settings.\n5. **Formatting**: The paper's formatting is problematic, with certain sections being hard to read, affecting the overall readability of the work.\n\nIn conclusion, while the proposed method shows strong intuition and detailed experimentation, there are concerns about project completeness, readability, and potential improvements in both writing and generalization when compared to existing work. \n\nI would like to change the rate after discussion, but at least you should finish the site you provide.", "questions": "1. Would you please update the sites in the paper?\n2. Could you please add some experiments with other imitation learning methods?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "1. **Innovative Approach**: The authors present a compelling intuition, demonstrating robustness in solving multitask generalization challenges.\n2. **Efficient Data Usage**: The method shows improvement in the way data is leveraged for robotics tasks, particularly in sub-trajectory retrieval.\n3. **Thorough Experiments**: The experiments are detailed and show promising results for sub-trajectory retrieval and policy creation.", "weaknesses": "1. **Project Incompleteness**: There is no accessible website or supplementary information, suggesting the project might still be unfinished.\n2. **Visual Readability**: The images in the paper are difficult to interpret, potentially detracting from the clarity of the results.\n3. **Writing Quality**: The paper's writing needs improvement, especially in terms of clarity and readability.\n4. **Generalization**: Some imitation learning methods, such as [Sparse Diffusion Policy](https://forrest-110.github.io/sparse_diffusion_policy/), [HPT](https://liruiw.github.io/hpt/), [RDT-Robotics](https://rdt-robotics.github.io/rdt-robotics), and [Humanoid Manipulation](https://humanoid-manipulation.github.io/), appear to show more generalization in similar settings.\n5. **Formatting**: The paper's formatting is problematic, with certain sections being hard to read, affecting the overall readability of the work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)