Title: Thinker: Learning to Plan and Act
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement learning

model-based reinforcement learning

planning

Monte Carlo Tree Search

Markov Decision Process
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

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.

1 Introduction

Model-based reinforcement learning (RL) has significantly enhanced sample efficiency and performance by employing world models, or simply models, to generate additional training data [1, 2], provide better estimates of the gradient [3, 4], and facilitate planning [5]. 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 remarkable success in recent years [7, 8, 5], underscoring the importance of planning.

However, a significant research gap persists in creating methods that enable an RL agent to 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 a handcrafted planning algorithm can handle these tasks, they pose a significant learning challenge for an RL 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, a novel approach that enables the agent to interact with a learned model as an integral part of the environment. The Thinker algorithm wraps a Markov Decision Process (MDP) with a learned model and introduces a new set of actions that enable agents to interact with the model. 1 The manner in which an agent uses a 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

real state

rollout 1

rollout 3

rollout 4

rollout 2

step 1

g:6.73

step 2

g:6.36

step 3

g:6.41

step 4

g:6.05

step 5

g:6.39

step 6

g:6.73

step 7

g:6.36

step 8

g:6.11

step 9

g:6.41

step 10

g:6.45

step 11

g:7.25

step 12

g:7.09

step 13

g:7.48

step 14

g:7.43

step 15

g:7.56

step 16

g:5.47

step 17

g:5.46

step 18

g:5.59

step 19

g:6.01

new real state

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 is 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 [11, 12]. According to this hypothesis, both model-free and model-based behaviors utilize a 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-augmented MDP, yield state-of-the-art performance in Sokoban. Notably, this approach attains a solving rate of 94.5% within 5e7 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 a significant improvement over those using the raw MDP. Visualization results, as shown in Fig 1, reveal the agent's effective use of the model. To summarize, the Thinker algorithm provides the following advantages:

Flexibility: The agent learns to plan on its own, without any handcrafted planning algorithm, allowing it to adapt to different states, environments, and models.

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-based RL algorithm by switching to the Thinker-augmented MDP.

Interpretability: We can visualize the agent's plan prior to its execution, as depicted in Fig 1. This provides greater insight compared to scenarios where the neural network internalizes the model and plans in a black box fashion [14].

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 a novel combination of architecture and feature loss for the model, which is 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], AlphaGo [15], and AlphaZero [7] employ MCTS. VPN [16], TreeQN and ATreeC [17] explore all possible action sequences up to a predefined depth. I2A [18] 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 to imitate actions from expert trajectories generated from MCTS. VIN [20] and DRC [14] are model-free approaches that utilize a special neural network architecture. VIN's architecture aims to facilitate value iteration, while DRC's architecture aims to facilitate general planning. IBP [21], 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, a method that is 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 E). 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.

Another research direction employs gradient updates to refine rollouts [22, 23, 24] prior to taking each real action. These methods start with a 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 a real action. For instance, previous works such as [22], [23], and [24] have proposed to maximize the rollout return using REINFORCE [25], PPO [26], 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.

Thinker is the first work showing that an RL agent can learn to plan with a 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 a learned model or (iii) being evaluated in a complex environment. For example, MuZero [5], VPN [16], TreeQN [17], and I2A [18] do not satisfy (i). VIN [20], DRC [14] 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 [29, 30], in Appendix H.

3 Background and Notation

We consider a Markov Decision Process (MDP) defined by a tuple ( S , A , P, R, γ, d 0 ) , where S is a set of states, A is a finite set of actions, P : S × A × S → [0 , 1] is a transition function representing the dynamics of the environment, R : S × A → R is a reward function, γ ∈ [0 , 1] is a discount factor, and d 0 : S → [0 , 1] is an initial state distribution. Denoting the state, action, and reward at time t by s t , a t , and r t respectively, P ( s, a, s ′ ) = Pr ( s t +1 = s ′ | s t = s, a t = a ) , R ( s, a ) = E [ r t +1 | s t = s, a t = a ] , and d 0 ( s ) = Pr ( s 1 = s ) , where P and d 0 are valid probability mass functions. An episode is a sequence of ( s t , a t , r t +1 ) , starting from t = 1 and continuing until reaching the terminal state, a special state where the environment ends. Letting g t = ∑ ∞ k = t γ k -t r k denote the infinite-horizon discounted return accrued after acting at time t , we are interested in finding, or approximating, a policy π : S ×A → [0 , 1] , such that for any time t ≥ 1 , selecting actions according to π ( s, a ) = Pr ( a t = a | s t = s ) maximizes the expected return E [ g t +1 | π ] . The value function for policy π is v π where for all s ∈ S , v π ( s ) = E [ g t +1 | s t = s, π ] . We adopt the notation ( s t , a t , r t +1 , s t +1 ) for representing transitions, as opposed to ( s t , a t , r t , s t +1 ) , which facilitates a clearer description of the algorithm. In this work, we only consider an MDP with a discrete action set.

✁

Agent

World Model

Imaginary steps (

mes)

Real step

Agent

Environment

Real

Reward

Real State

Augmented State

Real

Ac

/a0

on

4 Algorithm

The Thinker algorithm transforms a given MDP, denoted as M = ( S , A , P, R, γ, d 0 ) , into an augmented MDP, denoted as M aug = ( S aug , A aug , P aug , R aug , γ aug , d aug 0 ) . We call M the real MDP, and M aug the augmented MDP. Besides the real MDP, we assume that we are given (i) a model in the form of ( ˆ P, ˆ R ) that is an approximate version of ( P, R ) and is deterministic, and (ii) a base policy ˆ π : ( S , A ) → [0 , 1] and its value function ˆ v : S → R . 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.

4.1 Augmented Transitions

To facilitate interaction with the model, we add K -1 steps before each step in the real MDP, where K ≥ 1 and denotes the stage length . We typically use K = 20 in the experiments. The first K -1 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 to 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 a stage is shown in Fig 2. Each episode in the augmented MDP is 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 k ∈ { 1 , 2 , ..., K } to denote the augmented step within a stage, and t ∈ { 1 , 2 , ... } to denote the current stage. We omit the subscript t 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 s . We consider a simple method of traversing in a tree. At each imaginary step, the agent decides which child node to visit by imaginary action ˜ a ∈ A and also whether to reset the model to the root node by reset action ˜ a r ∈ { 0 , 1 } . 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 a maximum search depth, L , in the algorithm. This ensures that the agent is forced to reset once its search depth exceeds L , as a learned model may not be accurate for a large search depth. An illustration of the tree traversal is shown in Figure 3.

After the imaginary steps, the agent chooses a real action ˜ a ∈ A in a 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 aug := A×{ 0 , 1 } , and an augmented action by the tuple a aug := (˜ a, ˜ a r ) . During imaginary steps, ˜ a becomes the imaginary action, while ˜ a r decides whether to reset the current node. During real steps, ˜ a becomes the real action, while ˜ a r is not used. The pseudocode can be found in Algorithm 1.

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 -1 predicted

R se

triggered

Forc

d

r

s

t

t igger d

Figure 3: Illustration of tree traversal in imaginary steps, where the stage length K = 6 and the maximum search depth L = 2 . Two actions, a and b , are available. The root node corresponds to the real state s underlying the stage, and the child nodes correspond to the predicted states ˆ s a 1 ,a 2 ,... following the selection of actions a 1 , a 2 , . . . from state s . The augmented actions in the K -1 = 5 imaginary steps are depicted beneath the tree. The first action unrolls the model, whereas the second action determines whether the current node is reset to the root node. The yellow and the orange nodes indicate the current node before and after unrolling (before reset is applied), respectively. A forced reset is triggered on the 3 rd 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 a compact way of representing the tree in a low-dimensional vector. First, we consider how to encode a node compactly, followed by how to use these node representations to represent the tree.

To encode a node, we include the reward, the base policy and its 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 0 and its n -level descendant node i n , we define the rollout return from node i 0 to node i n through the path ( i 0 , i 1 , i 2 , ..., i n ) as:

where ˆ v i denotes the value function applied on the node i and ˆ r i denotes the predicted reward upon reaching node i . That is, ˆ v i := ˆ v (ˆ s i ) , and ˆ r i := ˆ R (ˆ s parent_node ( i ) , a parent_edge ( i ) ) .

For a given node i , we can compute g i,j for all its descendant nodes j that are visited in the tree. We also define g i,i := ˆ r i + γ ˆ v i . We define the mean rollout return g mean i and the maximum rollout return g max i as the mean of g i,j and the maximum of g i,j over all descendant nodes j (including itself) that are visited respectively. The mean and maximum rollout returns of a node's children are useful in guiding actions. For example, the mean rollout return of the child i of a node k 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 in the UCB formula of MCTS. We use g mean child ( i ) and g max child ( i ) to denote the vector of mean and maximum rollout returns for each child of a node i .

The visit count of a node, which represents the number of times the agent traverses that node, may also be useful. For example, a child node that has never been visited may be worth exploring more. We use n child ( i ) ∈ R |A| to denote the vector of visit count for each possible child of a node i .

Altogether, we represent a node i by a vector u i ∈ R 5 |A| +2 , called node representation , that contains:

ˆ r i , ˆ v i ∈ R , ˆ π i ∈ R |A| : The core statistics, composed of the reward, the value, and the base policy of the node i . Here, ˆ π i := (ˆ π ( s i , a 1 ) , ˆ π ( s i , a 2 ) , ..., ˆ π ( s i , a |A| )) .

one_hot ( a i ) ∈ R |A| : The one-hot encoded action that leads to node i .

g mean child ( i ) , g max child ( i ) , 1 K n child ( i ) ∈ R |A| : 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 representation u root and the current node's representation u cur to represent the tree (we use the subscript root to denote the root node, which is unchanged throughout a

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 u root and u cur, 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 u tree ∈ R 10 |A| + K +9 , called tree representation . We may also include the real state s or the current node's predicted state ˆ s cur in the augmented state, providing a 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 a terminal state, value estimates can be utilized to assess the potential of a 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 a compact summary of all rollouts. Compared to the concatenated predicted states, which have a dimension of K |S| , the tree representation, typically with fewer than 100 dimensions, is much easier to 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 is given for imaginary steps, and the reward for real steps is identical to the real reward from the real MDP. In other words, r aug k = 0 for k ≤ K and r aug K +1 = r , where r denotes the real reward. We set the augmented discount rate γ aug := γ 1 K to account for the K -1 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 g max root , 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 RL agents can indeed learn to improve the base policy (see Appendix H). As such, we only need to 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, ˆ R ) and (ii) a base policy ˆ π and its value ˆ v . To remove both assumptions, we propose employing a single, unified model that assumes the roles of ˆ P, ˆ R, ˆ π , and ˆ v , learning directly by fitting the real transitions. Specifically, we introduce a novel model, the dual network , designed for this purpose.

The dual network consists of two RNNs which we refer to as the state-reward network and the value-policy network . The state-reward network's input consists of the root state s t and an action sequence a t +1: t + L , and the network predicts future states ˆ s t +1: t + L , rewards ˆ r t +1: t + L , and termination probability ˆ p d t +1: t + L . Meanwhile, the value-policy network takes both the state-reward network's input and its predicted states ˆ s t +1: t + L as inputs, and the network predicts future policies ˆ π t : t + L and values ˆ v t : t + L . We can also view the dual network as an RNN with h = ( h sr , h vp ) being the hidden state, where h sr and h vp denote the hidden states of the two sub-RNNs. The relevant statistics required in the augmented state can be computed by unrolling the dual network using the root node and the action sequence leading to the current node.

To train the model, we first store all real transitions ( s t , a t , r t +1 , d t +1 ) in a replay buffer, which corresponds to the state, action, reward, and termination indicator. We then sample a sequence of real transitions with length L +1 from the buffer. The two sub-RNNs in the dual network are trained separately. For the state-reward network, the following loss is used:

where c r , c d , c s ≥ 0 are hyperparameters that modulate the loss strength. While a straightforward choice for the state prediction loss metric, L state , is the L2 loss - defined as L state ( s ′ , s ) = || s ′ -s || 2 2 -this approach could cause the model to allocate capacity towards predicting non-essential features. To address this issue, we suggest using the feature loss instead of the L2 loss. Specifically, L state ( s ′ , s ) = || g ( s ′ ) -g ( s ) || 2 2 , where g represents the encoding function of the value-policy network. As this encoding function concentrates on features relevant to action, the feature loss encourages the model to prioritize predicting the most crucial features.

For the value-policy network, the following loss is used:

where c v , c π ≥ 0 are hyperparameters and ˆ v target t + l denotes the multi-step return. One may also use the real action distribution π t place of one_hot ( a t + l ) for a better learning signal, though this requires passing the real action distribution instead of only the actions to the augmented MDP.

The design choice of the dual network aims to improve data efficiency while also ensuring inaccurately predicted states do not negatively affect the learning of the base policy and values. If the predicted states are inaccurate, the value-policy network, being an RNN, can learn to ignore these states, relying solely on the root state and action sequence. Conversely, if the predicted states are accurate, the valuepolicy network can effectively learn to make predictions based on these states. Consider, for instance, the task of predicting the value after executing several actions in Sokoban. If the predicted state suggests the level is close to completion, the predicted value should naturally be higher. Although, in theory, a standalone RNN might autonomously learn this intermediary step, in practice, it would likely require much more data. We explore the relationship of this dual network with models from other model-based RL algorithms in Appendix I.

With the dual network, the augmented state can also include the hidden state of the model at either the current node, h cur, or the root node, h root. These hidden states, usually in much lower dimensions, might be easier to process than the real state s or the current node's predicted state ˆ s cur . In our experiments, we use ( u tree , h cur ) as the augmented state, based on the intuition that the current hidden state is likely the most informative among the others. 2 This completes the specification of the Thinker algorithm, which takes an MDP as input and transforms it into another MDP.

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

Thinker's performance at 5e7

0

1

2

3

4

5

Frames

1e8

0%

20%

40%

60%

80%

100%

Solving Rate

Thinker

ResNet

DRC

I2A

DRC (original paper)

ATreeC

VIN

MuZero

5 Experiments

For all our experiments, we train a standard actor-critic algorithm, specifically the IMPALA algorithm [31], on the Thinker-augmented MDP. Although other RL algorithms could be employed, we opted for the IMPALA due to its computational efficiency and simplicity. The actor-critic's network uses an RNN for encoding the tree representation and a convolutional neural network for encoding the model's hidden state. These encoded representations are concatenated, then passed through a linear layer to generate actions and predicted values. We set the stage length, K , to 20 , and the maximum search depth or model unroll length, L , to 5 . More experiment details can be found in Appendix D.

5.1 Sokoban

We selected the game of Sokoban [32, 18], a classic puzzle problem, as our primary testing environment due to its inherent complexity and requirement for extensive planning. In Sokoban, the objective is to push all four boxes onto the designated red target squares, as depicted in Figure 1. We used the unfiltered dataset comprising 900,000 Sokoban levels from [14]. We train our algorithm for 5e7 frames, while baselines in the literature typically train for at least 5e8 frames. Despite this discrepancy, we present learning curves plotted against frame count for direct comparison.

The learning curve for the actor-critic algorithm applied to the Thinker-augmented MDP is depicted in Fig. 5, with results averaged across five seeds. We include seven baselines: DRC [14], Dreamer-v3 [33], MuZero [5], I2A [18], ATreeC [17], VIN [20], and IMPALA with ResNet [31]. In Sokoban, DRC stands as the current state-of-the-art algorithm, outperforming others by a significant margin, likely due to its well-suited prior for the game. The 'DRC (original paper)' result is sourced directly from the paper, whereas 'DRC' denotes our replicated results.

Thinker surpasses the other baselines in terms of performance in the first 5e7 frames. At 2e7 frames, Thinker solves 88% of the levels, while DRC (original paper) solves 80% and the other baselines solve at most 21%. At 5e7 frames, Thinker solves 95 % of the levels, while DRC (original paper) solves 93% and the other baselines solve at most 45%. These results underscore the enhanced planning capabilities afforded by the Thinker-augmented MDP to an RL agent.

To understand the benefits of planning, we evaluate a policy trained with K = 20 on the ThinkerAugmented MDP, by testing its performance with different K values in the same augmented environment during the testing phase. This variability allows us to control the degree of planning and the result is shown in Figure 6. The figure also depicts the performance of the base policy, which can be regarded as a distilled policy devoid of planning. We observe a significant performance improvement attributable to planning, even at the end of training.

The agent's behavior, visualized in Figure 1, illustrates that it learns to use the model for selecting better actions. Interestingly, it appears that the agents learn a planning algorithm that diverges from traditional n -step-exhaustive search and MCTS. For instance, the agent chooses real actions based on

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

K at testing

0%

20%

40%

60%

80%

100%

Solving Rate

Agent policy

Base policy

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

Dual net & feature state loss

Dual net & L2 state loss

Dual net & no state loss

Single net & L2 state loss

Single net & no state loss

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

K=20

K=5

K=2

K=1

Raw MDP

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Frames

1e8

0%

50%

100%

150%

200%

250%

Median h man-normalized score

Thinker-a gmented MDP

Raw MDP

Rainbow

both visit counts and rollout returns, and the agent learns to reset upon entering an irreversible state, contrasting the MCTS strategy of resetting at a leaf node. Further analysis on the trained agent's behavior can be found in Appendix F.

We conducted an ablation analysis to examine the impact of planning duration, varying the stage length K across { 1 , 2 , 5 , 10 , 20 } during both training and testing. When K = 1 , the agent does not interact with the model at all. The outcomes of this ablation analysis are presented in Figure 7. The results suggest that K = 10 already gives the optimal performance. This is in stark contrast to the 2000 simulations required by MCTS to achieve good performance in Sokoban [19].

As a comparison, we train the same actor-critic algorithm on the raw MDP, employing a similar network architecture used by the model to encode the raw state. Surprisingly, the result for K = 1 is worse than the raw MDP. A detailed explanation of the reasons behind this phenomenon can be found in Appendix D. However, employing a marginally larger K , such as 2, greatly improves the performance on the Thinker-augmented MDP, surpassing that of the raw MDP.

We use the dual network as the world model, trained with feature loss in state prediction. We consider several variants of the model, including (i) a dual network that uses L2 loss in state prediction, (ii) a dual network that has no loss in state prediction, (iii) a single network that uses L2 loss in state prediction, (iv) a single network that does not predict state. The single network refers to a single RNN that predicts the relevant quantities, and (iv) is equivalent to MuZero's model. The results of these four variants are shown in Figure 8. We observe that both the dual network architecture and the state loss are critical to performance. However, the use of feature loss instead of L2 loss gives marginal performance improvement, likely because the L2 loss' prior that each pixel is equally important suits well in Sokoban. Nonetheless, this prior may not be suitable for other games like Breakout.

Further ablation studies in Appendix E offer insights into the various components of Thinker, including the role of hints in tree representation, the agent's memory, base policy, planning reward, etc. For example, either the hints or an RNN agent is sufficient to achieve good performance, as the hints summarize the rollouts, and an agent with memory can learn to summarize rollouts on its own. However, if both hints and memory are omitted, the agent fails to learn to plan, as it cannot recall past rollouts. In addition, if the base policy and its value are also omitted, then an RNN agent still fails to learn, suggesting the necessity of a heuristic in learning to plan.

5.2 Atari 2600

Finally, we test our algorithm on the Atari 2600 benchmark [34] using a 200M frames setting. The performance of the actor-critic algorithm on both the Thinker-augmented MDP and raw MDP is evaluated using the same hyperparameters as in the Sokoban experiment, but with a deeper neural network for the model and an increased discount rate from 0.97 to 0.99.

The learning curve, measured in the median human-normalized score across 57 games, is displayed in Figure 9. We evaluate the agent in the no-ops starts regime at the end of the training, and the learning curve is shown in Fig 9. The median (mean) normalized score for the actor-critic algorithm applied on the Thinker-augmented MDP and the raw MDP are 261% (1372%) and 102% (514%), respectively, underscoring the advantages of the Thinker-augmented MDP over the raw MDP. For context, the median (mean) normalized score of Rainbow [35], a robust baseline, is 223% (868%).

The advantages of the Thinker algorithm are particularly pronounced in certain types of games, especially shooting games such as chopper command , seaquest , and phoenix . It is plausible that the agents benefit from predicting the trajectories of fast-moving objects in these shooting games. However, in some games, particularly those where a single action does not produce significant effects, the Thinker algorithm appears to offer no advantages. For instance, in the game qbert , the decision to move to a different box is made only once every five or more actions. This effectively transforms any 5-step plan into a single-step plan, diminishing the algorithm's benefits. Detailed experiment results on the Atari 2600 are available in Appendix D.

Many interesting behaviors learned by the agent are visualized in https://youtu.be/ 0AfZh5SR7Fk . We also observe that the predicted frames in the video are of high quality, showing that the feature loss can lead to the emergence of interpretable and visualizable predicted frames, despite the absence of an explicit loss in the state space. We attribute this capability to using a convolutional mapping g in the feature loss, as convolutions can preserve local features.

6 Future Work and Conclusion

While the Thinker algorithm offers many advantages, it also presents several limitations that provide potential areas for future research. Firstly, the Thinker algorithm carries a large computational cost. A single step in the original MDP corresponds to K steps in the augmented MDP, in addition to the overhead of model training. Secondly, the algorithm currently enforces rigid planning, requiring the agent to roll out from the root state and restricts it to planning for a fixed number of steps. Thirdly, the algorithm employs a deterministic model to facilitate tree traversal, which may not work well for a stochastic environment.

Future work could focus on addressing these limitations, such as enabling more flexible planning steps and integrating uncertainty into the model. Additionally, exploring the application of the algorithm in multi-environment settings is a promising direction. We hypothesize that the advantages of a learned planning algorithm over a handcrafted one will become more pronounced in such scenarios. Exploring how other RL algorithms perform within the Thinker-augmented MDP is an additional direction for future work.

The history of machine learning research tells us that learned approaches often prevail over handcrafted ones. This transition becomes especially pronounced when large amounts of data and computational power are available. In the same vein, we surmise that learned planning algorithms will eventually surpass handcrafted planning algorithms in the future.

References

Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations , 2021.

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.

Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. Advances in neural information processing systems , 28, 2015.

Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11) , pages 465-472, 2011.

Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

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.

David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature , 550(7676):354-359, 2017.

David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science , 362(6419):1140-1144, 2018.

Fan-Ming Luo, Tian Xu, Hang Lai, Xiong-Hui Chen, Weinan Zhang, and Yang Yu. A survey on model-based reinforcement learning. arXiv preprint arXiv:2206.09328 , 2022.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

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.

MDGyörgy Buzsáki. The brain from inside out . Oxford University Press, 2019.

Henry H Yin and Barbara J Knowlton. The role of the basal ganglia in habit formation. Nature Reviews Neuroscience , 7(6):464-476, 2006.

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.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in neural information processing systems , 30, 2017.

Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning. arXiv preprint arXiv:1710.11417 , 2017.

Théophane Weber, Sébastien Racaniere, David P Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al. Imagination-augmented agents for deep reinforcement learning. arXiv preprint arXiv:1707.06203 , 2017.

Arthur Guez, Théophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wierstra, Rémi Munos, and David Silver. Learning to search with mctsnets. In International conference on machine learning , pages 1822-1831. PMLR, 2018.

Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural information processing systems , 29, 2016.

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:1707.06170 , 2017.

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.

Mikael Henaff, William F Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint arXiv:1705.07177 , 2017.

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.

Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning , 8(3-4):229-256, 1992.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv:1611.05763 , 2016.

Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. Rl2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779 , 2016.

Richard Bellman. Dynamic Programming . Princeton University Press, Princeton, NJ, USA, 1957.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain 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.

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.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

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-279, 2013.

Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence , volume 32, 2018.

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.

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 Representations , 2018.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8):1735-1780, 1997.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, 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.

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 arXiv:1901.02860 , 2019.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Jessica B Hamrick, Abram L. Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Holger Buesing, Petar Veliˇ ckovi´ c, and Theophane Weber. Onthe role of planning in model-based deep reinforcement learning. In International Conference on Learning Representations , 2021.

Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations , 2017.

Simon Schmitt, Matteo Hessel, and Karen Simonyan. Off-policy actor-critic with shared experience replay. In International Conference on Machine Learning , pages 8545-8554. PMLR, 2020.

Matthew Aitchison, Penny Sweetser, and Marcus Hutter. Atari-5: Distilling the arcade learning environment down to five games. arXiv preprint arXiv:2210.02019 , 2022.

Pawel Ladosz, Lilian Weng, Minwoo Kim, and Hyondong Oh. Exploration in deep reinforcement learning: A survey. Information Fusion , 85:1-22, 2022.

Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028 , 2023.

Tom Zahavy, Zhongwen Xu, Vivek Veeriah, Matteo Hessel, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh. A self-tuning actor-critic algorithm. Advances in neural information processing systems , 33:20913-20924, 2020.

Zhongwen Xu, Hado P van Hasselt, Matteo Hessel, Junhyuk Oh, Satinder Singh, and David Silver. Meta-gradient reinforcement learning with an objective discovered online. Advances in Neural Information Processing Systems , 33:15254-15264, 2020.

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.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 15750-15758, 2021.

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems , 31, 2018.

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.

Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based reinforcement learning: A survey. Foundations and Trends® in Machine Learning , 16(1):1-118, 2023.

Russell Stuart and Norvig Peter. Artificial intelligence a modern approach third edition, 2010.

A Algorithm

Algorithm 1 implements the step function of the Thinker-augmented MDP, assuming we have access to a given model ( ˆ P, ˆ R ) , a base policy ˆ π , and its value ˆ v . We assume ˆ P is deterministic and use s ′ = ˆ P ( s, a ) to denote the predicted next state s ′ from state s and action a .

To compute the hints (i.e. mean rollout return, max rollout return, and visit counts) in the tree representation, we need to maintain a tree with each node storing the predicted reward ˆ r i , value ˆ v i , and visit count n i , but the details of updating the tree are omitted in the pseudocode for clarity.

It is important to note that since ˆ P is deterministic, one can store the computed node in a buffer to avoid re-computing it when the agent traverses the same node again within the same stage.

In the full Thinker-augmented MDP, both the model and the base policy with its values are learned by an RNN model. The code only needs slight modifications as follows:

After line 7, we insert the real transition in the form of ( s old , ˜ a, r, d ) , where s old denotes the previous real state, into a replay buffer. A parallel thread trains the model by sampling from the replay buffer.

Anode is represented by the hidden state of the RNN model, not the real state. Consequently, line 8 should initialize the root node with the first hidden state h of the RNN model. Line 17 updates the new current node to 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.

Implementation details Several key implementation details are worth noting. First, when transitioning into a new stage, we carry 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

subsequent stages. However, our experiments show that carrying the tree has minimal impact on performance. Second, for clarity, the planning reward is not shown in the pseudocode, but details can be found in Appendix G. Lastly, if the model predicts episode termination at node i , defined by ˆ p d i > 0 . 5 , we manually set the content of all its descendant nodes j to ˆ r j = 0 , ˆ π j = ˆ π i , ˆ v j = 0 , ˆ s j = 0 , and h j = h i , 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 u cur and the root node representation u root are included in the tree representation. Denoting the current search depth by d ∈ { 0 , 1 , ..., L } , the full tree representation u tree ∈ R 10 |A| + K +9 contains:

u root , u cur ∈ R 5 |A| +2 : The root node representation and the current node representation.

g root , cur ∈ R : Current rollout return.

g root , cur -γ d +1 ˆ v cur ∈ R : Sum of discounted rewards collected in the current rollout.

g max root ∈ R : Maximum rollout return of the root node.

γ d +1 ∈ R : Trailing discount factor.

one_hot ( k ) ∈ [0 , 1] K : Current step in a stage, represented in one-hot encoding.

i reset ∈ [0 , 1] : Indicator of whether the current node has just been reset to the root node.

And the following five inputs can be included in the augmented state s aug t,k :

u tree ∈ R 10 |A| + K +9 : The tree representation.

s ∈ S : The real state underlying the stage.

ˆ s cur ∈ S : The predicted real state of the current node.

h root ∈ H : The hidden state of the RNN model at the root node.

h cur ∈ H : The hidden state of the RNN model at the current node.

We used s aug t,k = ( u tree , h cur ) in our main experiments. Our code enables users to decide which inputs to include as part of the augmented state in the Thinker-augmented MDP.

A.2 Code

The complete code is accessible at https://github.com/stephen-chung-mh/thinker . We implemented the Thinker-augmented MDP in batches using Cython [36], a Python language extension that optimizes performance by interfacing with C. This ensures that the primary computational burden is attributed to 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-augmented MDP to share the same model.

Our implementation of the Thinker-augmented MDP aligns with the OpenAI Gym interface [10], facilitating easy integration with existing RL frameworks. The code allows customization of various parameters such as stage depth K , maximum search depth L , and the model's learning rate, among others. Included in the repository is the code for the IMPALA algorithm applied to the Thinkeraugmented MDP, reproducing the experimental results presented in the paper.

B Model Training Procedure

B.1 Collection of Training Samples

A transition tuple x t := ( s t , a t , r t +1 , d t +1 ) is stored in a replay buffer during each real step, which corresponds to the state, action, rewards, and termination indicator. The replay buffer is a first-in-firstout buffer with a capacity of 200,000 transitions. To train the model, a batch of sequences x i : i +2 L is randomly sampled from the buffer, where L represents the model rollout length. We use a batch size of 128 and a model unroll length L = 5 .

Consistent with MuZero [5], prioritized sampling [37] is used for sampling the sequence x i : i +2 L . Each transition tuple is assigned a priority p i . Newly added transition tuples are given the maximum priority value in the buffer. During training, the priority p i of the first transition in a sampled sequence x i : i +2 L is updated to the L1 loss of the value, | ˆ v i -v target i | . Transitions are sampled with a probability of P ( i ) = p α i ∑ j p α j , where α = 0 . 6 . Each sampled transition sequence's loss is weighted by the importance weight w i = ( 1 NP ( i ) ) β , where N is the total number of transition tuples in the buffer, and β is annealed from 0 . 4 to 1 . 0 throughout training.

We employ a replay ratio of 6 , implying that model training is paused if the ratio of 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 are trained separately, which leads to the following implications. First, when training the state-reward network, the g within feature loss for state prediction is not optimized since it belongs to the value-policy network. This prevents g from degeneration. Second, when training the value-policy network, gradients cannot flow through ˆ s t + l towards the state-reward network since ˆ s t + l is considered as the input to the value-policy network. This ensures that the predicted state aligns with the actual future state.

The multi-step return ˆ v target t + l in Equation 3 is defined by:

Instead of using the predicted value ˆ v t +2 L , we opted for the mean rollout return (at the end of the stage) g mean root ,t +2 L as the bootstrap target 3 . We found that it yields more stable training in our experiments. This stability likely stems from averaging across multiple rollouts starting from s t +2 L , rather than relying directly on the predicted value.

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, s t , 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 s t to the encoder, we concatenate it with the one-hot encoding of the action a t -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 a ReLu activation.

D residual blocks, each with 64 output channels.

Convolution with 128 output channels and stride 2, followed by a ReLu activation.

D residual blocks, each with 128 output channels.

Average pooling operation with stride 2.

D residual blocks, each with 128 output channels.

Average pooling operation with stride 2.

All convolutions use a kernel size of 3. The resulting output shape for both Sokoban and Atari 2600 is (128, 6, 6). We use D = 1 for Sokoban and D = 2 for Atari 2600.

We denote g sr ( s t ) and g ( s t ) as the encoder's output of s t for the state-reward network and the value-policy network respectively, and the dependency on a t -1 is omitted for clarity.

For the state-reward network, we define the hidden state h sr t := g sr ( s t ) . To obtain h sr t + l +1 from ( h sr t + l , a t + l ) with l ≥ 0 , we first concatenate h sr t + l with the one-hot encoding of the action a t + l , tiled to (6, 6), channel-wise. This concatenated output is passed to the unrolling function, which consists of 4 D +1 residual blocks, each having an output channel of 128. The output of the unroll function is h sr t + l +1 .

For the value-policy network, in addition to the previous hidden state, the unrolling function should also take the encoded predicted frame g (ˆ s t + l ) as input, so the unrolling procedure is slightly modified. First, we define h vp t -1 as a zero vector of shape (128, 6, 6). To obtain h vp t + l from ( h vp t + l -1 , a t + l -1 , g (ˆ s t + l )) with l ≥ 0 (replace ˆ s t + l with s t + l when l = 0 here), we first concatenate all three together channel-wise, with a t + l -1 tiled to (6, 6). This concatenated output is passed to the unrolling function, which consists of 4 D +1 residual blocks, each having an output channel of 128. The output of the unroll function is h vp t + l .

To compute the predicted quantities, we apply a prediction function to h t + l for both state-reward and value-policy networks. The predicted quantities for the state-reward network are rewards ˆ r t + l and termination probability ˆ p d t + l , while the predicted quantities for the value-policy network are values ˆ v l and policies ˆ π l . The architecture of this prediction function is the same for both networks and is 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 1D vector.

Separate linear layers for each predicted quantity.

In the case of probability outputs, namely the predicted termination probability ˆ p d t + l and the predicted policy ˆ π t + l , we pass the output of the linear layer through a softmax layer to ensure a valid probability distribution. In addition, the weights and biases of the final linear layers are initialized to zero to stabilize initial training.

For the state-reward network, we compute the predicted frame, ˆ s t + l for l ≥ 1 , by applying a decoder on h sr t + l for l ≥ 1 . The decoder's structure is as follows:

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 a transpose convolution with 64 output channels and stride 2.

D residual blocks, each with 64 output channels.

ReLu activation followed by a transpose convolution with 3 output channels and stride 2.

All convolutions use a kernel size of 4. The decoder outputs a tensor of shape (3 , 84 , 84) . In Atari, we concatenate this output with the three most recent true or predicted frames in s t + l -1 for l = 1 or ˆ s t + l -1 for l > 1 , resulting in a 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 are already present in the input or have been predicted in preceding steps.

C.2 Actor-critic

The augmented state consists of the model's hidden states h = [ h sr k , h vp k ] , a 3D vector of shape (256 , 6 , 6) , and the tree representation u tree, a flat vector of shape (10 |A| + K +9) . The model's hidden state is processed by a convolutional neural network (CNN), structured as:

Convolution with 64 output channels and stride 1, followed by a ReLu activation.

2 residual blocks, each with 64 output channels.

Convolution with 64 output channels and stride 1, followed by a 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 3D vector to a 1D 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-augmented MDP, which relies on long-term memory capabilities for plan execution. For instance, a 5-step plan in the augmented MDP requires the agent to maintain the plan in memory for 5 K = 100 steps, a demand that exceeds the capabilities of a traditional LSTM. However, using Transformers solely could lead to instability in RL [40] due to their disregard of the Markov property in MDP as they process inputs across all steps in the same manner. Therefore, our approach merges the LSTM with the attention module of the Transformer to allow for extended memory access while preserving the stability of the training process.

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 in 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:

The output from the CNN processing the model's hidden state.

The output from the RNN processing the tree representation.

The one-hot encoding of the last actions (real action, imaginary action, and reset action).

The last rewards (both the real reward and the planning reward).

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 is a real step or an imaginary step. To generate valid probabilities, the logits predicted by the linear layers are fed into a softmax layer. Additionally, an L2 regularization cost of 1e-9 is added for the final layer's input and 1e-6 for the real action's logits.

LSTM-Attention Cell This section details the 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 LSTM-Attention cell is denoted by:

where x t denotes the input to the cell, h t denotes the hidden state, c t denotes the input activation, K t and V t denotes the key and value vector for the attention module, respectively. The function LSTM_Attn is defined as:

where { W f , W ic , W ia , W o , W c , U f , U ic , U ia , U o , U c , b f , b ic , b ia , b o , b c } are parameters with output shape d m = 128 and ⊙ denotes element-wise product. The only differences compared to a standard LSTM cell are Equations 8, 10, 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 ( x t , K t -1 , V t -1 ) is based on the multi-head attention module used in Transformer, with the following modifications: (i) Attention is restricted to the last d t = 40 steps, equivalent to two stages. This is because maintaining an excessively long memory might complicate learning due to the Markov assumption of MDP. (ii) A learned relative positional encoding, akin to Transformer XL [41], is used to facilitate attention to steps relative to the current step.

For clarity, we first describe a single-head attention module. We denote q t ∈ R d e , k t ∈ R d e , v t ∈ R d e as the query, key, and value at step t , respectively, where d e = 16 is the dimension of the embedding. The input K t -1 ∈ R d t ,d e and V t -1 ∈ R d t ,d e represent the preceding d t keys and values, defined as follows:

We calculate the current q t , k t , v t via a linear projection on input x t :

where { W q , W k , W v } are parameters with output shape d e . This enables us to compute K t and V t by stacking the corresponding k and v . Subsequently, we compute the attention vector p t ∈ R d t as:

where U r ∈ R d t ,d e and b r ∈ R d t are learnable parameters representing the relative position encoding. This differs subtly from the relative position encoding used in Transformer XL, as the fixed attention

window in our case allows us to utilize fixed-size parameters for encoding relative positions. We initialize U r and b to be zero vectors, with the last entry of b set to 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 single-head attention module, denoted as ˜ a t ∈ R d e , is computed as follows:

For the multi-head attention module, we repeat the above single-head attention module d h = 8 times and stack the outputs to yield ˜ a stack t = stack ([˜ a t, 1 , ˜ a t, 2 , ..., ˜ a t,d h ]) , where ˜ a stack t ∈ R d m and ˜ a t,i denotes the output of the i th single-head attention module. The final output of the attention module a t is given by applying a linear layer and layer normalization [42] on this stacked output:

where x t represents the skip connection, and { W a , b a } are parameters with output size d m . This concludes the description of the LSTM-Attention cell.

LSTM-Attention Network An LSTM-Attention Network is constructed by stacking d n = 3 LSTM-Attention Cells together. We denote the network's input as x t , and use the superscript 1 ≤ n ≤ d n to represent the n th cell in the network. We denote the LSTM-Attention network as:

defined by the iterative equation for 1 ≤ n ≤ d n :

and we define h 0 t := h n t -1 , the top LSTM-Attention cell's output from the previous step. The final output of the network is the top LSTM-Attention cell's output, h n t . The input to each LSTM-Attention cell comprises both the network'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 in 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 it 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 LSTM-Attention network is also recommended to stabilize training further.

A significant drawback of the proposed LSTM-Attention network is the sequential computation it requires. Unlike the Transformer, which can compute attention across all steps simultaneously, the proposed network has to compute each step one at a time due to the recurrent computation of LSTM. Consequently, this leads to 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-augmented MDP, such as RL tasks that require long-term memory.

D Experiment Details

D.1 Hyperparameters

The hyperparameters used in all our experiments are shown in Table 4. We tune our hyperparameters exclusively on Sokoban based on the final solving rate. The specific hyperparameters we tune are: learning rates, model batch size, model loss scaling, maximum search depth or model unroll length, planning reward scaling, actor-critic clip global gradient norm, and actor-critic unroll length. In terms of preprocessing for Atari, we adhere to the same procedure as outlined in the IMPALA paper [31], with the exception that we do not apply grayscaling.

For the actor-critic baseline implemented on the raw MDP, we maintain the same hyperparameters, with two exceptions: (i) we adjust the actor-critic unroll length to 20 to ensure the same number of real transitions in an unrolled trajectory; (ii) we reduce the learning rate by half to 3e-4, after observing that the original learning rate led to unstable learning. We found that this set of hyperparameters provides optimal performance for the baseline in Sokoban, based on the final solving rate. Regarding the network architecture, we employ an encoder that mirrors the model's encoder architecture in Thinker, followed by a CNN that mirrors the actor-critic's CNN in Thinker.

D.2 Detailed Results on Sokoban

The performance of Thinker and the baselines on Sokoban can be found in Table 1, corresponding to the learning curves shown in Figure 5. We use our own implementation for DRC [14] and open-source implementation for Dreamer-v3 [33], with the result being averaged over three seeds. MuZero's result is taken from [43]. Results of DRC (original paper), I2A, ATreeC, VIN and IMPALA with Resnet are taken from [14].

As Thinker augments the MDP by adding K -1 steps for every step in the real MDP, there is an increased computational cost due to the extra K-1 imaginary steps, in addition to the added computational cost of training the model. On our workstation equipped with two A100s, the training durations for the raw MDP, DRC, and Thinker on Sokoban are around 1, 2, and 7 days, respectively.

However, when benchmarked against other model-based RL baselines, Thinker's training time is similar or even shorter. For instance, Dreamer-v3 (the open-sourced implementation) requires around 8 days for a single Sokoban run on our hardware, primarily due to its high replay ratio. MuZero, when conducting a standard number of simulations (e.g., 100), requires 10 training days. Notably, despite the Thinker's additional training overhead for RL agents, the algorithm compensates by requiring significantly fewer simulations. Furthermore, if one replaces the RNN in the actor-critic network with an MLP (as detailed in Appendix E), without compromising the performance, the computational time for a single Thinker run can be reduced to 3 days.

In Figure 7, we observed that the result for K = 1 , which does not use any imaginary transitions, underperforms compared to the result on the raw MDP. Initially, we hypothesized that this was because the model's hidden state, rather than the real state, was included in the augmented state. Although this hidden state might be more effective for predicting future states and rewards, it potentially complicates the learning of an optimal policy. However, subsequent experiments revealed that substituting the real state for the model's hidden state did not impact performance. This outcome suggests that the additional information given by the tree representation is, in fact, detrimental to performance. We speculate that this counterintuitive effect is due to the tree representation providing information, such as value estimation, that are so insightful that the critic in the actor-critic algorithm overly relies on this estimation, thereby neglecting the real state and impeding the learning of a useful state representation for the actor. To investigate this speculation, we experimented with masking the tree representation at a random rate (with a probability p = 0 . 20 ), which restored the performance to the level of the raw MDP. This propensity of the agent to depend excessively on a single type of input shown in these experiments points to a potential avenue for subsequent research.

D.3 Detailed Results on Atari 2600

The performance of Thinker and the baselines on Atari 2600 can be found in Table 2, corresponding to the learning curves shown in Figure 9. We include two additional baselines, UNREAL [44] and LASER [45], for comparison. The results of the baselines are taken from their original papers.

We note that Thinker's performance on Atari 2600 lags behind that of state-of-the-art algorithms. There are several possible reasons: (i) we did not tune hyper-parameters on Atari 2600; (ii) learning to plan may require a much larger number of frames to show significance, similar to how MuZero requires 20 billion frames. As the goal of the experiments on Atari 2600 is to demonstrate the general advantage of the Thinker-augmented MDP over the raw MDP, we leave the problem of mastering Atari 2600 using Thinker for future study. In principle, more advanced RL algorithms, including the baselines, could be applied to the Thinker-augmented MDP.

The final scores and the corresponding learning curves for each individual Atari game are presented in Table 5 and Figure 11, respectively. The Rainbow results are taken from the original paper [35], where the reported score is obtained from the best agent snapshots throughout training. The learning curves of Rainbow in Figure 11 are also taken from the original paper.

To validate the reproducibility of our results, we carried out multiple runs of five selected games with three different seeds. These games include battle zone , double dank , name this game , phoenix , and qbert . The selection of these games was influenced by the recommendations of [46], which suggests that these five games are pivotal in predicting the median human-normalized score. The final scores for these five games are displayed in Table 3, and their individual learning curves can be found in Figure 10. The results for these five games, as presented in other figures and tables in this paper, are derived from the average of these three seed values.

D.4 Discussion on Atari 2600's results

As stated in the main paper, Thinker demonstrates superior performance in shooting games like chopper command , seaquest , phoenix , beam rider , and space invader . A more general observation is that Thinker excels in environments with fast-moving objects, such as breakout . As depicted in https://youtu.be/0AfZh5SR7Fk , Thinker's model is capable of predicting the trajectories of these swift objects with near-perfect precision. We hypothesize that such predictions, essential for planning future actions, pose a challenge for an agent in the raw MDP to learn. For example, in Breakout , the video illustrates the agent exploring different methods of deflecting the ball at various angles, sometimes missing the ball in its imagination but performing adeptly in the

1e4

BattleZone

1e1

Double Dunk

1e4

Name This Game

1e5

Phoenix

1e4

8

2

3

4

6

1

4

3

2

4

-1

0

2

1

2

2

1

0

0

0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

1e8

1e8

1e8

1e8

1e8

Seed2

Seed3

actual game. Upon identifying a sequence of actions that might result in missing the ball, the agent can opt to avoid it. This adaptive decision-making and the significant performance gap between the agent policy and the base policy in these games underscore the benefits of planning in games with fast-moving objects.

It is important to highlight that applying a direct L2 loss on future states instead of a feature loss can undermine these advantages. For example, we observed that in Breakout , a model trained with direct L2 loss is able to predict frames of high quality, except that the ball, arguably the most important object in the game, is missing in the predicted frame. Given that these fast-moving objects typically inhabit only a few pixels and are challenging to predict, the model is prone to neglect these elements when trained using a direct L2 loss. In contrast, by using feature loss, these important objects, such as the ball in Breakout , are usually learned early in training.

We also observe that the agent's performance falls short in more demanding exploration games, such as montezuma revenge and venture . This is because the Thinker-augmented MDP lacks any mechanism to encourage exploration in the real environment, which results in challenges when dealing with these complex exploration games, much like an actor-critic applied directly to a raw MDP. Future research could explore methods to utilize Thinker to foster exploration, such as using model prediction loss as a curiosity reward. As exploration is a well-studied area in RL [47], many existing exploration methods could also be directly applied to the Thinker-augmented MDP.

Thinker also performs poorly in games where each action carries little impact, such as in qbert . Games in which the agent moves slowly, such as bank heist and ms pacman , also share this characteristic. These games are in stark contrast to games like Sokoban, where a single action can move the agent to a different grid. These games highlight a fundamental limitation of the Thinkeraugmented MDP that cannot be avoided by switching RL algorithms-planning is conducted in the primitive action space rather than the abstract action space, thereby necessitating the significance of each primitive action. Future research could investigate modifications to Thinker that would allow model interaction through abstract actions, thereby allowing high-level and long-term planning.

Interestingly, the actor-critic applied to the raw MDP occasionally outperforms the one applied to the Thinker-augmented MDP. This could be attributed to the fact that we provide the agent with the model's hidden state instead of the true or predicted frames. It is conceivable that the actual frame is simpler to learn from, given that the model learns the encoding of frames by minimizing supervised learning loss rather than maximizing the agent's returns. Future studies could consider allowing the agent to directly observe the true and predicted frame, thus ensuring that the agent possesses a strictly greater capability than an agent on the raw MDP. Another factor could be the impact of unoptimized hyperparameters. As we increased the network's depth while maintaining the same learning rate during the transition from Sokoban to the Atari environment, the training process became more unstable, as demonstrated in some learning curves. It is plausible that the poor performance in certain games stems from these unoptimized hyperparameters.

Finally, as evidenced by the individual learning curves illustrated in Figure 11, and the median human-normalized score in Figure 9, we observe that the slope of the learning curve is generally steeper at the end of training compared to other baselines. We posit that mastering the use of a model represents an additional skill to be acquired, as opposed to an agent acting on a raw MDP. This necessitates extended training to learn this skill, however, once this skill is acquired, the final performance generally surpasses that of an agent on the raw MDP. As such, future work could experiment with a larger training budget than 200 million frames.

1e8

1e4

Alien

1e3

Amidar

1e4

Assault

1e5

Asterix

1e5

Asteroids

8

1.5

1.0

4

6

1.0

0.5

2

0.5

0.0

0

0-

0.0L

1.00

1e6

Atlantis

1.5

1e3

Bank Heist

1e4

BattleZone

6

1e4

BeamRider

1e3

0.75

6

0.50

1.0

4

Y

0.25 -

0.5

2

2

0.00

0.0

0

le1

Bowling

1.00

1e2

Boxing

8

1e2

Breakout

8

1e3

Centipede

8

1e5ChopperCommand

4

0.75

0.50

6

4

6

6

4

2

0.25

2

21

0.00

0

2

0

1e5

CrazyClimber

1e5

1e5

DemonAttack

le1

DoubleDunk

1e3

Enduro

2

1.5

3

1.0

2

1.0

2

0.5

0

1

0.5

1

-1

0.0

0.0

-2

0.5

1e2

3

1e1

Freeway

1.00

1e4

Frostbite

1e4

Gopher

4

1e3

Gravitar

0.0

2

0.75 -

0.50

2

3

0.5

1

0.25

2

1

1.0 -

1e4

Hero

le1

IceHockey

0.00

1e4

Jamesbond

0

1e4

Kangaroo

0-

1e3

Krull

2

2.0

1.5

8

4

1.5

1.0

6

1.0

2

0

0.5

0.5

2

1

0.0

0.0

1e4

1e2MontezumaRevenge

1e3

MsPacman

1e4

NameThisGame

1e5

Phoenix

4

6

3

3

4

2

2

2

1

2

0

0

1e3

Pitfall

le1

Pong

1e2

PrivateEye

1e4

1e4

Riverraid

0.00

2.5 -

4

2.0

0.25

1

0.0

3

1.5

-0.50

0

2.5

2

1.0

0.75

-1

5.0

1

rw

0.5

1.00

1e4

RoadRunner

le1

Robotank

7.5

1e5

Seaquest

0

1e4

Sking

0.0

1e3

Solaris

61

9

4

0

3

4

4

2

-1

2

2

2

1

0

1e4

Space Invaders

0

1e5

StarGunner

0

le1

Surround

-3

le1

Tennis

0

1e4

Time Pilot

1.0

2

6

0.5

2

0.0

0

1

0.5

2

0

1e2

Tutankham

0

1e5

Up N Down

1.0 L

1e-2

Venture

-2

1e5

VideoPinball

0

1e4

WizardOfWor

5.0

3

2.5

2

00

1

2.5

2

1e5

YarsRevenge

0

1e4

Zaxxon

5.0

0.0

0.5

1.0

1.5

2.0

0L

0.0

0.5

1.0

1.5

2.0

0

0.0

0.5

1.0

1.5

2.0

2.0

1e8

1e8

1e8

1.0

1.5

Thinker-augmented MDP

1.0 -

0.5

0.5

Raw MDP

Rainbow

0.0

0.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

1e8

3

2

6

E Ablation Analysis

In this section, we investigate the significance of various elements of the Thinker-augmented MDP. All ablation experiments are performed in Sokoban under identical settings as those of the primary experiments, with results shown as the solving rate over the last 200 episodes.

Hints and memory We examine the effect of removing all hints and other auxiliary statistics from the tree representation. Specifically, we remove the following statistics:

from the augmented state. These statistics are deemed auxiliary since, theoretically, the agent can compute these from the remaining statistics in the previous and current steps. The result of removing the hints is shown in Figure 12.

The removal of hints causes the agent to learn at a slower pace while slightly reducing the final performance. We theorize that the agent requires more training to learn to summarize the rollouts across multiple steps, leading to slower learning. With prolonged training, it could be feasible for the agent to learn this summarization and achieve the same final performance, thereby rendering the provision of these handcrafted hints unnecessary.

A specialized RNN architecture is employed to endow the agents with long-term memories while preserving training stability, as elaborated in Appendix C. We then consider the impact of replacing this RNN with a three-layer MLP with 200 hidden units, effectively removing the agents' memory. In the MLP, we add a skip connection from the input to each hidden layer, to ensure easier flow of information. The result of this adjustment is presented in the same figure. Interestingly, the removal of memory does not hamper performance in any notable way.

We hypothesize that as we have already summarized the rollouts for the agent by providing the hints, the agent is not required to perform this summarization itself and hence, does not require memory. To verify this hypothesis, we consider the scenario of removing both memory and hint, and the outcome is demonstrated in the same figure. The effect of eliminating memory proves to be highly significant when the hints are not provided. Consequently, we conclude that the presence of either the hints or memory is necessary for learning to plan.

Base policy and imaginary actions We evaluate the effects of removing the base policy and its value. This leads to the tree representation being mostly empty as a number of statistics, such as the hints, are computed based on the node's value. Nonetheless, the agent still receives the model's hidden state as input. The result of this removal is shown in Figure 13.

The result shows that the agent fails to learn in such a scenario, indicating that the base policy is a critical component for learning to plan. We offer two potential explanations for this observation. Firstly, predicted values provide an agent with the ability to evaluate a leaf node more easily. This node evaluation allows an agent to evaluate a rollout without simulating the rollout till the end of the episode, thus allowing more shallow search. Even though the predicted value could theoretically be computed based on the hidden state of the node, which is provided to the agent, it is likely that this prediction is too difficult for the agent to learn by itself. Secondly, providing the base policy allows the agent to continuously improve upon the current policy. The base policy acts as a 'distilled' policy without planning, and the search corresponds to a learned policy improvement operator applied to this distilled policy. With the base policy and its value as inputs, the agent can learn different search algorithms as this policy improvement operator, completing the cycle of generalized policy iteration.

In addition, we consider the scenario where all imaginary actions are replaced with random actions, and we do not train the RL agent on these imaginary actions. The results are shown in Figure 13. This is intended to evaluate the benefits of learning imaginary actions. We find that performance significantly deteriorates, suggesting that the imaginary actions play a pivotal role in guiding real actions, consistent with the behavioural analysis in Appendix F.

Maximum search depth We evaluate the effect of increasing the maximum search depth and the model unroll length L from 5 to 10, thereby allowing the agents to search deeper. The result is shown in Figure 14. The initial learning speed is significantly slower, presumably due to the added time required for the model to learn to unroll over an increased number of steps. However, the final performance remains largely unaffected. This suggests that a deeper search does not necessarily lead

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

Default

Raw MDP

No hint

No mem

No hint & mem

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

L=5

L=10

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

Default

Raw MDP

No base policy

Random imaginary actions

0

1

2

3

4

5

Frames

1e7

0%

20%

40%

60%

80%

100%

Solving Rate

Default

Raw MDP

No planning rewards

No action dist.

to improved final performance. This finding aligns with previous studies indicating that a shallow search suffices to obtain satisfactory outcomes in MuZero [43].

Planning rewards and action target Lastly, we discuss some non-essential components of the algorithm. To mitigate the sparse rewards in the augmented MDP, we added a heuristic reward, which we refer to as the planning reward, during the imaginary steps. The details of the planning reward are discussed in Appendix G. Figure 13 shows the results of excluding planning rewards. The removal of the planning rewards does not have any significant impact on the performance. This suggests that the planning reward may not be a critical component for learning to plan. Additional behavior analysis of a trained agent in Appendix F demonstrates that an agent trained without planning rewards seems to learn a similar planning behavior to that of an agent trained with planning rewards.

We used the action distribution instead of the sampled action for training the base policy in the model. However, this approach necessitates that the augmented MDP receive the action distribution from the agent, a procedure that is not consistent with the definition of an MDP. We alternatively experimented with using the sampled action as the target for training the base policy in the model, with the results depicted in Figure 13. This modification did not significantly impact overall performance, though it led to a slightly slower initial learning. Nonetheless, we recommend employing the action distribution in lieu of the sampled action, as the former encompasses more information and has a lower variance. The adherence to the conventional definition of an MDP, while theoretically pertinent, is of marginal consequence in practical applications.

F Agent Behavior Analysis

In this section, we delve into the behavior learned by the trained agents in the Thinker-augmented MDP of Sokoban. Our inquiry aims to understand the nature of the search policies acquired by the agents. Due to the use of the dual network, it is easy to visualize what the agent is 'planning' by plotting the predicted states from the model on each imaginary step. Figure 1 shows a typical stage, where the agent plans different ways of solving the task. Videos of the whole game and other Atari games are available at https://youtu.be/0AfZh5SR7Fk .

Though the video is valuable in discerning the plan of an agent, a more quantitative analysis is warranted. Does the agent adhere to handcrafted planning algorithms, such as the n -step exhaustive search or MCTS, or does it learn novel planning algorithms? Given that the agents learn three types of actions-reset, imaginary, and real actions-we will analyze these actions individually.

Reset Action In a n -step exhaustive search, which involves iterating through every possible multistep action sequence and choosing the one with the highest rollout return, a reset is triggered when the search depth hits n . Similarly, MCTS initiates a reset at a leaf node. However, the trained agent seems to learn different behavior of resetting.

As visualized in the video, the agents often chose to reset when envisioning a poor rollout, for instance, when imagining a box being moved to an irreversible state. To illustrate this further, we can examine the distribution of Temporal Difference (TD) errors of the rollout from the root node to a node i n , defined as:

where { root , i 1 , i 2 , ..., i n } denotes the path from the root node to the node i n . This TD error can be interpreted as the quality of the rollout relative to the current policy. We compare the distribution of TD errors of the current rollout δ cur where the agent decided to reset with those where it did not, and the result is visualized in the left and middle panels of Figure 16.

We observe that the TD error distribution when agents opt to reset tends to lean more towards the left as compared to when they choose not to reset. This indicates a higher propensity for the agent to reset when the quality of the rollout is inferior. The thin tail on the left side of the distribution represents instances where the agent envisions a poor plan, such as reaching an irreversible state. Under such circumstances, the agent exhibits a high probability of resetting. This suggests that the agents learned to factor in the current rollout's quality in their decision to reset. This approach diverges from planning strategies like MCTS, which mandates a reset upon reaching a leaf node.

An additional dimension worth investigating is the expansion length , defined as the number of steps the agent executes beyond the last expanded node. In the case of MCTS, this length is consistently one as the algorithm requires a reset at every leaf node. However, as shown in the right panel of Figure 16, the trained agents generally exhibit a longer expansion length. In fact, when the maximum search depth L is increased to 10, the agent learns to search much deeper, as depicted in the same figure.

10.0

7.5

5.0

2.5

0.0

2.5

5.0

TD error

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Probability density

With planning rewards

not reset

reset

10.0

7.5

5.0

2.5

0.0

2.5

5.0

7.5

TD error

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Probability density

Without planning rewards

not reset

reset

1

2

3

4

5

6

7

8

9

10

Expansion length

0

20

40

60

80

100

Percentage (%)

With planning rewards

Without planning rewards (L=5)

Without planning rewards (L=10)

MCTS

Coupled with the aforementioned insights into the TD error distribution, we see that the agents prefer to persist with their existing plan unless they encounter an unfavorable leaf node or reach the set maximum search depth. These observations suggest that the agents tend to formulate plans that are both deep and narrow. Arguably, these learned behaviors mimic our planning processes more closely-we typically continue to think about a good plan, rather than simultaneously considering multiple plans and expanding each by a single step.

Imaginary Action In a n -step exhaustive search, imaginary action is chosen sequentially, while in MCTS, imaginary action is chosen according to the upper confidence bound (UCB). As for the trained agents, it is difficult to deduce the formula used to select the imaginary action as it is the output of a deep neural network. However, the maximum TD errors in a stage, defined as δ max := max k δ k where 1 ≤ k < K denotes the step within a stage, could be used as a proxy to evaluate the effectiveness of the search 4 . The distribution of maximum TD errors of search using the trained agent, UCB, and random search is shown in Figure 17. Note that we use the learned model of the trained agent when applying UCB.

We observe that the distribution of maximum TD errors has a slightly heavier tail on the right side when searching with a trained agent. This can be better visualized in log scale as shown in Figure 18. This tail mostly corresponds to the case of solving or being closer to solving a level in the imagination where the current policy has a high chance of not solving it. This scenario is relatively rare given that the current policy is capable of solving approximately 95% of levels, which explains the thin tails in all cases. Thus, the heavier tail suggests that the agents learn to search more effectively than MCTS within the same search budgets. This finding is consistent with previous studies on replacing MCTS with neural networks in supervised learning settings [19].

We conjecture that the difference between an RL agent and MCTS in this distribution would be much larger if the model were also trained with MCTS, given the large number of simulations required for MCTS to achieve good performance in Sokoban [19]. It is also noteworthy that even without planning rewards, the agent is still able to collect large maximum TD errors, which implies that the agent can learn to search effectively without the guidance of heuristics.

Real Action In a n -step exhaustive search, the real action is taken as the child node with the highest rollout return, while in MCTS, the real action is usually selected based on the visit count of each child node. In contrast, our trained agents appear to learn to use a combination of both.

To further investigate how real action is selected given the same search result, we used the imaginary action and reset action outputs from the trained agent and considered two different ways of selecting the real action: (i) the real action output by the trained agent, (ii) the real action output by MCTS, which is sampled based on visit counts with different temperatures T ∈ { 0 . 25 , 1 } as in MuZero [5]. For each method, we calculated the portion of the real action that corresponded to the highest visit count, the second highest visit count, and so on. The result is shown in the left panel of Figure 19. Additionally, we computed the same for the mean rollout return and the maximum rollout return of child nodes, as shown in the middle and right panels of the figure.

The figure shows that most of the real action selected by the agents corresponds to the child node with the highest visit counts or rollout returns. The figures also suggest that the maximum rollout return is the most indicative of the agent's real action among the three statistics. In contrast, MCTS tends to select the real action with the highest visit count. Notably, the action with the highest visit count does not necessarily correspond to the action with the highest rollout return, as evidenced by MCTS with T = 0 . 25 , which chooses the action with the highest visit count with over 80% probability, yet only around half of these actions correspond to the one with the highest rollout return. The simultaneous consideration of both rollout returns and visit counts by the trained agent seems more intuitively reasonable-it is unwise to follow a poor plan even if it is simulated numerous times.

Comparison with MCTS Lastly, to compare the efficiency of the trained agents across all three action sets with MCTS, we assess the solving rates of both the trained agents and MCTS. In MCTS, we utilize the model learned by the agent at the end of training, and we test a range of simulation

1

0

1

2

3

4

5

6

7

8

Maximum TD error

0.0

0.5

1.0

1.5

2.0

Probability density

With planning reward

Trained Agent

UCB

Random

1

0

1

2

3

4

5

6

7

8

Maximum TD error

0.0

0.5

1.0

1.5

2.0

Probability density

Without planning reward

Trained Agent

UCB

Random

1

0

1

2

3

4

5

6

7

8

Maximum TD error

10

6

10

5

10

4

10

3

10

2

10

1

10 0

Probability density

With planning reward

Trained Agent

UCB

Random

1

0

1

2

3

4

5

6

7

8

Maximum TD error

10

6

10

5

10

4

10

3

10

2

10

1

10 0

Probability density

Without planning reward

Trained Agent

UCB

Random

1

2

3

4

5

Visit count's ranking

0

20

40

60

80

100

Percentage of real action (%)

1

2

3

4

5

Mean rollout return's ranking

0

20

40

60

80

100

1

2

3

4

5

Max rollout return's ranking

0

20

40

60

80

100

Agent with planning reward

Agent without planning reward

MCTS (T=1)

MCTS (T=0.25)

10 1

10 2

10 3

Number of simulations

65%

70%

75%

80%

85%

90%

Solving Rate

MCTS

Agent Policy

Base Policy

steps 5 spanning over {5,10,20,40,80,160,320,640,1280,2560}. The results are depicted in Figure 20. We observe that MCTS needs 1280 simulation steps to match the performance of a trained agent, which achieves the same results with only 20 stage length. This notable difference highlights the potential benefits of trained agents compared to handcrafted algorithms-they exhibit the capacity to conduct searches with increased efficiency, particularly within limited computational resources.

In conclusion, the analysis in this section suggests that the trained agents appear to learn a planning algorithm that differs from common handcrafted planning algorithms. These agents learn to plan deeply and reset upon encountering a poor node, while being able to search for a good plan within a limited number of model simulations. It is likely that the agents also take advantage of the model's hidden state, which contains rich information about the root and predicted states, when deciding upon the three types of actions. However, a comprehensive investigation of this particular dynamic lies beyond the scope of our current analysis.

G Planning Rewards

In the main paper, we use the extended real reward as the augmented reward:

where r denotes the real reward obtained from the real step. However, this sparse reward may make learning difficult. We explore adding an auxiliary reward, which we call the planning reward , to help guide the learning of imaginary and reset actions. We define the planning reward ˜ r p k on step k to be:

where g max root ,k denotes the maximum rollout return of the root node on step k within a stage, as defined after Equation 1, and c p ≥ 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:

where { root , i ∗ 1 , i ∗ 2 , ..., i ∗ n } refers to the rollout with the highest return, g root ,i , in the entire stage. As ˆ v root is not influenced by imaginary and reset actions, the return of a stage is equivalent to ˆ r i ∗ 1 + γ 2 ˆ r i ∗ 2 + ...γ n ˆ r i ∗ n + γ n +1 ˆ v i ∗ n , the maximum rollout return. Consequently, the planning reward is a heuristic directing the agent to maximize the maximum rollout return, while not penalizing for searching paths with a low rollout return to encourage exploration in imagination.

However, this planning reward should not be given to real actions to prevent incentive misalignment. For example, collecting a large amount of planning rewards is possible by continually doing nothing in real steps but imagining a good plan. As such, the planning reward signal has to be separated from the real reward signal ˜ r k instead of being added together. We also decay c p from 1 to 0 throughout training to further ensure that the imaginary action is aligned with the real reward function. Combined, the augmented reward contains both the real reward and the planning reward: r aug := (˜ r, ˜ r p ) .

The incorporation of planning rewards necessitates certain modifications to the actor-critic algorithm in our experiments. Specifically, we separately learn the value of these two distinct rewards when training the critic, resulting in two different TD errors corresponding to the real and the planning rewards. In training the critic for the planning reward, we treat each stage as a single episode. This approach ensures that the planning reward incentives the agent to optimize the maximum rollout returns exclusively for the current stage. Furthermore, when using the TD errors to train the actor, we impose a mask on the planning rewards' TD error, setting it to zero during a real step.

Experiments in Appendix E show that planning rewards only provided a minimal increase in initial learning speed, indicating that the real rewards are sufficient for learning to plan. We conjecture that sparse real rewards are not a problem, as the TD errors induced by real rewards may not be sparse. For example, if the critic learns to use the maximum rollout return g max root ,k as the estimated value, which is provided in the tree representation, then the TD errors in an imaginary step equal the planning rewards, except for the adjustment due to the discount rate.

H Relationship with Generalized Policy Iteration and Meta-Reinforcement Learning

H.1 Generalized Policy Iteration Algorithm

Applying an RL algorithm to the Thinker-augmented MDP can be viewed as an instance of a generalized policy iteration algorithm [29, 30]. Policy improvements are carried out by an agent that conducts searches within the model. Given the base policy, these searches yield an improved policy. Simultaneously, the policy evaluation is done by the learning of the value-policy network. This involves training the network on trajectories generated by the improved policy, so as to learn the values of this improved policy and project the improved policy back to the base policy. A noteworthy divergence from a typical generalized policy iteration algorithm is the learning of the policy improvement operator, instead of employing a handcrafted one such as MCTS.

Given this perspective, a crucial question arises: how effective is a learned policy improvement operator in comparison to a handcrafted one? To address this question, we evaluate the efficiency of different policy improvement operators, such as MCTS and an RL agent, during a single step of policy improvement.

Specifically, we consider the Thinker-augmented MDP with a fixed base policy ˆ π and its value ˆ v . We do not update the base policy that is described in Section 4.4. The base policy is a given policy that performs slightly better than a random policy. For simplicity, we employ the true environment dynamics, represented as ( P, R ) , as the model in the augmented MDP and do not impose a limit on the maximum search depth, given that we are using the true environmental dynamics. Sokoban serves as the real environment in our experiments.

We experiment with three different types of policy improvement operators: (i) n -step exhaustive search, (ii) MCTS, and (iii) RL agents. In (iii), we trained the agent from scratch in the Thinkeraugmented MDP. Only the tree representation u tree is provided as input to the agent, so the agent cannot learn to act based on the real states directly. The results are shown in Figure 21. Due to the low solving rate, we report the return of an episode here for better illustration. The maximum return in Sokoban is around 14, and a return of 4 gives at most 30% solving rate.

From the figure, we see that the base policy achieves a return of 0.12, and all policy improvement operators are able to improve on this policy and achieve a better return. We observe that the policy improvement gap for RL agents widens as the value of K increases. With K =10, the agent's performance is similar to that of a 3-step exhaustive search, which requires 155 model simulations. Increasing K to 20, the performance is slightly worse than a 4-step exhaustive search, which requires 780 model simulations. At K =40, the agent performance surpasses a 5-step exhaustive search, which requires 3905 model simulations. In addition, the agent at K =10 also shows comparable performance to MCTS that uses 60 model simulations. It is also worth noting that the number of model simulations used by the agent is typically less than K , given that imaginary steps also include the traversal through expanded nodes.

The results here demonstrate that increasing K to 40 can still enhance performance, whereas in the main paper, the performance improvement gained by increasing K above 10 is insignificant. We conjecture that once the value-policy network is trained, the model can learn to predict the values of leaf nodes more accurately, reflecting the value of the current policy instead of the base policy. This eliminates the need for a deep search to evaluate a rollout more precisely, rendering a large K unnecessary.

The results also underscore the importance of a learned base policy. The agent achieves a maximum solving rate of less than 30%, whereas a 95% solving rate is attainable when the base policy undergoes training as well. This outcome is intuitive, considering that the agent gives an improved version of the base policy, and thus the agent policy will not significantly outperform it. If a fixed poor base policy is used, the agent policy will not perform well.

Nonetheless, these findings suggest that an agent can learn to perform policy improvement and achieve similar performance to handcrafted algorithms while requiring significantly fewer simulations. However, it requires a moderate number of transitions for an agent to learn this policy improvement operator.

0

1

2

3

4

5

Frames

1e7

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Trained agent K=10

Trained agent K=20

Trained agent K=40

Base policy

1-step exhaustive search (5 simulations)

2-step exhaustive search (30 simulations)

MCTS (60 simulations)

3-step exhaustive search  (155 simulations)

4-step exhaustive search  (780 simulations)

5-step exhaustive search  (3905 simulations)

H.2 Meta-Reinforcement Learning

In this section, we explore the relationship between meta-reinforcement learning (meta-RL) algorithms [48] and the Thinker algorithm. Meta-RL is a method designed to learn how to perform reinforcement learning. Its objective is to develop a policy that can quickly adapt to a variety of new tasks derived from a distribution of tasks, using minimal data. While meta-RL is often used in scenarios involving a task distribution, it is also applicable to single-task scenarios [49, 50], which is our primary focus here.

In the meta-RL framework, a trial typically comprises L episodes. During a trial, the agents initially encounter K episodes-termed burn-in episodes . Notably, the return within these episodes does not contribute to the main objective. Instead, the primary goal is to maximize the return of the remaining L -K episodes, known as evaluation episodes . Consequently, the purpose of a meta-RL algorithm is to use the knowledge gained throughout the entire trial to maximize returns during the evaluation episodes. One common class of meta-RL algorithm applies an RL algorithm with an RNN that can access previous actions and rewards, and the hidden state is not reset across episodes. As such, the knowledge of the trial is embedded in the RNN's hidden states [27, 28].

If the model perfectly represents the true environment dynamics, we can align the Thinker-augmented MDP with meta-RL by considering imaginary rollouts as (truncated) burn-in episodes, and real transitions as the evaluation episode in meta-RL. As agents only aim to maximize the return in real transitions, the objectives of meta-RL and RL agents on the Thinker-augmented MDP coincide. Moreover, the purpose of burn-in episodes in meta-RL or imaginary rollouts in planning is the same - to guide better actions in the evaluation episodes or real transitions, possibly through risk-taking behavior in them.

However, unlike conventional meta-RL, which uses a fixed number of burn-in episodes followed by evaluation episodes, the Thinker-augmented MDP interleaves the burn-in episodes within a single evaluation episode. In other words, for each step in the evaluation episode with state s t , we allow agents to have a few burn-in episodes starting from s t .

In meta-RL, it is generally assumed that the burn-in episodes and evaluation episodes are drawn from the same MDP, but this assumption does not hold when the burn-in episodes are generated by a learned model. But if the MDP underlying the burn-in episodes, ˆ 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 to how the RL agent in the Thinker-augmented MDP learns to use the imaginary rollout for better action selection, despite the possible inaccuracies of the model.

In summary, there is a close relationship between learning to plan and learning to learn. The former employs a model on each step to decide better actions at that state, while the latter uses episodes within a trial to determine better actions in the evaluation episodes. This highlights the converging ideas in meta-learning and planning, both aiming to make efficient use of available information for decision-making.

I Miscellaneous

I.1 Related Work on the Dual Network

The dual network is similar to the model in MuZero [5], except for (i) the use of the dual network architecture and (ii) the prediction of future states trained by feature loss. Improvements to MuZero's model have been proposed. EfficientZero [51] deviates from MuZero by adding SimSiam loss [52] to match future representations. We briefly experimented with it, but the result was only slightly better than that of MuZero's model, possibly because this loss works better in a limited data regime where EfficientZero was designed and evaluated. [43] also proposes to predict future states as an auxiliary loss but does not use the dual network architecture. In our experiments, this addition led to better performance, as shown in Figure 4, but the result was still significantly worse than that of the dual network.

Most other types of models used in model-based RL do not predict values and policies, as future values and policies are often not required in the algorithm. Among these models, various loss functions have been proposed for training the model to acquire meaningful representations for the task. [53] employs a VAE [54] to encode the state while Dreamer [33] utilizes a sequential VAE. Given the extensive body of research in model-based RL, we direct readers to the review [55] for a detailed exploration of this field.

I.2 Planning Capacity

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 ( a 1 , reset), ( a 2 , reset), ( a 1 , not reset), ( a 1 , reset), ( a 1 , not reset), ( a 2 , 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 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 [56, Ch 3.4] and partial-order planning [56, Ch 10.4], cannot be implemented by the agent. Agents are also unable to 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-augmented MDP can implement any forward-based 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 F and the video visualization. This is because, unlike traditional 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- **[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.\n\nWEAKNESSES:\n- **[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\nQUESTIONS:\n1. 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)\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\nThere 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\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\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- 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).\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-19"}

comments: 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.

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

2. 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.

2. The contribution is unclear to me:

    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.

    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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 1
DATE: 2023-07-19

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThis 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).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}