SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The Psychlab environment suite (Leibo et al., 2018) consists of 8 complex tasks and was designed to isolate various cognitive faculties of RL agents including vision and memory. The continuous recognition (CR) task of Psychlab was specifically designed to target the memory capacity of an agent. In this task the agent is placed in front of a monitor which displays a stream of objects. The agent then needs to identify whether it has already seen a particular object by swiping either left or right. If the agent correctly identifies an object it receives a reward of  $+1$ . A policy that always swipes in one direction achieves a reward of around 30, which we refer to as the random baseline. Every reward higher than that indicates that an agent actually utilizes its memory. This task is ideal to evaluate the memory capacity of an agent, since episodes last for about 2000 timesteps and usually require remembering up to 50 different objects.
We train all methods for  $10\mathrm{M}$  interaction steps on the CR task (Figure 5). We neglect the memoryless baseline because this task is unsolvable without memory. SHELM indeed learns to effectively utilize its memory within  $10\mathrm{M}$  interaction steps and significantly outperforms all competitors. Other approaches require interaction steps in the range of billions until convergence (Parisotto et al., 2020; Fortunato et al., 2019). These works report human normalized scores which are not publicly available, therefore we cannot compare SHELM to those. HELMv2 is the only competitor that attains a performance better than random, but reaches significantly lower scores than SHELM. Surprisingly, both variants of Dreamer do not exceed performance of the random baseline. Finally, we inspect the memory of SHELM and show the token mappings that are passed on to the memory module for some sample observations in Figure 5. In most cases SHELM assigns semantically correct tokens to the displayed objects. However, we also show some cases where the token retrieval of SHELM conflates different objects in Figure 13. We find that this can mostly be attributed to the low resolution of observations. SHELM can recover from these failure cases when using higher resolutions (see Figure 14).
3.6 Ablation studies
Are semantics important? We slightly alter the HELMv2 implementation from (Paischer et al., 2022b,c) by retrieving the closest tokens in the language space after their frozen batch normalization layer, and finally passing those to the LM. This setting is similar to SHELM in that the LM receives tokens in the form of text. However, semantics are not preserved since visual features are merely shifted to the language space. We call this setting HELMv2-Ret and train it on the Memory environment (see Figure 6). We find that if the mapping from observation to language is arbitrary, the performance decreases drastically.
Is it important to learn task-specific features? SHELM would be even more interpretable if not only the memory module but also the branch processing the current observation (CNN in Figure 2) could utilize language tokens. Therefore, we substitute the CNN with a CLIP vision encoder (SHELM-CLIP). An important consequence of this methodological change is that even features from the current observation can be interpreted by our retrieval mechanism. However, Figure 6 suggests that it is vital to learn task-specific features.
Should the history branch operate on task-specific features? To answer this question we substitute the CLIP encoder in the history branch with the CNN encoder learned at the current timestep (SHELM-CNN). An immediate drawback of that is the loss of readability, since we do not represent the observations in text form anymore. Further, we again observe a drastic drop in performance. The reason for this is that the task-specific features must be learned before they actually provide any benefit to the history branch. Thus, we prefer to keep abstract semantic features in text form since that does not require any training and has the additional benefit of being human-readable.
Figure 6: Ablation study on the effect of semantics, the influence of task-specific features in the history and the current timestep, and the importance of the pretrained LM. We report IQM and  $95\%$  CIs across 30 seeds on the MiniGrid-Memory task.
Is a pretrained language encoder required? We replace the LM with an LSTM operating on the text tokens (SHELM-LSTM). This setting resulted in performance equivalent to randomly choosing an object at the end of the corridor. However, we believe that after longer training SHELM-LSTM can eventually learn to solve the task, since the LSTM baseline also solved the task after longer training. Thus, the simpler
and more sample efficient method is to maintain the frozen pretrained encoder instead of learning a compressed history representation.
We show additional results for an ablation study where we exchange the ViT-B/16 vision encoder with a RN50 in Appendix E. SHELM appears to be quite sensitive to the selection of vision encoder, which corroborates our qualitative findings in Section 3.1.
RL and partial observability RL with incomplete state information can necessitate a memory for storing the history of observations an agent encountered. However, storing the entire history is often infeasible. History Compression tries to answer the question of what information to store in a stream of observations (Schmidhuber, 1992; Schmidhuber et al., 1993; Zenke et al., 2017; Kirkpatrick et al., 2016; Schwarz et al., 2018; Ruvolo & Eaton, 2013). A plethora of prior works have used history compression to tackle credit assignment (Arjona-Medina et al., 2019; Patil et al., 2022; Widrich et al., 2021; Holzleitner et al., 2021), and partial observability (Hausknecht & Stone, 2015; Vinyals et al., 2019; Berner et al., 2019; Pleines et al., 2022). The memory maintained by an agent can either
be external (Hill et al., 2021; Wayne et al., 2018), or directly integrated into the feature extraction pipeline via the network architecture. An orthogonal approach for history compression is training recurrent dynamics models (Ha & Schmidhuber, 2018; Pasukonis et al., 2022; Hafner et al., 2020, 2021). We believe language is very well suited as medium for compression to summarize past events, as suggested by Paischer et al. (2022a).
Language in RL Language provides useful abstractions for RL. These abstractions can be leveraged to represent high-level skills (Sharma et al., 2022; Jiang et al., 2019; Jacob et al., 2021). LMs have been used to improve exploration in text-based environments (Yao et al., 2020), or in visual environments via a language oracle (Mu et al., 2022). Pretrained vision-language models (VLMs) provide abstract embedding spaces that can be used for exploration in visual environments (Tam et al., 2022). We leverage VLMs to spatially compress visual observations to language tokens. Furthermore, pretrained LMs were leveraged for (i) initializing policies in text-based environments (Li et al., 2022), (ii) grounding to various environments (Andreas et al., 2018; Huang et al., 2023; Carta et al., 2023; Hill et al., 2021), (iii) sequence modeling in the offline RL setup (Reid et al., 2022), and (iv) generating high-level plans (Huang et al., 2022a,b; Wang et al., 2023; Singh et al., 2022; Ichter et al., 2022; Liang et al., 2022a; Dasgupta et al., 2023; Du et al., 2023; Shah et al., 2022; Ahn et al., 2022; Zeng et al., 2022). Other works train agents to generate descriptions of virtual scenes (Yan et al., 2022), or thought processes of humans (Hu & Clune, 2023). Additionally, language has been used for augmenting the reward function (Wang et al., 2019; Bahdanau et al., 2019; Goyal et al., 2019; Carta et al., 2022; Kwon et al., 2023), or learning a dynamics model (Zhong et al., 2022, 2021, 2020; Wu et al., 2023). To manage a language-based memory, Park et al. (2023) stores a record of an agent's memory which comprises different levels of abstraction. Importantly, all agent-environment interactions in their work are scripted, while our agent enables a memory based on language for visual inputs.
Language for interpretability In the realm of supervised learning, a plethora of prior works had used language as human-interpretable medium to explain classification decisions in computer vision (Hendricks et al., 2016, 2018; Park et al., 2018; Zellers et al., 2019; Hernandez et al., 2022), or in natural language processing (Andreas & Klein, 2017; Zaidan & Eisner, 2008; Camburu et al., 2018; Rajani et al., 2019; Narang et al., 2020). Interpretability methods in RL are scarce and usually follow a post-hoc approach (Puiutta & Veith, 2020). Intrinsically interpretable models are designed to be inherently interpretable even during training time and are preferable over post-hoc approaches (Rudin et al., 2021). They often restrict the complexity of the model class, which in turn results in reduced performance of the agent. Therefore, (Glanois et al., 2021) propose to adopt a modular approach to interpretability. To this end, our work focuses on the memory module. This enables us to provide some extent of intrinsic interpretability while exceeding performance of existing (non-interpretable) methods on tasks that necessitate memory.
Foundation models The advent of the Transformer architecture (Vaswani et al., 2017) gave rise to foundation models (FMs, Bommasani et al., 2021), such as GPT-3 (Brown et al., 2020). As shown by Petroni et al. (2019); Talmor et al. (2020); Kassner et al. (2020); Mahowald et al. (2023), pretrained LMs can learn abstract symbolic rules and show sparks of reasoning. We leverage their abstraction capabilities for history compression in RL. Further, vision FMs have been demonstrated to be well adaptable to foreign domains (Adler et al., 2020; Evci et al., 2022; Ostapenko et al., 2022; Parisi et al., 2022). Our approach combines language-based FMs with vision-language models, such as CLIP (Radford et al., 2021) or ALIGN (Jia et al., 2021). We use CLIP to obtain language tokens that semantically correspond to concepts present in synthetic environments.
Token-level abstraction One potential shortcoming of our method is that our retrieval is based on single tokens. However, we have shown that for environments where one or a few objects are important, this is sufficient. Further, our approach is very flexible and the semantic database can easily be augmented with more detailed descriptions of objects and their relationships. We aim to investigate more in this direction in future work.
Wall clock time A current limitation of SHELM is wall clock time. The rollout phase is particularly expensive since each timestep needs to be propagated through the LM. Despite that, it is still more efficient than, e.g., Dreamerv2. This is due to the fact that the memory mechanism is kept frozen and need not be updated durint the training phase. A potential solution for decreasing the complexity of the rollout phase would be distillation of the LM into smaller actor networks as in (Parisotto & Salakhutdinov, 2021).
Modality gap Our semantic mapping is limited by the inherent ability of CLIP vision encoders to extract semantically meaningful features in synthetic environments. However, CLIP suffers from the modality gap (Liang et al., 2022b), i.e., a mis-alignment between image and text embedding spaces. In the future, we aim at incorporating methods that mitigate the modality gap (Fürst et al., 2022; Ouali et al., 2023).
Distribution shift We use a pretrained CLIP model to retrieve tokens that are similar to visual observations. However, CLIP was pretrained on large-scale data crawled from the web. We have shown that it can still extract semantics of synthetic environments when those are sufficiently photorealistic, i.e., MiniWorld, Avalon, or Psychlab. Further, prompting can enhance the quality of the retrieval. For environments such as MiniGrid, the retrieval yields tokens that do not correspond to aspects in the image. However, Gupta et al. (2022) has shown that CLIP can handle task-specific texts for MiniGrid, thus, we would like to investigate the effect of augmenting our semantic database with such texts in the future.
In many real-world scenarios an agent requires a memory mechanism to deal with partial observability. Current memory mechanisms in RL act as a black box where it is not comprehensible for humans what pieces of information were stored. To solve this problem, we proposed a new method called Semantic HELM that represents past events in form of human-readable language by leveraging pretrained vision-language models. We showed compelling evidence that even for synthetic environments our memory mechanism can extract semantics from visual observations. Further, SHELM outperforms strong baselines on photorealistic memory-dependent environments through its human-readable memory module. In cases where SHELM fails to extract semantics from observations, we can investigate the cause by inspection of the memory module. Even in such cases, SHELM mostly performs on-par with other memory-based approaches.
We believe that we can further enhance our method by generating full captions from history observations instead of only a small number of words. This could enable (i) a long-term memory that textually summarizes all captions of corresponding observations, (ii) a potential for planning in form of text from visual observations similar to Patel et al. (2023), and (iii) modeling dynamics of an environment in language space.
Yee, E., Jones, M., and McRae, K. Semantic Memory. 10 2017. doi: 10.1002/9781119170174. epcn309.
Zaidan, O. and Eisner, J. Modeling annotators: A generative approach to learning from annotator rationales. In 2008 Conference on Empirical Methods in Natural Language Processing, EMNLP 2008, Proceedings of the Conference, 25-27 October 2008, Honolulu, Hawaii, USA, A meeting of SIGDAT, a Special Interest Group of the ACL, pp. 31-40. ACL, 2008.
Zellers, R., Bisk, Y., Farhadi, A., and Choi, Y. From recognition to cognition: Visual commonsense reasoning. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pp. 6720-6731. Computer Vision Foundation / IEEE, 2019. doi: 10.1109/CVPR.2019.00688.
Zeng, A., Wong, A., Welker, S., Choromanski, K., Tombari, F., Purohit, A., Ryoo, M. S., Sindhwani, V., Lee, J., Vanhoucke, V., and Florence, P. Socratic Models: Composing Zero-Shot Multimodal Reasoning with Language. CoRR, abs/2204.00598, 2022. doi: 10.48550/arXiv.2204.00598. arXiv: 2204.00598.
Zenke, F., Poole, B., and Ganguli, S. Continual Learning Through Synaptic Intelligence. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 3987-3995. PMLR, 2017.
Zha, D., Ma, W., Yuan, L., Hu, X., and Liu, J. Rank the episodes: A simple approach for exploration in procedurally-generated environments. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Zhong, V., Rocktäschel, T., and Grefenstette, E. RTFM: generalising to new environment dynamics via reading. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Zhong, V., Hanjie, A. W., Wang, S. I., Narasimhan, K., and Zettlemoyer, L. SILG: the multi-domain symbolic interactive language grounding benchmark. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 21505-21519, 2021.
Zhong, V., Mu, J., Zettlemoyer, L., Grefenstette, E., and Rocktäschel, T. Improving policy learning via language dynamics distillation. CoRR, abs/2210.00066, 2022. doi: 10.48550/arXiv.2210.00066.
Aström, K. J. Optimal control of Markov processes with incomplete state information. Journal of Mathematical Analysis and Applications, 10:174-205, 1964.
We choose 8 diverse 3D environments of the MiniWorld benchmark suite:
- CollectHealth: The agent spawns in a room filled with acid and must collect medikits in order to survive as long as possible.
- FourRooms: The agent must reach a red box that is located in one of four interconnected rooms.
- MazeS3Fast: A procedurally generated maze in which the agent needs to find a goal object.
- PickupObjs: Several objects are placed in a large room and must be collected by the agent. Since the agent receives a reward of 1 for each collected object, the reward is unbounded.
- PutNext: Several boxes of various colors and sizes are placed in a big room. The agent must put a red box next to a yellow one.
- Sign: The agent spawns in a U-shaped maze containing various objects of different colors. One side of the maze contains a sign which displays a color in written form. The aim is to collect all objects in the corresponding color.
- Tmaze: The agent must navigate towards an object that is randomly placed at either end of a T-junction.
- YMaze: Same as TMaze, but with a Y-junction.
We neglect the OneRoom and the Hallway environments, since those are easily solved by all compared methods. Further, we neglect the Sidewalk environment since it is essentially the same task as Hallway with a different background. Since the rewards of PickupObjs and CollectHealth are unbounded, we normalize them to be in the range of [0, 1], which is the reward received in all other environments. For a more detailed description of the MiniGrid environments we refer the reader to Paischer et al. (2022a).
B Token retrieval for synthetic environments
Since we perform retrieval on a token level we investigate the effect of augmenting tokens with different prompts, and the effect of different vision encoders on retrieval performance. We analyse the former at the example of the Avalon environment. Figure 7 shows some examples. We observed that simply encoding single tokens in CLIP space results in a bias toward abstract tokens such as biome, or screenshot. The same occurs for using the prompts originally used for zero-shot classification on the ImageNet dataset $^{4}$  (Deng et al., 2009). However, one can alleviate this bias by including these frequently retrieved tokens in the prompt itself (see Figure 7, bottom). However, we found this strategy to be effective only for Avalon and Psychlab environments. The sets of prompts for both environments can be found in Table 1. For MiniGrid and MiniWorld we retrieve single tokens without any prompt.
Next, we investigate the influence of the choice of vision encoder architecture, e.g., RN50 vs ViT-B/16. We only consider those encoders since they induce the least complexity on our history compression pipeline. We show the closest tokens in CLIP space for observations of MiniWorld (see Figure 9) and Avalon environments (see Figure 8). For MiniWorld, CLIP can extract shapes and colors. However, the retrievals are very noisy, which is mirrored in the attained score of SHELM. For Avalon, however, the token retrievals are more convincing. Generally, we find that retrievals using the RN50 encoder of CLIP tend to be more noisy than retrieval utilizing a ViT-B/16 encoder.
Instead of using CLIP and our token retrieval we may be able to use pretrained image captioning systems, such as BLIP-2 (Li et al., 2023) directly. To investigate this, we use BLIP-2 to generate captions for MiniGrid-Memory, Avalon and Psychlab (Figure 10). The captions generated for the MiniGrid-Memory in the left column contain correctly detected colors (green and red), and shapes (square), but incorrectly combines them, e.g. there is neither a green square, nor a green arrow pointing to the left. A similar trend can be observed for Avalon in the middle comlumn, where BLIP-2 recognizes a frog, but incorrectly assumes that it stands on top of a rock, while the rock is
${}^{4}$  available at https://github.com/openai/CLIP/tree/main
Cosine similarity
Figure 7: Token rankings for ViT-B/16 encoder of CLIP on Avalon observations. Tokens are encoded with the CLIP language encoder without prompt (top), with ImageNet specific prompts (middle), or prompts designed for Avalon (bottom).
RN50
ViT-B/16
Figure 8: Token rankings for RN50 and ViT-B/16 encoders of CLIP on Avalon observations.
clearly beside the frog in the foreground. The second observation for Avalon is very accurate though. For the two Psychlab observations in the right column, BLIP-2 correctly describes the objects on the screen, but also mentions redundant information, such "the objects are on a screen". Overall, BLIP-2 correctly recognizes shapes, colors, and different objects, but is often incorrect on their compositionality. Further, a drawback of using off-the-shelf captioning engines is that we cannot control for the context length for each observation. This can quickly result in excessive context lengths, which in turn, results in a drastic increase of computational complexity.
C Qualitative analyses
We show token mappings for the memory mechanism of SHELM to identify potential failures and successes. Figure 11 shows a few sample episodes from a trained policy on the Memory environment. Clearly, CLIP is not capable of extracting semantics of the abstract 2D gridworld environments. Thereby, it maps the ball to the token miner and the key to the token narrow. For minimalistic 3D environments, CLIP is capable of extracting colors and shapes of objects as can be seen in Figure 9. However, these results are still uninspiring since the token retrievals are very noisy.
We also visualize episodes collected by a random policy from the Avalon environment in Figure 12. As opposed to the minimalistic MiniGrid and MiniWorld environments, CLIP successfully extracts semantically meaningful concepts which are subsequently passed to the memory. There are still some failure cases though. For example, in the second episode the agent is attacked by a bee. The feelers of the bee are mistaken for a sword as the bee moves into close proximity of the agent. Further after the agent successfully defends itself against the bee, it mistakes the dead bee for a crab. Alleviating
Figure 9: Token rankings for RN50 and ViT-B/16 encoders of CLIP on MiniWorld observations.
Figure 10: Captions generated via BLIP-2 for observations of the MiniGrid-Memory task, Avalon and Psychlab.
Table 1: Prompts used for token retrieval for the Avalon and Psychlab environments.
Environment	Prompts
Avalon	a screenshot of a screenshot of a a screenshot of many a biome containing a biome containing a a biome containing many a biome full of
Psychlab	a render of a render of a a screenshot of a screenshot of a a screen showing a screen showing a
such failure cases would require grounding CLIP to the respective environment as was done in (Fan et al., 2022).
Further, we visualize token mappings for objects the agent encounters in the continuous recognition task of Psychlab in Figure 13. The majority of token retrievals are semantically correct. However, sometimes the agent confuses different objects or conflates to the same token for different objects. An example for that are two middle objects in row 4 which are both mapped to the token icon, or the first two objects in row 6 that are mapped to the token tuber. We suspect that this occurs due to downscaling to a lower resolution which is conducted within the DMLab engine. Indeed, when taking a closer look at token retrievals at a higher resolution, they are mapped to different tokens (see Figure 14). Therefore, we consider two different aspects on how to alleviate this issue, (i) increasing the resolution of the observations and (ii) retrieving more than one token for an observation. The former results in increased computational complexity, which we aim to avoid, since the task is computationally very expensive already. The latter assumes that the retrieval differs at least in their top- $k$  most likely tokens and is a viable solution.