XLVIN: eXecuted Latent Value Iteration Nets
Authors: Andreea Deac
Petar Veličković
Ognjen Milinkovic
Pierre-Luc Bacon
Jian Tang
Mladen Nikolic
Venue:  (2021)
Keywords: value iteration
graph neural networks
reinforcement learning

Abstract: Value Iteration Networks (VINs) have emerged as a popular method to perform implicit planning within deep reinforcement learning, enabling performance improvements on tasks requiring long-range reasoning and understanding of environment dynamics. This came with several limitations, however: the model is not explicitly incentivised to perform meaningful planning computations, the underlying state space is assumed to be discrete, and the Markov decision process (MDP) is assumed fixed and known. We propose eXecuted Latent Value Iteration Networks (XLVINs), which combine recent developments across contrastive self-supervised learning, graph representation learning and neural algorithmic reasoning to alleviate all of the above limitations, successfully deploying VIN-style models on generic environments. XLVINs match the performance of VIN-like models when the underlying MDP is discrete, fixed and known, and provide significant improvements to model-free baselines across three general MDP setups.

Under review as a conference paper at ICLR 2021
XLVIN: EXECUTED LATENT VALUE ITERATION NETS
Anonymous authors
Paper under double-blind review
ABSTRACT
Value Iteration Networks (VINs) have emerged as a popular method to perform implicit planning within deep reinforcement learning, enabling performance im- provements on tasks requiring long-range reasoning and understanding of envi- ronment dynamics. This came with several limitations, however: the model is not explicitly incentivised to perform meaningful planning computations, the un- derlying state space is assumed be discrete, and the Markov decision process (MDP) is assumed fixed and known. We propose eXecuted Latent Value Itera- tion Networks (XLVINs), which combine recent developments across contrastive self-supervised learning, graph representation learning and neural algorithmic rea- soning to alleviate all the above limitations, successfully deploying VIN-style models on generic environments. XLVINs match the performance of VIN-like models when the underlying MDP is discrete, fixed and known, and provide sig- nificant improvements to model-free basclines across three general MDP setups.
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 (also called model-free planning) (Tamar et al. 2016 Oh et al., 2017]
Racanière al., 2017 Silver et al.) 2017: Niu et al.. 2018 Guez et al. 2018 2019). Instead of training explicit environment models and leveraging planning algorithms, such approaches propose inductive biases in the policy function to enable planning to emerge, while training the policy in a model-free manner. A notable example of this line of research which observe that the value iteration (VI) algorithm on a grid-world can be understood as a convo- lution of state values and transition probabilities followed by max-pooling, which inspired the use of a CNN-based VI module (Tamar et al. (016). Generalized value iteration networks (GVINs), based on graph kernels (Yanardag & Vishwanathan 2015), lift the assumption that the environment is grid-world and allow planning on irregular discrete state spaces (Niu et al. 2018). such as graphs.
value iteration networks (VINs),
While such models can learn perform VI, they are in no way constrained explicitly incen- tivised to do so. Policies including such planning modules might exploit their capacity for different purposes, potentially finding ways to overfit the training data instead of learning how to perform planning. Further, both VINs and GVINs assume discrete state spaces, incurring loss of information for problems with naturally continuous state spaces. Finally, and most importantly, both approaches require the graph specifying the underlying Markov decision process (MDP) to be known in advance and are inapplicable if it is too large to be stored in memory, or in other ways inaccessible.
In this paper, we propose the eXecuted Latent Value Iteration Network (XLVIN), implicit planning policy network which embodies the computation of VIN-like models while addressing all of the above issues. As a result, are able to seamlessly run XLVINs with minimal configuration changes on discrete-action environments from MDPs with known structure (such as grid- worlds,,,)) through pixel-based ones (such as Atari), all the way towards fully continuous-state environments, consistently outperforming or matching baseline models which lack XLVIN's inductive biases.
To achieve this, we have unified recent concepts from several areas of representation learning:
Using contrastive relf-supervised representation learning, fer dynamics of the MDP, even when it is not provided. In particular, we leverage the work of Kipt al. (2000): van der Pol et al. (2020), which uses the TransE model (Bordes al.,
able to meaningfully in-

---

Under review as a conference paper at ICLR 2021
2013) embed states and actions into vector spaces, in such way that the effect of action embeddings onto the state embeddings are consistent with the true environment dynamics.
By applying recent advances in graph representation learning (Battaglia et al. 2018 Bron- stein et al.) [017]] Hamilton et al.] 2017). we designed a message passing architecture (Gilmer et al., 2017) which can traverse our partially-infereed MDP, without imposing strong structural constraints (i.e., our model is not restricted to grid-worlds)
We better align our planning module with VI by leveraging recent advances from neural learn polynomial-time graph algo- rithms (Xu et al., 2019 Veličković et al., 2019; Yan et al.] 2020 Geerriev & Lio 2020) Veličković et al. 2020). by supervising them structure their problem solving process ac- cording to target algorithm. Relevantly, it was shown that GNNs are capable of executing value iteration in supervised learning settings (Deac et al., 2020). To the best of our knowl- edge, our work represents the first implicit planning architecture powered by concepts from neural algorithmic reasoning, expanding the application space of VIN-like models.
algorithm execution, which has shown that GNNs
While we focus our discussion on VINs and GVINs, which we directly generalise and with which we share key design concepts (like the VI-based differentiable planning module), there are other
lines of research that our approach could be linked to. Significant work has been done on repre-
sentation learning in RL (Givan et al.] [2003; Ferns et al.] [2004] | [aderberg et al.) 2017 Ha &
Schmidhuber 2018h
Gelada et al.] 2019), often exploiting observed state similarities.
Regarding work in planning in latent spaces (Oh et al. 2017, Farquhar et al. 2018 Hafner et al,,
2019 van der Pol et al. 2020), Value Prediction Networks (Oh et al. 2017) and TreeQN (Farquhar et al. 2018) explore some ideas similar to our work, with important differences; they use explicit planning algorithms, while XLVINS do fully implicit planning in the latent space. However, due to the way in which value estimates are represented, the policy network is capable of melding both model-free and bosel-based cues robustly. Furthermore, while our VI executor provides a repre- sentation that aligns with the predictive needs of value iteration, it can also incorporate additional information it benefits the performance of the model.
BACKGROUND
Value iteration Value iteration is successive approximation method for finding the optimal value function of a discounted Markov decision processes (MDPs) as the fixed-point of the so-called Bell- man optimality operator (Puterman 20114) A discounted MDP is tuple (S,A, 7, where s E S are states, a E A are actions, R : S x A + IR is a reward function, P P $ x A + Dist(S) transition function such that P(s'|s, a) is the conditional probability of transitioning to state s' when the agent executes action a in state s. and E [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 stationary Markov policy : S x A - [0, 1] such that [[ass) is a conditional probability distribution over actions given a state. The return is defined as G2 (=(s,a) = Ez [2 [s = s, ag = a] 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 set- ting, we know that there exists an optimal stationary Markov policy n" such that for any policy æ it holds that V=r (s) Z (s) for all s E S. Furthermore, such optimal policy can be deterministic greedy with respect to the optimal values. Therefore, in order to find x² it suffices to find the unique optimal value function V as the fixed-point of the Bellman optimality operator. Value iteration is in fact the instantiation of the method of successive approximation method for finding the fixed-point of a contractive operator. The optimal value function V* is such a fixed-point and satisfies the Bellman optimality equations (Bellman 1966):
o y R(ac+&+ St++). Value functions (=(s,a) E,, [G,|s: = s] and
R(s,a) + +y Σ ((s)ls,))V'(5)) V (s) = max
(1)
a
s'ES
TransE The TransE (Bordes et al. 2013) loss for embedding objects and relations can be adapted to RL. State embeddings are obtained by an encoder 2 : S - R² and the effect of an action in given state is modelled by translation model T R x A + R Specifically, T(=(s), a) is
2

---

Under review as a conference paper at ICLR 2021
Encoder
Transition
Executor
Tail
a1)
2(6)
Thh up)
T(h as)
Construct latent graph using T
Execute (latent) Value Iteration using X"
Figure 1: XLVIN model summary. The individual modules are explained (and colour-coded) in Section 3.1 and the dataflow is outlined in Section 3.2
translation vector to be added to =(s) in order to obtain an embedding of the resulting state when taking action in state s Such embeddings should also be far away from all other states 2()). Therefore, the embedding functions are optimized using the following variant the triplet loss:
L d(=(s) + T(=(s), a), z(s')) + max(0. d(z(5),
where & is the hyperparameter of the hinge loss, (s, a, s') is transition negative sample state. Trajectories sampled from the MDP can be used as a training set of positive triplets (s, s'), whereas negative triplets (s, 8) are obtained by sampling & uniformly at random.
given MDP and is
GNNs and algorithmic executors Graph neural networks (GNNs) have been intensively studied as tool to process graph-structured inputs, and were successfully applied to various RL tasks (Wang al., 2018 Klissarov & Precup 2019; Adhikari et 2020 GNN operations are commonly ex- pressed through the message passing framework (Gilmer et al. 2017). For each node i in the graph, set of messages is computed one message for each neighbouring node j. derived by applying message function M to the relevant node and edge features (h,, hy, E.y). Incoming messages in a neighbourhood (i) are then aggregated through permutation invariant operator + (such as sum or max), obtaining a summarised message rī..
m
A JEN(()
M, M, ky.
The features node i are then updated through function U applied to these summarised messages: h
Un (h, m.
(4)
An important research direction explores the of neural networks for learning to execute algorithms- which was recently extended to algorithms on graph-structured data (Veličković et al., 2019). In particular, et al. (2019) establishes algorithmic alignment between GNNs and dynamic programming algorithms. Furthermore, Veličković et al.) (2019) show that supervising the GNN on the algorithm's ' intermediate results highly beneficial for out-of-distribution generalization. As VI is, in fact, a dynamic programming algorithm, a GNN executor is suitable choice for learning it, and good results on executing VI have emerged on synthetic graphs (Deac et al. 2020).
XLVIN ARCHITECTURE
With all of the above components in place, we can fully specify the computations of the eXecuted Latent Value Iteration Network (XLVIN). Throughout this section, we recommend referring to Fig- ure I for a visualisation of the model dataflow (more compact overview in Appendix A).
We propose policy network-a function, To (als), which for a given state s E S specifies a prob- ability distribution of performing each of the actions E A in that state. Here, 0 are the neural network parameters, to be optimised using gradient ascent.
For stochastic MDPs, the resulting emhedding should he understood an expectation under P(s'la,
3

---

Under review as a conference paper at ICLR 2021
3.1 XLVIN MODULES
Encoder The encoder function, ;: S RR consumes state representations E S and produces flat embeddings, h2 2(s) Rt. 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.
Transition The transition function, T R A + R* models the effects of taking actions, in the latent space. Accordingly, it consumes state embedding =(x) and an action a and produces the appropriate translation of the state embedding, to match embedding of the successor state (in expectation). That is, it is desirable that T satisfies the following property:
=s)) ((=(s), a) E (())a))))))
(5)
As 7 operates in flat embedding space, it commonly realised as an ML..
R RlAi*k + Rh generalises the planning modules proposed VIN (Tamar et al. 2016 and GVIN (Niu et al. 2018). It processes an embedding hig of given state s, alongside a neighbourhood set (hha), which contains (expected) embeddings of states that immediately neighbour 8 8-for example, through taking actions. Hence,
Executor The executor function. X
((h.) ~ { ((ee)P()))) (())) )EEA
The executor combines the neighbourhood set features produce an updated embedding of state X. X ((h)))(hh)), 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 to meaningfully plan from state 8 by stacking several layers of X (with K layers allowing for exploring length-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 X + [0, 1| consumes the state embed-
ding Fi and the updated state embedding X.s producing action probabilities To(a|s) A [ Xx specifying the policy to be followed by our XLVIN agent. Lastly, note that we may also have addi- tional tail networks which have the same input as A. We train XLVINs using proximal policy optimi- sation (PPO) (Schulman et al. 2017), which necessitates value function: V R R R. Hence, we also include as a component of our model.
3.2 OVERVIEW OF XLVIN
Putting all of the above components together, now present step-by-step algorithm used by XLVINs to derive a policy Ta (als), for a given input state and executor depth K:
1. Embed the input state by using the encoder function: hi 2(s).
2. Initialise the depth-O set of state embeddings to the singleton set, $0 {h.,,, containing only the embedding of s. Accordingly, initialise the set of edges to E = 0.
3. For each graph depth, A [0, K):
(a) Initialise the depth-(k: + 1) embedding set to Sk-1 = 0.
(b) For each state embedding in the previous depth, h Sh and action, a E A:: Compute (expected) neighbour embedding of h upon taking action a using the transition model: N/' T + T(),
ii Attach H'' to the graph, by gh+1 = g U = BE
(c) For each state embedding h. E st. define its neighbourhood as the set of all adjacent embeddings: (λA = {h' 3a E A. K', a) E E).
4. Run the execution model over the graph specified by the nodes $ = =K @l and edges E, by repeatedly applying h = =X(h,NNhh)) for every embedding E S, for K steps. Denote the updated embedding of obtained by the above procedure, X.-
U
4

---

Under review as a conference paper at ICLR 2021
5. Finally, use the actor and tail to predict the policy and value functions from the state em- bedding and updated state embedding of 8: (o( ) A(h, Xe): V()) V(h,,.a Xe).
Discussion
The entire procedure end-to-end differentiable, does impose any assumptions on the structure of the underlying MDP, and has the capacity to perform computations directly aligned with value iteration. Hence our model can be considered as a generalisation of VIN-like methods to settings where the MDP is not provided or otherwise difficult to obtain.
Our tree expansion strategy is breadth-first, which expands every action from every node, yielding O(IAA|)) time and space complexity. While this is prohibitive for scaling up to larger values of K, we have empirically found that performance tends to plateau by K I 4 for all environments we have considered, mirroring prior work on implicit planning (Farquhar et al. 2018).
We defer allowing for deeper expansions and large action spaces to future work, but note that will likely require a rollout policy, selecting actions to expand from given state. For example, 12A (Racanière et al., 2017) obtains a rollout policy by distilling the agent's policy network. Extensions of our model to continuous actions could also be achieved by rollout policies, although discretising the action space through binning (Tang & Agrawal 2020) is also a viable route.
3.3 XLVIN TRAINING
As discussed in Section 3.1 our transition function, T, and executor function, X, should ideally respect certain properties (e.g. Equations 5H We pre-train both of them using established methods
in the literature: (Kipf al. 2020 van der Pol al. 2020) for TransE and (Deac et al. 2020 for
the executor. Due to space constraints, we report detailed pre-training information in Appendix B
To optimise the neural network parameters 0 we use proximal policy optimisation (PPO Schul- man et al., 2017). Note that the PPO gradients also flow into parameters of T and X. which has the potential to displace them from the properties required by the above, leading to either poorly constructed graphs (that don't respect the underlying MDP) or lack of VI-aligned computation.
For the transition model, we resolve this by periodically re-training on the TransE loss during train- ing (with multiplicative factor of 0.001 to the loss in Equation 2 As we have no such easy way of retraining the executor, X, without knowledge of the underlying MDP, we instead opt to freeze the parameters of X after pre-training, treating them constants rather than parameters to be optimised.
While we found sufficient to fine-tune on TransE using only the trajectories obtained from our pol- icy network, we note that, especially once the policy becomes more exploitative, the data collected in this way may strongly bias the transition model and make unusable for exploration. Hence, we an- ticipate that some environments will require careful tradeoff between exploration and exploitation for the data collection strategy for training the transition model.
4 EXPERIMENTS
We now proceed to evaluating the proposed XLVIN architecture across three distinct kinds of envi- ronments, and variety of training and testing modes. Our work is centered around the following three research questions, which we will repeatedly refer to throughout this section:
Q1 How do XLVINs compare to (G)s when they are applicable (i.e. on fixed, known and discrete MDPs)? Is the graph built within XLVINs as useful as the underlying MDP?
Q2 Do XLVINs effectively generalise the VIN model? Do they outperform baseline architec- tures (without the transition and executor models) across general environments?
Q3 Does the XLVINS' environment model help them robustly plan? Are XLVINs amenable to dynamically changing environments (e.g. continual learning) or low-data regimes?
We use the available Py PPO implemmmmtnmnemnnnnnmmmmmttt hndeeprrntteeeeeeeeeeeereeeettt hrgn from the following repos- -ac- ppo-ppppo--epppopppp -cctr-gaaa
itory by Ilya Kostrikov: httpp:://ggttto...oo////ettteeet
5

---

Under review as a conference paper at ICLR 2021
4.1 EXPERIMENTAL SETUP
We will now categorise our environments into three types-for further details, see Appendix C
Known MDP In order to compare XLVINS performance against VINs and GVINs (Q1) we evalu- ate on an established environment with a known, fixed and discrete MDP. We use the 8 8x8 8 world mazes proposed by Tamarr al. 2016). The observation for this environment consists of the maze image, the starting position of the agent and the goal position. Every maze is associated with a difficulty level, equal to the length of the shortest path between the start and the goal.
We utilise this difficulty to formulate the continual maze task: the agent is, initially. trained to solve only mazes of difficulty 1. Once the agent reaches 95% success rate on the last 1,000 sampled episodes of difficulty d, it automatically advances to difficulty d + (without observing difficulty d again). If the agent fails to reach 95% within 1,000,000 trajectories, is considered to have failed. At each passed difficulty, the agent is evaluated by computing its success rate on held-out test mazes.
Setting up the task in this way allowed us evaluate the resilience of the agents to catastrophic for- getting (Q3)-as the difficulties get harder, the models get further inspired to memorise the training set (see Appendix D for further statisties) We also test out-of-distribution generalisation (training on 8 x 8 mazes, testing on 16 x 16), allowing for another way to test the agents' robustness (Q3).
Given the grid-world structure, our encoder for the maze environment is a three-layer CNN comput- ing 128 latent features and 10 outputs we describe in Appendix E It important to note that the VIN model of Tamar et al. (2016) performs the aggregation by slicing directly the features on the agent's coordinates. This assumes upfront knowledge of where the agent actually is the map (rather than having to infer it from data) and hence puts VINs in a privileged position. For a possibly fairer comparison (Q1), we also attempted to train VIN-mean and VIN-max-VIN models where the slicing is replaced by global average pooling or global max pooling.
The transition function is a three-layer MLP with layer normalisation (Ba et al. 2016) after the second layer, for all environments. For mazes, it computes 128 hidden features and applies ReLU.
The executor is, for all environments, identical the message passing executor Deac al. (2020). For mazes, we exploit the fact that the MDP is known: we train the executor exactly on the graph structures generated from the training mazes; please see Appendix F for further details. We apply the executor until depth K = 4. with layer normalisation (Ba et al. 2016) applied after every step.
Continuous-space observations The latent-space execution of XLVINS allow us to deploy it in generic -espenennsssseciillll........ continuous environments are now supported without hav- ing to discretise them. We investigate the performance of XLVIN across these environments (Q2), focusing on three classical control environments from the OpenAl Gym: CartPole-vO, Acrobot-vl and MountainCar-vO.. We also address Q3 by presenting extremely-limitedddan scenarios.
The encoder function is now three-layer MLP with ReLU activations, computing 50 output fea- tures and F hidden features, where F' = 64 for CartPole, F = 32 for Acrobot, and F = 16 for MountainCar. The same hidden dimension is also used in the transition function, which otherwise matches the one used for mazes.
The executor has been trained from synthetic graphs which are designed to imitate the dynamics of CartPole very crudely-the the MDP graphs being binary trees where only certain leaves carry zero re- ward and are terminal. We also attempt training the executor from completely random deterministic graphs-making no assumptions on the underlying environment. More details on the graph con- struction, for both of these approaches, given in Appendix F The same trained executor is then deployed across all environments, to demonstrate robustness to synthetic graph construction (Q3). In all cases, the XLVIN uses K 2 executor layers.
Before moving on, it is worthy to note that CartPole offers dense and frequent 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 beyyod 100 epochs of training on this dataset. Conversely, Acrobot and MountainCar are both sparse-reward, and known to be challenging for policy gradient. For these environments, we sample 5 trajectories at a time, twenty times during training (for a total of 100 trajectories).

---

Under review as a conference paper at ICLR 2021
Table 1: Mean scores for CartPole-vO, Acrobot-vl and MountainCar-vO after training, averaged over 100 episodes and five seeds. Bascline CartPole results reprinted from van der Pol ct al. 2020))
CartPole-vO 100 trajectories
Only 10 trajectories
Acrobot-v1
Score
REINFORCE
23.84 + 0.88
PPO
-500.0
WM-AE
114.47 17.32
XLVIN-R
-353.1 = 120.3
LD-AE
154.73 50.49
XLVIN-CP
-245.4 = 48.4
DMDP-H
72.81 = 20.16
PRAE, 5
171.53 = 34.18
MountainCar-vo
Score -200.0
PPO
104.6
1 48.5
PPO
XLVIN-R
199.2 1 1.6
XLVIN-R
-185.6 8.1
XLVIN-CP
195.2
XLVIN-CP
-168.9 + 24.7
Pixel-space Lastly, investigate how XLVINS perform on high-dimensional pixel-based observations (Q2), using the Atari-2600 (Bellemare al., (013). We focus on three games:
Freeway, Alien and Enduro. These environments encompass various aspects of complexity: sparse rewards (on Freeway), larger action spaces (18 actions for Alien), visually rich obser- vations (changing time-of-day on Enduro). Further, we suc- cessfully re-use (Q3) the executor trained on random determin- istic graphs. showing additionally that that it is comparable to using the CartPole-basee graphs Freeway (whose "up-and- down" structure aligns it somewhat to such environments). We evaluate the agents' low-data performance by allowing only 1,000,000 observed transitions (Q3). We re-use exactly the
KWW VN GAM
on M61 mam
environment and encoder from here and run the executor for
K 2 layers for Freeway and Enduro and K = 1 for Alien.
4.2 RESULTS
In results that follow, we use "PPO" to denote our baseline model-free agent; it has no transition/executor model, but oth- erwise matches the XLVIN hyperparameters. When applica- ble, we use "XLVIN-CP" to denote XLVIN executors pre- trained using strte synthetic graphs, and "XLVIN- R" for pre- vre-training them on random deterministic graphs.
Panaed
Figure 2: Success rate on 8 x 8 (top) and 16 16 (bottom) held- out mazes for XLVIN, PPO, VIN and GVIN, obtained after passing each level of 8 x 8 train mazes. VIN-mean and VIN-max failed to pass difficulty 1.
Continual maze The results of our continual maze eval- uation are summarised in Figure 2 We observe that, in- distribution. the XLVIN is competitive with all other models (including VIN and GVIN that fails to pass difficulty 6). It is also the most resilient to level 10, which has only one maze, and is hence very easy to overfit to.
The XLVIN remains competitive with the baseline model when evaluated out-of-distribution. There gap to VINs, which can be attributed to the previously mentioned slicing corroborated by the fact both VIN-mean and -max failed to pass even level of the slicing operation will underperform whenever some level of global context needs to be taken into account when determining the grid-world's transitioning or reward rules. In Appendix G we illus- trate explicitly how XLVINs outperform VINs in contextual maze: a the continual maze grid world, where such global context can modify the environment dynamics.
mazes. Generally, the
version of
CartPole, Acrobot and MountainCar Results for the continuous-space control environments are provided in [a We find that, for CartPole, the XLVIN model solves the environment from only 10 trajectories, outperforming all the results given van der Pol 2020 (incl. REINFORCE
a2c 2___poo_acktr/mode11 169-

---

Under review as a conference paper at ICLR 2021
7.5
6-8
26
Chserwe
Figure 3: Average reward on Freeway, Alien and Enduro over 1,000,000 processed transitions.
i
0.75
0.00
0.50
0.5
0.0
0.5
0.5
-1.0-0.5p. n.
A-022
0.5
0.5
10
-5
Figure 4: Top: A test maze (leff) 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 embed- dings after passing the first (left), second (middle), and ninth (right) level of the continual maze.
(Williams 1992).
Autoencoders, World Models (Ha & Schmidhuber 2018a), DeepMDP (Gelada al. 2019) and PRAE (van der Pol al. 2020)) while using 10x fewer samples.
Further, our model is capable of solving the Acrobot and MountainCar environments from only 100 trajectories, spite of sparse rewards. Conversely, the baseline model unable to get off the ground at all, remaining stuck at the lowest possible score in the environment until timing out. The above results still hold when XLVIN is trained on random deterministic graphs, demonstrating that the executor training need not be dependent on knowing the underlying MDP specifics.
Freeway, Alien and Enduro Lastly, the average reward of the Atari agents across the first million transitions can be visualised in Figure 3 From the inception of the training, the XLVIN model explores and exploits better, consistently remaining ahead of the baseline model in the low-data regime (matching it in the latter stages of Enduro). The fact that its executor was transferred from randomly generated graphs (Appendix F further statement the model's robustness.
4.3 QUALITATIVE RESULTS
We qualitatively study the embeddings learnt by the encoder (and transition model), hoping to elu- cidate the mechanism in which XLVIN organises its plan (and hence provide further insight on Q3).
At the top row of Figure 4 we (left) colour-coded a specific 8 x 8 test maze by proximity to the goal state, and (right) visualised the 3DD PCA projections of the "pure-TransE" embeddings of these states
8

---

Under review as a conference paper at ICLR 2021
61340
120
XUW
0.2
64
12
0.
04
6.5
02
6.5
19
Figure 5: Comparison of XLVIN and ATreeC update rules over the first 1,000,000 Atari transitions.
(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, visualise how these embeddings and transitions evolve the agent keeps solving levels; at levels one (leff) 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 (rightt), grid structure emerges again, but now the states become partitioned by proximity: nearer neighbours of the goal are closer to the goal embedding. In way, the XLVIN agent learning to reorganise the grid; this time in way that respects shortest-pati proximity.
ABLATION STUDIES
Our XLVIN model is powered by a synergy of learned latent-space models and latent-space graph algorithm executors. To probe the influence of these components, we conclude with two ablation studies, directly comparing our architecture incorporating pixel-based world models, or implicit planners that rely on explicitly modelling scalar quantities, such as ATreeC (Farquhar et al.) 2018).
Comparison to pixel-based world models Aligned with prior studies of VAE losses on Atari (Anand al., ) our attempts use world model based on pixel reconstruction (e.g. Ha & Schmidhhberr (2018)) for XLVINs were unsuccessful. This supports using transition models optimised in the low-dimensional latent space. For details and visualisations, Appendix H
Comparison to scalar value modelling We compare our results with the ATreeC architecture (Farquhar al., (2018). but still using the PPO loss. ATreeC relies on explicitly invoking TD(A) backup to combine scalar-level predictions in every node of the expanded tree and compute the final policy logits. Since the ATreeC policy is directly tied the result of applying TD(A), its ultimate performance is closely tied to the quality of its scalar value predictions.
The ATreeC model failed solve any of the continuous state environments (with its CartPole reward marginally exceeding the PPO baseline, at 117.1 + 56.2, and Acrobot and MountainCar staying at -500 and -200, respectively). On all three Atari environments, it consistently trailed behind XLVIN during the first 500,000 transitions (Figure and on Enduro, it underperformed even compared to the PPO baseline, indicating overreliance on scalar predictions may damage low-data performance.
CONCLUSIONS
We presented eXecuted Latent Value Iteration Networks (XLVINs), combining recent progress in self-supervised contrastive learning, graph representation learning and neural algorithm execution to generalise Value Iteration Networks to irregular, continuous or unknown MDPs. Our results have shown that XLVINs have matched or outperformed appropriate baselines, often at low-data or duttof-distribution regimes. The learnt executors are robust and transferable across environments, despite being trained on synthetic graphs that need not align with the underlying MDP. XLVINs represent, to the best of our knowledge, one of the first times neural algorithmic executors are used for implicit planning, and we believe that this is very promising direction for future RL research.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Ashutosh Adhikari, Xingdi Yuan, Marc-Alexandre Côté, Mikuláš Zelinka, Marc-Antoine Rondeau, Romain Laroche, Pascal Poupart, Jian Tang, Adam Trischler, and William L Hamilton. Learning dynamic knowledge graphs to generalize on text-based games. arXiv preprint 0rXi::2002.09127, 2020.
Réka Albert and Albert-Lászle Barabási. Statistical mechanics of complex networks. Reviews of modern physics, 74(1):47, 2002.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, pp. 8769-8782, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXivv1607.064500 2016.
Andrew Barto, Richard S Sutton, and Charles Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE transactions systems, man, and cybernetics, (5):834-844 846, 1983.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanzaez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806 01261, 2018.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Richard Bellman. Dynamic programming. Science, 153(3731):34-37, 1966.
Antoine Bordes, Nicolas Usunier, Alberto Gareía-Durán, Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi-relational data. In NIPS, pp. 2787-2795, 2013.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geomet- ric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18-42, 2017.
Andreea Deac, Pierre-Luc Bacon, and Jian Tang. Graph neural induction of value iteration. arXiv preprint arXiv:2009. /2604, 2020.
Paul Erdōs and Alfréd Rényi. On the evolution of random graphs. Publ. Math. Inst. Hung. Acad. Sci, 5(1):76000 1960.
Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning. In ICLR, 2018.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In AAAI, pp. 950-951, 2004.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM J. Comput., 1066):662--114,, 2011.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In ICML, volume 97, pp- 2170-2179, 2019.
Dobrik Georgiev and Pietro Lió. Neural bipartite matching. arXiv preprint arXiv:2005.11304 2020.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. arXiv preprint /70.00/22 2017.
R. Givan, T. Dean, and M. Greig. Equivalence notions and model minimization Markov decision processes. Artificial Intelligence Journal, 147:163-224, 2003.
10

---

Under review as a conference paper at ICLR 2021
Arthur Guez, Theophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wier- stra, Rémi Mumos, and David Silver. Learning to search with metsnets. In JCML, volume 80, pp. 1817-1826. PMLR,
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, pp. 2464- 2473, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprini
2018. 2018a.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NeurIPS, pp. 2455-2467, 2018b.
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha. Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML, volume 97, pp- 2019.
William L Hamilton, Rex Ying, and Jure Leskovec. Representation learning graphs: Methods and applications. arXiv preprint (.X.vv:7700055500 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167 2015.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In ICLR, 2017.
Thomas N. Kipf. Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In ICLR, 2020.
Martin Klissarov and Doina Precup. Graph convolutional networks as reward shaping functions, In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforeement learning. nature, 51 ((754))):5 533, 2015.
Andrew William Moore. Efficient memory-based learning for robot control. 1990.
Sufeng Niu, Siheng Chen, Hanyu Guo, Colin Targonski, Melissa Smith, and Jelena Kovacevic. Generalized value iteration networks: Life beyond lattices. In AAAI, pp. 6246-6253. AAAI Press, 2018.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In NIPS, pp. 6118-6128, 2017.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
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. Imagination-augmented agents for deep reinforcement learning. In NIPS, pp. 5690-5701, 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.065. 06347, 2017.
David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Armodd David P. Reichert, Neil C. Rabinowitz, André Barreto, and Thomas Degris. The predictron: End-to-end learning and planning. ICML, volume 70, pp. 3191-3199, 2017.
Jost Tobias Springenberg, Alexey Dosovitskiy, Thomas Brox, and Martin Riedmiller. Striving for simplicity: The all convolutional net. arXiv preprint arXiv: 14/. 2806 2014.

---

Under review as a conference paper at ICLR 2021
Richard S Sutton. Generalization in reinforcement learning: Successful examples using sparse coarse coding. In Advances in neural information processing systems, pp. 1088,044, 1996.
Aviv Tamar, Sergey Levine, Pieter Abbeel, Yi Wu, and Garrett Thomas. Value iteration networks. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), NIPS, pp. 2146-2154, 2016.
Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5981-5988. 2020.
Elise der Pol, Thomas Kipf, Frans A. Oliehoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In AAMAS 2020, 2020.
Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execu- tion of graph algorithms. arXiv preprint arXiv:1910.10599 2019.
Petar Veličković, Lars Buesing, Matthew Overlan, Razvan Pascanu, Oriol Vinyals, and Charles Blundell. Pointer graph networks. arXiv preprint arXiv:2006.065: 06380, 2020.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with
graph neural networks. In International Conference on Learning Representations, 2018.
Ronald J Williams. Simple statistical gradient-olllwingg algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):22--256, 1992.
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? arXiv preprint arXiv: /32/1, 2019.
Yujun Yan, Kevin Swersky, Danai Koutra, Parthasarathy Ranganathan, and Milad Hashemi. Neural 7JA77
Execution Engines, 2020. URL http::: https://ppenrevieww
unrii-mJiirrg BB4YYDr
Pinar Yanardag and SVN Vishwanathan. Deep graph kernlss. In Proceedings of the 21th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, pp. 1365-1374, 2015.
12

---

Under review as a conference paper at ICLR 2021
T Th, a1)
hig,a
2(s)
ht
X
hs,02
h,saa
X (hs,))(hs))
Figure 6: XLVIN model summary with compact dataflow. The individual modules are explained
(and colour-coded) in Section 3.1 and the dataflow is outlined in Section 3.2
A ALTERNATE RENDITION OF XLVIN DATAFLOW
See Figure for an alternate visualisation of the dataflow of XLVIN-whidl is more compact, but does not explicitly sequentialise operations the transition model with the operations of executor.
B PRE-TRAINING OF TRANSITION AND EXECUTOR FUNCTIONS
The transition condition of Equation 5 perfectly aligns with the TransE loss (Bordes et al. 2013), hence we pre-train T by optimising the loss function of Equation using transitions sampled the environment using random policy-not unlike the prior work by Kipf al. (2020): van der Pol et al. (020) that also trains transition functions in this way.
The desirable behaviour of the executor is align with VI, and hence pre-train X to be predictive of the value updates performed by VI, following the work of Deac et al. (2020). Note that standard VI procedure requires access to fully-specified MDP, to use as training data. When an MDP is not available, following the remarks of Deac et al. (2020), generate synthetic discrete MDPs that align with the target MDP much as possible
which can generate VI trajectories
finding
that useful transfer still occurs, echoing the findings of Deac et al. (2020).
To summarise in brief, the executor function training proceeds follows:
1. First, generate a dataset of MDPs which, as much as possible, mimics in some way the characteristics of the MDP we'd like deploy on (e.g. determinism, number of tions, etc.). If no such knowledge is available upfront, resort to generic graph distribution like Erdös-Réó or Barabási-Allee Alberr.
Then, execute the Value Iteration algorithm on these MDPs, keeping track of intermediate values Vi(s) at each iteration I, until convergence.
3. Supervise the executor network-a graph neural network operating over the transitions of the MDP as edges- to receive Vi(s)- and and all other parameters of the MDP as inputs, and predict V (s) outputs (optimise using mean-squared error).
13

---

Under review as a conference paper at ICLR 2021
Figure 7: The seven environments considered within our evaluation: x and 16x 16 mazes Tamar al. 2016) (known grid-like MDP), continuous control environments (CartPole-vO,, Acrohot-vI, MountainCar-vO) and pixel-based environments (Atari Freeway, Alien and Enduro).
C ENVIRONMENTS UNDER STUDY
We provide visual overview of all seven environments considered in Figure 7
Maze The maze environment with randomly generated obstacles from Tamar et al. (0016). Obser- vations include map of the maze (pointing out all obstacles), an indicator of the agent's location, and an indicator of the goal location. Actions correspond to moving in one of the eight principal directions, incurring a reward of -0.01 every move (to encourage shorter solutions). an obstacle (which terminates the episode), and for hitting the goal (which also terminates the episode).
for hitting
CartPole The CartPole environment is classic example of continuous control, first proposed by Barto et al.] (1983). The goal is to keep the pole connected by an un-actuated joint to a cart in an upright position. Observations velocity as well as pole's angle from vertical and pole's velocity at the tip. Actions correspond to staying still, pushing the engine forwards backwards. The agent receives fixed reward of +1 for every timestep that the pole remains upright. The episode ends when the pole is more than 15
four-dimensional vectors indicating the cart's position and
degrees from the vertical, the cart moves more than 2.4 units from the center or by timing out (at 200 steps), at which point the environment is considered solved.
Acrobot The Acrobot system includes two joints and two links, where the joint between the two links is actuated. Initially, the links are hanging downwards, and the goal is to swing the end of the lower link up to a given height. The environment was first proposed by SSttton 1996). The ceevvations-spcifyiiii
in full the Acrobot's corfguration-coosiiuuee six-dimensional vector, and the agent is able to swing the Acrobot using three distinct actions. The agent receives a fixed negative reward of -1 until either timing out (at 500 steps) swinging the acrobot up, when the episode terminates.
MountainCar The MountainCar environment is an example of challenging, sparse-reward, continuous-space environment first proposed by Moore (990).. The objective is to make car reach the top of the mountain, but its engine is too weak to go all the way uphill, so the agent must use gravity to their advantage by first moving in the opposite direction and gathering momentum. Obser- vations are two-dimensional vectors indicating the car's position and velocity. Actions correspond to staying still, or pushing the engine forward or backward. The agent receives fixed negative reward until either timing out (at 200 steps) or reaching the top, when the episode terminates.
Freeway Freeway is game for the Atari 2600, published by Activision in 1981, where the goal is to help the chicken cross the road (by only moving vertically upwards or downwards) while avoiding cars. is a standard part of the Atari Learning Environment and the OpenAI Gym.
Observations in this environment are the full framebuffer of the Atari console while playing the game, which has been appropriately preprocessed in Mnih al. (2013). Actions correspond to staying still, moving upwards or downwards. Upon colliding with car, the chicken will be set back a few lanes, and upon crossing road, will be teleported back at the other side to cross the road again (which is also the only time when it receives positive reward of +1). The game automatically times out after fixed number of transitions.
'If no prior knowledge about the environment is known, one might resort to generic graph distributions, such Brdös-Rényi (Erdós Rényi 1960) Barahási-Allert (Albert & Barabási 2002).
14

---

Under review as a conference paper at ICLR 2021
100
1000
D00
Figure 8: Number mazes with the optimal path of given length from: 8 x train dataset (left), 8 x 8 test dataset (middle) and 16 x 16 test dataset (right).
Enduro Enduro is game for the Atari 2600, published by Activison in 1983. The goal of the game is to complete an endurance race, overtaking a certain number of cars each day of the race to continue the next day. It is standard part of the Atari Learning Environment and the OpenAl Gym.
Observations in this environment the full framebuffer of the Atari console while playing the game, which has been appropriately preprocessed as in Mnih et al. (2015). This game is one of the first games with day/night cycles as well weather changes which makes particularly visually rich. There are nine different actions we can take in this environment corresponding to staying still as well accelerating, decelerating, moving left/right and combinations of of them.
Alien Alien is game for the Atari 2600, published by 20th Century Fox in 1982. The goal of the game is to destroy the alien eggs laid in the hallways (similar to the pellets in Pac-Man) while running away from three aliens on the ship. is standard part of Atari Learning Environment and the OpenAl Gym.
Observations in this environment are the full framebuffer of the Atari console while playing the game, which has been appropriately preprocessed as in Mnih et al. (2015). There are 18 different actions we can take in this environment corresponding to staying still, firing the flamethrower and moving firing the flamethrower in eight directions.
D DATA DISTRIBUTION OF MAZES
We provide an overview of simple count-based statistics of the maze datasets, stratified by difficulty, in Figure [ Namely, we can observe that the distribution of 8 x 8 maze difficulties follows power-law, whereas the x 16 maze difficulty counts decay linearly. This poses an additional challenge when extrapolating out-of-distribution, the distributions of the two testing datasets vary drastically-and what worked well for one's performance measure may not necessarily work well for the other.
ENCODER ON CONTINUAL MAZE
Given the grid-world structure, our encoder for the maze environment is CNN. We stack three con- volutional layers computing 128 features, of kernel size 3 3, each followed by batch normalisation (Toff & Szegedy 2015) and the ReLU activation. We then aggregate all positions by performing global average pooling (Springenberg et al., 2014). Finally, the aggregated representation is passed to a two-layer MLP with hidden dimension 128 and output dimension 10, with a ReLU activation in-between. The output dimension was chosen be comparable with VIN (Tamar al. 2016).
15

---

Under review as a conference paper at ICLR 2021
Figure 9: Synthetic graphs constructed for pre-training the GNN executor: Maze (left), CartPole (middle) and random deterministic (20 states, 8 actions) (right)
F SYNTHETIC GRAPHS
Figure presents the three kinds of synthetic graphs used for pretraining the GNN executor. The one for mazes (left) emphasises the terminal node as a node to which all nodes are connected; all other nodes have maximum eight neighbours, corresponding to the possible action types.
In all other cases, we pre-train the executor using randomly generated deterministic graphs (right): for [ = 20 states and |A| = actions, create [S]-node graph. For each state-action pair we select, uniformly at random, the state it transitions to, deterministically. We sample the reward model using the standard normal distribution. Overall, the graphs are sampled as follows:
a(s,a) Uniform((S1)
P(s' (8,) [I ='
(8)
lo otherwisc
R(s,a)
10)
(9)
These k--N style graphs do not assume upfront any structural properties of the underlying MDP, and are good prior distribution for evaluating the performance XLVIN.
For CartPole-style environments, attempt third type of graph (middle). It is binary tree, where red nodes represent nodes with reward 0, and blue nodes have reward This aligns with the idea that going further from the root, which is equivalent with taking repeated left (or right) steps, leads to being more likely to fail the episode.
We also attempt using the CartPole graph for pre- pre-trainn training the executor for the other two continuous- observation environments (MountainCar, Acrobot) and for Freeway. Primarily, the similar action space of the environments is a possible supporting argument of the observed transferabilityy More- over, MountainCar and Acrobot can be related to a inverted reward graph of CartPole, with more aggressive combinations left/right steps often bringing a higher chance of success.
G CONTEXTUAL MAZE ENVIRONMENT
The grid- world based VIN architecture relies on the assumption that the agent is in fixed and known environment, and that the current agent coordinate within this environment is known. This allows value iteration to be aligned with weight-shared convolutional neural networks, followed by a slicing operation in the agent coordinate. While this gives the final embedding higher relevance compared pooling all of the pixels together, it also makes the final embedding tightly localised around an agent's 's current state. Should the environment possess any global context that dictates the transition and reward dynamics, the VIN model may lack ways easily integrate this necessary contextual information.
To illustrate this issue, propose slight modification to continual maze environment: the contextual maze. We randomly sample two scalars, 6 ~ (00 1), and condition the reward model depending on their sum; if a + < the agent must proceed to the goal as usual; otherwise, the environment's 's reward model is inverted: reaching the goal now provides negative reward, and hitting (any) wall brings positive reward (and terminates the episode).
Note that could not have performed the usual flatten operation, in order generalise 16 16
16

---

Under review as a conference paper at ICLR 2021
GELL L
Figure 10: Procedurally generated contextual mazes. The green cell represents the goal, and the blue cell represents the agent. Black cells represent walls. The upper and left borders' intensity modified to mirror the sampled values of and b, determining the environment dynamies.
75
Accuracy on contextial mazes
66
58
PPO
50
XLVIN
VIN
Level
Figure 11: Success rate of the XLVIN model against VIN and PPO on the contextual maze environ- ment, after training on each level, on held-out test mazes.
This context is provided the agent through the colour intensity of the left and top border, match the values of and b, respectively see Figure I0 for an illustration of several procedurally gener- ated contextual mazes. Successful agents now must meaningfully extract the contextual information before committing to a plan, which may be unfavourable for the VIN architecture.
Organising the environment into levels (based on the agent's 's initial distance the goal), as before, and learning in continual setting, we present the results on held-out test mazes in Figure I1 On contextual mazes, clear gap opens up between the slicing- slicing-based VIN architecture and context- based models (XLVIN and PPO). Further, XLVINs once again demonstrate stronger resilience to catastrophic forgetting compared to the baselines, by not overfitting to the hardest levels, which also have the fewest mazes.
H PIXEL-BASED WORLD MODELS
We attempt replacing Atari transition model with variant that learns representations through pixel-based reconstructions (using a VAE objective, as done by Ha & Schmidhuber (2018b)). We found that representations obtained in this way encodings converged to a fixed-point, and that the pixel-space reconstructions completely ignored the foreground observations (see Figure 12) This aligns with prior investigations of VAE-style losses on Atari, which found they tend to overly focus on reconstructing the background and were less predictive of RAM state than latent-space models, well randomly-initialisdd CNNs (Anand) et al. 2019). This comparison stands in favour of our approach to using transition model optimised purely in the liver-dimeniionll latent space.
not useful, observed that most of state
17

---

Under review as a conference paper at ICLR 2021
Figure 12: Freeway frames (above) and reconstructions (below) using a VAE-style world model.
18


Reviews:
[{"review_id": "iEfpwdrRchz", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe paper proposed a framework that combines TransE-style world model learning and value iteration networks. Section 3.2 gives a nice and clear overview of the algorithm: generating trees by calling the learned world model, run value iteration networks on the tree, and use the value iteration algorithms to compute the value and the policy at each state.\nSince the output of the algorithm is a policy, the model was trained using PPO.\n\nComments:\nThe overall presentation of the algorithm is quite clear, and the proposed model is a nice combination of value iteration networks with learned world models. My comments will thus focus on the claims made by the authors and the completeness of the experiment section.\n\n- Claims\n1. The authors have made very strong arguments about the generality of the proposed method: \"As a result, we are able to seamlessly run XLVINs with minimal configuration changes on environments from MDPs with known structure (such as grid-worlds), through pixel-based ones (such as Atari), all the way towards fully continuous-state environments, consistently outperforming or matching baseline models which lack XLVIN’s inductive biases.\" The authors have made several arguments about non-discrete and non-deterministic environments in the paper, however, they haven't shown any results on challenging continuous and stochastic environments. For example, how will the model perform in environments with high-dimensional action spaces: such as torques for robots? How will the model perform in partially-observable environments? If the proposed algorithm does not work for continuous action spaces (Page 4, footnote), the authors should make their claims weaker: finite and discrete action spaces + finite planning horizon makes the number of states visited also finite.\n\n- Experiments\n1. Overall the experiments are relatively weak. Adding more results on high-dimensional continuous environments (such as humanoids) or more difficult atari games could possibly help.\n2. The contribution of the proposed algorithm comes from two parts: the world-model part and the VIN+Policy Network part. However, there are only few fine-grained ablation studies. For example, there have been a great number of models that learn the forward dynamics of the environment. How will other forward dynamics models perform in the same framework (e.g., Ha & Schmidhuber 2018a)? Does the TransE-style low-dimensional prior help?\n3. Related to my previous comment: is there any way that the authors can make a quantitative evaluation of the learned world models?\n4. The authors should make comparisons with other models that do explicit/implicit planning based on the same learned world models, e.g., Value Prediction Networks (Oh et al., 2017) and TreeQN (Farquhar et al., 2018).\n5. There naturally exists a chicken-and-egg problem between learning world models and learning task-specific policies: learning good world models in complex environments requires good policies to collect useful data, while a better world model can improve task-specific policies. I think the authors should be more clear about how they are collecting data for training their TransE models and discuss the trade-off between exploration and exploitation in their model.", "extra_scores": {}, "extra_text": {}}, {"review_id": "_JgiyPWZ5m7", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\n\nThe authors propose a generalization of Value Iteration Networks to unknown, potentially continuous state spaces. They describe a framework for leveraging a learned graph embedding model (TransE) in combination with a deep RL model and an execution model based on graphical message passing to perform a VI-like operation. The authors show improved performance compared to baselines on a grid-world task with a known MDP, as well as several simple continuous control environments and the Atari game Freeway.\n\n**General**\n\nOverall the approach is interesting, but the current results are somewhat limited and the model, in particular the executor, seems to depend critically on the availability of a suitable graph for pre-training. The authors leverage a \"CartPole-inspired\" graph for the majority of the tasks, and the tasks seem to have been chosen to be consistent with the statistics of this synthetic graph (from the text, Freeway's \"'up-and-down' structure aligns it somewhat to environments like CartPole\"). While the authors illustrate a limited ability to generalize from graph/task misalignment on the maze task (8x8 to 16x16), it's still not clear that the model will generalize to more complex tasks where it's not straightforward to pre-train on a representative graph. For instance, all of the non-maze tasks appear to use simple ternary actions (e.g. left, right and stay for CartPole; up, down and stay for Freeway). As well, working from a task-relevant pre-training graph appears to give the agent foreknowledge of the task itself, which makes comparisons to baselines like PPO unfair.\n\nThe main question here for me is whether the paper reaches the threshold on originality and significance. The paper builds on the VIN/GVIN work and particularly depends on Deac et al. 2020 for illustrating VI with GNNs. It provides the first extension of these works to more complex tasks with unknown MDPs. However, the limited results and dependence on the availability of a suitable graph makes the significance of these extensions unclear. Nonetheless, it could lay the framework for future work that pushes the VI framework into new domains.\n\nSection 4.3 provides a nice clarification on what the model is doing under the hood.\n\n**Specific**\n\n- I would like to see more exploration of the executor depth, i.e. do we see gradual improvement as the depth is increased?\n\n- Are the Freeway with respect to the first 1 million transitions (text) or 100000 (figure)?\n\n- \\alpha is used for both the hinge loss in Eq. 2 and 3.2 3.(c) which is somewhat confusing at first glance. \n\n**Pros:**\n- Extends VI framework to more complex, unknown tasks.\n- Well-written and clear, excluding the specific issues noted above.\n- Nice qualitative results\n\n**Cons:**\n- Results are provided on a limited number of simple tasks, primarily using PPO as a baseline.\n- Tasks seem to have been chosen such that they have a representative graph available.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ACNtmeSDiea", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposed a novel policy prediction model that combines self-supervised contrastive learning, graph representation learning and neural algorithm execution to generalize the Value Iteration Networks to MDPs. The method described in the paper is a combination of existing works in the literature but seems to work well in practice. The experiments evaluate multiple aspects of the proposed model (E.g. number of executor layers, etc.) and show significant performance improvement over the existing approaches.\n\nThe latent representation used for policy prediction implicitly incorporates two-step aggregation 1) from initial representation extraction through encoding network and 2) from message passing, which seems helpful for the generalization.\n\nAs I am not actively tracking RL literature, I am not sure if there is a similar approach has already been proposed or not. My comments are based on the assumption that no existing work using GNN to do further representation learning on top of the individual encoded information. It would be better if the authors clearly state this novelty at the end of related works.\n\nPros\n===\n1. The paper is very well written and provided sufficient background knowledge to let the reader follow the description.\n2. While it appears to be a simple combination of existing techniques, the proposed model shows the benefit of obtaining better latent state representations for the policy prediction task. \n\nCons\n===\n1. The novelty of the proposed model is a bit weak in terms of a lack of specialization for this particular task. \n2. Figure 1 in the paper is not quite meaningful. A better demonstrative figure would help improve this paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "2LRGzVQTi_x", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper tackles an open problem of the value-iteration-network-paradigm. The proposed method (XLVIN) has a conceptual edge over traditional value iteration networks in that it can be applied to continuous problems and problems where the state space is either too big or not fully known in advance. The experiments mostly succeed in making the case that XLVINs:\n\n* are a drop-in replacement for traditional VINs in discrete settings\n* are able to improve on models without explicit planning models in general\n\nOne issue with the first point is that the experiments on the $16\\times16$ environments are in favor of VINs against XLVINs. The explanation provided in the paper makes sense but on the whole I get the impression that more experiments are necessary to asses whether XLVINs can truly replace VINs in discrete environments. \n\nThere is also quite a bit of algorithmic complexity that results from using the graph neural network paradigm. This is justified by the fact that XLVINs are also applicable to a much wider range of problems. However, as a reader, I'd appreciate it a lot if the paper would spend time on discussing some of that complexity in greater detail (I go into specifics under \"Cons\"). Regardless, I recommend accepting the paper as the ideas presented here are relevant to the VIN-community.\n\n### Pros\n\n* The method is applicable to continuous problems and large or not fully-known state-spaces, where VINs are not.\n* The experiments on continuous problems show an edge over methods without a planning component.\n* The method has an edge in the low-data regime.\n\n### Cons\n\n* Some important questions are not addressed:\n  * How does the method scale with the planning horizon $K$?\n  * How feasible is it to produce synthetic MDPs on real-world problems?\n  * How feasible is it to use generic graphs instead?\n\n  It's understandable that some content needs to be left out due to space constraints, but at the same time, these questions would be at the heart of future work on this topic and the paper would benefit from addressing them.\n\n* The experiments on the $16\\times16$ mazes no longer support the claim that XLVINs can replace VINs gracefully. An explanation is provided for this, though it's still unclear how the two models compare in discrete environments other than the $8\\times 8$ mazes. I feel like more experiments are needed here, perhaps in scenarios that don't give an unfair advantage to VINs. The mean- or max-pooled VINs are not convincing for that because they fail to solve even the simplest tasks.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The work extends the line of work based on value iteration networks. The main goal is to extend VINS to continuous and partially observable state spaces. The approach combines self-supervised contrastive learning and graph representation learning with VINs to address these issues. Reviewers liked the premise of the paper and had several follow-up clarifications.  The authors provided the rebuttal and addressed some of the concerns. However, upon post rebuttal discussion, the reviewers decided to maintain their score. While everyone recommended weak acceptance, no one championed the paper. This was primarily due to the concerns in the empirical analysis. It is not clear that XLVINs are clearly outperforming VINs and Graph-VINs in all settings. All baselines are not present in all the environments, so it is difficult to draw a consistent conclusion. The paper is in a good state but not fully polished to infer clear conclusions about the effectiveness of the proposed approach. Please refer to the feedback below for more details. We believe strengthening the experimental results section will turn this paper into a very strong submission.
Decision: Reject