Accelerating exploration and representation learning with offline pre-training

Bogdan Mazoure 1 Jacob Bruce 1 Doina Precup 1 Rob Fergus 1 Ankit Anand 1

Abstract
Sequential decision-making agents struggle
with long horizon tasks, since solving themrequires multi-step reasoning. Most reinforcement
learning (RL) algorithms address this challenge byimproved credit assignment, introducing memory
capability, altering the agent’s intrinsic motivation(i.e. exploration) or its worldview (i.e. knowledge
representation). Many of these components
could be learned from offline data. In this work,
we follow the hypothesis that exploration andrepresentation learning can be improved byseparately learning two different models from asingle offline dataset. We show that learning a state
representation using noise-contrastive estimation
and a model of auxiliary reward separately froma single collection of human demonstrations cansignificantly improve the sample efficiency on thechallenging NetHack benchmark. We also ablate
various components of our experimental setting
and highlight crucial insights.

1. Introduction
Sequential decision-making tasks with long horizon form an
important class of problems, but can be notoriously difficult
to solve (Bellemare et al., 2013; Vinyals et al., 2019; Küttler
et al., 2020). For example, NetHack is a popular terminalbased rogue-like video game with a multitude of sub-tasksneeding to be completed in order to win. So far, only around
1% of human players can fully solve the game (Hambro
et al., 2022). How does an agent prioritize conducting an
invocation over finding a dungeon key, or fully explore amaze? Undirected exploration is hopeless in this scenario,
since the agent has to try out a number of actions exponential
in trajectory length, quickly becoming intractable even forsimple sub-tasks. However, as recent works(Anonymous,
2023) show, even more powerful exploration techniques
1
Google DeepMind. Correspondence to: Bogdan Mazoure (nowat Apple) <bogdan.mazoure@apple.com>.

Interactive Learning with Implicit Human Feedback Workshop at
ICML 2023.

such as Burda et al. (RND, 2018) do not make significant
progress on harder tasks within the NetHack Learning
Environment, or NLE (Küttler et al., 2020).
Often, this problem is addressed by determining the action expansion order through either modifying the agent’s intrinsic
motivation (i.e., exploration, Dudik et al., 2011; Bellemare
et al., 2016; Burda et al., 2018; Ecoffet et al., 2019; Guo et al.,2022) or structuring its internal world representation (i.e.,
representation learning, Jaderberg et al., 2016; Anand et al.,2019; Srinivas et al., 2020; Mazoure et al., 2020; 2021;
Eysenbach et al., 2022). Both exploration and representation
learning approaches to long horizon RL problems have seen
their respective successes in complex domains e.g. Atari’sMontezuma’s Revenge (Ecoffet et al., 2019) and maze-likeenvironments (Raileanu & Rocktäschel, 2020). Some
attempts were made to combine both exploration and representation learning (Misra et al., 2020; Yarats et al., 2021), butthey are limited to specific families of MDP (e.g. BlockMDP)
which might not hold in complex real-world scenarios.
How can one leverage the representation learning paradigm
in order to achieve more sample-efficient exploration in
long-horizon RL problems? While exploration cannot
be performed without online interactions, it is possible toamortize the sample complexity of representation learningby pre-training the agent on offline data. This approach ofoffline pre-training followed by online fine-tuning has leadto important advances in complex, multi-modal control tasks
such as Minecraft (Fan et al., 2022; Baker et al., 2022). In
these approaches, human curators were used as a labelingmechanism on incomplete trajectory data, which was thenused in the representation learning phase. Does executingclassical exploration strategies on top of state representations
with desirable properties (i.e., capturing the progress ofhuman demonstrations towards solving the downstreamtask) make them more sample-efficient?
In this work, we posit the hypothesis that combining representation learning with exploration improves performanceon long-horizon tasks. We test our hypothesis in the NetHack
Learning Environment (Küttler et al., 2020), a challenging
domain with high-dimensional state and action spaces,
strong notion of forward progress and organic definitionof multiple subgoal tasks. While algorithms based on a
combination of RL and imitation learning can solve easierAccelerating exploration and representation learning with offline pre-training

subtasks such as the NeurIPS’21 challenge, completing thegame is only possible for symbolic agents or human experts.
Therefore, NLE provides an excellent opportunity bridging
the performance gap between RL agents and humans, unlike
other common benchmarks such as the Arcade Learning
Environment (Bellemare et al., 2013).
Previous work, Explore-Like-Exports (ELE, Anonymous,
2023) showed that many of the sparse reward tasks in
NetHack can be solved by learning a simple scalar function
which predicts the expert progress or temporal distance between two observations in any trajectory from the expert data.
While sparser tasks are solved by introducing this additional
reward, the performance on dense reward tasks like NeurIPS’
21 challenge and scout suffer in comparison to baseline.
Secondly, we also hypothesize that ELE discards useful
information contained in expert trajectories by compressing
the dataset into a single scalar-valued function. We addressthese issues by using the same expert data to learn representations by contrastive pre-training and in conjunction with
ELE, develop an agent that not only has significantly bettersample efficiency but also improved performance than ELE
and standard imitation learning baselines on a spectrum oftasks ranging from sparse to dense in a challenging domain
like NetHack. This work illustrates how one can use thesame dataset to learn representations and auxiliary reward,
thereby achieving better sample efficiency and performance.
Specifically, we show that a simple offline pre-training
scheme based on contrastive learning (Eysenbach et al., 2022;
Mazoure et al., 2022) can be used in conjunction with ELE i.elearning a progress reward (Anonymous, 2023) on the same
expert data. This not only improves the sample efficiency andbase performance of Muesli, a strong RL baseline (Hessel
et al., 2021a) but also using representation learning or
progress reward alone on NetHack in a wide variety of tasks.

2. Related works
2.1. Auxiliary tasks in RL
While one of the main goals in RL problems is to find a policy
which maximizes expected reward, it can often be challenging due to a multitude of factors, e.g. sparse reward signal,
untractably large policy space, long task horizon, etc. Since
the problem is extremely hard in its current formulation,
it is possible to augment it with external learning signals,
which are notably specified via auxiliary downstream tasks.
Auxiliary learning objectives have been widely studied in theliterature, in both online (Jaderberg et al., 2016; Stooke et al.,2021) and offline settings (Schwarzer et al., 2021; Yang &Nachum, 2021). They can be used to equip RL agents withdesirable inductive biases, e.g. disentanglement (Higgins
et al., 2017), alignment and uniformity (Wang & Isola, 2020)or predictivity of future observations (Jaderberg et al., 2016;

Mazoure et al., 2020).
World models provide one natural pre-training objective for
RL agents, allowing it to capture crucial parameters of the environment such as transition dynamics, reward function andinitial state distribution. Single-step world models such as
DreamerV3 (Hafner et al., 2023) and Director (Hafner et al.,2022) equip RL agents with single-step transition and reward
models that can then be used for planning. However, training
such models from offline data is non-trivial and costly; using
them in online settings is computationally inefficient as it requires unrolling the sequence of latent states and actions in an
autoregressive manner. On the other hand, infinite-horizon
models such as γ-models (Janner et al., 2020) or contrastive
value functions (Eysenbach et al., 2022; Mazoure et al., 2022)
are harder to learn, but directly capture the probability of observing a future state when rolling out from the current state.
2.2. Exploration
Some of the inductive biases for challenging tasks can belearned from offline demonstrations, e.g. human interactions
with the environment (Reid et al., 2022; Fan et al., 2022;
Baker et al., 2022). In hard tasks with sparse rewards andlong horizons, agents need to rely on other forms of supervision, i.e. intrinsic motivation. Intrinsic motivation for guided
exploration has been an active area of research in the past
years, encompassing count-based exploration (Bellemare
et al., 2016; Tang et al., 2017), knowledge gathering (Kimet al., 2018; Zhang et al., 2021) and curiosity (Burdaet al., 2018; Raileanu & Rocktäschel, 2020). However,curiosity-based exploration from tabula rasa is still a hard
problem in some tasks (e.g. NetHack), and hence warrants
the use of learned auxiliary rewards from data.
2.3. Learning from demonstrations
In domains where RL agents have not yet achieved humanlevel performance, learning can be accelerated by training
on demonstrations of experts (symbolic agents, humans, etc).
Classical imitation learning methods like Behavior Cloning(Pomerleau, 1988), is one of the most effective and popularly used methods in presence of large quantities of data incomplex domains like Minecraft(Baker et al., 2022), computer control(Humphreys et al., 2022) etc. Other approaches
like GAIL (Ho & Ermon, 2016) learns a discriminator to
distinguish expert trajectories from agent trajectories which
could be modeled as a reward. These methods have beenfurther extended to work on expert trajectories without actions as BCO (Torabi et al., 2018a) and GAIfO (Torabi et al.,2018b). Another generative approach FORM (Jaegle et al.,2021) augments the environment reward by an additional
reward by learning a forward generative model of transition
dynamics from offline data and rewarding transitions under
the learned model. In scenarios of unlabeled(that containAccelerating exploration and representation learning with offline pre-training

Figure 1. Experimental setting studied in this work. An offline pre-training phase of the agent’s representations f as well as the progressmodel g (top) is followed by an online fine-tuning phase, during which the agent uses both f and g to collect information about theenvironment.

no actions) datasets like NetHack , experts can be used toannotate existing datasets without action information, e.g.
add action information based on external sources such as in
MineDojo or VPT (Fan et al., 2022; Baker et al., 2022). However, such labeling schemes can involve collecting data fromhuman experts or training complex RL agents from scratch,both of which are prohibitively expensive in many scenarios.
Alternatively, demonstrations can be used to guide RL agentsthrough intrinsic motivation using learned heuristic functions.
For example, ELE (Anonymous, 2023) learns a heuristic
function quantifying temporal progress in expert trajectories.
It outperforms prior state-of-the-art on 7 NetHack tasks withsparse rewards, but still does not solve the game itself. We
hypothesize that the main drawback of ELE is that it reduces
the pre-training dataset to a single scalar-valued function, anddoes not extract the most information out of the data. Specifically, the degree to which a dataset heuristic is beneficial fora given online task depends on its alignment with the optimalvalue function in that MDP (Cheng et al., 2021). In this work,
we focus on using offline data which does not contain any actions and hence, limit ourselves to compare with ELE, BCO,
GAIfO and FORM as standard imitation learning baselines.

collected by policy π over an episode of length H:
V π (st ) = EPπt:H [∆(X ) denotes the entire set of distributions over the space X .

(1)where Pπt:t+K denotes the joint distribution of{st+k , at+k }K
k=1 obtained by deploying π in the environment M from timestep t to timestep t + K. The
state-action value function is defined analogously as
Qπ (st ,at ) = EPπt:H [H−t
X

γ k r(st+k ,at+k )|st ,at ],

(2)

k=0

such that Qπ (st ,at ) = r(st ,at )+γET (st ,at ) [V π (st+1 )].
The reinforcement learning problem consists in finding a
Markovian policy π ∗ that maximizes the state value function
over the set of initial states:π ∗ = maxEP[S0 ] [V π (s0 )],
π∈Π

(3)

for s0 ∼ P[S0 ] and set of policies Π. Alternatively, thevalue function can also be re-written as the expectation of
the reward over the geometric mixture of k-step forward
transition probabilities:
V π (st ) =3.1. Reinforcement learning1

γ k r(st+k ,at+k )|st ],k=0

3. Preliminaries
The classical reinforcement learning setting assumes that theenvironment follows a Markov decision process M definedby the tuple M = ⟨S,S0 ,A,T ,r,γ⟩, where S is the state space,
P[S0 ],S0 ∈ S is the distribution of starting states, A is theaction space, T = P[·|st ,at ] : S ×A → ∆(S) is the transition
kernel1 , r : S × A → [rmin ,rmax ] is the reward function andγ ∈ [0,1) is a discount factor. The environment is initialized
in s0 ∼ P[S0 ]. At every timestep t = 1, 2, 3, .., the policy
π : S → ∆(A), samples an action at ∼ π(·|st ). The environment then transitions into the next state st+1 ∼ T (·|st ,at )and emits a reward rt = r(st ,at ). The state value function
is defined as the cumulative per-timestep discounted rewards

H−t
X

1
[r(s,a)],E
1−γ (s,a∼ρπ (st )×π(s))

(4)

where
ρπ (s|st ) = (1−γ)

HXγ ∆t−1 P[St+∆t = s|st ;π]

∆t=1

(5)

= E∆t∼GeoH[P[St+∆t |st ,∆t;π]],t (1−γ)and GeoHt (1−γ) denotes a truncated geometric distribution
with probability mass re-distributed over the interval [t,H].
This decomposition of the value function is useful in
scenarios in which the environment’s rewards are delayed,leaving the learner only with access to states. It has beenused in previous works based on the successor representation (Dayan, 1993; Barreto et al., 2016) and, more recently,
explicit (Janner et al., 2020) and implicit (Eysenbach et al.,2022; Mazoure et al., 2022) infinite-horizon models.

Accelerating exploration and representation learning with offline pre-training

3.2. Explore-Like-Experts
Exploration in long-horizon problems with large action
spaces and sparse rewards is hard: an uninformed agent
would have to try |A|H actions for a horizon length H, whichis infeasible in NetHack, where |A| = 121 and H can be close
to 106 (Küttler et al., 2020). Augmenting the uninformative
extrinsic reward with an intrinsic signal which drives theagent to visit rare state-action pairs can directly translate
into higher overall returns, as the learner uncovers more of
the extrinsic reward structure. More formally, it is achievedby constructing an auxiliary MDP M ′ , where the reward
function at timestep t is a combination of the extrinsic reward
from M as well as some heuristic function h : S ×A → R:r′ (st ,at ) := r(st ,at )+λh(s0:t ,a0:t )(6)

When h(s0:t ,a0:t ) = γET (st ,at ) [v(st+1 )] for some functionv : S → R, then solving Equation (3) in M ′ is equivalent
to finding π in M if v approximates V ∗ (original valuefunction), but with a lower discount factor γ ′ = γ(1 − λ)when λ < 1 (Cheng et al., 2021)2 . While setting v = V ∗leads to maximal sample efficiency, V ∗ is not accessible inpractice, and h has to be selected based on the structure of M .
In particular, good heuristics can be constructed from datausing existing pessimistic offline RL algorithms, e.g. improvable heuristics lead to small estimation bias of V ∗ , since
they are smaller than the maximum of the Bellman backup.
Intuitively, given two states, how to prioritize one over theother during the exploration process? If we had a systematic
way to evaluate which state is closer to the goal under the
optimal policy, then we could force the agent to expandthat state during the exploration phase. Maximizing the
progress in the task can be captured through a monotonicallyincreasing function of states learned from optimal data(where, by definition, progress is maximal). Specifically, theExplore Like Experts algorithm (ELE) (Anonymous, 2023),
first trains a function g : S → R by solvingg ∗ = minED [ℓELE (g,st ,st+∆t )]
g∈F

(7)

where
ℓELE (g,st ,st+∆t ) = {g(st ,st+∆t )−sgn(∆t)log(1+|∆t|)}2 ,(8)
D is a set of expert human demonstrations and∆t ∼ LogUniform(0, 104 ). Specifically, Equation (8)does mean-squared error regression in the signed log-space
to predict the time offset ∆t from states st and st+∆t .
In the second step, ELE uses the pre-trained progress modelin place of the h heuristic in Equation (6)rELE (st ,at ) := r(st ,at )+λg(st−∆t ,st ),(9)2
If v is an improvable heuristic aligned with the optimal valuefunction, then the discount factor in M ′ is lowered.

an approximation of the local progress from st−∆t to st .
While ∆t was sampled by LogUnform distribution while
training g, it was kept fixed during the online phase in ELE.
In other words, auxiliary reward always computed progresswith respect to a state ∆t steps behind from current state.
3.3. Contrastive representation learning
The conditional probability distribution of st+∆t givenst can be efficiently estimated using an implicit modelf : S × S → R trained via contrastive learning (Oord et al.,2018) on offline demonstrations D by solving:−f ∗ = minED [ℓContrastive (f,st ,s+t ,st )]
f ∈F

(10)

where
−
ℓContrastive (f,st ,s+t ,st ) = −log+

ef (st ,st ).
P′
ef (st ,st )

(11)

+
−
s′t ∈st ∪st

To approximate the occupancy measure defined in Equation (5), positive samples are sampled from s+t ∈{st+∆t ; ∆t ∼ GeoH
(1
−
γ)}
for
timestept.
Specifically,t
they are constructed by first sampling the interval ∆t fromGeoH
t (1−γ) and subsequently querying st+∆t in the same
episode. The negative samples s−t are uniformly sampledfrom any timestep within the current or any other episode.
Minimizing Equation (11) over D yields a function f ∗which, at optimality, approximates the future state visitation
probability under π up to a multiplicative term (Ma &Collins, 2018; Poole et al., 2019).
P[st+∆t |st ;π].
(12)f ∗ (st ,st+∆t ) ∝ log
P[st+∆t ;π]It should be noted that the time offsets in both ELE’s
progress model and in the contrastive pre-training phaseare sampled from similar distributions (see Appendix A.1).
In the following section, we show how f can be used foraccelerating exploration in the online setting.

4. Methodology
In this section, we provide details of both the pre-training
phase on offline human demonstrations and using these state
representations in Muesli (Hessel et al., 2021a), a strong RL
baseline. We also describe how to use same offline data fortraining progress model as well as training ELE’s progressmodel.
Pre-training state representations with contrastive training The idea behind pre-training offline representations isfairly straightforward: learn fundamental inductive biases
required for exploration from existing demonstrations (e.g.
state connectivity structure, action effects, perceptual local invariance, sequential dependence of states), therefore improving the sample-efficiency of the agent during the online phase.

Accelerating exploration and representation learning with offline pre-training

Figure 1 and Algorithm 1 (see Appendix) outline the generalparadigm of offline pre-training with online learning usedin all of our experiments, and which relies on finding ϕ thatminimizes Equation (11) over the set of possible encoders.
The pre-trained encoder is kept frozen or fixed through out
the training so that even if agent explores a new part of state
space, it does not drift away from the pre-trained representation. But it should be noted that we use a standard LSTM and
MLP on top of the frozen encoder which are trained throughout the training. In our experiments, we observe that thesepre-trained representations themselves are very useful for
improving sample efficiency of dense tasks but fail to solve
the sparse version of tasks themselves in NetHack where asingle reward is provided in the whole episode(more details
on sparsity of tasks in experimental section). To address this,we add the ELE progress reward (Anonymous, 2023) learnt
from the same data to the environment reward. The main
hypothesis is that the signal from the progress model solves
the problem of hard exploration in sparse reward tasks while
pre-training helps for faster learning and hence, improving
sample efficiency. Hence, one can use the same dataset tolearn signals of representation as well as additional reward
to assist exploration providing orthogonal benefits.
Why contrastive pre-training? Why should one pick thecontrastive pre-training scheme over any other objective?
First, as mentioned in Section 5, we test our hypothesis on
NetHack data containing only sequences of states (i.e. noactions nor rewards), which prevents the use of inverse RL
and reward prediction objectives such as SGI (Schwarzer
et al., 2021). Second, strong empirical evidence from prior
works suggests that contrastive learning is optimal forvalue function representation and outperforms latent andreconstruction based forward losses (Eysenbach et al., 2022).
Finally, latent-space prediction losses are known to be prone
to representation collapse, i.e. when bootstrapped prediction
targets are being exactly matched by the online networkand become independent of the state (Gelada et al., 2019).
Contrastive learning avoids representation collapse due tothe presence of negative samples which ensure uniformityof coverage on the unit sphere (Wang & Isola, 2020).

5. Experiments
In this section, we conduct a series of experiments tovalidate our central hypothesis: combining representation
learning with exploration improves the agent’s sample
complexity more than representation learning or explorationby themselves.
5.1. Experimental Details
Baselines: Our main results are obtained by comparing theperformance of tabula rasa Muesli and ELE (Anonymous,

2023) with their counterparts using pre-trained state representations. In addition, we also compare with standard baselinesthat use action-free expert demonstrations: GAIfO (Torabi
et al., 2018b), BCO (Torabi et al., 2018a) and FORM (Jaegle
et al., 2021), the same baselines as ELE. All these baselines
learn from the same offline data and are implemented on topof Muesli agent (Hessel et al., 2021a) for fair comparison and
we use the same hyperparameters as provided in ELE (Anonymous, 2023). Since previous work (Eysenbach et al., 2022)
demonstrated that contrastive pre-training performs muchbetter than other representation learning techniques, we limit
our comparison in this work to contrastive pre-training. We
performed preliminary investigations for other types of pretraining using forward prediction and latent models but they
performed inferior to contrastive pre-training. The motivation for using contrastive pre-training of state representations
is two-fold: 1) it allows Muesli to predict value functions
using a linear layer, making the task simpler, and 2) it was
shown to perform significantly better than latent-space orreconstruction-based objectives (see Eysenbach et al., 2022).
Tasks We use 7 different tasks from NetHack ranging fromdense to sparse rewards which were proposed in (Küttler et al.,2020) and ELE (Anonymous, 2023). On the dense side of thespectrum, we use the Score and Scout tasks, which reward
the agent for increasing the game score and revealing tiles
in the dungeon, respectively. At the sparse end, the Depth
N, Level N, and Oracle tasks deliver a reward of zero at all
timesteps until a trigger condition is met: reaching a particular dungeon depth, achieving a particular experience level, or
finding and standing next to the Oracle character (found between dungeon depths 5 and 9 inclusive); these sparse tasks
terminate when the condition is met. We believe that the wide
range of sparsity levels exhibited by this task collection represents a good selection of conditions under which to evaluate
the sample complexity of the algorithms we compare.
Dataset We use the NAO Top 10 dataset proposed inprevious work (Anonymous, 2023) which consists of humangames from top 10 players on nethack.alt.org. These
trajectories are useful for pre-training our contrastive
representations as this dataset provides a good balance oftrajectory quality and diversity(Anonymous, 2023) to learn
representations. This dataset consists of approximately 16K
trajectories of expert play, with a total of 184M transitions.
Frame Budget As we want to compare algorithms on
sample efficiency, we use 200M actor steps inspired from the
Atari benchmark (Bellemare et al., 2013) on all these tasks,
with the exception of the Oracle task. As this task poses a
significantly harder exploration challenge, we allow a larger
budget of 500M actor steps.

Accelerating exploration and representation learning with offline pre-training

Figure 2. Examples of observations generated by the NetHack Learning environment.

(a) Score (Dense)(b) Scout (Dense)(e) Level-2 (Sparse)(c) Depth-2 (Sparse)(f) Level-4 (Sparse)(d) Depth-4 (Sparse)(g) Oracle (Sparse)

Figure 3. Episode returns of Muesli and ELE, with and without pre-trained state representations. Dense reward tasks like Score and Scout
benefit immensely from contrastive pre-training in both performance and sample efficiency. While ELE’s exploration reward is needed tosolve sparse reward tasks, contrastive pre-training augments ELE by improving sample efficiency for all the sparse tasks. All curves arereported over 5 random seeds ± one standard deviation.

Architecture Inspired from previous work (Anonymous,
2023), we use a Residual Network (ResNet, He et al., 2016)architecture which encodes 80 × 24 TTY arrays (shown in
Figure 2) with a series of 2d convolutional block. This modelacts as encoder which is used by contrastive pre-training,
ELE’s progress model as well as Muesli agent. During online
phase, we pass the generated representation with a recurrent
network (LSTM) and MLP to predict policy and value heads
in Muesli. In case of contrastive pre-training, we simply passthe ResNet encoder through an MLP in order to project thestate features into a latent space. The ELE’s progress modelfuses the two states given as inputs, which are then passed
through a similar ResNet followed by an MLP to predicta scalar value in logarithmic space, that corresponds to thetemporal distance between both input states.

5.2. Results
We state the main results followed by ablation for different
components. All our experimental results are ran with 5
random seeds and are plotted with ± standard deviation.
Comparison of Progress Model and baseline with andwithout Pre-training Figure 3 shows that equippingstrong RL algorithms such as Muesli and ELE withhuman demonstrations via offline pre-training significantly
improves the sample complexity of the underlying method.
While ELE significantly outperformed Muesli on the sparser
tasks, the performance did not improve on the denser Score
and Scout tasks and in fact was inferior to Muesli. Using
contrastive pre-training with both Muesli and ELE, however,significantly improves its performance in the sample regime
under investigation in this work. On the sparse tasks DepthAccelerating exploration and representation learning with offline pre-training(a) Score (Dense)(b) Scout (Dense)(e) Level-2 (Sparse)(c) Depth-2 (Sparse)(f) Level-4 (Sparse)(d) Depth-4 (Sparse)(g) Oracle (Sparse)

Figure 4. Episode returns of ELE with pre-trained representation in contrast to standard imitation learning baselines (without access toactions) like BCO, GAIfO and FORM. Contrastive Pre-training + ELE outperforms all baselines on sparse as well as dense tasks. All curves
are reported over 5 random seeds ± one standard deviation.

(a) Score (Dense)(b) Scout (Dense)(e) Level-2 (Sparse)(c) Depth-2 (Sparse)(f) Level-4 (Sparse)(d) Depth-4 (Sparse)(g) Oracle (Sparse)

Figure 5. Episode returns of encoder extracted from ELE to Encoder separately trained by contrastive pre-training. All the tasks demonstratethat training a separate encoder by contrastive pre-training is much more useful on all the tasks illustrating ELE and contrastive pre-training
capture two different dimensions of exploration and representation learning respectively.

2, Depth 4, Level 2, Level 4 and Oracle, pre-training
with ELE significantly improves performance in the low
sample regime. It should be noted that the contrastive
pre-training without an exploration bonus struggles tosolve the sparser tasks, and the progress reward is clearly
beneficial in this case. This illustrates that the same dataset

can be used for both pre-training representations as well as
learning an exploration reward, and that these two different
applications of human data target orthogonal problems:
exploration bonuses help the agent discover the reward, andrepresentation learning improves its ability to exploit it.

Accelerating exploration and representation learning with offline pre-training

Comparison with Standard Imitation Learning Baselines: Figure 4 show the comparison of ELE + Pre-training
with other standard imitation learning baselines like GAIL
from Observations (GAIfO) (Torabi et al., 2018b), Behavior
Cloning from Observations (BCO) (Torabi et al., 2018a)
and FORM (Jaegle et al., 2021). On the sparser tasks,
only ELE and ELE + Pre-training are able to solve themat all, and contrastive pre-training improves convergence
speed significantly. Dense tasks like Score and Scout arelearned by many of the baselines, but contrastive pre-training
significantly improves sample efficiency.
5.3. Ablations
Is using progress model’s representation encoder as goodas contrastive pre-training? An interesting question
which stems from this work is that if the ELE’s progressmodel is useful for exploration, could we use the trained
progress model’s torso as an encoder for initializing
representations as well ? We experiment with extracting
the torso of trained progress model and use it to initialize
the representation encoder (instead of using encoder fromcontrastive pre-training). Figure 5 shows the comparisonof using progress model representations v/s training theseseparately by contrastive pre-training on the same data. We
observe that using ELE’s torso as encoder and additional
reward from ELE takes off faster but eventually achieves
poor performance than ELE itself and is significantly poor
than using ELE with contrastive pre-training encoder.
Do pre-trained state representations need to be finetuned
during the online phase? Next, we study the effect offreezing3 the pre-trained representations of observation
encoder from pre-training phase. We observe no significant
difference on most tasks with or without freezing representations. Figure 6 shows 1 sparse task and 1 dense task for thisablation (more details in appendix). We stick with freezing
representations for the online phase as our default setting
through out the paper.
How does encoder architecture impact its performance?
A natural question which arises when pre-training the state
representations offline is how well does the model capture
the future states. In all of our experiments, we use a simple
ResNet (He et al., 2016) which takes as inputs an array of TTY
characters. However, recent works have shown that the local
invariance biases from the 2d convolutions can be learned
through a vision transformer model (ViT, Dosovitskiy et al.,2020; Raghu et al., 2021), which positioned ViT as a competitive alternative to standard convolution-based architectures.
The main drawback of ViTs is their need to be trained on vast
amounts of data, which is abundant in NetHack. We have conducted pre-training experiments comparing the contrastive
3

Fixing a set of weights during the online phase.

(a) Score (Dense)(b) Oracle (Sparse)

Figure 6. Episode returns of with/without freezing representationsof contrastive pre-training during online phase. There is no significant difference between freezing or not freezing the state representations. We show 1 dense task (Score) and 1 sparse task (Oracle) as
representatives. All other tasks are shown in the Appendix

prediction accuracy of convolution-based models with that of
ViTs. Results shown in Figure 7 hint that ResNet-like models
are better suited for NLE, as they obtain better training set
and test set categorical accuracy as compared to ViTs.

(a) State prediction accuracy on (b) State prediction accuracy on
the training episodes.
the test episodes.
Figure 7. Ablation on different Model Architectures for contrastive
pre-training. We observe that Vision transformer performs muchworse than ResNet architecture.

6. Discussion
In this work, we posited that same offline data could be
used for learning representations as well as learning an
auxiliary reward to aid exploration and training these models
separately provide orthogonal benefits. We show thatpre-training state representations using contrastive learningand then using this network to initialize the representations
provides a large sample-efficiency improvement. However,using pre-training alone fail to solve sparse tasks. We
address the problem by adding a learned auxiliary reward
and observe that pre-training helps in representation
learning and auxiliary reward aids exploration. We validate
our hypothesis in the NetHack, a challenging rogue-liketerminal-based game with large state and action spaces, longtask horizon and strong notion of forward progress.

Accelerating exploration and representation learning with offline pre-training

References
Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-A., andHjelm, R. D. Unsupervised state representation learningin atari. Advances in neural information processing
systems, 32, 2019.
Anonymous. Learning about progress from experts. In
Submitted to The Eleventh International Conference
on Learning Representations, 2023. URL https://openreview.net/forum?id=sKc6fgce1zs.
under review.
Baker, B., Akkaya, I., Zhokhov, P., Huizinga, J., Tang, J.,
Ecoffet, A., Houghton, B., Sampedro, R., and Clune, J.
Video pretraining (vpt): Learning to act by watching unlabeled online videos. arXiv preprint arXiv:2206.11795,
2022.
Barreto, A., Dabney, W., Munos, R., Hunt, J. J., Schaul,
T., Van Hasselt, H., and Silver, D. Successor features
for transfer in reinforcement learning. arXiv preprint
arXiv:1606.05312, 2016.
Bellemare, M., Srinivasan, S., Ostrovski, G., Schaul,
T., Saxton, D., and Munos, R. Unifying count-based
exploration and intrinsic motivation. Advances in neural
information processing systems, 29, 2016.
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M.
The arcade learning environment: An evaluation platform
for general agents. Journal of Artificial Intelligence
Research, 47:253–279, 2013.
Burda, Y., Edwards, H., Storkey, A., and Klimov, O.
Exploration by random network distillation. arXiv
preprint arXiv:1810.12894, 2018.
Cheng, C.-A., Kolobov, A., and Swaminathan, A. Heuristicguided reinforcement learning. Advances in Neural
Information Processing Systems, 34:13550–13563, 2021.
Dayan, P. Improving generalization for temporal difference learning: The successor representation. NeuralComputation, 5(4):613–624, 1993.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn,D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer,
M., Heigold, G., Gelly, S., et al. An image is worth 16x16
words: Transformers for image recognition at scale. arXiv
preprint arXiv:2010.11929, 2020.
Dudik, M., Hsu, D., Kale, S., Karampatziakis, N., Langford,J., Reyzin, L., and Zhang, T. Efficient optimal learning forcontextual bandits. arXiv preprint arXiv:1106.2369, 2011.
Ecoffet, A., Huizinga, J., Lehman, J., Stanley, K. O., andClune, J. Go-explore: a new approach for hard-exploration
problems. arXiv preprint arXiv:1901.10995, 2019.

Eysenbach, B., Zhang, T., Salakhutdinov, R., and Levine,
S. Contrastive learning as goal-conditioned reinforcement
learning. arXiv preprint arXiv:2206.07568, 2022.
Fan, L., Wang, G., Jiang, Y., Mandlekar, A., Yang, Y.,
Zhu, H., Tang, A., Huang, D.-A., Zhu, Y., and Anandkumar, A. Minedojo: Building open-ended embodied
agents with internet-scale knowledge. arXiv preprint
arXiv:2206.08853, 2022.
Gelada, C., Kumar, S., Buckman, J., Nachum, O., andBellemare, M. G. Deepmdp: Learning continuous latent
space models for representation learning. In International
Conference on Machine Learning, pp. 2170–2179. PMLR,
2019.
Guo, Z. D., Thakoor, S., Pı̂slar, M., Pires, B. A., Altché,F., Tallec, C., Saade, A., Calandriello, D., Grill, J.-B.,
Tang, Y., et al. Byol-explore: Exploration by bootstrapped
prediction. arXiv preprint arXiv:2206.08332, 2022.
Hafner, D., Lee, K.-H., Fischer, I., and Abbeel, P. Deephierarchical planning from pixels.
arXiv preprint
arXiv:2206.04114, 2022.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. Mastering
diverse domains through world models. arXiv preprint
arXiv:2301.04104, 2023.
Hambro, E., Mohanty, S., Babaev, D., Byeon, M.,
Chakraborty, D., Grefenstette, E., Jiang, M., Daejin, J.,
Kanervisto, A., Kim, J., et al. Insights from the NeurIPS
2021 NetHack challenge. In NeurIPS 2021 Competitions
and Demonstrations Track, pp. 41–52. PMLR, 2022.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual
learning for image recognition. In Proceedings ofthe IEEE conference on computer vision and patternrecognition, pp. 770–778, 2016.
Hessel, M., Danihelka, I., Viola, F., Guez, A., Schmitt,S., Sifre, L., Weber, T., Silver, D., and Van Hasselt, H.
Muesli: Combining improvements in policy optimization.
In International Conference on Machine Learning, pp.
4214–4226. PMLR, 2021a.
Hessel, M., Kroiss, M., Clark, A., Kemaev, I., Quan, J., Keck,T., Viola, F., and van Hasselt, H. Podracer architectures
for scalable reinforcement learning. arXiv preprint
arXiv:2104.06272, 2021b.
Higgins, I., Pal, A., Rusu, A., Matthey, L., Burgess, C.,
Pritzel, A., Botvinick, M., Blundell, C., and Lerchner,
A. Darla: Improving zero-shot transfer in reinforcement
learning. In International Conference on Machine
Learning, pp. 1480–1490. PMLR, 2017.

Accelerating exploration and representation learning with offline pre-training

Ho, J. and Ermon, S. Generative adversarial imitation
learning. Advances in neural information processing
systems, 29, 2016.

rich-observation reinforcement learning. In International
conference on machine learning, pp. 6961–6971. PMLR,
2020.

Humphreys, P. C., Raposo, D., Pohlen, T., Thornton, G.,
Chhaparia, R., Muldal, A., Abramson, J., Georgiev, P.,
Santoro, A., and Lillicrap, T. A data-driven approach forlearning to control computers. In Chaudhuri, K., Jegelka,
S., Song, L., Szepesvari, C., Niu, G., and Sabato, S.
(eds.), Proceedings of the 39th International Conference
on Machine Learning, volume 162 of Proceedings of
Machine Learning Research, pp. 9466–9482. PMLR,
17–23 Jul 2022. URL https://proceedings.mlr.
press/v162/humphreys22a.html.

Oord, A. v. d., Li, Y., and Vinyals, O. Representation
learning with contrastive predictive coding. arXiv preprint
arXiv:1807.03748, 2018.

Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo,
J. Z., Silver, D., and Kavukcuoglu, K. Reinforcement
learning with unsupervised auxiliary tasks. arXiv preprint
arXiv:1611.05397, 2016.
Jaegle, A., Sulsky, Y., Ahuja, A., Bruce, J., Fergus, R.,
and Wayne, G. Imitation by predicting observations.
In International Conference on Machine Learning, pp.
4665–4676. PMLR, 2021.
Janner, M., Mordatch, I., and Levine, S. Generative temporal
difference learning for infinite-horizon prediction. arXiv
preprint arXiv:2010.14496, 2020.
Kim, H., Kim, J., Jeong, Y., Levine, S., and Song, H. O. Emi:
Exploration with mutual information. arXiv preprint
arXiv:1810.01176, 2018.
Küttler, H., Nardelli, N., Miller, A., Raileanu, R., Selvatici,
M., Grefenstette, E., and Rocktäschel, T. The nethacklearning environment. Advances in Neural Information
Processing Systems, 33:7671–7684, 2020.
Ma, Z. and Collins, M. Noise contrastive estimation and negative sampling for conditional models: Consistency and statistical efficiency. arXiv preprint arXiv:1809.01812, 2018.
Mazoure, B., Tachet des Combes, R., Doan, T. L., Bachman,
P., and Hjelm, R. D. Deep reinforcement and infomax
learning. Advances in Neural Information Processing
Systems, 33:3686–3698, 2020.
Mazoure, B., Ahmed, A. M., MacAlpine, P., Hjelm,
R. D., and Kolobov, A. Cross-trajectory representation
learning for zero-shot generalization in rl. arXiv preprint
arXiv:2106.02193, 2021.
Mazoure, B., Eysenbach, B., Nachum, O., and Tompson, J.
Contrastive value learning: Implicit models for simple
offline rl. arXiv preprint arXiv:2211.02100, 2022.
Misra, D., Henaff, M., Krishnamurthy, A., and Langford,
J. Kinematic state abstraction and provably efficient

Pomerleau,
D.
A.
Alvinn:Anautonomous
landvehicleinaneuralnetwork.
In
Touretzky,D.
(ed.),
Advancesin
NeuralInformation
Processing
Systems,
volume1.
Morgan-Kaufmann,
1988.
URL
https://
proceedings.neurips.
cc/paper/1988/file/812b4ba287f5ee0bc9d43bbf5bbe87fb-Paper.
pdf.
Poole, B., Ozair, S., Van Den Oord, A., Alemi, A., andTucker, G. On variational bounds of mutual information.
In International Conference on Machine Learning, pp.
5171–5180. PMLR, 2019.
Raghu, M., Unterthiner, T., Kornblith, S., Zhang, C.,
and Dosovitskiy, A. Do vision transformers see likeconvolutional neural networks? Advances in Neural
Information Processing Systems, 34:12116–12128, 2021.
Raileanu, R. and Rocktäschel, T. Ride: Rewarding
impact-driven exploration for procedurally-generated
environments. arXiv preprint arXiv:2002.12292, 2020.
Reid, M., Yamada, Y., and Gu, S. S.
help offline reinforcement learning?
arXiv:2201.12122, 2022.

Can wikipedia
arXiv preprint

Schwarzer, M., Rajkumar, N., Noukhovitch, M., Anand, A.,
Charlin, L., Hjelm, R. D., Bachman, P., and Courville,
A. C. Pretraining representations for data-efficient
reinforcement learning. Advances in Neural Information
Processing Systems, 34:12686–12699, 2021.
Srinivas, A., Laskin, M., and Abbeel, P. Curl: Contrastive
unsupervised representations for reinforcement learning.
arXiv preprint arXiv:2004.04136, 2020.
Stooke, A., Lee, K., Abbeel, P., and Laskin, M. Decoupling
representation learning from reinforcement learning.
In International Conference on Machine Learning, pp.
9870–9879. PMLR, 2021.
Tang, H., Houthooft, R., Foote, D., Stooke, A., Xi Chen,O., Duan, Y., Schulman, J., DeTurck, F., and Abbeel, P. #exploration: A study of count-based exploration for deep

Accelerating exploration and representation learning with offline pre-training

reinforcement learning. Advances in neural information
processing systems, 30, 2017.
Torabi, F., Warnell, G., and Stone, P. Behavioral cloning fromobservation. arXiv preprint arXiv:1805.01954, 2018a.
Torabi, F., Warnell, G., and Stone, P.
Generative
adversarial
imitationfromobservation.
CoRR, abs/1807.06158, 2018b.
URL
http://arxiv.org/abs/1807.06158.
Vinyals, O., Babuschkin, I., Czarnecki, W. M., Mathieu, M.,
Dudzik, A., Chung, J., Choi, D. H., Powell, R., Ewalds, T.,
Georgiev, P., et al. Grandmaster level in starcraft ii using
multi-agent reinforcement learning. Nature, 575(7782):
350–354, 2019.
Wang, T. and Isola, P. Understanding contrastive representation learning through alignment and uniformity on thehypersphere. In International Conference on Machine
Learning, pp. 9929–9939. PMLR, 2020.
Yang, M. and Nachum, O. Representation matters:
offline pretraining for sequential decision making. In
International Conference on Machine Learning, pp.
11784–11794. PMLR, 2021.
Yarats, D., Fergus, R., Lazaric, A., and Pinto, L. Reinforcement learning with prototypical representations.
In International Conference on Machine Learning, pp.
11920–11931. PMLR, 2021.
Zhang, J., Wang, J., Hu, H., Chen, T., Chen, Y., Fan, C.,
and Zhang, C. Metacure: Meta reinforcement learningwith empowerment-driven exploration. In International
Conference on Machine Learning, pp. 12600–12610.
PMLR, 2021.

Accelerating exploration and representation learning with offline pre-training

A. Appendix
A.1. Experimental details
Distribution of ∆t. For both the contrastive pre-training scheme, as well as ELE’s progress model, the distribution of
offsets between current and future states has a similar shape, albeit not being exactly equal: our current work chooses∆t ∼ GeomHt (1 − γ), which has three parameters, while ELE uses ∆t ∼ LogUniform(t,H). The advantage of using a
truncated geometric distribution is two-fold: 1) the support is discrete and therefore allows to directly sample time offsets
without truncation, and 2) the extra parameter γ is directly tied to the environment and allows to control for the shape of the
distribution, making the approach more flexible. Specifically, increasing γ → 1 brings the distribution closer to Uniform(t,H),
while decreasing γ → 0 brings it closer to δ(t).

Figure 8. Comparison of a log-uniform distribution with truncated geometric distributions with various success probabilities. It can be seenthat both coincide for some success probability.

Computational resources

We ran all of our experiments on 8 V100 GPUs for 6 hours.

A.2. Additional results
Freezing the pre-trained representation. Figure 9 shows the performance difference between using a frozen pre-trained
state representation in the online phase, vs allowing the gradient of the online loss propagate through the representation.
A.3. Algorithm
Algorithm pseudocode is provided in Algorithm 1.
Algorithm 1: Pre-training with online learning
Input :Dataset Dµ ∼ Pµ0:H , environment M , state embedding ϕ, policy π/* Pre-train ϕ using Equation (11)µ1 for minibatch B ∼ D do
2
Update ϕ using ∇ϕ ℓContrastive (ϕ)(B) ;/* Learn forward progress model (ELE)µ3 for minibatch B ∼ D do
4
Update g to minimize Equation (8) ;5 Initialize the online encoder with ϕ;
6 for iteration i = 1 → N do
7
Generate the trajectory τ using acting policy π(θ)/* Update reward in τ by adding progress reward8
for (st ,at ,st+1 ,rt ) in τ do
9Update rt ← rt +λg(st−k ,st )10Update θ using Muesli Loss function keeping ϕ frozen.

*/

*/

*/

Accelerating exploration and representation learning with offline pre-training(a) Score (Dense)(b) Scout (Dense)(e) Level-2 (Sparse)(c) Depth-2 (Sparse)(f) Level-4 (Sparse)(d) Depth-4 (Sparse)(g) Oracle (Sparse)

Figure 9. Episode returns of with/without freezing representations of contrastive pre-training during online phase. There is no significant
difference with/without freezing representations of encoder during online training for most of the tasks. We report all other results withfreezing the representations from the pre-training phase . All curves are reported over 5 random seeds ± one standard deviation.

A.4. Implementation details
Hyperparameters for our approach are shown in Table 1. All other hyperparameters are identical to ELE (Anonymous, 2023),
and the architecture of the encoder is the same as the one used in ELE. Each experiment was run on 8 TPUv3 accelerators
using a podracer configuration (Hessel et al., 2021b).
Hyperparameter

Value

Contrastive discount γ
Batch size

0.95
576 sequences

Table 1. Muesli hyperparameters that differ between our approach and ELE.