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 work extends the value iteration network family of reinforcement learning agents to a much less constrained form to avoid the "algorithmic bottleneck". The proposed model XLVIN creates a plannable representation of the state space by 1) first learning an embedding space for the state transitions via a TransE (or contrastive) approach and 2) then aggregating several layers of look-ahead state embeddings via a pretrained GNN. The policy and value functions are parameterized on top of the plannable representation. Empirically XLVINs provide a significant boost on the stochastic control tasks and Atari games in the low-data regime.