Value Explicit Pretraining for Learning Transferable Representations

## Abstract

Abstract We propose Value Explicit Pretraining (VEP), a method that learns generalizable representations for transfer reinforcement learning. VEP enables learning of new tasks that share similar objectives as previously learned tasks, by learning an encoder for objective-conditioned representations, irrespective of appearance changes and environment dynamics. To pre-train the encoder from a sequence of observations, we use a self-supervised contrastive loss that results in learning temporally smooth representations. VEP learns to relate states across different tasks based on the Bellman return estimate that is reflective of task progress. Experiments using a realistic navigation simulator and Atari benchmark show that the pretrained encoder produced by our method outperforms current SoTA pretraining methods on the ability to generalize to unseen tasks. VEP achieves up to a 2 × 2\times 2 × improvement in rewards on Atari and visual navigation, and up to a 3 × 3\times 3 × improvement in sample efficiency. For videos of policy performance visit our website .

## 1 Introduction

Learning generalizable representations for control is an open problem in visual sequential decision-making.
Typically, an encoder $\phi$ is learned using a large offline dataset via a predetermined objective function. Subsequently, $\phi$ is used for control by mapping high-dimensional visual observations from the environment $\mathbf{o}_{:t}$ into a lower-dimensional latent representation $\mathbf{z}_{t}$. The representation $\mathbf{z}_{t}$ is fed into a policy $\pi(\cdot|\mathbf{z}_{t})$ to generate an action $\mathbf{a}_{t}$ to solve a task. The key question in visual representation learning is: *what should $\phi$ be?*

The challenge in learning $\phi$ mainly lies in discovering the correct inductive biases that yield representations that can be used to learn a variety of downstream tasks in a sample efficient manner.

It is unclear, however, what such useful inductive biases are. Initial approaches (Shah & Kumar, ; Yuan et al., ; Parisi et al., ) to this problem included simply reusing pretrained vision models trained to solve computer vision tasks like image recognition, zero-shot for control. Subsequently, works like R3M and VIP (Nair et al., ; Ma et al., ), tried to utilize temporal consistency, enforcing that images that are temporally close in a video demonstration are embedded close to each other.
Other directions include works like Voltron (Karamcheti et al., ) and Masked Visual Pretraining (Radosavovic et al., ; Xiao et al., ; Seo et al., ) which attempt to use image reconstruction as one such inductive bias.

While biases induced by pretraining objectives like image reconstruction and temporal consistency have been shown to greatly improve downstream policy performance, these pretraining objectives used to learn $\phi$ are *distinct* from the downstream usage of $\phi$, e.g., the task of image reconstruction is very different from that of action prediction. Thus, there exists an unmet need for representation learning approaches that *explicitly* encode information directly useful for downstream control during the process of learning $\phi$.

This is, of course, challenging — how do we encode control-specific information without actually training online on a control task? Our crucial insight is that encoding control-specific information in the representations generated by $\phi$ is possible by harnessing the power of Monte Carlo estimates of control heuristics computed offline on demonstration datasets.

Our key contribution is *Value Explicit Pretraining (VEP)*, a contrastive learning approach that utilizes offline demonstration datasets with reward labels to learn a representation for visual observations. Our method utilizes the insight that observations with similar Monte Carlo Bellman return estimates across multiple tasks share a similar propensity of success and in tasks with related goals, also share a similar optimal policy. For example, in shooter games on Atari, despite differences in the visual appearances of adversaries, the strategy to effectively shoot them is similar.

VEP utilizes this intuition to learn an encoder using a contrastive loss which embeds observations with similar value function estimates across a set of training tasks near each other. We study the performance gains obtained by utilizing the VEP representation for policy learning, both on the train set of tasks and visually distinct yet related held-out tasks. We experiment on the Atari benchmark and on a visual navigation benchmark comparing VEP to state-of-the-art methods like VIP(Ma et al., ) and SOM(Eysenbach et al., ). We find up to a $2\times$ improvement in the rewards obtained on both Atari games and visual navigation and $3\times$ improvement in sample efficiency of online RL algorithms trained on VEP.

Figure: Figure 1: High-level overview of our problem statement The encoder $f_{\phi}$ is pretrained using expert videos from a set of train tasks, that is then reused for an unseen task. We evaluate pretrained encoders produced by our method and the baselines on the Atari and Navigation benchmarks.
Refer to caption: x1.png

## 2 Related Work

Representation learning for Robotics. Besides the general idea that the representation has the role of encoding the essential information for a given task, while discarding irrelevant aspects of the original data, typical state representation learning methods attempt to embed an observation into a latent representation that could be utilized by the downstream task (Lesort et al., ). It is also important that these methods produce a low dimensional representation that allows the control policy to efficiently learn the downstream task. Traditionally, unsupervised methods like VAE (Kingma & Welling, ), can learn disentangled representations that can be used to correlate with observation-specific appearance factors (Higgins et al., ), and are used for policy learning (Ha & Schmidhuber, ). However, in many environments, these representations prove difficult to learn a more optimal policy, since there the temporal structure is missing in these representations. (Anand et al., ) explores this direction and learns representations by enforcing temporal structure through contrastive loss.

Pretraining for RL. Pretraining for representation learning, in the context of RL, involves learning transferable knowledge, typically in the form of good representations, that helps the agent perceive the world better (Xie et al., ). Compared with traditional unsupervised methods, the objective of pretraining for RL is to learn representations that allow efficient downstream RL. Majority of the earlier Online pretraining works learn representations that model the task dynamics that can be learnt through a sequence of observations during the RL procedure (Pathak et al., ). More recent Offline pretraining methods like (Schwarzer et al., ) build on the prior work (Anand et al., ) by pretraining an encoder using unlabeled data and then finetune on a small amount of task-specific data. In comparison with these approaches, our method, focuses on learning representations that not only aid solving in-domain tasks, but also is able to generalizes to out-of-domain tasks (tasks that the model has not encountered, through unlabeled data).

Transfer after Pretraining. Transfer learning for reinforcement Learning is an ongoing research problem. Early works like Progressive networks (Rusu et al., ) attempt to solve it by reusing features learnt from source tasks through adapters. (Gamrian & Goldberg, ) performs image-to-image translation using GANs. However these method are limited to predefined source or target domains. More recent works focus on a more challenging problem of using only diverse, expert videos for offline pretraining that could later be transferred to solve a downstream task. These methods have gained popularity in RL for their use of self-supervised based pretraining (Sermanet et al., ) based on Contrastive Learning. Recent works like (Eysenbach et al., ; Ma et al., ) propose methods for learning an embedding space that could be used for Zero-shot reward specification.

Figure: Figure 2: Description of our method (VEP). We compute value estimates (Bellman returns), as denoted by $G$, for each frame. We then use a contrastive learning-based pretraining method that learns task-agnostic representations based on $G$. The above figure is a pictorial representation of a training scenario where the sampling batch size $b_{T}$ is 2 and the training batch size $b_{G}$ is 1. This results in anchor, positive and negative sampled from two sequences in each batch.
Refer to caption: x2.png

## 3 Problem setting and Preliminaries

Let $\mathcal{T}_{train}=\{\mathcal{T}_{1},\mathcal{T}_{2},...\mathcal{T}_{m}\}$ be a set of train tasks, that have associated expert video demonstration datasets $\mathcal{D}_{train}=\{\mathcal{D}_{1},\mathcal{D}_{2},\mathcal{D}_{3},...
\mathcal{D}_{m}\}$. During pretraining, we assume that the Encoder model $f_{\phi}$ parameterized by $\phi$ has access to $\mathcal{D}_{i}$ for each of the train-tasks $\mathcal{T}_{train}$. Each $\mathcal{D}_{i}$ consists of a sequence of frames $\{o_{t}^{i}\}$. The Encoder $f_{\phi}$ learns to encode images/observations $o_{t}$ into an embedding $z_{t}$, which is taken as an input by the policy $\pi$ to perform a test-task. The set of test-tasks are denoted by $\mathcal{T}_{test}=\{\mathcal{T}_{m+1},\mathcal{T}_{m+2},...\mathcal{T}_{n}\}$. Note that although $\mathcal{T}_{train}\cap\mathcal{T}_{test}=\varnothing$, all the tasks in $\mathcal{T}_{train}\cup\mathcal{T}_{test}$ share a similar objective, as defined by the value function.

As our approach aims to facilitate and accelerate transfer learning for tasks with a similar objective, such similarity must be evaluated. This is an open problem in general (Zhu et al., ). Several methods have been proposed to quantify similarity across RL tasks (Lazaric et al., ). This, however, is not our primary objective here, and instead, we manually selected tasks that are intuitively similar, in two domains: 1) for urban visual-based navigation, we use maps of several cities and photographs taken along the available streets; to make tasks similar, we use a similar goal location relative to the agent’s starting point, which eliminates the need to explicitly mention goal location. 2) In Atari games, we select the shooter games, which all highly resemble the Space Invaders concept, albeit with graphical and other variations: an army of alien enemies descends towards the bottom of the screen, where the agent’s ship is, which can move left or right or shoot straight up.

For both our method and the baselines, the encoder $f_{\phi}$ is trained only using the expert videos $\{\mathcal{D}_{i},\mathcal{D}_{2},..\mathcal{D}_{m}\}$ without any fine-tuning. The objective of our method is to efficiently learn the encoder using a sequence of observations $\mathcal{D}$ from the source tasks $\mathcal{T}_{train}$, such that the embeddings from $f_{\phi}$ could be zero-shot transferred to unseen test tasks $\mathcal{T}_{test}$.

### 3.1 Contrastive Representation Learning

Typically, contrastive representation learning methods for RL utilize offline demonstration datasets. These methods typically input a batch of anchors $\mathbf{o}_{an}$, positives $\mathbf{o}_{ps}$, and negatives $\mathbf{o}_{ng}$ and minimize a predetermined similarity metric that enables an encoder model to learn consistent and meaningful representations that can be used for downstream tasks. The earliest known formulation (Schroff et al., ) uses an $||\cdot||_{2}$ distance to embed the positives and the anchor close to each other and the negatives far away from the anchor.

$$ $\begin{multlined}\mathcal{L}_{triplet}=\\ \;\sum_{\mathbf{z}\in\mathcal{X}}\mathbf{max}\Big{[}\mathbf{0},||\mathbf{z}_{ an}-\mathbf{z}_{ps}||_{2}^{2}-||\mathbf{z}_{an}-\mathbf{z}_{ng}||_{2}^{2}+ \epsilon\Big{]}\end{multlined}\mathcal{L}_{triplet}=\\ \;\sum_{\mathbf{z}\in\mathcal{X}}\mathbf{max}\Big{[}\mathbf{0},||\mathbf{z}_{ an}-\mathbf{z}_{ps}||_{2}^{2}-||\mathbf{z}_{an}-\mathbf{z}_{ng}||_{2}^{2}+ \epsilon\Big{]}$ (1) $$

In the above equation $\mathbf{z}_{an}$, $\mathbf{z}_{ps}$ and $\mathbf{z}_{ng}$ represent the embeddings that are obtained after passing observations $\mathbf{o}_{an}$, $\mathbf{o}_{ps}$ and $\mathbf{o}_{ng}$ (anchors, positives and negatives) through the encoder network $f_{\phi}$ and are of shape $[B,D]$, where $B$ is the batch size and $D$ is the embedding size. Other metrics like cosine similarity could also be used instead of $||\cdot||_{2}$ distance, to compute the similarity between embeddings.

Similar to recent methods like (Ma et al., ), the InfoNCE (van den Oord et al., ) objective can also be used to optimize the encoder parameters. Unlike Triplet loss from Equation 1, InfoNCE permits utilizing multiple negative examples for calculating the loss (via to the expectation term in the denominator of Equation 2). As depicted below, InfoNCE aims to maximize mutual information of the anchors and positives.

$$ $\mathcal{L}_{InfoNCE}=\mathbb{E}_{\mathbf{z}_{ps}}\left[-\log\frac{\mathcal{S} _{\phi}\left(\mathbf{z}_{an},\mathbf{z}_{ps}\right)}{\mathbb{E}_{\mathbf{z}_{ ng}}\mathcal{S}_{\phi}\left(\mathbf{z}_{an},\mathbf{z}_{ng}\right)}\right]$ (2) $$

In the above equation, $\mathcal{S}_{\phi}$ is a distance function in the $\phi$-representation space that is used to compute the similarity between a pair of embeddings. In our experiments, that use InfoNCE, $\mathcal{S}$ takes the form of cosine similarity. More recently, Soft-Nearest Neighbor loss (Frosst et al., ) is a generalization of InfoNCE which also allows multiple positive examples (Weng, ) to be utilized in the computation of the objective.

$$ $\begin{multlined}\mathcal{L}_{SNN}=\\ \;\sum_{\mathbf{z}\in\mathcal{X}}-\log\frac{\exp(-\mathcal{S}_{\phi}\left( \mathbf{z}_{an},\mathbf{z}_{ps}\right)/\tau)}{\exp(\mathcal{S}_{\phi}\left( \mathbf{z}_{an},\mathbf{z}_{ng}\right)/\tau)}\end{multlined}\mathcal{L}_{SNN}= \\ \;\sum_{\mathbf{z}\in\mathcal{X}}-\log\frac{\exp(-\mathcal{S}_{\phi}\left( \mathbf{z}_{an},\mathbf{z}_{ps}\right)/\tau)}{\exp(\mathcal{S}_{\phi}\left( \mathbf{z}_{an},\mathbf{z}_{ng}\right)/\tau)}$ (3) $$

The Soft-Nearest Neighbors loss is especially useful for VEP. VEP utilizes it during a sampled batch of multiple positives (sampled from multiple training tasks from $\mathcal{T}_{train}$).

Figure: Figure 3: Pretraining results on Atari. Performance of different pretraining methods on the respective games as mentioned above. The encoder is pretrained only on the first 2 games (Demon Attack and Space Invaders) and is evaluated on the other out-of-domain games.
Refer to caption: x3.png

### 3.2 Discounted Returns and Value Functions

We operate using the formalism afforded by a POMDP (Partially Observable Markov Decision Process) with ($\mathcal{O}$, $\mathcal{S}$, $\mathcal{A}$, $p$, $\theta$, $r$, $T$, $\gamma$) representing an observation space $\mathcal{O}$, state space $\mathcal{S}$, action space $\mathcal{A}$, transition function $p$, emission function $\theta$, reward function $r$, time horizon $T$, and discount factor $\gamma$. An agent in state $\mathbf{s}_{t}$ takes an action $\mathbf{a}_{t}$ and consequently causes a transition in the environment through $p(\mathbf{s}_{t+1}\mid\mathbf{s}_{t},\mathbf{a}_{t})$. The agent receives the next state $\mathbf{s}_{t+1}$ and reward $r_{t}=r(\mathbf{o}_{t},\mathbf{a}_{t})$ calculated using the observation $\mathbf{o}_{t}$.
The goal of the agent is to learn a policy $\mathbf{\pi}$ which maximizes the expected discounted sum of rewards. The discounted sum of rewards for a sub-trajectory starting from state $\mathbf{s}_{t}$ is given by:

$$ $G(\mathbf{s}_{t})=r_{t+1}+\gamma r_{t+2}+\cdots+\gamma^{3}r_{t+4}=\sum_{k=t}^{ T}\gamma^{(k-t)}r_{k}$ (4) $$

The expectation of this discounted return under the trajectory distribution $p(\tau)$ and the policy $\pi$ where $\tau$ is a trajectory of the form $(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1},\cdots)$ is often defined as the value of the state $\mathbf{s}_{t}$ under policy $\pi$, denoted by $\mathcal{V}^{\pi}(\mathbf{s}_{t})$.

## 4 Method

### 4.1 Intuition

The value $\mathcal{V}^{\pi}(\mathbf{s}_{t})$ of a state $\mathbf{s}_{t}$ under a policy $\pi$ intuitively defines the propensity for the success of solving a task by following policy $\pi$. Intuitively, if two states have similar value estimates, they likely have a similar expected return under $\pi$.

With this in mind, we now motivate VEP with an example. Consider the task of shooting an adversary in the Atari game of Space Invaders. Assume that there exists an optimal policy for this task denoted by $\pi^{*}(\cdot|\mathbf{o}_{t})$ which operates on image observations and the associated optimal value function $\mathcal{V}^{\pi^{*}}(\cdot)$. Now, consider a slightly perturbed version of this game in which all the adversaries are colored orange. If policy $\pi^{*}$ must solve this perturbed task, it must be invariant to the color of the adversary. One way to achieve this invariance is to enforce that the value estimates of states with similar propensities for success are similar, e.g., the value estimate of a state containing a bullet very close to an adversary should be the same regardless of whether the adversary is yellow or orange.

VEP utilizes this exact intuition by learning representations that induce such an invariance. We assume access to expert demonstrations consisting of only observations for the set of training tasks $\mathcal{T}_{train}$. This kind of data can be obtained from YouTube videos of gameplay or demonstration videos and does not contain any action labels. Further, it is assumed to be generated by an expert and that each of the trajectories is successful. Subsequently, we define the end of each demonstration as a goal observation $\mathbf{o}^{goal}$. We also do not have access to the true reward function, so we operate under a sparse reward setting, assuming that a reward of $1$ is obtained at $\mathbf{o}^{goal}$ and 0 0 everywhere else. We now label each observation with a value estimate computed using Equation 4. Ideally, this label would have been computed using $\mathcal{V}^{\pi^{*}}(\cdot)$, but since we do not have access to the true value function of the optimal policy, we utilize a Monte Carlo estimate of this using Equation 4.

Having labeled the demonstration datasets for tasks in $\mathcal{T}_{train}$ with $G(\cdot)$ from Equation 4, we now train the encoder $\phi$ using a contrastive learning objective which first samples a scalar value estimate $g$ between 0 0 and $1$ and then samples multiple observations from $\mathcal{D}_{1}\cup\mathcal{D}_{2}\cdots$ with $G(\cdot)$ values within an $\epsilon$ of $g$. Subsequently, an encoder $\phi$ is learned which embeds these observations close to each other. Consequently, observations with a similar propensity for success have similar embeddings.

Figure: Figure 4: Pretraining results on Navigation. Performance of different pretraining methods on the respective cities as mentioned above. Similar to the Atari experiments, for all the baselines, expert videos from the first two tasks (Wall Street and Union Square) were used for pretraining. VEP representations improve PPO policy performance by up to $2\times$.
Refer to caption: x4.png

### 4.2 Implementation

To make the training computationally efficient, we preprocess $\mathcal{D}_{1}\cup\mathcal{D}_{2}\cdots$ and save a dictionary that maps sorted values $G(\cdot)$ to the indices of corresponding observations with the same Monte Carlo value estimate. This speeds-up the value look-up subroutines through binary search (see supplementary material for implementation details).

We first sample a batch of value estimates from the dataset determined by train batch size $b_{G}$. Next, we sample a $b_{T}$ number of training tasks. For a given value estimate and a given task, we load the observation with this value estimate as an anchor image. The positive observation corresponding to the anchor is an image with a value estimate within an $\epsilon$ of the value estimate of the anchor. In our experiments, we only sample $2$ training tasks during pretraining, i.e., $b_{T}=2$. Subsequently, the pretraining objective becomes the following:

$$ $\begin{multlined}\max_{\phi}\mathbb{E}_{g\sim[0,1]}[\mathcal{S}_{\phi}(\mathbf {z}_{an}^{\mathcal{T}_{1}},\mathbf{z}_{ps}^{\mathcal{T}_{1}})+\mathcal{S}_{ \phi}(\mathbf{z}_{an}^{\mathcal{T}_{2}},\mathbf{z}_{ps}^{\mathcal{T}_{2}})+\\ \mathcal{S}_{\phi}(\mathbf{z}_{an}^{\mathcal{T}_{1}},\mathbf{z}_{ps}^{\mathcal {T}_{2}})+\mathcal{S}_{\phi}(\mathbf{z}_{an}^{\mathcal{T}_{2}},\mathbf{z}_{ps} ^{\mathcal{T}_{1}})]\end{multlined}\max_{\phi}\mathbb{E}_{g\sim[0,1]}[\mathcal {S}_{\phi}(\mathbf{z}_{an}^{\mathcal{T}_{1}},\mathbf{z}_{ps}^{\mathcal{T}_{1}} )+\mathcal{S}_{\phi}(\mathbf{z}_{an}^{\mathcal{T}_{2}},\mathbf{z}_{ps}^{ \mathcal{T}_{2}})+\\ \mathcal{S}_{\phi}(\mathbf{z}_{an}^{\mathcal{T}_{1}},\mathbf{z}_{ps}^{\mathcal {T}_{2}})+\mathcal{S}_{\phi}(\mathbf{z}_{an}^{\mathcal{T}_{2}},\mathbf{z}_{ps} ^{\mathcal{T}_{1}})]$ (5) $$

where $g$ corresponds to a particular scalar value of the Bellman return estimate $G(\cdot)$. $\mathbf{z}_{an}^{\mathcal{T}_{1}}$ corresponds to the embedding of the anchor, i.e., observation with a value estimate equal to $g$ sampled from task $\mathcal{T}_{1}$, and $\mathbf{z}_{an}^{\mathcal{T}_{2}}$ corresponds to the embedding of the anchor sampled from task $\mathcal{T}_{2}$. Similarly, $\mathbf{z}_{ps}^{\mathcal{T}_{1}}$ corresponds to a positive, i.e., an observation with a value estimate within $\epsilon$ of $g$ sampled from task $\mathcal{T}_{1}$ and $\mathbf{z}_{ps}^{\mathcal{T}_{2}}$ corresponds to a positive sampled from task $\mathcal{T}_{2}$.

Intuitively, this objective encourages the positives and anchors from all the sampled tasks to embed near each other consequently, using the a value function estimate to organize the latent space of the learned encoder $\phi$.
For full implementation details like batch sizes etc., please refer to the supplementary material.

Figure: Algorithm 1 Value Explicit Pretraining

Figure: Figure 5: Comparision of our method with End-to-end trained method for Navigation task. Note that in each of the above training curves, the end-to-end baseline has the entire model trained on each of the above tasks, whereas our method (VEP) is pretrained only on expert videos from Wall Street and Union Square. Compared to any pretrained method, End-to-end training baseline takes significantly longer time ($2.1\times$ for Navigation and $3.3\times$ for Atari). Since both the methods were trained for the same number of iterations (20M), our method finished earlier and the dotted line is only for comparison
Refer to caption: x5.png

## 5 Experimental Setup

We study whether utilizing VEP as a pretraining objective to learn an encoder improves (1) policy learning on in-distribution tasks, i.e., those tasks for which data was available to pre-train the encoder and (2) whether the learned encoder aids transfer learning of new tasks.
We performed our experiments using the benchmark specified in the next paragraph. We used the RLLib library (Liang et al., ) under the Ray ecosystem for all our RL experiments. We used PPO (Schulman et al., ) for training the policy. Additional details for our experimental setup are mentioned in the Supplementary material.

### 5.1 Environments

Atari. We perform experiments by selecting a set of Shoot’em up games. Although all the games share a common objective of shooting enemies that spawn from above, there are significant differences in appearances and dynamics across games. We select a set of six games and split them into $\mathcal{T}_{train}$ and $\mathcal{T}_{test}$. For pretraining the encoder, we use a sequence of observations obtained from the D4RL datasets (Fu et al., ). We truncate large sequences of raw trajectories into smaller trajectories by setting a state to terminal, when the agent gets a reward. The terminal states across all the games are unified with the objective, as all the Shoot’em up games deal with shooting the enemy objects. The value estimates of each frame in a sequence are then computed using Equation 4. The list of all games and the comparison across each of them is part of the supplementary material.

Figure: Figure 6: Reward functions for Navigation. For a specific map, the agent spawns at a predetermined starting location (red), with the flexibility to initiate at a random location within a $r$-step to the fixed starting point. The sparsity of the rewards (brown lines) that enable the agent to navigate to the goal (green) can be adjusted through the parameter $L$.
Refer to caption: extracted/5454854/pics/reward.png

Navigation. We build an engine (Henghui Bao*, ) that loads the StreetLearn dataset (Mirowski et al., ) to perform visual navigation, based on gym (Brockman et al., ). In a typical Navigation task, the agent is designed to randomly respawn within a radius $r$ of a predetermined location $(src_{x},src_{y})$, with the objective reaching a goal location that is sampled within a radius $r$ of a location $(d_{x}+src_{x},d_{y}+src_{y})$. Reward acquisition is structured through a linear distribution of $L$ reward points (including the reward obtained upon reaching the target) uniformly spanning the starting point to the goal. The agent only earns rewards as it moves closer to the goal as depicted through the yellow lines in Figure 6. We have six cities for this benchmark, and we established consistent horizontal and vertical displacements $(d_{x},d_{y})$ between the starting and target points across all cities, avoiding the need for any explicit goal information (details are provided in the supplementary material). The agent is then expected to transfer to an unseen test city after learning from the expert videos obtained from a set of cities. Tasks across all the cities are all solvable within a predefined horizon. Lastly, using a planner, we obtain the expert videos by sampling all the paths under a distance limit and truncate all the paths such that the terminal state of the truncated path has a reward of $1$. For all the tasks, we set $L=15$, and $r=5$.

We use the same encoder architecture for both Atari and Navigation to embed pixel observation into vector space. To enable to temporal understanding of the state, all the embeddings in the past four timesteps are concatenated together and passed onto the policy. For the Navigation task, apart from the four stacked embeddings, we also obtain the subsequent four stacked odometry information $(odom_{x},odom_{y})$, of the agent, that is concatenated with the four stacked embedding and passed into a linear layer. This enables the agent to understand its ego-centric post that is crucial for navigating to the goal.
We first pretrain the encoder $f_{\phi}$ using the method described in the previous section and visually shown in Figure 2. This is achieved by using a sequence of unlabelled trajectories from both the games. Once we obtain the pretrained encoder, we use an online RL algorithm, in our case PPO (Schulman et al., ), to train a policy. We summarize results in Figure 3 and Figure 4.

Figure: Figure 7: We compared by equating the batch size and the number of iterations to match those of VEP by combining sample and train batch size, to show that the learning ability of our method is due to value estimates amidst tasks.
Refer to caption: x6.png

### 5.2 Results

Baselines. Value Implicit Pretraining (VIP) (Ma et al., ) learns temporally smooth embeddings by learning by encoding the goal (positive) and the start (anchor) image close and the middle images (negatives) within a sampled sub-trajectory. By learning this objective, through multiple sampled sub-trajectories, the encoder recursively learns continuous embeddings across sequential frames in a trajectory.
Time Contrastive Learning (TCN) involves sampling the positive within a certain margin distance $d_{thresh}$ from the anchor and a negative anywhere from the positive to the end of the trajectory (Sermanet et al., ). If the anchor is sampled at time instant $t_{a}$, positive is sampled at $t_{p}$ and the negative sampled at $t_{n}$, then $|t_{n}-t_{a}|>|t_{p}-t_{a}|$. We then use the standard triplet loss for optimization, although other contrastive losses could also be used. Unlike TCN, the positives could also be sampled from State Occupancy Measure (SOM) (Eysenbach et al., ) that could be embedded close to the anchor. The negative, on the other hand, is sampled anywhere from the other episodes under a specific task or other task. State occupancy measure from a specific instant $t$ is sampling within a truncated geometric distribution $Geo_{t}^{H}(1-\gamma)$ with probability mass re-distributed over the interval $[t,H]$, where $H$ is the horizon. (Mazoure et al., ):

Online RL experiments on Atari. For experiments involving Atari games, we trained the policy by freezing the pretrained encoder, without any additional fine-tuning. The encoder is pretrained using expert videos from Demon Attack and Space Invaders, and evaluated on a set of in-domain and out-of-domain environments. We find that the pretrained encoder is able to outperform baselines on the in-domain by $\sim 25$ percent. This margin is increased in the transfer experiments, most notably on Phoenix, with nearly $2\times$ improvement over baselines.

Online RL experiments on Navigation. Similar to the above, we froze the encoder and trained a layer policy. As mentioned before, unlike the model that was used for Atari, we also had the odometry information for a specific image that had to part of the embedding for the policy to perform the task. The embedding that was obtained from the CNN was concatenated with the 2D odometry information and was passed through another fully connected layer to obtain an embedding. All these parameters were used to pretrain. VEP outperforms all of our baselines by a larger margin in the navigation set as seen in Figure 4. VEP also outperformed the End-to-End trained baseline by achieving the same performance $2.1\times$ faster (Figure 5). We hypothesize that the better performance of our method in the Navigation tasks was due to a similar distribution of value estimate across the cities in the Navigation task, than the Atari games. Detailed specifications of the value estimates for all the Atari games and the cities in Navigation are described in the supplementary material.

Larger batch size and more iterations All the baseline approaches we compared against had a fixed train batch size that is used for computing gradients. For VEP, we are required to use a larger batch size — $b_{G}\times b_{T}$. To ensure that gains demonstrated by VEP cannot be attributed to larger batch sizes, we doubled the batch size for TCN as seen in Figure 7. The larger batch size for TCN still does not match the performance of VEP.

Early stopping to prevent overfitting. For the Navigation task, we increased the number of training tasks. We observed that the performance degraded in this setting. As shown in Figure 8, when we reduce the number of iterations, the model retains the performance, which suggests that our method learns much faster with an increase in data diversity and early stopping can prevent overfitting.

Figure: Figure 8: Comparison of our method with different early stop iterations. Notice that with an increase in number of games, our method performs better with less number of training iterations
Refer to caption: x7.png

## 6 Conclusion

Transferring policies to novel but related tasks is an important problem that needs to be addressed. We formulated a method to learn representations of states from different tasks solely based on the temporal distance to the goal frame. This way, the skills learnt from the train tasks could be transferred to unseen related tasks. We show the efficacy of our method by performing comprehensive evaluations on Atari and Visual Navigation.

## 7 Impact Statement

Our work opens new avenues for efficient training of new RL tasks by leveraging what was previously learned on similar tasks. We believe that has the potential to enable a much broader use of RL in real-life scenarios, as it eliminates the major hurdle of long and tedious training from scratch for each new task. We do not believe this work has particular ethical concerns. Its potential transformative societal impact is high as it makes RL for sequential tasks more achievable than previously possible.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 5
presentation: 2
questions: In general, I have these two core questions regarding the reliance on expert data in similar tasks and the experimental evaluation. I hope the author can add corresponding experiments or give necessary clarifications in their rebuttal.
soundness: 3
strengths: - This paper is well-written and easy to understand. It offers a clear background for the proposed method and compares it with recent competitive models, such as VIP.
- The fundamental idea that observations with similar value estimates should share similar representations is interesting. Following this line, the paper introduces an innovative contrastive pretraining objective. This approach differs from earlier methods that depend on image reconstruction or temporal consistency.
- The experiments conducted in visual navigation tasks are especially noteworthy, showing that the pretraining method significantly improves generalization to previously unseen tasks.
summary: This paper presents a novel transfer learning method. The aim is to learn generalizable representations from offline expert data that can be effectively transferred to similar reinforcement learning tasks. 

This approach's main contribution is a new form of self-supervised contrastive loss for training the visual encoder of the RL agent, which is based on an explicit state value function. Specifically, the positive pairs in this objective function are generated by sampling states with similar expected values. 

Experiments were conducted in two different environments, including Atari games and visual navigation tasks. The proposed method was shown to improve the averaged episodic returns.
weaknesses: Main concerns:
- Regarding the methodology, the proposed pretraining method relies heavily on high-quality expert data, which should be closely related to the actual tasks being performed. This limitation affects the applicability of the method. I am curious whether the proposed approach would still be viable when the source and target tasks share only visual similarities but have significantly different reward definitions.
- The experiments could be strengthened further. The Atari environment used in this paper is becoming somewhat "outdated" in current visual RL research, and the navigation tasks are not commonly employed in mainstream studies. I would suggest that the authors consider incorporating a more widely used environment, such as DeepMind Control, Meta-World, or Minecraft. It can facilitate direct comparisons with existing methods and highlight the effectiveness of the proposed approach.

Minor issues:
- Some figures (such as Figures 3-5) are not positioned properly in the paper, which can affect the overall reading experience.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 1/ It is unclear how the offline data is used for pre-training. The paper mentions that a suboptimal agent was used to collect the data, meaning that not all episodes are successful. Computing the Bellman return estimates appears to require successful trajectories. My question is: Are the failure trajectories excluded from pre-training, or are they used in some way?

2/ The message conveyed by Fig. 5 is unclear. Is it suggesting that VEP requires less wall time for online downstream task training compared to end-to-end training, despite having the same number of interactions, due to the use of a frozen encoder? If that is the case, for a fair comparison, the wall time for pre-training should also be taken into account. Additionally, it is worth noting that in Fig. 5, the Allegheny curve starts to trend upward toward the end. This raises the possibility that end-to-end training might perform better if extended over a longer period.

3/ [1] appears similar in that it uses contrastive learning to bring behaviorally similar trajectories closer together. This paper would benefit from a comparison and discussion with [1].

[1] Amy Zhang, et al. "Learning invariant representations for reinforcement learning without reconstruction." ICLR (2021)
soundness: 2
strengths: The problem setup is both interesting and challenging. The goal is to learn a generalizable representation that can transfer to out-of-distribution tasks with varying dynamics and appearances. The offline data used for training contains only sparse rewards, with no action labels or dense rewards, and the dataset is collected using a suboptimal agent. Given these challenges, the paper is valuable as it addresses a difficult and realistic problem.
summary: The paper introduces a new method designed to learn generalizable representations that enhance transferability in RL. The approach aims to efficiently learn new tasks with similar objectives by utilizing a contrastive loss that relates states across different tasks based on Bellman return estimates, which indicate task progress. Experiments conducted on navigation and shooting tasks in Atari demonstrate that their method (VEP) outperforms baseline methods on a few previously unseen tasks.
weaknesses: 1/ Although the high-level idea of using contrastive learning based on Bellman return estimates sounds interesting, I am not fully convinced that task progress is a suitable criterion for contrastive loss when aiming to transfer across different tasks. Let’s say we have task 1 (T1) and task 2 (T2). To accomplish each task, four sub-tasks need to be completed: S1, S2, S3, and S4, where each sub-task is semantically and visually distinct. For T1, the sub-tasks must be completed in the order of S1, S2, S3, and S4, while for T2, the order is S2, S3, S1, and S4. Based on the task progress, S1 and S2, S2 and S3, and S3 and S1 would be encouraged to become closer, which would ignore the semantics.

  This is problematic because the contrastive loss would fail to respect the semantic and structural differences between the sub-tasks, potentially causing confusion in the learned representation. For example, grouping semantically distinct sub-tasks together solely based on their task progress could degrade the model’s ability to generalize effectively to tasks that require a different order or strategy.

  I do believe that using task progress could be beneficial for the same tasks with different appearances. It makes sense to apply the proposed method for learning representations in visually different scenarios of the same task, but I am still unsure if this approach is effective for transferring across entirely different tasks. Additionally, the experiments presented in the paper to demonstrate transferability seem to focus more on in-distribution tasks rather than truly out-of-distribution tasks. 


2/ The paper requires significant improvements in writing, particularly in the following senses:

- Citation Style: Use \citep instead of \cite for parenthetical citations to ensure proper formatting.

- Figure References: Some figures, such as Fig. 1, are not referred to in the text.

- Figure Alignment: Ensure that all figures are properly aligned with the surrounding text to enhance readability.

- The paragraph titled 'Baselines for VEP' in the Related Work section appears incomplete.

- Missing Elements in Figures: In Fig. 6, the brown line mentioned in the text is missing or not clearly visible. This needs to be addressed either by correcting the figure or updating the text to accurately describe the content of the figure.

Overall, these revisions are crucial for improving the clarity, coherence, and readability of the paper.


3/ A proper discussion on the failure cases is necessary to provide better insights for the readers, particularly in relation to Fig. 3. This discussion should detail the specific scenarios where the proposed method (VEP) underperforms, the reasons behind these failures, and the potential limitations of the approach. Providing this context will help readers understand the conditions under which VEP may not be effective and inform future work on addressing these challenges.

  Additionally, I suggest tempering the claims about the method’s improvement over the baseline, especially in statements such as the last sentence of the abstract. The current language may give an impression of a substantial performance leap, whereas, in reality, the gap between VEP and the baselines is not very large in most cases. Adjusting this statement to more accurately reflect the experimental results will make the paper’s conclusions more credible and align the claims with the presented data.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: What are the significant differences between the proposed method and other SSL pertaining methods and what would make the proposed significantly better than state-of-the-art pretraining methods when the pretrain data scales up?
soundness: 3
strengths: + The paper proposes a novel way to annotate the final step of the play data as reward 1. to denote the progress of the policy, thus saving annotation cost and not requiring actions/rewards in the offline dataset;

+ Experiments on Atari and navigation datasets show that the proposed method effectively improves sample efficiency.
summary: Value Explicit Pretraining (VEP) aims to learn generalizable representations for transfer reinforcement learning, which utilizes contrastive learning to learn a transferrable encoder from offline play datasets. Experiments show that VEP achieves up to 3x improvement in sample efficiency.
weaknesses: - The novelty of the proposed method is limited, the idea of pretraining encoder with self-supervised learning methods is common both in general AI domains and in RL applications;

- From Figure 3, the improvements in Atari domain are not very significant, and it performs similarly to other self-supervised pretraining methods.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - The current submission appears to involve only simple behaviors, i.e. "Fire" in Atari and "Move & Turn" in navigation tasks. It remains unclear whether more complex and practical behaviors, such as obstacle-avoidance and path-finding, can be reflected in either testbed.
- It is also unclear whether odometry information plays a critical role in the navigation tasks. In practice, odometry readings are often noisy and unreliable in real-world scenarios. Recent work on 3D visual navigation (e.g., [1]) has emphasized learning navigation policies purely from egocentric RGB information. Could the authors provide more ablations to address this?
- About batch size. It is unclear why a fixed training batch size of $b_G\times b_T$​ is not used for a fair comparison across all baselines. Further clarification on this point would be helpful.

[1] NOLO: Navigate Only Look Once. Arxiv 2024.
soundness: 2
strengths: 1. VEP presents a novel approach by focusing on pretraining representations from actions and sparse rewards for downstream control tasks, without requiring additional fine-tuning. This has promising potential for enabling more efficient representation learning.
2. The paper offers a comprehensive evaluation of VEP on both Atari games and egocentric visual navigation tasks, demonstrating its superiority over baseline methods.
3. The paper is well-structured and easy to follow, with a clear explanation of the methodology.
summary: The paper proposes Value Explicit Pretraining (VEP), a method for learning transferable representations for downstream online reinforcement learning. VEP learns a visual encoder that is invariant to changes in environment dynamics and appearance, which is pre-trained using a self-supervised contrastive loss that relates states across different tasks based on the Bellman return Monte Carlo estimate. Extensive experiments on both Atari games and visual navigation tasks demonstrate the effectiveness. However, I have some concerns about this paper. My detailed comments are as follows.
weaknesses: 1. Originality: The novelty of the paper seems limited, as pretraining representations from video data for downstream reinforcement learning has already been well explored in existing works, especially for tasks with discrete actions. This includes methods like latent inverse dynamics [1], latent action mining with downstream fine-tuning [2], reward function learning [3], and temporal representation learning & reconstruction [4]. These approaches have also demonstrated generalization capabilities to novel or similar tasks. A more thorough comparison with these methods, along with additional ablations and visualizations, would be crucial to substantiate the claim that state contrastive learning from returns across similar tasks is a key contribution. While the proposed method in this submission is plausible, its effectiveness and novelty are not sufficiently demonstrated, particularly given that contrastive learning has been widely applied in control tasks.
2. Quality: The reported performance does not appear to significantly surpass existing methods. In in-distribution environments Demon Attack and Space Invaders, VEP performs comparably or even falls short of TCN. In Air Raid, pretrained representations even show a negative impact on RL performance compared to random representations. Could the authors give further explanation about this?
3. Significance: 
 - The current submission focuses only on short-horizon tasks with discrete action spaces. More evidence is needed to demonstrate whether VEP can be applied to other domains, such as long-horizon tasks with continuous action spaces.
 - Many large multi-modal models (e.g., vision-language models) have already pretrained semantic generalizable visual representations. Could the authors compare the quality of visual representations and final performance between VEP and these pretrained (or finetuned on limited play data) vision-language models?

[1] Learning to Act without Actions. ICLR 2024.

[2] Genie: Generative Interactive Environments. ICML 2024.

[3] Video Prediction Models as Rewards for Reinforcement Learning. NeurIPS 2023.

[4] Learning from Visual Observation via Offline Pretrained State-to-Go Transformer. NeurIPS 2023.