Paper: Robot Fleet Learning via Policy Merging (ICLR-2024)

Reviewer comment (Experiments):
The proposed algorithm was only demonstrated for behavior cloning, which is not what we generally consider as the setup where we benefit from distributed learning. If we have a large amount of static expert data, it is not too difficult to just merge the datasets and train a policy on the joint dataset. When we consider fleet policy learning, it is much more useful to consider a reinforcement learning setup where a collection of robots are collecting data for various tasks, and we wish to collectively use these data sources to learn optimal policies. But this setup is much more complicated due to non-stationarity and other challenges with online RL. It would be much more interesting to see if these merging strategies could work in an online RL setting.

Author response:
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.

Reviewer comment (Novelty):
Could the authors motivate distributed policy learning for behavior cloning? If static expert datasets are already present, it does not seem useful to train separate policies and merge them. Does this approach work in the online reinforcement learning setting?

Author response:
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.

Reviewer comment (Novelty):
The reviewer is leaning slightly towards rejection since behavioral cloning does not seem to be an appropriate task to motivate fleet learning, and the novelty of the FLEET-MERGE algorithm is not clear, though they might raise the score if the authors provide a strong response to these questions and clarify the contributions.

Author response:
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. 

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;

Reviewer comment (Presentation):
The experiment section can be better structured by having subsections for each of the evaluation setting instead of each benchmark; right now, all evaluation settings are presented upfront without concrete contextualization and results that interleave between them.

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

Reviewer comment (Experiments):
The new benchmark Fleet-Tools is not well motivated. The section on "Keypoint transformation & trajectory optimization" within Section 4 is too dense while not directly related to the algorithmic components of the paper. The authors should better motivate why Fleet-Tools is particularly well-suited for studying the policy merging problem.

Author response:
>  (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.

Reviewer comment (Evaluation):
The performance improvement of Fleet-Merge compared to simple baselines is moderate at best. Fleet-Merge does not seem to be substantially better than Git-Rebasin or FedAvg on the Metaworld and Fleet-Tool benchmark.

Author response:
**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.

Reviewer comment (Experiments):
Considering that either baseline is missing key components in Fleet-Merge, it is worth asking whether Fleet-Merge is truly necessary given its added algorithmic complexity as well as computational cost.

Author response:
**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.

Reviewer comment (Experiments):
While the paper emphasizes the multi-task aspect of policy merging, most of the experiments are conducted in the single-task setting. In the multi-task setting, simple averaging appears to be very competitive.

Author response:
**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.

Reviewer comment (Writing):
The definition of Loss Barrier should be included somewhere early. In Section 5, the concept of "performance barrier" is introduced assuming readers' familiarity with the equivalent concept of loss barrier in the supervised learning model merging literature.

Author response:
>  (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.

Reviewer comment (Writing):
While the method proposed is clear to me, I do not understand how the merged policy can determine what task to solve during test time.

Author response:
**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.

Reviewer comment (Evaluation):
In the results section I did not see how well the policies perform on a single task. In figure 4, it is clear that FleetMerge performs better than Federated Averaging, however I am curious to see how much of a decrease in performance is incurred by merging the policies at all.

Author response:
**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.

Reviewer comment (Reproducibility):
At test time, how does the merged policy know what task to solve? Could there be an ambiguity on what the task specification is?

Author response:
**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.

Reviewer comment (Experiments):
What would the performance be of a policy trained on all the data collected during the fleet learning? I am curious to know what the upper bound on multi-task training is.

Author response:
**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.

Reviewer comment (Reproducibility):
When choosing the first permutation matrix, does it matter which RNN policy you align the rest of the policies? I am unclear on how this initialized permutation matrix is chosen/calculated.

Author response:
**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.