Neural Algorithmic Reasoners are Implicit Planners

1 Introduction . Planning is an important aspect of reinforcement learning ( RL ) algorithms , and planning algorithms are usually characterised by explicit modelling of the environment . Recently , several approaches explore implicit planning [ 40 , 30 , 33 , 37 , 29 , 17 , 16 ] . Such approaches propose inductive biases in the policy function to enable planning to emerge , while training the policy in a model-free manner . Accordingly , implicit planners combine the effectiveness of large-scale neural network training with the data efficiency promises of planning , making them a very attractive research direction . Many popular implicit planners attempt to align with the computations of the value iteration ( VI ) algorithm within a policy network [ 40 , 30 , 29 , 13 , 26 ] . As VI is a differentiable algorithm , guaranteed to find the optimal policy , it can combine with gradient-based optimisation and provides useful theoretical guarantees . We also recognise the potential of VI-inspired deep RL , hence it is our primary topic of study here . However , applying VI assumes that the underlying RL environment ( a ) is tabular , and that its ( b ) transition and ( c ) reward distributions are both fully known and provided upfront . Such assumptions are unfortunately unrealistic for most environments of importance to RL research . Very often the dynamics of the environment will not be even partially known , and the state space may ⇤Work performed while the author was at DeepMind . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . either be continuous ( e.g . for control tasks ) or very high-dimensional ( e.g . for pixel-space observation in Atari ) , making a tabular representation hard to realise from a storage complexity perspective . Accordingly , VI-based implicit planners often offer representation learning based solutions for alleviating some of the above limitations . Impactful early work [ 40 , 29 , 26 ] showed that , in tabular settings with known transition dynamics , the reward distribution and VI computations can be approximated by a ( graph ) convolutional network . While highly insightful , this line of work still does not allow for RL in generic non-tabular environments with unobserved dynamics . Conversely , approaches such as ATreeC [ 13 ] and VPN [ 30 ] lift the remaining two requirements , by using a latent transition model to construct a “ local environment ” around the current state . They then use learned models to predict scalar rewards and values in every node of this environment , applying VI-style algorithms directly . While such approaches apparently allow for seamless VI-based implicit planning , we discover that the prediction of scalar signals represents an algorithmic bottleneck : if the neural network has observed insufficient data to properly estimate these scalars , the predictions of the VI algorithm will be equally suboptimal . This is limiting in low-data regimes , and can be seen as unfavourable , particularly given that one of the main premises of implicit planning is improved data efficiency . In this paper , we propose the eXecuted Latent Value Iteration Network ( XLVIN ) , an implicit planning policy network which embodies the computation of VI while addressing all of the limitations mentioned previously . We retain the favourable properties of prior methods while simultaneously performing VI in a high-dimensional latent space , removing the requirement of predicting scalars and hence breaking the algorithmic bottleneck . We enable this high-dimensional VI execution by leveraging the latest advances in neural algorithmic reasoning [ 43 ] . This emerging area of research seeks to emulate iterations of classical algorithms ( such as VI ) directly within neural networks . As a result , we are able to seamlessly run XLVINs with minimal configuration changes on a wide variety of discrete-action environments , including pixel-based ones ( such as Atari ) , fully continuous-state control and navigation . Empirically , the XLVIN agent proves favourable in low-data environments against relevant model-free baselines as well as the ATreeC family of models . Our contributions are thus three-fold : ( a ) we provide a detailed overview of the prior art in value iteration-based implicit planning , and discover an algorithmic bottleneck in impactful prior work ; ( b ) we propose the XLVIN implicit planner , which breaks the algorithmic bottleneck while retaining the favourable properties of prior work ; ( c ) we demonstrate a successful application of neural algorithmic reasoning within reinforcement learning , both in terms of quantitative analysis of XLVIN ’ s data efficiency in low-data environments , and qualitative alignment to VI . 2 Background and related work . We will now present the context of our work , by gradually surveying the key developments which bring VI into the implicit planning domain and introducing the building blocks of our XLVIN agent . Planning has been studied under the umbrella of model-based RL [ 36 , 34 , 20 ] . However , having a good model of the environment ’ s dynamics is essential before being able to construct a good plan . We are instead interested in leveraging the progress of model-free RL [ 35 , 27 ] by enabling planning through inductive biases in the policy network—a direction known as implicit planning . The planner could also be trained to optimise a supervised imitation learning objective [ 38 , 2 ] . This is performed by UPNs [ 38 ] in a goal-conditioned setting . Our differentiable executors are instead applicable across a wide variety of domains where goals are not known upfront . Diff-MPC [ 2 ] leverages an algorithm in an explicit manner . However , explicit use of the algorithm often has issues of requiring a bespoke backpropagation rule , and the associated low-dimensional bottlenecks . Throughout this section we pay special attention to implicit planners based on VI , and distinguish two categories of previously proposed planners : models which assume fixed and known environment dynamics [ 40 , 29 , 26 ] and models which derive scalars to be used for VI-style updates [ 13 , 30 ] . Value iteration ( VI ) is a successive approximation method for finding the optimal value function of a discounted Markov decision process ( MDPs ) as the fixed-point of the so-called Bellman optimality operator [ 32 ] . A discounted MDP is a tuple ( S , A , R , P , ) where s 2 S are states , a 2 A are actions , R : S⇥A ! R is a reward function , P : S⇥A ! Dist ( S ) is a transition function such that P ( s0|s , a ) is the conditional probability of transitioning to state s0 when the agent executes action a in state s , and 2 [ 0 , 1 ] is a discount factor which trades off between the relevance of immediate and future rewards . In the infinite horizon discounted setting , an agent sequentially chooses actions according to a stationary Markov policy ⇡ : S⇥A ! [ 0 , 1 ] such that ⇡ ( a|s ) is a conditional probability distribution over actions given a state . The return is defined as Gt = P1 k=0 k R ( at+k , st+k ) . Value functions V ⇡ ( s , a ) = E⇡ [ Gt|st = s ] and Q⇡ ( s , a ) = E⇡ [ Gt|st = s , at = a ] represent the expected return induced by a policy in an MDP when conditioned on a state or state-action pair respectively . In the infinite horizon discounted setting , we know that there exists an optimal stationary Markov policy ⇡⇤ such that for any policy ⇡ it holds that V ⇡ ⇤ ( s ) V ⇡ ( s ) for all s 2 S. Furthermore , such optimal policy can be deterministic – greedy – with respect to the optimal values . Therefore , to find a ⇡⇤ it suffices to find the unique optimal value function V ? as the fixed-point of the Bellman optimality operator . The optimal value function V ? is such a fixed-point and satisfies the Bellman optimality equations [ 7 ] : V ? ( s ) = maxa2A R ( s , a ) + P s02S P ( s 0|s , a ) V ? ( s0 ) . Accordingly , VI randomly initialises a value function V0 ( s ) , and then iteratively updates it as follows : Vt+1 ( s ) = max a2A R ( s , a ) + X s02S P ( s0|s , a ) Vt ( s0 ) ! . ( 1 ) VI is thus a powerful technique for optimal control in RL tasks , but its applicability hinges on knowing the MDP parameters ( especially P and R ) upfront—which is unfortunately not the case in most environments of interest . To make VI more broadly applicable , we need to leverage function approximators ( such as neural networks ) and representation learning to estimate such parameters . Value iteration is message passing Progress towards broader applicability started by lifting the requirement of knowing R. Several implicit planners , including ( G ) VIN [ 40 , 29 ] and GPPN [ 26 ] , were proposed for discrete environments where P is fixed and known . Observing the VI update rule ( Equation 1 ) , we may conclude that it derives values by considering features of neighbouring states ; i.e . the value V ( s ) is updated based on states s0 for which P ( s0|s , a ) > 0 for some a . Accordingly , it tightly aligns with message passing over the graph corresponding to the MDP , and hence a graph neural network ( GNN ) [ 15 ] over the MDP graph may be used to estimate the value function . Graph neural networks ( GNNs ) have been intensively studied as a tool to process graph-structured inputs , and were successfully applied to various RL tasks [ 45 , 24 , 1 ] . For each state s in the graph , a set of messages is computed—one message for each neighbouring node s0 , derived by applying a message function M to the relevant node ( ~hs , ~hs0 ) and edge ( ~es0 ! s ) features . Incoming messages in a neighbourhood N ( s ) are then aggregated through a permutation-invariant operator L ( such as sum or max ) , obtaining a summarised message ~ms : ~ms = M s02N ( s ) M ( ~hs , ~hs0 , ~es0 ! s ) ( 2 ) The features of state s are then updated through a function U applied to these summarised messages : ~h 0 s = U ( ~hs , ~ms ) ( 3 ) GNNs can then emulate VI computations by setting the neighbourhoods according to the transitions in P ; that is , N ( s ) = { s0 | 9a 2 A. P ( s0|s , a ) > 0 } . For the special case of grid worlds , the neighbours of a grid cell correspond to exactly its neighbouring cells , and hence the rules in Equations 2–3 amount to a convolutional neural network over the grid [ 40 ] . While the above blueprint yielded many popular implicit planners , the requirement of knowing upfront the transition neighbourhoods is still limiting . Ideally , we would like to be able to , on-the-fly , generate states s0 that are reachable from s as a result of applying some action . Generating neighbouring state representations corresponds to a learned transition model , and we present one such method , which we employed within XLVIN , next . TransE The TransE [ 8 ] loss for embedding objects and relations can be adapted to RL [ 23 , 42 ] . State embeddings are obtained by an encoder z : S ! Rk and the effect of an action in a given state is modelled by a translation model T : Rk ⇥A ! Rk . Specifically , T ( z ( s ) , a ) is a translation vector to be added to z ( s ) in order to obtain an embedding of the resulting state when taking action a in state s. This embedding should be as close as possible to z ( s0 ) , for the observed transition ( s , a , s0 ) , and also far away from negatively sampled state embeddings z ( s̃ ) . Therefore , the embedding function is optimised using the following variant of the triplet loss ( with hinge hyperparameter ⇠ ) : LTransE ( ( s , a , s0 ) , s̃ ) = d ( z ( s ) + T ( z ( s ) , a ) , z ( s0 ) ) + max ( 0 , ⇠ d ( z ( s̃ ) , z ( s0 ) ) ) ( 4 ) Having a trained T function , it is now possible to dynamically construct N ( s ) . For every action a , applying T ( ~hs , a ) yields embeddings ~hs , a which correspond to one neighbour state embedding in N ( s ) . We can , of course , roll out T further from ~hs , a to simulate longer trajectory outcomes—the amount of steps we do this for is denoted as the “ thinking time ” , K , of the planner . With neighbourhoods constructed , one final question remains : how to apply VI over them , especially given that the embeddings ~hs are high-dimensional , and VI is defined over scalar reward/value inputs ? If we also train a reward model , R ( ~hs , a ) , and a state-value function , V ( ~hs ) from state embeddings , we can attach scalar values and rewards to our synthesised graph . Then , VI can be directly applied over the constructed tree to yield the final policy . As the VI update ( Equation 1 ) is differentiable , it composes nicely with neural estimators and standard RL loss functions . Further , R and V can be directly trained from observed rewards and returns when interacting with the environment . This approach inspired a family of powerful implicit planners , including VPN [ 30 ] , TreeQN and ATreeC [ 13 ] . However , it remains vulnerable to a specific bottleneck effect , which we discuss next . Algorithmic bottleneck Through our efforts of learning transition and reward models , we reduced our ( potentially highly complex ) input state s into an abstractified graph with scalar values in its nodes and edges , so that VI can be directly applied . However , VI ’ s performance guarantees rely on having the exactly correct parameters of the underlying MDP . If there are any errors in the predictions of these scalars , they may propagate to the VI operations and yield suboptimal policies . We will study this effect in detail on synthetic environments in Section 4.3 . As the ATreeC-style approaches commit to using the scalar produced by their reward models , there is no way to recover from a poorly predicted value . This leaves the model vulnerable to an algorithmic bottleneck , especially early on during training when insufficient experience has been gathered to properly estimate R. Accordingly , the agent may struggle with data efficiency . With our XLVIN agent , we set out to break this bottleneck , and do not project our state embeddings ~hs further to a low-dimensional space . This amounts to running a graph neural network directly over them . Given that these same embeddings are optimised to produce plausible graphs ( via the TransE loss ) , how can we ensure that our GNN will stay aligned with VI computations ? Algorithmic reasoning An important research direction explores the use of neural networks for learning to execute algorithms [ 10 , 43 ] —which was recently extended to algorithms on graphstructured data [ 44 ] . In particular , [ 47 ] establishes algorithmic alignment between GNNs and dynamic programming algorithms . Furthermore , [ 44 ] show that supervising the GNN on the algorithm ’ s intermediate results is highly beneficial for out-of-distribution generalization . As VI is , in fact , a dynamic programming algorithm , a GNN executor is a suitable choice for learning it , and good results on executing VI emerged on synthetic graphs [ 11 ] —an observation we strongly leverage here .

Summary: This submission proposes eXecuted Latent Value Iteration Networks (XLVINs) to relax the assumptions of Value Iteration Networks. The method infers the dynamics of the MDP using known representation learning techniques and uses the inferred dynamics to construct a graph representation which is used for planning. The method allows continuous state spaces and unlike VINs, not restricted to grid-world.