STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning

• Task-specific data collection is costly, while multi-task generalist policies can underperform specialists because of negative transfer.Zero-shot generalization to unseen tasks is challenging, and collecting large amounts of in-domain fine-tuning data can also be expensive.Source: §1 · p. 2
• STRAP improves robustness and generalization over prior retrieval methods, zero-shot multi-task policies, and policies trained only on test-time data.Source: §1 · p. 2
• STRAP uses a small amount of test-time in-domain data to retrieve relevant sub-trajectories from a large multi-task corpus for training.The retrieved data is intended to guide models toward desired behavior.Source: §1 · p. 2
• Foundation-model features and dynamic time warping support matching across visual differences and variable-length sequences, while automatic extraction avoids manually segmenting the training corpus.Source: §1 · p. 2

Retrieval for Behavior Replay

• Behavior replay methods retrieve past demonstrations using varied embeddings, but assume offline data resembles test demonstrations or that intermediate representations bridge the gap.The authors say these assumptions limit use of large multi-task datasets collected across domains.Source: §2 · p. 3

Retrieval for Few-shot Imitation Learning

• BehaviorRetrieval and FlowRetrieval learn custom representations and retrieve individual states, whereas STRAP retrieves sub-trajectories to share data across tasks while retaining temporal information.The authors characterize custom representations as difficult to scale with growing offline datasets and significant visual or semantic differences.Source: §2 · p. 3

Learning from Sub-trajectories

• STRAP presents robot sub-trajectory retrieval as a mechanism for partitioning large offline datasets and enabling cross-task positive transfer during policy learning.Source: §2 · p. 3

3.1 DYNAMIC TIME WARPING

• DTW uses dynamic programming to find a minimum-cost alignment between sequences while satisfying boundary and continuity conditions.Its cumulative cost matrix adds each point-pair cost to the minimum of three neighboring cumulative costs.Source: §3.1 · pp. 3–4
• STRAP encodes target demonstrations and offline data in a shared vision-feature space, slices sub-trajectories, matches them with S-DTW, and trains on retrieved and target data.Source: Figure 2 · p. 4
• S-DTW matches a shorter query to a potentially different-length subsequence of a longer reference, allowing alignment to start and end within that reference.Backtracking from the minimum value in the last row identifies the best-matching subsequence.Source: §3.1 · p. 4
• Figure 3 depicts target sub-trajectories matched to relevant segments in prior trajectories, without requiring their starts and ends to align.Source: Figure 3 · p. 5

4.1 PROBLEM SETTING: RETRIEVAL-AUGMENTED POLICY LEARNING

• STRAP retrieves a relevant subset of offline data and combines it with few-shot target demonstrations to train a language-conditioned policy for the target environment.The offline dataset is assumed to contain expert trajectories from a matching embodiment, but its tasks and scenes may differ from the target.Source: §4.1 · p. 5

4.2 SUB-TRAJECTORIES FOR DATA RETRIEVAL

• Variable-length sub-trajectories preserve temporal dynamics while allowing behavior shared across different tasks to be retrieved.STRAP segments target demonstrations, not the much larger offline dataset; its proposed heuristic marks transitions when end-effector velocity falls below a threshold.Source: §4.2 · p. 6

4.3 FOUNDATION MODEL-DRIVEN RELEVANCE METRICS FOR RETRIEVAL

• Vision foundation-model embeddings provide semantic and visual similarity robust to low-level scene variation, while DTW-based matching retains sub-trajectory dynamics.The method uses pairwise embedding costs rather than temporal averaging, which would lose actions and dynamics.Source: §4.3 · pp. 6–7

4.4 EFFICIENT SUB-TRAJECTORY RETRIEVAL WITH SUBSEQUENCE DYNAMIC TIME WARPING

• Subsequence dynamic time warping matches variable-length target chunks to aligned segments in the offline data, and the lowest-cost matches are selected uniformly across target chunks.Repeated matches may be retained, treating recurring chunks as important by up-weighting them in training.Source: §4.4 · p. 7

4.5 PUTTING IT ALL TOGETHER: STRAP

• STRAP trains on the union of target and retrieved data, using language conditioning to distinguish misleading retrieved actions from useful cross-task examples.The retrieval procedure embeds observations, chunks target demonstrations, selects top matches, and limits S-DTW backtracking step size to [1, 2].Source: §4.5 · pp. 7–8

5.1 EXPERIMENTAL SETUP

• Evaluation covers 10 LIBERO tasks and real-world vegetable pick-and-place in three kitchen environments, with an additional 5000-trajectory DROID prior dataset.LIBERO uses five target demonstrations and 4500 prior trajectories; the real-world settings include three target demonstrations per environment.Source: §5.1 · p. 8, Table 2 · p. 9

5.2 EXPERIMENTAL EVALUATION

• STRAP outperforms BR and FR by +24.7% and +25.0% across all 10 LIBERO tasks, respectively.BC is strong on LIBERO but performs much worse in real-world experiments with increased evaluation randomization.Source: §5.2 · p. 9, Table 1 · p. 9
• Fine-tuning is the most competitive baseline, trailing STRAP by −6.4%, while multi-task performance improves when prior data overlaps the target setting.Greater randomization challenges both methods; multi-task policies sometimes ignore language instructions, which the authors hypothesize may reflect data imbalance.Source: §5.2 · p. 9
• STRAP generalizes to real-world poses unseen in target demonstrations and shows recovery when an initial grasp fails and changes the object’s pose.Its policy-training stage is independent of prior-dataset size, allowing larger datasets such as DROID to be added while maintaining performance.Source: §5.2 · p. 10
• Sub-trajectory retrieval improves performance by +4.1% across all 10 LIBERO tasks compared with full-trajectory retrieval.The authors hypothesize that irrelevant segments in full trajectories can hurt performance and cumulative-cost accuracy; the best retrieval count K varies by task.Source: §5.2 · p. 10
• DINOv2 and CLIP differ by only +0.7% across all 10 tasks, while no representation method outperforms the others on every task.In state-based retrieval, BR and FR outperform D-S by +5.0% and +4.7%, respectively.Source: Table 1 · p. 9, §5.2 · p. 10
• STRAP retrieves data from five of 90 semantically relevant tasks, and each retrieved task shares at least one sub-task with the target.The task-distribution figure shows retrieval proportions by method; the top five tasks are shown and the remainder grouped as “others.”Source: §5.2 · p. 10, Figure 13 · p. 10

6 CONCLUSION

• STRAP combines visual foundation models with S-DTW sub-trajectory retrieval, avoiding training on the entire retrieval dataset while maintaining minimal compute overhead.The authors report improved data utilization and more effective capture of dynamics, along with strong performance in challenging real-world scenarios.Source: §6 · p. 11

A.1 SIMULATION EXPERIMENTS

• Table 3 presents different methods’ performance on LIBERO-10 simulation tasks, marking the best results in bold and runner-up results with underlining.State-based baselines retrieve single-state action pairs and pad them with neighboring states for transformer-policy compatibility.Source: Table 3 · p. 14, §A.1 · p. 14
• Results are reported over three training and evaluation seeds, with 100 retrieved segments and five expert demonstrations.Retrieval uses agent-view images; policies train on agent-view and in-hand observations for 300 epochs.Source: §A.1 · p. 14
• Both FR and BR outperform STRAP on Cream-Cheese-Butter; the authors hypothesize overly long chunks mix semantic tasks and worsen retrieval matches.They suggest this may decrease downstream performance.Source: §A.1 · p. 14

A.2 REAL-WORLD EXPERIMENTS

• Retrieval averages embeddings across left, right, and in-hand camera observations and retrieves 100 segments.Policies train on all three image observations for 200 epochs, using batch size 32 and an epoch every 100 gradient steps.Source: §A.2 · p. 14
• The policy’s ResNet-18 vision encoders are initialized with ImageNet-pretrained weights.Source: §A.2 · p. 14
• Demonstrations are collected and policies deployed on the DROID setup, which runs at 15 Hz.Source: §A.2 · p. 14

A.2.1 FRANKA-PEN-IN-CUP

• The target pen-in-cup task is absent from the 100 prior demonstrations across 10 tasks, while the target dataset contains three demonstrations.The experiment tests retrieval for an unseen pose and object and target appearance.Source: §A.2.1 · p. 14
• Prior tasks include potentially shared subtasks, such as placing a pen beside markers, and unrelated tasks, such as making a hotdog or pulling a cable.The appendix includes figures labeled chess, cube stacking, hotdog, and knocking over a box.Source: §A.2.1 · p. 14, Figure 14 · p. 15, Figure 15 · p. 15
• Table 4 presents Franka-Pen-in-Cup real-world results, while Table 5 lists task and language instructions for the prior dataset.Source: Table 4 · p. 15, Table 5 · p. 15

A.2.2 DROID-KITCHEN

• Figure 28 marks task-relevant objects in the table, sink, and stove environments and notes randomized object pose and type during data collection.The listed objects include a carrot, pepper, and chili, respectively.Source: Figure 28 · p. 16
• The table environment includes a carrot, chips bag, and can; the sink has pepper, soap dispenser, and sponge; the stove has chili, chicken, and utensil.Source: Figure 28 · p. 16
• Table 6 lists DROID-Kitchen task and language instructions; each prior-data task combines two randomly sampled instructions joined by “then.”Source: Table 6 · p. 16

A.3 ABLATIONS

• Tuning K improves reported success rates on 8/10 tasks, while retrieval ablations compare full trajectories with 30-step sliding-window sub-trajectories.The optimal K varies by task; the retrieval-seed study uses 10 expert demonstrations and evaluates combinations of three training and three retrieval seeds.Source: Table 7 · p. 16, Figure 29 · p. 17, Table 8 · p. 18

A.4 ADDITIONAL BASELINES

• The diffusion-policy baseline experiments evaluate BehaviorRetrieval and FlowRetrieval on LIBERO-10 without language conditioning, replicating their training setup.Source: Table 10 · p. 18

A.5.1 ENCODING TRAJECTORIES WITH VISION FOUNDATION MODELS

• Encoding one image takes 2.83 ms ± 0.08 on average, and encoding DROID’s 18.9M single-view timesteps takes 26 hours.Embedding scales as O(T * V), and the embeddings can be precomputed and reused for later retrieval.Source: §A.5.1 · p. 18

A.5.2 SUBSEQUENCE DYNAMIC TIME WARPING

• For an offline dataset the size of DROID (76k), S-DTW retrieval takes approximately 300 seconds; distance-matrix computation scales as O(n * m * E).The distance-matrix computation can be expressed as matrix multiplications and may use GPU deployment or custom CUDA kernels.Source: §A.5.2 · p. 19

A.5.3 POLICY TRAINING

• Training a single policy takes 35 minutes ± 4 on average across 10 trials.The reported setup trains for 200 epochs, with 100 gradient steps per epoch and batch size 32.Source: §A.5.3 · p. 19

A.5.4 INCREASING THE DATASET SIZE

• Retrieval scales linearly with Dprior and completes within 5 minutes even on the largest available datasets, such as DROID.Policy learning depends on retrieved data K rather than the size of Dprior.Source: Figure 30 · p. 19, §A.5.4 · p. 19

A.6 QUALITATIVE RESULTS

• For the black-bowl-to-bottom-drawer task, S-DTW finds matches regardless of start and end points or trajectory length, producing varied retrieved segment lengths.The match distribution also varies in its start and end points.Source: Figure 31 · p. 20