Robot Fleet Learning via Policy Merging

Thank you for finding the topic of our work is of “useful contribution to the robot manipulation community” and “increasingly important” , and appreciating our contribution. We address your questions below.

Behavior Cloning:

Thank you for this comment. Although there are a plethora of policy-training techniques, we focus on behavior cloning (BC) for the following main reasons:

1. Its simplicity helps us disambiguate the effects of our method from the intricacies of the underlying training methodology’’.

2. BC is arguably the state-of-the-art approach for the industry-scale robot policy training settings where our method is most useful [1].

3. Even with static/offline datasets, as in BC, data-pooling is not always feasible. In particular, the datasets for policy learning in robotics are usually videos and embodied robotic data (such as observation-action trajectories), compared to those for supervised learning (as texts), can be much heavier to communicate and centralize. Moreover, robots belonging to different service providers/companies/client industries/households are not motivated/have the liberty to share the local data that can contain private information of the clients (analogous to the privacy concerns of federated learning for mobile devices).

[1] Open X-Embodiment Collaboration. Open X-Embodiment: Robotic learning datasets and RT-Xmodels. https://robotics-transformer-x.github.io, 2023

We do agree that it is interesting to see if our method extends to online RL approaches, and we acknowledge the reviewer’s point that the benefits of distributed learning are even more pronounced for online RL training. As per your suggestion, we have included some preliminary experiments in the end of the revised paper (Sec. E in Appendix), demonstrating that yes, indeed, model-merging approaches in general are effective in the RL setting as well and are robust even with little communications.

Algorithmic Novelty:

We want to emphasize that the proposed Fleet-Merge algorithm is different from a simple extension of Pena et al., in the following important aspects:

1. We are dealing with the “policy merging” setting with multiple models rather than just aligning two models, as considered in Pena et al.;

2. Because of 1), our algorithm maintains a central model as the “anchor policy” in merging – as in the seminal FedAvg algorithm – which allows for an “interpolation” between “pure model aligning (with one-shot merging)” and “federated learning (with multiple rounds of merging)”. We consider this as one of our key algorithmic and conceptual contributions. In contrast, a naive application of Pena et al. would require pairwise aligning and merging, as in the GitRebasin paper, which may depend on the order of the merging sequence;

3. We include a “hard projection” step in our algorithm, which takes on a new interpretation. Rather than permuting the weights to evaluate mode connectivity, we use hard projections as a way to compute optimal mode alignment for merging. We show via our ablations (Figure 12 in the appendix) that this is necessary;

4. Our algorithm allows iterative merging, i.e., merging multiple times during the course of training, compared to other (feedforward) neural network aligning approaches as in Pena et al., which allows tradeoffs between communication cost and merged-model performance;

5. Unlike the two-model-merging setting of Pena et al., we do not have access to a common dataset $\mathcal{D}$ for the aligning updates. Instead, each model is updated by sampling trajectories from its local dataset $\mathcal{D}_{local,i}$, obviating the need for dataset sharing (the key consideration in our setup).

6. Finally, because we consider policy learning from partial observations as images/point-clouds, the models are parameterized by recurrent neural networks, forcing us to handle permutation symmetries not only between layers, but also between timesteps which was not dealt with in Pena et al., which only handled feedforward neural networks;

Thank you very much again for your feedback, and please do let us know if you have any other questions and/or comments.

Thank you for finding our work very novel, and appreciate the potential impact in the robot learning community. We also appreciate your comments on presentation and experiments, and have addressed them in the revision and the responses below.

Presentation. We appreciate all your constructive feedback regarding our presentation, and have incorporated them in the revision. We summarize our changes as follows:

(a) We have updated the charts in Figures 3 and 4 as suggested, making them clearer.

(b) We have revised the Experiments section by restructuring it according to Experimental Settings instead of Experimental Benchmark, as per the reviewer’s suggestion.

(c) Thanks for bringing up the confusion about the definition of loss barrier. We have added the definition earlier in Sec. 3, as the reviewer suggested.

(d) We believed the design of Fleet-Tools, as a novel benchmark for robotic manipulation tasks that allows easy construction of compositional and contact-rich tasks, is an important contribution to robot (fleet) learning in its own right (in addition to our algorithmic components). This is why we endeavored to provide a detailed introduction to the benchmark, though we recognize that the writing was somewhat dense, so we have revised it to be clearer. We believe that this benchmark is well suited to testing fleet policy learning for the following reasons:

(i) the ability to adjust proportions of combinations of objects and tools provides a systematic and scalable way to create local/distributed datasets in fleet learning

(ii) the Drake simulator allows for physically realistic dynamics, giving us a better sense of the efficacy of merging in the real world.

(iii) Our benchmark allows for rapid construction of local demonstration data that readily enables high-fidelity policy learning.

(iv) The tool-use tasks considered in Fleet-Tools are more challenging than common tasks in most existing robotic manipulation simulators, e.g., MetaWorld (for which we have tested our framework also). To the best of our knowledge, Fleet-Tools is the first benchmark designed for imitation learning that involves tasks with intricate and challenging dynamics.

As per the reviewer’s feedback, we have simplified the introduction of Fleet-Tools, and added more discussions on its motivation for fleet learning, in the main body of the paper.

Fleet-Merge performance. Thank you for this valuable feedback. We have contextualized the relative increases in performance in the revision. In short, robotic behavior cloning tasks are quite challenging and obtaining any increase in performance should be considered a significant achievement. Notice, for example, that an increase from a score of .55 to .65, as observed in Figure 3(d) with alpha = .01 corresponds to a 20% gain in performance, which we consider to be significant. Moreover, while there is indeed some additional algorithmic complexity introduced by Fleet-Merge, the additional computational overhead is minor. This is because merging operations are infrequent in our overall setup, and thus the cost of merging is dominated by the training time. Beyond the performance improvements, we believe that the conceptual contribution of this paper -- that of model merging over data sharing — will inspire future algorithms to achieve even greater performance improvements as others in the community consider this important problem of fleet policy learning.

Multi-Task. We acknowledge that on the Meta-World benchmark, simple averaging is surprisingly competitive. Notice that this too is an interesting finding, and we view it as a contribution of our empirical study – merging models instead of data can work in multi-task robot learning, rather than a drawback of our proposed approach. However, we believe that the Meta-World experiments exactly motivate our Fleet-Tools benchmark for fleet learning: whereas simple merging may suffice for the former, Fleet-Merge is necessary to handle the additional dexterity for robotic manipulation demanded by the latter. We further note that Fleet-Tools can be viewed as another form of multi-task learning, wherein the combination of each object to be manipulated and the tool used corresponds to a task. We have made such a connection clearer in the revision. We believe that task- or language-conditioning policies will be interesting future work.

Thank you very much again for your constructive feedback. In sum, we have revised the paper significantly by adding more motivations of our setup and benchmark, as well as restructuring the sections and introducing additional explanations to improve the paper presentation. Please do not hesitate to let us know if you have any other questions and/or comments.

Thanks again for your feedback, and we will update the draft to reflect these changes.

Figure 3(d) doesn't seem to suggest an improvement from 0.55 to 0.65. Could you clarify?

Sorry for the typos in the response due to the reorganization of the paper (and subfigures). The improvement should refer to Figure 4 (a,b more specifically) where Fleet-Merge outperforms other methods (and apologize that the exact number should be from 0.44 to 0.65, which is even better). We also highlight that in the supervised learning setting in Appendix F Figure 12, Fleet-Merge can have over 100% improvements compared to FedAvg in some cases.

I am still not sure ... in the context of this paper.

While we agree MetaWorld is a useful benchmark in its own right, we argue that it is necessary to test on a second suite of tasks to confirm that the benefits of Fleet-Merge do not overfit to the specificities of any single environment. We took care to make our methodology for the construction of Fleet-Tools transparent so that it is clear that we did not in any sense tune or overfit this environment to improve the relative performance of Fleet-Merge. Rather, it is because Fleet-Tools tasks are intrinsically more challenging (as described below, and as in our previous response) than those on Meta-World that we have room to see greater relative gains in performance for Fleet-Merge over comparator methods. Also, as mentioned before, the combinatorial composition of tasks in Fleet-Tools makes it scalable and easy to create local datasets in fleet learning.

We acknowledge that Fleet-Tools is not limited to only study fleet learning, but is more general for robotic manipulation. Nevertheless, we believe this benchmark provides a number of advantages:

(1) Tool-use is a challenging task that requires high dexterity and generalization. This allows us to test methods on more challenging tasks.

(2) While there may ultimately be abundant data for these tasks in real robots and human demonstrations in tool-use, we still lack simulation data for effective benchmarking of robotic tool-use tasks. Fleet-Tools provides a purely in-simulation benchmark.

(3) Tool-use is particularly well-suited for studying multitask policies or policy learning under heterogeneity, because it requires object-object affordance, and tool-object compositionality (humans can use one tool for many tasks). This multitask flavor is useful for studying the ability of Fleet-Merge to combine diverse policies, and becomes especially relevant in fleet-learning scenarios.

(4) Fleet-Tools is the first robotic benchmark built with Drake. Drake offers numerous unique benefits as a simulation environment that complements other simulators such as Mujoco or IsaacGym: these features include trajectory optimization, model-based control, optimization, and the code is open source. Our Fleet-Tools thus inherits these advantages as a simulation benchmark.

The set of baselines ... It's not clear the justification behind this choice.

All benchmarks in the paper serve slightly different purposes. We made sure to run all experiments which “made sense” on either benchmark. Specifically, we study linear mode connectivity for both Meta-World and Fleet-Tools in Figure 3 (a,b), and robustness to communication frequencies in Figure 4 (a,d). However, there were some experiments that were better suited to one benchmark versus the other. The choice of results reported in Figures 3 and 4 is purely a matter of presentation. We will explain these decisions here and then be sure to clarify further in the revision:

(1) Meta-World is a relatively simple standard benchmark and is used as a proof-of-concept. It can run faster and in larger-scale experiments; it also has many tasks to choose from, which means policies training on these separate tasks can be very diverse. Therefore, it is well-suited to testing the effects of changing the fraction of policies selected for merging in Fleet-Merge at each round (Algorithm LIne 6). This is what has been reported in Figure 4c.

(2) Fleet-Tools is more challenging and closer to the real-world setup. Thus, we were able to include real-world figures and some preliminary sim-to-real results in Appendix Figure 7. Because of this, we used Fleet-Tools to evaluate experiments that would be more indicative of real-world setups. This motivated our use of Fleet-Tools in the fine-tuning experiments (Figure 3c). We will be sure to include fine-tuning results on Meta-World as well for the camera-ready.

As described in the paper, we investigate other ablations with our Linear control benchmark and supervised learning experiments.

Thank you very much again for the feedback! Hope our responses and revisions have addressed your remaining concerns, especially regarding "practical utility of the method" and "more clearly present the paper" as you have commented before. Feel free to let us know if there is anything else that can help improve the paper.

Thank you for finding the topic of our work very interesting and novel, and the paper well-written. We address your questions as follows.

Which task to solve: This is a good point and we have added clarification in the revision. We consider multi-task settings in which the task can be inferred from observations. For example, in Fleet-Tools, the task is always to apply a tool to the given object, and the tool and object can be inferred from the point cloud, which is the input of the policy. For both Fleet-Tools and Metaworld, we verify that there exists a single policy which executes all the tasks with proficient performance. We have included plots of this centralized trained policy in the Appendix. We also note that our method is entirely consistent with task- or language-conditioned policies, and this remains an exciting direction for future work. If we do not know what the test-task is, we can envision Fleet-Merge as a tool for training a large-pretrained model, which can be refined further via finetuning. Indeed, our experiments (Figure 3(b)) show that Fleet-Merge is effective for this purpose.

Performance of training on all data: Thanks for bringing this up. We have added the upper-bound performance of the joint policy training to figures and mentioned in captions as “the performance is upper-bounded and lower-bounded by the joint policy success rates and one-shot policy merging.” We observe that there is some degradation of the performance depending on the frequencies of (thus the number of times of) averaging, but overall the weight merging works surprisingly well. More specifically, in Figure 3 describing the wrench task in Fleet-Tools, the upper bound is around 0.82 success rate if we pool all the task data together. For Metaworld, the joint training success rates often can be really high (around 0.9). We also like to emphasize that we do have multitask results and one can also check the Metaworld experiments in the Appendix for larger-scale experiments.

Choice of initial permutation: The permutation matrices are initialized as identity and the initial anchor policy is the “naive averaged” policy. We attempted to improve performance by experimenting with more advanced initialization schemes, but none improved performance to a degree warranting the additional complexity. We describe this detail in Algorithm description in Sec. 3.1.

Thank you very much again for the positive comments, and hope our responses have fully addressed your questions.

I am satisfied with the authors' responses to my questions and maintain my score. Thank you for the detailed response!