Thinker: Learning to Plan and Act
Authors: Stephen Chung
Ivan Anokhin
David Krueger
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement learning
model-based reinforcement learning
planning
Monte Carlo Tree Search
Markov Decision Process
TL;DR: We suggest wrapping the environment with a world model, thereby enabling a reinforcement learning agent to interact with the world model in this augmented environment and autonomously learn to plan.

Abstract: We propose the Thinker algorithm, a novel approach that enables reinforcement learning agents to autonomously interact with and utilize a learned world model. The Thinker algorithm wraps the environment with a world model and introduces new actions designed for interacting with the world model. These model-interaction actions enable agents to perform planning by proposing alternative plans to the world model before selecting a final action to execute in the environment. This approach eliminates the need for handcrafted planning algorithms by enabling the agent to learn how to plan autonomously and allows for easy interpretation of the agent's plan with visualization. We demonstrate the algorithm's effectiveness through experimental results in the game of Sokoban and the Atari 2600 benchmark, where the Thinker algorithm achieves state-of-the-art performance and competitive results, respectively. Visualizations of agents trained with the Thinker algorithm demonstrate that they have learned to plan effectively with the world model to select better actions. Thinker is the first work showing that an RL agent can learn to plan with a learned world model in complex environments.

Thinker: Learning to Plan and Act
Stephen Chung University of Cambridge mhc48@cam. ac. uk
Ivan Anokhin Mila, Université de Montréal ivan anokhin@mila. quebec
David Krueger University of Cambridge dsk30&cam ac. uk
Abstract
We propose the Thinker algorithm, novel approach that enables reinforcement learning agents to autonomously interact with and utilize a learned world model. The Thinker algorithm wraps the environment with world model and introduces new actions designed for interacting with the world model. These mndee-interrattion actions enable agents perform planning by proposing alternative plans to the world model before selecting a final action to execute in the environment. This approach eliminates the need for handcrafted planning algorithms by enabling the agent to learn how to plan autonomously and allows for easy interpretation the agent's 's plan with visualization. We demonstrate the algorithm's effectiveness through experimental results in the game of Sokoban and the Atari 2600 benchmark, where the Thinker algorithm achieves state-of- the-art performance and competitive results, respectively. Visualizations of agents trained with the Thinker algorithm demonstrate that they have learned to plan effectively with the world model select better actions. Thinker is the first work showing that an RL agent can learn plan with learned world model in complex environments.
Introduction
Model-based reinforcement learning (RL) has significantly enhanced sample efficiency and perfor- mance by employing world models, or simply models, to generate additional training data [ provide better estimates of the gradient 4) and facilitate plamning [.] Here, planning refers to the process of interacting with the world model to inform the subsequent selection of actions. Handcrafted planning algorithms, such as Monte Carlo Tree Search (MCTS) (6), have achieved underscoring the importance of planning.
remarkable success in recent years 7 g 5
However, significant research gap persists in creating methods that enable RL agent learn to plan [9]. meaning to interact autonomously with a model, eliminating the need for handcrafted planning algorithms. We suggest this shortfall stems from the inherent complexities of mastering several skills concurrently required by planning: searching (exploring new potential plans), evaluating (assessing plan quality). summarizing (contrasting different plans), and executing (implementing the optimal plan). While handcrafted planning algorithm can handle these tasks, they pose significant learning challenge for an RI. agent that learns to plan. Introducing a learned world model further complicates the issue, as predictions from the model might not be accurate.
To address these challenges, we introduce the Thinker algorithm, novel approach that enables the agent to interact with a learned model as an integral part of the environment. The Thinker algorithm wraps Markov Decision Process (MDP) with learned model and introduces new set of actions that enable agents to interact with the model The manner in which an agent uses model is not predetermined. In principle. an agent can learn various common planning algorithms, such as n-step exhaustive search or MCTS, or ignore the model if planning is not beneficial. Importantly, we
Full code is available the Thinker-augmentedd MDP with the interface as OpenAl Gym 10
wtt.//oo////aeeeee cton//stepenn-chune mh/thinker which allows for using
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

rollout
real state
g:0.73
step
step 2 g:6.36
step 3 g:8.41
step 4 g:6.05
9:0.39
Step 5
step 9:8.73
step. 7 g:0.36
step 8 g:6.11
Step 9
g:6.41
step 10 9:6.45
step
step
9:7.43
9:7.56
step 9:5.47
step
step 18 g:5.59
step g:6.01
real state
P
rollout 3
rollout
Figure 1: Illustration of trained agent during single stage in Sokoban, where the goal is push all four boxes to oed-bordered squares. From the root state, the agent generates four rollouts using the learned model. Except for the two real states, all frames generated by the model, showcasing near- perfect frame prediction. Predicted rollout return (equivalent to gmot, in Equation [ is displayed. The four rollouts reasonable plans with agents pushing different boxes towards targets. The agent follows the third rollout for the next five real steps (not fully shown in the figure), likely due to its high rollout return. Crucially, plan generation, plan execution, and the model are all learned. Videos of the whole game and other Atari games are available at /httts://youut be/OAfZh5SR7FK
have designed the augmented MDP in a way that simplifies the processes of searching, evaluating, summarizing, and executing, making it easier for agents to learn to use the model.
Drawing from neuroscience, our work inspired by the influential hypothesis that the brain conducts planning via internalized actions-actions processed within our experience-based internal world model rather than externally [1] 12
According to this hypothesis, both model-free and model-based behaviors utilize similar neural mechanism. The main distinction lies in whether actions are directed toward the external world or the internal world model. This perspective is consistent with the structural similarities observed in brain regions governing model-free and model-based behaviors 13. This hypothesis offers insights into how an RL agent might learn to plan-by unifying both imaginary and real actions and leveraging existing RL algorithms.
Experimental results show that actor-critic algorithms,, when applied to the Thinker- aakkkketed MDP, yield state-of-the-arr performance in Sokoban. Notably. this approach attains solving rate of 94.5% within 5c7 frames, a significant improvement over the 56.7% solving rate achieved when the same algorithm is applied to the raw MDP. On the Atari 2600 benchmark, actor-critic algorithms using the Thinker-augmented MDP show significant improvement over those using the raw MDP. Visualization results, as shown in Fig [ the Thinker algorithm provides the following advantages:
reveal the agent's effective use of the model. To summarize,
The agent learns to plan on its own, without any handcrafted planning algorithm, allowing it to adapt to different states, environments, and models.
Flexibility:
Generality: The Thinker algorithm only dictates how an MDP is transformed, making it compatible with any RL algorithm. Specifically, one can convert any model-free RL algorithm into a model-basee based RL algorithm by switching to the Thinker-augmented MDP.
Interpretability: We can visualize the agent's 's plan prior to its execution, as depicted in Fig 1 This provides greater insight compared scenarios where the neural network internalizes the model and plans in a black box fashion [1].
Aligned objective: Both the real and imaginary actions are trained using the same rewards from the environment, ensuring that the objectives of planning and acting align.
Improved learned model: We introduce novel combination of architecture and feature loss for the model, which designed to prioritize the learning of task-relevant features and enable visualization at the same time.

---

2 Related Work
Various approaches exist for incorporating planning into RL algorithms,, with the majority relying on handcrafted planning algorithms. In other words, the policy underlying the interaction with the world model, termed here as the imaginary policy, is typically handcrafted. For example, MuZero [5], AlphaGoo [55, and AlphaZero 7 employ MCTS. VPN [6), TreeQN and ATreeC [7 explore all
possible action sequences up to a predefined depth. 12A
[8]] generates one rollout per action, with
the first action iterating over the action set and subsequent actions following the current policy.
Learning to plan is a challenging domain, and a few works exist in this area. MCTSnets [19] is a supervised learning algorithm that learns imitate actions from expert trajectories generated from MCTS. VIN [20 and DRC [14] are model-free approaches that utilize special neural network architecture. VIN's architecture aims to facilitate value iteration, while DRC's architecture aims to facilitate general planning. IBP [2I], which is closely related to our work, allows the agent to engage with both the model and the environment. However, several critical differences distinguish IBP from Thinker. First, in IBP, the imaginary policy is trained by backpropagating through the model, method that primarily suitable for continuous action sets, whereas our approach focuses on environments with discrete action sets and treats the model as a black-box tool external to the agent. Second, unlike our method, IBP does not predict future values and policies, which we have identified as crucial for complex environments (see Appendix ED. Third, we have designed an augmented state representation that significantly simplifies the process of using the model by providing more than just the predicted states to the agent, as is the case with IBP.
prior to taking each real action. These methods start with preliminary imaginary policy and use this policy to gather rollouts from the model. Subsequently, these rollouts are used to calculate gradient updates for the imaginary policy by maximizing the imaginary rewards, thus incrementally refining the rollouts' quality. In the final step, a handcrafted function is used to select real action. For instance, previous works such as [22] 223 and (24] have proposed to maximize the rollout return using REINFORCE (25]. PPO 226). and backpropagation through the model, respectively. In contrast to these works, Thinker does not employ such gradient updates to refine rollouts, as the agent has to learn how to refine rollouts and select a real action on its own. Moreover, unlike these works, Thinker trains the imaginary policy to maximize the real rewards, thus preventing model exploitation when the model is learned.
Another research direction employs gradient updates to refine rollouts [2
23 24
Thinker is the first work showing that an RL agent can learn to plan with learned world model in complex environments. Prior related works either do not satisfy (i) learning to plan, i.e., a learned imaginary policy, (ii) using learned model or (iii) being evaluated in complex environment. For
example, MuZero [5). VPN [16), TreeQN [17. and I2A [18) do not satisfy (i). VIN 20), DRC [4]
and [22 (24 do not satisfy (ii). IBP [21] and [23] do not satisfy (iii).
We describe the connection between our algorithm and Meta-RL
27
28
as well as generalized
policy iteration 22 30), in Appendix H
3 Background and Notation
where S set of states, A is a finite set of actions, P S x A S - [0. is a transition function representing the dynamics of the environment, R : S x A + R is a reward function, E [0.1] a discount factor, and dg S + [0, 1] is an initial state distribution. Denoting the state, action, and reward at time L by 8t, des and respectively. P(s,a,8')
We consider Markov Decision Process (MDP) defined by tuple (S.A,P,R,
s 8'1s 8, 0, a). R(s,a) EErr+1 8t 8, 0z = a), and do(s) Pr(81 = 8), where P and dg are valid probability mass functions. An episode is sequence of (st, ae, 11) starting from f = 1 and continuing until reaching the terminal state, a special state where the environment ends. Letting 9: EE=, >*-=r denote the infinite- horizzon discounted return accrued after acting at time t, we are interested in finding, or approximating, a policy π : S x A + [0, 1], such that for any time 2 1, selecting actions according to π π(,, a) = Pr(a, = alst = s) maximizes the expected return Elgh+1| The value function for policy π is v" where for all E S, v (s) E[ge+1|st 8, m). We adopt the notation (st, Gt. T 11+1) for representing transitions, as opposed to (St, a,, rr, 8(+1). which facilitates a clearer description of the algorithm. In this work, we only consider an MDP with a discrete action set.
Pr(st+1
3

---

Imaginary steps (K times)
Real step Augmented State aug StK
Agent
Agent
Augmented State aug $ L,ke
Augmented Action Gg A
Real
Real
Reward
Action
World Model
Environment
Roal State
St
Figure 2: Illustration of stage in the augmented MDP, which is composed of K - 1 imaginary steps and one real step. The environment outputs a state 8, that sets the root state of the world model. The agent interacts with the world model K times with augmented actions. In each imaginary step, .uug which contains the latest outputs from the world model.
the agent receives an augmented state *t,k Finally, the agent acts in the real environment for single step based on 8t.K*
4 Algorithm
The Thinker algorithm transforms given MDP, denoted as . = augmented MDP, denoted as _MSe (Smng A pam. R_mn aug d. We call M the real MDP. and MM"gg the augmented MDP. Besides the real MDP, we assume that we are given (i) model in the form of (P, It) that is an approximate version of (P, R) and is deterministic, and (ii) a base policy (S,A) + [o, 1] and its value function S + IR. The base policy can be any policy, which helps to simplify the search and evaluation of rollouts for the agent. We will show how both the model and base policy can be learned from scratch in Section 4.4 but for now we assume that both are given.
do), into an
4.1 Augmented Transitions
To facilitate interaction with the model, we add K 1 steps before each step in the real MDP, where K 2 and denotes the stage length. We typically use K = 20 in the experiments. The first K - steps are called imaginary steps, as actions in these steps are passed to the model instead of the real MDP. The step that follows the imaginary step is called the real step, as the action selected is passed the real MDP. These K steps in the augmented MDP, including K 1 imaginary steps and one real step, constitute a single stage in the augmented MDP. An illustration of stage is shown in Fig 2 Each episode in the augmented MDP composed of multiple stages, with a single stage corresponding to one step in the real MDP. The augmented MDP terminates if and only if the underlying real MDP terminates. We use ie E {1, 2..., K) denote the augmented step within a stage, and E 21.2....} to denote the current stage. We omit the subscript where it is not pertinent.
Our next step is to determine the form of the augmented transition. We can view the model search as tree traversal, with the tree's root node corresponding to the current real state 8. We consider a simple method of traversing a tree. At each imaginary step, the agent decides which child node to visit by imaginary action à E A and also whether to reset the model to the root node by reset action ā" {0, 13. In other words, the imaginary action unrolls the model by one step while the reset action sets the model back to the root node. We also impose maximum search depth, L. in the algorithm. This ensures that the agent is forced to reset once its search depth exceeds L, as learned model may not be accurate for a large search depth. An illustration of the tree traversal is shown in Figure
After the imaginary steps. the agent chooses a real action à A in real step, where this real action is passed to the real MDP. We can merge both the imaginary and real actions, while interpreting the reset action as an additional one, thereby defining the augmented action space by A == A {0, 1}, and an augmented action by the tuple au == (ā, a"). During imaginary steps, ă becomes the imaginary action, while â' decides whether to reset the current node. During real steps, becomes the real action, while ä" is not used. The pseudocode can be found in Algorithm
4.2 Augmented State
We then consider the design of the augmented state that is passed to the agent. One straightforward method involves concatenating all predicted states in the tree. At each real step, the K predicted

---

8
Reset triggered
Forced reset triggered
Sb
a
8b
$
Sb
6
Sb
$
$b
Sa.b
Sa.0
Sa.b
Step 1:(b, reset) Step 2:(a, not reset) Step 3:(b, not reset) Step 4:(a, not reset)
Step 5:(a, reset)
Figure 3: Illustration of tree traversal in imaginary steps, where the stage length K = 6 and the maximum search depth L 2. Two actions, and b, are available. The root node corresponds to the real state underlying the stage, and the child nodes correspond to the predicted states &n aa..... following the selection of actions 01 d2.... from state The augmented actions in the K - = imaginary steps are depicted beneath the tree. The first action unrolls the model, whereas the second action determines whether the current node is reset the root node. The yellow and the orange nodes indicate the current node hefore and after unrolling (before reset is applied), respectively. A forced reset triggered on the 3rd imaginary step, as the maximum search depth has been reached.
states serve as input, informing the selection of the real action. However, this method presents two significant challenges: (i) the high dimensionality of the predicted states makes learning to search and evaluate rollouts difficult, and (ii) the concatenation effectively multiplies the state's dimension by K, further complicating the learning of rollout summarization. To mitigate these issues, we propose compact way of representing the tree in a wiwddimesional vector. First, we consider how to encode a node compactly, followed by how to use these node representations to represent the tree.
To encode node, we include the reward, the base policy and value at the node. Moreover, we provide certain hints that summarize the rollouts. Taking inspiration from MCTS, we consider three hints: mean rollout return, maximum rollout return, and visit count.
For a node i and its n-level descendant node in. we define the rollout return from node in to node in through the path (io, 21, .., in) as:
9io.á. = Pro + rf, + 7 " + +
where úg denotes the value function applied on the node i and A, denotes the predicted reward upon reaching node i. That is, 0; == 0($:), and A; R( (2(pprent_nodeeleee
edne(t))
For given node we can compute Ji, for all its descendant nodes that are visited in the tree. We also define Ji, := n; + 70; We define the mean rollout return g gluesn and the maximum rollout return g as mean of gi..j and the maximum gio over all descendant nodes j (including itself) that are visited respectively. The mean and maximum rollout returns of node's children are useful in guiding actions. For example, the mean rollout return of the child i of a node & gives a rough estimate of the average estimated return that can be collected following action i from node k. The mean rollout return is also equivalent to Q(s, a) used the UCB formula of MCTS. We use Schild i) and (child(s) to denote the vector of mean and maximum rollout returns for each child of node i.
mean
The visit count of a node, which represents the number of times the agent traverses that node, may also be useful. For example, child node that has never been visited may be worth exploring more. We use (elid(ss) E R |AA to denote the vector of visit count for each possible child of node i.
Altogether, represent node i by vector u; RSAA|+ called node representation, that contains:
1. fr, 0; E R, n. E R|AI; The core statistics, composed of the reward, the value, and the base policy of the node i. Here, मं (f a,), á (s;, az), ..., 5 (s,, alA))).
2. one hot(ho)(a;) E R R|AI; The one-hot encoded action that leads to node i.
3. meean
()hild(j) Fahild(s) K (emiaas))) E R1- The hints, composed of mean rollout return, maximum rollout return and visit counts of the child nodes normalized by K.
We only include the root node's 's representation Unat and the current node's representation Uear to represent the tree (we use the subscript roor denote the root node, which is unchanged throughout
5

---

stage, and the subscript cur to denote the current node before reset). The former guides the selection of real actions, while the latter guides the selection of imaginary actions. Besides Umot and Heur. we also add auxiliary statistics, such as the current step in a stage, in the representation. Further details can be found in Appendix A Combined, we represent the tree with a fixed-size vector Unee E [R]0]A]-K++9 called free representation. We may also include the real state or the current node's predicted state Sour in the augmented state, providing richer context for the agent.
The tree representation is designed to simplify various planning procedures. In searching. the base policy and visit count help the identification of promising, yet unvisited rollouts. In evaluation, rather than simulating until terminal state, value estimates can be utilized assess the potential of given state. In summarization, mean and maximum rollout returns consolidate the outcomes for the current node. Finally, in execution, the hints at the root node provide compact summary of all rollouts. Compared to the concatenated predicted states, which have a dimension of K K|S), the tree representation, typically with fewer than 100 dimensions, is much easier learn from.
The tree representation enables an agent to execute various common planning algorithms, such as n-step exhaustive search and MCTS, without requiring memory. Employing an agent with memory, such as an actor-critic algorithm with an RNN, may facilitate learning a broader class of planning algorithms, as the agent can access previous node representations. In principle, an agent with memory can learn to compute the hints or other ways of summarizing rollouts by itself.
4.3 Augmented Reward
To align the agent's goals in the augmented MDP with those in the real MDP, we adopt the following augmented reward function: no reward given for imaginary steps. and the reward for real steps is Jug identical to the real reward from the real MDP. In other words, 'xx = 0 for k < K and "K+1"
aug
where r denotes the real reward. We set the augmented discount rate - := 7 * to account for the K T-11 imaginary steps preceding one real step, ensuring that the return in the augmented MDP equals the return in the real MDP.
We also experimented with the use of an auxiliary reward, called the planning reward, to mitigate the sparse reward in the augmented MDP. The core idea is to encourage the agent to maximize the maximum rollout return guer but we later found that it only provided a minimal increase in initial learning speed, indicating that the real rewards are sufficient for learning to plan. The details of the planning reward can be found in Appendix G
This completes the description of the augmented MDP. One may view an RL agent acting on this augmented MDP with a fixed base policy as performing one step of policy improvement in generalized policy iteration, but the policy improvement operator is learned instead of handcrafted. Experiments show that RI. agents can indeed learn to improve the base policy (see Appendix [H). As such, we only need project the improved base policy back to the base policy space to complete the full cycle of generalized policy iteration, which will be discussed next.
4.4 Learning the World Model and the Base Policy
Until now, we assume that we are provided (i) the model (P, Rt and (ii) a base policy á and its value à. To remove both assumptions, we propose employing a single, unified model that assumes the roles of P., R., it, and i, 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; and an action sequence 04+1:1+L. and the network predicts future states $141:+L. rewards f1+1:+ and termination probability Pi+1-1--- Meanwhile, the value-policy network takes both the state-reward network's 's input and its predicted states $1+11+L as inputs, and the network predicts future policies f and values t. We can also view the dual network as an RNN with = (hær, hPP) being the hidden state, where 1 and h 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.

---

1, Pt+1
2x P6+2
h
h
St
St+2
g(st)
(
g( ((2+2)
Model Inputs
h
R
up
State-reward
Network
0e, Rt
B Mt+1
8442 #c42
Value-policy
Network
Figure 4: Illustration of the dual network's architecture, composed of the state-reward and the value- policy network. See Equations and for the loss functions and Appendix B for training details.
To train the model, we first store all real transitions (st, Ge. re+1. d in a replay buffer, which corresponds to the state, action, reward, and termination indicator. We then sample sequence of real transitions with length L + from the buffer. The two sub-RNNs in the dual network are trained separately. For the state-reward network, the following loss is used:
Σ (f - rr+1)² + ddes log 14 + c² Lotate (8e42, Be+1))
L
where c², c,t 2 are hyperparameters that modulate the loss strength. While a straightforward choice for the state prediction loss metric, L is the 1,2 loss - defined as L ()) x)= ||
x1
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 1.2 loss. Specifically, L (s', s)
[lg(s') - g(s)] where 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:
L
(c"(
6441 target )² + c con_hho((a++)) log R++1).
(3)
1=0
where c", c. are hyperparameters and alunget denotes the multi-step return. One may also use the real action distribution ne place of one meehot((a++) 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 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 RI algorithms in Appendix
With the dual network, the augmented state can also include the hidden state of the model at either the current node, /uur. or the root node, Amort. These hidden states, usually in much lower dimensions, might be easier process than the real state s or the current node's predicted state S In our experiments, we use (Nuee heur) 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.
"Strictly speaking. using (umse hewr) as the augmented state leads to a POMDP instead of an MDP. One can recover an MDP by omitting the auxiliary hints in Unse and adding Scce to the augmented state, but both do not help performance in experiments.
7

---

Thinker's performance at 5e7
80%
Solving Rate
60%
Solving
40%
20%
2
3
4
5
Frames
1e7
Frames
1e8
Thinker
DRC
DRC (original paper)
Dreamervv3
MuZero
ResNet
12A
ATreeC
VIN
Figure 5: Running average solving rate over the last 200 episodes in Sokoban, compared with other baselines.
Experiments
For all our experiments, we train 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 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 concatenated, then passed through linear layer to generate actions and predicted values. We set the stage length, K, to 20, and the maximum search depth model unroll length, L, to 5. More experiment details can be found in Appendix D
model's hidden state. These encoded representations
5.1 Sokoban
a classic puzzle problem, as our primary testing envi- ronment due 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 I We used the unfiltered dataset comprising 000,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.
We selected the game of Sokoban [32. 18
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
331. MuZero [5], I2A [18), ATreeC [17], VIN [20]. and IMPALA with ResNet [31 In Sokoban,
DRC stands as the current state-of-the-ar 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 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 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, evaluate policy trained with K 20 the Thinker- Augmented MDP, by testing its performance with different values in the same augmented environ- ment during the testing phase. This variability allows control the degree of planning and the result is shown in Figure The figure also depicts the performance of the base policy, which can be regarded distilled policy devoid of planning. We observe significant performance improvement attributable to planning, even at the end training.
The agent's s behavior, visualized in Figure 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

---

100%
100%
Agent policy
policy
00%
80%
Rate
Solvin Rate
40%
40%
K-10
K=55
Km2
K=1
0%
0%
Raw MDP
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
1
Figure Final solving rate on Sokoban test levels with varying stage length K during test time but trained with = 20.
Figure 7: Running average solving rate over the last 200 episodes in Sokoban with varying stage length K.
Thinker-su gmeented
Median humannnnrmllzzee score
250%
Raw MDP
80%
wwwh
200%
Soling Rate
feature state Dual net & L2 state loss
60%
state
150%
Single not state loss
Single
1000%
20%
0%
0%
0.00
0.25
0.50
1.00 1.25
1.75
2.00
frames
Frames
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 leaf node. Further analysis the trained agent's behavior can be found in Appendix F
We conducted ablation analysis to examine the impact of plamning duration, varying the stage length K across {1, 2, 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 presented in Figure The results suggest that = 10 already gives the optimal performance. This is 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 similar network architecture used by the model to encode the raw state. Surprisingly, the result for K = 1 worse than the raw MDP. detailed explanation of reasons behind this phenomenon can be found in Appendix(D [ However, employing 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 dual network that uses L2 loss in state prediction, (ii)
several variants of the model, including dual network that has no loss in state prediction, (iii) a single network that uses L. loss in state prediction, (iv) single network that does not predict state. The single network refers 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 X We observe that both the dual network architecture and the state loss are critical to performance. However, the use of feature loss instead of 1.2 loss gives marginal performance improvement, likely because the 1.2 loss' prior that each pixel 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 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-augnn augmented MDP and the raw MDP are 261% % (1372%) and 102% (5 144%), respectively, underscoring the advantages of the thinker-augmentte MDP over the raw MDP. For context, the median (mean) normalized score of Rainbow [3). a robust baseline, is 223% (868%).
The advantages of the Thinker algorithm 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 '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://youtt: be/ OAf2h5SR7Fk] 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 deterministic model facilitate tree traversal, which may not work well for stochastic environment.
Future work could focus on addressing these limitations. such enabling more flexible planning steps and integrating uncertainty into the model. Additionally, exploring the application of the algorithm in multi-environment settings is 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-augment MDP an additional direction for future work.
The history of machine learning research tells us that learned approaches often prevail over hand- crafted ones. This transition becomes especially pronounced when computational power are available. In the same vein, we surmise that learned planning algorithms will eventually surpass handcrafted planning algorithms in the future.
large amounts of data and
10

---

Acknowledgement
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.cm cat.ll.aaaa....mmmmm mommaac..kk uk), provided by Dell EMC and Intel using Tier-2 funding from the Engineering and Physical Sciences Research Council (capital grant EP/T02219//), and DiRAC funding from the Science and Technology Facilities Council ..w.d.
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.
References
[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 informa- tion processing systems, 28, 2015.
[4] Mare 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 (ICMI,-1),, pages 465 662, 2011.
[5] Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-663 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, al. Mastering the game of go without human knowledge. nature. 550(7676):354-355 -59, 2017.
[8] David Silver, Thomas Hubert. Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419):1140-11444 2018
[9] Fan-Ming Luo, Tian Xu, Hang Lai, Xiong-Hur 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 Wojeicch Zaremba. Openai gym. arXiv preprint 160iv:1606.011
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 Knowlton. The role of the basal ganglia in habit formation. Nature Reviews Neuroscience, 766):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.
11

---

[15] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre. George Van Den Driess- che, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
[16] Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in neural information processing systems, 30, 2017.
[17] Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atrecc: Differentiable trre-structured models for deep reinforcement learning. arXiv preprint arXiv:1710.114177 2017.
[18] Théophane Weber, Sébastien Racaniere, David P Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yu- jia Li, et al. Imagination-augmae arXiv: 1707.0620. 2017.
agents for deep reinforcement learning. arXiv preprint
[19] Arthur Guez, Théophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wierstra, Rémi Munos, and David Silver. Learning search with metsnets. In International conference on machine learning, pages 1822-1831. PMLR, 2018.
[20] Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks Advances in neural information processing systems, 29, 2016.
[21] 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 arXiv:/707.06/70, 2017.
[22] Thomas Anthony, Robert Nishihara, Philipp Moritz, Tim Salimans, and John Schulman. Policy gradient search: Online planning and expert iteration without search trees. arXiv preprint 1rXiv:1904003666, 2019.
[23] Mikael Henaff, William Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint arXiv: 10X5:07705.07777 2017.
[24] Arnaud Fickinger, Hengyuan Hu, Brandon Amos, Stuart Russell, and Noam Brown. Scalable online planning via reinforcement learning fine-tuning. Advances in Neural Information Processing Systems, 34:16951-16663 2021.
[25] Ronald J Williams. Simple statistical sradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8(3-4):229 256, 1992.
[26] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 17Xiv:1707.06347 2017.
[27] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel 7. Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv: 16l1: 1.55663, 2016.
[28] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arX1v:1611.02279, 2016.
[29] Richard Bellman. Dynamic Programming. Princeton University Press, Princeton, NJ, USA, 1957.
[30] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[31] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning, pages 1407-1416. PMLR, 2018.
12

---

[32] Adi Botea, Martin Müller, and Jonathan Schaeffer. Using abstraction for planning in sokoban. In Computers and Games: Third International Conference, CG 2002, Edmonton, Canada, July 25-27, 2002. Revised Papers 3, pages 360-375. Springer, 2003.
[33] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 2201::2330.00004 2023.
[34] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-2117 279, 2013.
[35] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improve- ments in deep reinforcement learning. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018.
[36] Stefan Behnel, Robert Bradshaw, Craig Citro, Lisandro Dalcin, Dag Sverre Seljebotn, and Kurt Smith. Cython: The best of both worlds. Computing in Science & Engineering, 13(2):31-39, 2011.
[37] Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Hasselt, and David Silver. Distributed prioritized experience replay. In International Conference on Learning Representattons, 2018
[38] Sepp Hochreiter and Jürgen Schmidhuber: Long short-term memory. Neural computation, 18):17351 1997.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[40] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning, pages 7487-7498. PMLR, 2020.
[41] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint 1901.0/9000288000 2019.
[42] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607.06550 2016.
[43] Jessica B Hamrick, Abram L. Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Holger Buesing. Petar Veličković, and Theophane Weber. On the role of planning in model-based deep reinforcement learning. In International Conference on Learning Representations, 2021.
[44] Max Jaderberg, Volodymyr Mnih, Wojeiech Marian Czamecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017.
[45] Simon Schmitt, Matteo Hessel, and Karen Simonyan. Off-policy actor-critic with shared experience replay. In International Conference on Machine Learning, pages $545-8554. PMLR, 2020.
[46] Matthew Aitchison, Penny Sweetser, and Mareus Hutter. Atari-5: Distilling the arcade learning environment down to five games. arXiv preprint arXiv:2210.02019, 2022.
[47] Pawel Ladosz, Lilian Weng, Minwoo Kim, and Hyondong Oh. Exploration in deep reinforce- ment learning: A survey. Information Fusion, 85:1-22, 2022.
[48] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. survey of meta-reinforcemmnn learning. arXiv preprint arXiv:2301.0002. 2023.
13

---

[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 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-gradien 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 computer vision and pattern recognition, pages 15750- 2021.
[53] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Ad- vances 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, Catholija M Jonker, al. Model-based reinforcement learning: A survey. Foundations and TrendsO in Machine Learning, 16(1):1-118,
2023.
[56] Russell Stuart and Norvig Peter. Artificial intelligence a modern approach third edition, 2010.
14

---

A
Algorithm
Algorithm Transition in Thinker-Augmented MDP
1: Input: real MDP: env, model: P, base policy and its value function: At 0.;
2: Algorithm Parameter: stage length K 2 1, maximum search depth L 2 1.
3: Initialize & + K
Ensures that the first step is a real step à denotes imaginary/real action, ä" denotes reset action
4: procedure STEP(ă, â')
5: if K then
> Real step Initialize step number and search depth > Compute reward, state and termination for the real MDP Record the root node D Set the current node to be the root node No predicted rewards for the root node Imaginary step
6: k, < 1,0
7: 8, d t env.slep(a)
8: 8root
9:
Seur f- Smot
Fooo
11:
else
k.
k +11++1
13:
if l 2 L then ā" 1- end if
D Check if search depth exceeds maximum search depth L Override reset
15:
fear R(Seurs a)
b Compute the predicted reward B Compute the new current node by unrolling the model & No rewards for imaginary steps No termination for imaginary steps
17:
Sour + P(scur: a)
18:
0
19:
- False
20:
end if
21:
D Compute the base policy and its value Compute gu using S1, Pi, te, i. where i E {root, eur) described in Section A.1 A.1 Reset to the root node if reset is triggered
Teur; Uzur +- (Seur); ((Serr)
22:
23:
if ã" then
24:
1 0
25:
Scur + Sroot
end if
27:
return r, s end procedure
>> Return rewards, state, and termination for the augmented MDP
Algorithm implements the step function of the Thinker-augmented MDP, assuming we have access to given model (P, R), base policy n, and its value é. We assumee P is deterministic and use P(s,a) a) to denote the predicted next state s' from state and action
To compute the hints (i.e. mean rollout return, max rollout return, and visit counts) in the tree representation, we need to maintain tree with each node storing the predicted reward f. value ú, and visit count n,, but the details of updating the tree are omitted in the pseudocode for clarity.
is important to note that since P is deterministic, one can store the computed node in buffer to avoid 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:
After line 7 we insert the real transition in the form of (Sola à, r, d), where sou denotes the previous real state, into replay buffer. A parallel thread trains the model by sampling from the replay buffer.
2. A node is represented by the hidden state of the RNN model, not the real state. Consequently, line Ameuld initialize the root node with the first hidden state h of the RNN model. Line [7 updates the new current node the new hidden state by unrolling the RNN model with the current node and the latest action. Lines 16 and 21 compute the reward, policy, and value as predicted by the RNN model.
and
Implementation details Several key implementation details are worth noting. First, when transition- ing into new stage, we earry forward the tree from the previous stage and prune the non-descendant nodes of the new root node. This approach ensures the continuity of the previous plan into the
15

---

subsequent stages. However, our experiments show that carrying the tree has minimal impact on performance. Second, for clarity, the planning reward is not shown the pseudocode, but details can be found in Appendix G Lastly, if the model predicts episode termination at node i, defined by 1 > 0.5, we manually set the content of all descendant nodes to i O, f T.. új = 8j 0, and hy = h4g which represents the terminal state that is absorbing and has zero rewards and values.
A.1 Augmented State
As discussed in Section, both the current node representation and the root node representation root are included in the tree representation. Denoting the current search depth by d {0, 1, L). the full tree representation Uuue [R contains:
1. Urott Ulur E +| The root node representation and the current node representation.
2. 9root,cur E Current rollout return.
3. Gront.cur
d+ Dour E I: Sum of discounted rewards collected in the current rollout.
4. max
E I: Maximum rollout return root node.
5. dd+ E R: Trailing discount factor.
6. one mne_hot(A:) E [0, 1] Current step a stage, represented one-hot encoding.
7. reset
[0, 1]: Indicator of whether the current node has just been reset the root node,
And the following five inputs can be included in the augmented state st,k"
1. Mtoe E + The tree representation.
2. 8 E S: The real state underlying the stage.
3. Seur S: The predicted real state of the current node.
4. hroot H: The hidden state of the RNN model the node.
5. heur E H:: The hidden state of the RNN model the current node.
We used auy- = (Uurees heur) in our main experiments. Our code enables users to decide which inputs to include as part of the augmented state the Thinker-augmented MDP.
A.2 Code
The complete code is accessible at https://ggtthn com/ www. com/stephen- chung mh/ thinker) We implemented the Thinker-augmenteer MDP in batches using Cython [36], Python language extension that optimizes performance by interfacing with C. This ensures that the primary computational burden attributed model training and the RL algorithm, rather than the logic flow within the augmented MDP. The code supports multi-thread processing, allowing multiple threads of Thinker-augmenter MDP to share the same model.
Our implementation of the Thinker-augmente augmented MDP aligns with the OpenAI Gym interface [10), facilitating easy integration with existing RL frameworks. The code allows customization of various parameters such stage depth K, maximum search depth L, and the model's learning rate, among others. Included the repository is the code for the IMPALA algorithm applied the Thinker- augmented MDP, reproducing the experimental results presented in the paper.
16

---

B Model Training Procedure
B.1 Collection of Training Samples
A transition tuple == (st, de. Te+1, de+1) is stored in a replay buffer during each real step, which corresponds the state, action, rewards, and termination indicator. The replay buffer a first-in-first- out buffer with a capacity of 200,000 transitions. To train the model, batch of sequences 245-21 is randomly sampled from the buffer, where represents the model rollout length. We use batch size of 128 and model unroll length L = 5.
Consistent with MuZero [5 prioritized sampling (37] is used for sampling the sequence 2. Each transition tuple is assigned a priority Newly added transition tuples are given the maximum priority value in the buffer. During training, the priority ps of the first transition in sampled sequence Tei+2L is updated the L1 loss of the value, |
target Transitions sampled with probability U;
of P(i) = Py where a = 0.6. Each sampled transition sequence's loss is weighted by the importance weight is annealed from 0.4 to 1.0 throughout training.
(NP()) where N is the total number of transition tuples in the buffer, and
We employ replay ratio of 6, implying that model training paused if the ratio the total number of transitions trained on the model to the total collected transitions exceeds 6.
B.2 Additional Details for Training the Model
The state-reward network and the value-policy network trained separately, which leads the following implications. First, when training the state-reward network, the g within feature loss for state prediction is not optimized since belongs to the value-policy network. This prevents g from degeneration. Second, when training the value-policy network, gradients cannot flow through $ towards the state-reward network since $141 considered as the input the policy network. This ensures that the predicted state aligns with the actual future state.
The multi-step return target
Equation defined by:
0241 1:44+ + 7T+142 + ... + 2- ' T1+2L + 2L-1
Instead of using the predicted value 01421, we opted for the mean rollout return (at the end of the stage) experiments. This stability likely stems from averaging across multiple rollouts starting from 8422. rather than relying directly on the predicted value.
meam F0ot,1+21 as the bootstrap targef" We found that it yields more stable training in our
Snvet does not contain the reward upon arriving the root node, froce manually 0 in Algorithm
17

---

C
Network Architecture
In this section, we describe the neural networks that we use for the model and the actor-critic network.
C.1
Model
For Sokoban, the shape of the real state, is (3, 80, 80), while for Atari 2600, the real state's shape is (12, 80, 80) due to the stacking of four frames together. Before inputting the real state to the encoder, we concatenate it with the one-hot encoding of the action d;-1, tiled to (80, 80), channel-wise. The encoder architecture is the same for both the state-reward and value-policy networks, structured as follows:
Convolution with 64 output channels and stride 2, followed by ReLu activation.
D residual blocks, each with 64 output channels.
Convolution with 128 output channels and stride 2, followed by ReLu activation.
residual blocks, each with 128 output channels.
Average pooling operation with stride 2.
residual blocks, each with 128 output channels.
Average pooling operation with stride
All convolutions use kernel size of 3. The resulting output shape for both Sokoban and Atari 2600 is (128, 6, 6). We use D = for Sokoban and D = 2 for Atari 2600.
We denote gd (8x) and g(s+) as the encoder's output of 8t for the state-reward network and the value-policy network respectively, and the dependency on 0;-1 is omitted for clarity.
For the state-reward network, define the hidden state h := gs To obtain h from (hit, G(+1) with 1 2 0, we first concatenate h with the one-hot encoding of the action d,4z, tiled to (6, 6), channel-wise. This concatenated output is passed to the unrolling function, which consists of 4D + 1 residual blocks, each having output channel of 128. The output of the unroll function is hir141
For the value-policy network, in addition to the previous hidden state, the unrolling function should also take the encoded predicted frame g($e+1) as input, so the unrolling procedure is slightly modified. First, we define h1_, 1 as a zero vector of shape (128, 6, 6). To obtain hTP from (h 00-1-1,1111 (8t++)) with / (replace Se-1 with 8:44 when / = 0 here), first concatenate all three together channel-wise, with d:41-1 tiled to (6, 6). This concatenated output is passed the unrolling function, which consists of 4D + residual blocks, each having output channel of 128. The output the unroll function up
To compute the predicted quantities, we apply prediction function to h&+1 for both state-reward and value- caluee-poicyy networks. The predicted quantities for the state-reward network are rewards f:41 and termination probability Pf.. while the predicted quantities for the palue-policy network are values i and policies मं. The architecture of this prediction function the same for both networks and structured as follows:
Convolution with 64 output channels and stride 1, followed by a ReLu activation.
Convolution with 32 output channels and stride 1, followed by a ReLu activation.
Flatten the 3D vector to a ID vector.
Separate linear layers for each predicted quantity.
In the case of probability outputs, namely the predicted termination probability Pi1, and the predicted policy Reti. we pass the output of the linear layer through softmax layer to ensure valid probability distribution. In addition, the weights and biases of the final linear layers are initialized to zero to stabilize initial training-
For the reward network, we compute the predicted frame, 6 for I 2 1, by applying a decoder on her, for / > 1. The decoder's structure is as follows:
18

---

D residual blocks, each with 128 output channels.
ReLu activation followed by a transpose convolution with 128 output channels and stride 2.
D residual blocks, each with 128 output channels.
ReLu activation followed by a transpose convolution with 128 output channels and stride 2.
ReLu activation followed by transpose convolution with 64 output channels and stride 2.
D residual blocks, each with 64 output channels.
ReLu activation followed by transpose convolution with 3 output channels and stride 2.
All convolutions use kemne size of 4. The decoder outputs tensor of shape 83,84, we concatenate this output with the three most recent true or predicted frames in S141-1 for ( = 1 or 8141-1 for / 1, resulting in final decoded predicted state with shape (12, 84, 84). This approach is adopted to prevent the model from redundantly replicating previous inputs, as the prior three frames already present in the input have been predicted in preceding steps.
84, 84). In Atari,
C.2 Actor-critic
The augmented state consists of the model's hidden states & [hg h. 3D vector of shape (256, 6, 6), and the tree representation Uhree a flat vector of shape (10 + K + 9). The model's hidden state is processed by convolutional neural network (CNN), structured as:
Convolution with 64 output channels and stride 1, followed by ReLu activation.
2 residual blocks, each with 64 output channels.
Convolution with 64 output channels and stride 1, followed by ReLu activation.
2 residual blocks, each with 64 output channels.
Convolution with 32 output channels and stride 1, followed by a ReLu activation.
2 residual blocks, each with 32 output channels.
Flatten the vector to a ID vector.
A linear layer with output size 256, followed by a ReLu activation.
The tree representation is processed by a special recurrent neural network (RNN) architecture, which combines the Long Short-Term Memory (LSTM) (38) with the attention module of the Transformer [39]. This RNN architecture is necessitated by the requirements of the Thinker-augmentted MDP, which relies long-term memory capabilities for plan execution. For instance, 5-step plan in the augmented MDP requires the agent to maintain the plan in memory for 5K = 100 steps, a demand that exceeds the capabilities of traditional LSTM. However, using Transformers solely could lead to instability in RL. 400 due to their disregard of the Markov property in MDP as they process inputs manner. Therefore, our approach merges the LSTT with the attention module of the Transformer to allow for extended memory access while preserving the stability of the training process.
across all steps in the
In detail, the RNN processes the tree representation as follows:
A linear layer with output size 128, followed by a ReLu activation.
A LSTM-Attention network (explained detail below) with output size 128.
A linear layer with output size 128, followed by a ReLu activation.
Finally, the input for the final layer is obtained by concatenating the following:
1. The output from the CNN processing the model's 's hidden state.
2. The output from the RNN processing the tree representation.
3. The one-hot encoding of the last actions (real action, imaginary action, and reset action).
4. The last rewards (both the real reward and the planning reward).
19

---

This input is passed to separate linear layers to generate the final output of the actor-critic network, which includes predicted values and logits of the policies for real actions, imaginary actions, and reset actions. We use a separate head for computing the policies for real and imaginary actions, and select the respective logits based on whether the current step a real step or an imaginary step. To generate valid probabilities, the logits predicted by the linear layers are fed into softmax layer. Additionally, an L2 regularization cost of le-9 is added for the final layer's input and le-6 for the real action's s logits.
LSTM-Atfention Cell This section details the LSTM- LSTM-Attention cell, which serves as the building block of the RNN. The notation here is distinct and should not be conflated with the notation in other sections. The LSTM-Attention cell, while akin to a conventional LSTM,, incorporates an additional input from the attention module, multiplied by a gate variable. The LST cell is denoted by:
(he, Ct, Ke, Vi) LSTMSMMMA LSTMMAAtt(ra,, 1 K V2-1),
(5)
where Iy denotes the input to the cell, h, denotes the hidden state, Gr denotes the input activation, K2 and V2 denotes the key and value vector for the attention module, respectively. The function
LSTM Attn is defined as:
ft (W + Uphe-1 + by);
(6)
4 (Wicst + Uzch bic),
(7)
4 = o (WiaTz + Uzahi-1 + dsa),
(8)
Ot a (WoTe + Uphi-1 + bo),
(9)
oq, K₁, V2 = attn K,-1, V/-1);
(10)
è. tanh(Wa ct + Uch + be);
(11)
Q = ft @ c-11 + E @ E + ih o tanh(ar)
(12)
h; 0r @ tanh(c:),
(13)
where ( (Wy. W Wa,, Wo. We, Uy, Use: Ulai Un. Uc, by, bicr bia b,, be) are parameters with output
shape dm, = 128 and @ denotes element-wise product. The only differences compared to a standard LSTM cell are Equations 8 TO and 12 Equation 8 computes the input gate for the attention module outputs, Equation 10 represents the attention module which will be described below, and Equation 12 adds the gated output from the attention module to the input activation.
The attention function attn K1- 11-1) is based on the multi-hea attention module used in Transformer, with the following modifications: (i) Attention is restricted to the last de 40 steps, equivalent to two stages. This is because maintaining an excessively long memory might complicate learning due the Markov assumption of MDP. (ii) A learned relative positional encoding, akin to Transformer XL. [11, is used to facilitate attention to steps relative to the current step.
For clarity, we first describe single-head attention module. We denote 4: E IRd,, ke E Rdd, v E RRd. as the query, key, and value at step f, respectively, where de = 16 is the dimension of the embedding. The input K E de and V Rd..de represent the preceding d, keys and values, defined as follows:
K,-1 := [ki-d, hi- 1 lz- 177
(14)
Vi-1 := [Dz d, di +1 +.... Ut 1] 177
(15)
We calculate the current qr k2 Ve via linear projection input 2,:
qe WyNe.
(16)
k; WA
(17)
Ur Woree
(18)
where (W,, WR. W.) are parameters with output shape da. This enables us to compute K, and V, by stacking the corresponding is and v. Subsequently, we compute the attention vector Pt E Rd as:
= softmax (Ke + Ur)4: + by
(19)
dd
where U, E de and by Rd are learnable parameters representing the relative position encoding. This differs subtly from the relative position encoding used in Transformer XL, as the fixed attention
20

---

window in our case allows us to utilize fixed-size parameters for encoding relative positions. We initialize U, and b be zero vectors, with the last entry of b set 5, reflecting the prior that the current step should be attended to more. In addition, the softmax here should be masked for steps in the previous episode.
The final output of a singlc-head attention module, denoted as ã E Rd is computed as follows:
ā₂ = VV Pc.
(20)
For the multi-head attention module, we repeat the above single-head attention module d, = 8 times and stack the outputs to yield attuck stack( 2r.2:.... (r,dal), 1), where a, R and d denotes the output of the Hth single-head attention module. The final output of the attention module as is given by applying a linear layer and layer normalization [42] on this stacked output:
LayerNorm (Waa + ba + xx),
(21)
where Zt represents the skip connection, and {Wa, ba] are parameters with output size dm- This concludes the description of the LSTM-Attention cell.
LSTM-Attention Network An LSTM-Attention Network is constructed by stacking d, = LSTM-Attention Cells together. We denote the network's input 24, and use the superscript 1 S n S dn to represent the nh cell in the network. We denote the L.STM-Attention network as:
(ht lida K K.I'd V:lidn) = LSTM Attn Bett__ee(zz=1111 Lida Lida Klddn Villd",, A
(22)
defined by the iterative equation for 1 < n S du: KT, ((") = LSTLLSSMAAA (ATTM_AAt((tac(((r--11..
K2
(23)
and we define hoP : h2-1. the top LSTM-Attention cell's output from the previous step. The final output the network is the top LSTM-Attention cell's output, hg. The input to each LSTM-Autention cell comprises both the network's 's input and the output from the previous LSTM-Attention cell. This design choice was inspired by DRC [14, which suggests that this recurrence method may facilitate better planning.
The rationale for integrating LSTM with the attention module this manner stems from the following reasoning. The LSTM's hidden states can be interpreted as a form of short-term memory, primarily focusing on the current step. In contrast, the attention module acts as a long-term memory unit, as maintains an equal computational path from all past steps. By fusing these two forms of memory, the LSTM is granted the capability to control access to long-term memory, including the ability to disregard it entirely by setting the gate variable to zero. This flexibility enables a focus on the most recent few steps while still preserving access to long-term memory, thus combining the advantages of both types of RNN. Using gradient clipping when applying the SSTM-Attention network is also recommended to stabilize training further.
A significant drawback of the proposed LSTM-Attention network is the sequential computation requires. Unlike the Transformer, which can compute attention across all steps simultaneously. the proposed network has to compute each step one at time due to the recurrent computation of LSTM. Consequently, this leads a much higher computational cost. Nonetheless, this sequential computation limitation may not be as significant in an RL setting, An RL task inherently requires the agent to generate action sequentially, and this sequential computation cannot be avoided even for Transformer. Future studies could explore the potential of the proposed LSTM-Attention network in contexts beyond the Thinker- aakkeerre MDP, such as RL tasks that require long-tern terg memory.
21

---

D Experiment Details
D.1 Hyperparameters
The hyperparameters used in all our experiments are shown in Table We tune 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 model unroll length, planning reward scaling, actor-critic clip global gradient norm, and actor-critic unroll length. In terms of preprocessing for Atari, adhere to the same procedure as outlined in the IMPALA paper [31. with the exception that we do not apply grayscaling.
For the actor-critic baseline implemented on the raw MDP, 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 unrolled trajectory; (ii) 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 I. corresponding to the learning curves shown in Figure(3] 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 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 A A100s, the training durations for the raw MDP, DRC, and Thinker Sokoban are around 1, 2, and 7 days, respectively,
1 steps for every step in the real MDP, there is
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 single Sokoban run on our hardware, primarily due its high replay ratio. MuZero, when conducting 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 MIP (as detailed in Appendix E}. without compromising the performance, the computational time for single Thinker run can be reduced to days.
In Figure we observed that the result for K 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 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 this estimation, thereby neglecting the real state and impeding the learning of useful state representation for the actor. To investigate this speculation, we experimented with masking the tree representation at random rate (with probability 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 potential avenue for subsequent research.
which does not use any imaginary transitions,
0.20), which restored the performance to
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 We include two additional baselines, UNREAL 44 and LASER [45), for comparison. The results of the baselines are taken from their original papers.
22

---

Table Atari scores after 200M frames training for the five selected games with different seeds. Up to 30 no-ops at the beginning of each episode.
Seed
Seed 2
Seed
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 state-of-the-art algorithms. There several possible reasons: (i) did not tune myper-parametten
Atari 2600; (ii) learning to plan may require much larger number of frames to show significance, similar to how MuZero requires 20 billion frames. As the goal of the experiments advantage of the Thinker-augmented MDP over the raw MDP, we leave the problem of mastering advanced RI algorithms, including the
Atari 2600 is to demonstrate the general
Atari 2600 using Thinker for future study. In principle, baselines, could be applied to the Thinker-augmented MDP.
The final scores and the corresponding learning curves for each individual Atari game presented Table and Figure [1] respectively. The Rainbow results are taken from the original paper (35), where the reported score obtained from the best agent snapshots throughout training. The learning curves of Rainbow in Figure II 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], score.
which suggests that these five games are pivotal predicting the median numan-normalln The final scores for these five games are displayed in Table and their individual learning curves can be found Figure IO The results for these five games, as presented 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. more general observation is that Thinker excels in environments with fast-moving objects, such as breakout. As depicted in https://youtu. be/ 0f5255F77 Thinker's model capable of predicting the trajectories of these swift objects with near-perfect precision. We hypothesize that such predictions, essential for planning future actions, pose 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 various angles, sometimes missing the ball its imagination but performing adeptly in the
Table 1:
Solving rate on Sokoban.
Table 2: Final human- normalized score on Atari 2600.
At 2e7 frames
At Se7 frames
Median
Mean
Thinker (Agent policy)
261%
1372%
Thinker
88%
95%
DRC [14
68%
92%
Thinker (Base policy)
160%
807%
Raw MDP
102%
514%
DRC (original paper) 14
80%
93%
Rainbow 35
223%
868%
Dreamer-v3 33
12%
32%
MuZero 5
0%
21%
UNREAL
250%
880%
LASER 45
431%
IMPALA with ResNet
14%
45%
I2A 18
21%
43%
ATrecC
1%
3%
VIN 20
12%
21%
23

---

Bottik Zone
Double Dunk
Name This Game
Phoonix
Obert
0.0
0.5
10
1.3
20
0.0
25
0.0
0.5
10
15
2.2
0.2
0.2
1.0
13
20
160
368
563
Seed 1
Seed 2
Seed 3
Figure 10: Learning curves for Thinker in the five selected games with different seeds. Fach curve is computed as the running average of episode returns over the last 500 episodes.
actual game. Upon identifying 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 games with fast-moving objects.
It is important to highlight that applying direct L2 loss future states instead of a feature loss can undermine these advantages. For example, we observed that in Breakout, model trained with direct L. loss is able 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 few pixels and are challenging predict, the model prone neglect these elements when trained using direct L2 loss. In contrast, by using feature loss, these important objects, such 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-augmenti MDP lacks any mechanism to encourage exploration the real environment, which results in challenges when dealing with these complex exploration games, much like actor-critic applied directly to raw MDP. Future research could explore methods to utilize Thinker to foster exploration, such as using model prediction loss curiosity reward. As exploration is well-studied existing exploration methods could also be directly applied to the Thinker- augmented MDP.
in RL [47]. many
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 in stark contrast games like Sokoban, where single action 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, therehy necessitating the significance of each primitive action. Future research could investigate modifications 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 attributed to the fact that we provide the agent with the model's ' hidden state instead of the true 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 agent on the raw MDP. Another factor could be the impact of unoptimized hyperparameters. As we increased the network's '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. is plausible that the poor performance in certain games stems from these unoptimized hyperparameters.
Finally, as evidenced by the individual learning curves illustrated Figure TT and the median normalzee score in Figure 9 we observe that the slope of the learning curve is generally steeper at the end training compared to other baselines. We posit that mastering the use a model represents additional skill to be acquired, opposed to agent acting This necessitates extended training to learn this skill, however, once this skill is acquired, the final performance generally surpasses that of agent the raw MDP. As such, future work could experiment with a larger training budget than 200 million frames.
raw MDP.
24

---

Table 4: Hyperparameters used in experiments.
Parameter
Value
Thinker-augmentted MDP
Stage length K
20
Maximum search depth L
Augmented discount rate 7
ry k
Planning reward scaling cº
Anneal linearly from to 0
Model
Learning rate
Anneal linearly from le-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
Batch size
128
Model unroll length L
5
Reward loss scaling cx
Termination indicator loss scaling cd
Feature loss scaling c*
10
Value loss scaling e²
0.25
Policy loss scaling e
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
le-3
Entropy regularizer for imaginary and reset actions
5e-5
Environment specific-Sokoban
Discount rate ? Input shape
0.97
(3, 80, 80)
Environment specific-Atari 2600
Discount rate ?
0.99
Grayscaling
No
Action repetitions
Max-pool over last N action repeat frames
2
Input shape
(12, 80, 80)
Frame stacking
4
End the episode when life lost
Yes
Reward clipping
[-1, 1]
Sticky action
No
25

---

1.0
1.
0.5
0.0
Atfentis
Bottle Zone
Beam Rider
Serzerk
1.00
Bowling
Booing
Breakout
Concipode
Chopper Commind
0.50
223
100
1.5
1.0
1.0
0.5
Pshiin Derby
Freewsy
Froshbilte
Gopher
Srav far
MV
0.-
6:29
0:50
0.75
-1.0
8-
0.00
1nl
Hockey
124
Kargarco
160
1.-
3.0
Nunig Master
Marisame Bevenge
Ms. Pecmam
Natte This Garnir
Phosnit
2.3
Pont
Prizate Eye
Obert
Rivernald
/
2.0
-150
-3.5
-3.75
-10
Robottak
Seaquest
Skinng
Solens
M
Space Iivaacess
Star Gurner
Suroound
Nanris
Time Hilo
-05
502
165
16-2
16t
2.5
Yars Revenge
Zacopn
0.0 0.5
1.5
0.2
0.5
3.0
0.3
1.0
35
2.-
3.0
Thinker--ugmeetedd MDP
0.5
Raw MDP
Rainbow
0.0
2.0
Figure 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.
26

---

Table 5: Atari scores after 200M frames of training. Up to 30 no-ops at the beginning of cach 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.00 .00
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
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.66
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 Rye
4234.00
48.41
93.61
94.00
Qhert
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 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
$1.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
Vidco Pinball
533936.50
281210.04
444212.40
512567.33
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
27

---

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 statisties from the tree representation. Specifically, we remove the following statistics:
max [9childrzont> Schildrmot) Techildrrot): Genidteury Schildicury Plshildrcur) groet, groet
mat meam
maten
Ocur,
mas
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 [2
The removal of hints causes the agent to learn at 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 [3]
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 leaf node more easily. This node evaluation allows an agent to evaluate a rollout without simulating the rollout till the end 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 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 [ This is intended to evaluate the benefits of learning imaginary actions. We find that performance significantly deteriorates, suggesting that the imaginary actions play pivotal role in guiding real actions, consistent with the behavioural analysis in Appendix
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 [4] 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
28

---

80%
80%
Solving Rate
Default
Raw MDP
palicy
40%
Default
4.%
Random ima ginary actions
Raw
20%
No hint
mom
0%
No hint & mem
0%
Fromes
Frames
Figure 12: Ablation analysis on hints and mem- ory.
Figure 13: Ablation analysis on base policy and imaginary actions.
L-5
80%
80%
Solving Rate
60%
Solving
40%
40%
Default
20%
Row MDP
planning rewards
0%
0%
No action dist.
Frames
187
Frames
Figure 14: Ablation analysis on maximum search depth model unroll length L..
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 [3)
Planning rewards and action target Lastly, we discuss some non-essential components of the algorithm. To mitigate the sparse rewards the augmented MDP, we added heuristic reward, which 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 the performance. This suggests that the planning reward may not be a critical component for learning to plan. Additional behavior analysis demonstrates that agent trained without planning rewards seems to learn a similar planning behavior to that of an agent trained with planning rewards.
of trained agent in Appendix F
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, 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 [3] 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, 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.
29

---

F Agent Behavior Analysis
In this section, we delve into the behavior learned by the trained agents in the Thinker- aannkered MDP of Sokoban. Our inquiry aims to understand the nature 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 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 hhtpss: //yo:t//uuuuuuu be00A2255R
Though the video is valuable in discerning the plan of an agent, more quantitative analysis is warranted. Does the agent adhere 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 n-step exhaustive search, which involves iterating through every possible multi- step action sequence and choosing the one with the highest rollout return, reset is triggered when the search depth hits n. Similarly, MCTS initiates a reset at 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 box being moved to irvevesibble state. To illustrate this further, examine the distribution of Temporal Difference (TD) errors of the rollout from the root node to node 2n defined as:
8 :-:- +,:,,7771 +
".. + 7 2 0., Uroot
(24)
in) denotes the path from the root node to the node in. This TD error be interpreted as the quality of the rollout relative to the current policy. We compare the distribution of TD errors of the current rollout 6ur 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 [6
where (root, i1, f2, .....
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 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 poor plan, such as reaching an irreversible state. Under such circumstances, the agent exhibits high probability of resetting. This suggests that the agents learned to factor in the current rollout's 's quality in their decision to reset. This approach diverges from planning strategies like MCTS, which mandates reset upon reaching leaf node.
An additional dimension worth investigating 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 T6 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.
With planning remants
Without planning rewards
1.2
not reset reset
not reset resot
vitt planning rewards
1.0
1.0
Without planningg rewards (L=S1
Pubbabity dennhy
Pyy densky
Persettige (10)
Without planning (L-13)
MCTS
50
40
0.2
0.2
20
0.0
0.0
10.0
5.0
2.5
2.5
5.0
10.0
5.0
5.0
TD
Exparsion
Figure 16: Left and Middle: Probability density of TD eтror 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 different maximum search depths L. All statisties are averaged over 20000 stages, each starting from random initial states.
30

---

Coupled with the aforementioned insights into the TD enror 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 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 stage, defined as 6mux == maxq Og where 1 S. k < K denotes the step within a stage, could be used as a proxy to evaluate the effectiveness of the search 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 shown in Figure [8] This tail mostly corresponds to the case of solving or heing closer to solving a level in the imagination where the current policy has high chance of not solving 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 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 rt-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 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 E (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. 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 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 leameed by the agent at the end of training, and we test a range of simulation
'This maximum TD error is proportional to the undiscounted of planning rewards in stage.
31

---

planning Agent
Without planning
Trained Agent
2.0
UCB
2.0
UCB
Probability denulyy
Probabiiity densty
1.5
1.0
1.0
0.5
0.0
0.0
Maximum TD est
Maximum TD error
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.
planning
Without planning
10a
Agent
UCB
UCB
Probability dersity
10"4
Random
Probability density
10-1
Random
10"2
10-2
10
10:4
10
10
10-1
10-4
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.
Percentage action 189
80
60
40
40
40
20
20
20
0
Visit count's ranking
Mean rollout return's ranking
Max rollout return's ranking
Agent with planning reward
Agent without planning reward
MCTS (T=1)
MCTS (T=0.25)
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.
32

---

90%
85%
80%
Solving
75%
70%
Agent Policy
65%
Base Policy
10A
102
10
Number simulations
Figure 20: Final solving rate on Sokoban training levels achieved by MCTS with varying numbers of stage length, compared to the performance of the trained agent policy and the base policy.
steps spanning over [5, 15,10,20,40 40, 160,120, 640, 2280,260).. The results are depicted in Figure 20 We observe that MCTS needs 1280 simulation steps to match the performance of 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- 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 planning algorithm that differs from common handcrafted planning algorithms. These agents learn plan deeply and reset upon encountering poor node, while being able to search for a good plan within limited number of model simulations. 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, 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.
33

---

G Planning Rewards
In the main paper, we use the extended real reward as the augmented reward:
k k
(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 call the planning reward, to help guide the learning of imaginary and reset actions. We define the planning reward f% on step l: to be:
cP
Oroot,
Droot, (-iol,k-1)
if S K if = 1+
(26)
where I denotes the maximum rollout return of the root node step & within stage, as defined after Equation] I and d² 2 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:
K+1
Σ 2=
P c²l (g Froot,0)
(27)
; 2 + ....."""" Frz + n 02g
(28)
where (root, i 22. i5......, } refers to the rollout with the highest return, groo,, in the entire stage. As Ürone is not influenced by imaginary and reset actions, the return of a stage is equivalent to Fit + 7"ig + ff; + * 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 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 good plan. As such, the planning reward signal has to be separated from throughout
the real reward signal Fx instead of being added together. We also decay from 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: , == (f, FP).
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 eтrоrs 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 mask on the planning rewards' TD error, setting it to zero during real step.
Experiments in Appendix 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 not problem, the TD errors induced by real rewards may not be sparse. For example, if the critic learns to use the maximum rollout return goax 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.
34

---

H Relationship with Generalized Policy Iteration and Meta-Reinforcement Learning
H.1
Generalized Policy Iteration Algorithm
Applying an RL algorithm to the thinker-augmentn MDP can be viewed as an instance of 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, crucial question arises: how effective is 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, consider the Phinker-augmenter MDP with fixed base policy @ and its value é 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 random policy. For simplicity, employ the true environment dynamics, represented as (P, It), 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 (ii), we trained the agent from scratch in the Thinker- augmented MDP. Only the tree representation Uiee 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, 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 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 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 K00 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 rollout more precisely, rendering 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 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, requires a moderate number of transitions for an agent to learn this policy improvement operator.
35

---

4.0
3.0
Epikode Ratur
2.0
agent
15
Trained agent K-20
Trained agent K-60
Base policy
1.0
step uxhaustive search (5 simulations
ushaustive
simulations)
MCTS (60 s mulationss
0.5
ststep
4 step exha ustive search (780 simulations)
5-step exhe ustive search (3905 simulatioos)
Frames
Figure 21: Comparison of different policy improvement operators applied to 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.
36

---

H.2 Meta-Reiinfreement Learning
In this section, we explore the relationship between meta-reinforcement learning (meta-RL) algo- rithms [48] and the Thinker algorithm. Meta-RL is a method designed to learn how to perform reinforcement learning. Its objective is develop policy that can quickly adapt to a variety of new tasks derived from a distribution of tasks, using minimal data. While meta-RI. is often used in scenarios involving task distribution, it is also applicable to single-task scenarios our primary focus here.
which
In the meta-RL framework, a trial typically comprises episodes. During trial, the agents initially encounter K spisodes-termed burn-in episodes. Notably, the return within these episodes does not contribute 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-RA RL algorithm 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 not reset across episodes. As such, the knowledge of the trial is embedded in the RNN's hidden states [2 28
If the model perfectly represents the true environment dynamics, 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 maximize the return in real transitions, the objectives of meta-RL and RL agents on the Thinker-augmented MDP coincide. Moreover, the purpose 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-RI., which uses a fixed number of burn-in episodes followed by evaluation episodes, the Thinker-augmenter MDP interleaves the burn-in episodes within a single evaluation episode. In other words, for each step in the evaluation episode with state 8t, we allow agents have few burn-in episodes starting from St.
In meta-RL, 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 learned model. But if the MDP underlying the burn-in episodes, M, is similar to that of the evaluation episode, M, then the meta-RL. algorithm may also learn how to utilize the collected data while accounting for the difference. This is analogous 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 a close relationship between learning to plan and learning to learn. The former employs model each step to decide better actions at that state, while the latter uses episodes within a trial to determine better actions the evaluation episodes. This highlights the converging ideas in meta-learning and planning, both aiming to make efficient use of available information for decision-min makiin.
37

---

I Miscellaneous
L.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 [3]] 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 's model, possibly because this loss works better in a limited data regime where EfficientZero was designed and evaluated. [43] also proposes 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 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 [35] utilizes a sequential VAE. Given the extensive body research in model-basedd RL, we direct readers to the review [55] for detailed exploration of this field.
1.2 Planning Capacity
In this section, we discuss the planning algorithms that can be implemented in the Thinker-augmented MDP. Uninformed search strategies [56] Ch 3.4] that do not require backward search, including Breadth-first search (BFS), Depth-first search (DFS), MCTS can be implemented by the RL agent in the Thinker-augmented MDP. For example, to implement BFS with two actions available, the agent can select (a1, reset), (a2. reset), (a1, not reset), (a1, reset), (01, not reset), (0g, reset), ... as the imaginary actions. If we treat the values as heuristics, then informed search strategies [56] Ch 3.5] that do not require backward search, can also be implemented. For example, A could be implemented by expanding the unexpanded node that has the highest rollout return in the tree, which can be achieved by keeping a record of unexpanded nodes along with the rollout return in the memories.
Nonetheless, as an MDP does not necessarily involve goal state as in planning problems and our trained model does not support backward unrolling, any backward search from goal state, such as bidirectional search [56] Ch 3.4 and partial-order planning [56] Ch 10.4], cannot be implemented by the agent. Agents are also unable implement planning algorithms that involve heuristics beyond value function, such as GRAPHPLAN [56] Ch 10.3] that employs heuristics based on planning graphs. In summary, an RL agent under the thinker-augmentee MDP can implement any forward-base planning algorithms that do not involve heuristics beyond the value functions.
In practice, the learned planning algorithm is very different from the aforementioned handcrafted planning algorithms, as illustrated in Appendix[Fann the video visualization. This is because, unlike traditional planning algorithms whose goal solve the task in the search, the RL agent's primary incentive is to provide useful information for selecting the next immediate real action.
38


Reviews:
[{"review_id": "2H7yCFdq4S", "reviewer": "Reviewer_WUsk", "summary": "This paper presents the Thinker algorithm, a method that augments an MDP such that an agent takes 'imaginary actions' in a learned model prior to executing actions in the environment. This allows the agent to incorporate planning strategies into a policy learned with RL. State-of-the-art performance is achieved on Sokoban and competitive results are shown on Atari.", "questions": "I would appreciate clarification on the items described in weaknesses; some additional questions:\n- Is it correct that the augmentation increases the cardinality of the action space from |A| to 2|A|^2? Does this adversely affect the difficulty of learning a policy?\n- Is there an intuition for why the K=1 case does worse than in the raw MDP in Figure 5?\n- An ablation on the unroll length L is provided in Figure F.3 in the Appendix. Have the authors tested additional values <5 or >10 -- i.e., is there a clear trend here? I am also curious how sensitive the values of L and K are to a particular task.", "limitations": "- The authors have noted some limitations in Section 6, including computational cost and rigid planning for a fixed number of steps, including determinism of the MDP.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written and polished.\n- The appendix is thorough and documents implementation details and architecture details. Source code is provided for reproducibility.\n- The proposed algorithm for augmenting an MDP to induce planning behavior in an agent seems to be original, as is the architecture of the dual network for learning the model. The results on Sokoban and certain Atari tasks also appear to be quite strong.\n- An analysis of the learned behavior of agents (e.g., when resetting happens during planning, what types of imaginary and real actions are taken) is provided in the Appendix. This is useful for characterizing the behaviors of the learned policies.", "weaknesses": "- Comparison is given on Sokoban to DRC and to Rainbow on Atari. Most significantly, I think this work would benefit from comparison to model-based RL baselines to help contextualize how well the Thinker algorithm performs.\n- The method has multiple moving parts and several involved design choices (e.g., what elements to include in the augmented state, non-stationary planning reward, dual network architecture), though ablations are provided for some of these choices in Appendix F.\n- Related to the above point, regarding the dual network. The paper notes that \"one could adopt the same architecture and training methodology used by MuZero\" but that \"[MuZero] suffers from sample inefficiency since it discards information from the future state.\" It would be useful to provide a quantitative comparison/ablation here, where the encoder/unrolling function/prediction function are learned in the manner of MuZero, as described, in order to justify the proposed architecture.\n- The augmented MDP seems to increase the complexity of the MDP (in the state-action space as well as the time-horizon). An expanded discussion on the computational cost of both training and evaluation, compared to using the raw MDP and also to other model-based methods, would be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for the constructive comments and detailed feedback! The weaknesses are addressed as follows:\n\n**Insufficient baselines** - Please refer to (G1) in the global response.\n\n**Multiple moving parts** - Please refer to (G2) and (G3) in the global response for more ablation analysis.\n\n**Ablation on the dual network** - Please refer to (G2)  in the global response.\n\n**Computation cost** - We acknowledge the need for a more thorough discussion on the computation cost, and will add the following discussion in the appendix when discussing the new baselines: \n\n*As discussed in the main paper, Thinker augments the MDP by adding K-1 steps for every step in the underlying MDP. Consequently, 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.*\n\n*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 7 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 days of training. It is noteworthy that 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 F), without compromising the performance, the computational time for a single Thinker run can be reduced to 3 days.*\n\nThe questions raised are addressed as follows:\n\n>Is it correct that the augmentation increases the cardinality of the action space from |A| to 2|A|^2? Does this adversely affect the difficulty of learning a policy?\n\nThe cardinality of the action space is 2|A|, as we have (real action / imaginary action, reset) as the new action set. Note that the real actions on imaginary steps (the first K-1 step in a planning stage) and the imaginary actions on a real step (the K step in a planning stage) are not used (or as Reviewer U2R9 points out, they do not exist), so they can be considered to be the same type of action. As the new action space’s cardinality is still linear in |A|, this shouldn’t adversely affect the difficulty of learning a policy.\n\n>Is there an intuition for why the K=1 case does worse than in the raw MDP in Figure 5?\n\nThe reason for this is that, instead of providing the raw image as input to the actor-critic network for all K steps, we feed in the model's hidden state. While this hidden state might be more effective for predicting future states and rewards of the MDP, it can be more difficult for learning a good policy. This is particularly the case because the model undergoes updates throughout the training process. This approach potentially explains why in a small number of Atari games, the actor-critic applied to the raw MDP occasionally outperforms the one applied to the Thinker-augmented MDP, as highlighted in Appendix E3:\n\n*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.*\n\nFuture research could consider providing both the raw image and the model's hidden state as inputs to the RL agent to potentially address this issue.\n\n>Have the authors tested additional values <5 or >10 -- i.e., is there a clear trend here?\n\nWe have not conducted tests for scenarios where L < 5. However, we have tested the case of L=20 (implying no restriction on search depth given that the number of planning steps, K, is 20) and conducted a brief behavioural analysis. Notably, the learned maximum search depth within a planning stage averages between 16 and 17. This suggests that RL agents typically prefer to unroll a single, extended rollout. In particular, for some simpler Sokoban levels, the RL agent can imagine solving the level in the first real step. Nevertheless, the initial learning speed is notably slower, while the final performance observed is marginally worse. Our conjecture for this trend is that longer training is required for the model to learn to unroll over a span of 20 steps. We did not include this result in the paper as the results are from the early version of Thinker, where the MDP was augmented in a slightly different manner.\n\nWe hope that the answers provided above adequately address the issues raised by the reviewer, and we kindly ask the reviewer to consider adjusting the score based on our responses. We are grateful for the reviewer’s comments and welcome any further questions!\n"}, {"review_id": "pdSSEOPpiC", "reviewer": "Reviewer_aBLR", "summary": "This work presents Thinker, a new model-based reinforcement learning algorithm that achieves state-of-the-art performance on the puzzle video game, Skoban. The main contribution made by the authors is using a dual network architecture which addresses the sample inefficiency problem that other algorithms such as MuZero have. This dual-network setup consists of two sub-networks, namely a stationary and non-stationary network.  The stationary network takes as input the current state and a sequence of raw actions and predicts future states, rewards and episode termination probabilities. The non-stationary network takes as input both the stationary network’s inputs and its predicted next states. \n\nIn addition to the supervised training loss over the four predicted quantities, the authors proposed another L2 loss that encourages the encoded representation of the stationary network’s predicted state to match the encoded representation of the true state observed. The encoder helps the non-stationary network focus on only encoding task-specific information. The static network encodes static policy-independent information, while the non-static network encodes information that changes as the policy changes. Lastly, the encoder is only updated using the loss of the non-static network. The static network uses this encoder, without updating it, and updates its future state predictions to minimise the L2 squared distances in latent space.\n", "questions": "1. On line 245 you state, “ In contrast to these methods, our proposed dual network fits the next-state prediction with a feature loss that enables visualization and prioritizes the learning of task-relevant features.” How does a feature loss enable visualisation? Especially if the agent only needs to remember task-relevant information. When the stationary network predicts the next state, would it not be able to ignore parts of the state that the encoder does not use? Therefore the latent space distance is still close to zero, but the predicted state might have some missing information.\n\n\n2. On line 210 you state, “It is important to note that g is not being optimized when minimizing this loss, and ŝt+l does not receive gradients from the non-stationary network, as the two sub-networks are separately optimized. Further details regarding the model’s training can be found in Appendix B.” If it uses the encoder from the non-stationary network would that not make the stationary network’s updates non-stationary as well? As the policy changes the encoder (g)’s parameters might change. Which in turn changes the next state representation that gets generated.\n\n\n3. On line 192 you state, “Nonetheless, this approach suffers from sample inefficiency since it discards information from the future state, which carries a rich supervised learning signal.” In what way does the dual network approach address the inefficiency problem that you state other methods such as MuZero and DreamerV3 have?\n", "limitations": "This proposed solution currently only works for environments with deterministic dynamics. Furthermore, the planning component is quite computationally expensive, which might make it too expensive to use in many scenarios. \n\nThe main limitation I see is that there has not been enough evidence provided that the dual network setup performs better than alternatives. Especially since the stationary network actually uses components from the non-stationary setup. It would be of great benefit to have additional experiments that show the superiority of this architecture when compared to alternatives (such as predicting every value using a single network).\n", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors show that their proposed system achieves a new state-of-the-art result on Sokoban. Furthermore, they proposed using an L2 distance in a latent representation instead of using the full state. This makes the entire system much more scalable. They further provide a detailed breakdown of the algorithm in the appendices and open source their code. ", "weaknesses": "The main weakness I see is that the use of these dual networks is not well motivated. Along with a better motivation an additional experiment or two is needed to demonstrate that they outperform a single network setup.\n\nFurthermore, it would be helpful to directly compare against open-source versions of models such as DreamerV3 and MuZero. This can help motivate in what ways their approach is superior.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for the constructive comments and detailed feedback! The summary gives a clear description of the proposed dual network, but we would like to point out that the main focus of the paper is the augmented MDP and learning to plan (thus the paper’s title), and dual network is only a part of the overall algorithm. The weaknesses are addressed as follows:\n\n**Insufficient baselines** - Please refer to (G1) in the global response. \n\n**Ablation on the dual network** - Please refer to (G2) in the global response. \n\nThe questions raised are addressed as follows:\n\n> How does a feature loss enable visualisation?\n\nPrevious work [1] has shown that autoencoder training can rely solely on a feature (or perceptual) loss, omitting pixel-space losses, while still achieving good pixel-level reconstruction. We conjecture the reasons for that are (i) feature loss retains the spatial structure as convolutional layers are used, (ii) usually activations from early layers (1-5) are employed so the features are still close to the image itself.\n\nIt is common in computer vision that the network used to compute the feature (or perceptual) loss is pre-trained and fixed. However, Thinker deviates from this norm by utilizing features derived from a network that is concurrently trained to predict policy and values. Nevertheless, we believe the intuition is the same. Moreover, arguably because the feature loss should be satisfied in the dynamic of a concurrently trained network, it incentivizes the network to also predict well in the image space. Finally, it’s worth noting the domain of games is much simpler than the domain of natural images, possibly another reason why feature loss reconstructs images so well in our experiments.\n\nWe agree that the feature loss does not necessarily lead to accurate reconstruction. Indeed one can see some artefacts in the video on the main paper, especially in Atari games. We believe these artefacts are typical for feature loss. For instance, Fig 6 in [1] shows that by using L2 loss, the network omits the purple ball in the reconstructed image, but the feature loss reconstructs it with a dotted grid-like object. The same dotted grid-like object can be seen at the 30:44 timestamp of the video in the paper. Nonetheless, we observe that reconstructed images from the video are of high quality and can be easily interpretable.\n\n> If it uses the encoder from the non-stationary network would that not make the stationary network’s updates non-stationary as well?\n\nYes, the loss function is non-stationary, which changes as the non-stationary network is updated. But we did not find any issues with the non-stationary loss. An ablation analysis of directly using L2 loss on state space is also included in Fig A3, showing that the non-stationary loss yields better performance than the vanilla L2 loss on the raw state space.\n\nWe agree that the name *stationary network* may be confusing as it uses a non-stationary loss. As such, we will revise the naming from *stationary network* to *state-reward network* and *non-stationary network* to *value-policy network* to better reflect their nature.\n\n> In what way does the dual network approach address the inefficiency problem that you state other methods such as MuZero and DreamerV3 have?\n\nWe include the feature loss such that the learned model has to predict future features of the environment. This provides an additional supervised learning signal for the model. The stationary network (or the state-reward network) is designed to predict future states, which then act as intermediary steps or hints for estimating future output statistics. Take, for instance, the task of predicting the last reward after making the moves (up, up, up, right) in Sokoban. Predicting the reward becomes easier based on projected future frames. For example, if the character is anticipated to be on the left of a box after three upward moves, then the likely reward prediction would be one. The primary role of the stationary network (or the state-reward network) is to make these future frame predictions. While, in theory, a neural network could autonomously learn this intermediary step, doing so would necessitate significantly more data.\n\nAs seen in Figure A2, MuZero generally requires a much larger amount of frames to reach a competitive performance. As for DreamerV3, we believe its model is data efficient (it performs well in the Atari 100k benchmark) as it also includes image reconstruction loss in training the model. However, DreamerV3 still performs a lot worse than Thinker in Sokoban - this is likely due to the extensive planning imposed on the Thinker algorithm. Again, Thinker differs from both MuZero and DreamerV3 in both how a model is used (augmented MDP in Thinker v.s. MCTS in MuZero v.s. Dyna-like simulation in DreamerV3) and the model itself (dual network in Thinker v.s. single network in MuZero v.s. VAE in DreamerV3, which does not predict future values or policies that are necessary for an agent to learn in the augmented MDP). \n\nWe hope that the answers provided above adequately address the issues raised by the reviewer, and we kindly ask the reviewer to consider adjusting the score based on our responses. We are grateful for the reviewer’s comments and welcome any further questions!\n\n**References**\n\n[1] Gustav Grund Pihlgren, Fredrik Sandin, Marcus Liwicki. Improving Image Autoencoder Embeddings with Perceptual Loss, arXiv preprint arXiv:2001.03444, 2020\n"}, {"review_id": "euxF1Yo77D", "reviewer": "Reviewer_sy9V", "summary": "This paper presents a model-based RL approach, Thinker, that learns a world model and plans to take actions by generating imaginary rollouts with the learned world model. The paper claims that Thinker is a general method for any RL algorithms and does not rely on any hand-crafted planning algorithm. Experiments were conducted on Sokoban and the Atari 2600 benchmark, showing better performance over model-free RL methods.", "questions": "Please address my questions in the review.", "limitations": "It is hard for me to judge whether limitations have been sufficiently addressed since I do not even understand what new technical advances have been proposed here.", "rating": 5, "confidence": 3, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "1. The basic idea of learning world models and training model-based RL policies with the learned models is technically solid. Experiments indeed show better performance against model-free baselines\n\n2. The method description is clear and it is good to see performance in different domains.", "weaknesses": "1. I am very surprised that Thinker is not compared against model-based RL baselines.\n\n2. The contribution is unclear to me:\n\n    a) I do not understand the claim about not having a hand-crafted planning algorithm. Based on the method described in the paper, you did introduce an MCTS-type of planning algorithm. This claim needs to be explained and justified.\n\n    b) What is the difference between your method and other model-based RL methods? Particularly, it looks similar to MuZero, except that you have a different network architecture for learning the world model. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the questions. The questions raised are addressed as follows:\n\n**Insufficient baselines** - Please refer to (G1) in the global response. \n\n**The difference with MuZero** - Please refer to the first three paragraphs of the global response and reviewer #fWg3’s summary of the algorithm. In particular, *we did not introduce any MCTS in our algorithm*. An actor-critic is trained to select imaginary action (instead of the UCB formula in MCTS) and real action (instead of based on visit counts in MCTS) by maximizing the return in the augmented MDP. An RL agent can learn a variety of planning algorithms depending on its needs. Please refer to Appendix G for the analysis of the agents’ learned planning algorithm, which differs significantly from hand-crafted algorithms such as MCTS."}, {"review_id": "eVBpLw3Xio", "reviewer": "Reviewer_U2R9", "summary": "- This paper introduces a novel algorithm, Thinker, which learns how to plan by interacting with a learned world model using “imaginary actions”—negating the need for hand-designing a planner. The algorithm produces SOTA results on Sokoban, and its efficacy on Atari 2600 is demonstrated as well. ", "questions": "- In line 234, various approaches using gradient updates to plan within a planning stage are cited. These approaches do not appear to be used in evaluation as baselines/ablations—why not?\n- Could some of the improved results be due to the novel duel network architecture, rather than the new algorithm? I don’t see an ablation for this in the main body or supplementary material.\n- How were the output statistics chosen? Could alternatives have been used?\n- I am confused by the phrase “the real action at the first K-1 steps” (132), since I understand that the fist K-1 actions are imaginary. Is it possible the authors are trying to say that we’re not using real actions for the first K-1 steps? It seems like these real actions are not simply “not used”, but they don’t exist in the first place. The structure of the sentence also makes it unclear to me which “imaginary action” we are talking about (I would suggest maybe ditch the last comma for clarity, since this also corresponds to the last step, K).", "limitations": "- The authors note the limitations that (i) the algorithm carries a large computational cost, (ii) only rigid planning is supported, i.e. requiring agent to roll out from the root state and restricts it to planning for fixed number of steps, and (iii) a deterministic environment is assumed. I agree with this authors that these elements are appropriate for future work.\n- No other major limitations stand out to me.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The algorithm design is natural given the motivation of the work.\n- The introduced algorithm has a number of novel components, including the planning stage used over the augmented MDP, and the augmented reward, and the duel network architecture.\n- Thinker produces state-of-the-art results on Sokoban (over previous state-of-the-art, DRC), and in Atari 2600, the Thinker-augmented approach outperforms the non-Thinker-augmented baseline.\n- The writing is clear and relatively concise, the paper is well-structured, the significance of the work is well-communicated and compelling, and the related work section appears to be complete.", "weaknesses": "- The authors note “as the goal of the [Atari 2600] experiments is to investigate the benefits of the augmented MDP, we do not include other common baselines here” (line 307). While this point is understood, the results on Atari would have been more compelling if Thinker were to have outperformed comparable learned planning approaches, e.g. those mentioned in related work.\n- Figure 1 is a bit difficult to read—it is hard to note the difference between each of the successive images. Better visualization to highlight the differentials (and zooming in on the level, maybe) could have been employed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for the constructive comments and detailed feedback! We appreciate the reviewer’s recognition of our work. The weaknesses are addressed as follows:\n\n**Insufficient baselines** - Please refer to (G1) in the global response.\n \n**Clarity of Figure 1** - We will use a bordered window to emphasize the agent and enlarge the images. Additionally, we will update the border color of the rollout image to match the rollout summary on the left.\n\nThe questions raised are addressed as follows:\n\n> In line 234, various approaches using gradient updates to plan within a planning stage are cited. These approaches do not appear to be used in evaluation as baselines/ablations—why not?\n\nThis is because [1, 2] requires perfect information about the MDP, meaning one needs a perfect simulator of the environment, while [3, 4] requires a continuous action space. We focus on the conventional RL setting, where prior knowledge about the MDP is not available, and only consider MDPs with a discrete action space.\n\n> Could some of the improved results be due to the novel duel network architecture, rather than the new algorithm? I don’t see an ablation for this in the main body or supplementary material.\n\nPlease refer to (W3) in the global response for the ablation results on imaginary actions. If we replace the learned imaginary action with a random imaginary action, performance drops significantly. As such, the improved results are not only due to the novel dual network architecture but also the learning of imaginary actions in the augmented MDP. Section G in the appendix contains additional comparisons with MCTS.\n\n> How were the output statistics chosen? Could alternatives have been used?\n\nThe most important output statistics are values and policies, inspired by the following rationale: Values allow evaluating a rollout without the need to roll till the end of the episode, while policies allow the search to be more efficient by providing hints to the agents. Both of these quantities are critical - see Fig F4 in the appendix, which shows learning cannot occur without these two quantities. We do not know if there are any alternatives that can replace these two quantities.\n\n> It seems like these real actions are not simply “not used”, but they don’t exist in the first place.\n\nYes, one can interpret that the real actions at the first K-1 planning steps, or the imaginary actions at the K planning step, do not exist at all. However, since the actor-critic network outputs both real action and imaginary action with different heads on each planning step, so in terms of the code, those actions do exist (but they will be discarded). We acknowledge that it is clearer if we state the first K-1 actions are imaginary actions while the K action is a real action instead of saying the other actions are not used. We will revise the paper to reflect this. Thank you for the suggestion.\n\nWe hope that the answers provided above adequately address the issues raised by the reviewer. We are grateful for the reviewer’s comments and welcome any further questions!\n\n**References**\n\n[1] Thomas Anthony, Robert Nishihara, Philipp Moritz, Tim Salimans, and John Schulman. Policy gradient search: Online planning and expert iteration without search trees. arXiv preprint arXiv:1904.03646, 2019.\n\n[2] Arnaud Fickinger, Hengyuan Hu, Brandon Amos, Stuart Russell, and Noam Brown. Scalable online planning via reinforcement learning fine-tuning. Advances in Neural Information Processing Systems, 34:16951–16963, 2021.\n\n[3] Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal planning networks: Learning generalizable representations for visuomotor control. In International Conference on Machine Learning, pages 4732–4741. PMLR, 2018.\n\n[4] Mikael Henaff, William F Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint arXiv:1705.07177, 2017.\n"}, {"review_id": "iTSSOTF4qm", "reviewer": "Reviewer_9fjR", "summary": "The authors propose Thinker, a model-based RL algorithm that transforms a MDP into one where at each step, the agent is required to generate plans in imagination before selecting a real world action. Then through the RL gradient, the agent learns to plan in order to act. The authors get good results in Atari and Sokoban benchmarks.", "questions": "I think there is a lot of potential in this paper. It needs a good amount of polishing on the writing side to make the idea digestible to a general reader. Next, the experiment section can be improved in a variety of ways - more baselines, ablations, and seeds.\n\nSome additional questions:\n\nWhy is the non-stationary network needed? Couldn't we get value and policy distributions directly by evaluating the policy and value functions on the predicted states?\n\nWhy does K=1 ablation perform poorly? I thought this would be equivalent to the Raw-MDP baseline.\n\nWhat is the agent policy vs. the model policy? Is the agent policy the actor in IMPALA and the model policy the world model's predicted action probabilities?", "limitations": "The limitations section is already quite fair. However, there are some missing pieces. \n1. Computational Cost. MBRL algorithms can be very slow to run in comparison to model free RL.  I would like to see more information about the computational and resource cost of Thinker, and how it compares to the model-free baseline. How many GPUs does it require, how long does it take to run, and what are the main performance bottlenecks? \n2. Hyperparameter tuning.  There seems to be quite a bit of hyperparameters.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- An interesting way of combining planning and learning by integrating planning into the MDP itself. By viewing planning as part of the MDP itself, an RL agent maximizing this MDP will learn how to plan rather than relying on a given planning algorithm to maximize reward. \n- Good results on Sokoban, a hard planning task, and analysis (some issues, see below)\n", "weaknesses": "There are two major concerns I have, one in experimentation and one in writing.\n\n**Experimentation issues**.\n\n- **Lack of baselines** In Sokoban, the work currently compares to one baseline, DRC in Sokoban. DRC is a good choice for a baseline, since it also is a \"learning to plan\" baseline. \n\n  - However, it is worth comparing the learning to plan strategy against hand-crafted planning and other ways of using the world model for RL. For example, Dyna and model-based value expansion are two alternative ways to leverage a world model for RL. I would like to see competitive baselines from these families of MBRL agents (e.g. DreamerV3 and STEVE).\n\n  - The authors claim to omit common baselines since they want to investigate the benefit of the Augmented MDP. If so, I would like to see the Augmented MDP be applied to a few other RL algorithms (Rainbow, DreamerV3, etc.). Then the focus isn't on raw performance, but rather how much Augmented MDP improves over the vanilla MDP. \n\n- **Low number of seeds (3)** reduces my confidence in the results. Please run at least 5 seeds before reporting metrics, and even better, use IQM and 95% bootstrapped confidence intervals. \n\n- **Dual network ablation:** There is no ablation on the dual network itself. What happens if we just use 1 big RNN with the same amount of parameters to predict all quantities? \n\n- **Parameters for Thinker vs. baselines:** Is it possible the gains from Thinker are just from the increase in world model parameters compared to baselines? \n\n\n**Poor Writing and Presentation.** \n\nThis paper unfortunately suffers heavily from writing and presentation issues, which I believe will make it inaccessible to readers. The authors tend to write dense paragraphs about the important  ideas and mechanisms, rather than summarizing points succinctly and relying on visual figures. It seems like the paper is easy to understand for those deep in the subfields of \"learning to plan\" and MuZero style of MBRL. However, readers not in this particular area will find it hard to grasp. \n\nL88 - The paragraph on explaining the planning stage is verbose and hard to understand. The provided Figure 2 hides a lot of the details in the planning stage. \n\nI was initially a bit confused between K and L. Thought that K was the max length of 1 rollout. But instead, it seems like K is the # nodes of a planning tree, which is composed of several rollouts all starting from the same root node. Would have been nice to just have a figure of the planning tree to avoid this confusion, perhaps Figure 2 can be updated to have these details? \n\nSome more explanation of planning reward is needed. How does this avoid penalizing searching for paths with low rollout? Is this because we are using Max instead of mean return?\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments and detailed feedback! We appreciate the reviewer recognizing the potential of our paper while also raising valid concerns regarding the experiments. The weaknesses are addressed as follows:\n\n**Insufficient baselines** - Please refer to (G1) in the global response. \n\n**Number of seeds** - We have added two more seeds to the Sokoban run, bringing the total to five seeds. The results can be found in Fig A4. In general, Thinker’s performance on Sokoban has a very low variance, so having more seeds won’t affect the conclusion. \n\nRegarding Atari, we acknowledge that using more seeds can enhance result reliability. However, the significant costs associated with it (55 runs, with each run requiring 7 days of training on two A100 workstations) prevent us from incorporating a larger number of seeds. As our code is fully public, one can also reproduce our results easily. Moreover, we have not claimed that our results would achieve state-of-the-art performance in Atari, where a more number of seeds might be necessary. Lastly, we kindly ask the reviewer to consider the limited budget we have.\n\n**Dual network ablation** - Please refer to (G2) in the global response.\n\n**Parameters for Thinker vs baselines** - Please refer to (G3) in the global response.\n\n**Poor writing**- We thank the reviewer for pointing out the unclear aspects of the paper, and we will improve the clarity of the paper as follows. \n\n**The paragraph explaining the planning stage** - The paragraph in line 88, along with Figure 2, is intended to give an overview of the planning stage. We acknowledge that it may be too verbose. We will start a new paragraph after *the new real state* on line 96 and remove the line *In each of the K steps, the agent sees....*, which should be better explained in Section 3.2. We also include a better version of Figure 2 that should be clearer, which can be found in Fig A1 of the attached pdf.\n\n**About K and L** - In our paper, K represents the total number of steps taken during a planning stage. Of these steps, K-1 correspond to the 'imaginary actions', which is the same as the number of nodes the agent traverses within a tree. It's worth noting that this doesn't equate to the total number of nodes present in the tree, because the agent can traverse the same node multiple times. On the other hand, L denotes the maximum depth the agent can reach during a rollout. If the agent's search hits this depth (L), it will be reset to the starting point or root node. Typically, the value of L (for instance, L=5) is smaller than K (for instance, K=20). \n\n**Planning reward** - Due to page constraints, we have detailed the planning reward in Section H of the appendix. The reviewer is right: using `max` over `mean` avoids penalizing rollouts with low returns. As elaborated in the appendix, the planning reward acts as a heuristic, guiding the agent to maximize the highest possible rollout return. This gives the agent incentive to try rollouts that have both a low expected return and high variance. If we were to use `mean` instead of `max`, the agent would not have incentives to try rollouts with low expected returns.\n\nThe questions raised are addressed as follows:\n\n> Why is the non-stationary network needed?\n\nFirst, the predicted states may not be perfect, especially during early training, so directly applying a value or policy function to predicted states may lead to inaccurate outputs. As the non-stationary network itself is an RNN, one can recover the MuZero network's way of predicting future policies and values, even when the predicted states are all wrong. Secondly, the RL agent is an RNN that sees both the past imaginary transitions and real transitions. The RL agent receives the augmented state as input in each recurrent step instead of the raw state. Thus, the value and policy function from the RL agent cannot be directly applied to predicted states. Another way of seeing this is that the RL agent's policy acts like a system 2 policy that requires heavy planning, while the model's policy acts like a system 1 policy that only requires a single feedforward pass. The system 1 policy serves as the distilled version of the system 2 policy and guides the search in the model, so they cannot be replaced by one another.\n\n> Why does K=1 ablation perform poorly?\n\nPlease refer to our response to Reviewer WUsk (the second question).\n\n> Is the agent policy the actor in IMPALA and the model policy the world model's predicted action probabilities?\n\nYes.\n\nThe limitations raised are addressed as follows:\n\n**Computational Cost**\n\nWe acknowledge the need for a more thorough discussion on the computation cost and will add a discussion in the appendix. Please refer to our response to Reviewer WUsk regarding computation cost (the last weakness).\n\nThe main bottleneck of the algorithm is training the actor-critic network, and that is why replacing the RNN in the actor-critic network with an MLP can significantly reduce training time. The bottleneck becomes training the model if we use an MLP actor-critic network.\n\n**Hyperparameter tuning**\n\nYes, our method does introduce additional hyperparameters, specifically the number of planning steps, K, and the maximum search depth, L. However, we found that our selected hyperparameters are largely transferable across different tasks. As evidence, our Atari results were achieved by directly applying the hyperparameters from the Sokoban task, with the only modifications being an increase in network depth and a decrease in the learning rate. The same set of hyperparameters is used on 56 environments. This suggests that extensive hyperparameter tuning is not necessary for most environments.\n\nWe hope that the answers provided above adequately address the issues raised by the reviewer, and we kindly ask the reviewer to consider adjusting the score based on our responses. We are grateful for the reviewer’s comments and welcome any further questions!\n"}, {"review_id": "B79x1Y4sOG", "reviewer": "Reviewer_fWg3", "summary": "This paper presents *Thinker*, a method for augmenting an MDP with a learned model so that model-free algorithms can perform “planning” on the augmented MDP. Specifically, Thinker creates a cross-product MDP between the “real world” and a learned world model. This augmented MDP has a union action set that enables an RL agent to either take an action in the real world at a given step, or take an action in the world model. Decision-making is broken down into K-1 steps of acting in the world model (“planning”), followed by a single step of acting in the real model, potentially exploiting information from the K-1 steps of simulation. The authors argue that this allows model-free algorithms to learn how to execute various common planning algorithms.\n\nEvaluation is performed on Sokoban and 57 ALE games. The authors find that IMPALA applied to the Thinker-augmented MDP achieves a higher solve rate for a fixed number of world steps than IMPALA applied to the base MDP, as well as better results than one pre-existing baseline. Analogous results hold on the ALE benchmarks.", "questions": "1. Please explain in more detail why you believe the benefits of this algorithm come from it doing planning (systematic search), as opposed to some other benefit of the Thinker augmentation (i.e. address W1 above).\n2. Is there a crisp way of separating the class of planning algorithms that Thinker can emulate from those that it can’t? (W3)", "limitations": "I felt the limitations were pretty thorough, although I would have liked it if this section listed the expressiveness concern (W3) as well, or at least made more precise claims about the expressiveness of Thinker-augmented model-free RL algorithms.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- **[S1]** The fact that this method merely augments the MDP with some extra transitions (and a learned world model) means that it is both conceptually simple and broadly applicable to many existing reinforcement learning algorithms.\n\n- **[S2]**  Experimental evaluation is thorough. Sokoban is a genuinely difficult benchmark for planning (it’s PSPACE hard in general), and the paper also evaluates across 57 Atari games, which greatly reduces the probability that results were cherry-picked. The ablations in the appendix were also quite thorough.\n\n- **[S3]**  The paper is tackling a problem of significance for the community (planning in reinforcement learning) and was clearly written.", "weaknesses": "- **[W1]** It’s not clear how much the empirical benefits of Thinker come from actually doing search as opposed to just giving the model a more expressive feature space (since the policy can depend on the rollout history at each node), or something else of that nature. The fact that the algorithm’s benefits seem to plateau after K=10-20 planning steps is particularly surprising: traditional planning algorithms might take tens of thousands of node expansions to find a plan of tens to hundreds of steps.\n\n  The ablations go some way towards resolving this issue, particularly the ablations that decrease K and the ablation that removes some of the redundant features from the state space. However, more analysis of what the algorithm is actually doing during the planning phase would be helpful. For instance, the paper could include a qualitative analysis of Sokoban rollouts that tries to explain what kind of planning the algorithm is doing. It could also include an extra ablation in which actions are randomly chosen during the planning phase, rather than being chosen by Impala—does learning the planner with Impala actually help?\n\n- **[W2]** As I understand it, the main contribution of the paper is a way of augmenting MDPs, as opposed to a specific algorithm. Indeed, the paper uses this as a justification for the fact that it only has one third-party baseline for Sokoban and one for Atari. However, the paper only applies the Thinker augmentation to one RL algorithm (Impala). This is not a fatal problem for the paper, but it would be valuable to see the Thinker augmentation applied to other model-free RL algorithms to verify that it actually does work well for algorithms other than Impala.\n\n- **[W3]** The paper appeals to the “generality” and “universality” of the Thinker augmentation, applying that it can express any (?) planning algorithm. I think this is over-claiming: the Gym-like interface available during planning steps makes it hard to implement even a simple algorithm like BFS, and I’m not sure how Thinker would express more complex algorithms like regression (backwards search from the goal) or partial-order planning, not to mention planning algorithms that require complicated heuristics. It would be helpful if the paper was more precise about what planning algorithms Thinker can and can’t express.\n\nNote that W1 is my biggest objection, since I feel it cuts against the core claim of the paper. I would really like to see this addressed in the rebuttal. W3 is a more minor objection but is pretty easy to address, so it would be good to see the paper updated to reflect that (or to see an argument in the rebuttal for why it can’t be done). W2 is a big request for more experiments & so I understand if the authors don’t have time or compute to do this; it’s more of a nice-to-have than something I consider essential to publication.\n\nI’m marking this a weak accept for now due to the issues listed above, but I’m overall pretty happy with the paper, and expect to upgrade my rating later so long as the rebuttal is reasonable & no other reviewers identify major flaws in the experiments or conceptual basis of the idea.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments and detailed feedback! We appreciate the reviewer’s concise summary of the proposed algorithm and acknowledgement of our experiments. The weaknesses are addressed as follows:\n\n**(R1) Algorithm’s benefits seem to plateau after K=10-20 planning steps is particularly surprising**\n\nPlease refer to (G3) in the global response for the ablation results on imaginary actions and section G in the appendix for the analysis of learned agents’ behaviour. We found that RL agents can learn to search much more efficiently compared to other handcrafted planning algorithms, as shown in Fig G5 in the appendix. [1] also shows that a trained neural network can search much more efficiently than MCTS in a supervised learning setting.\n\nNote that if the better performance of Thinker stems merely from the expanded feature spaces of the model, and the imaginary actions from the RL agent don't enhance performance, then one would expect the rollouts in runs without planning rewards to be completely random. This is because, if the planning reward is disabled, the RL agent would only be incentivized to maximize true returns, and the entropy regularization would push all non-useful actions towards a uniform distribution. However, we observe that the rollouts in runs without planning rewards remain specific and meaningful (e.g., at 8:35 timestamp of the video in the paper, the agent imagines solving the task in various ways). The new experimental results in Fig A3 provide more direct evidence of the benefits of the learned imaginary actions.\n\n**(R2) Thinker augmentation applied to other model-free RL algorithms**\n\nWe aimed to demonstrate the general improvement offered by the augmented MDP by running a standard actor-critic algorithm across a wide range of domains (56 different environments). While it would be interesting to explore how the augmented MDP interacts with various RL algorithms, rather than different environments, the substantial computational cost associated with the algorithm prohibits us from doing so. We leave this for future research. Moreover, given that we employed a straightforward actor-critic algorithm without modification, it's reasonable to expect that the augmented MDP would be compatible with more advanced RL algorithms. We greatly appreciate the reviewer’s understanding on this matter.\n\n**(R3) Planning capacity under the augmented MDP**\n\nWe recognize that our claim that an RL agent can express most planning algorithms is misleading, and we apologize for it. We will revise the paper (e.g., replace line 30 *any planning algorithm* with *common forward-based planning algorithms*)  and include the following discussions in the appendix:\n\n*An RL agent under the Thinker-augmented MDP can implement any forward-based planning algorithms that do not involve heuristics beyond the value functions. Uninformed search strategies [2, Ch3.4] that do not require backward search, including Breadth-first search (BFS), Depth-first search (DFS), MCTS can be implemented by the RL agent. For example, to implement BFS with two actions available, the agent can select (action_1, reset),  (action_2, reset), (action_1, not reset), (action_1, reset), (action_1, not reset), (action_2, reset), … as the imaginary actions\\*. If we treat the values as heuristics, then informed search strategies [2, Ch3.5] that do not require backward search, can also be implemented. For example, A\\* could be implemented by expanding the unexpanded node that has the highest rollout return in the tree, which can be achieved by keeping a record of unexpanded nodes along with the rollout return in the memories.*\n\n*Nonetheless, as an MDP does not necessarily involve a goal state as in planning problems and our trained model does not support backward unrolling, any backward search from a goal state, such as bidirectional search [2, Ch3.4] and partial-order planning [2, Ch10.4], cannot be implemented by the agent. Agents are also unable to implement planning algorithms that involve heuristics beyond value function, such as GRAPHPLAN [2, Ch10.3] that employs heuristics based on planning graphs.*\n\n*In practice, the learned planning algorithm is very different from the aforementioned handcrafted planning algorithms, as illustrated in Appendix G and the video visualization. This is because, unlike planning algorithms whose goal is to solve the task in the search, the RL agent's primary incentive is to provide useful information for selecting the next immediate real action.*\n\n\\*We note that searching in this manner is not efficient, as the agent must begin expanding from the root node every time instead of directly starting from the unexpanded nodes (or the *frontier nodes*). Also, implementing BFS requires the agent to have memory, as whether one should reset at a specific node depends on whether all nodes at the same level have been expanded, which is beyond the auxiliary statistics.\n\nThe questions raised are addressed as follows:\n\n> Please explain in more detail why you believe the benefits of this algorithm come from it doing planning (systematic search)...\n\nPlease refer to our answer (R1) above.\n\n> Is there a crisp way of separating the class of planning algorithms that Thinker can emulate from those that it can’t?\n\nPlease refer to our answer (R3) above.\n\nWe hope that the answers provided above adequately address the issues raised by the reviewer, and we kindly ask the reviewer to consider adjusting the score based on our responses. We are grateful for the reviewer’s comments and welcome any further questions!\n\n**References**\n\n[1] Guez, A., Weber, T., Antonoglou, I., Simonyan, K., Vinyals, O., Wierstra, D., ... & Silver, D. (2018, July). Learning to search with MCTSnets. In International conference on machine learning (pp. 1822-1831). PMLR.\n\n[2] Stuart, R., & Peter, N. (2010). Artificial Intelligence A Modern Approach Third Edition.\n"}]

Author rebuttal:
We thank the reviewers for the constructive comments and detailed feedback. We appreciate that most reviewers recognize the novelty and value of our proposed algorithm, Thinker. We note that some reviewers seem to underestimate the novelty, viewing Thinker as merely a variant of MuZero, or focusing overly much on the architecture as opposed to the core novelty of our approach, which is to unify planning and acting within a single augmented MDP (hence the paper’s title). For reference, Reviewer #fWg3 provided a succinct summary of the algorithm, and a clearer illustration of the algorithm (replacing Fig 2 in the paper) is shown in Fig A1 of the attached pdf. 

To clarify the novelty with respect to MuZero,  MuZero can be understood as (i) MCTS and (ii) a learned model, while Thinker can be understood as (i) an RL agent in an augmented MDP and (ii) a learned model with the dual network architecture. Thus, *Thinker does not employ MCTS in any capacity.* While the specific architecture of the model is not our central concern, we do note that the use of a dual network enhances performance, and consider it a significant contribution, just not the main one.

**Thinker is the first work showing that an RL agent can learn to plan with a learned world model in complex environments.** The area of 'learning to plan' is still nascent, largely due to its intrinsic challenges. Prior works either do not satisfy (i) learning to plan, i.e., training a neural network to select imaginary actions, (ii) using a learned model or (iii) being evaluated in a complex environment. IBP [20] shares similar ideas with Thinker, but their models do not predict values and policies, which we found necessary in complex domains (see Fig F4 in the appendix). As such, their work has only been shown to work in simple domains, and (iii) is not satisfied. TreeQN [14] and I2A [15] do not satisfy (i). VIN [19] and DRC [10] do not satisfy (ii). Another line of work [21-24] corresponds to algorithms that perform gradient updates within a planning stage. This requires either a perfect model [21, 24] or simple environments [23], so either (ii) or (iii) is not satisfied. ([22] is an imitation learning algorithm instead of an RL algorithm.)

Common issues raised by reviewers are addressed as follows:

**(G1) Additional Baselines**

We added six new baselines in the Sokoban domain, including Dreamer-v3 [52], MuZero [1], I2A [15], ATreeC [14], VIN [19], and IMPALA with ResNet [35]. We use open-source implementation for Dreamer-v3, with the result being averaged over three seeds. MuZero’s result is taken from [53]. Results of DRC (original paper), I2A, ATreeC, VIN and IMPALA with Resnet are taken from [10]. The results are shown in Fig A2 and Table A1 in the attached pdf. In hard planning domains, planning-orientated RL algorithms, such as DRC, can generally achieve a much better performance than other RL algorithms. And within planning-orientated RL algorithms, DRC usually outperforms others by a wide margin [10]. As such, we only included DRC as the baseline in our submission, but we agree with the reviewers that more baselines would enrich our paper. As for Atari, we added the new baselines of UNREAL [54] and LASER [55] from previous papers, and the results are shown in Table A2. Readers who are interested in more baselines can refer to other papers, as Atari 200m is a common benchmark. Note that the goal of the Atari experiments is to evaluate the benefits of the Thinker-augmented MDP in other domains beyond planning domains, instead of competing with baselines.

**(G2) Ablation on the Dual Network**

We conducted three additional ablation experiments: 

(i) Single network that predicts the relevant quantities but not future states, which is the same as MuZero's model.

(ii) Same as (i), but with the addition of predicting states using L2 loss (though the reviewers did not request this run, we believe it provides valuable insights. The experiment is still running, and we expect to complete the training by Aug 13).

(iii) Dual network with L2 state prediction loss (instead of feature loss).

The result is shown in Fig A3. 

**(G3) Ablation on Imaginary Action**

To address reviewers' concerns that the performance may be due to the model's increased parameter count or additional features, we included an ablation experiment by using random imaginary actions while training the RL agent only on real actions. The result is presented in Fig A3. We also note that the IMPALA with ResNet shown in Fig A2 has 52.5m parameters, whereas Thinker has 8.5m parameters, so increasing the parameter count does not necessarily improve performance.

All aforementioned figures will be incorporated into the revised appendix. With the new experiments included, *our paper now encompasses over 70 distinct runs, spanning evaluations across 56 environments and 16 ablation studies, in addition to the extensive analysis in the appendix.* Most related works on learning to plan have fewer experiments. For example, [24], published in NeurIPS 2021, evaluated their algorithm on two environments with a single baseline. We hope that our new experiments, along with consideration of the computational cost, could address the reviewers’ concerns regarding experiments.

**References**

[1-51] See the main paper’s references.

[52] Hafner, D., Pasukonis, J., Ba, J., & Lillicrap, T. (2023). Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104.

[53] Hamrick, J. B., Friesen, A. L., Behbahani, F., Guez, A., Viola, F., Witherspoon, S., ... & Weber, T. (2020). On the role of planning in model-based deep reinforcement learning. arXiv preprint arXiv:2011.04021.

[54] Jaderberg, Max, et al. "Reinforcement learning with unsupervised auxiliary tasks." arXiv preprint arXiv:1611.05397 (2016).

[55] Schmitt, S., Hessel, M., & Simonyan, K. (2020). Off-policy actor-critic with shared experience replay. In ICML (pp. 8545-8554). PMLR.



Meta-review: After introducing the additional baselines, the reviewers unanimously recommended the acceptance of the paper.
The authors should explicitly compare to model-based approaches, as was done in the rebuttal, and relate to the algorithms mentioned in the rebuttal, including [56].

[56] https://arxiv.org/abs/2202.01108
Decision: Accept (poster)