SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fabian Paischer  $^{1}$ , Thomas Adler  $^{1}$ , Markus Hofmacher  $^{2}$ , Sepp Hochreiter  $^{1}$
$^{2}$  JKU LIT SAL eSPML Lab, Institute for Machine Learning,
paischer@ml.jku.at
Reinforcement learning agents deployed in the real world often have to cope with partially observable environments. Therefore, most agents employ memory mechanisms to approximate the state of the environment. Recently, there have been impressive success stories in mastering partially observable environments, mostly in the realm of computer games like Dota 2, StarCraft II, or MineCraft. However, existing methods lack interpretability in the sense that it is not comprehensible for humans what the agent stores in its memory. In this regard, we propose a novel memory mechanism that represents past events in human language. Our method uses CLIP to associate visual inputs with language tokens. Then we feed these tokens to a pretrained language model that serves the agent as memory and provides it with a coherent and human-readable representation of the past. We train our memory mechanism on a set of partially observable environments and find that it excels on tasks that require a memory component, while mostly attaining performance on-par with strong baselines on tasks that do not. On a challenging continuous recognition task, where memorizing the past is crucial, our memory mechanism converges two orders of magnitude faster than prior methods. Since our memory mechanism is human-readable, we can peek at an agent's memory and check whether crucial pieces of information have been stored. This significantly enhances troubleshooting and paves the way toward more interpretable agents.
In reinforcement learning (RL) an agent interacts with an environment and learns from feedback provided in the form of a reward function. In many applications, especially in real-world scenarios, the true state of the environment is not directly accessible to the agent, but rather approximated via observations that reveal mere parts of it. In such environments, the capability to approximate the true state by virtue of an agent's perception is crucial (Åström, 1964; Kaelbling et al., 1998). To this end, many approaches track events that occurred in the past. The brute-force strategy is to simply store all past observations. However, this is often infeasible and it is much more efficient to store more abstract representations of the history. Thus, many RL algorithms use memory mechanisms such as LSTM (Hochreiter & Schmidhuber, 1997) or Transformer (Vaswani et al., 2017) to compress sequences of high-dimensional observations. This has led to impressive successes mostly in the realm of mastering computer games on a human or even super-human level. Some examples are Dota 2 (Berner et al., 2019), StarCraft II (Vinyls et al., 2019), or MineCraft (Baker et al., 2022; Patil et al., 2022).
Most state-of-the-art methods dealing with partial observability in RL employ a memory mechanism that is not intelligible for humans. In this regard, we draw inspiration from the semantic memory
Figure 1: We add a semantic and human-readable memory to an agent to tackle partially observable RL tasks. We map visual observations  $\boldsymbol{o}_t$  to the language domain via CLIP retrieval. The memory component, a pretrained language encoder, operates on text only and compresses a history of tokens into a vector  $\boldsymbol{h}_t$ . The agent takes an action  $\boldsymbol{a}_t$  based on the current observation  $\boldsymbol{o}_t$  and the compressed history  $\boldsymbol{h}_t$ .
present in humans (Yee et al., 2017) and propose to represent past events in human language. Humans memorize abstract concepts rather than every detail of information they encountered in the past (Richards & Frankland, 2017; Bowman & Zeithamova, 2018). Their ability to abstract is heavily influenced by the exposure to language in early childhood (Waxman & Markow, 1995). Further, humans use language on a daily basis to abstract and pass on information. Therefore, language is a natural choice as a representation for compounding information and has the key advantage of being human-readable. This enables analyzing if critical pieces of information have entered the memory or not. Based on this data, it becomes clear which parts of the system require refinement. Moreover, natural language has been shown to be effective as a compressed representation of past events in RL (Paischer et al., 2022a).
Our proposed method, Semantic HELM (SHELM), leverages pre-trained foundation models to construct a memory mechanism that does not require any training. We use CLIP (Radford et al., 2021) to associate visual inputs with language tokens. Thereby, the vocabulary of the CLIP language encoder serves as a semantic database of concepts from which we retrieve the closest tokens to a given observation. These tokens are passed to a pretrained language model that serves as memory and provides the agent with a coherent representation of the past.
We illustrate the benefits of a human-readable and semantic memory in partially observable RL problems. First, we conduct a qualitative analysis on whether a CLIP vision encoder is capable of extracting semantics out of synthetic environments. Then, we test SHELM on a set of partially observable 2D MiniGrid (Chevalier-Boisvert et al., 2018), and 3D MiniWorld (Chevalier-Boisvert, 2018) environments. We find that even though these environments are only partially observable, they often do not necessitate a memory mechanism as they are solvable by a memory-less policy. The MiniGrid-Memory task, however, clearly requires memory and SHELM reaches state-of-the-art performance. On more realistic 3D environments such as Avalon (Albrecht et al., 2022) and Psychlab (Leibo et al., 2018), SHELM successfully assesses the semantics of visual observations. In turn, it requires approximately 100 times fewer interaction steps than prior methods on Psychlab's continuous recognition task that explicitly evaluates for memory capacity. On Avalon, we find that the addition of our semantic memory performs on-par with the current state-of-the-art for a limited interaction budget, while adding interpretability to the memory.
Our goal is to express visual observations in language space such that the resulting representations become comprehensible for humans. To this end, we instantiate a mapping from images to text in the form of pretrained components which are entirely frozen during training. This way the available computational resources are invested into performance rather than interpretability. Before we describe SHELM, we briefly review the HELM framework, which serves as a starting point for our work.
HELM (Paischer et al., 2022a) was proposed as a framework for RL in partially observable environments. It utilizes a pretrained language model (LM) as memory mechanism that compresses past
(c)
Figure 2: Architecture of SHELM. We compile a semantic database  $S$  by encoding prompt-augmented tokens from the overlapping vocabularies of CLIP and the LM (a). Given an observation  $o_t$  we retrieve the top- $k$  embeddings present in  $S$  and select their corresponding text tokens (b). These tokens are passed to the LM which represents the memory module of SHELM (c).  $C_{t-1}$  represents the memory cache of the LM which tracks past tokens.
observations. To map environment observations  $\pmb{o}_t \in \mathbb{R}^n$  to the LM, it introduces the FrozenHopfield (FH) mechanism, which performs a randomized attention over pretrained token embeddings  $\pmb{E} = (e_1, \dots, e_k)^\top \in \mathbb{R}^{k \times m}$  of the LM, where  $k$  is the vocabulary size and  $m$  is the embedding dimension. Let  $P \in \mathbb{R}^{m \times n}$  be a random matrix with entries sampled independently from  $\mathcal{N}(0, n / m)$ . The FH mechanism performs
$$
\boldsymbol {x} _ {t} = \boldsymbol {E} ^ {\top} \operatorname {s o f t m a x} (\beta \boldsymbol {E P o} _ {t}), \tag {1}
$$
where  $\beta$  is a scaling factor that controls the dispersion of  $x_{t}$  within the convex hull of the token embeddings. This corresponds to a spatial compression of observations to a mixture of tokens in the LM embedding space. Since  $\pmb{P}$  is random, the association of observations  $\pmb{o}_t$  with token embeddings  $e_i$  is arbitrary, i.e., not meaningful. That is, the FH mechanism does not preserve semantics. For temporal compression, HELM leverages a pretrained LM. At time  $t$ , HELM obtains a compressed history representation by
$$
\boldsymbol {h} _ {t} = \operatorname {L M} \left(\boldsymbol {c} _ {t - 1}, \boldsymbol {x} _ {t}\right) \tag {2}
$$
where  $c_{t}$  represents the context cached in the memory register of the LM up to timestep  $t$ .
More recent work has shown that the FH mechanism is prone to representation collapse if observations are visually similar to each other (Paischer et al., 2022b,c). They propose a new method, namely HELMv2, which substitutes the random mapping with a pretrained CLIP encoder. Subsequently, they adopt a batch normalization layer (Ioffe & Szegedy, 2015) with fixed shifting and scaling parameters to transfer the image embedding to the language space. Consequently, HELMv2 computes the inputs to the LM as
$$
\boldsymbol {x} _ {t} = \operatorname {B N} _ {\beta = \boldsymbol {\mu} _ {E}, \gamma = \boldsymbol {\sigma} _ {E}} \left(\operatorname {C L I P} _ {\mathrm {V M}} \left(\boldsymbol {o} _ {t}\right)\right), \tag {3}
$$
where  $\mathrm{CLIP}_{\mathrm{VM}}$  denotes the CLIP vision model and  $\mu_{E}$  and  $\sigma_{E}$  denote mean and standard deviation of the embedded vocabulary  $E$ . This effectively fits the statistics of the image embeddings to those of the LM embeddings. Since the embedding spaces of CLIP and the LM were trained independently they are not semantically aligned. Therefore, also HELMv2 fails to preserve semantics of observations and, consequently, the memory mechanism of HELMv2 is not human-readable.
2.2 Semantic HELM
Semantic HELM (SHELM) inherits the high-level architecture from HELM but introduces some changes to the memory module. Similarly to HELMv2, we also use CLIP to embed environment observations. However, we replace the batch normalization layer of HELMv2 with a token-retrieval mechanism. We retrieve tokens that are similar to an observation in CLIP space and pass them to the LM in the form of text so they can be regarded as textual descriptions of environment observations.
In a first step, we determine the overlap of the token vocabularies of CLIP and the LM. This is necessary to (i) to control for the number of tokens the LM receives per observation, and (ii) to avoid loss of information due to different tokenizers used by CLIP and the LM. Thereby, we obtain a set of tokens  $\mathcal{V}$ . Since CLIP was pretrained on image-caption pairs, we augment each token  $v\in \mathcal{V}$  with a set of pre-defined prompts  $\mathcal{P} = \{p_1,p_2,\ldots \}^1$ . The contents of  $\mathcal{P}$  are hyperparameters of our method and can be designed to specifically target certain aspects of the different environments. We embed a token  $v$  in the CLIP output space by computing the average embedding of its prompt augmentations. That is, we define the function
$$
\operatorname {e m b e d} (v) = \frac {1}{| \mathcal {P} |} \sum_ {p \in \mathcal {P}} \operatorname {C L I P} _ {\mathrm {L M}} (\operatorname {c o n c a t} (p, v)). \tag {4}
$$
We do this for every  $v \in \mathcal{V}$ , which results in a set  $S$  of CLIP-embedded tokens
$$
\mathcal {S} = \{\operatorname {e m b e d} (v) \} _ {v \in \mathcal {V}}. \tag {5}
$$
We denote by  $\arg \max^k$  an extension of the arg max operator that returns the subset of the  $k$  largest elements in a set. For each observation we retrieve the  $k$  most similar tokens in terms of cosine similarity by
$$
\mathcal {S} ^ {*} = \arg \max  _ {\boldsymbol {s} \in \mathcal {S}} ^ {k} \operatorname {c o s s i m} (\boldsymbol {s}, \operatorname {C L I P} _ {\mathrm {V M}} (\boldsymbol {o} _ {t})), \tag {6}
$$
$$
\operatorname {c o s s i m} (\boldsymbol {x}, \boldsymbol {y}) = \frac {\boldsymbol {x} ^ {\top} \boldsymbol {y}}{\| \boldsymbol {x} \| \| \boldsymbol {y} \|}. \tag {7}
$$
Note that  $k = |\mathcal{S}^*|$  is another hyperparameter of our method, namely the number of tokens that represent an observation. Effectively,  $k$  controls the degree of compression in the memory.
Finally, we embed single tokens  $v$  corresponding to the set of tokens in  $S^*$  in the LM embedding space and pass them to the LM. In this manner, we provide the LM with a textual representation of the current observation. While HELM and HELMv2 also leverage the latent structure of a pre-trained LM, they do not explicitly represent the observations in the form of text, thus, do not provide a human-readable form of past observations. Another improvement over HELMv2 is that SHELIM removes the restriction of HELMv2 that the embedding spaces of CLIP and LM must have the same dimensionality. In turn, any CLIP-like encoder can be used as semantic extractor for SHELIM. Figure 2 shows an illustration of the methodology of SHELIM.
First, we investigate in Section 3.1 whether CLIP can extract semantics of artificial scenes. Next, we train SHELM on four different environments, namely MiniGrid (Chevalier-Boisvert et al., 2018), MiniWorld (Chevalier-Boisvert, 2018), Avalon (Albrecht et al., 2022), and Psychlab (Leibo et al., 2018). We compare SHELM to HELMv2, LSTM (a recurrent baseline based on the LSTM architecture), and the popular Dreamerv2 (Hafner et al., 2021) and Dreamerv3 (Hafner et al., 2023). We show that a semantic memory boosts performance in environments that are both heavily dependent on memory and photorealistic. Finally, in Section 3.6 we perform ablation studies on the benefit of semantics and the trade-off between interpretability and performance.
We train all HELM variants and LSTM with Proximal Policy Optimization (PPO, Schulman et al., 2017) on RGB observations. Following Paischer et al. (2022a), we instantiate the LM with a pretrained TransformerXL (TrXL, Dai et al., 2019) model. For training Dreamerv2 and Dreamerv3, we use the respective codebases $^{23}$  and train on RGB observations. We report results via IQM (Agarwal et al., 2021) and  $95\%$  bootstrapped confidence intervals (CIs) unless mentioned otherwise. We follow (Colas et al., 2019) and perform a Welch's t-test with a reduced significance level of  $\alpha = 0.025$  at the end of training to test for statistical significance. We elaborate on the architectural design and hyperparameter sweeps in Appendix F.
$^{1}$ Following the original implementation at https://github.com/openai/CLIP/blob/main/ notebooks/Prompt_Engineering_for/ImageNet.ipynb
Dreamerv2: https://github.com/danijar/dreamerv2
Dreamerv3: https://github.com/danijar/dreamerv3
Figure 3: Left: Accumulated reward for different methods on six MiniGrid environments, Middle: on the MiniGrid-Memory task, Right: on eight MiniWorld tasks. We report IQM and  $95\%$  CIs across 30 seeds for each method.
3.1 Extracting semantics of virtual scenes
First, we analyse whether CLIP vision encoders are able to extract semantics from artificial scenes that are typically encountered in RL environments. We compare the two smallest ResNet (He et al., 2016) and ViT (Dosovitskiy et al., 2021) architectures, namely RN50 and ViT-B/16, since those add the least amount of computational overhead to SHELM. In this regard, we look at observations sampled via a random policy, provide them to the CLIP vision encoder and retrieve the closest tokens in the CLIP embedding space as explained in Section 2.2. We find that the retrieval of tokens strongly varies between vision encoder architectures for MiniWorld and Avalon (see Figure 9 and Figure 8 in the appendix). The differences are especially prominent for MiniWorld environments, where the ViT-B/16 encoder recognizes shapes and colors, whereas RN50 ranks entirely unrelated concepts highest. More photorealistic observations from the Avalon benchmark show a similar, but less pronounced trend. We also observe a strong bias toward abstract tokens such as screenshot, biome, or render. We alleviate the retrieval of these tokens by including them in our prompts for retrieval. Therefore, instead of using the prompt "An image of a <tok>", we consider prompts such as "A screenshot of a <tok>", or "A biome containing a <tok>", where <tok> stands for a token in  $\mathcal{V}$ . We only use prompts for Avalon and Psychlab, since the effect of this prompting scheme was negligible on the other environments. Table 1 in the appendix features the full list of prompts for the two environments. We also add an analysis on using off-the-shelf captioning engines, such as BLIP-2 (Li et al., 2023) in Appendix B. We find that, while BLIP-2 correctly recognizes shapes, colors, and objects, it lacks accuracy in their compositionality. Based on this analysis we use the ViT-B/16 encoder in combination with our designed prompts as semantics extractor for SHELM.
3.2 MiniGrid
We compare all methods on a set of six partially observable grid-world environments as in as (Paischer et al., 2022a). Additionally, we train on the MiniGrid-MemoryS11-v0 environment, which we refer to as Memory. The Memory task requires the agent to match the object in the starting room to one of the two objects at the end of a corridor. The agent's view is restricted so it cannot observe both objects at the same time. Therefore, it needs to memorize the object in the starting room. If the agent chooses the wrong object at the end of the corridor, it receives no reward and the episode ends.
Figure 3 (left and middle, respectively) shows the results across the six MiniGrid environments and the Memory environment. On the MiniGrid environments, Dreamerv3 excels and converges much faster than any other method. However, final performance is approximately equal for Dreamerv3, HELMv2, and SHELIM. Interestingly, although all MiniGrid environments are partially observable, they are solvable with a memory-less policy. In the Memory environment, SHELIM performs on par with Dreamerv3 but exhibits faster convergence and more stable training. Figure 11 visualizes the tokens passed to the memory of SHELIM after sampling episodes from a trained policy. SHELIM primarily maps the ball to the token miner, and the key to the token narrow. Although the retrieved tokens do not represent the semantically correct objects in the observations, they are still stored as two different language tokens in the memory. This enables SHELIM to discriminate between them and enables faster learning as mirrored in the learning curves. The LSTM baseline suffers from poor
Method	Mean Human Normalized Score
HELM	0.143 ± 0.010
HELMv2	0.145 ± 0.010
PPO	0.146 ± 0.010
SHELM	0.155 ± 0.011
Dreamerv2	0.168 ± 0.012
Figure 4: Left: Mean and standard deviation of human normalized score across all tasks on the Avalon test worlds after  $10\mathrm{M}$  timesteps. Right: Two sample episodes and their corresponding token mappings for episodes sampled from Avalon with a random policy.
sample efficiency and does not learn to utilize its memory within the budget of  $2\mathrm{M}$  interaction steps. LSTM requires approximately 5 M steps to solve the task.
3.3 MiniWorld
The MiniWorld benchmark suite provides a set of minimalistic 3D environments. In line with Paischer et al. (2022b), we select eight different tasks from MiniWorld and train our method and all baselines on those. The tasks comprise collecting and placing objects, finding objects, or navigating through mazes. A detailed description of each task can be found in Appendix A. Prior work on the MiniWorld benchmark mostly used single environments to evaluate specialized approaches (Liu et al., 2021; Venuto et al., 2019; Zha et al., 2021; Khetarpal et al., 2020), or handcraft novel environments (Sahni et al., 2019; Hutsebaut-Buysse et al., 2020). To the best of our knowledge, HELMv2 is the only method that was evaluated on all eight tasks and has been state of the art so far (Paischer et al., 2022b).
Figure 3 (right) shows the results for all methods. Interestingly, SHELM reaches performance on par with HELMv2, even though semantics can be extracted from the 3D observations to a certain degree (as explained in Section 3.1). Further, Dreamerv2 outperforms both and reaches state-of-the-art performance. Surprisingly, Dreamerv3 attains a significantly lower score and performs on par with HELM and the memory-less PPO. This might be due to suboptimal hyperparameters, even though Hafner et al. (2023) claim that the choice of hyperparameters transfers well across environments. LSTM again suffers from poor sample efficiency reaching the lowest scores out of all compared methods. Finally, we again observe that PPO can in principle solve all tasks, which yields further evidence that partial observability does not automatically imply necessity for memory. Our results suggest that a memory mechanism can result in enhanced sample efficiency (SHELM vs PPO). However, memory is not imperative to solve these tasks.
3.4 Avalon
Avalon is an open-ended 3D survival environment consisting of 16 different tasks. The aim for the agent is to survive as long as possible by defending against predators, hunting animals, and eating food in order to restore energy. An episode ends if the agent has no energy left. This can happen if the agent receives environmental damage (e.g. from falling), is being killed by a predator, or does not eat frequently. The agent receives a dense reward as the difference in energy between consecutive timesteps. Additionally, it receives a delayed reward upon successful completion of a task, e.g., eating food. The observation space are RGBD images as well as proprioceptive input that comprise, e.g., the agent's energy. We adopt the same architecture and training strategy as Albrecht et al. (2022) for all HELM variants. Specifically, we add the history compression branch to the baseline PPO agent and train on all 16 tasks including their difficulty curriculum. The history compression branch is only based on RGB images and does not receive the proprioceptive state. We also train Dreamerv2, and a memory-less baseline (PPO), which is identical to the PPO baseline in Albrecht et al. (2022). Due to computational constraints and since we observed superior performance of Dreamerv2 on 3D environments, we neglect Dreamerv3 and LSTM and train on a limited budget of  $10\mathrm{M}$  interaction steps. We elaborate on training details and hyperparameter selection in Appendix F. The final performance of an agent is measured in terms of mean human normalized scores on a curated set of 1000 test worlds.
Figure 5: Performance for all methods on CR task from Psychlab. We report IQM and  $95\%$  CIs across 5 seeds (left). Observation and corresponding tokens for SHELM on CR environment (right).
The results are shown in Figure 4, left. For detailed results per task see Table 2 in the appendix. SHELM and Dreamerv2 yield the highest performance on average after  $10\mathrm{M}$  interaction steps. However, the difference to the memory-less PPO is not statistically significant. To further investigate this finding we train PPO for  $50\mathrm{M}$  interaction steps and compare it to the baseline results reported in Albrecht et al. (2022) in Appendix D. Indeed, we find that our memory-less baseline attains scores on-par with memory-based approaches trained for  $50\mathrm{M}$  interaction steps. This yields further evidence that Avalon does not necessitate the addition of a memory mechanism. Finally, we can glance at SHELM's memory to identify failure cases. We show observations of two episodes and their token correspondences for SHELM in Figure 4, right. The observations are mostly mapped to semantically meaningful tokens that represent parts in the image observations. However, we also observe some cases where CLIP retrieves tokens that are semantically related, but not correct, which we show in Figure 12 in the appendix.
3.5 PsychLab