Semantic Exploration from Language Abstractions and Pretrained Representations
Abstract
Effective exploration is a challenge in reinforcement learning (RL).
Novelty-based exploration methods can suffer in high-dimensional state spaces, such as continuous partially-observable 3D environments.
We address this challenge by deﬁning novelty using semantically meaningful state abstractions, which can be found in learned representations shaped by natural language.
In particular, we evaluate vision-language representations, pretrained on natural image captioning datasets.
We show that these pretrained representations drive meaningful, task-relevant exploration and improve performance on 3D simulated environments.
We also characterize why and how language provides useful abstractions for exploration by considering the impacts of using representations from a pretrained model, a language oracle, and several ablations.
We demonstrate the beneﬁts of our approach with on- and off-policy RL algorithms and in two very different task domains— one that stresses the identiﬁcation and manipulation of everyday objects, and one that requires navigational exploration in an expansive world.
Our results suggest that using language-shaped representations could improve exploration for various algorithms and agents in challenging environments.
Introduction
Exploration is one of the central challenges of reinforcement learning (RL).
A popular way to increase an agent’s tendency to explore is to augment trajectories with intrinsic rewards for reaching novel environment states.
However, the success of this approach depends critically on which states are considered novel, which can in turn depend on how environment states are represented.
The literature on novelty-driven exploration describes several approaches to deriving state representations [7].
One popular method employs random features and represents the state by embedding the visual observation with a ﬁxed, randomly initialized target network [Random Network Distillation; 6].
Another method uses learned visual features, taken from an inverse dynamics model [Never
Give Up; 3].
These approaches work well in classic 2D environments like Atari, but it is less clear whether they are as effective in high-dimensional, partially-observable settings such as 3D environments.
For instance, in 3D settings, different viewpoints of the same scene may map to distinct visual states/features, despite being semantically similar.
The difﬁculty of identifying a good mapping between visual state and feature space is exacerbated by the fact that useful state abstractions are highly task dependent.
For example, a task involving tool use requires object-affordance abstractions, whereas navigation does not.
Thus, acquiring state representations that support effective exploration is a chicken-and-egg problem—knowing whether two states should be considered similar requires the type of understanding that an agent can only acquire after effectively exploring its environment.
To overcome these challenges, we propose giving agents access to prior knowledge during training, in the form of abstractions derived from large vision-language models [e.g. 41] that are pretrained on image captioning data.
We use these pretrained models to derive a intrinsic reward that reﬂects meaningful novelty.
We hypothesize that representations acquired by vision-language pretraining drive effective, semantic exploration in 3D environments, because the representations are shaped by the unique abstract nature of natural language.
Several aspects of natural language suggest that it could be useful to direct novelty-based exploration.
First, language is inherently abstract: language links superﬁcially distinct, but causally-related situations by describing them similarly, and contrasts between causally-distinct states by describing them differently, thus outlining useful concepts [29, 28].
Second, humans use language to communicate important information efﬁciently, without overspecifying [20, 21].
Thus, human language omits distracting irrelevant information and focuses on important aspects of the world.
For example, it is often observed that an agent rewarded for seeking novel experience would be attracted forever to a TV with uncontrollable and unpredictable random static [7].
However, a human would likely caption this scene “a TV with no signal” regardless of the particular pattern; thus an agent exploring with language abstractions would quickly leave the TV behind.
Figure 1 shows another conceptual example of how language abstractions can accelerate exploration.
We ﬁrst perform motivating proof-of-concept experiments using a language oracle.
We show that language is a useful abstraction for exploration not only because it coarsens the state space, but also because it coarsens the state space in a way that reﬂects the semantics of the environment.
We then demonstrate that our results scale to environments without a language oracle using pretrained vision encoders, which are only supervised with language during pretraining.
This work strives to enhance the representations used in novelty-based exploration, rather than compare various exploration methods.
We consider two popular novelty-based exploration methods from the literature, Never Give Up (NGU;
Badia et al.
[3]) and Random Network Distillation (RND; Burda et al.
[7]), and compare them to their language-augmented variants, Lang-NGU/LSE-NGU and Lang-RND.
We evaluate performance and sample efﬁciency on object manipulation, search, and navigation tasks in two challenging 3D environments simulated in Unity: Playroom (a house containing toys and furniture) and City (a large-scale urban setting).
Our results show that language-based exploration with pretrained visionlanguage representations improves sample efﬁciency on Playroom tasks by 18-70%.
It also doubles [Figure] the visited areas in City, compared to baseline methods.
We show that language-based exploration is effective for both on-policy (IMPALA [17]) and off-policy (R2D2 [25]) agents.
Related Work
Exploration in RL Classical exploration strategies include (cid:15) -greedy action selection [51], statecounting [49, 4, 32, 30, 3], curiosity driven exploration [44], and intrinsic motivation methods [36].
Our work is part of this last class of methods, where the agent is given an intrinsic reward for visiting diverse states over time [35].
Intrinsic rewards can be derived from various measures: novelty [43, 55, 6, 56], prediction-error [38, 3], ensemble disagreement [11, 39, 48, 46, 18, 50], or information gain [23].
One family of methods gives intrinsic reward for following a curriculum of goals [8, 12, 40].
Others use novelty measures to identify interesting states from which they can perform additional learning [16, 54].
These methods encourage exploration in different ways, but they all rely on visual state representations that are learned jointly with the policy.
Although we focus on novelty-based intrinsic reward and demonstrate the beneﬁts of language in NGU and RND, our methodology is relatively agnostic to the exploration method.
We suggest that many other exploration methods could be improved by using language abstractions and pretrained embeddings to represent the state space.
Pretraining representations for RL Pretraining has been used in RL to improve the representations of the policy network.
Self-supervised representation learning techniques distill knowledge from external datasets to produce downstream features that are helpful in virtual environments [15, 53].
Some recent work shows beneﬁts from pretraining on more general, large-scale datasets.
Pretrained
CLIP features have been used in a number of recent robotics papers to speed up control and navigation tasks.
These features can condition the policy network [26], or can be fused throughout the visual encoder to integrate semantic information about the environment [37].
The goal of these works is to improve perception in the policy.
Pretrained language models can also provide useful initializations for training policies to imitate ofﬂine trajectories [42, 27].
These successes demonstrate that large pretrained models contain prior knowledge that can be useful for RL.
While the existing literature uses pretrained embeddings directly in the agent, we instead allow the policy network to learn from scratchm and only utilize pretrained embeddings to guide exploration during training (Figure S2).
We imagine that future work may beneﬁt from combining both approaches.
Language for exploration Some recent works have used language to guide agent learning, by either using language subgoals for exploration/planning or providing task-speciﬁc reward shaping [47, 33, 13, 19].
Schwartz et al.
[45] use a custom semantic parser for VizDoom and show that representing states with language, rather than vision, leads to faster learning by simplifying policy inputs.
Chaplot et al.
[10] tackle navigation in 3D by constructing a semantic map of the environment from pretrained SLAM modules, language-deﬁned object categories, and agent location.
This approach lends itself to navigation, but it is unclear how it would extend easily to more generic settings or other types of tasks, such as manipulation.
Work concurrent to ours by Mu et al.
[34] shows how language, in the form of hand-crafted BabyAI annotations, can help improve exploration in 2D environments.
These works demonstrate the value of language abstractions: the ability to ignore extraneous noise and highlight important environment features.
However, these prior methods rely on environment-speciﬁc semantic parsers or annotations, which may limit the settings to which they can be applied.
In contrast, by exploiting powerful pretrained vision-language models, our approach can be applied to any visually-naturalistic environment, including 3D settings, which have not been widely studied in prior exploration work.
We additionally do not require any language from the environment itself.
Our method could even potentially improve exploration for physical robots, but we leave that for future work.
Method
We consider a goal-conditioned Markov decision process deﬁned by a tuple ( S , A , G , P, R e , γ ) , where S is the state space, A is the action space, G is the goal space, P : S ×
A → S speciﬁes the environment dynamics, R e : S × G → R e is the extrinsic reward, and γ is the discount factor.
State s t is presented to the agent as a visual observation O V .
In some cases, in order to calculate intrinsic reward, we use a language oracle O : S
→ L that provides natural language descriptions of the state,
O L .
Note that O L is distinct from the language instruction g ∈ G , which is sampled from a goal distribution at the start of an episode—the agent never observes O L .
We later remove the need for a language oracle by using pretrained models.
training may incorporate pretrained language (pictured above) or image encoder, which may respectively require O L or O V .
The latter does not rely on language provided by the environment.
See Figure S2 for more details. We use goal-conditioned reinforcement learning to produce a policy π g ( · | O V ) that maximizes the expected reward E [ (cid:80)
Ht =
γ t ( r et + βr it )] , where H is the horizon, r et is the extrinsic reward, r it is the intrinsic reward, and β is a tuned hyperparameter.
The intrinsic reward is goal-agnostic and is computed with access to either O V or O L .
Note that neither O L nor pretrained embeddings are used by the policy, and thus we only use them during training to compute the intrinsic reward (Figure 3).
Our approach builds on two popular exploration algorithms: Never Give Up (NGU; Badia et al. [3]) and Random Network Distillation (RND; Burda et al.
These algorithms were chosen to demonstrate the value of language under two different exploration paradigms.
While both methods reward visiting novel states, they differ on several dimensions: the novelty horizon (episodic versus lifetime), how the history of past visited states is retained (non-parametric versus parametric), and how states are represented (learned controllable states versus random features).
3.1 Never Give Up (NGU)
To more clearly isolate our impact, we focus only on the episodic novelty component of the NGU agent [3].
State representations along the trajectory are written to a non-parametric episodic memory buffer.
The intrinsic reward reﬂects how novel the current state is relative to the states visited so far in the episode.
Novelty is a function of the L2 distances between the current state and the k -nearest neighbor representations stored in the memory buffer.
Intrinsic reward is higher for larger distances.
Full details can be found in the original paper; however, we make two key simpliﬁcations.
While
Badia et al.
[3] proposes learning a family of policy networks that are capable of different levels of exploration, we train one policy network that maximizes reward r = r e + βr i for a ﬁxed hyperparameter β .
We also ﬁx the long-term novelty modulator α to be 1, essentially removing it.
The published baseline method, which we refer to as Vis-NGU , uses a controllable state taken from an inverse dynamics model.
The inverse dynamics model is trained jointly with the policy, but the two networks do not share any parameters.
[Table] [Table] The intrinsic reward relies on directly comparing state representations from the buffer, so our approach focuses on modifying the embedding function to inﬂuence exploration (Table 1).
LangNGU uses a frozen pretrained language encoder to embed the oracle caption O L .
We compare language embeddings from BERT [14], CLIP [41], Small-ALM, and Med-ALM.
The ALMs (ALign
Models) are trained with a contrastive loss on the ALIGN dataset [24].
Small-ALM uses a 26M parameter ResNet-50 image encoder [22]; Med-ALM uses a 71M parameter NFNet [5].
The language backbones are based on BERT and are all in the range of 70-90M parameters.
We do not ﬁnetune on environment-speciﬁc data; this preserves the real world knowledge acquired during pretraining and demonstrates its beneﬁt without requiring any environment-speciﬁc captions.
LSE-NGU does not use the language oracle.
Instead, it uses a frozen pretrained image encoder to embed the visual observation O V .
We use the image encoder from CLIP or ALM, which are trained on captioning datasets to produce outputs that are close to the corresponding language embeddings.
The human-generated captions structure the visual embedding space to reﬂect features most pertinent to humans and human language [31], so the resulting representations can be thought of as Language Supervised Embeddings (LSE).
The primary beneﬁt of LSE-NGU is that it can be applied to environments without a language oracle or annotations.
CLIP and ALM are trained on real-world data, so they would work best on realistic 3D environments.
However, we imagine that in future work the pretraining process or dataset could be tailored to maximize transfer to a desired target environment.
3.2 Random Network Distillation (RND)
Our RND-inspired family of methods rewards lifetime novelty.
Generically, the intrinsic reward is derived from the prediction error between a trainable network and some target value generated by a frozen function (Table 2).
The trainable network is learned jointly with the policy network, although they do not share any parameters.
As the agent trains over the course of its lifetime, the prediction error for frequently-visited states decreases, and the associated intrinsic reward consequently diminishes.
Intuitively, the weights of the trainable network implicitly store the state visitation counts.
For clarity, we refer to the baseline published by Burda et al.
[7] as Vis-RND .
The trainable network f V :
O V → R k maps the visual state to random features.
The random features are produced by a ﬁxed, randomly initialized network ˆ f V .
Both f V and ˆ f V share the same architecture: a ResNet followed by a MLP.
The intrinsic reward is the mean squared error (cid:107) f V ( O V ) − ˆ f V
( O V ) (cid:107) 2 .
In network distillation ( ND ), the target function is not random, but is instead a pretrained text or image encoder from CLIP/ALM.
The trainable network f learns to reproduce the pretrained representations.
To manage inference time, f is a simpler network than the target (see Appendix A.2).
The intrinsic loss is the mean squared error between f and the large pretrained network.
Like the respective Lang-NGU and LSE-NGU counterparts, text-based ND requires a language oracle, but image-based ND does not.
In Section 5.1 we compare against two additional methods to motivate why language is a useful abstraction.
The ﬁrst, Lang-RND , is a variant in which the trainable network f L : O L → R k maps the oracle caption to random features.
The intrinsic reward is the mean squared error between the outputs of f L and ﬁxed ˆ f L with random initialization.
Both f L and ˆ f L networks are of the same architecture.
The second method, language distillation ( LD ), is loosely inspired by RND in that the novelty signal comes from a prediction error.
However, instead of learning to produce random features, the trainable network learns to caption the visual state, i.e. f C :
O V → O L .
The network architecture consists of a CNN encoder and LSTM decoder.
The intrinsic reward is the negative log-likelihood of the oracle caption under the trainable model f C .
In LD, the exploration dynamics not only depend on how frequently states are visited but also the alignment between language and the visual world.
We test whether this caption-denoted alignment is necessary for directing semantic exploration by comparing
LD to a variant with shufﬂed image-language alignment ( S-LD ) in Section 5.1.
Experimental Setup 4.1 Environments
Previous exploration work benchmarked algorithms on video games, such as 2D grid-world MiniHack and Montezuma’s Revenge, or 3D ﬁrst-person shooter Vizdoom.
In this paper, we focus on ﬁrst-person
Unity-based 3D environments that are meant to mimic familiar scenes from the real world (Figure 2).
Playroom Our ﬁrst domain, Playroom [1, 52], is a randomly-generated house containing everyday household items (e.g. bed, bathtub, tables, chairs, toys).
The agent’s action set consists of 46 discrete actions that involve locomotion primitives and object manipulation, such as holding and rotating.
We study two settings in Playroom.
In the ﬁrst setting, the agent is conﬁned to a single room with 3-5 objects and is given a lift or put instruction.
At the start of an episode, the set of objects are sampled from a larger set of everyday objects (i.e. a candle, cup, hairdryer).
Object colors and sizes are also randomized, adding superﬁcial variance to different semantic categories.
The instructions take the form: "Lift a <object>" or "Put a <object> on a {bed, tray}" .
With a lift goal, the episode ends with reward 1 or 0 whenever any object is lifted.
With a put goal, the episode ends with reward 1 when the condition is fulﬁlled.
This setting tests spatial rearrangement skills.
In the second setting, the agent is placed in a house with 3-5 different rooms, and is given a find instruction of the form "Find a {teddy bear, rubber duck}" .
Every episode, the house is randomly generated with the teddy and duck hidden amongst many objects, furniture, and decorations.
The target objects can appear in any room— either on the ﬂoor, on top of tables, or inside bookshelves.
The agent is randomly initialized and can travel throughout the house and freely rearrange objects.
The episode ends with reward 1 when the agent pauses in front of the desired object.
The find task requires navigation/search skills and tests the ability to ignore the numerous distractor objects.
City Our second domain, City, is an expansive, large-scale urban environment.
Each episode, a new map is generated; shops, parks, and buildings are randomly arranged in city blocks.
Daylight is simulated, such that the episode starts during the morning and ends at nighttime.
The agent is randomly initialized and is instructed to “explore the city.”
It is trained to maximize its intrinsic reward and can take the following actions: move_{forward,backward,left,right} , look_{left,right} , and move_forward_and_look_{left,right} .
We divide up the map into a 32 × 32 grid and track how many unique bins are visited in an episode.
Additionally, City does not provide explicit visual or verbal signage to disambiguate locations.
As such, systematic exploration is needed to maximize coverage.
In contrast to Playroom, City tests long horizon exploration.
A Playroom episode lasts only 600 timesteps, whereas a City episode lasts 11,250 and requires hundreds of timesteps to fully traverse the map even once.
The City covers a 270-by-270 meter square area, which models a 2-by-2 grid of real world blocks.
4.2 Captioning Engine
We equip the environment with a language oracle that generates language descriptions of the scene,
O L , based on the Unity state, s (Figure 2).
In Playroom, the caption describes if and how the agent interacts with objects and lists what is currently visible to it.
In City, O L generally describes the object that the agent is directly looking at, but the captions alone do not disambiguate the agent’s locations.
Since these captions are generated from a Unity state, these descriptions may not be as varied or rich as a human’s, but they can be generated accurately and reliably, and at scale.
4.3 Training Details
At test time, the agent receives image observation
O V and language-speciﬁed goal g .
The policy network never requires caption
O L to act.
During training, the exploration method calculates the intrinsic reward from O L or O V .
We show that language-based exploration is compatible with both policy gradient and Q-learning algorithms.
We use Impala [17] on Playroom and R2D2 on City [25].
Q-learning is more suitable for the City, because the action space is more restricted compared to the one needed for Playroom tasks.
For both environments, the agent architecture consists of an image ResNet encoder and a language
LSTM encoder that feed into a memory LSTM module.
The policy and value heads are MLPs that receive the memory state as input.
If the exploration method requires additional networks, such as the trainable network in RND or inverse dynamics model in NGU, they do not share any parameters with the policy or value networks.
Figure S2 is a visualization of an Impala agent that uses languageaugmented exploration.
Hyperparameters and additional details are found in Appendix A.
Results 5.1 Motivation: Language is a Meaningful Abstraction
[Figure]
We share the intuition with other work [e.g. 34] that language can improve exploration.
We design a set of experiments to show how and why this may be the case.
Our analysis follows the desiderata outlined by Burda et al. [6]—prediction-error exploration ought to use a feature space that ﬁlters irrelevant information (compact) and contains necessary information (sufﬁcient).
Burda et al.
[6] speciﬁcally studies RND and notes that the random feature space, the outputs of the random network, may not fully satisfy either condition.
As such, we use the language variants of RND to frame this discussion.
We hypothesize that language abstractions are useful, because they (1) create a coarser state space and (2) divide the state space in a way that meaningfully aligns with the world (i.e. using semantics).
First, if language provides a coarser state space, then the random feature space becomes more compact, leading to better exploration.
We compare Lang-RND to Vis-RND to test this claim.
Lang-RND learns the lift task 33% faster and solves the put task as Vis-RND starts to learn (Figure 4a).
The states are grouped together based on similarities in the visual random feature space and assigned a label.
Exploring in this shufﬂed space is less effective than exploring with the semanticallymeaningful abstractions shown in Figure 1.
Second, we ask whether semantics – that is how language divides up the state space – is critical for effective exploration.
We use LD to test this hypothesis, precisely because the exploration in LD is motivated by modeling the semantic relationship between language and vision.
We compare LD to a shufﬂed variant SLD, where we replace the particular semantic state abstraction that language offers with a statistically-matched randomized abstraction (Figure 5).
S-LD is similar to LD; the intrinsic reward is the prediction error of the captioning network.
However, instead of targeting the language oracle output, the S-LD trainable network produces a different target caption (cid:102)
O L that may not match the image.
(cid:102)
O L is produced by a ﬁxed, random mapping ˆ f S :
O V → (cid:102) O L .
ˆ f S is constrained such that the marginal distributions
P ( O L ) ≈ P ( (cid:102) O L ) are matched under trajectories produced by policy π LD .
See Appendix A. for full details on the construction of S-LD.
Thus, whereas the LD captions parcel up state space in a way that reﬂects the abstractions that language offers, the randomized mapping ˆ f S parcels up state space in a way that abstracts over random features of the visual space (Figure 5).
We control for the compactness and coarseness of the resulting representation by maintaining the same marginal distribution of captions.
If semantics is crucial for exploration, then we expect to see LD outperform S-LD.
This indeed holds experimentally (Figure 4b).
We can also view these results under the Burda et al.
[6] framework.
The
S-LD abstractions group together visually similar, but semantically distinct states.
A single sampled caption likely fails to capture the group in a manner that is representative of all the encompassing states.
In other words, ˆ f S produces a compact feature space that may not be sufﬁcient.
This may explain why S-LD learns faster than Vis-RND on the simpler lift task but fails on the more complex put and find tasks.
The S-LD experiments imply that language abstractions are helpful for exploration because they expose not only a more compact, but also a more semantically meaningful state space.
5.2 Pretrained Vision-Language Representations Improve Exploration
Having shown how language can be helpful for exploration, we now incorporate pretrained visionlanguage representations into NGU and RND to improve exploration.
Such representations (e.g. from the image encoder in CLIP/ALM) offer the beneﬁts of explicit language abstractions, without the need to rely on a language oracle.
We also compare language-shaped representations to pretrained
ImageNet embeddings to isolate the effect of language.
To keep the number of experiments tractable, we only perform a full comparison on the Playroom tasks.
City We ﬁrst compare how representations affect performance in a pure exploration setting.
With no extrinsic reward, the agent is motivated solely by the NGU intrinsic reward to explore the City.
We report how many unique areas the agent visits in an episode in Figure 6.
While optimizing coverage only requires knowledge of an agent’s global location rather than generic scene understanding, vision-language representations are still useful simply because meaningful exploration is inherently semantic.
Lang-NGU, which uses text embeddings of O L , visits an area up to 3 times larger.
LSENGU achieves 2 times the coverage even without querying a language oracle (Appendix Figure S5). N G U B a s e li n e L S E - N G U ( I m a g e N e t )
L a n g - N G U ( C L I P ) L a n g - N G U A L M ) L S E - N G U
( C L I P )
L S E - N G U
( A L M ) 0.
C o v e r a g e
Playroom We next show that pretrained vision-language representations signiﬁcantly speed up learning across all Playroom tasks (Figure 7).
The LSE-NGU and Lang-NGU agents improve sample efﬁciency by 50-70% on the lift and put tasks and 18-38% on the find task, depending on the pretraining model used
The ND agents are signiﬁcantly faster than VisRND, learning 41% faster on the find task.
We also measure agent-object interactions.
Nearly all LSE-NGU and Lang-NGU agents learn to foveate on and hold objects within 40k learning updates, whereas Vis-NGU agent takes at least 60k updates to do so with the same frequency (Appendix Figure S7).
Although LSENGU and image-based ND agents do not access a language oracle, they are similarly effective as their annotation-dependent counterparts in the Playroom tasks (Appendix Figure S6), suggesting that our method could be robust to the availability of a language oracle.
To demonstrate the value of rich language, we compare LSE-NGU agents to a control agent that instead uses pretrained ImageNet embeddings from a 70M NFNet [5].
ImageNet embeddings optimize for single-object classiﬁcation, so they confer some beneﬁt to the most objectfocused tasks, lift and put .
However, ImageNet embeddings hurt exploration in the find task, where agents encounters more complex scenes (Figure 7b).
By contrast, the language-shaped representations are well-suited for not only describing simple objects, but also have capacity for multi-object, complex scenes.
Of course, current CLIPstyle models can be further improved in their ability to understand multi-object scenes, which may explain why the beneﬁts are less pronounced for the find task.
However, as the performance of pretrained vision-language models improve, we expect to see those beneﬁts transfer to this method and drive even better exploration.
Discussion
We have shown that language abstractions and pretrained vision-language representations improve the sample efﬁciency of existing exploration methods.
This beneﬁt is seen across on-policy and off-policy algorithms (Impala and R2D2), different exploration methods (RND and NGU), different 3D domains (Playroom and City), and various task speciﬁcations (lifting/putting, searching, and intrinsically motivated navigation).
Furthermore, we carefully designed control experiments to understand how language contributes to better exploration.
Our results are consistent with cognitive perspectives on human language—language is powerful because it groups together situations according to semantic similarity.
In terms of the desiderata that Burda et al.
[6] present, language is both compact and sufﬁcient.
Finally, we note that using pretrained vision-language representations to embed image observations enables more effective exploration even if language is not available during agent training.
This is vital for scaling to environments that do not have a language oracle or annotations.
Limitations and future directions
We highlight several avenues for extending our work.
First, additional research could provide a more comprehensive understanding of how language abstractions affect representations.
This could include comparing different types of captions offering varying levels of detail, or task-dependent descriptions.
These captions could be dynamically generated at scale by prompting a large multimodal model [2].
Second, it would be useful to investigate how to improve pretrained vision-language representations for exploration by ﬁnetuning on relevant datasets.
The semantics of a dataset could even be tailored to task-speciﬁc abstractions to increase the quality of the learnt representations. Such approaches would potentially allow applying our method to virtual [Figure] [Figure]
[Figure] [Figure] [Figure] environments that are farther from the pretraining distribution, such as Atari.
In contrast, compared to our experiments, we believe that the current pretrained representations would deliver even more beneﬁt for entirely photorealistic, visually rich environments, such as Matterport3D [9]. Finally, we note that a limitation of this approach is that current pretrained vision-language models may be less effective on multi-object scenes.
Future pretraining innovations or larger models would presumably produce more robust representations and thus lead to even more effective exploration. Acknowledgments and Disclosure of Funding
We would like to thank Iain Barr for ALM models and Nathaniel Wong and Arthur Brussee for the
Playroom environment.
For the City environment, we would like to thank Nick Young, Tom Hudson,
Alex Platonov, Bethanie Brownﬁeld, Sarah Chakera, Dario de Cesare, Marjorie Limont, Benigno
Uria, Borja Ibarz and Charles Blundell.
Moreover, for the City, we would like to extend our special thanks to Jayd Matthias, Jason Sanmiya, Marcus Wainwright, Max Cant and the rest of the Worlds
Team.
Finally, we thank Hamza Merzic, Andre Saraiva, and Tim Scholtes for their helpful support and advice.
A Training Details
We use a distributed RL training setup with 256 parallel actors.
For Impala agents, the learner samples from a replay buffer that acts like a queue.
For R2D2 agents, the learner samples from a replay buffer using prioritized replay.
Training took 8-36 hours per experiment on a 2 × 2 TPUv2.
All agents share the same policy network architecture and hyperparameters (Table S1).
We use an
Adam optimizer for all experiments.
The hyperparameters used to train Impala [17] and R2D2 [25] are mostly taken from the original implementations.
For Impala, we set the unroll length to 128, the policy network cost to 0.85, and state-value function cost to 1.0.
We also use two heads to estimate the state-value functions for extrinsic and intrinsic reward separately.
For R2D2, we set the unroll length to 100, burn-in period to 20, priority exponent to 0.9, Q-network target update period to 400, and replay buffer size to 10,000.
[Table]
A.1 Training Details for NGU Variants
We use a simpliﬁed version of the full NGU agent as to focus on the episodic novelty component.
One major difference is that we only learn one value function, associated with a single intrinsic reward scale β and discount factor γ .
The discount factor γ is 0.99 and we sweep for β (Table S2).
Another major difference is that our lifetime novelty factor α is always set to 1.
The NGU memory buffer is set to 12,000, so that it can always has the capacity to store the the entire episode.
The buffer is reset at the start of the episode.
The intrinsic reward is calculated from a kernel operation over state representations stored in the memory buffer.
We use the same kernel function and associated hyperparameters (e.g. number of neighbors, cluster distance, maximum similarity) found in Badia et al.
For Vis-NGU, the 32-dimension controllable states come from a learned inverse dynamics model.
The inverse dynamics model is trained with an Adam optimizer (learning rate 5e-4, β 1 = 0 . 0 , β 2 = 0 .
95 , (cid:15) is 6e-6).
For the variants, we use frozen pretrained representations from BERT, ALM, or CLIP.
The
ALM pretrained embeddings are size 768 and the CLIP embeddings are 512.
Med-ALM comprises a 71M parameter NFNet image encoder and 77M BERT text encoder.
Small-ALM comprises a 25M Resnet image encoder and 44M BERT text encoder.
Using Small-ALM can help mitigate the increased inference time.
To manage training time, we use Small-ALM in the City environment, where the episode is magnitudes longer than Playroom.
For the LSE-NGU ImageNet control, we use the representations from a frozen 71M parameter NFNet pretrained on ImageNet (F0 from Brock et al.
This roughly matches the size of the CLIP and Med-ALM image encoders.
We notice that it is crucial for the pretrained representations to only be added to the buffer everytimesteps.
Meanwhile, Vis-NGU adds controllable states every timestep, which is as or more effective than every 8.
This may be due to some interactions between the kernel function and the smoothness of the learned controllable states.
We also ﬁnd that normalizing the intrinsic reward, like in RND, is helpful in some settings.
We use normalization for Lang-NGU and LSE-NGU on the Playroom tasks.
[Table]
A.2 Training Details for RND-Inspired Agents
For the RND-inspired exploration agents
[Table]
the hyperparameters for training the trainable network are taken from the original implementation [7].
We perform a random hyperparameter search over a range of values for the intrinsic reward scale, V-Trace entropy cost
[Table]
and the learning rate for the trainable network.
The values can be found in Table S3.
We also normalize all intrinsic reward with the rolling mean and standard deviation.
[Table]
The various RND-inspired agents differ in the input and output space of the trainable network and target functions (Figure S1).
Some trainable networks and target networks involve a convolutional network
[Table]
which consists of (64, 128, 128, 128) channels with (7, 5, 3, 3) kernel and (4, 2, 1, 1) stride.
For ND, we balance the model capacity of the trainable network with the memory required for learning larger networks.
We notice that using larger networks as the target function can make distillation harder and therefore, requires more careful parameter tuning.
Our experiments use the 26M parameter Small-ALM vision encoder for generating target outputs
[Table]
which minimizes the discrepancy in complexity between the trainable network and target function
This also managed the memory requirements during training
[Table]
A.3 Engineering Considerations
Integrating large pretrained models into RL frameworks is nontrivial.
High quality pretrained models are orders of magnitudes larger than policy networks, introducing challenges with inference speed.
[Table]
[Figure]
Slow inference not only increases iteration and training times but also may push learning further off-policy in distributed RL setups [e.g. 17, 25].
We mitigate this issue in Lang-NGU and LSENGU by only performing inference computations when it is necessary.
We compute the pretrained representations only when they are required for adding to the NGU buffer (i.e. every 8 timesteps).
A.4 S-LD Construction
As explained in Section 5.1, we compare the efﬁcacy of exploration induced by the language distillation (LD) method with that induced by a shufﬂed variant, S-LD.
LD employs an exploration bonus based on the prediction error of a captioning network f C :
O V → O L , where the target value is the oracle caption.
Our goal with S-LD is to determine whether the efﬁcacy of the exploration is due to the speciﬁc abstractions induced by f C , or whether it is due to some low level statistical property (e.g. the discrete nature of O L , or the particular marginal output distribution).
We sample a ﬁxed, random mapping ˆ f S :
O V → O L while trying to match the low-level statistics of f C .
To do this, we construct a (ﬁxed, random) smooth partitioning of O V into discrete regions, which in turn are each assigned to a unique caption from the output space of f C .
The regions are sized to maintain the same marginal distribution of captions, P π LD ( L )
≈ P π S − LD
( L ) .
More precisely, our procedure for constructing ˆ f S is as follows.
We build ˆ f S as the composition of three functions, ˆ f S =
g 3 ◦ g 2 ◦ g 1 :
• g 1 : O V → R is a ﬁxed, random function, implemented using a neural network as in VisRND.
• g 2 : R → Cat ( K ) is a one-hot operation, segmenting g 1 ( O V ) into K non-overlapping, contiguous intervals.
• g 3 : Cat ( K )
→ O L is a ﬁxed mapping from the discrete categories to set of captions produced by the language oracle, f C .
To ensure that the marginal distribution of captions seen in f C was preserved in f S , we ﬁrst measure the empirical distribution of oracle captions encountered by π LD , a policy trained to maximize the LD intrinsic reward.
We denote the observed marginal probability of observing caption l i as
P π LD ( O L = l i ) =
q i (where the ordering of captions l
i ∈
O L is ﬁxed and random).
We then measure the empirical distribution of g 1 under the same action of the same policy, P π LD ( g 1 ( O V )) , and deﬁne the boundaries of the intervals of g 2 by the quantiles of this empirical distribution so as to match the CDF of q , i.e. so that P π LD ( g 2 ◦ g 1 ( O V )
= i ) =
q i .
Finally, we deﬁne g 3 to map from the i th category to the corresponding caption l i , so that P π LD ( g 3 ◦ g 2 ◦ g 1 ( O V )
= l i ) =
q i .
B Additional ablation: Pretrained controllable states
The state representations used by Vis-NGU are trained jointly with the policy, whereas the pretrained representations used by Lang-NGU and LSE-NGU are frozen during training.
To isolate the effect of the knowledge encoded by the representations, we perform an additional experiment where we pretrain the controllable states and freeze them during training.
We use the weights of the inverse dynamics model from a previously trained Vis-NGU agent.
Figure S3 shows that pretrained VisNGU learns at the same rate as Vis-NGU (if not slower).
Thus, the increased performance in LangNGU and LSE-NGU agents is due to the way the vision-language embeddings are pretrained on captioning datasets.
This furthermore suggests that the converged controllable states do not fully capture the knowledge needed for efﬁcient exploration and in fact may even hurt exploration at the start of training by focusing on the wrong parts of the visual observation.
[Figure]
[Table]
(cid:111)(cid:186)(cid:214)(cid:303)(cid:121)(cid:200)(cid:143)(cid:303)(cid:174)(cid:186)(cid:186)(cid:172)(cid:161)(cid:181)(cid:156)(cid:303)(cid:121)(cid:210)(cid:303)(cid:121)(cid:303)(cid:133)(cid:174)(cid:214)(cid:143)(cid:303)(cid:156)(cid:200)(cid:143)(cid:232)(cid:303)(cid:133)(cid:200)(cid:161)(cid:134)(cid:172)(cid:303)(cid:186)(cid:155)(cid:242)(cid:134)(cid:143)(cid:204)(cid:303)(cid:161)(cid:181)(cid:303)(cid:121)(cid:303)(cid:186)(cid:155)(cid:242)(cid:134)(cid:143)(cid:303)(cid:133)(cid:214)(cid:161)(cid:174)(cid:139)(cid:161)(cid:181)(cid:156)(cid:265) (cid:111)(cid:186)(cid:214)(cid:303)(cid:134)(cid:121)(cid:181)(cid:303)(cid:204)(cid:143)(cid:143)(cid:303)(cid:121)(cid:181)(cid:303)(cid:121)(cid:200)(cid:180)(cid:134)(cid:160)(cid:121)(cid:161)(cid:200)(cid:266)(cid:303)(cid:121)(cid:181)(cid:303)(cid:186)(cid:210)(cid:210)(cid:186)(cid:180)(cid:121)(cid:181)(cid:266)(cid:303)(cid:121)(cid:303)(cid:204)(cid:160)(cid:143)(cid:174)(cid:155)(cid:266)(cid:303)(cid:121)(cid:181)(cid:139)(cid:303)(cid:121)(cid:303)(cid:210)(cid:143)(cid:139)(cid:139)(cid:232)(cid:265)
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Figure]
[Figure]
[Figure]
N u m b e r o f O
b j e c t s
[Figure]
N u m b e r o f O
b j e c t s
[Figure]
s
[Figure]
r o f O
b
j e c t s
[Figure]
object interaction.