Active Fine-Tuning of Generalist Policies

We would like to thank the reviewer for the thorough review, and their suggestions. We believe that we can address the outlined weaknesses, and provide additional empirical support to clarify our contribution.

Parameter isolation

Parameter isolation is critical because we tackle a relatively novel, more challenging setting. In the standard multi-task behavior cloning setting, all data that the policy is trained on remains available at all times, and the data distribution does not shift. In our setting, the policy is pre-trained on a sizeable pre-training dataset, but this dataset is not available during fine-tuning (see common response). This induces a large distribution shift between pre-training and early rounds of fine-tuning. If this distribution shift is not addressed, performance suffers significantly, as shown in Figure 4 and Appendix G. Thus, the necessity of some form of continual learning arises from our problem setting itself, rather than from particular algorithms.

An effective technique to prevent forgetting in a continual setting is rehearsal, which implicitly occurs in the standard multi-task setting, as the dataset is fixed and always accessible. Unfortunately, this is simply not possible when part of the training data becomes unavailable, as in our setting. For this reason, we benchmarked different continual learning techniques, and found parameter isolation to perform best (see Appendix G, including updates). We also note that, at the end of fine-tuning, it is possible to distill the sets of weights into a single, multi-task policy.

Importance weights

We appreciate this suggestion, and we are happy to add an analysis of importance weights as Appendix J. In short, for both AMF-GP and AMF-NN we confirm that they might be inaccurate in early parts of training. Nevertheless, we find that they become more reliable estimates in later rounds. We refer to Appendix J for a more complete discussion.

Other sampling schemes

Thank you for suggesting this baseline: we think it highlights a nice feature of the problem we tackle. We included a comparison in our revision. Please, see the common response for a detailed response, and further empirical evidence. In short, we find that our method matches or outperforms the performance of the suggested baseline, despite not having access to the pre-training data distribution.

We hope that we were able to clarify the challenging fine-tuning setting we tackle, and why particular techniques might become desirable or not applicable to this case. In light of this, we would ask the reviewer to consider raising their score. We are of course happy to discuss and address any further questions or concerns!

Dear reviewer, We greatly appreciate your time and effort in reviewing our work, and we believe that our response addresses the points you raised. As the discussion period is nearing its end, we wanted to kindly check if there is anything else we can clarify. We remain available for any further discussion.

We thank the reviewer for the detailed review and questions. We’ll comment on the weaknesses, and answer questions one by one.

Benchmarks

We appreciate the suggestion to explore more complex benchmarks, and agree that this represents a promising avenue for future research. We find that the underlying learning algorithm (in this case, behavior cloning (BC)) may not be able to solve very complex tasks, even in the single-task setting. To the best of our knowledge, the best BC baseline across the suggested environments only solves (>80% success rate) less than a third of available tasks. Thus, these environments largely test the behavior cloning component, which makes them less informative of the data selection technique. Our evaluation focuses on showing that, as long as behavior cloning can solve a set of tasks, AMF can substantially improve sample efficiency.

Performance after uniform pre-training

As the reviewer correctly points out, in the uniform settings, the performance gap between AMF and baselines shrinks. However, we argue that this does not occur because the performance of AMF worsens, but rather because this setting is particularly favorable to uniform task selection. The fact that AMF recovers its performance without any knowledge of the pre-training distribution confirms its ability to adapt across scenarios, and is a desirable trait. Moreover, we believe that a scenario in which the pre-training covers all tasks equally and sufficiently is rather unrealistic. In this case, fine-tuning might simply not be necessary. For any practical use of generalist policies, we can expect some substantial distribution shift, and this is the setting where AMF shows promising performance gains.

Runtime

We find the computational requirements for AMF to be modest. Even with a generous sampling budget, data selection with AMF-NN requires at most <8 minutes at each round, which represent a small fraction of the time required for the entire fine-tuning process. We have added details on runtime to Appendix N.4.

Parameter isolation

We agree with the reviewer: unfortunately, parameter isolation prevents positive transfer, but we found it to work best across several continual learning techniques, as shown in Appendix G. In our setting, the impact of positive transfer is overall lower than that of catastrophic forgetting. At the end of fine-tuning, the parameter count may be reduced by distilling the multiple copies of the weights into one.

Scope of OCTO experiments

To the best of our knowledge, experiments with OCTO include all currently available environments for Real-to-Sim. The evaluation is thus currently limited by lack of available simulated benchmarks. We agree that further experiments in this setting would be interesting, and look forward to the development of broader benchmarks for Real-to-Sim.

Figure 2

Thanks for pointing out the misalignment in Figure 2! As it was generated by plotting actual noisy demonstrations, the trajectories are not perfectly straight. We’re happy to tune down the noise for plotting purposes, and have replaced the Figure.

Figure 4

In Figure 4, early drops in the red curve are a symptom of catastrophic forgetting in neural networks. Across all tasks, networks are pre-trained to reach non-trivial performance in at least some tasks. This performance is shown by the first point for each curve. On the first round of fine-tuning, the fine-tuning dataset only contains the first selected batch of demonstrations, thus inducing a large distribution shift from the pre-training distribution. For this reason, the first round of fine-tuning “undoes” some of the pre-training, and performance for some tasks drops significantly. Hence, average performance also drops. This is a general issue in continual and lifelong reinforcement learning, which arises from the setting itself. In the context of this work, we adopt parameter separation to mitigate this issue, as described in line 348 and Appendix G.

We thank the reviewer for the discussion, which highlighted some important features of our problem setting, and allowed us to clarify relevant points. Please, let us know if you have any further questions or concerns.

Dear reviewer, We greatly appreciate your time and effort in reviewing our work, and we believe that our response addresses the points you raised. As the discussion period is nearing its end, we wanted to kindly check if there is anything else we can clarify. We remain available for any further discussion.

We thank the reviewer for the detailed review and questions. As suggested, we have edited our submission to further highlight AMF-NN, namely by editing Algorithm 1. For the time being, we have kept a concise introduction to GPs, as AMF-NN largely builds upon a GP approximation. We hope that this solution can improve the readability of our paper for readers that have a stronger background in both GPs and NNs. All other questions are answered one by one in the following.

Gaussian Processes in AMF-NN

In AMF-NN, no Gaussian Processes need to be explicitly trained. Instead, we leverage embeddings learned by the policy network, and a GP approximation to estimate the policy’s conditional entropy. If the pretrained policy also captured epistemic uncertainty, then entropy reduction could be estimated by comparing fine-tuned and pretrained policies. However, this process would require several gradient steps for each of the tasks to consider, while our current method remains more computationally efficient.

Task-space partitioning

In practice, we found that task-space partitioning for AMF-NN works best when having a set of weights for each task. As the partitioning gets coarser, performance gradually decreases. This limitation is entirely due to forgetting issues in neural networks, and we expect that improvements in continual learning would eventually remove it. Generally speaking, the AMF framework scales to continuous sets of tasks, as displayed by AMF-GP (which does not require any partitioning).

Model merging

We appreciate the suggestion regarding model merging. We find that its applications for continual learning [1,2] rely on the availability of a replay buffer of past data. In our case, as we further discuss in the common response, this is not possible: fine-tuning algorithms do not have access to pre-training data. However, we found that one of the methods at the core of model-merging (Git-ReBasin [3]) is also described in a “weight-matching” variant, which does not require data, and is therefore directly applicable in our setting. In the context of this rebuttal, we have implemented a continual learning algorithm based on this idea, as detailed in Appendix G. Unfortunately, we found it to be empirically ineffective in our fine-tuning setting.

AMF with no parameter isolation

Without parameter isolation, all methods we evaluated in Kitchen and Metaworld suffer from catastrophic forgetting, and perform equally. We only report curves for the naive baseline, as this drop in performance is independent of the data selection algorithm. Catastrophic forgetting is thus a challenge arising from our setting (in which the pre-training data is unavailable), and not from specific algorithms. As suggested by the reviewer, we have added complete plots as Figure 15.

Tasks selected over fine-tuning

We thank the reviewer for suggesting an analysis of task selection frequencies over training. We added this analysis in Appendix I, where we observe qualitatively interesting patterns. In short, we find that AMF is capable of rebalancing the training data distribution, but is also consistently selecting tasks based on learning progress or correlation.

We thank the reviewer for suggesting several additional experiments. We believe that these additional results have been very helpful in improving our paper. We hope that our clarifications also address the remaining comments. We would be happy to discuss any further questions!

References

[1] Wang et al., “Robot Fleet Learning via Policy Merging”, ICLR 2024

[2] Pena et al., “Re-basin via Implicit Sinkhorn Differentiation”, CVPR 2023

[3] Ainsworth et al. “Git Re-Basin: Merging Models modulo Permutation Symmetries”, ICLR 2023

Dear reviewer, We greatly appreciate your time and effort in reviewing our work, and we believe that our response addresses the points you raised. As the discussion period is nearing its end, we wanted to kindly check if there is anything else we can clarify. We remain available for any further discussion.

Thank you for supporting our submission, and providing more detailed and informative feedback. We understand your comments and would like to provide some more context.

We would like to stress that parameter isolation is adopted as a solution to a problem, that is catastrophic forgetting. Unfortunately, without addressing catastrophic forgetting, fine-tuning a pre-trained model loses much of its nuances, as the pre-training may quickly be undone. As forgetting constitutes a bottleneck, it is unlikely that active data selection can be effective. We find this issue to arise particularly in the setting of Kitchen and Metaworld: in this case, parameter isolation happens to be the best performing solution. We agree that increasing model and dataset size can alleviate this issue, but we believe it is still important to highlight its occurrence with networks and dataset of modest size, and report the solution that was found to be most effective.

One disadvantage of parameter isolation is that it prevents positive interference across tasks, and specializes a set of weights to each task, as the reviewer points out, and we acknowledge in lines 352-353. However, we would like to stress that the uniform baseline (e.g. Fig. 4, in yellow) also adopts parameter isolation. Thus, the better performance of AMF-NN cannot be explained by parameter isolation, and is rather induced by a different data selection strategy.

As requested, we have further updated Appendix G to include results for OCTO without parameter isolation (new Fig. 15). In this setting, we find that AMF-NN without parameter isolation performs within confidence intervals of AMF-NN with parameter isolation. It is competitive with baselines, but does not significantly outperform them. As reported in lines 508-510, this is likely due to the fact that all fine-tuning tasks belong to the pre-training distribution, and a uniform sampling strategy is thus very desirable. However, these results suggest that parameter isolation is not strictly necessary for AMF-NN with large models, as scale partially alleviates catastrophic forgetting. This confirms two important points: (1) catastrophic forgetting is more prominent at smaller scales, and (2) it is a general issue in fine-tuning, and does not affect any strategy in particular.

Finally, we thank the reviewer for stressing the importance of proposing principled fine-tuning techniques, and working towards practical instantiations, which we fully agree with.

We would like to thank the reviewer for their thorough assessment. We believe that we can address all comments directly.

Assumptions

The assumptions in Section 4.1 are used to formally prove convergence with high probability, but they are not be necessary conditions for good performance. For instance, convergence still occurs under much weaker assumptions, albeit in expectation (see Appendix B). In high-dimensional cases, even these weaker assumptions may be violated. This is indeed the case in both Kitchen and Metaworld, particularly when learning from images. In these settings, we observe that AMF still performs well. While our theoretical results provide guarantees on AMF's performance in regular settings, we find it encouraging to empirically observe similar behavior in complex and high-dimensional tasks.

Rebalancing baseline

Thank you for this suggestion, we believe it highlights an important feature of our problem setting. Please, see the common response for a detailed response, and further empirical evidence. In short, we find that our method matches or outperforms the performance of this rebalancing baseline, despite not having access to the pre-training data distribution. Moreover, it displays data selection strategies that are more nuanced than simply rebalancing counts (e.g., it accounts for learning progress and correlations).

Shared structure across tasks

We agree that multi-task benchmarks generally involve relatively disjoint tasks. We selected evaluation environments specifically to include tasks with a significant amount of shared structure. In high-dimensional settings, while standard multi-task benchmarks often involve one scene per task, all tasks within our Kitchen or Metaworld evaluation share the same scene and initial state distribution (i.e., all objects are always present, irrespectively of the task). Thus, state distributions for different tasks are largely overlapping. Moreover, the expert demonstrations for several pairs of tasks will share some sequences (e.g., part of the demonstration for grabbing a cup and moving it to the left is also valid for moving it to the right). For this reason, we would argue that the amount of shared structure is significant, and larger than in many multi-task benchmarks. Additionally, experiments on the 2D integrator involve a continuous task space, in which two tasks might be arbitrarily close to each other, and their demonstration can be very correlated. The final evaluation on Simpler features arguably more disjoint tasks, although most of them involve picking and placing objects.

We hope that our response clarifies our setting, and provides additional empirical evidence demonstrating that AMF goes beyond rebalancing the count of demonstrations. If this is the case, we would ask the reviewer to consider raising their score. We of course remain available for further discussion!

Dear reviewer, We greatly appreciate your time and effort in reviewing our work, and we believe that our response addresses the points you raised. As the discussion period is nearing its end, we wanted to kindly check if there is anything else we can clarify. We remain available for any further discussion.

In light of the discussion phase, we would like to comment on a few important points:

• We are glad to see the reviewers confirm their appreciation for the theoretical guarantees of our method (e.g. "the paper is commendable for theoretically motivating an algorithm", gsEP), as well as its relevance to practical problems ("this demonstrates its practical applicability in realistic scenarios", Qd9g).
• We acknowledge that the particular instantiation of AMF-NN is limited by the issue of catastrophic forgetting, which we could only handle with hard parameter separation. This is not ideal in multi-task settings, as it prevents positive transfer for AMF-NN. Unfortunately, catastrophic forgetting remains a fundamental issue in neural networks, and is still far from being solved. As our main focus is proposing a principled algorithm, and studying how it can scale beyond regular setting, addressing catastrophic forgetting is not our main focus. We sincerely appreciate suggestions on this topic from Reviewer gsEP, and hope that further research in continual learning can provide more scalable approaches.
• With the caveat of catastrophic forgetting, we provide evidence that AMF is applicable across policy classes (GP, MLP, OCTO), and can largely improve adaptation of policies to new tasks that were not primarily represented in the training distribution, with further evidence added as Appendix H and I.

We again thank the reviewers for the discussion and remain available for further comments.