D ON ’ T F REEZE Y OUR E MBEDDING : L ESSONS FROM P OLICY F INETUNING IN E NVIRONMENT T RANSFER
A BSTRACT
A common occurrence in reinforcement learning RL research is making use of a pretrained vision stack that converts image observations to latent vectors
Using a visual embedding in this way leaves open questions , though: should the vision stack be updated with the policy
In this work, we evaluate the effectiveness of such decisions in RL transfer settings
We introduce policy update formulations for use after pretraining in a different environment and analyze the performance of such formulations.
Through this evaluation, we also detail emergent metrics of benchmark suites and present results on Atari and AndroidEnv
I NTRODUCTION
Imagine an agent completing a task when the background changes color.
The underlying task is identical; only the observation space has changed.
One might wonder: in this setting, might we want to update the embedding network while keeping the policy the same?
This is counter to a common strategy of freezing the embedding while ﬁnetuning the policy head.
Much work in RL uses a pretrained visual embedding network, which may or may not be ﬁnetuned during policy learning.
In this work, we aim to evaluate the effectiveness of these paths.
To this end, we present a transfer framework that allows for evaluation of zero-shot generalization, embedding ﬁnetuning (frozen policy head), head ﬁnetuning (frozen embedding), and full policy ﬁnetuning.
Comparing these approaches gives us a better understanding of how an agent makes use of the different parts of its policy network.
The effectiveness of these paths depends on environment similarity at transfer time.
Exploiting this reasoning, we introduce environment metrics that evaluate an environment along multiple dimensions based on ﬁnetuning success.
We present such metrics for Atari ﬂavors in the Arcade Learning
R ELATED W ORK 2. M ULTITASK T ESTBEDS IN RL
In recent years, researchers have identiﬁed a growing need for better evaluation in RL (Machado et al., 2018; Henderson et al., 2018).
One such dimension is a generalization component of RL testbeds.
Machado et al.
(2018) examined determinism in Atari environments and proposed the addition of sticky actions and variants on Atari games they call ﬂavors.
Others have created RL testbeds with generalization in mind using procedural generation of environments (Chevalier-Boisvert et al., 2018; Cobbe et al., 2020).
Multitask benchmarks have also cropped up in robot learning (Yu et al., 2020b).
In our work, we make use of the Arcade Learning Environment as formulated by Machado et al.
We also introduce task variants for transfer in AndroidEnv (Toyama et al., 2021).
2. M ULTITASK M ETHODS
An entire body of research is dedicated to multitask learning, which approaches ranging from metalearning (Finn et al., 2017) to gradient surgery (Yu et al., 2020a).
Many such approaches are not necessarily speciﬁc to RL and can be applied to a range of supervised tasks like vision.
Hierarchical approaches that can exploit common structures across tasks are also well-suited to multitask learning (Sutton et al., 1999).
Our work, in contrast, does not optimize for multiple tasks at the same time; tasks are seen sequentially, and there is a clear target task in mind at any given point in training.
2. T ASK T RANSFER AS AN E VALUATION P ROTOCOL
Some recent work has created evaluation protocols for task transfer.
Farebrother et al.
(2018) looks at DQN generalization with experiments on pretraining and transfer between Atari game ﬂavors.
For consistency, we use the same set of ﬂavors in our Atari experiments, though our focus is on types of policy updates rather than generalization techniques like dropout and regularization.
Parisi et al.
(2021) also crafts a transfer evaluation framework in RL.
The framework involves pretraining a self-supervised exploration policy which is used to bias actions of a task policy trained from scratch at transfer time.
The framework is implemented in 2D navigation settings (MiniGrid and AI Habitat).
Our work, in contrast, trains a single policy using extrinsic rewards across pretraining and transfer, and we present results in more dynamic settings (Atari and AndroidEnv).
M ETHOD
Our approach has two primary components: ﬁrst, a transfer framework for training a policy (Section 3.1), and second, a policy update formulation used after transfer (Section 3.2)
. 3. T RAINING
Training consists of 2 phases.
First, pretrain an RL agent in the default version of an environment.
Second, put the agent in a different ﬂavor of the environment while training with modiﬁed policy updates (described in Section 3.2).
This sequence of events yields a number of evaluation points, from random to zero-shot generalization to ﬁnetuning.
Such evaluation is beneﬁcial for understanding the performance not only of the agent but also of the environment itself.
We use IMPALA (Espeholt et al., 2018) for our RL agent, which enables the decoupling of action and evaluation.
3. P OLICY U PDATES
[Figure]
(Espeholt et al., 2018).
At transfer time, we independently update or freeze these components.
We apply a set of modiﬁed policy updates during transfer time.
As depicted in Figure 1, we split the IMPALA policy network into 2 parts: an embedding and a policy head.
The embedding is a deep stack of convolutional layers which takes in an image observation and outputs a 256-dimensional vector.
The head represents the temporal and action aspects of the agent, consisting of an LSTM (Hochreiter &
Schmidhuber, 1997) and fully connected layer before producing a policy and value function.
The architecture is based on the large network from Espeholt et al.
(2018), which containsconvolutional layers and 1.6 million parameters.
Our only change is the addition of a fullyconnected layer to the head after the LSTM.
At transfer, we activate our modiﬁed policy updates, which independently learn or freeze each model component.
This yields 4 formulations (freeze or not freeze head, freeze or not freeze embedding), one of which is trivial (freeze entire policy, zero-shot evaluation).
E XPERIMENTS
In our experiments, we set out to answer 2 overarching questions:
• How does freezing the policy embedding affect performance, on both the current environment ﬂavor and the pretrained one?
• What kind of generalization are Atari game ﬂavors and AndroidEnv variants testing?
The ﬁrst question is agent-centric, while the second is environment-centric.
In this section, we present experiments tackling these areas on Atari (Section 4.1) and AndroidEnv (Section 4.2).
All experiments are run with 5 random seeds, and plots show mean and conﬁdence across these seeds.
4. A TARI E XPERIMENTS
In our Atari experiments, we use the set of games and ﬂavors from Farebrother et al.
(2018), comprised of 9 transfer tasks across 4 games (Figure 2).
Experiments on the game Freeway, in line with prior work, did not yield signiﬁcant learning due to the sparsity of the reward, so we focus our analysis on the 3 remaining games.
We pretrain the agent on the default game ﬂavor (mode 0 difﬁculty 0, abbreviated m0d0) for 50M steps, followed by an additional 50M steps of training in another ﬂavor.
The top row shows default ﬂavors used for pretraining and the bottom row shows variants seen during transfer: Freeway’s mode 1 adds buses, more vehicles, and increases velocity; Hero’s mode 1 starts the agent at level 5; Breakout’s mode 12 hides all bricks unless the ball has recently collided with a brick; Space Invader’s difﬁculty 1 widens the agent’s spaceship and oscillates the red shields.
4.1. A TARI T RANSFER R ESULTS
Figure 3 shows reward curves after transfer to a new Atari ﬂavor.
Updating the entire policy performs best while freezing the entire policy shows the ﬁxed performance after pretraining.
The most interesting ﬁnding from these results is that across the 3 Space Invaders ﬂavor transfers, freezing the policy head performs better than freezing the embedding.
Freezing the embedding is commonly seen in RL, especially in settings where we expect the observation space to be similar.
From Figure 2 right, Space Invaders difﬁculties 0 and 1 appear to have similar observations (only the size of the agent differs).
This makes it even more surprising that updating the embedding is so crucial to performance ( especially in difﬁculty 1 transfer, Figure 3 bottom left).
From this result, we hypothesize that the embedding is learning task-centric representations and conclude that RL researchers should be more wary of freezing an embedding network.
[Figure]
4.1. A TARI E NVIRONMENT F LAVOR M ETRICS
Comparing the policy update approaches yields a general framework for environment difﬁculty and similarity scores.
We compare episode rewards from 5 points: 1) trained from scratch on target ﬂavor for 50M actor steps, 2) random policy, 3) pretrain then freeze the embedding, 4) pretrain then freeze the head, and 5) pretrain then update everything.
In Table 1, we display 2-5), all normalized by 1).
Freeway agents rarely yield nonzero rewards, so the Freeway metrics do not shed much light.
The normalized random policy reward is a metric for environment difﬁculty.
Random policies on the Breakout and Hero ﬂavors receive 2% of the rewards of the policy trained on the target; random policies on Space Invaders receive 10-30% (suggesting that Space Invaders is an easier game).
The normalized pretrain then update everything (‘update all’) is a metric for ﬂavor similarity – if the ﬂavors are more similar, then pretraining should provide a boost in target performance.
Pretraining provides a 2-4x boost in performance over training from scratch, with the exception of Hero Mode 2 – here pretraining actually hurts performance.
The normalized metrics for pretraining and freezing either the embedding or head tell us how two environments are similar or different.
There are 2 distinct ﬁndings here.
First, as seen on Hero mode 1, freezing the embedding is twice as effective as freezing the head.
Second, on Space Invaders difﬁculty 1, freezing the head is three times as effective as freezing the embedding.
We present these metrics not only to illustrate our particular results but also as a general framework for evaluating multitask and transfer benchmarks.
These metrics could simply be computed once upon release of a benchmark to show what kinds of challenges it presents.
[Table]
Presented at the Generalizable Policy
[Figure]
D ODGE
C ATCH
We also evaluate our policy learning strategies and evaluation framework on AndroidEnv, a new testbed that is challenging for RL due to its lack of determinism and complex action space.
We present results on the Catch task as depicted in Figure 4.
We introduce 4 variants of this task as ﬂavors.
We use an action wrapper that discretizes the pointer location into a 9x6 grid, as done in prior work.
The image observations are downsampled by 4x in x and y for a resolution of 120x80.
We run IMPALA with 5 actors.
[Figure]
Figure 5 shows the reward curves during transfer to Catch variants after 5M steps of pretraining in the default ﬂavor.
Reinforcing a ﬁnding from the Atari experiments, on the large ball variant (Figure 5 left), freezing the embedding performs signiﬁcantly worse than freezing the policy head or updating the policy altogether.
In fact, freezing the policy head does not seem to degrade performance at all.
Other variants do not show as much distinction between policy update strategies.
We hypothesize that this is due to the simplicity of the Catch task and the binary nature of its reward.
In future work, we hope to expand these experiments to more nuanced tasks in AndroidEnv.
[Figure]
[Equation]
[Figure]
than freezing the policy head.
A NDROID E NV E NVIRONMENT F LAVOR M ETRICS
Table 2 shows the metrics for AndroidEnv across the 4 Catch ﬂavors.
Note that since the catch rewards can be positive and negative, we do not normalize the metrics as we did for Atari ﬂavors.
Each number is the mean episode reward for that strategy, averaged over 100 episodes from that point in training.
‘Freeze all’ shows the initial performance of the pretrained policy right after transfer, which gives us a similarity metric: ‘small paddle’ does almost as all initially as after ﬁnetuning (the other ﬂavors yield lower initial performance).
Updating all policy parameters performs well except in the large ball ﬂavor, where freezing the head surprisingly performs the best.
[Table]
D ISCUSSION
Our approach has two goals: ﬁrst, to analyze agent effectiveness in cases where part of the policy is frozen, and second, to analyze the underlying environment testbeds.
Notably, our results show that freezing a policy embedding sometimes performs worse than freezing the policy head, opening doors for future work to explore the role of a policy embedding.
A clear takeaway from this result is that researchers should think twice before using a pretrained embedding without updating it in concert with the rest of the policy.
Next, on the environment-centric side of our analysis, we presented a collection of metrics to analyze environment ﬂavors, their difﬁculties, and their similarities to one another.
Such metrics hold promise for designers of future benchmarks or anyone hoping to gain a better understanding of existing environments.
In a ﬁeld struggling with sample efﬁciency as we move towards more complex applications, RL researchers must increase focus on generalization and ﬁnetuning to make these problems tractable; our work makes simple steps towards these goals