Foundation Models for History Compression in Reinforcement Learning
Abstract
Agents interacting under partial observability require access to past observations via a memory mechanism in order to approximate the true state of the environment.
Recent work suggests that leveraging language as abstraction provides benefits for creating a representation of past events.
History Compression via Language Models (HELM) leverages a pretrained Language Model (LM) for representing the past.
It relies on a randomized attention mechanism to translate environment observations to token embeddings.
In this work, we show that the representations resulting from this attention mechanism can collapse under certain conditions.
This causes blindness of the agent to subtle changes in the environment that may be crucial for solving a certain task.
We propose a solution to this problem consisting of two parts.
First, we improve upon HELM by substituting the attention mechanism with a feature-wise centering-and-scaling operation.
Second, we take a step toward semantic history compression by leveraging foundation models, such as CLIP, to encode observations, which further improves performance.
By combining foundation models, our agent is able to solve the challenging MiniGrid-Memory environment.
Surprisingly, however, our experiments suggest that this is not due to the semantic enrichment of the representation presented to the LM, but rather due to the discriminative power provided by CLIP.
We make our code publicly available at https://github.com/ml-jku/helm .
Introduction
In Reinforcement Learning (RL) an agent interacts with an environment and learns from feedback provided in the form of a reward function.
RL agents that are deployed in the real world often have to cope with partial observability.
Therefore, the capability to approximate the true state of their surrounding environment by virtue of an agent’s perception is crucial (Åström, 1964; Kaelbling et al., 1998).
To this end, many agents employ a memory mechanism to track events that occurred in the past.
In this memory, it is much more efficient to store abstract representations of the past rather than every detail the agent encountered.
Thus, memory mechanisms such as LSTM (Hochreiter &
Schmidhuber, 1997) or Transformer (Vaswani et al., 2017) compress sequences of high-dimensional observations.
We propose to use foundation models (FM; Bommasani et al., 2021) as memory mechanism.
Since
FMs come pretrained, they allow for a much higher sample efficiency than memory mechanisms trained from scratch.
Moreover, FMs have demonstrated remarkable few-shot capabilities (Brown et al., 2020) and can learn abstract symbolic rules and perform reasoning (Petroni et al., 2019; Talmor et al., 2020; Kassner et al., 2020).
In this work, we aim to leverage these properties, which FMs
Deep Reinforcement Learning Workshop at Neural Information Processing Systems, 2022.
[Figure]
learned on large-scale text corpora, for RL.
We believe that language is well suited as domain for a memory mechanism, since it appears to be an essential component in forming meaningful abstractions.
In humans, the ability to abstract is heavily influenced by the exposure to language in early childhood (Waxman & Markow, 1995).
Recently, there has been a surge of complex models that combine multiple modalities, in particular images and text, as in CLIP (Radford et al., 2021).
Since many RL environments rely on visual inputs, we use CLIP in our proposed memory mechanism to map visual inputs meaningfully into the language domain.
Prior work has illustrated that foundation models pretrained on language can efficiently compress sequences of observations and facilitate agent learning in partially observable RL environments (HELM, Paischer et al., 2022).
The key challenge hereby is to map image-based observations to language representations.
HELM tackles this problem using a randomly chosen mapping called
FrozenHopfield (FH) that is — due to its randomness — inherently unable to form meaningful abstractions.
In fact, we show that under certain conditions the representations are prone to collapse, rendering the agent unable to distinguish between different inputs (see Fig. 1).
We build upon HELM by (i) substituting FH with a feature-wise centering-and-scaling operation and (ii) incorporating a
CLIP image encoder (Radford et al., 2021), which is pretrained in a multimodal fashion on web data consisting of images and text.
We term the resulting new method HELMv2.
Using HELMv2, we are able to distinguish even minute differences in the input when necessary, which drastically enhances downstream performance.
We demonstrate the effectiveness of HELMv2 on a set of partially observable environments.
Concretely, we train on 2D MiniGrid (Chevalier-Boisvert et al., 2018), and 3D MiniWorld environments (Chevalier-Boisvert, 2018).
HELMv2 yields significant improvements over HELM in all environments.
Further, we conduct ablation studies, which show that the improvements are not only due to the CLIP image encoder but also due to the replacement of the FH.
Finally, we construct a mapping that successfully conveys the semantics extracted by CLIP to the LM.
Surprisingly, however, using this mapping does not further improve the results in the selected environments.
Methods
HELM has demonstrated that pretrained LMs are well suited for compressing past observations that are randomly mapped to language tokens.
In this regard, HELM performs two different forms of compression: (i) spatial compression, and (ii) temporal compression.
The former is realized with the
FH mechanism, while the latter is performed with a pretrained TransformerXL (TrXL, Dai et al., 2019).
The spatial compression via FH consists of a random projection matrix P and an attention mechanism over pretrained token embeddings (VocabAttn).
More formally, let E =
( e 1 , . . .
, e k ) ⊤ ∈ R k × m be the token embedding matrix of the pretrained LM consisting of k embeddings e i ∈
R m .
At every timestep t , we obtain inputs x t ∈ R m for the LM from observations o t ∈ R n via the FH mechanism by
[Equation]
where σ is the softmax function and P ∈ R m × n has entries sampled from N (0 , n/m ) .
The resulting x t lies in the convex hull of the token embeddings of the LM.
The parameter β is a scaling factor that controls the dispersion of x t within that convex hull.
HELMcs Our aim is to avoid representation collapse caused by the FH mechanism.
We illustrate in
Fig.
1 that FH is prone to representation collapse even for higher values of β , especially if observations tend to be visually similar.
To sidestep this issue, we substitute the attention mechanism in FH with a feature-wise centering-and-scaling operation.
Let B be a batch of observations o t and let µ B ∈ R n be its mean feature vector and σ B ∈ R n the vector of standard deviations both taken over B .
Likewise, let µ E ∈ R m and σ E ∈ R
m be means and standard deviations of the token embeddings E .
Then we compute x t as
[Equation]
where diag( · ) takes a vector and constructs a diagonal matrix from it.
Intuitively, the centering and scaling in the observation space pronounces differences between single observations.
Subsequently,
P approximately preserves the distances between the centered-and-scaled observations according to the Johnson-Lindenstrauss lemma (Johnson & Lindenstrauss, 1984).
Finally, to obtain x t we shift the observations to the language space by adding µ E and scaling by σ E .
We refer to this setting as
HELMcs.
HELMv2 We obtain another setting by encoding the observations with the ResNet-50 CLIP image encoder, which has the same output dimension as the embeddings e i used by TrXL.
Therefore we eliminate the need for the random mapping P .
Intuitively, performing the centering-and-scaling operation in the abstract CLIP space pronounces differences between different concepts encoded by
CLIP.
We believe this is more effective than looking at differences in the raw pixel space.
That is, we let z t =
CLIP( o t ) ∈
R m and construct B from z t and, consequently, µ B ϕ ∈ R m and σ B ϕ ∈ R m .
We compute the LM inputs as
[Equation]
The complexity imposed by the ResNet-50 CLIP encoder is negligible since it is kept frozen and only utilized during inference.
Experimental Results
We investigate the limitations of HELM on the MiniGrid-MemoryS11-v0 environment (Memory,
Fig. 1, left).
The task for the agent is to remember the object in the room it spawns in, after navigating through a corridor.
The corridor ends at a T-junction showing two objects in each direction, one of which is equivalent to the object in the starting room.
The agent then has to choose the direction towards the object in the starting room.
If the agent chooses the wrong direction the episode ends yielding no reward.
However, the objects only slightly differ in shape (green ball vs. green key).
The agent receives a partially observable egocentric view of the environment in the form of RGB images.
We demonstrate that FH collapses to the same representation for both objects by measuring the distances between the softmax distribution σ ( · ) over token embeddings (Fig. 1, right).
Indeed, we observe that even for high values of β there is very little chance that HELM can discriminate between the two objects.
Moreover, we take a closer look at the resulting representations for the two different observations and measure their distance in terms of cosine similaritiy.
Compared to
HELM and HELMcs, HELMv2 improves separability of the two observations by a large margin (see Fig.
2, right).
Based on these findings we add higher values for β to the gridsearch for HELM.
Also, we show the performance of a recurrent agent based on the LSTM architecture (Hochreiter &
Schmidhuber, 1997).
HELM, HELMcs and LSTM do not achieve better performance than randomly choosing a path at the end of the corridor (see Fig. 2, left) after 2M interaction steps.
On the contrary,
HELMv2 is able to distinguish and memorize the objects and solve the Memory task.
Additionally, we compare HELMv2 to HELMcs, HELM and an LSTM baseline on a set of six diverse partially observable gridworld environments.
Particularly, we select the same MiniGrid environments as Paischer et al.
Moreover, we add eight 3D environments from the MiniWorld benchmark suite (Chevalier-Boisvert, 2018).
For more details about our selected MiniWorld environments see
Appendix A.1.
Results are shown in Fig.
3 for MiniGrid (left) and MiniWorld (right).
HELMv significantly outperforms HELM on both benchmark suites ( p = 2 . 24 e-
10 , p = 0 .
012 , for MiniGrid,
[Figure]
A cc u m u l a t ed
R e w a r d
[Figure]
HELMv2 HELM
[Figure]
Method
Cosine Similarity ( ↓ )
HELM, β =0.99 ± 1.68eHELM, β =0.99 ± 5.65eHELM, β =0.99 ± 8.2eHELM, β =0.97 ± 0.
[Figure]
0.99 ± 3.65e
[Figure] [Figure]
0.83 ±
[Figure] [Figure]
[Figure]
and MiniWorld, respectively).
Also, HELMcs significantly outperforms HELM on the MiniGrid environemnts ( p = 8 .
7 e- 4 ).
There is no significant difference between the performance of HELM and
HELMcs on MiniWorld environments.
As in Paischer et al.
(2022), the LSTM baseline is consistently outperformed by the HELM variants.
In the following we analyse what components lead to the improved performance of HELMv2.
We conduct ablation studies to isolate the effect of the different components.
In this regard, we add two additional settings: (i) HELMshift, and (ii) HELMclip.
HELMshift adds the centeringand-scaling operation from HELMcs to the original HELM setting.
HELMclip uses the same
ResNet-50 CLIP image encoder with the centering-and-scaling operation followed by VocabAttn.
We compare the performance of HELMclip and HELMshift to HELMv2 and HELM on all MiniGrid, and MiniWorld environments.
Fig.
4 shows that adding the centering-and-scaling operation to HELM (HELMshift) does not lead to improved performance.
In fact, on MiniWorld environments it even leads to worse performance than HELM.
We suspect this is due to the frequent pixel changes in the observation space that do not correspond to significant changes in the environment state.
While we also use image observations in MiniGrid the observation space is much more abstract since entire tiles change at once.
Substituting P in HELMshift with the CLIP encoder (HELMclip) however,
Number of Interaction Steps 0 .0 .0 .0 .0 .1 .
A cc u m u l a t ed
R e w a r d
MiniGrid
HELMv2 HELMclip
HELMshift HELM
Number of Interaction Steps 0 .0 .0 .0 .0 .0 .0 .0 .
A cc u m u l a t ed
R e w a r d
MiniWorld
HELMv2 HELMclip
HELM HELMshift results in an immediate improvement on MiniGrid environments.
Since HELMclip still uses the
VocabAttn, it is prone to representation collapse if changes in the CLIP space are subtle (see Fig. 1).
Finally, discarding VocabAttn (HELMv2) drastically improves performance on both, MiniGrid and
MiniWorld environment suites.
While VocabAttn can help in some environments, its proneness to collapse hurts performance overall.
To show the effect of different image encoders, we perform an ablation where we substitute the
CLIP image encoder with a ResNet (He et al., 2016, HELMv2-RN34-IN), and a Vision Transformer (Dosovitskiy et al., 2021, HELMv2-ViT-L16-IN), pretrained in a supervised manner on the popular
Imagenet dataset (Deng et al., 2009).
Additionally, we compare to a Vision Transformer version of CLIP (HELMv2-ViT-B/16).
First, we measure cosine similarities to quantify the ability of the different vision encoders to separate the two observations (see Fig.
5, right).
The HELMv2-ViT-B/ variant exhibits the lowest cosine similarity between both observations.
However, better separability does not correlate with improved downstream performance (see Fig.
5, left).
In fact, the best performing agent uses the ViT-L16 pretrained on ImageNet.
An explanation for this might be the difference in scale (Table 5) or the different pretraining paradigms (i.e. supervised vs. contrastive).
We aim at answering this question in future work.
All our methods are trained with Proximal Policy Optimization (PPO, Schulman et al., 2017) on
RGB observations.
We evaluate all methods via the interquartile mean (IQM) and 95% bootstrapped confidence intervals (CIs) as suggested by Agarwal et al.
To test for statistical significance, we perform a Wilcoxon test (Wilcoxon, 1945) at the end of training.
We follow the architectural design of Paischer et al.
but extend their hyperparameter search according to Appendix A.4.
Investigation on Semantic History Compression
We want to pave the way toward a semantic history compression for RL, that is, to preserve the semantic concepts of an observation when mapping it into the LM space.
In this regard, we provide compelling evidence that it is indeed feasible to create such a semantic mapping between CLIP and the LM (see Fig. 6).
When using this semantic mapping in the two selected minimalistic RL environments, however, we do not observe a significant improvement over HELMv2.
We can compute a lightweight mapping using the vocabularies of the CLIP language encoder
V CLIP and the LM encoder V LM alone.
First, we identify the overlap between the two vocabularies
V OV = V CLIP ∩V
LM with size l = |V OV | .
Next, we embed V OV in the CLIP output space and the LM input space yielding embedding matrices F = ( f 1 , . . .
, f l ) ∈ R n × l and E l ∈ R m × l , respectively.
In general, the mapping matrix W can be computed using any linear model.
Algorithm 1 outlines a procedure for computing such a semantic mapping.
[Figure]
A cc u m u l a t ed
R e w a r d
Memory
HELMv2­ViT­L16­IN HELMv2­RN34­IN
HELMv2­RN50 HELMv2­ViT­B
Method
Cosine Similarity ( ↓ )
HELMv
[Figure]
[Figure] [Figure]
IN
[Figure]
IN
[Figure] [Figure]
Algorithm 1 Associating CLIP output space with LM input space
Require: CLIP language encoder CLIP LM , Vocabulary of CLIP V CLIP , Language Model Embedding
Layer LM ( v i ) , Language Model vocabulary V LM
V OV ← V CLIP ∩ V LM ▷ Search for overlapping vocabulary f i = CLIP LM ( v i ) for v
i ∈ V OV
▷ Embed tokens in CLIP output space e i =
LM ( v i ) for v
i ∈ V OV
▷ Embed tokens in LM input space
W ← create_mapping ( F , E )
▷ Compute mapping between embeddings
A common choice for creating a mapping between monolingual embedding spaces is the Procrustes method (Artetxe et al., 2018; Hoshen & Wolf, 2018; Smith et al., 2017; Lample et al., 2018; Zhang et al., 2016;
Xing et al., 2015; Minixhofer et al., 2022).
In total, we compare four different linear mappings:
• Linear: Ordinary Least Squares, as in (Mikolov et al., 2013) • Ridge: Least Squares with Thikonov regularization • Procrustes: Least Squares with orthogonality constraint (Schönemann, 1966; Gower & Dijksterhuis, 2005))
RobProc: The Robust Procrustes algorithm, that iteratively refines the Procrustes method based on its predicion error ∥ F W − E ∥ 2 F 1 .
We perform a 5-fold cross validation and measure the accuracy, considering each token as its own class.
Table 2 in Appendix A.2 shows the average train and test accuracy for the different linear mapping methods for various CLIP backbones and TrXL.
We observe a strong overfitting effect for
Linear, while constrained optimization (Ridge, Procrustes, and RobProc) generalize better.
The best method in terms of generalization is the Procrustes method.
RobProc does not provide significant gains over Procrustes.
Due to the alignment of the CLIP image encoder and the CLIP language encoder, our mapping can be used to project images to the LM space while preserving its semantics.
We demonstrate that by identifying the closest tokens in the LM space after applying our mapping to project natural images (see Fig. 6).
Further, we conduct a quantitative analysis using publicly available image captioning datasets.
In this regard, we draw a random subset of 1000 image-caption pairs of the popular
MSCOCO dataset (Lin et al., 2014).
The subset is filtered to contain only image-text pairs where the captions contain at least 5 tokens of our computed vocabulary overlap above.
For preprocessing of the captions we remove stop words and apply stemming.
Next, we propagate the image through various
CLIP backbones, to obtain an image embedding and map it to the LM space using our pre-computed mappings.
Finally, we rank tokens in the LM space based on their cosine similarity to the mapped image.
Based on the obtained ranking we compute the Mean Reciprocal Rank (MRR, Craswell, 2009) and the Normalized Discounted Cumulative Gain (NDCG, Järvelin & Kekäläinen, 2002).
Table 1 shows the NDCG for various CLIP backbones mapping to the embedding space of TrXL relative to ranking in the CLIP space (rNDCG).
We observe that the Procrustes mapping consistently outperforms the Linear and Ridge mapping.
Furthermore, there is no improvement in iteratively refining the Procrustes method as in RobProc.
Also, we observe an improvement in rNDCG for larger model architectures.
Table 3 in Appendix A.2 shows the MRR for the various CLIP backbones.
[Table]
Fig.
6 shows promising results for natural images from the MSCOCO dataset and their corresponding token rankings.
Following our analysis we propose a new setting, namely SHELM (for semantic
HELM), that utilizes a CLIP vision backbone in combination with the semantic mapping.
We select the best backbone-mapping combination by ranking all combinations according to the average absolute NDCG (aNDCG), and select the top 5 settings.
Furthermore, we perform a Wilcoxon test for statistical significance between the combinations after bonferroni correction (Bonferroni, 1936).
The two combinations, ViT-B/16+Procrustes, and ViT-L/14*+Procrustes, significantly outperform all competitors in terms of absolute NDCG.
Due to the imposed complexity of ViT-L/14* (see Table 5), we choose ViT-B/16+Procrustes to instantiate SHELM.
[Figure]
We show preliminary results of SHELM on the Memory environment, all MiniGrid, and all MiniWorld environments in Fig.
Additionally, we compare with HELMv2+RandOrtho, which samples random orthogonal matrices from the Haar distribution (Stewart, 1980).
Surprisingly, we observe no statistically significant differences in performance between the different methods.
We believe this is due to the fact that our selected environments are artificial 2D and 3D scenes for which the
CLIP image encoder is unable to sufficiently extract semantically meaningful features.
This is in line with findings of Fan et al.
(2022) who minimally finetune CLIP for the Minedojo environment.
We validate this finding by taking a closer look at the token rankings for the two observations of the
Memory environment in Fig.
8 in the Appendix.
Indeed, the tokens exhibiting the highest similarity in the CLIP space do not describe the semantics of the image, but rather similar concepts, i.e., the
[Figure]
l a t
[Figure]
r d
[Figure]
d
[Figure]
token pong is ranked higher when the ball is present in the observation.
Still, it is worth mentioning that the ranked tokens differ, which aligns with our finding of CLIP’s discriminative power.
The ability of CLIP to extract semantics may improve in more realistic environments as shown in Tam et al.
We aim to apply SHELM to similar environments in the future.
Finally, we conduct an additional analysis on how much the different methods preserve the semantics of natural images.
In this regard, we perform the same quantitative analysis as above for the different methods presented in this work.
We show the MRR and aNDCG for our MSCOCO subset of image-caption pairs in Table 4 in Appendix A.3.
Methods such as HELMv2, HELMclip, or
HELMv2+RandOrtho yield on-par or worse results compared to a random ranking.
This is due to the fact, that for lower values of β the VocabAttn tends to output the mean over token embeddings.
In this case, the resulting ranking is equal across images, and thus, results in lower MRR and NDCG than random rankings.
Also, substituting the RobProc mapping with a random mapping drawn from the Haar distribution (HELMv2+RandOrtho) leads to random rankings.
On the contrary, SHELM preserves more of the semantics resulting in much better rankings.
Adding VocabAttn after the
Procrustes mapping of SHELM again results in loss of information and a worse ranking.
Related Work
Reinforcement Learning with incomplete state information necessitates compression of past events.
A variety of prior works has 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 question of what information to store given a stream of observations was investigated in prior works (Schmidhuber, 1992; Zenke et al., 2017;
Kirkpatrick et al., 2016; Schwarz et al., 2018; Ruvolo &
Eaton, 2013).
Typical choices for network architectures that are capable of compressing a stream of data are LSTM (Hochreiter & Schmidhuber, 1997), and Transformer (Vaswani et al., 2017).
The Transformer architecture has blossomed in the realm of Natural Language Processing.
Many works have focused on scaling pretraining of Transformers to create LMs that generate realistic human-like text (Devlin et al., 2019; Radford & Narasimhan, 2018;
Radford et al., 2019; Brown et al., 2020; Thoppilan et al., 2022; Rae et al., 2021; Zhang et al., 2022; Hoffmann et al., 2022; Chowdhery et al., 2022; Dai et al., 2019).
More recently, interest has sparked in leveraging the Transformer for decision making (Parisotto et al., 2020; Sukhbaatar et al., 2021;
Chen et al., 2021; Janner et al., 2021;
Zheng et al., 2022; Melo, 2022).
Moreover, it has been found that pretrained LMs are very well suited for creating abstractions of a sequence of visual observations (Paischer et al., 2022).
Language provides useful abstractions for Reinforcement Learning.
Hill et al.
(2021) illustrate that language abstraction enables compositional generalization and intrinsic motivation in embodied environments.
LMs have been used for exploration in text-based environments (Yao et al., 2020), or in visual environments that provide a language oracle (Mu et al., 2022).
Multimodal models pretrained with language supervision provide abstract embedding spaces that can be used for visual exploration (Tam et al., 2022).
Further, language pretraining has been leveraged to initialize policies in text-based environments (Li et al., 2022), and for sequence modeling in the offline RL setup (Reid et al., 2022).
Moreover, language has been used for reward shaping (Wang et al., 2019;
Bahdanau et al., 2019; Goyal et al., 2019).
We leverage pretrained multimodal models and map them into a semantically meaningful region of the embedding space of a LM.
Foundation models (Bommasani et al., 2021), such as GPT-3 (Brown et al., 2020), demonstrated remarkable few-shot capabilities.
Pretraining on large-scale text corpora enables such complex models to abstract relevant knowledge for solving math-word problems (Griffith & Kalita, 2019), symbolic math problems (Noorbakhsh et al., 2021), and even university-level math exercises (Drori et al., 2021).
As shown by Petroni et al.
; Talmor et al.
; Kassner et al.
(2020), pretrained
LMs can learn abstract symbolic rules and perform reasoning.
Recently, there has been a surge of complex models trained on large scale data that combines multiple modalities, such as images and text, as in CLIP (Radford et al., 2021).
Moreover, vision FMs have been demonstrated to be well adaptable to foreign domains (Adler et al., 2020; Evci et al., 2022).
The computation of our mapping network is reminiscent of “model-stitching” (Lenc & Vedaldi, 2019;
Bansal et al., 2021; Csiszárik et al., 2021;
Scialom et al., 2020).
In model-stitching an encoder is stitched via a sparse linear transformation to a compatible decoder.
Moschella et al.
(2022) use relative representation spaces to avoid the training of stitching layers, however requires decoder to be trained on these relative spaces.
Merullo et al.
; Scialom et al.
(2020) optimize a linear mapping from a vision encoder to a LM on the downstream task of image captioning.
In contrast, our mapping can be computed with a closed form and does not require optimization of a vast amount of data, while preserving the symmetry between embedding spaces due to the orthogonality constraint.
A plethora of works have focused on conditioning LMs on the popular CLIP model.
Alayrac et al.
(2022) use cross-attention combined with the Perceiver architecture (Jaegle et al., 2021) trained on vast web data to combine a frozen pretrained CLIP (Radford et al., 2021) with a generative language model.
Yu et al.
(2022) adds a contrastive loss between CLIP and a LM encoder for image captioning.
Similarly, (Mokady et al., 2021) trains a mapping model from the CLIP output space to a LM input space for image captioning.
In contrast, (Su et al., 2022) proposes to condition the decoding step of a finetuned LM on image captioning on a frozen CLIP model.
Finally, (Tewel et al., 2021) optimizes the context cache during inference to mazimize a similarity score computed with CLIP for zero-shot image captioning.
In contrast to these methods, our mapping can be efficiently computed offline, and allows conditioning a LM on visual information without any additional optimization procedure.
Zeng et al.
(2022) use language as communication interface between various foundation models to solve certain tasks.
Further, CLIP has demonstrated its robustness and versatility across various tasks in the
RL setup (Ostapenko et al., 2022; Parisi et al., 2022).
We aim to investigate extensions to the CLIP model, such as CLOOB (Fürst et al., 2021), in our framework as well.
Conclusion
Reinforcement Learning with incomplete state information requires an agent to remember past events.
Recent work had illustrated that large pretrained LMs are suitable for creating abstractions of past events (Paischer et al., 2022).
HELM randomly maps observations to language tokens before compressing the compounding information via a LM.
We presented certain conditions under which
HELM is incapable of remembering information that is essential to solve a given task.
Further, we proposed an incremental improvement over HELM comprising a feature-wise centering-and-scaling operation and a pretrained vision encoder, which we call HELMv2.
On the challenging MiniGridMemory environment HELMv2 overcomes the limitations of HELM and successfully solves the task.
Furthermore, HELMv2 yields significant improvements over HELM on minimalistic 2D and 3D environments.
Additionally, we pave the way toward a semantic history compression mechanism, that does not require training in the RL context.
By leveraging multi-modal pretrained foundation models we demonstrated how to obtain a semantic mapping between a vision encoder and a LM.
First results on
MiniGrid and MiniWorld environments, however, showed that a semantic alignment between CLIP and LM does not significantly improve performance.
Inherently, the semantic mapping is limited by the ability of the vision encoder to extract semantically meaningful features of the corresponding simulations.
We surmise that this restriction is the reason why we did not observe improvements.
In the future, we aim at improving the semantic mapping between the foundation models and demonstrating that it can provide benefits for more realistic scenes.
Acknowledgements
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the
Federal State Upper Austria.
IARAI is supported by Here Technologies.
We thank the projects AIMOTION
(LIT-2018-6-YOU-212), AI-SNN (LIT-2018-6-YOU-214), DeepFlood (LIT-2019-8-YOU213), Medical Cognitive Computing Center (MC3), INCONTROL-RL (FFG-881064), PRIMAL (FFG-873979), S3AI (FFG-872172), DL for GranularFlow (FFG-871302), AIRI FG 9-N (FWF36284, FWF-36235), ELISE (H2020-ICT-2019-3 ID: 951847).
We thank Audi.
JKU Deep Learning
Center, TGW LOGISTICS GROUP GMBH, University SAL Labs initiative of Silicon Austria Labs (SAL), FILL Gesellschaft mbH, Anyline GmbH, Google, ZF Friedrichshafen AG, Robert Bosch
GmbH, UCB Biopharma SRL, Merck Healthcare KGaA, Verbund AG, Software Competence Center
Hagenberg GmbH, TÜV Austria, Frauscher Sensonic and the NVIDIA Corporation.
A Supplementary Material A.1 Environments
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 our methods.
Further, we neglect the Sidewalk environment since it is essentially the same task as Hallway with a different background.
Since the reward 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.
A.2 Mapping between Language Embedding Spaces
We consider all publicly available CLIP backbone variants and align their output spaces with the
TrXL used in HELM.
According to Algorithm 1 we determine the overlap in the CLIP and TrXL vocabularies.
In total, there are 5285 tokens that appear in both, V CLIP , and V LM .
Prior work has found that as little as ten word correspondences are sufficient to train an orthogonal mapping between monolingual embedding spaces of closely related languages (Zhang et al., 2016).
This assumes a certain degree of isomorphism between the embedding spaces.
Since we train a mapping between embedding spaces of the same language, we expect this assumption to hold in our setting as well.
The
CLIP output space and the TrXL input space differ greatly in their statistics, therefore we perform centering and scaling as preprocessing.
This way, during inference in the RL experiments we can simply re-center and re-scale by the statistics of LM space to obtain a suitable input for the TrXL.
Table 2 shows the average training and test accuracy across 5 folds for different linear projection methods.
We find that generally there is a tendency that the mapping improves for larger backbones.
The Procrustes variants consistently exhibit the highest average accuracy on the test set across 5 folds.
We do not show variance estimates, because those are negligibly small.
A.3 Mapping Images to the Language Space
We complement our findings on the ranking experiment in Section 4 with results for additional vision backbones of CLIP.
Table 1 shows the relative NDCG to ranking in CLIP space, while Table 3 shows the MRR for the different image encoders.
A.4 Hyperparameter Search
We adapt the hyperparameter search conducted in Paischer et al.
Particularly, we search for learning rate in { 5 e- 4 , 3 e- 4 , 1 e- 5 , 5 e- 5 } , entropy coefficient in { 0 . 05 , 0 . 01 , 0 . 005 , 0 . 001 } , rollout
[Table]
[Table]
0.2 Cosine similarity miner narrow cross
[Figure]
length in { 32 , 64 , 128 } for HELMv2, and HELMcs.
Since our analysis in Section 3 showed that for β = 1 , 10 the spatial compression collapses to the mean over pretrained token embeddings, we alter the grid for β of HELM to { 100 , 500 , 1000 , 5000 } .
To decrease wall-clock time of HELM variants, we vary the size of the memory register of TrXL such that it can fit the maximum episode length (Table 6).
We lower the number of interaction steps for the gridsearch if we observe convergence before the 500k interaction steps.
If no convergence is observed within the 500K interaction steps, we tune for the entire duration.
We apply the same scheme for tuning the LSTM baseline and tune the same hyperparameters as in Paischer et al.
[Table]
[Table] [Table]