FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning
Authors: Li-Heng Lin
Yuchen Cui
Amber Xie
Tianyu Hua
Dorsa Sadigh
Venue: CoRL 2024 (2024)
Keywords: Data Retrieval
Few-shot Learning
Imitation Learning
TL;DR: We propose a few-shot imitation learning method that uses optical flow to retrieve data with similar motion from prior dataset and augment downstream policy learning.

Abstract: Imitation learning policies in robotics tend to require an extensive amount of demonstrations. It is critical to develop few-shot adaptation strategies that rely only on a small amount of task-specific human demonstrations. Prior works focus on learning general policies from large scale dataset with diverse behaviors. Recent research has shown that directly retrieving relevant past experiences to augment policy learning has great promise in few-shot settings. However, existing data retrieval methods fall under two extremes: they either rely on the existence of exact same behaviors with visually similar scenes in the prior data, which is impractical to assume; or they retrieve based on semantic similarity of high-level language descriptions of the task, which might not be that informative about the shared behaviors or motions across tasks. In this work, we investigate how we can leverage motion similarity in the vast amount of cross-task data to improve few-shot imitation learning of the target task. Our key insight is that motion-similar data carry rich information about the effects of actions and object interactions that can be leveraged during few-shot adaptation. We propose FlowRetrieval, an approach that leverages optical flow representations for both extracting similar motions to target tasks from prior data, and for guiding learning of a policy that can maximally benefit from such data. Our results show FlowRetrieval significantly outperforms prior methods across simulated and real-world domains, achieving on average 27% higher success rate than the best retrieval-based prior method. In the Pen-in-Cup task with a real Franka Emika robot, FlowRetrieval achieves 3.7x the performance of the baseline learning from all prior and target data.

FLOWRETRIEVAL: Flow-Guided Data Retrieval for Few-Shot Imitation Learning
Li-Heng Lin', Yuchen Cui', Amber Xie*, Tianyu Hual, Dorsa Sadigh t*
Few-shot imitation learning relies on only a small amount of task- specific demonstrations to efficiently adapt a policy for given downstream tasks. Retrieval-bass methods come with promise of retrieving relevant past experi- ences augment this target data when learning policies. However, existing data retrieval methods fall under two extremes: they either rely on the existence of exact behaviors with visually similar scenes in the prior data, which is impractical to assume; they retrieve based on semantic similarity of high-level language descriptions of the task, which might not be that informative about the shared low-level behaviors or motions across tasks that is often more important factor for retrieving relevant data for policy learning. In this work, we investigate how we can leverage motion similarity in the vast amount of cross-task data to improve few-shot imitation learning of the target task. Our key insight that notion-similar data car- ries rich information about the effects of actions and object interactions that can be leveraged during few-shot adaptation. We propose FLOW RETRIEVAL, approach that leverages optical flow representations for both extracting similar motions to target tasks from prior data, and for guiding learning of policy that can maxi- mally benefit from such data. Our results show FLOW RETRIEVAL significantly outperforms prior methods across simulated and real- world domains, achieving on average 27% higher success rate than the best retrieval-bart prior method. In the Pen-in-Cup task with real Franka Emika robot, FLOW RETRIEVAL achieves 3.7x the performance of the baseline imitation learning techniques that learn from all prior and target data. Website: https: htttss:/fllowrettie http///lwowrrrreaalllll ..l..giito github.is
Abstract:
Introduction
Imitation learning has shown significant promise in robotics with the rise of deep end-to-end models, but a considerable challenge is overcoming its high data demands. Typically, substantial number of demonstrations, often ranging from hundreds to thousands depending on the complexity and dexterity of the task, are needed to achieve a policy with reliable performance across the expected range of scenarios. This reliance on extensive amounts of data makes it impractical for the policy to adapt quickly to slightly different or out-of aut-of-iistriuttion environments tasks. Therefore, is crucial to develop methods that facilitate rapid adaptation with small amounts of demonstrations.
One promising approach is to retrieve relevant data from collected experiences, which enables more efficient policy learning for given downstream target task by incorporating valuable information from this prior data. Prior works identify reusable skills from pretraining data by training latent representations that capture state-action similarity between the target and prior data [1, 2]. However, these methods often assume that the exact visual observation present in the target task exists in the robot's prior experiences, and thus primarily retrieve based on such visual similarities. practice, this assumption rarely met in real-world applications. For instance, the task of "rotating a doorknob" appears to be visually different from "rotating a faucet". Nevertheless, the movement involved in one task ought to be applicable and beneficial for the other, regardless of their visual disparities.
Our key insight is that prior datasets serve a broader purpose than merely retrieving the same skills of visually similar states. Target task data may in fact exhibit similarities to prior data in
Computer Science Department, Stanford University, Stanford, CA, USA.
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

Prior Date
Representation Pretraining
Targel Task Demo S,,
Auxiiirr
Co-trainingg Batck
Flow Prediction
INS
Flow Decoder
VAE
Optical Aool
Visual Encoder
line
Optial Flow
Imitation
Action
Motion Space
Latent
Policy Learning
Policy
Opk
Data Retrieval
learn latent motion space
2. retrieve from prior data
3. co-train with retrieved data and auxiliary loss
Figure I: Overview of FLOW VETRIEVAL. We first learn motion-centric latent space through training VAE for embedding optical flows; then retrieve prior data similar to target data by measuring pairwise distances in the learned latent space: and augment policy co-training with supervised prediction loss for optical flow.
terms of low-level motion, offering an opportunity for knowledge transfer of motions. For example, both turning doorknob turning faucet necessitate similar rotating behavior, despite task dissimilarity. We posit that retrieving the robot's past experiences with similar motion patterns to the downstream task are particularly valuable for learning. The idea of retrieving and using data with similar behaviors is not new. Recent works often leverage language descriptions of tasks for retrieving similar behaviors [3-5]; however, language often operates at much higher level of abstraction and cannot readily capture low-level motions such turns and twists that could in fact be beneficial and transferable across prior and target tasks. Building upon this insight, we introduce an approach for retrieving and learning from motion-similat data from prior datasets.
We illustrate an overview of FLOW RETRIEVLL in Fig. 1. Our approach leverages optical flow as representation for motion-similarity based retrieval. FLOW RETRIEVAL first acquires a motion-centric latent space by computing optical flow between the current frame and a future frame the robot's RGB visual observations, and employing a variational autoencoder (VAE) to embed the optical flow data, During the data retrieval process, we rank prior data points based on their distance target task data in the latent space and retrieve the closest ones. We then train the imitation policy network using an augmented dataset through co-training with the target-specific and retrieved data. During policy learning, FLOWRETRIES RETRIEVAL leverages auxiliary loss of predicting the optical flow as additional guidance for representation learming, encouraging the model to encode the image with enough details to predict optical flow alongside predicting the action.
We demonstrate VETRIVELL achieves higher performance than prior approaches both in simulation and real robot environments. We experiment with two simulation tasks: Square Assembly task with a mixture of similar and adversarial tasks as prior data [2], and pick-and-place Can task from the LIBERO benchmark, with random subset of the LIBERO tasks as prior data [6]. Our real robot experiments spans two settings with pick-and-place tasks as the target task: using Widow-X robot with the Bridge [7] dataset as the prior data, and ii) using a Franka Panda robot with either past data from other similar tasks on the same robot or sub-sampled DROID [8] dataset as the prior data. Our experiments demonstrate FLOWREETIEVAL. achieves an average of 14% higher success rate than the best baseline each domain (+ 10% in simulation, +1% in real). BOORETRIENT also achieves on average 27% higher success rate than the best prior retrieval method.
2 Problem Statement
We model our imitation learning problem via a Markov decision process (MDP), where an agent moves between states s E S by taking actions a E A. The transition between states is governed by the dynamics function P P:Sx + S. An underlying reward function the target task but is unobserved by the learning agent. policy π (aas) is the conditional distribution of actions at each state. We assume the human demonstrator's 's policy gg is near-optimal under the target task reward function. The goal of imitation learning is to find a policy itaas) that imitates
- defines achieving

---

expert behavior me with dataset Durget = ((s ( r)} containing demonstrations of the target task, where r denotes the reward of achieving the target task.
In our few-shot imitation setting, Duaget is assumed to contain only a few demonstrations not enough to be able to learn the task from scratch. We additionally assume access to prior dataset Dprior = {(s, a)la a ~ TT(als,.) }} that consist of data of the robot performing diverse tasks. The goal of data retrieval is to extract the most relevant subset of prior data (6% of Prrioe) augment Pasget and improve the learning of ft.
In this work, decouple the problem of few-shot imitation learning via data retrieval into two sub problems: 1) the data retrieval problem focuses constructing similarity function S to extract the desired data from prior dataset ( ("what to retrieve?") "), and 2) the policy learning problem is concerned with how to maximally utilize the retrieved data D retrieved for improving performance on the target task ( ("how to use retrieved data?"). Decoupling the two problems allows us to independently evaluate the performance of data retrieval irrespective of the choice of policy learning, and further evaluate policy learning performance under the assumption of having access to the most relevant and valuable retrieved data. At the core of the data retrieval problem is to define measure of similarity ((p, t) that compares data points (sp, ap) Dptior to (st, a1) Dungee. The top 8% of Ppior ranked by S, will be retrieved. With the retrieved data, the problem of policy learning aims to improve performance on target task by maximally capitalizing on the common patterns and regularities in the retrieved data. Given Dreuseved and Durget. the goal is to find policy λ that can outperform learning only from Prrget or even learning from naively combining Pprior and Praget-
3 Related Work
Learning from demonstrations in data-efficient way has been long-sought goal of robot learning. The full literature is beyond the scope of this work. Our work is closely related to the body of work that leverages prior data for few-shot imitation of visuomotor policy. Pretraining and finetuning is the most common scheme of few-shot imitation learning that leverages prior datasets [9-11]. However, this two-stage scheme often faces challenges in balancing the extent of freezing and tuning parameters to achieve optimal performance. Co-training has emerged as a popular approach by directly incorporating both prior and target data in the same training batch [1, 12, 13]-albeit requiring balancing act between prior and target data -achieving competent performance when target data is scarce. In this work, we employ the co-training scheme for training the downstream policy network. Additionally, the choice of prior data in the training batch can significantly affect policy performance [11, 14]. Our approach, FLOW RETRIVVAL carefully filters the prior dataset based on measuring motion similarity of the provided demonstrations from the target task.
Extracting telk-relevant data from prior datasets can be achieved with different similarity metrics- from high-level language descriptions to dense visual dynamics (Fig. 2). In the rest of this section, we examine these different types of motion guidance and discuss how they are used solve the two sub problems in our problem setting discussed in Section 2.
Guidance Type
Language
Feature Traces
Optical Flow
Visual Dynamics
PPick up the
Example
"Insert square pex"
Granularity
Coarse
Medium
Medium
Fine
Attributes
Label Source
Human Annotation
Otfche-Shell Vision Model Motion
OHf-th--Shilf Vision Model
Ground-Truth Video Testure Shape Motion
Information
Semantic
Shape Mation
Figure 2: Different types of motion guidance for imitation learning, Here, we present a spectrum of different types of motion guidance used in the literature for either retrieval dircctly learning the policy. This spectrum varies the level of granularity of motion starting from coarse semantic information via language from the left to guidance about feature traces or optical flow capturing motion and shape objects all the way to more fine-grained guidance such as visual dynamics that also capture the texture of the scene.
3

---

Language Descriptions. On one end of the spectrum in Fig. 2, we can retrieve base on high-level language descriptions to reuse semantically similar prior experiences. Recent work has explored using text to retrieve similar past experience for novel target tasks [3, 5, 15]. At the same time, a body of work leverages language for improving policy learning through explicit representation learning [16-22] or using language prediction as an auxiliary task [3, 23-25J. However, we argue that language descriptions may not capture the detailed information for retrieving the most relevant data or guiding policy learning to properly reuse prior data. For example, the same description of "open the door" can apply opening doors with either a knob or handle, each requiring a different low-level motion - hence not all semantically similar motions are useful for the target task. At the same time, the rotating motion "turn on the faucet" might also be useful for learning how to turn a door knob, and focusing on semantic similarity will oversee such data points. In this work, we focus on leveraging fine-grained low-level similarity metries that can capture motion similarity.
Visual Dynamics. On the other end of the spectrum in Fig. we can consider similarity based on visual dynamics, which compares tasks at the pixel level. Behavior Retrieval [2] pretrains variational auto outo-encoder over state-action pairs and uses similarity in that latent space for filtering prior data. Similarly, SAILOR [1] focuses on identifying reusable skills from prior data by pretraining a latent space via an inverse dynamics loss on actions with sequences of images as input. Recent work of Palo and Johns [26] directly compares extracted visual features at annotated bottleneck states for retrieving and replaying recorded trajectories. However, in all of these works, directly embedding RGB observations couples visual similarity with task similarity and hinders the latent space's 's ability extract data from visually different tasks that are similar at the motion level.
Flow-based Motion Guidance. There are intermediate motion representations that do not suffer from the shortcomings of language or visual dynamics similarity. Feature traces and optical flow (shown in the middle range of Fig. 2) are some examples that consider motion similarity at an intermediate level. However, most prior work often use them directly within the policy learning paradigm and not as a retrieval method [27-31]. FLOWRETTIEVAL. instead attempts to use these intermediate representations for retricval enabling a more policy-agnostee approach when tapping into prior data. Additionally, FLOWRETRIEVA, leverages an auxiliary reconstruction loss on optical flow for policy learning to encourage the model to learn motion-centric representations. Such guidance as auxiliary task allows RETRRRERIL to be agnostic the backbone imitation learning algorithm and therefore FLOWRETRIEVAL can be applied to different existing policy learning models unlike prior work that use intermediate motion representations.
4 FLOWRETRIEVAL: Flow-Guided Data Retrieval for Imitation Learning
LOWRETRIEVVL consists of three stages: 1) motion representation pretraining, 2) data retrieval, and 3) policy learning. We discuss each in details below.
4.1 Motion-Centric Pretraining
To retrieve data with similar motion patterns, it is essential to have a latent space encoding observations based on their motion similarity. In particular, existing robot datasets, which contains similar motion patterns, often are collected diverse environments requiring the motion-centrie representation to exclude irrelevant features such background information from the images. Optical flow presents an ideal solution it encodes the motion of pixels without being influenced by other features such as object textures. However, optical flow data is inherently high-dimensional, and direct distances between raw optical flows isn't t meaningful. Hence, we employ straightforward variational autoencoder ( (VAE) to embed optical flow information, creating a latent space that facilitates the assessment of similarity between optical flows.
We process the entire Pprior with flow operator Enow to find optical flow f. for each state s; in Dprioe. In our experiments, we leverage GMFlow [32] for computing the optical flow between visual state 8, and future frame St-k- The value of k is determined by downstream learning policy's action prediction length. For standard behavioral cloning (BC) algorithms that predict single action as output, we set k = 1. For models that predict sequence of actions such as diffusion policy [33, 34],

---

we set k to match the action length the model predicts at each step (k = 16 in our experiments). We denote the matching ordered set of optical flow for prior as priorr
F Fpicr = (Enow (Se, $1.41)) &t+-k E prior
(1)
The VAE consists an encoder (f() and decoder 4s (-), and is trained with reconstruction loss:
V (0,E(0, v)) = Be 1lg (p)( - fille
prior is large enough, we expect the VAE trained only on Fprior to generalize to the target task data Flargets as the domain difference in optical flow much smaller than that in the RGB space.
4.2 Data Retrieval with the Learned Motion-Centric Latent Space
The data retrieval process involves assessing every data point in the prior dataset based on its relative similarity to the target task data, subsequently extracting the highest-ranked portion of these similar data points. Given Dunget we process every state find the optical flow the same way as we process Prrior We then compute the similarity score for each state in prior data by computing pairwise distances target task data and taking the minimum value. The similarity function is:
S(s:)
min YfJEFuge
[Ipo(fi) po(ff)ll fi F Fprice()
We select the top 8% of the prior dataset ordered by the similarity score. Specifically, the similarity threshold can be found at the [8N] -th element of the sorted dataset, where N =====pppiic I/Dprioll::
= sorted S
(s;, [s, E Proice) [[
(4)
The retrieved dataset thereby contains all data with similarity score higher than 7y:
Dretrieved = ((5, ()..a)+(++((s a) Dprior and S(st) > n)
(5)
In practice, it may not be feasible or necessary to scan through all prior data. One can randomly sample a large enough candidate set of data serve as D Presr as suggested by Nasiriany et al. [1].
4.3 FloweGuided Learning
To train the downstream policy, it is not sufficient to only include the retrieved data with similar motion. To ensure the policy pays attention to the retrieved motion, we train the policy with both an action prediction loss and an auxiliary loss for predicting the optical flow. contrast to approaches that require the model to predict dense visual guidance or directly reconstruct the action from optical flow [27, 29], FLOWRETRIEVAL treats optical flow prediction only auxiliary task. This is due to the fact that requiring the model predict such dense visual guidance leads to encoding unnecessary details for predicting end-effector motion. Thus, we employ optical flow prediction as an auxiliary task to provide dense guidance for policy learning without imposing the challenge of learning a significantly more complex task.
bottleneck layer before the action head
FLOWRETRIEVAL is agnostic to the choice of imitation learning algorithm for policy learning. In our experiments, we implemented FETREVALL with diffusion-ban based policy [7, 33]. We assume the BC backbone has some visual encoder dene that embeds visual inputs into a vector form and a policy head &BC that predicts the action output. The visual encoder can optionally start from pretrained representation (e.g. ImageNet-pretraaiee
weights). We incorporate additional flow decoder (Φa)) from the bottleneck representation layer to enforce the auxiliary task loss. For co-training, sample equal amounts of data from Drenioved and Diarger into each batch B. Concretely, the loss for FLOORRETIEEVAL policy learning is:
Lichene &BC, ) LBBc))..eeBLnCCee (as+k) 2) + Aau (feenc(s))) - Enow (8i, 8++4) (6)
where À is hyperparameter (we use À = 0.01 in all of our experiments), and LBC is a supervised Ta loss on action prediction. Algorithm presents the pseudo code of FI.WRETTIEVAL.. line 1-3 pretrain the representations, line 4-8 retrieve the relevant data, and line 9-13 train policy using both the target and retrieved data.
5

---

Algorithm FLOWRETRIEVAL (Drargets D Presce k, 8, &now)
Require: demos Duspet. prior data Paiior. action chunking length k. retrieval threshold 6, flow model Eaow
1: Representation Pretraining (can be done only once fixed prior data)
2: Fpriar {Erow(St, (s+ E priio:;
compute optical flow for prior data f
update 0, 1) by minimizing LVAE(0,
(0, (b) = Teginepann
(po www.:mmmlllllppppp
b. train VAE with Eq. (2)
4: / 2. Data Retrieval */
Furget = {Eiio (8e, 8txx) Dirgai:
> compute optical flow for target data
6: S (-miny, EFuen [Ipo(fi) poff)llleff Fptae);
compute similarity scores with Eq. (3)
7: " = ([
& find retricval threshold Eq. (4)
8: Drenival {(s, ())=:44(((s a)) price and 811)
retrieve data Eq. (5)
9: f= 3. Policy Learning $/
10: randomly initialize policy madel Orne, dBc, Owx (optionally load pretrained weights for ven))
11: repeat
sample batch Bo.b Durget, B.41:2b Dratieved to update with loss L(den ODC, dax)):
Dr Eq. (6)
13: until & converged: return
Square Assembly
LIBERO-Can
Bridge-Pot
Goddge-iirrowen
Franka-Pen-in--up
square pey
Put soup
Thanefer put
Open
up red morker and place eup
enstrconments
2 taoks
environments 18 taoks
envirenment
cur badks Wildet tn-thescild divense taeks
Figure 3: Experimental Setup. We experiment with 5 different manipulation tasks for evaluating FLOwRRr TRIEVAL. Bottom shows the prior dataset used in each experiment and their meta data.
5
Experiments
In this section, we first introduce the task domains and the baselines we compare with, and then present the quantitative results of policy learning and qualitative analysis of retrieval.
Task Domains. We evaluate FLOWRETRIEVAL table-top manipulation tasks of varying difficulty and different compositions of prior datasets (a total of 5 tasks shown in Fig. 3);
Square Assembly is peg insertion task from the Robomimic benchmark [35]. Behavior Retrieval [2] adapted this task to evaluate data retrieval. Diarget consists of 10 demonstrations.
D prirr consists of two different types of data: 200 episodes of useful demonstrations placing the nut into the goal peg, and 200 episodes adversarial demonstrations that place the nut into wrong peg. We further adapted this task emulate an extreme case where the target task is in the same environment as the adversarial data while the useful data has a different background.
LIBERO-Can is pick-and place task from the LIBERO benchmark [6]. D Duaget consists of 10 human demonstrations from the LIBERO benchmark, and DD Ppiior consists of 18 selected tasks from LIBERO-90 and each with 50 trajectories, resulting in a total of 900 trajectories. Half of the prior data are pick-and- -lace motions that are intuitively useful for learning the target task, and the other half are tasks involving other diverse motions such as opening a drawer.
Bridge-Pot and Bridge-Microwave are two real robot tasks a toy kitchen setting. We leverage Bridge-V3 V2 [7] as Pprior and collected Piager in our own setup with ViperX arm [36], which is similar to some environments in Bridge- V2 but does not exist in the prior dataset. We collected 10 demonstration for each of the target tasks and use a subset of Bridge 2 as prior data.
Franka-Pen-in-Cuu -up is real task with the Franka Emika Panda robot arm. Drange consists of 10 human demonstrations. We tested two instances of D Dprice: 1) PnP (Pick-and-Place): 105 trajectories of leck-and-place tasks from the same robot, and 2) Wild: 400 randomly sampled trajectories from the DROID [8] dataset (filtered to roughly match the viewpoint target task). PnP contains semantically similar data target task while Wild contains diverse range of tasks.

---

Square Assembly
LIBERO-Can
Microwave
Bridge-Pot
Pranka Pen-in-Cup
RowRetteeal (PhP))
FlowBeriival (WH22)
255
64
CF
SC
FoowC (taggt))
S4
(hrrpe))
FlowRetrieaal (owns)
Prupriolletrieval
BR
5R
BC
GCBC
Pick
Figure 4: Quantitative Results. We plot success rates (%) of learned policies observe that FL.Ww RETRIEVAL. outperforms baselines across tasks. Simulation results are averaged over 2 training seeds and 3 evaluation seeds (50 rollouts each). Real results are from 25 rollouts of best-of-last-3 checkpoints.
Baselines. We leverage co-training for downstream policy learning for all baselines that use prior data, and each training batch contains mixture of half target task data and half prior data. In the simulation domains and real- world Bridge tasks, we compare FLOw RETRIEVAL with several baseline imitation learning algorithms including prior data retrieval methods (full details in Appendix):
BC is standard behavior cloning with diffusion policy [33] trained only on target data;
GCBC is goal-conditioned diffusion policy trained on all data;
BR (Behavior Retrieval) a prior work [2] that trains a VAE for state-action pairs for retrieval;
SR (SAILOR Retrieval) is the pretraining algorithm from SAILOR [I]: we use the pretrained latent skill space for retrieving data but keep policy learning consistent with other methods; ProprioRetrieval is naive method where we use proprioceptive data for computing similarity.
In Franka-Pen-in-Cul Cup task, we focus on evaluating FLOW with different types of prior dataset (PnP and Wild in Fig. 3) and compare it with co-training-onl baseline BC-Co co-trained well as ablation FlowBC that applies flow loss during
on all of prior data and target data policy learning from target task data without retrieval.
Quantitative Results. Fig. 4 presents the success rates of different methods for each task. FLOW RE- TRIEVAL outperforms the best baseline method across different tasks, achieving an average of 14% higher success rate than the best baseline method in each domain (+10% simulation,+199 in real). FL.ORRETRIVVAL also achieves on average 27% higher success rate than the best retrieval-based prior method. We observe that BR and SR (shown in blue) can perform worse than BC (shown in grey), which is potentially due to retrieving harmful data from the prior dataset. PeeprioRetrieval (shown in green) performs competently in the two simulated tasks. However, it can fail in more complex settings (e.g. when camera viewpoints in prior data vary), due to the fact that it is ignorant of the visual state (more in Appendix). In the real Pen-in-Cup task with a Franka Emika robot (right of Fig. 4), FLOWRETRIEVAL (shown in orange) achieves 3.7 x the performance of the imitation baseline, learning from all prior and target data (shown in grey).
Ablation of Data Composition. Fig. (right) shows the evaluation and ablation of FL.WRRETRIEALL with different dataset compositions Franka Pen-in-Cae Cup task. this task, the pick stage of the target task shares similarity with prior dataset (PnP, which consists of mostly pick and place tasks) while the place stage of the task employs unique strategy that is very different from this prior dataset. We see that BC-Co (in grey) with PnP improves pick performance compared to BC (in black) but performs worse for the overall pick and place task. FLOW R ETRIEVAL (in orange) instead successfully filters out the harmful data in PnP and achieves the highest performance. F..OWRETRIEVA. (in dark orange) also achieves performance higher than FlowBC (+16%) (in yellow) in this task by retrieving from in-the-wild prior data, Note that due the large discrepancy between viewpoints prior and target data, do not use the action retrieved from Wild for action prediction.
Qualitative Analysis. Fig. 5 visualizes example retrieved data points by each method FLOWRE- TRIEVAL retrieves similar motion from prior data while BR and SR retrieve based on visual similarity of the state and consequently can retrieve adversarial data. In the Square Assembly example, both baselines end up retrieving data points where the robot moves towards the round peg. FLOWRE- TRIEVAL. on the other hand selects a very similar motion toward the square peg even when the
Additional visualization of retrieved data for each task can be found in Appendix and project website.
7

---

Most Similar Data Retrieved by Different Method
Task Domain
Targel Query fnsss
Retrieeii (Chrrs)
BR
Squere Acsembly
LIBERO-CHI Can
Figure 5: Qualitative Analysis. We visualize most similar motion prior dataset to queries from target task demonstrations. For each example data point, we show the image observations and 814-66 then overlay the optical flow top of We see that FLOw FLOWRETRIEWS RETRIEVAL focuses on retrieving similar motions to target queries while baselines may retrieve visually similar states with different motions.
Usefill Deta fenter (adigrunn)
Ubgfil Data Gifferent bazkgronud)
Afterarrial Dits fead scfice frrking point)
Non-hamyfii Det
resooal das
RowRRttiival tours
ProprioRetrivval
SALLOR
RetiriccRttiiral
(Bach 72
Trsnsder
1lace
Souch Pickup
Truneer
11ace
Noach Pick
Transfer
11ac
Raach 1966
Tran-der
7lace
Figure Retrieval Distribution in Square Assembly. visualize the distribution and composition of retrieved data by different retrieval methods. BLOWRETRIEVAL mostly retrieves useful data (green) for each stage of the task especially during the challenging parts of picking up and transferring the square to the correct peg and filters the majority of adversarial data (red).
background color completely different from that in target data. In LIBERO-Cam, BR and SR again focus mainly on visual similarity instead of the picking up motion and retrieved data points that have very different motion (moving downwards). However, FEOWRETRIEVAL retrieves the correct motion of picking up, even if it is picking up different object than the one in the target task.
Additionally, consider the Square Assembly task, where have curated prior dataset. Specifi- cally, we know the data right after the robot picks up the nut is crucial: we label the actions that move towards the wrong goal after picking up as adversarial data. We label the data points after the forking point for transferring the square on the decided peg as non-harmful. We can therefore analyze the quality of retrieved data under different retrieval methods in the Square Assembly task by plotting the amount of different types of data retrieved from each stage of the task (split into 10 bins) as shown in Fig. 6. RETRIEVAL uniformly retrieves from prior useful data and retrieves little adversarial data, while baseline models either cannot effectively filter out the adversarial data or do not retrieve enough useful data at the bottleneck stage (between pick up and transfer) of the task.
6 Conclusion
Summary. We propose FL.WWRETRIEVALL, method for few-shot imitation learning that retrieves motion-similar data from heterogeneous prior datasets via optical flow. FLOWRETRIEVAL addresses the limiting assumption of prior data retrieval methods that often require visually similar experi- ences to exist in prior dataset. In our experiments, FLLOWRETRIEVA achieves significantly better performance (+14% on average) than other retrieval-based methods.
Limitations and Future Work. Pretraining with large-scale datasets requires extensive computation but done only once. In contrast, FLOWRETRIEVAL processes prior data for each new target task, calculating and sorting pair-wise distances, which can add overhead and doesn't scale well with large datasets. Caching embeddings and sub-sampling can help, but practical deployment remains challenging, limitation left for future work. Meanwhile, the optimal retrieval threshold task- dependent and needs manual tuning for best performance. Future work could automate threshold discovery using annotated boundary data active learning methods.

---

Acknowledgments
This work is supported by NSF projects #1941722, #2125511, #2218760, ONR project N00014-21-1- 2298, DARPA #W91 INF2210214. This research was supported in part by Other Transaction award HR001 12490375 from the U.S. Defense Advanced Research Projects Agency (DARPA) Friction for Accountability in Conversational Transactions (FACT) program.
References
[1] S. Nasiriany, T. Gao, A. Mandlekar, and Y. Zhu, Learning and retrieval from prior data for ball-bbsed imitation learning. In Conference on Robot Learning (CoRL), 2022.
[2] M. Du, S. Nair. D. Sadigh, and C. Finn. Behavior retrieval: Few-shot imitation learning by querying unlabeled datasets. In Robotics: Science and Systems (RSS), 2023.
[3] S. Belkhale, T. Ding. T. Xiao, P. Sermanet, Q. Vuong, J. Tompson, Y. Chebotar, D. Dwibedi, and D. Sadigh. Rt-h: Action hierarchies using language. arXiv preprint arXiv:2403.0182 2024.
[4] L. Shao, T. Migimatsu, Q Zhang, K. Yang, and J. Bohg. Concept2robot: Learning manipulation concepts from instructions and human demonstrations. The International Journal of Robotics Research, 40(12-14): 1419-1434, 2021.
[5] I.. Zha, Y. Cui, I.-H. Lin, M. Kwon, M. G. Arenas, A. Zeng, F. Xia, and D. Sadigh. Distilling and retrieving generalizable knowledge for robot manipulation via language corrections. In International Conference on Robotics and Automation, 2024.
[6] B. Liu, Y. Zhu, C. Gao, Y. Feng, Q. Liu, Y. Zhu, and P. Stone. Libero: Benchmarking knowledge transfer for lifelong robot leamning. Conference on Neural Information Processing Systems, 2023.
[7] H. Walke, K. Black, A. Lee, M. J. Kim, M. Du, C. Zheng. T. Zhao, P. (lansnn-Estruch Q. Vuong, A. He, V. Myers, K. Fang, C. Finn, and S. Levine. Bridgedata v2: A dataset for robot learning at scale. In Conference on Robot Learning (CoRL), 2023.
[8] A. Khazatsky, K. Pertsch, S. Nair, A. Balakrishna, S. Dasari, S. Karamcheti, S. Nasiriany, M. K. Srirama, L. Y. Chen, K. Ellis, P. D. Fagan, Hejna, M. Itkina, M. Lepert, J. Ma, T. Miller, J. Wu, S. Belkhale, S. Dass, H. Ha. A. Jain, A. Lee, Y. Lee. M. Memmel. S. Park, I. Radosavovic, K. Wang. A. Zhan, K. Black,
C. Chi, K. B. Hatch, S. Lin, J. Lu, J. Mercat, A. Rehman, P. R. Sanketi, A. Sharma, C. Simpson, Q. Vuong. H. R. Walke, B. Wulfe. T. Xiao, J. H. Yang. A. Yavary. T. Z. Zhao, C. Agia, R. Baijal. M. G. Castro. D. Chen, Q. Chen, T. Chung, J. Drake E. P. Foster, J. Gao, D. A. Herrera, M. Heo, K. Hsu, J. Hu, D. Jackson, C. Le, Y. Li, K. Lin, R. Lin, Z. Ma, A. Maddukuri, S. Mirchandani, D. Morton, T. Nguyen,
A. O' O'Neill, R. Scalise, D. Scalc, V. Son, S. Tian, E. Tran, A. E. Wang, Y. Wu, A. Xic, J. Yang, P. Yin, Y. Zhang, O. Bastani, G. Berseth, J. Bohg, K. Goldberg, A. Gupta, A. Gupla, D. Jayaraman, J. J. Lim, J. Malik, R. Martín-Martn, S. Ramamoorthy, D. Sadigh, S. Song. J. Wu, M. C. Yip. Y. Zhu, T. Kollar, S. Levine, and C. Finn. Droid: A large-scale in-the-wild robot manipulation dataset. 2024.
[9] S. Dasari, M. K. Srirama, U. Jain, and A. Gupta. An unbiased look at datasets for visuo-motor pre-training. In Conference Robot Learning, pages 1183-1198. PMLR, 2023.
[10] I. Radosavovic, T. Xiao, S. James, P. Abbeel, J. Malik, and T. Darrell. Real-world robot learning with masked visual pre-training. In Conference on Robor Learning, pages 416 426. PMLR, 2023.
[11] Octo Model Team, D. Ghosh, H. Walke, K. Pertsch, K. Black, O. Mees, S. Dasari, J. Hejna, C. Xu, J. Luo, T. Kreiman, Y. Tan, D. Sadigh, C. Finn, and S. Levine. Octo: An open-source generalist robot policy. gi github. io, 2023.
https: http:////oooeeeeeee
[12] Z. Fu, T. Z. Zhao, and C. Finn. Mobile aloha: Learning bimanual mobile manipulation with low-cost whole-body teleoperation. In arXiv, 2024.
[13] O. X.-E. Collaboration, A. O'Neill, A. Rehman, A. Maddukuri, A. Gupta, A. Padalkar, A. Loe, A. Pooley, A. Gupta, A. Mandlekar, A. Jain, A. Tung, A. Bewley, A. Herzog, A. Irpan, A. Khazatsky, A. Rai, A. Gupta, A. Wang. A. Singh, A. Garg, A. Kembhavi, A. Xie, A. Brohan, A. Raffin, A. Sharma, A. Yavary, A. Jain, A. BBlakkrihha,, A. Wahid, B. Burgess-Limerick, B. Kim, B. Schölkopf, B. Wulfe, B. Ichter, C. Lu, C. Xu, C. Le. Finn, C. Wang, C. Xu, C. Chi, C. Huang, C. Chan, C. Agia, C. Pan, Fu, C. Devin, D. Xu, D. Morton, D. Driess, D. Chen, D. Pathak, D. Shah, D. Büchler, D. Jayaraman, D. Kalashnikov, D. Sadigh, E. Johns, E. Foster, F. Liu, F. Ceola, F. Xia, F. Zhao, F. Stulp, G. Zhou, G. S. Sukhatme, G. Salhotra. G. Yan, G. Feng, G. Schiavi, G. Berseth, G. Kahn, G. Wang. H. Su, H.-S. Fang, H. Shi, H. Bao, H. B. Amor, H. I.

---

Christensen, H. Furuta, H. Walke, H. Fang, H. Ha, I. Mordatch, I. Radosavovie, I. Leal, J. Liang, J. Abou- Chakra, J. Kim, J. Drake, J. Peters, J. Schneider, J. Hsu, J. Bohg. J. Bingham, J. Wu, J. Gao, J. Hu, J. Wu, Wu, J. Sun, J. Luo, J. Gu, J. Tan, J. Oh, J. Wu, J. Lu, J. Yang. Malik, J. Silvério, Hejna, J. Booher, J. Tompson, J. Yang, J. Salvador, J. J. Lim, J. Han, K. Wang, K. Rao, K. Pertsch, K. Hausman, K. Go, K. Gopalakrishnan, K. Goldberg, K. Byrne, K. Oslund, K. Kawaharazuka, K. Black, K. Lin, K. Zhang, K. Ehsani, K. Lekkala, K. Ellis, K. Rana, K. Srinivasan, K. Fang. K. P. Singh, H. Zeng, K. Hatch, K. Hsu, L. Itti, I. Y. Chen, I.. Pinto, I., Fei-Fei, L.. Tan, I.. J. Fan, I.. Ott, I., Lee, L. Weihs, M. Chen, M. Lepent, M. Memmel, M. Tomizuka, M. Itkina, M. G. Castro, M. Spero, M. Du, M. Ahn, M. C. Yip, M. Zhang, M. Ding, M. Heo, M. K. Srirama, M. Sharma, M. J. Kim, N. Kanazawa, N. Hansen, N. Heess, N. J. Joshi, N. Suenderhauf, N. Liu, N. D. Palo, N. M. M. Shafiullah, O. Mees, O. Kroemer, O. Bastani, P. R. Sanketi, P. T. Miller, P. Yin, Wohlhart. P. Xu, P. D. Fagan, P. Mitrano, P. Sermanet, P. Abbeel, Sundaresan. Q. Chen. Q. Vuong, R. Rafailov, Tian, R. Doshi, R. Mart in-MMrt in, R. Baijal, R. Scalise, R. Hendrix. R. Lin, R. Qian, R. Zhang, R. Mendonca, R. Shah, R. Hoque, R. Julian, S. Bustamante, S. Kirmani, S. Levine, S. Lin, S. Moore, S. Bahl, S. Dass, S. Sonawami, S. Song, S. Xu, S. Haldar, S. Karamcheti, S. Adebola, S. Guist, S. Nasiriany, S. Schaal, S. Welker, S. Tian, S. Ramamoorthy, S. Dasari, S. Belkhale, S. Park, S. Nair, S. Mirchandani, T. Osa. T. Gupta, T. Harada, T. Matsushima, T. Xiao, T. Kollar, T. Yu, Zhan, T. Ding, Davchev, Z. Zhao, T.. Armstrong, T. Darrell, Chung, V. Jain, V. Vanhoucke,
W. Zhou, W. Burgard, X. Chen, X. Wang. X. Zhu, X. Geng, X. Liu, X. Liangwei, X. Li, Y. Lu, Y. J. Ma, Y. Kim, Y. Chebotar, Y. Zhou, Y. Zhu, Y. Wu, Y. Xu, Y. Wang. Y. Bisk, Y. Cho, Y. Lee, Y. Cui, Y. Cao. Y.-. Wu, Y. Tang, Y. Zhu, Y. Zhang, Y. Jiang. Y. Li, Y. Li, Y. Iwasawa, Y. Matsuo, Z. Ma. Z. Xu, Z. J. Cui, Zhang, and Z. Lin, Open X- X-Embodiment: Robotic learning datasets and RT-X models. https:// httpa////awxvvvo/aeeeeee
2abb/2310 088644, 2023.
[14] S. Belkhale, Y. Cui, and D. Sadigh. Data quality in imitation learning. Advances in Neural Information Processing Systems, 36, 2024.
[15] G. Wang, Y. Xie, Y. Jiang, A. Mandlekar, C. Xiao, Y. Zhu, L. Fan, and A. Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv: (e0.-2235..629 2023.
[16] C. Lynch and P. Sermanet. Language conditioned imitation learning over unstructured data. Proceedings of Ruboticx: Science and Systems (RSS), 2021.
[17] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar. J. Dabis, C. Finn, K. Gopalakrishnan, K. Hausman, A. Herzog. J. Hsu, et al. Rt-1: Robotics transfommer for real-world control at scale. arXiv preprint arXiv:2212.068// 2022.
[18] C. Lynch, A. Wahid, J. Tompson, T. Ding, J. Betker, R. Baruch, T. Armstrong, and P. Florence, Interactive language: Talking to robots in real time. IEEE Robotics and Automation Letters, 2023.
[19] A. Rashid, Sharma, C. M. Kim, J. Kerr, L. Y. Chen, A. Kanazawa, and K. Goldberg. Language embedded radiance fields for zero-shot task-oriented grasping. In Conference on Robot Learning, pages 178-200. PMIRR 2023.
[20] Y. J. Ma, W. Liang. V. Som, V. Kumar, A. Zhang, O. Bastani, and D. Jayaraman. Liv: Language-image representations and rewards for robotic control. arXiv preprint urXiv:2306.0095. 2023.
[21] S. Karamcheti, S. Nair, A. S. Chen, T. Kollar, C. Finn, D. Sadigh, and Liang. Language-driven representation learning for robotics. In Robotics: Science and Systems (RSS), 2023.
[22] S. A. Sontakke, Zhang, S. Amold, K. Portsch, E, Biyik, D. Sadigh, C. Finn, and L. Itti. Roboclip: One demonstration is enough to learn robot policics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[23] S. Mirchandani, S. Karamcheti, and D. Sadigh. Ella: Exploration through learned language abstraction. Advances in neural information processing systems, 24:29529-29540. 2021.
[24] J. Hejna, P. Abbeel, and L. Pinto. Improving long-horizon imitation through instruction prediction. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 7857-7865, 2023.
[25] M. Srivastava, C. Colas, D. Sadigh, and J. Andreas, Policy learning with a language bottleneck. arXiv preprint 2405:22000...188,, 2024.
[26] N. D. Palo and E. Johns. Dinobot: Robot manipulation via retrieval and alignment with vision foundation models. In IEEE International Conference on Robotics and Automation (ICRA), 2024.
[27] C. Wen, X. Lin, J. So. K. Chen, Q. Dou, Y. Gao, and P. Abbeel. Any-point trajectory modeling for policy learning. arXiv preprint arXiv:2401.000225 2023.
10

---

[28] M. Vecerik, C. Doersch, Y. Yang, T. Davchev, Y. Aytar, G. Zhou, R. Hadsell, L. Agapito, and J. Scholz. Robotap: Tracking arbitrary points for few-shot visual imitation. arXiv, 2023.
[29] P.-C. Ko, J. Mao, Y. Du, S.-H. Sun, and J. B. Tenenhaum. Learning to Act from Actionless Video through Dense Correspondences. arXiv:2310.08576, 2023.
[30] H. Bharadhwaj. R. Mottaghi, A. Gupta, and S. Tulsiani. Track2act: Predicting point tracks from internet videos enables diverse zero-shot robot manipulation, 2024.
[31] Y. Zhu, A. Lim, P. Stone, and Y. Zhu. Vision-hased manipulation from single human video with open-world object graphs. arXiv preprint arXiv:2405.20:3 2024.
[32] H. Xu, J. Zhang, J. Cai, H. Rezatofighi, and D. Tao. Gmflow: Learning optical flow via global matching. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 8121-8130, 2022.
[33] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. Proceedings of Robotics: Science and Systems (RSS), 2023.
[34] M. Reuss, M. Li, X. Jia, and R. Lioutikov. Goal-conditioned imitation learning using score-hased diffusion policies. arXiv preprint arXiv:2304.0255: 025322. 2023.
[35] A. Mandlekar, D. Xu, J. Wong. S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese, Y. Zhu, and
R. Martin-Martin. What matters learning from offline human demonstrations for robot manipulation. In arXiv preprint arXiv:2108.0333 2021.
[36] T. Robotics. Viperx 300 robot arm, viperx-300 -00000000-0-t-aee arm.arr.ppppp..... aspx.
URL http:::/ww..troseerootto...
[37] T. Robotics. Widowx 250 robot arm, widowx- (adowx-250--a ro0otoott--aaaae-- arm- wdw..ddofffoe.ooeo aspx.
URL httpp://www.troseeerobbbt
[38] Y. Cui. Niekum, A. Gupta, V. Kumar, and A. Rajeswaran. Can foundation models perform zero-shot task specification for robot manipulation? In Learning for dynamics and control conference, pages 893-905. PMLR, 2022.
[39] S. Nair, A. Rajeswaran, V. Kumar, C. Finn, and A. Gupta. R3m: A universal visual representation for rohot manipulation. In Conference Robot Learning, pages 892-909. PMLR, 2023.
[40] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, pages 8748 8763. PMLR, 2021.
[41] M. Oquab, T. Darcet, T. Moutakanni, H. Vo. M. Szafraniec, V. Khalidov, P. Fernandez, D. Haziza. F. Massa, A. El-Nouby, al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193 2023.

---

Appendix
Table of Contents
A Experimental Setup
12
A.1. Experimental Domain Details
12
A.2 Implementation Details of Baseline Methods and Ablations
13
A.3 Retrieval Threshold Selection
13
B Additional Baselines and Ablation Results
14
B.1 Additional Baselines
14
B.2 Using Proprioception for Retrieval
14
Retrieval Visualization
15
Ablating Retrieval Strategies
D.1. KNN-based Retrieval Strategy
16
D.2 Retrieving with Pre-trained Representations
16
Experimental Setup
A.1 Experimental Domain Details
Below provide full details about the experimental domain we used in experiments:
Square Assembly: The goal is to pick up and place square nut into the square peg. Durget consists of 10 demonstrations. Pprior consists of two different types of data: 200 episodes of useful demonstrations placing the nut into the goal peg, and 200 episodes of adversarial demonstrations that place the nut into the wrong peg. Note that the initial phase of the adversarial data also consists of useful motion for learning to pick up the nut. In setup, the target task is in the same environment as the adversarial data while the useful data has a different background. We generated optimal demonstrations using scripted policies.
The goal is to pick up a can and place it into a basket. Corresponding task description is "pick the alphabet soup and place it in the basket" and the environment is LIV- ING ROOM SCENE1 (from the LIBERO-90 suite). D Duarger consists of 10 human demonstrations from the LIBERO benchmark, and D prior consists of 18 selected tasks from LIBERO-90 and each with 50 trajectories, resulting in total of 900 trajectories. The list of task IDs in LIBERO-90 that was selected to form the prior dataset is:
LIBERO-Can:
"LIVING BOETIN___OOO_SEE _0ON_SCENELL____
thethe. the wha._kaccupppaee _n___annddd_______eeeeee_______e ptt__ttt_t____t_ttttt______te__e sa__thh_baaaaee
"LIVING ROOM, SCENE2, pick_up.pppppoppppppppppkpppkkkp up. the
SsaNE2_pick_uu
the_thtterr butter_and and_put,____pppppppp put_it,ttttttt
it_iiitti_nni in_in_tth the, basket"
"LiVING_ R0OM_SCENE2 ScENE2_pick.ue ptck_ppppppppppppppppppppppppppp the_tee__ee cheeoorngee__.. oraaooaanneau..e.. jurge__jucc__am aeee_and__uu_tt_ and_put_it_ff
it_iitnn....... in_
up_the the tonato sautt___aeeeeeeeeee aaddeeeendd___eee and__at___tt_tt it, th__the_bakktt "LIVING, ROOM, SCENE2, ppck_up_
"LIVING RCOM_
ScENE2
pick_up-
the milk put__ut_ii. in_
Ste_tie________eeeeeeeeeeeeeee__ the basket"
"LIVING_AA" BEVWING_DOO_____ _tW__HOCCOEE____________2 ppc__pp__ppppppppppppppppp__pppp sp_the__t tho___lbabet__eee sappabet_eoup__d___p___ppp_ soup_and_putt___ ant_aadttti_iin________a_
"LIVING ROOM sew___CEEE__pp__
pick_pp_thhe_peepppppppppppppppe the
spp_the__lphabet alphabat,
"LIVING ROOM_ R0ON_SCENE,,____ set.EecE_epppp.
but.eutterr... wtt.aanddttt....i... in___itt____tee thay"
"LIVING ROOM, SCENEA_ stack the the_left, bowl_oolooool on, bowl_on_the_rrte rew_trigghbboo bowb_oadd_pieeeeeeeeaaa and_place, Silace_then_inn in the_tray"
"LIVING_ ROOM _tw___CCOOOO_____________ Battie__pt___heeeee
"KITCHEN SCENE1 ooeen.tta the botton, draver_ of_ the_cabinet"
"KITCHEN BEITCHEN_SCENE__CCC Be__e2_put_tteeeeeeeeeee_ee_eeee beceeennddlle__eeeeeeeeeeeeeeeee black blackbbbbcllalaaae pool_ea_the_plate
"KITCHEN_ SITCCEEN_SCENES, wwrn_tuu_ttttteeoooetttooo saar__oo_the_teeee____ee cece_stto__aad__
scenes_turn_ot ofrn_off satr_oof_thh_tttt the_atove" "KiTCHEn_ Baceeessseee____
"KITCHEN_ SCENE1O
top_draver_of of cheinetb
"STuDY SCEXE1 pick up. the book, and_ppaa___t_...._oaa stceee_ttii_____ee it_in_the wtwiiiteefroottteeeeeeeeeeeeeeee stont_compaataentt of the caddy" the
btt..bcccccopptteeettt... theettt________ee_ee___eeeeeeee_ the__addy" "STUDY ScExe2_picckeeccceeeeecec pick_up_thh btt_thhh_ooo..oo bea_book_and_ plt_pcl_t___ittttttte
"STUDY_ SRRDY_SSEEE___.E ppckpppcppppppp the white.
thite__ngg_eee sagg.aaaana www_edddde_it___.eaaaaeaa ttt._ittttttttrtgtteeeeee of__oof__t___a__________________
Bridge-Pot and Bridge-Mierooave: In Bridge-Pot the robot picks up pot on the burner and places into the sink. In
the robot grasps the handle of the microwave door and
12

---

pulls it open. We leverage B2idge-V2| [7] as Dprior and collected Durget in our own setup with ViperX arm [36], which is similar to some environments Bridge- V2 but does not exist in the prior dataset. We use a ViperX arm [36] instead of the WidowX [37] as in Bridge- V2 dataset, which introduces additional domain difference for transferring useful pattern in prior data. We collected 10 demonstration for each of the target tasks using VR teleoperation and use subset of Bridge- V2 as prior data. The full list of environments in the prior dataset is:
batacoll_ten toykitchenn
Saacclll_oykttthhem
[hacoll__ooddng... table
datacol2_robot. rooo___ooo___e
Sttaco2_.toykttcchee
Salacol2_toykitcttn
datacol2_ toysink2
deepthought_robtt..dd
weepthought_toykktchhee
segpthought_toykitccheee
The goal is to pick up marker and put it into a cup. Durget consists of 10 human demonstrations. We tested two instances of D Ppia:: 1) PnP: 105 trajectories of pick-and- place tasks from the same robot that we collected ourselves through VR teleoperation, and 2) Wild: 400 randomly sampled trajectories from the DROID [8] dataset (filtered to roughly match the viewpoint in target task). Pen-in-Cup is a task that exists DROID but the demonstrations were collected in very different environments.
Prania--en-in-Cup::
A.2 Implementation Details of Baseline Methods and Ablations
We have two sets of implementations of experiments in simulation as well as the Franka-Pent Pen-in- task, we adapted the diffusion policy implementation of Chi et al. [33]; we use the U-Net variant of diffusion policy and load pretrained ImageNet weights for initializing the encoder network; 2) for Bridge-* tasks, we use the diffusion model implementation provided by Walke et al. [7];; ; also load pretrained ImageNet weights for the encoder of the policy.
and baselines for our experiments: 1) for
For BR, we re-implemented the pretraining logic for the VAE and reused the hyperparameters as documented in Du et al. [2]. For SR results in the paper, we re-implemented the pretraining logic for the latent skill space with the provided hyperparameters in Nasiriany et al. [1]. We also ran our simulation experiments using the full implementation of SAILOR provided by Nasiriany et al. []], but fail to obtain non-zero success rates. One potential reason is that the tasks we consider in our experimental setup is relatively short, while SAILOR was designed for tackling long-horizon tasks (4-6 steps per task), such that our experimental domains may require a completely different set of hyperparameters than those used in the original work.
ProprioRetrieval is a method where we use proprioceptive data for computing similarity at the retrieval stage. We also apply flow loss during policy learning, and therefore it can be considered as an ablation of the retrieval space of FLOW RETRIEVAL. Specifically, we leverage end-effector Cartesian position difference between 8; and 8t+-& to represent motion. While it is desirable match the rotational motion as well, we find using the full end-effector Cartesian pose to compute similarity requires tuning scaling factors for balancing positional state versus rotational state (orientation) and not only does that require manual evaluations, but also existing datasets often use different conventions for orientation and therefore introduces additional challenges for comparing proprioceptive states across datasets directly. Hence in our implementation of ProprioRetrieval, we use only delta of the position of the end effector to retrieve from prior data, The feature e, for datapoint s, is computed as:
et (st, St+ [ 3] -s[[ 3])
The similarity score is then computed as the negative &a distance between feature vectors.
A.3 Retrieval Threshold Selection
In FLOWRETRIEVAL, we retrieve the top 8% from the prior data. However, the optimal threshold can be different for different target task and prior dataset. Intuitively, if the similarity metric truly ranks the prior datapoints by usefulness to target task, we want to retrieve just the right amount of
phttps://giiee cattchuuu.ommm/aiiiaaaa www/rail-berkeleeeeeoeeeeeeoeeee bww.bbeeeeeelbbbrre.dda.....
13

---

Method
SquareAssembly
LIBERO-Can
Bridge-Mien
Bridge-Por
Pen-in- -n-
BC
7%
64%
56%
20%
23%
BC-Co
3%
16%
56%
32%
12%
FlowBC
7%
71%
48%
8%
40%
ProrioRetrieval-
44%
73%
24%
40%
ProrioRetrieval
54%
81%
16%
12%
44%
FLOWRETRIEVAL
90%
68%
64% %
56%
Table Success rates of baselines and ablations of FLOWRETRIEVAL.
data that supports the learning of target task. Du et al. [2] observed bell-shaped curve between the relationship of retrieval threshold and policy performance. Due constraints computing resources, we only sparsely searched over a small range of threshold values for 3 tasks (2 in sim, in retrieve 35% from
real): Square Assembly, LIBERO-Cam and Bridge-Pot. In experiments, prior data in Square Assebmly, 10% in LIBERO-Car and 1% in Bridge tasks. We then reuse the same threshold from bridge-Pot for the Tridge-Microwaye task. For Franka Pen-in-Cup, reuse the threshold of found in the Square Assembly task.
B Additional Baselines and Ablation Results
B.1
Additional Baselines
We present the full evaluations BC-Co and FlowBC Table BC-Co directly cotrains with the entire prior dataset. FlowBC applies auxiliary flow loss to BC. We see that while each could improve the performance from vanilla BC in certain domains, they can also hurt performance in others, depending the composition of prior dataset.
B.2
Using Proprioception for Retrieval
Low level motion features from proprioception are simple compute and be effective for extracting similar actions from past experience. Therefore we ablate our retrieval method with using propriceptive information, and denote this method ProprioRetrieval. However, proprioceptive features are entirely agnostic to the visual observation and therefore may retrieve irrelevant data such those from very different viewpoints if the prior dataset has different viewpoints, e.g. OXE [13] and DROID [8].
ProprioRetrieval still applies the auxiliary flow loss during policy learning. We additionally evaluate baseline version that does not use the flow loss as a comparison point, and use ProprioRetrieval- to denote this baseline. The success rates both proprioception-bant based methods are reported in Table We see that ProirioRetrieval can lead to high success rates when camera viewpoints are consistent (e.g. in Square Assembly and LIBERO-Can Can), but does not retrieve enough useful data from Wild to match the performance of FI.Ow RETRIEVAL in the Franka Pen-in-Cup task, and performs poorly in the Bridge tasks that have a large variety of prior tasks. Note that the downstream policy learning does not use the low-level action data to update the policy branck in the Franka Pen-in-Cup task when retrieving from the Wild dataset, due to the large variation in camera viewpoints. However, we still use the low level actions from the retrieved prior data of Bridge since the viewpoints are better aligned, but may introduce additional multimodality in policy learning.
14

---

C Retrieval Visualization
Teok Donsils
AcessAy
LIBERO Cae
Quiry
Quar2
Query
Qany
Targel
Most Similer Data
FluuRorica! (Curn)
SALLOR
Figure 7: Visualization paired retrieved data query larget datapoint by different methods in Square Assembly and LIBERO-Can.
Fig. 7 shows additional queries and retrieved data by different retrieval method in the two simulated tasks. BR and SR both encodes visual observations and low level actions the latent space used for retrieval. We that BR often overly focuses visual scene similarity, and motion similarity is likely a second order feature (only effective the visual scene is similar in the first place). Therefore in all cases, it cannot ignore the effect of the environment (background) and end up retrieving potentially adversarial data. SR can latch on either of the similarities, sometimes focusing on visual scene (Sqaure Assembly), and sometimes retrieving based on action similarity (LIBERO-Can). In contrast,
consistently focuses on visual motion similarity.
ProprioRetrievall
PnP Dalasel
Wild Datasel
Dalase!
Figure 8: Visualization of example data points retrieved by FLOWRETRIEVAL and ProprioRetrieval in Franka
Pen-in-Cup task.
Fig. 8 shows the example data points retrieved by FLOWRETRIEVAL and ProprioRetrieval in Franka Pen-in-Cup task. When retrieving from the PnP dataset, FLOWRETRIEVAL focuses on the pick- up and transfer stage of the target task, and does not retrieve the placing motions from the prior dataset, effectively filtering adversarial prior data. When retrieving from the Wild dataset, we see that FLOWRETRIEVE ProprioRetrieval retrieves very different viewpoints (sometimes robot is not even in the view see example in rightmost column, second from bottom). These data points are less informative for the downstream learning of the target task but may still provide additional regularization learning diverse set of visual features.
retrieves viewpoints better aligned with that in the target task, while
15

---

top
nearest neighbor
Figure Illustration two retrieval strategies. Depending on distribution target and prior data points. top-% may retrieve data only close to certain part the target trajectory while KNN would retricve uniformly.
Top 15$
10-NN
Figure 10: Visualization of the retrieved datapoints in Bridge tasks using different strategies.
D Ablating Retrieval Strategies
D.1 KNN-based Retrieval Strategy
In our implementation of FLOWRETRIEVAL follow the practice of prior works and retrieve base on threshold of the similarity score, taking the top 8% closest datapoints from prior dataset to any point target data. However, this does not augment the target task trajectory uniformly. As result, when there are pauses in the dataset (which is true in most human- demonstrated datasets if post-processed), we observe that the retrieved data contain large portion of small motion data where the robot arm barely moves (Fig. 10 left).
One intuitive way to solve this issue to retrieve top-k data points for each state in the target task data. See illustration of the two different retrieval strategies in Fig. 9. We visualize samples of the data retrieved by 1O-NN in Bridge tasks in Fig. 10 (right) and see that able to retrieve meaningfully similar data to target task. However, find that, when retrieving similar amount of total data as top 1 1,, this approach surprisingly does not lead to as performant policies as the existing approach, achieving 40% in Square Assembly (-15%), 80% in LIBERO-Can (-10%), and 60% (-8%) Bridge-Microwan
One potential reason is that when we force each datapoint to have at least some datapoints retrieved from prior data, might end up retrieving dissimilar and potentially adversarial data if the & is not carefully selected for each datapoint.
D.2 Retrieving with Pre-trained Representations
We evaluated retrieval in the Square Assembly task with pretrained visual representations to see if off-the-shelf models could be leveraged retrieve motion-similar data. Specifically we take the delta between the features of and 8 and use that to represent motion as proposed in prior work [38]. Fig. 11 shows the detailed analysis data retrieved by Voltron [21], R3M [39], CLIP [40], and Dino-v2 [41]. We see that in general these models focus visual features
than the motion
itself and retrieving with such embeddings cannot bypass adversarial data in this task, with motion Linuuage-aligned models (Voltron and R3M) suffering the most.
Lingfui Bazkground)
Lisefai Data ldifferent backgowss)
Advonianial
Mos-laameiul Dats
5900 redeved dho
Voltoon
RSM
CLIP
ORodd up
Teaooter
Tbxc
Roudh
Toooter
Tlace
Boadh 15
Transfer
Fboo
Reath Pxck
Thocster
Place
Figure 11: Retrieval Analysis for pretrained visual representations.
16


Reviews:
[{"review_id": "lyza6y8pgd", "reviewer": "Reviewer_JhhG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 4, "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": "The authors propose a way of retrieving useful prior data when trying to complete a target task. They accomplish this by selecting the parts of the prior data which represent similar motions to the ones present in the target data. Additionally, in order to preserve the importance of motion during the learning and not only the retrieval, they include an auxiliary task of predicting the flow when training the policy. they finally provide a plethora of experiments both in simulation and real world scenarios.", "summary_of_recommendation": "The paper addresses the limitations of the previous work and is open about the limitations of the presented work. They provide many experiments that are tailored to their work and therefore prove to be interesting.", "technical_quality": 4, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "I would appreciate if the authors could comment on the few things I have mentioned above in the review.\n\nA suggestion that I have is to add an interesting ablation to the Franka pen in cup experiment. More specifically, I would evaluate also FlowRetrieval with both PnP data and wild data. If your retrieval works well it should basically only retrieve PnP and therefore have a very similar performance to FlowRetrieval with only PnP.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "51oI0bvbcU", "reviewer": "Reviewer_GJ3C", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 4, "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": "This work proposes a novel retrieval-based few-shot imitation learning method that leverages motion similarities, specifically using optical flow between observation frames. The key insight of this research is that motion-based imitation retrieval offers a more robust approach compared to pure visual similarity methods, while providing more detailed information than high-level language-based semantic retrieval.", "summary_of_recommendation": "Recommendation: Accept with minor revisions. This paper presents a novel and promising approach to few-shot imitation learning using motion similarity-based retrieval. The work's key strength lies in its innovative use of optical flow to capture motion similarities, bridging the gap between visual and semantic retrieval methods. The comprehensive experiments demonstrate the potential advantages of this approach over existing visual and language-based methods.", "technical_quality": 4, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "The paper is well-written, and the experiments are comprehensive, supporting the authors' proposal that motion similarity retrieval-based imitation is more beneficial for few-shot skill adaptation compared to visual similarity or language semantic similarity-based approaches.\nRegarding the method and its potential, I have the following questions:\n\n1) Motion Similarity and Affordance Learning:\nThe performance of motion similarity retrieval depends on the quality of the learned representations. Given that there are many affordance learning methods (including optical flow-based approaches) trained on large datasets, is there a way to leverage these existing methods to enhance the motion similarity retrieval process?\n\n2) Dataset Size Impact:\nHow does the dataset size affect the imitation retrieval performance? It would be valuable to explore:\na) What is the minimum amount of training data required to learn a robust motion representation?\nb) How does the dataset size correlate with the test performance? Is there a clear trend or diminishing returns past a certain point?\n\n3) Cross-Embodiment Learning:\nHow does this method contribute to cross-embodiment learning scenarios? For instance:\na) If the robot executing the real task is significantly different from those in the dataset, can this method still effectively guide the unseen robot to accomplish the task?\nb) What would be the expected performance of a novel robot with a different embodiment?\nc) Are there any limitations or additional considerations when applying this method to robots with vastly different morphologies?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZzA8tIQQ3m", "reviewer": "Reviewer_Z3WV", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "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": "The authors propose a few-shot manipulation learning method based on flow similarity. The proposed methods augment the small few-shot training set from a much larger prior dataset, by mining trajectories with similar motion patterns. The motion patterns are determined through representation learning with optical flow.", "summary_of_recommendation": "The flow retrieval few-shot learning method is sound but lacks some ablation studies to further justify the design choice. I list the current draft as \"weak accept\".", "technical_quality": 4, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 2, "questions_for_rebuttal": "1. (Corresponds to W1) How much performance contribution comes from the flow auxiliary loss, instead of the proposed flow retrieval strategy?  If you haven't included this (in case I missed anything), I suggest \"random retrieval + flow loss\" as another comparison, instead of just FlowBC (flow loss without retrieval). \n\n2. Is it possible to add some qualitative examples of why flow retrieval performs better than Proprio Retrieval? \n\n3. Do you train one VAE for all environments, or each one separately? \n\n4. Optical flow lacks real groundtruth and is usually very noisy when obtained from off-the-shelf models instead of simulation environments. I feel like the flow-aware representation learning part can be tricky, it would be better if more details are provided. \n\n5. There are some recent works on one-shot imitation learning that share a very similar setting to this work, such as\n\n     - One-shot Imitation Learning via Interaction Warping (CoRL2023) \n     - One-Shot Imitation Learning- A Pose Estimation Perspective (CoRL2023)\n     - One-Shot Imitation Learning with Invariance Matching for Robotic Manipulation (RSS2024)\n   \n   The authors are suggested to discuss them in the related work.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces a new method for few-shot imitation learning in robot manipulation. The main idea is to use optical flow estimation from the few-shot demonstration videos to retrieve similar data from prior datasets to improve the policy learning.

Strengths

- The reviewers agrees on the novelty of the proposed method by using optical flow for data retrieval.

- The paper is well-written with solid experiments to demonstrate the effectiveness of the proposed method.

Weaknesses

The reviewers have also raised several concerns of the paper.

- Scalability of the proposed method on the size of the dataset.
- The effect of the quality of the retrieved data on the policy. 
- Missing citations and some ablation study of the proposed method.

Post-Rebuttal
- The authors have successfully addressed the concerns from the reviewers.
Decision: Accept