XLVIN: eXecuted Latent Value Iteration Nets

1 INTRODUCTION . Planning is an important aspect of reinforcement learning ( RL ) algorithms , and planning algorithms are usually characterised by explicit modelling of the environment . Recently , several approaches explore implicit planning ( also called model-free planning ) ( Tamar et al. , 2016 ; Oh et al. , 2017 ; Racanière et 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 are value iteration networks ( VINs ) , which observe that the value iteration ( VI ) algorithm on a grid-world can be understood as a convolution of state values and transition probabilities followed by max-pooling , which inspired the use of a CNN-based VI module ( Tamar et al. , 2016 ) . Generalized value iteration networks ( GVINs ) , based on graph kernels ( Yanardag & Vishwanathan , 2015 ) , lift the assumption that the environment is a grid-world and allow planning on irregular discrete state spaces ( Niu et al. , 2018 ) , such as graphs . While such models can learn to perform VI , they are in no way constrained or explicitly incentivised 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 ) , an implicit planning policy network which embodies the computation of VIN-like models while addressing all of the above issues . As a result , we 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 self-supervised representation learning , we are able to meaningfully infer dynamics of the MDP , even when it is not provided . In particular , we leverage the work of Kipf et al . ( 2020 ) ; van der Pol et al . ( 2020 ) , which uses the TransE model ( Bordes et al. , 2013 ) to embed states and actions into vector spaces , in such a 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. , 2017 ; Hamilton et al. , 2017 ) , we designed a message passing architecture ( Gilmer et al. , 2017 ) which can traverse our partially-inferred 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 algorithm execution , which has shown that GNNs can learn polynomial-time graph algorithms ( Xu et al. , 2019 ; Veličković et al. , 2019 ; Yan et al. , 2020 ; Georgiev & Lió , 2020 ; Veličković et al. , 2020 ) , by supervising them to structure their problem solving process according to a 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 knowledge , our work represents the first implicit planning architecture powered by concepts from neural algorithmic reasoning , expanding the application space of VIN-like models . 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 representation learning in RL ( Givan et al. , 2003 ; Ferns et al. , 2004 ; 2011 ; Jaderberg et al. , 2017 ; Ha & Schmidhuber , 2018b ; 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 model-based cues robustly . Furthermore , while our VI executor provides a representation that aligns with the predictive needs of value iteration , it can also incorporate additional information if it benefits the performance of the model . 2 BACKGROUND . Value iteration Value iteration is a successive approximation method for finding the optimal value function of a discounted Markov decision processes ( MDPs ) as the fixed-point of the so-called Bellman optimality operator ( Puterman , 2014 ) . A discounted MDP is a tuple ( S , A , R , P , γ ) where s ∈ S are states , a ∈ A are actions , R : S × A → R is a reward function , P : S × A → Dist ( S ) is a 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 γ ∈ [ 0 , 1 ] is a discount factor which trades off between the relevance of immediate and future rewards . In the infinite horizon discounted setting , an agent sequentially chooses actions according to a stationary Markov policy π : S × A → [ 0 , 1 ] such that π ( a|s ) is a conditional probability distribution over actions given a state . The return is defined as Gt = ∑∞ k=0 γ kR ( at+k , st+k ) . Value functions V π ( s , a ) = Eπ [ Gt|st = s ] and Qπ ( s , a ) = Eπ [ Gt|st = s , at = a ] represent the expected return induced by a policy in an MDP when conditioned on a state or state-action pair respectively . In the infinite horizon discounted setting , we know that there exists an optimal stationary Markov policy π∗ such that for any policy π it holds that V π ∗ ( s ) ≥ V π ( s ) for all s ∈ S . Furthermore , such optimal policy can be deterministic – greedy – with respect to the optimal values . Therefore , in order to find a π∗ 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 ) : V ? ( s ) = max a∈A ( R ( s , a ) + γ ∑ s′∈S P ( s′|s , a ) V ? ( s′ ) ) . ( 1 ) 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 z : S → Rk and the effect of an action in a given state is modelled by a translation model T : Rk × A → Rk . Specifically , T ( z ( s ) , a ) is a translation vector to be added to z ( s ) in order to obtain an embedding of the resulting state when taking action a in state s.1 Such embeddings should also be far away from all other states z ( s̃ ) . Therefore , the embedding functions are optimized using the following variant of the triplet loss : L = d ( z ( s ) + T ( z ( s ) , a ) , z ( s′ ) ) + max ( 0 , ξ − d ( z ( s̃ ) , z ( s′ ) ) ) ( 2 ) where ξ is the hyperparameter of the hinge loss , ( s , a , s′ ) is a transition in a given MDP and s̃ is a negative sample state . Trajectories sampled from the MDP can be used as a training set of positive triplets ( s , a , s′ ) , whereas negative triplets ( s , a , s̃ ) are obtained by sampling s̃ uniformly at random . GNNs and algorithmic executors Graph neural networks ( GNNs ) have been intensively studied as a tool to process graph-structured inputs , and were successfully applied to various RL tasks ( Wang et al. , 2018 ; Klissarov & Precup , 2019 ; Adhikari et al. , 2020 ) . GNN operations are commonly expressed through the message passing framework ( Gilmer et al. , 2017 ) . For each node i in the graph , a set of messages is computed—one message for each neighbouring node j , derived by applying a message function M to the relevant node and edge features ( ~hi , ~hj , ~eij ) . Incoming messages in a neighbourhood N ( i ) are then aggregated through a permutation-invariant operator⊕ ( such as sum or max ) , obtaining a summarised message ~mi : ~mn+1i = ⊕ j∈N ( i ) Mn ( ~h n i , ~hnj , ~e n ij ) ( 3 ) The features of node i are then updated through a function U applied to these summarised messages : ~hn+1i = Un ( ~hni , ~m n+1 i ) ( 4 ) An important research direction explores the use of neural networks for learning to execute algorithms—which was recently extended to algorithms on graph-structured data ( Veličković et al. , 2019 ) . In particular , Xu et al . ( 2019 ) establishes algorithmic alignment between GNNs and dynamic programming algorithms . Furthermore , Veličković et al . ( 2019 ) show that supervising the GNN on the algorithm ’ s intermediate results is highly beneficial for out-of-distribution generalization . As VI is , in fact , a dynamic programming algorithm , a GNN executor is a suitable choice for learning it , and good results on executing VI have emerged on synthetic graphs ( Deac et al. , 2020 ) . 3 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 Figure 1 for a visualisation of the model dataflow ( more compact overview in Appendix A ) . We propose a policy network—a function , πθ ( a|s ) , which for a given state s ∈ S specifies a probability distribution of performing each of the actions a ∈ A in that state . Here , θ are the neural network parameters , to be optimised using gradient ascent . 1For stochastic MDPs , the resulting embedding should be understood as an expectation under P ( s′|s , a ) .

Summary: The 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.