SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Though the video is valuable in discerning the plan of an agent, a more quantitative analysis is warranted. Does the agent adhere to handcrafted planning algorithms, such as the  $n$ -step exhaustive search or MCTS, or does it learn novel planning algorithms? Given that the agents learn three types of actions—reset, imaginary, and real actions—we will analyze these actions individually.
Reset Action In a  $n$ -step exhaustive search, which involves iterating through every possible multi-step action sequence and choosing the one with the highest rollout return, a reset is triggered when the search depth hits  $n$ . Similarly, MCTS initiates a reset at a leaf node. However, the trained agent seems to learn different behavior of resetting.
As visualized in the video, the agents often chose to reset when envisioning a poor rollout, for instance, when imagining a box being moved to an irreversible state. To illustrate this further, we can examine the distribution of Temporal Difference (TD) errors of the rollout from the root node to a node  $i_n$ , defined as:
$$
\delta_ {i _ {n}} := \hat {r} _ {i _ {1}} + \gamma^ {2} \hat {r} _ {i _ {2}} + \dots \gamma^ {n} \hat {r} _ {i _ {n}} + \gamma^ {n + 1} \hat {v} _ {i _ {n}} - \hat {v} _ {\text {r o o t}}, \tag {24}
$$
where  $\{\mathrm{root}, i_1, i_2, \dots, i_n\}$  denotes the path from the root node to the node  $i_n$ . This TD error can be interpreted as the quality of the rollout relative to the current policy. We compare the distribution of TD errors of the current rollout  $\delta_{\mathrm{cur}}$  where the agent decided to reset with those where it did not, and the result is visualized in the left and middle panels of Figure 16.
We observe that the TD error distribution when agents opt to reset tends to lean more towards the left as compared to when they choose not to reset. This indicates a higher propensity for the agent to reset when the quality of the rollout is inferior. The thin tail on the left side of the distribution represents instances where the agent envisions a poor plan, such as reaching an irreversible state. Under such circumstances, the agent exhibits a high probability of resetting. This suggests that the agents learned to factor in the current rollout's quality in their decision to reset. This approach diverges from planning strategies like MCTS, which mandates a reset upon reaching a leaf node.
An additional dimension worth investigating is the expansion length, defined as the number of steps the agent executes beyond the last expanded node. In the case of MCTS, this length is consistently one as the algorithm requires a reset at every leaf node. However, as shown in the right panel of Figure 16, the trained agents generally exhibit a longer expansion length. In fact, when the maximum search depth  $L$  is increased to 10, the agent learns to search much deeper, as depicted in the same figure.
Figure 16: Left and Middle: Probability density of TD error upon choosing to reset and not to reset. These statistics are computed for agents trained with (Left) and without planning rewards (Middle). Right: Percentage of expansion length for agents trained with and without planning rewards at different maximum search depths  $L$ . All statistics are averaged over 20000 stages, each starting from random initial states.
Coupled with the aforementioned insights into the TD error distribution, we see that the agents prefer to persist with their existing plan unless they encounter an unfavorable leaf node or reach the set maximum search depth. These observations suggest that the agents tend to formulate plans that are both deep and narrow. Arguably, these learned behaviors mimic our planning processes more closely—we typically continue to think about a good plan, rather than simultaneously considering multiple plans and expanding each by a single step.
Imaginary Action In a  $n$ -step exhaustive search, imaginary action is chosen sequentially, while in MCTS, imaginary action is chosen according to the upper confidence bound (UCB). As for the trained agents, it is difficult to deduce the formula used to select the imaginary action as it is the output of a deep neural network. However, the maximum TD errors in a stage, defined as  $\delta^{\mathrm{max}} \coloneqq \max_k \delta_k$  where  $1 \leq k < K$  denotes the step within a stage, could be used as a proxy to evaluate the effectiveness of the search<sup>4</sup>. The distribution of maximum TD errors of search using the trained agent, UCB, and random search is shown in Figure 17. Note that we use the learned model of the trained agent when applying UCB.
We observe that the distribution of maximum TD errors has a slightly heavier tail on the right side when searching with a trained agent. This can be better visualized in log scale as shown in Figure 18. This tail mostly corresponds to the case of solving or being closer to solving a level in the imagination where the current policy has a high chance of not solving it. This scenario is relatively rare given that the current policy is capable of solving approximately  $95\%$  of levels, which explains the thin tails in all cases. Thus, the heavier tail suggests that the agents learn to search more effectively than MCTS within the same search budgets. This finding is consistent with previous studies on replacing MCTS with neural networks in supervised learning settings [19].
We conjecture that the difference between an RL agent and MCTS in this distribution would be much larger if the model were also trained with MCTS, given the large number of simulations required for MCTS to achieve good performance in Sokoban [19]. It is also noteworthy that even without planning rewards, the agent is still able to collect large maximum TD errors, which implies that the agent can learn to search effectively without the guidance of heuristics.
Real Action In a  $n$ -step exhaustive search, the real action is taken as the child node with the highest rollout return, while in MCTS, the real action is usually selected based on the visit count of each child node. In contrast, our trained agents appear to learn to use a combination of both.
To further investigate how real action is selected given the same search result, we used the imaginary action and reset action outputs from the trained agent and considered two different ways of selecting the real action: (i) the real action output by the trained agent, (ii) the real action output by MCTS, which is sampled based on visit counts with different temperatures  $T \in \{0.25, 1\}$  as in MuZero [5]. For each method, we calculated the portion of the real action that corresponded to the highest visit count, the second highest visit count, and so on. The result is shown in the left panel of Figure 19. Additionally, we computed the same for the mean rollout return and the maximum rollout return of child nodes, as shown in the middle and right panels of the figure.
The figure shows that most of the real action selected by the agents corresponds to the child node with the highest visit counts or rollout returns. The figures also suggest that the maximum rollout return is the most indicative of the agent's real action among the three statistics. In contrast, MCTS tends to select the real action with the highest visit count. Notably, the action with the highest visit count does not necessarily correspond to the action with the highest rollout return, as evidenced by MCTS with  $T = 0.25$ , which chooses the action with the highest visit count with over  $80\%$  probability, yet only around half of these actions correspond to the one with the highest rollout return. The simultaneous consideration of both rollout returns and visit counts by the trained agent seems more intuitively reasonable—it is unwise to follow a poor plan even if it is simulated numerous times.
Comparison with MCTS Lastly, to compare the efficiency of the trained agents across all three action sets with MCTS, we assess the solving rates of both the trained agents and MCTS. In MCTS, we utilize the model learned by the agent at the end of training, and we test a range of simulation
4This maximum TD error is proportional to the undiscounted sum of planning rewards in a stage.
Figure 17: Maximum TD error collected by different search methods, using the model learned by the agent. All statistics are averaged over 20000 stages, each starting from random initial states.
Figure 18: Maximum TD error collected by different search methods in log scale, using the model learned by the agent. All statistics are averaged over 20000 stages, each starting from random initial states.
Figure 19: Percentage of chosen real action grouped by the ranking of visit count, mean rollout return, and maximum rollout return. All statistics are averaged over 20000 stages, each starting from random initial states.
Figure 20: Final solving rate on Sokoban training levels as achieved by MCTS with varying numbers of stage length, compared to the performance of the trained agent policy and the base policy.
steps<sup>5</sup> spanning over {5,10,20,40,80,160,320,640,1280,2560}. The results are depicted in Figure 20. We observe that MCTS needs 1280 simulation steps to match the performance of a trained agent, which achieves the same results with only 20 stage length. This notable difference highlights the potential benefits of trained agents compared to handcrafted algorithms—they exhibit the capacity to conduct searches with increased efficiency, particularly within limited computational resources.
In conclusion, the analysis in this section suggests that the trained agents appear to learn a planning algorithm that differs from common handcrafted planning algorithms. These agents learn to plan deeply and reset upon encountering a poor node, while being able to search for a good plan within a limited number of model simulations. It is likely that the agents also take advantage of the model's hidden state, which contains rich information about the root and predicted states, when deciding upon the three types of actions. However, a comprehensive investigation of this particular dynamic lies beyond the scope of our current analysis.
Note that the stage length can be smaller than the number of simulations, as the former includes traversing expanded nodes while the latter does not.
G Planning Rewards
In the main paper, we use the extended real reward as the augmented reward:
$$
\tilde {r} _ {k} := \left\{ \begin{array}{l l} 0 & \text {i f} k \leq K, \\ r, & \text {i f} k = K + 1, \end{array} \right. \tag {25}
$$
where  $r$  denotes the real reward obtained from the real step. However, this sparse reward may make learning difficult. We explore adding an auxiliary reward, which we call the planning reward, to help guide the learning of imaginary and reset actions. We define the planning reward  $\tilde{r}_k^p$  on step  $k$  to be:
$$
\tilde {r} _ {k} ^ {p} := \left\{ \begin{array}{l l} c ^ {p} \left(g _ {\text {r o o t}, k} ^ {\max } - g _ {\text {r o o t}, k - 1} ^ {\max }\right), & \text {i f} k \leq K \\ 0, & \text {i f} k = K + 1, \end{array} \right. \tag {26}
$$
where  $g_{\mathrm{root},k}^{\mathrm{max}}$  denotes the maximum rollout return of the root node on step  $k$  within a stage, as defined after Equation 1, and  $c^{p} \geq 0$  is a scalar hyperparameter that controls the strength of the planning reward.
Note that the total undiscounted sum of planning rewards in a stage is:
$$
\begin{array}{l} \sum_ {k = 2} ^ {K + 1} \tilde {r} _ {k} ^ {p} = c ^ {p} \left(g _ {\text {r o o t}, K} ^ {\max } - g _ {\text {r o o t}, 0} ^ {\max }\right) (27) \\ \propto \hat {r} _ {i _ {1} ^ {*}} + \gamma^ {2} \hat {r} _ {i _ {2} ^ {*}} + \dots \gamma^ {n} \hat {r} _ {i _ {n} ^ {*}} + \gamma^ {n + 1} \hat {v} _ {i _ {n} ^ {*}} - \hat {v} _ {\text {r o o t}}, (28) \\ \end{array}
$$
where  $\{\mathrm{root},i_1^*,i_2^*,\dots,i_n^*\}$  refers to the rollout with the highest return,  $g_{\mathrm{root},i}$ , in the entire stage. As  $\hat{v}_{\mathrm{root}}$  is not influenced by imaginary and reset actions, the return of a stage is equivalent to  $\hat{r}_{i_1^*} + \gamma^2\hat{r}_{i_2^*} + \dots \gamma^n\hat{r}_{i_n^*} + \gamma^{n + 1}\hat{v}_{i_n^*}$ , the maximum rollout return. Consequently, the planning reward is a heuristic directing the agent to maximize the maximum rollout return, while not penalizing for searching paths with a low rollout return to encourage exploration in imagination.
However, this planning reward should not be given to real actions to prevent incentive misalignment. For example, collecting a large amount of planning rewards is possible by continually doing nothing in real steps but imagining a good plan. As such, the planning reward signal has to be separated from the real reward signal  $\tilde{r}_k$  instead of being added together. We also decay  $c^p$  from 1 to 0 throughout training to further ensure that the imaginary action is aligned with the real reward function. Combined, the augmented reward contains both the real reward and the planning reward:  $r^{\mathrm{aug}}\coloneqq (\tilde{r},\tilde{r}^p)$ .
The incorporation of planning rewards necessitates certain modifications to the actor-critic algorithm in our experiments. Specifically, we separately learn the value of these two distinct rewards when training the critic, resulting in two different TD errors corresponding to the real and the planning rewards. In training the critic for the planning reward, we treat each stage as a single episode. This approach ensures that the planning reward incentives the agent to optimize the maximum rollout returns exclusively for the current stage. Furthermore, when using the TD errors to train the actor, we impose a mask on the planning rewards' TD error, setting it to zero during a real step.
Experiments in Appendix E show that planning rewards only provided a minimal increase in initial learning speed, indicating that the real rewards are sufficient for learning to plan. We conjecture that sparse real rewards are not a problem, as the TD errors induced by real rewards may not be sparse. For example, if the critic learns to use the maximum rollout return  $g_{\mathrm{root},k}^{\mathrm{max}}$  as the estimated value, which is provided in the tree representation, then the TD errors in an imaginary step equal the planning rewards, except for the adjustment due to the discount rate.
H Relationship with Generalized Policy Iteration and Meta-Reinforcement Learning
H.1 Generalized Policy Iteration Algorithm
Applying an RL algorithm to the Thinker-augmented MDP can be viewed as an instance of a generalized policy iteration algorithm [29, 30]. Policy improvements are carried out by an agent that conducts searches within the model. Given the base policy, these searches yield an improved policy. Simultaneously, the policy evaluation is done by the learning of the value-policy network. This involves training the network on trajectories generated by the improved policy, so as to learn the values of this improved policy and project the improved policy back to the base policy. A noteworthy divergence from a typical generalized policy iteration algorithm is the learning of the policy improvement operator, instead of employing a handcrafted one such as MCTS.
Given this perspective, a crucial question arises: how effective is a learned policy improvement operator in comparison to a handcrafted one? To address this question, we evaluate the efficiency of different policy improvement operators, such as MCTS and an RL agent, during a single step of policy improvement.
Specifically, we consider the Thinker-augmented MDP with a fixed base policy  $\hat{\pi}$  and its value  $\hat{v}$ . We do not update the base policy that is described in Section 4.4. The base policy is a given policy that performs slightly better than a random policy. For simplicity, we employ the true environment dynamics, represented as  $(P,R)$ , as the model in the augmented MDP and do not impose a limit on the maximum search depth, given that we are using the true environmental dynamics. Sokoban serves as the real environment in our experiments.
We experiment with three different types of policy improvement operators: (i)  $n$ -step exhaustive search, (ii) MCTS, and (iii) RL agents. In (iii), we trained the agent from scratch in the Thinker-augmented MDP. Only the tree representation  $u_{\mathrm{tree}}$  is provided as input to the agent, so the agent cannot learn to act based on the real states directly. The results are shown in Figure 21. Due to the low solving rate, we report the return of an episode here for better illustration. The maximum return in Sokoban is around 14, and a return of 4 gives at most  $30\%$  solving rate.
From the figure, we see that the base policy achieves a return of 0.12, and all policy improvement operators are able to improve on this policy and achieve a better return. We observe that the policy improvement gap for RL agents widens as the value of  $K$  increases. With  $K = 10$ , the agent's performance is similar to that of a 3-step exhaustive search, which requires 155 model simulations. Increasing  $K$  to 20, the performance is slightly worse than a 4-step exhaustive search, which requires 780 model simulations. At  $K = 40$ , the agent performance surpasses a 5-step exhaustive search, which requires 3905 model simulations. In addition, the agent at  $K = 10$  also shows comparable performance to MCTS that uses 60 model simulations. It is also worth noting that the number of model simulations used by the agent is typically less than  $K$ , given that imaginary steps also include the traversal through expanded nodes.
The results here demonstrate that increasing  $K$  to 40 can still enhance performance, whereas in the main paper, the performance improvement gained by increasing  $K$  above 10 is insignificant. We conjecture that once the value-policy network is trained, the model can learn to predict the values of leaf nodes more accurately, reflecting the value of the current policy instead of the base policy. This eliminates the need for a deep search to evaluate a rollout more precisely, rendering a large  $K$  unnecessary.
The results also underscore the importance of a learned base policy. The agent achieves a maximum solving rate of less than  $30\%$ , whereas a  $95\%$  solving rate is attainable when the base policy undergoes training as well. This outcome is intuitive, considering that the agent gives an improved version of the base policy, and thus the agent policy will not significantly outperform it. If a fixed poor base policy is used, the agent policy will not perform well.
Nonetheless, these findings suggest that an agent can learn to perform policy improvement and achieve similar performance to handcrafted algorithms while requiring significantly fewer simulations. However, it requires a moderate number of transitions for an agent to learn this policy improvement operator.
Figure 21: Comparison of different policy improvement operators applied to a fixed base policy in Sokoban. For trained agents, the running average of returns from the last 2000 episodes is shown, whereas for others, an average return from 2000 episodes is shown.
H.2 Meta-Reinforcement Learning
In this section, we explore the relationship between meta-reinforcement learning (meta-RL) algorithms [48] and the Thinker algorithm. Meta-RL is a method designed to learn how to perform reinforcement learning. Its objective is to develop a policy that can quickly adapt to a variety of new tasks derived from a distribution of tasks, using minimal data. While meta-RL is often used in scenarios involving a task distribution, it is also applicable to single-task scenarios [49, 50], which is our primary focus here.
In the meta-RL framework, a trial typically comprises  $L$  episodes. During a trial, the agents initially encounter  $K$  episodes—termed burn-in episodes. Notably, the return within these episodes does not contribute to the main objective. Instead, the primary goal is to maximize the return of the remaining  $L - K$  episodes, known as evaluation episodes. Consequently, the purpose of a meta-RL algorithm is to use the knowledge gained throughout the entire trial to maximize returns during the evaluation episodes. One common class of meta-RL algorithm applies an RL algorithm with an RNN that can access previous actions and rewards, and the hidden state is not reset across episodes. As such, the knowledge of the trial is embedded in the RNN's hidden states [27, 28].
If the model perfectly represents the true environment dynamics, we can align the Thinker-augmented MDP with meta-RL by considering imaginary rollouts as (truncated) burn-in episodes, and real transitions as the evaluation episode in meta-RL. As agents only aim to maximize the return in real transitions, the objectives of meta-RL and RL agents on the Thinker-augmented MDP coincide. Moreover, the purpose of burn-in episodes in meta-RL or imaginary rollouts in planning is the same – to guide better actions in the evaluation episodes or real transitions, possibly through risk-taking behavior in them.
However, unlike conventional meta-RL, which uses a fixed number of burn-in episodes followed by evaluation episodes, the Thinker-augmented MDP interleaves the burn-in episodes within a single evaluation episode. In other words, for each step in the evaluation episode with state  $s_t$ , we allow agents to have a few burn-in episodes starting from  $s_t$ .
In meta-RL, it is generally assumed that the burn-in episodes and evaluation episodes are drawn from the same MDP, but this assumption does not hold when the burn-in episodes are generated by a learned model. But if the MDP underlying the burn-in episodes,  $\hat{\mathcal{M}}$ , is similar to that of the evaluation episode,  $\mathcal{M}$ , then the meta-RL algorithm may also learn how to utilize the collected data while accounting for the difference. This is analogous to how the RL agent in the Thinker-augmented MDP learns to use the imaginary rollout for better action selection, despite the possible inaccuracies of the model.
In summary, there is a close relationship between learning to plan and learning to learn. The former employs a model on each step to decide better actions at that state, while the latter uses episodes within a trial to determine better actions in the evaluation episodes. This highlights the converging ideas in meta-learning and planning, both aiming to make efficient use of available information for decision-making.
I Miscellaneous
I.1 Related Work on the Dual Network
The dual network is similar to the model in MuZero [5], except for (i) the use of the dual network architecture and (ii) the prediction of future states trained by feature loss. Improvements to MuZero's model have been proposed. EfficientZero [51] deviates from MuZero by adding SimSiam loss [52] to match future representations. We briefly experimented with it, but the result was only slightly better than that of MuZero's model, possibly because this loss works better in a limited data regime where EfficientZero was designed and evaluated. [43] also proposes to predict future states as an auxiliary loss but does not use the dual network architecture. In our experiments, this addition led to better performance, as shown in Figure 4, but the result was still significantly worse than that of the dual network.
Most other types of models used in model-based RL do not predict values and policies, as future values and policies are often not required in the algorithm. Among these models, various loss functions have been proposed for training the model to acquire meaningful representations for the task. [53] employs a VAE [54] to encode the state while Dreamer [33] utilizes a sequential VAE. Given the extensive body of research in model-based RL, we direct readers to the review [55] for a detailed exploration of this field.
I.2 Planning Capacity