Neural Algorithmic Reasoners are Implicit Planners
Authors: Andreea Deac
Petar Veličković
Ognjen Milinković
Pierre-Luc Bacon
Jian Tang
Mladen Nikolic
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: graph neural networks
value iteration
implicit planning
algorithmic bottleneck

Abstract: Implicit planning has emerged as an elegant technique for combining learned models of the world with end-to-end model-free reinforcement learning. We study the class of implicit planners inspired by value iteration, an algorithm that is guaranteed to yield perfect policies in fully-specified tabular environments. We find that prior approaches either assume that the environment is provided in such a tabular form---which is highly restrictive---or infer "local neighbourhoods" of states to run value iteration over---for which we discover an algorithmic bottleneck effect. This effect is caused by explicitly running the planning algorithm based on scalar predictions in every state, which can be harmful to data efficiency if such scalars are improperly predicted. We propose eXecuted Latent Value Iteration Networks (XLVINs), which alleviate the above limitations. Our method performs all planning computations in a high-dimensional latent space, breaking the algorithmic bottleneck. It maintains alignment with value iteration by carefully leveraging neural graph-algorithmic reasoning and contrastive self-supervised learning. Across seven low-data settings---including classical control, navigation and Atari---XLVINs provide significant improvements to data efficiency against value iteration-based implicit planners, as well as relevant model-free baselines. Lastly, we empirically verify that XLVINs can closely align with value iteration.

Neural Algorithmic Reasoners are Implicit Planners
Andreea Deac" Mila Québec Al Institute Université de Montréal
Petar Veličković DeepMind
Ognjen Milinković Faculty Mathematics University of Belgrade
Pierre-Lue Bacon
Jian Tang
Mladen Nikolié
Mila Québec Al Institute
Mila Québec Al Institute
Faculty of Mathematics University of Belgrade
Université de Montréal
HEC Montréal
Abstract
Implicit planning has emerged as elegant technique for combining learned mod- els of the world with end-to-end model-free reinforcement learning. We study the class of implicit planners inspired by value iteration, algorithm that is guaran- teed to yield perfect policies in fully-specified tabular environments. We find that prior approaches either assume that the environment is provided in such tabular form- whinh is highly restrictive infer "local neighbourhoods" of states to value iteration over for which discover algorithmic bottleneck effect. This effect caused by explicitly running the planning algorithm based on scalar predictions in every state, which can be harmful data efficiency if such scalars are improperly predicted. We propose eXecuted Latent Value Iteration Networks (XLVINs). which alleviate the above limitations. Our method performs all planning computations a high-dimensional latens space, breaking the algorithmic bottle- neck. It maintains alignment with value iteration by carefully leveraging neural graph-algorithmic reasoning and contrastive self-supervised learning. Across eight low-data settings-including classical control, navigation and Atari- Atari-XLVINs XLVINs provide significant improvements to data efficiency against value iteration implicit planners, as well relevant model-free baselines. Lastly, we empirically verify that XLVINs can closely align with value iteration.
Introduction
Planning is an important aspect of reinforcement learning (RL) algorithms, and planning algorithms are usually characterised by explicit modelling the environment. Recently, several approaches 1 [6]. 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 planning, making them a very attractive research direction.
explore implicit planning 20 30
33
37
29
Many popular implicit planners attempt to align with the computations of the value iteration (VI) algorithm within a policy network 44 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 Very often the dynamics of the environment will not even partially known, and the state space may
20 13 26). As VI is differentiable algorithm, guaranteed
unfortunately unrealistic for most environments of importance to RL research.
Work performed while the author DeepMind.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

either be continuous (e.g. for control tasks) or very sigh-dimensional (e.g. for pixel-space observation in Atari), making a tabular representation hard to realise from a storage complexity perspective.
Accordingly, VI-hased implicit planners often offer representation learning based solutions for allevi- ating some of the ahove limitations. Impactful early work 4 29 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 [ and VPN BO 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 rewards and values in every node of this environment, applying VI-style algorithms directly.
20 showed that, in tabular settings
ATreeC
While such approaches apparently allow for seamless VI-hased 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), 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 nigh-dimensions and hence breaking the algorithmic bottleneck. We enable this high-dimensional VI execution by leveraging the latest advances in neural a algorithmic reasoning 4 This emerging area of research seeks emulate iterations of classical algorithms (such VT) directly within neural networks. As 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.
latent space, removing the requirement of predicting scalars
Our contributions are thus fhree-fold: (a) we provide a detailed overview of the prior art in value dieration-base 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) 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 VT 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
However, having a good model of the environment's 's dynamics is essential before being able to construct good plan. We are instead interested in leveraging the progress of model-free RL
34
20
27 by enabling planning through inductive biases in the policy network- a direction known implicit planning. The planner could also be trained to optimise a supervised imitation learning objective 38 2]. This is performed by UPNs 38] in goal-conditinn conditiodd setting. Our differentiable executors 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 bespoke backpropagation rule, and the associated low-dimensional bottlenecks.
335
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 20 26
and models which derive scalars to be used for VI-style updates [ 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 A,, R. P,7) where & E S are states, a E AA are actions, R: S x A- - RR is a reward function, P: S x A + Dist(S) is transition function such that P(s (os'ss,a) a) is the conditional probability of transitioning to state s' when the agent executes action in state and y [0, 1] is discount factor which trades off between the relevance of immediate and

---

future rewards. In the infinite horizon discounted setting, an agent sequentially chooses actions ac- cording to a stationary Markov policy S x A - [0. 1] such that ग s)) is conditional probability distribution over actions given state. The return is defined as Gt- - ? R[ac+kk St+k). Value functions (s, a) = Exx [G,] St = s] and (s, a) = E,, [Gr|s; = 5, = aj represent the expected return induced by a policy in an MDP when conditioned on state or state-action pair respectively. In the infinite horizon discounted setting, we know that there exists an optimal stationary Markov policy x¹ such that for any policy и it holds that V (s) 2 (s) for all & E S. Furthermore, such optimal policy can be deterministic- greedy with respect to the optimal values. Therefore, to find a πT" 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 (s) = m (R(ss, a) + + + P(s' ss, a) (s')). Accordingly, VI randomly initialises a value function Vo(s), and then iteratively updates it as follows:
V1+1(s) = max
.EAÀ R(s,a)+ Σ P(5))s,))V((()) SEES
(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 cnvironments where P is fixed and known. Observing the VI update rule (Equation D. we may conclude that it derives values by considering features of neighbouring states; ie. the value (s) is updated based on states for which P(s [[s'|s, a) > for some a. Accordingly, tightly aligns with message passing over the graph corresponding to the MDP, and hence graph neural network (GNN) [13] 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 wraph-structured inputs, and were successfully applied to various RL tasks 45 24 M. For each state in the graph, a
set of messages is computed-one message for each neighbouring node 8', derived by applying
message function M to the relevant node (h., h. and edge (Egr+s) features. Incoming messages
in neighbourhood (s)) are then aggregated through a permutation-invan invariint operator + (such as sum or max), obtaining a summarised message m,,:
m
① s'EN(ss) (he, Fi (o)
(2)
The features of state are then updated through a function U applied to these summarised messages:
Wy U(h,, (..)
GNNs can then emulate VI computations by setting the neighbourhoods according to the transitions in P; that is, (sss) = (s' Ea E A. P(s'|s, a) > 0}. For the special case of grid worlds, the neighbours of grid cell correspond to exactly its neighbouring cells, and hence the rules in Equations 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 s' that are reachable from as result of applying some action. Generating neighbouring state representations corresponds
learned transition model, and we present one such method, which we employed within XLVIN, next.
loss for embedding objects and relations can be adapted RL State embeddings are obtained by an encoder 2 : S + R² and the effect of an action in a given state modelled by translation model T R& x A + R Specifically, ((x(s), a) is translation vector to be added to 2(s) in order to obtain an embedding of the resulting state when taking action a in state This embedding should be as close as possible to =(s'), for the observed transition (s, a, s'), and also far away from negatively sampled state embeddings 233.. Therefore, the embedding function is optimised using the following variant of the triplet loss (with hinge hyperparameter E)
TransE
The TransE g
CTransE((s, ((,,,,(sss,)))) (s,,,,'), s'), d(=(s) + T(z(s), a), 2(s')) + max(0, & - d(z(š), z(s')))
(4)
3

---

Encoder
Transition
Exeentor
Tail
oi)
z(+)
Beans
Thh oa)
Construct latent graph using? T
Exccute (latent) Value Iteration using. X
Figure 1: XLVIN model summary. Its modules are explained (and colour-coded) in Section 3.1
Having a trained T function, is now possible to dynamically construct (ss). For every action applying Thi,, a) yields embeddings ha, which correspond to one neighbour state embedding in (s). We can, of course, roll out T further from hi 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 h2 are
and VI is defined over scalar reward/value inputs?
we also train a reward model, R((,,a) a), and a state-value function, V V(ha) from state embeddings, we can attach scalar values and rewards to synthesised graph. Then, VI can be directly applied over the constructed tree to yield the final policy. As the VI update (Equation D 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 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 efforts of learning transition and reward models, reduced our (potentially highly complex) input state s into an abstractified graph with scalar values in its nodes and edges, that VI can be directly applied. However, VI's performance guarantees rely 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
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 Th, 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 ensure that our GNN will stay aligned with VI computations?
loss), how
Algorithmic reasoning An important research direction explores the use of neural networks for learning to execute algorithms [CC
43 -which was recently extended to algorithms on graph- structured data [4! In particular, 4 establishes algorithmic alignment between GNNs and dynamic programming algorithms. Furthermore, [44] show that supervising the GNN on the algorithm's 's intermediate results is highly beneficial for out-of-distribution generalization. As VI is, in fact, dynamic programming algorithm, a GNN executor is suitable choice for learning it, and good results on executing VI emerged on synthetic graphs
an observation we strongly leverage here.
3 XLVIN Architecture
Next, we specify the computations of the eXecuted Latent Value Iteration Network (XLVIN). We recommend referring to Figure I for a visualisation of the model dataflow (more compact overview in Appendix A) and to Algorithm for step-by-step description of the forward pass.

---

We propose policy network- function, iaa [a(a)s), which for a given state E S specifies probability distribution of performing each action E A in that state. Here, A are the policy parameters, to be optimised with gradient ascent.
3.1 XLVIN modules
The encoder function, 2 S IR= consumes state representations 8 E S and produces flat embeddings. Ai zs)) E I$ The design of this component is flexible and may be dependent on the structure present in states. For example, pixel-based environments will necessitate CNN encoders. while environments with flat observations are likely to be amenable to MLP encoders.
Encoder
Transition The transition function, T R X A - RR models the effects of taking actions, in the latent space. Accordingly, it consumes a state embedding z(s) and an action a and produces the appropriate translation of the state embedding, expectation). That is, it is desirable that 7 satisfies Equation S
match the embedding of the successor state (in and is commonly realised as an MLP.
2(s) T(z(s), Es ((r
(5)
Executor The executor function, X R x TA RR. processes embedding h2 of given state s, alongside a neighbourhood set (h.)), which contains (expected) embeddings of states that immediately neighbour s for example, through taking actions. Hence,
N(h..) (B)=)(o1ss)=((")))))
The executor combines the neighbourhood set features produce an updated embedding of state 8, X. = X ((h)),(h))) which is mindful of the properties and structure of the neighbourhood. Ideally, X would perform operations in the latent space which mimic the one-step behaviour of VI, allowing for the model meaningfully plan from state s by stacking several layers of X (with K layers allowing for exploring length-K K trajectories). Given the relational structure of state and its neighbours, the executor is commonly realised as graph neural network (GNN).
Actor & Tail components The actor function, A R R [o, 1|1-1 consumes the state embedding hig and the updated state embedding Xs, producing action probabilities To(als) A [i,,, X.) sspccifying the policy to be followed by our XLVIN agent. Lastly, note that we may also have additional tail networks which have the same input A. For example, train XLVINs using proximal policy optimisation (PPO) K55 which necessitates state-value function: V(h,, Xs).
Algorithm 1: XLVIN forward pass
Input : Input state s, executor depth K"
Output:Polig function no (als), state-value function V(s)
h, = =
// Embed the input state with the encoder
so [h1), E 0
for & E [0, K) do
g++1 =0;
// Initialise depth-(k + 1) embeddings
for h E Sk, a do R' = h a(h gh+1 = sh+1 u ( (B,
If (expected) neighbour embedding // Attach to the graph
end end
Run execution model over the graph specified by the nodes UK 3 and edges E, by repeatedly applying h = X (R),)(()). for every embedding i E S. for K steps.
X. = EXECUTE(F, U Sh E, X, K); // See Appendix B for details on the EXECUTE function Use the actor and tail to predict the policy and value functions from the (updated) state embedding s */ (s, A(T., X.). (s)) = V(h.. Xx)
5

---

Discussion The entire procedure is end-to-end differentiable, does not impose any assumptions on the structure of the underlying MDP, and has the capacity to perform computations directly aligned with value iteration, while avoiding the algorithmic bottleneck. This achieves all of our initial aims.
The transition function produces state embeddings that correspond the expectation of the successor state embedding over all possible outcomes (Equation 5.. While taking expectations is an operation that aligns well with VI computations, it can pose limitations in the case of non-deterministic MDPs. This is because the obtained expected latent states may not be trivially further expandable, should we wish to plan further from them. One possible remedy we suggest is employing probabilistic (e.g. variational) transition model from which we could repeatedly sample concrete next-state latents.
Our tree expansion strategy is breadth-ffisst which expands every action from every node, yielding O([A|K) time and space complexity. While this is prohibitive for scaling up K, we empirically found that performance plateaus by K Σ 4 for all studied environments, mirroring prior findings 1 Even if these are shallow trees of states, we anticipate a compounding effect from optimising TransE together with PPO's value/policy heads. We defer allowing for deeper expansions and large action spaces to future work, but note that it will likely require a rollout policy, selecting actions to expand from given state. For example, 12A 3K obtains a rollout policy by distilling the agent's policy network. Extensions to continuous actions could also be achieved by rollout policies, or discretising the action space by binning 41
3.2 XLVIN Training
As discussed, the success of XLVIN relies on our transition function, T, constructing plausible graphs, and our executor function, X, reliably imitating VI steps in high dimensional latent space. Accordingly, we train both of them using established methods: TransE for T and [11] for X.
To optimise the neural network parameters 0, we use proximal policy optimisation (PPO i Note that the PPO gradients also flow into and X, which could displace them from the properties required by the above, leading to either poorly constructed graphs or lack of VI-aligned computation.
35
Without knowledge of the underlying MDP, we have no easy way of training the executor, X, online. We instead opt to pre-train the parameters of X and freeze them, treating them as constants rather than parameters be optimised. In brief, the executor pre-training proceeds by first generating a dataset of synthetic MDPs, according to some underlying graph distribution. Then, we execute the VI algorithm on these MDPs by iterating Equation I. keeping track of intermediate values (s) at each step f, until convergence. Finally, we supervise GNN (operating over the MDP transitions as edges) to receive and all other parameters of the MDP-as as inputs, and predict V++1 (6) (optimised using mean-squared error). Such a graph neural network has three parts: an encoder, mapping Vs( to a latent representation, a processor, which performs a step of VI the latent space, and decoder, which decodes back Vs+1(s) from the latent space. We only retain the processor as our exccutor function X, in order to avoid the algorithmic bottleneck in our architecture.
For the transition model, we found it sufficient to optimise TransE (Equation 4 using only on-policy trajectories. However, we do anticipate that some envirouments will require careful tradeoff between exploration and exploitation for the data collection strategy for training the transition model.
Thus, after pre-training the GNN predict one-step value iteration updates and freezing the processor, a step of the training algorithm corresponds to:
1. Sample on-policy rollouts (with multiple parallel actors acting for a fixed number of steps).
2. Based the transitions in these rollouts, evaluate the PPO B5 and TransE (Equation 4 losses. Negative sample states for TransE, 3, are randomly sampled from the rollouts.
3. Update the policy network's 's parameters, 6, using the combined loss. is defined, for a single rollout, T {(s,, ae Te, Se4 st41)).. )e, as follows:
L(T)= L (T) Σ CTram 81,11, 8:)
8 P())T)
where we set À = 0.001, and ((s) is the empirical distribution over the states in T.
We use the PPO implementation and hyperparameters from Kostrikov 25

---

It should be highlighted that our approach can easily be modified to support value- based methods such as DQN 12 merelly by modifying the tail component the network and the RL loss function.
Experiments
We now deploy XLVINS in generic discrete-action environments with unknown MDP dynamics (further details in Appendix D. and verify their potential as an implicit planner. Namely, we investigate whether XLVINS provide gains in data efficiency, by comparing them in low-data regimes against a relevant model-free PPO baseline, and ablating against the ATreeC implicit planner [I3] which executes an explicit TD(A) backup instead of latent-space executor, and is hence prone to algorithmic bottlenecks. All chosen environments were previously explicitly studied in the context of planning within deep reinforeement learning, and were identified as environments that benefit from planning computations in order to generalise [2] 1 30
4.1 Experimental setup
Common elements On all environments, the transition function, T, three-layer MLP with layer normalisation 4 after the second layer. The executor, X, is, for all environments, identical to the message passing executor of [Л. We train the executor from completely random deterministic graphs-make making no assumptions on the underlying environment's topology.
classical continuous-state control CartPol,,,, Acrobot and MountainCar, and a continuous-state spaceship navigation task, Lunarl In all cases, we study data efficiency by presenting extremely limited data scenarios.
Continuous-space We focus four OpenAl Gym environments
The encoder function a threc-layer MLP with ReLU activations, computing 50 output features and F hidden features, where F' = 64 for CartPole, F = 32 for Acrobot, F' 16 for MountainCar and F = 64 for LunarLander. The same hidden dimension is also used in the transition function MLP.
As before, we train our executor from random deterministic graphs. In this setting only, we also attempt illustrate the potential benefits when the graph distribution is biased by our beliefs of the environment's topology. Namely, we attempt training the executor from synthetic graphs that imitate the dynamics of CartPole very crudely-the MDP graphs being binary trees where only certain leaves carry zero reward and are terminal. More details on the graph construction, for both of these approaches, is given in Appendix E The same trained executor then deployed across all environments, to demonstrate robustness to synthetic graph construction. For LunarLander, the 2.
XLVIN uses K = 3 executor layers; in all other cases,
It is worthy to note that CartPole offers dense and frequent rewards rewards-making! it easy for policy gradient methods. We make the task challenging by sampling only 10 trajectories at the beginning, and not allowing any further interaction- beyond 100 epochs of training on this dataset. Conversely, the remaining environments are all sparse-reward, and known to be challenging for policy gradient methods. For these environments, we sample 5 trajectories at a time, twenty times during training (for a total of 100 trajectories) for Acrobot and MountainCar, and fifty times during training (for a total of 250 trajectories) for Lundrr.ander.
Pixel-space Lastly, we investigate how XLVINS perform on high-dimensional pixel-based obser- vations, using the Atari-2600 [ We focus on four games: Freeway, Alien, Enduro and H.E.R.O.. These environments encompass various aspects of complexity: sparse rewards (in Freeway), larger action spaces (18 actions for Alien and H.E.R.O.) and visually rich observations (changing time-of- day on Enduro) and long-ram rangerane credit assignment (on H.E.R.O.). Further, we successfully re-use the executor trained on random deterministic graphs, showing its transfer capability across vastly different settings. We evaluate the agents' low-data performance by allowing only .000,000 observed transitions. We re-use exactly the environment and encoder from Kostrikov (25), and run the executor for K 2 layers for Freeway and Enduro and K for Alien and H.F.R.O..
4.2 Results
In our results, we use "XLVIN-CP" to denote XLVIN executors pre-trained using CartPole-style synthetic graphs (where applicable), and "XLVIN-R" for pre-training them on random deterministic graphs. "PPO" denotes our baseline model-free agent; it has transition/executor model, but
7

---

Table 1: Mean scores for low-data CartPole-vO, Mo and LunarLander-v2, averaged over 100 episodes and five seeds.
CartPole-v0
Acrobot-vl 100 trajectories
MountainCar-vo 100 trajectories
LunarLander-v2 250 trajectories
Agent
10 trajectories
PPO
104.6 = 48.5
-500.0 0.0
-200.0 0.0
90.52 = 9.54
ATreeC
117.1
4 56.2
-500.0
+ 0.0
-200.0 + 0.0
84.04 + 5.35
XLVIN-R
199.2
1.6
-353.1
120.3
+ -185.6
99.34 +: 6.77
XLVIN-CP
195.2
5.0
-245.4
+ 48.4
-168.9 9 24.7
N/A
FPD
PP0
x VVNAR
R/VIN Rt
AfrenC
AfreeC
D.0
Avange roodet
Awsrage
1.5
100
62
04
0.5
86
1.
52
64
2.6
0.0
10
Obscreed Tarraloor
Obsirred traralarra
Endura
HFF R.G.
1913
10
80
Oboorved
1ed
02
tranatione
0.6
0.5
Figure 2: Average clipped reward on Freeway, Alien, Enduro and H.E.R.O. over 1,000,000 transitions and ten seeds.
otherwise matches the XLVIN hyperparameters. As planning-like computation was shown to emerge in entirely model-free agents [16] this serves as a check for the importance of the VI computation.
To analyse the impact of the algorithmic bottleneck, we use "ATreeC" [K] as of our baselines, capturing the behavior of a larger class of VI-based implicit planners (including TreeQN and VPN (3O). For maximal comparability, we make ATreeC fully match XLVIN's hyperparameters, except for the executor model. Since ATreeC's policy is directly tied to the result of applying TD(A), its ultimate performance is closely tied to the quality of its scalar value predictions. Comparing against ATreeC can thus give insight into negative effects of the algorithmic bottleneck at low-data regimes.
CartPole, Acrobot, Car and LunarLander Results for the continuous- control environments are provided in Table [ We find that the XLVIN model solves CartPole from only 10 trajectories, outperforming all the results given in 42 (incl. REINFORCE [46]. Autoencoders, World Models [8) DeepMDP details, see Appendix D
14 and PRAE
42
while using 10x fewer samples. For more
Further, our model is capable of solving the Acrobot and MountainCar environments from only 100 trajectories, in spite of sparse rewards. Conversely, the baseline model, as well as ATrecC, are unable to get off the ground at all, remaining stuck at the lowest possible score in the environment until timing out. This still holds when XLVIN is trained on the random deterministic graphs, demonstrating that the executor training need not be dependent on knowing the underlying MDP specifics.
Mirroring the above findings, XLVIN also demonstrates clear low-data efficiency gains on Lu- ander, compared to the PPO baseline and ATreeC. In this setting, ATreeC even underperformed compared to the PPO baseline, highlighting once again the issues with algorithmic bottlenecks.
8

---

+
-04
0.50
0.5
1.5
2
-1.0_0.pp.
-2 4
0
0.5
10
-5
15
Figure 3: Top: A test maze (left) and the PCA projection of its TransE state embeddings (right). colour-coded by distance to goal (in green). Bottom: PCA projection of the XLVIN state embeddings after passing the first (left), second (middle), and ninth (right) level of the continual maze.
Freeway, Alien, Enduro and H.E.R.O. Lastly, the average clipped reward of the Atari agents across the first million transitions can be visualised in Figure 2 From the inception of the training, the XLVIN model explores and exploits better, consistently remaining ahead of the baseline PPO model in the low-data regime (matching it in the latter stages of Enduro). In H.E.R.O.. XLVIN is also the first agent to break away from the performance plateau, towards the end of the 1,000,000 transitions. The fact that the executor was transferred from randomly generated graphs (Appendix B a further statement to XIVIN's robustness.
On all four games, ATreCC consistently trailed behind XLVIN during the first half of the training, and on Enduro, it underperformed even compared the PPO baseline, indicating that overreliance on scalar predictions may damage low-data performance. It empirically validates our observation of the potential negative effects of algorithmic bottlenecks at low-data regimes.
4.3 Qualitative results
The success of XLVIN hinges the appropriate operation of its two modules; the transition model and the executor GNN X. In this section, we qualitatively study these two components, hoping to elucidate the mechanism in which XLVIN organises and executes its plan.
To faithfully ground the predictions of T and X on an underlying MDP, we analyse pre-trained XLVIN agent with a CNN encoder and K = 4 executor steps randomly-generde 8 x 8 world environments. We chose grid- worlds because V can be computed exactly. We train on mazes of progressively increasing difficulty (expressed in terms of their level: the shortest-path length from the start state to the goal). We move on to the next level once the agent solves at least 95% of the mazes from the current level. On these environments, implicit planners that are aware of the grid-world structure. See Appendix for details, including the hyperparameters of the XLVIN architecture.
generally found that XLVIN is competitive with
Projecting the embeddings We begin by qualitatively studying the embeddings learnt by the encoder and transition model. At the top row of Figure we (left) colour-coded specific held-out 8 maze by proximity to the goal state, and (right) visualised the 3D PCA projections of the "pure-TraanEF" embeddings of these states (prior to any PPO training), with the edges induced by the transition model. Such model merely seeks to organise the data, rather than optimise for returns: hence, grid-like structure emerges.

---

At the bottom row, we visualise how these embeddings and transitions evolve as the agent keeps solving levels; at levels one (left) and two (middle). the embedder learnt to distinguish all 1-step and 2-step neighbours of the goal, respectively, by putting them on opposite parts of the projection space. This process does not keep going. because the agent would quickly lose capacity. Instead, by the time it passes level nine (right), grid structure emerges again, but now the states become partitioned by proximity: nearer neighbours of the goal are closer to the goal embedding. In a way, the XLVIN agent is learning to reorganise the grid; this time in way that respects shortest-path proximity.
V predictibility We hypothesise that the encoder func- tion 2 is tasked with learning to map raw states to a latent space where the executor GNN can operate properly, and then the GNN performs VI-aligning computations in this latent space. We provide qualitative study to validate this: for all positions in held-out 8 x 8 test mazes, we computed the ground-truie truth values (s) (using VI), and performed linear regression to test how accurately they are decodable from the XLVIN state embeddings before (hg) and after (Xx) applying the GNN. We computed the R² goodness- of-fit measure, after passing each of the ten training levels. Our results (Figure g are strongly in support of the hypoth- esis: while the embedding function (=(s) = has in green) already reasonably predictive of (s) (R² as 0.85 on average), after the GNN computations are performed, the recovered state embeddings (Xs: in red) are consistently almost-perfectly linearly decodable to VI outputs (s) (R² s 1). Hence, the encoder maps s to a latent space from which the executor can perform VI.
6.
6.
0.7
04
0.
Passed
Figure 4: Coefficient of determination from linearly regressing on the state embeddings obtained from the encoder (green) and from the executor (red).
Algorithmic bottleneck In formulating the algorithmic bottleneck, we assume that inaccuracies in the scalar val- ues used for VI will have a larger impact on degrading the performance than perturbations in gigh-dimensionall state embeddings inputted to the executor. To faithfully evaluate policy accuracy. we study this sensitivity on the randomly generated MDPs used to train the executor. Here, ground- truth values V* can be explicitly computed, and we can compare the recovered policy directly to the greedy policy over these values. Hence, we study the effect of introduc- ing Gaussian noise in the scalar inputs fed to VI compared to introducing Gaussian noise in the high- dighesiinnnii la- tents fed into the XLVIN executor. In Figure 5 we plot the policy accuracy as function of noise standard de- viation, showing that, while XLVIN is unable to predict policies perfectly zero noise, it quickly dominates the VI's policy predictions once the noise magnitude increases. Besides indicating that imperfections in TransE outputs can be handled with reasonable grace, this experiment provides direct evidence of the algorithmic bottleneck: errors in scalar inputs an algorithm can impact its predictions substantially, while high-dimensions latent space executor is able to more gracefully handle such perturbations.
100
W
00
XVV
Pulay ecursey
70
60
80
0.8
02
04 Nakus abencaard clcceaion
0.
08
Figure 5: Policy accuracy from introduc- ing Gaussian noise in the scalar input fed into VI (red) and in the embeddings fed into the XLVIN executor (green).
5 Conclusions
We presented eXecuted Latent Value Iteration Networks (XLVINs), combining recent progress in self-supervised contrastive learning, graph representation learning and neural algorithm execution for implicit planning on irregular, continuous or unknown MDPs. Our results showed that XLVINs match or outperform appropriate baselines, often at low-data or out-of-distribution regimes. The learnt executors are robust and transferable across environments, despite being trained on purely random graphs. XLVINS represent, to the best of our knowledge, one of the first times neural algorithmic exccutors used for implicit planning, and they successfully break the algorithmic bottleneck.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank the developers of PyTorch B1 We specially thank Charles Blundell and Jess Hamrick for the extremely useful discussions. Additionally, we would like to thank all of the anonymous reviewers of our work while it was under submission to ICLR ICML'2I and NeurlPS'21 substantially, and also broadened our own perspective about XLVIN's importance along the way.
every iteration, received insightful comments that strengthened the paper
AD and JT declare support from the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AT Chair Program, collaboration grants between Mierosoft Research and Mila, Samsung Electronics Co., Ldt., Amazon Faculty Research Award, Tencent AI Lab Rhino-Bo Bird Gift Fund and a NRC Collaborative R&D Project A14D-CORE-06 This project was also partially funded by IVADO Fundamental Research Project grant PRF-2019-358313977
PV Research Scientist at DeepMind. AD was Research Scientist Intern at DeepMind while completing this work.
Andreea Deac and Petar Veličković wish to dedicate this paper to their family for always being by their side, through thick and thin.
References
[1] Ashutosh Adhikari, Xingdi Yuan, Marc-Alexandre Côté, Mikuláš Zelinka, Marc-Antoine Rondeau, Romain Laroche, Pascal Poupart, Jian Tang, Adam Trischler, and William L Hamil- ton. Learning dynamic knowledge graphs to generalize on text-based games. arXiv preprint arXiv:2002.09127, 2020.
[2] Brandon Amos, Ivan Dario Jimenez Rodriguez, Jacob Sacks, Byron Boots, and J Zico Kolter. Differentiable mpe for end-to-end planning and control. arXiv preprint arXiv: 18 10. 13400, 2018.
[3] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare-Alexandre Côté, and Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, pages 8769-8782, 2019.
[4] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton, Layer normalization. arXiv preprint 107.00560 2016.
[5] Andrew G Barto, Richard S Sutton, and Charles W Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE transactions on systems, man, and cybernetics, (5):834-846, 1983.
[6] Mare Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[7] Richard Bellman. Dynamic programming. Science, 153(3731):34-77, 1966.
[8] Antoine Bordes, Nicolas Usunier, Alberto Garefa-Durán, Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi-relational data. In NIPS, pages 2787-2795, 2013,
[9] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jic Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:
2016.
[10] Quentin Cappart, Didier Chételat, Elias Khalil, Andrea Lodi, Christopher Morris, and Petar Veličković. Combinatorial optimization and reasoning with graph neural networks. arXiv preprint
2021.
[11] Andreea Deac, Pierre-Luc Bacon, and Jian Tang. Graph neural induction of value iteration. arXiv preprini 20t.::20009//66044 2020.
[12] Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor. Yuhuai Wu, and Peter Zhokhov. Openai baselines. nttps://gii tht.b/gtttuu.../m com/ comm/openaaa//aslllness 2017.
[13] Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning. In ICLR, 2018.

---

[14] Carles Gelada. Saurabh Kumar, Jacob Buckman, Ofrr Nachum, and Mare G. Bellemare. Deep- mdp: Learning continuous latent space models for representation learning. In ICML, volume 97, pages 2170-2179, 2019.
[15] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry, arXiv preprint arXiv:1704.01221 2017.
[16] Arthur Guez. Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Theophane Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, Greg Wayne, David Silver, and Timothy Lillicrap. An investigation of model-free planning. In ICML, volume 97, pages 2464-2473, 2019.
[17] Arthur Guez, Theophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wierstra, Rémi Munos, and David Silver. Learning to search with metsnets. In ICML, volume 80, pages 1817-1826. PMLR, 2018.
[18] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122 2018.
[19] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NeurlPS, pages 2455-2467, 2018.
[20] Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In JCML, volume 97, pages 2555-2565, 2019.
[21] Jessica B Hamrick, Abram L Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Buesing, Petar Veličković, and Théophane Weber. On the role of planning in model- deep reinforcement learning. arXiv preprint ar arXiv:2011.0402 2020.
[22] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model- based reinforcement learning for atari. arXiv preprint arXiv:1903.00331 2019.
[23] Thomas N. Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In ICLR, 2020.
[24] Martin Klissarov and Doina Precup. Reward propagation using graph convolutional networks. arXiv preprint arXiv:2010.02474. 2020.
[25] Ilya Kostrikov. Pytorch implementationns of reinforcement learning algorithms. https://
cithub..oom/kootriiikooo http.///o//ooor.ooooooooooodddda
www.-aacppppp-p-----aaa------ 2018.
[26] Lisa Lee, Emilio Parisotto, Devendra Singh Chaplot, Eric Xing, and Ruslan Salakhutdinov. Gated path planning networks. International Conference on Machine Learning, pages 2947-2955. PMLR, 2018.
[27] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. nature, 518(7540):529-53, 2015.
[28] Andrew William Moore. Efficient memory-based learning for robot control. 1990.
[29] Sufeng Niu, Siheng Chen, Hanyu Guo, Colin Targonski, Melissa C. Smith, and Jelena Kovacevic. Generalized value iteration networks: Life beyond lattices. In AAAI, pages 6246-6253 AAAI Press, 2018
[30] Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In NIPS, pages 6118 0128, 2017.
[31] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, perffrreffem deep learning library. In Advances in neural information processing systems, pages 8026-8037 2019.
[32] Martin Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[33] Sébastien Racanière, Theophane Weber, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia
Li, Razvan Pascanu, Peter W. Battaglia, Demis Hassabis, David Silver, and Daan Wierstra.
12

---

imagination-auggeent agents for deep reinforcement learning. In NIPS, pages 5690-5701, 2017.
[34]
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609 2020.
[35]
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: /tf.:/777.000 007.06647, 2017.
[36] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International Conference on Machine Learning, pages 8583-8592. PMLR, 2020.
[37] David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Amold, David P. Reichert, Neil C. Rabinowitz, André Barreto, and Thomas Degris. The predictron: End-to-end learning and planning. In ICML,, volume 70, pages 3191-3199, 2017.
[38] Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal plan- ning networks: Learning generalizable representations for visuomotor control. In International Conference on Machine Learning, pages 4732-4741. PMLR, 2018.
[39] Richard S Sutton. Generalization in reinforcement learning: Successful examples using sparse coarse coding. In Advances in neural information processing systems, pages 1038-1044, 1996.
[40] Aviv Tamar, Sergey Levine, Picter Abbeel, Yi Wu, and Garrett Thomas. Value iteration networks. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, NIPS, pages 2146-2154, 2016.
[41] Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy opti- mization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5981-5988, 2020.
[42] Elise van der Pol, Thomas Kipf, Frans A. Olichoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In AAMAS 2020, 2020.
[43] Petar Veličković and Charles Blundell. Neural algorithmic reasoning.
arXiv preprint
arXiv:2105.02761, 2021.
[44] Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execution of graph algorithms. arXiv preprint arXiv: 10593, 2019.
[45] Tingwu Wang. Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International Conference on Learning Representations, 2018.
[46] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8(3-4):22--256, 1992.
[47] Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? arXiv preprint arXiv: 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] We study value iteration-based implicit planners, identify algorithmic bottleneck issue with prior work in the area, and propose novel agent to rectify this issue.
(b) Did you describe the limitations of your work? [Yes] See Section 3.1. ("Discussion" paragraph) and 3.2. (w.rtt exploration/exploittiion tradeoff when collecting data).
(c) Did you discuss any potential negative societal impacts of your work? [Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
13

---

3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- URL)? [No] We share links to mental results (either in the supplemental material environments and PPO optimiser the base implementations of large parts of our agents,
2
We aim open source our full implementation a later point.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section 4.1. and relevant appendices for details of our agents' hyperparameters.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] All performance metrics and plots are swept over multiple seeds with error bars clearly marked.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix.
4. If you using existing assets (e-g., code, data, models) curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite both the base implementation we used (by Kostrikov) and the environment suite we used (OpenAI Gym).
(b) Did you mention the license of the assets? [Yes] Licenses included in the Appendix.
(c) Did you include any new assets either in the supplemental material or as a URL? [No] We aim to open source our full implementation at later point.
(d) Did you discuss whether and how consent obtained from people whose data you'le 're
using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "z7vF6mxCCg8", "reviewer": "Reviewer_wygb", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- This paper provides a solution to the issues presented in previous work along this line:\n  - Not all high-dimensional state spaces can be mapped properly to a lower-dimensional space for VIN-like updates.\n  - The implicit planning in previous work is prone to error, especially from the incorrect estimation of the reward and state values.\n- The improvements compared with existing baselines in low-data regimes are significant.\n- The paper is well written, and the experiments are solid with several extra studies beyond standard ones.\n\nCons:\n- Since the complexity of the proposed method scales with O(|A|^K), the model can only look ahead for a very small number of steps.\n  - Why would a small K help build a better state representation? For atari games, it seems that looking 1-2 steps ahead into the future won't provide a decent gain.\n- The discussion of how the quality of the learned TransE state embeddings will affect the final performance is not included in the evaluation section. Such sensitivity analysis will be helpful in seeing how XLVINs can generalize to more complicated state spaces (e.g., molecule space) where TransE style embeddings are hard to learn.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "_J0KL-Z_3K", "reviewer": "Reviewer_UvSf", "summary": "In the context of implicit planning algorithms, the paper proposes to replace exact value iteration (VI) computations with a learned graph neural network that is trained to perform VI computations in an abstract latent space.\n\nThe paper shows competitive performance on several benchmark datasets using this approach, against a PPO baseline, and an implicit planning algorithm using exact VI.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper argues that performing exact VI computations in an implicit planning setup suffers from an \"algorithmic bottleneck\". They argue that this algorithmic bottleneck occurs when the exact VI computations are performed on poor initial estimates of the value function. \n\nTo remedy this \"algorithmic bottleneck\" the paper proposes to instead perform (approximate) VI computations in a latent space using a pre-trained Graph Neural Network (GNN). The GNN consists of 3 parts, an encoder, an update function, and a decoder, and is pre-trained to perform approximate VI on a dataset of randomly sampled Markov Decision Processes, after which only the update function is kept and frozen.\n\nGiven this GNN, the \"executor\", that can perform approximate VI in a latent space, the paper uses a fairly standard implicit planning setup. A learned encoder maps the environment state into a latent space, and a learned transition function maps latent states and actions to the next (expected) latent state. Given this tree of latent states, the executor performs its approximate VI computations and finally, a learned policy and value network outputs the policy and value function.\n\nThe paper is exceptionally well written. Especially section 2 gives the reader an excellent overview of the relevant material needed to understand the contributions of the paper. The paper shows compelling results demonstrating that it's better than a solid baseline (PPO) and more importantly, ATreeC, which is the method they directly built on by replacing the exact VI computations.\n\nI have a few comments that I think could make the paper even stronger. \n\n1) The paper argues that the algorithmic bottleneck is a problem, but doesn't really offer any evidence of this, except in the full experiments where they show their proposed method is better than the one using the exact VI computations. Further, the paper doesn't give any intuition as to why performing the VI computations in a high dimensional latent space might avoid the algorithmic bottleneck. I would greatly appreciate a small-scale experiment that directly shows the effect of the \"algorithmic bottleneck\" and how doing (approximate) VI in a higher dim latent space avoids it. The authors could use the simple MDPs they use to train the executor, and inject noise in the value functions and show how their learned executor is more robust to noise than the direct VI equations. This along with a few sentences on the intuition why doing approximate VI in a higher dim space is more robust would greatly improve the paper in my opinion. \n\n2) It was not entirely clear to me how the whole system was trained. In particular, I find it hard to understand how the encoding function and transition functions are trained. The loss in eq. 4 seems to indicate that they're trained entirely offline on (S,A,S') triplets, but that can't be true since then the learned embeddings wouldn't align with the one used by the executor. It would be very helpful if the authors included the joint loss that is optimized with PPO, and/or an algorithm describing the training procedure (training the executor first, then freezing, etc.).\n\nA few minor things:\n - I'd have liked a small discussion of how well (or not) the approach handles multi-modal MDPs. Computing an expected state over a multi-modal outcome seems like a poor approach.\n - Line 156: \"... and do not project our state to a low-dimensional space\": well, you do, with your encoder function. I know what you mean, but this sentence is kind of confusing.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "7_FdYG0iLzM", "reviewer": "Reviewer_8TgE", "summary": "This paper proposes a new implicit planner module, XLVINs, for end-to-end policy learning, which mimics the value iteration planning procedure with a pre-trained graph convolution network and overcomes the existing limitations of previous VI-style planning modules. This paper carefully investigates the existing literature with a thorough analysis of the algorithmic limitations of previous VI modules. The experiments also provide a complete understanding of the proposed method. ", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is extremely well written with lots of insightful discussions and explanations. It is really a great pleasure to read this paper and everything is so easy to follow under its current form. The discussion of algorithmic bottleneck is insightful. \n\nFrom a technical perspective, the paper proposes to use a TransE style latent model and adopt a GCN-based VI module. These two techniques directly address the existing limitations of previous works, which look clearly novel to me. I like the idea of pre-training a VI-module using random graphs and was surprised to see it is really working ---- particularly the experiment where the ground truth information is recovered. \n\nIn general, I like this paper and feel happy to see it gets accepted. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There are some minor comments.\n\nRegarding the TransE, it seems a bit unjustified since it seems that even if we simply train a neural network to predict the next state (e.g., a transition model, or a world model), things would still work without any changes? Is there any specific reason for using TransE other than its simplicity?\n\nFinally, I definitely feel the experiment scenarios are sufficient in the current paper. If possible, the paper can be even stronger if results on more Atari games can be presented. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "1KmXbaDOv0u", "reviewer": "Reviewer_Li46", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nNote:\nI reviewed this paper for ICML 2021. The method is identical as far as I can tell. The writing quality has been improved at some places. In terms of experiments, the authors have removed known MDP maze experiments and have added one continuous space environment, LunarLander-v2. \n\n\nStrengths:\n- The paper tackles an important problem of planning under uncertain maps which is of interest to the community.\n- The method incorporates several prior techniques to tackle planning from pixel space which makes it complicated, but it is described in sufficient detail at most places.\n\n\nWeakness:\n- The paper aims to improve over VINs by tackling continuous space and pixel space environments. However, while doing so it seems the proposed method isn't as effective as VINs at discrete state planning under known MDPS. Instead of adding larger maze size experiments, the authors chose to remove all maze experiments. If the method isn't effective at known MDP discrete space planning, this should be acknowledged in the paper.\n- For continuous state space tasks, the authors only conduct experiments in only 4 simple gym games and for pixel space observations, only 3 Atari games. The motivation of using these benchmarks for testing planning algorithms is unclear. It would make more sense to test the proposed method for planning in high dimensional continuous state space 3D manipulation or navigation tasks. In any case, the experimental evaluation needs to be more comprehensive in my opinion to establish the effectiveness of XLVINs.\n- The paper misses some relevant related work, for example Universal planning networks (Srinivas et al. ICML 2018), Differentiable MPC for End-to-End Planning and Control (Amos et al., Neurips 2018). It would be useful to discuss how the proposed method differs from these methods.\n- The baselines used are weak in my opinion. The proposed model learns a transition model between latent states. It seems more similar to model-based RL methods than to VINs. I would have liked to see empirical comparisons to Universal planning networks (Srinivas et al. ICML 2018) or PlaNet (Hafner et al. ICML 2019).\n\n\nUpdate after author response:\nThe author response has addressed my major concerns. I am increasing my rating accordingly. I still believe larger mazes and robotics applications are better environments to test planning methods rather than a few Atari games.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes proposes eXecuted Latent Value Iteration Networks (XLVINs) to relax the assumptions of Value Iteration Networks.  While most neural architectures in deep RL are not well motivated, this paper builds on an important line of work that treats value representations as implicit planners.  This is very nice work that generalizes previous work.  Well done!
Decision: Accept (Spotlight)