Combining Q-Learning and Search with Amortized Value Estimates

1 INTRODUCTION . Model-based methods have been at the heart of reinforcement learning ( RL ) since its inception ( Bellman , 1957 ) , and have recently seen a resurgence in the era of deep learning , with powerful function approximators inspiring a variety of effective new approaches ( Silver et al. , 2018 ; Chua et al. , 2018 ; Hamrick , 2019 ; Wang et al. , 2019 ) . Despite the success of model-free RL in reaching state-of-the-art performance in challenging domains ( e.g . Kapturowski et al. , 2018 ; Haarnoja et al. , 2018 ) , model-based methods hold the promise of allowing agents to more flexibly adapt to new situations and efficiently reason about what will happen to avoid potentially bad outcomes . The two key components of any such system are the model , which captures the dynamics of the world , and the planning algorithm , which chooses what computations to perform with the model in order to produce a decision or action ( Sutton & Barto , 2018 ) . Much recent work on model-based RL places an emphasis on model learning rather than planning , typically using generic off-the-shelf planners like Monte-Carlo rollouts or search ( see Hamrick ( 2019 ) ; Wang et al . ( 2019 ) for recent surveys ) . Yet , with most generic planners , even a perfect model of the world may require large amounts of computation to be effective in high-dimensional , sparse reward settings . For example , recent methods which use Monte-Carlo Tree Search ( MCTS ) require 100s or 1000s of model evaluations per action during training , and even upwards of a million simulations per time step at test time ( Anthony et al. , 2017 ; Silver et al. , 2018 ) . These large search budgets are required , in part , because much of the computation performed during planning—such as the estimation of action values—is coarsely summarized in behavioral traces such as visit counts ( Anthony et al. , 2017 ; Silver et al. , 2018 ) , or discarded entirely after an action is selected ( Bapst et al. , 2019 ; Azizzadenesheli et al. , 2018 ) . However , large search budgets are a luxury that is not always available : many real-world simulators are expensive and may only be feasible to query a handful of times . In this paper , we explore preserving the value estimates that were computed by search by amortizing them via a neural network and then using this network to guide future search , resulting in an approach which works well even with very small search budgets . We propose a new method called “ Search with Amortized Value Estimates ” ( SAVE ) which uses a combination of real experience as well as the results of past searches to improve overall performance and reduce planning cost . During training , SAVE uses MCTS to estimate the Q-values at encountered states . These Q-values are used along with real experience to fit a Q-function , thus amortizing the computation required to estimate values during search . The Q-function is then used as a prior for subsequent searches , resulting in a symbiotic relationship between model-free learning and MCTS . At test time , SAVE uses MCTS guided by the learned prior to produce effective behavior , even with very small search budgets and in environments with tens of thousands of possible actions per state—settings which are very challenging for traditional planners . 2 BACKGROUND AND MOTIVATION . Unifying the complementary approaches of learning and search has been of interest to the RL and planning communities for many years ( e.g . Gelly & Silver , 2007 ; Guo et al. , 2014 ; Gu et al. , 2016 ; Silver et al. , 2016 ) . SAVE is motivated in particular by two threads in this body of work : one which uses planning in-the-loop to produce experience for Q-learning , and one which learns a policy prior for guiding search . As we will describe next , both of these previous approaches can suffer from issues with training stability which are alleviated by SAVE by simultaneously using MCTS to strengthen an action-value function , and Q-learning to strengthen MCTS . 2.1 LEARNING FROM PLANNED ACTIONS . A number of methods have explored learning from planned actions . Guo et al . ( 2014 ) trained a model-free policy to imitate the actions produced by an MCTS agent . Other methods use planning in-the-loop to recommend actions , which are then executed in the environment to gather experience for model-free learning ( Silver et al. , 2008 ; Gu et al. , 2016 ; Azizzadenesheli et al. , 2018 ; Shen et al. , 2018 ; Lowrey et al. , 2018 ; Bapst et al. , 2019 ; Kartal et al. , 2019 ) . However , problems can arise when learning with actions that were produced via planning , even with off-policy algorithms like Q-learning . As noted by both Gu et al . ( 2016 ) and Azizzadenesheli et al . ( 2018 ) , planning avoids suboptimal actions , resulting in a highly biased action distribution consisting of mostly good actions ; information about suboptimal actions therefore does not get propagated back to the Q-function . As an example , consider the case where a Q-function recommends taking action a . During planning , this action is explored and is found to yield lower reward than expected . The planner will end up recommending some other action a′ , which is executed in the environment and later used to update the Q-function . However , this means that the original action a is never actually experienced and thus is never downweighed in the Q-function , resulting in poorly approximated Q-values . One way to deal with this problem is to use a mixture of both on-policy and planned actions ( Gu et al. , 2016 ) . However , this throws away information about poor actions which is acquired during the planning process . In SAVE , we instead make use of this information by using the values estimated during search to help fit the Q-function . If the search finds that a particular action is worse than previously thought , this information will be reflected by the estimated values and will thus ultimately get propagated back to the Q-function . We explicitly test and confirm this hypothesis in Section 4.2 . 2.2 USING PRIOR KNOWLEDGE IN SEARCH . Much research has leveraged prior knowledge in the context of MCTS ( Gelly & Silver , 2007 ; 2011 ; Silver et al. , 2016 ; Segler et al. , 2018 ; Silver et al. , 2017b ; 2018 ; Anthony et al. , 2017 ; 2019 ) . Some of the most successful methods ( Anthony et al. , 2017 ; Silver et al. , 2018 ) use a prior policy to guide search , the results of which are used to further improve the policy . However , such methods use information about past behavior to learn a policy prior—namely , the visit counts of actions during search—and discard other search information such as inferred Q-values . We might anticipate one potential failure mode of such “ count-based policy learning ” approaches . Consider an environment with sparse rewards , where most actions are highly suboptimal . In the limit of infinite search , actions which have highest value will be visited most frequently , resulting in a policy that guides search towards regions of high value . However , in the regime of small search budgets , the search may very well end up exploring mostly suboptimal actions . These actions have higher visit counts , and so are reinforced , leading to the agent being more likely to explore poor actions . Rather than implicitly biasing search towards value through the use of visit counts , SAVE relies on a prior that explicitly encodes knowledge about value . If SAVE ends up searching poor actions , it will learn that they have low values and this knowledge will be reflected in future searches . Thus , in contrast to count-based approaches , a SAVE agent will be less likely to visit poor actions in the future despite having frequently visited them in the past . We explicitly test and confirm this hypothesis in Section 4.1 . 2.3 OTHER RELATED WORK . Finding effective ways of combining model-based and model-free experience has been of interest to the RL community for decades . Most famously , the Dyna algorithm ( Sutton , 1990 ) proposes using real experience to learn a model and then using the model to train a model-free policy . A number of more recent works have explored how to incorporate this idea into deep architectures ( Kalweit & Boedecker , 2017 ; Feinberg et al. , 2018 ; Buckman et al. , 2018 ; Serban et al. , 2018 ; Kurutach et al. , 2018 ; Kaiser et al. , 2019 ) , with an emphasis on dealing with the errors that are introduced by approximate models . In these approaches , the policy or value function is typically trained using on-policy rollouts from the model without using additional planning . Another way to combine model-free and model-based approaches is “ implicit planning ” , in which the computation of a planner is built into the architecture of a neural network itself ( Weber et al. , 2017 ; Buesing et al. , 2018 ; Pascanu et al. , 2017 ; Silver et al. , 2017b ; Oh et al. , 2017 ; Guez et al. , 2018 ; Farquhar et al. , 2018 ; Hamrick et al. , 2017 ; Srinivas et al. , 2018 ; Yu et al. , 2019 ; Tamar et al. , 2016 ; Karkus et al. , 2017 ) . While SAVE is not an implicit planning method , it shares similarities with such methods in that it also tightly integrates planning and learning . 3 METHOD SAVE features two main components ( Figure 1 ) . First , we use a search policy that incorporates the Q-function Qθ ( s , a ) as a prior over Q-values that are estimated during search . Second , to train the Qfunction we rely on an objective function that combines both the TD-error from Q-learning with an amortization loss that amortizes the value computation performed by the search . The amortization loss , combined with the prior over Q-values , thus enables future searches to build on previous ones , resulting in stronger search performance overall . 3.1 STANDARD MCTS . Before explaining how SAVE leverages search , we briefly describe the standard MCTS algorithm ( Kocsis & Szepesvári , 2006 ; Coulom , 2006 ) . While we focus here on the single-player setting , we note that the formulation of MCTS ( and by extension , SAVE ) is similar for two-player settings . MCTS uses a simulator or model of the environment to explore possible future states and actions , with the aim of finding a good action to execute from the current state , s0 . In MCTS , we assume access to a budget of K iterations ( or simulations ) . The kth iteration of MCTS consists of three phases : selection , expan- sion , and backup . In the selection phase , we expand a search tree beginning with the current state and taking actions according to a search policy : πk ( s ) = argmax a ( Qk ( s , a ) + Uk ( s , a ) ) , ( 1 ) whereQk is the currently estimated value of taking action awhile in state s , which will be explained further below . Uk ( s , a ) is the UCT exploration term : Uk ( s , a ) = cUCT √ log ( ∑ aNk ( s , a ) ) Nk ( s , a ) , ( 2 ) where Nk ( s , a ) is the number of times we have explored taking action a from state s and cUCT is a constant that encourages exploration . This selection procedure is repeated for T − 1 times , until a new action aT−1 that had not previously been explored is chosen from state sT−1 . This begins the expansion phase , during which aT−1 is executed in the simulator , resulting in a reward rT−1 and new state sT . The new state sT is added to the search tree , and its value V ( sT ) is estimated either via a state-value function or ( more traditionally ) via a Monte-Carlo rollout . At this point the backup phase begins , during which the value of sT is used to update ( or “ back up ” ) the values of its parent states earlier in the tree . Specifically , for state st , the ith backed up return is estimated as : Ri ( st , at ) = γ T−tV ( sT ) + T−1∑ j=t γj−trj , ( 3 ) where γ is the discount factor and rj was the reward obtained after executing aj in sj when traversing the search tree . These backups are then used to estimate the Q-function in Equation 1 as Qk ( s , a ) =∑Nk ( s , a ) i=1 Ri ( s , a ) /Nk ( s , a ) .

Summary: This paper proposes Search with Amortized Value Estimates (SAVE), which combines Q-learning and Monte-Carlo Tree Search (MCTS). SAVE makes use of the estimated Q-values obtained by MCTS at the root node (Q_MCTS), rather than using only the resulting action or counts to learn a policy. It trains the amortized value network Q_theta via the linear combination of Q-learning loss and the cross-entropy loss between the softmax(Q_MCTS) and softmax(Q_theta). Then, SAVE incorporates the learned Q-function into MCTS by using it for the initial estimate for Q at each node and for the leaf node evaluation by V(s) = max_a Q_theta(s,a). Experimental results show that SAVE outperforms the baseline algorithms when the search budget is limited.