M-Walk: Learning to Walk over Graphs using Monte Carlo Tree Search
Learning to walk over a graph towards a target node for a given query and a source node is an important problem in applications such as knowledge base completion (KBC). It can be formulated as a reinforcement learning (RL) problem with a known state transition model. To overcome the challenge of sparse rewards, we develop a graph-walking agent called M-Walk, which consists of a deep recurrent neural network (RNN) and Monte Carlo Tree Search (MCTS). The RNN encodes the state (i.e., history of the walked path) and maps it separately to a policy and Q-values. In order to effectively train the agent from sparse rewards, we combine MCTS with the neural policy to generate trajectories yielding more positive rewards. From these trajectories, the network is improved in an off-policy manner using Q-learning, which modifies the RNN policy via parameter sharing. Our proposed RL algorithm repeatedly applies this policy-improvement step to learn the model. At test time, MCTS is combined with the neural policy to predict the target node. Experimental results on several graph-walking benchmarks show that M-Walk is able to learn better policies than other RL-based methods, which are mainly based on policy gradients. M-Walk also outperforms traditional KBC baselines.
This work tackles the problem of knowledge base completion (KBC) by walking through a knowledge graph. The graph walk is treated as an MDP with a known state transition model. MCTS is used to plan a trajectory through the graph, and these trajectories are used to train a value function with Q-learning. The Q-function is used to form a policy prior for the MCTS. A KBC-specific RNN structure is also proposed. Results on a range of KBC benchmarks show small but consistent gains over several baselines, including on-policy RL approaches and non-MCTS ablations.  Quality MCTS is a strong planning algorithm for delayed-reward problems and seems a good fit for KBC, so this is an at the least an interesting application paper. The authors also extend previous work on combining MCTS with learned value functions and policy priors by leveraging the MCTS trajectories to learn the optimal value function of the KBC graph-traversal policy with Q-learning. This seems like a reasonable approach to improve the sample efficiency of learning the prior, compared to just updating the policy prior to match the MCTS visitation frequencies at the root node. The experiments seem to present competitive baselines and show some promise on a number of tasks.  However, I would appreciate some more discussion or clarity of the exact roles of pi, Q, and V. Firstly, the background section 2 does not make the import distinction between an arbitrary policy and the optimal policy (and their corresponding value functions). In particular, note that Q-learning does not estimate Q^pi but Q* for the MDP.  In M-Walk, Q is trained off-policy, and so learns about the optimal state-action value function Q* for graph-traversal. However, V is trained on-policy (wrt the MCTS tree policy, not pi_theta), but only on terminal MCTS states (not terminal MDP states, as I understand it). Could the authors elaborate on this choice (as opposed to for example using V = max_a (Q) )?  Meanwhile, pi is simply a soft transformation of Q into a probability distribution (the authors note that simply using a softmax works similarly to their final implementation). Traditionally, Q-learning with Boltzmann exploration as the behaviour policy would anneal its entropy (ie temperature) to converge to a good solution. I am concerned that a fixed temperature would restrict the ability of the algorithm to develop strong confidence in its prior policy, even if it were justified by the data. In this particular setting with strongly bounded rewards and sufficient MCTS samples to correct for even an uncertain prior, it seems to work fine, but the exact mechanism by which the Q-learning transfers to the MCTS tree policy deserves some further discussion.   Clarity The paper is quite clearly written on the whole. Some of the details of the architecture and intermediate state representation are a little hard to follow, and some design decisions regarding the neural architecture (s3.1) are a bit opaque. It should also be emphasised that the policy pi_theta is neither used to traverse the MDP (except indirectly as a prior) nor is its value learned by any function, which is a departure from the expectations set by section 2.  Originality The combination of MCTS with learned policy and value priors has been explored in a number of domains, but this application to KBC appears novel and includes a number of algorithmic and architectural innovations.  Significance This work addresses a significant and well-studied problem, as well as a significant class of solutions (MCTS with learned value/policy priors) with methods that may transfer to other problem settings. The experimental results are promising.  However, more care with the formulation and explanation of the approach wrt the exact roles of the learned value functions and consequent policy prior would be very useful for solidifying the theoretical foundations of the approach.