SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1. After line 7, we insert the real transition in the form of  $(s_{\mathrm{old}},\tilde{a},r,d)$ , where  $s_\mathrm{old}$  denotes the previous real state, into a replay buffer. A parallel thread trains the model by sampling from the replay buffer.
2. A node is represented by the hidden state of the RNN model, not the real state. Consequently, line 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  $\hat{p}_i^d > 0.5$ , we manually set the content of all its descendant nodes  $j$  to  $\hat{r}_j = 0$ ,  $\hat{\pi}_j = \hat{\pi}_i$ ,  $\hat{v}_j = 0$ ,  $\hat{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_{\mathrm{cur}}$  and the root node representation  $u_{\mathrm{root}}$  are included in the tree representation. Denoting the current search depth by  $d \in \{0,1,\dots,L\}$ , the full tree representation  $u_{\mathrm{tree}} \in \mathbb{R}^{10|\mathcal{A}| + K + 9}$  contains:
1.  $u_{\mathrm{root}}, u_{\mathrm{cur}} \in \mathbb{R}^{5|\mathcal{A}| + 2}$ : The root node representation and the current node representation.
2.  $g_{\mathrm{root,cur}} \in \mathbb{R}$ : Current rollout return.
3.  $g_{\mathrm{root,cur}} - \gamma^{d + 1}\hat{v}_{\mathrm{cur}}\in \mathbb{R}$  : Sum of discounted rewards collected in the current rollout.
4.  $g_{\mathrm{root}}^{\mathrm{max}} \in \mathbb{R}$ : Maximum rollout return of the root node.
5.  $\gamma^{d + 1}\in \mathbb{R}$  : Trailing discount factor.
6. onehot  $(k)\in [0,1]^K$  : Current step in a stage, represented in one-hot encoding.
7.  $i^{\mathrm{reset}} \in [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_{t,k}^{\mathrm{aug}}$ :
1.  $u_{\mathrm{tree}} \in \mathbb{R}^{10|\mathcal{A}| + K + 9}$ : The tree representation.
2.  $s\in S$  : The real state underlying the stage.
3.  $\hat{s}_{\mathrm{cur}} \in S$ : The predicted real state of the current node.
4.  $h_{\mathrm{root}} \in \mathcal{H}$ : The hidden state of the RNN model at the root node.
5.  $h_{\mathrm{cur}} \in \mathcal{H}$ : The hidden state of the RNN model at the current node.
We used  $s_{t,k}^{\mathrm{aug}} = (u_{\mathrm{tree}}, h_{\mathrm{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 Thinker-augmented MDP, reproducing the experimental results presented in the paper.
B Model Training Procedure
B.1 Collection of Training Samples
A transition tuple  $x_{t} \coloneqq (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-first-out buffer with a capacity of 200,000 transitions. To train the model, a batch of sequences  $x_{i:i+2L}$  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+2L}$ . 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+2L}$  is updated to the L1 loss of the value,  $|\hat{v}_i - v_i^{\mathrm{target}}|$ . Transitions are sampled with a probability of  $P(i) = \frac{p_i^\alpha}{\sum_j p_j^\alpha}$ , where  $\alpha = 0.6$ . Each sampled transition sequence's loss is weighted by the importance weight  $w_i = (\frac{1}{NP(i)})^\beta$ , where  $N$  is the total number of transition tuples in the buffer, and  $\beta$  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  $\hat{s}_{t + l}$  towards the state-reward network since  $\hat{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  $\hat{v}_{t + l}^{\mathrm{target}}$  in Equation 3 is defined by:
$$
\hat {v} _ {t + l} ^ {\text {t a r g e t}} := r _ {t + l + 1} + \gamma r _ {t + l + 2} + \dots + \gamma^ {2 L - l} r _ {t + 2 L} + \gamma^ {2 L - l} g _ {\text {r o o t}, t + 2 L} ^ {\text {m e a n}}. \tag {4}
$$
Instead of using the predicted value  $\hat{v}_{t + 2L}$ , we opted for the mean rollout return (at the end of the stage)  $g_{\mathrm{root},t + 2L}^{\mathrm{mean}}$  as the bootstrap target<sup>3</sup>. We found that it yields more stable training in our experiments. This stability likely stems from averaging across multiple rollouts starting from  $s_{t + 2L}$ , rather than relying directly on the predicted value.
$g_{\mathrm{root}}^{\mathrm{mean}}$  does not contain the reward upon arriving the root node, as  $\hat{r}_{\mathrm{root}}$  is manually set to 0 in Algorithm 1.
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_t^{sr} \coloneqq g^{sr}(s_t)$ . To obtain  $h_{t + l + 1}^{sr}$  from  $(h_{t + l}^{sr}, a_{t + l})$  with  $l \geq 0$ , we first concatenate  $h_{t + l}^{sr}$  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  $4D + 1$  residual blocks, each having an output channel of 128. The output of the unrolling function is  $h_{t + l + 1}^{sr}$ .
For the value-policy network, in addition to the previous hidden state, the unrolling function should also take the encoded predicted frame  $g(\hat{s}_{t + l})$  as input, so the unrolling procedure is slightly modified. First, we define  $h_{t - 1}^{vp}$  as a zero vector of shape (128, 6, 6). To obtain  $h_{t + l}^{vp}$  from  $(h_{t + l - 1}^{vp}, a_{t + l - 1}, g(\hat{s}_{t + l}))$  with  $l \geq 0$  (replace  $\hat{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  $4D + 1$  residual blocks, each having an output channel of 128. The output of the unroll function is  $h_{t + l}^{vp}$ .
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  $\hat{r}_{t + l}$  and termination probability  $\hat{p}_{t + l}^d$ , while the predicted quantities for the value-policy network are values  $\hat{v}_l$  and policies  $\hat{\pi}_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  $\hat{p}_{t + l}^{d}$  and the predicted policy  $\hat{\pi}_{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,  $\hat{s}_{t + l}$  for  $l\geq 1$  , by applying a decoder on  $h_{t + l}^{sr}$  for  $l\geq 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  $\hat{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_k^{sr}, h_k^{vp}]$ , a 3D vector of shape (256, 6, 6), and the tree representation  $u_{\mathrm{tree}}$ , a flat vector of shape  $(10|\mathcal{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  $5K = 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:
1. The output from the CNN processing the model's hidden state.
2. The output from the RNN processing the tree representation.
3. The one-hot encoding of the last actions (real action, imaginary action, and reset action).
4. The last rewards (both the real reward and the planning reward).
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:
$$
\left(h _ {t}, c _ {t}, K _ {t}, V _ {t}\right) = \operatorname {L S T M} _ {-} \operatorname {A t t n} \left(x _ {t}, h _ {t - 1}, c _ {t - 1}, K _ {t - 1}, V _ {t - 1}\right), \tag {5}
$$
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:
$$
f _ {t} = \sigma \left(W _ {f} x _ {t} + U _ {f} h _ {t - 1} + b _ {f}\right), \tag {6}
$$
$$
i _ {t} ^ {c} = \sigma \left(W _ {i c} x _ {t} + U _ {i c} h _ {t - 1} + b _ {i c}\right), \tag {7}
$$
$$
i _ {t} ^ {a} = \sigma \left(W _ {i a} x _ {t} + U _ {i a} h _ {t - 1} + b _ {i a}\right), \tag {8}
$$
$$
o _ {t} = \sigma \left(W _ {o} x _ {t} + U _ {o} h _ {t - 1} + b _ {o}\right), \tag {9}
$$
$$
a _ {t}, K _ {t}, V _ {t} = \operatorname {a t t n} \left(x _ {t},, K _ {t - 1}, V _ {t - 1}\right), \tag {10}
$$
$$
\tilde {c} _ {t} = \tanh  \left(W _ {c} x _ {t} + U _ {c} h _ {t - 1} + b _ {c}\right), \tag {11}
$$
$$
c _ {t} = f _ {t} \odot c _ {t - 1} + i _ {t} ^ {c} \odot \tilde {c} _ {t} + i _ {t} ^ {a} \odot \tanh  (a _ {t}) \tag {12}
$$
$$
h _ {t} = o _ {t} \odot \tanh  \left(c _ {t}\right), \tag {13}
$$
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  $\odot$  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  $\mathrm{attr}(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} \in \mathbb{R}^{d_{e}}$ ,  $k_{t} \in \mathbb{R}^{d_{e}}$ ,  $v_{t} \in \mathbb{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} \in \mathbb{R}^{d_{t},d_{e}}$  and  $V_{t-1} \in \mathbb{R}^{d_{t},d_{e}}$  represent the preceding  $d_{t}$  keys and values, defined as follows:
$$
K _ {t - 1} := \left[ k _ {t - d _ {t}}, k _ {t - d _ {t} + 1} \dots , k _ {t - 1} \right] ^ {T}, \tag {14}
$$
$$
V _ {t - 1} := \left[ v _ {t - d _ {t}}, v _ {t - d _ {t} + 1},..., v _ {t - 1} \right] ^ {T}. \tag {15}
$$
We calculate the current  $q_{t}, k_{t}, v_{t}$  via a linear projection on input  $x_{t}$ :
$$
q _ {t} = W _ {q} x _ {t}, \tag {16}
$$
$$
k _ {t} = W _ {k} x _ {t}, \tag {17}
$$
$$
v _ {t} = W _ {v} x _ {t}, \tag {18}
$$
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} \in \mathbb{R}^{d_{t}}$  as:
$$
p _ {t} = \operatorname {s o f t m a x} \left(\frac {\left(K _ {t} + U _ {r}\right) q _ {t}}{\sqrt {d _ {e}}} + b _ {r}\right), \tag {19}
$$
where  $U_{r} \in \mathbb{R}^{d_{t},d_{e}}$  and  $b_{r} \in \mathbb{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  $\tilde{a}_t \in \mathbb{R}^{d_e}$ , is computed as follows:
$$
\tilde {a} _ {t} = V _ {t} ^ {T} p _ {t}. \tag {20}
$$
For the multi-head attention module, we repeat the above single-head attention module  $d_h = 8$  times and stack the outputs to yield  $\tilde{a}_t^{\mathrm{stack}} = \mathrm{stack}([\tilde{a}_{t,1},\tilde{a}_{t,2},\dots,\tilde{a}_{t,d_h}])$ , where  $\tilde{a}_t^{\mathrm{stack}}\in \mathbb{R}^{d_m}$  and  $\tilde{a}_{t,i}$  denotes the output of the  $i^{\mathrm{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:
$$
a _ {t} = \operatorname {L a y e r N o r m} \left(W _ {a} \tilde {a} _ {t} ^ {\text {s t a c k}} + b _ {a} + x _ {t}\right), \tag {21}
$$
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 \leq n \leq d_{n}$  to represent the  $n^{\mathrm{th}}$  cell in the network. We denote the LSTM-Attention network as:
$$
\left(h _ {t} ^ {1: d _ {n}}, c _ {t} ^ {1: d _ {n}}, K _ {t} ^ {1: d _ {n}}, V _ {t} ^ {1: d _ {n}}\right) = \text {L S T M} _ {-} \operatorname {A t t n} _ {-} \operatorname {N e t} \left(x _ {t - 1}, h _ {t - 1} ^ {1: d _ {n}}, c _ {t - 1} ^ {1: d _ {n}}, K _ {t - 1} ^ {1: d _ {n}}, V _ {t - 1} ^ {1: d _ {n}}\right), \tag {22}
$$
defined by the iterative equation for  $1 \leq n \leq d_{n}$ :
$$
\left(h _ {t} ^ {n}, c _ {t} ^ {n}, K _ {t} ^ {n}, V _ {t} ^ {n}\right) = \operatorname {L S T M} _ {-} \operatorname {A t t n} \left(\operatorname {s t a c k} \left(\left[ x _ {t - 1}, h _ {t} ^ {n - 1} \right]\right), h _ {t - 1} ^ {n}, c _ {t - 1} ^ {n}, K _ {t - 1} ^ {n}, V _ {t - 1} ^ {n}\right), \tag {23}
$$
and we define  $h_t^0 \coloneqq h_{t-1}^n$ , 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_t^n$ . 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.