SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Example 1 (Passive T-Maze). T-Maze has a 4-direction action space  $\{\mathcal{L},\mathcal{R},\mathcal{U},\mathcal{D}\}$ . It features a long corridor of length  $L$ , starting from (oracle) state  $O$  and ending with (junction) state  $J$ , connecting two potential goal states  $G_{1}$  and  $G_{2}$ .  $O$  has a horizontal position of  $0$ , and  $J$  has  $L$ . The agent observes null  $N$  except at states  $J, O, G_{1}, G_{2}$ . The observation space  $\{N, J, O, G_{1}, G_{2}\}$  is two-dimensional. At state  $O$ , the agent can observe the goal position  $G \in \{G_{1}, G_{2}\}$ , uniformly sampled at the beginning of an episode. The initial state is  $S$ . The transition is deterministic, and the agent remains in place if it hits the wall. Thus, the horizontal position of the agent  $x_{t} \in \mathbb{N}$  is determined by its previous action sequence  $a_{1:t-1}$ .
In a Passive T-Maze  $\mathcal{M}_{\text{passive}}$ ,  $O = S$ , allowing the agent to observe  $G$  initially, and  $L = T - 1$ . Rewards are given by  $R_{t}(h_{1:t},a_{t}) = \frac{\mathbf{1}(x_{t + 1}\geq t) - 1}{T - 1}$  for  $t\leq T - 1$ , and  $R_{T}(h_{1:T},a_{T}) = \mathbf{1}(o_{T + 1} = G)$ . The unique optimal policy  $\pi^{*}$  moves right for  $T - 1$  steps, then towards  $G$ , yielding an expected return of 1.0. The optimal Markovian policy can only guess the goal position, ending with an expected return of 0.5. The worst policy has an expected return of -1.0.
Example 2 (Active T-Maze). In an Active T-Maze  $\mathcal{M}_{\text {active }}$ ,  $O$  is one step left of  $S$  and  $L = T - 2$ . The unique optimal policy moves left to reach  $O$ , then right to  $J$ . The rewards  $R_{1} = 0$ ,  $R_{t}(h_{1:t}, a_{t}) = \frac{\mathbf{1}(x_{t + 1} \geq t - 1) - 1}{T - 2}$  for  $2 \leq t \leq T - 1$ , and  $R_{T}(h_{1:T}, a_{T}) = \mathbf{1}(o_{T + 1} = G)$ . The optimal Markovian and worst policies have the same expected returns as in the passive setting.
$\mathcal{M}_{\mathrm{passive}}$  is specially designed for testing memory only. To make the final decision, the optimal policy must recall  $G$ , observed in the first step. All the memory lengths are  $T$  as the terminal reward depends on initial observation and observation at  $J$  depends on the previous action sequence. The credit assignment length is 1 as immediate penalties occur for suboptimal actions. In contrast,  $\mathcal{M}_{\mathrm{active}}$  also requires a credit assignment length of  $T$ , as the initial optimal action  $\mathcal{L}$  only affects  $n$ -step rewards when  $n = T$ . Passive and Active T-Mazes are suitable for evaluating the ability of pure memory, and both memory and credit assignment in RL, as they require little exploration due to the dense rewards before reaching the goal.
Memory and credit assignment have been extensively studied in RL, with representative works shown in Table 1. The work most closely related to ours is bsuite (Osband et al., 2020), which
Figure 2: Pixel-based tasks, Passive Visual Match (left) and Key-to-Door (right), evaluated in our experiments. Each task has three phases, where we adjust the length of Phase 2 and keep the lengths of the others constant. The agent (in beige) only observes nearby grids (within white borders) and cannot pass through the walls (in black). In Passive Visual Match, the agent observes a randomized target color (blue in this case) in Phase 1, collects apples (in green) in Phase 2, and reaches the grid with the matching color in Phase 3. In Key-to-Door, the agent reaches the key (in brown) in Phase 1, collects apples in Phase 2, and reaches the door (in cyan) in Phase 3.
also disentangles memory and credit assignment in their benchmark. bsuite is also configurable with regard to the memory and credit assignment lengths to evaluate an agent's capability. Our work extends Osband et al. (2020) by providing formal definitions of memory and credit assignment lengths in RL for task design and evaluation. Our definitions enable us to analyze existing POMDP benchmarks.
Several prior works have assessed memory (Parisotto et al., 2020; Esslinger et al., 2022; Chen et al., 2022; Pleines et al., 2023; Morad et al., 2023) in the context of Transformers in RL. However, these prior methods use tasks that require (relatively short) context lengths between 50 and 831. In addition, there are two separate lines of work using Transformers for better credit assignment. One line of work (Liu et al., 2019; Ferret et al., 2020) has developed algorithms that train Transformers to predict rewards for return redistribution. The other line of work (Chen et al., 2021; Zheng et al., 2022) proposes return-conditioned agents trained on offline RL datasets. Both lines are beyond the scope of canonical online RL algorithms that purely rely on TD learning, which we focus on.
5 Evaluating Memory-Based RL Algorithms
In this section, we evaluate the memory and credit assignment capabilities of memory-based RL algorithms. Our first experiment aims to evaluate the memory abilities by disentangling memory from credit assignment, using our Passive T-Maze and the Passive Visual Match (Hung et al., 2018) tasks in Sec. 5.1. Next, we test the credit assignment abilities of memory-based RL in our Active T-Maze and the Key-to-Door (Raposo et al., 2021) tasks in Sec. 5.2. Passive Visual Match and Key-to-Door tasks (demonstrated in Fig. 2) can be viewed as pixel-based versions of Passive and Active T-Mazes, respectively. Lastly, we evaluate memory-based RL on standard POMDP benchmarks that only require short-term dependencies in Sec. 5.3, with a focus on their sample efficiency.
We focus on model-free RL agents that take observation and action sequences as input, which serves as a simple baseline that can achieve excellent results in many tasks (Ni et al., 2022). The agent architecture is based on the codebase from Ni et al. (2022), consisting of observation and action embedders, a sequence model, and actor-critic heads. We compare agents with LSTM and Transformer (GPT-2 (Radford et al., 2019)) architectures. For a fair comparison, we use the same hidden state size (128) and tune the number of layers (varying from 1, 2, 4) for both LSTM and Transformer architectures. Since our tasks all require value memory lengths to be the full horizon, we set the context lengths to the full horizon  $T$ .
For the T-Maze tasks, we use DDQN (van Hasselt et al., 2016) with epsilon-greedy exploration as the RL algorithm to better control the exploration strategy. In pixel-based tasks, we use CNNs as observation embeddings and SAC-Discrete (Christodoulou, 2019) as RL algorithm following Ni et al. (2022). We train all agents to convergence or at most 10M steps on these tasks. We provide training details and learning curves in Appendix C.
5.1 Transformers Shine in Pure Long-Term Memory Tasks
First, we provide evidence that Transformers can indeed enhance long-term memory in tasks that purely test memory, Passive T-Mazes. In Fig. 3 (left), Transformer-based agents consistently solve the task requiring memory lengths from 50 up to 1500. To the best of our knowledge, this achievement
Figure 3: Transformer-based RL outperforms LSTM-based RL in tasks (purely) requiring long-term memory. Left: results in Passive T-Mazes with varying memory lengths from 50 to 1500; Right: results in Passive Visual Match with varying memory lengths from 60 to 1000. We also show the performance of the optimal Markovian policies. Each data point in the figure represents the final performance of an agent with the memory length indicated on the x-axis. All the figures in this paper show the mean and its  $95\%$  confidence interval over 10 seeds.
Figure 4: Transformer-based agent can solve long-term memory low-dimensional tasks with good sample efficiency. In Passive T-Maze with long memory length (from left to right: 500, 750, 1000, 1500), Transformer-based agent is significantly better than LSTM-based agent.
may well set a new record for the longest memory length of a task that an RL agent can solve. In contrast, LSTM-based agents start to falter at a memory length of 250, underperforming the optimal Markovian policies with a return of 0.5. Both architectures use a single layer, yielding the best performance. We also provide detailed learning curves for tasks requiring long memory lengths in Fig. 4. Transformer-based agents also show much greater stability and sample efficiency across seeds when compared to their LSTM counterparts.
On the other hand, the pixel-based Passive Visual Match (Hung et al., 2018) has a more complex reward function: zero rewards for observing the color in Phase 1, immediate rewards for apple picking in Phase 2 requiring short-term dependency, and a final bonus for reaching the door of the matching color, requiring long-term memory. The success rate solely depends on the final bonus. The ratio of the final bonus w.r.t. the total rewards is inversely proportional to the episode length, as the number of apples increases over time in Phase 2. Thus, the low ratio makes learning long-term memory challenging for this task. As shown in Fig. 3 (right), Transformer-based agents slightly outperform LSTM-based ones $^{8}$  in success rates when the tasks require a memory length over 120, yet their success rates are close to that of the optimal Markovian policy (around  $1/3$ ). Interestingly, Transformer-based agents achieve much higher returns with more apples collected, shown in Fig. 5. This indicates Transformers can help short-term memory (credit assignment) in some tasks.
5.2 Transformers Enhance Temporal Credit Assignment in RL, but not Long-Term
Next, we present evidence that Transformers do not enhance long-term credit assignment for memory-based RL. As illustrated in Fig. 6 (left), in Active T-Mazes, by merely shifting the Oracle one step to the left from the starting position as compared to Passive T-Mazes, both Transformer-based and LSTM-based agents fail to solve the task when the credit assignment (and memory) length extends to just 250. The failure of Transformer-based agents on Active T-Mazes suggests the bottleneck of credit assignment. This is because they possess sufficient memory capabilities to solve this task (as evidenced by their success in Passive T-Mazes), and they can reach the junction state without exploration issues (as evidenced by their return of more than 0.0). A similar trend is observed in pixel-based Key-to-Door tasks (Raposo et al., 2021), where both Transformer-based and LSTM-
It is worth noting that at the memory length of 1000, LSTM-based agents encountered NaN gradient issues in around  $30\%$  of the runs, which were not shown in the plots.
Figure 6: Transformer-based RL improves temporal credit assignment compared to LSTM-based RL, but its advantage diminishes in long-term scenarios. Left: results in Active T-Mazes with varying credit assignment lengths from 20 to 500; Right: results in Key-to-Door with varying credit assignment lengths from 60 to 500. We also show the performance of the optimal policies that lack (long-term) credit assignment – taking greedy actions to maximize immediate rewards. Each data point in the figure represents the final performance of an agent in a task with the credit assignment length indicated on the x-axis, averaged over 10 seeds.
Figure 5: Transformer-based agent is more sample-efficient in long-term memory high-dimensional tasks. In Passive Visual Match with long memory length (750 on the left; 1000 on the right), Transformer-based agents show a slight advantage over LSTM-based agents in the memory capability (as measured by success rate), while being much higher in total rewards (as measured by return).
based agents struggle to solve the task with a credit assignment length of 250, as indicated by the low success rates in Fig. 6 (right). Nevertheless, we observe that Transformers can improve credit assignment, especially when the length is medium around 100 in both tasks, compared to LSTMs.
On the other hand, we investigate the performance of multi-layer Transformers on credit assignment tasks to explore the potential scaling law in RL, as discovered in supervised learning (Kaplan et al., 2020). Fig. 7 reveals that multi-layer Transformers greatly enhance performance in tasks with credit assignment lengths up to 100 in Active T-Maze and 120 in Key-to-Door. Yet they fail to improve long-term credit assignment. Interestingly, 4-layer Transformers have similar or worse performance compared to 2-layer Transformers, suggesting that the scaling law might not be seamlessly applied to credit assignment in RL.
Figure 7: Increasing the number of layers (and heads) in Transformers aids in temporal, but not long-term, credit assignment. The left two figures show the final results from Active T-Maze and Key-to-Door, varying the number of layers and heads in Transformers from 1, 2 to 4. The right two figures show the associated learning curves in Active T-Maze with different credit assignment lengths.
5.3 Transformers in RL Raise Sample Complexity in Certain Short-Term Memory Tasks
Lastly, we revisit the standard POMDP benchmark used in prior works (Ni et al., 2022; Han et al., 2020; Meng et al., 2021). This PyBullet benchmark consists of tasks where only the positions (referred to as “-P” tasks) or velocities (referred to as “-V” tasks) of the agents are revealed, with dense rewards. As a result, they necessitate short-term credit assignment. Further, the missing velocities can be readily inferred from two successive positions, and the missing positions can be
Figure 8: Transformer-based RL is sample-inefficient compared to LSTM-based RL in most PyBullet occlusion tasks. The top row shows tasks with occluded velocities, while the bottom row shows tasks with occluded positions. Results are averaged over 10 seeds.
approximately deduced by recent velocities, thereby requiring short-term memory. Adhering to the implementation of recurrent TD3 (Ni et al., 2022) and substituting LSTMs with Transformers, we investigate their sample efficiency over 1.5M steps. We find that in most tasks, Transformer-based agents show worse sample efficiency than LSTM-based agents, as displayed in Fig. 8. This aligns with the findings in POPGym paper (Morad et al., 2023) where GRU-based PPO is more sample-efficient than Transformer-based PPO across POPGym benchmark that requires relatively short memory lengths. These observations are also consistent with the known tendency of Transformers to thrive in situations with large datasets in supervised learning.
We provide quantitative definitions of memory and credit assignment lengths, but developing programs capable of calculating them for any given task remains an open challenge. This can often be desirable since sequence models can be expensive to train in RL agents. The worst compute complexity of  $c^{\mathcal{M}}$  can be  $O(|S||\mathcal{A}|T^{2})$  for a finite MDP, while the complexity of  $m^{\mathcal{M}}$  can be  $O((|O| + |\mathcal{A}|)^T|\mathcal{O}||\mathcal{A}|T)$  for a finite POMDP. In addition,  $c^{\mathcal{M}}$  does not take the variance of intermediate rewards into account. These reward noises have been identified to complicate the learning process for credit assignment (Mesnard et al., 2021). On the other hand,  $m^{\mathcal{M}}$  does not take into account the memory capacity (i.e., the number of bits) and robustness to observation noise (Beck et al., 2019).
On the empirical side, we evaluate Transformers on a specific architecture (GPT-2) with a relatively small size in the context of online model-free RL. As a future research direction, it would be interesting to explore how various sequence architectures and history abstraction methods may impact and potentially improve long-term memory and/or credit assignment.
In this study, we evaluate the memory and credit assignment capabilities of memory-based RL agents, with a focus on Transformer-based RL. While Transformer-based agents excel in tasks requiring long-term memory, they do not improve long-term credit assignment and generally have poor sample efficiency. Furthermore, we highlighted that many existing RL tasks, even those designed to evaluate (long-term) memory or credit assignment, often intermingle both capabilities or require only short-term dependencies. While Transformers are powerful tools in RL, especially for long-term memory tasks, they are not a universal solution to all RL challenges. Our results underscore the ongoing need for careful task design and the continued advancement of core RL algorithms.
This research was enabled by the computational resources provided by the Calcul Québec (calculquebec.ca) and the Digital Research Alliance of Canada (alliancecan.ca). We thank Pierluca D'Oro and Zhixuan Lin for their technical help. We thank Jurgis Pasukonis and Zhixuan Lin for correcting some factual errors. We also thank anonymous reviewers, Pierluca D'Oro, Siddhant Sahu, Giancarlo Kerg, Guillaume Lajoie, and Ryan D'Orazio for constructive discussion.
D. Yarats, I. Kostrikov, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021a. 21, 22
D. Yarats, A. Zhang, I. Kostrikov, B. Amos, J. Pineau, and R. Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, 2021b. 21, 22
Q. Zheng, A. Zhang, and A. Grover. Online decision transformer. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, 2022. 1, 7
E. A. Zilli and M. E. Hasselmo. Modeling the role of working memory and episodic memory in behavioral tasks. Hippocampus, 18(2):193-209, 2008. 2
A Proofs 15
B Details on Environments Related to Memory and Credit Assignment 17
B.1 Abstract Problems 17
B.2 Concrete Benchmarks for Memory 18
B.3 Concrete Benchmarks for Credit Assignment 20
CExperiment Details 21
C.1 Memory-Based RL Implementation 21
C.2 Additional Results 23
C.3 Plotting Details 24
C.4 Training Compute 24
D Broader Impacts 24
Lemma 1 (Upper bound of value memory length). For any policy  $\pi$ ,
$$
l _ {\text {v a l u e}} ^ {\mathcal {M}} (\pi) \leq \max  \left(m _ {\text {r e w a r d}} ^ {\mathcal {M}}, m _ {\text {t r a n s i t}} ^ {\mathcal {M}}, l _ {\text {m e m}} (\pi) - 1\right) \tag {3}
$$
Proof of Lemma 1. We show this by expanding the sum of rewards in  $Q$ -value function:
$$
\begin{array}{l} Q ^ {\pi} \left(h _ {1: t}, a _ {t}\right) = \mathbb {E} _ {\pi , \mathcal {M}} \left[ \sum_ {i = t} ^ {T} \gamma^ {i - t} r _ {i} \mid h _ {1: t}, a _ {t} \right] (4) \\ = \underbrace {\mathbb {E} \left[ r _ {t} \mid h _ {1 : t} , a _ {t} \right]} _ {\text {r e w a r d m e m o r y}} + \sum_ {i = t} ^ {T - 1} \gamma^ {i + 1 - t} (5) \\ \end{array}
$$
$$
\begin{array}{l} \int \underbrace {P \left(o _ {i + 1} \mid h _ {1 : i} , a _ {i}\right)} _ {\text {t r a n s i t i o n m e m o r y}} \underbrace {\pi \left(a _ {i + 1} \mid h _ {1 : i + 1}\right)} _ {\text {p o l i c y m e m o r y}} \underbrace {\mathbb {E} \left[ r _ {i + 1} \mid h _ {1 : i + 1} , a _ {i + 1} \right]} _ {\text {r e w a r d m e m o r y}} d o _ {t + 1: i + 1} d a _ {t + 1: i + 1} (6) \\ \stackrel {(A)} {=} \mathbb {E} \left[ r _ {t} \mid h _ {t - m _ {\text {r e w a r d}} ^ {\mathcal {M}} + 1: t}, a _ {t} \right] + \sum_ {i = t} ^ {T - 1} \gamma^ {i + 1 - t} (7) \\ \end{array}
$$
$$
\begin{array}{l} \int P \left(o _ {i + 1} \mid h _ {i - m _ {\text {t r a n s i t}} ^ {\mathcal {M}} + 1: i}, a _ {i}\right) \pi \left(a _ {i + 1} \mid h _ {i - l _ {\text {m e m}} (\pi) + 2: i + 1}\right) \mathbb {E} \left[ r _ {i + 1} \mid h _ {i - m _ {\text {r e w a r d}} ^ {\mathcal {M}} + 2: i + 1}, a _ {i + 1} \right] d o _ {t + 1: i + 1} d a _ {t + 1: i + 1} (8) \\ = \mathbb {E} _ {\pi , \mathcal {M}} \left[ \sum_ {i = t} ^ {T} \gamma^ {i - t} r _ {i} \mid h _ {\min  (t - m _ {\text {r e w a r d}} ^ {\mathcal {M}} + 1, t - m _ {\text {t r a n s i t}} ^ {\mathcal {M}} + 1, t - l _ {\text {m e m}} (\pi) + 2, t - m _ {\text {r e w a r d}} ^ {\mathcal {M}} + 2): t}, a _ {t} \right] (9) \\ = \mathbb {E} _ {\pi , \mathcal {M}} \left[ \sum_ {i = t} ^ {T} \gamma^ {i - t} r _ {i} \mid h _ {t + 1 - \max  \left(m _ {\text {r e w a r d}} ^ {\mathcal {M}}, m _ {\text {t r a n s i t}} ^ {\mathcal {M}}, l _ {\text {m e m}} (\pi) - 1\right): t}, a _ {t} \right] (10) \\ \end{array}
$$
where line  $(A)$  follows from the definition of memory lengths in policy (Def. 1.C), reward (Def. 1.A), and transition (Def. 1.B). By the definition of value memory length (Def. 1.D), the shortest history to match the value, we have the upper bound:  $l_{\mathrm{value}}^{\mathcal{M}}(\pi) \leq \max(m_{\mathrm{reward}}^{\mathcal{M}}, m_{\mathrm{transit}}^{\mathcal{M}}, l_{\mathrm{mem}}(\pi) - 1)$ .
Lemma 2 (Lower bound of value memory length). For an optimal policy  $\pi^{*} \in \operatorname{argmin}_{\pi \in \Pi_{\mathcal{M}}^{*}} \{l_{\mathrm{mem}}(\pi)\}$ , we have  $l_{\mathrm{mem}}(\pi^{*}) \leq l_{\mathrm{value}}^{\mathcal{M}}(\pi^{*})$ .
Proof. Proof by contradiction. Suppose  $l_{\mathrm{value}}^{\mathcal{M}}(\pi^{*}) < l_{\mathrm{mem}}(\pi^{*})$ . We can construct a deterministic policy  $\pi'$  that has context length of  $l_{\mathrm{value}}^{\mathcal{M}}(\pi^{*})$ :
$$
\begin{array}{l} \pi^ {\prime} \left(h _ {t - l _ {\text {v a l u e}} ^ {\mathcal {M}} (\pi^ {*}) + 1: t}\right) = \underset {a _ {t}} {\operatorname {a r g m a x}} Q _ {l _ {\text {v a l u e}} ^ {\mathcal {M}} (\pi^ {*})} ^ {\pi^ {*}} \left(h _ {t - l _ {\text {v a l u e}} ^ {\mathcal {M}} (\pi^ {*}) + 1: t}, a _ {t}\right) (11) \\ = \underset {a _ {t}} {\operatorname {a r g m a x}} Q ^ {\pi^ {*}} \left(h _ {1: t}, a _ {t}\right), \quad \forall t, h _ {1: t} (12) \\ \end{array}
$$
This is a greedy policy derived from  $\pi^{*}$ 's value function. The policy memory length of  $\pi'$  is strictly shorter than  $l_{\mathrm{mem}}(\pi^{*})$ , by definition and assumption:
$$
l _ {\mathrm {m e m}} \left(\pi^ {\prime}\right) \leq l _ {\mathrm {c t x}} \left(\pi^ {\prime}\right) = l _ {\text {v a l u e}} ^ {\mathcal {M}} \left(\pi^ {*}\right) <   l _ {\mathrm {m e m}} \left(\pi^ {*}\right) \tag {13}
$$
Similar to policy improvement theorem (Sutton and Barto, 2018, Chap 4.2) in Markov policies, we can show that  $\pi'$  is at least as good as  $\pi^*$ , thus contradicting the assumption that  $\pi^*$  is an optimal policy that has the shortest policy memory length. To see this, for any  $t$  and  $h_{1:t}$ ,
$$
\begin{array}{l} V ^ {\pi^ {*}} \left(h _ {1: t}\right) \leq \max  _ {a _ {t}} Q ^ {\pi^ {*}} \left(h _ {1: t}, a _ {t}\right) (14) \\ \stackrel {(A)} {=} \max  _ {a _ {t}} Q ^ {\pi^ {*}} \left(h _ {1: t}, \pi^ {\prime} \left(h _ {t - l _ {\text {v a l u e}} ^ {\mathcal {M}} \left(\pi^ {*}\right) + 1: t}\right)\right) (15) \\ \stackrel {(B)} {=} \mathbb {E} _ {\mathcal {M}} \left[ r _ {t} + \gamma Q ^ {\pi^ {*}} \left(h _ {1: t + 1}, a _ {t + 1}\right) \mid h _ {1: t}, a _ {t} \sim \pi^ {\prime}, a _ {t + 1} \sim \pi^ {*} \right] (16) \\ \stackrel {(C)} {\leq} \mathbb {E} _ {\mathcal {M}} \left[ r _ {t} + \gamma Q ^ {\pi^ {*}} \left(h _ {1: t + 1}, a _ {t + 1}\right) \mid h _ {1: t}, a _ {t: t + 1} \sim \pi^ {\prime} \right] (17) \\ \dots (18) \\ \leq \mathbb {E} _ {\mathcal {M}} \left[ \sum_ {i = t} ^ {T} \gamma^ {i - t} r _ {i} \mid h _ {1: t}, a _ {t: T - 1} \sim \pi^ {\prime} \right] (19) \\ = \mathbb {E} _ {\pi^ {\prime}, \mathcal {M}} \left[ \sum_ {i = t} \gamma^ {i - t} r _ {i} \mid h _ {1: t} \right] = V ^ {\pi^ {\prime}} \left(h _ {1: t}\right) (20) \\ \end{array}
$$
where  $(A)$  follows the definition of  $\pi^{\prime}$ ,  $(B)$  uses the Bellman equation and  $(C)$  uses the property of greedy policy  $\pi^\prime$ .
Proof of Theorem 1. By Lemma 2,  $l_{\mathrm{value}}^{\mathcal{M}}(\pi^{*}) \geq l_{\mathrm{mem}}(\pi^{*})$ . On the other hand, by Lemma 1,  $l_{\mathrm{value}}^{\mathcal{M}}(\pi^{*}) \leq \max(m_{\mathrm{reward}}^{\mathcal{M}}, m_{\mathrm{transit}}^{\mathcal{M}}, l_{\mathrm{mem}}(\pi^{*}) - 1)$ . Thus,