Active Fine-Tuning of Multi-Task Policies

We would like to thank the reviewer for the thorough evaluation of our submission, and the insightful comments.

Additional baselines and prior work

We 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.

Adding a topline score

While 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.

Role of the adaptive prior

We 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.

Number of demonstrations in Figure 2

Thank you for catching this! We only collect 40 demonstrations, as performance plateaus afterwards. We will correct this in the text.

Data efficiency with 20 demonstrations

We 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.

Additional baseline evaluating policy performance

Evaluating 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.

Average task success curves and standard deviations

As 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). Asymptotic 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.

---

We 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.

References

1. Neiswanger et al., Generalizing bayesian optimization with decision-theoretic entropies, NeurIPS 2022

2. Mehta et al., An experimental design perspective on model-based reinforcement learning, ICLR 2022

We would like to thank the reviewer for thoroughly reviewing our work. We are happy to address each comment below.

Comparison with other methods in the domain

It would be great if the authors could include comparisons with other methods in the domain

To 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].

Number of iterations in Figure 4

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?

In 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.

While 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.

New environment or new tasks

[...] will this work even when the pretrained policy is introduced in a new environment with similar tasks?

Thanks 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.

---

We again thank the reviewer for their insightful feedback. If there are any additional comments, we would be very happy to engage in further discussion.

References

[1] Hejna et al., Robot Data Curation with Mutual Information Estimators, arXiv:2502.08623

[2] Mehta et al., An experimental design perspective on model-based reinforcement learning, ICLR 2022

[3] Kim et al., Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success, arXiv:2502.19645

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.

MT10/50

I would have expected to evaluate the proposed approach on the standard Metaworld scenarios (MT10 and MT50).

The 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. In 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.

AMF 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). 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.

Other baselines

[...] active (fine-tuning) learning, but I believe methods from this domain can be considered as baselines.

To 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.

Other open source models

I believe we already have so many open-sourced generalist policies.

We overall agree with the reviewer: open-source generalist policies are indeed not scarce anymore. We 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.

However, 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. We are happy to rephrase the discussion in Section 6 to make this clear.

---

We 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.

References

[1] Mehta et al., An experimental design perspective on model-based reinforcement learning, ICLR 2022

[2] Li et al., Evaluating Real-World Robot Manipulation Policies in Simulation, CoRL 2024

[3] Open X-Embodiment Collaboration, Open X-Embodiment: Robotic Learning Datasets and RT-X Models, CoRL 2023 TGR Workshop

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

[5] Liu et al., Benchmarking Knowledge Transfer for Lifelong Robot Learning, NeurIPS 2023 D&B

We thank the reviewer for a (second) evaluation of our submission. We will address each comment individually.

Additional environments/policy architectures

Future evaluations would benefit from testing on more challenging robotic benchmarks

We 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.

Uncertainty estimation

The effectiveness of AMF, especially with neural networks (AMF-NN), hinges on accurate uncertainty estimation

We 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.

Finally, 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.

References

[1] Mandlekar et al., What Matters in Learning from Offline Human Demonstrations for Robot Manipulation, CoRL 2021

[2] Chi et al., Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, RSS 2023