Robot Fleet Learning via Policy Merging
Authors: Lirui Wang
Kaiqing Zhang
Allan Zhou
Max Simchowitz
Russ Tedrake
Venue: ICLR 2024 poster (2024)
Keywords: Fleet Learning
Weight Merging
Multi-task Policy Learning
Primary area: applications to robotics, autonomy, planning
TL;DR: We investigated a "bottom-up" approach to learn robotic policies from a fleet of robots, by leveraging weight merging, and the method demonstrates strong results on settings such as linear control, Meta-World, and a novel robotic tool-use benchmark.

Abstract: Fleets of robots ingest massive amounts of heterogeneous streaming data silos generated by interacting with their environments, far more than what can be stored or transmitted with ease. At the same time, teams of robots should co-acquire diverse skills through their heterogeneous experiences in varied settings. How can we enable such fleet-level learning without having to transmit or centralize fleet-scale data? In this paper, we investigate policy merging (PoMe) from such distributed heterogeneous datasets as a potential solution. To efficiently merge policies in the fleet setting, we propose FLEET-MERGE, an instantiation of distributed learning that accounts for the permutation invariance that arises when parameterizing the control policies with recurrent neural networks. We show that FLEET-MERGE consolidates the behavior of policies trained on 50 tasks in the Meta-World environment, with good performance on nearly all training tasks at test time. Moreover, we introduce a novel robotic tool-use benchmark, FLEET-TOOLS, for fleet policy learning in compositional and contact-rich robot manipulation tasks, to validate the efficacy of FLEET-MERGE on the benchmark.

Published as a conference paper at ICLR 2024
ROBOT FLEET LEARNING VIA POLICY MERGING
Lirui Wang', Kaiqing Zhang", Allan Zhou', Max Simchowitz', Russ Tedrake' MIT CSAIL University of Maryland, College Park *Stanford
ABSTRACT
Fleets of robots ingest massive amounts of heterogeneous streaming data silos generated by interacting with their environments, far more than what can be stored or transmitted with ease. At the same time, teams of robots should co-acquire diverse skills through their heterogeneous experiences in varied settings. How can we enable such flect-level learning without having to transmit or centralize fleet-scale data? In this paper, we investigate policy merging (PoMe) from such distributed heterogeneous datasets as a potential solution. To efficiently merge policies in the fleet setting, we propose GLEET-MERCEE an instantiation of dis- tributed learning that accounts for the permutation invariance that arises when parameterizing the control policies with recurrent neural networks. We show that FLEET- MERG consolidates the behavior of policies trained on 50 tasks in the Meta-Wore World environment, with good performance on nearly all training tasks at test time. Moreover, we introduce novel robotic tool-use benchmark, FLEET-TOOLS, for fleet policy learning in compositional and contact-rich robot manipulation tasks, to validate the efficacy of FLEET-MERGE on the benchmark."
INTRODUCTION
With the fast-growing scale of robot fleets deployed in the real world, learning policies from the diverse datasets collected by the fleet (Osa et al., 2018; Bagnell, 2015; Kumar et al., 2020; 2022) becomes increasingly promising approach to training sophisticated and generalizable robotic agents (Jang et al., 2022; Levine et al., 2020). We hope that both the magnitude of the data streamed and actively generated by robots interacting with their surroundings the environments and tasks around which the data are collected, will allow robot fleets to acquire rich and varied sets of skills. However, the data heterogeneity and total amount of the data are becoming as much of a challenge as a benefit. Real-world robot deployments often run on devices with real-time constraints and limited network bandwidth, while generating inordinate volumes of data such video streams. Hence, a "top-down" scheme of centralizing these data (Grauman et al., 2022; Collaboration, 2023), and training single policy to handle all the diverse tasks, can be computationally prohibitive and violate real-world communication constraints. At the same time, we wish to consolidate the skills each robot acquires after being trained on its local datasets via various off-the off-the-shelf robot learning approaches. Thus, it is natural to ask: How can the entire fleet efficiently acquire diverse skills, without having to transmit the massive amount of heterogeneous data that is generated constantly in silos, when each one of the robots has learned some skills from its own interactions?
and the diversity of
To answer this question, we propose policy merging (Figure 1), PoMe, a "bottom-up" approach for fleet policy learning from multiple datasets. Specifically, consider policies that are already trained separately on different datasets and tasks, and seek to merge their weights form one single policy, while preserving the learned skills of the original policies. Policy merging acquires skills efficiently with drastically reduced communication costs, by transmitting only the trained weights of neural networks but the training data. Such bottom-up merging scheme is agnostic to and thus compatible with any local training approaches used in practice.
Merging the weights of neural networks has been studied in various contexts, including finetuning foundation models (Wortsman et al., 2022a;b), multi-task learning (He et al., 2018; Stoica et al., 2023). and the investigation of linear mode connectivity hypothesis for (feedforward) neural networks (Frankle et al., 2020; Entezari et al., 2021; Ainsworth et al., 2022). Distributed and federated learning
"oode available http::/ggthah..com/iiiii/ieettt

---

Published as a conference paper at ICLR 2024
Data Pooling and Centralized Training
-TB/day
D
D2
D D
DN
2
-MB/day
Local Training and Policy Merging
We consider the problem of merging multiple policies trained on potentially distinct and diverse tasks, which can be more computation and communication efficient than pooling all data together for joint training. Instead of acquiring astronomical size of data from the top-down (red arrow, requiring terabytes-per-day worth of data transfer), we demonstrate that the bottom-up approach (green arrow, megabytes-per-dan policies that incorporate skills learned by the individual constituent policies. Moreover, local training and sharing weights are more suitable for agents that actively generate data, which is especially the case in robotic control, and are more efficient in communicating with the other agents. We aim to achieve the following objective in fleet learning: One robot learns, the entire fleet learns.
Figure 1:
merging from locally trained policies, can also produce general
(McMahan et al., 2017; Wang et al., 2020a; Konečný et al., 2016), which iteratively update central model weights from decentralized updates, can be viewed as an iterative form of model merging; these approaches have achieved tremendous success in learning from diverse datasets, especially in solving supervised learning tasks (McMahan et al., 2017; Wang al., 2020a; Mansour et al., 2020). However, such an approach has not yet demonstrated its power in solving robot learning and control tasks, which are generally more challenging due to their dynamic and sequential nature, and the richness of the tasks and environments. Notably, with the commonly used sensors in robot learning, e.g., cameras, one has to handle the partial observability in learning such visuomotor policies (Levine et al., 2016), which necessitates the use of latent state in parameterizing the policies, and is usually instantiated by recurrent neural networks (RNNs) (Elman, 1990). Indeed, policies with latent state dynamics are known to be theoretically necessary for partially-observed linear control problems (Bertsekas, 2012), and outperform policies parameterized by stateless policies in practice, i.e., those using feedforward neural networks, in perception based robotic tasks (Andrychowicz et al., 2020).
In this work, we address policy merging in robotic fleet learning, with a focus on ENN-parameterizer policies. A naive approach of averaging the weights would easily fail because multiple configurations of network weights parametrize the same function. One reason behind this is the known permutation invariance of neural networks, i.e., one can swap any two units of hidden layer in network, without changing its functionality (Hecht-Nielsen, 1990; Entezari et al., 2021). We have to account for such invariance in merging multiple policies. Indeed, such a fact has been accounted for recently in merging the weights of trained neural networks, with extensive focuses on aligning the weights of feedforward neural networks, and solving supervised learning tasks (Entezari al., 2021; Ainsworth et al., 2022; Peña et al., 2022). We generalize these insights to the RNN setting, where permutation symmetries not only appear between layers, but also between timesteps, in solving robotic control tasks. Compared with one of the few federated learning methods that also explicitly account for permutation symmetry in merging RNNs (Wang al., 2020a), we develop a new merging approach based on "soft" permutations (see Section 3 for formal introduction) of the neurons, with more efficient implementation and an application focus on robotic control tasks. We detail our main contributions as follows, and defer detailed related work overview in Appendix A.
We design a new policy-merging approach, FLEFT-MERGE, that outperforms baselines by over 50%, by accounting for the permutation symmetries in RNN-parameterizee policies, and also extend the approach to the training stage, by allowing multiple rounds of merging between each training update, and also extend to merging multiple (more than two) models.
We evaluate our proposed approach with different input modalities such as states, images, and pointelouds, in linear control and the Meta- Meta-World benchmark (Yu et al., 2020) settings.
We develop and evaluated on novel robotic tool-use benchmark, FLEET-TOOLS, for policy merging and fleet robot learning in compositional and contact-rich manipulation tasks.
2

---

Published as a conference paper at ICLR 2024
2 PRELIMINARIES
Model. We consider the general setting of policy learning for controlling dynamical system, when the system state may not be fully observed. Specifically, consider a sequential decision-makin setting with time index / 2 1, where time I, robot makes observations or E O of latent state 8, E S, and then selects a control action a, E A. Dynamics and observations are characterized by probability distributions Pobs and Payn. such that or ~ Poba(- st) and 8141 ~ Payu (- $t, at). In the special case with full state observability, we have O = S and = st. In the special setting of linear control, both the dynamics and observations can be characterized by linear functions (with
additional noises), see Appendix G for formal and detailed introduction.
Feedforward and recurrent policies. For simplicity, we consider the case where the robot agent exe- cutes deterministic recurrent policies æ. Specifically, we parameterize these policies by maintaining policy state l4, updated as h4g = (0t: 11-1). As a special case, we consider static feedback policies O - AA that select a = x (ot) as function of only the instantaneous observation. Given non- linear activation function a(-). which is applied to vectors entry-wise, we can parameterize the static feedback policies by L-layer feedforward neural networks with parameter ( (W ft:
b" Josessss L-1+ where W and be denote the weight and bias at layer f, respectively. The action a, tT (o:) given observation or is then given by
h OF, at h4, h (W + +b),, 0 0
(2.1)
where above he are the hidden layer activations. For the general case with recurrent policies, parameterize them with Elman recurrent neural networks (Elman, 1990), with parameter 0 = (W rec W bd) bosesss-.-- Let h4 = (Mfssegs be a sequence of hidden states, such that h4 = #t
To(or, (tt 14-1) 1) given by
h Ob. a₂ = h4,
hy o (W hat + With + b'), 0
(2.2)
The presence of the recursive term W h,+1 that incorporates the hidden state at the previous time t distinguishes the RNN architecture in Eq. (2.2) from the feedforward architecture in Eq. (2.1). Notice here that at time t, the action a is just the last layer of the hidden state h,"
Permutation invariance. As well-understood for supervised learning (Frankle et al., 2020; Wortsman et al., 2022b; Ainsworth et al., 2022), policies trained separately even on the same dataset - can exhibit similar behavior whilst having very different weights. This is in large part due to the invariances of neural network architectures to symmetry transformations. For an L-layer neural network with layer-dimensions do, d1, ..... : dt, let Gperm denote the set of hard permutation operators. These sequences of matrices p = (P", p1 ..... PL), where we always take p Idg p Ldz and take p as a de x de permutation matrix for S EéS L 1. We let Güin 3 9perm denote the set of linear transformation operators that are sequences of matrices P = (Po P!...... PL), where we still have p Id., pL Idz, but now we allow (P1 ,..., ....,,,.111.
be general invertible matrices.
Elements of Gin (and thus Sperm) act on feedforward models 0 via
(Wé, f:: b²) + ( is known that the feedforward architecture Eq. (2.1) is invariant, in terms of input-output behavior, to all hard permutation transformations P E Sperm, but not to general P E Gin- When the activation function is an identity mapping, i.e., the neural networks are linear, becomes invariant to Gin-
(2.3)
Imitation learning. As a basic while effective imitation learning method, we here focus on behavior cloning (Osa et al., 2018; 8; Bagnell, 2015) for the purpose of introducing the policy-merging framework next. Note that our merging framework and algorithms will be agnostic to, and can be readily applied to other imitation learning algorithms. In behavior cloning, one learns a policy ne. parameterized by some 0 E RRd. that in general maps the observation- action trajectories to actions, by imitating trajectories generated by expert policies. Let D (7()) Isis denote a set of M trajectories, with (+) = (o(), a. JusEST denoting the i-th trajectory of length T. As an example, we study behavior cloning with the Eg-imitation loss, instantiated by (0c D) = EM Lhe ((9), where for given 7 = (ot, () aI)SSTTT
(0: == Σtz ||āg,1 aell2, where har == To(or,ho,c11) ho,c-)) āg.1 = h4g
(2.4)
3

---

Published as a conference paper at ICLR 2024
where h denotes the hidden state that arises from executing the recurrent policy Ta on the observa-
tion sequence 01, O2...., Ot, via the Elman recurrent updates in Eq. (2.2), and the action is part of the hidden state corresponding to the last layer. Note that in the special case where xo is static feedback policy, we can drop hgge from the above display, and generate each āø, using Or based on Eq. (2.1).
Policy merging framework. We now introduce policy merging, our framework for fleet policy learning from diverse datasets. Consider N datasets collected by fleet of N robots from possibly different tasks and environments, and can potentially be highly heterogeneous and non-i..d. Each robot agent i 1,2, .... N only has access the dataset D; of itself, in the form given in Section 2. Ideally, if the robot designer can pool all the data together, then the objective is to minimize the following imitation loss across datasets
min Σ (.) 0
(2.5)
Let Pual denote the solution to Eq. (2.5), i.e., the best policy parameter one can hope for when seeing all the data, and will provide an upper bound for the performance we will compare with later. Let 0; denote the policy parameter of robot by minimizing the loss associated with D;, i... 0 E argming (0: Dz). The goal of policy merging is to find a single policy parameter marg- as some aggregation of the local policy parameters (01,... ,......,,,,,, ,6N), without sharing the datasets.
An example of this aggregation direct averaging, i.e., 0mry = - E 0,. We propose more advanced policy merging methods in Section 3 by accounting for the symmetries of RNN weights. Note that the above merging process can also occur multiple times during training: at round we first merge the trained policies (0f").. ................................ xx)) ) to obtain Olrgg and send it back to the robots to either conduct more training and/or collect more data. At round 2, the newly trained local policy parameters
(o(2),... ..... (22) are then merged to obtain 0 This iteration can proceed multiple times, and we will refer to it as the iterative merging setting. When the merging is instantiated by direct averaging. this iterative setting exactly corresponds to the renowned FedAvg algorithm in federated learning (McMahan et al., 2017). When there is only one such iteration, we refer as the one-shot merging setting. Fewer iterations lead to fewer communication rounds between the robots and the designer, and note that no data is transmitted between them.
3 METHODOLOOGY
Given the invariance properties introduced in Section 2, merging by naive averaging the parameters Omrg + + + 0g may not perform well. Prior work has instead proposed merging aligned feedforward neural network models Oorg + + M. (,,),, where the elements PI..... PN are weight transformations which are often, but not necessarily, hard permutation operators (i.e., elements of Gperm). The GITREBASIN algorithm (Ainsworth et al., 2022) iteratively computes the weight permutations (Pt). At each step, agent index i is drawn uniformly from [N], and one constructs - 0% = N- Ejjt Oj by averaging the parameters of indices j & i. It then solves a series of linear assignment problems (LAPs) (Kuhn, 1955; Jonker & Volgenant, 1988; Bertsekas, 1998) for each layer l to find some Pe, which is derived by matching the activations between two models via ordinary least squares regression. The algorithm then repeats the sampling from (0I..... 0x) and the computation of 0% until convergence.
Peña et al. (2022) instead propose a gradient-based variant to merge two models by relaxing the rigid constraint of using hard permutation matrix. The direct extension of their algorithm to our setting is as follows: given two models (0, 0'), iteratively trajectories T from common dataset D, and update the aligning parameters P by following the gradient of (a)(aP()) (1 a)0'; T)), where Lbe is as given in Eq. (2.4). Thus, for each iteration 2
P. + P. 1/ (P))(( + (1 - a)0': P=P,' a ~ Unif 00, 1]. T~D D
(3.1)
with some stepsize 17 > O. Note that the updated matrices in P. - (P. P.:) are not necessarily (even close to) permutation matrices. We define soft permutation projection with regularization
For convenience, speak heuristically of exact minimizers in this section. In practice, we understand "minimizer" "model trained minimize the given loss".
4

---

Published as a conference paper at ICLR 2024
Algorithm 1 FLEET-MERGER Fleet Learning of Policies via Weight Merging
1: Input: Models ...... NN. datasets Dlocal. for each i [N]
2: Parameters: Epoch length E, iteration number S, soft-projection parameter T > 0, stepsize
1)
3: Initialize: Permutations Phard,1,..... Nhard,N, N Psoft.1,..... Psoft, + Identity
4: for Epoch ........ do
5: Average models: 4- E (0ara,
6: Sample indices I C [M]
7: for E do
8: for Iteration I....TTuu do
9:
Sample data pair (o, a) ~ Diocal,i to form a trajectory т, and sample interpolation parameter Unif(0,
Update with gradient: Psott.i Psott //V (aP(0) +1 a)); p
11:
Update Psoft,e & Proj, (Psott:)) by applying the soft projection step Eq. (3.2)
12:
Update Phard, f- Projhard (Psoft,))
13: Return: 0 2/ Em= (0,)
0 as:
Proj,( = - I---1 where P argmax (P', P), + T.((()
(3.2)
PEAA,
and the associated hard permutation projection Projhard == Proj- r= where Bd is the Birkhoff polytope of doubly-stochastic matrices, (P) 2013; Mena et al., 2018), and т > 0 is some hyperparameter that weights the strength of the entropy regularization. Computation of Projhard can implemented efficiently via solving linear assignment problem, and the solution with 7 0 can be solved approximately via Sinkhorn iteration (Eisenberger et al., 2022; Peña et al., 2022), which also allows gradient computation on any differentiable objective. The operators are then updated as P.+1 t Proj, (P.). Note that for / and I., we directly set P to be identity matrices. At the final step, Peña et al. (2022) conducts hard projection onto the space of permutation matrices.
Pty log(P (1)) is the matrix entropy (Cuturi,
To measure the performance of model alignment, we study the (imitation) loss barrier, defined previously the supervised learning setting (Frankle et al., 2020; Ainsworth et al., 2022). Specifically, given two policy parameters 0, 0' such that Lhe (0; D) R3 Lb.e (0'; DD, the loss barrier of the policies, defined as maxxe[o.1] (bc A)0 + Ad'; D) ((O) D) (D)+ + (8'; (B)) evaluates the worst performing policy linearly interpolating between 0 and 0', where we recall the definition of Zbc above Eq. (2.4). More amenable the control and policy learning setting, task performance barrier, which replaces the behavior cloning loss Loc with any suitable measure T of task performance (e.g.. the accumulated rewards maxxe(o.1) (T()) + T(6))) T((1 A)Ø + AO'); ); the sign is flipped to model the rewards achieved in accomplishing the tasks. These metrics will be used in our experiments Section 5.
also define the
the success rates of task completion):
3.1 MERGING MANY RECURRENT POLICIES
In this section, we describe our new algorithm for merging many RNN-parameterized policies.
Permutation invariance of RNNs. Given recurrent neural network with parameter 0 (W W b D'sssss-11 as parameterized in Eq. (2.2), we let weight transformations ff
(P" Pl p²) Iln act on through
(Wé ree: ft 1--1bb611) + (P'Wé (P¹)-1 p' ( pbb
(3.3)
for each layer with 1 < L-LL- 1. In Appendix B, we verify that RNNs are invariant to the above operation when E 9perm Giin are hard permutation operators.
Proposition 3.1. Any recurrent neural network given by Eq. (2.2) is invariant to any transformation of P Sperm-
5

---

Published as a conference paper at ICLR 2024
In Appendix B, also expand upon the group structure of Sperm to the larger sets of invariance groups, for the special architecture with ReLU activations. Moreover, we also argue that permutations in essence the only generic invariances of ReLU and polynomial networks whose weights minimize the lz norm, which might be of independent interest.
Merging many models with a single reference. Rather than sequentially merging N = 2 models, we merge all models to common reference model 0. Inspired by the update rules in federated learning, this approach has the following advantages: (a) it removes the dependence of sequential merging on the order of the merging sequence; (b) it allows better pooling of weights from many models to guide each merging step; (c) we align only a subsel of models per iteration, which becomes more efficient when N is very large. We show the benefits of our approach over sequential merging in ablation studies (Appendix F). In addition, unlike the two-model- -oeeeeg setting of Peña et al. (2022), we do not have access to a common dataset D for the aligning updates. Instead, each model is updated by sampling trajectories from its local dataset Piocal, obviating the need for dataset sharing.
Algorithm description. Our algorithm, FLEET-MERGE, is depicted in Algorithm 1. We maintain hard transformation operators ... Phard.N E Gperm. initialized by identity matrices. At each epoch. we compute the reference model 0 by averaging each model under the associated transformation (Line 5). We then select subset of models I, and initialize the "soft" permutation Psoft, +- P as the hard permutation operator. For each E I, we update the "soft" permutation Psoft. for T steps. Importantly, because each 0g corresponds to a recurrent neural network model, the action of (B.) in the gradient step in Line 10 is given by Eq. (3.3). Of equal significance (and as noted above), the trajectories T in Line 9 are sampled not from a common dataset, but rather from a local dataset i associated with the i-th agent. We conclude by re-projecting each Psoft., onto Operm to obtain new Phard,; (Line 12), which are used to update 0 accordingly in the next epoch. In particular, our algorithm allows iterative merging, i.e., merging multiple times during the course of training, compared to other (feedforward) neural network merging approaches (Ainsworth et al., 2022; Peña et al., 2022; Stoica et al., 2023), which allows tradeoffs between communication cost and model performance. At test time, the merged policy Ty is deployed for the test task. In the multi-task setting, we consider the cases where the task identity can be inferred from observations.
4 FLEET-TOOLS: A ROBOTIC TOOL--SE SIMULATION BENCHMARK
To validate the performance of our policy-merging framework, we develop a new robotic tool-use benchmark, FLEET-TOOLS, in the robotic simulator Drake (Tedrake & the Drake Development Team, 2019). The benchmark FLEET- focuses on robotic manipulation tasks with rich contact- dynamics, and category-level composition, in several existing literature in robotic manipulation for tool-use (Toussaint et al., 2018; Holladay et al., 2019).
Expert tasks. To scale the generation of expert demonstrations for the tool-use tasks, we specify the tasks via keypoints (Manuelli et al., 2022) when generating the expert trajectories. We mainly consider four skills examples (or task families): (wrench, hammer, spatula, a, and knife) consisting of tasks involving the eponymous tool-type. Specifically, we consider use a spanner to apply wrenches, use hammer to hit, use a spatula to s scoop, and use kni fe to split. All of these tasks can be solved to reasonable performance by specifying a few keypoints. For instance, in the wrench tasks as shown in Figure 5, the wrench needs to reach the nut first and then rotate by 45 degrees while maintaining contact. The same strategy can be applied to a set of nuts and wrenches. We choose these tasks because they provide a balance of common robotic tasks that require precision, dexterity, and generality in object-object affordance. We use the segmented tool and object point-cloud as the observation o₁, and the translation and rotation of the robot end effector (6-dof) the action a,- At test time of multi-task setting, the task would be identified by the observations, i.e., the (combination of) a tool and an object, which can be inferred from the point could as the input of the test policy, identifies the task being tested on.
Keypoint transformation & trajectory optimization. We denote the end effector frame and world frame as E and W. respectively. For each tool, we specify n keypoints, for example, when n we can specify one point on the tool head, one point on the tool tail, and one point on the side. We define rigid transformation X SE(3) R the current locations of the keypoints on the tool in the world frame W as W " "2001
W too1.1 ,.., "'p ptoe1.n], and the poses in the

---

Published as a conference paper at ICLR 2024
spatula
knife
hammer
wrench
Figure 2: FLEET-TOOLS Benchmark. We develop several tool-use tasks that focus on contact-rich motions and category-level compositions in the Drake simulator.
3 =p f0t respectively. For the object, specify m keypoinnss in the world frame W poobj [W pobs.1 W "poj..m Then, one can solve an optimization problem, referred to as kPAM-Opt (where kPAM stands for KeyPoint Affordance-basee Manipulation) (Manuelli al., 2022), to find the transformation X (see Eq. (C.1)-Eq. (C.3) for the detailed formulation). However, recovering the joint angle from the transformation X via inverse kinematics can be computationally challenging. Hence, we propose to solve the kPAM optimization in the joint space, i.e., optimize over the joint angles q that parameterize the transformation through forward kinematics (i.e., forward-ke nemat is (q)), and only add constraints to the keypoints in the transformed frame. Specifically, solve the following joint-space optimization problem:
end-effector frame as "pt = =poool,1
Joint-Space-kPAM-OP
min X
X E P tool
W
p tool
(4.1)
S.t.
X P tool P
bbbs. S E, Vi.j
(4.2)
Oij S 3 X E 100l F pptool.j ai.j te, Vi,j
(4.3)
X =
(4.4)
where we choose i E {1.... n2 and E (1,... ... m). and Bi,j E IR' is vector the unit sphere with the fourth dimension (4) 0, aij is some constant that represents the target angle alignment, and >> 0 is some relaxation level. We provide detailed explanation of the constraints in $C.1. After finding the optimal joint angle, one can find pre-actuation and the post-actuation trajectories by solving trajectory planning problem (see Eq. (C.8)-Eq. (0. for more details).
The benchmark can be easily extended to new tool-use tasks by labeling few keypoints and providing configuration for constraints and costs for the desired task. The overall pipeline also easily represent "category-level" composition and generalization, across various (combinations of) tools and objects. providing a systematic and scalable way create "local and distributed demonstration datasets". This makes FLEEE--OOO TOOLS great fit for evaluating large-scale fleet policy learning. More details of the benchmark can be found in Appendix C.
EXPERIMENTS: POLICY MERGING
We evaluate the performance of various algorithms in several benchmark environments: FLEET- TooLs, described in Section 4, manipulation tasks, and linear control. This section summarizes results on FLEET-TOOLS and Meta- World, deferring additional experimental details with different network architectures and results to Appendix C.2 and Appendix D. Details for our results on linear control problems can be found in Appendix G. Ablation studies have also been conducted for supervised learning settings including classification (Deng, 2012) and language generation (Caldas et al., 2018); these are deferred to Appendix F. To summarize, the ablations validate the benefits of: (a) merging to a common reference model (Line 5) and (b) performing hard projections (L.ine 12) in Algorithm 1, as well as other algorithm-design choices.
(Yu et al., 2020), which has 50 distinct robotic
Baselines. We compare our FLEET-MERSE against the following baselines [NAIEEVEERAGE.
GITREBASIN, SINGLEDATASET): NAIVEAVERAGE and GITREBASIN denote the method of naive averaging and that in Ainsworth et al. (2022), respectively, i.e., for NAIVEAVERAGE, it directly averages 0 = X0 + (1 - 2)0', and for GITREBASIN, it computes an aligning permutation P. from the ReBasin algorithm in Ainsworth et al. (2022) and merges 0 20 + (1 (O") SINGLEDATASET trains on a single dataset which differs from the test dataset, and thus measures the performance of "zero shot adaptation" in contrast model merging. Each merging method is applied only once after training is complete. We study how the performance of these methods varies as the source distributions of datasets which the different models are trained become more diverse.
7

---

Published as a conference paper at ICLR 2024
(a) Mode Connectivity on Fleet-Tools
(b) Mode Connectivity Meta-Worke
Hammen-To
tettom koodem
moveroperst
0.83
0.
04
Sealerange ressonn
0.52
0.35
0.5
0.37
035
100
0.50
(.5
0.39
035
3.00
nahe
00
8.5
1.6
3.
05
10
06
05
30
14
(c) Merging as Warmstarting on Fleet-Tools
(d) Multi-task Merging Performance on Mcta-Worke
0.50
0.45
500
0.40
Method
902
0.42
0,42
0.51
0.35
Dectreegge
0.35
0.30
2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0
Epochs For Finctuning
singlecallee
exerage
flettrregge
Figure 3: Results on One-Shot Merging. (a & b) FLEET-MERGE attains better mode connectivity for policies in both FLEET-TOOLS and Meta-World. (c) One-shot merging performance is predictive of relative performance after finetuning. (d) FLEET-MERGE succeeds in multitask settings with varying data heterogeneity for 5 data sources, measured by the Dirichlet parameter a.
Data heterogeneity. We also validate the performance of policy merging by varying the levels of data heterogeneity, an important metric in distributed and federated learning (McMahan al., 2017; Hsieh et al., 2020). We adopt popular approach in federated learning (Hsich et al., 2020) to create data heterogeneity: for each experiment, we begin with K component distributions Di...., D (e.g., imitation learning from a hammer task v.s. spatula task). We then sample N source distributions by sampling a mixture weights pl, .PN & E (K) ~ Dirichlet x) from an evenly- weighted Dirichlet distribution with parameter in dimension K. Smaller biases toward "peakier" probability distributions, encouraging greater dataset heterogeneity. By default, we use = 5 and each data source has the same number of data points.
5.1 ONE-SHOT MERGING
We first study one-shot model merging, and validate the presence of mode connectivity between the policy models, as in the weight-merging literature for feedforward NNs (Frankle et al., 2020; Entezari et al., 2021). We consider two policy models 0, I' trained on the same dataset, and measure the performance of NAIVEAVERAGE: and GITREBASIN (as described above), and SOFTUPDATE, which computes the aligning perturbation using the sroteprojention updates as described in Section 3.
FLFET- TOOLS. As preliminary test, we first train multiple feedforward policies single dataset with distinet tasks with behavior cloning. In Figure (a), we evaluate the performance barrier for every pair of the trained policies. We observe that the algorithm that accounts for permutation alignment has a significantly smaller performance barrier when interpolating along the segment in parameter space between two aligned weights. This implies that it is possible to acquire a policy by merging the policies trained on separate datasets, without sharing data or changing the onput-output behavior of each individual model. We extend this merging setting to multiple datasets and multiple tasks. In Figure (c), we show that the merged model can also serve as a warm start for policy finetuning, which be further finetuned downstream tasks with held-out small dataset. In this case, FLEET-MERGE enjoys superior performance compared to naive averaging, git-rebasin Ainsworth et al. (2022)'s extension multiple models, and Sinkhorn rebasin Peña al. (2022)'s.
Meta-World. We use the frozen ResNet features on the images as policy inputs. In Figure (b), we compare different merging algorithms by measuring the mode connectivity in the one-shot merging regime. We observe that there are almost no barriers for certain tasks. See Appendix D for more experiments with different network architectures. outperforms baseline methods in the challenging multitask setting, across all settings of the Dirichlet parameter a (which, we recall, controls the local dataset heterogeneity).
Figure 3 (d), we show FLEET-MERT MERGE
Linear control. Lastly, we evaluate on Linear Quadratic Gaussian (LQG) control tasks with high- dimensional observations. We briefly summarize our findings here and defer detailed description of the setting, algorithms, and experiments to Appendix G, together with the implications of policy merging on the topology the LQG landscape. averaging of policy parameters, permutation- based averaging computed via linear assignment, and a gradient- based alignment approach that aligns over all invertible matrices. First, we discover that framework still be effective for merging linear control policies, in both single-task and
short, we compare three baselines: direct
8

---

Published as a conference paper at ICLR 2024
(d) Robusiness to Merging Frequencics on Flect-Tools
(b) Robustness to Heterogencity on Flect-Tools
0750.33
0.6
0.6
0.52
0.4
0..4
Matied
0.2
festmeegg
Festererge prenom-ccta
Epochs Between Policy Marging
100
200
eeerag
0.0
[c) Robustmess to Merging Ratios on Meta- Meta-World
(d)
Menging Frequencciss on Meta World
0.73
0.67
0.78
0.70
0.71
0.85
0.86
0.81
0.75
0.68.
0:74
feet-merge ge 02
fleet merge
Feettug
FodAvg
0.2
044 04
100
Figure 4: Results on Iterative Merging. We showed that FedAvg with our merged algorithm achieves better performance both FLEET-TOOLS and Meta-World benchmarks, when changing (b) Dirichlet parameters, (c) partial participation ratios (25 tasks), and (a,d) communication epochs. The perfor- mance is upper-bounded by joint training (which pools all the data together for training).
certain multi-task settings. Additionally, because control policies exhibit symmetries to general invertible linear transformations, we find that the gradient-basee alignment approach performs the best. This directly contrasts what we find when merging RNN policies with nonlinear activations. Indeed, in the ablations depicted in Figure 12 in Appendix F, we find that failing to project the "soft permutation" matrices back to hard permutations leads to non-trivial degradation performance when merging variety of neural network architectures. Thus, linear experiments demonstrate that the group structure of policy invariances can be essential for proper algorithm design.
5.2 ITERATIVE MERGING
We also evaluate the performance of our framework when merging models over the course of training, i.e. iterative merging. We focus on multi-task problems in this setting and try minimize the merging frequencies and ratios under high non-IIDness. Specifically, at every M epoch, we merge models according to different methods, and then reinitialize all the individual models 0, the merged model 0 farg. This scheme allows more communication rounds between the robots and the designer, trading-off the biases caused by local training when there is high data heterogeneity across robots. This scheme connects to the popular distributed learning protocol of federated learning (McMahan et al., 2017). See also Section 3.1 and Algorithm for more details of this setting.
FLEET-TOOLS. We first evaluate the performance of the full FLEET-MERCH algorithm in FLEET- TOOLS, as plotted Figure (b)) We find that the performance of our method degrades gradually the frequency of iterative merging decreases (i.c., the number of epochs between merging increases), which is essential in flect-learning applications. We also have real-world demos for tool- using the merged policy from FLEET-MERGE in simulation, which are deferred to Appendix C.2. In Figure 4 (b), we show that FLEET-MERGE also has the best absolute performance and resilience to dataset heterogeneity in the Meta-World benchmark.
Meta- World. Figure 4 (c)&(d) study iterative merging in Meta- World. We find that FLEET- MERGE resilient lower frequency of weight merging (which allows
epochs of local training between merges) and to a smaller fraction of models being merged per round (which we refer to as the "participation ratio" of the agents in the fleet). These two together show that communication between agents in the fleet can be reduced without significantly harming the performance of the merged model. We observe similar results across different neural network architectures, different inputs and metrics, and in parge-scale settings. We defer the detailed results to Appendix D. Notably, we can learn multi-task learning policy to solve all 50 manipulation tasks jointly in Meta- Meta-World, without pooling the data together.
CONCLUSION
We studied policy merging, framework for fleet learning of control policies from distributed and potentially heterogeneous datasets, by aggregating the parameters of the (trained) policies. We developed new algorithms merge multiple policies by taking into account their parameterization ambiguity using recurrent neural networks. Finally, we proposed a novel robotic manipulation benchmark, FLEEE- in generalizable tool-use tasks, which might be of broader interest.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGEMENT
The authors would like to thank many helpful discussions from Eric Cousineau and Hongkai Dai at Toyota Research Institute well as Tao Pang, Rachel Holladay, and Chanwoo Park MIT. We thank MIT Supercloud for providing computing cluster resources for running the experiments. This work is supported part by Amazon Greater Boston Tech Initiative and Amazon PO No. 2D-06310236.
REFERENCES
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation. and distribution shift. The Journal of Machine Learning Research, 22(1):4431-4500 2021. 18
Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836 2022. 1, 4, 6, 7, 8, 18, 20, 28
OpenAl: Marcin Andrychowicz. Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-20, 2020. 2
Karl Åström. Introduction to stochastic control theory. Courier Corporation, 2012. 18
Karl Johan Åström and Richard M Murray. Feedback systems: an introduction for scientists and engineers. Princeton university press, 2021. 18, 29
J Andrew Bagnell. An invitation to imitation. Technical report, Camegic-Mellon Univ Pittsburgh Pa Robotics Inst, 2015. 3,3
Dimitri Bertsekas. Network optimization: continuous and discrete models, volume 8. Athena Scientific, 1998. 3.
Dimitri Bertsekas. Dynamic programming and optimal control: Volume volume 1. Athena scientific, 2012. 2
Adam Block, Ali Jadbabaie, Daniel Pfrommer, Max Simchowitz, and Russ Tedrake. Provable guarantees for generative behavior cloning: Bridging low-level stability and high-level behavior. In Thirty-sevnnth Conference Neural Information Processing Systems, 2023. 19
Johanni Brea, Berfin Simsek, Bernd Illing, and Wulfram Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv:1907.0299 2019. 18
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint urXiv:22/2.068/7 2022. 18
Jingjing Bu, Afshin Mesbahi, and Mehran Mesbahi. On topological and metrical properties of stabilizing feedback gains: the mimo case. arXiv preprint arXiv: (tt.:::944..2233 2019. 18
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.1277 12712. 2023. 18
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097, 2018. 29
Open X-Embodiment Collaboration. Open X-Embodiment: Robotic learning datasets and RT-X models. https: // wt t.u...u.. io, 2023.
10

---

Published as a conference paper at ICLR 2024
Ronan Collobert and Jason Weston. A unified architecture for natural language processing: Deep neural networks with multitask learning. In Proceedings of the 25th international conference on Machine learning, pp. 160-1677 167, 2008.
2008. 18
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013. 5, 20
Li Deng. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012. 7
Lior Deutsch, Erik Nijkamp, and Yu Yang. A generative model for sampling segh-perffrmance and diverse weights for neural networks. arXiv preprint arXiv:1905.00 02899, 2019. 18
Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint arXiv:2303.03375 2023. 18
Marvin Eisenberger, Aysim Toker, Laura Leal-Taixé, Florian Bernard, and Daniel Cremers, A unified framework for implicit sinkhorn differentiation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 509-518, 2022. 5
Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2):19-211, 1990. 2,
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. In International Conference Learning Representations, 2021. 1, 2, 8, 18, 20
Maryam Fazel, Rong Ge, Sham Kakade, and Mehran Mesbahi. Global convergence policy gradient methods for the linear quadratic regulator. In International conference machine learning, pp-
1467-1476. PMLR, 2018. 18
Han Feng and Javad Lavaci. Connectivity properties of the set of stabilizing static decentralized controllers. SIAM Journal on Control and Optimization, 58(5):2790-25)) 2820, 2020. 18
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agmostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1135-1135, PMLR, 2017. 18
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference Machine Learning. pp. 1259-32661 PMLR, 2020. 1, 3, 5, 8,
Wei Gao and Russ Tedrake. kpam 2.0: Feedback control for category-level robotic manipulation. IEEE Robotics and Automation Letters, 6(2):2962-2966 2021. 24
Philip E Gill, Walter Murray, and Michael A Saunders. SNOPT: An sqp algorithm for large-scale constrained optimization. SIAM review, 17(1)999-131, 2005. 23
Kristen Grauman, Andrew Westbury, Eugene Byrne, Zachary Chavis, Antonino Furnari, Rohit Girdhar, Jackson Hamburger, Hao Jiang, Miao Liu, Xingyu Liu, al. Ego4d: Around the world in 3,000 hours of egocentric video. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18995-19012, 2022. 1
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprini arXiv: 18/ 1.00604 2018. 18
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016. 25
Xiaoxi He, Zimu Zhou, and Lothar Thiele. Multi-task zipping via layer-wise neuron sharing. Advances in Neural Information Processing Systems, 31, 2018. 1, 18
Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pp. 135-135. Elsevier, 1990. 2, 18

---

Published as a conference paper at ICLR 2024
Alexander Herzog, Kanishka Rao, Karol Hausman, Yao Lu, Paul Wohlhart, Mengyuan Yan, Jessica Lin, Montserrat Gonzalez Arenas, Ted Xiao, Daniel Kappler, et al. Deep rl at scale: Sorting waste in office buildings with fleet of mobile manipulators. arXiv preprint arXiv:2305.032700 2023. 18
Rachel Holladay, Tomás Lozano- Pérez, and Alberto Rodriguez. Force-and-motion constrained planning for tool use. In 2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 7409-7416. IERE, 2019. 6, 24
J Hong, N Moehle, and S Boyd. Lecture notes in "introduction matrix methods", 2021. 33
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of decentralized machine learning. In International Conference on Machine Learning, pp. 4387-4398. PMLR, 2020. 8
Bin Hu, Kaiqing Zhang, Na Li, Mehran Meshahi, Maryam Fazel, and Tamer Başar. Toward theoretical foundation of policy optimization for learning control policies. Annual Review of Control, Robotics, and Autonomous Systems, 6: 6:123-158, 2023. 18
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Conference on Robot Learning, pp. 991-1002. PMLR, 2022. 1
Roy Jonker and Ton Volgenant. A shortest augmenting path algorithm for dense and sparse linear assignment problems. In DGOR/SOO:: Papers of the 16th Annual Meeting of DGOR in Coopera- tion with NSOR/Vorträge der 16. Jahrestagung der DGOR zusammen mit der NSOR, pp. 622-622. Springer, 1988.
Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski. Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic rein- forcement learning at scale. arXiv preprint arXiv:2104.08212. 2021. 18
Oussama Khatib. A unified approach for motion and force control of robot manipulators: The operational space formulation. IEEE Journal on Robotics and Automation, 3(1):43-53, 1987. 24
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14X2::4412.69880 2014. 24
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, al. Wilds: benchmark of in-the-wild distribution shifts. In International Conference Machine Learning, pp. 5637-5664. PMLR, 2021. 18
Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint 2, 18
2016.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009. 29
Harold W Kuhn. The hungarian method for the assignment problem. Naval Research Logistics quarterly, 2(1-2):83-97, 1955. 4
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179-1191, 2020.
Aviral Kumar, Anikait Singh, Frederik Ebert, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre- training for robots: Offline RL enables learning new tasks from a handful trials. arXiv preprint 10.i:2210.051778 2022.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17())13344-333, 2016. 2
12

---

Published as a conference paper at ICLR 2024
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint 0rXiv:2055.01643 2020.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint 1xXXx:2002.106199 2020.
Lucas Manuelli, Wei Gao, Peter Florence, and Russ Tedrake. kpam: Keypoint affordances for category-level robotic manipulation. In Robotics Research: The 19th International Symposium ISRR, pp. 157. Springer, 2022. 6, 7, 22, 23
Joseph Masterjohn, Damrong Guoy, John Shepherd, and Alejandro Castro. Velocity level approxima- tion of pressure field contact patches. IEEE Robotics and Automation Letters, 7(4): 11600, 2022. 24
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effin efficiiet learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017. 2, 8, 9,
H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private recurrent language models. In International Conference on Learning Representations, 2018. 18
Gonzalo Mena, David Belanger. Scott Linderman, and Jasper Snock. Learning latent permutations with gumbel-sinkhorn networks. arXiv preprint arXiv: 1802.06655, 2018. 5
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equiv- ariant architectures for learning in deep weight spaces. arXiv preprint 0rXiv:2301.12780, 2023.
18
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614 2014. 20
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint arXiv: Tt03:8803.0990 2018. 18
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 7(1-2): 1-179, 2018. 3,
Fidel A Guerrero Peña, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti. Eric Granger, and Marco Pedersoli. Re-basin via implicit sinkhorn differentiation. arXiv preprint 1rXiv:2212. 12042, 2022. 2, 4, 5, 6, 8, 18, 28
Aubrey B Poore. Multidimensional assignment formulation of data association problems arising from multitarget and multisensor tracking. Computational Optimization and Applications, 3(1): 27-57, 1994. 18
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 18
Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint
(rXi::/70.005090 2017. 18
Daniel Seita, Yufei Wang, Sarthak J Shetty, Edward Yao Li, Zackory Erickson, and David Held. Toolflownet: Robotic manipulation with tools via predicting tool flow from point clouds. In Conference Robot Learning, pp. 1038-1049. PMLR, 2023. 24
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning, pp. 785-799. PMLR, 2023. 18
Max Simchowitz, Abhishek Gupta, and Kaiqing Zhang. Tackling combinatorial distribution shift: A matrix completion perspective. In The Thirty Sixth Annual Conference Learning Theory, pp- 3356-3468. PMLR, 2023. 24
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045-22055. 2020. 18
13

---

Published as a conference paper at ICLR 2024
Samuel L Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. arXiv preprint arXiv:2101.1217,, 2021. 20
Trevor Standley, Amir Zamir, Dawn Chen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? In International Conference Machine Learning, pp. 9120-9132. PMLR, 2020. 18
George Stoica, Daniel Bolya, Jakob Bjorner, Taylor Hearn, and Judy Hoffman. Ziplt! Merging models from different tasks without training. arXiv preprint arXiv:2305.0305. 2023. 1, 6, 18
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999. 28
Julia Tan, Ransalu Senanayake, and Fabio Ramos. Renaissance robot: Optimal transport policy fusion for learning diverse skills. In 2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 7052-7059. IEEE, 2022. 18
Norman Tatro, Pin- Yu Chen, Payel Das, Igor Melnyk, Prasanna Sattigeri, and Rongjie Lai. Optimizing mode connectivity via neuron alignment. Advances in Neural Information Processing Systems, 33: 15300-15311, 2020. 18
Russ Tedrake and the Drake Development Team. Drake: Model-based design and verification for robotics, 2019. URL https: / drake meiiteddu 6, 22
Mare Toussaint. Robot trajectory optimization using approximate inference. In Proceedings of the 26th annual international conference on machine learning, pp. 1049-1056, 2009. 24
Mare A Toussaint, Kelsey Rebecca Allen, Kevin Smith, and Joshua B Tenenbaum. Diffcrentiable physics and stable modes for tool-use and manipulation planning. 2018. 6
Jack Umenberger, Max Simchowitz, Juan Perdomo, Kaiqing Zhang, and Russ Tedrake. Globally convergent policy search for output estimation. Advances in Neural Information Processing Systems, 35:22778-22777 2022. 18, 31
Ricardo Vilalta and Youssef Drissi. A perspective view and survey of meta-learning. Artificial intelligence review, 18:77-95, 2002. 18
Han Wang, Leonardo F Toso, and James Anderson. Fedsysid: federated approach to sample- efficient system identification. arXiv preprint arXiv:2211.14393 2022a. 18
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papaitiopoulos, and Yasaman Khazaeni. Fed- crated learning with matched averaging. In International Conference on Learning Representations, 2020a. 2, 18
Lirui Wang, Yu Xiang, Wei Yang, Arsalan Mousavian, and Dieter Fox. Goal-auxiliary actor-critic for 6d robotic grasping with point clouds. In Conference on Robot Learning, pp. 70-80. PMLR, 2022b. 24
Lirui Wang, Kaiqing Zhang, Yunzhu Li, Yonglong Tian, and Russ Tedrake. Does learning from decentralized non-iid unlabeled data benefit from self supervision? In The Eleventh International Conference on Learning Representations, 2023. 18
Yaqing Wang, Quanming Yao, James T Kwok, and Lionel M Ni. Generalizing from a few examples: A survey few-shot learning. ACM computing surveys (csur), 53(3):1-34, 2020b. 18
Mitchell Wortsman, Suchin Gururangan, Shen Li, Ali Farhadi, Ludwig Schmidt, Michael Rab- bat, and Ari S Morcos. lo-fi: distributed fine-tuning without communication. arXiv preprint arXiv:2210. 1948, 2022a. 1, 18
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lpess, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups:
averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pp. 23965-23998 PMLR, 2022b. 1, 18
14

---

Published as a conference paper at ICLR 2024
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.
In Conference robot learning. pp. 1094-1100. PMLR, 2020. 2, 7, 18, 25
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International
conference on machine learning. pp. 7252-7261. PMLR, 2019. 18
Thomas T Zhang, Katie Kang, Bruce D Lee, Claire Tomlin, Sergey Levine, Stephen Tu, and Nikolai Matni. Multi-task imitation learning for linear dynamical systems. arXiv preprint arXiv:2212.001886 2022. 18, 33
Yu Zhang and Qiang Yang. An overview of multi-task learning. National Science Review, 5(1): 30-43, 2018. 18
Yang Zheng, Yujie Tang, and Na Li. Analysis of the optimization landscape of linear quadratic gaussian (lgg) control. arXiv preprint arXiv:2102.04393, 2021. 18, 34
Allan Zhou, Kaien Yang, Kaylee Burns, Yiding Jiang, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. arXiv preprint arXiv:2302. 14040, 2023. 18
Kemin Zhou, JC Doyle, and Keither Glover. Robust and optimal control. Control Engineering Practice, 4(8):1189- 190, 1996. 34
15

---

Published as a conference paper at ICLR 2024
CONTENTS
1
Introduction
2 Preliminaries
3 Methodology
3.1 Merging Many Recurrent Policies
FLEET-TOOLS: A Robotic Tool-Use Simulation Benchmark
5 Experiments: Policy Merging
5.1 One-shot Merging
8
5.2 Iterative Merging
6
Conclusion
7
Acknowledgement
10
Detailed Related Work
18
B Permutation Invariance for RNNs
19
B.1 Proof of Proposition 3.1
19
B.2 Invariance for ReLU Activations
20
C FLEET-TOOLS Benchmark
22
C.I Implementation Details
22
C.2 Environment Setup and Implementation Details
24
Meta-World Experiments
25
D.I Further Experimental Details
25
Online Reinforcement Learning Ablation
28
F Algorithm Ablation Supervised Learning Tasks
28
G Linear Policy Merging
29
G.1 Problem Setup
29
G.2 Symmetries in Nonlinear Control
31
G.3 Details for Training Imitation Learning...
................................
G.4 Methods for Linear Policy Merging
32
G.4.1 Method 1: Permutation-Base Alignment
32
G.4.2 Method 2: Unconstrained Gradient Descent
33
G.5 Experiment Setup Details
33
16

---

Published as a conference paper at ICLR 2024
G.6
Results and Discussions
33
G.6.1
Basic Experimental Results
33
G.6.2
Gradient-based approach v.s. LQG Landscape
34
G.6.3
Insufficiency of Naive and Permutation Based Merging
.tw..n..........................
34
G.6.4
Avoiding
of Stabilizing Controllers
35
17

---

Published as a conference paper at ICLR 2024
A DETAILED RELATED WORK
Parameter invariance in neural networks and control. It has long been known that neural networks are invariant under certain transformations symmetries, of their weights (Hecht-Nielsen, 1990). In particular, permutation symmetries arise because the ordering of hidden neurons in neural networks does not affect their input-output behavior. This has been an important topic of recent studies of neural networks loss landscapes (Brea et al., 2019; Entezari et al., 2021; Tatro et al., 2020; Ainsworth et al., 2022) and of methods that process the weights of other networks (Deutsch et al., 2019; Navon et al., 2023; Zhou et al., 2023). On the other hand, in the controls literature, it is known that the input-output behavior of linear dynamical system is invariant under any similarity transformations the system matrices (Aström & Murray, 2021; Åström, 2012). Recent works have further studied the optimization landscape of classical linear controllers (Fazel et al., 2018; Zheng al., 2021: Umenberger et al., 2022; Hu et al., 2023), especially for dynamic controllers that have recurrent latent states (Zheng et al., 2021; Umenberger et al., 2022). In our work, we consider the symmetry of general nonlinear policies parametrized by neural networks, particularly recurrent neural networks, and how to leverage this property to merge multiple policies in robotic fleet learning.
Model merging & Mode connectivity. There have been variety of approaches to merging neural networks by interpolating their weights. For deep neural networks, this typically requires aligning hidden neurons using techniques similar to those in sensor fusion (Poore, 1994). One line of work uses optimal transport to align the weights of distinct models prior to merging (Singh & Jaggi, 2020; Tan et al., 2022), while related works have investigated removing permutation symmetries to align neurons before interpolation (Tatro et al., 2020; Entezari et al., 2021; Ainsworth et al., 2022; Peña et al., 2022). The connectivity property of the landscape when optimizing linear controllers has also been a resurgent research interest (Feng & Lavaei, 2020; Zheng et al., 2021; Bu et al., 2019). Notably, Bu et al. (2019) shows that the parameters of stabilizing static state fete-ffe linear controllers form single connected component, and Zheng et al. (2021) shows that heeputtfeedback dynamical controllers can form at most two connected components. Merging neural networks has also been studied in recent years in the context of finetuning foundation models (Wortsman al., 2022a;b), federated learning (McMahan et al., 2017; Wang et al., 2020a; Yurochkin et al., 2019; Konečný et al., 2016; Wang et al., 2023), multi-task learning (He al., 2018: Stoica et al., 2023), and studying linear mode connectivity hypothesis for (feedforward) neural networks (Frankle et al., 2020; Entezari et al., 2021; Ainsworth et al., 2022). Most of these works focused on the feedforward NN architecture. A few works have studied model merging for recurrent NNs in the context of federated learning (Hard et al., 2018; McMahan et al., 2018; Wang et al., 2020a). However, they either used direct averuging without accounting for the permutation symmetries (Hard et al., 2018; McMahan et al., 201 8), or focused on supervised learning tasks (Wang et al., 2020a). In our work, we aim to study how well these insights can be applied to policy learning and control settings, with recurrent neural network parameterization, and awareness of permutation invariance. Compared with the most related work Wang et al. (2020a), we also develop a new merging approach based on "soft" permutations (see Section 3 for a formal introduction) of the neurons, allowing a more efficient implementation and focus on robotic control tasks.
Multi-task policy learning. Multi-task learning and models that exhibit multi-task behaviors have shown impressive successes in computer vision (Zhang & Yang, 2018; Standley et al., 2020) and natural language processing (Radford et al., 2019; Collobert & Weston, 2008; Bubeck et al., 2023). Indeed, multi-task learning (Ruder, 2017: Yu et al., 2020), meta-learning (Vilalta & Drissi, 2002: Nichol et al., 2018; Finn et al., 2017), and few-shot learning (Wang et al., 2020b) have long lines of literature. Despite some promising recent attempts robotics (Brohan al., 2022: Shridhar et al., 2023). the dominant paradigm in robotics is still to conduct single-task data collection and training on a single domain. Recently, such multi-task paradigms have also been extended to linear control settings, see e.g., Zhang et al. (2022); Wang et al. (2022a). Our work is also related to the studies on the effects of distribution shifts and diverse data distributions (Agarwal et al., 2021; Koh et al., 2021). Inspired by the increasing scale of robot fleets deployed in the real-world, building policy learning framework at scale requires us to handle distribution heterogeneity and communication efficiency (Kalashnikov et al., 2021; Herzog al., 2023; Driess et al., 2023). Also note that the setting consider, with multiple datasets, include but are not limited to multi-tasks settings, as the datasets may come from different experts in imitation learning, or different time operating conditions of the
18

---

Published as a conference paper at ICLR 2024
same agent/policy in the same task. While provable guarantees exist for imitation in the single-task setting Block et al. (2023), understanding the implications of behavior cloning in multi-task settings from theoretical angle is exciting direction for future work.
B PERMUTATION INVARIANCE FOR RNNS
B. PROOF OF PROPOSITION 3.1
In this section, we prove Proposition 3.1 in the main paper. (P", Pl p4) Sperm. it acts on RNN models via
Recall that when P
(W ree W b (P'W( rec (P'), P'W, f (Pi-1),, P'b (B.1)
where we note that we replace the matrix inverse in Eq. (B.1) with matrix transpose because the two coincide for permutation matrices. We now note that the group operation of on weights can be decomposed as the computation of applying operators as follows
,I) ) (0) (P",
1) (L, I)..............................
(I, P4)
(B.2) where for the second equality, we recall that Sperm is defined so that its elements P =
PI. - P1,
1)
..... II, P
.p......P! E Sperm. have p and p4 equal to the identity (i.e., we do not permute in- puts or outputs). It suffices to show that the RNN model is invariant to applying this operation specific layer, i.e.
P= P=
I,Pe I.....II.
(B.3)
We now prove the following claim.
Claim B.1. Let l be such that 0 < <<< L, and let P be of the form Eq. (B.3), i.e., P is the only non-identity. Let hy be the hidden state at time and layer 0 l' < L of model 6 starting at hy 0.3 Let hú denote the same, but for the permuted model P(0). Then, for all
E' +l El f=t
My'
(B.4)
P'M
Note that this claim implies Proposition 3.1 because (B.2) shows that any P can be decomposed into permutations of the form (B.3) for 0 < f < L. particular, because & + L, (B.4) means that h hf as needed.
Proof of Claim B.I. It is clear that he" = hf for layers l' < l and t 2 0. It remains to consider layers l' l and E' f.
Let's start with é' L. We argue by induction on that if P'h Let's start with t 0. By assumption we have that ñι = O P. This proves the base case, For general t 0, by specializing (B.1) to permutation of the form (B.3), we have
h ww'w' roc P() Thp p'
P'b
P (W (P) P W 1/6 +b--l)
because any entry-wise activation (-) commutes with P As noted above, we have 7 116-1 Moreover, by the inductive hypothesis, hf_ = ph Thus, we have
h (Wo rea (P') T p + w 7 + b b
= =P'o(W' rec 1 + W6-1i be-1)
=P'M
This proves what we need for =.
This condition can be generalized to requiring that Tg'' P' for 0 < E < T. and hg" =n6" for
8" E (0, T.).
19

---

Published as a conference paper at ICLR 2024
For E >> f", it suffices to prove the case when e = L 1, because the weights for all subsequent layers are not changed by (B.3) and all layers above + depend only on the + I's input. Again, we induct over f. By assumption h +++1 = 1,t1. Moreover, permutations of the form (B.3) mean that
7,+1 = a (W 1 7 + W ( +b)).
rec
Using the inductive hypothesis that h not1 and the fact that ñff P proven above, we have
++t = 0 (W rec /1-1 <+1 + (P') TP')) bb)
= oW ne 1-1 + W + b¹) =ht+1,
concluding the proof of the claim.
B.2 INVARIANCE FOR RELU ACTIVATIONS
For RNNs with common activation functions, can study more general group of symmetries beyond just the permutations of Proposition 3.1. Again, consider the action of a general p = (I,P), p2 ...., ...,,P.. p II Gin. acting as
(W! reas W bb-1) (P'W rec (P')-1 P'w (P--1)-1 P'b²-1).
(B.5)
Proposition B.1. (-) )) be any activation function, and consider P E Gia acting as above. Suppose that, for each f, P commutes with d, in the sense that P-)(-) = (P* ())) Then, the output of the RNN is invariant to the transformation induced by P. In particular, ( = ReLU(.), then the RNN output invariant the set P E Iscated, where Gscated C 9ii is the set of transformations such that P D are scaled permutations, with P being permutation matrix and D being diagonal matrix with strictly positive elements."
The proof of Proposition B. follows exactly as in the proof of Proposition 3.1 above, by replacing transposes with inverses for the general case. The special case of the ReLU network follows from checking that the ReLU operation commutes with scaled permutations.
Even though RNNs are invariant scaled permutations their weights, we follow previous model merging work and focus only on the permutations (i.e., ignoring the scaling) (Ainsworth et al., 2022; Cuturi, 2013). As argued by the linear mode connectivity hypothesis (Entezari al., 2021), scaling symmetries should not appear in practice because the implicit regularization of the training algorithms, e.g., stochastic gradient descent (Neyshabur et al., 2014; Smith et al., 2021), controls the scale of the resulting weights. Here, we formalize the argument by proving that implicit regularization uniquely determines the scaling matrix.
Proposition B.2. For an RNN 0, define the norm
| 2 W I - + b 11
Σ
W rec l /|
112 + 11b
(B.6)
OCEEL
Let any model for which every bias term be has strictly nonzero entries. Define the set
(0) argmin
110'1| s.t. 0' P(0),
(B.7)
0r, P.PEeccied
of minimal norm models which correspond to transforming 0 by scaled permutation. Then, for every 8', 8" E (0), there exists a transformation consisting only of (non-scaled) permutations P 9perm such that 0' P(O").
Proof. Note that if 0' E (0)) x(0)) = x((0) because Gseated a group. Thus, it suffices to prove a slightly simpler claim:
E x0) and 0' X(0). and let P be such that p (I, D'D'.. ...., 1 I), where P-matrices are permutation matrices and D matrices are diagonal matrices with positive diagonal elements. Then, we must have
Claim B.2. Suppose
D' = Iap 0<E<L
(B.8)
where de is the dimension of the D
NNoe that this is group because D'D' p'D' for other D
20

---

Published as a conference paper at ICLR 2024
We now prove the claim. We expand
- A
||"||2
+ Σ
11 + |(P''))>>>
D<E<L
rec
D<E<L
= - H
11
Σ DCEEL
D' rec (D')
D -ow2-1(10 (D 1D |
dr di-1
=11b4-112
W [i,) -- -
de
D Li,
W rea
"
DctcLin i,j=1
de
de Σ
ΣΣΣ ictc
D D'
D'
OcteLi=]
where j] indexes matrices and [i] indexes vectors, and where we use that D' are diagonal matrices. Next, let represent the diagonals of the matrices D i] exp(Te,s) exponentials, which is valid because they are strictly positive. this representation, we have
du di-1
[P(O) | 1
w L---1
[ 11 + Σ "
w rec ( (
LeLLi.j
di
ΣΣΣ KEL
W
22 +
Σ O<ECLE
2 b'
=F((rahhsssdoooeee))
We now observe that the function F()) is convex in its arguments, because it is the sum over with positive coefficients. fact, it is strictly exponentials of linear functions of (Teshsisd convex, because all but the last terms are convex, and the term Eo<ket Ede ( is strictly convex under our assumption that none of the bias weights zero.
Therefore, there a unique (riphsissa.cete which minimizes F(.). Thus, by assumption that 0' E (6), we must have that
[i i] = exp(Tt,), 1sis de, << <<<
But similarly, since 0 E (0), then by writing and = (I,I,...I] applied to itself, we must have that
1) a 0 as the identity transformation
exp(T(,) = (L 1.
Thus, we conclude that
D'ii 11
as needed.
21

---

Published as a conference paper at ICLR 2024
C FLEET-TOOLS BENCHMARK
C.I IMPLEMENTATION DETAILS
In this section, we introduce more details of our newly developed tool-use benchmark for fleet policy leamingg based on the Drake simulator (Tedrake & the Drake Development Team, 2019).
Expert tasks. To scale the generation of expert demonstrations for the tool-use tasks, we specify the tasks via key points (Manuelli et al., 2022) when generating the expert trajectories. We mainly consider four skills as examples (or task families): (wrench, hammer, spatula, and knife) fe) -consisting of tasks involving the eponymous tool-type. Specifically, we use use a spanner to apply wrenches, use a hammer to hit, use a spatula to scoop, and use knife to split. All of these tasks have common feature that the skills can be reasonably, certainly not perfectly, solved by specifying a few sparse keypoints. For instance, in the wrench tasks as shown in Figure 5, the wrench needs to reach the nut first and then rotate by 45 degrees while maintaining contact. The same strategy can be applied to set of nuts and wrenches. We pick these tasks because they provide balance of common robotic tasks that require precision, dexterity, and generality in object-object affordance.
Keypoint transformation optimization. We denote the end effector frame and world frame as E and W. respectively. For cach tool, we specify n keypoints, for example, when n = 3, we can specify one point on the tool head, one point the tool tail, and one point on the side. We define rigid body transformation X E SE(3) C R4x4, the current locations of the n keypoints on the tool in the world frame W as "yoool= = W " p0o0111 .....
and the poses in the end effector frame as Eyro01 = [Ep [[ppto0111 .... pprool,]] respectively. For the object, we specify m keypoints in the world 1 poos = [ pap..1
"y"oe..") ,... W
frame as
W ppoo..m Each point (either p = pppto or W woob1.) lies in 1 and is expressed in (at, y. 2, 1) >-homogeneous coordinates. Then, one can solve the following optimization problem, referred to as kPAM-Opt (where kPAM stands for KeyPoint Affordance- Manipulation) (Manuelli et al., 2022), to find the transformation X:
kPAM-Opt
min X
x p tool - W. tool
pr ||
(C.1)
s.t.
|| E pooo w poDsss |1 Se, Vi,j
(C.2)
0s.j - S BT, X (B p too1,1 too1..))) S aij te, Vi,j
(C.3)
where we choose E {1, {1...... n} and j E (1, (1..... ,m), and Bi, E R is a vector on the unit sphere with the fourth dimension (4) 0, oz, some constant that represents the target angle alignment, and > 0 is some relaxation level. Constraint Eq. (C.5) restricts the distance between the keypoints on the tool to be close some target ones on the object be e-small. For example, when 1 and m 1, where we use 1pt0ol,1 to denote the keypoint corresponds to the head of the tool, then this constraint ensures the head of the tool, e.g., the hammer head, to be close to the object, e.g., the pin. The composition of (at.j 3 c) in Eq. (C.6) describes some constraints related to the orientations of the tool (e.g., the hammer head needs to be vertical to the world and stay flat), up to some ereelaxation. The overall formulation in Eq. (C.1)-Eq. (C.3) instantiates the general one in Manuelli et al. (2022).
Although the kPAM procedure above can solve for the desired transformation X, it can be computa- tionally challenging to recover the joint angle from X via inverse kinematics. Hence, we propose to solve kPAM optimization in the joint space, i.c., optimize over the joint angles q that parameterize the transformation through forward kinematics (i.e., Forward-kinema t i cs (4)), and only add constraints to the keypoints in the transformed frame. Specifically, we solve the following joint-space optimization problem:
E t00l w, p too2 }
Siat-Space-kPAMM-O0
min 9. X
(C.4)
s.t.
x p coo0s W pobsuul1 Se, Vi,j
(C.5)
0i
S S x E ppools Eyooold S aiy tc, Vi,j
(C.6)
X forward-kknematicsll
(C.7)
22

---

Published as a conference paper at ICLR 2024
Specify keypoints
Pre-actuation Trijectory
Actuation Trajectory
MIMM
Figure 5: Visualization of the kPAM-based (Manuelli et al., 2022) expert generation pipeline. There are three keypoints defined on the tool and one keypoint defined on the object. The kPAM objective and constraints define an actuation pose for each specific tool-use task. We then optimize for the pre-actuation trajectory and the post-actuation trajectory jointly.
where constraint Eq. (4.4) relates the rigid body transformation X and the output of forward www--kiiiiiii(((e under a joint angle to better contrast the formulation in Eq. (C.1)- Eq. (C.3).
We create the optimization problem Eq. EC. (C.6) in Drake, which is then solved using nonlinear programming solvers as SNOPT (Gill et al., 2005). The intuition this joint-space optimization is to leverage the multiple solutions in the manipulation tasks to simplify the constraint satisfaction problem in inverse kinematics. For instance, when solving the wrench tasks, the entire rotation space around the z axis in the world frame is free, solving for joint angle directly these constraints defined on keypoints can be easier than solving the pose-space KPAM problem and then solve inverse kinematics.
Pre-/Post-actuation optimization. After finding the joint angle solution q", we can work out the pre-actua actuation and the post-actuation poses (Figure 5). For instance, for wrench tasks, there is standoff pose that is few centimeters in the negative axis of the end effector, which are treated as keyframes. Once we construct the trajectory of pose keyframes the pre-actuation and post-actuation trajectories, we can solve the joint-space trajectory as an optimization problem that further improves the smoothness (by minimizing the norm of joint difference with fixed endpoints), while satisfying the keyframe pose constraints.
Specifically, 5 = ((2)..........111 denote joint-space trajectory length T. including trajectory before actuation at some timestep 1* < T, and the post-actuation trajectory from 1* + until T. Choose m < T indices { Im C 2. ,T}, and specify the corresponding pose keyframe trajectory [PI,,PI,:- PI_. Let dı be the current joint position, and be the solution to the previous joint-space KPAM problem Eq. (C.4)-Eq. (C.7). We then solve the following optimization problem:
min
Σ 1=1
Left 1] Elt]l1
(C.8)
s.t.
{[1 42
E[t"] q*
(C.9)
forward-kinemttessee = Pie, A 1......,
(C.10)
which is also solved by nonlinear solvers SNOPT (Gill et al., 2005) in Drake. This final solution gives planned trajectory &* and we can then track this trajectory with low-level controllers to provide demonstrations.
Discussions. Our tool-use benchmark and the pipeline above provide natural and easy way to create (new) tasks in Drake, with only configuration file of the actuation pose optimization and few keypoints labels that can easily obtained via a 3D visualizer. The overall pipeline, which is deterministic by design, is also general enough to act the expert in the imitation learning settings to represent "category- composition and generalization. The convenience category- generalization be important application domain study out-of-support distribution shift,
23

---

Published as a conference paper at ICLR 2024
Figure 6: The real-world setup that the benchmark recreates in the simulation with segmented point cloud visualizations.
which can be formulated as combinatorial extrapolationggeeenriiziioooo problems, i.e., generalizing to the unseen tool-object (or tool-task) combinations at test time, despite only being able to see specific tools and objects separately (while not such joint combinations) training time. See more formal treatment in Simchowitz et al. (2023).
There certain limitations of our pipeline. For example, it creates only hand-scripted experts, and there no feedback in the execution loop, which can be necessary for some dynamic and contact-rich motions (Gao & Tedrake, 2021). Moreover, tool-object and tool-hand contact dynamics modeled in the simulation but are not considered in the current expert formulation. These are all important directions we will address in the future. We believe tool-use tasks, which often involve rich force and contact interactions (?Holladay et al., 2019; Toussaint, 2009), represent one of the key challenging tasks in dexterous manipulation.
C.2 ENVIRONMENT SETUP AND IMPLEMENTATION DETAILS
We set up cameras to mimic the real-world setting of Panda robotic arms (see Figure 6). One overhead camera is mounted on the table and the other camera is mounted on the wrist. We have tool-object pairs for each tool in the benchmark (more assets available). At each scene initialization, we randomize the camera extrinsics, robot initial joints, object frictions, mass, and the tool-in-hand pose. We also randomize object texture when using colors and support offline blender rendering. We weld the tool to the robot end effector which forms kinematic chain. We assume RGB-D images with masks as inputs the robots. Specifically, we can use this information fuse pointelouds (Seita et al., 2023; Wang et al., 2022b) for both the tools and the objects, which we resample to contain 1024 points. We also measure joint torques and end effector wrench in the environment. We use the bounded relative end effector motion as the action output. Notably, in the simulation we have a perfect dynamics model that can potentially improve the contact-rich motions in tool-use, which enables to use customized operation operation-space controller (OSC) (Khatib. 1987) with tuned gains to achieve high success rates in solving the tasks. We use hydroelastic contact model with SAP contact solver (Masterjohn al., 2022) in drake. The environment timesteps 6Hz and the simulation dynamics step is 250Hz. We use the Ray library to support multiple rollouts in parallel.
We combine the segmented point cloud from wrist camera and an overhead camera into the tool and object pointcloud, and feed them as observations to the policy. We train the policy using Adam optimizer (Kingma& & Ba, 2014) with 200 epochs. The dataset contains 50000 data points for each tool instance, and the batch size is 512. We use 5 users across the merging experiments. Each test trial has 30 different scenes and we take average over 10 trials.
24

---

Published as a conference paper at ICLR 2024
Figure We deployed the merged policies the real world for the hammering task.
D HETA-WORLD EXPERIMENTS
D.1 FURTHER EXPERIMENTAL DETAILS
In this section, we provide the experiment details on the Meta-World benchmark (Yu al., 2020).
Implementation. For state-based input, we use a 4-layer network architecture with 512 hidden dimensions, including multilayer perceptron (MLP), RNN, and Transformer, to parameterize the policies. For image-based input, we use ResNet18 (He et al., 2016) as the default vision encoder and similar 4-layer MI.P for the policy head. We use Adam optimizer with a learning rate of and batch size of 256. We use 3, 5, and 10 users respectively for the single- merging, merging while training, and merging with participation ratio experiments in Figure 10. For validating mode connectivity,
the same dataset for training two models, but we find that training on different datasets (generated by the same expert) will work as well.
Evaluation across tasks. We evaluate and compare the performance of methods large-scale experiments, notably, the MT-50 tasks in Meta- World. In Figure 10, we provide the comparison of our policy merging methods with baselines: joint multi-task training and task-specific training. In particular, in joint training, we pool the data from multiple tasks to train task-conditioned policy; in task-specific training, we train policy for each task separately. As shown Figure 10, in terms of success rates, joint training, and task- speccspec training can achieve similar performance. We also shown
observed that the more data is used, the gap between the two smaller. More importantly, that our policy merging methods can perform on with these two baselines, by outputting only one single policy (instead of in task -pee-iie training), while without sharing the training data (instcad of in joint multi-task training).
Evaluation across network architectures. We also evaluate our policy-merging methods and corresponding baselines across different neural network architectures. In Figure 9, observe similar performance connectivity behaviors across network architectures, including MLP, RNN, and Transformers. This indicates that our policy merging formulation and algorithms can be general, and not specifically tied to the neural network architectures.
Norm variations. Since RELU neural network's weights are also symmetric up to scaling in between layers, it's common question to ask whether need to explicitly account for that. In Figure 8, we show that the network's 's weights RNN policies and neural net models at each layer are relatively stable, across different runs training the same dataset. The y-axis shows the relative ratio the weight different norm compared to the weight norm, which less than 1% in general.
25

---

Published as a conference paper at ICLR 2024
Shakespeare Dutaset
Meta: World Button Press
Meta World Doar Open
Wegl Dutiiee
Meat to sud
1.0
Figure 8: The weight norm ratio (weight difference between runs, compared to weight norms of the first run) of trained RNNs multiple benchmarks.
MI.P
button prest topdowe x2
00o 03E0
SIAAY open
reach V2
rake
RNN
topdown-x2
Da mallley
1.0
03
05
10
Transformer
3.0
2.0
0.5
1.0
2.0
0.5
1.
Figure 9: Performance connectivity across architectures. Git rebasin and train matching work consistently across many tasks for every architecture.
26

---

Published as a conference paper at ICLR 2024
Task-Specific
0.6-
0.2
Multi-Task
env name
0.8
total Success
0.4
Skill Merging
total totallSucoo Success
0.G
0.0
envy_name
Figure 10: Comparison of joint multi-task training, task-specific training, and policy merging on 50 tasks in Meta-World (MT-50).
27

---

Published as a conference paper at ICLR 2024
Epoch=10
Epoch=100
Epoch=1000
312
18
3000
13000 2006 40:00
3008
0908
3000
Figure 11: We show that fleet learning can also apply to online reinforcement learning settings where agents continue to gather new data for training. The performance robust to the communication epochs ranging from 10 to 1000 trajectories between every averaging. The merged model also outperforms the average of each individual model in this case.
E ONLINE REINFORCEMENT LEARNING ABLATION
To apply fleet learning to online reinforcement learning, we use the REINFORCE (Sutton et al., 1999) algorithm in the CartPole environment in OpenAl Gym proof of concept. In this experiment, we train 5 different policies on 5 instances of the CartPole environment with different random seeds, and average the policies after certain number of episodes. The experiment (Figure 11) shows that simple averaging can be used when each agent collects more data online and uploads model weights to compose a policy. Note that no data is shared during the learning process.
F ALGORITHM ABLATION ON SUPERVISED LEARNING TASKS
In this section, we ablate on different components of our proposed methods of merging multiple models in the one-shot setting, i.e., we only merge the models once, on benchmark supervised learning tasks with different neural network architectures. Specifically, into local datasets, with Dirichlet parameters a to create data non-IIDness (see ?? for more details), and we L-layer MLPs to parameterize the models. We lay out the methods ablate as follows:
split the MNIST dataset
1. average: This method naively averages all the weights of the trained model on each local dataset and then tests on the MNIST test set.
2. git-rebasin: This method uses the MergeMany algorithm in Ainsworth et al. (2022) to merge all the locally trained models. Specifically, the MangeMany algorithm is an alternation- -based procedure: at each round, it first randomly samples model, and then aligns it with the average of the rest of the models, by using coordinate descent on the aligning loss; the procedure continues until convergence.
3. sinkhorn-data: This method similarly applies the alternation- based algorithm above, and replaces the alignment step from Ainsworth et al. (2022) to sinkhorn-rebasin in Peña et al. (2022). Specifically, we use the validation dataset for the alignment. The difference between this alternation- method and our FT.EE--EERS algorithm (Algorithm 1) is whether we reinitialize the permutation with the hard permutation matrix each round or
4. FLEET-MERGE (ours): This is our main algorithm FLEET-MERGE (Algorithm 1) where we use a FedAvg style algorithm to update the permutation parameters in sinkhorn rebasin.
5. fleet-merge-soft: This an ablation method on our algorithm where we replace the hard permutations at line 6 in Algorithm by soft permutation matrices that are computationally more tractable.
6. siee--merg--iimpee:: This is ablation method algorithm where instead of trying to interpolate between each local model and the averaged model, we randomly sample in the convex hull of N local models and compute the gradients.
28

---

Published as a conference paper at ICLR 2024
(A) MLP
(b) RNN
(c) CNN
0.40
Metroc
0.35
Methad
Accurao
0.8
Accuracy
0.33
0.6
0.3
0.20
0.4
0.2
0.15
0.2
0.05
0.0
0.01
0.1
alpha
0.5
1.0
0.0
0.01
0.1
alpha
0.5
1.0
0.00
0.01
0.1
alpha
0.5
1.0
(d)
(c) Ablation
(1) Model Number (RNN)
Accuraly
0.8
Accuraty
0.8
Mest-mangeet mag.rgg Bd
Teet-mengeeet
0.4
Fext merge
Vethod
0.6
Meat auaragge
0.3
0.4
0.2
0.2
0.0
5
7
10
20
30
0.0
0.01
1.0
0.0
Figure 12: Algorithm Ablation Study on MNIST. Each model owns a shard the total dataset and our Mani-Rebasin algorithm outperforms all baseline methods.
Single Task Dynamic Feedbackk With Disconnected Companents
Multi-Task Static Foodback
Single-Task Dynamic Feedback
Multi-Task Dynamic Foodback
tock000
10030
3030
200-
200
0002
200
4000
200
100
200.6
2030
Repua
average
singee Sestt
pedeet vedelt-esedd
overge
stomoton
gradenk
nethod
method
method
Figure 13: Linear Single Task Results. (Left) We observe that even averaging linear policies can often improve performance almost match joint training on the pooled data for single-task settings, and match task-specific training for multi-task settings. (Right) We observe that gradient-base based methods find general invertible matrices perform better than finding permutation matrices with alternations.
7. fleet-merge-averagee: This is an ablation method on our algorithm where instead of trying interpolate between each local model and the averaged model, we the average of these local models and compute the gradients.
As we show in the figure, our methods outperform all the baselines across various neural network architectures, such MIP on MNIST dataset, RNN Shakespear dataset from LEAF (Caldas et al., 2018) (sub-figure f), as well as CNN on CIFAR- 10 dataset (Krizhevsky et al., 2009) (sub-figure c). We also ablate against the design choices in the algorithm (sub-figure e) and ablate against both non-lIDness (sub-figure a,b,c) and model numbers in the multiple model merging setting (sub-figure d, f).
G LINEAR POLICY MERGING
G.1 PROBLEM SETUP
Linear quadratic control. A classical sub-setting of the policy setup setting described in Section and one whose optimal policy is recurrent, is that of linear quadratic Gaussian (LQG) control (Åström Murray, 2021). To respect the notational convention in control theory, we use (x, Yr. u,) to denote Ot, 0t, a:) in Section 2. Let n be the dimension of the state, the (state, observation, action) tuple of (st,0t,
29

---

Published as a conference paper at ICLR 2024
p the dimension of the observations, and m the dimension of the control input; that is, x, R", Yt E R" and u, E Rm. The LQG problem is defined by a linear time-invariant (LTI) dynamical system parameterized by dynamical matrices (A, B, C) of appropriate dimenion, evolving according to the following dynamics:
X++1 = Ax; + Bu; + +W
(G.1)
y: Cx2 Vt-
Above, for & 2 O, W, and are noise vectors drawn i.i.d. Gaussian from (o E..) and ( E.) respectively. The initial state Xy is also drawn from Gaussian distribution (O, Eg). The classical LQR cost is positive definite quadratic cost
c(x, u) :- x Qx + u Ru, 0, R 0.
The goal of LQG is to minimize the following long-term average cost over policies (Xia, U1a-1) + un:
1
x² E argmin
( = sup T-000
Σ 1=
x Qx + u Ru;
(G.2)
n
where above Ex denotes expectation under the closed loop dynamics induced by and the dynamies Eq. (G.1) In the special case with full state observability, we have C = / and E, 0, and the problem is also referred to linear quadratic regulator (I.QR) control.
now classical that under standard controllability and observability assumptions", the optimal x² in Eq. (G.2) is dynamic linear policy, i.e. the specialization of recurrent policy to this linear setting.
Policy parameterization.
On finite horizons, the optimal control law roughly correspond takes the form u, KK., where the time-varying gain matrices K., solves a finite-time LQR problem, and where x; = E[[x ht].
where h2 is the observations (y1.2, U1:--1) up time t. Passing to the infinite horizon, the optimal policy takes a form we now describe.
First, compute an optimal control gain matrix K, can be computed
P.B K. -B -BBBB + B
(G.3)
where P. solves the discrete algebraic Riccati equation:
P A PA + A PB(B PB + - Q.
(G.4)
Similarly, we compute matrices L, and Σ, are solve Ricatti equations dual to those for which K, and P. + are the solutions:
L. = E.C ECE, CT +E.)-1
(G.5)
Σ, AE,A C (CE,C +E.)- Σ..
Because of the similarities between the equations defining L,, and K., L, is usually referred to as the Kalman gain. Under the optimal policy, the latter is used to define an optimal reccurent estimate x, of the true state F,, obtained via the Kalman filter equation:
x A Bu, + (y:
-B,,)))
(G.6)
The optimal control policy can then be expressed as
*; A. x + BB. Dt.
u, K, =: Ca-
(G.7)
Ag. = (A BK, - C.A(A + BK,)), Bg. = I Co- := K,
which itself a linear dynamical system, with inputs (yr) states (8t) and outputs (u.)22 i.e. inputs of the control policy are outputs of the system, and vice versa.
This motivates considering a policy class of policies parametrized by == (Ag, [p, Cg) where Ag _ (i.e., the same dimensions A.), Cg has the same dimension as TT, and B3 the same dimensions as C Following the formula in Eq. (G.7):
*; -
u, Caai,
(G.8)
and we can regard 0* := (Ag.. Bg. Ca- ) as the optimum of such set of parameters. Thus, in our behavior cloning objective, we optimize over 9 of the form (G.8). Again, we stressed can be viewed
Mooe generally, under stahilizability and detectability assumptions.
30

---

Published as a conference paper at ICLR 2024
special case the RNN parameterization introduced in Section 2, with one hidden layer and no activation functions. We let Th denote the corresponding policy, i.c., (TT) ) the cost associated with executing the policy induced by 0.
Feedforward controller. In the special case where C and Σ, = 0, yr = X, optimal controller because the static (i.e., one-layer feedforward) control law u, K,,,,, where K, is exactly in (G.3) above. This corresponds to the case of (Ao, Bo, C.) = I0, K.). Notice that only with longer optimal, because better estimate & of the
C = I (but Σ, + 0), feedforward policies system state 2g can be obtained by taking the history into account. Nevertheless, we can still view the policy x, + K, u, as an imperfect expert that one would like imitate, studied by Zhang et al. (2022).
G.2 SYMMETRIES IN NONLINEAR CONTROL
is both well-known and straightforward that linear control policies are invariant under invertible linear transformations acting follows
(Ao: Bo: Co) + (TApT TBp, C&T det * O.
(G.9)
As example, the policies (Ag, Ba, Ca) and ( Cg) are equivalent, maintaining internal state which are sign-flips of the other. In the special case of imitating a feedforward expert, there no symmetries. Indeed, if Kx K'x for all then K K'.
G.3 DETAILS FOR TRAINING IMITATION LEARNING
In this section, describe the details for training the imitation learning policy, given in Algorithm below. We begin with the training of dynamic policies. We assume access to a dataset of expert trajectories (() An important point is initialization. the dynamics defining & are stable We also ensure that Bş, Ca have unit Gaussian entries. Note that Ag = is not necessarily a critical point of the loss. Indeed, when because when As = &; A + Bay: = Bayt. & does not vanish, that thus gradients with respect to As can be incorporated. Note that Line 10 in the update requires differentiating through the entire dynamics. Notice that while LQG policy optimization is still unsolved, gradient descent for imitation learning is dual to policy search over Kalman filters, and can be solved in model-free manner due to Umenberger al. (2022). Feedforward imitation learning is achieved by direct least squares Algorithm 3.
initialize A = which ensures that
Algorithm 2 Dynamic Linear Policy Learning
1: Input: Expert dataset (i( (y -T-1 1T
u (
1
2: Initialize A₁ + O, and Bg: Cg to have unit Gaussian entries
3: for iter 1, 1.., itermax do
4: Sample one trajectory T (y1, UI -T-1.
5: Initiate: x1 + 0, 2-
6: for 1. E (1,...,7) do
7:
Update step: &; 1o.tAnu Agxx-1 + Bay:
8: Compute control: ú; f- Cax:
9: Update loss: E E+ | û
Gradient Update: 0+0
17 a where gradients are taking through the entire dynamics.
di
11: Return: 0 = (Ag, Ba, Ca)
Algorithm 3 Statie Linear Policy Learning
1: Input: Expert dataset (((()) yy
2: Solve K E argming Σ; (1) Ku by least squares
3: Return: K
31

---

Published as a conference paper at ICLR 2024
G.4 METHODS FOR LINEAR POLICY MERGING
In the linear setting, we propose two different types of policy merging algorithms combine the policies. Both algorithms alternate between updating each model's transformation matrix and the merged policy parameters, in order avoid the problem caused by the nonconvexity of the loss. The first algorithm leverages closed-form solutions such as linear assignments (see this subsection, Appendix G.4.1), whereas the second algorithm based on gradient-descent updates (see Appendix G.4.2).
G.4.1 METHOD PERMUTATION BBSSED ALIGNMENT
Our first approach is to focus finding permulation transfornations merge the models, despite that for linear dynamic policies, they are invariant to general invertible transformation matrices. This type of approach has also been used in the nonlinear setting, we propose to minimize the sum of the distances of each individual model 0g = (Aa,, Be,, Co,) to the merged model 0 (Ag. B C.) through permutations (P: (Pi)ee(N)
introduced in Section 3.1. Specifically,
N 9 (0, (P))) = ElIA- P.
P.T Bo, BI, Co. Pill
min P. (.P.)-EES perm
(G.10)
Despite the loss in Eq. (G.10) being jointly nonconvex in the variables (0, (P.) we notice that the problem becomes tractable by fixing of the variables and solving for the other. Hence, propose a two-step alternating procedure to solve Eq. (G.10): the Merging step and the Aligning step.
Merging step. In the merging step, fix 1 and optimize for 0. This is equivalent to computing the mean of all individual transformed 0, to estimate the merged model Specifically, solve
AA, B, C) <- arg min 9 (0: (P:)), where
N Σ A. B
N
N Č= Σ P.
2 P.B..
N 1=1
N
N
Note that this is essentially the solution to some least-square problem.
Aligning step. In the aligning step, for each local dataset i, we fix the As, and A, and compute the permutation matrix P; by solving linear assignment problem to align the local model with the merged model. To avoid the two-sided linear assignment problem that is not tractable, we reuse the permutation matrix from the previous iteration P/ for the current iteration, that we can solve it regular linear assignment problem (Bertsekas, 1998):
min Ppegm
|| (P') Ao, P. 1|=
BB P.B.B.l. ||C Co. Poll
(G.11)
which can be simplified to solving
max P. P.E9parm
(P,, AT P.Ag + B BJ + C.)-a)-
(G.12)
We note that there are finite number (N pair) of pairs (P;, P(), and the objective (P,, non-decreasing. Therefore this alternating procedure is guaranteed to converge.
SSince there only one hidden layer, notational convenience, we use P: the permutation matrix for each local dataset i, instead of using P as in Section 2, which denotes a sequence of transformation matrices across layers. With a slight abuse of notation, we still use Gperm to denote the set of such permutation matrices and P. Sperm- A similar convention applies to P; E Gun later to denote P; belonging to the set of all invertible matrices.
32

---

Published as a conference paper at ICLR 2024
G.4.2 METHOD 2: UNCONSTRAINED GRADIENT DESCENT
Our second approach for policy merging is find solution the set of general invertible matrices, through gradient descent for the objective as in Eq. (G.10) directly, i.e. now we allow P, E Glin
N
min 8, EG
9 (0, (P.) = - Ao | +|C P
[ +
Ao, Pill
1=1
(G.13)
We consider simpler loss, which is obtained by right-multiplying the expression inside the first and second terms by P,;
min M.(P,ML
9 (0. = Σ P.As A | B B
Co Pill
(G.14)
Note that the losses (G.13) and (G.14) are equivalent in spirit the sense that, if there exists matrices ] which perfectly align all 9; then those same matrices are the global optima for both expressions. The second expression (G.14) has the advantage admitting simpler gradients, and is indvidually convex with respect to either {P.} or (A,, Bg, Cş), holding the other fixed. Therefore, for N models, we run alternating gradient descent to update those parameters to update the averaged model parameters 0 and the transformation matrices P (P1, .Nn).
Remark G.1 (Possibility of Degenerate Solutions). Note that Eq. (G.14), as long as Bg, are non-zero, the degenerate solution with P; = 0 will not be optimal even in Eq. (G.14). More generally, we find that gradient descent on Eq. (G. (G.14) works well in practice. An alternate approach would be to alternate gradient updates on {P,} with exactly solving for the optimal 0 (Ag. Bğ, Ca) in Eq. (G.14) (which is a well-conditioned least squares problem whenever the singular values of each Pt are bounded away from zero). In our experiments, find this is not necessary.
G.5 EXPERIMENT SETUP DETAILS
We set up a linear system in the experiment with state dimension n = 4, input dimension m 2, and
To model different tasks, we choose R 12, Q = a (+) with (A) E log space (-2, H + 1) and H 9.
We consider such 10 different system realizations, and depending on whether it is partially observable (LQG case) or fully observable (LQR case), choose the observation matrix C randomized matrix (but fixed for each task) by sample each entries from a unit gaussian or C = I, respectively. For the former, the entries in C [250 are independently sampled from Gaussian distribution. We run 10 trials for each experiment. The rollout horizon 100. The system dynamics follow from that in Hong et al. (2021), and the setup of R, Q follows from that in Zhang al. (2022) on multi-task imitation learning in linear control. We generate expert data with the optimal LQG/LQR controller, and train the policies with imitation learning. Policy performance is evaluated in terms of the closed-loop rollout cost. Finally, in both the settings of learning static and dynamic policies, the parameterized controllers are trained with a standard gradient descent algorithm.
For the last experiment, we set up the following systems with 10 systems where half of them are positive and half of them are negative, following the disconnected component constructions in the next section. We show that under disconnected component settings, we need general invertible matrices to successfully merge different linear policies.
0.01
0.99
1.1
0.03
0.02
33.44
1.66
A Aunstab
0.01
0.47
4.7
B
-0.83
0.44
[0.02
-0.06
0.40
0.47
0.25
G.6 RESULTS AND DISCUSSIONS
G.6.1 BASIC EXPERIMENTAL RESULTS
As shown in Figure 13, with static feedback policies (i.c., "feedforward" policies), average can sometimes match the performance of policies that are trained with shared data (multi-task) task-
33

---

Published as a conference paper at ICLR 2024
specific data (task-speciffic). We then compare the performance in cases with dynamic policies (i.e., "recurrent" policies), and observe that the gradient- method can outperform the alternation- beeeattion-bued method that searches in only the space of permutation matrices. These two methods, by accounting for the similarity invariance in merging, outperform the naive averaging method (average). They also outperform single-task, where policies only trained with local datasets, since more data has been (implicitly) used by merging multiple policies. More interestingly, when there are disconnected components in the landscape (as constructed in odd-dimensional systems, see more details in the next subsections), we observe that only searching for permutation will not mitigate the issues.
G.6.2 GRADIENT-BASED APPROACH V.S. LQG LANDSCAPE
In this subsection, we reconcile the effectiveness of gradient- based policy alignment with the topo- logical properties of the the LQG cost landscape. First, recall that we parameterize controllers by π (A,, B Bgg Co), where Ag maintains latent state of the same dimension as the system state. Let C. denote the set of policies
(Ao, Bo, Co) such that the cost J(xr is finite, or equivalently, the set of controllers which stabilize the dynamics (Zhou et al., 1996). Zheng et al. (2021) show that this set is either (a) fully connected, or (b) disconnected, but has only two connected components
( cl Case (a) occurs if either
(i) the open loop system is stable, i.e. the spectral radius p(A) < 1, that is i.c. maximum eigenvalue magnitude) of the system dynamic matrix is strictly less than 1. This is equivalent to saying that (O) is finite for the zero controller (Ag, Bg, Cq) (0,0, 0).
(ii) there exists a policy ı with rank (Ao) < n. such that π E Cn-
And, in case (b), Zheng et al. (2021, Theorem 3.2) shows that any two connected components can be realized by any mapping of the form
(As, Bg, Cş) (TAST TB2, C&T )) det(TT < 0.
(G.15)
For example, in odd dimensions, the matrix T -I has det(T) = -1. Thus, we find a simple corollary:
Proposition G.1. Let the state dimension n be odd. Then, C, is disconnected and 0 (Ae, BB, Ca) is stabilizing, then the policies (As, By, Ce) and (Aa. -B. -Ca) lie in different connected components of Cn. More generally, consider any block diagonalization of
A
(1,1)
A (1,2)
B (1.1)
A
9
2.1)
2.2)
B 2.11
c(
(G.16)
0
where the block dimensions are chose such that "1" dimension has odd dimension. Then, the above is stabilizing and C, is disconnected, then the following equivalent controller lies in the other connected component of C,:
A (1,1)
(1.2)
B (1.1) (1.1)
(1.1) c((.2)]
(G.17)
- - 12,1)
A (2.2)
B
(2.1)
0
As a second example, note that when T is an "odd" permutation, det(T) = -1, placing the transformed parameters into separate components.
G.6.3 INSUFFICIENCY NAIVE AND PERMUTATION BASED MERGING
Importantly, path connectedness does not mean connectedness by linear paths, the latter being pre- cisely the definition of convexity. Indeed, consider the equivalence transformation (Ag. Bo, Co) H> A C.)). The mid point between these controllers if (A,, O, 0) which, one can check, is equivalent to the zero controller (0, 0, 0). Not only is the zero controller essentially never optimal, if A is not stable, then this controller is not stabilizing. Similarly, permutations can be insufficient to merge as well: indeed, permutation based mergin coincides with naive merging in dimension one, which the above shows will fail if A is unstable (i.e., in dimension one, |A] 2 1). These examples highlight the important of aligning using all invertible matrices, and not simply permutation matrices.
34

---

Published as a conference paper at ICLR 2024
G.6.4 AVOIDING WATH-DISCONNECTNEE OF STABILIZING CONTROLLERS
In addition the insufficiency of linear and permutation-based merging, may seem that the possible path disconneetedness of the set Cn may pose challenge to weight merging. However, are experiments do not seem to suggest that this is issue for the unconstrained gradient descent method proposed in Appendix G.4.2. The reason for this is our choice of merging cost in (G. (G.44). This cost penalizes only differences between the controller parameters. does not, as in FLEET-MERGE (Algorithm 1), use cost depending on the policy error. Therefore, even if one of the iterates lies outside of the stabilizing set Cn (and therefore has infinitely large infinite-horizon LQQ cost J), we still compute gradients with respect to the parameter error (G.14). In other words, merging in parameter loss allows us to pass through disconnectedness the landscape of control cost J.
This can be seen in the following simplified variant of (G.14) for the merging of two LQG policies 81 = (A1. B1, C1) and 02- = (Ag. B2, C2). Define
((P) = |PA| A2P PI. + |BPP- B2ll + ||C, C2PII
(G.18)
Then note that this loss is convex in P. Under natural conditions, it is in fact strongly convex, and thus admits a unique global minimizer P. If 0, and 02 are equivalent, i.e., (A2, Bg, C2) TTA,
TC₁) for some transformation T, then this unique minimizer must be P T, as this has zero loss in (8.18). This is true regardless of whether or not 01, Bg lie in separate connected components of C, Hence, disconnectedness.
how merging based parameters alone circumvents possible loss
35


Reviews:
[{"review_id": "TDi87BRxEa", "reviewer": "Reviewer_FoVN", "summary": "This paper introduces a method, called Fleet-Merge, for merging neural network policies trained on different tasks into one neural network policy. When a large fleet of robots learn a diverse set of skills on diverse tasks and it is not possible to collect a centralized dataset of all environmental interactions, how can we learn singular policies that capture all the training done by individual robots? A common tool to merge policies in federated learning is to simply average the weights/biases of all the neural networks into one policy. However, this leads to a drop in performance because the internal representation space of these policies may be significantly different. This paper introduces a method for merging the policies using the permutation invariance property of RNNs. By applying a permutation matrix to each layer of weights and biases, the paper aligns the internal representation space of the RNNs and is able to perform better merging. The authors test on a series of domains including linear control tasks and realistic robotic control tasks. They also introduce a tool use benchmark called Fleet-Tool, a task suite for robotic tool usage (like hammering, screwing, cutting,...)", "questions": "1. At test time, how does the merged policy know what task to solve? Could there be an ambiguity on what the task specification is?\n2. 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.\n3. 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/calcualted", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is a very interesting paper and is well written. The problem that they investigate is well supported and they propose (to my knowledge) a novel algorithm that allows us to merge different neural networks by aligning their internal representation spaces. They do extensive experiments in simulation and the real-world and show that policies from diverse tasks can be merged together in a smart way.\n\nThey leverage known properties of RNNs (i.e.  internal permutation invariance and time-invariance) and find transformation matrices to match the parameters of different RNNs together. This is an interesting idea and very applicable to the problem that they are solving.", "weaknesses": "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. \n\nFurther, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "QYgCpkfHnm", "reviewer": "Reviewer_2Z2Z", "summary": "This paper studies policy merge, which seeks to merge policies that are trained on different datasets and tasks into a single set of weights, while preserving the learned skills of the original policies. In particular, taking account into the sequential and dynamic nature of robot policy learning, this paper focuses on merging RNN-parameterized policies. The proposed method, Fleet-Merge, outperforms prior model merging methods by over 50%, accounting for not only the permutation symmetries in RNN policies but also enabling multiple rounds of merging between each training update as well as merging multiple models.", "questions": "Several suggestions and questions are already included in the Weaknesses section. Here are some additional questions/suggestions:\n\n1. 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.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper is very novel in both its problem setting. Policy merging, to the best of my knowledge, has not been explicitly studied before. This paper will have impact in the robot learning community, where large-scale datasets as well as data/model sharing across institutions are becoming increasingly prevalent. However, most such effort has primarily focused on consolidating all data to train a single large model (e.g., the recent RT-X effort), but this paper offers a fresh perspective that opens up the possibility of different institutions sharing their pre-trained policies, which are much more manageable.\n\nThe proposed method, at a technical level, is largely derived from a prior work (Pena et al., 2022). That said, the paper does introduce several additional improvements that are particular fitting to the policy learning setting. For one, Fleet-Merge is designed to be able to merge multiple models and do so in multiple-round fashion, both of which are important features to have in sequential decision making fashion; in contrast, prior works in the model merging literature is primarily considering the simplified case of merging two models. Because of these changes, Fleet-Merge does diverge from prior works with several modifications such as grounding the model merging in each iteration to the current average of all models, reminiscent of common practices in federated learning. Therefore, the approach, though extensions of prior works, is well motivated.", "weaknesses": "The main weaknesses of the paper is in its presentation and the strength of the experimental results.\n\nFirst, the presentation of the paper can be improved. In Figure 3 and 4, several charts are hard to see because the legends block them.  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. The new benchmark Fleet-Tools is not well motivated and the section on \"Keypoint transformation & trajectory optimization\" within Section 4 is too dense while not directly related to the algorithmic components of the paper. I suggest the authors better motivate why Fleet-Tools is particularly well-suited for studying the policy merging problem. \n\nThe 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.  Considering that either baseline is missing key components in Fleet-Merge, it is worth asking the question whether Fleet-Merge is truly necessary given its added algorithmic complexity as well as computational cost. Furthermore, while the paper emphasizes the multi-task aspect of policy merging, most of the experiments are conducted in the single-task setting. Interestingly, in the multi-task setting, simple averaging appears to be very competitive. \n\nOverall, I think this paper is promising and I am willing to improve my score if the authors could resolve my reservations about the practical utility of the proposed method as well as more clearly present the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "3kHqRuFQ5Y", "reviewer": "Reviewer_DxkE", "summary": "The authors propose a method to merge policies learnt by a fleet of robots trained for different skills. The proposed FLEET-MERGE algorithm which allows for distributed learning of recurrent neural networks through a merging strategy that accounts for permutation invariance of neurons in layers. The authors also introduce a new robot manipulation benchmark for fleet policy learning, which they use to evaluate their method.", "questions": "1) 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?\n\n2) Could the authors summarize the main novel contributions of the FLEET-MERGE algorithm, especially compared to the work from Pena et al. [1]? I would be content if the contributions are mainly experimental validation for robot learning, but would appreciate the clarity since the exact contribution is difficult for me to parse.\n\nIt is possible I did not understand the contributions properly, so I would be willing to raise the score if the authors can provide a strong response to my questions. But as of now I am 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.\n\n[1] Re-basin via implicit Sinkhorn differentiation, arxiv.org/abs/2212.12042", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "1) The experimental results for policy merging are strong for merging RNN policies trained for different tasks. The proposed method outperforms naive averaging and other similar baselines. Distributed learning is becoming increasingly important for robotics as deployment of fleets of robots for data collection is more feasible.\n\n2) The introduced benchmark FLEET-TOOLS is a useful contribution to the robot manipulation community, for easy collection of expert trajectories useful for training policies with imitation learning.\n\n3) The algorithms and experiments are explained clearly, with good presentation.", "weaknesses": "1) 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.\n\n2) The novelty of the proposed method is not very clear. The algorithm used by the authors is based heavily on Pena et al. [1], as such it is hard to see what the main contribution is, apart from the FLEET-TOOLS benchmark. Is it the application of the algorithm for training RNNs, and the experimental demonstration of its usefulness in fleet robot learning?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (poster)