SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Until now, we assume that we are provided (i) the model  $(\hat{P},\hat{R})$  and (ii) a base policy  $\hat{\pi}$  and its value  $\hat{v}$ . To remove both assumptions, we propose employing a single, unified model that assumes the roles of  $\hat{P},\hat{R},\hat{\pi}$ , and  $\hat{v}$ , learning directly by fitting the real transitions. Specifically, we introduce a novel model, the dual network, designed for this purpose.
The dual network consists of two RNNs which we refer to as the state-reward network and the value-policy network. The state-reward network's input consists of the root state  $s_t$  and an action sequence  $a_{t+1:t+L}$ , and the network predicts future states  $\hat{s}_{t+1:t+L}$ , rewards  $\hat{r}_{t+1:t+L}$ , and termination probability  $\hat{p}_{t+1:t+L}^d$ . Meanwhile, the value-policy network takes both the state-reward network's input and its predicted states  $\hat{s}_{t+1:t+L}$  as inputs, and the network predicts future policies  $\hat{\pi}_{t:t+L}$  and values  $\hat{v}_{t:t+L}$ . We can also view the dual network as an RNN with  $h = (h^{sr}, h^{vp})$  being the hidden state, where  $h^{sr}$  and  $h^{vp}$  denote the hidden states of the two sub-RNNs. The relevant statistics required in the augmented state can be computed by unrolling the dual network using the root node and the action sequence leading to the current node.
Figure 4: Illustration of the dual network's architecture, composed of the state-reward and the value-policy network. See Equations 2 and 3 for the loss functions and Appendix B for training details.
To train the model, we first store all real transitions  $(s_t, a_t, r_{t+1}, d_{t+1})$  in a replay buffer, which corresponds to the state, action, reward, and termination indicator. We then sample a sequence of real transitions with length  $L + 1$  from the buffer. The two sub-RNNs in the dual network are trained separately. For the state-reward network, the following loss is used:
$$
\mathcal {L} ^ {s r} := \sum_ {l = 1} ^ {L} \left(c ^ {r} \left(\hat {r} _ {t + l} - r _ {t + l}\right) ^ {2} + c ^ {d} d _ {t + l} \log \hat {p} _ {t + l} ^ {d} + c ^ {s} \mathcal {L} ^ {\text {s t a t e}} \left(\hat {s} _ {t + l}, s _ {t + l}\right)\right), \tag {2}
$$
where  $c^r, c^d, c^s \geq 0$  are hyperparameters that modulate the loss strength. While a straightforward choice for the state prediction loss metric,  $\mathcal{L}^{\mathrm{state}}$ , is the L2 loss - defined as  $\mathcal{L}^{\mathrm{state}}(s', s) = ||s' - s||_2^2 -$  this approach could cause the model to allocate capacity towards predicting non-essential features. To address this issue, we suggest using the feature loss instead of the L2 loss. Specifically,  $\mathcal{L}^{\mathrm{state}}(s', s) = ||g(s') - g(s)||_2^2$ , where  $g$  represents the encoding function of the value-policy network. As this encoding function concentrates on features relevant to action, the feature loss encourages the model to prioritize predicting the most crucial features.
For the value-policy network, the following loss is used:
$$
\mathcal {L} ^ {v p} := \sum_ {l = 0} ^ {L} \left(c ^ {v} \left(\hat {v} _ {t + l} - \hat {v} _ {t + l} ^ {\text {t a r g e t}}\right) ^ {2} + c ^ {\pi} \text {o n e} _ {-} \operatorname {h o t} \left(a _ {t + l}\right) ^ {T} \log \hat {\pi} _ {t + l}\right), \tag {3}
$$
where  $c^v, c^\pi \geq 0$  are hyperparameters and  $\hat{v}_{t + l}^{\mathrm{target}}$  denotes the multi-step return. One may also use the real action distribution  $\pi_t$  place of one\_hot  $(a_{t + l})$  for a better learning signal, though this requires passing the real action distribution instead of only the actions to the augmented MDP.
The design choice of the dual network aims to improve data efficiency while also ensuring inaccurately predicted states do not negatively affect the learning of the base policy and values. If the predicted states are inaccurate, the value-policy network, being an RNN, can learn to ignore these states, relying solely on the root state and action sequence. Conversely, if the predicted states are accurate, the value-policy network can effectively learn to make predictions based on these states. Consider, for instance, the task of predicting the value after executing several actions in Sokoban. If the predicted state suggests the level is close to completion, the predicted value should naturally be higher. Although, in theory, a standalone RNN might autonomously learn this intermediary step, in practice, it would likely require much more data. We explore the relationship of this dual network with models from other model-based RL algorithms in Appendix I.
With the dual network, the augmented state can also include the hidden state of the model at either the current node,  $h_{\mathrm{cur}}$ , or the root node,  $h_{\mathrm{root}}$ . These hidden states, usually in much lower dimensions, might be easier to process than the real state  $s$  or the current node's predicted state  $\hat{s}_{\mathrm{cur}}$ . In our experiments, we use  $(u_{\mathrm{tree}}, h_{\mathrm{cur}})$  as the augmented state, based on the intuition that the current hidden state is likely the most informative among the others. This completes the specification of the Thinker algorithm, which takes an MDP as input and transforms it into another MDP.
${}^{2}$  Strictly speaking,using  $\left( {{u}_{\text{tree }},{h}_{\text{cur }}}\right)$  as the augmented state leads to a POMDP instead of an MDP. One can recover an MDP by omitting the auxiliary hints in  ${u}_{\text{tree }}$  and adding  ${s}_{\text{root }}$  to the augmented state,but both do not help performance in our experiments.
Figure 5: Running average solving rate over the last 200 episodes in Sokoban, compared with other baselines.
For all our experiments, we train a standard actor-critic algorithm, specifically the IMPALA algorithm [31], on the Thinker-augmented MDP. Although other RL algorithms could be employed, we opted for the IMPALA due to its computational efficiency and simplicity. The actor-critic's network uses an RNN for encoding the tree representation and a convolutional neural network for encoding the model's hidden state. These encoded representations are concatenated, then passed through a linear layer to generate actions and predicted values. We set the stage length,  $K$ , to 20, and the maximum search depth or model unroll length,  $L$ , to 5. More experiment details can be found in Appendix D.
5.1 Sokoban
We selected the game of Sokoban [32, 18], a classic puzzle problem, as our primary testing environment due to its inherent complexity and requirement for extensive planning. In Sokoban, the objective is to push all four boxes onto the designated red target squares, as depicted in Figure 1. We used the unfiltered dataset comprising 900,000 Sokoban levels from [14]. We train our algorithm for 5e7 frames, while baselines in the literature typically train for at least 5e8 frames. Despite this discrepancy, we present learning curves plotted against frame count for direct comparison.
The learning curve for the actor-critic algorithm applied to the Thinker-augmented MDP is depicted in Fig. 5, with results averaged across five seeds. We include seven baselines: DRC [14], Dreamer-v3 [33], MuZero [5], I2A [18], ATreeC [17], VIN [20], and IMPALA with ResNet [31]. In Sokoban, DRC stands as the current state-of-the-art algorithm, outperforming others by a significant margin, likely due to its well-suited prior for the game. The 'DRC (original paper)' result is sourced directly from the paper, whereas 'DRC' denotes our replicated results.
Thinker surpasses the other baselines in terms of performance in the first 5e7 frames. At 2e7 frames, Thinker solves  $88\%$  of the levels, while DRC (original paper) solves  $80\%$  and the other baselines solve at most  $21\%$ . At 5e7 frames, Thinker solves  $95\%$  of the levels, while DRC (original paper) solves  $93\%$  and the other baselines solve at most  $45\%$ . These results underscore the enhanced planning capabilities afforded by the Thinker-augmented MDP to an RL agent.
To understand the benefits of planning, we evaluate a policy trained with  $K = 20$  on the Thinker-Augmented MDP, by testing its performance with different  $K$  values in the same augmented environment during the testing phase. This variability allows us to control the degree of planning and the result is shown in Figure 6. The figure also depicts the performance of the base policy, which can be regarded as a distilled policy devoid of planning. We observe a significant performance improvement attributable to planning, even at the end of training.
The agent's behavior, visualized in Figure 1, illustrates that it learns to use the model for selecting better actions. Interestingly, it appears that the agents learn a planning algorithm that diverges from traditional  $n$ -step-exhaustive search and MCTS. For instance, the agent chooses real actions based on
Figure 6: Final solving rate on Sokoban test levels with varying stage length  $K$  during test time but trained with  $K = 20$ .
Figure 7: Running average solving rate over the last 200 episodes in Sokoban with varying stage length  $K$ .
Figure 8: Running average solving rate over the last 200 episodes in Sokoban with different models.
Figure 9: Median human-normalized score on Atari 2600, computed by the running average episode returns over the last 500 episodes.
both visit counts and rollout returns, and the agent learns to reset upon entering an irreversible state, contrasting the MCTS strategy of resetting at a leaf node. Further analysis on the trained agent's behavior can be found in Appendix F.
We conducted an ablation analysis to examine the impact of planning duration, varying the stage length  $K$  across  $\{1,2,5,10,20\}$  during both training and testing. When  $K = 1$ , the agent does not interact with the model at all. The outcomes of this ablation analysis are presented in Figure 7. The results suggest that  $K = 10$  already gives the optimal performance. This is in stark contrast to the 2000 simulations required by MCTS to achieve good performance in Sokoban [19].
As a comparison, we train the same actor-critic algorithm on the raw MDP, employing a similar network architecture used by the model to encode the raw state. Surprisingly, the result for  $K = 1$  is worse than the raw MDP. A detailed explanation of the reasons behind this phenomenon can be found in Appendix D. However, employing a marginally larger  $K$ , such as 2, greatly improves the performance on the Thinker-augmented MDP, surpassing that of the raw MDP.
We use the dual network as the world model, trained with feature loss in state prediction. We consider several variants of the model, including (i) a dual network that uses L2 loss in state prediction, (ii) a dual network that has no loss in state prediction, (iii) a single network that uses L2 loss in state prediction, (iv) a single network that does not predict state. The single network refers to a single RNN that predicts the relevant quantities, and (iv) is equivalent to MuZero's model. The results of these four variants are shown in Figure 8. We observe that both the dual network architecture and the state loss are critical to performance. However, the use of feature loss instead of L2 loss gives marginal performance improvement, likely because the L2 loss' prior that each pixel is equally important suits well in Sokoban. Nonetheless, this prior may not be suitable for other games like Breakout.
Further ablation studies in Appendix E offer insights into the various components of Thinker, including the role of hints in tree representation, the agent's memory, base policy, planning reward, etc. For example, either the hints or an RNN agent is sufficient to achieve good performance, as the hints summarize the rollouts, and an agent with memory can learn to summarize rollouts on its own. However, if both hints and memory are omitted, the agent fails to learn to plan, as it cannot recall past rollouts. In addition, if the base policy and its value are also omitted, then an RNN agent still fails to learn, suggesting the necessity of a heuristic in learning to plan.
5.2 Atari 2600
Finally, we test our algorithm on the Atari 2600 benchmark [34] using a 200M frames setting. The performance of the actor-critic algorithm on both the Thinker-augmented MDP and raw MDP is evaluated using the same hyperparameters as in the Sokoban experiment, but with a deeper neural network for the model and an increased discount rate from 0.97 to 0.99.
The learning curve, measured in the median human-normalized score across 57 games, is displayed in Figure 9. We evaluate the agent in the no-ops starts regime at the end of the training, and the learning curve is shown in Fig 9. The median (mean) normalized score for the actor-critic algorithm applied on the Thinker-augmented MDP and the raw MDP are  $261\%$  (1372%) and  $102\%$  (514%), respectively, underscoring the advantages of the Thinker-augmented MDP over the raw MDP. For context, the median (mean) normalized score of Rainbow [35], a robust baseline, is  $223\%$  (868%).
The advantages of the Thinker algorithm are particularly pronounced in certain types of games, especially shooting games such as chopper command, seaquest, and phoenix. It is plausible that the agents benefit from predicting the trajectories of fast-moving objects in these shooting games. However, in some games, particularly those where a single action does not produce significant effects, the Thinker algorithm appears to offer no advantages. For instance, in the game qbert, the decision to move to a different box is made only once every five or more actions. This effectively transforms any 5-step plan into a single-step plan, diminishing the algorithm's benefits. Detailed experiment results on the Atari 2600 are available in Appendix D.
Many interesting behaviors learned by the agent are visualized in https://youtu.be/0AfZh5SR7Fk. We also observe that the predicted frames in the video are of high quality, showing that the feature loss can lead to the emergence of interpretable and visualizable predicted frames, despite the absence of an explicit loss in the state space. We attribute this capability to using a convolutional mapping  $g$  in the feature loss, as convolutions can preserve local features.
6 Future Work and Conclusion
While the Thinker algorithm offers many advantages, it also presents several limitations that provide potential areas for future research. Firstly, the Thinker algorithm carries a large computational cost. A single step in the original MDP corresponds to  $K$  steps in the augmented MDP, in addition to the overhead of model training. Secondly, the algorithm currently enforces rigid planning, requiring the agent to roll out from the root state and restricts it to planning for a fixed number of steps. Thirdly, the algorithm employs a deterministic model to facilitate tree traversal, which may not work well for a stochastic environment.
Future work could focus on addressing these limitations, such as enabling more flexible planning steps and integrating uncertainty into the model. Additionally, exploring the application of the algorithm in multi-environment settings is a promising direction. We hypothesize that the advantages of a learned planning algorithm over a handcrafted one will become more pronounced in such scenarios. Exploring how other RL algorithms perform within the Thinker-augmented MDP is an additional direction for future work.
The history of machine learning research tells us that learned approaches often prevail over handcrafted ones. This transition becomes especially pronounced when large amounts of data and computational power are available. In the same vein, we surmise that learned planning algorithms will eventually surpass handcrafted planning algorithms in the future.
This work was performed using resources provided by the Cambridge Service for Data Driven Discovery (CSD3) operated by the University of Cambridge Research Computing Service (www.csd3.cam.ac.uk), provided by Dell EMC and Intel using Tier-2 funding from the Engineering and Physical Sciences Research Council (capital grant EP/T022159/1), and DiRAC funding from the Science and Technology Facilities Council (www.dirac.ac.uk).
We note that author Ivan Anokhin has since moved to the Mila, Université de Montréal, after completing the primary research reflected in this paper at the University of Cambridge.
[1] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
[2] Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pages 216-224. Elsevier, 1990.
[3] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. Advances in neural information processing systems, 28, 2015.
[4] Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11), pages 465–472, 2011.
[5] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609, 2020.
[6] Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In Computers and Games: 5th International Conference, CG 2006, Turin, Italy, May 29-31, 2006. Revised Papers 5, pages 72-83. Springer, 2007.
[7] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-359, 2017.
[8] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc 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.
[9] Fan-Ming Luo, Tian Xu, Hang Lai, Xiong-Hui Chen, Weinan Zhang, and Yang Yu. A survey on model-based reinforcement learning. arXiv preprint arXiv:2206.09328, 2022.
[10] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540, 2016.
[11] György Buzsáki, Adrien Peyrache, and John Kubie. Emergence of cognition from action. In Cold Spring Harbor symposia on quantitative biology, volume 79, pages 41-50. Cold Spring Harbor Laboratory Press, 2014.
[12] MD György Buzsáki. The brain from inside out. Oxford University Press, 2019.
[13] Henry H Yin and Barbara J Knowlton. The role of the basal ganglia in habit formation. Nature Reviews Neuroscience, 7(6):464-476, 2006.
[14] 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. In International Conference on Machine Learning, pages 2464-2473. PMLR, 2019.
[49] Tom Zahavy, Zhongwen Xu, Vivek Veeriah, Matteo Hessel, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh. A self-tuning actor-critic algorithm. Advances in neural information processing systems, 33:20913-20924, 2020.
[50] Zhongwen Xu, Hado P van Hasselt, Matteo Hessel, Junhyuk Oh, Satinder Singh, and David Silver. Meta-gradient reinforcement learning with an objective discovered online. Advances in Neural Information Processing Systems, 33:15254-15264, 2020.
[51] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34:25476-25488, 2021.
[52] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 15750-15758, 2021.
[53] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018.
[54] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
[55] Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based reinforcement learning: A survey. Foundations and Trends® in Machine Learning, 16(1):1-118, 2023.
[56] Russell Stuart and Norvig Peter. Artificial intelligence a modern approach third edition, 2010.
Algorithm 1 Transition in Thinker-Augmented MDP
1: Input: real MDP: env, model:  $\hat{P},\hat{R}$  , base policy and its value function:  $\hat{\pi},\hat{v}_{\cdot}$    
2: Algorithm Parameter: stage length  $K\geq 1$  maximum search depth  $L\geq 1$    
3: Initialize  $k\gets K$  Ensures that the first step is a real step   
4: procedure STEP(a,  $\tilde{a}^r$  )  $\tilde{a}$  denotes imaginary/real action,  $\tilde{a}^r$  denotes reset action   
5: if  $k = K$  then Real step   
6:  $k,l\gets 1,0$  Initialize step number and search depth   
7:  $r,s,d\gets env.\mathrm{step}(\tilde{a})$  Compute reward, state and termination for the real MDP   
8:  $s_{\mathrm{root}}\gets s$  Record the root node   
9:  $s_{\mathrm{cur}}\gets s_{\mathrm{root}}$  Set the current node to be the root node   
10:  $\hat{r}_{\mathrm{root}}\gets 0$  No predicted rewards for the root node   
11: else Imaginary step   
12:  $k,l\gets k + 1,l + 1$    
13: if  $l\geq L$  then Check if search depth exceeds maximum search depth  $L$    
14:  $\tilde{a}^r\gets 1$  Override reset   
15: end if   
16:  $\hat{r}_{\mathrm{cur}}\gets \hat{R} (s_{\mathrm{cur}},\tilde{a})$  Compute the predicted reward   
17:  $s_{\mathrm{cur}}\gets \hat{P} (s_{\mathrm{cur}},\tilde{a})$  Compute the new current node by unrolling the model   
18:  $r\gets 0$  No rewards for imaginary steps   
19:  $d\gets False$  No termination for imaginary steps   
20: end if   
21:  $\hat{\pi}_{\mathrm{cur}},\hat{v}_{\mathrm{cur}}\gets \hat{\pi} (s_{\mathrm{cur}}),\hat{v} (s_{\mathrm{cur}})$  Compute the base policy and its value   
22: Compute  $s^{\mathrm{aug}}$  using  $s_i,\hat{r}_i,\hat{\pi}_i,\hat{v}_i$  where  $i\in \{\mathrm{root,cur}\}$  as described in Section A.1   
23: if  $\tilde{a}^r = 1$  then Reset to the root node if reset is triggered   
24:  $l\gets 0$    
25:  $s_{\mathrm{cur}}\gets s_{\mathrm{root}}$    
26: end if   
27: return  $r,s^{\mathrm{aug}},d$  Return rewards, state, and termination for the augmented MDP   
28: end procedure
Algorithm 1 implements the step function of the Thinker-augmented MDP, assuming we have access to a given model  $(\hat{P},\hat{R})$ , a base policy  $\hat{\pi}$ , and its value  $\hat{v}$ . We assume  $\hat{P}$  is deterministic and use  $s' = \hat{P}(s,a)$  to denote the predicted next state  $s'$  from state  $s$  and action  $a$ .
To compute the hints (i.e. mean rollout return, max rollout return, and visit counts) in the tree representation, we need to maintain a tree with each node storing the predicted reward  $\hat{r}_i$ , value  $\hat{v}_i$ , and visit count  $n_i$ , but the details of updating the tree are omitted in the pseudocode for clarity.
It is important to note that since  $\hat{P}$  is deterministic, one can store the computed node in a buffer to avoid re-computing it when the agent traverses the same node again within the same stage.
In the full Thinker-augmented MDP, both the model and the base policy with its values are learned by an RNN model. The code only needs slight modifications as follows: