SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The hyperparameters used in all our experiments are shown in Table 4. We tune our hyperparameters exclusively on Sokoban based on the final solving rate. The specific hyperparameters we tune are: learning rates, model batch size, model loss scaling, maximum search depth or model unroll length, planning reward scaling, actor-critic clip global gradient norm, and actor-critic unroll length. In terms of preprocessing for Atari, we adhere to the same procedure as outlined in the IMPALA paper [31], with the exception that we do not apply grayscaleing.
For the actor-critic baseline implemented on the raw MDP, we maintain the same hyperparameters, with two exceptions: (i) we adjust the actor-critic unroll length to 20 to ensure the same number of real transitions in an unrolled trajectory; (ii) we reduce the learning rate by half to 3e-4, after observing that the original learning rate led to unstable learning. We found that this set of hyperparameters provides optimal performance for the baseline in Sokoban, based on the final solving rate. Regarding the network architecture, we employ an encoder that mirrors the model's encoder architecture in Thinker, followed by a CNN that mirrors the actor-critic's CNN in Thinker.
D.2 Detailed Results on Sokoban
The performance of Thinker and the baselines on Sokoban can be found in Table 1, corresponding to the learning curves shown in Figure 5. We use our own implementation for DRC [14] and open-source implementation for Dreamer-v3 [33], with the result being averaged over three seeds. MuZero's result is taken from [43]. Results of DRC (original paper), I2A, ATreeC, VIN and IMPALA with Resnet are taken from [14].
As Thinker augments the MDP by adding  $K - 1$  steps for every step in the real MDP, there is an increased computational cost due to the extra K-1 imaginary steps, in addition to the added computational cost of training the model. On our workstation equipped with two A100s, the training durations for the raw MDP, DRC, and Thinker on Sokoban are around 1, 2, and 7 days, respectively.
However, when benchmarked against other model-based RL baselines, Thinker's training time is similar or even shorter. For instance, Dreamer-v3 (the open-sourced implementation) requires around 8 days for a single Sokoban run on our hardware, primarily due to its high replay ratio. MuZero, when conducting a standard number of simulations (e.g., 100), requires 10 training days. Notably, despite the Thinker's additional training overhead for RL agents, the algorithm compensates by requiring significantly fewer simulations. Furthermore, if one replaces the RNN in the actor-critic network with an MLP (as detailed in Appendix E), without compromising the performance, the computational time for a single Thinker run can be reduced to 3 days.
In Figure 7, we observed that the result for  $K = 1$ , which does not use any imaginary transitions, underperforms compared to the result on the raw MDP. Initially, we hypothesized that this was because the model's hidden state, rather than the real state, was included in the augmented state. Although this hidden state might be more effective for predicting future states and rewards, it potentially complicates the learning of an optimal policy. However, subsequent experiments revealed that substituting the real state for the model's hidden state did not impact performance. This outcome suggests that the additional information given by the tree representation is, in fact, detrimental to performance. We speculate that this counterintuitive effect is due to the tree representation providing information, such as value estimation, that are so insightful that the critic in the actor-critic algorithm overly relies on this estimation, thereby neglecting the real state and impeding the learning of a useful state representation for the actor. To investigate this speculation, we experimented with masking the tree representation at a random rate (with a probability  $p = 0.20$ ), which restored the performance to the level of the raw MDP. This propensity of the agent to depend excessively on a single type of input shown in these experiments points to a potential avenue for subsequent research.
D.3 Detailed Results on Atari 2600
The performance of Thinker and the baselines on Atari 2600 can be found in Table 2, corresponding to the learning curves shown in Figure 9. We include two additional baselines, Unreal [44] and LASER [45], for comparison. The results of the baselines are taken from their original papers.
Table 3: Atari scores after 200M frames of training for the five selected games with different seeds. Up to 30 no-ops at the beginning of each episode.
Seed 1	Seed 2	Seed 3	Mean
Battle Zone	78740.00	31110.00	34460.00	48103.33
Double Dunk	12.20	14.54	21.54	16.09
Name This Game	34079.60	33400.20	31775.80	33085.20
Phoenix	543987.70	535685.10	578485.30	552719.37
Qbert	30734.50	36810.25	37003.25	34849.33
We note that Thinker's performance on Atari 2600 lags behind that of state-of-the-art algorithms. There are several possible reasons: (i) we did not tune hyper-parameters on Atari 2600; (ii) learning to plan may require a much larger number of frames to show significance, similar to how MuZero requires 20 billion frames. As the goal of the experiments on Atari 2600 is to demonstrate the general advantage of the Thinker-augmented MDP over the raw MDP, we leave the problem of mastering Atari 2600 using Thinker for future study. In principle, more advanced RL algorithms, including the baselines, could be applied to the Thinker-augmented MDP.
The final scores and the corresponding learning curves for each individual Atari game are presented in Table 5 and Figure 11, respectively. The Rainbow results are taken from the original paper [35], where the reported score is obtained from the best agent snapshots throughout training. The learning curves of Rainbow in Figure 11 are also taken from the original paper.
To validate the reproducibility of our results, we carried out multiple runs of five selected games with three different seeds. These games include battle zone, double dank, name this game, phoenix, and qbert. The selection of these games was influenced by the recommendations of [46], which suggests that these five games are pivotal in predicting the median human-normalized score. The final scores for these five games are displayed in Table 3, and their individual learning curves can be found in Figure 10. The results for these five games, as presented in other figures and tables in this paper, are derived from the average of these three seed values.
D.4 Discussion on Atari 2600's results
As stated in the main paper, Thinker demonstrates superior performance in shooting games like chopper command, seaquest, phoenix, beam rider, and space invader. A more general observation is that Thinker excels in environments with fast-moving objects, such as breakout. As depicted in https://youtu.be/0AfZh5SR7Fk, Thinker's model is capable of predicting the trajectories of these swift objects with near-perfect precision. We hypothesize that such predictions, essential for planning future actions, pose a challenge for an agent in the raw MDP to learn. For example, in Breakout, the video illustrates the agent exploring different methods of deflecting the ball at various angles, sometimes missing the ball in its imagination but performing adeptly in the
Table 1: Solving rate on Sokoban.
At 2e7 frames	At 5e7 frames
Thinker	88%	95%
DRC [14]	68%	92%
DRC (original paper) [14]	80%	93%
Dreamer-v3 [33]	12%	32%
MuZero [5]	0%	21%
IMPALA with ResNet [31]	14%	45%
I2A [18]	21%	43%
ATreeC [17]	1%	3%
VIN [20]	12%	21%
Table 2: Final human-normalized score on Atari 2600.
Median	Mean
Thinker (Agent policy)	261%	1372%
Thinker (Base policy)	160%	807%
Raw MDP	102%	514%
Rainbow [35]	223%	868%
UNREAL [44]	250%	880%
LASER [45]	431%	n.a.
Figure 10: Learning curves for Thinker in the five selected games with different seeds. Each curve is computed as the running average of episode returns over the last 500 episodes.
actual game. Upon identifying a sequence of actions that might result in missing the ball, the agent can opt to avoid it. This adaptive decision-making and the significant performance gap between the agent policy and the base policy in these games underscore the benefits of planning in games with fast-moving objects.
It is important to highlight that applying a direct L2 loss on future states instead of a feature loss can undermine these advantages. For example, we observed that in Breakout, a model trained with direct L2 loss is able to predict frames of high quality, except that the ball, arguably the most important object in the game, is missing in the predicted frame. Given that these fast-moving objects typically inhabit only a few pixels and are challenging to predict, the model is prone to neglect these elements when trained using a direct L2 loss. In contrast, by using feature loss, these important objects, such as the ball in Breakout, are usually learned early in training.
We also observe that the agent's performance falls short in more demanding exploration games, such as montezuma revenge and venture. This is because the Thinker-augmented MDP lacks any mechanism to encourage exploration in the real environment, which results in challenges when dealing with these complex exploration games, much like an actor-critic applied directly to a raw MDP. Future research could explore methods to utilize Thinker to foster exploration, such as using model prediction loss as a curiosity reward. As exploration is a well-studied area in RL [47], many existing exploration methods could also be directly applied to the Thinker-augmented MDP.
Thinker also performs poorly in games where each action carries little impact, such as in qbert. Games in which the agent moves slowly, such as bank heist and ms pacman, also share this characteristic. These games are in stark contrast to games like Sokoban, where a single action can move the agent to a different grid. These games highlight a fundamental limitation of the Thinker-augmented MDP that cannot be avoided by switching RL algorithms—planning is conducted in the primitive action space rather than the abstract action space, thereby necessitating the significance of each primitive action. Future research could investigate modifications to Thinker that would allow model interaction through abstract actions, thereby allowing high-level and long-term planning.
Interestingly, the actor-critic applied to the raw MDP occasionally outperforms the one applied to the Thinker-augmented MDP. This could be attributed to the fact that we provide the agent with the model's hidden state instead of the true or predicted frames. It is conceivable that the actual frame is simpler to learn from, given that the model learns the encoding of frames by minimizing supervised learning loss rather than maximizing the agent's returns. Future studies could consider allowing the agent to directly observe the true and predicted frame, thus ensuring that the agent possesses a strictly greater capability than an agent on the raw MDP. Another factor could be the impact of unoptimized hyperparameters. As we increased the network's depth while maintaining the same learning rate during the transition from Sokoban to the Atari environment, the training process became more unstable, as demonstrated in some learning curves. It is plausible that the poor performance in certain games stems from these unoptimized hyperparameters.
Finally, as evidenced by the individual learning curves illustrated in Figure 11, and the median human-normalized score in Figure 9, we observe that the slope of the learning curve is generally steeper at the end of training compared to other baselines. We posit that mastering the use of a model represents an additional skill to be acquired, as opposed to an agent acting on a raw MDP. This necessitates extended training to learn this skill, however, once this skill is acquired, the final performance generally surpasses that of an agent on the raw MDP. As such, future work could experiment with a larger training budget than 200 million frames.
Table 4: Hyperparameters used in experiments.
Parameter	Value
Thinker-augmented MDP
Stage length K	20
Maximum search depth L	5
Augmented discount rate γ	γ1/K
Planning reward scaling cp	Anneal linearly from 1 to 0
Model
Learning rate	Anneal linearly from 1e-4 to 0
Optimizer	Adam
Adam beta	(0.9, 0.999)
Adam epsilon	1e-8
Replay buffer size	200,000
Minimum replay buffer size for sampling	200,000
Replay ratio	6
Batch size	128
Model unroll length L	5
Reward loss scaling cr	1
Termination indicator loss scaling cd	1
Feature loss scaling cs	10
Value loss scaling cv	0.25
Policy loss scaling cπ	0.5
Target for policy	Action distribution
Actor-critic
Learning rate	Anneal linearly from 6e-4 to 0
Optimizer	Adam
Adam beta	(0.9, 0.999)
Adam epsilon	1e-8
Batch size	16
Actor-critic unroll length	20K + 1
Baseline scaling	0.5
Clip global gradient norm	1200
Entropy regularizer for real actions	1e-3
Entropy regularizer for imaginary and reset actions	5e-5
Environment specific-Sokoban
Discount rate γ	0.97
Input shape	(3, 80, 80)
Environment specific-Atari 2600
Discount rate γ	0.99
Grayscaling	No
Action repetitions	4
Max-pool over last N action repeat frames	2
Input shape	(12, 80, 80)
Frame stacking	4
End of the episode when life lost	Yes
Reward clipping	[-1, 1]
Sticky action	No
Figure 11: Learning curves for Thinker and baselines in Atari games. For all methods except Rainbow, each curve is computed as the running average of episode returns over the last 500 episodes. For the five selected Atari games, the shaded area represents the standard deviation across three seeds.
Thinker-augmented MDP
Raw MDP
Table 5: Atari scores after 200M frames of training. Up to 30 no-ops at the beginning of each episode.
Rainbow	Raw MDP	Thinker Base Policy	Thinker Agent Policy
Alien	9491.70	13868.80	229.20	230.10
Amidar	5131.20	1121.56	1338.75	3591.84
Assault	14198.50	6998.29	18856.51	24262.47
Asterix	428200.30	279440.00	238992.50	85254.00
Asteroids	2712.80	12570.40	40030.60	212805.50
Atlantis	826659.50	896017.00	808384.00	809160.00
Bank Heist	1358.00	378.50	1531.20	1528.90
Battle Zone	62010.00	2320.00	47553.33	48103.33
Beam Rider	16850.20	24202.46	15213.50	59533.14
Berzerk	2545.60	2664.00	1024.60	1038.30
Bowling	30.00	23.78	31.26	31.14
Boxing	99.60	8.17	99.51	99.86
Breakout	417.50	593.83	660.84	840.14
Centipede	8167.30	8201.71	5169.38	6507.05
Chopper Command	16654.00	1046.00	43691.00	843973.00
Crazy Climber	168788.50	97824.00	98755.00	116193.00
Defender	55105.00	385397.50	71378.50	123064.50
Demon Attack	111185.20	502.30	133410.20	135040.10
Double Dunk	-0.30	-1.72	15.27	16.09
Enduro	2125.90	26.03	0.00	0.00
Fishing Derby	31.30	32.46	38.58	43.68
Freeway	34.00	0.00	28.04	28.65
Frostbite	9590.50	269.70	303.30	303.10
Gopher	70354.60	5049.40	15871.40	55911.00
Gravitar	1419.30	267.00	3766.00	4226.00
Hero	55887.40	13117.30	2984.30	2987.80
Ice Hockey	1.10	-5.27	20.62	22.99
Jamesbond	N.A.	1523.00	4171.50	7612.50
Kangaroo	14637.50	390.00	12424.00	13218.00
Krull	8741.50	5767.00	4663.10	4557.60
Kung Fu Master	52181.00	35209.00	32639.00	36391.00
Montezuma Revenge	384.00	0.00	12.00	14.00
Ms Pacman	5380.40	7117.66	2685.40	4206.10
Name This Game	13136.00	13391.00	25131.43	33085.20
Phoenix	108528.60	18201.70	156647.23	552719.37
Pitfall	0.00	-38.17	-2.95	-2.51
Pong	20.90	20.97	20.96	20.84
Private Eye	4234.00	48.41	93.61	94.00
Qbert	33817.50	18491.75	20184.58	34849.33
Riverraid	N.A.	23983.60	11664.30	11696.00
Road Runner	62041.00	2988.00	37967.00	41950.00
Robotank	61.40	31.73	40.89	40.53
Seaquest	15898.90	1501.40	32992.90	432855.00
Skiing	-12957.80	-16076.10	-21515.92	-24218.14
Solaris	3560.30	2136.60	2427.00	2323.40
Space Invaders	18789.00	8423.85	12880.40	42486.95
Star Gunner	127029.00	183274.00	290441.00	692283.00
Surround	9.70	-9.86	1.89	4.28
Tennis	0.00	20.51	-1.25	-1.19
Time Pilot	12926.00	42851.00	37310.00	50767.00
Tutankham	241.00	51.60	262.00	268.14
Up N Down	N.A.	319317.80	337508.10	404120.10
Venture	5.50	0.00	0.00	0.00
Video Pinball	533936.50	281210.04	444212.40	512567.38
Wizard Of Wor	17862.50	6010.00	19304.00	31691.00
Yars Revenge	102557.00	72845.97	70000.15	124342.12
Zaxxon	22209.50	15.00	17716.00	19979.00
E Ablation Analysis
In this section, we investigate the significance of various elements of the Thinker-augmented MDP. All ablation experiments are performed in Sokoban under identical settings as those of the primary experiments, with results shown as the solving rate over the last 200 episodes.
Hints and memory We examine the effect of removing all hints and other auxiliary statistics from the tree representation. Specifically, we remove the following statistics:
$$
\left\{g _ {\text {c h i l d} (\text {r o o t})} ^ {\max }, g _ {\text {c h i l d} (\text {r o o t})} ^ {\max }, n _ {\text {c h i l d} (\text {r o o t})}, g _ {\text {c h i l d} (\text {c u r})} ^ {\max }, g _ {\text {c h i l d} (\text {c u r})} ^ {\max }, n _ {\text {c h i l d} (\text {c u r})}, g _ {\text {r o o t}, \text {c u r}}, g _ {\text {r o o t}, \text {c u r}} - \gamma^ {d + 1} \hat {v} _ {\text {c u r}}, g _ {\text {r o o t}} ^ {\max } \right\}
$$
from the augmented state. These statistics are deemed auxiliary since, theoretically, the agent can compute these from the remaining statistics in the previous and current steps. The result of removing the hints is shown in Figure 12.
The removal of hints causes the agent to learn at a slower pace while slightly reducing the final performance. We theorize that the agent requires more training to learn to summarize the rollouts across multiple steps, leading to slower learning. With prolonged training, it could be feasible for the agent to learn this summarization and achieve the same final performance, thereby rendering the provision of these handcrafted hints unnecessary.
A specialized RNN architecture is employed to endow the agents with long-term memories while preserving training stability, as elaborated in Appendix C. We then consider the impact of replacing this RNN with a three-layer MLP with 200 hidden units, effectively removing the agents' memory. In the MLP, we add a skip connection from the input to each hidden layer, to ensure easier flow of information. The result of this adjustment is presented in the same figure. Interestingly, the removal of memory does not hamper performance in any notable way.
We hypothesize that as we have already summarized the rollouts for the agent by providing the hints, the agent is not required to perform this summarization itself and hence, does not require memory. To verify this hypothesis, we consider the scenario of removing both memory and hint, and the outcome is demonstrated in the same figure. The effect of eliminating memory proves to be highly significant when the hints are not provided. Consequently, we conclude that the presence of either the hints or memory is necessary for learning to plan.
Base policy and imaginary actions We evaluate the effects of removing the base policy and its value. This leads to the tree representation being mostly empty as a number of statistics, such as the hints, are computed based on the node's value. Nonetheless, the agent still receives the model's hidden state as input. The result of this removal is shown in Figure 13.
The result shows that the agent fails to learn in such a scenario, indicating that the base policy is a critical component for learning to plan. We offer two potential explanations for this observation. Firstly, predicted values provide an agent with the ability to evaluate a leaf node more easily. This node evaluation allows an agent to evaluate a rollout without simulating the rollout till the end of the episode, thus allowing more shallow search. Even though the predicted value could theoretically be computed based on the hidden state of the node, which is provided to the agent, it is likely that this prediction is too difficult for the agent to learn by itself. Secondly, providing the base policy allows the agent to continuously improve upon the current policy. The base policy acts as a 'distilled' policy without planning, and the search corresponds to a learned policy improvement operator applied to this distilled policy. With the base policy and its value as inputs, the agent can learn different search algorithms as this policy improvement operator, completing the cycle of generalized policy iteration.
In addition, we consider the scenario where all imaginary actions are replaced with random actions, and we do not train the RL agent on these imaginary actions. The results are shown in Figure 13. This is intended to evaluate the benefits of learning imaginary actions. We find that performance significantly deteriorates, suggesting that the imaginary actions play a pivotal role in guiding real actions, consistent with the behavioural analysis in Appendix F.
Maximum search depth We evaluate the effect of increasing the maximum search depth and the model unroll length  $L$  from 5 to 10, thereby allowing the agents to search deeper. The result is shown in Figure 14. The initial learning speed is significantly slower, presumably due to the added time required for the model to learn to unroll over an increased number of steps. However, the final performance remains largely unaffected. This suggests that a deeper search does not necessarily lead
Figure 12: Ablation analysis on hints and memory.
Figure 14: Ablation analysis on maximum search depth or model unroll length  $L$ .
Figure 13: Ablation analysis on base policy and imaginary actions.
Figure 15: Ablation analysis on planning reward and action target.
to improved final performance. This finding aligns with previous studies indicating that a shallow search suffices to obtain satisfactory outcomes in MuZero [43].
Planning rewards and action target Lastly, we discuss some non-essential components of the algorithm. To mitigate the sparse rewards in the augmented MDP, we added a heuristic reward, which we refer to as the planning reward, during the imaginary steps. The details of the planning reward are discussed in Appendix G. Figure 13 shows the results of excluding planning rewards. The removal of the planning rewards does not have any significant impact on the performance. This suggests that the planning reward may not be a critical component for learning to plan. Additional behavior analysis of a trained agent in Appendix F demonstrates that an agent trained without planning rewards seems to learn a similar planning behavior to that of an agent trained with planning rewards.
We used the action distribution instead of the sampled action for training the base policy in the model. However, this approach necessitates that the augmented MDP receive the action distribution from the agent, a procedure that is not consistent with the definition of an MDP. We alternatively experimented with using the sampled action as the target for training the base policy in the model, with the results depicted in Figure 13. This modification did not significantly impact overall performance, though it led to a slightly slower initial learning. Nonetheless, we recommend employing the action distribution in lieu of the sampled action, as the former encompasses more information and has a lower variance. The adherence to the conventional definition of an MDP, while theoretically pertinent, is of marginal consequence in practical applications.
F Agent Behavior Analysis
In this section, we delve into the behavior learned by the trained agents in the Thinker-augmented MDP of Sokoban. Our inquiry aims to understand the nature of the search policies acquired by the agents. Due to the use of the dual network, it is easy to visualize what the agent is 'planning' by plotting the predicted states from the model on each imaginary step. Figure 1 shows a typical stage, where the agent plans different ways of solving the task. Videos of the whole game and other Atari games are available at https://youtu.be/0AfZh5SR7Fk.