STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning

Dear Reviewer,

Thank you for your feedback and for acknowledging our “compelling intuition” and “thorough experiments”.

Project incompleteness

We’ve updated the link to the website: https://strapaper.github.io/strap.github.io/ Furthermore, we’ve added more real-world experiments and improved baselines to the main paper and extended the appendix with experimental details and discussion on the computational complexity and choice of hyperparameters.

Visual readability and formatting

We acknowledge that some of the sections lacked proper formatting. We’ve improved the formatting by increasing the margins for figures, equations, and tables. Furthermore, we’ve updated [Figure 2] to be more interpretable and added pointers from the visual stages to the different sections in the paper. Could you please comment on whether these changes are sufficient and clarify what other changes to the writing quality and visual readability you would like to see?

Generalization to other imitation learning methods

We’ve run additional experiments using alternative architectures and imitation learning methods. Unfortunately, we found simpler architectures (MLP+GMM, LSTM+GMM) to struggle with achieving non-zero success rates on the LIBERO benchmark or limited by their capability to be conditioned on language (Diffusion Policies). While vision-language-action models experience much better language-conditioning, they require significantly more fine-tuning data than available in our setting [3]. Therefore, we chose transformer-based policies since they observe much better language-conditioning and multi-task capabilities as shown in [1,2].

[1] Libero: Benchmarking knowledge transfer for lifelong robot learning. Liu et al, 2024

[2] BAKU: An Efficient Transformer for Multi-Task Policy Learning. Haldar et al. 2024

[3] OpenVLA: An Open-Source Vision-Language-Action Model. Kim et al. 2024

We hope these additional modifications have addressed your previous questions. Please don’t hesitate to let us know if you have any additional comments or questions.

Choice of hyperparameter $K$

We extensively ablate STRAP for $K \in (100, 200, 500, 1000, 2000, 4000)$ on all LIBERO-10 tasks. We find tuning $K$ to improve our reported success rates ($K=100$) on 8/10 tasks by an average of $7.5$%. The results suggest that the optimal value for $K$ is task-dependent with some tasks benefiting from retrieving less (4/10) and some from retrieving more (4/10) data. We hypothesize that the optimal $K$ depends on whether tasks benefit from broader less relevant or more refined data. Thank you for this suggestion! We update our results in [Tables 1 & 3] with the improved results. You can find additional details and visualizations in [Table 8] and [Figure 29] in the Appendix.

The choice of $K$ has no impact on the computational complexity of STRAP as we always compute all matches between D_target and D_prior before selecting the top K. This also means that the hyperparameter search on $K$ can be reduced to policy learning and evaluation by storing the matches and training on $K$ segments.

Baselines ablations

a) We compare non-parametric (frozen DINOv2, denoted as D-S) and parametric (training a VAE as in Behavior Retrieval, denoted as BR) embeddings in a state-based retrieval setting. [Table 1] shows BR outperforming DINOv2 embeddings by $+5$% on average across all LIBERO-10 tasks. Inspecting the individual success rates in [Tables 1 & 3] it becomes clear that the optimal embedding is highly task-dependent. However, we expect future representation learning techniques to close this gap. DINOv2 only encodes a single observation and, therefore, does not have a notion of dynamics and task semantics compared to BR.

b) & c) We compare retrieving states, sub-trajectories, and full trajectory retrieval in [Table 1]. While state-based retrieval is based on cosine similarity, the variable length of trajectories requires S-DTW for sub-trajectory and full trajectory retrieval. We find retrieving sub-trajectories outperforms state-based retrieval by $+17.2$ and full trajectory retrieval by $+4.2$ on average across all LIBERO-10 tasks.

To summarize, the largest gains come from retrieving (sub-)trajectories enabled by S-DTW. While the choice of embedding is largely task-dependent, STRAP allows for using off-the-shelf models by encoding the dynamics and semantics of the trajectories in the sub-trajectory retrieval process instead of the embedding model.

We hope these additional modifications have addressed your previous questions. Please don’t hesitate to let us know if you have any additional comments or questions.

Dear Reviewer,

Thank you for your feedback on how we can improve our paper. We provide additional discussions on the computational complexity of our retrieval algorithm, an extensive ablation of hyperparameter $K$, and a detailed comparison to recent retrieval systems.

Computational cost to adapt to new scenes

Current generalist policies are unable to adapt zero-shot to the target task. A common approach is to fine-tune a pre-trained policy on a few target demonstrations (FT). While recent generalist policies suggest fine-tuning on 10-150 demonstrations [1] we only have access to 3-5 demonstrations in our few-shot setting. While the pre-training can provide faster adaptation, fine-tuning on a small number of demonstrations is unable to cover the randomization the policy is exposed to during evaluation (cf. added FT baselines in [Table 1,2 & 3]). By retrieving relevant data, we expose our expert to a much larger training distribution making it more robust to the deployment scenario.

STRAP and a standard fine-tuning approach differ in the retrieval process and longer policy training. Our discussion in [A.5] shows that the retrieval scales linearly with the number of (sub-)trajectories in D_prior and D_target. Greater speedups can be achieved by parallelizing the retrieval and utilizing GPUs to compute the DTW distance matrix. Overall, STRAP’s runtime is longer compared to fine-tuning (~5min) consisting of retrieval (~5min on a DROID-scale dataset) and training a policy (~30min) but provides significant robustness benefits shown by an average performance boost of $+6.4$% and $+25.7$% across all LIBERO-10 and real-world Kitchen tasks, respectively.

[1] OpenVLA: An Open-Source Vision-Language-Action Model. Kim et al. 2024

Computational cost of the retrieval process

To avoid excessive compute during test time, we precompute the embeddings for D_prior! Encoding a single image running DINOv2 on an NVIDIA L40S 46GB and batch size 32 takes $2.83ms\pm 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. Every dataset only has to be encoded once when added to D_prior and can be reused for all future deployment scenarios. In contrast to previous methods (cf. BehaviorRetrieval, FlowRetrieval), using an off-the-shelf vision foundation model also eliminates the need to re-train the encoder and re-encode the entire dataset when D_prior grows.

Retrieving data with S-DTW scales linearly with the size of D_prior, 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 D_prior and only depends on the amount of retrieved data $K$, making it 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_prior. We thoroughly discuss the complexity of our retrieval algorithm and benchmark the runtime of S-DTW in [A.5]. Note that there are several possible ways to improve our implementation, e.g., leveraging GPU deployment and custom CUDA kernels to compute the distance matrix or parallelizing retrieval across trajectories.

Computational cost, memory, and scalability to real-world datasets

STRAP’s retrieval stage is based on S-DTW which consists of two stages: computing the distance matrix $D$ and finding the shortest path via dynamic programming. These stages have to be run sequentially for each sub-trajectory in D_target and trajectory in D_prior but don’t depend on the other (sub-)trajectories. Therefore, STRAP has a runtime complexity of $\mathcal{O}(N*M)$ with N the number of sub-trajectories in D_target and M the number of trajectories in D_prior.

We thoroughly benchmark the runtime of S-DTW in [A.5]. For an offline dataset the size of DROID (76k), retrieval takes approximately $300sec$ using our unoptimized research code. Note that there are several possible ways to improve our implementation, e.g., leveraging GPU deployment and custom CUDA kernels to compute the distance matrix or parallelizing retrieval across trajectories.

Since we encode D_prior using vision foundation models, the memory footprint of STRAP stays fairly low. Loading the embeddings to compute the cost matrix represents the largest memory consumption but is much lower than the memory consumed by loading image or video sequences.

For runtime visualization and information regarding data encoding and policy learning complexity please refer to [A.5] in the appendix.

Choice of hyperparameter $K$

Thank you for this suggestion! We extensively ablate STRAP for $K \in (100, 200, 500, 1000, 2000, 4000)$ on all LIBERO-10 tasks. We find tuning $K$ to improve our previously reported success rates ($K=100$) on 8/10 tasks by an average of $7.5$%. The results suggest that the optimal value for $K$ is task-dependent with some tasks benefiting from retrieving less (4/10) and some from retrieving more (4/10) data. We hypothesize that the optimal $K$ depends on whether tasks leverage (positive transfer) or suffer (negative transfer) from multi-task training. We update our results in [Tables 1 & 3] with the improved results. You can find additional details and visualizations in [Table 8] and [Figure 29] in the Appendix.

The choice of $K$ has no impact on the computational complexity of STRAP as we always compute all matches between D_target and D_prior before selecting the top $K$. This also means that the hyperparameter search on $K$ can be reduced to policy learning and evaluation by storing the matches and retrieving $K$ segments for each iteration.

Cross-embodiment generalization

We leave cross-embodiment retrieval to future work but emphasize that the embedding choice allows for some steerability of the retrieved sub-trajectories. For instance, averaging the embeddings of all three cameras leads to retrieving very similar scenes [see “Does retrieval scale to DROID?” on our website] while using embeddings of only the in-hand camera focuses much more on the manipulated object and can be embodiment-agnostic. While the retrieval might scale to multiple embodiments, training cross-embodied policies is much more challenging, and leveraging cross-embodied data is an open research problem.

We hope these additional modifications have addressed your previous questions. Please don’t hesitate to let us know if you have any additional comments or questions.

Dear Reviewer,

Thank you for your feedback on how we can improve our paper. We provide additional results on three real-world tasks, improved baselines, an extensive ablation of hyperparameter $K$, and a discussion on our retrieval algorithm's computational and memory complexity.

Additional real-world tasks

We’ve added three real-world tasks (“pick and place the [carrot, pepper, chili]”) in three realistic kitchen environments ([table, sink, stove]) and collected 50 multi-task demonstrations in every scene. Object poses are randomized in a $20\times20cm$ grid during data collection and evaluation. D_target is specific to each environment and contains 3 demonstrations of the downstream task.

We find STRAP to experience surprising generalization behavior from the 3 poses seen in D_target to the poses in the $20\times20cm$ grid exposed during evaluation (Kitchen). The policy further shows recovery behavior, completing the task even when the initial grasp fails and alters the object pose [cf. robot attempting to pick up the red pepper in the first video on our website]. As shown in [Table 2] STRAP even maintains its high performance, retrieving the relevant data for much larger offline datasets (Kitchen+DROID).

Meanwhile, the large pose randomizations in D_prior and the evaluation are challenging for the baselines. Behavior cloning and fine-tuning (BC, FT) fit the target demonstrations failing to adapt to unseen object poses. The multi-task policy (MT) replays trajectories from the offline dataset instead of solving the prompted task most likely caused by an imbalanced training dataset. Increasing the dataset size amplifies these challenges by making it significantly harder to pre-train or learn a multi-task policy, leading to performance drops of $-6.09$% (FT) and $-26.3$% (MT).

Stronger baselines

We extend our evaluations by including two additional baselines for all LIBERO-10 tasks and the real-world tasks introduced above.

1. Pre-training + fine-tuning (FT) which represents a policy pre-trained on D_prior and fine-tuned on the few available demonstrations in D_target.
2. Multi-task policy (MT) trained on D_prior $\cup$ D_target in contrast to only D_prior in our original submission.

The fine-tuning baseline is more competitive than standard behavior cloning but still falls short by $-6.4$% and $-25.7$% compared to STRAP on the LIBERO-10 and real-world tasks. The updated multi-task policy performs well on some tasks where we expect positive transfer, i.e., where environments and tasks in D_prior and D_target overlap, e.g., LIBERO-10 and LIBERO-90 both contain the Book-Caddy task. We’ve added the FT baseline and replaced MT trained only on D_prior with MT trained on D_prior $\cup$ D_target in [Tables 1,2 & 3].

Few-shot demos and model training at test time

Current generalist policies are unable to adapt zero-shot to the target task. A common approach is to fine-tune a pre-trained policy on few target demonstrations (FT). While recent generalist policies suggest fine-tuning on 10-150 demonstrations [1] we only have access to 3-5 demonstrations in our few-shot setting. While the pre-training can provide faster adaptation, fine-tuning on 3 demonstrations is unable to cover the $20\times20cm$ grid of possible object poses the policy is exposed during evaluation (cf. added FT baselines in [Table 1,2 & 3]). By retrieving relevant data, we expose our expert to a much larger training distribution making it more robust to the deployment scenario.

STRAP and a standard fine-tuning approach differ in the retrieval process and longer policy training. Our discussion in [A.5] shows that the retrieval scales linearly with the number of (sub-)trajectories in D_prior and D_target. Greater speedups can be achieved by parallelizing the retrieval and utilizing GPUs to compute the DTW distance matrix. Overall, STRAP’s runtime is longer compared to fine-tuning (~5min) consisting of retrieval (~5min on a DROID-scale dataset) and training a policy from scratch (~30min) but provides significant robustness benefits shown by an average performance boost of $+6.4$% and $+25.7$% across all LIBERO-10 and real-world Kitchen tasks, respectively.

[1] OpenVLA: An Open-Source Vision-Language-Action Model. Kim et al. 2024

Cross-environment and -embodiment generalization

We provide examples of sub-trajectories retrieved from the DROID dataset on our website [see “Does retrieval scale to DROID?”]. STRAP retrieves trajectories collected in environments with similar appearance, e.g., camera pose, table orientation, and texture, and similar tasks, e.g., picking up cylindrical objects. The choice of embedding allows for further steerability of the retrieved sub-trajectories. For instance, averaging the embeddings of all three cameras leads to retrieval of very similar scenes, while using embeddings of only the in-hand camera focuses much more on the manipulated object and can be embodiment-agnostic. While the retrieval could scale to multiple embodiments, training cross-embodied policies is much more challenging, and leveraging cross-embodied data is an open research problem that we leave to future work.

We hope these additional modifications have addressed your previous questions. Please don’t hesitate to let us know if you have any additional comments or questions.

Dear Reviewer,

Thank you for your feedback on how we can improve our paper. We provide additional results on improved baselines, discussions on the runtime differences between STRAP and a fine-tuning approach, and generalization to new environments and embodiment.

Additional real-world tasks

We’ve added three real-world tasks (“pick and place the [carrot, pepper, chili]”) in three realistic kitchen environments ([table, sink, stove]) and collected 50 multi-task demonstrations in every scene. Object poses are randomized in a $20\times20cm$ grid during data collection and evaluation. D_target is specific to each environment and contains 3 demonstrations of the downstream task.

We find STRAP to experience surprising generalization behavior from the 3 poses seen in D_target to the poses in the $20\times20cm$ grid exposed during evaluation (Kitchen). The policy further shows recovery behavior, completing the task even when the initial grasp fails and alters the object pose [see the robot attempting to pick up the red pepper in the first video on our website]. As shown in [Table 2] STRAP even maintains its high performance, retrieving the relevant data for much larger offline datasets (Kitchen+DROID).

Meanwhile, the large pose randomizations in D_prior and the evaluation are challenging for the baselines. Behavior cloning and fine-tuning (BC, FT) fit the target demonstrations failing to adapt to unseen object poses. The multi-task policy (MT) replays trajectories from the offline dataset instead of solving the prompted task most likely caused by an imbalanced training dataset. Increasing the dataset size amplifies these challenges by making it significantly harder to pre-train or learn a multi-task policy, leading to performance drops of $-6.09$% (FT) and $-26.3$% (MT).

Stronger baselines

We extend our evaluations by including two additional baselines for all LIBERO-10 tasks and the real-world tasks introduced above.

1. Pre-training + fine-tuning (FT) which represents a policy pre-trained on D_prior and fine-tuned on the few available demonstrations in D_target.
2. Multi-task policy (MT) trained on D_prior $\cup$ D_target in contrast to only D_prior in our original submission.

The fine-tuning baseline is more competitive than standard behavior cloning but still falls short by $-6.4$% and $-25.7$% compared to STRAP on the LIBERO-10 and real-world tasks. The updated multi-task policy performs well on some tasks where we expect positive transfer, i.e., where environments and tasks in D_prior and D_target overlap, e.g., LIBERO-10 and LIBERO-90 both contain the Book-Caddy task. We’ve added the FT baseline and replaced MT trained only on D_prior with MT trained on D_prior $\cup$ D_target in [Tables 1,2 & 3].

We hope our responses have adequately addressed your concerns! If so, we kindly ask you to increase your score recommendation.