SRSA: Skill Retrieval and Adaptation for Robotic Assembly Tasks
Authors: Yijie Guo
Bingjie Tang
Iretiayo Akinola
Dieter Fox
Abhishek Gupta
Yashraj Narang
Venue: ICLR 2025 Spotlight (2025)
Keywords: Robotic Assembly Tasks; Skill Retrieval; Skill Adaptation; Sim-to-real Transfer; Reinforcement Learning Fine-tuning;
Primary area: applications to robotics, autonomy, planning
TL;DR: We introduce SRSA, a novel pipeline that retrieves relevant skills from a pre-existing skill library and adapts them to efficiently solve new robotic assembly tasks.

Abstract: Enabling robots to learn novel tasks in a data-efficient manner is a long-standing challenge. Common strategies involve carefully leveraging prior experiences, especially transition data collected on related tasks. Although much progress has been made for general pick-and-place manipulation, far fewer studies have investigated contact-rich assembly tasks, where precise control is essential. We introduce SRSA} (Skill Retrieval and Skill Adaptation), a novel framework designed to address this problem by utilizing a pre-existing skill library containing policies for diverse assembly tasks. The challenge lies in identifying which skill from the library is most relevant for fine-tuning on a new task. Our key hypothesis is that skills showing higher zero-shot success rates on a new task are better suited for rapid and effective fine-tuning on that task. To this end, we propose to predict the transfer success for all skills in the skill library on a novel task, and then use this prediction to guide the skill retrieval process. We establish a framework that jointly captures features of object geometry, physical dynamics, and expert actions to represent the tasks, allowing us to efficiently learn the transfer success predictor. Extensive experiments demonstrate that SRSA significantly outperforms the leading baseline. When retrieving and fine-tuning skills on unseen tasks, SRSA achieves a 19% relative improvement in success rate, exhibits 2.6x lower standard deviation across random seeds, and requires 2.4x fewer transition samples to reach a satisfactory success rate, compared to the baseline. In a continual learning setup, SRSA efficiently learns policies for new tasks and incorporates them into the skill library, enhancing future policy learning. Furthermore, policies trained with SRSA in simulation achieve a 90% mean success rate when deployed in the real world. Please visit our project webpage https://srsa2024.github.io/.

Published as a conference paper at ICLR 2025
SRSA: SKILL RETRIEVAL AND ADAPTATION FOR ROBOTIC ASSEMBLY TASKS
Yijie Guo Bingjie Tang², Iretiayo Akinola', Dieter F Abhishek Gupta! & Yashraj Narang' NVIDIA Corporation, UUiiersityy of Southern California, University of Washington
Dieter
ABSTRACT
Enabling robots to learn novel tasks in a data-efficient manner is long-standing challenge. Common strategies involve carefully leveraging prior experiences, es- pecially transition data collected on related tasks. Although much progress has been made for general pick-and-place manipulation, far fewer studies have inves- tigated contact-rich assembly tasks, where precise control is essential. We intro- duce SRSA (Skill Retrieval and Skill Adaptation), a novel framework designed to address this problem by utilizing a pre-existing skill library containing policies for diverse assembly tasks. The challenge lies identifying which skill from the library is most relevant for fine-tuning on new task. Our key hypothesis is that skills showing higher zero-shot success rates on a new task are better suited for rapid and effective fine-tuning on that task. To this end, we propose predict the transfer success for all skills in the skill library on a novel task, and then use this prediction to guide the skill retrieval process. We establish framework that jointly captures features of object geometry, physical dynamics, and expert actions to represent the tasks, allowing us to efficiently learn the transfer success predic- tor. Extensive experiments demonstrate that SRSA significantly outperforms the leading baseline. When retrieving and fine-tuning skills on unseen tasks, SRSA achieves 19% relative improvement in success rate, exhibits 2.6x lower standard deviation across random seeds, and requires 2.4x fewer transition samples to reach satisfactory success rate, compared to the baseline. In continual learning setup, SRSA efficiently learns policies for new tasks and incorporates them into the skill library, enhancing future policy learning. Furthermore, policies trained with SRSA in simulation achieve a 90% mean success rate when deployed in the real world. Please visit our project webpage https://srsaa// gi thub. i0/
INTRODUCTION
Humans excel at solving new tasks with few demonstrations or trial-and-error interactions. robot learning, key challenge is to similarly enable robots to learn control policies from sensory input in a data-efficient manner. Achieving data-efficient learning is crucial for deploying robots in diverse real-world environments, such the household and industry. A compelling approach to efficient policy learning is the development of a foundation model or generalist policy that spans multiple tasks, as the model policy can offer long-term efficiency gains by providing strong base for adaptation to novel tasks. Significant advancements have been made in manipulation tasks, particu- larly in visual pre-training (Parisi et al. 2022 Nair et al. 2022). multi-task policy learning (Shridhar et al. 2022 Goyal et al. 2024). and policy generalization (Jag et al. 2022. Ebert et al. 2021).
Despite this progress, efficiently solving new tasks in contact-rich environments, such robotic as- sembly, remains underexplored. Robotic assembly plays a critical role in industries like automotive, aerospace, and electronics, but learning assembly policies is uniquely difficult. These tasks require contact-rich interactions with high levels of precision and accuracy, compounded by the physical complexity of the environments, part variability, and strict reliability standards. Much of the exist- ing research focuses on training specialist (i.e., single-task) policies for individual assembly tasks (Spector & Di Castro 2021 Spector et al. 2022 Tang et al. 2023). Building on the strengths of these specialist approaches, we propose novel method for tackling new assembly tasks. Our approach leverages a skill library a collection of diverse specialist policies and associated infor- mation (such object geometry and task-relevant trajectories) for various assembly tasks. These policies and data, regardless of the training strategies or learning approaches used to develop them, can be harnessed to efficiently solve previously-unse assembly challenges.

---

Published as a conference paper at ICLR 2025
Skill Library
Predicted Transfer Success
Retriever
0.5
10
Adapter
Replay
Buffer
New Task
Retrieved Skill
Figure 1: Overview of SRSA. We address assembly tasks, where the goal is use robot arm to insert diverse plugs (i.e., the white parts) into or onto corresponding sockets (i.e., the green parts). Specifically, we propose predict the transfer success of applying prior skills (i.e., policies) new task, retrieve the skill with the highest predicted success rate, and fine-tune it on the new task. During fine-tuning, we accelerate and stabilize adaptation by incorporating imitation learning of high-rewarding transitions from the agent's own replay buffer.
To utilize prior task experiences, previous work on general bock-and-lland tasks has explored meth- ods such as imitating state-action pairs from expert demonstrations (Du et al. 2023] Lin et al. 2024] 2024) and encoding sub-task skills macro-action choices (Lynch et al. 2020; Pertsch
Kuang et al.
et al., 2021: Nasiriany et al. 2022). Unlike these approaches, which focus on reusing data or sub-
task skills, our approach centers adapting policies from previous tasks to solve novel tasks. These policies encapsulate essential task-solving knowledge in a generative form, making them a valuable starting point for further refinement. Despite having access to library of policies, identifying the most relevant ones for fine-tuning on new tasks is still an open question, and the success of fine- tuning hinges on making the right selection. In this paper, we introduce SRSA (Skill Retrieval and Skill Adaptation), a novel framework designed to retrieve policies for similar tasks and adapt them to new tasks, as illustrated in Fig. The key contributions of this paper are follows:
We propose a skill retrieval method that simultaneously and explicitly learns embeddings for three fundamental components of assembly tasks: part geometry. interaction dynamics, and expert action choices. We subsequently introduce a novel objective that leverages these embeddings to predict transfer success between any source policy and target task, implicitly capturing additional critical factors for policy transfer. This approach enables the effective retrieval of relevant skills, resulting in higher zero-shot transfer success when applied to new tasks.
(1) Skill Retrieval Method:
(2) Skill A Adptattion Method: We propose a skill adaptation method that fine-tunes retrieved skills on new tasks while incorporating a self- ooff-imitation learning method (Oh et al. 2018) to enhance performance and stability during fine-tuning. In simulation-based dense-reward setting explored in the leading assembly baseline (Tang et al., 2024p. SRSA achieves a relative improvement of 19% in success rate with 2.4x faster training and 2.6x lower standard deviation across random seeds. In simulation-basee based, sparse-reward settings without demonstrations or curricula (closely aligning with real-world fine-tuning scenarios), SRSA outperforms the baseline with a relative improvement of 135% in success rate. Furthermore, we demonstrate that policies fine-tuned in simulation can be directly transferred to real-world robots, achieving a 90% average success rate without the need for additional training. This capability of effectively fine-tuning policies in simulation on novel tasks, and transferring these policies to the real world in zero-shot, highlights the potential for deploying high-pertte hegh-performanc solutions in real-world assembly tasks.
(3) Continual Learning with SRSA: Instead of training numerous specialist (i.e., single-task) poli- cies from scratch, we propose gradually expanding a small set of initial skills via retrieval and adaptation to cover a broader range of tasks. This strategy improves sample efficiency by over 80% compared to (ang et al.] 2024) and stays consistently efficient as the skill library and target tasks evolve. Thus, SRSA provides an efficient solution for accumulating a large-scale collection of skills.
RELATED WORK
Robotic Assembly Tasks Robotic assembly is a critical manufacturing process in the automotive, aerospace, electronics, and medical device industries, but adaptive robotic assembly (e.g., robustness to part types, initial part poses, perceptual noise, control error, and environmental perturbations)
is largely unsolved. Research (Beltran (Beltran-Hernane Hernandez et al.] 2020 Luo et al. 2021 Narang et al.
2022 Tang et al. 2023 Zhang et al. 2023] Noseworthy et al.] 2024h on adaptive assembly has
significant growth in recent years. Despite advancements in datasets and real-world benchmarks
for assembling small, realistic parts (Kimble al. 2020
2022 Willis al. 2022
Tian et al.)
2

---

Published as a conference paper at ICLR 2025
2022).
the exploration of policy learning across wide variety of parts remains relatively limited. Many recent efforts in robotic assembly have concentrated on perception (Fu et al. 2022 Wen
et al. 2022) or planning (Tian et al. 2022 2024), rather than learning policies that are robust to disturbances and noise. Additionally, the policy-learning efforts that have addressed the widest range of assemblies have typically been restricted to <30 parts (Spector & Di Castro 202T Spector
et al.] 2022] Zhao et al.] [2022). The largest study, AutoMate (Tang et al., [2024) introduced a diverse dataset featuring 100 assembly tasks with simulation environments and 3D-printable parts,
and explores policy learning across these tasks. However, its approach primarily focuses on learning
specialist (i.e., single-task) policies from scratch without leveraging prior experience or knowledge from related tasks. In contrast, our goal is to solve novel assembly tasks by leveraging skills from previously-solved assembly tasks.
Retrieval-based Policy Learning Many studies have explored techniques for utilizing datasets from other tasks for pretraining, such as visual pretraining (Parisi at.
2022 Nair et 2022 Xiao et al.
2022) and multi-task imitation learning aang et al. 2022 Ebert et al. 2021: Shridhar et al. 2022
Recently, in robotic manipulation, some works have investigated how to selectively incorporate of-
fline data from other tasks during policy learning, i... retrieving prior data according to expert
demonstrations on the target task (Nasiriany et al. 2022 Belkhale et al. 2024) (Shao et al. 2021] Zha al., 2024). For instance, Du et al. (2023) selects pertinent state-action pairs based on visual and action similarity from offiine, unlabeled datasets and jointly trains a policy using a small amount of expert demonstrations and the queried data via imitation learning. Lin al. (222)), on the other hand, emphasizes motion similarity rather than semantic similarity by retrieving state-action pairs based on optical flow representations, followed by few-shot imitation learning with expert demon- strations and the retrieved data. Kuang et al. (222)) takes different approach by extracting a unified affordance representation from diverse data sources and hierarchically retrieving and transferring 2D affordance information based on language instructions to perform zero-shot robotic manipulation. These works primarily study data retrieval for general pick-and-place manipulation tasks. (Zhu tasks, which selects policy candidates
et al. 2024) introduces a policy retriever for pick-and from a memory bank to align closely with the current input, based on the cosine similarity between instruction and observation features. In contrast to these works, focus on challenging contact-rich manipulation tasks, especially investigating transfer success predictor for policy retrieval.
Embedding Learning for Task and Skills Task embedding learning has been extensively explored in nett-reiiforcemment learning and multi-task reinforcement learning problems, where shared knowledge across tasks can significantly enhance learning efficiency for new tasks. Most previous approaches focus on capturing task features related to visual appearance in 2D images or dynamics in transitions (James et al., 2018 Rakelly et al. 2019 Lee et al. 2020p. Contrastive learning is often employed to bring similar tasks closer together in the embedding space while pushing dissim- ilar tasks farther apart (James et al. 2018p. Skill embedding learning, on the other hand, leverages unstructured prior experiences (i.e., temporally extended actions that encapsulate useful behaviors) and repurposes them to solve downstream tasks. Existing methods typically train a high-level policy where the action space consists of the extracted skills (Pertsch et al. [02T]] Nasiriany et al.) 2022
Hausman et al., 2018 Sharma et al. 2019 Lynch et al.) 2020). Although most previous approaches use skills to solve subtasks and combine sequences of skills for long-horizon tasks, we focus on selecting and adapting a single relevant skill for a new task; our tasks of interest are assembly tasks, which are typically short-horizon but difficult train due to exploration challenges and precise con- trol requirements. Additionally, we integrate multiple embedding-leamiin capturing three fundamental components of assembly tasks: part geometry, interaction dynamics, and expert actions. We consolidate these perspectives for more robust task representation.
approaches by jointly
PROBLEM SETUP
In this work, we consider the problem setting of solving a new target task leveraging pre-existing skills from skill library. This library contains policies, each designed to solve a specific previously- encountered task. Our approach is motivated by situations (Rusu et al. 2016 Tirinzoni et al.] 2019 Huang et al.) (02T) where an agent can draw on knowledge from previously-Iearnee policies to adapt quickly to new task at hand. Similar to the multi-task reinforcement learning (RL) formulation [Borsa et al., 2016, Sodhani et al. 2021 Calandriello et al., 2014), we consider a task space T where each task T E T is defined as a Markov decision process (MDP) (S, A,,,,, p). In this formulation, S represents the state space, AA the action space, p(8 | a,) the transition dynamics, r(8e, at) the reward function, 7 E [0. 1) the discount factor, and p the initial state distribution.
3

---

Published as a conference paper at ICLR 2025
01063
01079
01129
/
00211
00319
CO548
00863
01041
(b) 3D-printed assembly parts in real world
00213
00320
00624
C0648
0:025
0:152
01003
00255
00319
00537
00652
C1029
01135
(+) Various assembly tasks simulation
[c) Keyframes assembly real-world deployment
Figure 2: Illustration of assembly tasks in AutoMate and SRSA. (a) Samples of assembly tasks in the AutoMate benchmark. (b) 3D- 3D-printed parts of corresponding real-world assembly tasks in SRSA. (c) Keyframes from video recordings of our real-world deployments of performant policies.
Our study focuses on two-part assembly tasks, as depicted in Fig. 2 Following the setup of Au- toMate (Tang et al. 2024). each environment includes a Franka robot, plug (i.e., part to be inserted), and a socket (i.e., the part that mates with the given plug). In the initial state, we random- ize the robot's joint configuration and socket pose, as well the pose of the plug within the robot's gripper. The goal of each task is to insert plug into its corresponding socket. (See Appendix A
The state space S consists of the robot arm's joint angles and velocities, the end-effector pose and its linear/angular velocities, the current plug pose, and the end-effector goal pose. The action space A consists of incremental pose targets for task-space impedance controller. As described in [Tng et al., al,, (024), although assembly trajectories are infeasible to procedurally generate, disassembly puths can be easily generated, serving as reverse demonstrations that can be used by an RL agent. Specifically, the RL reward function is composed of terms that penalize the distance to the goal, penalize simulation error, reward task difficulty in a curriculum, and imitate the reversed disassembly paths. The assembly tasks all share the same state space S and action space A, but vary in part geometries, transition dynamics p. and initial state distribution p.
Given a target task T T T, we assume access to a prior task set Tprior = (Ti, T2, ST.TT. With policy space II: + A. the skill library contains policies Iprior = [m1, n2. Tn CI that solve each of the prior tasks, respectively. To solve target task, the goal of RL is to find policy π (aa (8e) that produces an action for each state to maximize the expected return. We propose to first retrieve skill (i.e., policy) for the most relevant prior task (Sec. 4.T), and then rapidly and effectively adapt to the target task by fine-tuning the retrieved skill (Sec.
CT.
4 METHOD
4.1 SKILL RETRIEVAL
To effectively retrieve the skills from Uprior that are useful for a new target task T, we require a means to estimate the potential of applying source policy Tare E II prior to the task T. Concretely, we aim to obtain function F 11 x T + R, which takes as input a source policy and a target task, and produces a scalar score measuring how well the source policy can be adapted to the target task.
According to the simulation lemma (Agarwal et al. 2019), the difference in expected value when applying the same policy to different tasks partially depends the difference in their transition dynamics and initial state distributions. We execute a source policy Tere on both target task Torg and its original source task Tree- Let Fsre,trg denote the zero-shot transfer success of sre Terg and T Tere,are its success rate on Teee- These success rates reflect the expected value of Trre on Terg and Tares respectively. Notably, if Tsre,tr is similar Trrc,ares it suggests that the transition dynamics and initial state distributions of two tasks may be closely aligned. Since Tave is already an expert on Tare with a high success rate sre, a high zero-shot transfer success rate T Tsre,trg indicates strong task similarity. Thus, we use the high transfer success rate as a heuristic indicator of similar dynamics and initial state distributions between source and target tasks. Details are Appendix A.2
Subsequently, we hypothesize that fine-tuning a source policy on a target task with similar dynamics will be efficient, as only requires adaptation to small differences in dynamics. Therefore, we propose using zero-shot transfer success as metric to gauge the potential to efficiently adapt source policy to a target task. To identify a source policy with high zero-shot transfer success on a given target task, we propose to learn a function F' to predict zero-shot transfer success for any pair of source policy Tare and target task Terg- The prediction F(Tarey Terg) serves as an indicator of whether Tere is strong candidate to initiate fine-tuning for the target task Tirg. Below, we describe data collection (Sec. (S. featurization (Sec. 4.1.2), training (Sec. 4.1.3) and inference (Sec. 4.1.4p for the transfer success predictor F.
4

---

Published as a conference paper at ICLR 2025
Noin Co d
Disassemblly Tajectory
Disassemblly Trajedoryy
Se 1
5.a: Ses
Sen ... Sot
Courstry
Ad32s
Counerry
Encoder Ec
Encoder Es
Fannder Fa
Eeomatry Jeature
Feature
Eriedding
Embedding
Dacoder
Farword Model
d Iettuce
Transfer Success
Deconntmaciion
Figure 3: Illustration of skill retrieval approach. We decompose skill retrieval into task feature learning(a-c) and transfer success prediction(d). (a) Geometry features are learned from point-cloud input using a PointNet autoencoder. (b) Dynamics features are learned from transition segments using a state-prediction objective. (c) Expert-action features are learned from transition segments using action-recoosstructinn objective. (d) The zero-shot transfer success rate (of applying source policy to target task) is predicted using these task features from the source and target tasks.
4.1.1 DATASET FORMULATION
In order train the prediction function F, construct dataset of tuples (Taeey Tergs ('srettrg) ) We treat any two tasks from the prior task set Tprior as a source-target task pair. For each pair (Trre: Ttrg), we evaluate the source policy øare the target task Trg obtain the zero-shot trans- fer success rate Tare,trg. In cases where multiple distinct policies exist for the same source task, each solving it in different manner, policy-specific features would be necessary to capture nu- ances between different policies. However, in our setting, each policy in the skill library is trained as an expert for specific source task, with a one-to-one mapping between policies and their cor- responding training tasks. Consequently, we use the features of the source task Tare as a proxy to represent the source policy Tsre- This process enables us to collect a training dataset of tuples (Tere: Tirgs (rrre,tg) from the prior skill library.
4.1.2 LEARNING TASK FEATURES
Given the limited number of (Trres Tirg) pairs (specifically, during training, we have n x n pairs for a total of tasks in Tprioo), we need strong featurization of both the source policy and target task for efficient learning of F. For assembly tasks, each task differs along three fundamental axes: part geometry, interaction dynamics, and expert actions that solve the task. Thus, propose framework that jointly captures features of geometry, dynamics, and expert actions to represent the tasks, allowing us efficiently learn the transfer success predictor F (Fig.
When learning geometry features, we assume access to object meshes for both seen and novel tasks. This assumption well-grounded in industry, where CAD models are widely available, allowing us to learn embeddings of 3D geometry. However, learning features for dynamics and expert actions poses unique challenge. For new assembly tasks, assume that expert demonstrations not available, as these are typically tedious to obtain and often suboptimal for assembly tasks. This deficit prevents us from easily computing dynamics or action embeddings.
We draw insight from TTia et al. 2022 Tang et al. 2024h, which noted that, although procedu- rally generating assembly demonstrations for new tasks intractable (narrow marrow-passage problem), disassembly paths can be trivially generated by employing compliant low-level controller to lift an inserted plug from its socket and move to randomized pose. We propose learning features for dynamics and expert actions using these dixassembly paths and hypothesize that such features are useful for predicting transfer success for assembly. We later empirically support this hypothesis.
From each task, randomly sample certain number of points from parts' meshes the point cloud P and also randomly sample the transition segments т from the disassembly trajectories. Using point clouds P or transition sequences we learn encoders Eg, Ep. and EA to capture fea- tures 2G (representing geometry). 2D (representing forward dynamics) and 2 (representing expert actions). We also train decoders De. DD. and D conditioned cloud for geometry, the next state for dynamics, and the action sequence for expert action choices. In Appendix A.4 explain the implementation details for learning these features.
these features to predict the point
4.1.3 LEARNING TRANSFER SUCCESS PREDICTOR
We consolidate task features of source Tuee and target tasks Tery to develop the transfer success predictor F. We feed the sampled point cloud Pare, Ptrg and transition segments Tsres Ttrg from
5

---

Published as a conference paper at ICLR 2025
Tere and Terg: into the pre-trained and frozen encoders Ec, ED and EA. The geometry, dynamies, and expert action features are concatenated together to form task features Zere and Ztrg. We then pass the concatenated task features through an MLP to predict the transfer success trre,trg as illustrated in Fig. dd). Formally, we train the function F to minimize the objective function (Eq.
L= [[ Tirg) terc.tryll Irl = [|MLP Etrg) Trre,trgll I1? ||MLP(EG(Prre) ET(TTrC); (T(re))), (o(rr (b((T) Ea(Tirg)) T rrre,
(1)
4.1.4 INFERRING TRANSFER SUCCESS FOR RETRIEVAL
At test time, we use the well-trained function F to predict the transfer success of applying any prior policy to new task trg Frrse Trrg). As described in Sec. can randomly sample a certain number of points from parts' meshes as point clouds and randomly sample transition segments from the disassembly trajectories. For each pair of source and tar- get tasks, we sample the input data for m times and average the output of F to obtain a more robust prediction of transfer success. Specifically, we sample point clouds P1, P2. Pm and transition segments Ti, T2," ... .mm and then compute the averaged prediction for these samples, i.e. Fergre # Ei= MLP(Ea (ferre.))) Ea (.D Eatterr))) In this way, we infer the predicted transfer success F( Tirg) for any source policies Tarc in the prior skill library Iprior
for each task, we
{11, .T. ... "
Although the well-trained F provides transfer success prediction as an effective guidance for re- trieval, its predictions may not always be perfectly accurate. To mitigate this, we retrieve the top-k source skills ranked by the predictor F. Among these k candidates, we identify the most relevant skill by evaluating their zero-shot transfer success on the target task, and ultimately select the skill with the best transfer performance. This technique is grounded in the same intuition as introduced in Sec. zero-shot transfer success serves as a reliable metric for skill relevance. In the experiments we set k to 5. Details are in Appendix A.12
in Sec. 5.2 2
4.2
SKILL ADAPTATION
As mentioned in Sec. D] our ultimate goal is to solve the new task as an RL problem. The re- trieved skill is used to initialize the policy network no (o) and we subsequently use proximal policy optimization (PPO) (Schulman et al. 2017) to fine-tune the policy on the target task. Our initialization provides strong start for policy learning, as initial trials with the retrieved skills can achieve a reasonable success rate. Inspired by self-imitation learning (Oh et al. 2018), we fully exploit these positive experiences gained during the initial phase of fine-tuning. We maintain re- play buffer D = {(1t, 0g, 0,, R to store the transitions encountered throughout the training, where R2 >x " is the discounted sum of rewards. We prioritize the state-action pairs (st, a.) based on R2 and imitate those pairs with high rewards. The objective function is defined in Eq. 2 a (s),, ( + 3C
(2)
where Csil policy log log (a))) (s)) Caalue 2
(3) 12, and and Vy are the policy and value function (see the details in Appendix A.3
(++ = max(., (),
As training progresses, the agent collects higher rewards on the target task, leading to an expanding replay buffer filled with improved experiences. As analyzed in (Tagg 2020). this self-imitation mechanism accelerates the agent's convergence to encountered high-reward behavior, even though it may introduce some bias into the policy. In our case, the behavior derived from the retrieved skill is advantageous for the target task. We find that self-imitation learning significantly enhances and stabilizes policy fine-tuning, proving especially beneficial in sparse-ree reward scenarios.
4.3 CONTINUAL LEARNING WITH SKILL-LLIBRARY EXPANSION
Continual learning investigates learning various tasks in sequential manner. The primary objec- tive is to overcome the forgetting of previously learned tasks and to leverage earlier knowledge to achieve better performance and/or faster convergence on incoming tasks (Ring 1994 Xu & Zhu
2018 Abel et al,, 2024 We integrate SRSA in the continual-leamin the skill library. Specifically, we begin with an initial skill library Iprror corresponding to prior apply SRSA to
setup and gradually expand
tasks T Tprior. When faced with new batch of tasks To (TI, Ty. ... Tkk.
retrieve and fine-tune policies for each new task T, The learned policies are then incorporated as Tprior- Tprior [T:]; IT Uprior- II pprio {T.) This approach allows us to efficiently tackle new
6

---

Published as a conference paper at ICLR 2025
tasks by leveraging the skill library and simultaneously prevent the forgetting of all learned tasks by maintaining and expanding the skill library. See Appendix A.. for the algorithm pseudocode.
EXPERIMENTS
We design experiments to answer questions: (1) Can SRSA retrieve source policies that achieve a better zero-shot transfer success rate on test tasks compared bascline retrieval approaches? (2) Can SRSA policy fine-tuning improve learning performance, stability, and efficiency on test tasks? SRSA nigh-performing policies from simulation be deployed in zero-shot to the real world? (4) Can SRSA be applied in the continual-learning scenario to improve learning efficiency by gradually expanding skill library? We investigate these questions on the AutoMate benchmark (Tang et al., 2024). which consists of 100 two-part assembly tasks with diverse parts, enabling us to study challenging contact-rich assembly tasks in simulation and the real world.
(3) After fine-tuning,
5.1 SKILL RETRIEVAL
AutoMate provides meshes and disassembly trajectories for each task. We use these data to learn the task embedding for retrieval. We compare the following retrieval strategies. Signature: retrieve the task with the closest path signature (Barcelos et al. 2024 Chen 1958 Kidger et al. 2019). which represents the disassembly trajectories retrieve the task with the closest VAE embedding of state-action pairs on disassembly trajectories.
collection of path integrals (Tang et 2024). Behavior:
Forward: retrieve the task with the closest latent vector for the transition sequence on the disassem-
bly trajectories, where the latent vector was trained to predict the forward dynamics. Geometry: retrieve the task with the closest PointNet (Qi et al. 2017| Wang et al. 2023) encoding for the point clouds of the assembly assets. SRSA: retrieve the source task with the highest prediction of transfer success on the target task. Implementation details can be found in Appendix A.4
HWW.
Trarater Succens
Trander Success
SRAA
%
Overal
00023
00028
02243
73085
02079
08022
01302
011.0
03222
Figure 4: Zero-shot transfer success of retrieved skills when applied to test tasks. For each test task, we retrieve policy from the prior skill library using 5 different approaches (4 baselines and SRSA). If the approach involves training neural networks, represents the best transfer success rate on the target task among all source policies. Left: Mcan and standard deviation of transfer success rate, averaged over 10 test tasks with 3 seeds each. Right: Mean and standard deviation of success rate for each test task, averaged over 3 seeds. Overall, SRSA substantially outperforms baselines.
train on 3 random seeds. Optimal
Given the 100 tasks in the AutoMate benchmark, we split the task set into 90 prior tasks (to build the skill library) and 10 test tasks (as the new tasks to solve). For both SRSA and baseline methods, we train the retrieval model with three random seeds and report the average and standard deviation of transfer success across these seeds. Fig. 4 shows the result the set of test tasks. SRSA performs best or second-best on all test tasks, except for one very challenging assembly, where all methods perform poorly (01029). In Appendix A.5 we show additional comparisons for other splits of prior and test task sets. Overall, SRSA retrieves source policies that obtain around 20% higher success rates on the test tasks, compared to baselines.
(..5
5.2 SKILL ADAPTATION
In this section, rather than investigating zero-shot transfer, we study policy learning on test tasks. We compare AutoMate learning specialist policies from scratch (Tang et al. 2024) and SRSA fine- tuning the retrieved specialist policy with self-imitation learning. Details are Appendix A.4 We consider both the dense-reward setting (identical to AutoMate), which includes a reward term imitating disassembly demonstrations and curriculum, and the sparse-reward setting, which only provides a nonzero reward for task success. The rewrse setting is designed to emulate real- world RL fine-tuning, where dense-reward information is much more challenging to acquire
7

---

Published as a conference paper at ICLR 2025
02036
01041
01033
Sucses
Nanne nomurd) Ausodisse
Epoch
Epoch
Epoch
CL.079
Sucees
Epoch
Fpooch
Epoch
01125
01129
01186
Succes
Figure 5: Learning curves on test tasks. The rraxis represents training epochs, where each epoch consists of 128 environment steps over 256 parallel environments. The y-axis represents success rate. The solid line shows the mean success rate over 5 runs with different random seeds, and the shaded area denotes the standard deviation.
Fig. 5 shows the learning curves in the set of test tasks. In the dense-reward setting, SRSA achieves strong per- formance with fewer training epochs than AutoMate. In the sparse-rewand setting, AutoMate struggles to achieve reasonable success rate, whereas SRSA ben- efits from the retrieved skill initialization and self- imitation learning, enabling it to reach higher perfor- mance. Additionally, in both settings, the learning curves of AutoMate exhibit instability with fluctuat- ing success rates as the training goes on. Tab. g and Tab. in Appendix A.5 summarize the mean and stan- dard deviation of the success rate at the last epoch of training. These values are averaged across random seeds for each test task. In the dense-rewart setting, SRSA reaches an average success rate of 82.6% on 10 test tasks, outperforming AutoMate (69.4%), cor- responding to relative improvement of 19% in per- formance. Moreover, SRSA shows greater stability, as AutoMate exhibits a 2.6x higher standard deviation. In the sparse-reward setting, SRSA delivers notable 135% relative improvement in the average success rate compared to the baseline. Fig. 6 demon- strates the number of training epochs required to reach the desired success rate in the dense-reward setting. Averaged over 10 test tasks and 5 random seeds, SRSA requires far fewer training samples, i..,, at least 2.4 times fewer training epochs, to achieve arbitrary success threshold.
SRSA AutoMalee
Number of Trakiin Eqochss Roquiad
Success Rate
Figure 6: Sample efficiency on test set. To achieve desired success rate (0.70, 0.75, 0.80, 0.85, or 0.90), we identify how many training epochs required for each run. We illustrate the mean and standard deviation of required epochs across 5 runs with the points and error bars in the figure, averaged over 10 test tasks.
5.3 REAL-WORLD DEPLOYMENT
8

---

Published as a conference paper at ICLR 2025
We now deploy the trained specialist poli- cies in the real world. As in (Tang et al.) 2024). we place the robot in lead-through mode (a.k.a., manual guide mode), grasp plug, guide it into the socket, and record the pose as target pose. We then program- matically lift the plug until free from con- fact, apply perturbations to the position and
Asset ID
000099
01053
01079
01129
01136
Overall
AutoMate SRSA
7/10
1/10
7/10
4/10
8/10
54%
9/10
8/10
8/10
10/10
10/10
90%
Real-world evaluation. We take the best checkpoint of policies across 5 runs within 500 epochs and report the success rate over 10 trials for each task. rotation of the end effector, and deploy a policy to assemble the plug into the socket. Such conditions emulate the perceptual noise and control error that are experienced in full robotic assembly pipelines.
Figure 7:
In Tab. 7 we take the best checkpoint 500 training epochs in simulation and record its performance when deployed in the real world. In this relatively brief training time, SRSA reaches higher success rates than the baseline on real-world assembly tasks. We show keyframes of roal-world deployments in Fig. c).. For videos, please refer the project website
github. io/
5.4 CONTINUAL LEARNING
U200
Nimber of Towring Epecte Daggine
Auto Mote
400
Number of Epocha
40
6
40
SRSA
AutoNate
Success
Batch
(b)
Figure 8: (a) Overall sample efficiency. We report the number of training epochs required to reach desired success rates (0.5, 0.6, 0.7, 0.8). We calculate the mean and standard deviation of the required training epochs over 5 runs, and report the average across 90 tasks. (b) Sample efficiency in batches. We sequentially introduce 9 batches of new tasks for policy learning, with each batch containing 10 new tasks. For each batch, we show the mean and standard deviation of training epochs required to reach success rate of 0.8. SRSA consistently requires fewer training epochs.
We study the continual-leaminng setting to obtain policies for each of the 100 Mate tasks. Rather than training 100 policies from scratch in parallel, we start from skill library with 10 tasks, and obtain 10 new policies for 10 new tasks utilizing the skill library. For each new task, we fine-tune the retrieved policy over 5 runs with different random seeds. We pick the best checkpoint with the highest success rate over 5 runs as the specialist policy for this new task. We repeat this process for 9 iterations, eventually covering the entire AutoMate benchmark. Essentially, we have a skill library that is gradually expanded with an increasing number of specialist policies.
In Fig. R we compare the sample efficiency of SRSA and AutoMate when learning specialist policies for the 90 tasks outside the initial skill library. We consider different desired success rates and report the number of training epochs required to reach each success rate. Overall, SRSA requires fewer training epochs to reach the desired success rate, demonstrating an 84% relative improvement in sample efficiency average (Fig. 8 ())) For each batch of new tasks, SRSA is more efficient than the baseline, regardless of the skill library and test tasks (Fig. ( (b)) In Fig. [4| we show the success rates for the best checkpoints encountered in runs for each task. SRSA achieves an average success rate of 79% compared to AutoMate's ' 70% across 100 tasks, while also exhibiting better training efficiency. In Appendix A.55 we present learning results for another ordering of batches of tasks, showing that the advantage of SRSA is insensitive to the order of encountering new tasks.
6 ABLATION STUDY
Effect of Skill Retrieval To verify the effect of skill retrieval, we conduct skill adaptation with retrieved skills using only geometry embedding, i.e., the second-best skill retrieval approach eval- uated in Fig. S Fig. shows the performance of policy fine-tuning for both SRSA and the geometry- based skill retrieval (SRSA-Geom). One can observe that retrieving worse skill hinders learning
9

---

Published as a conference paper at ICLR 2025
01029
01053
01078
Sucoes
%
Succes
Gean
Epocc
Epoch
Epoac
Figure 9: Comparison across variants of SRSA. For each method, we fine-tune the policy with different random seeds. The learning curves show the mean and standard deviation of success rate over these seeds. We show learning curves for more tasks in the Appendix A.5
efficiency, starting from a lower success rate and requiring more training epochs to reach high per- formance. Our retrieval approach SRSA improves adaptation efficiency over SRSA-Geom.
Effeet of Self-imitation Learning To demonstrate the benefits of self-imitation learning (SIL) in policy fine-tuning, we compare SRSA to the variant without this component (SRSA-noSIL). In SRSA outperforms the variant in terms of learning stability. In partieular, SRSA-noSIL, suffers from more fluctuations during fine-tuning and larger standard deviation of success rate (shaded area) across runs with different seeds.
Fig. 2
Effect of Generalist Policy We analyze whether fine-tuning generalist policy outperforms fine- tuning a selected specialist policy. For policy initialization, we use the generalist policy for 20 training tasks from (Tang et al. 2024). Although it does not cover numerous tasks, it is the strongest generalist policy reported to date that can solve a diverse set of assembly tasks with an average success rate greater than 80%. Fig. g shows the learning curves of fine-tuning the generalist policy on unseen tasks (SRSA-Gen). We observe that SRSA-Gen provides weaker initialization compared to SRSA. likely because the generalist policy's knowledge from the training tasks is less specialized than the skills retrieved by SRSA. Furthermore, adaptation is less efficient, possibly due to the larger neural network in the generalist policy, which requires more fine-tuning to adapt to new tasks, As result, its asymptotic performance is also lower than that of SRSA.
7 CONCLUSION
Summary: In this work, we propose a pipeline to retrieve and adapt specialist policies to solve new assembly tasks. To learn a retrieval model, we jointly learn features from geometry, dynam- ies, and expert actions to represent tasks, and predict transfer success to implicitly capture other transfer-related factors from tasks. By combining skill retrieval with policy fine-tuning and self- imitation learning, our method efficiently learns nigh-performance simulation-based policies. We demonstrate that these policies are transferable to real-world robots. Additionally, we demonstrate that our approach can continuously expand a skill library through efficient learning of various skills.
First, although we train policies for all assembly tasks in a leading benchmark (Tang al. 2024). do not address assemblies requiring rotational helical motion (e.g., nut-and-bolt assembly). Second, we primarily concentrate on learning specialist (i.e., single-task) policies; fu- ture work could explore learning generalist (i.e., multi-task) policies, and, furthermore, incorporate knowledge from both specialist and generalist policies to solve novel tasks with even greater effi- ciency. Third, although our real-world success rates outperform the state-of-the-art in sim-to-real transfer for our examined tasks, they still fall short of the 99+% success rates required for industry- level deployment. We believe that RL fine-tuning directly in real-world settings could help bridge the sim-to-real gap and further improve success rates.
Limitations:
Future Extensions: How to utilize existing policies for new tasks (rather than training from scratch) is an open and general question in robotics. This question is relevant not just for insertion tasks, but also for general pick-and-place tasks, dexterous manipulation tasks, advanced assembly tasks, etc. Most robotics tasks are governed by geometry, dynamics, and behavior/action. We believe that our ideas of learning task features and predicting zero-shot transfer success for policy transfer generalize other domains. For instance, in tool-use tasks, the skill of using scissors may be beneficial for learning to operate pliers due to their similar shape and operating mechanism. We leave it as future work to extend SRSA these additional robotics applications.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
David Abel, André Barreto, Benjamin Van Roy, Doina Precup. Hado P van Hasselt, and Satinder Singh. A definition of continual reinforcement learning. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Alckh Agarwal, Nan Jiang, Sham M Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep, 32:96, 2019.
Lucas Barcelos, Tin Lai, Rafael Oliveira, Paulo Borges, and Fabio Ramos. Path signatures for diversity in probabilistic trajectory optimisation. The International Journal of Robotics Research, 43(11):1993-1710, 2024.
Suneel Belkhale, Tianli Ding, Ted Xiao, Pierre Sermanet, Quon Vuong. Jonathan Tompson, Yevgen Chebotar, Debidatta Dwibedi, and Dorsa Sadigh. Rt-h: Action hierarchies using language. arXiv preprint arXiv:2403.01823 2024.
Cristian C. Beltran-Hernandez, Damien Petit, Ixchel G. Ramirez-Alpizar, and Kensuke Harada.
Variable compliance control for robotic peg-in-hole assembly: A degp-reintrreementt-aarinng approach Applied Sciences, 2020.
Diana Borsa, Thore Graepel, and John Shawe-Taylor, Learning shared representations in multi-task reinforcement learning. arXiv preprint arXiv: arXiv: 1603.0204, 2016.
Daniele Calandriello, Alessandro Lazaric, and Marcello Restelli. Sparse multi-task reinforcement learning. Advances in neural information processing systems, 27, 2014.
Kuo-Tsai Chen. Integration of paths-a faithful representation of paths by noncommutative formal power series. Transactions of the American Mathematical Society, 89(2):395 407, 1958.
Maximilian Du, Suraj Nair, Dorsa Sadigh, and Chelsea Finn. Behavior retrieval: Few-shot imitation learning by querying unlabeled datasets. arXiv preprint arXii:2304.08777 2023.
Frederik Ebert, Yanlai Yang. Karl Schmeckpeper, Bernadette Bucher, Georgios Georgakis, Kostas Daniilidis, Chelsea Finn, and Sergey Levine. Bridge data: Boosting generalization of robotic skills with cross-domain datasets. arXiv preprint arXiv:2 13396, 2021.
Bowen Fu, Sek Kun Leong, Xiaocong Lian, and Xiangyang Ji. 6d robotic assembly based on rgb- only object pose estimation. In 2022 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 4736-4742. IEEE, 2022.
Ankit Goyal, Valts Blukis, Jie Xu, Yijie Guo, Yu-Wei Chao, and Dieter Fox. Rvt-2: Learning precise manipulation from few demonstrations, arXiv preprint arXiv:2406. 205.::2000008455, 2024.
Yijie Guo, Qiucheng Wu, and Honglak Lee. Learning action translator for meta reinforcement learning sparse-reward tasks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 7792,6800, 2022.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning embedding space for transferable robot skills. International Conference on Learn- ing Representations, 2018
Biwei Huang. Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. arXiv preprint arXiv:2107.02729, 2021.
Stephen James, Michael Bloesch, and Andrew J Davison. Task-embedded control networks for few-shot imitation learning. In Conference on robot learning, pp. 783-795. PMLR, 2018.
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Confer- ence on Robot Learning, pp. 991-1002, PMLR, 2022.
Patrick Kidger, Patric Bonnier, mmanol Perez Arribas, Cristopher Salvi, and Terry Lyons. Deep signature transforms. Advances in neural information processing systems, 32, 2019.

---

Published as a conference paper at ICLR 2025
Kenneth Kimble, Karl Van Wyk. Joe Falco, Elena Messina, Yu Sun, Mizuho Shibata, Wataru Ue- mura, and Yasuyoshi Yokokohji. Benchmarking protocols for evaluating small parts robotic as- sembly systems. IEEE robotics and automation letters, 5(2):883-888)) 2020.
Kenneth Kimble, Justin Albrecht, Megan Zimmerman, and Joe Falco. Performance measures to benchmark the grasping, manipulation, and assembly of deformable objects typical to manufac- turing applications. Frontiers Robotics and AI, 9:999348, 2022.
Yuxuan Kuang, Junjie Ye, Haoran Geng, Jiageng Mao, Congyue Deng, Leonidas Guibas, He Wang, and Yue Wang. Ram: Retrieval-based affordance transfer for generalizable zero-shot robotic manipulation. arXiv preprint arXiv:2407.046899 2024.
Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynam- ics model for generalization in model- based reinforcement learning. In International Conference on Machine Learning. pp. 5757-5766 PMLR, 2020.
Li-Heng Lin, Yuchen Cui, Amber Xie, Tianyu Hua, and Dorsa Sadigh. Flowretrieval: Flow- lowwguidedd data retrieval for few-shot imitation learning, arXiv preprint arXiv:2408. 16944, 2024.
Jianlan Luo, Oleg Sushkov, Rugile Pevceviciute, Wenzhao Lian, Chang Su, Mel Vecerik, Ning Ye, Stefan Schaal, and Jon Scholz. Robust multi-modal policies for industrial assembly via rein- forcement learning and demonstrations: A large-scale study arXiv preprint 1155 2021.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on robot learning, pp. 1113- 1132. PMLR. 2020.
Tongzhou Mu, Zhan Ling, Fanbo Xiang. Derek Yang, Xuanlin Li, Stone Tao, Zhiao Huang. Zhi- wei Jia, and Hao Su. Maniskill: Generalizable manipulation skill benchmark with large-scale demonstrations. arXiv preprint vrXiv:2107.14883 2021.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A univer- sal visual representation for robot manipulation. arXiv preprint arXiv:2203. 126011 2022.
Yashraj Narang, Kier Storey, Iretiayo Akinola, Miles Macklin, Philipp Reist, Lukasz Wawrzyniak, Yunrong Guo, Adam Moravanszky, Gavriel State, Michelle Lu, et al. Factory: Fast contact for robotic assembly. arXiv preprint arXiv:2205.03332 2022.
Soroush Nasiriany, Tian Gao, Ajay Mandlekar, and Yuke Zhu. Learning and retrieval from prior data for bkill-based imitation learning. arXiv preprint arXiv:2210.1 1435,, 2022.
Michael Noseworthy, Bingjic Tang, Bowen Wen, Ankur Handa, Nicholas Roy, Dieter Fox. Fabio Ramos, Yashraj Narang, and Iretiayo Akinola. Forge: Force-guided exploration for robust contact-rich manipulation under uncertainty. arXiv preprint arXiv:2408. 04588, 2024.
Junhyuk Oh, Yijie Guo, Satinder Singh, and Honglak Lee. Self-imitation learning. In International conference on machine learning, pp. 3878-3887. PMLR, 2018
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. In international conference on machine learning, pp. 17359-17371. PMLR, 2022.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on robot learning, pp. 188-204, PMLR, 2021.
Charles R Qi. Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 652-660, 2017.
12

---

Published as a conference paper at ICLR 2025
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy retf-reiinorrement learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
Mark Bishop Ring. Continual learning in reinforcement environments. The University of Texas at Austin, 1994.
Andrei A Rusu, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. arXiv preprint arXiv:1606.0400 04671, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06777 2017.
Lin Shao, Toki Migimatsu, Qiang Zhang, Karen Yang, and Jeannette Bohg. Concept2robot: Learn- ing manipulation concepts from instructions and human demonstrations. The International Jour- nal of Robotics Research,
1434, 2021.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. arXiv preprint 197.00997 2019.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Cliport: What and where pathways for robotic manipulation. In Conference on robot learning, pp. 894-906. PMLR, 2022.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pp. 9767-9779. PMIR, 2021.
Oren Spector and Dotan Di Castro. Insertionnet-a scalable solution for insertion. IEEE Robotics and Automation Letters, 6(3):5509-5516, 2021.
Oren Spector, Vladimir Tchuiev, and Dotan Di Castro. Insertionnet 2.0: Minimal contact multi- step insertion using multimodal multiview sensory input. In 2022 International Conference on Robotics and Automation (ICRA), pp. 6330-6336. IEEE, 2022.
Bingjie Tang, Michael A Lin, Iretiayo Akinola, Ankur Handa, Gaurav S Sukhatme, Fabio Ramos, Dieter Fox, and Yashraj Narang. Industreal: Transferring contact-rich assembly tasks from simu- lation to reality. arXiv preprint arXiv:2305. 17,10, 2023.
Bingjie Tang, Iretiayo Akinola, Jie Xu, Bowen Wen, Ankur Handa, Karl Van Wyk, Dieter Fox, Gaurav S Sukhatme, Fabio Ramos, and Yashraj Narang. Automate: Specialist and generalist assembly policies over diverse geometries. arXiv preprint arXiv:2407.08028 2024.
Yunhao Tang. Self-imitation learning via generalized lower bound q-learning. Advances in neural information processing systems, 33: -10065---17,,----- 2020.
Yunsheng Tian, Jie Xu, Yichen Li, Jieliang Luo, Shinjiro Sueda, Hui Li, Karl DD Willis, and Woj- ciech Matusik. Assemble them all: Physics- based planning for generalizable assembly by disas- sembly. ACM Transactions Graphics (TOG), 41(6):1-11, 2022.
Yunsheng Tian, Karl DD Willis, Bassel Al Omari, Jieliang Luo, Pingchuan Ma, Yichen Li, Farhad Javid, Edward Gu, Joshua Jacob, Shinjiro Sueda, et al. Asap: automated sequence planning for complex robotic assembly with physical feasibility. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 1380-4386 IEFF, 2024.
Andrea Tirinzoni, Mattia Salvini, and Marcello Restelli. Transfer of samples in policy search via multiple importance sampling. In International Conference on Machine Learning, pp. 6264 6274. PMLR. 2019.
Weikang Wan, Haoran Geng, Yun Liu, Zikang Shan, Yaodong Yang, Li Yi, and He Wang. Unidex- grasp++: Improving dexterous grasping policy learning via geometry-aware curriculum and iter- ative generalist-speccialis learning. In Proceedings of the IEÉE/CVF International Conference on Computer Vision, pp. 3891-3902, 2023.
13

---

Published as a conference paper at ICLR 2025
Ruicheng Wang, Jialiang Zhang, Jiayi Chen, Yinzhen Xu, Puhao Li, Tengyu Liu, and He Wang. Dex- graspnet: A scale robotic dexterous grasp dataset for general objects based on simulation. In 2023 IEEE International Conference Robotics and Automation (ICRA), pp. 11359-11366. IEEE, 2023.
Bowen Wen, Wenzhao Lian, Kostas Bekris, and Stefan Schaal. You only demonstrate once: Category-level manipulation from single visual demonstration. arXiv preprint arXiv:2201.12716, 2022.
Karl DD Willis, Pradeep Kumar Jayaraman, Hang Chu, Yunsheng Tian, Yifei Li, Daniele Grandi, Aditya Sanghi, Linh Tran, Joseph G Lambourne, Armando Solar-Lezama, et al. Joinable: Learn- ing bottom-up up assembly of parametric cad joints. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 15849-15860, 2022.
Tete Xiao, Iliaa Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173. 2022.
Ju Xu and Zhanxing Zhu. Reinforced continual learning. Advances in neural information processing systems, 31, 2018.
Kuo-Hao Zeng, Zichen Zhang, Kiana Ehsani, Rose Hendrix, Jordi Salvador, Alvaro Herrasti, Ross Girshick, Aniruddha Kembhavi, and Luca Weihs. Poliformer: Scaling on-policy rl with trans- formers results in masterful navigators. arXiv preprint urXiv:2406.2008. 2024.
Lihan Zha, Yuchen Cui, Li-Heng Lin, Minae Kwon, Montserrat Gonzalez Arenas, Andy Zeng, Fei Xia, and Dorsa Sadigh. Distilling and retrieving generalizable knowledge for robot manipulation via language corrections. 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 15172-15179. IEEE, 2024.
Xiang Zhang, Changhao Wang, Lingfeng Sun, Zheng Wu, Xinghao Zhu, and Masayoshi Tomizuka. Efficient sim-to-real transfer of contact-rich manipulation skills with online admittance residual learning. In Conference Robot Learning, pp. 1621-1639. PMLR, 2023.
Tony 7. Zhao, Jianlan Luo, Oleg Sushkov, Rugile Pevceviciute, Nicolas Heess, Jon Scholz, Stefan Schaal, and Sergey Levine. Offline meta-reinforcemmnt learning for industrial insertion. In 2022 international conference on robotics and automation (ICRA), pp. 6386-6393. IEEE, 2022.
Yichen Zhu, Zhicai Ou, Xiaofeng Mou, and Jian Tang. detrieval-augmentee embodied agents. In Proceedings of the IEFE/CVF Conference on Computer Vision and Pattern Recognition, pp- 119985-17995, 2024.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 ROBOT SETUP
Plug
Socket
Bench Vise
Figure 10: Real-world experimental setup. A Franka Panda robot and a bench vise are mounted tabletop. At the beginning of each episode, 3D-printed plug is grasped by the robot gripper and 3D-printed socket is haphazardly placed and clamped in the bench vise. The task is to control the robot arm to fully insert the plug into the socket.
A.2 MOTIVATION WITH THEORETICAL PERSPECTIVE
Transferring knowledge from a source task to a target task can improve training efficiency and asymptotic performance. Consider source task T, and target task T,, which MDPs that share state space S, action space A, and reward function r, but have distinct transition functions pt. and initial state distributions P.. Pj. To measure the transferability of policy, we apply the same policy on both tasks and examine the difference their expected values. Here we note that the value difference partially depends the difference in their transition functions pt. Pj and initial state distributions Ps. Pj (Proposition
Proposition 1. Let T; AS,A, Pi,T, pi} and T; (S,A. T, T.Pi be two MDPs in the task space T. Applying a policy क on T, and TT, we have a function to describe the value difference:
V (p: T,) V
V* (p5s T3) f(p: Ps. pz)
15

---

Published as a conference paper at ICLR 2025
Proof.
(p,, V (p)(.T))
= E Esppe Ba (s, E. (
E Esppi E B (s, T;)
+Esopi ( (s, - Esmp,() (1)
= E E Ba (Q" (s as,
+Bumpe() V (s, [2) (x) (s,
(2= (s, (s, Cia- (2)
For the Q-value difference (" refer to the simulation lemma in (Agarwal et al. 2019) :
= ((-->P"(T))) p
where PT((T) denotes the transition matrix on state-action pairs induced by the policy a on the task Tj. i.e., P ((aa),ss (r) ( (T₃) Ps(''ss, a)x(a'|s')) (s)).
Consequently, Q" (s, T3 Q" (s, TT) (" T3) can be expressed as a function of (p; P,)-
(s. a, T;) is the (s, a) item in the matrix (T.) Q" (Ty), and
Thus, the value difference V (p;, T,) V (pg. T3) partially depends on (p: P,) and (p. -P))-
Assume the reward function r is a sparse, binary term indicating task success at the end of an episode. The success rate of applying a policy (=(p) (rpp--prrrprr 12 000 r're]. Here, our success rate V* (pp.T;) will naturally be t
task T can be represented as
high, as the source policy π already an expert policy for the source task Tý. When the success rate of applying the source policy target task T, is also high (i.e., V (pt, Tū) is close to V (p), TT)), then Proposition implies that the transition functions p; and Pý. as well as the initial state distri- |
butions P; and Pj. might be similar. Consequently, if source policy transfer success on a target task, the target task might have a similar transition function and initial state distribution as the source task. Hence, hypothesize that fine-tuning the source policy on the target task will be efficient.
achieve high zero-shot
However, it is important to note that achieving similarly high success rate two tasks with single policy does not necessarily indicate similar dynamics between the tasks. Proposition | establishes that similar dynamics and initial state distributions lead to similar expected values for given policy, but the reverse is not guaranteed. We use the high transfer success rate as heuristic indicator of similar dynamics, serving as intuitive motivation rather than strict theoretical justification.
16

---

Published as a conference paper at ICLR 2025
A.3 METHOD
Algorithm Policy Fine-tuning with Self-imitation Learning
Initialize parameters for policy He and parameters & for value function V4 from retrieved skill Initialize replay buffer +-
Initialize episode buffer & 0 for each iteration do # Collect training samples for each step do
Execute an action a, ~ Tr (ag(()))) in the environment and transit the next state 8;41 Store transition & + U (st, Tt)) end for
if +1+1 1 is terminal then
# Update replay buffer
"rx for all f in D + D {[st, at, R&)} } for all in &
Compute returns Rs = DB 7th "
end if
# Update parameter using PPO objective with samples in
0 + 0 PPoo
(Schulman et al. 2017
+ ab V LPPo
# Perform self-imitation learning
Sample a mini-batch {(s, a, R}} from D according to advantages
0 - 0 Vo
(Equation
1 - 3 v/ Lal
Clear episode buffer 5-0
end for
Algorithm 2 Continual Learning with Skill Library Expansion
Require:
Prior tasks Tprior {Ti, T2, T Skill library Uprir { , Rn) I: while given new batch of tasks T (T7, "E..T..... ... TE,
do
each task T? do 3: Retrieve a policy Tave from the skill library Npriar
Fine-tune Ware to get policy TE& solving the task T7
5:
Expand the skill library, Tprior T Tprior U (T7 Ipprio Iprior U } 6: end for
end while
17

---

Published as a conference paper at ICLR 2025
A.4 IMPLEMENTATION DETAILS
A.4.1 TASK FEATURE LEARNING IN SRSA
Geometry Features As shown in Fig. a), employ PointNet-based QQi et al. 2017) autoen- coder EG and DG to minimize the difference between input point cloud P and reconstructed point cloud (P)()))) The autoencoder is trained using point clouds of parts from all tasks.
We follow the implementation details outlined in (Tagg et al. 2024). In a large set of meshes M for various assembly parts, each mesh m, E M consists of (Vi, P.), where V denotes the vertices and E represents the (undirected) edges. During each training iteration, we sample a batch of meshes B C .. For each m, B, generate point cloud P. from the mesh, with each point located on the surface of m,. The point cloud P. is passed through PointNet encoder (Qi et al. 2017) based on the implementation from (Mu al.] 202T) produce a latent vector. The latent vector 2G, is subsequently fed into fully-convolutional decoder, following the implementation from (Wan et al., 2023) to produce the reconstructed point cloud Qi DE(Ea((
The network is trained to minimize reconstruction loss, defined here as the Chamfer distance be- tween P; and Q;:
LCD
1 [
min |Ip 2112 + || Σ PEP:
Σ
min lp 411 PEP.
PEQ.
gEQ.
Across 100 two-parts assembly tasks, we utilize a total of 200 meshes for the plug and socket components with [M| = 200. Each sampled point cloud P. contains 2000 points and the dimension of learned embedding is |ac.u| = 32. The autoencoder is trained for a total of 23,000 epochs, using batch size of 64 and a learning rate of 0.001.
To represent the features of one task, we gather the geometry features for the meshes of the plug, socket, and their assembled state, where the plug is fully inserted into the socket.
Dynamics Features We build upon prior work in context-based meta-RI. (Rakelly et al. 2019
Lee et al. 2020) to utilize a context encoder ED that produces a latent vector from transition seg- ments T--1 (st at-k, 9t-h+1; 0c-h-1, St-1, (-1) }. as shown in Fig. Bb). We sample the transition segments from disassembly trajectories, compute the latent vector ED(T:-2), 1). and feed the latent vector from transition segments to a forward dynamics model DD across all tasks. For any transition samples from any task, the forward dynamics model is trained to predict the next state 141 DD (ED((E(TTT))) St, ax) be close to the ground-truth next state 81+11
As described in (Tang et al. 2024). for each task, we generate disassembly paths by initializing the robot hand to grasp the plug in the assembled state, where the plug is fully inserted in the socket. Using a low-level controller, we lift the plug from the socket and move it to a randomized pose. We repeat this process until collecting 100 successful disassembly trajectories. We store the state and action at each timestep in the disassembly trajectories. Each task has a total of 100 disassembly trajectories, with each trajectory spanning 128 timesteps.
We sample the transition segment T--1 = {s:=h, at-h: 10--h-1.0-A-11 timesteps. The context encoder is modeled as a multi-layer perceptron (MIP) with hidden layers of sizes (256, 128, 64), producing 32-dimensional vector ZD,t (TD(TL-). Then, the forward dynamics model Dp receives the context vector as an additional input, where the input consists of a concatenation of state 9;, action a, and context vector DD, The forward dynamics model comprises four fully-connected layers of sizes (200, 200, 200, 200) with ReLU activation functions, outputting the prediction of the next state $',' The objective is to minimize L2-distance between the ground- truth next state $141 and the predicted next state 5'41- For the entire set of disassembly trajectories across 100 tasks, train the encoder and forward dynamics model for 200 epochs, using batch size of 128 and a learning rate of 0.001.
S1--1, a for h = 10
Expert Action Features We utilize the disassembly trajectories as reverse expert demonstrations for assembly tasks and aim to capture expert action information in an embedding space. As illus- trated in Fig. Blc), we sample a transition segment 7, from the disassembly trajectories, map it to the action embedding E FT)T) and reconstruct the action sequence (a--h, (4-hhy4thhtt (tthtt,"" "at-)] using
18

---

Published as a conference paper at ICLR 2025
decoder DA- We train both the encoder and decoder with transition segments from all tasks. This embedding effectively extracts the strategy for solving the task by reconstructing the expert actions from the disassembly trajectories.
We sample the transition segment 7--1 = {s at- h, +1, at-h+1," 1t-1, for 10 timesteps (i.e., h 10). The action encoder EA is modeled as multi-layer perceptron (MLP) with three hidden layers of sizes (256, 128, 64), producing a 32-dimensional vector AA,-- The action decoder D is an MLP with four hidden layers of sizes (200, 200, 200, 200) that pre- dicts the sequence of expert actions 16 between input action sequence {az-A, dt--+1,""""" a and the reconstructed action sequence {a/_n: o4 "h-h+1:"". -.-1)- The encoder and decoder are trained for 200 epochs, using a batch size of 128 and a learning rate of 0.001.
d(1) We minimize the L2-distance
A.4.2 TRANSFER SUCCESS PREDICTION IN SRSA
We learn the function Fi Ttrg) to predict the transfer success. For any pair of source policy and target task in the skill library, we execute the source policy on the target task for 1000 episodes with randomized initial conditions and average the success rate obtain the ground-truth label for F. For any task T in the prior task set, we sample the point cloud P. of plug, socket, and their combined geometry (i.e. the plug fully inserted in the socket) to extract geometry features 2G, with a dimension of 96. Then we sample transition segment т; to obtain the dynamics features 2D.s with a dimension of 32 and action features 2A,; with dimension of 32. By concatenating these features, we create task feature 2 with a dimension of 160 for the sampled point clouds and transition segment. We feed the task features Exre,i and Etrg.i for the source and target tasks into an MLP with one hidden layer of size 128 to predict transfer success. We optimize the MLP to learn the transfer success prediction. The training is conducted for 50 epochs with batch size 64 across all tarcee-trreet pairs in the prior task set.
A.4.3 BASELINE APPROACHES FOR SKILL RETRIEVAL.
Signature: Path signatures represent trajectories as a collection of path integrals and also quantify distances between trajectories. Inspired by (Tang et al. 2024), we find the closest path signature for skill retrieval. For each disassembly trajectory 7 on the target task T, we calculate the path signature 24 and search all disassembly trajectories over all source tasks to identify source disassembly trajectory Ty with the path signature 2j closest to žg. The source disassembly trajectory T, belongs to source task in Tprior; thus we match the target trajectory T&e to this source task, denoted as Tk- We count the times that one source task Tare E Tprior is assigned as the source task for a target disassembly trajectory, C(T.re) = Σx=1 [Tke = Taro]. Then we retrieve the policy for the source task with the highest count, i.e., arg maxx,.. C(Tare). In the case of ties, we select one at random.
Inspired by (Du et al. 2023), we employ state-action pairs on disassembly trajectories across all tasks and learn state-action embedding with a VAE for skill retrieval. For any state- action pair (sk; 0x) on the target task, we infer the embedding Z and look for one state-action pair (s az) from the disassembly trajectories in source tasks with an embedding æsa.j closest to Esa,k- The target state-action pair (sk, a&) is matched to the one source task that (sj, a,) belongs to. We denote this source task as Tg. Sweeping through all N state-action pairs in the disassembly trajectories on target task, we count the times that one source task Tarc E Typior is assigned as the =Tsre.. Then we retrieve the policy
Behavior:
source task for a target state-action pair, C(T.re) = EEEn -Eir)) for the source task with the highest count, i.e., arg max C(Tsre)
We learn the latent vector for transition sequence on disassembly trajectories. In order to retrieve one source task according to the distances between task embeddings, we average the embeddings for all transition sequences from the same task to obtain the task embedding, similar to Guo et al. 2022). We retrieve the policy for the source task with the closest task embedding.
Forward:
Geometry: As explained above, learn an autoencoder for the point clouds of the assembly assets to minimize the reconstruction loss, as conducted in (Tang et al. 2024). We retrieve the policy for the source task with the closest point-cloud embedding.
19

---

Published as a conference paper at ICLR 2025
Hyperparameters
Value
Policy Network Architecture
[256, 128, 64] [256, 128, 64]
Value Function Architecture
LSTM network size
256
Horizon length (T)
32
Adam learning rate
le-4
Discount factor (?)
0,99
GAE parameter (A)
0.95
Entropy coefficient
0.0
Critic coefficient
2
Minibatch size
8192
Minibatch epochs
8
Clipping parameter (c)
0.2
LSTM network size
256
SIL update per iteration SIL batch size
8192
SIL loss weight
SIL value loss weight (3) Replay buffer size
0.01
10
Exponent for prioritization
0.6
Table 1: Hyperparameters in PPO and Self-imitation Learning
A.4.4 SKILL ADAPTATION IN SRSA
PPO to train the stochastic policy Ty (i.e., actor) and an approximation of the value function Vi (i.e., critic), parameterized by neural networks with weights 0. The policy is stochastic, following multivariate normal distribution with a learned mean and standard deviation; however, at evaluation and deployment time, the action output from the policy is deterministic.
Implementation Details Following (Tang al. 2024).
The input state for the policy network consists of the robot arm's joint angles, the end-effector pose, the goal mnd-effector pose, and the relative pose of the end effector the goal. The state has dimensionality of 24. Due to the asymmetric actor-critic strategy, the states provided to the value function include privileged information available to the policy. The states for the critic include joint velocities, end-effector velocities, and the plug pose, resulting in an input dimensionality of 44 for the value function.
The action space consists of incremental pose targets, representing the position and orientation dif- ferences between the current pose and the target pose. We use incremental targets instead of absolute targets to restrict selection to small, bounded spatial range. The action dimensionality is 6.
SRSA combines PPO with self-imitation learning (Oh et al. 2018) mechanism for policy fine- tuning. We maintain replay buffer D for transitions encountered during training. The data samples in the buffer are prioritized based on the discounted accumulated reward.
As shown in Algorithm I each iteration includes one PPO update for the policy and value function, along with a batch sampling from D to perform one self-imitation learning update. This update aims minimize the loss function Lau defined in Sec. 4.2 For details network architectures and hyperparameters, refer to Tab.
Input Modality We follow prior work (Tang et al., 2024 to use object poses rather than visual observations as input to the policy. Incorporating vision-based observations would introduce addi- tional challenges for zero-shot sim-to-real transfer, it requires camera. In contrast, the current policy only relies on the fixed socket pose and the robot's proprioceptive features (including the end-effector pose), making it more straightforward to execute the policy in real-world settings.
Using visual observations or object pose is orthogonal to our proposed method (i.e., SRSA is in- dependent of the observation modality). The high-level idea of retrieving relevant skill and fine- tuning the retrieved policy remains applicable in scenarios involving vision-basee policies. The geometry features derived from point clouds in task representation partially capture visual
20

---

Published as a conference paper at ICLR 2025
similarities between tasks. This enables the retrieval of source tasks that visually similar to the target task to some degree.
At the same time, SRSA may require modifications to better support vision-based policies, where each policy relies on a vision encoder to process high-dimensional visual observations. There is guarantee that retrieved source and target tasks are visually similar enough, and the features extracted by the vision encoder in the policy might differ significantly on the source and target tasks. This could pose challenges for fine-tuning the policy on the target task. To address this, we consider two distinct directions: 1. how to perform retrieval to better account for visual similarity; 2. how train specialist policies with visual encoders such that the current SRSA retrieval strategy still likely to work. Below, we propose specific approaches for these two directions.
Enhancing retrieval by incorporating features from visual observations: For example, integrating Variational Autoencoder (VAE) to extract features from visual observations (as in Behaviorketrieval (Du al. 2023 and combining these with other task representations might improve the retrieval process. Additionally, learning dynamics features, such as predicting future visual observations, could implicitly encode relevant visual information task features for retrieval.
2. Improving the robustness of the visual encoder in the policy: Training the source policy with significant data augmentation (e.g., randomizing colors, poses, backgrounds, etc.) could make the visual encoder in the source policy more robust to diverse visual observations. It is more likely to extract similar features from geometrically similar tasks. Alternatively, leveraging state-of-the- art visual foundation models (e.g., DINOv2 (Oquab et al. 2023 as visual encoders could further enhance generalization and robustness. These models have demonstrated strong performance in han- dling diverse observations and sim-to-real challenges, shown PoliFormer (Zeng et al. 2024 Consequently, we believe that features extracted by such visual encoders are likely to remain con- sistent for visual observations across geometrically similar tasks.
A.5 EXPERIMENTS
A.5.1 SKILL RETRIEVAL
We first replicate specialist policy learning for 100 assembly tasks as described in (Tang et al. 2024). These 100 tasks are then split into 90 prior tasks and 10 test tasks. For the prior tasks, we use the trained specialist policies to build the skill library.
We train the task retriever on the prior tasks (Sec. and evaluate its performance on the test tasks. In Fig. 4 in the main text and Fig. [ and Fig. three different ways of splitting the 100 tasks. Overall, SRSA demonstrates superior performance in identifying relevant policies from the skill library, achieving a higher success rate in zero-shot transfer.
2 in the Appendix, present the test results for
Farand Goorwey sitek
Tharster Sucossa
Transfer Successs
0.
0.4
6.
Overaall
Figure 11: Transfer success of retrieved skills applied to test tasks. For each of the test tasks, we retrieve policy from the prior skill library using 5 different approaches. For each approach, if it involves training neural networks, we train it for 3 random seeds. Left: we illustrate the mean result over 10 test tasks. Right: For each test task, we show the mean and standard deviation of transfer success over 3 seeds. Overall, SRSA clearly outperforms baselines.
21

---

Published as a conference paper at ICLR 2025
Forwed Gooraey aRtA
Toanster Sucosia
Transfer Success
04 0b
Overaall
Figure 12: Transfer success of retrieved skills applied to test tasks. For each of the test tasks, we retrieve policy from the prior skill library using 5 different approaches. For each approach, if it involves training neural networks, we train it for random seeds. Left: we illustrate the mean result over 10 tasks. Right: For each task, show the mean and standard deviation of transfer success over 3 seeds. Overall, SRSA clearly outperforms baselines.
A.5.2 SKILL ADAPTATION
We show the learning curves in Fig. At the end of 1000 training epochs, we record the success rate of the learned policies on 10 test tasks, as shown Tab. 2 and Tab. 3 For AutoMate, policies are learned from scratch using PPO. In contrast, SRSA initializes the policies with retrieved skills and fine- tunes them using PPO combined with self-imitation learning. The retrieval mechanism is trained on skill library of 90 prior tasks, where the skills were pre-trained by AutoMate.
Compared to the baseline success rate of 69.4%, SRSA achieves significantly higher success rate of 82.6%, corresponding to absolute improvement of 13.2% and relative improvement of ap- proximately 19.0%. By leveraging the knowledge from the skill library, SRSA also obtains 2.6x lower standard deviation compared to AutoMate (Tab. D This advantage becomes even more pro- nounced in sparse-reward scenarios, where SRSA shows an absolute improvement of 40.8% and
relative improvement of 135% in comparison with the baseline. (Tab. 3h.
Task ID
01029
01036
01041
01053
01079
01092
01102
01125
01129
01136
Average
53.4
89.0
79.1
49.1
74.3
59.4
76.4
49.6
76.0
87.3
69.4
AutoMate
(27.4)
(7.7)
(8.4)
(15.3)
(32.9)
(13.1)
(11.4)
(3.2)
(3.0)
(4.2)
(12.7)
97.3
91.3
78.9
75.4
90.5
78.3
86.6
48.5
82.3
96.9
82.6
SRSA
(1.3)
(6.0)
(7.7)
(6.4)
(2.5)
(6.3)
(4.6)
(5.7)
(5.6)
(2.0)
(4.8)
Table 2: Mean (standard deviation) of success rate (%) on each test task in dense-reward setting. We calculate the mean and standard deviation over 5 runs of different random seeds at the last training epoch (i.e., 1000 epochs).
Task ID
01029
01036
01041
01053
01079
01092
01102
01125
01129
01136
Average
61.3
37.2
14.4
0
81.7
1.4
9.8
55.6
39.7
30.1
AutoMate
(26.5)
(31.4)
(1.6)
(0.5)
(15.1)
(0.5)
(1.0)
(2.0)
(6.0)
(5.4)
(9.0)
95.1
72.4
33.7
87.4
96.1
51.4
70.7
51.2
90.3
60.5
70.9
SRSA
(1.1)
(8.9)
(6.4)
(3.6)
(1.7)
(5.5)
(2.9)
(9.3)
(7.2)
(2.6)
(4.9)
Table 3: Mean (standard deviation) of success rate (%) on each test task in sparse-reward setting. We calculate the mean and standard deviation over 5 runs of different random seeds at the last training epoch (i.e., 1000 epochs).
A.5.3 CONTINUAL LEARNING
We begin with an initial skill library containing 10 policies and expand its size by 10 policies per round over 9 rounds, eventually reaching 100 policies. When the skill library contains fewer than 40 policies, the number of source-target task pairs from the prior task set limited. During this phase, retrieve skills solely based geometry embeddings. That is to say, the retrieved skill
22

---

Published as a conference paper at ICLR 2025
from the skill library is the with the closest geometry embedding to the new task. Once the skill library reaches 40 or more policies, we train the transfer success prediction function to guide skill retrieval for new tasks.
In the continual -eaaninng setting, Fig. g in main text and Fig. 13 in Appendix show the efficiency of SRSA and AutoMate under two different task batch orderings. In both cases, SRSA demonstrates significantly better sample efficiency compared to AutoMate.
1000
800
AutoMate
SRSA
850
Number of Epochs
6
of
SRSA
Success Rate
Batch
(a)
(b)
Figure 13: (a) Overall sample efficiency. We report the number of training epochs required to reach desired success rates (0.5, 0.6, 0.7, 0.8). We calculate the mean and standard deviation of the required training epochs over 5 runs, and report the average across 90 tasks. (b) Sample efficiency in batches. We sequentially introduce batches of new tasks for policy learning, with each batch containing 10 new tasks. For each batch, we show the mean and standard deviation of training epochs required to reach success rate of 0.8. SRSA consistently requires fewer training epochs.
Additionally, compare SRSA and AutoMate based the best checkpoint, measured by the highest rewards achieved over runs for each task. In our replication of AutoMate, we achieved an average success rate of 70% across 100 assembly tasks, which is lower than the 80% reported in the original paper. This discrepancy may be due to differences simulator versions, asset meshes, implementation details, and other factors.
On average, SRSA achieves a success rate of 79% in Fig. 14 and 73% in Fig. 15 for two cases of task ordering, respectively. SRSA demonstrates higher success and better sample efficiency than the baseline AutoMate.
Avg SRSA
Aug AutoMate
AutoMate
0.8
Success
0.6
62
69
Figure 14: Comparison of SRSA and AutoMate success rate over 100 tasks. We replicate the specialist policy learning in the AutoMate paper over all tasks, and run SRSA with the continual- learning approach train 90 specialist policies with the initial skill library of 10 policies. For both approaches, for each task, we select the best checkpoint among runs with different random seeds. We compare the success rate on all the tasks. On average, SRSA achieves higher success rate.
23

---

Published as a conference paper at ICLR 2025
Avg SRSA
Aug AutoMate
AutoMase
SRSA
10
0.8
Success
0.8
0.4
o2
Figure 15: Comparison of SRSA and AutoMate success rate over 100 tasks. We replicate the specialist policy learning in the AutoMate paper over all tasks, and run SRSA with the continual- learning approach to train 90 specialist policies with the initial skill library of 10 policies. For both approaches, for each task, select the best checkpoint among runs with different random seeds. We compare the success rate on all the tasks. On average, SRSA achieves a higher success rate.
A.5.4 ABLATION STUDY
Implementation Details Fig. 16 illustrates the learning curves of different SRSA variations across
10 test tasks.
Skills retrieved based solely on geometry embeddings may face challenges during adaptation due to dynamic differences between the source and target tasks. As a result, the learning curves of SRSA- Geom tend to be less efficient and more unstable than SRSA. We further analyze this baseline in the following section.
When oef-immiaaiinn learning is removed (SRSA-noSIL) from SRSA, the learning curves show in- creased fluctuation and higher variance across runs.
For the generalist policy, which was trained on 20 tasks from AutoMate (including tasks 01036, 01041, 01 129, 01 136), fine-tuning on these tasks yields strong performance since the policy was already optimized for them. However, other test tasks, the generalist policy not effective for efficient policy learning compared to the skills retrieved by SRSA.
Generalist Policy Fine-tuning state-based generalist policy does not perform well because the generalist policy has limited capacity and cannot cover more than 20 training tasks.
As prior work AutoMate TTag et al. 2024h has shown, the training success rate of state-based generalist policy decreases significantly when the number training tasks exceeds 20, given fixed policy architecture of RNN and MLP. We believe that this may be because each task requires precise control across distinct geometric features, and a single policy cannot capture the strategies for all these challenging tasks.
While "increasing model capacity" or moving toward a "large data and large model" regime might help mitigate this problem, it might introduce other challenges. Simply scaling model capacity could result in a generalist policy that works well in-domain but operates more like a "switching cir- cuit." effectively storing task-specific strategies without generalizing to out-of-domain tasks. This approach is suboptimal as it prioritizes in-domain performance at the expense of out-of-distributtin generalization. Thus, we do not want increase the model capacity indefinitely. Instead, we may need a more advanced architecture (e.g., diffusion policy) or model-based RL approach with plan- ning to better handle diverse tasks.
That said, several open questions remain for the state-based generalist policy: How do you design policy architectures capable of high-precision control across many tasks? How you train the generalist policy efficiently many assembly tasks considering possible gradient conflicts? How many training tasks are needed to achieve strong out-of-distribution generalization performance on new assembly tasks?
Fine-tuning a vision-based generalist policy presents more challenges, such as effectively learning a generalist policy across multiple prior tasks with nighddimensional vision observations, fine-tuning on new tasks without forgetting prior ones, and addressing continual learning scenarios, including whether to fine-tune the original generalist policy one already fine-tuned other tasks. We made
24

---

Published as a conference paper at ICLR 2025
CLO29
01036
02042
Pwigt age
Sucoso
Sucoees
4
Lpoch
Epoch
Epocch 01092
01063
01079
Sucossa
Succees
Eppach
Epoch
Epoch
60
Sup00s
Successs
Sucsass
Epach
Epech
011.96
Success
80
Epoch
Comparison for variants of SRSA with different ablated components. For each method, we have 5 runs with different random seeds. The learning curves show mean and standard deviation of success rate over these runs.
Figure 16:
an initial attempt to train a vision-based generalist policy with PPO and fine-tune it. Given 90 prior tasks, can only reach around 10% average success rate. We expect such a generalist policy would perform better than random initialization when fine-tuned for new tasks. Vision-based RL for generalist policy on assembly tasks is a relevantly new topic, and the development such policies lies beyond the scope of SRSA. We leave this direction for future research.
25

---

Published as a conference paper at ICLR 2025
A.6 COMPARISON WITH GEOMETRY-BASED RETRIEVAL
During adaptation, the final performance of SRSA-geom looks close to SRSA in some cases (see Fig. However, is statistically worse than SRSA, especially when there is smaller number of training epochs. To provide more comprehensive evaluation, we run SRSA-geom and SRSA across additional target tasks with three random seeds. The table below summarizes statisties of success rate at different numbers of training epochs, showing that SRSA consistently achieves higher success rates with lower variance. In industrial settings, 2-9% difference in success rate can be highly substantial.
Test task set I
Test task set 2
Success rate (%)
Epoch 500
Epoch 1000
Epoch 500
Epoch 1000
SRSA-geom SRSA
73.6 (= 6.9)
81.0 (=7.7)
67.7 (#7.1)
71.4 (=8.1)
81.4 4+77
82.6 (=4.8)
76.2 (#3.0)
77.6 (#3.5)
Geometry-bade based retrieval alone is not always sufficient. When tasks share similar geometry but have different dynamics, SRSA-geom struggles to transfer as effectively as SRSA. For example, for the target task 01092, SRSA-geom retrieves source task 00686, achieving transfer success rate of only 61.1%, whereas SRSA retrieves task 00213 with a higher success rate of 76.7%. Although the overall shapes of 01092 and 00686 are similar (see below), the lower part of the plug in task 01092 is thinner than the upper part, and there is only a short distance to insert this lower part into the socket. These features closely resemble task 00213, i.e., a narrow plug to be inserted a short distance to accomplish assembly. These shared physical characteristics and similar task-solving strategies make 00213 better suited for transfer. In assembly tasks, the dynamies of the contact region are often more critical than overall geometry for task success. Therefore, source task 00213 works better than 00686 when transferring to the target task 01092.
010g2
on36
00213
(a) Assembly tasks example
(b) Assembly tasks the second example
Additionally, we examine assembly tasks with identical geometry but differing physical parameters. For instance, consider the target task 01 136 with a friction value of 10.0. One source task has the same geometry 01136 but significantly lower friction value of 0.5. SRSA-geon selects this source task due to its geometric similarity; however, the corresponding source policy achieves only 88.9% transfer success on the target task due the friction mismatch (compared rate on its original source task). In contrast, SRSA selects the source task 00213, which not only shares geometric similarity but also has a friction value closer that of the target task. As a result, SRSA retrieved policy achieve higher zero-shot transfer success rate of 93.2%.
99.3% success
A.7 ANALYSIS SOURCE POLICY SUCCESS AS INPUT FOR RETRIEVAL
The success rate of the source policy on the source task is meaningful information to represent the source policy. To see whether it is practically beneficial for retrieval, we modify our approach. We simply concatenate the source success rate information with the task features of the source and target tasks. We train the transfer success predictor with these features as inputs.
We consider three random splits between the prior task set (90 tasks) and test task set (10 tasks). For each split, we train on the prior task set over three random seeds. For each seed, we test the trained function F on the test task set for retrieval. We report the mean transfer success rate of the retrieved skills on 10 test tasks, with the standard deviation reported over three seeds. Empirically, the source success rate input F only slightly improves the retrieval results.
26

---

Published as a conference paper at ICLR 2025
Average transfer success (%)
Test task set
Test task set
Test task set
SRSA SSSA+source success rate
62.7 (=5.))
53.7 (=5.5)
44.9 (=2.4)
66.7 (-3.)
23.7
43.7 (#3.7)
A.8 ANALYSIS OF OUT-OF-DISTRIBUTION TEST TASKS
For out-of-distribution (OOD) tasks where no skill transfers zero-shot, SRSA may indeed struggle, and the initialization from retrieved skill might not help much. To tackle this, it is essential to build a skill library which as diverse as possible. When the target task falls outside the current library's 's distribution, we can use SRSA's continual learning approach (section 4.3 & 5.4) to expand the library with new tasks. By building a larger, more varied skill library, we increase the likelihood that this target task will align better with tasks in the skill library.
We run experiments for target tasks with IDs 00004, 00015, 00016, 00028, 00030. These tasks suffer from low transfer success rate given small skill library with only 10 prior tasks. However, when we have a larger and larger skill library, the retrieved skill has a higher transfer success rate on the target task.
Transfer success rate (%)
00004
00015
00016
00028
00030
10-task library
15,9
6.9
0.2
12.2
39.1
50-task library
12.7
8.4
0.3
27.5
49.4
90-task library
24.2
28.4
19.3
18.1
82.6
As demonstrated, continual learning to expand the skill library is a promising step; however, gener- alizing OOD tasks is longstanding challenge in robotics, and is still an open question how to optimally construct the curriculum that governs the expansion of the skill library.
A.9 ANALYSIS OF OTHER METRICS FOR RETRIEVAL
We acknowledge that zero-shot transfer success rate may not be a perfect proxy for retrieval. We can consider several other possible metrics for retrieval: (1) Ground-truth success rate after adapta- tion (2) Predicted success rate after adaptation (3) Predicted zero-shot success rate (i.e., SRSA) (4) Predicted zero-shot dense reward.
Option is the ideal metric to identify the best skill for retrieval, as our final goal is to obtain the highest success rate on the target task after adaptation. However, it introduces chicken-and-egee problem, as we cannot get this metric without fine-tuning all candidate policies on the target task.
Option 2 requires training a predictor for the success rate after adapting any source policy any target task. We need the training labels of the ground-trun -tuun success rate after adaptation. Unfor- tunately, collecting this training data would require extensive computational resources. For each source-target pair, we need at least 20 GPU hours to finish adaptation; given a skill library of 100 tasks, 200,000 GPU hours would be required to collect training data. Furthermore, will remain intractable as the skill library becomes larger.
Option 3 (SRSA) requires much less resources to collect training data for the predictor. We only need 20 minutes on a GPU to evaluate one source policy on a target task. It thus requires 3,000 GPU hours collect training labels. We conduct an experiment to compare the performance of Option and Option 3 on two test tasks. To collect experimental results for Option 1, for each test task, we sweep all 90 source policies in our skill library. We fine-tune each source policy with one random seed to adapt to the test task and identify the best success rate after adaptation. Below we report the success rate of Option and Option 3 on two test tasks, after fine-tuning for 1500 epochs.
Success rate after adaptation (%)
Test task 1036
Test task 1041
Option (SRSA)
95.9
89.1
Option 1
98.3
94.0
Option 1 is the perfect but intractable metric for retrieval. The difference in success rate between the SRSA-retrieved skill (Option 3) and the best source skill (Option 1) is less than 5% after adaptation. Therefore, although zero-shot transfer success rate is not a perfect metric for retrieval, it is a high- quality metric for retrieval in terms of both performance and computational efficiency.
27

---

Published as a conference paper at ICLR 2025
Furthermore, we consider using dense reward information to guide retrieval (Option 4). We learn to predict the accumulated reward rather than success rate on the target task when executing the source policies in zero-shot manner, then we retrieve the source policy with the highest predicted transfer reward. In the table below, we show the performance of retrieved skills when they are applied on the target tasks.
Test task set I
Test task set 2
Transfer reward Transfer success (%)
Transfer reward Transfer success (%)
Option (SRSA)
8134
62.7
7722
53.7
Option 4
7976
54.8
7935
32.6
On the AutoMate task set, Option 3 (SRSA) yields slightly better skill retrievals, especially with higher transfer success the target task. However, success rate may not accurately reflect the expected value for tasks with dense rewards; the higher transfer success rate does not mean higher transfer reward in test task set Therefore, if is critical prioritize the reward achieved on the target task, using the cransferreward predictor for retrieval is reasonable choice. Conversely, if the success rate on the target task is more critical (as in our assembly tasks), transfer success would be the preferred choice as a retrieval metric.
A.10 ANALYSIS OF DISTANCE METRICS FOR TASK FEATURES
In SRSA method, we jointly learn features from geometry, dynamics and expert actions to represent tasks, and predict transfer success to implicitly capture other transfer-related factors from tasks. To investigate the advantage SRSA, we compare it against baselines that uses simple distance metrics for task features to determine task retrieval. Specifically, we concatenate the features of geometry, dynamics and expert actions as the task features and apply 1.2 distance, L1 distance, and negative cosine similarity between the vectors the distance metrics for retrieval. We consider three different ways to split the prior task set (90 tasks) and test task set (10 tasks). For each test task, we retrieve the source task with the closest task feature to the target task. The table below shows that SRSA outperforms the baselines on all test task sets.
Transfer success rate (%) 1,2 distance LI distance Cosine similarity SRSA
Test task I Test task set 2
51.6
50.8
52.6
62.7
47.1
49.0
46.5
53.7
Test task set
35.3
35.0
36.1
44.9
We attribute SRSA's advantage to its transfer success predictor F, which capture additional infor- mation relevant to policy transfer. By explicitly learning to predict transfer success, F provides a more effective metric for selecting source tasks with higher zero-shot transfer success."
A.11 ABLATION STUDY ON POLICY INITIALIZATION AND SELF- LEARNING
For policy learning, AutoMate uses PPO from random policy initialization, and SRSA uses PPO with self-imitation learning (SIL.) after initialization with the retrieved skill. Thus, the main differ- ence between SRSA and AutoMate lies in (1) strong initialization from retrieval and (2) SIL.
In Sec. we compared SRSA and SRSA-noSIL to show the effect of SIL. Below, we additionally compare with SRSA with random initialization (SRSA-noRetr) to show the effect of initialization from retrieval. Comparing AutoMate with SRSA-noRetr, we see the difference between PPO and PPO+SIL when learning policy from scratch. Both approaches started from poor performance, but SIL has greater learning efficiency and stability. Comparing SRSA-noRetr and SRSA, we see the difference between random initialization and initialization from retrieval. Policy retrieval provides a good start with a reasonable success rate. As a result, SRSA more efficiently reaches higher performance on the target task.
A. 12 TOP-k RETRIEVAL SELECTION IN SRSA
Although the transfer success predictor F in SRSA effectively guides the retrieval of relevant skills, its predictions may not always be perfectly accurate. To mitigate this issue, retrieve the top-
28

---

Published as a conference paper at ICLR 2025
01036
01041
01053
Success
Sucvess
Success
0g
SRSA
SRSA
Epoch
Epoch
Epoch 01102
01079
01092
Succrs %
Success
Success
4
Epoch 01125
Epoch
Epoch
Success 06
Success
Success
Epoch
Epoch
Comparison for variants SRSA with different ablated components. For each method, we have 5 runs with different random seeds. The learning curves show mean and standard deviation success rate over these runs.
Figure 17:
skills ranked by the predictor F. With these candidates, we evaluate their zero-shot transfer success on the target task by running each candidate for 100 episodes. We then select the best candidate with the highest average reward, ensuring that we identify the most relevant skill based on we set k to 5.
actual performance rather than just the predicted ranking. In Sec. 5
To illustrate the impact this selection process, we compare SRSA with SRSA-topl. Since the predictor F is not always precise, the top-1 skill based on predicted transfer success may not be the best in terms ground-trutt truth performance. The learning curves in Fig. 18 demonstrate that SRSA and SRSA-topl perform similarly in most cases. However, for certain tasks (c.g., 01 125, 01 029, 01136), SRSA benefits from selecting among the top-5 candidates and retrieves better-performing skills compared to SRSA-topl. Here SRSA-topl works well in our setting because the pre-trained predictor F is reliable given a large 90-task prior skill library. However, with a smaller skill library, F may be prone overfitting and top-k selection is probably more advantageous in this case.
29

---

Published as a conference paper at ICLR 2025
9
Sucodss
SRSA-1o1
SRSA
Fipech
01063
01079
01092
Su0005
Sucoees
Sucossa
%
6
Epoch
Epach
Epoch
01102
01125
01129
Success
%
5%
80
6%
Epoch
Epoch
Epoch
Success
80
Figure 18: Learning curves on test tasks. The --axis and y-axis represent training epochs (where each epoch consists of 128 environment steps over 256 parallel environments) and success rate, respectively. The solid line shows the mean success rate over runs with different random seeds, and the shaded area denotes the standard deviation. SRSA takes the top-5 relevant skills based on transfer success prediction and selects one skill for policy initialization based on the ground-truth zero-shot success rate of applying the skill the target task. SRSA-topl directly retrieves the skill with the highest transfer success prediction.
30


Reviews:
[{"review_id": "IhOiLdV1gE", "reviewer": "Reviewer_YitF", "summary": "Skill retrieval and Skill Augmentation (SRSA) tackles the problem of quickly learning new tasks in task family. In the paper, the authors consider various insertion tasks from the AutoMate benchmark which share the same observation space, action space and reward function. The tasks differ in the geometries and thereby the transition dynamics and expert actions required.\n\nThe authors assume that an availability of a skill library containing expert policies for different insertion tasks. Skill retrieval is based on geometry similarity, transition dynamics and expert actions. Transition dynamics and expert actions are based on reverse trajectory generation. The similarity is predicted using a learnt function F, which is a 0-shot policy transfer success rate predictor trained on existing skill library. Given a new task, the policy with the highest 0-shot transfer prediction is used extracted and fine-tuned on the new task using Reinforcement Learning (PPO) with additional Self Imitation Learning (SIL) objective.", "questions": "Should F use as input the success rate of the expert policy in the original task?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Strong simulation results\n2. Insightful that we can train a 0-shot transfer predictor function using geometry, transition and expert action encodings\n3. Results transfer to real robot", "weaknesses": "1. Assumption that that all tasks have the same observation space, action space and reward function\n2. Based on Fig. 6 and Appendix Fig.16, it seems SRSA-geom is very comparable to SRSA. Hence, it seems geometry based skill retrieval might be sufficient. It's also easier as it doesn't require to collect disassembly trajectories and training a 0-shot transfer predictor function", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fbQp94dmM9", "reviewer": "Reviewer_49Jv", "summary": "This paper focuses on the problem of learning to solve new assembly tasks efficiently while utilizing solutions to previously solved tasks (previously solved assembly problems). This paper achieves this using a skill retreival and adaptation approach. First, the paper assumes that there exists N set of skills where each skill can solve a particular assembly task (particular CAD model).  To solve a new task, it tries to retrieve the most similar previous CAD from the previous set of skills. This retrieval problem is solved by using a heuristic which assumes that two CADs are similar if the policy from source zero-shot solves the second task (or gives high reward in dense reward settings). This retrieval problem is used to construct a dataset to find similar tasks. This dataset is used to learn an embedding which can predict if the CAD models are similar or not (using labels from the reward heuristic).\n\nNow for a new incoming task (or new CAD model), the best existing skill is extracted (using the learned embedding). This skill is then refined using RL. This library can then be used in a naive continual learning setting.", "questions": "see above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and easy to follow. Most of the proposed approach makes sense, the experimental results along with the ablations show good performance of the proposed approach.", "weaknesses": "*Generality:* Unfortunately, I don’t think this approach is very general. This seems a bit too tailored for assembly tasks. This  paper might fit a robotics venue much better and would probably be appreciated more in that setting. However, this is not my decision and I don’t think this should in any way count as a negative for the paper.\n\n*Sparse reward setting:* The retrieval approach cannot work if we are in sparse reward settings, since there can be some CAD model for which no other skill transfers 0 shot. How would this approach solve this problem? This is definitely very relevant for real-world setting where dense rewards can be quite hard.\n\n*Generalist approach (ablation):* Do you have any intuition on why the generalist approach fails? Would better feature learning or better multi-task policies help a generalist approach. Clearly, a generalist policy trained on all tasks is way more appealing. Also, curious to hear any thoughts on why the generalist policy does not perform well on all tasks. Is this because some of the assembly tasks can be very different?\n\n*Learning embeddings for retrieval:* Given that there are less than 100 samples to learn the embedding function which uses high-dimensional input (such as point cloud embeddings). Did you find any challenges in learning this function. \n\n*Reward heuristic:* While the simulation lemma is used as a motivation for the reward heuristic to be used in a zero-shot setting, I think there is a difference here. If two MDPs share same transition dynamics and initial state distributions then same policy will give high rewards for both. But this does not necessarily go the other way around, i.e., one policy can gives high rewards for two MDPs *does not* mean that they share similar T and S_0. I think this is an important point that should be clarified in the paper. Thus, the use of transfer should be seen as a heuristic and not a theoretical justification.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oAD0G6gize", "reviewer": "Reviewer_F4Pz", "summary": "This paper introduces a method for more effective learning of RL policies for robotic assembly tasks. It first uses Skill Retrieval to retrieve the most relevant policy from the skill library, a set of available policies that solve similar types of tasks. The retrieval is done based on the predicted zero-shot transfer success rate using object geometry, dynamics and expert-action features. The retrieved policy is then fine-tuned on the new task, facilitating more efficient learning when compared to learning to solve this task from scratch. The authors also showcase how their method could be used for continual learning by progressively expanding the skill library. The experiments show improved performance when compared to baselines, while the ablation studies provide valuable insights about the proposed method.", "questions": "See the Weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Paper is well written and motivated.\n- The proposed method is simple, elegant and easy to understand.\n- Using geometry, dynamics and expert-action features for retrieval is a nice idea that incorporates different aspects of the underlying tasks.\n- Experiments are conducted well and the ablation studies provide valuable insights into the method.\n- Real-world experiments are appreciated.", "weaknesses": "- Line 217 - what does being similar with a high value mean? Maybe this sentence could be rephrased for clarity.\n- Equations could be formatted better (e.g. not to overflow) and also it would be useful if they were numbered and referred to in the text when needed.\n- References to the Appendix are broken as it's a separate file. This, I believe, will be addressed for the camera-ready version, just wanted to point it out.\n- Using zero-shot transfer success rate as a proxy for retrieval is intuitive, but, I believe, it doesn’t accurately reflect the expected value for tasks with dense rewards. Could authors provide more details as to why using a zero-shot transfer success rate is sufficient? Would it be beneficial to use additional metrics as well?\n- If the difference in expected value when applying the same policy to different tasks depends only on their difference in transition dynamics and initial state distributions that are captured by geometry, dynamics and expert-action features, why do you need to approximate the zero-shot transfer success rate at all? Wouldn't comparing these features using some distance metric be enough? Maybe authors could do some experiments on this or justify why such an approach is not appropriate.\n- Relying on disassembly paths seems to restrict the method’s applicability to simple insertion tasks on a plane. Could the authors more clearly point out such a limitation and/or propose solutions to it in the paper?\n- In a supervised learning setting, the loss function is usually positive. Why the objective function (line 287) is negative?\n- Referring to different tasks in the experiment section using their ID number (e.g. 01125) doesn’t provide any information as to why they are hard or easy. Adding some text or image information about these tasks would benefit the understanding of the experimental results.\n- In the experiments, improvements in the success rate over the baselines are provided in percentage points, which can be hard to understand given that success rate can also be expressed in percentages. For example, does a 22% increase mean a +22% or 22% increase when compared to the base? This is especially strange when seeing an 181% increase (given that the success rate is capped at 100%). Maybe the authors can clarify that, or (as I would encourage) use a different metric.\n- It would be beneficial to discuss the main difference between the proposed method and the AutoMate baseline (not including the retrieval part). It would then be easier to understand the experimental results (e.g. AutoMate’s instability, line 427).\n- The paper would benefit from a discussion on using visual observations, rather than object poses, and whether the proposed method would be better than fine-tuning a generalist policy that would have the capacity to generalise implicitly. \n- It would also be interesting to see whether fine-tuning or learning a residual on top of the retrieved policy is more efficient. Although, this might be outside the scope of the current approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)