Policy-Induced Self-Supervision Improves Representation Finetuning in Visual RL
Authors: Séb Arnold
Fei Sha
Venue: Submitted to ICLR 2023 (2023)
Keywords: visual reinforcement learning
representation learning
transfer learning

Abstract: We study how to transfer representations pretrained on source tasks to target tasks in visual percept based RL. We analyze two popular approaches: freezing or finetuning the pretrained representations. Empirical studies on a set of popular tasks reveal several properties of pretrained representations. First, finetuning is required even when pretrained representations perfectly capture the information required to solve the target task. Second, finetuned representations improve learnability and are more robust to noise.
Third, pretrained bottom layers are task-agnostic and readily transferable to new tasks, while top layers encode task-specific information and require adaptation. Building on these insights, we propose a self-supervised objective that \emph{clusters representations according to the policy they induce}, as opposed to traditional representation similarity measures which are policy-agnostic (\eg Euclidean norm, cosine similarity). Together with freezing the bottom layers, this objective results in  significantly better representation than frozen, finetuned, and self-supervised alternatives on a wide range of benchmarks.

POLICY-INDUCED SELF-SUPERVISION IMPROVES REP- RESENTATION FINETUNING IN VISUAL RL
Anonymous authors Paper under double-blind review
ABSTRACT
We study how to transfer representations pretrained on source tasks to target tasks visual percept based RL. We analyze two popular approaches: freezing or finetuning the pretrained representations. Empirical studies on of popular tasks reveal several properties of pretrained representations. First, finetuning is required even when pretrained representations perfectly capture the information required solve the target task. Second, finetuned representations improve learnability and are more robust to noise. Third, pretrained bottom layers are aask-agnostic and readily transferable to new tasks, while top layers encode task-specific information and require adaptation. Building on these insights, we propose self- supersuee objective that clusters representations according to the policy they induce, opposed to traditional representation similarity measures which are policy-agnosie (e.g. Euclidean norm, cosine similarity). Together with freezing the bottom layers, this objective results in significantly better representation than frozen, finetuned, and self-supervised alternatives on wide range of benchmarks.
INTRODUCTION
Learning representations via pretraining is staple of modern transfer learning. Typically, a feature encoder is pretrained on one or few source task(s). Then either frozen (i.e., the encoder stays fixed) or finetuned (i.e., the parameters of the encoders are to be updated) when solving new downstream task (Yosinski et al., 2014). While whether to freeze or finetune is application-specifa finetuning outperforms in general freezing when there are sufficient (labeled) data and compute. recipe has led to many success stories in vision (Razavian et al., 201 4; Chen et al., 2021), speech (Amodei et al., 2016; Akbari et al., 2021), and NLP (Brown al., 2020; Chowdhery et al., 2022).
This eeetrain-then-tranee
For reinforcement learning (RL), however, finetuning is costly option as the learning agent needs to collect its own data specific to the downstream task. Moreover, when the source tasks are very different from the downstream task, the first few updates from finetuning destroy the representations learned on the source tasks, cancelling all potential benefits of transferring from pretraining. For those reasons, practitioners often choose to freeze representations, thus completely preventing finetuning.
But representation freezing has its own shorteomings, especially pronounced in visual RL, where the (visual) feature encoder can be pretrained on existing image datasets such as ImageNet (Deng et al., 2009) and even collections of web images. Such generic but easier-to-annotate datasets are not constructured with downstream (control) tasks in mind and the pretraining does necesarily capture important attributes used to solve those tasks. For example, the downstream embodied AI task of navigating around household items (Savva al., 2019; Kolve et al., 2017) requires knowing the precise size of the objects in the scene. Yet this information is not required when pretraining on visual objection categorization tasks, resulting in what called negative transfer where frozen representation hurts downstream performance.
More seriously, even when the (visual) representation needed for the downstream task known priori, it is unclear that learning it from the source tasks then freezing it should be preferred to finetuning, as shown in Figure On the left two plots, freezing representations (Frozen) underperforms learning representations using only downstream data (De Novo). On the right two plots, we observe the opposite outcome. Finetuning representations (Finetuned) performs well overall, but fails to unequivocally outperform freezing the rightmost plots.

---

Evaluation Tioks MSR Jume
Walke--Sun
DMC
Givoon
Hahrot
Matterport 3D
Habiaa
10
300
Fruen
Frocen
Revodts (1 bette)
Recoads (1 batee)
#30
Finstnaa
SPL bette)
Finstuned
0.6
330
Now
20-
800
0.2
Fintural
40
Upper Blard
100
Firetored
00
00
10.
3-
1.5
34
Steps
xPP
Soeps
XMP
Steps
Stapp
(a)
(b)
(c)
(d)
Figure I: When should we freeze or finetune pretrained representations in visual RL? Reward and success weighted by path length (SPL) transfer curves on MSR Jump, DeepMind Control, and Habita tat tasks. Freezing pretrained representations can underperform no pretraining at all (Figures la and 1b) or outperform it (Figures le and ld). Finetuning representations is always competitive, but fails to significantly outperform freezing visually complex domains (Figures Ic and Id). Solid lines indicate mean over 5 random seeds, shades denote 95% confidence interval. See Section 3.2 for details.
Contributions When should we freeze representations, when do they require finetuning, and why? This paper answers those questions through several empirical studies on visual RL tasks ranging from simple game and robotic tasks (Tachet des Combes et al., 2018; Tassa et al., 2018) to photo-realistic Habitat domains (Savva et al., 2019).
Our studies highlight properties of finetuned representations which improve learnability. First, they are more consistent in clustering states according to the actions they induce on the downstream task; second, they more robust to noisy state observations. Inspired by these empirical findings, we propose PiSCO, a representation finetuning objective which encourages robustness and consistency with respect to the actions they induce.
We also show that visual percept feature encoders first compute task-agnosse information and then refine this information to be task-specific (i.e., predictive of rewards and or dynamics) - well- known lesson from computer vision (Yosinski et al., 2014) but, to the best of our knowledge, never demonstrated for RL so far. We suspect that finetuning with RL destroys the task-agnostic (and readily transferrable) information found in lower layers of the feature encoder, thus cancelling the benefits of transfer. To retain this information, we show how to identify transferrable layers, and propose to freeze those layers while adapting the remaining ones with PiSCo. This combination yields excellent results on all testbeds, outperforming both representation freezing and finetuning.
2 RELATED WORKS AND BACKGROUND
Learning representations for visual reinforcement learning (RL) has come a long way since its early days. Testbeds have evolved from simple video games (Koutník et al., 2013; Mnih et al., 2015) to self-driving simulators (Shah et al., 2018; Dosovitskiy et al., 2017), realistic robotics engines (Tassa al., 2018; Makoviychuk et al., 2021), and embodied AI platforms (Savva al., 2019; Kolve et al., 2017). Visual RI. algorithms have similarly progressed, and can now match human efficiency and performance on the simpler video games (Hafner et al., 2021; Ye et al., 2021), and control complex simulated robots in a handful of hours (Yarats et al., 2022; Laskin et al., 2020; Hansen et al., 2022).
In spite of this success, visual representations remain challenging to transfer in RL (Lazaric, 2012). Prior work shows that learned representations can be surprisingly brittle, and fail to generalize to minor changes in pixel observations (Witty et al., 2021). This perspective is often studied under the unbrella of generalization (Zhang et al., 2018; Cobbe et al., 2020; Packer et al., 2018) or adversarial RL (Pinto et al., 2017; Khalifa al., 2020). In part, those issues arise due to our lack of understanding in what can be transferred and how to learn it in RI. Others have argued for a plethora of representation
pretraining objectives, ranging from capturing policy values (Lchnert et al., 2020; Liu et al., 2021) and summarizing states (Mazoure et al., 2020; Schwarzer et al., 2021; Abel et al., 2016; Littman & Sutton, 2001) to disentangled (Higgins et al., 2017) and bi-simulation metrics (Zhang et al., 2021; Castro & Precup, 2010). Those objectives can also be aggregated in the hope of learning more generic representations (Gelada et al., 2019; Yang & Nachum, 2021). ). Nonetheless, it remains unclear which of these methods should be preferred to learn generic transferrable representations.
2

---

Jumps Move right,
i 14
otherwase.
Obstacle Coordinate
Augmentation T
21
grance 1
KL kulfulre sal) ma (mn) KL (Lfrt sill at mn))
2
Data Augmentttion
T
Encodar
Projector
p2
Obstacle æ Coordinate
(b)
Figure 2: (a) Annotated observation from MSR Jump. On this simple game, the optimal policy consists of moving right until the agent (white) is 14 pixels away from the obstacle (gray), which point it should jump. (b) Diagram for our proposed PiSCo consistency objective. It encourages the policy's actions to be consistent for perturbation of state s, as measured through the actions induced by the policy. See Section 4 for details.
Those issues are exacerbated when we pretrain representations on generie vision datasets (e.g., ImageNet (Deng et al., 2009)), as in our Habi tat experiments. Most relevant to this setting are the recent works of Xiao et al. (2022) and Parisi et al. (2022): both pretrain their feature encoder on real-world images and freeze it before transfer to robotics embodied AI tasks. However, neither reports satisfying results as they both struggle to outperform a non-visual baseline built on top of proprioceptive states. This observation indicates that more work is needed to successfully pipeline to visual RL. potentially simple source of gain comes
apply the setrrannthhen-taaee finetuning those generic representations the downstream task, as suggested by Yamada et al. (2022). While this is computationally expensive, prior work has hinted at the utility of finetuning in visual RL (Schwarzer et al., 2021; Wang et al., 2022).
In this work, we zero-in on representation freezing and finetuning in visual RL. Unlike prior represen- tation learning work, our goal is not to learn generic representations that transfer: rather, we would like to specialize those representations to specific downstream task. We contribute in Section an extensive empirical analysis highlighting important properties of finetuning, and take the view that finetuning should emphasize features that easily discriminate between good and bad action on the downstream task - in other words, finetuning should ease decision making. Building on this analysis, we propose novel method to simplify finetuning, and demonstrate its effectiveness on challenging visual RL tasks where naive finetuning fails to improve upon representation freezing.
3 UNDERSTANDING WHEN TO FREEZE AND WHEN TO FINETUNE
We presents empricial analysis of representation freezing and finetuning. In Sections 3.1 and 3.2 we detail the results from Figure I mentioned above. Then, we show in Section 3.3 a surprising result on the MSR Jump tasks: although freezing the pretrained representations results in negative transfer, those representations are sufficiently informative to perfectly solve the downstream task. This result indicates that capturing the relevant visual information through pretraining is not always effective RL transfer. Instead, Section 3.4 shows that representations need to emphasize task-relevant information to improve learnability insight we explore in Section 3.5 and build upon in Section 4 to motivate simple and effective approach finetuning.
3.1 EXPERIMENTAL TESTBEDS
Our testbeds are standard and thoroughly studied RL tasks in literature. Our description here focuses on using study RL transfer. We defer details and implementations to Appendix B.
MSR Jump (Tachet des Combes et al., 2018) The agent (a gray block) needs to cross the screen from left to right by jumping over an obstacle (a white block). The agent's observations are video game frames displaying the entire world state (see Figure 2a) based on which can choose to "move right" or "jump". We generate 1 140 tasks by changing the position of the obstacle vertically and horizontally; half them are used for pretraining, half for evaluation. We train the convolutional actor-critic agent from Mnih et al. (2016) with PPO (Schulman et al., 2017) for 500 iterations on all pretraining tasks, and transfer its layer feature encoder to jointly solve all evaluation tasks. Like pretraining, we train the actor and critic heads of the agent with PPO for 500 iterations during transfer.
3

---

Evaluation Tasks
MSR
Evaluation Tasks
MSR
Evaluation Tasks
MSR Jump
Distane Estimation Error (L [eette)
4.3
Oszification Error (1 better)
10
(1 (bbtte)
0.8
Cluster Puriyy (1 better
06
0.5
0.6
20-2
0.4
10-3
Random Franen
022
0.2
1017
Firetuned
00
0.3
0.
From
Fineturdd
Fuzeen
Fineuneed
Random
Frzee
Fintuned
10
Nase Magnitude
(c)
Figure 3: (a) Transfer can fail even though frozen representations are informative enough. On MSR Jump tasks, we can perfectly regress the optimal actions (accuracy = 1.0) and agent-obstacle distance (mean square error 1.0) with frozen or finetuned representations. Combined with Figure la. those results indicate that capturing the right information not enough for RL transfer. (b) Finetuned representations yield purer clusters. Given a state, we measure
expected purity of the cluster consisting the closest neighbours of that state in representation space. For Finetuned representations, this metric is significantly higher (98. 75%) than for Frozen (91. (91.1%) or Random (82.98% showing that states which beget the same actions are closer to one another and thus easier to learn. (c) Finetuned representations are more robust to perturbations. For source and downstream tasks, the classification
(1 action accuracy) degrades significantly Fine tuned representations than Frozen ones under increasing data augmentation noise, suggesting that robustness to noise improves learnability. See Section 3.3 for details.
slowly for
DeepMind Control (DMC) (Tassa et al., 2018) This robotics suite is based on the MuJoCo (Todorov et al., 2012) physics engine. We use visual observations as described by Yarats et al. (2021), and closely replicate their training setup. The learning agent consists of action-value and policy heads, together with shared convolutional feature encoder whose representations image observations. We pretrain with DrQ-v2 (Yarats et al., 2022) on a single task from given domain, and transfer only the feature encoder to different task from the same domain. For example, our Walker agent is pretrained on Walker-Wa wa lk and transfers to Walker-Run.
used in lieu of the
Note: The main text Wa laee to illustrate results DeepMind Control, and more tasks included in Appendix C (including Cartpole salance-SSiinge and Hopper Stand-Hop).
Habitat (Savva al., 2019) In this setting, we pretrain ConvNeXt (tiny) (Liu et al., 2022) to classify images from the ImageNet (Deng et al., 2009) dataset. We use the "trunk" of the Xt as a feature encoder (i.e., discard the classification head) and transfer it to either Gibson (Shen et al., 2021) or Matterport3D (Chang et al., 2017) scenes simulated with the Habitat embodied AI simulator. The tasks in Habitat consist of navigating to point from an indoor scene (i.e., PointNav) from visual observations, aided with GPS and compass sensing. Our implementation replaces the residual network from "Habitat-Lab" (Szot al., 2021) with our ConvNeXt' and uses the provided DDPPO (Wijmans et al., 2020) implementation for transfer training. We train for 5M steps on Gibson and 7M steps on Matterport3D. include all scenes from the respective datasets, and report reward weighted by path length (SPL).
3.2 WHEN FREEZING FINETUNING WORKS
As first experiment, we take the pretrained feature encoders from the previous section and transfer them to their respective downstream tasks. We consider three scenarios.
De Novo: where the feature encoder randomly initialized (not pretrained) and its repre- sentations are learned on the downstream task only. The feature encoder might use different
hyper-parameters (e.g., learning rate, momentum, etc) from the rest of the agent.
Frozen: where the feature encoder is pretrained and its representations are frozen (i.e., never update the weights of the feature encoder).
Finetuned: where the feature encoder is pretrained and its representations (i.e.. we update the weights of the feature encoder using the data and the algorithm of the downstream task). As for De Novo, the apprr-parametess for the feature encoder are tuned separately from the other hyper- parameters of the model.
finetuned
We choose Cony because it removes hatch normalization and dropout layers after pretraining.
4

---

Figure displays those results. Comparing De Novo and Frozen shows that freezing representations works well on Habitat (see Figures Ic and Id), likely because these tasks demand significantly richer visual observations. However, freezing representations completely fails the simpler MSR Jump and DeepMind Control tasks (see Figures 1a and lb). The pretrained feature encoders don't t provide any benefit over randomly initialized ones, and even significantly hurt transfer MSR Jump. The latter is especially surprising because there is very little difference between pretraining and evaluation tasks (only the obstacle position is changed). Those results beg the question: why does transfer fail so badly on MSR Jump?
On the other hand, Finetuned performs reasonably well across all benchmarks but those results hide an important caveat: finetuning can sometimes collapse to representation freezing. This was the case on Hab tat, where our typer-parameter search resulted in tiny learning rate for the feature encoder (Ie-6, the smallest in our grid-search), effectively preventing the encoder to adapt. We hypothesize that on those large-scale and visually challenging tasks, Finetuned offers no improvement over Frozen because RL updates are too noisy for finetuning. Thus, we ask: how we help stabilize finetuming in RL?
We answer both questions in the remaining of this paper. We take a closer look at MSR Jump tasks and explain why transfer failed. We show that a core role of finetuning is to emphasize task- specific information, and that layers which compute task-agnostic information can be frozen without degrading transfer performance. In turn, this suggests an instinctive strategy to help stabilize RL training: freezing task-agnostic layers, and only finetuning task-specific ones.
3.3 FREEZING FAILS EVEN WHEN LEARNED REPRESENTATIONS ARE USBFUL
We now show that, on MSR Jump, the pretrained representations are informative enough to perfectly solve all evaluation tasks.
Due to the simplicity of the dynamics, we can manually devise the following optimal policy because the size of the agent, the size of the obstacle, and the jump height are shared across all tasks:
Jump, Move right,
if the agent is 14 pixels to the left of the obstacle. otherwise.
π
We empirically confirmed this policy reaches the maximal rewards (81) on all pretraining and evalua- tion tasks. Immediately, we see that the distance from the agent to the obstacle is the only information that matters to solve MSR Jump tasks. This observation suggests the following experiment.
To measure the informativeness of pretrained representations, we regress the distance between the agent and the obstacle from pretrained representations. If the estimation error (i.e., mean square error) is lower than 1.0 pixel, the representations are informative enough to perfectly solve MSR Jump tasks. To further validate this scenario, we also regress the optimal actions from the simple policy above given pretrained representations. Again, if the acurracy is perfect the representations are good enough. We repeat this experiment for finetuned representations, and compare the results against one another.
Figure 3a reports results for pretrained and finetuned (i.e., after transfer) representations. We regress the distance and optimal actions with linear models from evaluation task observations, using a mean square error and binary cross-entropy loss. Surprisingly, both pretrained and finetuned representations can perfectly regress the distance and optimal actions: they get sub-pixel estimation error (0.94 and 0.79, respectively) and perfect accuracy (both 100%). Combined with the previous section, these results indicate that capturing the right visual information is not enough for successful transfer in RL. They also put forward the following hypothesis: one role of finetuning is to emphasize information that eases decision making. This hypothesis is supported by the lower distance error for Finetuned in the right panel of Figure 3a.
3.4 FINETUNING IMPROVES LEARNABILITY AND ROBUSTNESS TO NOISE
How is information refined for decision making? Our next set of experiments answers this question by highlighting differences between pretrained and finetuned representations.
First, we measure the ability of each feature encoder to cluster states assigned to identical actions together. Intuitively, expect states that require similar actions to have similar representations.
5

---

Evaliation Tasks
kimp
Walker-Run
DMC
Foshiation Tasks
Jump
Waller-Ran
DMC
Valce Ediration Eror E1 bente)
Value Edimation Ereer El bente)
Area Undaa Reward Care [1 Sette]
Area Under Roward Cave Sette)
9.8
30-2
0.6
21
Bantoon
Fictured
Firtured
Olter
Using Ape
Wher Using sye
Layers
When ayee
Figure 4: Frozen layers that retain the same information finetuned layers can be kept frozen. (a) Mean squared error when linearly regressing action values from representations computed by a given layer. Some early layers are equally good at predicting action values before or after finetuning. This suggests those layers can be kept frozen, polentially stabilizing training. See Section 3.5 for details. (b) Area under the reward curve when freezing up given layer and finetuning the rest. Freezing early layers does not degrade performance, and sometimes improves it (e.g., up Conv2 on DeepMind Control).
To that end, we compute the average purity when clustering a state with its closest neighbors in representation space. In other words, given state s, find the 5 closest neighbors to in representation space and count how many of them are assigned the same optimal action as s. Figure 3b the average cluster purity for random, pretrained, and finetuned representations. We find that Finetuned representations yield clusters with significantly higher purity (98. 75%) than Frozen (91.41%) or random (82. (82.98%) ones, which explains why PPO converges more quickly when representations are finetuned: it is easier to discriminate between the states requiring "jump" v.s. "move right", the representations for those states cluster homogenously.
Second, we uncover another property of finetuned representations, namely, that they are more robust to perturbations. We replicate the optimal action classification setup from Figure 3a but perturb states randomly rotating them. For each image, the degree of the rotation is randomly sampled from range (-x, 27. Figure shows how classification accuracy degrades as we let æ increase from 0 to 12 Both Frozen and Finetuned degrade with more noise, but Finetuned is more robust and degrades more slowly.
Those results point to finetuned representations that consistently map similar actions and are robust to perturbations They set the stage for Section 4 where we introduce PiSCo, a policy consistency objective which builds these insights. In the next section, continue analysis and further investigate how information gets refined in each layer of the feature encoder.
3.5 WHEN AND WHY IS REPRESENTATION FINETUNING REQUIRED?
As the previous section uncovered important pitfalls of representation freezing, finetuning in the hope of understanding why it succeeds when it does. To that end, we dive into the representations of individual layers in Fine feature encoders and compare them to the Frozen and random representations. We take inspiration from the previous section and hypothesize that "purpose drives adaptation"; that is, a layer only needs finetuning if it computes seee-sccc information that eases decision making on the downstream task. Conversely, layers that compute task- agnostic information can be frozen thus potentially stabilizing RL training. To this hypothesis, we conduct two by-er-byllyyer experiments: linear probing and incremental freezing.
turn to
With layer-by- by-leyer linear probing, we aim measure how relevant each layer is when choosing an action (Ding et al., 2021; Zhang et al., 2022). To start, we collect a dataset of 10, 000 state-action pairs and their corresponding action values with the Finetuned policy on the downstream task. The action values are computed with an action value function trained to minimize the temporal difference error on the collected state-action pairs (more details in Appendix C). We then proceed as follows for each layer/ / in the feature encoder. First, we store the representations computed by layer I, and extract their most salient features by projecting them to a -0-dimenniooal vector with PCA. Second, regress action values from the projected representations and their corresponding actions using linear regression model. Finally, we report the mean squared error for this linear model at layer l. Intuitively, the lower the linear probing error the easier it is predict how good an action is for

---

given state". We repeat this entire process for Frozen and Finetuned feature encoders, as well as one with a randomly initialized weights (i.e., Random).
Results for this experiment shown in Figure 4a. curves trend downwards and lie significantly lower than Random, indicating, as expected, that representations specialize for decision making that the first few layers of Frozen and Finetuned track each other closely before Frozen starts to stagnate, On MSR Jump, this happens after the last convolutional layer (Conv3) and on DeepMing nd Control after the second one (Conv2). This evidence further supports our previous insight that early layers compute task -agsostic information, and suggests a new hypothesis: can we freeze the early layers without degrading finetuning performance? If so, this strategy might help stabilize finetuning in visually complex tasks like Habí tat.
On all testbeds, both Frozen and Finetuned
move up the feature encoder. We also notice
We confirm this hypothesis by combining freezing and finetuning in our second layer-by-layer experiment. For each target layer 1, we take the pretrained feature encoder and freeze all layers up to and including &; the remaning layers are finetuned using the same setup as Figure 1.
Figure 4b summarizes this experiment with the (normalized) area under the reward curve (AURC). We preferred this metric over "highest reward achieved" since the latter does not consider how quickly the agent learns. On MSR Jump, the pretrained layers that have similar value estimation error in Figure 4a can be frozen without degrading adaptation. But, freezing yields lower AURC when value estimation error stagnates (as for Línear). Similarly, freezzing the last two layers on DeepMind Control (Conv3 and Conv4, which did not match Finetuned's value estimation error) also degrades performance. We also see that adapting too many layers (iee, when Conv1 is frozen but not Conv2) reduces the AURC. The training curves show this is due slower convergence, suggesting that Conv2 already computes useful representations which can be readily used for the downstream task.
Merging the observations from our layer-by-layer experiments, we conclude that pretrained layers which extract cask-agnostic information can be frozen. We show in Section 5 that this conclusion significantly helps stabilize training when finetuning struggles, as in Habitat.
4 FINETUNING WITH A POLICY-INDUCED SELF-SUPERVISED OBJECTIVE
Section 3.4 suggests that important outcome of finetuning is that states which are assigned the same actions cluster together. This section builds on this insight and introduces P PiSCo, a self-supervised objective which attempts accelerate the discovery of representations which cluster together for similar actions. At high level PiSCO works as follows: given a policy æ, it ensures that π T(a 81) and K (a $2) are similar if states $1 and $2 should yield similar actions (e.g., $1, $2 are perturbed version of state s). This objective is self-supervised because it applies to any policy π (not just the optimal one) and doesn't 't require knowledge of rewards nor dynamics.
More formally, we assume access to a batch of states B from which we can sample state Further, we compute two embedding representations for state s. The first, z = ((s). obtained by passing through encoder <)-); the second, p = lu(z), applies projector le(-) top of representation z. Both represenations and have identical dimensionality and can thus be used to condition the policy æ. The core contribution behind PiSCO is to measure the dissimilarity between 2 and p in terms of the distribution they induce through ..:
D(z,p) KL (1 (L( 2)) || [[- [p)),
where KL(-II-) is the Kullback-Leibler divergence, and (x) is the stop-gradient operator, which sets 0. The choice of the Kullback-Leibler is arbitrary other
all partial derívatives with respect statistical divergence are valid alternatives.
The final ingredient in PNO a distribution T(s') s) over perturbed states s'. This distribution is typically implemented by randomly applying set of data augmentation transformations on the state Then, the PiSCO objective (for Policy-induced Self-Consistency Objective) is to minimize the
We found action values to be better targets to measure case of decision making than, say, action accuracy which is upper bounded by 1.0 as in Figure 3a.
7

---

Eveluation Tasks
Walker-Run
DMC
Gibson
Matterpont50
Haliber
001
10
0.5
Rowndb ["bott]]
70
Reoordh better)
200
batter
0.8
M.b.
802
807
0.4
Frezen
Froeet
Froucn Findtaned
202
Festundl
Fiectared
Fiazsuncd
Frucs RISCO
Frucx Fraciane)
0.2
Frozen Findunnal
Frocent Fintunaa
Froze PSC0
Frozes PISCO
FrocersPiECO
XIP
X10
Steps
X30
Stops
XEP
Figure 5: Partial freezing improves convergence; adding our policy-induced consistency objective im- proves further. As suggested by Section 3.5, we freeze the early layers of the feature extractor and finetune the rest of the parameters without (Frozen-Finetunee PPCCO). On challenging tasks (e-g-, Habitat), partial freezing dramatically boosts downstream performance, while Frozen+PiSCo further improves upon Frozen+ F Finetuned across the hoard.
or with policy-induced consistency objective
dissimilarity between the representation of states sampled from T:
1 2
(D(z,,p2) + D(22.Pi))
E 8
-1.82/-((-1s)
where 81 and 82 two different perturbations obtained from state 8, and the objective uses symmetrized version of the dissimilarity measure D. In practice, the PiSCo objective is added as an auxiliary term to the underlying RL objective and optimized for the encoder and the projector h. Pseudocode is available in Appendix D.
Remarks PiSCO is reminiscent of SimSiam (Chen & He, 2020) and only differs in how dissimilarity between embeddings is measured. Our proposed policy-induced similarity measure crucial for best performance, shown in Section 5.2. In fact, similar representation learning algorithms for RL could be derived by replacing the embedding similarity measure in existing self-supervised algorithms, such SimCLR (Chen et al., 2020), MoCo (He et al., 2019), VICReg (Bardes et al., 2022). We chose SimSiam for its simplicity no target encoder nor negative samples required - which suits the requirements of RL training.
Alone, PiSCO is not a useful objective to solve RL tasks; rather, its utility stems from assisting the underlying RL algorithm in learning representations that are robust to perturbations. Could we obtain the same benefits by learning the policy with augmented data (and thus side-stepping the need for Pisco)? In principle, yes. However, the policy objective typically a function of the rewards which is known to be excessively noisy to robustly learn representations. For example, DrQ avoids learning representations through the policy objective, and instead relies on Bellman error minimization. We hypothesize that PiSCO succeeds in learning robust representations because its self-supervised objective is less noisy than the policy objective.
EXPERIMENTS
We complete our analysis with closer look at PiSCO and partially frozen feature encoders. First, we check whether partial freezing and policy-induced supervision can help accelerate finetuning with RL; second, we compares policy-induced supervision as an RI. finetuning objective against more common similarity measures in self-supervised learning.
5.1 PARTIAL FREEZING AND POLICY-INDUCED SUPERVISION IMPROVE RL FINETUNING
This section shows that partially freezing a pretrained feature encoder can significantly help stabilize downstream training. We revist the experimental setup of Figure 1, this time including two new transfer variations. The first one is Frozen+Finetuned, where we freeze early layers and finetune the remaining ones with Finetuned. The second, Frozen+ additionally includes the SCC objective as an auxiliary loss to help finetuning representations. Frozen+PiSe involves one extra syper-parameter. namely, the auxiliary loss weight, which we to 0.01 for all benchmarks.
8

---

Waller--uu
DMC
Hepper-Hip
DMC
Figure 6: Policy-induced supervision is better ob- jective than representation alignment for finetuning. We compare our policy-induced consistency objective as measure of representation similarity against pop- ular representation learning objectives (e.g. SimSiam, CURL). Policy supervision provides more useful similar- ities for RL finetuning, which accelerates convergence and reaches higher rewards.
Reveets better)
Roceeds (1 bedte)
600
150
1101
(00)
Frouert San Frozess RSCO
Fruust SinSaa
110
0.5
10 Steps
1.5
20
0.5
18
1.5
Steps
In both cases, we identify which layers to freeze and which to finetune following similar linear probing setup as for Figure 4a. For each layer, we measure the action value estimation error of pretrained and finetuned representations, and only freeze the first pretrained layers that closely match the finetuned ones, On MSR Jump, we freeze up to Conv3; on DeepMing ind Control (Wa (Walker) up to Conv2; and on Habi we freeze upto Conv8 for Gibson and upto Conv7 for Matterport3D. (See the Appendix C for results on Habitat.)
We report convergence curves in Figure 5. As suggested by our analysis of Figure 4b, freezing those early layers does not degrade performance; in fact, we see significant gains on Habitat both in terms of convergence rate and in asymptotic rewards. Those results are particularly noticeable since Finetunnd struggled to outperform Frozen on those tasks. We also note that Frozen+Pi SCO improves upon Frozen+ Finetuned across the board: it accelerates convergence on MSR Jump, and also improves asymptotic performance on DeepMin nd Control and Habi tat. These results tie all our analyses together: they show that (a) freezing early layers can help stabilize transfer in visual RL, (b) which layer to freeze is predicted by how well it encodes action value information, and (c) policy-induced supervision (and p SCO) is a useful objective to accelerate RL. finetuning.
5.2 POLICY-INDUCED SUPERVISION IMPROVES UPON CONTRASTIVE PREDICTIVE CODING
As final experiment, answer how important it is to evaluate similarity through the policy rather than with representation alignment as in contrastive predictive coding. In other words, could we swap PiSCo for SimSiam and oblain similar resulls in Figure 5? Our ablation studies focus on DeepMind Control as it is more challenging than MSR Jump yet computationally tractable, unlike labitat.
We include Frozen+SimSian and Frozen+CURL in addition to Frozen+ Finetuned and Frozen+PiSCo. Frozen+ SimSiam is akin to Frozen+PiSCO but uses the negative cosine sim- ilarity to measure the dissimilarity between embeddings z and p. For Frozen+ we implement the contrastive auxiliary loss of CURL (Laskin et al., 2020), self-supervised method specifically designed for RL. Both of these alternatives use DrQ as the underlying RL algorithm.
Figure 6 reports convergence curves on Walker Walk-> Run and Hopper Stand->Hop' Including the SimSiam objective typically improves upon vanilla finetuning, but including CURL does not. We hypothesize that most of the benefits from CURL are already captured by DrQ, as shown in prior work (Yarats et al., 2021). PiSCO significantly outperforms the standard self-supervised methods thus demonstrating the efficacy of using the policy when measuring similarity between embeddings. We hypothesize that representation alignment objectives naturally map two similarly looking states to similar representations, even when they shouldn't 't (e.g., when the agent is 15 and 14 pixels away from the obstacle in MSR Jump). Instead, PiSCO is free to assign different representations since those two states might induce different actions ("move right" and "jump", respectively).
6 CONCLUSION
Our analysis of representation learning for transfer reveals several new insights on the roles of finetuning: similarity between representations should reflect whether they induce the same similar distribution of actions; not all layers in encoders of states need to be adapted. Building on those insights, we develop a hybrid approach which partially freezes the bottom (and readily transferrable) layers while finetuning the top layers with policy- self-supervised objective. This approach is especially effective for hard downstream tasks, it alleviates the challenges of finetuning rich visual representations with reinforcement learning.
Hopper-HHop is especially difficult; egg, DrQ-v2 fails solve it after 3M iterations (Yarats et al., 2022).

---

REFERENCES
David Abel, David Hershkowitz, and Michael Littman. Near optimal behavior via approximate state abstraction. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 2915-2923, New York, New York, USA, 20-22 Jun 2016. PMLR. URL httts:://proceedings/ mlr. paress/v48// ///////////bbeeeee 116.11eel11.111 html. Cited on p. 2.
Hassan Akbari, Liangzhe Yuan, Rui Qian, Wei-Hong Chuang, Shih-Fu Chang, Yin Cui, and Boqing Gong. VATT: Transformers for multimodal self-supervised learning from raw video, audio and text. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: / wwwppppeevvieveeeeee net/ fort//ormm?tt dd= RzYrn625bu8. Cited on p. 1.
Dario Amodei, Sundaram Ananthanarayanan, Rishita Anubhai, Jingliang Bai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Qiang Cheng, Guoliang Chen, Jie Chen, Jingdong Chen, Zhijie Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Ke Ding, Niandong Du, Erich Elsen, Jesse Engel, Weiwei Fang, Linxi Fan, Christopher Fougner, Liang Gao, Caixia Gong, Awni Hannun, Tony Han, Lappi Vaino Johannes, Bing Jiang, Cai Ju, Billy Jun, Patrick LeGresley, Libby Lin, Junjie Liu, Yang Liu, Weigao Li, Xiangang Li, Dongpeng Ma, Sharan Narang, Andrew Ng, Sherjil Ozair, Yiping Peng. Ryan Prenger, Sheng Qian, Zongfeng Quan, Jonathan Raiman, Vinay Rao, Sanjeev Satheesh, David Seetapun, Shubho Sengupta, Kavya Srinet, Anuroop Sriram, Haiyuan Tang, Liliang Tang. Chong Wang, Jidong Wang, Kaifu Wang, Yi Wang. Zhijian Wang, Zhiqian Wang. Shuang Wu, Likai Wei, Bo Xiao. Wen Xie, Yan Xie, Dani Yogatama, Bin Yuan, Jun Zhan, and Zhenyao Zhu. Deep speech 2: End-to-end speech recognition in english and mandarin. In Proceedings of the 33rd International Conference on International Conference on Machine Learning- Volume 48, ICML'16, pp. 173-182. JMLR.org, 2016. Cited on p. 1.
Ba, Kiros, and Hinton. Layer normalization. arXiv preprint 1607 06450, 2016. URL http:/ htap://aaxxtt www. /66.. /111//6677766..... Cited on p. 17.
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Yarianceiivaarance--vvaaaaa regularization for self-supervised learning. In International Conference on Learning Representations, 2022. URL
Cited on p. 8.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler. Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL /procee/////cceee/eeee www.eeeddddnnnnnnisssseoee neurips.cc/ p2per/2020/ fiile//ee////
Paper pdd. Cited on p. 1.
Pablo Samuel Castro and Doina Precup. Using bisimulation for policy transfer in MDPs. In Twenty- Fourth AAAI Conference on Artificial Intelligence, July 2010. URL https:////// hopp:://www... www.aaati .org/
ocs/ fess/indxx.. php/ /hpp/AAII/AA /AAAIAAAAII paper wwtp:/viwwwrrreeee 1197. Cited on 2.
A. Chang, A. Dai, T. Funkhouser, M. Halber, M. Niebner, M. Savva, S. Song, A. Zeng, and Y. Zhang, Matterport3d: Learning from rgb-d data in indoor environments. In 2017 International Conference on 3D Vision (3DV). pp. 667-676, Los Alamitos, CA, USA, oct 2017. IEEE Computer Society. URL https:://of teeecomputersociety Cited on p. 4, 18.
doi: 10. 10.1109//VV.001.000088 3DV. 2017. 00081.
1109/
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé li and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020. URL https: http///pooco/i/eeee/ iggs mress/ chen2j html. Cited on p. 8.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. November 2020.
URL http:://axxvv.org// http///bb/ee....../ 2011.1111111111 10566. Cited on p. 8, 21, 22.
10

---

Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pp. 9640-9649, 2021. URL https:/ / github. com facebookresearchh/monn
v3. Cited on p. 1.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer. Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontack Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Ceier-Hellstem, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Ma Scaling language modeling with pathways. April 2022. URL http: / ar xvv or g/ abt/a/ss//222 62311. Cited on
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman, Leveraging procedural generation to benchmark reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2048-2056 2056. PMLR, 13-18 Jul 2020. URL /propee//proceed ing. mmlr. press/v119 cowbe20a..b.ecoooooooooo html. Cited on p. 2.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255. June 2009. URL http: on p. 1,3, 4, 17.
org/10. 1109/ / CVPR. 2009 5206848. Cited
Frances Ding, Jean -tan-SSanislas Denain, and Jacob Steinhardt. Grounding representation similarity through statistical testing. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: //pp:////peeevweeeee Cited on p.
net/ met/forum?ion ddw
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA:
An open urban driving simulator. In Sergey Levine, Vincent Vanhoucke, and Ken Goldberg (eds.), Proceedings of the 1st Annual Conference on Robot Learning. volume 78 Proceedings of Machine Learning Research, pp. 1-16. PMLR, 13-15 Nov 2017. URL https: /ttpss////oooeeeeeeee mlr. P..r.press/vi (o/////oovvii wtwiiiiskiiiii html. Cited on p. 2.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume of Proceedings of Machine Learning Research, pp. 2170-2179, PMLR, 2019. URL http: press/ htep://9/ggeeeee lada19a. html. Cited on p.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL n wrtrruumdd-BoabbyZb000 Cited on p. 2.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv [cs.LG], March 2022. URL http::/aarxiv.org/abb.... 0495. Cited on p.
2.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. November 2019. URL http::t/a/axxxi.rrgg/////// 1911. 05722. Cited on p.
D Hendrycks, K Gimpel arXiv preprint arXiv: 100.::6666...... 606108415, and 2016. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. arxivorrg 2016. URL https: /arxiv. org/abs /106. 08415. Cited on p. 16, 17.
11

---

Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference Machine Learning Volume 70, ICML'17, pp. 1480-1490. Jorg. 2017. Cited on p. 2.
Ahmed Khalifa, Philip Bontrager. Sam Earle, and Julian Togelius. Pegrl: Procedural content generation via reinforcement learning. In Proceedings of the Sixteenth AAAI Conference Artificial Intelligence and Interactive Digital Entertainment, AlIDE' AAAI Press, 2020. ISBN 978-1-57735-849--7 Cited on p. 2.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. December 2014.
URL http: http:///aaxavao/eeaeeeeo /112../11.... 1412.6980. Cited on p. 17.
Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt. Luca Weihs, Alvaro Herrasti, Daniel Gordon, Yuke Zhu, Abhinav Gupta, and Ali Farhadi. A12-THOR: An Interactive 3D Environment for Visual AI. arXiv, 2017. Cited on p. 1,
Jan Koutník, Giuseppe Cuccu, Jürgen Schmidhuber, and Faustino Gomez. Evolving large-scale neural networks for vision- based reinforcement learning. In Proceedings of the 15th annual conference on Genetic and evolutionary computation, GECCO '13, pp. 1061-1066 New York, NY, USA, July 2013. Association for Computing Machinery. URL https: doii org.orgg/o/o. 10. 1145/ 2463372. 2463569. Cited on p. 2.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised repre- sentations for reinforcement learning. In Hal Daumé IIII and Aarti Singh (eds.), Proceed- ings of the 37th International Conference on Machine Learning, volume 119 of Proceed- ings of Machine Learning Research, pp. 5639-5650. PMIR, 13-18 Jul 2020. URL https: html. Cited on p. 2.
//proceedings.mit wtw.press/119//assssssss
Alessandro Lazaric. Transfer in reinforcement learning: A framework and a survey. In Mareo Wiering and Martijn van Otterlo (eds.), Reinforcement Learning: State-of-the-Art, pp. 143-173. Springer Berlin Heidelberg, Berlin, Heidelberg. 2012. URL https://doi.orgg 1007/
978-3-642-27645-3 Cited on p. 2.
Lucas Lehnert, Michael L Littman, and Michael Frank. Reward preeurctrreeiiivveeee representations generalize across tasks reinforcement learning. PLoS Comput. Biol., 16(10);e 1008317, October 2020. URL http: //t:://xxxxx doi. or g 1371/journal pcbi. 1008317. Cited on 2.
Michael Littman and Richard S Sutton. Predictive representations of state. Adv Neural Inf. Process. Syst., 14, 2001. URL https://webs eecs. wwwch. hwwp:ichhhdubbbbbejj///paa ppr. pwtpppf.....ppp.p pdff Cited
on p. 2.
Guoqing Liu, Chuheng Zhang, Li Zhao, Tao Qin, Jinhua Zhu, Li Jian, Nenghai Yu, and Tie-Yan Liu. Return-based contrastive representation learning for reinforcement learning. In International Conference on Learning Representations, 2021. URL https: /wwp///veeeee...coooooo/oo net/ forum? id= _TM6rT7tXke. Cited on p. 2.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022. Cited p. 17.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. Isaac gym: High performance GPU-Based physics simulation for robot learning. August 2021. URL http: //arxiv.org/abee www.oo/g/////////oooo//oo 2108. 10470. Cited on p. 2.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and De- von Hjelm. Deep reinforeement and infomax learning. In NeurIPS 2020. ACM, De- cember 2020. URL https:: : mi crosoft. wom/en-us///eeeeooo /teeeu//rreeeacchppultttttttte Cited on p.
12

---

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518 (7540):529-553) February 2015. URL http: /ttt:://dx.co doi on p. 2, 16.
1038 wtt../0///////te11115500 Cited
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference machine learning, pp. 1928-19377 2016. Cited on p. 3.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, V Koltun, and D Song. Assessing gener- alization deep reinforcement learning. ArXiv, 2018. URL https: /www. semant icschola Cited on p. 2.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effecttivnness of Pre-Trained vision models for control. March 2022. URL
abs/2203. 03580. Cited on p.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorrc:: imperative style, High-Performance deep learning library. In H Wallach, H Larochelle, A Beygelzimer, F d' Alche Buc, E Fox, and R Garnett (eds.), Advances Neural Information Processing Systems 32, pp. 8026-8037. Curran Associates, Inc., 2019, URL -eeeaaineeede--eeeennnn--irrrrrr 2o15-pytorch ch- an-cch------- an- pdf. Cited p. 24.
Lerrel Pinto, James Davidson, Rahul Sukthankar, and Abhinav Gupta. Robust adversarial reinforce- ment learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp. 2817-2826. PMLR, 06-11 Aug 2017. URL https:/ https:/pproceciinggss pintol7a html. Cited on p. 2.
Ali Sharif Razavian, Hossein Azizpour, Josephine Sullivan, and Stefan Carlsson. Cnn features off fff-the-shelf: An astounding baseline for recognition. In 2014 IEEE Conference Computer Vision and Pattern Recognition Workshops, pp. 512-519, 2014. doi: 10. 109/CVPRW.2014.111
Cited on p. 1.
Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, Devi Parikh, and Dhruv Batra. Habitat: A Platform for Embodied AI Research. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), 2019. Cited on p. 1,2, 4, 17.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. July 2017. URL htw.ttt://axxxi.or//aaaaa wtt.or//////////.//.ooooooe 3, 16, 23.
https//177..0000 06347. Cited on p.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations, 2021. URL hittps://opents htwp:://e///eeeeee/eeeeeee wwwpe/rrwiiieeennefrrreeeee mottfforum?d== uCQfPZwRaUu. Cited on p. 2, 22.
Shital Shah, Debadeepta Dey, Chris Lovett, and Ashish Kapoor. Airsim: High-fidelity visual and physical simulation for autonomous vehicles. In Marco Hutter and Roland Siegwart (eds.). Field and Service Robotics, pp. 621-635, Cham, 2018. Springer International Publishing. ISBN 978-3-319-6736-1 Cited on p. 2.
Bokui Shen, Fei Xia, Chengshu Li, Roberto Martín-Martin Linxi Fan, Guanzhi Wang, Claudia Pérez-D's DérrzzD'Arpino, Shyamal Buch, Samjana Srivastava, Lyne Tchapmi, Micael Tchapmi, Kent Vainio, Josiah Wong, Li Fei-Fei, and Silvio Savarese. igibson 1.0: A simulation environment for interactive
13

---

tasks in large realistic scenes. In 202/ JEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 7520-7527, 2021. doi: 10.1109/RROSS1666 10211110.22.166611 Cited on 4.
Andrew Szot, Alex Clegg, Eric Undersander, Erik Wijmans, Yili Zhao, John Turner, Noah Maestre, Mustafa Mukadam, Devendra Chaplot, Oleksandr Maksymets, Aaron Gokaslan, Vladimir Vondrus, Sameer Dharur, Franziska Meier, Wojciech Galuba, Angel Chang, Zsolt Kira, Vladlen Koltun, Jitendra Malik, Manolis Savva, and Dhruv Batra. Habitat 2.0: Training home assistants rearrange their habitat. In Advances in Neural Information Processing Systems (NeurIPS), 2021. Cited on p. 17.
Remi Tachet des Combes, Philip Bachman, and Harm van Seijen. Learning invariances for policy generalization, 2018. URL https: http:::://///peevieeeeee wwwpeerwwwweeeeeeeoeeeeeeee net/forum? wett/oorumidd-BHHHBB... id-BDHRaK1PG. Cited on p. 2, 3, 16.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. Timothy Lillicrap, and Martin Riedmiller. Deep.Mind control suite. January 2018. URL http:://axiv...// http://xxvvvrrreee abs/ 00690. Cited on p. 2, 4, 17.
Emanuel Todorov, Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In 2012 JEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, October 2012. URL http: /tpp//xxoooo/oo wrg.org///oooo 1109 IROS. 6386109.. Cited on p. 4. 17.
Che Wang, Xufang Luo, Keith Ross, and Dongsheng Li. VRL3: data-driven framework for visual deep reinforcement learning. arXiv CVJ, February 2022. URL http:://rxxvv.rrg abs/ 2202. 10324. Cited on p. 3.
Erik Wijmans, Abhishek Kadian, Ari Morcos, Stefan Lee, Irfan Essa, Devi Parikh, Manolis Savva, and Dhruv Batra. Dd-ppo: Learning near-perfect pointgoal navigators from 2.5 billion frames. In International Conference on Learning Representations, 2020. URL https://openrett /openevviww
net/forum?11 ttttfoorummd-HllxSCAYYY
Cited on p. 4, 18.
Sam Witty, Jun K. Lee, Emma Tosch, Akanksha Atrey, Kaleigh Clary, Michael L.. Littman, and David Jensen. Measuring and characterizing generalization in deep reinforcement learning. Applied A/ Letters, 2(4):455, 2021. doi: hups://doi.org wtwp//..///////..... wiley. com/ doi abs 10. 1 a http//al...44. 45. Cited on p. 2.
URL https:: //o////nnnloo wnn//nnnnnnnnnnnne ieliibary....
Fei Xia, Amir R. Zamir, Zhiyang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gibson Env: real-world perception for embodied agents. In Computer Vision and Pattern Recognition (CVPR), 2018 IEEE Conference on. IEEE, 2018. Cited on 18.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv [cx. CV], March 2022. URL htt.ht::/aaaii..gg//aaaaaa wbtp:////..a//./....////// 06173. Cited on p. 3.
Jun Yamada, Karl Pertsch, Anisha Gunjal, and Joseph Lim. Task-induced representation learning. In International Conference on Learning Representations, 2022. URL /tttps://opereevieet net/ forum? detrrumm?dd=OxxxxIzzZzt.
Cited on p.
Mengjiao Yang and Ofir Nachum. Representation matters: Offline pretraining for sequential decision making. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 11784- 11794. PMLR, 18-24 Jul 2021. URL https: http://pooeeeed.../eee// wwweeedinng.m..........m html. Cited on p. 2.
wtt.sss..rrr.......mm. press/ /v13///1////////// /aagg//
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representations, 2021. URL https: /wwpreveeeeeeeeeee net/ nwttffoumm forum? dd= 6sTvGaf. Cited on p. 4, 9, 17.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data- augmented reinforcement learning. In International Conference Learning Representations, 2022. URL https: wtwprrrwwww.eeeefoouuunideeee sestrruuumiiddddddeeeee-- SJ__yyes8. Cited on p. 2, 4, 9, 17, 19, 23.
14

---

Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data, Adv. Neural Inf. Process. Syst. 34, 2021. URL htttps://proceeee Paper pdf. Cited on 2.
neur ips.
Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances in Neural Information Processing Systems, 27:3320-3328,, 2014. Cited on p. 2,
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. June 2018. URL URL.htp:///rxxtv.org//aas/ att/ 1806. 07937. Cited on p. 2.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning
invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Reprexentations, 2021. URL http:ps:///peeeeeeeeeneeeee/e/ meme/fooumm?ii -2FCwDKRREu. Cited p. 2.
Wancong Zhang, Anthony GX-Chen, Vlad Sobal, Yann LeCun, and Nicolas Carion. Light-weight probing of unsupervised representations for reinforcement learning. arXiv [cs.LG], August 2022.
URL http::t/aaxxi..rr/////// hatp//aabb//.....e//ee 1238.12345. Cited on p. 6.
15

---

Walker Run
DMC
Cartpole Swingup
DMC
Hopper Hop
DMC
900
500
250
800
Frozen
Rewards better)
600
Rowards (1 better)
Rewards (11 better)
200
Finctuned
6000
300
500
100
200
400
Frozen
Frozen
100
Finetuned
300
Finetuned
200
0
Steps
Steps
Steps
Figure 7: When should we freeze or finetune pretrained representations in visual RL? Reward transfer curves on Cartpole and Hopper DeepMind Control domains; replicates Figure 1. Freezing pretrained representations can underperform pretraining at all (Figures la and 1b) or oulperform it (Figures Ic and Id). Fineetuing representations is always competitive, but fails significantly outperform freezing visually complex domains (Figures and Id). Solid lines indicate mean over 5 random seeds, shades denote 95% confidence interval. See Section 3.2 for details.
A
APPENDIX
TESTBED DETAILS B
This section details the pretraining and transfer setups for the MSR Jump, DeepMind Control, and Habitat domains.
B. MSR Jump (TACHET DES COMBES AL., 2018)
As discussed in Section 3, the goal of MSR Jump tasks is for the agent to reach the right-hand side of the screen. The agent always starts on the left-hand side and has to jump over a single gray obstacle. The agent observes the entire screen from pixel values, and the only factors of variation are the æ and y coordinates possible actions are to move right jump. The jump dynamics are shared across tasks and mid-air actions don't 't affect the subsequent states until the agent has landed (iee. all jumps go up by 15 pixels diagonally, and down by 15 pixels diagonally). Our experiments carefully replicate the reward and dynamics of Tachet des Combes et al. (2018), save for one aspect: we increase the screen size from 64 pixels 84 pixels, as this significantly accelerated training with the DQN (Mnih et al., 2015) feature encoder. We provide a fast GPU-accelerated implementation of MSR Jump tasks in our code release (see Appendix E).
the obstacle (i.e., its horizontal position and the floor height). Its only
B.1.1 PRETRAINING SETUP
Source tasks We pretrain on 70 source tasks, varying both the obstacle (from 15 50 pixels) and y (from to 50) coordinates, every pixel. We don't 't use frame stacking nor action repeats.
Learning agent We train with PPO (Schulman al., 2017) for 500 iterations, where each iteration collects 80 episodes from randomly sampled source tasks. We use those episodes to update the agent 320 times with batch size of 64, discount factor of 0.99, and policy and value clip of 0.2. The learning rate starts at 0.0005 and is decayed by 0.99 once every 320 update. The agent consists of the 4-layer DQN feature encoder and linear policy and value heads. All models use GELU (Hendrycks et al., 2016) activations.
1.1 TRANSFER SETUP
At transfer time, we replace the 70 source tasks with 70 unseen downstream tasks. Those downstream tasks obtained by shifting the obstacle position (z and y coordinates) of each source task by 2 pixels. For example, if source task has an obstacle at = 20 and y = 30 there is corresponding downstream task at 22 and 32.
Downstream tasks
16

---

Learning agent The learning agent is identical to the source tasks, but rather than randomly initializing the weights of the feature encoder we use the final weights from pretraining. As mentioned in the main text, we also introduce one additional syper-parameter for the feature encoder and policy value heads. For Finetuned these learning rates are 0.001 for the heads and 0.0001 for the feature encoder; both learning rates are decayed during finetuning.
decouple the learning rate
B.2 DeepMind Control (TASSA ET AL., 2018)
The DeepMind Control tasks build on top of the MuJoCo physiscs engine (Todorov et al., 2012) and implement a suite of robotic domains, Our experiments focus on Walker Cartpole, and Hopper (see Appendix C.2 for the latter two). Walker and Hopper are robot locomotion domains, while tert is the classic cartpole environment. On all domains, the agent receives an agent- centric visual 84 RGB observation, and controls its body with continuous actions. Our task implementation exactly replicate the one from DrQ (Yarats et al., 2021).
B.2.1 PRETRAINING SETUP
Source task On Walker, the source task is Walk, where the agent is tasked walk in the forward direction. On Cartpole is Balance, where the pole starts in a standing position and the goal is to balance it. And on Hopper it is Stand, where the agent needs reach (still) standing position. For all tasks, we keep stack of the last 3 images and treat those as observations.
Learning agent Our learning agent is largely inspired by DrQ-v2 (Yarats et al., 2022), and com- prises 4-layer convolutional feature encoder, a policy head, and twin action value head. Both the policy and action value heads consist of a projection layer, followed by LayerNorm (Ba et al., 2016) normalization, and a 2-layer MLP with 64 hidden units and GELU (Hendryeks et al., 2016) activations. We train with DrQ-v2 for 200k iterations, starting the policy standard deviation at 1.0, decaying it by 0.999954 every iteration until reaches 0.1. We use 512 batch size and don't use n-steps bootstrapping nor delayed updates. For pretraining, we always a learning rate of 3e-4 and default Adam (Kingma Ba, 2014) hyper-parameters.
B.2.2 TRANSFER SETUP
Downstream task On Walker, we transfer to Run, where the agent needs to run in the forward direction. On Cartpole, the downstream task is Swingup where the pole starts in downward position and should be swung and balanced upright. On Hopper, we use Hop, where the agents move the forward direction by hopping.
Learning agent We the same learning agent as for pretraining, but initialize the feature encoder with the weights obtained from pretraining. We also add a learning rate hyper-parameter when finetuning the feature encoder, which is tuned independently for each setting and method.
B.3 Habitat (SAVVA ET AL., 2019)
On Habitat, the goal of the agent is navigate indoor environments to reach desired location. To solve this point navigation task, it relies on two sensory inputs: a GPS with compass and visual observations. The visual observations are 256 256 RGB images (n.b., without depth information). Our codebase builds on the task and agent definitions of "Habitat-Lab" (Szot et al., 2021), I), and only modifies them to implement our proposed methods.
B.3.1
PRETRAINING SETUP
Source task We use the ImageNet (Deng et al., 2009) dataset as source task to pretrain our feature encoder.
Learning agent The learning agent is ConvNeXt-tiny (Liu et al., 2022) classifier, trained to discriminate between ImageNet-IIe classes. We directly use the pretrained weights provided by the ConvNeXt authors, available at: https://giihaee com www.hh..ooaaccbboooo.eeeoooo http://aaooooeerrac/ccooooo
17

---

Walker-Run
DMC
Swtpple-Swingup
DMC
Hopper-Hop
DMC
Value Estimation Emor better)
Value Czimation Ener (1 better
Value Essimation Eneor (1 bettrr
Randon
Ransiom
Randam Frozen
Frezen
Frazen
Tinetuned
linetuned
Tinetured
Obgam.
Come
Comed
Gond
Comed
Olearv
Comil
Comut
Cane
Comet
Obeanv
Coned
Comvi
Camd
Come
When Using Features From Layer
When Using Features From Layer
When Using Features From Layer
Walker-Run
DMC
Cartpole Swingup
Sarigole-Swinup
DMC
Hopper-Hop
DMC
Ansa Undler Reward Curve (1 better)
Ane Under Reword Cure (1 better)
Under Reward (1 better)
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.2
0.2
0.2
0.1
0.0
8.0
Observ.
Comu
Cana
Cand
Camu
Observ.
Comd
Came
Cana
Cand
Obsmm.
Cardl
Camv
Cand
Come
When Frenzing Up To Layer
When Freezing Up To layer
When Freezing Up To Laye
Figure 8: Frozen layers that retain the same information as finetuned layers can be kept frozen. Replicates Figure 4a on Car tpole and Hopper DeepM ind Control domains. (top) Mean squared error when linearly regressing action values from representations computed given layer. Some early layers equally good predicting action values hefore or after finetuning. This suggests those layers be kept frozen, potentially stabilizing training. See Section 3.5 for details. (hottom) Area under the reward curve when freezing up to given layer and finetuning the rest. Freezing carly layers does not degrade performance, and sometimes improves it (e.g., up to Conv2 on DeepMind Control).
B.3.2 TRANSFER SETUP
Downstream tasks For transfer, we load maps from the Gibson (Xia et al., 2018) and Matter- port3D (Chang et al., 2017) datasets. Both are freely available for academic, non-commercial use at:
Gibson: https: http:///stttiiiiio////////oaaaat
Matterport3D: httph: http://////sttttooooo/oeeeee///e
Learning agent We take the pretrained ConvNeXt classifier, discard the linear classification head, and only keep its trunk as feature encoder. We train 2-layer L.TTM policy head and linear value head on top of this feature encoder with DDPPO (Wijmans et al., 2020) for 5M iterations on Gibson and 10M on Matterport3D. For both tasks, we collect 128 steps with workers for every update, which consists of 2 PPO epochs with a batch size of 4. We set PPO's entropy coefficient to 0.01, the value loss weight 0.5, GAE's norm to have 0.2 12-norm. The learning rate for the policy and value heads is set to 2.4c -4 and the feature encoder's learning rate to le for Frozen+F Finettnnee.ned. For Frozen-Varnishe we use single learning rate set to 1 for the feature encoder and policy value heads (but 2.5e gave similar results). We never decay learning rates.
0.95, the discount factor to 0.99, and clip the gradient
C ADDITIONAL EXPERIMENTS
This section provides additional details on the experiments from the main text, and also includes additional results further support conclusions.
18

---

Walker Run
DMC
Cartpole Swingup
DMC
Hopper Hop
DMC
600
900
300
800
250
Rewrrds (1 [etter)
Rowards (1 (etter)
Rewards (11 better)
400
200
6000
150
5000
Frozen
Frozen
100
Frozen
200
Finetured
400
Finetured
Frozen PISCO
Frazen-FFnneuned PISCo
900
Frozen FFinetundd
Frazen PISCO
50
Frozen + PISCO
Finetuned
200
0
1.0 Steps
Steps
Steps
Figure 9: Partial freezing improves convergence; adding our policy-induced consistency objective im- proves further. As suggested by Scction 3.5, we freeze the carly layers of the feature extractor and finetune the rest of the parameters without (Frozen-Finetunel) or with policy-induced consistency objective (Frozen+PiSCO), On challenging tasks (e.g., Hopper), partial freezing dramatically boosts downstream perfor- mance, while Frozen-PiSCO further improves upon Frozen+Finetuner Finetuned across the board.
Walker-Run
DMC
Cartpole-Swingup
DMC
Hopper-Hop
DMC
900
300
300
800
260
Rewords better)
Rewards (1 better)
700
Rewards (ft better)
200
6000
160
300
500
Frozon Finetuned
Frozen Finetuned
Frozzn i Finetuned
100
200
Frozen: CURL
800
Frozen: CURL
Frozen j CURL
Frozen-- Frozen-Simmiaa
Frozen SimSSam
Frozen: Simmiam
300
Frozen-
Frazen-PiSCC
200
0.5
1.0
1.5
2.0
0.25
0.50
0.75
1.00
0.5
1.0
1.5
200
Steps
x10
Steps
X10
Steps
x10P
Figure 10: Policy-induced supervision is a better objective than representation alignment for finctuning. We compare policy-induced induced consistency objective measure of representation similarity against popular representation leamnigg objectives (e.g. SimSiam, CURL). Policy supervision provides more useful similarities for RL finetuning, which accelerates convergence and reaches higher rewards.
C.1 COMPUTING ACTION VALUES
EEVER-BY-LAYER LINEAR PROBING EXPERIMENTS
We follow similar protocol for all tasks to obtain action values for mayer-by-layer linear probing experiments. stated in Section 3.5, we first collect 10, 000 pairs using policy trained (or finetuned) on the downstream task. We then compute the returns for each state-action pair by discounting the rewards observed in the collected trajectories. Then, we fit a 64-hidden unit 2-layer MLP with GELU activations to predict action values from the collected state-action pairs, and store the predicted action-values. On MSR Jump and Habitat we use the one-hot encoding of the actions rather than their categorical values. Finally, linear probing models.
the MLP- Ale--rrddeted action values regression targets for
C.2 FULL RESULTS ON DeepMind Control
As mentioned in Section 3.1, we also ran our analysis on two DeepMind Control domains other than Walker: Cartpole and Hopper. Both tasks use the same pretraining and transfer setups as Walker. In Cartpole, we pretrain on the Balance task, where the pole starts in an upright position and our goal to keep blanced. We then transfer Swingup, where the pole starts upside down and the goal is swing and balance upright. With Hopper, the pretraining task is Stand where the agent is asked to stand up still; the downstream task is Hop where the agent should move in the forward directino fast as can. Note that Hopper-Hop is considered more challenging than Cartpole and Walker tasks, and DrQ-v2 fails to solve even after 3M iterations (Yarats et al., 2022).
For both domains, all analysis results are qualitatively similar to the Wa lalker ones from the main text. Figure shows once more that freezing representations is not sufficient for transfer. Figure
19

---

Gibson
Habitat
Matterport3D
Habitat
Volue Estimation Eiror (1 beterr
Random Frazen
Value Extimation Emror (J better)
Random Frozen
2.
Comtl Cone
Conne
Comes Conde Conu?
Cami AugPoed
LayeNeam
Comu Comed
Conut
Conus Corud Comu
Conva ArpPoul
LayerHor
When Using Features From Layer
When Using Features From Layer
Figure 11: Identifying which layers freeze on Habitat tasks. replicate by-er-byllayer linear probing experiments on Habitat with the ImageNet pretrained feature extractor. Although the downward trend is less evident than in Figure 4a, we clearly see that layers Conv8 and Conv7 yield lowest value prediction error on Gibson and Matterpont3D, respectively.
replicates Figure 4a. On both Cartpole and Hopper, freezing beyond layer Conv2 increases the value estimation error (top panels); conversely, freezing up to those layers yields the best AURC (bottom panels). Note that Cartpole the degradation in value error is minimal and so is the difference in AURC. We see in Figure that partially freezing the feature encoder yields better transfer results, especially when combined with PiSCo. Finally, Figure 10 reconfirms our hypothesis that a policy-driven similarity metric improves RL finetuning.
C.3 SELECTING FROZEN LAYERS IN Habitat
We decided freeze up to layer Conv8 and Conv7 for the Gibson and Matterport3D experiments in Section 5.1. This section explain how our analysis informs the choice of those layers, without resorting to a finetuned policy nor layer-by-layer freezing experiments as Figure 4b.
Building on our analysis, we obtained 10, 000 state-action pairs from policy trained directly on Gibson and Matterport3D. This policy is provided by "Habitat-lab" but in practice couple expert demonstrations might suffice. As in Figure 4a, we fit action value function and save the action values for the observation-action pairs from the collected state-action pairs. Then, we replicate the layer-by-layer linear probing experiment of Section 3.5 with a randomly initialized feature encoder and the one pretrained on ImageNet. Two differences from Section 3.5 stem from the large and visually rich observations on tabitat. First, we project each layer's representations to PCA embedding of dimensions 600 rather than 50"; second, we omit the first convolutional layer as PCA prohibitively slow on such large representations.
Figure 11 displays those results. As expected, the downward trend less evident than in Figure 4a since the transfer gap between ImageNet and Habitat is larger than between tasks in MSR Jump DeepMind Control. Nonetheless, we can easily identify the representations layers Conv8 and Conv7 as the most predictive of action value on Gibson and Matterport3D, thus justifying our choice.
C.4 DeepMind Control TRANSFER FROM IMAGENET
Given the promising results of large feature extractors pretrained on large amounts of data Figures Ic and 1d, we investigate whether similarly trained feature extractors would also improve transfer on DeepMind Control tasks. To that end, we use the same ConvNeXt as our Habitat experiments, freeze its weights, learn policy and action-value heads on each one of the three DeepMind Control tasks (replicating the setup of Figure 1).
'Dimensions ranging from 500 to 750 work equally well. Smaller dimensions yield underfitting as PCA doesn't retain enough information from the representations. Larger dimensions yield overfitting and pretrained representations don't 't do much hetter than random ones.
20

---

Walker Run
DMC
Cartpole Swingup
DMC
Hopper Hop
DMC
600
300
Frazen (Walk)
800
(Stand)
500
(ImageNet)
250
(ImageNed)
Rewards (1 better)
Rowards {1 better)
Rewards (11 better)
600
600
200
Freezn (Balance)
300
150
Frozen (ImageNett
200
100
200
100
50
0
0
1.0
Steps
Steps
Steps
Figure 12: Pretraining on large and diverse data can hurt transfer when the generalization gap is too large. When transferring representations that are pretrained on ImageNct to DeepMind Control tasks, we see signficant decrease in convergence rate. We hypothesize this is due to the lack of visual similarity between ImageNet and DeepMind Control.
Walker-Run
DMC
Cartpole-Swingup
DMC
Hopper-Hop
DMC
De Nove
500
800
250
De Novo PISCO
Rewards (1 beterr
Rewards (1 better)
200
Rewards (1 better)
Frozon PISCO
600
150
300
500
De Nevo
-
De Nova
100
De Novo PISCO
800
De Novo PISCO
Frozen
Fropen
300
Frozen- PISCO
Frozen+FPISC
0.5
1.0
1.5
2,0
0.25
0.50
0.75
1.00
0.5
1.0
1.5
2.0
Steps
x10
X10
X10P
Figure 13: PiSCO also improves representation learning source tasks but shines with pretrained representations. When benchmarking PiSCO combined with fully finetunable features (De Novo+PisCO), we observe that it marginally outperforms DrQ-v2 on the downstream tasks (De Novo). However, the best performance is obtained when also transferring task-agnostic features (Frozen+PiSCO).
Figure 12 provides convergence curves comparing the 4-layer CNN pretrained on the source Deep- Mind Control task (i.e., either Walker-Walk Wa, Cartpole- aapoleeBBllance Hopper-Stand) against the ConvNeXt pretrained on ImageNet. The CNN drastically outperforms the ConvNeXt despite having fewer parameters and trained on less data. We explain those results as follows. The smaller network is trained on data that is more relevant to the downstream task, thus has smaller generaliza- tion gap to bridge. In other words, more data is useful insofar More parameters might help (i.e., pretraining the ConvNeXt on DeepMind nd Control tasks) but it notoriously difficult to train very deep architectures from scratch with reinforcement learning, as shown with De Novo Habitat tat tasks.
is relevant to the downstream task.
C.5
De Novo FINETUNING WITH PiSCO
As an additional ablation, we investigate the ability of PiSCO to improve reinforcement learning from scratch. Figure 13, we compare benefits of using PiSCO with when representations are (partly) frozen v.s. fully finetunable (De Novo). We see that PiSCO always improves upon learning from scratch (De novo+PiSCO outperforms Novo) but that is most beneficial when the task-agnostic features are frozen Prrozen+PiSce outperforms De Novo+PiSCO).
C.6 PiSCO WITHOUT PROJECTION LAYERS
The original SimSiam formulation includes a projector layer h(-), and Chen & He (2020) show that this projector is crucial to prevent representation collapse. Does this also holds true for RL?
To answer this question, we compare the performance a (partially) frozen feature extractor finetuned with PisCo, with and without projection layer, in Figure 14. The results clearly answer our question
21

---

Walker Run
DMC
Cartpole Swingup
DMC
Hopper Hop
DMC
600
900
300
Frozen
800
250
Froeen-FiSCO (no proj.)
Rewrrrs (1 [ebter)
Rowards (1 (eetter
700
Rewards (11 better)
Frozen PISCO (w/ proj.)
400
200
6000
150
500
200
Frozen
400
Frozen
100
Frozn-PPSSCO (w/ proj)
Frozen-PISCO (no proj.)
900
De Novo PISC+ (no pro..
Froeen+FiSCO (w/ proj.)
50
200
0
Steps
Steps
Steps
the projector h(-) necessary PisCO's formulation? Yes. Removing projection layer when finetuning task-specific layers (and freezing task-agnostic layers) drastically degrades performance on all DeepMind Control tasks.
Figure 14:
Walker Run
DMC
Cartpole Swingup
DMC
Happer Hop
DMC
000
9000
300
800
260
Rewards better
Rewards [11 better)
700
Rewards [f better)
400
200
600
160
100
Frozen-FFnetuned
Propen--Firttune
100
Frozen Finet
200
Frozen-PISCO
Frocen+FS5CO
Frocen+/1550
F
300
Propen--SPP
50
Frozen: SPR
Frozen- PISPR
Frozen-FFSSPH
Frozen+ PISP
200
0
1.5
0.25
1,00
1.0 Steps
Steps
Steps
x
x10P
Figure 15: Combining SPR with PiSCo significantly improves the performance of SPR. Swapping the cosine similarity objective in SPR (Schwarzer et al., 2021) for the policy-induces objective suggested by analysis significantly improves finetuning. Still, finetuning with PiSCO (based on SimSiam (Chen & Hc, 2020)) yields the best performance, while remaining easier to implement and faster in terms wall-clock time.
in the affirmative: the projector is also required and removing it drastically degrades performance the downstream task.
C.7 COMPARING AND COMBINING WITH SPR
To illustrate future avenues for extending the ideas motivating PiSCO, combine them with the self-predictive representation (SPR) objective of Schwarzer et al. (2021). We replace SimSiam for SPR the auxiliary self-supervised objective used during finetuning. Concretely, reimplemented SPR for DeepMind Control tasks and compare the original implementation (which uses cosine similarity cos(p1,P2 P2) ) to compare representations) against variant (PiSPR) which uses the policy-induced objective 2 (D(pi.P2) + D(p2. P1))
In terms of implementation details, we found that setting the number of contiguous steps sampled from the replay buffer to K = 3 yielded the best results on Dnd Control tasks. The transition model is 2-layer MLP (1024 hidden units) with layer normalization and ReLU activations. We used the same projector as as for SimSiam (which does not reduce dimensionality), and linear predictor with layer-normalized inputs.
Figure 155 reports the results for partially frozen feature extractors with different finetuning objectives. We observe that using SPR typically performs as well finetuning the task -ask-specific layers with the RL objective. In that sense, finetuning with PiSCO offers compelling alternative: it is substantially cheaper (3x faster wall-clock time), doesn't 't require learning a transition model nor keeping track of momentum models, and is simpler tune (fewer hyper-parametter parameters). More interestingly, we find that replacing the cosine similarity in SPR for the policy-induce objective (i.e., PiSPR) significantly improves performance SPR. Those results further validate the generality of our analysis insights.
NNote: Pi and P2 correspond ieta and det* in the original SPR notation.
22

---

D PSEUDOCODE
Here we provide pseudocode for implementing PiSCO on top of two popular reinforcement learning algorithms. Algorithm adds PiSCO PPO (Schulman al., 2017), while Algorithm 2 adds it to DrQ-v2 (Yarats et al., 2022).
Algorithm PPO with PiSCO
sample transitions from replay
2 a, s' hapllay.samplebbachhsiizzz
3
((s)
5 # compute value loss
C² 0.5 (V()) drscount(r))²
8 # compute policy loss
Ax log 2) - log Ta z)
10 GAE (V)(), T, Y, T)
_ min(exp ( A, clip(exp(Ar)- A,1 - +.
12
13 # compute PiSCO loss
14 2,, é(data_ (s)) o(data_ augment(s))))
pl, p2 h(z),, (x22)
16 [ 0.5 (L ( ( pp))) p2)) 0.5 (L(L(M ( ( p:))
18 adam.optimize +v. LV + A 1isco B H(*{
z))) 2)))
optimizes V, n, h, and s
Algorithm 2 DrQ-v2 with PiSCO
sample transitions from replay
8, a, r, s replay. wawpleeaatchhhhsiiiiiii..e.. size, n_steps, 7)
# compute policy loss
5 z = ((d(data__uggig auggmntts)))
6 à = )( rsample()
min â), Q2 à))
adam. osamooeiimize(C)
only optimizes x
10 # compute action-value loss
11 @(data__aument(())
12 2' = didata_augmett (s'))
13 a' # ( samplee)
q. f2 Qi a), Q2 a)
15 of = L(r 7 min (Qi (z', a'), Q2(2', d')))
16 _ 0.5 (q) - 0.5 (42 - q
17
compute PiSCO loss
19 21, 22 ó(data_ (dddta__ugment(s)) a(dataaaugmentte
20 pl, h(21), h(22)
C 0.5 (L [ p2)) 8.5 (L a2) [[ p.))
22
23
adam. www.:ittiiiieeeeemm (C + X (1000
only optimizes Q1 Q2, lo, and o
E CODE RELEASE
Together with our submission, release example code implement PiSCO in the code/ directory of the Supp. Material. We plan to publicly release complete implementations, together with examples for the pretrain- -nnerafeend- pipeline. We describe the main files below.
23

---

msr_jump. py: our GPU implementation of MSR Jump tasks.
ppo ppo_pisco.py: py: our PyTorch (Paszke et al., 2019) implementation of PiSCO on top of PPO.
drqv2_pisco. py: our PyTorch (Paszke et al., 2019) implementation of PiSCO on top of DrQ-v2.
24


Reviews:
[{"review_id": "nWpsK8vSBOb", "reviewer": "Reviewer_ZmQR", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies fine-tuning vs freezing in the context of reinforcement learning. The authors argue that the layers that represent general-purpose features(mainly the early layers) could be frozen while those that represent task-specific features(mainly the last layers) should be fine-tuned to achieve higher performance in the downstream task.", "summary_of_the_review": "In my opinion, this is a well-written paper that suffers from a lack of novelty and convincing arguments. Please see the above for more detailed feedback.", "strength_and_weaknesses": "Strength:\n\n- The paper is well-written and easy to follow.\n\nWeaknesses:\n\n- I am not an expert in reinforcement learning. However, the argument in the paper in terms of representation learning does not seem convincing and sufficient as an ICLR paper to me.  First, the author's argument about freezing vs finetuning in the introduction(namely the second and third paragraphs) holds for any neural network and downstream task. Posing it specifically as an RL problem does not seem rigorous to me.  Second, as it is mentioned by the authors themselves, the fact that the early layers of a convent capture general-purpose features while the later layers represent more task-specific features, is known and widely accepted. In fact, the pioneering work on fine-tuning such as Fast-RCNN achieves higher performance when the early layers are frozen, and the last layers are fine-tuned.\n\nI can not evaluate the performance, but  I am not entirely convinced about the sufficiency of the contributions in. this paper. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "h_HpILOU1E-", "reviewer": "Reviewer_nmYq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies transfer learning in the context of Reinforcement Learning. Using 3 diverse domains as a testbed (a toy platform game, some tasks from the DeepMind Control Suite, and some navigation tasks from habitat), it makes the following contributions:\n- Contribution #1: Shows that when transferring representations, the best strategy involves freezing the deepest layers and fine-tuning the others\n- Contribution #2: Provides a strategy to select which layers to freeze and which to finetune\n- Contribution #3: Introduces an auxiliary self-supervised objective to further enhance the transfer performance.\n\n\n", "summary_of_the_review": "My main critiques for each of the contributions are as follows:\n\n- The analysis of finetuning vs freezing representations misses the discussion on the capacity and pre-training diversity of the network involved, which in my opinion makes the picture incomplete\n- The proposed strategy to choose which layers to freeze seems too costly, in its current form, to be broadly applicable\n- The auxiliary objective is not studied in depth, be it with respect to its own design choices (eg the projection layer) or with respect to current state of the art auxiliary objectives (eg SPR)\n\nAs such, I don't recommend the paper in its current form for publication. That being said, all the aforementioned limitations seem fixable, and I am willing to increase my score if they are.", "strength_and_weaknesses": "All 3 contributions outlined above are more or less independent from each other, as such this section will be divided into three parts discussing the strength and weaknesses of each of them.\n\n## Contribution 1\n\nThe paper starts by analyzing the performance of two main transfer strategies, namely freezing the representations or finetuning them. This is compared to the baseline of training the representations from scratch on the target task (strategy nicknamed \"de novo\"). In a nutshell, the results can be summarized as follows:\n- For the toy jump task as well as the Deepmind control suite, finetuning outperforms \"de novo\" while freezing representations substantially decreases performance\n- For the Habitat tasks, both strategies outperform \"de novo\" but none seems strictly better than the other.\n\nThe paper provides a thorough and solid analysis of the failure of the freezing of the representations for the jump task, linking it to robustness issues.\nHowever, in my opinion, this falls short of making a broader conclusion: the results above suggest a meaningful qualitative difference between the Habitat task and the other ones (in Habitat, freezing works while it doesn't in the others). The reason for this difference is not discussed in the paper, yet it seems to be somewhat at the crux of the matter. Does this difference comes from intrinsic characteristics of the environments or from other factors?\nAs for the other factors, we can list 2 main ones: \n- Network capacity/architecture: The jump task uses the original DQN network, which is rather primitive (no skip connections for example) and has a tiny capacity (around 100k trainable parameters). By contrast, the Habitat task uses a state-of-the art CNN, with 28M parameters, ie two orders of magnitude larger.\n- Pre-training tasks: The convnext for Habitat is pre-trained on ImageNet, while in the other experiments the backbones are pretrained on related control tasks.\n\nIn a nutshell, the networks used for Habitat have received a much more varied pre-training and have a much higher capacity. It is well known in the vision community (see [1] for example) that both these properties improve robustness and transferability of the features, making the results presented in the paper somewhat unsurprising.\n\nOverall, it is not clear to me if the \"frozen\" transfer strategy would still fail in the two tasks where it currently does, if the same convnext, with the same pretraining, was used as the feature extractor. Depending on the answer to that question, a slightly different conclusion could be made: similarly to trends in other domains (Vision, NLP, speech, ...) does RL benefit from larger models pre-trained on large, general datasources?\n\nFinally, I believe the paper could discuss a bit the alternate method to full freezing and finetuning that is being privileged more and more both for vision and NLP: freezing the pre-trained weights but adding a small amount of trainable parameters to allow \"adpatation\" of the learned representation to the new task. See [2], [3] and references therein.\nAnother thread of work of interest is the idea of using a different learning rate per layer (typically lower for the deepest layers and increasing from there). See [4] for an example.\n\n\n## Contribution 2\n\nThe paper proposes to resort to linear probing of the feature of each intermediate layer to determine where the cut-off between frozen and fine-tuned should be made.\nWhile the method seems sound, several design choices are not really explained:\n- What is the justification behind reducing the dimensionality with PCA first, as opposed to simply using a linear layer directly on the features?\n- The paper states that the probing is using the output of a action-value function network. Would it be possible to alleviate the need to train such a model by simply predicting an empirical estimate (eg the mean) of the expected discounted return in each state?\n\n\nFigure 4.b (and similar figures in the appendix) would benefit from error bars since it's not clear what the variance is. Why does Fig. 11 not contain the estimates for the fine-tuned features?\n\nFinally, the main weakness of the approach is its computational cost. Section 5.1 states \"For each layer, we measure the action value estimation error of pretrained and finetuned representations, and only freeze the first pretrained layers that closely match the finetuned one\". This implies access to the said fine-tuned representations, which in-turn implies that a full RL training is required as a first step. As such, this doubles the cost of the method (both in wall-clock time as well as sample complexity), significantly reducing its practicality. For proper comparison, it would be necessary to see whether the \"Finetuned\" baseline from figure 5 is still lower than the \"Frozen+Finetuned\" if the former is trained for twice the number of steps.\nPerhaps, a \"softer\" scheme akin to [4] would do the job equally well while side-stepping the computational cost entirely.\n\n## Contribution 3\n\nThe last contribution is a self-supervised auxiliary objective aiming at improving the quality of the representation. Though it is inspired from findings from contribution #1, it is a contribution that is not specific to the transfer setting and could potentially be applied elsewhere.\n\nHere are some questions and comments related to this contribution:\n- What is the motivation for the projection layer in the SSL loss? The choice is not ablated in the paper. I understand the need for such a layer in methods that do direct embedding supervision (either contrastive loss or simple L2 loss for example), but in this case the loss directly operates on the output distribution. Similar vision method like DINO [5] don't use such projector\n- It would be valuable to provide \"De Novo + Pisco\" as a baseline to disentangle the gains from the transfer scheme from those due to the additional objective\n- Though comparison to CURL is provided, SPR [6] has been shown to outperform it significantly, so a comparison seems warranted.\n\n\n\n\n[1] \"Big Transfer (BiT): General Visual Representation Learning\", Kolesnikov et al, ECCV 2020\n\n[2] \"AdapterHub: A Framework for Adapting Transformers\", Pfeiffer et al, EMNLP 2020\n\n[3] \"Learning multiple visual domains with residual adapters\", Rebuffi et al, Neurips 2017\n\n[4] \"Universal Language Model Fine-tuning for Text Classification\" Howard et al, ACL 2018\n\n[5] \"Emerging Properties in Self-Supervised Vision Transformers\", Caron et al, ICCV 2021\n\n[6] \"Data-Efficient Reinforcement Learning with Self-Predictive Representations\", Schwarzer et al, ICLR 2021\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-WK1WGfRGqL", "reviewer": "Reviewer_b3ER", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors compare finetuning and freezing pretrained representations on various RL tasks.\nThey seek to understand under what conditions each of these can be effective, and when one should be preferred to the other.\nOn a simple task where the optimal policy is known, they show that the pretrained representation does contain enough information to regress the optimal actions, but finetuning still fails to recover the optimal policy in spite of this.\nThey use insights from these experiments to propose a new self-supervised objective for RL, and evaluate this objective on a variety of tasks.\n", "summary_of_the_review": "I think that the proposed approach shows promise, but the motivation is unclear (mainly per Weaknesses #1). I would feel comfortable accepting this paper if this were addressed (and the experiments/evaluation may need to change accordingly).", "strength_and_weaknesses": "Strengths:\n1. The proposed objective, PiSCO, seems elegant and promising. I am curious if the authors could say more about the projector $h(\\cdot)$ that is used in the PiSCO loss. The intuition for PiSCO makes sense, but I am not sure why why should compare $\\pi(\\cdot|h_1)$ and $\\pi(\\cdot|z_2)$ instead of simply $\\pi(\\cdot|z_1)$ and $\\pi(\\cdot|z_2)$?\n2. The paper is well-written, and the authors attempt to be insightful in their discussion.\n\nWeaknesses:\n1. The narrative of the paper feels a bit incoherent. It feels to me that the freezing vs finetuning investigation did not well-motivate PiSCO, and that PiSCO does not address the original problem statement that motivated said investigation. From what I can tell, PiSCO is an auxiliary loss to be used during RL, not a pretraining objective. I think this paper's contributions would be more clear if either: (a) the authors proposed a pretraining objective, or (b) the authors more directly motivated PiSCO, e.g. by discussing robustness or data augmentation/efficiency.\n2. Related to the previous point, if the authors were to consider option (b), then I feel that PiSCO could benefit from a more thorough experimental evaluation (a greater variety of tasks/domains, with more emphasis paid to sample efficiency).\n3. The authors make broad statements about finetuning vs freezing solely based on results from a toy environment, MSR jump. I understand the attractiveness of that domain (in that the optimal policy is known), but I'm wondering if it would be possible to conduct a similar study in other domains as well?\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject