Neuro-algorithmic Policies for Discrete Planning

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 of the spectrum lie approaches that require interaction with the environment to obtain data for policy extraction problem , also known as 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 algorithms require exhaustive amounts of exposure to environmental variability before starting to generalize Cobbe et al . ( 2019 ) . 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 . This has been argued to be the best strategy in the presence of an accurate model and sufficient computation time ( Daw et al. , 2005 ) . However , learning a precise transition model is often harder than learning a policy . This , in turn , makes them more general , but comes at a cost of increasing the problem dimensionality . The transition model has a much larger dimensionality and 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 Schrittwieser et al . ( 2019 ) . 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 specifically , 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 relates us to model-based approaches . 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 orders of magnitude faster 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 potential to endow artificial agents with the main component of intelligence , the ability to reason . Our contributions can be summarized as follows : • We identify that poor generalization is caused by lack of structural and combinatorial inductive biases and can be alleviated by introducing the correct inductive biases through neuro-algorithmic policies . • We show that architectures embedding TDSP solvers are applicable beyond goal-reaching environments . • 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 ( Gammell et al. , 2014 ; Burget et al. , 2016 ; Kuo et al. , 2018 ) . Similarly , along this direction , Kumar et al . ( 2019 ) 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 . ( 2019 ) employ a latent multi-step transition model . Savinov et al . ( 2018 ) propose a semi-parametric method for mapping observations to graph nodes and then applying a shortest path algorithm . Asai & Fukunaga ( 2017 ) ; Asai & Kajino ( 2019 ) use an autoencoder architecture in order to learn a discrete transition model suitable for classical planning algorithms . Li et al . ( 2020 ) learn compositional Koopman operators with graph neural networks mapping to a linear dynamics latent space , which allows for fast planning . Chen et al . ( 2018 ) ; Amos et al . ( 2017 ) perform efficient planning by using a convex model formulation and convex optimization . Alternatively , the replay buffer can be used as a non-parametric model in order to select waypoints ( Eysenbach et al. , 2019 ) or in an MPC fashion ( Blundell et al. , 2016 ) . None of these methods perform differentiation through the planning algorithm in order 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 ( Amos & Yarats , 2020 ; Bharadhwaj et al. , 2020 ) . Wu et al . ( 2020 ) use a ( differentiable ) recurrent neural network as a planner . Tamar et al . ( 2016 ) use a differentiable approximation of the value iteration algorithm to embed it in a neural network . Silver et al . ( 2017b ) differentiate through a few steps of value prediction in a learned MDP to match the externally observed rewards . Srinivas et al . ( 2018 ) use a differentiable forward dynamics model in latent space . Karkus et al . ( 2019 ) 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 . ( 2020 ) 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 objective function from demonstrations has been a central topic in reinforcement learning ( Ng & Russell , 2000 ) . 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 ( Abbeel & Ng , 2004 ; Neu & Szepesvári , 2012 ; Aghasadeghi & Bretl , 2011 ) . 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 ( Ross et al. , 2011 ) or other online reinforcement learning methods augmented with expert datasets ( Reddy et al. , 2019 ; Ho & Ermon , 2016 ) . Offline model-based reinforcement learning Model-based methods have shown promise by facilitating better generalization ( Janner et al. , 2019 ) . Approaches employing models fall into two camps : using models to extract a policy in a Dyna-style approach ( Sutton , 1991 ; Janner et al. , 2019 ; Sutton et al. , 2008 ; Yao et al. , 2009 ; Kaiser et al. , 2019 ) , or incorporating the model in a planning loop , i.e . model-predictive control ( Finn & Levine , 2017 ; Racanière et al. , 2017 ; Oh et al. , 2017 ; Silver et al. , 2017a ) . 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 ( Elmachtoub & Grigas , 2017 ; Demirovic et al. , 2019 ; Mandi et al. , 2019 ) . Also , specializations for concrete cases such as sparse structured inference ( Niculae et al. , 2018 ) , logical satisfiability ( Wang et al. , 2019 ) , submodular optimization ( Djolonga & Krause , 2017 ) or mixed integer programming ( Ferber et al. , 2020 ) 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 ( Vlastelica et al. , 2020 ; Berthet et al. , 2020 ) . Vlastelica et al . ( 2020 ) introduce an efficient implicit piece-wise linear interpolation scheme , while Berthet et al . ( 2020 ) introduce Monte Carlo technique for estimating the Jacobian of a Gaussian smoothing of the piecewise constant function . The approach from Vlastelica et al . ( 2020 ) is especially appealing , since it allows for uses in which the solver is the computational bottleneck . Therefore , we follow it in this work . By formulating the control problem as a time-dependent shortest path problem ( TDSP ) , we show that the framework from Vlastelica et al . ( 2020 ) is applicable in specific control settings . 3 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 specific terminal state has to be reached . Definition 1 A goal-conditioned Markov Decision Process ( gcMDP ) , M is defined by the tuple ( S , A , p , g , r ) , where S is the state space , A the action space , p ( s′ | a , s ) the probability of making the transition s → s′ when taking the action a , g is the goal , r ( s , a , s′ , g ) the reward obtained when transitioning from state s to s′ 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 . The goal of reinforcement learning is to maximize the return G = ∑T t=0 rt of such a process . In gcMDPs the reward is such that the maximal return can be achieved by reaching the goal state g. Given access to the transition probabilities and rewards , an optimal policy can be extracted by dynamic programming Bertsekas et al . ( 1995 ) . In a graph representation of an gcMDP , the set of vertices V corresponds to the set of states S , traversing an edge corresponds to making a transition between states . We assume a deterministic process , such that the optimal policy can be extracted by standard shortest path algorithms , such as Dijkstra ’ s algorithm . 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 , M̂ . Assuming that we have access to the topology of the gcMDP , by applying blackbox-differentiation theory Vlastelica et al . ( 2020 ) we are able to learn the underlying costs ( instead of rewards ) of M̂ such that the optimal policy on M̂ is also optimal inM . Although the MDP Definition 1 yields itself nicely towards learning the time-dependent edge costs cet , this can increase the problem dimensionality considerably with the out-degree of the vertices ( here |A| ) . Thus , we consider cases where the reward function only depends on the current state and the goal : r ( s , a , s′ , g ) = r ( s , g ) . In this case , vertex costs cvt are sufficient for finding the optimal solution to the gcMDP . Accordingly , we rely on a vertex-based version of the shortest path algorithm .

Summary: This work proposes a novel neuro-algorithmic policy architecture for solving discrete planning tasks. It takes a high-dimensional image input and processes it through modified ResNet encoders to obtain a graph cost map and a start/goal heatmap. This is fed into a differentiable Dijkstra algorithm to obtain the shortest trajectory prediction which is trained using an expert-annotated trajectory via a Hamming distance loss. This module is evaluated in two dynamic game environments demonstrating generalization to unseen scenes.