U NDERSTANDING THE G ENERALIZATION G AP IN V ISUAL R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep Reinforcement Learning (RL) agents have achieved superhuman performance on several video game suites.
However, unlike humans, the trained policies fail to transfer between related games or even between different levels of the same game.
Recently several techniques have been proposed to close the generalization gap such as data augmentation, domain invariant feature learning, separation of actor and critic networks, etc.
However, the transfer performance of RL agents still remains unsatisfactory.
In this work, we conduct a large scale empirical study using procedurally generated video games to understand why the generalization gap still exists.
We also show that simple auxiliary tasks can improve generalization of policies.
Furthemore, contrary to the belief that adaptation to new levels requires ﬁnetuning all layers of the policy network, we ﬁnd that features in the visual trunk can be kept ﬁxed and only the parameters that use visual features to predict actions require ﬁnetuning.
Finally, to inform fruitful avenues for future research, we construct simple oracle methods that close the generalization gap.
I NTRODUCTION
Deep Reinforcement Learning (RL) has achieved tremendous success on several video game suites (Mnih et al., 2013; Vinyals et al., 2019; Berner et al., 2019).
However, state-of-the-art (SOTA) RL agents are usually trained and tested on the same game.
It turns out that the learned policy fails to transfer to different games or even unseen levels of the same game (Cobbe et al., 2018; 2019).
In contrast, humans are remarkable at transferring to new tasks and environments.
Behavior studies (Dubey et al., 2018) show that humans make extensive use of visual priors such as object permanence, sub-goals, and intuitive physics when learning new games.
Without these priors, data efﬁciency of humans drops an order of magnitude (Dubey et al., 2018).
(Sax et al., 2018) learns visual priors with supervised learning on datasets with manual annotations of several semantically meaningful features.
In practice however, acquiring such datasets is challenging.
In the absence of labelled datasets and human supervision, past works have explored three main ways to incorporate visual priors to improve generalization.
The ﬁrst paradigm is data augmentation (Michael et al., 2020;
Kostrikov et al., 2020) where invariance in the learned representations is induced by training the agent on a large set of task-irrelevant variations of the observed data.
Since the distinction between what is relevant vs irrelevant is task-dependent, the optimal choice of data augmentation is also task dependent.
To avoid manual selection of data augmentations,
UCB-DrAC (Raileanu et al., 2020) used upper conﬁdence bound (UCB) algorithm (Auer, 2002) for choosing optimal augmentations.
However, it remains unclear whether choosing one out of many task-agnostic data augmentations, such as cropping, helps in closing the generalization gap.
In this paper, we show that set of transformations used in state-of-the-art methods are insufﬁcient at closing the generalization gap.
On the other hand, we also show that task-informed data augmentations close the generalization gap, but require task variation information during training.
How to construct task sensitive data augmentations remains an open question and an avenue of future research.
We discuss this ﬁnding in Section 4.
The second paradigm for transfer learning is to leverage domain confusion (Tzeng et al., 2014; 2015; Hoffman et al., 2013), technique commonly used in computer vision to make the source data distribution indistinguishable from the target distribution.
The main idea is to discourage learning of spurious domain/level-speciﬁc features by enforcing the constraint that it should not be possible
[Figure]
to predict the identity of the domain/level from the policy features.
iDAAC (Raileanu & Fergus, 2021) used this technique to improve policy generalization.
However, our investigation reveals that even policies that generalize well contain domain/level speciﬁc information.
In hindsight, this is not surprising. On image classiﬁcation tasks, the identity of the object doesnot change based on factors such as background clutter, size, etc. However, small changes in the size of obstacles or platforms can drastically change the optimal policy. Therefore, the policy must encode information about level-speciﬁc game layout, making application of domain invariance to policy learning challenging.
Details of this investigation are provided in Section 5.
The third paradigm of improving generalization involves using auxiliary prediction tasks (such as depth prediction, image reconstruction, contrastive learning) to regularize the policy features to prevent learning spurious features (Jaderberg et al., 2016; Srinivas et al., 2020).
In this work, we choose the auxiliary task of inverse model prediction (Agrawal et al., 2016) due to its simplicity.
This improves policy generalization.
We detail this ﬁnding in Section 6.
It is well-known that choice of hyper parameters has a substantial impact on the performance of
RL agents.
When it comes to policy generalization, discussions around hyperparameter selection have been limited.
We show that state-of-the-art (SOTA) methods that improve policy generalization require careful hyperparameter selection for each task.
When these methods are constrained to choose one set of hyperparameters across tasks, they donot outperform the base PPO algorithm.
We detail this ﬁnding in Section 7.
Finally, we probe if deep RL algorithms learn spurious features that deters performance of the same task in new environments.
Our experiments reveal that visual features learned by training on a limited set of levels of one game donot require adaption to achieve good performance on new levels.
Infact, only the layers of the policy network that convert visual features into actions require ﬁnetuning.
Details of this investigation are provided in Section 8.
In summary, our experiments on Procgen (Cobbe et al., 2019), a suite of procedurally generated games reveal that (i) task-informed data augmentation closes the generalization gap; (ii) learning levelinvariant (or domain-invariant) features is not necessary for good generalization; (iii) regularizing the policy features using simple auxiliary tasks can improve generalization; (iv) SOTA methods that improve policy generalization rely on careful hyperparameter selection; (v) training on a limited number of levels of the same game doesnot result in learning of level-speciﬁc spurious features. When operating on new levels, good performance can be achieved by only ﬁne-tuning two layers that transform these features into actions.  R ELATED W ORK
Prior works (Cobbe et al., 2018; 2019; Justesen et al., 2018; Zhang et al., 2018a;b; Juliani et al., 2019; Rajeswaran et al., 2017; Raileanu & Rocktäschel, 2020; Grigsby & Qi, 2020; Kuttler et al., 2020; Farebrother et al., 2018) have established the problem of overﬁtting in reinforcement learning. In response, multiple methods have been proposed to mitigate overﬁtting.
(Cobbe et al., 2018; 2019) showed that classical techniques such as dropout (Srivastava et al., 2014), L2 regularization, and batch normalization (Ioffe & Szegedy, 2015) – originally developed for supervised learning – reduce the generalization gap in RL.
Other works (Michael et al., 2020; Cobbe et al., 2018; Ye et al., 2019;
Raileanu et al., 2020; Wang et al., 2020) have made use of data augmentation to learn policies that generalize well.
Another popular approach is to use representation learning techniques such as variational information bottleneck (Igl et al., 2019) and bisimulation metrics (Agarwal et al., 2021;
Zhang et al., 2020) to help with the problem of overﬁtting in RL.
(Cobbe et al., 2020; Raileanu &
Fergus, 2021) learned a separate network for policy and value function and showed that it helps reduce the generalization gap in RL.
(Sax et al., 2018) showed that using mid-level visual features optimized for segmentation, depth, keypoints, surface normal prediction and others improve transfer.
In our work, rather than proposing a novel approach to solve overﬁtting, we analyze the limitations of the current techniques and propose avanues for improvements.
P RELIMINARIES
We consider a distribution of Partially Observable Markov Decision Process (POMDPs) p ( M ) such that M i =
( S i , O i , A
i , T i , Ω i , R i , γ ) ∼ p ( M ) .
We can think of different M i as instances of same task in different environments.
S i is the state space,
O i is the observation space,
A i is the action space, T i ( s (cid:48) | s, a ) is the transition function, Ω i
( o (cid:48) | s, a ) is the observation function,
R i ( s, a ) is the reward function and γ is the discount factor.
The goal is to ﬁnd a policy π θ that maximizes the expected sum of discounted rewards over the distribution of POMDPs, J ( π θ ) =
E π,p ( M ) (cid:104)(cid:80) T − 1 t =
γ t R i ( s t , a t ) (cid:105) .
During training, we only have access to a limited number of POMDPs ˆ M =
{M i } N i =1 .
Our goal is to ﬁnd a policy π θ that generalizes to new POMDPs sampled from p ( M ) .
3. E NVIRONMENT S ETUP
We conduct experiments using the ProcGen suite (Cobbe et al., 2019), a collection of 16 procedurally generated video games.
Each game contains multiple levels wherein the agent needs to perform the same task, but the environment varies due to changes in position of objects, their textures and the background texture.
We refer to variation in levels resulting solely from changes in texture of object, agent and the background as theme variation.
We collectively refer to variation in levels resulting from changes in positions of the objects as layout variation.
Figure 1 illustrates theme and layout variation for two games, Jumper and Climber . For each game, the ﬁrst and the second level differ in theme whereas the ﬁrst and the third level differ in both the theme and the layout.
3. M ETHODS
We use standard neural network architecture to represent the policy, a ∼ π θ ( a | s ) , consisting of: a visual encoder z
= π encθ 1 ( s ) and a policy head a ∼ π hθ 2 ( a | z ) with parameters θ 1 , θ 2 respectively.
Collectively, they are referred as θ =
[ θ 1 , θ 2 ] .
The architecture of the visual encoder is borrowed from IMPALA (Espeholt et al., 2018) and is a ResNet with 15 convolutional layers.
It outputs a ﬂattened feature map ( z ) that we refer to as the visual feature .
These features are passed into a two 2-layer fully-connected neural networks with 256 hidden units for predicting the actions (i.e., the policy head) and the value function (cid:0) i.e., the value head, V h ψ ( z ) (cid:1) .
We analyze the generalization performance of PPO (Schulman et al., 2017) and three state-of-the-art (SOTA) methods:
• UCB-DrAC (Raileanu et al., 2020) uses Upper Conﬁdence Bound (UCB) to automatically select game speciﬁc data augmentations out of eight possible ones: crop , grayscale , cutout , cutout-color , ﬂip , rotate , random convolution and color jitter .
For brevity, we sometimes refer to UCB-DrAC as simply DrAC.
• DAAC (Raileanu & Fergus, 2021) trains two separate networks for representing the policy and the value function. The intuition is that in POMDPs, the value network will overﬁt, but the policy can still generalize.
The seperation of networks prevents interference in visual features learned for representing the value function and the policy. PPG (Cobbe et al., 2020) also trains separate networks for actor and critic but mainly focuses on improved training of policies and not their generalization performance. • iDAAC (Raileanu & Fergus, 2021) adds an auxiliary loss to the training objective of DAAC to encourage domain invariance. [Figure]
We use PPO’s pytorch implementation (Kostrikov, 2018), the ofﬁcial implementation for UCBDrAC (Raileanu et al., 2020), and the ofﬁcial implementation for DAAC/iDAAC (Raileanu & Fergus, 2021).
See Appendix B for the hyperparameters used.
3. E VALUATION M ETRICS
Different procgen games have different score ranges which makes it difﬁcult for us to average the scores across games.
Recent works (Raileanu et al., 2020; Raileanu & Fergus, 2021) used percentage improvement over PPO for averaging across games.
However, this metric doesn’t properly take reward range of a game into account and is also biased by PPO’s performance on that game, thereby being more susceptible to large variations.
To remedy this, as suggested in (Cobbe et al., 2019), we use the normalized score
R − R min R max − R min to compute mean scores (MNS) across games.
Here, R max and R min are game-dependent constants provided in (Cobbe et al., 2019).
To calculate the ﬁnal train and test performances of a policy, we evaluate the policy 100 times on train levels (1-200) and test levels for each game.
We report the mean and standard deviation of ﬁnal train and test performances computed across 5 different seeds. We train the policy on the ﬁrst two hundred levels and evaluate it on a subset of remaining levels (i.e. approximately two million unique levels). [Table]
I NVESTIGATING POLICY GENERALIZATION WITH DATA AUGMENTATION
A policy can fail to generalize to a new level due to inability to deal with either: (i) layout variation or (ii) theme variation.
To tackle the second issue, prior work (Raileanu et al., 2020; Michael et al., 2020) has employed data augmentation.
If these methods are successful at achieving theme in-variance, then their performance on test-levels should be same as a policy that is trained and evaluated on
ProcGen games with ﬁxed themes.
To test if this is true, we trained UCB-DrAC, a state-of-the-art data augmentation method on ProcGen games with ﬁxed themes. We refer this to oracle method as
DrAC+FT, where FT stands for ﬁxed theme . Results in Figure 2 show a signiﬁcant test performance difference between UCB-DrAC and DrAC+FT which suggests that UCB-DrAC does not achieve theme invariance.
Moreover, as shown in Figure 3, UCB-DrAC does not outperform a well-tuned PPO baseline when averaged across all the games. The above results raise a question: is it possible to close the generalization gap by adopting a better data augmentation scheme?
To answer this, we constructed game levels that had a ﬁxed layout, but varied only in theme.
We refer to a PPO agent trained on this baseline as DrAC-TR for theme randomization .
Our intent is to use this domain randomization scheme, that makes use of privileged access to the game engine, as a proxy of better data augmentation schemes.
Figure 2 shows that DrAC+TR greatly improves over UCB-DrAC and gets closer to the test performance of DrAC+FT.
Our result demonstrates the viability of pushing for stronger data augmentation schemes for producing theme invariance and thus policy generalization.
We hope to encourage the development towards this venue of research with our reported ﬁndings.
I NVESTIGATING POLICY GENERALIZATION WITH DOMAIN CONFUSION
Another way to learn generalizable policies is by using domain confusion (Tzeng et al., 2014; 2015;
Hoffman et al., 2013) which discourages policies from predicting task irrelevant properties of the environment, thereby encouraging them to focus on task relevant properties.
In addition to separating the policy and the value network, (Raileanu & Fergus, 2021) used this technique to prevent the policy from learning level-speciﬁc features and obtained SOTA results.
Given these ﬁndings, we might think that we should make policies level invariant to obtain better generalization.
However, a level (in a game) is not only deﬁned by its theme but also by its layout which contains task relevant information.
Therefore, learning level-invariant policies might remove layout features important for solving the task and thus hurt the performance of these policies.
Hence, it’s not clear if the policies, that generalize well, are independent of level-speciﬁc features.
To answer this, we ﬁrst take the policies trained by PPO on 100 k levels, by DAAC on 200 levels and by iDAAC on 200 levels on different games of procgen.
All these policies have good generalization capabilities.
We want to test if the visual features coming from these policies contain level-speciﬁc features.
Therefore, we collect 1 trajectory with random exploration for each level from 1 to(per game) and label each state in the trajectory with the corresponding level id.
We then mix all the (state, level id) tuples in the collected trajectories and call the resulting dataset D .
The ﬁrst 80% of D is labelled as train set and the remaining as the test set.
For each of the policies described above, we use the train set to learn a linear level classiﬁer on top of their visual features and evaluate the performance of the learned classiﬁer on the test set.
Table 1 shows that each of the classiﬁer learned obtained a good train and test performance averaged across all games.
We further provide the performance of these classiﬁers on individual games in the Appendix.
This implies that all the above mentioned policies, with good generalization capabilities, contain level-speciﬁc features.
To further bolster our claim, we use the train set to learn a linear classiﬁer on top of visual features of [Figure] policies trained by PPO on 200 levels and evaluate its performance on the test set.
We see that its performance is similar to that of the previous classiﬁers.
This shows that policies have level speciﬁc features regardless of whether it generalizes or not.
DAAC/iDAAC (Raileanu & Fergus, 2021) performs well due to less overﬁtting as a result of separation of policy and value networks and the regularization of policy features from its auxillary losses.
However, the insight that these methods don’t learn level-speciﬁc features and discouraging the policies to learn level-speciﬁc features leads to better generalization, isn’t empirically grounded.
I NVESTIGATING POLICY GENERALIZATION WITH AUXILIARY TASKS
Data augmentation and domain confusion can be thought of as two separate methods of constraining the learned visual features to potentially improve generalization.
A third way to constrain features is to optimize them for auxiliary objectives in addition to the policy and value prediction losses.
(Jaderberg et al., 2016) showed that depth prediction improves both data efﬁciency and the generalization of learned policies on ﬁrst-person navigation tasks.
However, depth is an additional sensory modality that is not always available.
Therefore, we propose a simpler auxiliary task of encouraging the visual features coming from two consecutive states to predict the corresponding action.
This inverse model regularization has been shown to be an effective regularization technique in prior works (Agrawal et al., 2016; Pathak et al., 2017; 2018).
Let z t
=
π enc θ 1 ( s t ) .
If f φ is the inverse model and ( s t , a t , s t +1 ) is the state action state [Figure]
tuple, then this regularization minimizes CE ( a t , f φ ( z t , z t +1 )) where CE is the cross-entropy loss.
Figure 3 shows that this regularization is competitive with UCB-DrAC in terms of test performance despite being simpler.
E FFECTS OF HYPERPARAMETER SELECTION ON POLICY GENERALIZATION
Until now we have investigated several methods for learning policies that generalize across levels.
However, with exception of UCB-DrAC, we found that all the other methods (PPO inv m, DAAC, iDAAC) require careful hyperparameter selection to improve over the base PPO algorithm.
To quantify the impact of careful hyperparameter selection, we chose a single set of hyperparameters for each method that maximizes the mean normalized score (MNS) evaluated on validation levels 201 − 400 .
We then compare these policies on the test levels.
Figure 4 shows that average test performance across games of all methods is similar to PPO. Additional details about hyperparameter selection are provided in Appendix B.
These results suggest that improvements over PPO heavily depends on careful hyperparameter selection on each game.
One could argue that different games require different hyperparameters due to their unique properties.
However, careful hyperparameter tuning on each game is compute-intensive, and there’s no guidance on how to set these hyperparameters other than performing a grid search.
Hence, automated tuning of hyperparameter to save compute costs is an important avenue for future research.
I NVESTIGATING POLICY ADAPTATION TO NEW LEVELS
Until now we have investigated zero-shot evaluation of policies on test levels.
The other possibility is to ﬁnetune the policies on new levels.
First, we train a policy on a limited number of levels (i.e., levels 1 − 200 ).
However, rather than directly evaluating the trained policy, we ﬁnetune the trained policy on new levels (i.e., levels 201 − 400 ) to achieve the train-time performance quickly.
In this setting, it is natural to wonder if ﬁnetuning of the entire policy is required to achieve the train-time performance on new levels.
To answer this question, we ﬁx the visual encoder π encθ 1 of a policy trained using PPO on levels 1 − 200 of a game and ﬁnetune only the policy head π hθ 2 (as well as the value head V ψ ) on levels 201 − 400 of the game.
If ﬁnetuning the entire policy on levels 201 − 400 is required to achieve the train-time performance, then ﬁnetuning only the policy head on those levels should lead to a sub-optimal performance.
Figure 5 shows that the average returns of the policy on levels 201 − 400 become similar to the average returns of the policy on levels 1 − 200 after ﬁnetuning the policy head.
Furthermore, ﬁnetuning the entire policy on levels 201 − 400 levels gives a similar average return.
These observations imply that ﬁnetuning only the policy head (on levels 201 − 400 ) is sufﬁcient to recover the average return of levels 1 − 200 .
Additionally, we see that end-to-end training of policy from scratch on levels 201 − 400 is less sample-efﬁcient than ﬁnetuning the policy head as well as ﬁnetuning the entire policy.
Moreover, training of policy head with randomly initialized visual encoder as well as imagenet (Krizhevsky et al., 2012) trained visual encoder on levels 201 − 400 leads to poor average returns.
These observations show that the visual encoder features learned from levels 1 − 200 are meaningful and help ﬁnetune both the policy head and the entire policy.
Overall, this implies that the visual features learned on training levels ( 1 − 200 ) are meaningful and can be kept ﬁxed when [Figure] [Figure] adapting to new levels ( 201 − 400 ).
Instead, it is the policy head feature that fails to transfer to new levels and requires further ﬁnetuning.
However, the strategy of training a policy on training levels ( 1 − 200 ) and then ﬁnetuning the policy head on new levels has limitations.
As the number of new levels increases, it becomes harder to ﬁnetune the policy head on those levels and recover the average policy returns of the training levels.
From Figure 6, we can see that ﬁnetuning the policy head on 200 and 1 k new levels recovers the average policy returns of the training levels.
But, it fails to do the same for 10 k and 100 k new levels.
D ISCUSSION
How to use visual features?
Most of the work on visual transfer in RL has focused on learning generalizable visual features.
However, another related problem setting is policy adaptation to new levels.
Section 8 shows that visual features learned by PPO can be kept ﬁxed when adapting to new levels but the policy head, tasked to combine visual features to predict actions, requires ﬁnetuning on new levels.
This suggests a fruitful avenue for future work: instead of learning better visual features, focus on how to leverage the learned features to solve a new task.
Towards better data augmentation.
UCB-DrAC performs similar to PPO baseline when averaged across all the games and doesn’t fully address the problem of learning invariant features.
Speciﬁcally, it is signiﬁcantly worse as compared to training PPO with ﬁxed texture.
This indicates that better data augmentation techniques are required to make the policy invariant to irrelevant visual appearances.
We show that one way to overcome this challenge is by combining UCB-DrAC with theme randomization.
However, theme randomization requires task-speciﬁc knowledge of how the theme varies across different variants of the task.
In general, such knowledge may not available to the agents.
An exciting area of future investigation is to develop data augmentation method that can automatically discover task-relevant augmentations.
Further investigation of auxiliary losses.
We show PPO, with inverse model regularization, signiﬁcantly improves upon PPO despite being a simple change.
Similarly, iDAAC made use of auxiliary losses to focus on task-relevant features of the environment.
This indicates the need for further investigation into auxiliary losses for better policy generalization.
Towards compute-efﬁcient hyperparameter tuning.
We show that state-of-the-art (SOTA) methods that improve policy generalization heavily rely on careful hyperparameter selection for each game.
However, hyperparameter tuning via grid search on each game is compute-intensive.
Therefore, automatically tuning hyperparameters in a compute-efﬁcient manner is a fruitful avenue for future works.
A PPO WITH I NVERSE M ODEL R EGULARIZATION
PPO (Schulman et al., 2017) is an actor-critic algorithm that alternates between sampling data with environment interaction and optimizing objective function with stochastic gradient descent.
Let the policy being optimizes be π θ and the associated value function be V θ,ψ .
PPO maximizes
[Equation]
where S π θ is the entropy for aiding exploration,
[Equation]
J V is the value function loss,
[Equation]
and J π θ is the policy objective term
[Equation]
Here, r θ = π θ ( a | s )
π θ old ( a | s ) is the importance weight for estimating the advantage function and ˆ A
=
ˆ
A θ old ( s, a ) is estimated advantage function for π θ old using returns and value function V θ,ψ ( s ) .
We make use of inverse models to regularize policy features for improved performance and better generalization.
Let z t = π encθ 1 ( s t ) .
If f φ is the inverse model and ( s t , a t , s t +1 ) is the state action state tuple coming from π θ old , then this regularization minimizes J inv =
CE ( a t , f φ ( z t , z t +1 ))
where
CE is the cross-entropy loss.
Overall, PPO with inverse model regularization maximizes
[Equation]
where α inv is the regularization constant associated with inverse model regurlarization loss.
B H YPERPARAMETERS
We use the default hyperparameters from PPO (Schulman et al., 2017), UCB-DrAC (Raileanu et al., 2020) and DAAC/iDAAC (Raileanu & Fergus, 2021).
DrAC-FT and DrAC-TR uses same hyperparameters as UCB-DrAC (DrAC).
PPO with inverse model regularization only introduces an extra hyperparameter α inv .
We write all the hyperparameters in table 2 and 3.
While table 2 describes the best single set of hyperparameters for different methods selected using validation levels 201 − 400 , table 3 details game-speciﬁc changes to some of those hyperparameters for improved performance.
C P ROCGEN R ESULTS
Table 4 and 5 describe scores on train levels (1-200) and test levels across all procgen games for
PPO, UCB-DrAC, DrAC-FT (UCB-DrAC with ﬁxed theme games), DrAC-TR (UCB-DrAC with randomized theme games), DAAC, iDAAC, and PPO
inv-m (PPO with inverse model regularization).
Figure 7 places games next to the RL algorithm on which they perform the best.
Some games are excluded as they perform the best on two or more RL algorithm.
Figures 8, 9, and 10 show the learning curves for ﬁnetuning the policy head as well as the full policy on 1 k, 10 k and 100 k new levels respectively.
The policy was originally trained on levels 1-200 with PPO.
End-to-end training of policy and training of policy head with randomly initialized visual encoder on new levels are also included as baselines.
Finally, table 6 and 7 give train and test accuracy of level classiﬁer trained on the encoder features of policies learned with PPO on 200 levels, with PPO on 100 k levels, with
DAAC on 200 levels, and with iDAAC on 200 levels on all procgen games.
[Figure] [Figure]
[Figure] [Figure]
[Table] [Table]
fruitbot  0.
bossﬁght  0.
jumper  0.
ninja  0.
starpilot  0.
coinrun  0.
heist  0.
caveﬂyer  0.
[Table] [Table]
[Table] [Table]