Discrete Planning with Neuro-algorithmic Policies
Abstract
Although model-based and model-free approaches to learning the control of systems have achieved impressive results on standard benchmarks, most have been shown to be lacking in their generalization capabilities.
These methods usually require sampling an exhaustive amount of data from different environment conﬁgurations.
We introduce a neuro-algorithmic policy architecture with the ability to plan consisting of a model working in unison with a shortest path solver to predict trajectories with low way-costs.
These policies can be trained end-to-end by blackbox differentiation.
We show that this type of architecture generalizes well to unseen environment conﬁgurations.
[Figure]
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, i.e. imitating an expert.
On the other end of the spectrum lie approaches that require interaction with the environment in order to obtain data for policy extraction, also known as the problem of exploration.
Most Reinforcement Learning (RL) algorithms fall into this category of approaches.
In this work, we concern ourselves primarily with the setting where limited data is available, and a policy needs to be extracted from it.
This is particularly true for learning in problems with high-dimensional input space, such as imagebased inputs.
In order to alleviate this problem, learning specialized or partial models has shown to be a viable alternative, e.g. in MuZero Schrittwieser et al.
We propose to use recent advances in making combinatorial algorithms differentiable in a blackbox fashion as proposed by Vlastelica et al.
(2020) to train neuro-algorithmic policies with embedded planners end-to-end.
More speciﬁcally, we use a time-dependent shortest path planner acting on a temporally evolving graph generated by a deep network from the inputs.
This enables us to learn the time-evolving costs of the graph and connects us in a narrow sense to model-based approaches.
We demonstrate the effectiveness of this approach in an 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 orders of magnitude faster than naive learners.
Using neuro-algorithmic architectures facilitates generalization by shifting the combinatorial aspect of the problem to efﬁcient algorithms, while using neural networks to extract a good representation for the problem at hand.
They have potential to endow artiﬁcial agents with the main component of intelligence, the ability to reason.
Our contributions can be summarized as follows:
• A general differentiable policy architecture embedding shortest path algorithms.
• Demonstration of learning generalizing policies in dynamic game environments from images. Showing that the policies enhanced with shortest path algorithms exhibit superior generalization performance in comparison to standard methods.
Markov Decision Processes and Shortest Paths
We follow the MDP framework Puterman (2014) in a goal-conditioned setting Schaul et al. (2015).
This is used in sequential decision making problems where a speciﬁc terminal state has to be reached.
Deﬁnition
A goal-conditioned Markov Decision Process (gcMDP), M is deﬁned by the tuple ( S ,
A , p , g , r ), where S is the state space, A the action space, p ( s (cid:48) | a, s ) the probability of making the transition s → s (cid:48) when taking the action a , g is the goal, r ( s, a, s (cid:48) , g ) the reward obtained when transitioning from state s to s (cid:48) while taking action a and aiming for goal g .
Concretely, we concern ourselves with fully observable discrete MDPs, in which the Markov assumption for the state holds and where the state and action-space are discrete.
In this work, we imitate expert trajectories by training a policy with an embedded time-dependent shortest path solver end-to-end.
Although the actual gcMDP solved by the expert may be stochastic, we learn a deterministic latent approximate gcMDP, (cid:99) M .
Assuming that we have access to the topology of the gcMDP, by applying blackbox-differentiation theory Vlastelica et al.
we are able to learn the underlying costs (instead of rewards) of (cid:99) M such that the optimal policy on (cid:99) M is also optimal in M .
Shortest Path Algorithm and its Differentiation
We will employ an efﬁcient implementation of Dijkstra’s algorithm for computing the shortest path.
For differentiation, we rely on the framework for blackbox differentiation of combinatorial solvers Vlastelica et al.
This framework has already been successfully used for ranking Rolínek et al.
and keypoint matching Rolínek et al.
(2020) problems.
The purely combinatorial setup can be formalized as follows.
Let G =
( V, E ) be a graph.
For every v i ∈ V , let c 1 i , . .
c
Ti be non-negative real numbers; the costs of reaching the vertex v i at time-points 1 , 2 , ..., T , where T is the planning horizon.
The TIME - DEPENDENT - SHORTEST - PATH problem (TDSP) has as input the graph G , a pair of vertices s, e ∈ V (start and end) and the matrix
C ∈ R | V |× T of the costs c ti .
This version of the shortest path problem can be solved by executing the
Dijkstra shortest path algorithm 2 Dijkstra (1959) on an augmented graph.
In particular, we set
[Equation]
where the cost of vertex ( v i , t ) ∈ V ∗ is simply c ti and E (cid:44) → is the original edge set E appended with all self-loops.
This allows to “wait” at a ﬁxed vertex v from timestep t to timestep t + 1 .
In this graph, the task is to reach vertex ( e, T ) from ( s, 1) with the minimum traversal cost.
We can take advantage of this formulation of the time-dependent shortest path problem in a modelpredictive control fashion in that we predict the time dependent costs.
We learn this by applying blackbox differentiation from Vlastelica et al.
Neuro-algorithmic Policy Architecture
We propose the Neuro-algorithmic Policy (NAP) framework, that is an end-to-end trainable deep policy architecture embedding an algorithmic component using the afore-mentioned techniques.
In this paper we consider the concrete architecture consisting of two main components: a backbone
ResNet18 (without the ﬁnal fully connected layers, similarly as in Vlastelica et al.
architecture and the shortest path solver, see Fig.
At each time step the policy receives two images concatenated channel-wise from which it predicts the cost matrix C for the planning horizon T with the costpredictor and the start vertex s and end vertex e with the goal predictor , explained in the appendix.
The cost matrix C is given to the solver along with the start vertex s and end vertex e to compute the time-dependent shortest path Y .
The cost-predictor is trained using the Hamming distance between the predicted plan Y and the expert plan Y ∗ that we use for supervision.
Such a cost-prediction ca be seen in Fig.
It is necessary to know the start and goal vertex when applying the shortest-path solver for which we learn respective predictors, which we describe in appendix Sec.
C.
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 ﬁrst action from the plan.
[Figure]
Experiments
To validate our hypothesis that embedding planners into neural network architectures leads to better generalization, we consider several procedurally generated environments (from the ProcGen suite Cobbe et al.
and C RASH J EWEL H UNT ) with considerable variation between the realizations, called levels .
A more detailed description of the environments can be found in section F of the appendix.
We compare to two baselines: a standard imitation learning baseline using a ResNet18 architecture trained with a cross-entropy classiﬁcation loss on the same dataset as our method; and a reinforcement learning baseline using the PPO algorithm.
We note that a comparison to PPO is to some extent unfair, since the algorithm also needs to handle the problem of exploration.
Nevertheless, the PPO agent observes much more diverse data than our method, yet fails to generalize in comparison.
For the experimental validation, we aim to answer the following questions:
• Can NAP be trained to perform well as policies in procedurally generated environments?
• Can NAP generalize in a low data regime,
i
e. after seeing only few different levels?
C RASH 5 × C
[Figure]
×
[Figure]
5.1 Results
We train our method (NAP) and the imitation learning baseline until saturation on the training set, resulting in virtually 100% success rate when evaluating on train conﬁgurations in the environment.
For the PPO baseline we use the code from Cobbe et al.
and provide also two subsequent frames and 200M time steps for training.
For our method we also provide a version with access to the true start and end-point prediction (NAP+ oracle)
In Fig.
3 we show the performance of the methods when exposed to different number of levels at training time.
As reported in Cobbe et al.
(2019), the baselines have a large generalization gap, and also poor performance when < 10 000 levels are seen.
We ﬁnd that NAP shows strong generalization performance, already for < 500 levels.
In some environments, such as the M AZE we obtain near 80% success rate already with just 100 levels which is reached by PPO after seeing 200 000 levels.
For the C RASH J EWEL H UNT 5 × 5 already with 30 trajectories a third of the 1000 test-levels can be solved, the baseline manages less than 50 out of the 1000.
In the C HASER env, see Sec.
G, subgoals need to be predicted, which is illustrated in Fig.
We show strong generalization with very few training levels – reward of ∼ 6 on unseen test-levels when using 100 training levels.
Discussion
We have shown that hybrid neuro-algorithmic policies consisting of a deep feature extraction and a shortest path solver – made differentiable via blackbox differentiation – enables learning policies that generalize to unseen environment settings in the low-data regime .
These policies can be trained end-to-end by employing blackbox differentiation theory Vlastelica et al.
Although there is a clear beneﬁt in using NAP, the methods comes with certain caveats.
We assume that the topological structure of the latent planning graph is known a priori, relaxing this assumption we leave for future work.
Furthermore, we assume that the structure of the latent graph is ﬁxed and not dynamically changing over time, i.e. that each available action at a vertex corresponds to the same edge.
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)
A Related Work
Planning in reinforcement learning Differentiable planning has been proposed in previous works, e.g. in the continuous case with CEM Amos & Yarats (2020); Bharadhwaj et al.
Learning a representation for planning was considered, among others, in Hafner et al.
but without differentiating though the planner.
Silver et al.
(2017) differentiate through a few steps of value prediction in a learned MDP to match the externally observed rewards.
In Srinivas et al.
(2018) a differentiable planner was obtained via relaxation and used for learning a representation.
Alternatively to learning the representation, a planning graph can be constructed from the agents’ experience using a learned value function for the edge costs, as done by Eysenbach et al.
An orthogonal research direction is to optimize the planners with prior experience, e.g. Chen et al.
(2020), that can be combined with our approach.
Discrete algorithms in end-to-end trainable networks Incorporating expert discrete solvers into end-to-end trainable architectures is a topic with both rich history and exciting recent developments.
The simpler setup, in which the output of the non-differential discrete algorithm is directly compared to ground truth values has been ﬁrst addressed by max-margin techniques (Tsochantaridis et al., 2005).
Since then, various adaptations have been proposed such as the “predict-and-optimize” framework and its variants (Elmachtoub & Grigas, 2017; Demirovic et al., 2019; Mandi et al., 2019).
Also, specializations for cases of concrete types of discrete optimization problems such as sparse structured inference (Niculae et al., 2018), logical satisﬁability (Wang et al., 2019), submodular optimization (Djolonga & Krause, 2017) or mixed integer programming (Ferber et al., 2020).
The harder task of purely providing a gradient and therefore allowing entirely hybrid architectures containing a discrete algorithm as an intermediate layer has only been addressed recently.
First, by (Vlastelica et al., 2020) who introduce an efﬁcient implicit piece-wise linear interpolation scheme of the true underlying piece-wise constant function.
On the other hand, Berthet et al.
(2020) offer a
Monte Carlo technique for estimating the Jacobian of a Gaussian smoothing of the piecewise constant function at hand.
B Blackbox Differentiation
The framework presented in Vlastelica et al.
(2020) 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. 4.
The exact gradient of this linear interpolation is computed efﬁciently via evaluating the solver on only one more instance (see
Algorithm 1) in appendix.
In order 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 ) with time-dependent costs C ∈ R | V |× T we deﬁne Y ∈ { 0 , 1 } | V |× T an indicator matrix of visited vertices.
In particular, Y ti =
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 ∗ ) will be denoted as
Adm( G ) .
The time-dependent shortest path optimization problem can be then rewritten as
[Equation]
This is an inner-product objective and thus the theory from Vlastelica et al.
(2020) applies.
In effect, the deep network producing the cost tensor C can be trained via supervision signal from ground truth shortest paths.
C Goal Prediction, Static and Dynamic
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 end vertices ( s, e ).
To this end, we employ a second
ResNet18 – the goal-predictor – similar to the cost-predictor that learns to extract the agent start position and a suitable target position.
The training of this predictor is using a Cross-Entropy loss and is independent of learning the costs of the latent graph representation.
function F ORWARD P ASS ( C , s , e )
Y :=
TDSP ( C, s, e ) //
Run Dijkstra’s algorithm save Y , C , s , e // Needed for backward pass return Y function B ACKWARD P ASS ( ∇ L ( Y ) , λ ) load Y , C , s , e
C λ :=
C + λ ∇ L ( Y ) //
Calculate modiﬁed costs
Y λ :=
TDSP ( C λ , s, e ) //
Run Dijkstra’s algo.
return 1 λ (cid:0) Y λ − Y (cid:1)
At training time, given the expert trajectories Y ∗ we have access to the current position of the agent and its position in the future.
Thus, for predicting s we have a simple supervision signal, namely the current position.
For the goal prediction e we extract a set of suitable goal locations from the expert.
Here, we distinguish between global and local planning.
In the global setting, the last position of the expert is the goal e , corresponding to, for instance, the jewel in C RASH J EWEL H UNT , see 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 + 2 T .
The local setting allows to limit the complexity of our method, which grows with the planning horizon.
It also is a trade-off between the combinatorial complexity solved by the TDSP 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 -ends in a
This formulation makes our architecture a hierarchical method similar to Blaes et al.
; Nachum et al.
(2018), and allows for solving tasks that are not typical goal-reaching problems, such as the
C HASER environment, see below.
D Sensitivity to the Planning Horizon
We provide a sensitivity analysis of the performance with different planning horizons.
Our results indicate that longer horizons beneﬁt environments with increased dynamical interactions.
As apparent from Fig. 5, our method outperforms the imitation baseline greatly in the crash environment, the gap between the methods being correlated with the complexity of the environment ( 5 × 5 vs 5 × 10 ).
It can be seen also that making the planning horizon smaller in these environments hurts performance.
On the other hand, for environments with no dynamics, such as the maze environment, there is no beneﬁt in using time-dependent costs, as expected.
Nevertheless, there is still strong performance gain in generalization when using NAP oppose to vanilla imitation learning from expert trajectories.
C RASH 5 X
C RASH
[Figure] [Figure]
cc
[Figure]
E Cost Margin
Our work is related to the problem of metric learning in the sense that we learn the distance metric between the current position of the agent (state) and the target position in the underlying gcMDP, allowing us to solve it with a shortest path algorithm.
It has been shown that inducing a margin on the metric can be beneﬁcial for generalization.
Similarly to (Rolínek et al., 2020) in the context of rank-based metric learning, we induce a margin α on the costs of the latent gcMDP, increasing the cost of ground truth path and decreasing the rest in the training stage of the algorithm:
[Equation]
During evaluation, the cost margin is removed from the shortest path calculation.
F Environments
[Equation]
[Equation]
We ﬁrst consider an environment we constructed to test NAP, called C RASH J EWEL H UNT which can be seen in Fig.
The environment corresponds to a grid-world of dimensions h
× 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 conﬁgurations 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 initial and goal positions.
F.2 ProcGen Benchmark
In addition to the jewel hunt environment, we evaluate our method on the M AZE , L EAPER and
C HASER environments from the ProcGen suite Cobbe et al.
We have chosen these environments because their structure adheres to our assumptions.
For the L EAPER we modiﬁed the environment such that a grid-world dynamics applies (L EAPER (G RID )).
Based on performance of the baselines, the resulting (L EAPER (G RID )) is not an easier environment.
The M AZE and the L EAPER (G RID ) tasks have a static goal whose position only varies across levels, whereas the C HASER requires collection of all orbs without contact with the spiders, so the local goals need to be inferred on the ﬂy.
The C HASER environment is also particularly challenging as even the expert episodes require on average 150 steps, each of which carries a risk of dying.
For this reason, we used three expert trajectories per level.
G Results on C HASER Environment
Here we provide additional results of NAP in comparison to PPO on the C HASER environment.
This environment is speciﬁc in that there are not explicit goals needed to be reached, but the agent needs to pick up all of the green orbs while avoiding the enemies.
This poses a challenge, since we are reduced to a very short-horizon in comparison to the episode length of the expert.
Furthermore, the environment doesn’t have independent dynamics but the dynamics depend on the movement of the agent.
Experiments show that NAP still outperforms PPO considerably in the low data regime (Fig 6)
[Figure]