# DiNNO: Distributed Neural Network Optimization
for Multi-Robot Collaborative Learning

Javier Yu,1 Joseph A. Vincent,1 Mac Schwager1 *This work was funded by NASA
ULI grant 80NSSC20M0163, NSF NRI grant 1925030, and NSF NRI grant 1830402. The
first author was also supported on a NSF Graduate Research Fellowship, and the
second author was supported on a Dwight D. Eisenhower Transportation
Fellowship.1Department of Aeronautics and Astronautics, Stanford University,
Stanford, CA 94305, USA, {javieryu, josephav<EMAIL_ADDRESS>
###### Abstract

We present a distributed algorithm that enables a group of robots to
collaboratively optimize the parameters of a deep neural network model while
communicating over a mesh network. Each robot only has access to its own data
and maintains its own version of the neural network, but eventually learns a
model that is as good as if it had been trained on all the data centrally. No
robot sends raw data over the wireless network, preserving data privacy and
ensuring efficient use of wireless bandwidth. At each iteration, each robot
approximately optimizes an augmented Lagrangian function, then communicates
the resulting weights to its neighbors, updates dual variables, and repeats.
Eventually, all robots’ local network weights reach a consensus. For convex
objective functions, we prove this consensus is a global optimum. We compare
our algorithm to two existing distributed deep neural network training
algorithms in (i) an MNIST image classification task, (ii) a multi-robot
implicit mapping task, and (iii) a multi-robot reinforcement learning task. In
all of our experiments our method out performed baselines, and was able to
achieve validation loss equivalent to centrally trained models. See
https://msl.stanford.edu/projects/dist_nn_train for videos and a link to our
GitHub repository.

##  I Introduction

A group of collaborating robots has the ability to explore, interact with, and
experience their environment as a collective much faster than a single robot
acting alone. This ability to rapidly gather a large volume and variety of
data makes multi-robot systems especially well suited for tasks that involve
training deep neural networks using data gathered by the robots. In a cloud
robotics scenario, one can imagine thousands of robots networked over a cloud
server, able to collectively gather and process vast volumes of data for a
common task (e.g. manipulation, autonomous driving, or human behavior
prediction). In a mesh network scenario, one can similarly imagine a team of
robots collaborating to map an environment, learn a control policy, or learn
to visually recognize threats in the environment. A central unsolved problem
in collaborative robotics, therefore, is how to train neural network models on
the robots through local communication such that each robot benefits from the
data collected by the entire multi-robot system.

Figure 1: DiNNO allows robots to cooperatively optimize local copies of a
neural network model without explicitly sharing data. In this figure
(representative of Section V-B), three robots use DiNNO to cooperatively
optimize a building occupancy map represented as a neural network. Each robot
only sees part of the building, collecting a local lidar data set (colored
cylinders). The robots communicate over a wireless network (dashed lines) to
cooperatively optimize their local neural network copies. The resulting model
is as good as if it were trained centrally with all data at once.

To solve this problem, we propose Distributed Neural Network Optimization
(DiNNO), an algorithm built on the alternating direction method of multipliers
(ADMM) [1], and that integrates easily with standard deep learning tool sets
like PyTorch. Using DiNNO, robots alternate between local optimization of an
objective function, and communication of intermediate network weights over the
wireless network. The robots eventually reach a consensus on their network
weights, with each robot learning a neural network that is as good as if it
had been trained centrally with the data from all robots, as illustrated
graphically in Figure 1. DiNNO inherits the strong convergence properties of
ADMM—for convex objective functions we prove that all robots obtain globally
optimal parameters. However, neural network training is rarely convex. Using
standard deep learning tools within DiNNO we are able to retrieve state-of-
the-art deep learning performance, but in a distributed, multi-robot
implementation. DiNNO can accommodate time-varying communication graphs (e.g.,
due to robot motion), and also suits streaming data or online learning, if
robots train their network while they move and collect data. Finally, DiNNO
operates by sharing network weights over the communication network, not raw
data. Therefore, robots using DiNNO preserve the privacy and integrity of
their own local data set. This is crucial in scenarios where user data or
observations of humans are involved, or when different robot manufacturers
need to preserve the privacy of their own data sets.

A naive approach to solving the multi-robot deep learning problem is to use a
mesh network routing protocol to aggregate the data gathered by all of the
robots in the system to a single “leader” robot which then optimizes a deep
neural network model, and sends a copy of that trained model back to all of
the other robots in the system. We refer to this approach as a “centralized”
solution, and it has a number of distinct drawbacks. First, depending on the
size of the gathered data, algebraic connectivity of the communication graph,
bandwidth of the communication links, and efficiency of the routing protocol
it can take a significant amount of time to aggregate the gathered data at the
leader node. A centralized approach is not robust to failure of the leader
node, and there are some applications for which it may not be possible to
transmit data due to privacy considerations, for instance, due to the European
Union General Data Protection Regulation article 46 [2]. DiNNO overcomes all
of these limitations by enabling leaderless distributed neural network
training through local communication among the robots.

The paper is organized as follows. We give related work in Section II and
introduce the distributed collaborative learning problem in Section III. In
Section IV we derive DiNNO starting from a well-known variant of ADMM called
Consensus ADMM. In Section V we present three example robotic deep learning
tasks that showcase our method.

##  II Related Work

Learning has been used to address a variety of problems for multi-robot
systems and is becoming more popular as deep learning tools are made more
accessible. A deep learned controller is used to model multi-quadrotor
interactions in [3] and [4], and is shown to considerably outperform
traditional non-linear controllers. Reinforcement learning can also be a
useful tool in multi-robot contexts and [5] shows that an offline learned
value function can be used to perform real-time and uncertainty aware
collision avoidance. The collective training methods in [6] and [7]
demonstrate how experience aggregation from multiple robots can speed up
policy optimization. While not deep learning, Gaussian processes are another
popular learning tool that have been used for various regression tasks with
data collected online by multi-robot systems [8, 9]. These works all showcase
applicability of collaborative learning models for multi-robot systems, but in
general, aside from [9], do not provide a distributed framework from which to
perform this learning.

The problem of training neural networks in a distributed way using data
aggregated from individual robots can be viewed as a specific instantiation of
a more general class of problems referred to as distributed optimization
problems. Distributed optimization is the study of algorithms for solving
optimization problems where a sum of individual objective functions, which
correspond in this case to the individual robots, is optimized using local
computation and message passing. This formulation was first proposed in [10],
and has been of renewed interest since the seminal work [11] which presented
distributed subgradient descent for convex distributed optimization problems.
Subsequent research has focused on improving convergence rates [12] and
extending analysis to a broader range of problems including time-varying
communication graphs [13], and online or streaming convex objectives [14]. For
an overview of the broader distributed optimization literature we direct the
reader to the surveys [15, 16, 17].

Our algorithm is designed for networked robots collaboratively optimizing a
neural network model using data local to each robot. Specifically, this task
requires distributed optimization of nonconvex deep learning loss functions in
situations where the communication graph can be time-varying, and where data
may be streaming as it is collected by the robots. While, to our knowledge, no
existing works formally address all of these requirements, there are some
works in the distributed optimization literature that address general
nonconvex distributed optimization objectives, and some use simple distributed
neural network training problems as motivating examples.

In [18] the distributed subgradient descent algorithm is extended to
stochastic gradients as the distributed stochastic gradient descent (DSGD)
algorithm, and uses training of a CIFAR-10 classification model as a benchmark
problem. One approach to improving convergence rates in distributed
optimization is to introduce an auxiliary variable that estimates the global
gradient. Several works make use of this mechanism, and extend it to the
domain of nonconvex optimization with stochastic gradients [19, 20]. The
Choco-SGD algorithm for distributed deep learning [21] is another algorithm
similar to DSGD with the variations that it uses a gossip mechanism for
improved consensus, and incorporates a quantization step for reducing
communication bandwidth. The authors of [22] propose using local convex
approximations, based on global gradient estimates found through gradient
tracking, for distributed neural network optimization and show basic results
on small networks.

The edge consensus learning algorithm proposed in [23] is similar to our
approach in that it is derived from ADMM, but instead of addressing the
nonconvex primal update directly it uses a linearization similar to that
proposed in [24], which results in a gradient descent like update. While still
technically a primal-dual method, the update equations do not include a local
optimization procedure, and are more similar to primal domain methods like
DSGD (see [24], Remark 1). A number of other nonconvex distributed
optimization methods are discussed in the survey [25].

Our proposed algorithm for distributed multi-robot deep learning takes full
advantage of existing deep learning tools and optimizers, and we demonstrate
superior performance compared to benchmarks on robotics related deep learning
tasks such as neural implicit mapping and deep multi-agent reinforcement
learning.

##  III Problem Formulation

We consider deep learning problems where portions of a data set,
$\mathcal{D}$, are sent to, or collected by, $N$ robots that operate in a
connected communication graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$. Let
$\mathcal{D}_{i}$ be the portion of local data that belongs to robot
$i\in\mathcal{V}$ where the union of all the local data sets,
$\mathcal{D}_{i}$, is the joint data set, $\mathcal{D}$. In some cases
$\mathcal{D}_{i}$ can represent access to a time-varying data set gathered
from an incoming private data-stream (as in Section V-B).

The model we would like to optimize has the form $y=f(x;\>\theta)$.
Specifically we consider $f$ to be a deep neural network that implements a
continuous function $f(x):\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$ to give a
map from inputs $x\in\mathbb{R}^{n}$ to outputs $y\in\mathbb{R}^{m}$. The
neural network is parameterized by network weights $\theta\in\mathbb{R}^{d}$.
We make no special assumptions about the architecture (e.g. feed-forward,
convolutional, residual, etc.). We can then formalize our distributed learning
optimization problem as

$\displaystyle\underset{\theta\in\mathbf{R}^{d}}{\text{minimize}}\quad\sum_{i\in\mathcal{V}}\ell(\theta;\mathcal{D}_{i})$
(1)

where $\ell(\cdot)$ is the objective function (loss function) which is
generally nonconvex and often nonsmooth (due to ReLU activation). Common deep
learning tasks such as classification, regression, and unsupervised learning
have different objective functions and a distributed deep learning optimizer
should be general enough to achieve good performance across all of these
problems.

Suppose that the decision variable, $\theta$, is separated such that each
robot maintains their own instance of it, $\theta_{i}\in\mathbb{R}^{d}$. This
yields the equivalent optimization problem

$\displaystyle\underset{\theta\in\mathbf{R}^{d}}{\text{minimize}}$
$\displaystyle\quad\sum_{i\in\mathcal{V}}\ell(\theta_{i};\mathcal{D}_{i})$
(2a) subject to
$\displaystyle\quad\theta_{i}=\theta_{j}\quad\forall(i,j)\in\mathcal{E}.$ (2b)

This optimization problem is amenable to a distributed solution in which
robots minimize local objective functions, and take additional steps to come
to agreement (consensus) on the value of the decision variable. Replacing the
data defined loss functions in (2) with arbitrary objective functions yields
the general formulation of a distributed optimization problem. Though we
consider distributed deep learning which is typically unconstrained, a
constrained formulation of (2) covers many other robotics problems including
distributed target tracking, coordinated package delivery, and cooperative
multi-robot mapping [15].

##  IV Distributed Training

A standard method for solving convex distributed optimization problems is the
consensus alternating direction method of multipliers (CADMM). CADMM is an
ADMM-based optimization method where compute nodes (robots) alternate between
updating their primal and dual variables and communicating with neighboring
nodes. To achieve a distributed primal-dual update, CADMM introduces auxiliary
primal variables (i.e. $\theta_{i}=z_{ij}$ and $\theta_{j}=z_{ij}$ instead of
$\theta_{i}=\theta_{j}$). CADMM works by first optimizing the auxiliary primal
variables, followed by the original primal variables, then the dual variables,
as in the original formulation of ADMM [1]. Implementations of CADMM then
perform minimization with respect to the primal variables and gradient ascent
with respect to the dual on an augmented Lagrangian that is fully distributed
among the robots:

$\displaystyle\mathcal{L}_{a}=\sum_{i\in\mathcal{V}}\ell(\theta_{i})+p_{i}^{\top}\theta_{i}+\frac{\rho}{2}\sum_{j\in\mathcal{N}_{i}}||\theta_{i}-z_{ij}||_{2}^{2}$
(3)

where $p_{i}$ represents the dual variable that enforces agreement between
node $i$ and its communication neighbors, and $\mathcal{N}_{i}$ is the set of
indices for neighboring nodes of $i$. The parameter $\rho$ that weights the
quadratic terms in $\mathcal{L}_{a}$ is also the step size in the gradient
ascent of the dual variable. Furthermore, the algorithm can be simplified by
noting that the auxiliary primal variable update can be performed implicitly
($z_{ij}=\frac{1}{2}(\theta_{i}+\theta_{j})$). Initializing the dual variables
at zero then yields the following distributed update equations for CADMM:

$\displaystyle p_{i}^{k+1}$
$\displaystyle=p_{i}^{k}+\rho\sum_{j\in\mathcal{N}_{i}}(\theta_{i}^{k}-\theta_{j}^{k})$
(4a) $\displaystyle\theta_{i}^{k+1}$
$\displaystyle=\underset{\theta}{\text{argmin}}\>\>\ell(\theta;\mathcal{D}_{i})+\theta^{\top}p_{i}^{k+1}$
(4b)
$\displaystyle\qquad\qquad+\rho\underset{j\in\mathcal{N}_{i}}{\sum}\left\lVert\theta-\frac{\theta_{i}^{k}+\theta_{j}^{k}}{2}\right\rVert_{2}^{2}.$

Typically, the primal variables are initialized uniformly to an initial guess
$\theta_{i}^{0}=\theta_{\text{initial}}$. This derivation of CADMM is
addressed in much finer detail in [26].

In our multi-robot learning problem, the objective function for the primal
update (4b) if performed by each robot using CADMM is composed of three terms:
a neural network loss on the robot’s local data, a linear term from the dual
variable, and a regularization term. It is obvious that applying CADMM
directly to the neural network training problem results in intractable primal
updates due to the neural network loss component. The key insight, which we
use in our algorithm DiNNO, is that this primal optimization can be performed
approximately, stopping well before convergence to a local optimum. Formally,
we propose replacing the exact minimization of the primal update, (4b), with
an approximate solution found by taking a small number of steps, $B$
(typically between 2 and 10), of a stochastic mini-batch first order method
(SFO) on the entire primal objective function. Our proposed algorithm is shown
in Algorithm 1 with the approximate primal update performed in lines 12 - 16.
We replace the current primal iterate with $\psi$ in order to avoid including
two iteration count super scripts, and let
$G(\psi^{\tau};\rho,p_{i}^{k+1},\theta_{i}^{k},\\{\theta_{j}^{k}\\}_{j\in\mathcal{N}_{i}},\mathcal{D}_{i})$
represent the step taken by a SFO on the objective in the primal update. To be
clear, $G$ computes a stochastic gradient from a mini-batch in
$\mathcal{D}_{i}$ not a gradient on the entire local data set.

In some of our experiments we found it beneficial to also add a “scheduled”
increase (Algorithm 1, line 9) for the penalty parameter $\rho$ in similar
fashion to the learning rate schedules used in deep learning. For notational
simplicity, we overload the variable $\rho$ to also mean this schedule of
parameter values, and make explicit note of all cases where one is used.
Although generally we leave this term constant, it can be useful to gradually
increase it when faster consensus is desired. This schedule can be provided to
robots prior to optimization, and does not compromise the distributed nature
of the DiNNO algorithm.

An added benefit of our approach for distributed deep learning is that it
pairs well with existing deep learning libraries (e.g. PyTorch [27]) because
the approximate primal minimization can be performed with minimal changes to
the typical training loops used to optimize individual neural networks. We
find that this is beneficial because automatic differentiation and state-of-
the-art neural network optimizers, like Adam [28], can be used to perform the
approximate primal update, and practitioner knowledge from experience training
individual neural networks is transferable.

Algorithm 1 Distributed Neural Network Optimization (DiNNO)

1:Require: $\ell(\cdot),\>\theta_{initial},\>\mathcal{G},\>\mathcal{D},\>\rho$

2:for $i\in\mathcal{V}$ do $\triangleright$ Initialize the iterates

3: $p_{i}^{0}=0$ $\triangleright$ Dual variable

4: $\theta_{i}^{0}=\theta_{initial}$ $\triangleright$ Primal variable

5:end for

6:

7:for $k\leftarrow 0\text{ to }K$ do $\triangleright$ Main optimization loop

8: Communicate: send $\theta_{i}^{k}$ to neighbors $\mathcal{G}$

9: for $i\in\mathcal{V}$ do $\triangleright$ In parallel

10:
$p_{i}^{k+1}=p_{i}^{k}+\rho\sum_{j\in\mathcal{N}_{i}}(\theta_{i}^{k}-\theta_{j}^{k})$

11: $\psi^{0}=\theta_{i}^{k}$

12: for $\tau\leftarrow 0\text{ to }B$ do $\triangleright$ Approximate primal

13:
$\psi^{\tau+1}=\psi^{\tau}+G(\psi^{\tau};\rho,p_{i}^{k+1},\theta_{i}^{k},\\{\theta_{j}^{k}\\}_{j\in\mathcal{N}_{i}},\mathcal{D}_{i})$

14: end for

15: $\theta_{i}^{k+1}=\psi^{B}$ $\triangleright$ Update primal

16: end for

17:end for

18:return $\\{\theta_{i}^{K}\\}_{i\in\mathcal{V}}$

###  IV-A Convergence Properties

######  Theorem 1 (Optimality of Algorithm 1)

Let each local objective function $l(\theta,\mathcal{D}_{i})$ be strongly
convex and $L$-smooth. Furthermore, let

$\displaystyle
G=-\frac{1}{L}\nabla_{\theta}\big{(}l(\theta,\mathcal{D}_{i})+\theta^{\top}p_{i}^{k+1}+\rho\underset{j\in\mathcal{N}_{i}}{\sum}\left\lVert\theta-\frac{\theta_{i}^{k}+\theta_{j}^{k}}{2}\right\rVert_{2}^{2}\big{)}.$

Then Algorithm 1 converges to the unique global solution with linear
convergence rate.

###### Proof:

Given strongly convex and $L$-smooth local objectives, each primal update
(Algorithm 1 line 13) converges to the global solution with a linear
convergence rate as shown in [29]. Given globally optimal primal updates and
the stated assumptions, Algorithm 1 is a special case of the decentralized
ADMM algorithm studied in [30] where it was shown to have linear convergence
to the global solution. ∎

Clearly, in deep learning problems global solutions and linear convergence
rates are not ensured due to nonlinear neural networks creating nonconvex, and
often nonsmooth, objective functions. Nevertheless, we find that in practice
Algorithm 1 performs well for solving distributed deep learning problems, and
converges to solutions similar in quality to those found by centralized
optimization.

###  IV-B Baseline Algorithms

In Section V we show that DiNNO is an extremely effective method for
distributed training of neural network models. We compare DiNNO against two
other commonly referenced stochastic first order distributed optimization
methods: distributed stochastic gradient descent (DSGD) [18], and distributed
stochastic gradient tracking (DSGT) [19]. Like DiNNO, both methods have each
robot maintain a local copy of the optimization variable (neural network
weights), and use message passing and locally computed stochastic gradients to
collaboratively optimize the neural network. DSGD uses the following simple
update equation

$\displaystyle\theta_{i}^{k+1}=\underset{j\in\mathcal{V}}{\sum}w_{ij}\theta_{j}^{k}-\alpha^{k}g(\theta_{i}^{k})$
(5)

where $w_{ij}$ is an element of a doubly stochastic matrix $W$ that has a
sparsity pattern matching that of the graph Laplacian of $\mathcal{G}$,
$\alpha^{k}$ is a decaying step size, and $g(\theta_{i}^{k})$ is a stochastic
(or mini-batch) gradient of $\ell(\theta_{i}^{k};\mathcal{D}_{i})$. While (5)
may not at first appear to be a distributed algorithm, the sparsity pattern of
$W$ means that each node only needs $\theta_{j}^{k}$ from its immediate
neighbors to compute its update step.

The updates for DSGT are similar to those of DSGD, but an additional auxiliary
variable is added to estimate the gradient of the joint loss:

$\displaystyle\theta_{i}^{k+1}$
$\displaystyle=\underset{j\in\mathcal{V}}{\sum}w_{ij}(\theta_{j}^{k}-\alpha
y_{j}^{k})$ (6) $\displaystyle y_{i}^{k+1}$
$\displaystyle=\underset{j\in\mathcal{V}}{\sum}w_{ij}y_{j}^{k}+g(\theta_{i}^{k+1})-g(\theta_{i}^{k}).$
(7)

An important point to note is that for DSGT the message size sent at each
communication round is double that of both DSGD and DiNNO which only send
$\theta_{i}^{k}$. For DSGT and DSGD we use the Metropolis-Hastings weights as
$W$.

Alternative benchmark algorithms include ([20, 21, 22, 23, 25]) but, in
general, they share many core characteristics with the proposed baselines DSGT
and DSGD.

Figure 2: Each plot shows the Top-1 accuracy of the neural network models on
the validation set of the MNIST problem. For each algorithm we evaluate the
local neural networks stored by each robot on the validation set. Solid lines
show the average validation loss across all robots at the current
communication round, and filled areas are upper and lower bounded by the best
and worst performing networks for each particular algorithm. The four plots
show algorithm performance on four different communication graphs (using the
same hyperparameters across all instances). In all cases our method, DiNNO,
converges to the same accuracy as the centralized solution, and outperforms
both of the baseline methods.

###  IV-C Data Distributions

The way in which local data is partitioned between the robots strongly
influences the convergence rate of distributed deep learning methods. For
example, in classification, problems where each robot has access to a subset
of examples from all classes are easier to solve with distributed optimization
than problems in which each robot only has access to labelled data for a
single class. We refer to these two data distributions as homogenous and
heterogenous respectively. In homogeneous classification a robot which
optimizes directly on its local data set without communication may be able to
achieve a relatively high validation accuracy. Whereas a robot in the
heterogeneous case is unlikely to achieve a high accuracy on any class other
than what it observed.

##  V Experiments

In the following three examples we demonstrate that DiNNO can be applied to a
wide range of multi-robot learning applications, and demonstrates a
substantial improvement against baseline distributed optimization algorithms.
In each example we compare with two other common distributed optimization
methods: DSGD and DSGT. We implement DiNNO, DSGD, and DSGT in a general
framework such that for each experiment the optimization algorithm is
unchanged, but a different objective function and data set are provided.
Hyperparameter values are reported in Section Appendix.

###  V-A MNIST Classification

To clearly illustrate the potential for Algorithm 1 to train a shared neural
network from disparate data observers we first consider the well known MNIST
classification problem [31]. In this problem a neural network must learn to
classify images of handwritten digits to their correct integer values. We
train a network composed of a convolutional layer with three 5x5 filters
followed by 2 hidden linear layers of width $576,64$ with ReLU activation and
an output layer with log-softmax activation. We use the negative log-
likelihood loss function. Each robot only has access to labelled digits from a
single class. In this experiment, we perform distributed training with this
heterogeneous data distribution on four different communication topologies: a
complete graph, a cycle graph, and two random graphs.

In Figure 2 we show the average, worst, and best Top-1 accuracy for each
method on the distributed MNIST classification problem. Also included is the
centralized result (98.5% validation accuracy), and the individual results,
which as expected have roughly 10% validation accuracy. We observe that while
DSGT quickly trains to good accuracy in these problems, DiNNO achieves much
better final accuracy as training progresses with lower variance in later
iterations. DSGD has relatively poor performance with high variance. The
Fiedler value (algebraic connectivity) of each graph is reported to emphasize
that DiNNO is capable of training to centralized accuracy with highly
heterogeneous data distributions, and low graph connectivity.

###  V-B Neural Implicit Mapping

In robotics and computer vision there is growing interest in using neural
networks to represent functions which implicitly define the geometry and
visual features of an environment [32, 33]. In their basic form, implicit
density field networks take as input an $(x,y,z)$ spatial coordinate and
output a single density value between 0 and 1. Such networks are able to
represent complicated 3D scenes in a single memory-efficient function. In this
example we use DiNNO to learn the density field of a two dimensional
environment where data collection and computation is distributed across
multiple robots. The robots also have access to a global coordinate frame
which enables cooperative mapping, but a future line of research would be to
implement this same pipeline in conjunction with a distributed pose
optimization algorithm.

The environment we seek to map is a 2D building floorplan environment from the
CubiCasa5K data set [34]. Seven robots are deployed, and each robot gathers
data from the environment by collecting lidar scans as it traverses a closed
loop, precomputed trajectory. To simulate data streaming the robots update
their local networks at regular intervals from datasets of their last 400
collected lidar scans (an entire trajectory has 3000-4000 scans). Figure 3
shows the ground truth environment with seven robot paths and one lidar scan.
There is some overlap in the locations traversed by each robot, but many
locations, especially at the edges and corners of the map, are only viewed by
a single robot.

Figure 3: Ground truth map. Highlighted is a single robot’s trajectory and a
single lidar scan is shown with high (gold) and low (blue) density points.

We train a feedforward network with four hidden layers of size $256,64,64,64$
where the first hidden layer has sinusoidal activation, the remaining hidden
layers have ReLU activation, and the output layer has sigmoid activation to
restrict our density estimates to $(0,1)$. The sinusoidal activation of the
first hidden layer is common in implicit mapping networks and inspired by
Fourier Feature Networks [35]. For a loss function we use the binary cross
entropy between the sampled and predicted density.

Our validation set is composed of novel lidar scans from uniformly sampled
locations across the entire map, and this ensures that the validation loss
reflects loss only on areas where the robots have been able to gather data
(e.g. not inside walls). For the communication graph at each round an edge
between two robots exists if their pairwise distance is below a certain
threshold (1500 in this experiment). The constant motion of the robots results
in a communication graph that is dynamic (in terms of edges), but always
connected.

Figure 4 shows the validation loss for our method as well as DSGD and DSGT. Of
the tested methods DiNNO best minimizes the validation loss, once again
approaching the performance of centralized training. DSGD and DSGT train less
effectively, and converge to a poor quality solution.

Figure 4: Validation loss versus communication iteration for the neural
implicit mapping experiment. Both baseline algorithms DSGT and DSGD appear to
converge to a poor quality minima while DiNNO (ours) converges to a model with
validation loss matching that of the centralized solution.

Figure 5 shows the map learned by each method as well as learned maps from
individual robots training on only their own data. Each of the robots when
using DiNNO is able to provide a faithful reconstruction of the ground truth
environment while never having traversed its entirety or having received any
raw data from other robots. When tested with DSGD and DSGT robots converge to
similar local minima which result in incoherent maps.

To verify the performances of DSGT and DSGD we reran this experiment several
times, and both methods always converged to poor performing local minima.
Additionally, we emphasize that the simulation code for these two methods is
unchanged between this experiment and the MNIST one above where both methods
are able to learn acceptable classifiers. We speculate that this is a
challenging problem where only a small amount of suboptimality is allowable to
achieve a useful representation, and DSGT and DSGD may be unable to either
fine tune their neural network weights, escape poor local minima, or handle
streaming data.

Figure 5: The left most plot shows the ground truth density map, and moving
right the next three plots are the reconstructions from the neural implicit
maps found by the three tested distributed algorithms. Here DiNNO is the only
method that is able to learn a coherent map. The reconstructions were produced
by querying the optimized (and agreed upon) networks on a grid mesh of points
on the map. The last three plots show reconstructions produced from three of
the seven robots when communication is not used (training exclusively on local
data with Adam for 10 epochs). Since these robots do not have information from
other areas on the map they are only able to reconstruct regions which those
robots have traversed.

###  V-C Multi-Agent Reinforcement Learning

For the final example we show that our algorithm can be used for deep multi-
agent reinforcement learning (deep MARL). Specifically we use DiNNO for
distributed learning of a decentralized policy applied to a standard
continuous state and action, multi-robot, predator-prey problem that was first
introduced in [36]. MARL is known to be an especially hard learning task due
to the inherent nonstationarity of the environment. That is, the environment
changes during learning because the other learning agents also have evolving
policies. For background on MARL and a recent survey of deep MARL algorithms
see [37] and [38].

In our learning environment three robots must work together to pursue a faster
evader robot in the presence of stationary randomly placed obstacles, as shown
in Figure 6. Implemented in PettingZoo [39], the environment operates
according to the Actor Environment Cycle (AEC) Game model in which pursuers
make observations, take actions, and receive rewards sequentially before the
environment as a whole is updated. The pursuers are homogeneous with actions
$a=$ [none, right, left, up, down] $\in\mathbb{R}^{5}:\mathcal{A}$ and
observations $o=$[self_vel, self_pos, other_pursuers_rel_pos, evader_rel_pos,
evader_rel_vel] $\in\mathbb{R}^{12}:\mathcal{O}$. Actions are clipped to be on
the interval $[0,1]$. The evader obeys a heuristic policy, moving in the
direction opposite of the position of the nearest pursuer. However, to prevent
unfair evasion, the evader cannot propel itself outside a square of radius
$1.2$. The reward function penalizes pursuing robots based on their distance
from the evader, and pursuers receive a large positive reward for tagging
(capturing) the evader.

Figure 6: A decentralized policy rollout in the predator-prey environment.
Pursuers (red), using policies learned with DiNNO applied to PPO, attempt to
capture a faster evader (green) in the presence of obstacles (black).

To solve this problem we extend the PPO algorithm [40] with DiNNO to train a
shared, decentralized policy. PPO uses an actor-critic scheme, and with DiNNO
the policy (composed of two networks) can be optimized in a distributed way.
At consensus the robots all converge to the same policy in accordance with a
parameter sharing approach which has been shown to be effective for many MARL
problems [41, 37]. Typically the policies for parameter sharing are trained by
some centralized compute node that aggregates the experiences of each of the
robots. Applying DiNNO to PPO results in a relatively unexplored paradigm for
MARL where both training and execution are fully distributed.

In this example the actor and critic networks are feedforward ReLU networks
with 3 hidden layers of 64 neurons each. The robots simulate communication
through a fully connected graph, and robots update their policies using
individually collected data every 10 episodes. Results from this experiment
are shown in Figure 7. For each training scenario (DiNNO, DSGD, DSGT,
centralized) we show the mean of the average episodic reward achieved by the
multi-robot predator team as training progresses. To verify training quality
of each algorithm, we run the training scenario five times for each algorithm
to generate the shown statistics.

As it does in previous experiments DiNNO achieves the same average episodic
reward as a policy trained using PPO with experiences aggregated from all
three of the robots. Both DSGD and DSGT seem unable to learn a policy that
results in positive episodic reward.

Figure 7: Episodic reward (averaged across 10 episodes per network update) vs
environment time steps (summed across all episodes) for DiNNO, DSGD, DSGT, and
centralized. Shown results are averaged across 5 training runs to show
training robustness. DiNNO is the only algorithm to achieve good performance,
matching centralized performance after 10 million time steps.

##  VI Conclusion

We presented the DiNNO algorithm that enables high performance distributed
training of deep neural networks, and showcase its versatility on three
diverse multi-robot learning tasks. In comparisons to existing distributed
learning methods our algorithm consistently achieves better validation
performance and converges to performance of centrally trained models.
Directions for future work include learning neural implicit density functions
from real 3D depth data, and exploring the capabilities of DiNNO for more
complex distributed reinforcement learning tasks.

## Appendix

###  VI-A MNIST Experiments

Hyperparameters used across all four graphs were the same. DiNNO uses $B=2$,
$\rho_{0}=0.5$ increasing 0.3% per communication round, and Adam as its primal
optimizer with a log learning rate schedule (0.005 - 0.0005) for the primal
update. DSGT uses $\alpha=0.005$. DSGD uses a decaying stepsize following
$\alpha^{k+1}=\alpha^{k}(1-\mu\alpha^{k})$ where $\alpha^{0}=0.005$ and
$\mu=0.001$. All methods use batch size 64.

###  VI-B Neural Implicit Mapping Experiment

DiNNO uses $B=5$, $\rho^{0}=0.1$ increasing 0.3% per communication round,
batch size of 10,000, and Adam with a log learning rate schedule of (0.001 -
0.0001). DSGT uses $\alpha=0.001$, and a batch size of 20,000. DSGD uses
$\alpha^{0}=0.001$, $\mu=0.001$, and a batch size of 20,000.

###  VI-C Multi-agent Reinforcement Learning

For each algorithm we use the following hyperparameters: 200 steps per
episode, 2000 steps between actor/critic network updates, reward discount
factor $\gamma=0.99$, and PPO clipping parameter $0.2$. We allow each
algorithm 5 gradient steps ($B=5$) per batch of data to update actor and
critic networks. The actor learning rates for DiNNO, DSGD, and DSGT are
$0.0003$, $0.001$, and $0.01$, respectively. Only DSGT has a separate critic
learning rate, $0.00001$, due to exploding gradients otherwise. For DiNNO we
set a constant $\rho=1.0$.

## References

  * [1] S. Boyd, N. Parikh, and E. Chu, _Distributed optimization and statistical learning via the alternating direction method of multipliers_. Now Publishers Inc, 2011.
  * [2] B. Custers, A. M. Sears, F. Dechesne, I. Georgieva, T. Tani, and S. Van der Hof, _EU Personal Data Protection in Policy and Practice_ , 03 2019.
  * [3] G. Shi, W. Hönig, Y. Yue, and S.-J. Chung, “Neural-swarm: Decentralized close-proximity multirotor control using learned interactions,” in _International Conference on Robotics and Automation_. IEEE, 2020.
  * [4] G. Shi, W. Hönig, X. Shi, Y. Yue, and S.-J. Chung, “Neural-swarm2: Planning and control of heterogeneous multirotor swarms using learned interactions,” _IEEE Transactions on Robotics_ , 2021.
  * [5] Y. F. Chen, M. Liu, M. Everett, and J. P. How, “Decentralized non-communicating multiagent collision avoidance with deep reinforcement learning,” in _IEEE International Conference on Robotics and Automation_. IEEE, 2017.
  * [6] A. Yahya, A. Li, M. Kalakrishnan, Y. Chebotar, and S. Levine, “Collective robot reinforcement learning with distributed asynchronous guided policy search,” in _IEEE/RSJ International Conference on Intelligent Robots and Systems_. IEEE, 2017, pp. 79–86.
  * [7] G. Sartoretti, Y. Wu, W. Paivine, T. S. Kumar, S. Koenig, and H. Choset, “Distributed reinforcement learning for multi-robot decentralized collective construction,” in _Distributed autonomous robotic systems_. Springer, 2019, pp. 35–49.
  * [8] W. Luo and K. Sycara, “Adaptive sampling and online learning in multi-robot sensor coverage with mixture of gaussian processes,” in _2018 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2018, pp. 6359–6364.
  * [9] G. Habibi and J. P. How, “Human trajectory prediction using similarity-based multi-model fusion,” _IEEE Robotics and Automation Letters_ , vol. 6, no. 2, pp. 715–722, 2021.
  * [10] J. N. Tsitsiklis, “Problems in decentralized decision making and computation.” Massachusetts Inst of Tech Cambridge Lab for Information and Decision Systems, Tech. Rep., 1984.
  * [11] A. Nedic and A. Ozdaglar, “Distributed subgradient methods for multi-agent optimization,” _IEEE Transactions on Automatic Control_ , vol. 54, no. 1, pp. 48–61, 2009.
  * [12] W. Shi, Q. Ling, G. Wu, and W. Yin, “Extra: An exact first-order algorithm for decentralized consensus optimization,” _SIAM Journal on Optimization_ , vol. 25, no. 2, pp. 944–966, 2015.
  * [13] A. Nedić and A. Olshevsky, “Distributed optimization over time-varying directed graphs,” _IEEE Transactions on Automatic Control_ , vol. 60, no. 3, pp. 601–615, 2014.
  * [14] S. Hosseini, A. Chapman, and M. Mesbahi, “Online distributed convex optimization on dynamic networks,” _IEEE Transactions on Automatic Control_ , vol. 61, no. 11, pp. 3545–3550, 2016.
  * [15] T. Halsted, O. Shorinwa, J. Yu, and M. Schwager, “A survey of distributed optimization methods for multi-robot systems,” _arXiv preprint arXiv:2103.12840_ , 2021.
  * [16] A. Nedić, A. Olshevsky, and M. G. Rabbat, “Network topology and communication-computation tradeoffs in decentralized optimization,” _Proceedings of the IEEE_ , vol. 106, no. 5, pp. 953–976, 2018.
  * [17] T. Yang, X. Yi, J. Wu, Y. Yuan, D. Wu, Z. Meng, Y. Hong, H. Wang, Z. Lin, and K. H. Johansson, “A survey of distributed optimization,” _Annual Reviews in Control_ , 2019.
  * [18] X. Lian, C. Zhang, H. Zhang, C.-J. Hsieh, W. Zhang, and J. Liu, “Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent,” in _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , 2017, pp. 5336–5346.
  * [19] S. Pu and A. Nedić, “Distributed stochastic gradient tracking methods,” _Mathematical Programming_ , vol. 187, no. 1, pp. 409–457, 2021.
  * [20] S. Lu, X. Zhang, H. Sun, and M. Hong, “Gnsd: A gradient-tracking based nonconvex stochastic algorithm for decentralized optimization,” in _2019 IEEE Data Science Workshop (DSW)_. IEEE, 2019, pp. 315–321.
  * [21] A. Koloskova, T. Lin, S. U. Stich, and M. Jaggi, “Decentralized deep learning with arbitrary communication compression,” in _International Conference on Learning Representations_ , 2019.
  * [22] P. Di Lorenzo and S. Scardapane, “Parallel and distributed training of neural networks via successive convex approximation,” in _2016 IEEE 26th International Workshop on Machine Learning for Signal Processing (MLSP)_. IEEE, 2016, pp. 1–6.
  * [23] K. Niwa, N. Harada, G. Zhang, and W. B. Kleijn, _Edge-Consensus Learning: Deep Learning on P2P Networks with Nonhomogeneous Data_. New York, NY, USA: Association for Computing Machinery, 2020, p. 668–678. [Online]. Available: https://doi.org/10.1145/3394486.3403109
  * [24] Q. Ling, W. Shi, G. Wu, and A. Ribeiro, “Dlm: Decentralized linearized alternating direction method of multipliers,” _IEEE Transactions on Signal Processing_ , vol. 63, no. 15, pp. 4051–4064, 2015.
  * [25] T.-H. Chang, M. Hong, H.-T. Wai, X. Zhang, and S. Lu, “Distributed learning in the nonconvex world: From batch data to streaming and beyond,” _IEEE Signal Processing Magazine_ , vol. 37, no. 3, pp. 26–38, 2020.
  * [26] T.-H. Chang, M. Hong, and X. Wang, “Multi-agent distributed optimization via inexact consensus admm,” _IEEE Transactions on Signal Processing_ , vol. 63, no. 2, pp. 482–497, 2014.
  * [27] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, “Pytorch: An imperative style, high-performance deep learning library,” in _Advances in Neural Information Processing Systems 32_ , H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, Eds. Curran Associates, Inc., 2019, pp. 8024–8035.
  * [28] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _arXiv preprint arXiv:1412.6980_ , 2014.
  * [29] H. Karimi, J. Nutini, and M. Schmidt, “Linear convergence of gradient and proximal-gradient methods under the polyak-łojasiewicz condition,” in _Joint European Conference on Machine Learning and Knowledge Discovery in Databases_. Springer, 2016, pp. 795–811.
  * [30] W. Shi, Q. Ling, K. Yuan, G. Wu, and W. Yin, “On the linear convergence of the admm in decentralized consensus optimization,” _IEEE Transactions on Signal Processing_ , vol. 62, no. 7, pp. 1750–1761, 2014.
  * [31] Y. LeCun, “The mnist database of handwritten digits,” _http://yann. lecun. com/exdb/mnist/_ , 1998.
  * [32] B. Mildenhall, P. P. Srinivasan, M. Tancik, J. T. Barron, R. Ramamoorthi, and R. Ng, “Nerf: Representing scenes as neural radiance fields for view synthesis,” in _European conference on computer vision_. Springer, 2020, pp. 405–421.
  * [33] E. Sucar, S. Liu, J. Ortiz, and A. J. Davison, “imap: Implicit mapping and positioning in real-time,” 2021.
  * [34] A. Kalervo, J. Ylioinas, M. Häikiö, A. Karhu, and J. Kannala, “Cubicasa5k: A dataset and an improved multi-task model for floorplan image analysis,” in _Scandinavian Conference on Image Analysis_. Springer, 2019, pp. 28–40.
  * [35] M. Tancik, P. P. Srinivasan, B. Mildenhall, S. Fridovich-Keil, N. Raghavan, U. Singhal, R. Ramamoorthi, J. T. Barron, and R. Ng, “Fourier features let networks learn high frequency functions in low dimensional domains,” _NeurIPS_ , 2020.
  * [36] R. Lowe, Y. Wu, A. Tamar, J. Harb, P. Abbeel, and I. Mordatch, “Multi-agent actor-critic for mixed cooperative-competitive environments,” _Neural Information Processing Systems (NIPS)_ , 2017.
  * [37] J. K. Terry, N. Grammel, A. Hari, L. Santos, and B. Black, “Revisiting parameter sharing in multi-agent deep reinforcement learning,” 2021.
  * [38] G. Papoudakis, F. Christianos, L. Schäfer, and S. V. Albrecht, “Benchmarking multi-agent deep reinforcement learning algorithms in cooperative tasks,” in _Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1)_ , 2021.
  * [39] J. K. Terry, B. Black, N. Grammel, M. Jayakumar, A. Hari, R. Sulivan, L. Santos, R. Perez, C. Horsch, C. Dieffendahl, N. L. Williams, Y. Lokesh, R. Sullivan, and P. Ravi, “Pettingzoo: Gym for multi-agent reinforcement learning,” _arXiv preprint arXiv:2009.14471_ , 2020.
  * [40] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” 2017.
  * [41] J. K. Gupta, M. Egorov, and M. Kochenderfer, “Cooperative multi-agent control using deep reinforcement learning,” in _International Conference on Autonomous Agents and Multiagent Systems_. Springer, 2017, pp. 66–83.