secAdditional References

# Neuro-algorithmic Policies enable Fast Combinatorial Generalization

Marin Vlastelica, Michal Rolínek, Georg Martius
Max Planck Institute for Intelligent Systems
{mvlastelica, mrolinek<EMAIL_ADDRESS>
###### Abstract

Although model-based and model-free approaches to learning the control of
systems have achieved impressive results on standard benchmarks,
generalization to task variations is still lacking. Recent results suggest
that generalization for standard architectures improves only after obtaining
exhaustive amounts of data. We give evidence that generalization capabilities
are in many cases bottlenecked by the inability to generalize on the
combinatorial aspects of the problem. Furthermore, we show that for a certain
subclass of the MDP framework, this can be alleviated by neuro-algorithmic
architectures.

Many control problems require long-term planning that is hard to solve
generically with neural networks alone. We introduce a neuro-algorithmic
policy architecture consisting of a neural network and an embedded time-
dependant shortest path solver. These policies can be trained end-to-end by
blackbox differentiation. We show that this type of architecture generalizes
well to unseen variations in the environment already after seeing a few
examples.

##  1 Introduction

One of the central topics in machine learning research is learning control
policies for autonomous agents. Many different problem settings exist within
this area. On one end of the spectrum are imitation learning approaches, where
prior expert data is available and the problem becomes a supervised learning
problem. On the other end lie approaches that require interaction with the
environment to obtain data for policy extraction, posing the problem of
exploration. Most Reinforcement Learning (RL) algorithms fall into the latter
category. In this work, we concern ourselves primarily with the setting where
limited expert data is available, and a policy needs to be extracted by
imitation learning.

Independently of how a policy is extracted, a central question of interest is:
how well will it generalize to variations in the environment and the task?
Recent studies have shown that standard deep RL methods require exhaustive
amounts of exposure to environmental variability before starting to generalize
[14].

There exist several approaches addressing the problem of generalization in
control. One option is to employ model-based approaches that learn a
transition model from data and use planning algorithms at runtime or to
improve value-learning. This has been argued to be the best strategy in the
presence of an accurate model and sufficient computation time [15].
Furthermore, one can use the transition model alongside a reward model to
generate offline data to improve value function learning [49; 26]. However,
learning a precise transition model is often harder than learning a policy.
The transition model often has a much larger dimensionality than the policy
since it needs to model aspects of the environmental dynamics that are perhaps
irrelevant for the task. This is particularly true for learning in problems
with high-dimensional inputs, such as raw images. In order to alleviate this
problem, learning specialized or partial models has shown to be a viable
alternative, e.g. in MuZero [45].

We propose to use recent advances in making combinatorial algorithms
differentiable in a blackbox fashion[52] to train neuro-algorithmic policies
with embedded planners end-to-end. More specifically, we use a time-dependent
shortest path (TDSP) planner acting on a temporally evolving graph generated
by a deep network from the inputs. By learning the time-evolving costs of the
graph our method builds a specific model of the system that is sufficient for
planning. We demonstrate the effectiveness of this approach in an offline
imitation learning setting where a few expert trajectories are provided. Due
to the combinatorial generalization capabilities of planners, our learned
policy is able to generalize to new variations in the environment out of the
box and needs orders of magnitude fewer samples than naive learners. Using
neuro-algorithmic architectures facilitates generalization by shifting the
combinatorial aspect of the problem to efficient algorithms, while using
neural networks to extract a good representation for the problem at hand. They
have the potential to endow artificial agents with the main component of
intelligence, the ability to reason.

Our contributions can be summarized as follows: (i) We show that combinatorial
inductive biases implemented through neuro-algorithmic policies can be used to
tackle the generalization problem in reinforcement learning. (ii) We show that
architectures embedding TDSP solvers are applicable beyond goal-reaching
environments. (iii) We demonstrate learning neuro-algorithmic policies in
dynamic game environments from images.

##  2 Related Work

#### Planning

There exist multiple lines of work aiming to improve classical planning
algorithms such as improving sampling strategies of Rapidly-exploring Random
Trees [23; 12; 30]. Similarly, along this direction, Kumar et al., [29]
propose a conditional VAE architecture for sampling candidate waypoints.
Orthogonal to this are approaches that learn representations such that
planning is applicable in the latent space. Hafner et al., [24] employ a
latent multistep transition model. Savinov et al., [43] propose a semi-
parametric method for mapping observations to graph nodes and then applying a
shortest path algorithm. Asai and Fukunaga, [6]; Asai and Kajino, [7] use an
autoencoder architecture to learn a discrete transition model suitable for
classical planning algorithms. Li et al., [31] learn compositional Koopman
operators with graph neural networks mapping to a linear dynamics latent
space, which allows for fast planning. Chen et al., [13]; Amos et al., [3]
perform efficient planning by using a convex model formulation and convex
optimization. Alternatively, the replay buffer can be used as a nonparametric
model to select waypoints [20] or in an MPC fashion [11]. None of these
methods perform differentiation through the planning algorithm to learn better
latent representations.

#### Differentiation through planning

Embedding differentiable planners has been proposed in previous works, e.g. in
the continuous case with CEM [4; 9]. Wu et al., [54] use a (differentiable)
recurrent neural network as a planner. Tamar et al., [51] use a differentiable
approximation of the value iteration algorithm to embed it in a neural
network. Silver et al., 2017b [47] differentiate through a few steps of value
prediction in a learned MDP to match the externally observed rewards. Srinivas
et al., [48] use a differentiable forward dynamics model in latent space.
Karkus et al., [28] suggest a neural network architecture embedding MDP and
POMDP solvers and during the backward pass, they substitute the algorithms by
learned approximations. In comparison, we do not perform any relaxation or
approximation of the planner itself and we learn interpretable time-dependent
costs of the latent planning graph based on expert demonstrations by
differentiating through the planner. Similarly to our work, Yonetani et al.,
[56] embed an A∗ algorithm into a neural network, but in comparison, their
method does not operate with time-dependent costs, subgoal selection and does
not provide a policy for closed-loop control.

#### Inverse reinforcement learning and imitation learning

Uncovering the expert’s objective function from demonstrations has been a
central topic in reinforcement learning [35]. Our method is connected to
inverse reinforcement learning in the sense that we learn the objective
function that the expert optimizes to extract an optimal policy, also called
apprenticeship learning [1; 34; 2]. What separates our approach is that the
inferred costs are inherently part of the learned neuro-algorithmic policy in
conjunction with the applied planner on the costs.

Our method is an offline imitation learning method, but since we propose an
end-to-end trainable policy, it is naturally extendable to the online case
with a method such as DAgger [42] or other online reinforcement learning
methods augmented with expert datasets [40; 25].

#### Offline model-based reinforcement learning

Model-based methods have shown promise by facilitating better generalization
[26]. These approaches fall into two camps: using models to extract a policy
in a Dyna-style approach [49; 26; 50; 55; 27], or incorporating the model in a
planning loop, i.e. model-predictive control [22; 39; 37; 46; 38]. In this
work, we consider the latter case where an implicit transition model is
“hidden” within the predicted time-dependent costs.

#### Combinatorial algorithms in end-to-end trainable networks

We suggest a hybrid policy consisting of a neural network and an accompanying
expert (shortest path) discrete solver that is trainable end-to-end.
Incorporating expert discrete solvers into end-to-end trainable architectures
is a topic with exciting recent developments. For the simpler setup of
comparing to ground-truth values on the solver output, numerous frameworks
have been suggested such as the “predict-and-optimize” framework and its
variants [19; 16; 32]. Moreover, specializations for concrete cases such as
sparse structured inference [36], logical satisfiability [53], submodular
optimization [18], or mixed integer programming [21] have been proposed.

We are interested in the harder case of providing an entirely hybrid
architecture which may use the solver at intermediate levels and is trainable
end-to-end. For this case, two approaches have recently emerged [52; 8].
Vlastelica et al., [52] introduce an efficient implicit piecewise linear
interpolation scheme, while Berthet et al., [8] introduce an estimation of a
Gaussian smoothing of the piecewise constant function. The approach from
Vlastelica et al., [52] is especially appealing, since it allows for uses in
which the solver is the computational bottleneck. By formulating the control
problem as a time-dependent shortest path problem (TDSP), we show that the
framework of Vlastelica et al., [52] is applicable in specific control
settings.

##  3 Markov Decision Processes and Shortest Paths

We consider the problem formulation in the context of a Markov Decision
Process (MDP). We concern ourselves with policies that solve a time-dependent
shortest path problem on a latent graph representation related to that MDP. We
start with a class of MDPs that can be directly mapped to a shortest path
problem and construct this mapping. Then we consider conditions that allow for
a reduced latent graph where the optimal policy follows a time-dependent
shortest path problem.

First, we consider discrete MDPs with deterministic transitions. In addition,
we follow a goal-conditioned setting [44]. This is used in sequential decision
making problems where a specific terminal state has to be reached.

######  Definition 1.

A goal-conditioned Markov Decision Process (gcMDP) ${\mathcal{M}}$ is defined
by the tuple (${\mathcal{S}}$, ${\mathcal{A}}$, $p$, $g$, $r$), where
${\mathcal{S}}$ is the state space, ${\mathcal{A}}$ the action space,
$p(s^{\prime}\mid a,s)$ the probability of making the transition
$s\in{\mathcal{S}}\rightarrow s^{\prime}\in{\mathcal{S}}$ when taking the
action $a\in{\mathcal{A}}$, $g\in{\mathcal{S}}$ is the goal,
$r(s,a,s^{\prime},g)$ the reward obtained when transitioning from state $s$ to
$s^{\prime}$ while taking action $a$ and aiming for goal $g$.

In episodic reinforcement learning, the objective is to find a policy that
maximizes the return $G=\sum_{t=0}^{T}r_{t}$ of such a process. In gcMDPs, the
reward is such that the maximal return can be achieved by reaching the goal
state $g$, which is also the terminal state.

######  Definition 2.

A discrete and deterministic goal-conditioned Markov Decision Process
(ddgcMDP) $\mathcal{\ddot{M}}$ is a gcMDP with discrete and finite state space
${\mathcal{S}}$ and action space ${\mathcal{A}}$ with deterministic
transitions, i.e. $p(s^{\prime}\mid a,s)$ is one-hot.

Let us consider the following graph representation $(G,v_{s},v_{g})$: a
weighted graph $G=(V,E,C)$ together with start vertex $v_{s}$ and goal vertex
$v_{g}$, where $V$ is the vertex set, $E\subset(V\times V)$ is the edge set,
and $C\in\mathbb{R}_{+}^{|E|}$ is the cost matrix with positive entries. We
write $C(e)$ for the cost of edge $e$. In addition, we consider an inverse
model $\psi\colon E\to\mathcal{A}$, associating an edge to an action.

A direct translation of a ddgcMDP to a graph is given by a bijective
correspondence $\phi\colon\mathcal{S}\to V$ between states and vertices, for
each possible transition $(s,a,s^{\prime})$ there is an edge
$e=(\phi(s),\phi(s^{\prime}))$ with $\psi(e)=a$, and the goal vertex is
$v_{g}=\phi(g)$. The cost matrix $C$ takes the role of the reward function
with $C(\phi(s),\phi(s^{\prime}))=c_{\mathrm{max}}-r(s,a,s^{\prime},g)$ where
$c_{\mathrm{max}}$ is an upper bound on the reward ensuring positive cost
values. Due to the deterministic nature of the ddgcMDP, the optimal policy
yields a path with maximal return (sum of rewards) from start $s_{0}$ to goal
$g$, which now coincides with the shortest path according to $C$ by
definition.

###  3.1 Factorized MDPs and Time-dependent Shortest Path

In many interesting environments, the state space is exponentially large, e.g.
due to independent dynamics of entities. As an illustrative example, consider
the game environment shown in Fig. LABEL:fig:pipeline and 2 – an agent with
moving obstacles. For such environments, the corresponding graph would become
intractably large. We now define conditions under which the size of the graph
can be drastically reduced, lwithout losing the property that its shortest
path solves the MDP. To this end, we assume the state space can be factorized
as ${\mathcal{S}}={\mathcal{S}}_{A}\times{\mathcal{S}}_{E}$, where
${\mathcal{S}}_{A}$ is affected only by the actions and ${\mathcal{S}}_{E}$ by
the environment dynamics independent of the actions. We write the
decomposition of each state as $s=(s^{a},s^{e})$. The mapping from state space
to graph vertices $\phi\colon\mathcal{S}\to V$ is bijective only w.r.t.
${\mathcal{S}}_{A}$ and ignores ${\mathcal{S}}_{E}$, i.e. $\forall
s^{a}\in{\mathcal{S}}_{A}\;\exists\,v\in
V\colon\phi(s^{a},s^{e})=v\,,\forall\,s^{e}\in{\mathcal{S}}_{E}$. For brevity,
we write $\phi(s^{a})$. For instance, using ${\mathcal{S}}_{A}$ as the agent’s
positions on a discrete grid results in just as many vertices as there are
positions of the agent.

Next, we show how a solution to this factorized ddgcMDP can be found by
solving a time-dependent shortest path (TDSP) problem on the reduced graph.
The cost matrix $C$ is now a time-dependent quantity, i.e.
$C\in\mathbb{R}_{+}^{{T\times|E|}}$ that assigns every edge $e$ a positive
cost value for each time $t\in\\{1,2,\dots,T\\}$ of the planning
horizon/episode. To ease the notation, we write $C^{t}(e)$ for the cost of
edge $e$ at time $t$. The TDSP problem is defined as reaching the goal vertex
$v_{g}$ within at most $T$ steps when starting at time step $1$ in the start
vertex $v_{s}$.

Two situations need to be modeled by the cost:111Note that learning the
_exact_ costs is not necessary, since we are only interested in optimal
trajectories. a) the environment dynamics can make a transition along an
unavailable edge $e$ (for instance, an obstacle is moving in the way), then
the cost should be infinite: $C^{t}(e)=\infty$, and b) the environment
dynamics change the reward, thus for all other edges we have

$\displaystyle
C^{t}(\phi(s^{a}),\phi(s^{\prime}{{}^{a}}))=c_{\mathrm{max}}-r((s^{a},s_{t}^{e}),a,(s^{\prime}{{}^{a}},s_{t}^{e}))$
(1)

where $a=\psi(\phi(s^{a}),\phi(s^{\prime}{{}^{a}}))$. Again, with this
construction, the time-dependent shortest path solution coincides with the
optimal policy of the specific ddgcMDP. In fact, we can also deal with
stochastic environment dynamics as long as the action’s effect on the state
stays deterministic. This changes the reward term in Eq. 1 to an expectation
over environment dynamics:

$\displaystyle\underset{s_{t}^{e}}{\mathbb{E}}[r((s^{a},s_{t}^{e}),a,(s^{\prime}{{}^{a}},s_{t}^{e}))].$

In our architecture, described below, the policy generates a latent graph at
every observation/state and solves a ddgcMDP to optimality every time step
following a model predictive control approach using receding horizon planning.

##  4 Shortest Path Algorithm and its Differentiation

We aim to predict the cost matrix $C$ via a deep neural network and train the
entire system end-to-end via gradient descent on expert trajectories, as
illustrated in Fig. LABEL:fig:pipeline. We will employ an efficient
implementation of Dijkstra’s algorithm for computing the shortest path. For
differentiation, we rely on the framework for blackbox differentiation of
combinatorial solvers [52].

#### Time-dependent shortest path with vertex costs.

In our neuro-algorithmic policy, we use the time-dependent-shortest-path
(TDSP) formulation based on vertex-costs instead of edge-costs, due to the
reduction in cost matrix size by a factor of $|{\mathcal{A}}|$. The TDSP
solver has as input the graph $G(V,E,C)$ with time-dependent cost matrix
$C\in\mathbb{R}_{+}^{{T\times|V|}}$ and a pair of vertices $v_{s},v_{g}\in V$
(start and goal). We write $C^{t}(v)$ for the cost of vertex $v$ at time $t$.

This version of the shortest path problem can be solved by executing the
Dijkstra’s shortest path algorithm222Even though the classical formulation of
Dijkstra’s algorithm is edge-based, all of its properties hold true also in
this vertex-based formulation. [17] on an augmented graph. In particular, we
set

$\displaystyle V^{*}$ $\displaystyle=\\{(v,t)\colon v\in V,t\in[1,T]\\}\quad$
$\displaystyle E^{*}$
$\displaystyle=\\{((v_{1},t),(v_{2},t\\!+\\!1))\colon(v_{1},v_{2})\in
E^{\circlearrowleft},t\in[1,T\\!-\\!1]\\},$

where the cost of vertex $(v_{i},t)\in V^{*}$ is simply $C^{t}(i)$ and
$E^{\circlearrowleft}$ is the original edge set $E$ appended with all self-
loops. This allows to “wait” at a fixed vertex $v$ from timestep $t$ to
timestep $t+1$. In this graph, the task is to reach the vertex $(v_{g},T)$
from $(v_{s},1)$ with the minimum traversal cost.

###  4.1 Applicability of Blackbox Differentiation

The framework presented in [52] turns blackbox combinatorial solvers into
neural network building blocks. The provided gradient is based on a piecewise
linear interpolation of the true piecewise constant (possibly linearized) loss
function, see Fig. 1. The _exact_ gradient of this linear interpolation is
computed efficiently via evaluating the solver on only one additional instance
(see Algo. 1).

function ForwardPass($C$, $v_{s}$, $v_{g}$)

$Y:=\text{TDSP}(C,v_{s},v_{s})$ // Run Dijkstra’s algo.

save $Y$, $C$, $v_{s}$, $v_{e}$ // Needed for backward pass

return $Y$

function BackwardPass($\nabla L(Y)$, $\lambda$)

load $Y$, $C$, $v_{s}$, $v_{e}$

$C_{\lambda}:=C+\lambda\nabla L(Y)$ // Calculate modified costs

$Y_{\lambda}:=\text{TDSP}(C_{\lambda},v_{s},v_{g})$ // Run Dijkstra’s algo.

return $\frac{1}{\lambda}\bigl{(}Y_{\lambda}-Y\bigr{)}$

Algorithm 1 Forward and backward Pass for the shortest-path algorithm Figure
1: Differentiation of a piecewise constant loss resulting from incorporating a
combinatorial solver. A two-dimensional section of the loss landscape is shown
(left) along with two differentiable interpolations of increasing strengths
(middle and right).

To apply this differentiation scheme, the solver at hand needs to have a
formulation in which it minimizes an inner-product objective (under arbitrary
constraints). To that end, for a given graph $G=(V,E,C)$ as describe above, we
define $Y\in\\{0,1\\}^{T\times|V|}$ an indicator matrix of visited vertices.
In particular, $Y_{i}^{t}=1$ if and only if vertex $v_{i}$ is visited at time
point $t$. The set of such indicator matrices that correspond to valid paths
in the graph $(V^{*},E^{*})$ from start to goal will be denoted as
${\mathcal{Y}}\coloneqq\mathrm{Adm}(G,v_{s},v_{g})$. The time-dependent
shortest path optimization problem can be then rewritten as

$\displaystyle\mathrm{TDSP}(C,v_{s},v_{g})$
$\displaystyle=\operatorname*{arg\,min}_{Y\in{\mathcal{Y}}}\sum_{(i,t)}Y_{i}^{t}C_{i}^{t},$
(2)

where ${\mathcal{Y}}=\mathrm{Adm}(G,v_{s},v_{g})$. This is an inner-product
objective and thus the theory from [52] applies. and allows us to learn the a
cost-matrix generating network with gradient descent, as described below.

###  4.2 Cost Margin

The supervision-signal for the costs $C$ of the latent graph is only
indirectly given via the shortest path solutions. At training time there is no
incentive to make these costs robust to small misestimations. Thus, inducing a
margin on costs can be beneficial for generalization. Similarly to Rolínek et
al., [41], where it was used in the context rank-based metric learning, we
induce a margin $\alpha$ by increasing the cost of the ground-truth path and
decreasing the rest in the training stage of the algorithm:

$c_{i}^{t}=\begin{cases}c_{i}^{t}+\frac{\alpha}{2}&\text{if $(v_{i},t)\in
Y^{*}$}\\\ c_{i}^{t}-\frac{\alpha}{2}&\text{if $(v_{i},t)\notin
Y^{*}$}\end{cases}\;\forall\,(v_{i},t)\in V^{*}.$ (3)

##  5 Neuro-algorithmic Policy Framework

We propose the Neuro-algorithmic Policy (NAP) framework, which is an end-to-
end trainable deep policy architecture embedding an algorithmic component
using the afore-mentioned techniques. Following the definitions from Sec. 3,
we concern ourselves with learning the mapping
$\varphi_{\theta}:\mathcal{S}\mapsto\mathbb{R}^{|V|\times T}\times V\times V$,
i.e. mapping from MDP states to cost matrices and respective start and end
vertices for the $\mathrm{TDSP}$ problem 333We hold the graph structure fixed,
namely the set of edges $E$ and learn the costs $C$. Therefore we replace $G$
with costs $C$ in the $\mathrm{TDSP}$ solver to simplify the notation. of
planning horizon $T$. This enables us to construct the policy

$\pi_{\theta}\coloneqq\psi\circ\mathrm{TDSP}\circ\varphi_{\theta}.$ (4)

The mapping $\varphi_{\theta}$ can be decomposed into $\varphi_{\theta}^{c}$
(_cost-predictor_), $\varphi_{\theta}^{s}$ (_start-predictor_),
$\varphi_{\theta}^{g}$ (_goal-predictor_), i.e. mappings from state to costs,
start vertex and goal vertex. In practice, instead of learning
$\varphi_{\theta}^{s}$ and $\varphi_{\theta}^{g}$ directly, we learn the
conditional probability densities $p_{\theta}^{s}(v|s)$ and
$p_{\theta}^{g}(v|s)$.

For this work, we examine the application of neuro-algorithmic policies to the
imitation learning setting, where we have access to trajectories $\tau$
sampled from the expert policy distribution $\eta(\pi^{*})$. Given a fixed
planning horizon $T$, the objective that we consider consists of three main
parts, the latent cost term, start vertex term and goal vertex term. The
latent cost objective is defined as

$J^{C}(\theta,T)=\operatorname{\mathop{\mathbb{E}}_{\tau_{t:t+T}\sim\eta({\pi^{*}})}}\big{[}\mathrm{d^{H}}(\mathrm{TDSP}(\varphi_{\theta}(\tau_{t})),\phi^{\prime}(\tau))\big{]},$
(5)

where $\mathrm{d^{H}}(\cdot,\cdot)$ denotes the Hamming distance between
predicted and expert paths in the latent graph, and
$\phi^{\prime}\colon\mathcal{S}\mapsto V$ the mapping of the expert-visited
states to latent graph vertices. The second part of the objective is a
standard cross-entropy term for the start and goal vertices that allows us to
train a _start-_ and _goal-predictor_ :

$\displaystyle\begin{split}J^{P}(\theta,T)=\operatorname{\mathop{\mathbb{E}}_{\tau_{t:t+T}\sim\eta({\pi^{*}})}}\big{[}&-\log
p_{\theta}^{s}(\phi^{\prime}(\tau_{t})|\tau_{t})\\\ &-\log
p_{\theta}^{g}(\phi^{\prime}(\tau_{t+T})|\tau_{t})\big{]}.\end{split}$ (6)

We assume access to $\phi^{\prime}$ at training time in order to map the
expert to the latent graph for calculation of $J^{C}$ and $J^{P}$. Finally, we
optimize for the sum of $J^{C}$ and $J^{P}$.

We utilize a concrete architecture consisting of two main components: a
backbone ResNet18 architecture (without the final fully connected layers (a
detailed description is available in Sec. C in the supplementary) and the
shortest path solver, see Fig. LABEL:fig:pipeline. At each time step, the
policy receives two images concatenated channel-wise from which it predicts
the cost tensor $C$ for the planning horizon $T$ with the _cost-predictor_ ,
the start vertex $v_{s}$ and goal vertex $v_{g}$ with the _goal-predictor_ ,
explained below.

The policy is used in a model-predictive control fashion, i.e. at execution
time, we predict the plan $Y$ for horizon $T$ at each time step and execute
the first action from the plan.

###  5.1 Global and Local Goal Prediction

In order to apply the solver to the learned latent graph representation, we
need to map the current state of the environment to appropriate start and goal
vertices ($v_{s},v_{g}$). To this end, we employ a second ResNet18 similar to
the _cost-predictor_ that approximates $p^{s}(v|s)$ and $p^{g}(v|s)$, i.e. the
start and goal conditional densities.

At training time, given the expert trajectories, we have access to the mapping
$\phi^{\prime}$ that maps the expert trajectory to the latent graph. In the
_global_ setting, the last position of the expert is the goal $v_{g}$,
corresponding to, for instance, the jewel in Crash Jewel Hunt which is also
the terminal state, see Fig. Fig. 2.

In the _local_ setting, we expect the end vertex to be an intermediate goal
(“collect an orb”), which effectively allows for high-level planning
strategies, while the low-level planning is delegated to the discrete solver.
In this case, the positively labeled supervision at time $t$ are all locations
of the (expert) agent between step $t+T$ and $t+2T$.

The local setting allows to limit the complexity of our method, which grows
with the planning horizon. This is also a trade-off between the combinatorial
complexity solved by the $TDSP$ solver and the goal predictor. Ideally, the
planning horizon $T$ used for the cost-prediction is long enough to capture
the combinatorial intricacies of the problem at hand, such as creating detours
towards the goal in the case of future dangerous states, or avoiding dead-ends
in a maze.

This formulation makes our architecture akin to hierarchical methods similar
to Blaes et al., [10]; Nachum et al., [33], and allows for solving tasks that
are not typical goal-reaching problems, such as the Chaser environment.

##  6 Experiments

To validate our hypothesis that embedding planners into neural network
architectures leads to better generalization in control problems, we consider
several procedurally generated environments (from the ProcGen suite [14] and
Crash Jewel Hunt) with considerable variation between _levels_.

We compare with the following baselines: a standard behavior cloning (BC)
baseline using a ResNet18 architecture trained with a cross-entropy
classification loss on the same dataset as our method; the PPO algorithm as
implemented in Cobbe et al., [14] and data-regularized actor-critic (DrAC)
[5]. A comparison to DrAC is especially interesting, since the method claims
to improve generalization by applying optimized data augmentations. As there
are multiple variations of data augmentation suggested by anu et al., [5], we
run all of them and select the best result as DrAC∗ for performance comparison
to our method444We defer a more detailed description of DrAC along with
performance plots to Sec. F. We also ablate the _start-_ and _goal-predictor_
and replace with the ground truth vertices, this serves as an upper baseline
for NAP which we denote with NAP∗. More details on the training procedure and
the hyperparameters can be found in the supplementary Sec. D.

(a) Frame | (b) Game logic
---|---
|
(c) Cost prediction example
Figure 2: The Crash Jewel Hunt environment. The goal for the fox, see (a), is
to obtain the jewel in the right most column, while avoiding the moving wooden
boxes (arrows in (b)). When the agent collides with a wooden box it instantly
fails to solve the task. We observe that the predictions of the costs in (c)
are highly interpretable, corresponding to (future) movements of the boxes.

For the experimental validation, we aim to anwser the following questions: (i)
Can NAP be trained to perform well in procedurally generated environments?
(ii) Can NAP generalize in a low data regime, i.e. after seeing only few
different levels? (iii) Can we also solve non-goal-reaching environments?

Crash $5\times 5$ | Crash $5\times 10$ | Maze | Leaper(Grid)
---|---|---|---
|  |  |
Figure 3: Performance on unseen test levels for a different number of training
levels. We observe that NAP already shows signs of generalization after only
being trained on 100 levels.

###  6.1 Crash Jewel Hunt

For proof of concept, we first consider an environment we constructed to test
NAP, called Crash Jewel Hunt which can be seen in Fig. 2. The environment
corresponds to a grid-world of dimensions $h\times w$ where the goal is to
move the agent (Fox) from an arbitrary start position in the left-most column
to the goal position (jewel) arbitrarily positioned in the right-most column.
Between the agent and the goal are obstacles, wooden boxes that move downwards
(with cyclic boundary conditions) with velocities that vary across levels but
not within a level, see Fig. 2 (right). At each time step, the agent can
choose to move horizontally or vertically in the grid by one cell or take no
action.

To make the task challenging, we sample distinct environment configurations
for the training set and the test set, respectively. More concretely, we vary
the velocities, sizes and initial positions of the boxes as well as the start
and goal positions.

###  6.2 ProcGen Benchmark

In addition to the jewel hunt environment, we evaluate our method on the hard
version Maze, Leaper and Chaser environments from the ProcGen suite [14]. We
have chosen these environments because their structure adheres to our
assumptions. For Leaper, we modified the environment such that grid-world
dynamics apply (Leaper(Grid)).

The Maze and the Leaper(Grid) tasks have a static goal whose position only
varies across levels, whereas the Chaser requires collection of all orbs
without contact with the spiders, so the local goals need to be inferred on
the fly. The Chaser environment is also particularly challenging as even the
expert episodes require on average 150 steps, most of which carry the risk of
dying. For this reason, we used three human expert trajectories per level.

###  6.3 Results

We train our method (NAP) and the imitation learning baseline until saturation
on a training set, resulting in virtually $100\%$ success rate when evaluating
on train configurations in the environment. For the PPO baseline we use the
code from [14] and provide also two subsequent frames and 200M time steps for
training. For the DrAC baselines we use the code from anu et al., [5]. For our
method, we also report performance of a version with access to the true start
and end-point prediction (NAP∗), with the exception of the Chaser where true
goals are not well-defined.

(a) Performance on Chaser

(b) Short-horizon plans
---
|

Figure 4: Performance of NAP against PPO on the Chaser environment (a) trained on 10, 20, 50, 100 and 200 levels. In (b) we show the short-horizon plans (white) of the agent (blue) at step 5 and 110 in the environment. Crash 5x5 | Crash 5x10 | Maze | Leaper(Grid)
---|---|---|---
|  |  |
Figure 5: Test success rate of our method with different horizon lengths. The
solver assumes that the last horizon step costs remain to infinity. In this
sense, the horizon of length 1 corresponds to a static solver.

In Fig. 3, we show the performance of the methods when exposed to a different
number of levels at training time. As reported in [14], the baselines have a
large generalization gap, and also poor performance when $<10\,000$ levels are
seen. We find that NAP shows strong generalization performance already for
$<500$ levels. In some environments, such as Maze we obtain near $80\%$
success rate already with just $100$ levels which is never reached by PPO or
DrAC even after seeing a $1000$ levels. Our NAP trained on $1000$ levels
reaches the same performance as PPO trained on $200\,000$ levels of the Maze
environment. For Crash Jewel Hunt $5\times 5$, already with 30 trajectories a
third of the $1000$ test-levels can be solved, while the behavior cloning
baseline manages less than $50$ out of the $1000$.

###  6.4 Sensitivity to the Planning Horizon

We provide a sensitivity analysis of the performance with different planning
horizons. Our results indicate that longer horizons benefit environments with
increased dynamical interactions. As apparent from Fig. 5, our method
outperforms the behavior cloning baseline in all of the environments, the gap
between the methods being correlated with the complexity of the environment.
It can also be seen that making the planning horizon smaller in environments
with dynamics hurts performance.

On the other hand, for environments with no dynamics, such as the maze
environment, there is no benefit in using a longer planning horizon, as
expected. Nevertheless, there is still strong performance gain in
generalization when using NAP as opposed to vanilla imitation learning from
expert trajectories thanks to the introduced combinatorial inductive bias.

###  6.5 Path Optimality

Leaper(Grid) | Maze
---|---
|
Figure 6: Distributions of test level episode lengths after training on $500$
levels for the Leaper(Grid) and Maze environments over 3 seeds. We observe
that NAP tends to take shorter paths than DrAC∗ and PPO. We also observe that
the median episode length for the baselines in the Maze task is located in the
upper part, which corresponds to unsolved levels.

We have observed that introducing combinatorial structural biases with a
method such as NAP can benefit generalization greatly. Another interesting
question is how optimal the paths taken by NAP are. Even though NAP has an
advantage in having access to expert optimal or near-optimal trajectories,
this doesn’t necessarily translate to being able to act optimally in unseen
situations.

For this analysis we have plotted the episode length distribution across runs
for NAP, DrAC∗ and PPO, which can be seen in Fig. 6. This shows us that NAP
generalizes with more optimal (shorter) paths in comparison to DrAC∗ and PPO,
even when compared to levels solved by all of the methods. Even in the
Leaper(Grid) environment, where DrAC∗ outperforms NAP by measure of success
rate, NAP still outperforms DrAC∗ in terms of the path lengths.

##  7 Discussion

We have shown that hybrid neuro-algorithmic policies consisting of deep
feature extraction and a shortest path solver – made differentiable via
blackbox differentiation [52] – _enable learning policies that generalize to
unseen environment settings in the low-data regime_. Hybrid architectures are
a stepping stone towards use of better inductive biases that enable stronger
generalization. In NAP, the inductive bias that we impose is the topology of
the latent planning graph in conjunction with a planning algorithm.
Introducing the shortest-path solver as a module shifts the combinatorial
complexity of the planning problem to efficient algorithmic implementations
while alleviating the learning of good representations for planning.

Although there is a clear benefit in using NAP, the method comes with certain
caveats. We assume that the topological structure (i.e. that there is an
underlying grid structure with a set of 5 actions) of the latent planning
graph is known a priori. Furthermore, we assume that the structure of the
latent graph is fixed and not dynamically changing over time, i.e. that each
available action at a vertex corresponds to the same edge. Any results
allowing to abandon some of these assumption will vastly increase
applicability of this method and should be of immediate interest.

##  8 Acknowledgment

We thank the International Max Planck Research School for Intelligent Systems
(IMPRS-IS) for supporting Marin Vlastelica. We acknowledge the support from
the German Federal Ministry of Education and Research (BMBF) through the
Tübingen AI Center (FKZ: 01IS18039B).

## References

  * Abbeel and Ng, [2004] Abbeel, P. and Ng, A. Y. (2004).  Apprenticeship learning via inverse reinforcement learning.  In Proceedings of the twenty-first international conference on Machine learning, page 1.
  * Aghasadeghi and Bretl, [2011] Aghasadeghi, N. and Bretl, T. (2011).  Maximum entropy inverse reinforcement learning in continuous state spaces with path integrals.  In 2011 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 1561–1566. IEEE.
  * Amos et al., [2017] Amos, B., Xu, L., and Kolter, J. Z. (2017).  Input convex neural networks.  In International Conference on Machine Learning, pages 146–155.
  * Amos and Yarats, [2020] Amos, B. and Yarats, D. (2020).  The differentiable cross-entropy method.  In International Conference on Machine Learning, ICML.
  * anu et al., [2020] anu, R., Goldstein, M., Yarats, D., Kostrikov, I., and Fergus, R. (2020).  Automatic data augmentation for generalization in deep reinforcement learning.  arXiv preprint arXiv:2006.12862.
  * Asai and Fukunaga, [2017] Asai, M. and Fukunaga, A. (2017).  Classical planning in deep latent space: Bridging the subsymbolic-symbolic boundary.  arXiv preprint arXiv:1705.00154.
  * Asai and Kajino, [2019] Asai, M. and Kajino, H. (2019).  Towards stable symbol grounding with zero-suppressed state autoencoder.  In Proceedings of the International Conference on Automated Planning and Scheduling, volume 29, pages 592–600.
  * Berthet et al., [2020] Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.-P., and Bach, F. (2020).  Learning with differentiable perturbed optimizers.  arXiv preprint arXiv:2002.08676.
  * Bharadhwaj et al., [2020] Bharadhwaj, H., Xie, K., and Shkurti, F. (2020).  Model-predictive control via cross-entropy and gradient-based optimization.  arXiv preprint arXiv:2004.08763.
  * Blaes et al., [2019] Blaes, S., Vlastelica Pogančić, M., Zhu, J., and Martius, G. (2019).  Control what you can: Intrinsically motivated task-planning agent.  In Advances in Neural Information Processing Systems 32, pages 12541–12552. Curran Associates, Inc.
  * Blundell et al., [2016] Blundell, C., Uria, B., Pritzel, A., Li, Y., Ruderman, A., Leibo, J. Z., Rae, J., Wierstra, D., and Hassabis, D. (2016).  Model-free episodic control.  arXiv preprint arXiv:1606.04460.
  * Burget et al., [2016] Burget, F., Bennewitz, M., and Burgard, W. (2016).  Bi 2 rrt*: An efficient sampling-based path planning framework for task-constrained mobile manipulation.  In 2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 3714–3721. IEEE.
  * Chen et al., [2018] Chen, Y., Shi, Y., and Zhang, B. (2018).  Optimal control via neural networks: A convex approach.  In International Conference on Learning Representations.
  * Cobbe et al., [2019] Cobbe, K., Hesse, C., Hilton, J., and Schulman, J. (2019).  Leveraging procedural generation to benchmark reinforcement learning.  arXiv preprint:1912.01588.
  * Daw et al., [2005] Daw, N. D., Niv, Y., and Dayan, P. (2005).  Uncertainty-based competition between prefrontal and dorsolateral striatal systems for behavioral control.  Nature neuroscience, 8(12):1704–1711.
  * Demirovic et al., [2019] Demirovic, E., Stuckey, P. J., Bailey, J., Chan, J., Leckie, C., Ramamohanarao, K., and Guns, T. (2019).  Predict+optimise with ranking objectives: Exhaustively learning linear functions.  In Kraus, S., editor, Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI 2019, Macao, China, August 10-16, 2019, pages 1078–1085. ijcai.org.
  * Dijkstra, [1959] Dijkstra, E. W. (1959).  A note on two problems in connexion with graphs.  Numer. Math., 1(1):269–271.
  * Djolonga and Krause, [2017] Djolonga, J. and Krause, A. (2017).  Differentiable learning of submodular models.  In Advances in Neural Information Processing Systems, pages 1013–1023.
  * Elmachtoub and Grigas, [2017] Elmachtoub, A. N. and Grigas, P. (2017).  Smart "predict, then optimize".  ArXiv, abs/1710.08005.
  * Eysenbach et al., [2019] Eysenbach, B., Salakhutdinov, R. R., and Levine, S. (2019).  Search on the replay buffer: Bridging planning and reinforcement learning.  In Advances in Neural Information Processing Systems 32, pages 15246–15257. Curran Associates, Inc.
  * Ferber et al., [2020] Ferber, A., Wilder, B., Dilkina, B., and Tambe, M. (2020).  Mipaal: Mixed integer program as a layer.  In AAAI, pages 1504–1511.
  * Finn and Levine, [2017] Finn, C. and Levine, S. (2017).  Deep visual foresight for planning robot motion.  In 2017 IEEE International Conference on Robotics and Automation (ICRA), pages 2786–2793. IEEE.
  * Gammell et al., [2014] Gammell, J. D., Srinivasa, S. S., and Barfoot, T. D. (2014).  Informed rrt*: Optimal sampling-based path planning focused via direct sampling of an admissible ellipsoidal heuristic.  In 2014 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 2997–3004. IEEE.
  * Hafner et al., [2019] Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. (2019).  Learning latent dynamics for planning from pixels.  In International Conference on Machine Learning, volume 97 of ICML’19, pages 2555–2565, Long Beach, California, USA.
  * Ho and Ermon, [2016] Ho, J. and Ermon, S. (2016).  Generative adversarial imitation learning.  In Advances in neural information processing systems, pages 4565–4573.
  * Janner et al., [2019] Janner, M., Fu, J., Zhang, M., and Levine, S. (2019).  When to trust your model: Model-based policy optimization.  In Advances in Neural Information Processing Systems, pages 12519–12530.
  * Kaiser et al., [2019] Kaiser, L., Babaeizadeh, M., Milos, P., Osinski, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., et al. (2019).  Model-based reinforcement learning for atari.  arXiv preprint arXiv:1903.00374.
  * Karkus et al., [2019] Karkus, P., Ma, X., Hsu, D., Kaelbling, L. P., Lee, W. S., and Lozano-Pérez, T. (2019).  Differentiable algorithm networks for composable robot learning.  In Bicchi, A., Kress-Gazit, H., and Hutchinson, S., editors, Robotics: Science and Systems XV, University of Freiburg, Freiburg im Breisgau, Germany, June 22-26, 2019.
  * Kumar et al., [2019] Kumar, R., Mandalika, A., Choudhury, S., and Srinivasa, S. (2019).  Lego: Leveraging experience in roadmap generation for sampling-based planning.  2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 1488–1495.
  * Kuo et al., [2018] Kuo, Y.-L., Barbu, A., and Katz, B. (2018).  Deep sequential models for sampling-based planning.  In 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 6490–6497. IEEE.
  * Li et al., [2020] Li, Y., He, H., Wu, J., Katabi, D., and Torralba, A. (2020).  Learning compositional koopman operators for model-based control.  In International Conference on Learning Representations.
  * Mandi et al., [2019] Mandi, J., Demirovic, E., Stuckey, P. J., and Guns, T. (2019).  Smart predict-and-optimize for hard combinatorial optimization problems.  CoRR, abs/1911.10092.
  * Nachum et al., [2018] Nachum, O., Gu, S. S., Lee, H., and Levine, S. (2018).  Data-efficient hierarchical reinforcement learning.  In Advances in Neural Information Processing Systems, pages 3303–3313.
  * Neu and Szepesvári, [2012] Neu, G. and Szepesvári, C. (2012).  Apprenticeship learning using inverse reinforcement learning and gradient methods.  arXiv preprint arXiv:1206.5264.
  * Ng and Russell, [2000] Ng, A. Y. and Russell, S. (2000).  Algorithms for inverse reinforcement learning.  In in Proc. 17th International Conf. on Machine Learning. Citeseer.
  * Niculae et al., [2018] Niculae, V., Martins, A. F., Blondel, M., and Cardie, C. (2018).  Sparsemap: Differentiable sparse structured inference.  arXiv preprint arXiv:1802.04223.
  * Oh et al., [2017] Oh, J., Singh, S., and Lee, H. (2017).  Value prediction network.  In Advances in Neural Information Processing Systems, pages 6118–6128.
  * Pinneri et al., [2021] Pinneri, C., Sawant, S., Blaes, S., and Martius, G. (2021).  Extracting strong policies for robotics tasks from zero-order trajectory optimizers.  In International Conference on Learning Representations, ICLR.
  * Racanière et al., [2017] Racanière, S., Weber, T., Reichert, D., Buesing, L., Guez, A., Jimenez Rezende, D., Puigdomènech Badia, A., Vinyals, O., Heess, N., Li, Y., et al. (2017).  Imagination-augmented agents for deep reinforcement learning.  Advances in neural information processing systems, 30:5690–5701.
  * Reddy et al., [2019] Reddy, S., Dragan, A. D., and Levine, S. (2019).  Sqil: Imitation learning via reinforcement learning with sparse rewards.  arXiv preprint arXiv:1905.11108.
  * Rolínek et al., [2020] Rolínek, M., Musil, V., Paulus, A., Vlastelica, M., Michaelis, C., and Martius, G. (2020).  Optimizing ranking-based metrics with blackbox differentiation.  In Conference on Computer Vision and Pattern Recognition, CVPR’20.
  * Ross et al., [2011] Ross, S., Gordon, G., and Bagnell, D. (2011).  A reduction of imitation learning and structured prediction to no-regret online learning.  In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627–635.
  * Savinov et al., [2018] Savinov, N., Dosovitskiy, A., and Koltun, V. (2018).  Semi-parametric topological memory for navigation.  In International Conference on Learning Representations.
  * Schaul et al., [2015] Schaul, T., Horgan, D., Gregor, K., and Silver, D. (2015).  Universal value function approximators.  In International Conference on Machine Learning, volume 37 of ICML, pages 1312–1320.
  * Schrittwieser et al., [2019] Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. (2019).  Mastering atari, go, chess and shogi by planning with a learned model.  arXiv preprint arXiv:1911.08265.
  * [46] Silver, D., Hasselt, H., Hessel, M., Schaul, T., Guez, A., Harley, T., Dulac-Arnold, G., Reichert, D., Rabinowitz, N., Barreto, A., et al. (2017a).  The predictron: End-to-end learning and planning.  In International Conference on Machine Learning, pages 3191–3199. PMLR.
  * [47] Silver, D., van Hasselt, H., Hessel, M., Schaul, T., Guez, A., Harley, T., Dulac-Arnold, G., Reichert, D., Rabinowitz, N., Barreto, A., and Degris, T. (2017b).  The predictron: End-to-end learning and planning.  In International Conference on Machine Learning, volume 70 of ICML’17, pages 3191–3199. PMLR.
  * Srinivas et al., [2018] Srinivas, A., Jabri, A., Abbeel, P., Levine, S., and Finn, C. (2018).  Universal planning networks: Learning generalizable representations for visuomotor control.  In International Conference on Machine Learning, volume 80 of ICML’18, pages 4732–4741, Stockholmsmässan, Stockholm Sweden.
  * Sutton, [1991] Sutton, R. S. (1991).  Dyna, an integrated architecture for learning, planning, and reacting.  ACM Sigart Bulletin, 2(4):160–163.
  * Sutton et al., [2008] Sutton, R. S., Szepesvári, C., Geramifard, A., and Bowling, M. (2008).  Dyna-style planning with linear function approximation and prioritized sweeping.  In Proceedings of the Twenty-Fourth Conference on Uncertainty in Artificial Intelligence, pages 528–536.
  * Tamar et al., [2016] Tamar, A., Wu, Y., Thomas, G., Levine, S., and Abbeel, P. (2016).  Value iteration networks.  In Advances in Neural Information Processing Systems, pages 2154–2162.
  * Vlastelica et al., [2020] Vlastelica, M., Paulus, A., Musil, V., Martius, G., and Rolínek, M. (2020).  Differentiation of blackbox combinatorial solvers.  In International Conference on Learning Representations, ICLR’20.
  * Wang et al., [2019] Wang, P.-W., Donti, P. L., Wilder, B., and Kolter, Z. (2019).  Satnet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver.  arXiv preprint arXiv:1905.12149.
  * Wu et al., [2020] Wu, G., Say, B., and Sanner, S. (2020).  Scalable planning with deep neural network learned transition models.  Journal of Artificial Intelligence Research, 68:571–606.
  * Yao et al., [2009] Yao, H., Bhatnagar, S., Diao, D., Sutton, R. S., and Szepesvári, C. (2009).  Multi-step dyna planning for policy evaluation and control.  In Advances in neural information processing systems, pages 2187–2195.
  * Yonetani et al., [2020] Yonetani, R., Taniai, T., Barekatain, M., Nishimura, M., and Kanezaki, A. (2020).  Path planning using neural a* search.  arXiv preprint arXiv:2009.07476.

Supplemetary Material for Neuro-algorithmic Policies enable Fast Combinatorial
Generalization

##  Appendix A Data Generation

To do imitation learning, we require expert data. For Crash Jewel Hunt (Crash
$5\times 5$ and $5\times 10$), Leaper(Grid) and Maze we can determine the
exact ground truth costs leading to optimal behavior. As an example, Crash
$5\times 5$ contains moving boxes that when encountered lead to instant death,
meaning infinite costs and otherwise the fixed cost of moving around in the
environment.

Since the environments become deterministic for a fixed random seed, we first
unrolled their dynamics for each level. After obtaining the underlying grid
structure and entities, we labeled them with costs and constructed a graph
that reflects the grid structure. An expert trajectory is constructed by
applying Dijkstra’s algorithm on this graph and the human-labeled costs and
then executing in simulation.

For the Crash Jewel Hunt experiments, we randomly sampled $2000$ solvable
levels by varying number of boxes per column, their speed, the agent start
position and the jewel position. The training levels were taken from the first
half and the second half of levels was used for testing. For the ProcGen
environments Leaper(Grid) and Maze we have taken the levels determined by
seeds 0-1000.

For Chaser, we applied a similar procedure but additionally, we recorded two
sets of human trajectories, as we observed benefits in performance by
incorporating more different expert trajectories for the same level. Since
both the search procedure and human labeling are time consuming for this
environment, we collected fewer expert trajectories for the Chaser than for
the other environments, $3\times 100$, two-thirds of which are from human
players.

Level seeds 1000000-1001000 were taken for testing in the ProcGen experiments.

##  Appendix B Environments

| Crash $5\times 5$ | Crash $5\times 10$ | Leaper(Grid) | Maze
---|---|---|---|---
learning rate | $10^{-3}$ | $10^{-3}$ | $10^{-3}$ | $10^{-3}$
$\alpha$ | 0.2 | 0.2 | 0.15 | 0.15
$\lambda$ | 20 | 20 | 20 | 20
resnet layers | 4 | 4 | 4 | 4
kernel size | 4 | 4 | 6 | 6
batch size | 32 | 32 | 16 | 16
Table 1: Training hyperparameters, where $alpha$ denotes the margin that was
used on the vertex costs and $\lambda$ the interpolation parameter for
blackbox differentiation of Dijkstra’s algorithm. We vary the kernel size of
the initial convolution for ResNet18.

Our method is applicable in discrete environments, therefore we evaluated on
environments from the ProcGen benchmark and the Crash Jewel Hunt environment.

We created the Crash Jewel Hunt environment to evaluate our method, where the
goal is for the fox (Crash) to reach the jewel. We found this environment
convenient since we can influence the combinatorial difficulty directly, which
is not true for the ProcGen benchmark where we are limited to the random seeds
used in the OpenAI implementation. The sources of variation in the Crash Jewel
Hunt are the box velocities, initial positions, sizes, as well as the agent
initial position and the jewel position.

We modified the Leaper environment to make grid steps for our method to be
applicable. This involved making the logs on the river move in discrete steps
as well as the agent. Moreover, in our version, the agent is not transported
by the logs as they move, but has to move actively with them. For an
additional description of the ProcGen environmnets, we refer the reader to
Cobbe et al., [14].

##  Appendix C Network Architecture and Input

For all of our experiments, we use the PyTorch implementation of the ResNet18
architecture as the base of our model. All approaches receive two stacked
frames of the two previous time steps as input to make dynamics prediction
possible. For the PPO baseline, we did not observe any benefit in adding the
stacked frames as input and we used stable-baselines implementation from
OpenAI to train it on the ProcGen environments.

In the case of the behavior cloning baseline, the problem is a multi-class
classification problem with the output being a multinomial distribution over
actions.

For the variant NAP∗, we train a cost prediction network on top of which we
run Dijkstra’s algorithm on the output costs of the planning graph. This
requires modifications to the original ResNet18 architecture. We remove the
linear readout of the original ResNet18 architecture and replace it with a
convolutional layer of filter size 1 and adaptive max pooling layer to obtain
the desired dimensions of the underlying latent planning graph. More
concretely, the output $x$ of the last ResNet18 block is followed by the
following operation (as output by PyTorch) to obtain the graph costs:

Seq | uential(
---|---
| Conv2d(256, 2, kernel_size=(1, 1), stride=(1, 1))
| Abs()
| AdaptiveMaxPool2d(output_size=(grid_height, grid_width))
)

Where grid_{height,width} denotes the height and width of the planning grid.
For the full variant of NAP with goal prediction and agent position prediction
we have a separate position classifier that has the same base architecture as
the cost prediction network with 2 additional linear readouts for the
likelihoods of the latent graph vertices, more concretely (as output by
PyTorch):

Seq | uential(
---|---
| Conv2d(256, 2, kernel_size=(1, 1), stride=(1, 1))
| Abs()
| AdaptiveMaxPool2d(output_size=(grid_height, grid_width))
| Flatten()
|  Linear(grid_height $\times$ grid_width, grid_height $\times$ grid_width)
)

For training the position classifier, we use a standard cross-entropy loss on
the likelihoods. For NAP with position classification, we use the ground-truth
expert start and goal positions to calculate the Hamming loss of the predicted
path by the solver. At evaluation time, NAP uses the position classifier to
determine the start and end vertices in the latent planning graph.

##  Appendix D Training Procedure

For Crash $5\times 5$, Crash $5\times 10$, Leaper(Grid) and Maze we train the
models on the same #levels, namely 1, 2, 5, 10, 20, 50, 100, 200, 500 and
1000. We evaluate on unseen 1000 levels in order to show that NAP exhibits
superior generalization. The levels are generated as per description in
section A. Each dataset is normalize to be zero mean and unit variancesize.
For each dataset size (#levels) we run experiments with 3 random restarts
(seeds for network initialization). For all experiments, we make use of the
Adam optimizer.

We determine the number of epochs for training depending on each dataset size
as $\min(150000/\text{\\#levels},15000)$ to have roughly the same number of
gradient updates in each experiment. We take the minimum over the 2 values
because for smaller number of levels a large number iterations is not
necessary to achieve good performance, but for a larger number of levels it is
necessary. If we observe no error on the training set, we stop the training.

For the Chaser, the training conditions were analogous to the other
environments, only of slightly smaller scale due to its higher complexity.
Models were trained on 10, 20, 50, 100 and 200 levels and evaluated on 200
unseen levels.

| Chaser
---|---
learning rate | $1e^{-3}$
$\alpha$ | 0.2
$\lambda$ | 40
resnet layers | 3
kernel size | 4
batch size | 16
Table 2: Training hyperparameters for the Chaser experiment, where $\alpha$
denotes the margin that was used on the vertex costs and $\lambda$ the
interpolation parameter for blackbox differentiation od Dijkstra.

###  D.1 PPO Training Procedure

The training of the PPO baseline is exactly the same as described in Cobbe et
al., [14] using the official code from https://github.com/openai/train-
procgen, see Table 3 for the used parameters. The network architecture is the
IMPALA-CNN. The algorithm is trained on the specified number of levels for 200
million environments interactions gathered from 256 (instead of 64 as in Cobbe
et al., [14]) to compensate for not having access to 4 parallel workers. We
report numbers for 3 independent restarts.

learning rate | $5e^{-4}$
---|---
$\alpha$ | 0.2
discount $\gamma$ | 0.999
entropy coefficient | 0.01
steps per update | $2^{16}$
Table 3: PPO hyperparameters, as used in Cobbe et al., [14].

###  D.2 DrAC Training Procedure

For the DrAC algorithm, we run all versions introduced by anu et al., [5]
(Meta-DrAC, RL2-DrAC, UCB-DrAC and DrAC-Crop) and choose the best one in the
main plots, denoted as DrAC∗. We used the original hyperparameters from anu et
al., [5] and the implementation from https://github.com/rraileanu/auto-drac.
As with the other experiments, we report numbers from 3 different random
seeds.

##  Appendix E On Comparing Imitation Learning to Reinforcement Learning

We compare our method to Data Augmented Actor Critic, PPO and a behavior
cloning baseline. Arguably, since NAP is used in an imitation learning
setting, it reaps benefits from having access to expert trajectories.
Nevertheless, it is not straight forward that embedding a solver in a neural
architecture leads to better generalization in comparison to reinforcement
learning methods.

Behavior cloning with a standard neural architecture has access to the same
amount of data, whereas reinforcement learning agents have access to orders of
magnitude more data for inference ($2\cdot 10^{8}$ transitions in comparison
to $\sim 10^{5}$ max). This would lead us to believe that reinforcement
learning agents are able to generalize well because of the sheer amount of
data that they have at their disposal, but we show that nevertheless it is
possible to extract meaningful policies even with a small number of training
levels seen with expert trajectories, that are more optimal.

In addition, NAP is a general architecture paradigm that may be composed with
various different objective functions, including a reinforcement learning
formulation. It would be interesting to see how NAP behaves when used in such
a formulation and if this would lead to even better generalization properties
with more data. We provide training performance curves in Fig. 7 and density
plots for different numbers of training levels in Fig. 8.

Crash $5\times 5$ | Crash $5\times 10$ | Maze | Leaper
---|---|---|---
|  |  |
Figure 7: The dotted lines denote the training performance of the methods. We observe that the behavior cloning baseline and NAP have fitted the training set almost perfectly. |  10 |  20 |  50 |  100 |  200 |  500
---|---|---|---|---|---|---
Maze |  |  |  |  |  |
Leaper |  |  |  |  |  |
Figure 8: Density plots of performance on the test set (1000 unseen levels)
after different number of training levels for the Maze and Leaper
environments, the white point denotes the median performance on the test set.

##  Appendix F Data Regularized Actor-Critic

The DrAC algorithm [5] attacks the problem of generalization in reinforcement
learning from a different angle, namely applying (in some versions optimized)
data augmentations to the PPO algorithm. The main insight is that naively
applied data augmentations result in faulty policy gradient estimates because
the resulting policy after applying the augmentation is too far from the
behavior policy.

To alleviate this, a policy regularization term $G_{\pi}$ and value function
regularization term $G_{V}$ are added to the PPO objective:

$J_{\text{DrAC}}=J_{\text{PPO}}-\alpha_{r}(G_{\pi}+G_{V}).$ (7)

Furthermore, various augmentations and augmentation selection schemes were
proposed. We ran all of the proposed selectin schemes on our environments,
Meta-DrAC, RL2-DrAC, UCB-DrAC and DrAC. Meta-DrAC meta-learns the weights of a
convolutional neural network used for data augmentation. RL2-DrAC meta-learns
a policy that selects an augmentation from a pre-defined set of augmentations.
UCB-DrAC is a bandit formulation of the augmentation selection problem with
application of a upper confidence bound algorithm for selection strategy. DrAC
denotes the version with the crop augmentation, which has been shown to work
well with more than half of the environments in the ProcGen benchmark. For
more details, we refer the reader to anu et al., [5].

DrAC’s approach to improving generalization is orthogonal to NAP and the
approaches may be composed in order to achieve even better generalization
capabilities.

##  Appendix G Additional Related Work

#### Generalization in reinforcement learning

In addition to the work of anu et al., [5], there is a plethora of approaches
that attempt to improve generalization in reinforcement learning by
considering various data augmentation techniques while mainly drawing
motivation from supervised learning approaches kostrikov2020image,
laskin2020reinforcement. Other approaches combine unsupervised learning with
data augmentation srinivas2020curl, zhan2020framework.

Notably, the problem of sim-to-real transfer can be seen as a problem of
generalization to different system dynamics. Domain randomization
tobin2017domain, i.e. augmenting system dynamics in a structured way, has
emerged as one of the main techniques for tackling this problem.

main apalike