Active Fine-Tuning of Multi-Task Policies
Authors: Marco Bagatella
Jonas Hübotter
Georg Martius
Andreas Krause
Venue: ICML 2025 poster (2025)
Keywords: imitation learning
deep reinforcement learning
multi-task reinforcement learning
active learning
fine-tuning
Primary area: reinforcement_learning
TL;DR: We actively query demonstrations when fine-tuning a multi-task policy through behavioral cloning.

Abstract: Pre-trained generalist policies are rapidly gaining relevance in robot learning due to their promise of fast adaptation to novel, in-domain tasks.
This adaptation often relies on collecting new demonstrations for a specific task of interest and applying imitation learning algorithms, such as behavioral cloning.
However, as soon as several tasks need to be learned, we must decide *which tasks should be demonstrated and how often?*
We study this multi-task problem and explore an interactive framework in which the agent *adaptively* selects the tasks to be demonstrated.
We propose AMF (Active Multi-task Fine-tuning), an algorithm to maximize multi-task policy performance under a limited demonstration budget by collecting demonstrations yielding the largest information gain on the expert policy.
We derive performance guarantees for AMF under regularity assumptions and demonstrate its empirical effectiveness to efficiently fine-tune neural policies in complex and high-dimensional environments.

Active Fine- Tini-Tnnig of Multi- Task Policies
Marco Bagatella 12
Jonas Hübotter'
Georg Martius 23
Andreas Krause'
Abstract
Pretrained policy
Agent selects task
Pre-trained generalist policies are rapidly gaining relevance in robot learning due to their promise of fast adaptation
novel, in-domain tasks. This adaptation often relies on collecting new demonstrations for a specific task of interest and applying imitation learning algorithms, such as behavioral cloning. However, as soon as several tasks need to learned, we must decide which tasks should be demonstrated and how often? We study this multi-task problem and explore an interactive framework in which the agent adap- tively selects the tasks to be demonstrated. We propose AMF (Active Multi-task Fine-tuning), an algorithm to maximize multi-tasi policy per- formance under a limited demonstration budget by collecting demonstrations yielding the largest information gain on the expert policy. We derive performance guarantees for AMF under regular- ity assumptions and demonstrate its empirical ef- fectiveness to efficiently fine-tune neural policies in complex and high-dimensions environments.
Expert provides demonstration
Dataset
Behavior cloning
Figure 1: Interactive loop between agent and expert. We consider scenario where we receive pre-trained policy, and are able to obtain expert demonstrations of tasks. We study how to select tasks (in blue) to obtain the best- performing policy after as few demonstrations as possible.
ing (Ceron et al., 2024b;a; Nauman et al., 2024) suggest the possibility of leveraging large pre-trained policies. These "generalist" policies have decent performance on many tasks, and can be fine-tuned on particular set of tasks while lever- aging their previously learned representations and skills. We investigate whether representations of pre-trained policies can be used to significantly bootstrap learning progress.
As a motivating example, consider a household robot that is delivered with a pre-trained "generalist" policy, and deployed in different conditions than those observed in its training data. While the robot may achieve some tasks in a zero-shot fashion (e.g. simple pick-and-place). other tasks might necessitate further fine-tuning (e.g. cooking an omelett)). The robot should be able to interactively request demonstrations to compensate for its shortcomings. We seek to answer which demonstrations should be requested from the user to achieve the best performance, as quickly as possible.
1. Introduction
The availability of large pre-trained models has trans- formed entire areas of machine learning, from computer vision (Krizhevsky et al., 2012; He et al., 2016; Dosovitskiy et al., 2021; Radford et al., 2021), to natural language processing (Radford et al., 2019; Brown et al., 2020) and generative modeling in general (Ho et al., 2020; Esser et al., 2024). This paradigm has started to extend to robotics and control (Collaboration, 2023; Ma et al., 2024), in readily
If the agent only needs perform well in a single task, the fine-tuning process conventionally relies behavioral cloning (Chen et al., 2021; Reed et al., 2022; Bousmalis et al., 2024) of expert demonstrations. As collecting demon- strations is in general costly, the number of demonstrations required, and thus the expert's effort, should be minimized. However, as each demonstration should solve the same task, the allocation of the expert's 's effort is straightforward, The multi-tase case presents the more nuanced problem of select- ing which tasks to demonstrate, and when. This motivates the main focus of this work: provided a pre-trained pol- icy, how can we maximize multi-task performance with a minimal number of additional demonstrations?
particular for systems for which demonstrations available (Octo Model Team et al., 2024), or can be easily collected (Zhao et al., 2023). Even when demonstrations are not easily obtained, scaling laws in reinforcement learn-
'Department of Computer Science, ETH Zürich, Zürich, Switzerland "Max Planck Institute for Intelligent Systems, Tübingen, Germany University of Tübingen, Tübingen, Germany, Correspondence to: Marco Bagatella <mbagatella@ ethz.ch>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
To address this problem, we propose AMF (Active Multi-

---

Active Multi-task Policy Finc-tuning
task Fine-luning), which selects maximally informative demonstrations. AMF parallels recent work on active super- vised fine-tuning of neural networks (Hübotter et al., 2024). To this end, AMF relies on estimates of the demonstrations' information gain on the expert policy. We prove that in suffi- ciently regular Markov decision processes, AMF converges to the expert policy. We then focus on practical scenarios where policies are represented as neural networks. We show that, despite additional challenges, AMF can effectively guide active task selection in such settings, leading to better policies after fewer demonstrations. Our contributions are:
models (Brown al., 2020) has been mirrored by explo- ration of multi-task, or generalist policies, often trained via imitation learning (Reed et al., 2022; Bousmalis et al., 2024; Collaboration, 2023). These recent works mostly build upon algorithms developed for the single-task case, and simply integrate cask-condiiioning as part of the state. While sev- eral works hand-select parts of large, open-source roboties datasets for pre-training (Octo Model Team et al., 2024), active data selection for multi-task fine-tuning has not been addressed. Prior work on meta-learning has studied how one can explicitly meta-learn the ability to adapt to task demonstrations (Finn et al., 2017). We find this capability to emerge even from models that are not explicitly trained in this way, and focus on which demonstrations to obtain.
We propose AMF, an algorithm for multi-task policy fine-tuning that maximizes the information gain of demonstrations about the expert policy.
We prove statistical guarantees for AMF, which ex- tend the results of Hübotter et al. (2024) to dynamical systems.
Data Selection The idea of directing a sampling process to gather information has been central to machine learning research and studied extensively in experimental de- sign (Chaloner & Verdinelli, 1995), active learning (Settles, 2009) and reinforcement learning (Mehta et al., 2022). Most work on active data selection summarizes data without focusing on a particular task (e.g., Sener & Savarese, 2017: Ash et al., 2020; Holzmüller al., 2023; Lightman et al.. 2023), which has been predominantly applied pre-training.
We empirically scale AMF to high-dimensional tasks involving pre-trained neural policies.
We additionally propose a practical approach to alle- viate catastrophic forgetting in pre-trained neural poli- cies, which benefits arbitrary data selection strategies.
2. Related Work
Recently, adapting models after pre-training and during deployment has gained interest. Several works, mostly in computer vision, focus on unsupervised fine-tuning on a test instance (Jain & -Mille, 2011; Krause et al., 2018; Sun et al., 2020; Wang et al., 2021b; Chen et al., 2022). We focus instead on supervised fine-tuning of learning- based controllers in dynamical systems. Our approach extends work task-directed data selection (Kothawade et al., 2020; Wang et al., 2021a; Kothawade et al., 2022; Bickford Smith et al., 2023), which has recently been applied to the supervised fine-tuning of large-scale neural networks in vision (Hübotter et al., 2024) and language (Rotman & Reichart, 2022; Xia et al., 2024; Hübotter al., 2025).
Learning- baaming-based control and active data selection are both well-established research directions. This section discusses some topical works in either direction, and clarifies the novelty and placement of this work with respect to them.
Behavioral Cloning Numerous imitation learning ap- proaches have been developed with the goal of distilling knowledge from high-quality demonstrations to a control policy (Osa et al., 2018). Within this family of techniques, behavioral cloning (BC, Bain & Sammut, 1995; Ross & Bag- nell, 2010) aims to maximize policy performance by mini- mizing the distance of its actions to demonstrated actions, simply through supervised learning. While BC may suffer from accumulating errors (Ross et al., 2011) 1), its empirical effectiveness has seen increasing support when high-quality demonstrations are readily available (Kumar et al., 2022). Next to recent empirical successes (Chi et al., 2023), formal analysis has also advanced (Spencer et al., 2021; Block et al., 2024a; Belkhale et al., 2024; Foster et al., 2024), and estab- lished provable performance guarantees for BC policies (Xu et al., 2020; Maran et al., 2023; Block et al., 2024b).
3. Background
3.1. Multi-Tass Reinforcement Learning
The multi-task setting can be modeled by casting the en- vironment as a contextual Markov decision process (MDP) M = (S, AS,AC C, P, R, Y, Ho) where S E R and A E R-NA are possibly continuous state and action spaces. is a (potentially infinite) set of tasks, with each task represented by an Ne-dimensional vector e E Ne P -> A(S) models the transition probabilities (A(S) represents the set of probability distributions over S), R : S x C -> R is a scalar reward function, E (0,1) is a discount factor and Mo E A(S) is the initial state distribution. In this setting, a policy is simply a sall-and-task-condiiin
Multi-task and Generalist Policies Traditionally, behay- ioral cloning has mostly been deployed in a single-task setting. Multi-task learning in sequential decision-makin has largely been investigated in the context of reinforcement learning (Teh et al., 2017; Sodhani et al., 2021; Yu et al., 2021; Sun al., 2022; Cho et al., 2022; Hendawy et al., 2023). Moreover, the recent rise of multi-task generative
2

---

Active Multi-task Policy Finc-tuning
action distribution π SS + A(A)'. Any given policy induces task-conditional distribution over trajectories:
which can be used to fine-tune the policy, and condition the agent's query at step + 1. The process is repeated for N rounds, with the goal of producing fine-tuned policy that maximizes the expected returns for the desired task distribution Pe- We note that the fine-tuning process does not assume access to the pre-training data distribution. This setting is both realistic, as large pre-training datasets are rarely publicly available, and challenging, as prohibits any naive data rebalancing strategy.
Tir ((00,,000,)) 81, (I,.a.....) Ic)
0c
(o(5)) II T(, I sc. P(s I a;)
The discounted returns for a specific task e E C or a task distribution je ACC) are, respectively,
Modeling assumptions We take a Bayesian perspective on active multi-task fine-tuning, by assuming a Bayesian model π over policies. We assume that demonstrations follow a noisy expert: a c) = H rs c) + c(s, c) where e(s, c) is independent noise. We remark, however, that AMF can also be understood from a non-Bayesian perspective as selecting tasks that most quickly minimize the size of frequentist confidence sets around the optimal policy.
8C Σn R(((s c) and
E (80 (NNT(()) (c)
J MC
J c-pe
Reinforcement learning algorithms traditionally aim directly at maximizing J/. which is notoriously chal- lenging. In the scope of this work, we instead consider an imitation learning setting, in which expert demon- strations from an optimal policy x² are provided. In particular, we focus on behavioral cloning algorithms, which reduce control to a supervised learning problem. Given a set of N task-conditioned, H-length trajectories fi.N = (s), ab..... - -) with task labels C1:N. behavioral cloning proposes proxy objective for the policy π: an empirical estimate of the log-likelihood under the data distribution: prxxy = N E log ग ((a st, es). If trajectories T1:N are obtained from the optimal policy, cover the support of the desired task distribution Nee and the searched policy class is sufficiently rich, the maximizer of proxy will also maximize Jhh as N and H increase. How- ever, in general, there is a clear mismatch between proxy and J.H (Xu et al., 2020; Maran et al., 2023). Nonetheless, the optimization of poxx is relatively straightforward su- pervised learning problem, while the full RL problem raises several convergence issues, particularly in the offline setting (Levine et al., 2020). Thus, we use J only for evaluation, and carry out optimization through the proxy objective.
4. Method
The active multi-task fine-tuning problem outlined so far re- quires active data selection for sample-efficien learning. We thus build top of principled active learning approaches for non-sequential domains (Hübotter et al., 2024), and propose AMF, which selects queries that maximize the expected information gain about the expert policy over its occupancy:
#-1
Gn = arg max E C'EC
Σ (T((s,)): cin- F1mm-1),
1=U
with e Mes (so....) ~ T()), Tim-1 ~ 7 (C 1). (1)
We show in Section 4.1 that, under certain regularity assumptions, the policy learned by AMF converges to the expert policy and matches its performance. These results constitute a first-of-it-kinn performance guarantee for active multi-task fine-tuning. The main novelties of this guarantee are the extension of prior work to sequential do- mains where the visited trajectory (Su, ag $1,...) ) ~ ((2,) is unknown when selecting the task en for a demonstration, and the connection to imitator performance. we discuss the design choices that make AMF amenable to optimization in practical settings.
3.2. Active Policy Fine-Tunine TaninT
In this work, we consider an active fine- mee-tuning scheme for multi-task policies. The goal is to fine-tune a pre-trained pol- icy to perform well on a desired task distribution Me using as few expert demonstrations as possible. The agent is allowed N sequential queries for demonstrations according to the fine-tuning budget. The n-thh query should consist of task On E C. Once the agent selects a task, feedback is received from the optimal policy n²= S x C + A (i.e., an optimal demonstrator). At each round the agent receives an H Htstep demonstration conditioned on the chosen task cn This can be seen as a single measurement from a stochastic process over trajectories T C + A(S x A) Each observed trajectory up to round n is stored in dataset (Ci., f,m)),
Section 4.2,
4.1. Performance Guarantees
We begin by presenting the performance guarantees for AMF. Our proof builds upon rates for uncertainty reduction, then ties these to probabilistic convergence guarantees to n², finally resulting in performance guarantees within the MDP. We summarize the main result here, and include a formal proof in Appendix A.
Informal Assumption 4.1. We make these assumptions:
We use and to denote stochastic and deterministic policies, respectively, and a(2, e) for realizations.
1. The expert policy n² is deterministic, Lipschitz-
3

---

Active Multi-task Policy Finc-tuning
smooth, lies in the reproducing kernel Hilbert space (S C) of the kernel & with norm ||=" ||ae oo and induces Lipschitz-smooth Q-function.
as it does not depend on c', and rewrite the second entropy term as an expectation over =(c)). )- As long as the task space C' is finite and its cardinality is tractable, the arg min opera- tor can be evaluated exhaustively, and the expectation over the task distribution je can be computed exactly. When this is not the case, the arg min can be optimized through discretization, or with sampling- -aasldnng-eedddddd optimizers. The ex- pectation over pe is also not particularly problematic, as can be computed in closed form (if C is discrete) or esti- mated empirically through sampling, as pe is assumed to be known. However, two issues need to be resolved: (i) computing the expectation over the noisy expert's 's trajectory distribution т, and (ii) estimating the conditional entropy term H(-1 ().
2. The noise c(s, c) affecting demonstrations is condition- ally p-sub-Gaussian and bounded.
3. The dynamics of the contextual MDP M are Lipschitz- smooth with bounded support, the initial state dis- tribution No has bounded support, and the reward is Lipschitz-smooth
Under these assumptions, we prove the following perfor- mance guarantee for active multi-task behavioral cloning.
Informal Theorem 4.2 (Performance guarantees for active multi-task BC). Let all regularity assumptions hold. If each demonstrated task of length H is selected according to the criterion in Equation 1, then with probability 1 - ô the performance difference between the expert policy x and the imitator policy Tn after n demonstrations can be upper bounded:
Occupancy estimation Computing the expectation over policy's occupancy over states or trajectories is in general intractable in continuous state spaces. Fortunately, coarse empirical estimate can be obtained as soon as few expert demonstrations become available. The expectation can be estimated through a single sample, which is always available in the form of the trajec- tories f collected so far, as they have effectively been sampled from 7 (C):n-1). However, the remaining two ex- pectations (i.e., (r (-) and )uT(a)))><(())
P. JH S O(T(F)))
where nn is the mean of π at round n and ((H)) is the maximum information gain about the expert policy from Hn samples, and is sublinear for a large class of problems. The O-) notation suppresses all multiplicative terms that do not depend on n.
(-) involve the distribution over trajectories for an arbilrary task, which might not have been demonstrated yet. We observe that, at round n, the tasks demonstrated so far induce the empirical distribution (a(+) = ب Σn' de, (.), while the trajecto- ries collected similarly induce +()= Ei (-). where 8 indicates the Dirac delta distribution. We can show that expectations over the trajectory distribution for an arbitrary task e E C can be estimated through importance sampling (i.e., by sampling trajectories from +()) instead
Intuitively, this theorem proves that the imitator will eventually achieve the demonstrator's performance in smooth, regular MDPs with sublinear KHH (for a formal definition, we refer to Lemma A.6 in the Appendix). We can also prove a more general result under weaker assumptions: as long as the policy is regular, the imitator will reach the noisy expert performance in arbitrary, non-smooth MDPs, albeit only in expectation. A full derivation of this further result can be found in Appendix B.
of T(-le):: Er-r(+le) f(r) E.)+(()
T(r c)
*(r) f(r)- The
importance weights can then be estimated as
4.2. Practical Algorithms
((TCc) ((T)
((Tcc
(T|C)
Theorem 4.2 guarantees that, under regularity assumptions, the adaptive demonstration sampling scheme leads to con- vergence of the imitator's performance to the optimal one. However, this criterion involves state occupancies and a con- ditional entropy term, which are hard to access or estimate in practice. Thus, here we derive a practical objective to be deployed in general settings. We first rephrase the objective from Equation I in its entropy form:
Seee (e(c)
n- ΣE T(Tlo:)
(n 1)H0(50) H-1 (a,
1==
E (80) )
I A c)) 186, Uz)
(n
11 [ ((cess
= w(T,c)
Σ
H (aclst.
(=0
where T = (s0, ao. .... and मे can be approximated with the current estimate of æ. Intuitively, the likelihood ratio of a trajectory under two different tasks only depends on the like- lihood of actions under the policy, and thus does not require knowledge of the MDP. As the estimate may be inaccurate for small numbers of samples, in practice the algorithm can invest the first few rounds to query a single demonstration for each of the tasks (in case they are countable and few) or to sample the task space uniformly. On the other hand, the
H-1
Cn = arg min C'EC 1=
E.
(() d','''' C1m- 11
with c~ Mo: (So.... ~ r'c), T(d).
Tim-1 T()
(2)
where we use the definition of mutual information [(-)T(c)) = H(-) - ((-)T(d)), drop the first entropy term
4

---

Active Multi-task Policy Finc-tuning
high-variance of the estimate can be controlled by practical solutions such as clipping. We present the resulting empir- ical estimate for Equation 2 in full in Appendix C, and qualitative analysis of importance weights in Appendix K.
Even if the pre-trained model was perturbed and ensembled for fine-tuning, the ensemble disagreement would not cap- ture the pre-training data distribution. Second, access to pre-training data is in general unrealistic, or hard to manage due to size and ownership of large robotic datasets.
Entropy estimation The estimation of conditional en- tropy terms such as H(( .) has been widely researched in the literature. When the policy is represented through a Gaussian process GP k) (Williams & Rasmussen, 2006) with known mean function j and kernel k, the entropy can be directly quantified by the predicted variance. Let us denote state-task tuple as z = c),, and let X be the sample vector obtained from concatenating states and tasks from previous trajectories (c.g., Cim--1 Tim--, c', T'). The unconditional entropy can be measured in closed form as (()()))) 1/ (or(2rk((z x)) + 3. and the conditional en- tropy can be obtained by simply replacing the kernel & with kx (z, .z) k(x,z) (zz, ()X,(XXXX + +2
Nevertheless, we can leverage the approximation of neural networks as linear functions over an embedding space (Ts, c; 0) = B (e(s,c), c), where both weights B and em- beddings (- exist in p-dimensional latent space (Lee et al., 2019; Khan et al., 2019). This technique does not violate any of the practical constraints listed above, and allows us to adapt the machinery introduced in GP set- tings. While several embedding strategies exist (Jacot et al., 2018; Devlin et al., 2019; Holzmüller et al., 2023), we adopt loss gradient embeddings (Ash et al., 2020). Assuming the prior 33 ~ NO, I), the policy π T(S, 8) can be mod- eled by Gaussian Process with kernel kg ((8, c), (s',c) = (Øo(s, c), áo c'))- When coupled with this approxima- tion, the conditional entropy objective in Equation 2 can be reformulated:
(xX, z),,
where of is the variance of the observation noise c(s, c), as- suming it is distributed according to a zero-mean Gaussian.
Algorithm AMF
H
Input: initial policy nu. budget N, desired task distr. Pe Output: fine-tuned policy TT
Cn arg min d'EC
E cope (sg.... (or(e)
((B,,c,)),
(4)
t=
Thin TI
Initialize dataset Da = 0
{'wt(d))
for n [0...., N - 1] do
where K (. (X, x.), 2t = (st, and X is the vector of states and tasks in (c', 11m As the collected dataset grows, the conditioning on previous trajectories Timn-1 can instead be addressed by fine-tuning the network's 's parameters (e.g.. through conventional gradient descent), resulting in updates in the embedding function de-
Compute c, as the solution to Eq. 2 Collect new demonstration on for task Cn.
if +1 1 % B= = 0 then
Dn+1= = D P=+1-B U (cn-B+I:n Ta- B+1mm) Update 1 from T-+1-B B with Dn+1
end if
end for
Dealing with forgetting Catastrophic forgetting is an in- herent challenge to neural function approximation under shifts to the training distribution, as optimization over new data is non-local, and may undo learning progress (Mc- Closkey & Cohen, 1989; French, 1999). This issue is critical in our setting, as actively guiding the fine-tuning distribution will necessarily accentuate the distribution shift from pre- training. Common strategies for its mitigation often involve rehearsal (Atkinson et al., 2021; Verwimp et al., 2021) or regularization (Kirkpatrick et al., 2017). Unfortunately, the former is not possible in this setting due to lack of access to pre- trainiiiing data, and the latter was not found to be em- pirically effective (see Appendix H). Scale and a diverse pre-training dataset can also mitigate forgetting (Ramasesh et al., 2022), but neither can be controlled during fine-tuning
Thus, when the policy can be modeled as a GP, the only approximation needed concerns occupancy estimation. We refer to this first, practical instantiation as AMF-GP, and present a general algorithmic framework in Algorithm Application of the method to policies parameterized by neu- ral networks will adopt the same scheme. It will also require additional care on two distinct topics: kernel approximtions and mitigation of catastrophic forgetting.
Kernel approximations When the policy is parameter- ized through a neural network, estimation of the conditional entropy is far less straightforward. First, we cannot as- sume the availability of ad-hoc techniques for uncertainty estimation (e.g., Dropout (Srivastava et al., 2014; Gal & Ghahramani, 2016) or ensembles (Lakshminarayanan et al., 2017)). as they might not be featured in pre-trained models.
We thus propose a practical algorithmic solution to alleviate catastrophic forgetting when fine-tuning multi-task policies. Intuitively, we would like the policy to retain the skills it mastered during pre-training, while focusing on novel tasks during fine-tuning. Inspired by previous works in behavioral
For simplicity, we consider single-output GP, but general- ize to multi--dimessoonal policies with multi-output GPs in both experiments and formal proofs.
5

---

Active Multi-task Policy Fine-tuning
Average return
Average entropy
Average return
Average return
Average entropy
-7.00
@ 40 demonstrations
-7.75
0.610
-7.5
0.80
-8.00
0.605
-10
0.75
8.25
0,600
0.70
8.50
0.505
-12
0.65
-875
-9.0
9.00
12
10
2
1
-14
0.60
Rounds 0
25
of pretraining tasks
0 25 Rounds 0
25
No mismatch (12/12 tasks demonstrated)
AMF-GP
Uniform
Mismatch (6/12 tasks demonstrated)
Figure 2: Experiments in GP settings for a 2D integrator (see Figure 3). AMF-GP selects tasks that minimize the policy's 's posterior entropy and improves the agent's returns faster than uniform task sampling. the middle, the improvement in final return over the baseline is greater when the pre-training distribution does not perfectly match the evaluation distribution pe and only demonstrates fewer tasks. We report return and entropy eurves for non-mismatched and mismatched pre-training (left and right, respectively). We plot means and 90% simple bootstrap confidence intervals over 10 random seeds; dots and crosses mark corresponding measurements.
priors (Bagatella et al.. 2022) and offline RL (Kumar et al., 2020), we propose retain a copy of the pre-trained policy, which we refer to as prior (TT)). For a given state E S and task e E C, we can then linearly combine actions sampled from the fine-tuned policy π with those sampled from the prior:
faster empirically faster convergence with respect to naive approaches to data collection. To do so, we compare AMF to uniform i.i.d. sampling from the set of tasks C. First, we consider a classic 2D integrator as a benchmark environ- ment (see Figure 3). The agent a pointmass initialized in the origin, and can directly control its 2D velocity, which is integrated over the past trajectory return the current state. We can define a continuous task space, in which each task consists of reaching a point on a circle centered on the ori- gin, and the agent is rewarded with the negative Euclidean distance to it. The evaluation distribution Me assigns equal probability to 12 points in different directions. The initial state distribution deterministie, dynamics are both deter- ministic and smooth, while the expert policy is smooth and corrupted with i.i.d. Gaussian noise. We model the policy as a Gaussian Process with RBF kernel, and we condition it on a pre-training dataset of 12 noisy demonstrations. We then collect 40 additional demonstrations by running both AMF-GP and uniform sampling.
( (1 - ac)),, where à c-)s,,)), Tp(-|s.c). Crucially, a(c) E [o, 1] is a task-dependent weight trained through gradient descend on an proxy behavior cloning loss, with an additional conservative penalty that encourages closeness to 0. As a result, actions will drift towards the fine-tuned policy's 's output as soon as it robustly outperforms pre-training performance on a given task. Tasks which are not improved during fine-tuning will rely instead on samples from the prior, and will not be forgotten. We refer to this technique as Adaptive Prior, and present a detailed description and comparison to continual learning techniques in Appendix H.
By combining the approximations required by AMF-GP with the described solutions for estimating entropy and pre- venting catastrophic forgetting, we obtain a method for ac- tive multi-task fine-tuning of policies parameterized via neural networks, which we refer to as AMF-NN.
As a sanity check, we first con- sider perfectly uniform pre- training regime, in which each evaluation task is demonstrated exactly once. The pre-training task distribution juD E ACC) thus perfectly matches the eval- uation task distribution pe (Fig- ure 2, left). As it actively minimizes the policy's entropy. AMF-GP increases the policy's returns slightly faster when compared to uniform sampling of demonstrations. We then ex- tend this evaluation to more re-
5. Experiments
The experiment section is designed to evaluate active multi- task fine-tuning and provide an empirical answer to several questions. We thus reserve section to each of them. Addi- tional evaluations are reported in Appendix D, E and F.
Figure 3: 2D integrator. Starting from the origin, each task involves reaching given point on a circle, as shown by differently colored trajectories.
5.1. When is AMF beneficial?
When none of the assumptions listed in Section 4.1 is vio- lated, AMF is guaranteed to converge to the optimal policy. We furthermore investigate whether AMF also results in
6

---

Active Multi-task Policy Finc-tuning
State-based
Visual
Kitchen, mismatch
Kitchen, no mismatch
Kitchen, mismatch
Kitchen, no mismatch
0.5
Muli-tasf success rate
rate
0.85
Muki-tus suocess rate
Multitask success rate
0.80
0.4
066
0.75
0.3
0:35
0.70
0.30
0.4
0.2
0.65
0.20
0.60
0.1
0
10
20
5
10
Rounds
Rounds
Rounds Metaworld. mismatch
Rounds
Metaworld, mismatch
Metaworld, no mismatch
Metaworld. mismatch
Multi-usk success rate
0.7
Multi-isk -ask-iskk success rale
Multi took success rate
0.6
066
0.6
0.6
0.
0.5
0.5
0.5
0.4
Multi-tak
0.4
0.4
044
0.3
03
0.2
0
10 20
20
10
5 Rounds
10
Rounds
Rounds
Rounds
AMF-NN Uniform, with adaptive prior
Uniform
Figure 4: AMF with neural policies in Frankakitchen (top) and Metaworld (bottom). Experiments are repeated for state and RGB inputs (left and right). We evaluate both mismatched and monmmmsatt settings. AMF-NN is overall desirable, and highly beneficial for mismatched pre-training distributions. We report means and 90% simple bootstrap confidence intervals over 10 seeds.
alistic pre-training distributions, characterised by a mis- match with respect to the evaluation distribution: Pe # pr (Figure 2, middle), and compare the final performance of the two methods as the pre-training budget is allocated to a decreasing number of tasks. As the pre-training distribution diverges from #e (e.g., when only 6/12 tasks are demon- strated in Figure right), we observe that the performance gap between uniform task sampling and AMF-GP grows larger. This is to be expected, as in this case the informa- tion gain from the next demonstration heavily depends on the queried task, and taking the arg max of the criterion in Equation 1 is significantly better than choosing random task. Intuitively, in this case, uniform sampling of tasks fails to reliably provide demonstrations for tasks that were observed less often during pre-training.
benchmarks for multi-task learning, both with finite set of tasks.
In Metaworld (Yu et al., 2020) we create a scene with a robotic arm, a cup and a faucet, defining tasks: moving the cup to two distinet positions, opening and closing the faucet.
In FrankaKitchen (Fu et al., 2020), consider 5 tasks, namely turning a knob on or off, opening pivoting or a sliding cabinet, or opening the microwave door.
In both environments, we evaluate AMF-NN when learning from state measurements, as well as from raw pixels. In the first case, the policy simply parameterized through a MLP, while in the second the ML. receives the embedding of pre-trained visual encoder (Nair et al., 2022). The policy is pre-trained on ~ 15 total demonstrations, which we allocate either uniformly across all tasks, or only on half of them, reproducing the non-mismatched and mismatched regimes from the previous experiments. Afterwards, when learning from state measurements, we apply AMF-NN for 20 iterations, collecting one demonstrations at each iteration. To compensate for the increased complexity, in visual settings we instead provide 2 demonstrations per each task selected and only evaluate 10 iterations.
5.2. Can AMF scale high-dimensional tasks?
In realistic settings, the assumptions enabling formal anal- ysis of AMF are soon violated. As the complexity of the environments of interest increases, most modern behavior cloning applications rely on neural networks for policy pa- rameterization (Reed et al., 2022; Chi et al., 2023). Moti- vated by this pattern, we now study a second version of our method, AMF-NN, and evaluate its ability to scale to com- plex, high-dimensional tasks. We consider two common
Figure 4 reports average multi-task success rates at each iter-
7

---

Active Multi-task Policy Finc-tuning
09
Kitchen, mismatch
Kitchen, no mismatch
Robomimic, mismatch
Robomimic, bo mismatch
Multi Saltit-nsk success rate
Multi-aak saccesss rate
0.85
Multi-tukk seccess rale
0.7
Molis-tukk ralee
07
0.80
06
0.6
06
0.75
0.5
0.5
0.5
0.70
04
04
0.
0.63
0.3
0.3
0.3
0.00
10
20
Rounds
Rounds
Rounds
Rounds
Last-layer gradient Last-layer
Dropout
AME-NN
Uniform, with adaptive prior
Uniform
Figure 5: AMF performance with alternative
Figure 6: Evaluation on Robomimic. AMF-NN widely applicable (e.g,, to diffusion policies).
uncertainty estimation schemes.
ation compared to a random uniform task selection scheme. For the baseline, we additionally report performance with an adaptive prior, highlighting how this solution prevents per- formance degradation due to catastrophic forgetting of pre- training demonstrations, As reported in the previous section, AMF beneficial under distribution mismatch, in which the pre-training dataset does not exactly cover the evalua- tion distribution Mic. As sanity check, we also observe that, when pre-training demonstrates all tasks equally and a uniform task allocation would be very effective, AMF's performance matches this naive baseline.
Robomimic benchmark, which involves four long-horizon, precise manipulation tasks (up to =000 steps). While exper- iments in previous sections rely on demonstrations collected by scripted policies or RL agents, expert trajectories are in this case provided by humans; due to the increased scale, we sample 20 demonstrations for the task selected each iteration. Finally, we fine-tune larger generative model, namely diffusion policy (Chi et al., 2023), which remains compatible with AMF-NN.
Despite the change in data source and architectures, the results we observe in Figure 6 are consistent with those reported in previous settings: active fine-tuning and an adaptive prior are overall helpful, especially when the pre- training distribution does not match the evaluation distribu- tion PCC-
These trends are consistent across both environments, and both modalities. For qualitative analysis of the strategy induced by AMF, we refer Appendix and J.
5.3. How do uncertainty estimates for AMF compare?
6. Discussion
As entropy estimation is at the core of AMF-NN, we addi- tionally compare the adopted GP approximation with loss- gradient embeddings to other approaches from the liter- ature. In particular, we also consider an alternative GP approximation using last-layer embeddings (Holzmüller et al., 2023), as well as test-time Dropout (Loquercio et al., 2020). The latter simply selects the task maximizing prior entropy, that arg max E E (((T(Se,C) T1mm-1), with 71:n-1 (Ciin-1) and (su....) these schemes are in practice desirable, as they do not re- quire access to action labels. However, we observe that these two schemes are less effective in driving task selec- tion. Hence, as shown in Figure 5 (and in Appendix G), multi-task performance in general lower, suggesting that the entropy estimation technique is important for AMF-NN.
As generalist robotic policies gain prominence, a new set of challenges and opportunities emerge. This work responds this trend by investigating active multi-task fine-tuning scheme, which adaptively selects the task to be demonstrated for sample-efficient multi-task behavioral cloning. This approach is developed from first principles, extending a formally-motivvteed motivated, information-base criterion to trajectories over dynamic systems. The resulting method is both formally supported by novel performance guarantees and widely applicable. Moreover, practical instantiation enables sample-efficient multi-task fine-tuning across GP and neural network policy classes.
T(c). Both of
Naturally, active multi-task fine-tuning has several limita- tions. When coupled with neural networks, the algorithm relies on uncertainty estimation techniques, which remain an open problem. While the approximation we leverage is informative in experiments, AMF could benefit if large pre-trained policies would allow other off-the-shelf uncertainty quantification techniques (e.g., through model ensembling during pre-training), Second, we found the per- formance of AMF to depend naturally on the pre-training
5.4. Can AMF be applied to off-the-shelf models?
As AME-NN has minimal requirements (essentially, access to differentiable pre-trained prior is sufficient), it should be widely applicable. In this section we investigate further scaling evaluation to more complex tasks, multi-moda demonstrators and modern policy classes. We consider the
8

---

Active Multi-task Policy Finc-tuning
data distribution. While AMF induces efficient learning for mismatched pre-training distributions, it naturally brings more modest gains when the pre-trained policy is equally capable for all tasks, and uniform task sampling is sufficient.
Bagatella, M., Christen, S. J., and Hilliges, O. Sfp: State- free priors for exploration in off-policy reinforcement learning, In TMLR, 2022.
Bain, M. and Sammut, C. A framework for behavioural cloning, In Machine Intelligence, volume 15, 1995.
On top of addressing the current limitations, this work sug- gests multiple interesting directions. An extensive empirical evaluation of active fine-tuning with large-scall generalist policies is clearly desirable, but remains infeasible at the moment due to the scarce availability of open-source cali- brated benchmarks. Another future research direction would involve direct estimation of the RL objective, thus removing the dependence on non-equivalent BC proxy objectives.
Belkhale, S., Cui, Y., and Sadigh, D. Data quality in imita- tion learning. In NeurIPS, volume 36, 2024.
Bickford Smith, F., Kirsch, A., Farquhar, S., Gal, Y., Foster, A., and Rainforth, T. Prediction-orienter bayesian active learning. In AISTATS, 2023.
Block, A., Foster, D. J.. Krishnamurthy, A., Simchowitz, M. and Zhang, C. Butterfly effects of sgd noise: Error amplification in behavior cloning and autoregression. In ICLR, 2024a.
Acknowledgments
We thank Nico Gürtler, Ji Shi, Andreas René Geist and Us- man Namjad for the valuable discussions. Marco Bagatella is supported by the Max Planck ETH Center for Learn- ing Systems. Georg Martius is a member of the Machine Learning Cluster of Excellence, EXC number 2064/1 Project number 390727645. We acknowledge the support from the German Federal Ministry of Education and Re- search (BMBF) through the Tübingen AI Center (FKZ: 011S18039B), from the European Research Council (ERC) under the European Union's 's Horizon 2020 research and Innovation Program Grant agreement no. 815943, and from the Swiss National Science Foundation under NCCR Au- tomation, grant agreement 51NF40 180545.
Block, A., Jadbabaic, A., Pfrommer, D., Simchowitz, M., and Tedrake, R. Provable guarantees for generative behav- ior cloning: Bridging low-level stability and high-level behavior. In NeurlPS, volume 36, 2024b.
Bogunovic, I., Scarlett, J., Krause, A., and Cevher, V. Trun- cated variance reduction: A unified approach to bayesian optimization and level-set estimation. NeurIPS, 29, 2016.
Bousmalis, K., Vezzani, G., Rao, D., Devin, C. M., Lee, A. X., Villalonga, M. B., Davchev, T., Zhou, Y., Gupta, A., Raju, A., Laurens, A., Fantacci, C., Dalibard, V., Zambelli, M., Martins, M. F., Pevceviciute, R., Blokz- ijl, M., Denil, M., Batchelor, N., Lampe, T., Parisotto, E., Zolna, K., Reed, S., Colmenarejo, S. G., Scholz, J.. Abdolmalcki, A., Groth, O., Regli, J.-B., Sushkov, O., Rothörl, T., Chen, J. E., Aytar, Y., Barker, D., Ortiz, J., Riedmiller, M., Springenberg, J. T., Hadsell, R., Nori, F., and Heess, N. Robocat: self-improving generalist agent for robotic manipulation. TMLR, 2024.
Impact Statement
This paper formally introduces framework for active data selection in MDPs, and investigates its applicability in di- verse settings. Although this direction as whole might eventually have an effect on labor and automation, we do not foresee a direct impact of our work in terms of societal consequences.
References
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A. Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. In NeurIPS, volume 33, 2020.
Abbasi-Yadkonl Y. Online learning for linearly parametrized control problems. PhD thesis, University of Alberta, 2013.
Ceron, J. S. O., Courville, A., and Castro, P. S. In value- based deep reinforcement learning, a pruned network is good network. In ICML, 2024a.
Ash, J. T., Zhang, C., Krishnamurthy, A., Langford, J., and Agarwal, A. Deep batch active learning by diverse, uncertain gradient lower bounds. In ICLR, 2020.
Atkinson, C., McCane, B., Szymanski, L., and Robins, A. Pseudo-rehearsal: Achieving deep reinforcement learning without catastrophic forgetting. Neurocomputing 428, 2021.
Ceron, J. S. O., Sokar, G., Willi, T., Lyle, C., Farebrother, J., Foerster, J. N., Dziugaite, G. K., Precup, D., and Castro, P. Mixtures of experts unlock parameter scaling for deep RL. In ICML, 2024b.
Ba, J. L. Layer normalization. NIPS 2016 Deep Learning Symposium, 2016.
Chaloner, K. and Verdinelli, I. Bayesian experimental de- sign: review. Stotistical science, 1995.

---

Active Multi-task Policy Finc-tuning
Chaudhry, A., Rohrbach, M., Elhoseiny, M., Ajanthan, T., Dokania, P. K., Torr, P. H., and Ranzato, M. On tiny episodic memories in continual learning, arXiv preprint WrXiv::902.10480 2019.
Gal, Y. and Ghahramani, Dropout as bayesian approxi- mation: Representing model uncertainty in deep learning. In ICML, 2016.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In CVPR, 2016.
Chen, D., Wang, D., Darrell, T., and Ebrahimi, S. Con- trastive test-time adaptation. CVPR, 2022.
Hendawy, A.. Peters, J., and D Eramo, C. Multi-task rein- forcement learning with mixture of orthogonal experts. In ICLR, 2023.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. In NeurIPS, volume 34, 2021.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. In NeurlPS, volume 33, 2020.
Chi, C., Feng, S., Du, Y., Xu, Z., Cousineau, E., Burchfiel, B., and Song, S. Diffusion policy: Visuomotor policy learning via action diffusion. In RSS, 2023.
Holzmüller, D., Zaverkin, V., Kästner, J., and Steinwart, framework and benchmark for deep batch active learning for regression. JMLR, 24, 2023.
Cho, M., Jung. W., and Sung, Y. Multi-task reinforcement learning with task representation method. In ICLR Work- shop on Generalizable Policy Learning in Physical World, 2022.
Hübotter, J., Sukhija, B., Treven, L., As, Y., and Krause, A. Transductive active learning: Theory and applications. In NeurIPS, 2024.
Hübotter, J., Bongni, S., Hakimi, I., and Krause, A. Effi- ciently learning at test-time: Active fine-tuning of llms. In ICLR, 2025.
Collaboration, O. X.-E.. Open x-embodiment: Robotic learn- ing datasets and RT-x models. In Towards Generalist Robots: Learning Paradigms for Scalable Skill Acquisi- tion @ CoRL2023, 2023.
https:// / wtw.://www....eeeeeeeeeeeaaaaaaa all ail-MiniLM-16-- 2026.
HuggingFace.
all-minilm-l6-v2. huggingface.com
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. In NAACL, 2019.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In NeurIPS, volume 31, 2018.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x 16 words: Transformers for image recognition at scale. In ICLR, 2021.
Jain, V. and Learned- E. Online domain adaptation pre-trained cascade of classifiers. In CVPR, 2011.
Kakade, S. and Langford, J. Approximately optimal approx- imate reinforcement learning. In ICML, 2002.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, Podell, D., Dockhorn, T., English, Z., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In ICML, 2024.
Khan, M. E. E., Immer, A., Abedi, E., and Korzepa, M. Approximate inference turns deep networks into gaussian processes. In NeurlPS, volume 32, 2019.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In ICML, 2017.
Kirkpatrick, J., Pascanu, R., Rabinowitz, N., Veness, J., Des- jardins, G., Rusu, A. A., Milan, K., Quan, J., Ramalho, T., Halbska-Barwinkke A., et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 1 114, 2017.
Foster, J., Block, A., and Misra, D. Is behavior cloning all you need? understanding horizon in imitation learning. arXiv preprint arXiv:2407.1500 2024.
Kothawade, S., Beck, N., Killamsetty, K., and Iyer, Simi- lar: Submodular information measures based active learn- ing in realistic scenarios. In NeurIPS, 2020.
French, R. M. Catastrophic forgetting in connectionist net- works. Trends in cognitive sciences, 3, 1999.
Kothawade, S., Kaushal, V., Ramakrishnan, G., Bilmes, J., and lyer, R. Prism: A rich class of parameterized submodular information measures for guided data subset selection. In AAAI, 2022.
Fu. J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072/1 072/9 2020.
10

---

Active Multi-task Policy Finc-tuning
Krause, B., Kahembwe, E., Murray, I., and Renals, S. Dy- namic evaluation of neural sequence models. In ICML, 2018.
Mandlekar, A., Xu, D., Wong, J., Nasiriany, S., Wang. C., Kulkarni, R.,, Fei-Fei, L., Savarese, S., Zhu, Y., and Martín-Martún, R. What matters in learning from offline human demonstrations for robot manipulation. In CORL, 2021.
Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In NeurlPS, volume 25, 2012.
Maran, D., Metelli, A. M., and Restelli, M. Tight perfor- mance guarantees of imitator policies with continuous actions. In AAAI, volume 37, 2023.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Conser- vative q-learning for offline reinforcement learning. In NeurlPS, 2020.
McCloskey, M. and Cohen, N. J. Catastrophic interference in connectionist networks: The sequential learning prob- lem. Psychology of learning and motivation, 24, 1989.
Kumar, A., Hong, J., Singh, A., and Levine, S. Should run offline reinforcement learning behavioral cloning? In ICLR, 2022.
Mchta. V.. Paria, B., Schneider, J., Ermon, S., and Neiswanger, W. An experimental design perspective on model-based reinforcement learning. In ICLR, 2022.
Kumar, V., Shah, R., Zhou, G., Moens, V., Caggiano, V., Gupta, A., and Rajeswaran, A. Robohive: unified framework for robot learning. NeurlPS, 36, 2024.
Nair, S.,, Rajeswaran, A., Kumar, V., Finn, C.,, and Gupta, A. R3m: A universal visual representation for robot manipulation. In CORL, 2022.
Lakshminarayaman, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. In NeurlPS, volume 30, 2017.
Nauman, M., Ostaszewski, M., Jankowski, K., Milos, P., and Cygan, M. Bigger, regularized, optimistic: scaling for compute and sumple-efficien continuous control. arXiv preprint arXiv:2405. 16158, 2024.
Lee, J., Xiao, L., Schoenholz. S., Bahri, Y., Novak, R., Sohl- Dickstein, J., and Pennington, J. Wide neural networks of any depth evolve as linear models under gradient descent. In NeurIPS, volume 32, 2019.
Octo Model Team, Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Xu, C., Luo, J., Kreiman, T., Tan, Y., Chen, L. Y., Sanketi, P., Vuong, Q., Xiao, T., Sadigh, D., Finn, C., and Levine, S. Octo: An open-source generalist robot policy, In RSS, 2024.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005. ..00000..663,... 2020.
Li, X., Hsu, K., Gu, J., Pertsch, K., Mees, O., Walke, H. R., Fu, C., Lunawat, I., Sieh, I., Kirmani, S., Levine, S., Wu, J., Finn, C., Su, H., Vuong, Q., and Xiao, T. Evaluat- ing real-world robot manipulation policies in simulation. arXiv preprint WwX.:2244500044111 2024.
Osa, T., Pajarinen, J., Neumann, G., Bagnell, J. A., Abbeel, P., Peters, J., al. An algorithmic perspective on imita- tion learning. Foundations and Trends@ in Robotics, 2018.
Rachelson, F. and Lagoudakis, M. G. On the locality of action domination in sequential decision making. Inter- national Symposium on Artificial Intelligence and Mathe- matics, 2010.
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let's verify step by step. arXiv preprint arXiv:2305.20050. 2023.
Radford, A., Wu, J., Child, R.. Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Loquercio, A., Segu, M., and Scaramuzza, A general framework for uncertainty estimation in deep learning. IEEE Robotics and Automation Letters, 5, 2020.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. In ICLR, 2019.
Lyle, C., Zheng, Z., Nikishin, E., Pires, B. A., Pascanu, R., and Dabney, W. Understanding plasticity in neural networks. In ICML, 2023.
Ramasesh, V. V., Lewkowycz, A., and Dyer, E. Effect of scale catastrophic forgetting in neural networks. ICLR, 2022.
Ma, Y., Song, Z., Zhuang, Y., Hao, J., and King, I. A survey on lasioo--anguage-ace action models for embodied ai. arXiv preprint Arvv:24405 14093, 2024.
Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-maron, G., Giménez, M., Sulsky,

---

Active Multi-task Policy Finc-tuning
Y., Kay, J., Springenberg. J. T., Bccles, T., Bruce, J., Razavi, A., Edwards, A., Heess, N., Chen, Y., Hadsell, R., Vinyals, O., Bordbar, M., and de Freitas, N. A generalist agent. TMLR, 2022.
Wang, C., Sun, S., and Grosse, R. Beyond marginal uncer- tainty: How accurately can bayesian regression models estimate posterior predictive correlations? In AISTATS, 2021a.
Riemer, M., Cases, I., Ajemian, R., Liu, M., Rish, L., Tu, Y., and Tesauro, G. Learning to learn without forgetting by maximizing transfer and minimizing interference. In ICLR, 2019.
Wang, D., Shelhamer, E., Liu, S., Olshausen, B., and Darrell, T. Tent: Fully test-time adaptation by entropy minimiza- tion. In ICLR, 2021b.
Wang, I.., Zhang, X., Su, H., and Zhu, J. A comprehen- sive survey of continual learning: theory, method and application. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Ross, S. and Bagnell, D. Efficient reductions for imitation learning. In AISTATS, 2010.
Ross, S., Gordon, G., and Bagnell, D. A reduction of imitation learning and structured prediction to no-regret online learning. In AISTATS, 2011.
Williams, K. and Rasmussen, C. Gaussian processes for machine learning. MIT press Cambridge, MA, 2006.
Xia, M., Malladi, S., Gururangan, S., Arora, S., and Chen, D. Less: Selecting influential data for targeted instruction tuning. ICML, 2024.
Rotman, G. and Reichart, R. Multi-task active learning for pre-trained transformer-based models. ACL, 2022.
Sener, O. and Savarese, S. Active learning for convolutional neural networks: A core-set approach. In ICLR, 2017.
Xu, T., Li, Z., and Yu, Y. Error bounds of imitating policies and environments. In NeurlPS, volume 33, 2020.
Settles, B.
Active learning literature survey. Technical report, University of Wisconsin-Madison Department of Computer Sciences, 2009.
Yu, T., Quillen, D., He, Z., Julian, Hausman, K., Finn, C., and Levine, S. Meta-world A benchmark and evaluation for multi-task and meta reinforcement learning. In CORL, 2020.
Sodhani, S., Zhang, A., and Pineau, J. Multi-task rein- forcement learning with context-based representations. In ICML, 2021.
Yu, T. Kumar, A., Chebotar, Y., Hausman, K., Levine, S., and Finn, C. Conservative data sharing for multi-task offline reinforcement learning. In NeurlPS, volume 34, 2021.
Spencer, J., Choudhury, S., Venkatraman, A., Ziebart, B., and Bagnell, J. A. Feedback in imitation learning: The three regimes of covariate shift. arXiv preprint arXiv:2102.02872 2021.
Zhao, Z., Kumar, V., Levine, S., and Finn, C. Learning fine ime-grained bimanual manipulation with low-cost hard- ware. In ICML Workshop New Frontiers in Learning, Control, and Dynamical Systems, 2023.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: simple way to prevent neural networks from overfitting. JMLR, 15, 2014.
Sun, L., Zhang, H., Xu, W., and Tomizuka, M. Paco: Parameter-comppstionna multi-task reinforcement learn- ing. In NewrlPS, volume 35, 2022.
Sun, Y., Wang, X., Liu, Z., Miller, J., Efros, A., and Hardt, M. Test-time training with self-supervision for general- ization under distribution shifts. In ICML, 2020.
Teh, Y., Bapst, V., Czarnecki, W. M., Quan, J., Kirkpatrick, J., Hadsell, R., Heess, N., and Pascanu, R. Distral: Robust multitask reinforcement learning. In NeurIPS, volume 30, 2017.
Verwimp, E., De Lange, M., and Tuytelaars, Rehearsal revealed: The limits and merits of revisiting samples in continual learning. In /CCV, 2021.
12

---

Active Multi-task Policy Finc-tuning
A. Performance guarantees under regularity assumptions
This sections retrieves guarantees on the performance of the imitator policy as a function of the number provided demonstrations n. At first, this analysis focuses on policies over single-dimensonaa action space. An extension multi-dimensional outputs is introduced later on. The general sketch of the proof can be informally described as follows:
we first introduce the regularity assumptions required for the guarantees;
then show that, in Lipschitz, bounded MDPs, the effect of stochasticity information gain at each round can be controlled;
we show how the variance over the imitator's policy shrinks according to the maximum information gain at each round, which in turn depends on the maximum information gain over of queries to the expert;
starting from the previous result, we leverage a well-known theorem (Abbasi- Yadukii,, 2013) to retrieve probabilistic, anytime guarantee on the error of the imitator;
we quantify the relationship between the imtator's error and its performance, thus retrieving our main theoretical result.
A.1. Assumptions
It is clear that bounding imitation performance would be hopeless without any regularity assumption, as slight errors in the imitator's 's policy could result in arbitrary differences in return. We thus introduce the following:
Assumption A.1. (Regular, noisy policy) We assume that the optimal policy x² ~ GP(p, k) with known mean function and kernel k. Furthermore the noise e(s, c) is mutually independent and zero-mean Gaussian, with known variance (3 c) for all (s x C.
In order to motivate further assumptions, let us recall the criterion from Equation 1:
H-1
Cn = arg max d' EC
E
Σ
I(m T))) cin-11 T1m-1).
(5)
Thin- m( coeper (90
(or(())
1=
Through this section, we will use slightly more precise formulation:
On arg max
E Then (eLa (d)) oretes (ea (orr(e)
Σ ();
cim-1, Tim-1),
(6)
e'EC
in which we clarify that the mutual information is only computed with respect to the actions of the noisy expert, and overload the notation with #(r, d') (#(s,, (())) for (st, (s))")))" The criterion selects the task e, with the greatest expected mutual information between the policy and the trajectory associated with the task. We note that, the objective produces a fully deterministic sequence of tasks, as all stochasticity is resolved in the expectation. Nevertheless, the actual sequence of states at which the demonstrator is queried remains stochastic. For this reason, we require the following two sets of assumptions to ensure that information gained along empirical trajectories is not arbitrarily smaller than the expected one.
Assumption A.2. (Lipschitz, bounded MDP and policy) Given the contextual MDP M = (S, A, C, P. R, 7, Ho) and the noisy expert हे we assume that, for every {(s, c, a), (d',''a'') C x A:
the support of the initial state distribution μn is bounded by Eno -ball
max 81,8 ESupP(μo)
||5 -5ll22 S Epo
the transition kernel P Lm-shhhth
W(P(Hss.)) S Lp ((((((s))),((d)))
13

---

Active Multi-task Policy Finc-tuning
where d((s, a), (8', a')) |ss s|12 ||a 222,, and W(:-)) is the Wasserstein 1-distance with respect d(;;)) furthermore, the support of P(-|s, a) is bounded by er-ball
max $1.68 (uupp((()ss.o))
- - S CP,
the reward function R is Ln-smooth
[ ,'''' < La d(ss,,,))
where di((s, (((ss,,,))) (, (,','' ||s - 8||2 ||c ll + lla - d'l12.
the noisy expert गे is Smooth
[(((ss....) d')) s Lx d((s,e,a), (s,,,
where (s', a')) l l|a aa'l2 and w(,,) is the 1-Wasserstein distance with respect to d(.,)) )) furthermore, the support of म (-18, is bounded by ex
max Esuppl (-(1,,))
||an - a1lla I cp.
finally, the Q-function for the expert x² --msmooth
Q (8,c,a) - Q (d)) Lo d(s,,a))
We note that smoothness of the noisy expert is guaranteed by construction if the expert x² is -mooth..h.
Assumption A.3. (Smooth MI) For every pair of sequences of trajectories {Tim- 1 S (S x A) (s, C. Cin-1 E C=-1, (S x A)" and e₂ E C, we assume that the mutual information at step is L--smooth with respect the mean square deviation of collected trajectories:
Tim I(T((S c); #(, Cn) C1mm (imm-1)) S Lr d(Ti:n 11m--1)
where d(( d..1............................ ... 1H 1)) (66 d6, . d.m- = + Dmml (>
HH-1 || s 113 +/
is the mean square deviation over the concatenation of trajectories.
A.2. Proof
We first prove that, under Assumptions A.2 and A.3, the effect of stochasticity on the mutual information at step n is bounded.
Lemma A.4. Let Assumptions A.2 and A.3 hold. Fix sequence tasks Cin-1 and consider two arbitrary sequences of trajectories Tim--1 and Tim-1 sampled from T(Cin-1). Fix one state-task pair (s, c) E S x C, one task E C and one trajectory 7 TC..). Let 6n SH + max(Ir, L,)) max(co, Cn Cr). The difference in mutual information when conditioning on the two sequences of trajectories can be bounded:
(F. (( S en.
Proof. Under Assumption A.3 it is sufficient to show that stochasticity in the MDP does not cause the demonstrator's trajectories to deviate excessively. This is a direct consequence of smoothness and boundedness, which we assume in
Assumption A.2, and can be shown by induction. Let us fix task c₂ E C and consider two trajectories T, T' ~ ((2)). For the two initial states (So. *6),, boundedness of the initial state distribution A4 implies that ||so - %6ll2 I Cpo- Now, assuming
14

---

Active Multi-task Policy Finc-tuning
that the distance between two states (5t; 8)) is bounded ||s, - - S eg. we have that
4+1 == || 41121l2
2 ((P(--5r, 0r); P(-]s), a() + 2Ep
(8)
$ Lp. (1|s - s/1l2 + ||: - d,l12) + 22p
(9)
-pp. (a + ||a, d'll2) 2cp
(10)
(
S Lp (4 + (( a.) k(-1s/1, ())) + 2cx) + 22p
(11)
S S (ez + La ||5 - 5/112 + 22x) ) + 22p
(12)
Lp (EL + Li 6, + 22,) + 2pp
(13)
Lp ((Lx + +1 44 2(,) + 2cp
(14)
Lr(1 + [e)c 2[ + (p)
(15)
== Ac + B,
(16)
where Lemma N.2 used in (i) and (ii); Assumption A.2 and the fact that C2 de were used through the rest of the derivation. The recurrence relation can be easily unrolled as
e S A' +
ABB
(17)
i=0
Ao + max (A 1)Bt
(18)
S max(A (A, (c Bt)
(19)
$ (1 + (1 1) + 2t(pcc + EP)) + 21((1 + Lpder +ep)) S (1 + Lp) (1 (1+ + L. ) (2t + Lp) max(eo, cp)))
L1)
(20)
(21)
(22)
(1 2t(1 + 1 +
44,² max Cn,
(23)
thus bounding the L2 distances between states at each step of the trajectory a ||s: st 21l2- We note that the distance between actions can also be easily bound by Lemma N.2: ||a d'lla S L + 2t,, This can in turn be related to distances
over trajectories. Let us fix C1a--1 E C and consider Tim 1, 1 T (cinn-)). We have that
n-1
((1m-11771m-1
|| |sm. s + Ila., dl 112)
(24)
n
H-1
2 + (Toze + 2)x)
(25)
0=
#-1
2 + (Laee + tz)
(26)
0
2 2 +1 +z))
(27)
2(1 + (22)
(28)
H-1
= (2(1 + Iun) 12
Σ
2)
(29)
1
H-1
-2(11 + Lz)(
(30)
1=
15

---

Active Multi-task Policy Finc-tuning
L.)
(31)
V2H(1 + [H
(32)
s V2H(1 LT) 2(H - 1) Lp) (1 + -a cx
(33)
4H² (1+ L.) max(Co Cx₁ Ep)
(34)
s 8H $ (1 map L+)) max(E0, Ex,E))
(35)
Having obtained an upper bound on the distance between sequences of trajectories, the result follows naturally from smoothness of mutual information according to Assumption A.3.
We focus the main result. We start by introducing an important measure, quantifying the maximum information gain at each round:
I max v (c)
max C'EC
E
Σ ((: C1m-1, Tin-1)
(36)
d'EC
N (c)
7 MoT (c')
coupee (Sp.... (CoT(c)
We note that the criterion in Equation 6 takes the arg max of the same quantity I, maximizes As common in the literature (Bogunovic et al., 2016; Kothawade et al.. 2020; Hübotter et al., 2024), we make standard assumption on diminishing informativeness.
Assumption A.5. For each n, i E N with < the maximum information gain at round n is not greater than the maximum information gain at round i:
In TT
This be leveraged show that the expected mutual information is sublinear in the number of rounds From this point, overload the notation and allow policies (e.g., т)) to map vector to random vectors, that is π E (S x C)". Lemma A.6. Under Assumptions A. and A.5, if Co) follows the criterion in Equation 6, then In S w ((Hm)
(2n-1)) for (zo....,
H
where Y(Hn) = maxxçsxC I(T(II x C); (X)).
Proof.
Fn
(37)
(38)
H-1
max e'Ec
:):
cin-1: 7im-1)
(39)
Tha ((pp--
1=
(c)(
coffte (sg.... lw((c)
H-1
ci:-11 7in-1)
(40)
Tin N (1)
1=0
7' (T(()))
oriiles (c0... (o)(()
H-1
E
E
[): C1m T1m--1)
(41)
(so.... (oT(C)
1=el
7in-1 7
1=0
16

---

Active Multi-task Policy Finc-tuning
(11
E
[(T((00 c): 7 (T
(42)
Ten) Tin- 1o (c1m-1)
(80
(()
1=0
H-1
E
Σ
max I( #X:+
(43)
n
XCSC X Hnn
180
(or(())
1=
H-1
E
max I(T(S x ();+(X))
(44)
eNc
XCSX
(50.... (NT(()
X|=Hn
max I(TT(S x (X))
(45)
XCse
[
H
7(H))
(46)
where (i) follows from Assumption A.5, (ii) follows from Equation 6, (iii) is due to the chain rule of mutual information. We note that 7n2 = maxxgsxc, |X|sn [( C); (X)) is sublinear for a large class of GPs. this cases, looser upper bound would be 1/20=
This bound on expected round-wise mutual information can then be leveraged describe how the total variance shrinks over rounds.
Lemma A.7. (Uniform convergence of marginal variance, following Hübotter et al. (2024)) Under Assumption A. A.2 and 3.3.for any n 2 0 and (s,c) E S x C.
2 n
(s, S (1 e,,
min
where 6² (xx c)) (2 c) and Tiin min oess RyNT(C) ler.
Proof.
(s. = Var(TT(s,) c)
Ci:m: riml]
(47)
(Var(T((s,) c)
Cim Timl + (s. c)) - (s),c)
(48)
Var(#(s,c) SlmmiTion
Var (8, (), Tim)
(49)
(i)
S の² log
Vari#(s, c)
Clan, Tim
(50)
Var(# π (s, c), Cl:ns Timn
20 20² (O'II((S,,)))
A ((, Cl:n, 7im)
(51)
-2"
E ((T, (s. cimi Iim) ETAT(C) 1Eer (~)(c)
(52)
(1) 2
E. ILETI(T(S,C)) 7 (s, Ti:))
(53)
Tmin TNT(C)
20²
<
E (TT,I c) Cims T1m)
(54)
Tmin F~J(()
2
E I(Tr(s,)); (T, Ci:1 T1m))
(55)
Tmin T~r(()
2g²
E erukee
E For(()
c); Cisn, Tim)
(56)
Tiin E.
Coitae (sp.... (oT(()
1 1s0.......) (so.... l~T(e)
20
<
E
1sE(so...) )
E (N)((C)
I(T((S,) ((T,c) Cim, I
(57)
2
(No.... )NT((c
17

---

Active Multi-task Policy Fine-tuning
2ō
E ToJ((c) (fo }or(())
1 ss IsE(so...IIIT((s, [(T((, ( ((, Cim. Timn)
(58)
min
20
H-1
E oxpe TNTT(C) (#0 (CT(()
1 IE(
Σ
[(T (T. Cimi Tim)
(59)
1=0
25
#-1
E ompe TNT(CC
Σ
[(T (); Cis: Tim)
(60)
7 min
t=0
(fo (or((c)
(8))
2
H-1
S
(1 + nn
E
Σ 1=0
I(T((s, 6i Tim)
(61)
enppe TNT((C)
7
(so... (oN(())
20
= (1 + En)
E
(): cla: Tim)
(62)
2
min
Tin 1
TNT(C)
empler (80 (or((c)
20
S (1 + En
max C'EC
(): Cran; Tisa)
(63)
m
Tin- 101 (1)m P'NT(C') coppor (5 (NN((C)
20
(1 + Cn)
(64)
min
where (i) follows from Lemma N.1 and monotonicity of variance, (ii) holds as the state s is within the support of T(-cc) and (iii) follows from Lemma A.4 as the difference between the expected mutual information and the mutual information for a realized trajectory less than the difference in mutual information for two arbitrary realized trajectories.
This result can then be translated to the agnostic setting, for regular policy '², which we still model through the stochastic process π. Without loss of generality we will assume that the prior variance is bounded by Var (s, S 1.
Lemma A.8. (Well-calibrated confidence intervals, following Abbasi- Yadkori (2013)) Pick 8 E (0, 1). Assume that x² lies in the RKHS (C)(C) of the kernel & with norm ||= < oc, the noise c, is conditionally p-sub-Gaussin sublinear in n. Let B,( = || + P (((()))) + + log(1/6]). Then, for any n > and (s, c) E S x C. GP k) an all-time well-calbbrated model of n² Thus, jointly with probubility at least 1 5,
and Yn is
|m" (s, (s(s,c)) Bn(8)a
We note that (6) depends on (H)) as Hn samples from the demonstrator's policy are collected up to round n. Combining Lemmas A.7 and A.8 we easily get for all (s, e) E S x C and 2 0 with probability 1 d:
Lemma A.8 | -s,
Lemma A.7
Ba(d) ( (1+ Gv)
2 min
(6)0)
(65)
While the analysis has so far dealt with scalar n", a simple union bound can guarantee that
(s, un,s s (6))oa) (a (1+te))
21
(66)
m
with probability at least 1 for an action space of dimension [A]. where now (5) = | + (()(() + +1+ log( (A/6)))). From now on, we will refer to nn as nn. We are thus able to globally bound the L1 distance of the imitator policy with respect the expert policy with high probability under active fine-tuning.
18

---

Active Multi-task Policy Finc-tuning
It is clear that, even if this distance is small, the performance of an imitator which does not exactly match the expert (s,((8,)) c) A n² (s, c) for some (s, c) E 8 x C) can be arbitrarily low for arbitrary MDPs. It is however possible to show that, as long as the Q-function of the expert is smooth, the performance gap to the expert can be controlled. We note that, in case (L + ) < 1, then the Q-function Q is guaranteed to be Lo-Lipschitz continuous with LQ S [->Lp(I+La.) (Rachelson & Lagoudakis, 2010). If smoothness holds, it is casy to connect the divergences in action space to performance gaps (Maran et al., 2023).
Lemma A.9. Let and x' denote two deterministic policies. If the state-action value function Q"' is Laz -iischiit continuous, then:
J - Jm S Lqq 1-7 Exnd [|| - (Tss c)lla)-
Proof.
Given a function f: 7:A+> R, we denote the Lipschitz semi-norm [If(-))l= a'
[f(a)-((a))) 1a-aall
We have:
J
Bondd
E (awr(1 (s)
(67)
E BaNt (s, (o0 V (1,
Band
(68)
1
/ T(als)Q* (0,00((ss,,,, (8,c,a)- (s,c)
E Exnd
(69)
- ?
1
Egnd
JOEA Qt (s,c,a)tT s.c) - x (())s)))
(70)
- ?
1
Exedt
sup JaEA s,EESc
()( (.,) - x'(a s.c]]
(71)
7
(1)
Bend
sup s.cESxc
(((s s,,), (1(,,))))
(72)
(11)) L La Bannt WWT((- x' s, s,)))]
(73)
(1)) Lq Eg [|| s, - x' s, c)|l
(74)
where (i) follows from the performance difference lemma (Kakade & Langford, 2002), (ii) follows from the definition of La Wasserstein distance, (iii) holds LQ- 2 sups, Q" (s, e, l1 and (iv) follows from both policies being deterministic. The proof is completed by taking the absolute value on both sides.
So far, we have shown rates of convergence for the imitator, and connected its error performance. Our main formal result can be shown by coordinating the lemmas so far presented.
Theorem A.10. (Performance guarantees for active multi tasti-tasskae BC) Let Assumptions A.2, A.3 and A.5 hold. Pick E (0. 1). Assume that x² lies in the RKHS H4 (C) of the kernel k with norm ||< oc, the noise cn is conditionally o-sub-Gaussian and Yn is sublinear in n. If each demonstrated task is selected according to the criterion in Equation 1, then with probability at least - 8 the performance difference between the expert policy n² and the imitator policy w., after demonstrations can be upper bounded:
JR Jnn S V2L0- 11611 (1++e)))(())) 1
(7(H))) Vn,
Tmin (1 77)
where on = = + max Lx, Lp)) cax( Furthermore, if O(log n) (e.g., for linear kernels), then Ja Jrx "300 0.
19

---

Active Multi-task Policy Fine-tuning
Proof.
JT
J J
(75)
Lemma A.9
Lon*
S
Bendit ([) x (s,
(76)
1-ÿ
Lemma A.7,A.8 LIr
3,16610011 (a +u++nn en) 2r 3
S
(77)
1-9
m
q 1ō11 mmin (1 -7
(1++6))3(00T.)
(78)
where (i) is due to the fact that J" 2 J for any policy и, and the expectation fades due to uniform convergence. The only terms with a dependency notation in the Theorem. If 7/2 (log n), then J J = (log n /////// 3300 0. For summary of magnitudes of Yn for common kernels, we refer to Table 3 in Hübotter et al. (2024).
are 8, (6) (((())) and I'2 ((T(H))) (mn, which can be combined in the asymptotic
B. Guarantees in non-Lipschitz MDPs
The main result reported in Theorem A. provides anytime guarantees on the agent's performance, assuming smoothness in the MDP. However, is possible replace this assumption with weaker one, at the cost of only retaining guarantees in expectation. This weaker version of the theorem be retrieved by simply assuming smoothness on the noise, rather than on the MDP, and leveraging results recently presented by Maran et al. (2023). Assumption B.1. The noise distribution e is Ty-ishiis continuous.
This assumption is satisfied by large class of Gaussian and sub- Gaussian distributions (Maran al., 2023). We can build upon Assumption A.5 and Lemma A.6, and start by providing weaker version of Lemma A.7. Lemma B.2. (Uniform convergence of marginal variance in expectation) Under Assumption A.I, for any n 2 0 0nd (s,c) E S x C,
20
E
on
c) S
2
Ti ((
min
where 5² max(s,))sxx (s),c) + p² (s, c) and Timl ming,cesxe (Twr(() 1.Er.
Proof. We resume from Inequality 60 in the proof of Lemma A.7:
202
#-1
( c) <
E
Σ
I(m (s) (), ci:: Tim)
(79)
2
TNTlCe) (eo.... (oT((c)
1=
Therefore,
2 H-1
(8) S
E
Σ
[(T((), cimy
(80)
(()
Tin 10 (oT((C)
T
t=0
ompiee (s0.... l~T(()
202 H
< max CEc
Σ
[(T (): #(T,6)) Clm, T1m)
(81)
2
Tun ( 1) P'N((')
min
1=
enperr (No.... )or((c)
27 2r,
(82)
min
20

---

Active Multi-task Policy Fine-tuning
Having bounded variance at each round, this time in expectation, invoke Lemma A.8 bound the expected distance the optimal policy with high probability:
E Thin- Ti
(s, $
(83)
2
min
Instead of leveraging bounds for the imitator's performance in Supschizz-smool settings, we can instead use the fact that the expert's actions are corrupted by smooth noise. this setting, is instead possible to control the suboptimality the imitator with respect the noisy expert. We report the following Theorem from Maran et al. (2023), and refer to the original work for the proof.
Lemma B.3. Let nt, it and π denote the expert, noisy expert and imitator policy, respectively. If Assumption B.l holds, then:
J J 2LeQ 1-y Benua (W(= (- [8)]]
where Q maxis.c.aeSxxxx. as) and W represents the Wasserstein 1-distance.
As the expert x* and the imitator n, are both deterministic, this implies that
2L.Q E Evnue (1 s), T( s))
J J 1-7
(84)
By invoking this Lemma, we can thus conclude that, with probability at least 1
2 L Lmx 10 (0)r = =O( B).
J Ja <
(85)
m (1 - -7)
7mm- 1 (CLm-))
Therefore, if 7xx = Ologn)) then Eren-r-(((um-)) - -== O only hold in expectation, they arise from minimal assumptions, mostly regarding the policy class and the perturbation noise, and can thus be applied to arbitrary MDPs.
200 0. While these performance guarantees
vn
C. Practical Objective
Following up the approximations reported in Section 4.2, present the empirical estimate of the objective that is used through experiments. In particular, we show how the expectations in Equation may be approximated with finite samples. The original criterion is expressed as
H-1
C arg min (c)) = arg min clec
E The ( 1), ('m((e) enptos (S0.... (c)
H(T(s(, Cin ,,,,,einnn Tmn-1).
(86)
e'Ec
empirical estimate can be derived as follows:
21

---

Active Multi-task Policy Fine-tuning
H
on(c)=
H((T()),)) c) (e,.)),,,,.....m..... 1m--1,7mm))
(87)
Tim Com-107te 10 (Cim- r' (c))
t=
coupee (So... (c)
H 111
E
H(T((s,, T (,m-
(88)
P()
t=
Coiftee (80.... (o)
H-1
(66)
Σ
H(TT((,.)) 1' Tim-1)
(89)
[C)
OEC (so... (NT((C)
1=0
T(Tlc')) T( T(sp..... (c)
H-1
Σ
H(T(();,) |i 1im-1, cum-i 71m-1)
(90)
10| OEC (#D.... Jou
r'n
A(-') f((00....)
w(r (((T'c))((((00 ),
Σ t, Tin-1)
1mm--
(91)
C
r'nnt fout
CEC (so....
H-1
(i3)
Σ wo(T',c)u(( .,,c)
Σ H(I((), I 10m-,,,m1,,,, 7m--),
(92)
Cn - 1)²
eEC r'Etinnl (N0 JEfin-1
where (i) uses single sample to estimate the expectation over past trajectories, (ii) uses a sample-based approximation to the target task distribution Pe, and (iii) uses the importance sampling trick introduced in Section 4.2, with w(T,c)
(1-1)111 *(a.ls..)
Sn -1 #(as|ar,0)) This final approximate objective does not involve expectations, and can be efficiently computed. The complexity of evaluating the criterion for single task é' scales linearly with the number of samples in ć and quadratically with the number of rounds However, the dependency on the number of rounds can be removed by evaluating the second sum fixed number of trajectories sampled among T1:n-1. ensuring that the complexity does not depend on the round.
D. A Additional results for AMF-GP
Figure only reports full return curves for two representative pre-training settings, namely those involving 6/12 and 12/12 demonstrated tasks. We here report full results for each task allocation, spanning from 1/12 to 12/12 12 demonstrated tasks. For each setting, we report both average multi-task return and average policy entropy curves.
E. Additional results for AMF-NN
Results in Figure 4 are computed over two representative pre-training distributions: one allocating pre-training demon- strations uniformly over all tasks, the other one only demonstrating the first two tasks. We report these results again, and compare them with those for several other pre-training distributions which tasks have been shuffled. Results are reported for FrankaKitchen in Figure and for Metaworld in Figure and are consistent with patterns observed Figure 4.
22

---

Active Multi-task Policy Fine-tuning
Average return
Average x entropy
Average return
Average x entropy
Average return
Average π entropy
-8
1.0
-8
1.0
10
100
-10
0.9
0.9
0.9
12
12
14
0.8
14
0.8
14
16
0.7
16
0.7
0.7
16
-20
0.6
-18
0.6
18
0.6
25
25
0
25
25
0
25
0
25
Rounds
Rounds
Rounds
1/12 tasks
2/12 tasks
3/12 tasks
Average return
Average π entropy
Average return
Average x entropy
Average return
Average x entropy
8
0.80
0.80
0.80
0.75
-10
0.75
-10
0.75
-12
0,70
0.70
0.70
-12
-12
14
0.65
0.65
0.65
-14
16
0.60
0.60
-14
0.60
0.55
-16
0
25
0
25
0
25
0
25
0
25
25
Rounds
Rounds
Rounds
4/12 tasks
5/12 tasks
6/12 tasks
Average return
Average π entropy
Average return
Average T' entropy
Average return
Average x entropy
0.85
8
88
0.80
0.80
0.80
0.75
0.75
0.75
10
11
0.70
10
0.70
100
0.70
0.65
111
0.65
0.65
11
-33
0.60
-12
0.60
0.60
-14
-13
0.55
25
25
25
0
25
0
25
0
25
Rounds
Rounds
Rounds
7/12 tasks
8/12 tasks
9/12 tasks
Average return
Average π entropy
Average return
Average π entropy
Average return
Average π entropy
0.85
-7.
-7.50
88
0.610
0.610
0.80
-7.75
-8.0
0.75
0.605
8.00
0.605
-8.5
0.600
-825
0.600
-10
0.70
-9.0
0.595
88.0
0.595
0.65
-9.5
0.590
8.75
0.60
-12
-10.0
0.585
-9.00
0.585
25
25
0
25
0
25
0
25
0
25
Rounds
Rounds
Rounds
10/12 tasks
11/12 tasks
12/12 tasks
AMF-GP
Uniform
Figure 7: Additional results in GP settings for 2D integrator (see Figure 3). AMF-GP results in improved sample efficiency across all pre-training regimes, and is particularly effective for skewed pre-training distributions (e.g., when pre-training demonstrations have been allocated to 1/12 or 6/12 tasks).
23

---

Active Multi-task Policy Finc-tuning
Kitchen, mismateh
Kitchen, mismatch
Kitchen, mismatch
Kitchen, mismatch
Kitchen, mismatch
0.9
Multi-takk success rate
03
rate rate
Multi task success rate
0.8
0.8
Multi Multi-task suovess rate
0.85
Molti task soccess
Multi-tusk success
0.7
0.90
0.
0.6
0.6
0.6
0.75
04
0.4
0.4
0.5
0.70
0.63
0.2
0.3
0.2
20
Rounds
Rounds
Rounds
Rounds
Rounds
[8 0,
[0, 10,8,8,00 0,
0,,00088,8
33 3,
AMF-NN
Uniform, with adaptive prior
Uniform
Figure 8: Additional results for AMF-NN in FrankaKitchen with state inputs. We evaluate several allocations of the pre-training demonstrations, as labeled below each plot (e.g., the label [8, 8, O, O, 0] indicates that 8 demonstrations were provided for each of the first two tasks each, and none for the remaining tasks).
Metaworid, mismatch
Metaworld, mismatch
Metaworld, mismatch
Metaworld, no mismatch
0.8
08
0.7
Multi-tsk suocess rute
0.7
Multi-takk success rate
0.7
Multi-task success rate
Molti-task success rate
0.6
0.6
0.6
0.5
0.5
0.5
0.4
0,4
0.2
0.3
20
0
10
20
Rounds
Rounds
Rounds
Rounds
[8, 0
[0,8, 8,
[0, 00,0.80 8,
[2, 2, 2, 2]
AME-NN (a = 3) Uniform, with adaptive prior
Uniform
Figure 9: Additional results for AMF-NN in Metaworld with state inputs. We evaluate several allocations of the pre-training demonstrations, labeled below each plot (e.g., the label [8, 8, O, 0] indicates that 8 demonstrations were provided for each of the first two tasks each, and none for the remaining tasks).
WidowX
F. AMF with generalist policies
70
This section investigates scaling evaluation to recently published open-source generalist policies. For this pur- pose, we choose Octo (Octo Model Team et al., 2024). This model relies on a transformer backbone for integrat- ing multimodal information (in the form state sensors, camera images and text or RGB task descriptions), and uses diffusion-based policy head for action prediction (Chi et al., 2023). For computational reasons, we will focus on fine-tuning the action head alone. Octo is pre- trained on a large-scale real-world robotic dataset (Col- laboration, 2023), and thus designed for inference on physical hardware. Nonetheless, a recently proposed eval- uation suite enables simulated evaluations that statistically
Mulli-task success
60
50
40
AMF-NN
Uniform, with adaptive prior
Uniform
Rounds
Figure 10: Evaluation on life-like WidowX tasks. AME-NN can be applied to large-scale settings.
24

---

Active Multi-task Policy Finc-tuning
correlate with real-world results (Li et al., 2024). We thus collect rollouts from a pre-trained Octo agent on the WidowX tasks, and filter them only include successes, akin self-distillation schemes (Bousmalis et al., 2024). On availability of such self-supervised demonstrations, we then apply AMF-NN for 5 iterations, providing demonstrations in each round. The results are reported in Figure 10. As all evaluation tasks are largely demonstrated in the pre- pre-training dataset (Collaboration, 2023), and the evaluation is significantly noisier with respect to other benchmarks, we find that the performance of AMF-NN falls within confidence intervals of uniform task collection, confirming the trend we observed for uniform pre-training distributions in Figure 4. Nonetheless, we observe that it constitutes an effective method for data selection, and can be applied as a drop-in replacement for fine-tuning of off sheelf models."
G. Uncertainty ablation
Metaworld, mismatch
Metaworld, no mismatch
Section 5.3 evaluates alternative uncertainty quantification schemes in FrankaKitchen for two representative pre-training distributions. This Section extends these results to include re- sults for Metaworld (see Figure 11). Results are consistent with those so far reported, suggesting that loss gradient embeddings an important component for the empirical performance of AMF-NN.
0.8
Multi-taek successs rate
0.
0.7
05
Multi-rak
0.5
0.4
0.4
0.3
10
0
10
20
H. Mitigating forgetting
Rounds
Rounds
The ability of neural networks to adapt to shifts in training dis- tribution while retaining information is an important object of interest in lifelong and continual learning (Wang et al., 2024). In general, learned models display trade-off' between their ability to integrate novel information, and their memory of previously observed training samples. Arguably, common neural network architectures can easily fit new data (save for loss of plasticity (Lyle et al., 2023)), but are known to forget previous information, often catastrophically. This problem is of utmost relevance in our setting, which the pre-trained network not just leveraged as a useful initialization, but may already capable of solving some tasks. Hence, the fine-tuning procedure should be careful not disrupt this ability.
Last-layer gradient
-Last-layer
Dropout
Figure 11: Additional results for AMF-NN in Meta- world with state inputs and different uncertainty quan- tification techniques.
Several methods aimed at mitigating forgetting can be traced back to rehearsal (Riemer et al., 2019; Chaudhry et al., 2019) and regularization (Kirkpatrick et al., 2017) strategies. While rehearsal approaches are often effective, they also require access to pre-training data, which is unrealistic in our setting. Hence we consider two common regularization technique,
Thii evaluation also reports an interesting trend, that is vast reduction in catastrophic forgetting, to the point that an adaptive prior is almost not necessary. This anecdotal evidence can be seen instance of general trend of mitigated catastrophic forgetting in large models (Ramasesh et al., 2022).
Kitchen, mismatch
Kitchen, mismatch
Metaworld, mismatch
Metaworld, no mismatch
Multi-task sugcess rate
0.8
Multi-usk -usk-usk sMeeess rute
0.85
Multi-task success rate
Molti-task success rite
8.6
0.6
06
0.75
0.5
0.4
0.70
0.4
0.65
0.2
0.
0.2
0.60
10
20
10
10
20
Rounds
Rounds
Rounds
Rounds
Adaptive Prior
EWC
L2 regularization
Baseline
Figure 12: Performance of AMF-NN with several techniques to mitigate forgetting. Darker shades represent stronger regularization coefficients. L2, EWC regularization are not effective in this setting.
25

---

Active Multi-task Policy Finc-tuning
Sliding door
Microwave
Left door
Knob off
Knob on
1.0
1.0
0.4
1.0
1.0
1.0
1.0
1.0
1.0
Single-task success rate
0.8
0.6
Single-task suecess rale
08.8
0.3
0.8
0.8
Single-task success rafe
0.8
0.8
Single-task success rate
0.8
0.8
Single-task SUCCess
0.6
0.6
06
0.6
0.6
0.5
0.6
0.6
0:4
0.2
0.4
0.4
0.4
0.4
04
0.1
0.2
02
0.2
0.2
0.
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
20
0
20
20
20
0
20
Figure 13: Success rates (blue) and mixing coefficients a(c) (grey) over training in Kitchen with skewed pre- pre-training grows quickly for novel tasks, and more conservatively for those that were demonstrated during pre-training.
namely L2-regularization to the pre-trained weights, and EWC (Kirkpatrick et al., 2017). The latter can be seen as more nuanced version of the former, which adaptively scales the regularization strength according to the curvature of the loss landscape. Unfortunately, we find these methods to be insufficient in our setting, as reported in Figure 12. When coupled with large regularization weights, the asymptotical performance of I.2-regularization and EWC is significantly limited. When regularization weights are too low, they recover the performance of a naive baseline. Intermediate values were found to interpolate between the two behaviors, without addressing the forgetting issue. As result, the policy quickly loses information on its pre-training tasks, thus rendering adaptive data selection strategies ineffective.
This motivates our adoption of a novel technique to alleviate forgetting, which we refer to as Adaptive Prior. Inspired by previous works in behavioral priors (Bagatella et al., 2022) and offline RL (Kumar et al., 2020), we linearly combine the policy's 's output with that of a prior (in practice, a frozen copy Ty of the pre-trained policy). When an action needs to be sampled, we instead output a linear combination of actions sampled from the fine-tuned policy πT and the prior Tp. for given state-task pair (s,c) c) E S C, the action selected is a(c) + (1 oc(c)a )a where à ~ <-- e) and ~ Tp(-ss, c). The weight a<c) [0, 1] is task-dependent and learnable: ideally, it would be high when the fine-tuned policy is more suitable for the task, and low when the prior is instead preferable. In practice, it may be per-task learnable parameter for finite task spaces, or the output of a parametrized function (e.g., a neural network) in continuous task spaces. While this formulation departs from existing ones (Bagatella et al., 2022), in which the action is instead sampled from a mixture of policy and prior, it enables a simple update rule for a, which remains applicable when the policy class does not allow straightforward likelihood evaluation (e.g., Diffusion policies). Under the assumption that both л and p are isotropic Gaussians, the BC loss computed with respect to the linear combinations of actions simplifies a simple mean squared error:
CiC
||a (a(c))ā + (1 acc))a)ll2
(93)
N
0=
where à ~ T((-s),,;),a ā ~ TTp(-1s..c) c:) and iI. (s/, abj. , -1 -1 is the dataset of N task-conditioned, H H-length trajectories with task labels CI.N. As this loss function is differentiable with respect to a(c), the weight can be casily trained through gradient descent. Intuitively, this gradient pushes weights up for tasks in which the fine-tuned policy is more accurate than the pre-trained one. As π is fine-tuned on the dataset used for estimating Cgc. its performance can be easily overestimated. Thus, we propose to introduce a conservative penalty term, which encourages the weight to be updated only if the fine-tuned policy π significantly outperforms the prior नp: L CgC + B/ Σn a(c). We remark that, in our empirical evaluation, the Gaussian assumption is often violated; nevertheless, we found this update rule to be empirically effective across models and tasks. To illustrate, we visualize trajectories of mixing weights a(c) during fine-tuning in Figure 13. We consider a mismatched setting, in which the pre-training distribution largely demonstrates the first two tasks in Kitchen, and does not cover the remaining tree. During early training, AMF-NN mostly samples demonstrations for the last three tasks. Thus, π quickly improves, and the respective mixing coefficients converge to 1. On the other hand, ,, outperforms π on the first two tasks for the majority of training, a remains low and evaluation performance does not drop as actions are largely sampled from Tp-
26

---

Active Multi-task Policy Finc-tuning
Kitchen, mismatch
Kitchen, no mismatch
Metaworld, mismatch
Metaworld, no mismatch
0.8
Multi-task success rale
0.8
Multi-tash tasi success rate
0.85
Multi-task success rate
Multi-task success rate
0.6
0.80
0.6
0.6
0.75
0.5
0.4
0.70
0.4
0.4
0.65
0.3
0.2
0.2
0.60
10
20
10
20
10
20
0
10
20
Rounds
Rounds
Rounds
Rounds
AMF-NN
Rebalancing
Uniform, with adaptive prior
Uniform
Figure 14: Extended results from Figure 4, including a privileged "rebalancing" baseline.
I. Does AMF rebalance demonstration counts?
In discrete task spaces, counting the number of demonstrations for each task is possible. In this case, a naive data selection strategy would simply request demonstrations for tasks that have been demonstrated the least in the past. If all tasks require a similar amount of demonstrations, this would empirically perform very well. In our setting, however, data selection algorithms do not have knowledge of pre-training data. For this reason, count could only be kept with respect to the fine-tuning demonstrations: actively balancing this count would lead to near-uniform task selection, and recover the performance of uniform sampling in expectation.
Nevertheless, we implement this "rebalancing" criterion as a privileged baseline, which assumes access to the pre-training task distribution. We evaluate in the standard settings for AMF-NN from Figure 4. In Figure 14, we observe that AMF-NN is able to match the performance of this privileged baseline, despite having no knowledge of the pretraining distribution.
This implies that AMF can infer information on the pre-training phase through estimation of the policy's s uncertainty, and capable of automatically recovering "rebalancing" strategy. Moreover, AMF-NN considers the reduction in entropy across several tasks: hence, it can in contrast focus on tasks that are harder to learn or that could, in principle, lead to learning progress on other tasks. Further empirical evidence for these behaviors is shown in Appendix J.
J. Single-task performance
This Section presents a detailed look at the data selection strategies induced by AMF-NN. For this purpose, we consider the main experiments in Kitchen and Metaworld outlined in Figure 4, and plot single-task success rates, well as the amount of demonstrations collected over time. For reference, asymptotic single-task success rates on Metaworld and Kitchen converge between 70% and 100%, depending on the task.
In the case of skewed pre-training (Fig. 15 and 17), we observe that AMF samples tasks that were not present in the pretraining dataset more often, without having access to any direct information on the pre-training distribution. Moreover, even if multiple tasks have the same frequency in the pre-training distribution, AMF will prefer the ones that induce larger reduction in posterior uncertainty: for instance, in Metaworld, AMF selects Close Faucet more often than Open Faucet, despite the fact that both were similarly demonstrated during pre-training. We remark that these task selection strategies arise naturally from our information based criterion in Equation 1, without any direct information on the pre-training distribution, nor any explicit policy evaluation.
27

---

Active Multi-task Policy Finc-tuning
Cup left
Cup right
Open faucet
Close faucet
1.0
1.0
1.0
1.0
Sinlestask success rate
Single- tassask success rale
Single-task leeteekkkkkk stLccess rate
0.8
0.8
0.8
0.8
Single-task sticcsss
0.6
0.8
0.8
0.6
0.4
0.0
0.0
10
20
0
10
10
20
10
20
demonstrations
# of domonstration
# of demonsttratoon rations
# of demonstration
0
20
10 Rounds
10
20
0
10
20
Rounds
Rounds
Rounds
AMF-NN 1o=
Uniform, with adaptive prior -
# of pre-training demonstrations
Figure 17: Single-task curves for skewed pre-training in Metaworld. Dashed lines represent demonstrations counts, with grey lines displaying the (inaccessible) count of pre-training demonstrations.
Cup left
Cup right
Open faucet.
Close faucet
1.0
1.0
1.0
1.0
Single-task suoeess rate
Single-task strcesss rate
0.8
0.8
0.8
0.8
Single-tasi taak succeess
0.6
0.8
Single-tack
0.0
0.0
0.0
0.0
10
20
0
10
20
10
20
10
20
of demonstrotion
of demonst roatons
of dlemonsttrtion
of deemmmsstttionn
0
10
20
10
20
10
20
10
20
Rounds
Rounds
Rounds
Rounds
AMF-NN (a 13
Uniform, with adaptive prior
# of pre-training demonstrations
Figure 18: Single-task curves for uniform pre-training in Metaworld. Dashed lines represent demonstrations counts, with grey lines displaying the (inaccessible) count of pre- teaiiiinnng demonstrations.
28

---

Active Multi-task Policy Finc-tuning
Sliding door
Microwave
Left door
Knob off
Knob on
1.0
1.0
1.0
1.0
Single-tack suocess rale
Single-task snecess rate
Single- taak strepess rale
Single-tade tlsk success rase
0.8
0.8
0.4
Single-toet Moo-task
0.4
0.2
0.2
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.0
10
20
10
20
20
10
20
10
10
10
10
demonstration
of demonstrations
of demonstrations
demonstrations
demonstration
8
10
10
20
0
10
20
10
20
Rounds
Rounds
Rounds
Rounds
Rounds
AMF-NN (o le 3) Uniform, with adaptive prior of pre-training demonstrations
Figure 15: Single-task curves for skewed pre-training in Kitchen. Dashed lines represent demonstrations counts, with grey lines displaying the (inaccessible) count of pre-training demonstrations.
Sliding door
Microwave
Left door
Knob off
Knob on
1.0
1.0
1.0
Single-tasl suocess rate
Single-tasks shooess rate
Sing le-taek-tkkk stnntes rale
Single- Single-task s100ess rate
0.6
0.4
0.4
0.4
0.4
Single-tasle
0.4
0.2
0.2
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.0
10
20
20
10
10
10
10
demuosttrttin
demonstration
dlemonstration
of demonstrations
# demonneration
10 Rounds
20
Rounds
Rounds
Rounds
Rounds
AME-NN (a 33 Uniform, with adaptive prior of pre-training demonstrations
Figure 16: Single-task curves for uniform pre-training in Kitchen. Dashed lines represent demonstrations counts, with grey lines displaying the (inaccessible) count of pre-training demonstrations.
29

---

Active Multi-task Policy Finc-tuning
K. Analysis of Importance Weights
Importance weights (as introduced in Equation 3) allow estimating the expert's occupancy for arbitrary tasks. Naturally, the quality of importance weights depends on many factors, including the dimensionality of the trajectory space, and the density with which available data covers In this section, we report a qualitative evaluation of importance weights for both AMF-GP and AMF-NN (Figures 19 and 20, respectively). In both cases, find that informative weights can be retrieved eventually, given the proper amount of clipping (as described in Appendix M). While in early rounds of the algorithm, weights can be inaccurate, leading to poor estimate of the objective, we observe that the quality of importance sampling weights improves within handful of rounds.
Importance weights @ round
Importance weights @ round
Importance weights round 9
0.8
0.6
0.2
2
0.0
0
Task space
Task space
Task space
Figure 19: Analysis of importance sampling weights for AMF-GP. We consider the skewed pre-training setting from Figure and compute importance weights after 1, and rounds. We sample four tasks co:3, represented by vertical dashed lines of different colors. For each task c; we collect demonstration Tı and sweep over c' E C on the x-axis; we plot w(T,, d') with solid lines. We observe that importance weights uninformative in early parts of training, but converge more accurate values within few rounds.
Sliding door
Left door
Kanh
Sliling door
Milr
Left duar
Knob off
Knob
Sliding door
Mier owcave
Left
Knch off
Kieb
1.0
1.0
Sliding door
Sliding door
Sliding
0.8
0.8
0.8
Microwave
Microwave
Microwave
0.6
0.6
0.6
Left door
Left
0.4
0.4
0.4
Knob off
Knob
Knob
0.2
0.2
0.2
0.0
0.0
0.0
Kitchen, skowed & round
Kitchen, skowed & round 10
Kitchen, skewed d round 20
Cup
Crpp risht
Funcett oprn
Funent closs
Clup left
Chpp right
Fateett ope
Fascest clows
Cup luft
Cap right
Funvet open
Funent
Cup
0.8
Cup left
Cup right
0.6
Cup right
Cup right
0.4 Fancst
Fanost open
Fauost
0.2
0.2
close
close
Metaworld, & round
Metaworld, & round
Metaworld, skewee
Figure 20: Analysis of importance sampling weights for AMF-NN. We consider the skewed pre-training setting from Figure 4, and compute importance weights after I, 10 and 19 rounds. We visualize weights for both Kitchen (top) and Metaworld (bottom). As the task set is discrete, we consider all tasks (c: E C), and collect one demonstration т; for each. The entry of each colormap at row and column represents wr(T,, G.). Again, we observe that at the beginning of training importance weights can be inaccurate, particularly for tasks e; that have not been sufficiently demonstrated. However, more data collected and the policy specializes each task, the weights converge.
30

---

Active Multi-task Policy Finc-tuning
L. Criterion vs returns
As a didactic example, we evaluate the criterion optimized by AMF-GP in particular instance. We adopt the set- tings presented in Section 5.1, and pre-train GP policy by providing 50 demonstrations in the 2D integrator envi- ronment, uniformly sampled among tasks in the top half of the target circle. We represent the task space along one dimension, and plot the smoothed pre-training dis- tribution on the top of Figure 21. The second row of the Figure displays the evaluation of the criterion in Equation 2 for 100 tasks uniformly sampled across the entire task space. By comparison with the plot above, it is evident that the criterion is significantly lower for tasks that have not yet been demonstrated. These tasks are also those that, if demonstrated, would lead to a greater increase in multi-task performance after fine-tuning, as reported in the bottom row of Figure 21. In this instance, it's easy to see that the criterion leads to selection of tasks which have not been demonstrated sufficiently, and that will thus lead to greater policy performance.
Pretraining distribution
0.3
0.2
0.1
0.0
Estimated expected posterior entropy
1e4
1.5
1.0
0.5
0.0
Average return after finetuning
-12.5
M. Implementation Details
-13.0
In order to ease reproducibility, we open-source our code- base on the project's repo.' Furthermore, we describe several implementation details in the following sections.
13.5
Task space
M.1. Metrics
Figure 21: Didactic example on correlation between pre- training distribution over tasks (top), evaluations of the AMF criterion for each task (middle) and return after fine-tuning on demonstration for given task (bottom).
All metrics are reported in the form of their mean and the 90% simple bootstrap confidence intervals over 10 random seeds.
M.2. GP settings
In GP settings (5.1), each expert demonstration involves steps, is corrupted with Gaussian noise and collected by a scripted policy. As the task space is continuous, the criterion is simply optimized via uniform random shooting, with budget of 100. Multi-task returns are averaged over 20 episodes per task.
M.3. Neural network settings
1. ENVIRONMENTS
We evaluate AMF-NN across four environment suites, namely FrankaKitchen, Metaworld, Robomimic and WidowX. For the first two, demonstrations are 50 steps, while while for the latter two they involve up to 700 steps, and 100 steps respectively. In FrankaKitchen, demonstrations are provided by Kumar et al. (2024), and collected by trained SAC agents. In Metaworld, demostrations are instead collected by the scripted policies provided (Yu et al., 2020). In Robomimic, demonstrations are provided by proeficient human demonstrators (Mandlekar et al., 2021). Finally, in WidowX successful trajectories are collected by Octo-small (Octo Model Team et al., 2024) itself and filtered according to success labels, in an instance of self-supervised distillation. Furthermore, in the case of WidowX, the initial position of the object is not randomized, as found this to result in very inconsistent performance for the data collection policy. In the first two suites, 50 attempts for each task evaluated, while evaluation Robomimic and WidowX involves 25 attempts.
"githuu..ommin com/marbagg//mmm
31

---

Active Multi-task Policy Finc-tuning
M.3.2. BEHAVIOR CLONING
The MLP policy has with layers and 256 units per layer, with layer normalization (Ba, 2016). The Diffusion policy shares architecture and hyperparameters with the original implementation (Chi et al., 2023). Task conditioning word embedding extracted by sentence transformer (all-MmmiL--L6-22 (HuggingFace, 2026). Policies pre- trairained for 200 epochs with batch size of 256, learning rate of 10- using the AdamW optimizer[Loshchilov & Hutter, 2019).
M.3.3. IMPORTANCE SAMPLING WEIGHTS
In GP setting, importance weights are computed from the Gaussian policy distribution, and prg-oabbbillties are clipped to the range [-12, 0j. In NN settings, for deterministic policies, we interpret the policy's output as the mean of a Gaussian 1.0, and only clip log-probabilities for numerical stability. In experiments involving pre-trained Octo policies, evaluated two solutions. One option consisted of fitting Gaussian distribution through maximum likelihood methods to samples from the diffusion policy, and was found to underperform. We thus treat the Octo policy as strictly deterministic: with continuous action spaces, this simplifies importance sampling weights to w(c, r) 1 in case is a demonstration provided exactly for task c, and 0 otherwise. We note that this solution cannot be used to evaluate the criterion on yet unobserved tasks, but remains feasible when tasks are finite and few.
with fixed standard deviation
M.3.4. AMF
Each fine-tuning round involves 3000 gradient steps, each with a batch size of 256. We warm-start each algorithm by collecting the first |C] demonstrations uniformly, mentioned in Section 4.2. In the case of loss-gradient embeddings, we found it to be beneficial to use separate copy of the policy for task selection, which is not trained these initial trajectories (which thus can be seen as small "validation" set).
M.3.5. ADAPTIVE PRIOR
We found relatively high learning rates for the mixing weight a to work well: all experiments use learning rate of 1.0. The conservative coefficient was tuned for each suite, and is set 0.01 for all but Kitchen, in which takes the value of 0.001.
M.4. Runtime
Each experimental run for AMF-NN takes at most 5 hours with GPU acceleration. In this case, data selection itself requires less than second per round, and can be significantly sped up by reducing the sampling budget. AMF-GP experiments can be reproduced within 10 minutes CPU.
N. Useful inequalities
Lemma N.1. a E M) for some M > 0 and b 2 a then
b )
b S M. log
If additionally, a 2 M for some M' > 0 then
b
b b- M' log
Proof. Let f(x) del logrr By the mean value theorem, there exists E (a, b) such that
= (()) f(b) - f(a)
log log
log (8)
b-a
b b-a
Thus,
b-a a = c. c. log bb)
(0)
< M log
32

---

Active Multi-task Policy Finc-tuning
Under the additional condition that 2 M', obtain
M' log -))
b - e.c..oo c.log
Lemma N.2. Let us consider two spaces X E R", Y E R", and conditional distribution p X - A(Y) whose support supp(p((lz)) is bounded by a ball of radius e for all x X, that is
max 31 (su
1l 31l12 S E.
For all (I, S y ~ p(11z) and of ~ p(-|z') it holds that
- 1 S x((()z)), + 2c,
where K denotes the Wasserstein 1-distance.
Proof.
11y 1/1l
|1y 3/ 1|2 E'xup(p(((1a"
(94)
yEsupp(p( (ouupp(p(-aap)) [z])
WV3
min @Esupp(p(-1a)
I1y y1112 + 2e Expp(pp(112))
(95)
W(p(1=z). 2e,
(96)
where (i) follows from the triangle inequality, and (2) due to the fact that the integral of the distance between two points in Y for any coupling is greater than the minimum distance.
33


Reviews:
[{"review_id": "z8JFzDRhea", "reviewer": "Reviewer_9SjT", "summary": "*Note: I previously reviewed this paper during an earlier submission cycle. While I acknowledge the authors' efforts to enhance the manuscript, several concerns I raised earlier remain inadequately addressed in this revision. Thus, I incorporated some parts of my prior review, and adjusted the content according to the current submission.*\n\n\nThis paper introduces AMF (Active Multi-task Fine-tuning), an algorithm for efficiently fine-tuning pre-trained \"generalist\" robot policies to perform multiple tasks. Given a limited demonstration budget, AMF actively selects which tasks to request demonstrations for, aiming to maximize overall multi-task performance. It does this by selecting tasks that yield the largest information gain about the expert policy, focusing on areas where the current policy is most uncertain.\n\nThe authors provide theoretical performance guarantees for AMF under regularity assumptions, showing that it converges to the expert policy in sufficiently smooth MDPs. They also demonstrate AMF's effectiveness in practice, applying it to some robotic manipulation tasks with neural network policies. Experiments in simulated robotic environments like FrankaKitchen and Metaworld show that AMF significantly outperforms uniform task sampling, especially when the pre-training data is skewed towards a subset of tasks. The authors also demonstrated that AMF can be applied to off-the-shelf models like Octo, though the improvement over the naive baseline is marginal.\n\n## update after rebuttal\nI appreciate the rebuttal from the authors. I still recommend weak acceptance since generally this paper is solid.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for a (second) evaluation of our submission. We will address each comment individually.\n\n**Additional environments/policy architectures**\n\n> Future evaluations would benefit from testing on more challenging robotic benchmarks\n\nWe understand the significance of more challenging benchmarks. In fact, we have closely followed this helpful suggestion from your earlier review: we have included an evaluation in RoboSuite/RoboMimic [1], which ships with human demonstrations for long-horizon tasks up to 700 steps (Figure 6). As suggested, in this case we fine-tune a Diffusion Policy [2], as we found it to outperform ACT in single-task BC. The results are consistent with those on other benchmarks: AMF significantly outperforms a uniform strategy under distribution mismatch, and overall improves data efficiency. Therefore, AMF remains effective in more complex environments, and with fundamentally different policy parameterizations. We thank the reviewer for providing actionable feedback, and suggesting this extended evaluation, which we believe has substantially increased the breadth of our empirical support.\n\n**Uncertainty estimation**\n\n> The effectiveness of AMF, especially with neural networks (AMF-NN), hinges on accurate uncertainty estimation\n\nWe agree: uncertainty estimation is at the core of our method, and we acknowledge this limitation in Section 6. Fortunately, we find that loss gradient embedding perform reliably, as confirmed by the new evaluations with a Diffusion Policy in Figure 6. Moreover, AMF is rather agnostic to the particular choice of uncertainty quantification scheme, and we expect future developments to also benefit our framework.\n\nFinally, we would like to thank the reviewer for pointing out additional references, which we would be happy to integrate into our current related works. We of course remain available for further discussion, in particular if any major comment remains standing.\n\n**References**\n\n[1] Mandlekar et al., What Matters in Learning from Offline Human Demonstrations for Robot Manipulation, CoRL 2021\n\n[2] Chi et al., Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, RSS 2023", "claims_and_evidence": "I think most of the claims in this submission are supported by evidence.", "methods_and_evaluation_criteria": "Yes, they make sense to me.", "theoretical_claims": "The theoretical performance guarantees appear reasonable, though I did not verify the full derivation step by step.", "experimental_designs_or_analyses": "- FrankaKitchen and MetaWorld are relatively simple robotic benchmarks due to their narrow initial state distributions and short task horizons. Future evaluations would benefit from testing on more challenging robotic benchmarks such as RLBench [1], RoboSuite [2], ManiSkill [3], and BiGym [4], which offer greater complexity and variability. While naive BC + MLP/CNN approaches may not be sufficient for solving these benchmarks, many modern imitation learning methods (e.g., ACT, Diffusion Policy) are capable of achieving strong performance given high-quality demonstrations. Therefore, it should be possible to evaluate AMF on those benchmarks if the BC component is replaced with ACT or Diffusion Policy. *It is crucial to demonstrate that the proposed AMF method can also succeed on these more challenging benchmarks, as improving performance on overly simplified tasks provides limited insight*.\n\n- According to Figure 4, the performance of the proposed AMF-NN method is only marginally better than the uniform baseline. This raises concerns about the practicality of the method and the significance of the results, particularly given the relatively simple nature of the tasks.\n\n[1] James, Stephen, et al. \"Rlbench: The robot learning benchmark & learning environment.\" IEEE Robotics and Automation Letters 5.2 (2020): 3019-3026.\n\n[2] Zhu, Yuke, et al. \"robosuite: A modular simulation framework and benchmark for robot learning.\" arXiv preprint arXiv:2009.12293 (2020).\n\n[3] Mu, Tongzhou, et al. \"Maniskill: Generalizable manipulation skill benchmark with large-scale demonstrations.\" arXiv preprint arXiv:2107.14483 (2021).\n\n[4] Chernyadev, Nikita, et al. \"BiGym: A Demo-Driven Mobile Bi-Manual Manipulation Benchmark.\" arXiv preprint arXiv:2407.07788 (2024).", "relation_to_broader_scientific_literature": "AMF combines these two areas. Traditional active learning focuses on selecting the most informative data points to label in a single-task setting. Multi-task learning aims to improve performance on multiple tasks by sharing information between them. AMF extends active learning principles to the multi-task setting, specifically for fine-tuning a pre-trained policy. This is related to works like those by [1, 2] which apply task-directed data selection and active fine tuning, but AMF extends this to sequential decision-making in robotics. Other works have explored multi-task active learning, but not in the context of fine-tuning pre-trained Transformer-based models for NLP [3].\n\n[1] Smith, Freddie Bickford, et al. \"Prediction-oriented Bayesian active learning.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2023.\n\n[2] Hübotter, Jonas, et al. \"Transductive active learning: Theory and applications.\" Advances in Neural Information Processing Systems 37 (2024): 124686-124755.\n\n[3] Rotman, Guy, and Roi Reichart. \"Multi-task active learning for pre-trained transformer-based models.\" Transactions of the Association for Computational Linguistics 10 (2022): 1209-1228.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "### Strengths\n\n- The paper studies the timely problem of efficiently fine-tuning generalist robot policies, which are becoming increasingly important in robotics.\n\n- The authors provide performance guarantees for AMF under certain regularity assumptions, proving its convergence to the expert policy in smooth MDPs. This adds to the credibility and understanding of the algorithm's behavior.\n\n- AMF demonstrates improvements over uniform sampling, particularly when the pre-training data is biased towards a subset of tasks. This is an advantage as real-world pre-training datasets are sometimes unevenly distributed.\n\n  \n\n### Weaknesses\n\n- The effectiveness of AMF, especially with neural networks (AMF-NN), hinges on accurate uncertainty estimation. While the proposed loss-gradient embedding approach works well empirically, uncertainty quantification in neural networks remains a challenging open problem. The performance can degrade if the uncertainty estimates are unreliable.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See above sections.", "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": "zpmX3QsIue", "reviewer": "Reviewer_jtjj", "summary": "Pretrained generalist policies are becoming popular in the robot learning field for the gained capabilities by large-scale training. Nevertheless, deploying such policies in a zero-shot manner is still lacking. Hence, adaptation of generalist policies is a must to utilize the acquired representations and skills. A popular approach for adaptation is fine-tuning on some expert demonstrations. Since collecting such demonstrations is expensive, we should minimize the number of the demonstrations needed for the downstream fine-tuning. The problem is more prominent when the target is to learn multiple downstream tasks. In this work, the authors propose an algorithm, named AMF, which aims to actively fine-tune large-scale pretrained models with the minimum number of demonstrations possible, collected from more than one task. The approach aims to maximize the information gain of demonstrations about the expert policy by selecting the right task for the upcoming demonstration collection. Under some strong assumptions, this work offers theoretical guarantees for matching the performance of the expert policy used for collecting the demonstrations. In addition, the authors offer a practical version of the algorithm based on a GP policy, called AMF-GP, by tackling the occupancy and entropy estimation. Furthermore, a more realistic approach, named AMF-NN, has been proposed when the policy is modeled by a neural network pre-trained on a large-scale dataset. Both algorithms have been empirically evaluated against the naive approach of sampling tasks uniformly for demonstration collection.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the author for their thorough review and positive evaluation of our work. We are happy to provide an answer to each question and comment.\n\n**MT10/50**\n> I would have expected to evaluate the proposed approach on the standard Metaworld scenarios (MT10 and MT50).\n\nThe evaluation on Metaworld is designed such that all tasks may share the same initial state distribution. In the standard MT10/MT50 suite, most tasks can be identified through their initial state (e.g., if a button is present, the task is most likely ‘press-button’). This would allow one to potentially quantify policy uncertainty from the states alone, independently of the specified tasks, arguably simplifying the problem.\nIn our instantiation, all tasks share the same initial state distribution (faucet and cup are both always present). On one hand, this has the advantage of stressing uncertainty quantification, ensuring that task specification cannot be ignored. On the other hand, a shared state space may facilitate transfer learning, which is an important motivation for multi-task learning as a whole. These are the reasons motivating the adoption of the current Metaworld setup.\n\nAMF however remains fully compatible with disjoint initial state distributions: to demonstrate this, we evaluated AMF on standard MT10 tasks (see new Figure W1 on the [anonymous website](https://sites.google.com/view/active-multitask-finetuning)). Similarly to the existing setup, pre-training only focuses on half of the tasks in mismatched settings, and on all tasks in non-mismatched settings. Results are consistent with those in the current setup, confirming that AMF may be widely applied.\n\n**Other baselines**\n> [...] active (fine-tuning) learning, but I believe methods from this domain can be considered as baselines.\n\nTo the best of our knowledge, principled active selection of demonstrations for fine-tuning remains largely unexplored, and we could not find active fine-tuning methods that may be directly applied to this setting. There exists parallel work adopting related objectives, the closest being perhaps [1] in the context of model-based RL. The objective from [1] is however not applicable to this setting: it would involve estimation of a dynamics model from pre-training data, which is not available. Learning the dynamics model on fine-tuning data alone would result in similar behavior to the uniform baseline. \n\n**Other open source models**\n> I believe we already have so many open-sourced generalist policies.\n\nWe overall agree with the reviewer: open-source generalist policies are indeed not scarce anymore. \nWe have chosen Octo due to its good performance in the Simpler suite (see [2], Figure 7), which should make the evaluation more informative with respect to RT-1-X [3] models. OpenVLA [4] is also available as a base model; however, it was evaluated sparingly in simulation, in which it performed similarly to Octo ([4], Table 12). This motivated our adoption of Octo.\n\nHowever, we believe that benchmarks remain sparse. The few existing real-to-sim simulated benchmarks are currently designed to mainly support zero-shot evaluations [2], rather than data collection and subsequent fine-tuning. Other simulated environments exist [5], but the real-to-sim gap has not been properly quantified. For this reason, we believe that a comprehensive evaluation of AMF with open-source models would need to go beyond simulation and include extensive hardware experiments. As such, it would lie beyond the scope of this work, which aims at establishing a principled framework for fine-tuning multi-task policies across different parameterizations and environments.\nWe are happy to rephrase the discussion in Section 6 to make this clear. \n\n---\n\nWe hope that these clarifications are helpful in supporting experimental choices, considering the motivation and scope of this work. We would be happy to engage in further discussion or answer any additional questions. Thank you again for taking the time to review out work.\n\n**References**\n\n[1] Mehta et al., An experimental design perspective on model-based reinforcement learning, ICLR 2022\n\n[2] Li et al., Evaluating Real-World Robot Manipulation Policies in Simulation, CoRL 2024\n\n[3] Open X-Embodiment Collaboration, Open X-Embodiment: Robotic Learning Datasets and RT-X Models, CoRL 2023 TGR Workshop\n\n[4] Kim et al., OpenVLA: An Open-Source Vision-Language-Action Model, CoRL 2024\n\n[5] Liu et al., Benchmarking Knowledge Transfer for Lifelong Robot Learning, NeurIPS 2023 D&B", "claims_and_evidence": "- I believe the problem demonstrated in this work is valid. It is important to minimize the number of demonstrations collected for fine-tuning generalist policies on a task. The issue is even more challenging when we need to decide from which task we should collect demonstrations and how many demonstrations we should collect. \n- In my opinion, the motivation, claims, and evidence were clear. The flow of this work is consistently good.\n- The theoretical discussion strongly supports the claims presented in this work.", "methods_and_evaluation_criteria": "- The proposed method sounds. I have no comment on the incorrectness of the approach. \n- The benchmarks used are convincing. However, I believe that the chosen metaworld tasks are few and too similar. I would have expected to evaluate the proposed approach on the standard Metaworld scenarios (MT10 and MT50).", "theoretical_claims": "- I believe the theoretical claims provided in this work sound. I have checked the claims, and I have no issue with them.", "experimental_designs_or_analyses": "- The benchmarks used in this work are relevant. \n- As mentioned, I have a concern regarding the similarity between the tasks in the metaworld benchmark. I think this is an important weakness to highlight. Nevertheless, I appreciate the robomimic benchmark used.\n- I have concerns regarding the scarcity of the baselines.\n- I am not familiar with the literature of active (fine-tuning) learning, but I believe methods from this domain can be considered as baselines. \n- In addition, as stated in this work, the meta-learning literature has already approaches for a similar objective. I think it is important to illustrate the final performance obtained using the proposed approach in comparison to approaches designed for learning to adapt.", "relation_to_broader_scientific_literature": "- This work is highly related to the field of robot learning. Recently, many generalist policies have been introduced, and the need for fine-tuning them is high. \n- In addition, this work is related to meta-learning and active learning as discussed already by the authors.", "essential_references_not_discussed": "- I have no references to add.", "other_strengths_and_weaknesses": "I have to comment on the last paragraph in the discussion section. I believe we already have so many open-sourced generalist policies. In my opinion, the major weakness in this work is the lack of experiments that show the ability of the proposed approach in fine-tuning different large models (other than Octo) for robot learning. I think that benchmarking with more than one backbone model will strengthen the soundness of the proposed approach and validate the claims.", "other_comments_or_suggestions": "- No comments", "questions_for_authors": "- The chosen metaworld scenario consists of only 4 tasks; the tasks, in my opinion, are very similar. Why not consider many tasks from the standard set of MT10 and MT50 in metaworld?", "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": "OVoEKszWPp", "reviewer": "Reviewer_uqQh", "summary": "This paper investigates an active multi-task fine-tuning scheme, which adaptively selects the task to be demonstrated for sample-efficient fine-tuning of multi-task behavioral cloning policies. The authors provide a practical version of the proposed method and highly the efficacy of the proposed method through experiments on a controlled 2D integrator environment as well as more realistic scenarios such as Metaworld, Franka Kitchen, and the WidowX (Appendix F) environments.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for thoroughly reviewing our work. We are happy to address each comment below.\n\n**Comparison with other methods in the domain**\n> It would be great if the authors could include comparisons with other methods in the domain\n\nTo the best of our knowledge, the problem of active-finetuning for behavior cloning has not been explored in the past. For this reason, there are no established baselines for this specific problem. For instance, [1] studies a dataset may be “curated” according to mutual information, but it does not address online data selection or multi-task learning. [2] deploys a related objective to ours in the context of model-based RL; however, this objective cannot be evaluated in our setting, as it would require a dynamics model estimated on pre-training data, which are not available during fine-tuning. The uniform baseline is what is practically used in an offline setting [3].\n\n**Number of iterations in Figure 4**\n> In Sec. 5.2, when the authors allow 10 or 20 evaluation iterations, is this across the 4 or 5 tasks in the task set?[...] Does the performance in Figure 4 improve with more iterations?\n\nIn Sec. 5.2, 20 demonstrations are collected across all tasks for all plots: in non-visual settings, we collect 1 demonstration for 20 iterations, and in visual settings, we collect 2 demonstrations for 10 iterations. This is a relatively small number, but we must consider that the policy is already pre-trained on several demonstrations (~10-20 depending on the setting). Furthermore, significant implementation efforts were directed towards data-efficiency in the underlying (BC) algorithm to reduce runtime. Albeit more slowly, performance further improves with more iterations, and eventually converges to a value of ~92% in Metaworld, and ~89% in Kitchen.\n\nWhile Kitchen and Metaworld remain high-dimensional, challenging deep RL benchmarks, we find that they require relatively few demonstrations. Thus, we also provide an evaluation on Robomimic (Figure 6), which involves arguably harder tasks, requiring a larger number of demonstration is necessary to achieve comparable success rates (>100 demonstrations for some tasks). We find that results in this setting confirm the previous ones on Kitchen and Metaworld.\n\n**New environment or new tasks**\n> [...] will this work even when the pretrained policy is introduced in a new environment with similar tasks?\n\nThanks for the interesting question. Our formal analysis of the algorithm assumes a fixed MDP, and thus tackles the scenario in which the environment does not change between pre-training and evaluation. However, the practical algorithm may also be applied in case of slight shifts in the environment, as long as the policy generalizes across these changes. If the environment at evaluation is completely unrelated to pre-training data, then AMF would in practice recover a less informed strategy. However, if the changes in the environment are relatively small, and uncertainty estimates remain somewhat meaningful, AMF would still be able to leverage them to efficiently query demonstrations.\n\n---\nWe again thank the reviewer for their insightful feedback. If there are any additional comments, we would be very happy to engage in further discussion.\n\n**References**\n\n[1] Hejna et al., Robot Data Curation with Mutual Information Estimators, arXiv:2502.08623\n\n[2] Mehta et al., An experimental design perspective on model-based reinforcement learning, ICLR 2022\n\n[3] Kim et al., Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success, arXiv:2502.19645", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Not in detail. The empirical results look promising.", "experimental_designs_or_analyses": "The experimental setting makes sense for evaluating the sample-efficient fine-tuning of pretrained policies through the proposed method.", "relation_to_broader_scientific_literature": "I do not follow the relevant work in this area very closely. It would be nice if the authors could compare with existing methods in the area to show where the proposed method stands compared to them.", "essential_references_not_discussed": "I do not follow the relevant work in this area very closely.", "other_strengths_and_weaknesses": "Strengths\n- The authors provide theoretical guarantees of the method and provide a practical algorithm applicable to more realistic settings.\n- The paper includes experiments on a controlled 2D integrator environment as well as more realistic scenarios such as Metaworld, Franka Kitchen, and the WidowX (Appendix F) environments.\n- The authors study uncertain estimation choices for AMF and study its applicability to other off-the-shelf models such as Diffusion Policy (Sec 5.4) and Octo (Appendix F).\n- The paper includes the limitations of the proposed method.\n\nWeaknesses\n- It would be great if the authors could include comparisons with other methods in the domain in addition to the uniform sampling baseline. This would help the reader get a sense of where the proposed method stands in comparison with existing methods.\n- In Sec. 5.2, when the authors allow 10 or 20 evaluation iterations, is this across the 4 or 5 tasks in the task set? If yes, this seems like a small number of demonstrations. Does the performance in Figure 4 improve with more iterations?\n- Also, since the finetuned policies get around or greater than 50% success rate despite starting with a low success rate (Fig. 4), especially in the mismatched case, does this mean that the tasks are very simple? Since 10 or 20 demonstrations across 4 or 5 tasks lead to a significant improvement in performance, even with uniform sampling.", "other_comments_or_suggestions": "- Based on the experimental environments, does this method tackle the scenario where the environment remains the same and the policy on fine-tuned on new tasks that are introduced during evaluation? Or will this work even when the pretrained policy is introduced in a new environment with similar tasks? I am curious the hear about the authors’ thoughts on this.", "questions_for_authors": "It would be great if the authors could address the weaknesses as well as the question mentioned earlier.", "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": "HEpMYTohkM", "reviewer": "Reviewer_GTTR", "summary": "The paper tries to tackle the problem of maximizing the multi-task performance of a pre-trained policy with minimal additional demonstrations via active learning. The proposed algorithm builds upon existing active learning approaches in non-sequential domain. AMF selects queries that maximizes the expected info gain about expert policy over the occupancy with some guarantee that the resulting policy will converge to the expert policy under some assumptions. AMF proposes practical approaches to compute conditional entropy with Gaussian Process policies and occupancy estimation via importance sampling. Lastly, AMF proposes a simple technique using a prior policy to mitigate issues of catastrophic forgetting of learned behavior. The authors demonstrate the approach on a 2D didactic problem along with continuous control tasks in MetaWorld and Franka Kitchen. Further, they show that AMF can scale to work with pretrained policies such as diffusion policies.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the thorough evaluation of our submission, and the insightful comments.\n\n**Additional baselines and prior work**\n\nWe thank the reviewer for suggesting these two additional works. We found them to be very relevant to this direction, but not directly applicable as baselines. To the best of our knowledge, Bayesian active selection of demonstrations for fine-tuning remains relatively unexplored. [1] focuses on a standard, non-sequential BO setting, and [2] is designed for exploration in MBRL. Nevertheless, the objective from [2] could be applied to our setting, and would select tasks inducing trajectories minimizing posterior uncertainty over dynamics predictions along optimal trajectories. This criterion would however require access to a dynamics model estimated on pre-training data. In realistic settings, pre-training data is not available; if it was, other strategies (such as rebalancing, see Appendix I) would become promising. We focus on the realistic setting in which pre-training data is not available: in this case, a dynamics model trained from scratch would be unaware of the pre-trained policy’s uncertainty, and would fail to drive effective task selection. These two references remain very relevant; we are happy to include and extend this discussion in our related works.\n\n**Adding a topline score**\n\nWhile we are not allowed to submit a revision at the moment, we have run the experiments to estimate a topline score. Under no mismatch, a policy fine-tuned “at once” on 20 demonstrations uniformly distributed across tasks matches the score of the Uniform baseline at iteration 20: \\~64% in Metaworld and \\~85% in Kitchen. If the number of demonstrations is increased 5x to 100, the policy reaches a performance of \\~92% in Metaworld, and \\~89% in Kitchen.\n\n**Role of the adaptive prior**\n\nWe found the adaptive prior to be very helpful when pre-training and evaluation task distributions do not match. In this situation, when the agent requests the first demonstrations and fine-tunes on them, it may quickly forget tasks for which no demonstration was queried. As a consequence, the average success rate drops significantly in early iterations under mismatch (Figure 4, first and third columns, red line). This phenomenon is substantially alleviated if an adaptive prior is used (yellow line), as it can retain information about pre-training tasks. We hope this explains the issue, and would ask the reviewer whether anything else could be clarified.\n\n**Number of demonstrations in Figure 2**\n\nThank you for catching this! We only collect 40 demonstrations, as performance plateaus afterwards. We will correct this in the text.\n\n**Data efficiency with 20 demonstrations**\n\nWe directed significant efforts at making the implementation of the underlying algorithm as data efficient as possible, in order to reduce runtime. As a result, the policy performs well with only 20 additional fine-tuning demonstrations, as the reviewer points out. We must however remark that policies are already pre-trained on a number of demonstrations (\\~10-20), and that, while good performance is achieved quickly, a larger number of demonstrations is required to reach asymptotic performance.\n\n**Additional baseline evaluating policy performance**\n\nEvaluating the multi-task policy, and sampling demonstration according to the inverse of the policy performance on each task would indeed represent an interesting solution to active multi-task finetuning. This is however an online solution, which has additional requirements. First, it requires access to the environment in order to estimate the policy performance, potentially involving the execution of unsafe actions. Second, it requires designing a per-task reward to evaluate the rollouts. In settings in which these issues would not be relevant (e.g. a good simulator and reward function are available), this method would do very well. In contrast, AMF is a fully offline method, which avoids environment interaction and uses uncertainty as a proxy for performance. This relationship is formally described in Theorem A.10.\n\n**Average task success curves and standard deviations**\n\nAs the reviewer points out, plots are averaged across tasks. Shaded areas are 90% simple bootstrap confidence intervals of the average task performance (thus modeling stochasticity across seeds).\nAsymptotic single-task success rates on Metaworld and Kitchen converge between 70% and 100%, depending on the task. We are happy to report this in detail given a chance to update the paper.\n\n---\nWe would like to thank the reviewer for taking the time to evaluate our submission. We hope we were able to address all comments, and we would like to further discuss any point that remains unresolved.\n\n**References**\n\n1) Neiswanger et al., Generalizing bayesian optimization with decision-theoretic entropies, NeurIPS 2022\n\n2) Mehta et al., An experimental design perspective on model-based reinforcement learning, ICLR 2022", "claims_and_evidence": "The claims are supported with experiments in a didactic environment and continuous control experiments across several robot learning benchmarks. It would be nice to see other baseline methods that use Bayesian active selection in a decision-making context. Currently, the authors only compare against the naive baseline of uniform sampling for the next task id. The results make sense intuitively in that if there is a mismatch in distribution between the pretraining and deployment time task distributions then uniform sampling would be worse as you can query for task demonstrations suboptimally. The experimental results do support this as we observe that in the case of mismatched distributions, AMF learns quicker compared to the uniform sampling baseline.", "methods_and_evaluation_criteria": "Yes, the proposed methods make sense. The primary evaluation criteria is success rate of the learned policy over consecutive rounds of active querying of expert demonstrations. The benchmarks used in the experiments are also standard manipulation tasks used in the literature.", "theoretical_claims": "I did not check the correctness of the proofs.", "experimental_designs_or_analyses": "The experimental design seems reasonable. The authors evaluated on both state- and image- based settings. A policy is pretrained on set number of demonstrations and some tasks are heldout to simulate a distribution mismatch.", "relation_to_broader_scientific_literature": "This paper is relevant to the multi-task learning, active imitation learning, and imitation learning literature. AMF extends the ideas of active query selection in Bayesian optimization to the sequential decision-making domain and specifically for the setting of multi-task learning.", "essential_references_not_discussed": "There are some works in active learning for sequential decision-making and Bayesian experiment design that should have been referenced and used as potential baseline methods.\nThe second work also similarly proposes a Bayesian approach for query selection in the setting of model-based RL and has a similar objective for estimating occupancy measure between trajectory and expert policy. Also, they propose similar idea of using Gaussian process for modeling the policy. It would be nice to see some additional discussions contrasting to this prior work.\n\n[1] Neiswanger, Willie, et al. \"Generalizing bayesian optimization with decision-theoretic entropies.\" Advances in Neural Information Processing Systems 2022.\n[2] Mehta, Viraj, et al. \"An experimental design perspective on model-based reinforcement learning.\" International Conference on Representation Learning 2022.", "other_strengths_and_weaknesses": "Strengths:\n- Paper is well written and easy to follow\n- Method is technically sound \n- Experimental setting is reasonable and AMF is evaluated on a good suite of environments\n- Results on pretrained off-the-shelf policies seem to suggest that AMF is useful for improving policy performance quickly in a more realistic setting\n\nWeaknesses:\n- Additional references expected, see above section. There are some works in active query selection for RL that are not referenced. \n- It would be nice to have one additional baseline method that is a stronger baseline that uniform sampling \n- Some of the ideas in the AMF algorithm have already been proposed in prior work, e.g. occupancy estimation in [2] and prior regularization for catastrophic forgetting", "other_comments_or_suggestions": "- I would recommend putting a topline score for what an expert policy would have achieved if provided all the training demonstrations from the very beginning without applying any active learning.\n- It is not clear from the main results how much the adaptive prior is helping to prevent catastrophic forgetting.", "questions_for_authors": "- Why is Figure 2 showing the results at 40 demonstrations and not 50 demonstrations which is the number of demonstrations reported in the text?\n- For my understanding: in the state-based experiments, if one new demonstration is added at each iteration and if there are 20 iterations, does that mean 20 new demonstrations are added? I'm quite surprised that AMF is able to learn a multi-task policy for all those tasks with so few demonstrations.\n- Would this be a reasonable baseline: first evaluate the base policy to get a initial estimate of the policy's performance on each task and use that as your prior distribution to query for demonstrations? This baseline would aim to select tasks that the base policy performs poorly on.\n- Are the result plots averaged across each of the tasks in the benchmark? How do I interpret the standard deviation intervals? It would be interesting to see the individual task successes with more demonstrations.", "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:

Meta-review: The paper considers the setting of active task selection for policy finetuning. The paper is easy to read and the results are sound. Overall, I share the same concerns and probably the feelings of the reviewers. The paper is honest in its claims and results, and it is mostly a combination of existing techniques. From my perspective, the paper's novelty is to a certain extent limited, but I still think it is an interesting contribution to the RL community.
Decision: Accept (poster)