# PushWorld: A benchmark for manipulation planning with tools and movable
obstacles

Ken Kansky1 Skanda Vaidyanath2∗ Scott Swingle3 Xinghua Lou3 Miguel Lázaro-
Gredilla3 Dileep George3
1Intrinsic Innovation
2Stanford University
3DeepMind
https://deepmind-pushworld.github.io/play/

###### Abstract

While recent advances in artificial intelligence have achieved human-level
performance in environments like Starcraft and Go, many physical reasoning
tasks remain challenging for modern algorithms. To date, few algorithms have
been evaluated on physical tasks that involve manipulating objects when
movable obstacles are present and when tools must be used to perform the
manipulation. To promote research on such tasks, we introduce PushWorld, an
environment with simplistic physics that requires manipulation planning with
both movable obstacles and tools. We provide a benchmark of more than 200
PushWorld puzzles in PDDL and in an OpenAI Gym environment. We evaluate state-
of-the-art classical planning and reinforcement learning algorithms on this
benchmark, and we find that these baseline results are below human-level
performance. We then provide a new classical planning heuristic that solves
the most puzzles among the baselines, and although it is 40 times faster than
the best baseline planner, it remains below human-level performance.

**footnotetext: During author’s internship at DeepMind.

##  1 Introduction

Figure 1: Example PushWorld puzzles. Black regions are immovable walls, and
the green object is the agent. The agent can push blue and red objects, which
can transitively push other blue and red objects. To solve a puzzle, all red
objects must move to their outlined goal positions. (left) An obstacle
prevents the agent from directly pushing the objects into their goal
positions. Instead, the agent must arrange all red objects in their
corresponding slots in the obstacle and then use the obstacle as a tool to
simultaneously push them into place. There are three distracting objects that
do not help achieve the goals. (center) The goal object is enclosed by two
obstacles that prevent pushing the object into its outlined position. The
agent’s shape is unable to push the obstacles off of the goal object, so the
agent must push the enclosure onto the peg of the object on the right wall and
move the peg upward to pull apart the enclosure. All other objects are useless
distractors. (right) A challenging puzzle. The agent must remove the obstacle
that blocks the gap in the wall. The agent is too large to pull the obstacle
away from the wall, so it must use the lower left U-shaped object as a tool
that hooks behind the obstacle to help pull it away. However, the lower right
object is an obstacle in the way of the tool, so the agent must first remove
this second obstacle. Finally, the agent cannot fit through the gap in the
wall, so the agent must use the lower right object as a tool to push the red
object through the gap to its goal.

In everyday tasks, people manipulate objects to achieve their goals, like
grasping a coffee cup from a shelf and rotating it to lift it over other cups.
If there is not enough space to lift the cup over others, people will push the
other cups out of the way, or if the cup is out of reach, people might use a
spoon as a tool to pull the handle of the cup within reach. This task is an
example of manipulation planning in the presence of tools and movable
obstacles. People are adept at solving such tasks, but reaching human
performance on these tasks remains an active research area in artificial
intelligence and robotics [?; ?; ?].

One research approach is to construct simplified environments to simulate
real-world tasks, then develop algorithms in these simplified environments,
and then generalize these algorithms for use in the real world. For example,
grid-based environments such as Sokoban have been extensively used for
evaluating physical planning algorithms [?; ?], and these algorithms are now
applied in real-world robotic tasks [?]. Indeed, environments like Sokoban
have contributed to the development of a variety of intelligent algorithms,
including tree search with pattern database heuristics [?] and using features
as a high-level road map for search [?]. The deep reinforcement learning
community also continues to develop Sokoban solvers in both model-free [?] and
model-based [?] settings.

Despite its popularity as a planning benchmark, Sokoban does not involve
manipulation planning or tool use. Moreover, existing environments for robotic
manipulation planning typically do not include both tools and movable
obstacles [?; ?; ?; ?]. Therefore, we introduce PushWorld, a simplified
environment for testing manipulation planning in the presence of both tools
and movable obstacles. As shown in Figure 1, PushWorld is a two-dimensional
physical environment with finite, discrete states. Like Sokoban, it is fully
observable, deterministic, and static. However, PushWorld additionally allows
objects to have diverse shapes and sizes, and multiple objects can be pushed
simultaneously. These differences from Sokoban introduce many new challenges:
a PushWorld puzzle may require pushing multiple obstacles out of the way
simultaneously, or it may require assembling a tool to push another object.
Section 2 elaborates on the challenges that PushWorld offers.

PushWorld is also usable as a sparse-reward and hierarchical reinforcement
learning (RL) environment that is simpler than many existing environments that
involve robots [?] or continuous action spaces [?]. Three-dimensional RL
environments have been created to test tool use [?], but PushWorld
demonstrates that non-trivial tool use problems are plentiful in only two
dimensions.

Our contributions are as follows: First, we introduce PushWorld, a simple
physical environment suitable for research in classical planning,
reinforcement learning, combined task and motion planning for robotics, and
cognitive science. Second, we benchmark several state-of-the-art classical
planning and deep reinforcement learning algorithms on PushWorld as baseline
evaluations. Third, we provide a novel planning heuristic that outperforms all
baselines. Finally, we provide an open-source
package111https://github.com/deepmind/pushworld/ that includes an OpenAI Gym
environment, a benchmarking dataset of 223 hand-designed PushWorld puzzles
with scripts to generate PDDL and SAS+ [?], and the novel planning heuristic.

##  2 The PushWorld Environment

As shown in Figure 1, PushWorld is a simplified two-dimensional physical
environment in which the goal is to push one or more objects into target
positions. Objects have rigid shapes that occupy sets of discrete positions,
which cannot intersect with the occupied positions of other objects. Objects
can translate but cannot rotate.

One object is designated as the agent, which can perform one of four actions
at each discrete time step: moving up, down, left, or right by a single
discrete position. The agent can push other objects in the direction of its
movement, and there is no limit on the number of objects that the agent can
push simultaneously. Any object can also push any other, so the agent can use
objects as tools to indirectly push other objects. Objects have no momentum,
so they immediately stop moving after the agent pushes them, as if all objects
rest on a high-friction surface. Immovable walls also occupy discrete
positions, and actions are forbidden that would push any object into a wall’s
position. Figure 1 illustrates three example puzzles, and Appendix A.1
provides a formal definition of a PushWorld puzzle.

In spite of its simplified physics, solving PushWorld puzzles requires diverse
skills:

  * •

Path Planning involves finding a collision-free path to move an object from
one position to another, considering walls and object shapes.

  * •

Manipulation Planning involves exploring alternative ways to push an object
along a desired path. Some puzzles require the agent to preemptively position
itself so that it can move from pushing in one direction to pushing in a
different direction in the future.

  * •

Movable Obstacles require the agent to decide between finding a path around an
obstacle or moving the obstacle out of the way. Some obstacles introduce
choices in which freeing one path blocks another, and choices can be
irreversible. Due to object shapes and the inability to pull objects,
obstacles can also be “parasitic”: once pushed against an object, the agent
can never separate the obstacle from the object.

  * •

Tool Use. To move an object into a desired position, the agent may need to use
one or more objects as tools to indirectly push the target object: the agent
pushes a tool object, which simultaneously pushes the target or another tool.
In some puzzles the agent must assemble a tool by pushing multiple objects
together.

  * •

Prioritizing Multiple Goals. Moving an object into its goal position too soon
may prevent achieving another goal, and some puzzles require achieving
multiple goals simultaneously by preemptively arranging objects and pushing
them all at once.

Within a single puzzle, an object can change its role over time, sometimes
functioning as a tool, as an obstacle, or as a manipulation target. PushWorld
therefore requires flexibly applying the skills above to adapt to changing
situations.

As a reinforcement learning environment, PushWorld also poses several
challenges:

  * •

Sparse rewards: The agent only receives a reward when it pushes an object into
its goal position. Compared to Sokoban, PushWorld puzzles tend to require more
actions to achieve all goals, which makes learning from rewards more
challenging.

  * •

Credit assignment: Judging which actions contributed to causing an outcome can
hasten RL training [?] and can improve exploration [?]. In some PushWorld
puzzles it is possible to reach states from which it is no longer possible to
achieve all goals, so RL agents can avoid these states by identifying their
causes. A model-based approach can answer counterfactual questions that assist
with credit assignment [?].

  * •

Exploration: Efficient exploration is one strategy to counteract sparse
rewards [?; ?]. An intelligent exploration algorithm can identify which parts
of the state space to explore and which to ignore [?]. The agent must also
explore safely to avoid states from which goals are no longer achievable [?].

  * •

Hierarchy: Humans can decompose PushWorld puzzles into subgoals to remove
obstacles, retrieve tools, navigate to desired positions, and more. This
decomposition suggests that hierarchical approaches [?; ?; ?; ?] may be
necessary to achieve human-level performance.

##  3 Related Environments

| Sokoban | Sliding Block Puzzles | Robotic Manipulation | Grid Path Planning | PushWorld
---|---|---|---|---|---
Varied Object Shapes | No | Yes | Yes | Yes | Yes
Path Planning | Yes | No | Yes | Yes | Yes
Manipulation Planning | No | No | Yes | No | Yes
Movable Obstacles | Yes | Yes | Sometimes | No | Yes
Tool Use | No | No | Sometimes | No | Yes
Multiple Goals | Yes | Yes | Sometimes | No | Yes
Simplified Physics | Yes | Yes | No | Yes | Yes
Table 1: Properties of PushWorld vs. related environments as a planning
benchmark.

###  3.1 Sokoban

A Sokoban puzzle (Fig. 2A) requires an agent to push one or more boxes into
goal positions. Unlike PushWorld, the agent can only push one box at a time,
and all boxes have the same shape. By allowing objects to have different
shapes and allowing the agent to push multiple objects at once, PushWorld adds
complexity to path planning and obstacle interactions, and it introduces the
use of tools. Like Sokoban, PushWorld puzzles include the potential for
deadlocks, in which action effects are irreversible, and for conflicts, in
which achieving a goal in the wrong order may prevent achieving other goals.

###  3.2 Sliding Block Puzzles

Sliding block puzzles involve sliding one or more 2D objects into goal
positions, and well-known examples include Klotski (Fig. 2B) and the fifteen
puzzle. Objects can have varying shapes, and actions can directly move all
objects. Challenging sliding block puzzles typically confine all objects
without much free space to move. In PushWorld, actions can directly move only
one object, the agent, so limited free space is not necessary to construct
challenging puzzles.

###  3.3 Robotic Manipulation Among Movable Obstacles

A robotic manipulation task involves moving one or more objects into goal
states, typically by executing alternating phases of transit, during which the
robot moves without moving any other object, and transfer, during which the
robot controls the movement of one or more objects [?]. Tasks may require
moving objects out of the way to make space for moving the robot or for
manipulating another object [?; ?], and in some tasks the robot must push
objects out of the way because they are too heavy to lift or too large to
grasp [?]. While PushWorld has a smaller state space than these robotic tasks,
PushWorld nevertheless requires similar planning capabilities to move
obstacles and to manipulate objects via pushing. Some robotic tasks involve
tool use [?] (Fig. 2C), but robotic tasks involving both tool use and movable
obstacles have received less attention. PushWorld therefore presents a novel
integration of planning challenges.

###  3.4 Grid-Based Path Planning

Some robotic path planning tasks can be represented by discretizing the space
of robot configurations into a grid (Fig. 2D) and then annotating which
discrete configurations are collision-free [?; ?]. Planners must then solve
for a collision-free path in the grid from a starting configuration to a goal
configuration. Path planning is a sub-problem in PushWorld, so techniques for
grid-based path planning, like hierarchical path planning [?], are applicable.

Figure 2: Examples of environments related to PushWorld. (A) Sokoban [?], (B)
Klotski, (C) Robotic Manipulation [?], (D) Grid-Based Path Planning [?].

##  4 The Recursive Graph Distance Planner

Existing classical planning algorithms [?; ?; ?] are effective at logistics
problems involving the constrained movements of multiple objects, but they are
not optimized for planning motion paths. Simultaneously, existing algorithms
for robotics [?; ?; ?] are designed for efficient path planning but are not
typically applied in environments that require the diversity of obstacle and
tool manipulations present in PushWorld puzzles. To reduce the gap between the
specializations of existing algorithms, we introduce the recursive graph
distance (RGD) heuristic, which optimizes a variation of the context-enhanced
additive heuristic (CEA) [?] to achieve faster path planning. We then combine
RGD with the novelty heuristic [?] to explore efficiently in a greedy best-
first (GBF) search.

###  4.1 Planning Task Definition

A multi-valued planning task $\Pi=(V,s_{0},s_{*},O)$ contains a set of
variables $V$, an initial state $s_{0}$, a goal $s_{*}$, and a set of
operators $O$ that map states to other states as a means to achieve the goal
from the initial state. A state $s$ is a mapping from variables $v\in V$ to
values $d\in D_{v}$ in $v$’s domain $D_{v}$. For convenience we will write
$(v,d)\in s$ to indicate that $s(v)=d$. A partial state is a mapping of a
subset of variables in $V$ to values in their respective domains. In
PushWorld, variables are typically object positions, although other state
representations are possible. The goal $s_{*}$ is a partial state, and a state
$s$ satisfies the goal whenever $s(v)=s_{*}(v)$ for all variables $v$ in the
domain of $s_{*}$.

Operators can change one state into another, and each operator $o$ has a set
of preconditions $\text{pre}(o)$ and a set of effects $\text{eff}(o)$, both of
which are partial states. When a state $s$ satisfies an operator’s
preconditions, such that $s(v)=d$ for all $(v,d)\in\text{pre}(o)$, then
applying the operator changes the state to $s^{\prime}$ in which
$s^{\prime}(v)=d$ for all $(v,d)\in\text{eff}(o)$ and in which
$s^{\prime}(v)=s(v)$ for all other variables $v$ that are not in the domain of
$\text{eff}(o)$.

###  4.2 The Context-Enhanced Additive Heuristic

The CEA heuristic estimates the number of operators needed to change state $s$
into a state in which the goal is satisfied. For each variable $v$, the
heuristic uses a domain transition graph $\text{DTG}(v)$, which is a labeled
directed multigraph with vertex set $D_{v}$ and an edge $(d,d^{\prime})$
associated with each operator $o$ that has both $(v,d)\in\text{pre}(o)$ and
$(v,d^{\prime})\in\text{eff}(o)$. The graph may contain parallel edges if
multiple operators can change the value of $v$ from $d$ to $d^{\prime}$, and
edges are labeled with the preconditions $\text{pre}(o)$ of their associated
operators.

The CEA heuristic computes the sum of estimates of the number of operators
needed to change the value $s(v)$ of each goal variable $v$ into a state
$s^{\prime}$ in which $s^{\prime}(v)=s_{*}(v)$. Details of this computation
are available in [?] and are not repeated here, but to estimate the cost of
changing a variable $v$ from one value to a target value, the heuristic
searches for a minimum-cost path in $\text{DTG}(v)$.

Figure 3: The objects and the goal position remain in the corners as the
puzzle scales with the distance $N$ between the objects. A greedy best-first
search with the CEA heuristic [?] solves this puzzle in $O(N^{3})$ time, and
the same search with the RGD heuristic solves this puzzle in $O(N^{2})$ time.

If variables correspond to object positions in PushWorld, then in the puzzle
shown in Fig 3, computing the CEA heuristic for a single state has a worst-
case complexity of $O(N^{2})$ with respect to the distance $N$ between objects
in the puzzle. This complexity results from the search for minimum-cost paths
in DTGs, which each include $O(N^{2})$ nodes for all object positions. A GBF
search computes the heuristic in every explored state, resulting in $O(N^{3})$
time to solve the puzzle.

###  4.3 Recursive Graph Distance

The RGD heuristic modifies the CEA heuristic to cache approximate costs of
minimum-cost paths, and a GBF search with this heuristic requires $O(N^{2})$
time to solve the puzzle in Fig 3. Let $h^{\text{RGD}}(s)$ denote the
estimated number of operators needed to achieve the goal $s_{*}$ starting from
state $s$, and let $h_{v}^{\text{RGD}}(d|s)$ denote the estimated number of
operators needed to change state $s$ into any state $s^{\prime}$ in which
$s^{\prime}(v)=d$. The RGD heuristic estimates the cost to the goal by summing
independent estimates of the cost to reach the goal value of each variable in
the goal’s domain:

$h^{\text{RGD}}(s)=\sum_{v\in\text{dom}(s_{*})}{h_{v}^{\text{RGD}}(s_{*}(v)|s)}$
(1)

Next, let $L_{v}(d,d^{\prime})$ denote the number of edges in a minimum-length
path from $d$ to $d^{\prime}$ in $\text{DTG}(v)$, and let
$\text{Succ}_{v}(d)=\\{(d^{\prime},p)\\}$ denote the set of end nodes
$d^{\prime}$ and corresponding labeled preconditions $p$ of all outgoing edges
from $d$ in $\text{DTG}(v)$. To compute $h_{v}^{\text{RGD}}(d|s)$, we search
over all operators that can modify the value of $v$ in $s$ to find which
operator minimizes the sum of the remaining graph distance to $d$ after
applying the operator, plus $1$ for the cost of the operator, plus the
estimated cost of the operator’s preconditions:

$\begin{split}h_{v}^{\text{RGD}}(d|s)=\min_{(d^{\prime},p)\in\text{Succ}_{v}(s(v))}\biggl{[}L_{v}&(d^{\prime},d)+1+\\\
&\sum_{(v_{i},e_{i})\in
p}{h^{\text{RGD}}_{v_{i}}(e_{i}|s)}\biggr{]}\end{split}$ (2)

We also apply two exceptions: whenever $s(v)=d$, we return $0$ from
$h_{v}^{\text{RGD}}(d|s)$, and to avoid infinite recursion, we return infinity
whenever $v$ appears in any recursive parent call of
$h_{v}^{\text{RGD}}(d|s)$. Figure 4 illustrates calculating the RGD heuristic
in PushWorld. Domain transition graphs are computed once for all variables at
the beginning of each search that uses this heuristic.

This formulation allows $h^{\text{RGD}}(s)$ to cache the shortest path lengths
$L_{v}(d,d^{\prime})$. Continuing with the example in Fig 3, the minimum-
length path to move the agent to the red object requires $O(N^{2})$ time to
compute in the DTG for the agent’s position, and likewise for the path from
the red object’s position to its goal. In a GBF search, expanded states can
reuse the minimum-length paths from those computed while expanding preceding
states. Consequently, in this example the heuristic requires $O(N^{2})$ time
when it is first computed, and in all subsequent expanded states it requires
$O(1)$ time, resulting in a total time of $O(N^{2})$ to solve the puzzle in a
GBF search.

Appendix A.3 provides pseudocode for RGD and includes additional
optimizations. We also provide an open-source C++
implementation.222https://github.com/deepmind/pushworld/

Figure 4: An illustration of the RGD heuristic. The red object must move two
positions left to reach its goal position, but the actor (green) is unable to
push it due to the walls. RGD detects that the blue object can push the red
object, and it recursively discovers that the actor can move into a position
to push the blue object along a path toward the red object. RGD finds the
minimum-length paths to move objects into either a pushing position or a goal
position. The actual number of actions required to achieve the goal is 12, but
RGD estimates 4+3+2 = 9 actions because it only considers the cost of pushing
the blue object by one position; it ignores the cost of moving the actor to
continue pushing the blue object along its path. To estimate the minimum
number of actions to achieve a goal, RGD explores all pushing directions, all
possible pushers, and all relative pushing positions between objects.

###  4.4 Combining with the Novelty Heuristic

The novelty heuristic [?] promotes exploration from states that contain
previously unseen partial states. For some set of generated states
$S=\\{s_{i}\\}$, the novelty $w(s)$ of a newly generated state $s$ is the size
of the smallest non-empty set of variables $V$ such that there is no state
$s_{i}\in S$ for which $s(v)=s_{i}(v)$ for all $v\in V$. In other words, $V$
defines a set of variables whose combined values occur in $s$ but have not yet
occurred in any other generated state, and $w(s)$ measures the size of the
smallest such set.

When variables correspond to the positions of objects in a PushWorld puzzle,
then $w(s)=1$ if $s$ is the first state in which some object has a position
that has not occurred in any previously generated state. If $s$ is the first
state in which some pair of objects have positions that have not jointly
occurred in any previously generated state, then $w(s)=2$.

We define the Novelty+RGD planner as a GBF search that prioritizes states
using a lexicographic combination of heuristics: first prioritize states that
minimize $w(s)$, and whenever states have the same novelty, then prioritize
states that minimize $h^{\text{RGD}}(s)$. In PushWorld this planner defines a
state $s$ as a mapping from object identifiers to their 2D positions, and
because computing novelty is combinatorially expensive in the number of
objects, the planner limits the maximum novelty to 3 to reduce the novelty
computation to polynomial time.

##  5 Evaluation

###  5.1 Benchmarking Dataset

To compare how efficiently different algorithms can solve PushWorld puzzles,
we provide a set of 223 hand-designed puzzles, which are available to play
online and download.333https://deepmind-pushworld.github.io/play/ Puzzles are
organized into four levels of difficulty, such that people can typically solve
Level 1 puzzles within seconds, Level 2 puzzles within a few minutes, Level 3
puzzles within 15 minutes, and longer for Level 4. From levels 1 to 4, the
dataset contains 68, 74, 67, and 14 puzzles, respectively.

###  5.2 Classical Planning

To evaluate the performance of existing classical planners, we aimed to cover
multiple classes of search heuristics, including delete-relaxations [?],
causal graphs [?], landmarks [?], and state novelty [?], as well as portfolios
containing multiple planners [?; ?]. We preferred the winning planners from
International Planning Competitions (IPC), resulting in the following
selection:

  * •

Fast Forward (FF) [?] is a seminal planning algorithm that relies on a delete
relaxation heuristic.

  * •

Fast Downward (FD) [?] introduced DTGs and the causal graph heuristic (winner
of IPC 2004 classical track). We selected the most recent implementation,
which uses the CEA heuristic.

  * •

LAMA [?] introduced the landmark heuristic (winner of IPC 2008 sequential
satisficing track).

  * •

Best-First Width Search (BFWS) [?] uses the novelty heuristic, which
prioritizes states that contain novel substates (winner of IPC 2018 agile
track).

  * •

Fast Downward Stone Soup (FDSS) [?] is a portfolio of planners (winner of IPC
2018 satisficing track).

  * •

Saarplan [?] is another portfolio of planners (runner-up of IPC 2018 agile
track).

Multiple versions of these planners are available, and Appendix A.2 explains
the versions we selected. In particular, we found that among the available
variations of BFWS, Dual BFWS performed best in PushWorld. All planners are
implemented in C or C++.

Some planners are optimized for different problem representations, like SAS+
[?] or PDDL [?]. Translating between representations can exceed the search
time in PushWorld puzzles. Because these representations make no assumptions
that are specific to PushWorld’s dynamics, we omitted translation times from
the reported planning times to avoid penalizing planners for their problem
representations.

###  5.3 Model-Free Reinforcement Learning

We selected two seminal, diverse deep RL algorithms to evaluate on PushWorld:
Deep Q-Network (DQN) [?], an off-policy, value-based algorithm, and Proximal
Policy Optimization (PPO) [?], an on-policy, policy-gradient algorithm. We
chose these algorithms because they are widely used, easy to implement, and
have shown competitive performance in environments like Atari games and
continuous robotic control. Furthermore, recent work [?] has shown that model-
free approaches can be successful at solving tasks that may seem more suited
for model-based approaches. Appendix A.4 provides details on the network
architectures and hyperparameter settings we used in experiments.

We implemented PushWorld as an OpenAI Gym environment [?] with a reward
function similar to that in [?]: the agent receives +1 reward for moving an
object onto its goal position and -1 for moving off of its goal position. The
agent receives +10 reward for achieving all goals in a puzzle. In contrast to
[?], the agent receives -0.01 reward for each action, since solutions to
PushWorld puzzles tend to be longer than solutions to Sokobans.

DQN and PPO require more training data than the 223 hand-crafted puzzles can
provide, so we procedurally generated sets of puzzles to provide sufficient
training data. To measure how puzzle parameters affect RL algorithms, we
generated the following puzzle sets:

  * •

Base: The puzzle grid size is 5x5 discrete positions. Each puzzle contains
three walls, one object with a goal position, and one movable obstacle. All
objects have 1x1 shape, including the walls and the agent. All objects and
goal positions are placed randomly.

  * •

Larger Puzzle Sizes: Same as Base, but the puzzle width and height are
independently uniformly sampled between 5 and 10, inclusive.

  * •

More Walls: Same as Base, but the number of walls is uniformly sampled between
3 and 5, inclusive.

  * •

More Obstacles: Same as Base, but with two movable obstacle objects.

  * •

More Shapes: Same as Base, but the shape of the agent, obstacle, and target
object are independently uniformly sampled polyominoes with 1, 2, or 3
squares.

  * •

Multiple Goals: Same as Base, but with two objects with goal positions. To
distinguish these objects, one has shape 1x1 and the other 1x2.

  * •

All: All previous variations combined simultaneously. Each puzzle’s width and
height are uniformly sampled between 5 and 10, the number of walls is
uniformly sampled between 3 and 5, the number of movable obstacles is
uniformly sampled between 1 and 2, the number of goal objects is uniformly
sampled between 1 and 2, and object shapes are uniformly sampled polyominoes
with between 1 and 3 squares.

These puzzles are limited to small sizes to mitigate sparse reward; an RL
agent must initially be able to solve some puzzles by chance so that it can
learn how to achieve goals. Because these puzzle sets are generally easier
than the hand-designed Level 1 puzzles, we designate them as Level 0.

For each Level 0 variation, we generated a set of 2,000 puzzles for training
and 200 more for testing. This puzzle generating process employed a classical
planner to guarantee that all generated puzzles were solvable. To reduce bias
during both training and testing, we augmented the puzzles by including every
combination of 90 degree rotation and flipping, resulting in 16,000 augmented
puzzles per set. For variations involving puzzle size changes, we added walls
around all puzzles to pad them up to the largest puzzle’s size. The position
of each smaller puzzle was randomized within the padded walls, providing
additional augmentation.

To evaluate performance on the hand-designed puzzles, we trained each
algorithm on all puzzles in the Level 1 set and included the same
augmentations of rotation, flipping, and padding puzzles up to the largest
size. For this evaluation we did not pretrain on the Level 0 puzzle sets. To
give the algorithms the highest chance of learning from the training puzzles,
we did not use a train/test split because the results showed that few training
puzzles were solved. Due to their low performance on Level 1, we did not
evaluate PPO or DQN on Levels 2 – 4.

##  6 Experimental Results

Figure 5: Number of solved puzzles vs. planning time. Planning time is
measured independently for each puzzle.

###  6.1 Classical Planning Results

We evaluated all classical planners on the set of hand-designed puzzles. All
planners ran on a 2.25 GHz AMD CPU, subject to a 30 GB memory limit and a 30
minute time limit per puzzle. Results are shown in Figure 5. Among existing
classical planners, FDSS performed the best, solving 61.4% of the puzzles.
Despite its relative simplicity, Fast Forward was competitive with most other
planners. The RGD plot shows the performance of a greedy best-first search
using the RGD heuristic, which outperformed all existing classical planners
except FDSS. The Novelty+RGD planner performed best, solving 69.1% of the
puzzles. In fact, Novelty+RGD solved as many puzzles in 45 seconds as FDSS
solved in 30 minutes, which amounts to a 40x speed improvement.

###  6.2 Deep Reinforcement Learning Results

#### PPO and DQN on Level 0

Table 2 shows PPO and DQN results on Level 0. PPO’s training accuracy is
significantly higher than its test accuracy, indicating possible overfitting.
We found that this phenomenon persisted despite tuning the network capacity.
By contrast, DQN shows a smaller discrepancy between training and testing,
suggesting that DQN’s learned value function is more transferable across
puzzles than PPO’s learned policy.

Both PPO and DQN perform best on Larger Puzzle Sizes. We believe this is
because these puzzles have the lowest spatial density of walls and movable
obstacles. Both algorithms perform worst on Multiple Goals. We suspect this is
because puzzles with multiple goals are less likely to solve by chance and
because some puzzles become unsolvable if goals are achieved in the wrong
order. Both algorithms also have low performance on All, indicating that the
diversity in puzzles hinders both training and transfer to unseen puzzles from
the same distribution.

| PPO | DQN
---|---|---
Level 0 Puzzle Set | Train | Test | Train | Test
Base | 88.5% | 40.4% | 24.9% | 19.5%
Larger Puzzle Sizes | 93.2% | 71.5% | 61.2% | 61.8%
More Walls | 77.9% | 23.3% | 18.5% | 13.1%
More Obstacles | 64.7% | 21.4% | 17.6% | 13.1%
More Shapes | 74.5% | 24.2% | 23.9% | 17.1%
Multiple Goals | 5.3% | 1.1% | 0.9% | 0.2%
All | 21.2% | 5.7% | 11.4% | 10.1%
Table 2: The percentage of training and testing puzzles solved by PPO and DQN.

#### PPO and DQN on Level 1

Compared to the Level 0 puzzles, Level 1 puzzles are generally larger, include
more object shapes, and involve more complex arrangements of obstacles, walls,
and goals. Without pretraining on Level 0 puzzle sets, we trained both PPO and
DQN on all Level 1 puzzles for 350 million steps and measured their training
performance, the percentage of Level 1 puzzles they could solve. DQN converged
to solving less than 1% of the puzzles, and PPO converged to solving 6% of the
puzzles. We believe this low performance is in part due to the low probability
of solving most Level 1 puzzles with the initial policy, resulting in sparse
positive rewards.

##  7 Conclusions & Discussion

We present PushWorld, a manipulation planning environment involving tools and
movable obstacles. Compared to similar environments such as Sokoban, we
believe PushWorld presents many novel challenges relevant for research
communities such as classical planning, reinforcement learning, and robotics.
In our baseline evaluations we introduced the RGD+Novelty planner, which
outperformed many state-of-the-art classical planners and RL algorithms on
PushWorld, yet is still below human-level performance. To close this gap, we
hope PushWorld can become a standard benchmark to drive forward research in RL
and physical planning.

For future work, we believe in exploring the following directions: First,
hierarchical planning [?] and factorized task and motion planning [?] may be
effective for PushWorld. Second, model-based RL algorithms such as AlphaZero
[?] and MuZero [?] can potentially discover new and interesting heuristics.
Third, many RL research areas could improve both training efficiency and
generalization to unseen PushWorld puzzles: hierarchical RL, credit assignment
[?] using a causal model [?], efficient exploration, multi-task learning [?],
meta-learning [?], and curriculum learning [?]. Finally, a cognitive science-
inspired analysis of how people solve PushWorld puzzles may help develop
better algorithms.

The PushWorld environment and benchmark puzzles are open source, and we
encourage the community to expand on this benchmark. We hope that lessons
learned in PushWorld will translate into more intelligent systems in real-
world applications.

## References

  * [Andreas et al., 2016] Jacob Andreas, Dan Klein, and Sergey Levine.  Modular multitask reinforcement learning with policy sketches, 2016.
  * [Andrychowicz et al., 2017] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba.  Hindsight experience replay, 2017.
  * [Bäckström and Nebel, 1995] Christer Bäckström and Bernhard Nebel.  Complexity results for sas+ planning.  Computational Intelligence, 11(4):625–655, 1995.
  * [Bellemare et al., 2016] Marc G. Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos.  Unifying count-based exploration and intrinsic motivation, 2016.
  * [Brockman et al., 2016] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym, 2016.
  * [Dogar and Srinivasa, 2011] Mehmet Dogar and Siddhartha Srinivasa.  A framework for push-grasping in clutter.  Robotics: Science and systems VII, 1, 2011.
  * [Dor and Zwick, 1999] Dorit Dor and Uri Zwick.  Sokoban and other motion planning problems.  Computational Geometry, 13(4):215–228, 1999.
  * [Fickert et al., 2018] Maximilian Fickert, Daniel Gnad, Patrick Speicher, and Jörg Hoffmann.  Saarplan: Combining saarland’s greatest planning techniques.  IPC2018–Classical Tracks, pages 10–15, 2018.
  * [Finn et al., 2017] Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks, 2017\.
  * [Florensa et al., 2017] Carlos Florensa, Yan Duan, and Pieter Abbeel.  Stochastic neural networks for hierarchical reinforcement learning.  arXiv preprint arXiv:1704.03012, 2017.
  * [Frances et al., 2018] Guillem Frances, Hector Geffner, Nir Lipovetzky, and Miquel Ramirez.  Best-first width search in the ipc 2018: Complete, simulated, and polynomial variants.  IPC2018–Classical Tracks, pages 22–26, 2018.
  * [Garg et al., 2022] Divyansh Garg, Skanda Vaidyanath, Kuno Kim, Jiaming Song, and Stefano Ermon.  Lisa: Learning interpretable skill abstractions from language, 2022.
  * [Garrett et al., 2020] Caelan Reed Garrett, Tomás Lozano-Pérez, and Leslie Pack Kaelbling.  Pddlstream: Integrating symbolic planners and blackbox samplers via optimistic adaptive planning.  In Proceedings of the International Conference on Automated Planning and Scheduling, volume 30, pages 440–448, 2020.
  * [Ghallab et al., 1998] M. Ghallab, A. Howe, C. Knoblock, D. Mcdermott, A. Ram, M. Veloso, D. Weld, and D. Wilkins.  PDDL - The Planning Domain Definition Language.  Technical Report, Tech. Rep., 1998.
  * [Grinsztajn et al., 2021] Nathan Grinsztajn, Johan Ferret, Olivier Pietquin, Philippe Preux, and Matthieu Geist.  There is no turning back: A self-supervised approach for reversibility-aware reinforcement learning, 2021.
  * [Guez et al., 2019] Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Théophane Weber, David Raposo, Adam Santoro, Laurent Orseau, et al.  An investigation of model-free planning.  In International Conference on Machine Learning, pages 2464–2473. PMLR, 2019.
  * [Hamrick et al., 2020] Jessica B Hamrick, Abram L Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Buesing, Petar Veličković, and Théophane Weber.  On the role of planning in model-based deep reinforcement learning.  arXiv preprint arXiv:2011.04021, 2020.
  * [Harabor and Botea, 2008] Daniel Harabor and Adi Botea.  Hierarchical path planning for multi-size agents in heterogeneous environments.  In 2008 IEEE Symposium On Computational Intelligence and Games, pages 258–265. IEEE, 2008.
  * [Harutyunyan et al., 2019] Anna Harutyunyan, Will Dabney, Thomas Mesnard, Mohammad Gheshlaghi Azar, Bilal Piot, Nicolas Heess, Hado P van Hasselt, Gregory Wayne, Satinder Singh, Doina Precup, et al.  Hindsight credit assignment.  Advances in neural information processing systems, 32, 2019.
  * [Haslum et al., 2007] Patrik Haslum, Adi Botea, Malte Helmert, Blai Bonet, Sven Koenig, et al.  Domain-independent construction of pattern database heuristics for cost-optimal planning.  In AAAI, volume 7, pages 1007–1012, 2007.
  * [Helmert and Geffner, 2008] Malte Helmert and Héctor Geffner.  Unifying the causal graph and additive heuristics.  In ICAPS, pages 140–147, 2008.
  * [Helmert, 2006] Malte Helmert.  The fast downward planning system.  Journal of Artificial Intelligence Research, 26:191–246, 2006.
  * [Hoffman et al., 2020] Matt Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, et al.  Acme: A research framework for distributed reinforcement learning.  arXiv preprint arXiv:2006.00979, 2020.
  * [Hoffmann and Nebel, 2001] Jörg Hoffmann and Bernhard Nebel.  The ff planning system: Fast plan generation through heuristic search.  Journal of Artificial Intelligence Research, 14:253–302, 2001.
  * [Hong et al., 2018] Zhang-Wei Hong, Tzu-Yun Shann, Shih-Yang Su, Yi-Hsiang Chang, Tsu-Jui Fu, and Chun-Yi Lee.  Diversity-driven exploration strategy for deep reinforcement learning.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [Kaelbling and Lozano-Pérez, 2010] Leslie Pack Kaelbling and Tomás Lozano-Pérez.  Hierarchical planning in the now.  In Workshops at the Twenty-Fourth AAAI Conference on Artificial Intelligence, 2010.
  * [Kansky et al., 2017] Ken Kansky, Tom Silver, David A Mély, Mohamed Eldawy, Miguel Lázaro-Gredilla, Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George.  Schema networks: Zero-shot transfer with a generative causal model of intuitive physics.  In International conference on machine learning, pages 1809–1818. PMLR, 2017.
  * [Lecun et al., 1998] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324, 1998.
  * [Li et al., 2019] Alexander C. Li, Carlos Florensa, Ignasi Clavera, and Pieter Abbeel.  Sub-policy adaptation for hierarchical reinforcement learning, 2019.
  * [Lipovetzky and Geffner, 2017] Nir Lipovetzky and Hector Geffner.  Best-first width search: Exploration and exploitation in classical planning.  In Thirty-First AAAI Conference on Artificial Intelligence, 2017\.
  * [Liu et al., 2020] Evan Zheran Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn.  Decoupling exploration and exploitation for meta-reinforcement learning without sacrifices, 2020.
  * [Mesnard et al., 2020] Thomas Mesnard, Théophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Tom Stepleton, Nicolas Heess, Arthur Guez, Éric Moulines, Marcus Hutter, Lars Buesing, and Rémi Munos.  Counterfactual credit assignment in model-free reinforcement learning, 2020.
  * [Mnih et al., 2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  nature, 518(7540):529–533, 2015.
  * [Nachum et al., 2018] Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine.  Data-efficient hierarchical reinforcement learning, 2018.
  * [Narvekar et al., 2020] Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E Taylor, and Peter Stone.  Curriculum learning for reinforcement learning domains: A framework and survey.  arXiv preprint, 2020.
  * [Richter and Westphal, 2010] Silvia Richter and Matthias Westphal.  The lama planner: Guiding cost-based anytime planning with landmarks.  Journal of Artificial Intelligence Research, 39:127–177, 2010.
  * [Schrittwieser et al., 2020] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al.  Mastering atari, go, chess and shogi by planning with a learned model.  Nature, 588(7839):604–609, 2020.
  * [Schulman et al., 2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347, 2017.
  * [Seipp and Röger, 2018] Jendrik Seipp and Gabriele Röger.  Fast downward stone soup 2018.  IPC2018–Classical Tracks, pages 72–74, 2018.
  * [Shoham and Schaeffer, 2020] Yaron Shoham and Jonathan Schaeffer.  The fess algorithm: A feature based approach to single-agent search.  In 2020 IEEE Conference on Games (CoG), pages 96–103. IEEE, 2020\.
  * [Silver et al., 2017] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al.  Mastering chess and shogi by self-play with a general reinforcement learning algorithm.  arXiv preprint arXiv:1712.01815, 2017.
  * [Siméon et al., 2004] Thierry Siméon, Jean-Paul Laumond, Juan Cortés, and Anis Sahbani.  Manipulation planning with probabilistic roadmaps.  The International Journal of Robotics Research, 23(7-8):729–746, 2004.
  * [Stilman and Kuffner, 2008] Mike Stilman and James Kuffner.  Planning among movable obstacles with artificial constraints.  The International Journal of Robotics Research, 27(11-12):1295–1307, 2008.
  * [Stilman et al., 2007] Mike Stilman, Jan-Ullrich Schamburek, James Kuffner, and Tamim Asfour.  Manipulation planning among movable obstacles.  In Proceedings 2007 IEEE international conference on robotics and automation, pages 3327–3332. IEEE, 2007.
  * [Tamar et al., 2016] Aviv Tamar, Sergey Levine, and Pieter Abbeel.  Value iteration networks.  CoRR, 2016.
  * [Team et al., 2021] Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, et al.  Open-ended learning leads to generally capable agents.  arXiv preprint arXiv:2107.12808, 2021.
  * [Thrun, 2002] Sebastian Thrun.  Probabilistic robotics.  Communications of the ACM, 45(3):52–57, 2002.
  * [Todorov et al., 2012] Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026–5033. IEEE, 2012.
  * [Toussaint et al., 2018] Marc A Toussaint, Kelsey Rebecca Allen, Kevin A Smith, and Joshua B Tenenbaum.  Differentiable physics and stable modes for tool-use and manipulation planning.  Robotics: Science and Systems 2018, 2018.
  * [Yan et al., 2022] Dong Yan, Jiayi Weng, Shiyu Huang, Chongxuan Li, Yichi Zhou, Hang Su, and Jun Zhu.  Deep reinforcement learning with credit assignment for combinatorial optimization.  Pattern Recognition, 124:108466, 2022.

##  Appendix A Appendix

###  A.1 PushWorld Puzzle Definition

Formally, a PushWorld puzzle is a 4-tuple $(S,A,G,T)$ with the following
components:

$S$ is a finite set of states, where each state $s$ is a set of object states.
An object state $o_{i}$ is a pair $(p_{i},Y_{i})$ containing the 2D integer
position $p_{i}$ of the object and a set $Y_{i}$ of the 2D integer positions
of all cells that the object occupies relative to $p_{i}$. One object $o_{A}$
is designated as the agent, which is directly controlled by actions, and
another object $o_{W}$ is designated as an immovable wall. For convenience,
let $Z_{i}(p)=\\{y+p\mid y\in Y_{i}\\}$ denote the set of absolute positions
of all cells in object $i$ when the object is located at position $p$. No
object’s occupied cells are allowed to intersect the occupied cells of any
other object in any state: $Z_{i}(p_{i})\cap Z_{j}(p_{j})=\emptyset$ for all
object pairs $i$, $j$.

$A$ is a fixed set of actions
$\\{\textrm{Left},\textrm{Right},\textrm{Up},\textrm{Down}\\}$, which
correspond to moving the agent by a single cell in the corresponding
direction. For each action $a$, let $\hat{u}_{a}$ denote the unit vector in
the direction of movement.

$G$ defines the goal of the puzzle. It is a mapping from one or more object
indices $i$ to corresponding goal positions $g_{i}$. The goal of the puzzle is
achieved in a given state when $p_{i}=g_{i}$ for all $(i,g_{i})$ pairs in $G$.

$T$ is a state transition function that maps $S\times A\rightarrow S$. Let
$s^{\prime}=T(s,a)$, and likewise let $p_{i}^{\prime}$ denote the position of
object $i$ in $s^{\prime}$. For every action $a$, a subset of objects $M$ may
move in the direction $\hat{u}_{a}$. $M$ contains the agent object as well as
every object $m$ for which there exists a chain of two or more objects
${c_{1},c_{2},...,c_{n}}$ that satisfy these conditions:

  * •

$c_{1}$ is the agent, which is the root cause of the movement.

  * •

$c_{n}=m$

  * •

Every object in the chain pushes against the next object in the chain:
$Z_{c_{i}}(p_{c_{i}}+\hat{u}_{a})\cap Z_{c_{i+1}}(p_{c_{i+1}})\neq\emptyset$
for $i=1,\dotsb,n-1$.

The wall object $W$ cannot change position, so if $M$ contains $W$, then no
objects move, resulting in $s^{\prime}=s$. If not, then:

$p_{i}^{\prime}=\begin{cases}p_{i}+\hat{u}_{a}&\textrm{if}\ i\in M\\\
p_{i}&\textrm{otherwise}\end{cases}$

for every object $i$. $T$ copies the object cells $Y_{i}$ without
modification.

In one variation, the agent object may be constrained by additional walls that
all other objects can move through. These agent-specific walls can force the
agent to use tools to manipulate objects out of reach.

###  A.2 Classical Planner Implementations

We used the following implementations of these classical planners:

  * •

Fast Forward: We used the version included in the Fast Downward project at
https://www.fast-downward.org/.

  * •

Fast Downward: We used the “seq-sat-fd-autotune-1” version included in the
Fast Downward project, and this version uses the context-enhanced additive
heuristic [?] that succeeded the causal graph heuristic. This version
outperforms “seq-sat-fd-autotune-2” in PushWorld.

  * •

LAMA: We used the “seq-sat-lama-2011” version included in the Fast Downward
project.

  * •

Best-First Width Search: We used the version from
https://github.com/nirlipo/BFWS-public and only report results for the “DUAL-
BFWS” mode, which outperforms “BFWS-f5” and “k-BFWS” in PushWorld.

  * •

Fast Downward Stone Soup: We used the “seq-sat-fdss-2018” version included in
the Fast Downward project.

  * •

Saarplan: We used the version from IPC 2018 at
https://ipc2018-classical.bitbucket.io/.

###  A.3 Recursive Graph Distance Pseudocode

Algorithm 1 defines the RGD heuristic in the context of PushWorld. This
implementation uses movement graphs, which are equivalent to object position
DTGs that omit parallel edges. Instead of labeling graph edges with operator
preconditions, the implementation uses the RelativePushingPositions function
to return all relative position offsets from which one object can push another
in a given direction based on the objects’ shapes. The movement graphs omit
all positions of objects that would collide with walls, so the implementation
does not include explicit collision checks.

In the reported experimental results, the RGD implementation includes the
following optimizations, which are not shown in Algorithm 1:

  * •

The heuristic ignores negative preconditions of operators; it does not check
for whether movable obstacles would block a movement. Such a check could be
exponentially expensive in the number of objects because the movement might be
blocked by an obstacle that is blocked by another obstacle, etc., which is
finally blocked by a wall. This choice does not affect the completeness of a
search that uses this heuristic.

  * •

Instead of allowing unbounded recursion depth in PushingCost, which can be
exponentially expensive in the number of objects, planning time is empirically
reduced by incrementally increasing the maximum recursion depth until the
heuristic cost is not infinite.

  * •

The calculation of $d_{\text{min}}$ within the for loop on line 30 is cached
for all pairs of objects in all positions.

  * •

ShortestPathLength computes the length of the shortest directed path between
two nodes in a graph. These path lengths are cached, and computing the
shortest path length between one pair of nodes can reduce the computation in
subsequent calls for different nodes. One possible implementation is to
maintain a breadth-first expansion from each target node and to continue the
expansion whenever a source node is not in the visited set whose shortest path
lengths to the target are known.

  * •

RelativePushingPositions is memoized to avoid repeated comparisons of the same
object shapes.

Algorithm 1 Recursive Graph Distance Heuristic

1:function
RecursiveGraphDistanceCost($\text{state},\text{goal},\text{movement\\_graphs}$)

2: $c\leftarrow 0$

3: for all $(\text{object\\_id},p_{\text{goal}})\in\text{goal}$ do

4: $c\leftarrow
c+\textsc{CostToReachPosition}(\text{object\\_id},p_{\text{goal}},\text{state},\text{movement\\_graphs})$

5: if $c=\infty$ then

6: break

7: return c

8:

9:function
CostToReachPosition($\text{object\\_id},p_{\text{goal}},\text{state},\text{
movement\\_graphs}$)

10: $p\leftarrow\text{state}[\text{object\\_id}].\text{position}$

11: if $p=p_{\text{goal}}$ then

12: return $0$

13: $c_{\text{min}}\leftarrow\infty$

14: $\text{used\\_object\\_ids}\leftarrow\\{\text{object\\_id}\\}$

15: for all
$p_{\text{next}}\in\text{movement\\_graphs}[\text{object\\_id}].\text{direct\\_successors}(p)$
do

16:
$d\leftarrow\textsc{ShortestPathLength}(\text{movement\\_graphs}[\text{object\\_id}],p_{\text{next}},p_{\text{goal}})$
$\triangleright$ Returns $\infty$ if no path exists.

17: if $d<c_{\text{min}}$ then

18: $c_{\text{min}}\leftarrow
d+\textsc{PushingCost}(\text{object\\_id},p_{\text{next}},\text{used\\_object\\_ids},\text{state},\text{movement\\_graphs},c_{\text{min}}-d)$

19: return $c_{\text{min}}$

20:

21:function
PushingCost($\text{object\\_id},p_{\text{next}},\text{used\\_object\\_ids},\text{state},\text{movement\\_graphs}$,
cost_upper_bound)

22: $c_{\text{min}}\leftarrow\text{cost\\_upper\\_bound}$

23: $p\leftarrow\text{state}[\text{object\\_id}].\text{position}$

24: $\hat{u}_{a}\leftarrow p_{\text{next}}-p$

25: for $\text{pusher\\_object\\_id}\leftarrow 1\textbf{ to }|\text{state}|$
do

26: if $\text{pusher\\_object\\_id}\in\text{used\\_object\\_ids}$ then

27: continue

28:
$p^{\text{pusher}}\leftarrow\text{state}[\text{pusher\\_object\\_id}].\text{position}$

29:
$\text{next\\_used\\_object\\_ids}\leftarrow\text{used\\_object\\_ids}\cup\\{\text{pusher\\_object\\_id}\\}$

30: $\triangleright$ For all pusher positions that are adjacent to the
pusher’s current position, compute the minimum graph

31: $\triangleright$ distance from each adjacent position to the position
where the pusher makes contact with the object_id.

32: for all
$p_{\text{next}}^{\text{pusher}}\in\text{movement\\_graphs}[\text{pusher\\_object\\_id}].\text{direct\\_successors}(p^{\text{pusher}})$
do

33: $d_{\text{min}}\leftarrow\infty$

34: for all $\Delta
p\in\textsc{RelativePushingPositions}(\text{state}[\text{object\\_id}].\text{shape},\text{state}[\text{pusher\\_object\\_id}].\text{shape},\hat{u}_{a})$
do

35: $p_{\text{start}}^{\text{pusher}}\leftarrow p+\Delta p$

36: $p_{\text{end}}^{\text{pusher}}\leftarrow
p_{\text{start}}^{\text{pusher}}+\hat{u}_{a}$

37: if
$(p_{\text{start}}^{\text{pusher}},p_{\text{end}}^{\text{pusher}})\notin\text{movement\\_graphs}[\text{pusher\\_object\\_id}].\text{edges}$
then

38: continue

39: if $p_{\text{start}}^{\text{pusher}}=p^{\text{pusher}}\textbf{ and
}p_{\text{end}}^{\text{pusher}}=p_{\text{next}}^{\text{pusher}}$ then
$\triangleright$ This is a simultaneous push, so there is no cost.

40: $d_{\text{min}}\leftarrow 0$

41: break

42: else

43:
$d\leftarrow\textsc{ShortestPathLength}(\text{movement\\_graphs}[\text{pusher\\_object\\_id}],p_{\text{next}}^{\text{pusher}},p_{\text{start}}^{\text{pusher}})$

44: $d\leftarrow d+1$ $\triangleright$ Add the cost of moving the pusher from
its start position to its end position.

45: $d_{\text{min}}\leftarrow\text{min}(d_{\text{min}},d)$

46: if $\text{pusher\\_id}=\text{AGENT\\_OBJECT\\_ID}$ then

47: $c_{\text{min}}\leftarrow\text{min}(c_{\text{min}},d_{\text{min}}+1)$
$\triangleright$ Moving the agent to its next position costs 1 action.

48: else if $d_{\text{min}}<c_{\text{min}}$ then

49: $c_{\text{min}}\leftarrow d_{\text{min}}+\textsc{PushingCost}($
$\text{pusher\\_object\\_id},p_{\text{next}}^{\text{pusher}},\text{next\\_used\\_object\\_ids},\text{state},\text{movement\\_graphs},c_{\text{min}}-d_{\text{min}})$

50: return $c_{\text{min}}$

###  A.4 Deep RL Network Architecture and Hyperparameters

We used the PPO and DQN implementations in DeepMind Acme [?]:
https://github.com/deepmind/acme. For both PPO and DQN, we used a
convolutional neural network (CNN) [?] as the vision processor, implemented in
JAX (https://github.com/google/jax). The CNN consisted of 3 convolutional
layers below 2 fully connected layers, and every layer used ReLU activation.
The three convolutional layers used filter sizes of 3x3, 3x3, and 5x5, with
stride lengths of 3, 1, and 1. The fully connected layers had sizes 256 and
128, respectively.

The PPO hyperparameters were: entropy cost 0.01, learning rate 0.0002, and
number of epochs 2. The DQN hyperparameters were: learning rate 0.0001,
epsilon 0.05, samples per insert 2, batch size 256, discount 1.0, and 1-step
updates. These parameters were determined via hyperparameter sweeps.

We set the episode length to 100 steps, considering that Level 0 puzzles can
be solved in fewer than 15 steps on average, Level 1 puzzles can be solved in
fewer than 30 steps on average, and puzzles with multiple goals can provide
positive rewards without solving the entire puzzle.

###  A.5 PPO and DQN Results

Figure 6 shows the performance of PPO and DQN on Level 0 puzzle sets.

Figure 6: PPO and DQN training and test curves on the Level 0 puzzle sets.


#  WAFFLE: Weighted Averaging for
Personalized Federated Learning

Martin Beaussart
EPFL
<EMAIL_ADDRESS>&Felix Grimberg
EPFL
<EMAIL_ADDRESS>&Mary-Anne Hartley
EPFL
<EMAIL_ADDRESS>&Martin Jaggi
EPFL
<EMAIL_ADDRESS>
###### Abstract

In federated learning, model personalization can be a very effective strategy
to deal with heterogeneous training data across clients. We introduce WAFFLE
(Weighted Averaging For Federated LEarning), a personalized collaborative
machine learning algorithm that leverages stochastic control variates for
faster convergence. WAFFLE uses the Euclidean distance between clients’
updates to weigh their individual contributions and thus minimize the
personalized model loss on the specific agent of interest. Through a series of
experiments, we compare our new approach to two recent personalized federated
learning methods—Weight Erosion and APFL—as well as two general FL
methods—Federated Averaging and SCAFFOLD. Performance is evaluated using two
categories of non-identical client data distributions—concept shift and label
skew—on two image data sets (MNIST and CIFAR10). Our experiments demonstrate
the comparative effectiveness of WAFFLE, as it achieves or improves accuracy
with faster convergence.

##  1 Introduction

Federated learning (FL) is a collaborative learning technique to build machine
learning (ML) models from data distributed among several participants
("agents"). The objective is to generate a common, robust model not by
exchanging the data between agents, but rather through exchanging parameter
updates for the common model that all agents share at a certain frequency.
This technique addresses some major problems of centralized data-sharing
approaches, can enable data privacy and security and is therefore attractive
to many applications with sensitive data.

We consider centralized FL, in which a single server orchestrates the
execution of the algorithms. Each iteration of learning can thus be divided
into four main steps. Firstly, the central server sends the current model to
all agents. Secondly, each agent trains the model with their data and,
thirdly, sends the updated model parameters back to the central server.
Finally, the central server aggregates these results and generates a new
global model. WAFFLE intervenes in this last step: instead of a global one-
size-fits-all model, the aggregation produces a personalized model for each
specific agent.

###  1.1 Control Variates for Heterogeneous Data

Federated learning develops a common model for all agents, typically under the
assumption that data is _independent and identically distributed (IID)_ across
agents. In real-world applications, this assumption rarely holds, and thus the
convergence rate and final performance of FL algorithms like Federated
Averaging (FedAvg) can vary significantly [McMahan et al., 2017, Karimireddy
et al., 2020]. Karimireddy et al. [2020] introduce SCAFFOLD, which uses
stochastic control variates (SCV) to tackle the _client-drift_ that results
from non-IID-ness across agents. Each agent maintains an SCV estimating the
sum of all other agents’ gradients, and uses it to “correct” each step of
stochastic gradient descent (SGD). Compared to FedAvg, SCAFFOLD converges
faster and towards a better model. WAFFLE extends SCAFFOLD to collaboratively
train a personalized ML model.

###  1.2 Model Personalization

A complementary solution for dealing with non-IID data is to train a unique
(personalized) model for each agent, rather than a single global model.
Recently, there have been several key works on personalization in federated
learning, see e.g. Kulkarni et al. [2020] for an overview. Some methods of
adapting global models for individual agents are summarized in Kairouz et al.
[2019], such as local fine-tuning [Deng et al., 2020, Kairouz et al., 2019],
multi-task learning [Smith et al., 2017], and model-agnostic meta-learning
(MAML) [Finn et al., 2017, Fallah et al., 2020].

The model personalization method most similar to WAFFLE is Weight Erosion (WE)
[Grimberg et al., 2020]. It is conceptually related to local fine-tuning,
which consists of training a common global model and subsequently
personalizing it for each agent by performing a small number of SGD steps on
the agent’s local data (resulting in one model per agent) [Kairouz et al.,
2019]. WE achieves a smoother and more differentiated transition from global
to local training by using a _weighted average_ of the agents’ gradients to
update the server model at each communication round. Thus, the server model
after running WE once is personalized for one specific agent—called _user_ or
_Alice_.111 The authors motivate this setting, in which all other agents
contribute selflessly to Alice’s objective, by constructing a fictive scenario
of cross-silo FL across hospitals. Alternatively, for cross-silo FL, WE (or
WAFFLE) can be run several times in parallel to obtain personalized models for
all agents.  Like WE, WAFFLE uses aggregation weights derived from the
Euclidean distance between gradients to update the global model, but it
combines this approach with the use of SCVs (Section 1.1). These different
approaches for transitioning from global to local training are illustrated in
Figure 3 (Section 3).

Instead of interpolating between gradients at each round, a slightly different
approach consists of interpolating between a local and global model [Grimberg
et al., 2021, Donahue and Kleinberg, 2020, Deng et al., 2020, Zhang et al.,
2021]. For instance, Deng et al. [2020] propose the adaptive FL algorithm APFL
that achieves personalization by learning a global model and a set of local
corrective models with which to interpolate. Another related approach is
personalized MAML, where a global model is optimized with respect to (w.r.t.)
the _loss after local fine-tuning_ on each agent’s local data [Finn et al.,
2017, Fallah et al., 2020]. Finally, regularization can be used (instead of
weighted averaging) to mix global and local training [Smith et al., 2017, Li
et al., 2021]. For instance, Li et al. [2021] recently showed very promising
results with Ditto, where the local model’s training is split across all
communication rounds and regularized by its distance from the current global
model.

###  1.3 Benchmark

While several categories of non-IID data have been identified by Kairouz et
al. [2019], we focus only on label skew and concept shift:

  * •

Label skew: The distribution of labels varies between agents, but the "true"
classification function is the same for all agents. This is, so far, the type
of non-IID-ness most commonly used in personalized FL benchmarks [Li et al.,
2021, Fallah et al., 2020, Deng et al., 2020].

  * •

Concept shift: The mapping from the features to the label varies across
agents. We include concept shift to model naturally partitioned data sets like
those used by Smith et al. [2017].

###  1.4 Contributions

  * •

We present a new personalized collaborative ML algorithm named WAFFLE
(Weighted Averaging For Federated LEarning), which builds on the existing
methods SCAFFOLD and Weight Erosion (WE) by using stochastic control (SC) and
weighted gradient aggregation. To the best of our knowledge, this is the first
method that uses SC in model personalization.

  * •

We build an image classification benchmark with label skew and concept shift.
We then use it to compare WAFFLE against baselines (Federated Averaging and
Local), its parent methods (SCAFFOLD and WE), and against other state-of-the-
art personalized FL methods (APFL), thus contributing to further empirical
evaluation of these methods.

  * •

Finally, we show that in most cases, WAFFLE converges faster than its
competitors WE and APFL and matches or improves their accuracy, while
requiring less hyperparameter tuning.

##  2 The WAFFLE Algorithm

To introduce WAFFLE, we will first summarize the known SCAFFOLD FL algorithm
on which it is based, and then detail how personalization can be achieved by
weighted averaging between agents, for suitable weight choices. SCAFFOLD
maintains an estimate of what all agents are learning, called a control
variate, so that agents can simulate having all the data and estimate the
direction in which they should update their gradient. Thus, it is more
efficient and addresses the problem of _client drift_. Indeed Karimireddy et
al. [2020] prove that SCAFFOLD converges to the globally optimal model instead
of converging to a weighted average of local models and does so much faster
and more accurately than FedAvg in the presence of inter-agent heterogeneity.

###  2.1 Personalization Using Weighted Update Aggregation

WAFFLE is essentially a personalized version of SCAFFOLD, where the goal is no
longer to get a global model for all agents but a personalized model for one
particular agent, (Alice). The idea is to start from global training
(SCAFFOLD) and to gradually move to local training. SCAFFOLD computes the
gradient for the server model ($\Delta\boldsymbol{x}$) and the control variate
$\Delta\boldsymbol{c}$ (the assumption of what other agents are learning) by
averaging each agent’s update $\Delta\boldsymbol{y}_{i}$ and local control
variates $\boldsymbol{c}_{i}$ at each round [Karimireddy et al., 2020,
Algorithm 1, Line 16]. It is at this step that WAFFLE intervenes: Instead of
averaging all agents at each round, WAFFLE assigns a weight $\alpha_{i}$ to
each agent and computes $\Delta\boldsymbol{x}$ and $\Delta\boldsymbol{c}$ by a
weighted combination as per Equation 1. Changes w.r.t. SCAFFOLD are
highlighted in red:

$\left(\Delta\boldsymbol{x},\Delta\boldsymbol{c}\right)\leftarrow\sum_{i\in{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\\{1,\dots,N\\}}}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\alpha_{i}}\left(\Delta\boldsymbol{y}_{i},\Delta\boldsymbol{c}_{i}\right)$
(1)

For $N$ agents, we recover SCAFFOLD (global training) by setting all weights
equal to $\frac{1}{N}$. In contrast, we recover local training by setting all
weights to $0$, except the weight of Alice.

WAFFLE is based on a smooth transition from global to local training. To do
this, the weight of each agent is updated according to the degree of
personalization desired for the round. Just as SCAFFOLD uses the control
variates to converge to the global model for the union of all agents’
datasets, WAFFLE uses them to converge to a personalized model based on a
weighted subset of agents (Figure 1). Each agent is included and weighted in
this subset based on the current degree of personalization and whether their
dataset is sufficiently similar to Alice’s, as measured by the distance
between their gradients at each round.

Figure 1: Evolution of the personalized model and agent models on WAFFLE in a
three-rounds scenario where agent 1 has a higher weight than agent 2.

The weight of each agent at a given round depends on the distance between its
and Alice’s gradient. As the method uses SGD (i.e., stochastic gradients), the
agent weights are also subject to stochastic noise. Indeed, at some rounds,
some agents may have an atypical gradient resulting in a bad weight, even if
their contribution is still needed for the global model. Therefore, we average
the current weight with the last two weights to smooth out strong random
effects (line 10 of Algorithm 1).

To compare each agent with Alice, we use the Euclidean distance between
gradients. Simply, we want an agent with a small distance to have a higher
weight than an agent with a larger distance. Like SCAFFOLD, we have a central
server that coordinates all the agents. At each starting round it gives the
current personalized model and control variate to all agents for training.
When the agents finish training the model with their data, they send the
updated model and their local control variate back to the server. The server
then calculates the Euclidean distance of each update from the selected agent
and modifies the weight assigned to each agent accordingly. Finally, it
calculates the new personalized model and the new control variate (Equation
1). The implementation of WAFFLE can be seen in the modified version of
SCAFFOLD (Algorithm 2, Appendix B) where changes w.r.t. to the original code
are written in red ink. Agent weights are obtained by CalcWeight (Algorithm
1), which is explained in the next subsection.

###  2.2 Computing Agent Weights for WAFFLE

WAFFLE uses a new data-dependent approach to select the contribution weight
$\alpha_{i}^{r}$ of agent $i$ at each round $r$. It relies on hyperparameters
$\Omega$ and $\Psi$: two functions defining the degree of personalization for
each round $r$ (where a value of $1$ corresponds to global training and $0$ is
personalized (local) training). As explained later, we reduce the functions
$\Omega$ and $\Psi$ to a single, positive real-valued hyperparameter for ease
of tuning.

The weight is calculated according to Equations 2 and 3, where

  * •

agent $i^{\star}$ is Alice

  * •

$d_{i}$ is the distance between the gradients of agents $i$ and $i^{\star}$

  * •

$\mathit{dM}$ is the largest such distance:
$\mathit{dM}\leftarrow\underset{i}{\max}\ d_{i}$

  * •

Analogously, $\mathit{dm}$ is the smallest such distance (excluding
$d_{i^{\star}}$)

  * •

$\Omega(r)$ and $\Psi(r)$ denote the degree of personalization for round $r$

The weight of an agent is based on the distance between its gradient and the
gradient of agent $i^{\star}$, but also on the distribution of the distances
of the agents. The distance of $i^{\star}$ would always be 0, we want to
change it to a similar value compared to the distances of other agents. The
further away Alice’s distance is from the others, the less weight we assume
the agents have compared to Alice. We need a “distance” $d_{i^{\star}}$ for
Alice in Equation 3. Thus, Equation 2 serves to assign her such a distance
based on the hyperparameter $\Omega$. When $\Omega\approx 1$ (i.e., more
global training), Alice’s assigned “distance” ($d_{i^{\star}}$) is close to
the next smallest distance $\mathit{dm}$, resulting in similar weights for
these two agents. The more $\Omega$ decreases towards 0, the smaller
($d_{i^{\star}}$) becomes compared to all other agents’ distances, resulting
in a large weight for Alice (i.e., more local training).

$d_{i^{\star}}\leftarrow\mathit{dm}\cdot\Big{(}1-{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\frac{\mathit{dM}-\mathit{dm}}{\mathit{dM}}}(1-\Omega(r))\Big{)}$
(2)
$\alpha_{i}^{r}\leftarrow\max\Big{\\{}\Psi(r)-\frac{d_{i}-d_{i^{\star}}}{\mathit{dM}-d_{i^{\star}}},\,0\Big{\\}}$
(3)

Equation 3 serves as a threshold of inter-agent utility. To pass the
threshold, the gradient of agent $i$ must be closer than the other agents’
gradients (based on the range of distances). With $\Psi$ close to 0 (i.e.,
more local training), only a small fraction of the range between the minimum
and maximum distance will map to a non-zero weight. Thus, learning is
concentrated in a subset of agents with the most similar gradients.

Using the two functions $\Omega$ and $\Psi$, we can define different
strategies for WAFFLE. We could, as an example, shift rapidly toward zero and
thus drastically limit larger steps in global training in the initial phase of
personalization.

As Algorithm 1 would otherwise have $2R$ positive real-valued parameters to
tune (values of $\Omega$ and $\Psi$ for each round), we restrict ourselves to
$\Omega(r)=\Psi(r)=$ a single-parameter function of $r$. However, more methods
should be investigated to see the full potential of WAFFLE in different
scenarios.

In this paper, we use a sigmoid function to give the general training a
gradual slope towards smaller values Equation 4 and Figure 8. The function has
a parameter $\Delta\Omega$ that controls the slope, a higher value will
steepen the slope and thus move faster towards the local training. According
to our experiments, a good value for $\Delta\Omega$ is around 3.2,
irrespective of the model used. With this value, WAFFLE transitions to mostly
local SGD after $70-90\%$ of the total number of rounds, depending on the
similarity of the agents’ gradient to that of Alice. A simple modification to
delay or accelerate the local learning time is to move the $\Omega$ function
horizontally by adding a value to $r$.

$\displaystyle\Omega(r)=\Psi(r)=\frac{1}{1+e^{\Delta\Omega\cdot(r/(R/2)-1)})}$
(4)

input : $N$, $R$, $i^{\star}$, $r$, $\\{\Delta\boldsymbol{y}_{i}\\}$,
$\mathbf{\alpha}^{r-1}$, $\mathbf{\alpha}^{r-2}$

hyperparameter : Functions $\Omega$ and $\Psi$ from $\\{1,\dots,R\\}$ to
$\mathbb{R}_{>0}$

output : Weight vectors $\bar{\mathbf{\alpha}}^{r},\mathbf{\alpha}^{r}$

1

2foreach  __agent_ $i\in\\{1,\dots,N\\}$ _ do

// Compute the Euclidean distance between updates of agents $i$ and
$i^{\star}$.

$d_{i}\leftarrow\left\|\Delta\boldsymbol{y}_{i}-\Delta\boldsymbol{y}_{i^{\star}}\right\|_{2}$

// $d_{i^{\star}}=0$

3

4$\mathit{dM}\leftarrow\underset{i\neq i^{\star}}{\max}\ \ d_{i}$,
$\mathit{dm}\leftarrow\underset{i\neq i^{\star}}{\min}\ \ d_{i}$

$d_{i^{\star}}\leftarrow\mathit{dm}\cdot\left(1-{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\frac{\mathit{dM}-\mathit{dm}}{\mathit{dM}}}(1-\Omega(r))\right)$

// $0\leq d_{i^{\star}}\leq\mathit{dm}$ if $\Omega(r)\in[0,1]$

5 foreach  __agent_ $i\in\\{1,\dots,N\\}$ _ do

6
$\alpha_{i}^{0}\leftarrow\max\big{\\{}\Psi(r)-\frac{d_{i}-d_{i^{\star}}}{\mathit{dM}-\underset{i}{\min}\
d_{i}},0\big{\\}}$

7 if  _$r\geq 0.95R$_ then

8 $\alpha_{i}^{0}\leftarrow 0$ if $i\neq i^{\star}$ else $1$

9

10

11$\mathbf{\alpha}^{0}\leftarrow\frac{1}{\sum\alpha^{0}_{i}}(\alpha^{0}_{1},\dots,\alpha^{0}_{N})^{\top}$

12
$\bar{\mathbf{\alpha}}^{r}\leftarrow\frac{1}{3}(\mathbf{\alpha}^{-2}+\mathbf{\alpha}^{-1}+\mathbf{\alpha}^{0})$

13

Algorithm 1 CalcWeights-$\Omega$

##  3 Benchmarking Personalized Collaborative Learning Methods

To compare the performance of WAFFLE with the state of the art, we evaluate it
against two recent personalized FL methods—APFL [Deng et al., 2020] and Weight
Erosion (WE) [Grimberg et al., 2020]—as well as two global methods—Federated
Averaging [McMahan et al., 2017] and SCAFFOLD [Karimireddy et al., 2020]—and
finally, local training using only Alice’s dataset (Local). SCAFFOLD and WE
were selected to validate whether WAFFLE outperforms the methods it builds
upon. We included APFL in the benchmark as an unrelated personalized FL method
which has an open-source implementation. The benchmark consists of training
standard image classification networks on standard IID and non-IID image
datasets described below.

Two 10-class image classification tasks are used based on either the MNIST or
CIFAR10 datasets. To generate confidence estimates, we run the MNIST benchmark
with five different seeds. However, using five seeds for CIFAR10 would be
prohibitive due to computational limitations.

Distributions and Models. We test five distributions (A,B,C,A*,B*), the first
one (A) is IID and serves as a baseline. A classic label skew distribution (B)
was selected for consistency with several other benchmarks [Collins et al.,
2021, Deng et al., 2020, McMahan et al., 2016], where only a small number of
agents have samples from any given class (uniformly distributed amongst these
agents). Distribution C exhibits more nuanced label skew, where the labels are
not uniformly distributed among the agents. Instead, each agent will have 40%
of one label, 20% of two labels, and 10% of two other labels. We also present
two other distributions—concept shift on the IID distribution (A*), and
concept shift on top of classic label shift (B*). Concept shift is achieved by
randomly swapping the class labels of all agents, except for Alice. We specify
each distribution as a list of ten numbers that sum to one (see Figure 2)).
For each agent $i$, we move the list to the right $i$ times to get non-IID
data, e.g. with distribution B, agent 0 will only have labels 0 to 3, agent 1
labels 1 to 4, etc.

For MNIST we use a LeNet-5 model [LeCun et al., 1989] and for CIFAR10 we use
DLA [Yu et al., 2017].

Figure 2: Repartition of labels between all agents for each distribution
benchmarked

Hyperparameters. We use an SGD optimizer with a learning rate of 0.1 for MNIST
and 0.01 for CIFAR10. Following [Yu et al., 2017] we add a momentum 0.9 and a
weight decay 5e-4 for CIFAR10 . Concerning the batch size, it is set to 32 for
MNIST and 128 for CIFAR10. Except for the APFL method, we define the
hyperparameters in a principled manner (based on the expected behaviour of the
algorithm). As WE and WAFFLE use a weight system, we attempt to set the
parameters in such a way that both algorithms reach fully local learning at
about the same time. In practice, this was difficult to achieve as some
hyperparameters are very sensitive. Fully local learning is obtained when all
weights reach zero except for Alice (Agent 0), who has a weight of one. Figure
3 shows an example of the evolution of the weights. Alice (agent 0) is placed
in the centre. For the hyperparameters of APFL, we use the same method mention
in the paper by [Deng et al., 2020], and adaptively update the hyperparameter
to guarantee the best generalization performance.

Figure 3: Evolution of weight for personalized FL with the distribution C on
MNIST (seed = 1)

The tuned hyperparameters are listed in Table 3 (Appendix B). Out of the two
methods, WAFFLE is the easiest to parametrize, as the hyperparameter does not
change in contrast to the weighting system used by WE. This is because WAFFLE
can work very well without tuning, as opposed to WE which requires
hyperparameter tuning to perform well. As explained in Section 2.2, we
purposefully restrict WAFFLE to a single hyperparameter, $\Delta\Omega$
(Equation 4) and we use the same value for all experiments. Nevertheless, we
show in the result section that WAFFLE is able to obtain a very good accuracy
even without tuning.

##  4 Results

We present the results of our benchmarking in Table 1 below. For each method
and each distribution, it shows the best accuracy reached at any of the 100
epochs. In the MNIST section, results are averaged across five random seeds
and listed along with the standard deviation as $avg\pm std$.

Figure 4: Evolution of accuracy with distribution C on MNIST, average over
five seeds of the best accuracy obtained up to each turn.

To see the evolution of the accuracy, we provide an example in Figure 4. The
evolution for each experiment is shown in the appendix (see Figures 6 and 7).

MNIST. On MNIST, WAFFLE and WE outperform the other methods for the label
skews (B, C), but not for IID distribution (A). The difference between WAFFLE
and WE lies in the speed of convergence: As expected WAFFLE converges faster
thanks to the use of SCVs.
For the concept shift distributions (A*, B*), all personalized FL methods
improve accuracy compared to global FL methods. Among all personalized FL
methods, WE performs slightly better in this setting. Here, WAFFLE is not more
efficient than WE (time taken to reach convergence). This may be related to
the fact that only small fractions of agents are useful in this distribution,
so until we reach this critical point, both algorithms are bound to the same
slow increase until a threshold. Regarding APFL, the method also performs
better than global FL methods except for the IID distribution (A), but it
struggles to achieve a better result than the other two personalized FL
methods and Local.

CIFAR10. Concerning CIFAR10, the lack of other seeds makes us less confident
about the interpretation of our results, however, the results are coherent
with our expectations. For example, like with MNIST, global FL methods
outperform personalized methods on the IID distribution and present a
significant improvement compared to Local. Distribution B and C (label skew)
produce results close to those of MNIST, where WAFFLE and WE outperform
others. Concerning concept shift, personalized FL methods are still a great
improvement compared to global FL or Local. The accuracy of WAFFLE and WE is
approximately the same but as for MNIST we don’t have a noticeable difference
between the speed of convergence to best accuracy. For APFL, we get similar
results as MNIST, APFL performs better than the global FL methods, but this
time it is better than Local. Unfortunately, compared to the other two
customized methods, it performs much worse.

Table 1: Comparison of average local test accuracy and standard deviation of
different algorithms given different data distributions. Algorithms
---
|  Local |  FedAvg |  SCAFFOLD |  WE |  WAFFLE |  APFL
MNIST |  |  |  |  |  |
IID (A) |  $96.88\%$ $\pm 0.47$ |  $98.99\%$ $\pm 0.08$ |  $\mathbf{99.03}\%$ $\pm 0.08$ |  $99.01\%$ $\pm 0.16$ |  $98.93\%$ $\pm 0.1$ |  $97.55\%$ $\pm 0.13$
Distr. B |  $99.38\%$ $\pm 0.09$ |  $98.97\%$ $\pm 0.18$ |  $99.24\%$ $\pm 0.17$ |  $\mathbf{99.69}\%$ $\pm 0.07$ |  $\mathbf{99.69}\%$ $\pm 0.03$ |  $99.36\%$ $\pm 0.09$
Distr. C |  $98.54\%$ $\pm 0.1$ |  $98.56\%$ $\pm 0.1$ |  $98.77\%$ $\pm 0.1$ |  $\mathbf{99.0}\%$ $\pm 0.08$ |  $\mathbf{99.0}\%$ $\pm 0.08$ |  $98.45\%$ $\pm 0.08$
Distr. A* |  $96.56\%$ $\pm 0.27$ |  $13.37\%$ $\pm 3.21$ |  $19.93\%$ $\pm 5.95$ |  $\mathbf{98.09}\%$ $\pm 0.26$ |  $97.99\%$ $\pm 0.26$ |  $96.24\%$ $\pm 0.22$
Distr. B* |  $99.47\%$ $\pm 0.07$ |  $45.26\%$ $\pm 4.1$ |  $57.51\%$ $\pm 4.02$ |  $\mathbf{99.67}\%$ $\pm 0.09$ |  $99.63\%$ $\pm 0.11$ |  $99.44\%$ $\pm 0.07$
CIFAR10 |  |  |  |  |  |
IID (A) |  0.771% |  0.925% |  $\mathbf{0.927\%}$ |  0.912% |  0.907% |  0.776%
Distr. B |  0.888% |  0.884% |  0.85% |  $\mathbf{0.935}$% |  $0.931\%$ |  0.895
Distr. C |  0.832% |  0.857% |  0.868% |  0.907% |  $\mathbf{0.914}$% |  0.809%
Distr. A* |  0.763% |  0.248% |  0.259% |  0.877% |  $\mathbf{0.9}$% |  0.78%
Distr. B* |  0.876% |  0.782% |  0.476% |  $\mathbf{0.937}$% |  0.926% |  0.878%

##  5 Limitations and Future Work

The benchmark has some limitations that should be addressed in future work and
are discussed in detail in Appendix A. Firstly, the theoretical properties of
WAFFLE should be investigated, in particular whether the guarantees of
SCAFFOLD are maintained. Secondly, our benchmark should be expanded to cover
more types of non-IID distributions and naturally partitioned datasets, and to
include more personalized FL methods like the recently published Ditto.
Finally, the versatility of WAFFLE could also be further investigated by
optimizing the functions $\Omega$ and $\Psi$ for different contexts.

##  6 Conclusion

In this work, we propose WAFFLE, a personalized FL algorithm based on
SCAFFOLD, that can overcome client drift and accelerate convergence to an
optimal personalized model. WAFFLE uses the Euclidean distance between agent
updates to weigh their contributions and thus transition gradually from global
to more local training. We explored the performance of WAFFLE in two cases of
non-identical data—concept shift and label skew—on two standard image
datasets—MNIST and CIFAR10. We demonstrate that in most cases WAFFLE shows a
faster convergence and a possible improvement in accuracy compared with other
personalized FL methods. We also demonstrate that WAFFLE is easier to handle
than most personalized learning methods. Indeed, a fixed default value for its
single hyperparameter yields competitive results in all test cases. However,
by design, WAFFLE can be adapted to a task by modifying $\Omega$ and $\Psi$
functions and thus potentially reach even better accuracy.

## Acknowledgements

We thank David Roschewitz for sharing with us his implementation of APFL, and
Freya Behrens for inspiring us on how to create synthetic concept shift. We
are grateful to Celiane De Luca for her suggestions and comments on the
writing of this paper. We thank Vincent Yuan and Damien Gengler for their
collaboration with Martin Beaussart on adapting Weight Erosion to a neural
network for the first time, which lead directly to the present work [Gengler
et al., 2020].

## References

  * McMahan et al. [2017] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-Efficient Learning of Deep Networks from Decentralized Data.  In _AISTATS - Proceedings of the 20th International Conference on Artificial Intelligence and Statistics_ , volume 54, pages 1273–1282. PMLR, 2017.
  * Karimireddy et al. [2020] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh.  Scaffold: Stochastic controlled averaging for federated learning.  In _International Conference on Machine Learning_ , pages 5132–5143. PMLR, 2020.
  * Kulkarni et al. [2020] Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant.  Survey of personalization techniques for federated learning, 2020.
  * Kairouz et al. [2019] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al.  Advances and open problems in federated learning.  _arXiv preprint arXiv:1912.04977_ , 2019.
  * Deng et al. [2020] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi.  Adaptive personalized federated learning.  _arXiv preprint arXiv:2003.13461_ , 2020.
  * Smith et al. [2017] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar.  Federated multi-task learning.  _Advances in Neural Information Processing Systems_ , 2017-December:4425–4435, 5 2017.  URL http://arxiv.org/abs/1705.10467.
  * Finn et al. [2017] Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks.  _34th International Conference on Machine Learning, ICML 2017_ , 3:1856–1868, 3 2017.  URL http://arxiv.org/abs/1703.03400.
  * Fallah et al. [2020] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar.  Personalized federated learning: A meta-learning approach.  2020\.  URL http://arxiv.org/abs/2002.07948.
  * Grimberg et al. [2020] Felix Grimberg, Mary-Anne Hartley, Martin Jaggi, and Sai Praneeth Karimireddy.  Weight erosion: An update aggregation scheme for personalized collaborative machine learning.  In _Domain Adaptation and Representation Transfer, and Distributed and Collaborative Learning_ , pages 160–169. Springer, 2020.
  * Grimberg et al. [2021] Felix Grimberg, Mary-Anne Hartley, Sai P Karimireddy, and Martin Jaggi.  Optimal model averaging: Towards personalized collaborative learning.  In _FL ICML workshop_ , 2021.
  * Donahue and Kleinberg [2020] Kate Donahue and Jon Kleinberg.  Model-sharing games: Analyzing federated learning under voluntary participation.  _arXiv preprint arXiv:2010.00753_ , 2020.  URL http://arxiv.org/abs/2010.00753.
  * Zhang et al. [2021] Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M. Alvarez.  Personalized federated learning with first order model optimization.  In _ICLR - International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=ehJqJQk9cw.
  * Li et al. [2021] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith.  Ditto: Fair and robust federated learning through personalization.  2021\.  URL http://arxiv.org/abs/2012.04221.
  * Collins et al. [2021] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai.  Exploiting shared representations for personalized federated learning.  2 2021.  URL http://arxiv.org/abs/2102.07078.
  * McMahan et al. [2016] H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017_ , 2 2016.  URL http://arxiv.org/abs/1602.05629.
  * LeCun et al. [1989] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel.  Backpropagation Applied to Handwritten Zip Code Recognition.  _Neural Computation_ , 1(4):541–551, 1989.  ISSN 0899-7667.  doi: 10.1162/NECO.1989.1.4.541.
  * Yu et al. [2017] Fisher Yu, Dequan Wang, Evan Shelhamer, and Trevor Darrell.  Deep layer aggregation.  _Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition_ , pages 2403–2412, 7 2017.  URL http://arxiv.org/abs/1707.06484.
  * Gengler et al. [2020] Damien Gengler, Martin Beaussart, and Vincent Yuan.  Personalized federated image classification using weight erosion.  2020\.  URL https://www.epfl.ch/labs/mlo/wp-content/uploads/2021/05/crpmlcourse-paper854.pdf.

## Appendix

##  Appendix A Limitations and Future Work

Non-IID distributions.

This work only focuses on certain categories of non-identical data, which
limit our evaluation of WAFFLE. Further categories like Prior probability
shift, Concept drift, and Covariate shift should be tested.

Dataset selection.

The benchmark could also have been more realistic with data sets created for
this purpose, like FEMNIST or real world data sets, which would benefit from
private personalised learning such as medical images.

Hyperparameters.

Hyperparameters were also a limitation, as some can be very sensitive (Weight
Erosion and WAFFLE) and more rigorous testing methods are necessary to find
the optimal settings.

Computational cost.

A major limitation was the number of runs we were able to perform due to
computational time limitations, especially for CIFAR10. This only allowed us
to report the results of one seed which reduces the confidence of the results.

$\Omega$ and $\Psi$.

The WAFFLE benchmark is also limited by the functions we chose for $\Omega$
and $\Psi$. Both of these functions were designed to offer versatility in
WAFFLE, so that they could be optimized for various contexts. Depending on the
category of non-identical data, it might be that some type of $\Omega$ and
$\Psi$ functions gives better results in general. Therefore, more research
should be done to find better functions depending on the problem.

Benchmark.

WAFFLE must still be benchmarked against more personalized FL methods. This
work only shows two methods (Weight Erosion and APFL). However, the
performance of APFL was lower than anticipated. It is possible that this is
because APFL adapts poorly to the tasks proposed in our setting and further
experimentation is recommended.

Robustness.

Finally, we assume honest participation of agents which may not be the case in
reality. While tests should be done to verify the strength of WAFFLE against a
data poisoning attack, we assume that WAFFLE should not be vulnerable by
design as the weighting mechanism should mitigate vulnerability by ignoring
malicious participants.

##  Appendix B Figures, Tables and Algorithms

Table 2: Repartition of labels for the distribution C between our ten agents
Label
---
|  0 |  1 |  2 |  3 |  4 |  5 |  6 |  7 |  8 |  9
Agent 0 |  0 |  0 |  0 |  0.1 |  0.2 |  0.4 |  0.2 |  0.1 |  0 |  0
Agent 1 |  0 |  0 |  0.1 |  0.2 |  0.4 |  0.2 |  0.1 |  0 |  0 |  0
Agent 2 |  0 |  0.1 |  0.2 |  0.4 |  0.2 |  0.1 |  0 |  0 |  0 |  0
Agent 3 |  0.1 |  0.2 |  0.4 |  0.2 |  0.1 |  0 |  0 |  0 |  0 |  0
Agent 4 |  0.2 |  0.4 |  0.2 |  0.1 |  0 |  0 |  0 |  0 |  0 |  0.1
Agent 5 |  0.4 |  0.2 |  0.1 |  0 |  0 |  0 |  0 |  0 |  0.1 |  0.2
Agent 6 |  0.2 |  0.1 |  0 |  0 |  0 |  0 |  0 |  0.1 |  0.2 |  0.4
Agent 7 |  0.1 |  0 |  0 |  0 |  0 |  0 |  0.1 |  0.2 |  0.4 |  0.2
Agent 8 |  0 |  0 |  0 |  0 |  0 |  0.1 |  0.2 |  0.4 |  0.2 |  0.1
Agent 9 |  0 |  0 |  0 |  0 |  0.1 |  0.2 |  0.4 |  0.2 |  0.1 |  0

In Algorithm 2, changes w.r.t. SCAFFOLD are written in red ink [Karimireddy et
al., 2020].

server input : Initial global model $\boldsymbol{x}$ and global control
variate $\boldsymbol{c}$, global step-size $\eta_{g}$, and index $i^{\star}$
of the requesting agent

agent $i$’s input : Initial local control variate $\boldsymbol{c}_{i}$, and
local step-size $\eta_{l}$

output : Model $\boldsymbol{x}$ optimized for agent $i^{\star}$

1

2
$\left(\mathbf{\alpha}^{r-2},\mathbf{\alpha}^{r-1}\right)\leftarrow(\frac{1}{N}\mathbf{1}_{N\times
1},\frac{1}{N}\mathbf{1}_{N\times 1})$

3 for  __round_ $r\leftarrow 1$ to $R$_ do

4 communicate $\left(\boldsymbol{x},\boldsymbol{c}\right)$ to all agents
$i\in\mathcal{S}$

5 on each agent  _
$i\in{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\\{1,\dots,N\\}}$
_ in parallel do

6 initialize local model $\boldsymbol{y}_{i}\leftarrow\boldsymbol{x}$

7 for  _ $k\leftarrow 1$ to $K$_ do

8 compute mini-batch gradient $g_{i}(\boldsymbol{y}_{i})$

9
$\boldsymbol{y}_{i}\leftarrow\boldsymbol{y}_{i}-\eta_{l}(g_{i}(\boldsymbol{y}_{i}))-\boldsymbol{c}_{i}+\boldsymbol{c})$

10

11
$\boldsymbol{c}_{i}^{+}\leftarrow\boldsymbol{c}_{i}-\boldsymbol{c}+\frac{1}{K\eta_{l}}(\boldsymbol{x}-\boldsymbol{y}_{i})$

12
$\left(\Delta\boldsymbol{y}_{i},\Delta\boldsymbol{c}_{i}\right)\leftarrow(\boldsymbol{y}_{i}-\boldsymbol{x},\boldsymbol{c}_{i}^{+}-\boldsymbol{c}_{i})$

13 $\boldsymbol{c}_{i}\leftarrow\boldsymbol{c}_{i}^{+}$

14 communicate
$\left(\Delta\boldsymbol{y}_{i},\Delta\boldsymbol{c}_{i}\right)$ to the server

15

16 $(\bar{\mathbf{\alpha}}^{r},\mathbf{\alpha}^{r})\leftarrow$
CalcWeights(_$N$ , $R$, $i^{\star}$, $r$, $\\{\Delta\boldsymbol{y}_{i}\\}$,
$\mathbf{\alpha}^{r-1}$, $\mathbf{\alpha}^{r-2}$_)

17 $\mathbf{\alpha}^{r-2}\leftarrow\mathbf{\alpha}^{r-1}$ and
$\mathbf{\alpha}^{r-1}\leftarrow\mathbf{\alpha}^{r}$

18
$\left(\Delta\boldsymbol{x},\Delta\boldsymbol{c}\right)\leftarrow\sum_{i\in{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\\{1,\dots,N\\}}}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bar{\alpha}^{r}_{i}}\left(\Delta\boldsymbol{y}_{i},\Delta\boldsymbol{c}_{i}\right)$

19 $\boldsymbol{x}\leftarrow\boldsymbol{x}+\eta_{g}\Delta\boldsymbol{x}$ and
$\boldsymbol{c}\leftarrow\boldsymbol{c}+{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\Delta\boldsymbol{c}}$

20

21

Algorithm 2 WAFFLE: Weighted Averaging for Personalized Federated Learning Table 3: Hyperparameters for personalized FL methods |  |  |  MNIST |  |  |  |  |  CIFAR10 |  |
---|---|---|---|---|---|---|---|---|---|---
Distr. |  A |  B |  C |  A* |  B* |  A |  B |  C |  A* |  B*
WE $(p_{d})$ |  0.006 |  0.0067 |  0.006 |  0.0088 |  0.0088 |  1.1 |  1.1 |  0.9 |  0.6 |  0.6
WAFFLE $(\Delta\Omega)$ |  3.2 |  3.2 |  3.2 |  3.2 |  3.2 |  3.2 |  3.2 |  3.2 |  3.2 |  3.2
Figure 5: Comparison of the accuracy using MNIST between Weight Erosion with
or without a distance penalty factor following different distribution of
labels and size of samples. The definition of all distributions used can be
seen in Table 4.  Table 4: Label and size sample distribution for distance
penalty experiment, example Label distribution A with Size distribution C mean
that agent 4 will have two times more data than agents 0,1,2,3,7,8,9 and data
is IID. Label
---
|  0 |  1 |  2 |  3 |  4 |  5 |  6 |  7 |  8 |  9
Label distribution |  |  |  |  |  |  |  |  |  |
A |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1
C |  0.25 |  0.25 |  0.25 |  0.25 |  0 |  0 |  0 |  0 |  0 |  0
Size distribution |  |  |  |  |  |  |  |  |  |
A |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1 |  0.1
B |  0.21 |  0.06 |  0.11 |  0.06 |  0.11 |  0.06 |  0.16 |  0.06 |  0.11 |  0.06
C |  0.1 |  0.1 |  0.1 |  0.1 |  0.2 |  0.05 |  0.05 |  0.1 |  0.1 |  0.1
Figure 6: Evolution of the accuracy on MNIST, average over five seeds of the
best accuracy obtained up to each turn.  Figure 7: Evolution of the accuracy
on CIFAR10, best accuracy obtained up to each turn.  Figure 8: Evolution value
of omega with $\Delta\Omega$ = 3.2