SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Stephen Chung
University of Cambridge mhc48@cam.ac.uk
Ivan Anokhin
Mila, Université de Montréal  
ivan.anokhin@mila.quebec
David Krueger
University of Cambridge dsk30@cam.ac.uk
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.
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.<sup>1</sup> 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
$^{1}$ Full code is available at https://github.com/stephen-chung-mh/thinker, which allows for using the Thinker-augmented MDP with the same interface as OpenAI Gym [10].
Figure 1: Illustration of a trained agent during a single stage in Sokoban, where the goal is to push all four boxes to red-bordered squares. From the root state, the agent generates four rollouts using the learned model. Except for the two real states, all frames are generated by the model, showcasing near-perfect frame prediction. Predicted rollout return  $g$  (equivalent to  $g_{\mathrm{root,cur}}$  in Equation 1) is displayed. The four rollouts are reasonable plans with agents pushing different boxes towards targets. The agent follows the third rollout for the next five real steps (not fully shown in the figure), likely due to its high rollout return. Crucially, plan generation, plan execution, and the model are all learned. Videos of the whole game and other Atari games are available at https://youtu.be/0AfZh5SR7Fk.
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.
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.
We consider a Markov Decision Process (MDP) defined by a tuple  $(S, \mathcal{A}, P, R, \gamma, d_0)$ , where  $\mathcal{S}$  is a set of states,  $\mathcal{A}$  is a finite set of actions,  $P: S \times \mathcal{A} \times S \to [0,1]$  is a transition function representing the dynamics of the environment,  $R: S \times \mathcal{A} \to \mathbb{R}$  is a reward function,  $\gamma \in [0,1]$  is a discount factor, and  $d_0: S \to [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') = \operatorname*{Pr}(s_{t+1} = s'|s_t = s, a_t = a)$ ,  $R(s, a) = \mathbb{E}[r_{t+1}|s_t = s, a_t = a]$ , and  $d_0(s) = \operatorname*{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 = \sum_{k=t}^{\infty} \gamma^{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  $\pi: S \times \mathcal{A} \to [0,1]$ , such that for any time  $t \geq 1$ , selecting actions according to  $\pi(s, a) = \operatorname*{Pr}(a_t = a|s_t = s)$  maximizes the expected return  $\mathbb{E}[g_{t+1}|\pi]$ . The value function for policy  $\pi$  is  $v^\pi$  where for all  $s \in S$ ,  $v^\pi(s) = \mathbb{E}[g_{t+1}|s_t = s, \pi]$ . 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.
Figure 2: Illustration of a stage in the augmented MDP, which is composed of  $K - 1$  imaginary steps and one real step. The environment outputs a state  $s_t$  that sets the root state of the world model. The agent interacts with the world model  $K - 1$  times with augmented actions. In each imaginary step, the agent receives an augmented state  $s_{t,k}^{\mathrm{aug}}$ , which contains the latest outputs from the world model. Finally, the agent acts in the real environment for a single step based on  $s_{t,K}^{\mathrm{aug}}$ .
The Thinker algorithm transforms a given MDP, denoted as  $\mathcal{M} = (\mathcal{S},\mathcal{A},P,R,\gamma ,d_0)$ , into an augmented MDP, denoted as  $\mathcal{M}^{\mathrm{aug}} = (S^{\mathrm{aug}},\mathcal{A}^{\mathrm{aug}},P^{\mathrm{aug}},R^{\mathrm{aug}},\gamma^{\mathrm{aug}},d_0^{\mathrm{aug}})$ . We call  $\mathcal{M}$  the real MDP, and  $\mathcal{M}^{\mathrm{aug}}$  the augmented MDP. Besides the real MDP, we assume that we are given (i) a model in the form of  $(\hat{P},\hat{R})$  that is an approximate version of  $(P,R)$  and is deterministic, and (ii) a base policy  $\hat{\pi}:(S,\mathcal{A})\to [0,1]$  and its value function  $\hat{v}:S\to \mathbb{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 \geq 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 \in \{1, 2, \dots, K\}$  to denote the augmented step within a stage, and  $t \in \{1, 2, \dots\}$  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  $\tilde{a} \in \mathcal{A}$  and also whether to reset the model to the root node by reset action  $\tilde{a}^r \in \{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  $\tilde{a} \in \mathcal{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  $\mathcal{A}^{\mathrm{aug}} \coloneqq \mathcal{A} \times \{0,1\}$ , and an augmented action by the tuple  $a^{\mathrm{aug}} \coloneqq (\tilde{a},\tilde{a}^r)$ . During imaginary steps,  $\tilde{a}$  becomes the imaginary action, while  $\tilde{a}^r$  decides whether to reset the current node. During real steps,  $\tilde{a}$  becomes the real action, while  $\tilde{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
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, \dots, i_n)$  as:
$$
g _ {i _ {0}, i _ {n}} := \hat {r} _ {i _ {0}} + \gamma \hat {r} _ {i _ {1}} + \gamma^ {2} \hat {r} _ {i _ {2}} + \dots \gamma^ {n} \hat {r} _ {i _ {n}} + \gamma^ {n + 1} \hat {v} _ {i _ {n}}, \tag {1}
$$
where  $\hat{v}_i$  denotes the value function applied on the node  $i$  and  $\hat{r}_i$  denotes the predicted reward upon reaching node  $i$ . That is,  $\hat{v}_i \coloneqq \hat{v}(\hat{s}_i)$ , and  $\hat{r}_i \coloneqq \hat{R}(\hat{s}_{\text{parent_node}(i)}, a_{\text{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} \coloneqq \hat{r}_i + \gamma \hat{v}_i$ . We define the mean rollout return  $g_i^{\mathrm{mean}}$  and the maximum rollout return  $g_i^{\mathrm{max}}$  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_{\mathrm{child}(i)}^{\mathrm{mean}}$  and  $g_{\mathrm{child}(i)}^{\mathrm{max}}$  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_{\mathrm{child}(i)} \in \mathbb{R}^{|\mathcal{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}\in \mathbb{R}^{5|\mathcal{A}| + 2}$ , called node representation, that contains:
1.  $\hat{r}_i, \hat{v}_i \in \mathbb{R}, \hat{\pi}_i \in \mathbb{R}^{|\mathcal{A}|}$ : The core statistics, composed of the reward, the value, and the base policy of the node  $i$ . Here,  $\hat{\pi}_i \coloneqq (\hat{\pi}(s_i, a_1), \hat{\pi}(s_i, a_2), \dots, \hat{\pi}(s_i, a_{|\mathcal{A}|}))$ .
2. onehot  $(a_{i})\in \mathbb{R}^{|\mathcal{A}|}$  : The one-hot encoded action that leads to node  $i$
3.  $g_{\mathrm{child}(i)}^{\mathrm{mean}}, g_{\mathrm{child}(i)}^{\mathrm{max}}, \frac{1}{K} n_{\mathrm{child}(i)} \in \mathbb{R}^{|\mathcal{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_{\mathrm{root}}$  and the current node's representation  $u_{\mathrm{cur}}$  to represent the tree (we use the subscript  $\mathrm{root}$  to denote the root node, which is unchanged throughout a
stage, and the subscript  $\mathrm{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_{\mathrm{root}}$  and  $u_{\mathrm{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_{\mathrm{tree}} \in \mathbb{R}^{10|\mathcal{A}| + K + 9}$ , called tree representation. We may also include the real state  $s$  or the current node's predicted state  $\hat{s}_{\mathrm{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_k^{\mathrm{aug}} = 0$  for  $k \leq K$  and  $r_{K + 1}^{\mathrm{aug}} = r$ , where  $r$  denotes the real reward. We set the augmented discount rate  $\gamma^{\mathrm{aug}} \coloneqq \gamma^{\frac{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_{\mathrm{root}}^{\mathrm{max}}$ , 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