Combining Q-Learning and Search with Amortized Value Estimates
Authors: Jessica B. Hamrick
Victor Bapst
Alvaro Sanchez-Gonzalez
Tobias Pfaff
Theophane Weber
Lars Buesing
Peter W. Battaglia
Venue:  (2020)
Keywords: model-based RL
Q-learning
MCTS
search

Abstract: We introduce "Search with Amortized Value Estimates" (SAVE), an approach for combining model-free Q-learning with model-based Monte-Carlo Tree Search (MCTS). In SAVE, a learned prior over state-action values is used to guide MCTS, which estimates an improved set of state-action values. The new Q-estimates are then used in combination with real experience to update the prior. This effectively amortizes the value computation performed by MCTS, resulting in a cooperative relationship between model-free learning and model-based search. SAVE can be implemented on top of any Q-learning agent with access to a model, which we demonstrate by incorporating it into agents that perform challenging physical reasoning tasks and Atari. SAVE consistently achieves higher rewards with fewer training steps, and---in contrast to typical model-based search approaches---yields strong performance with very small search budgets. By combining real experience with information computed during search, SAVE demonstrates that it is possible to improve on both the performance of model-free learning and the computational cost of planning.

Published as a conference paper at ICLR 2020
COMBINING Q-LEARNING AND SEARCH WITH AMORTIZED VALUE ESTIMATES
Jessica B. Hamrick
Victor Bapst
Alvaro Sanchez-Gonzalez DeepMind
DeepMind chamrick@googl com
DeepMind @bapsttgoog com
vlavsg@googge....
Tobias Pfaff
Théophane Weber
Lars Buesing
DeepMind tpfaffegoogle.com
DeepMind
DeepMind
cheophane@google.comm
Lbuesing@google.com
Peter W. Battaglia
DeepMind
peterbattagglia@goolleccoo
ABSTRACT
We introduce "Search with Amortized Value Estimates" (SAVE), an approach for combining model-free Q-learning with model-based Monte-Carlo Tree Search (MCTS). In SAVE, Iearned prior over state-action values is used guide MCTS, which estimates an improved set state-action values. The new Q-estimates are then used in combination with real experience to update the prior. This effectively amortizes the value computation performed by MCTS, resulting in coopera- tive relationship between model-free learning and model-based search. SAVE can be implemented on top of any Q-learning agent with access to a model, which we demonstrate by incorporating into agents that perform challenging physical reasoning tasks and Atari. SAVE consistently achieves higher rewards with fewer training steps, and-in contrast to typical model-based search approaches-yieee strong performance with very small search budgets. By combining real experience with information computed during search, SAVE demonstrates that it is possible to improve on both the performance of model-free learning and the computational cost planning.
INTRODUCTION
Model-based methods have been 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 variety of effective
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-ar performance challenging domains (e.g. Kapturowski al., 2018 Haarnoja 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 decision or action (Sutton & Barto 2018).
Much recent work on model-based RL places an emphasis on model learning rather than plan- ning, typically using generic off-the-shelf planners like Monte-Carlo rollouts or search (see [Ham- rick (2019); Wang et al. (2019) 19) 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 Carto Tree Search (MCTS) require 100s 1000s of model evaluations per action during training, and even upwards of million simulations per time step at test time (Anthony et al. 2017] (Silver et al. 2018). These large search budgets required, in part, because much of the computation performed during planning-sud

---

Published as a conference paper at ICLR 2020
as the estimation of action values- is coarsely summarized in behavioral traces such 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 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 encoun- tered 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 prior for subsequent searches, resulting 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 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 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-Jearning to strengthen MCTS.
2.1 LEARNING FROM PLANNED ACTIONS
A number of methods have explored learning from planned actions. Guo al. (0014) trained 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 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 al.) 2019 Kartal al. 2019). However, problems can arise when learning with actions that were produced via planning, even with off- off-policy algorithms like Q-learning. As noted by both Gu al. (2016) and AAzizzaeenesheli 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 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 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. 20176] 2018: Anthony et 2017 2019). Some of the most successful methods (Anthony et al. [017 Silver et al. 2018) use prior policy to guide search, the results of which are used to further improve the policy. However, such methods use information about past behavior learn policy prior-namely, the visit counts of actions during
2

---

Published as a conference paper at ICLR 2020
search-ane discard other search information such as inferred Q-values. We might anticipate one potential failure mode of such "count-base" 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, will learn that they have low values and this knowledge will be reflected in future searches. Thus, in contrast to count-based approaches, 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 model-free policy. A number of more recent works have explored how to incorporate this idea into deep architectures (Kalweit & Boedecker [017]] [Finnberg 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- approaches is "implicit planning". in which the computation of a planner is built into the architecture of a neural network itself (Weber et al. [017]] [uesing et al.,
2018
Pascanu et al., 2017 Silver al., 201 7b Oh et al..
2017 Guez al., [018]] Farquhar al.
2018
Hamrick et al. 2017 Srinivas et al. 2018 Yu et al. 2019] Taaar et al., [2016] 6, Karkus et al.,
2017
While SAVE not an implicit planning method, it shares similarities with such methods in
that it also tightly integrates planning and learning.
METHOD
Acting
Learning
Qa
QK
sonrch experiuue {s, QUCCTS)
SAVE features two main components (Figure I). First, use search policy that incorporates the Q-function Qo a) as a prior over Q-values that estimated during search. Second, train the Q- function we rely on an objective function that com- bines both the TD-error from Q-learning with an amortization loss that amortizes the value computa- tion 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.
QMeTS
Qo
Qo
MCTS
Lg
zeal experienoe {8,5',ar))
W
Figure 1: Illustration of SAVE. When acting, the agent uses ffunnctio,, Qe. prior for the Q-values estimated during MCTS. Over K steps of search, Qo Qo is built up to Qk. which is returned as QMCTS (Equa- tions and 4 From QMCTS. an action is selected via epsilon-greedy and the resulting experience (s, r, s' QMers) is added to a replay buffer. When learning, the agent uses real experience to update Qa via Q-learning (Lq) as well as an amortization loss (C.) which regresses Qe towards the Q-values es- timated during search (Equation 6).
3.1 STANDARD MCTS
Before explaining how SAVE leverages search, we briefly describe the standard MCTS algorithm (Koc-) sis & Szepesvári 2006 Coulom 2006). While we Tocus here on the single-player setting, we note that the formulation MCTS (and by extension, SAVE) is similar for two-player settings. MCTS uses sim- ulator or model of the environment to explore possi- ble future states and actions, with the aim of finding a good action to execute from the current state, $p- In MCTS, we assume access to a budget of K iterations (or simulations). The tth iteration of MCTS consists of three phases: selection, expan-
3

---

Published as a conference paper at ICLR 2020
sion, and backup. In the selection phase, we expand search tree beginning with the current state and taking actions according to a search policy:
k (8)
arg max (s
(1)
where Qu is the currently estimated value of taking action a while in state 8, which will be explained further below. (s, a) is the UCT exploration term:
(v(s, a) CUCTTV log (E. Na(
(2)
Na(s,a)
where (s, a) is the number of times we have explored taking action from state and CUCT is constant that encourages exploration. This selection procedure is repeated for new action ar-1 that had not previously been explored is chosen from state $r-1- This begins the expansion phase, during which ar-1 is executed in the simulator, resulting in a reward ry-1 and new state sT. The new state sT is added the search tree, and its value V(8T) 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 9t, the zth backed up return is estimated as:
times, until a
T-1
Ra(se, = TV(((((s)) + Σ
rd Tg.
(3)
where y is the discount factor and was the reward obtained after executing ay in 8, when traversing the search tree. These backups are then used to estimate the Q-function Equation I as Qx(s,a)
EN (a..a) Rs(s, aN (s, a).
21=1
3.2 INCORPORATING A PRIOR DURING SEARCI
SAVE makes several changes to the standard MCTS procedure. First, it assumes it has visited every state and action pair once by initializing N(s, a) 1 for all states and actions Second, for each of these state-action pairs, it assumes a prior estimate of its value, Qa(s, a), and uses this as an initial estimate for Qk. similar to Gelly & Silver (2077] (01T):
Qs(5, Q0(s, a) + Σ Ne(sa)--1 Rx(s,a)
(4)
Ng(s,a)
where Qa(s,a) (s, a). Third, rather than using a separate state-value function or Monte-Carlo rollouts estimate the value of new states, SAVE uses the same state-action value function, i.e. V(s) max, Qa(s, a). These three changes provide a mechanism for incorporating Q-based prior knowledge into MCTS: specifically, SAVE acts if has visited every state-action pair once, with the estimated values being given by Qo. Roughly speaking, this can be interpreted as using MCTS to perform Bayesian inference over Q-values, with the prior specified by Qo with a weight equivalent to pseudocount of one. This set of changes contrasts with UCT, which does not incorporate prior knowledge, as well as PUCT (Rosin) 2011 Silver et al. 2017a 2018). which incorporates prior knowledge via policy in the exploration term U)(s,a).
After iterations, we return QMCTS(s, a) == Q (s, a) and select an action to execute in the envi- ronment via epsilon-greedy over QMCTS( a). After the action is executed, store the resulting experience along with a copy of QMCTSS process is illustrated in Figure I (left).
{QMCTs(*) a;)}; in the replay buffer. This
3.3 Q-LEARNING WITH AN AMORTIZATION LOSS
During learning, the results of the search are amortized into an updated prior Qa Figure I right). We impose an amortization loss La which encourages the distribution of Q-values output by the neural network to be similar to those estimated by MCTS. The amortization loss is defined to be
We could also consider initializing N(s. a) hased on an estimate of previous visit counts, which we leave an interesting direction for future work.
4

---

Published as a conference paper at ICLR 2020
(a] Budget =
(B) Budget=
(4) Budget =
(6) Function Approximation
1.0
1.0
SAVE
G-Learing
0.5
0.5
10.6
1.0-
1.0
1.0
0.6
0.44
0.2
0.0
85
90
05
50
86
90
65
90
85
90
96
89
0.0
% Terminal
Temmnal
N: Terminal
0.25
0.50
0.75
1.00
PUCT
Loaming
Egtsodos
Figure 2: Results on Tightrope. (a-c) Tabular results comparing SAVE, PUCT, UCT, and Q-learning (with MCTS at test time) for varying percentages of terminal actions on the z-axes and for different search budgets. The y-axes show reward for either the sparse or dense reward setting of Tightrope. Lines show medians across 20 seeds, with error bars showing 95% confidence intervals. (d) Results on Tightrope when using function approximation. comparing SAVE with PUCT and Q-learning. Lines show medians across 10 seeds, with shaded regions indicating min and max seeds.
the cross-entropy between the softmax of the Q-values before (Qp) and after (QMets) MCTS. This cross-entropy loss achieves better performance than alternatives like L22, described in [Section 4.2 Setting PMCTS = softmax, (MMCTS(s 1) and po = softmax, (Qo(s, (, 1)), where 7 1 is the softmax temperature, the loss is defined as:
CA D) = (PMOTT)) log Pe,
(5)
NNNNN D
where D is batch of N experience tuples (st: at, St+1. QMCTS .)) sampled from the replay buffer. This amortization loss is linearly combined with Q-learning loss,
[(, D) = $0, D) + AAC
(6)
where 3g and BA are coefficients to scale the loss terms. Lo may be any basue- loss function, such that based on 1-step TD targets, n-step TD targets, or X-returns (Sutton) 1988). The amorti- zation loss does make SAVE more sensitive to off-policy experience, as the values of QMCTS stored in the replay buffer will become less useful and potentially misleading Qo improves; however, we did not find this to be an issue in practice.
4 EXPERIMENTS
We evaluated SAVE in four distinct settings that vary in their branching factor, sparsity of rewards, and episode length. First, we demonstrate through new Tightrope environment that SAVE per- forms well in settings where count-based policy approaches struggle, as discussed in Section 2.2 Next, we show that SAVE scales to the challenging Construction domain (Bapst et al. 2019 and that it alleviates the problem with off-policy actions discussed in infeection" 2.1 We also perform sev- eral ablations to tease apart the details of SAVE. Finally, we demonstrate that SAVE dramatically improves over Q-learning in new and even more difficult construction task called Marble Run, as well as in more standard environments like Atari (Bellemare et al., (0013). In all our experiments we use SAVE with a perfect model of the environment, though expect approach would work with learned models as well.
4.1 TIGHTRRPP
In Section 2.2, we hypothesized that approaches which use count-based policy learning rather than value-based learning (e.g. Anthony al.) 2017: Silver et al. 2018) may suffer in environments with large branching factors, many suboptimal actions, and small search budgets. To test this hypothesis, developed toy environment called Tightrope with these characteristics. Tightrope is deter- ministic MDP consisting of 11 labeled states linked together in a chain. At each state, there are 100 actions take, M% of which are terminal (meaning that when taken they cause the episode to end).
5

---

Published as a conference paper at ICLR 2020
The other non-terminal actions will cause the state to transition to the next state in the chain. We considered two settings of the reward function: dense rewards, in which case the agent receives a reward of 0.1 when making it to the next state in the chain and 0 otherwise; and sparse rewards, in which case the agent receives a reward of only when making it to the final state. In the sparse re- ward setting, we randomly selected one state in the chain he the "final" state to form curriculum over the length of the chain. With the exception of the final state in the sparse reward setting, the transition functiin of the MDP is exactly the same across episodes, with the same actions always having the same behavior.
Tabular Results We first examined the behavior of SAVE on Tightrope in tabular setting to eliminate potential concerns about function approximation (see SScction B.2). We compared SAVE three other agents. UCT a pure-search agent which runs MCTS using UCT search policy with no prior. It uses Monte-Carlo Carlo rollouts following a random policy to estimate Vss). PUCT is based on AlphaZero (Silver et al., 2018) and uses a policy prior (which is learned from visit counts during MCTS) and -tauu function (which is learned from Monte-Carte returns). During search, the policy is used in the PUCT exploration term and the value function is used for bootstrapping. More details on PUCT in general are provided in [Section A.3 Q-Learning performs one-step tabular Q-lcarning during training, and MCTS at test time using the same search procedure as SAVE.
Figure 2 illustrates the results in the tabular setting after 500 episodes. UCT, which does not use any learning. illustrates the difficulty of using brute-force search. Q-learning, which does not use any search during training, is slow to converge in the sparse reward setting; additionally, adding search at test time does not substantially improve things. Although the incorporation of learning with PUCT does improve the results, we can see that with small search budgets and high proportions of terminal actions, PUCT struggles to remember which actions are safe (nonterminal), especially in the sparse reward setting. In contrast, SAVE solves the Tightrope environment in all of the dense reward settings and most of the sparse reward settings. As the search budget increases, we see that both PUCT and SAVE reliably converge to solution; thus, if a large search budget is available both methods may fare equally well. However, if only a small search budget is available, SAVE results in much more reliable performance.
solution within the 500 episodes, particularly
Function Approximation Results We also looked at the ability of SAVE and PUCT to solve the Tightrope environment when using function approximation, along with a model-free Q-learning baseline (see Section B.3). We evaluated all agents on the sparse reward version of Tightrope with 95% terminal actions, and used a search budget of 10 (except for Q-learning, which used a test budget of zero). The results, shown in Figure 21, follow the same pattern in the tabular setting.
4.2 CONSTRUCTION
We next evaluated SAVE in three of the Construction tasks explored by Bapst al. (2019). in which the goal is to stack blocks to achieve a functional objective while avoiding collisions with obstacles. In Connecting, the goal is connect target point in the sky the floor. In Covering, the goal is to cover obstacles from above without touching them. Covering Hard is the same as Covering, except that only a limited number of blocks may be used. The Construction tasks are challenging for model- free approaches because there is a combinatorial space of possible scenes and the physical dynamics are challenging to predict. However, they are also difficult for traditional search methods, as they have huge branching factors with up to tens of thousands of possible actions per state. Additionally, the simulator in the Construction tasks is expensive to query. making infeasible use with search budgets of more than 10-20.
To implement SAVE, we used the same agent architecture as Bapst et al. 2019 We compared SAVE to a baseline version of SAVE without amortization loss (i.e., L( D) BoLo(( (6, D)), similar to the MCTS agent described in Bapst et al. 2019). We also compared to a Q-learning baseline which performs pure model-free learning during training (but which may also utilize MCTS at test time using the same search procedure as SAVE), as well as a UCT baseline which did not any learning (but which did use pretrained value function for bootstrapping). For SaVe-based agents, we used a training budget of 10 simulations and varied the budget at test time; for UCT, we used a constant budget of 1000 simulations at test time (see Appendix C).

---

Published as a conference paper at ICLR 2020
Connecting
(b) Covering
Covering Hard
(8) Ablations
3.0
2.5
4-
2.0
Rowardd
1.5-
3-
1.0
2-
2-
G-Lamming
0.5
SAVE
0.0
0-
SAVZ
SANE SAVE win DL
SOVE WA AL
Test Budget
Test Budgee
Test Budges
SAVE PUCT
Figure 3: Results on Construction. (a-c) Each subplot shows results for SAVE, SAVE without amortization loss, Q-learning with MCTS at test time, and pure search (UCT). The r-axis shows the effect of increasing the number of MCTS simulations at test time. During training. SAVE with and without amortization loss used a search budget of 10 simulations. UCT used a search budget of 1000 simulations. Points show medians across 10 seeds, and emror bars indicate min and max seeds. (d) Ablation experiments on the Covering task. We compare SAVE to variants that do not have amortization loss, which use L2 amortization loss, which do not use the Q-Learning loss, and which use PUCT rather than UCT. Results are shown at the hardest level of difficulty for the Covering task with test budget of 10. The colored bars show median reward across 10 seeds, and error bars show min and max seed.
Results Figure 3 shows the results on the three construction tasks. The poor performance of UCT (dotted lines) highlights the need for prior knowledge to manage the huge branching factor in these domains. While model-free Q-learning improves performance, simply performing search on top of the learned Q-values only results in small gains in performance, if any. The performance of SAVE without amortization loss highlights exactly the issue discussed in [Section 2.1 Without the amortization loss, the Q-learning component of SAVE only learns about actions which have been selected via search, and thus rarely sees highly suboptimal actions, resulting in a poorly approxi- mated Q-function. Indeed, as performance falls off dramatically, suggesting that the underlying Q-values are poor. Using search test time can make up for this problem to some degree, but only when used with budget very close to that with which it was trained: large search budgets can actually result in worse search performance (e.g. in Covering and Covering Hard) because the poor Q-values are also being used for bootstrapping during the search. It is only by leveraging search during training time and incor- porating an amortization loss do we see a synergistic result: using SAVE results in higher rewards across all tasks, strongly outperforming the other agents.
in the case where the search budget is zero, the agent's
Ablation Experiments In the past two sections, we compared SAVE to alternatives which do not include amortization loss, which use count-based policy learning rather than value-based learning. However, a number of additional questions remain regarding the architectural choices in SAVE. To address these, results shown in infFiuure 31.. Specifically, we compared SAVE with versions that use an L2 loss (rather than cross entropy), that do not use the Q-learning loss, and that use the Q-values to guide search via PUCT rather than initializing Qo- Overall, we find that the choices made in SAVE result in the highest levels of performance. Of particular note the ablation that uses the L2 loss, indicating that the softmax cross entropy loss plays an important role in SAVE's performance. We speculate this is true for two reasons. First, because we use small search budgets, the estimated QMCTS is likely to be noisy, and thus it may be more robust to preserve just the relative magnitudes of action values rather than exact quantities. Second, the cross entropy loss means that Qo need not represent the values of poor actions exactly, thus freeing up capacity in the neural network to more precisely represent the values of good actions. Details and further discussion is provided in Section C.3 We also compared
number of ablation experiments
the Covering task, with the
policy-based PUCT agent like that described in Section 4.1, but found this did not achieve positive reward on the harder tasks like Covering. This result again highlights the same problem with count-based policy training and small search budgets, discussed in Section 2.2

---

Published as a conference paper at ICLR 2020
(a) Sticky
Sticky block cost 0.04
Sticky block cost
Sticky black cost
1.0-
IN
0.8
Reward
0.4-
0.2
SAVE Q-Learning
0.0
0
Difficllty
Difficulty
Figure 4: (a-b) Results on the Marble Run environment for model-free Q-Learning as well as SAVE as a function of curriculum difficulty level, for two different settings of the cost of "sticky" blocks. Points indicate medians across 10 seeds, and error bars show min and max seeds. (c-d) Structures built by SAVE which solve the same scene for two different costs of sticky blocks (difficulty 6). Ad- ditional videos showing agent behavior are available at ftthtts::/ttt /ttns:///inyuull com/ com/yxm4ma* 4xm4ma47
4.3 MARBLE RUN
SAVE is able to achieve near-ceiling levels of performance on the original Construction tasks. Thus, we developed a new task in the style of the previous Construction tasks called Marble Run which is even more challenging in that involves sparser rewards and a more complex reward function. Specifically. the goal in Marble Run is to stack blocks to enable marble to get from its original starting position to a goal location, while avoiding obstacles. At each step, the agent may choose from number of differently shaped rectangular blocks well ramp shapes, and may choose to make these blocks "sticky" (for price) so that they stick to other objects in the scene. The episode ends once the agent has created structure that would get the marble the goal. The agent receives a reward of one if it solves the scene, and zero otherwise.
We used the same agent architecture and training setup as with the Construction tasks, except for the curriculum. Specifically, we found it was important to train agents on this task using an adaptive curriculum over difficulty levels rather than a fixed linear curriculum. Under the adaptive curriculum, we only allowed an agent to progress to the next level of difficulty after it was able to solve at least 50% of the scenes at the current level of difficulty. Further details of the Marble Run task and the curriculum are given in Appendix D
Results Figure 4] shows the results for SAVE and Q-learning for the two different costs of sticky blocks, well as some example constructions. SAVE progresses more quickly through the curriculum and reaches higher levels of difficulty (see[Figure D.Ib and overall achieves much higher levels of reward at every difficulty level. Additionally, we found that the Q-learning agent reliably becomes unstable and collapses at around difficulty 4-5 (see Figure D.2), while SAVE does not have this problem. Qualitatively (Figure 4b-d), SAVE able to build structures which allow the marble to reach targets that are raised above the floor while also spanning multiple obstacles.
These results on Marble Run also allow us to address the trade-off between model-free experience versus planned experience. Specifically, with a search budget of 10, SAVE effectively sees 10 times as many transitions as model-free agent trained on the same number environment interactions. Would model-free agent trained for 10 times as long achieve equivalent performance? As can be seen in Figure D.2, this is not the case: the model-free agent sees more episodes but results in worse performance. We find the same result in other Construction tasks as well (see [Section C.4) This highlights the positive interaction that occurs when learning both from experience generated from planned actions and from the values estimated during search.
4.4 ATARI
To demonstrate that SAVE is applicable to more standard environments, also evaluated it subset of Atari games (Bellemare et al. 2013).. We implemented SAVE on top of R2D2, a distributed Q-learning agent that achieves ofte-offthheart results Atari (Kapturowski et al. 2018). To allow
8

---

Published as a conference paper at ICLR 2020
for a fair comparisor between purely model-free R2D2 and version with SAVE, controlled R2D2 to have the same replay ratio as SAVE and then tuned its hyperparameters to have approxi- mately the same level of performance the baseline version of R2D2 (see Appendix E). We find that SAVE outperforms or equals this controlled version of R2D2 in all games, with particularly high performance Frosthite, Alien, and Zaxxon (shown in Figure 5). SAVE also outperforms the baseline version of R2D2 (see Table E. and Figure E.I
5 DISCUSSION
change relotive to R2D2
1600%
We introduced SAVE, a method for combining model-free Q-Jearning with MCTS. During training. SAVE leverages MCTS to infer set of Q-values, and then uses combina- tion of real experience plus the estimated Q-values to fit Q- function, thus amorlizing the value computation of previous searches via a neural network. The Q-function is used as prior guide future searches, enabling even stronger search performance, which in turn is further amortized via the Q- function. test time, SAVE can be used to achieve high levels of reward with only very small search budgets, which we demonstrate across four distinct domains: Tightrope, Construction (Bapst et al. 2019). Marble Run, and Atari (Bellemare et al., 2013] Kapturowski et al. 2018). These results suggest that SAVEing the experience generated by search in an explicit Q-function, and initializing future searches with that information, offers impor- tant advantages for model-based RL.
200%
100%
0%
Froobrite
Allen
Zaaron
Ueam Rider
Ms. Pacman
Invaders
Nome This Came
Orvlipade
River Reid
Grayy Climber
Asteroids
Up tn Down
Gravitar
Hero
Figure 5: Results on Atari.
When combining Q-values estimated both from prior searches and real experience, it may also be useful to account for the quality confidence of the estimated Q-values. Count-based policy methods (Anthony et al., 2017, Silver et al.) 2018) do this by leveraging an estimate of confidence based on visit counts: actions with high visit counts should both have high value (or else they would not have been visited so much) and high confidence (because they have been explored extensively). However, as we have shown. relying solely on visit counts can result poor performance when using small search budgets (Section 4.J). A key future direction will be to amortize both the computation of value and of reliability, achieving the best of both SAVE and count-based methods. Encoding confidence estimates into the Q-values may also be helpful for applying SAVE to settings with learned models, which may have non-trivial approximation errors. particular, it may be helpful to attenuate the contribution of search-estimate Q-values to the Q-prior both when an action has not been sufficiently explored and when model error is high.
Our work demonstrates the value of amortizing the Q-estimates that are generated during MCTS. Indeed, we have shown that by doing so, SAVE reaches higher levels of performance than model- free approaches while using less computation than is required by other model-based methods. More broadly, we suggest that SAVE can be interpreted as a framework for ensuring that the valuable computation performed during search is preserved, rather than being used only for the immediate action or summarized indirectly via frequency statistics of the search policy. By following this philosophy and tightly integrating planning and learning, expect that even more powerful hybrid approaches can be achieved.
6 ACKNOWLEDGEMENTS
We would like to thank GB Parascandolo. George Papamakarios, Nicolas Heess, Ioannis Antonoglou, Thomas Hubert, Julian Schrittweiser, and David Silver for helpful comments and feed- back on this project.
R222 is very sensitive to the speed of the learners and actors. If the actors are slower (which they will be when performing search), the replay ratio will increase which thus affects performance.

---

Published as a conference paper at ICLR 2020
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: system for large- scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Imple- mentation (OSDI 16), pp. 265-283, 2016.
Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow with deep learning and tree search. In Advances in Neural Information Processing Systems, pp. 3360-5370, 2017.
Thomas Anthony, Robert Nishihara, Philipp Moritz, Tim Salimans, and John Schulman. Pol- icy gradient search: Online planning and expert iteration without search trees. arXiv preprint 1rXiv:1904003644 2019.
Kamyar Azizzadenesheli, Brandon Yang, Weitang Liu, Emma Brunskilland Zachary C Lipton, and Animashree Anandkumar. Surprising negative results for generative adversarial tree search. arXiv preprint wrXiv:/806.0577000 pp. 1-25, 2018.
Victor Bapst, Alvaro Sanchez-Gonzaler Carl Doersch, Kimberly L Stachenfeld, Pushmeet Kohli, Peter W Battaglia, and Jessica B Hamrick. Structured agents for physical construction. In Pro- ceedings of the International Conference on Machine Learning (ICML), 2019.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzan Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 1rXXvv:806601111 1806.0/26,, pp. 1-38, 2018.
Marc 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. Princeton University Press, 1957.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lec. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Advances Neural Information Processing Systems, pp. 8224 8224-82334 2018.
Lars Buesing, Théophane Weber, Sébastien Racanière, S. M. Ali Eslami, Danilo Rezende, David P. Reichert, Fabio Viola, Frédéric Besse, Karol Gregor, Demis Hassabis, and Daan Wierstra. Learning and querying fast generative models for reinforcement learning. arXiv preprint Ntf:::8802.00006 pp. 1-15.2015
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning handful of trials using probabilistic dynamics models. In Proceedings of the 32nd Conference on Neural Information Processing Systems (NeurIPS 2018), 2018.
Rémi Coulom. Efficient selectivity and backup operators in Monte-Carlo tree search. Interna- tional conference on computers and games, pp. 72-83. Springer, 2006.
Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. TreeQN and ATreeC: Differentiable tree planning for deep reinforcement learning. In Proceedings of the 6th Interna- fional Conference on Learning Representations (ICLR 2018), 2018.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I. Jordan, Joseph E. Gonzalez, and Sergey Levine. Model-based value expansion for efficient model-free reinforcement learning. In Pro- ceedings of the 35th International Conference on Machine Learning (ICML 2018), 2018.
Sylvain Gelly and David Silver. Combining online and offline knowledge in UCT. In Proceedings of the 24th international conference Machine learning, pp. 273-280. ACM, 2007.
Sylvain Gelly and David Silver. Monte-carlo tree search and rapid action value estimation com- puter go. Artificial Intelligence, 175(11): 1856-1855. 2011.
10

---

Published as a conference paper at ICLR 2020
Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep Q-learning with model-based acceleration. In Proceedings of the 33rd International Conference on Machine Learning (ICML 2016), 2016.
Arthur Guez, Théophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wier- stra, Rémi Munos, and David Silver. Learning to search with MCTSnets. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), 2018.
Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Théophane Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, et al. An investigation of model-free planning. arXiv preprint wwt.i::/011.335500 2019.
Xiaoxiao Guo, Satinder Singh, Honglak Lee, Richard I. Lewis, and Xiaoshi Wang. Deep learning for real-time Atari game play using offlin Monte-Carlo tree search planning. In Advances in neural information processing systems, pp. 1338-3346, 2014.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with stochastic actor. arXiv preprint arXiv:1801.01290 2018.
Jessica Hamrick. Analogues of mental simulation and imagination in deep learning. Current Opinion in Behavioral Sciences, 29:8-16, 11,8-1,220 2019.
Jessica B. Hamrick, Andrew J. Ballard, Razvan Pascanu, Oriol Vinyals, Nicolas Heess, and Peter W. Battaglia. Metacontrol for adaptive imagination-based optimization. In Proceedings of the 5th International Conference Learning Representations (ICLR 2017), 2017.
Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for Atari. arXiv preprint arXiv: 1903.004433. 2019.
Gabriel Kalweit and Joschka Boedecker. Uncertainty- driven imagination for continuous deep re- inforcement learning. In Proceedings of the 1st Conference on Robot Learning (CoRL. 2017), 2017.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent expe- rience replay in distributed reinforcement learning. In Proceedings of the International Confer- ence on Learning Representations, 2018.
Peter Karkus, David Hsu, and Wee Sun Lee. QMDP-Net: Deep learning for planning under partial observability. In Proceedings of the 31st Conference Neural Information Processing Systems (NeurIPS 2017), 2017.
Bilal Kartal, Pablo Hernandez-Leal, and Matthew E Taylor. Action guidance with mets for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment, volume 15, pp. 153-159, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 6980, 2014.
Levente Kocsis and Csaba Szepesvári. Bandit based Monte-Carlo planning. In European conference on machine learning, pp. 282-292 293. Springer, 2006.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. Proceedings of the 6th International Conference on Learning Representations (ICLR 2018), 2018.
Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offline: Efficient learning and exploration via model-based control. arXiv preprint arXiv:1811.01844 2018.
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 reinforcement learning. Nature, 518(7540):529, 2015.
11

---

Published as a conference paper at ICLR 2020
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In Proceedings of the 31st Conference on Neural Information Processing Systems (NeurIPS 2017), 2017.
Razvan Pascanu, Yujia Li, Oriol Vinyals, Nicolas Heess, Lars Buesing, Sebastien Racanière, David Reichert, Théophane Weber, Daan Wierstra, and Peter Battaglia. Learning model-based planning from scratch. arXiv preprint FrXXiv:707.06/// pp. 1-13, 2017.
Malcolm Reynolds, Gabriel Barth-Maron, Frederic Besse, Diego de Las Casas, Andreas Fidjeland, Tim Green, Adrià Puigdomènech, Sébastien Racanière, Jack Rae, and Fabio Viola. Open sourcing Sonnet library for constructing neural networks. https: https://depminnd.co com/bloc opeen-suurcing sonnet/ / 2017.
Christopher D Rosin. Multi-armed bandits with episode context. Annals of Mathematics and Artifi- cial Intelligence, 61()):203 230, 2011.
Marwin HS Segler, Mike Preuss, and Mark Waller. Planning chemical syntheses with deep neural networks and symbolic AL. Nature, 555 2018.
Iulian Vlad Serban, Chinnadhurai Sankar, Michael Pieper, Joelle Pineau, and Yoshua Bengio. The bottleneck simulator: model-based deep reinforcement learning approach. arXiv preprint 1r07.v/807.0472. 2018.
Yelong Shen, Jianshu Chen, Po-Sen Huang, Yuqing Guo, and Jianfeng Gao. M-Walk: Learning to walk over graphs using monte carlo tree search. In Advances in Neural Information Processing Systems, pp. 6786-6797, 2018.
David Silver, Richard S Sutton, and Martin Müller. Sample-based learning and search with perma- nent and transient memories. In Proceedings of the 25th international conference on Machine learning, pp. 968-975. ACM, 2008.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529:484-489, 2016.
David Silver, Julian Schritwieserr, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap. Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of Go without human knowledge. Nature, 550:354 359, 2017a
David Silver, Hado van Hasselt, Mattco Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Arnold, David Reichert, Neil Rabinowitz, Andre Barreto, and Thomas Degris, The predic- tron: End-to-end learning and planning. In Proceedings of the 34th International Conference on Machine Learning (ICML 2017), 2017b.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters Chess, Shogi, and Go through self-play. Science, 362(6419): 1140-1144, 2018
Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal planning networks. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), 2018
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, (1):9-44, 1988.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. In Proceedings of the 7th International Conference on Machine Learning (ICML 1990), 1990.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning. MIT Press, 2nd edition, 2018
12

---

Published as a conference paper at ICLR 2020
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. In Proceedings of the 30th Conference on Neural Information Processing Systems (NeurIPS 2016), 2016.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforce- ment learning. arXiv preprint arXiv:1907.0205: 02057, 2019.
Théophane Weber, Sébastien Racanière, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pascanu, Peter Battaglia, Demis Hassabis David Silver, and Daan Wierstra. Imagination-augmenue agents for deep reinforcement learning. In Proceedings of the 31st Conference Neural Information Processing Systems (NeurIPS 2017), 2017.
Tianhe Yu, Gleb Shevchuk, Dorsa Sadigh, and Chelsea Finn. Unsupervised visuomotor control through distributional planning networks. arXiv preprint 1rX:v:1902.05542, 2019.
13

---

Published as a conference paper at ICLR 2020
A FURTHER AGENT DETAILS
In all experiments except Tabular Tightrope (see (Section B.2) and Atari (see Appendix E), we use a distributed training setup with GPU learner and 64 CPU actors. Our setup was implemented using TensorFlow (Abadi et al. 2016) and Sonnet (Reynolds et al. (017), and gradient descent was performed using the Adam optimizer (Kingma & Ba 2014) with the TensorFlow default parameter settings (except learning rate).
A.1 Q-LEARNING
Except for in Atari (see Appendix E), we used 1-step implementation of Q-learning, with the standard setup with experience replay and a target network (Mnih et al. 2013). We controlled the rate of experience processed by the learner such that the average number times each transition was replayed (the "replay ratio") was kept constant. For all experiments, we used batch size of 16, learning rate 0.0002, a replay size of 4000 transitions (with a minimum history 100 transitions), replay ratio of 4, and updated the target network every 100 learning steps.
We used a variant of epsilon-greedy exploration described by Bapst et al. (2019) in which epsilon changed adaptively over the course of episode such that it is lower earlier in the episode and higher later in the episode, with an average value over the whole episode. We annealed the average value of from to 0.01 over 1e4 episodes.
A.2 SAVE
Algorithm A.1 Pseudocode for the SAVE algorithm.
1:
procedure SAVE(0)
2:
while true do
3: Begin episode at
4:
while acting do
5:
Estimate QMCTS(s, + MCTS(s, Qa)
6:
Select a using epsilon-greedy from QMeTs(s,) Execute in environment and receive s', Add (s, a, r. QMCTT(s,.) replay buffer +- s'
7:
8:
9:
10:
while learning do Sample minibatch of experience from the replay buffer Update to minimize Equation
11:
12:
13:
procedure MCTS(so, Qa)
15:
Qo(s,a) +- (s for all No(s,a) for all s.a
17:
k 0
18:
while search budget remains (k K) do
19:
Traverse the search tree with T.. (Equation Expand new state 8y and add it to the search tree Evaluate max,, (s a) and backup returns (Equation
20:
21:
22:
Set N (8, a) < NE(8, a) and then increment counts of visited states and actions Compute estimates for Q (8,a) a) (Equation
23:
24:
+ k +
25:
Return {Qk (So, a;));
The SAVE agent is implemented as described in Section and Algorithm A. provides additional pseudocode explaining the algorithm. Algoorri A. provide example of using SAVE in an episode setting where learning happens after every episode; however, SAVE can be used any Q-learning setup including in distributed setups where separate processes concurrently acting and learning. particular, in our experiments we use the distributed setup described in (Section A. Note that when performing epsilon-greedy exploration (Line 6 of Algorithm A. either choose
14

---

Published as a conference paper at ICLR 2020
an action uniformly at random with probability and otherwise choose the action with the highest value of QACTS out of the actions which were explored during search (i.e., we do not consider actions that were not explored, even if they have higher QMCTs). In all experiments (except tabular Tightrope), we use a UTC exploration constant of c = 2, though we have found SAVE's performance to be relatively robust this parameter setting.
A.3 PUCT
The PUCT search policy is based on that described by Silver et al. (2017a) and Silver et al. (2018). Specifically, we choose actions during search according to Equation 1| with:
Qk = Ei=1 Nalss,) (₂(,,a)
(x(
Us(s, a) (s))(s V. Na(s,a)
(s(s
where e is an exploration constant, I (s, a) is the prior policy, and M&(s, a) is the total number of times action a had been taken from state s at iteration k of the search. Like [iiver et al. (2017a 2018), we add Dirichlet noise to the prior policy:
(3, = e) ( a) + EM);
where ~ Dir(1 (ri(1//aann /naciooss) In our experiments we set = 0.25 and = 2. During training, after search is complete, we sample action execute in the environment from TMCTS (so, a) (s. a)/ Σ, (*) At test time, we select the action which has the maximum visit count (with random bie-breaking).
To train the PUCT agent, we used separate policy mo(s, a) and value (s)s) heads which were trained using a combined loss (Equation 60, with:
Lq
Vo(s) R|12
N
LA N
TMCTS(S,.) log (o(
where R. is the Monte-Carlo return observed from state s. We used fixed values of åQ = 0.5 and 8A 0.5 in all experiments with PUCT. We used the same replay and training setup used in the Q-learning and SAVE agents, with two exceptions. First, we additionally include episodic Monte-Carlo returns R and policies TMCTS in the replay buffer they be used during learning. Second, we did not use e-greedy exploration (because the Dirichlet noise in the PUCT term already enables sufficient exploration).
We tried several different hyperparameter settings and variants of the PUCT agent to attempt to improve the results. For example, we tried using I-step TD error for learning the values, which should have lower variance and thus result in more stable learning of values. We also tried reducing the replay ratio to 1 and the replay size to 400 in order to make the experience for training more on-policy. However, we did not find that these changes improved the results. We also tried different settings of e for the Dirichlet noise, but found that lower values resulted too little exploration, while higher values resulted in too much exploration.
B DETAILS ON TIGHTROPE
B. ENVIRONMENT
The Tightrope environment has 1I states which are connected together in a chain. Each state has 100 actions, M% of which will cause the episode to terminate when executed and the rest of which will cause the environment to transition to the next state. Each state is represented using a vector of 50 random values drawn from standard normal distribution, which are the same across episodes. The indices of terminal actions are selected randomly and are different for each state but are consistent across episodes. Agents always begin in the first state of the chain.
15

---

Published as a conference paper at ICLR 2020
In the sparse reward setting, randomly select one of the states in the chain to be the "final" state (excluding the first state), to enable the agent to sometimes train on easy problems and sometimes train on hard problems. If the agent reaches this final state, it receives reward of 1 and the episode terminates. If it takes a non-terminal action, it transitions to the next state in the chain and receives reward of O. Otherwise, if it takes terminal action, the episode terminates and the agent receives a reward of O.
In the dense reward setting, the "final" state is always chosen be the last state in the chain. If the agent reaches the final state in the chain, it receives a reward of 0. and the episode terminates. If it takes non-telminal action, it transitions to the next state in the chain and receives a reward of 0.1. Otherwise, if it takes a terminal action, the episode terminates with a reward of O.
B.2 TABULAR EXPERIMENTS
During training, we execute each tabular agent in the environment until the episode terminates. Then, we perform a learning step using the experience generated from the previous episode. This process repeats for some number of episodes (in our experiments, 500). After training, we execute each agent in the environment 100 times and compute the average reward achieved across these 100 episodes. For all cases in which search is used, we use a UCT exploration constant of c = 0.1.
Q-Learning Tabular Q-learning begins with a table of state-action values initialized to zero We perform -ereedyy exploration with e = 0.1, and add the resulting experience to replay buffer with maximum size of 1000 transitions. We perform episodic learning, where during each episode the Q-values are fixed and after the episode is complete we update the Q-values by performing a single pass through the experience in the replay buffer in random order. We use a learning rate of 90 = 0.01. At test time, the Q-learning agent uses MCTS in the same manner as SAVE.
SAVE Tabular SAVE begins with table of state-action values initialized to zero. During search, values are looked up in this table and used to initialize Qu- The values are also for bootstrapping. During learning, we perform both Q-learning (as described in the Q-learning agent) as well as an update based on the gradient of the cross-entropy amortization loss (Equaatio 6). We use Bo = 0.01 and AA-
PUCT Tabular PUCT begins with two tables; one with state values (initialized to zero) and one with action probabilities (initialized the uniform distribution). During search, action probabilities are looked and used in the PUCT term, while state values are looked up and used for bootstrapping. Search proceeds as described in Seccion A.3] During learning, TMCTS is copied back into the action probability table (this is equivalent to an L2 update with a learning rate of I); we also experimented with doing an update based on the cross entropy loss but found this resulted in worse performance. The value at episode is given by:
Vi(s) = -
-1-a) Vi- + acR:1(s),
(8)
where (s) is the return obtained after visiting state s during episode ( 1. In our experiments we used a = 0.5. We also experimented with using Q-learning rather than Monte- Carlo returns, but found that these resulted in similar levels of performance.
UCT The UCT agent is as described in Section 3.1 with Vs)) at unexplored nodes estimated via a Monte-Carlo rollout under uniform random policy. The only difference from regular UCT is that we did not require all actions to be visited before descending down the search tree; unvisited actions were initialized to a value of zero. For Tightrope, this the optimal setting of the default Q-values because all possible rewards are greater than or equal to zero. Once an action is found with non-zero reward the best option is to stick with it, so it would not make sense to set the values optimistically. Actions that cause the episode to terminate have reward of zero, so it would also not make sense to set the values pessimistically as this would lead to over-exploring terminal actions. Setting the values to the average of the parent would either have the effect of setting zero or setting optimistically (if the parent had positive reward).
To select the final action to execute in the environment, the UCT agent selects a visited action with the maximum estimated value. We could consider alternate approaches here, such as selecting uni- formly at random from unexplored actions if none of the visited actions have high enough expected
16

---

Published as a conference paper at ICLR 2020
(a) SAVE
(b) SAVE AL
(c) Q-Learning
Reward
(d) SAVE w/ L2
(e) SAVE w/o QL
(f) SAVE PUCT
Reward
0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
# Episodes
1e6
# Episodes
1e6
# Episodes
Figure C.I: Learning curves on the Covering task. Each plot shows median performance across 10 seeds, with shaded regions showing the min and max seed.
values. We experimented with this approach, using threshold value of zero (which is the expected value for bad actions in Tightrope), and find that this indeed improves performance (p = 0.02), though the effect size is quite small: on the dense setting with M 95% we achieve a median reward of 0.08 (using this thresholding action selection policy) versus 0.07 (selecting the max of visited actions).
B.3 FUNCTION APPROXIMATION EXPERIMENTS
We used the same learning setup for the Q-learnin, SAVE, and PUCT agents described in Ap- pendix A For the network architecture of our agents, we used shared multilayer perceptron (MLP) torso with two layers of size 64 and ReLU activations. To predict Q-values, we used an MLP head with two layers of size 64 and ReLU activations, with final layer of size 100 (the number of ac- tions) with linear activation. To predict policy in the PUCT agent, we used the same network architecture as the Q-value head. To predict state values in the PUCT agent, MLP head with two layers of size 64 and ReLU activations, and a final layer of size with linear activation. All network weights were initialized using the default weight initialization scheme in Sonnet (Reynolds et al. 2017). For both the SAVE and PUCT agents we used loss coefficients of 8q 0.5 and BA 0.5.
used separate
We trained each agent 10 times and report results after le6 episodes in a version of Tightrope that has 95% terminal actions Figure 2 right). During training, the SAVE and PUCT agents had access search budget of 10 simulations; the Q-learning agent did not use search. We also explored training agents with different numbers of terminal actions and different budgets. Qualitatively, same results as in the tabular setting: the PUCT agent can perform well for larger budgets (50+), but struggles with small budgets, underperforming performed well in all our experiments, even for small budgets like 5 or 10.
found the
model--fee Q-learning agent. In contrast, SAVE
C DETAILS ON CONSTRUCTION
C. AGENT DETAILS
SAVE For SAVE, we annealed PQ from to 0.1 and 3p1 from 0 to 4.5 over the course of 5e4 episodes. We found this allowed the agent to rely more on Q-learning early in training to build a good Q-value prior, and then more on MCTS later in training once good prior had already been established.
17

---

Published as a conference paper at ICLR 2020
(a) SAVE
(b) SAVE AL
(c) Q-Learning
Reward
(d) SAVE w/ L2
(e) SAVE w/o QL
(f) SAVE PUCT
Reward
0
5
10
20
50
0
5
10
20
50
0
5
10
20
50
Test Budget
Test Budget
Test Budget
Figure C.2: Detailed final results on the Covering task. Each plot shows median performance across 10 seeds, with error bars showing the min and max seed.
Q-Learning The Q-Learning agent is as described in Section A.1 In particular, we follow the same setup the GN-DQN agent deseribed in [Bapst et al. (2019). During training, Q-Jearning with no search. At test time, we may allow the Q-learning agent to additionally perform MCTS, using the same search procedure as that used by SAVE (i.e., initializing the Q-values using the trained Q-function and initializing the visit counts to one).
pure
SAVE without Amortization Loss The SAVE without an amortization loss the same as the basic SAVE agent, except that it includes no amortization loss (i.e., L(O, D) = (l D)). This equivalent the GN-DQN-MCTS agent described by Bapst et al. (2019).
with (ss) at unexplored nodes estimated via using pretrained aation-value function (trained using the same setup as the Q-learning agent). Additionally, unlike standard UCT we did not require all actions to be visited before descending down the search tree.
UCT UCT is described in Section 3.1,
SAVE with L2 SAVE with L2 loss is identical SAVE except that uses different amorti- zaton loss:
Ca EllMOcT((s,)) Qo(s, 1]1
Similar to the SAVE agent, anneal BQ from to 0.1 and A from 0 to 0.045 over the course of 5e4 episodes.
SAVE without Q-Learning SAVE without the Q-learning loss identical to SAVE except that we do not use Q-learning and we use the 1.2 amortization loss described in the previous paragraph:
C( D) (A,
where we set НΛ = 0.025. The reason we use the L2 loss rather than the cross-entropy loss is that otherwise the Q-values will not actually be real Q-values, in that they will have grounding in the actual scale of rewards. We did experiment with using only the cross-entropy loss with no Q-learning, and found slightly worse performance than when using the 1.2 loss and no Q-learning.
SAVE with PUCT SAVE with PUCT uses the same learning procedure as SAVE but different search policy. Specifically, we use the PUCT search policy described in (Section A.3 and Equation 7 To do this, set (s, a) a (QQo (Qo(s, a)), where is the softmax actions with a temperature of 1. We use the same settings for Dirchlet noise to encourage exploration during search. After search complete, select action using the same epsilon-greedy action procedure used by the SAVE
18

---

Published as a conference paper at ICLR 2020
agent rather than selecting based visit counts. We experimented with selecting based on visit counts instead, but found this resulted in the same level of performance.
C.2 EXPERIMENTAL SETUP
Observations are given as graphs representing the scene, with objects in the scene corresponding to nodes in the graph and edges between every pair of objects. All agents use the same network architecture (Battaglia et al. 2018)) described in [Bapst et al. (2019) to process these graphs. Briefly, we use a graph network architecture which takes graph as input and returns a graph with Q-values relative object-hased action like "pick up block B and put it on block D Each edge additionally has multiple actions associated with it which correspond to particular offset locations where the block should be placed, such as "on the top left".
on the edges of the graph. Each edge corresponds
Bapst et al. (2019) describe four Construction tasks: Silhouette, Connecting, Covering, and Covering Hard. We reported results on three of these tasks in the main text (Connecting, Covering, and Covering Hard). The agents in Bapst et al. (2019) already reached ceiling performance on Silhouette and thus we do not report results for that task here, except to report that SAVE also reaches ceiling performance.
The agents used 10 MCTS simulations during training and were evaluated on 0 to 50 simulations at test time, with the exception of the UCT agent, which always used 1000 simulations test time, and the Q-learning agent, which did not peform search during learning. We trained 10 seeds per agent and report results after le6 episodes. Figure C.T show details of learning progress for each of the agents compared in the ablation experiments on the Covering task (Section 4.2). Figure C.2 shows detailed final performances evaluated different test budgets. We evaluated all agents on the hardest level of difficulty of the particular task they were trained on for either 10000 episodes Figuure 3D-c) or 1000 episodes (Figure 3 and FFigure C.2). In general, while we find that search at test time can provide small boosts in performance, the main gains are achieved by incorporating search during training.
C.3 DISCUSSION OF ABLATION RESULTS
Here expand the results presented in the main text and in Figure 3 and Figure C.2
Cross-entropy vs. L2 loss While the L2 loss Figure C.2 orange) can result in equivalent perfor- mance as the cross-entropy loss (Figure C.2 green), this is the cost of higher variance across seeds and lower performance on average. This is likely because the L2 loss encourages the Q-function to exactly match the Q-values estimated by search. However, with a search budget of 10, those Q- values will be very noisy. In contrast, the cross-entropy loss only encourages the Q-function to match the overall distribution shape of the Q-values estimated by search. This is less strong con- straint that allows the information acquired during search to be exploited while not relying on it too strongly. Indeed, we can observe that the agent with L2 amortization loss actually performs worse than the agent that has amortization loss at all (Figure C.2, purple) when using a search budget of 10, suggesting that trying to match the Q-values during search too closely can harm performance.
consider interesting interaction between Q-learning and the amortization loss. Due to the search locally avoiding poor actions, Q-learning will rarely actually operate on low-valued actions, meaning most of its computation is spent refining the estimates for high-valued actions. The softmax cross entropy loss ensures that low-valued actions have lower values than high-valued actions, but does not force these values to be exact. Thus, in this regime we should have good estimates of value for high-valued actions and worse estimates of value for low-valued actions. In contrast, an L2 loss would require the values be exact for both low and high valued actions. By using cross entropy instead, we can allow the neural network to spend more of its capacity representing the high-valued actions and less capacity representing the low-valued actions, which we care less about in the first place anyway.
Additionally,
teal), the SAVE agent's 's perfor- mance suffers dramatically. As discussed in the previous section, the Q-values estimated during search are very noisy, meaning it is not necessarily a good idea to try to match them exactly. Ad- ditionally, QMeTS on-pol experience and can become stale if Qo changes too much between
With vs. without Q-learning Without Q-learning Figure C.2
19

---

Published as a conference paper at ICLR 2020
Reward
Epeilon-Greedy
3-
Calegarical (Visits)
2-
Categorical (Q-Values)
0.25
0.75
1.00
Episodes
Figure C.3: Performance of different exploration strategies the Covering task.
Reward
3-
2-
SAVE
G-Learning
Transitions
107
Figure C.4: Performance of SAVE and Q-learning on Covering, controlling for the same number of environment interactions (including those during search).
when QMTTS was computed and when is used for learning. Thus, removing the Q-learning loss makes the learning algorithm much more on-policy and therefore susceptible to the issues that come with on- on-policy training. Indeed, without the Q-learning loss, we can only rely on the Q-values estimated during search, resulting in much worse performance than when Q-learning is used.
UCT vs. PUCT Finally, we compared to a variant which utilizes prior knowledge by transforming the Q-values into a policy via softmax and then using this policy as a prior with PUCT, rather than using to initialize the Q-values (Figure C.2 brown). With large amounts of search, the initial setting of the Q-values should not matter much, but the case of small search budgets (as seen here), the estimated Q-values do not change much from their initial values. Thus, if the initial values are zero, then the final values will also be close to zero, which later results in the Q-function being regressed towards nearly uniform distribution of value. By initializing the Q-values with the Q- function, the values that are regressed towards may be similar to the original Q-function but will not be uniform. Thus, we can more effectively reuse knowledge across multiple searches by initializing the Q-values with UCT rather than incorporating prior knowledge via PUCT.
C.4 ADDITIONAL RESULTS
We performed several other experiments to tease apart the questions regarding exploration strategy and data efficiency.
Exploration strategy When selecting the final action perform in the environment, SAVE uses an epsilon-greeds, exploration strategy. However, many other exploration strategies might be con- sidered, such as UCB, categorical sampling from the softmax of estimated Q-values, or categorical sampling from the normalized visit counts. We evaluated how well each of these exploration strate- gies work, with the results shown in Figure C.3) We find that using apsilon-greedd works the best out of these exploration strategies by substantial margin. We speculate that this may be because important for the Q-function to be well approximated across all actions, so that it is useful during MCTS backups However, UCB and categorical methods will not uniformly sample the
20

---

Published as a conference paper at ICLR 2020
tion space, meaning that some actions are very unlikely to be ever learned from. The amortization loss will not help either, as these actions will not be explored during search either. The error in the Q-values for unexplored actions will grow over time (due to catastrophic forgetting), leading to poorly approximated Q-function that is unreliable. In contrast, epsilon-greedy consistently spends little bit of time exploring these actions, preventing their values from becoming too inaccurate. We expect this would be less of a problem if we were to use a separate state-value function for bootstrapping (as done by AlphaZero).
Data efficiency With a search budget of 10, SAVE effectively sees 10 times many transitions as a model-free agent trained on the same number of environment interactions. To more carefully compare the data efficiency of SAVE, we compared its performance to that of the Q-learning agent on the Covering task, controlling for the same number of environment interactions (including those seen during search). The results are shown in Figure C.4) illustrating that SAVE converges to higher rewards given the same amount of data. We find similar results in the Marble Run environment, shown in Figure D.2
D DETAILS ON MARBLE RUN
D.1 SCENE GENERATION
Scenes contain the following types of objects (similar Bapsi et al. (2019))
Floor (in black) that supports the blocks placed by the agent.
Available blocks (row of blue blocks at the bottom) that the agent picks and place in the scene (with replacement).
Blocks (blue blocks above the floor) that the agent has already placed. They may take lighter blue color to indicate that they are sticky. A sticky block gets glued to anything touches.
Goal (blue dot) that the agent has to reach with the marble.
Marble (green circle) that the agent has to route to the goal.
Obstacles (red blocks, including two vertical walls), that the agent has to avoid, by not touching them neither with the blocks or the marble.
All the initial positions for obstacles in the scene are sampled from tessellation (similar to the Silhouette task in [Bapst al. (2019) made of rows with random sequences of blocks with sizes of discretization unit in height and or 2 discretization units in width (a discretization unit corresponds to the side of the first available block). The sampling process goes as follows:
1. Set the vertical position of the goal to the specified discrete height (according to level) corresponding to the center of one of the tessellation rows, and the vertical position of the marble rows above that.
2. Uniformly sample a horizontal distance between the marble and the goal from predefined range, and uniformly sample the absolute horizontal positions respecting that absolute dis- tance.
3. Sample a number of obstacles (according to level) from the tessellation spanning up to the vertical position of the marble.
Obstacles are sampled from the tessellation sequentially. Before each obstacle is sampled, all objects in the tessellation that are too close (# 2 layers vertically and with less than 2 discretization units of clearance sideways) to the goal, the target, or previously placed obstacles, are removed from the tessellation in order to prevent unsolvable scenes. Then probabilities are assigned all of the remaining objects in the tessellation according to one of the following criteria (the criteria itself is also picked randomly with different weights) designed to avoid generating trivial scenes:
(Weight=4) Pick uniformly tessellation object lying exactly the floor and between the marble and the goal horizontally, since those objects prevent the marble from rolling freely the floor (only applicable if the tessellation still has objects of this kind available).
21

---

Published as a conference paper at ICLR 2020
(Weight=1) Pick tessellation object that is close (horizontally) to the marble. Probabilities proportional to (d/r)*40.1 (where d is the horizontal distance between each object and the marble scaled by the width of the scene and T is a temperature set to 0.1) are assigned to all objects left in the tessellation, and one of them is picked.
(Weight=1) Pick a tessellation object that close (horizontally) the goal. Identical the previous one, but using the distance to the goal.
(Weight=1) Pick tessellation object that is close (horizontally) to the middle point between the ball and the goal. Identical to the previous one, but using the distance to middle point, and temperature of 0.2.
(Weight=1) Pick any object remaining in the tessellation with uniform probability (to in- crease diversity).
D.2 CURRICULUM DIFFICULTY
We used a curriculum to sample scenes of increasing difficulty (Fig. D.I according to:
Level
Goal height
Marble/Goal distance (scene width fraction)
# obstacles
Max # steps
(discretization units)
0
[0.03,0.3]
20
0
[0.36,0.49]
20
2
0
[0.50,0.63]
2
20
0
[0.69,0.82]
20
4
0 1
[0.83,1]
3
20
5
[0.83,1]
25
2
[0.83.1]
30
During both training and testing, episodes at a certain curriculum level are sampled not only from that difficulty, but also from all of the previous difficulty levels, using truncated geometric distri- bution with a decay of 0.5. This means that each level, about half the episodes correspond to that level, half of the remaining episodes correspond to the previous level, half of the remaining the level before that, and so on. By truncated we mean that, because it is not possible to sample episodes for negative levels, so we truncate the probabilities there and re-normalize.
D.3 ADAPTIVE CURRICULUM
Given the complexity and the sparsity of rewards in this task, we trained agents using an adaptive curriculum to avoid presenting unnecessarily hard levels to the agent until the agent is able to solve the simpler levels. Specifically at each level of the curriculum results according to all possible combinations of scene properties consisting of:
keep track and bin past episode
Height of the target (discretized to tessellation rows).
Horizontal distance d between marble and goal (discretized to d < 1/3, 1/3 < d < 2/3,
> 2/3, where is normalized by the width the scene).
Number of obstacles.
Height of the highest obstacle (discretized to tessellation rows).
Height of the lowest obstacle (discretized tessellation rows).
and require the agents to have solved least 50% of scenes of the last 50 episodes in each bin
individually, but simultaneously in all bing"] before allow the agent to progress to the next level of difficulty. This is very strict criteria, which effectively means the agent has find solutions for all representative variations of the task at that level before is allowed to progress to the next level.
"ignoring underrepresented bins for which episodes are generated at less than 0.25 the expected generation rate according to a uniform prior across bins (which we only start estimating after the first 300 episodes, to discover sufficient number of distinct bins). This is to avoid delayed curriculum progress just due to lack of recent statisties for those bins.
22

---

Published as a conference paper at ICLR 2020
Sample
Sample
Sample
Sample
Sample
Sample
Level 0
Level 1
Level 2
Level 3
Level 4
Level
Level
Figure D.I: Scenes samples each curriculum level for the marble run task. During training, the n-th level of the curriculum consists of scenes sampled from the rows up to the n-th row with truncated geometric distribution with decay of 0.5.
D.4 AGENT STEP, ACTION AND REWARD EVALUATION
Each agent step consists of four phases:
1. Block placement phase: The agent picks one object from the available objects and places it into the scene. If the block placed by the agent was sticky the agent will receive a negative reward according to the cost (which may be either 0 or 0.04).
2. Block settlement phase: The physics simulation (keeping the marble frozen) is run until the placed blocks settle (up maximum of 20 s). During this phase the new block may affect the position of previously placed blocks.
3. Marble dynamics phase: The physics simulation including the marble is run until the mar- ble collides with 8 objects, with timeout of 10 at each collision, that is maximum 80s. This phase may terminate early the marble reaches the goal (task solved and episode terminated with reward of 1.), but also if the marble any of the blocks touch an obstacle.
4. Restore state phase: After the marble dynamics phase, the marble and of the blocks are moved back to the position where they were at the end of the block settlement phase. This is to prevent the agent from using the marble to indirectly move the blocks with persistent effect across steps.
The block placement phase and block settlement phase, as well as the action space is identical to those in Bapst al. (2099).
D.5 OBSERVATION
The observation is identical to the Construction tasks in Bapst et al. 2019), with additional one-hot encoding the object shape (e.g. rectangle vs triangle vs circle) and includes all blocks positions and the initial marble position at the end of the block settlement phase. Note that the agent
23

---

Published as a conference paper at ICLR 2020
Sticky block cost 0
Sticky block cost 0.04
1.0
SAVE
0.6
Difficulty
0.2
1.0
0.8
DLLarrin
0.6
Final
0.4
0.2
0.0
10a
10a
10
107
10
10*
10a
10
107
100
Episodes
Episodes
Figure D.2: Learning curves for the Marble Run environment. Each line shows the median reward across 10 seeds, and the shaded regions show min and max seed performance. Each color corre- sponds to a different level of curriculum difficulty. Difficulties less than the final difficulty are only evaluated while the agent is training at that curriculum level; the final level of difficulty is always evaluated.
never actually gets observe the marble's dynamics, and therefore does not get direct feedback about why the marble does or does not make it to the goal (such that is getting stuck in a hole). An interesting direction for future work would be to incorporate this information into the agent's learning as well.
D.6 TERMINATION CONDITION
There are several episode termination conditions that may be triggered before the task is solved:
An agent places a block in position that overlaps with an existing block or obstacle.
An agent has placed block that during the settlement phase touches obstacle.
An agent has placed block that, at the end of the block settlement phase overlaps with the initial marble position.
Maximum number of steps is reached.
Note that touching obstacles during the marble dynamics phase does not terminate the episode be- cause purely evaluating the reward function and, during the restore state phase, all objects are returned to there previous locations. This makes it possible for the agent to correct for any ob- stacle collisions that happened during the marble dynamics phase, by placing additional blocks that re-route the marble.
D.7 ADDITIONAL RESULTS
We used the same experimental setup in the other Construction tasks (Apendix C). In particular, during training. for each seed each agent we checkpoint the weights which achieve the highest reward on the highest curriculum level, and then use these checkpoints to evaluate performance in Figure 4 Figure D.2] additionally shows details of the training performance at each level of difficulty in the curriculum. We can see that at around difficulty level 4-5, the Q-learning agent becomes unstable and crashes, while the SAVE agent stays stable and continues to improve. Indeed, shown in Figure D.3] the Qelearning agent never makes it to difficulty level 6 (when sticky blocks are free) or even difficulty level (when sticky blocks have a moderate cost). The SAVE agent is able reach harder levels of difficulty, and does with fewer learning steps.
24

---

Published as a conference paper at ICLR 2020
Sticky
10'
Sticky block 0.04
8 Episodes
10a
10
SAVE Q-Leaming
5
Dificullyy
Dificully
Curriculum progress in Marble Run. Light lines show individual curriculum progress per seed, and dark lines are computed over the median of these seeds. The r-axis shows the particular curriculum level and the y-axis indicates at which episode that level of difficulty was reached.
Figure D.3:
Level
Baseline
Controlled
SAVE
% Change
Alien
71925.1
96013.5
280227.3
191.9%
Asteroids
251033.3
266306.7
274431.7
3.1 %
Beam Rider
96654.4
113930.6
195703.8
71.8%
Centipede
517332.2
562742.3
767206.6
36.3%
Crazy Climber
311203.8
271151.5
324726.4
19.8%
Frostbite
15814.2
11052.3
202744.2
1734.4%
Gravitar
7854.0
11314.3
11484.1
1.5%
Hero
30515.9
44574.3
44796.0
0.5%
Ms. Pacman
25377.4
27776.3
47186.0
69.9%
Name This Game
45027.1
40790.0
58621.1
43.7%
River Raid
33819.5
32720.8
41031.6
25.4%
Space Invaders
3639.2
42387.4
63684.7
50.2%
Up 'n' Down
563661.0
568735.6
585475.6
2.9%
Zaxxon
116892.6
73073.1
213370.4
192.0%
Median
58476.1
58823.7
199224.0
40.0%
Mean
149339.3
154469.2
222192.1
174.5%
Table E.1: Results on Atari. Scores are final performance averaged over seeds. "Baseline" the standard version of R2D2 (Kapturowski et al., 2018). "Controlled" is our version that is controlled to have the same replay ratio as SAVE. The rightmost column reports the percent change in reward of SAVE over the controlled version of R2D2. Bold scores indicate scores that are within 5% of the best score on a particular game. The last two rows show median and mean scores, respectively. The percentages in the last two rows show the median and mean across percent change, rather than the percent change of the median/mean scores.
E DETAILS ON ATARI
E. EXPERIMENTAL. SETUP
of 14 Atari games in the Arcade Learning Environment (Bellemare) et al. 2013). The games were chosen combination of classical action Atari games such as As- teroids and Space Invaders, and games with a stronger strategic component such as Ms. Paeman and Frostbite, which are commonly used evaluation environments for model-based agents (Buesing
We evaluated SAVE
al. 2018 Farquhar et al. 2018 Oh et al. 2017 Guez al. 2019).
SAVE was implemented on top of the R2D2 agent (Kapturowski et al. 2018) as described in Algo- rithm A.I Concretely, this means we evaluate the function QMCTS instead Qa to select an action in the actors, and optimize the combined loss function (Equation 6 instead of the TD loss in the learner. For hyperparameters, we used a search budget of 10, and Bo = 1, 3 = 10. We did very little tuning to select these hyperparameters, only sweeping found while both of these settings resulted in similar performance, 3 = 10 worked slightly better. It is likely that with further tuning of these parameters, even larger increases in reward be achieved,
two values of 3 E {1, 10). We
25

---

Published as a conference paper at ICLR 2020
1e5
allen
3.5 1e5
asterolds
_.0_1e5
3.0
1e5
@ederrider
3.0
2.5
Reward
2.5
2.0
1.5
1.5
1.0
0.5
0.5
0.0
0.0
Reward 6
1e5
centipede
1ee5
clay_cc
1e5
frostbite
5
5
1.4
1e4
gravitar
1e4
hero
6-
1e4
ms_pacman
5
5-
Reward
1.0
0.6
0.4
0.2
0.0
1e4
name this tame_gameeeee_a_e
5
1e4
riverraid
7-
1e4
space invaders
space_invaders
OUNUADO SwN-RTt-10
Reward
1e5
nown
4.5 _.__
zaxxon
Learner Steps
le6
4.0
3.5
Reward
Baseline
0.5
Controlled
0.0
Learnr Steps
le6
Learner Steps
1e6
Figure E.1: Learning curves on Atari games. Solid lines show the average over 3 seeds, and shaded regions show min and max seeds.
as Lq and LA will have very different relative magnitudes depending on the scale of the rewards in each game.
All plyee--prammete of R2D2 remain unchanged from the original paper, with the exception of actor speed compensation. By running MCTS, multiple environment interactions need to be evaluated for each actor step, which means transition tuples are added to the replay buffer slower rate, changing the replay ratio. To account for this, we increase the number of actors from 256 to 1024, and change the actor parameter update interval from 400 to 40 steps.
E.2 EVALUATION
The learning curves of our experiment are shown in Figure E. and Table E.I shows the final performance in tabular form. We ran three seeds for each of the Baseline, Controlled and SAVE agents for each game and computed final scores as the average score over the last 2e4 episodes of training. The Baseline agent represents the unchanged R2D2 agent from (Kapturowski The Controlled agent is R2D2 agent controlled to have the same replay ratio as SAVE, which we achieve by running MCTS in the actors but then discarding the results. As in SAVE, use 1024 actors with update interval 40 for the controlled agent.
al. 2018))
We can observe that in the majority of games, SAVE performs not only better than the controlled agent but also better than the original R2D2 baseline. While we see big improvements in the strategic games such Ms. Pacman, also notice gain in many of the action games. This suggests that
26

---

Published as a conference paper at ICLR 2020
model-based methods like SAVE can be useful even in domains that do not require much long- term reasoning.
27


Reviews:
[{"review_id": "Hkxs3EFlcr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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 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.\n\n\n- The idea of training Q-network using the result of MCTS planning is not new (e.g. UCTtoRegression in Guo et al 2014), but this paper takes further steps: the learned Q-network is again used for MCTS planning as Q initialization, the cross-entropy loss is used instead of L2-loss for the amortized value training, and the total loss combines Q-learning loss and the amortization loss.\n- In Figure 1, it says that the final action is selected by epsilon-greedy. Since SAVE performs MCTS planning, UCB exploration seems to be a more natural choice than the epsilon-greedy exploration. Why does SAVE use a simple epsilon-greedy exploration? Did it perform better than UCB exploration or softmax exploration? Also, what if we do not perform exploration at all in the final action selection, i.e. just select argmax Q(s,a)? Since exploration is performed during planning, we may not need exploration for the final action selection?\n- Can SAVE be extended to MCTS for continuous action space? SAVE trains Q-network, rather than a policy network that can sample actions, thus it seems to be more difficult to deal with continuous action space.\n- In Eq. (5), we may introduce a temperature parameter that trade-offs the stochasticity of the policy to further improve the performance of SAVE.\n- In Tightrope domain (sec 4.1), it says: \"The MDP is exactly the same across episodes, with the same actions always having the same behavior.\", but it also says: \"In the sparse reward setting, we randomly selected one state in the chain to be the “final” state to form a curriculum over the length of the chain.\" It seems that those two sentences are contradictive.\n- In the Tightrope experiment's tabular results (Figure 2), the performance of Q-learning is not reported. I want to see the performance of Q-learning here too.\n- In Figure 2, the search budgets for training and testing are equal, which seems to be designed to benefit SAVE than PUCT. Why the search budget should be very small even during training? Even if the fixed and relatively large search budget (e.g. 50 or 100) is used during training and the various small search budgets are only used in the test phase, does SAVE still outperform PUCT?\n- In Figure 2 (d), model-free Q-learning does not perform any planning, thus there will be much less interaction with the environment compared to SAVE or PUCT. Therefore, for a fair comparison, it seems that the x-axis in Figure 4-(d) should be the number of interactions with the environment (i.e. # queries to the simulator), rather than # Episodes. In this case, it seems that Q-Learning might be much more sample efficient than SAVE.\n- In Figure 3, what is the meaning of the test budget for Q-Learning since Q-Learning does not have planning ability? If this denotes that  Q-network trained by Q-learning loss is used for MCTS, what is the difference between Q-Learning and SAVE w/o AL?\n- In Figures 3, 4, 5, it seems that comparisons with PUCT are missing. In order to highlight the benefits of SAVE for efficient MCTS planning, the comparison with other strong MCTS baselines (e.g. PUCT that uses learned policy prior) should be necessary. A comparison only with a model-free baseline would not be sufficient.\n\n\n-----\nafter rebuttal:\n\nThank the authors for clarifying my questions and concerns. I feel satisfied with the rebuttal and raise my score accordingly.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "H1xECZ5_qS", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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 proposes an approach, named SAVE, which combines model-free RL (e.g. Q-learning) with model-based search (e.g. MCTS). SAVE includes the value estimates obtained for all actions available in the root node in MCTS in the loss function that is used to train a value function. This is in contrast to closely-related approaches like Expert Iteration (as in AlphaZero etc.), which use the visit counts at the root node as a training signal, but discard the value estimates resulting from the search. \n\nThe paper provides intuitive explanations for two situations in which training signals based on visit counts, and discarding value estimates from search, may be expected to perform poorly in comparison to the new SAVE approach:\n1) If a trained Q-function incorrectly recommends an action \"A\", but a search process subsequently corrects for this and deviates from \"A\", no experience for \"A\" will be generated, and the incorrect trained estimates of this action \"A\" will not be corrected.\n2) In scenarios with extremely low search budgets and extremely high numbers of poor actions, a search algorithm may be unable to assign any of the visit count budget to high-quality actions, and then only continue recommending the poor actions that (by chance) happened to get visits assigned to them. \n\nThe paper empirically compares the performance of SAVE to that of Q-Learning, UCT, and PUCT (the approach used by AlphaZero), on a variety of environments. This includes some environments specifically constructed to test for the situations described above (with high numbers of poor actions and low search budgets), as well as standard environments (like some Atari games). These experiments demonstrate superior performance for SAVE, in particular in the case of extremely low search budgets.\n\nI would qualify SAVE as a relatively simple (which is good), incremental but convincing improvement over the state of the art -- at least in the case of situations with extremely low search budgets. I am not sure what to expect of its performance, relative to PUCT-like approaches, when the search budget is increased. For me, an important contribution of the paper is that it explicitly exposes the two situations, or \"failure modes\", of visit-count-based methods, and SAVE provides improved performance in those situations. Even if SAVE doesn't outperform PUCT with higher search budgets (I don't know if it would?), it could still provide useful intuition for future research that might lead to better performance more generally across wider ranges of search budgets.\n\n\nPrimary comments / questions:\n\n1) Some parts of the paper need more precise language. The text above Eq. 5 discusses the loss in Eq. 5, but does not explicitly reference the equation. The equation just suddenly appears there in between two blocks of text, without any explicit mention of what it contains. After Eq. 6, the paper states that \"L_Q may be any variant of Q-learning, such as TD(0) or TD(lambda)\". L_Q is a loss function though, whereas Q-learning, TD(0) and TD(lambda) are algorithms, they're not loss functions. I also don't think it's correct to refer to TD(0) and TD(lambda) as \"variants of Q-learning\". Q-learning is one specific instead of an off-policy temporal difference learning algorithm, TD(lambda) is a family of on-policy temporal difference learning algorithms, and TD(0) is a specific instead of the TD(lambda) family.\n\n2) Why don't the experiments in Figures 2(a-c) include a tabular Q-learner? Since SAVE is, informally, a mix of MCTS and Q-learning, it would be nice to not only compare to MCTS and another MCTS+learning combo, but also standalone Q-learning.\n\n3) The discussion of Tabular Results in 4.1 mentions that the state-value function in PUCT was learned from Monte-Carlo returns. But I think the value function of SAVE was trained using a mix of the standard Q-learning loss and the new amortization loss proposed in the paper. Wouldn't it be more natural to then train PUCT's value function using Q-learning, rather than Monte-Carlo returns?\n\n4) Appendix B.2 mentions that UCT was not required to visit all actions before descending down the tree. I take it this means it's allowed to assign a second visit to a child of the root node, even if some other child does not yet have any visits? What Q-value estimate is used by nodes that have 0 visits? Some of the different schemes I'm aware of would involve setting them to 0, setting them optimistically, setting them pessimistically, or setting them to the average value of the parent. All of these result in different behaviours, and these differences can be especially important in the high-branching-factor / low-search-budget situations considered in this paper.\n\n5) Closely related to the previous point; how does UCT select the action it takes in the \"real\" environment after completing its search? The standard approach would be to maximise the visit count, but when the search budget is low (perhaps even lower than the branching factor), this can perform very poorly. For example, if every single visit in the search budget led to a poor outcome, it might be preferable to select an unvisited action with an optimistically-initialised Q-value.\n\n6) In 4.2, in the discussion of the Results of Figure 3 (a-c), it is implied that the blue lines depict performance for something that performs search on top of Q-learning? But in the figure it is solely labelled as \"Q-learning\"? So is it actually something else, or is the discussion text confusing?\n\n7) The discussion of Results in 4.3 mentions that, due to using search, SAVE effectively sees 10 times as many transitions as model-free approaches, and that experiments were conducted on this rather complex Marble Run domain where the model-free approaches were given 10 times as many training steps to correct for this difference. Were experiments in the simpler domains also re-run with such a correction? Would SAVE still outperform model-free approaches in the more simple domains if we corrected for the differences in experience that it gets to see?\n\n\nMinor Comments (did not impact my score):\n- Second paragraph of Introduction discusses \"100s or 1000s of model evaluations per action during training, and even upwards of a million simulations per action at test time\". Writing \"per action\" could potentially be misunderstood by readers to refer to the number of legal actions in the root state. Maybe something like \"per time step\" would have less potential for confusion?\n- When I started reading the paper, I was kind of expecting it was going to involve multi-player (adversarial) domains. I think this was because some of the paper's primary motivations involve perceived shortcomings in the Expert Iteration approaches as described by Anthony et al. (2017) and Silver et al. (2018), which were all evaluated in adversarial two-player games. Maybe it would be good to signal at an early point in the paper to the reader that this paper is going to be evaluated on single-agent domains. \n- Figure 2 uses red and green, which is a difficult combination of colours for people with one of the most common variants of colour-blindness. It might be useful to use different colours (see https://usabilla.com/blog/how-to-design-for-color-blindness/ for guidelines, or use the \"colorblind\" palette in seaborn if you use seaborn for plots).\n- The error bars in Figure 3 are completely opaque, and overlap a lot. Using transparant, shaded regions could be more easily readable.\n- \"... model-free approaches because is a combinatorial ...\" in 4.2 does not read well.\n- Appendix A.3 states that actions were sampled from pi = N / sum N in PUCT. It would be good to clarify whether this was only done when training, or also when evaluating.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "BklMPEJx5H", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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 proposes SAVE that combines Q learning with MCTS. In particular, the estimated Q values are used as a prior in the selection and backup phase of MCTS, while the Q values estimated during MCTS are later used, together with the real experience, to train the Q function. The authors made several modifications to ‘standard’ setting in both Q learning and MCTS. Experimental results are provided to show that SAVE outperforms generic UCT, PUCT, and Q learning.\n\nOverall the paper is easy to follow. The idea of the paper is interesting in the sense that it tries to leverage the computation spent during search as much as possible to help the learning. I am not an expert in the of hybrid approach, so I can not make confident judgement on the novelty of the paper.\n\n The only concern I have is that the significance of the result in the paper:\n1. The proposed method, including the modifications to MCTS and Q learning (section 3.2 and 3.3), is still a bit ad-hoc. The paper has not really justified why the proposed modification is a better choice except a final experimental result. Some hypotheses are made to explain the experimental results. But the authors have not verified those hypotheses. Just to list a few here: (a). The argument made in section 2.2 about count based prior; (b). the statement of noisy Q_MCTS to support the worse performance of L2 loss in section 4.2; (c). In the last paragraph of section 4.3, why would a model free agent with more episodes results in worse performance?\n2. The baselines used in this paper are only PUCT and a generic Q learning. What are the performances of other methods that are mentioned in section 2.1, like Gu 2016, Azizzadenesheli 2018, Bapst 2019? \n\nOther comments:\n1. What is the performance of tabular Q-learning in Figure 2 (a-c)?\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes Search with Amortized Value Estimates (SAVE) that combines Q-learning and MCTS.  SAVE uses the estimated Q-values obtained by MCTS at the root node to update the value network, and uses the learned value function to guide MCTS.

The rebuttal addressed the reviewers’ concerns, and they are now all positive about the paper. I recommend acceptance.
Decision: Accept (Poster)