Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
Learning generalizable policies that can adapt to unseen environments remainschallenging in visual Reinforcement Learning (RL).
Existing approaches try toacquire a robust representation via diversifying the appearances of in-domain observations for better generalization.
Limited by the specific observations of theenvironment, these methods ignore the possibility of exploring diverse real-worldimage datasets.
In this paper, we investigate how a visual RL agent would benefit from the off-the-shelf visual representations.
Surprisingly, we find that theearly layers in an ImageNet pre-trained ResNet model could provide rather generalizable representations for visual RL.
Hence, we propose P re-trained I mage
E ncoder for G eneralizable visual reinforcement learning (PIE-G), a simple yeteffective framework that can generalize to the unseen visual scenarios in a zeroshot manner.
Extensive experiments are conducted on DMControl Generalization
Benchmark, DMControl Manipulation Tasks, and Drawer World to verify theeffectiveness of PIE-G. Empirical evidence suggests PIE-G improves sample efficiency and significantly outperforms previous state-of-the-art methods in termsof generalization performance.
In particular, PIE-G boasts a 55 % generalizationperformance gain on average in the challenging video background setting.
Project
Page: https://sites.google.com/view/pie-g/home .
Introduction
Visual Reinforcement Learning (RL) has achieved significant success in learning complex behaviorsdirectly from image observations [51, 38, 40].
Despite the progress, RL agents are often plaguedby the overfitting problem [70], especially in high-dimensional observation space.
Previous studiesshow that it is difficult for the visual agents to generalize to unseen scenarios [10, 43], which severelylimits their deployment in real-world applications.
In general, visual RL methods rely on their encoders to learn a visual representation to perceivethe world.
Recent studies have found that data augmentation [68] leads to more generalizablerepresentations so that the agents can adapt to the unseen environments with different visual appearances [61, 77].
However, most of those approaches only augment the observations of the trainingenvironments [40, 42, 71], which is unable to provide enough diversity for generalization over largedomain gaps.
Furthermore, naively applying data augmentation may damage the robustness of learnedrepresentations and decrease training sample efficiency [31, 87].
To overcome these drawbacks, what we require is a universal representation that can generalize to avariety of unseen scenarios.
Recent works in representation learning demonstrate promising results in
[Figure]
Buffer
GENET IM
[Figure]
enabling pre-trained models to provide strong priors for downstream tasks [34, 13].
The pre-trainedmodels contain representations obtained from a wide range of existing real-world image datasets.
These representations are proved to be robust to noises and capable of distinguishing salient featuresdespite the diversity and the variability [14].
Based on the observations, we would like to ask thefollowing question: is it possible to train a visual RL agent that is augmented with pre-trained visualrepresentations so that it can better generalize to novel tasks?
Towards answering the question, the main contribution of this paper is a surprising discovery that theoff-the-shelf features of frozen models trained with ImageNet can be used as universal representationsfor visual RL.
Based on such findings, we present P retrained I mage E ncoder for G eneralizable visualreinforcement learning (PIE-G), a visual RL framework that allows agents to obtain enhanced trainingefficiency and generalization ability via integrating the extracted representations from a pre-trained ResNet [33] encoder into RL training.
Straightforward as the framework appears, PIE-G enjoysthoughtful details and nuanced design choices to acquire representations that are suitable for controland generalizable to novel scenarios.
Specifically, we show that the choice of early layer features andthe ever-updating Batch Normalization (BatchNorm) [36] are crucial for the performance gain.
To validate the effectiveness of our framework, we conduct a series of experiments on 3 benchmarks:
DMControl Generalization Benchmark (DMC-GB) [29], DMControl Manipulation Tasks [75], and
Drawer World [78] that is modified from Meta World [86].
Empirical studies have demonstratedthat PIE-G achieves better or competitive results in training sample efficiency, and significantlyoutperforms previous state-of-the-art methods in generalization capability without bells and whistles.
Our main contributions are summarized as follows: (i)
We find that pre-trained encoders from offthe-shelf image datasets provide generalizable representations in visual RL.
(ii)
We propose PIE-G,a simple yet effective framework with a pre-trained encoder that can boost the sample efficiencyand generalization ability in visual RL.
(iii) PIE-G outperforms state-of-the-art methods in 3 visualgeneralization benchmarks by a large margin, with a 55 % boost on average on the hardest setting in
DMC-GB.
Related Work
Representation learning in RL.
A large corpus of literature has sought to leverage representationlearning in the setting of RL [64, 24, 71, 63, 85, 72].
More recently, there are a branch of methodsthat perform unsupervised pre-training to encourage exploration for better sample efficiency [56,47, 19, 46, 41, 32].
A typical approach is to use a contrastive learning method to jointly incentivize exploration and acquire useful representations [17].
Liu et al.
[47] introduce a new type of pretraining techniques via entropy maximization in embedding space for better exploration.
In particular,inspired by the new representation algorithm in computer vision, Yarats et al.
[83] propose a SwAVlike architecture [6] for pre-training with an exploration scheme that maximizes the entropy of thestate visitation distribution.
However, all these methods require the data collected in the targetenvironments, resulting in additional sample cost.
We instead propose a cross-domain pre-trainingway to improve visual RL agents’ performance without any in-domain interaction during the pretraining time.
Pre-trained visual encoders for RL.
Applying the pre-trained vision model from other domainsto the control tasks has gradually attracted researchers’ attention [84, 69, 39, 65, 18, 54, 80].
Forexample, Shah et al.
[66] and Parisi et al.
[55] suggest that with the help of experts’ demonstrations,the pre-trained ResNet [33] representations can achieve competitive performance with state-basedinputs.
Moreover, human video datasets are introduced to pre-train a visual representation fordownstream policy learning [52].
The pre-trained model has also been proposed for goal-specificationvia behavior cloning [11].
However, few works explore the effectiveness of pre-trained models forgeneralization.
In contrast to prior approaches, PIE-G enables agents to generalize well to the unseenvisual scenarios with a large distributional shift in a zero-shot manner while achieving high sampleefficiency in a standard RL training paradigm.
Generalization in visual RL.
Researchers have investigated to improve visual RL agents’ generalization ability from various aspects [1, 30, 43, 76, 3, 53, 81].
Data augmentation [43, 77, 31, 87, 20, 50]and domain randomization [74, 59, 57, 62, 7] are effective ways for generalization in visually differentenvironments.
Notably, Hansen et al.
[29] employ a BYOL-like[26] architecture to decouple theaugmentation from policy learning for better generalization.
In order to control the high variancewhen implementing data augmentation, Hansen et al.
[31] add a regularization term of the Q functionbetween un-agumented and augmented data as an implicit variance reduction technique.
Meanwhile,
Yuan et al.
[87] propose a task-aware data augmentation method with the Lipschitz constant [21]for maintaining training stability.
Fan et al.
[20] apply data augmentation in the imitation learningparadigm.
Prior works rely on data augmentation to gain a robust representation for each task.
Inthis work, we tackle this challenge from a different perspective, utilizing a single pre-trained encoderwith the universal visual representations for all the tasks.
Preliminaries
Reinforcement learning.
Due to partial state observability from images in visual RL [37], weconsider learning in a Partially Observable Markov Decision Process (POMDP) [4] formulated bythe tuple ⟨S , O , A , r, P , γ ⟩ where S is the state space, O is the observation space,
A is the actionspace, r :
S × A
(cid:55)→ R is a reward function, P ( s t +1 | s t , a t ) is the state transition function, andγ ∈ [0 , 1) is the discount factor.
The goal is to find a policy π ∗ to maximize the expected cumulativereturn π ∗ = arg max π E a t ∼ π ( ·| s t ) , s t ∼P (cid:104)(cid:80)
Tt =1 γ t r ( s t , a t ) (cid:105) , starting from an initial state s 0 ∈ Sand obtained by following the policy π θ ( · | s t ) which is parameterized by learnable parameters θ .
Generalization.
In terms of generalization, we consider a set of POMDPs: M = {M 1 , M 2 , ..., M n }that shares the same dynamics and structures.
The only difference among them is the observationspace
O .
This setting is more formally described as “Block MDPs” [16].
During the training process,we only have the access to a fixed POMDP denoted M i .
Our purpose is to train an agent on a specificscenario M i to learn a policy π ∗ G which can maximize the expected cumulative return over the wholeset of POMDPs in a zero-shot generalization manner.
Method
In this section, we introduce PIE-G, a simple yet effective framework for visual RL which benefitsfrom the pre-trained encoders on other domains to facilitate sample efficiency and generalizationability.
4.1 Pre-trained Encoder
PIE-G explicitly leverages the pre-trained models as the representation extractor without any modification.
The pre-trained encoder projects high-dimensional image observations into compact,low-dimensional embeddings that are later used by RL policies.
Note that PIE-G is as simple asimporting a pre-trained ResNet model from the torchvision [49] library.
This avoids the design of anyauxiliary tasks to acquire useful representations.
For all the training tasks on different benchmarks, the encoder’s parameters are frozen to obtainuniversal visual representations.
Since the pre-trained model contains the priors from a wide rangeof real-world images, we hypothesize that the inherited power from a pre-trained model may helpto capture and distinguish the main components of different tasks’ observations regardless of thechanges of visual appearances or deformed shapes, and will further improve the sample efficiencyand generalization abilities of RL agents.
To validate our hypothesis, we first encode each observation independently to obtain embeddings.
Then, the embeddings from the second layer of the pre-trained model are fused as input features tothe policy networks [67, 54].
Moreover, we enable BatchNorm [36] to keep updating the runningmean and running standard deviation during the policy training.
The key findings are: 1) earlylayers of a neural network would provide better representations for visual RL generalization, whichresonates with prior works in imitation learning [55]; 2) the always updating statistics in BatchNormhelps better adapt to the shift in observation space and thus improve the generalization ability.
Moredetailed discussion can be found in Section 5.4.
4.2 Reinforcement Learning Backbone
We implement DrQ-v2 [82] as the base visual reinforcement learning algorithm.
DrQ-v2 is thestate-of-the-art method for visual continuous control tasks, which adopts DDPG [44] coupling withclipped Double Q-learning [22] to alleviate the overestimation bias of target Q-value.
The agents aretrained with two Q θ k value functions and their corresponding target network Q ¯ θ k .
The critic lossfunction is as follows, and the mini-batch of transitions τ = ( s t , a t , r t : t + n − 1 , s t + n ) is sampled fromthe replay buffer D :
[Equation]
with n-step TD target y :
[Equation]
The actor π ϕ is trained with the following objective:
[Equation]
where s t is augmented by random shift, a t = π ϕ ( s t )
+ ϵ , ϵ is sampled from clip (cid:0) N (cid:0) 0 , σ 2 (cid:1) , − c, c (cid:1)with a decaying exploration noise σ .
Thanks to the efficiency of DrQ-v2, PIE-G enjoys faster wall-clock training time and fewer computational footprints.
We emphasize that PIE-G does not need any other proprioceptive states and sensoryinformation as the inputs besides the representations extracted from original image observations.
Inprevious works [40, 82, 61, 31], different schemes of data augmentation are proposed to improvesample efficiency and generalization performance.
In practice, weak augmentation methods (e.g.,random shift) of DrQ [40] and DrQ-v2 [82] are found to be beneficial for sample efficiency.
In thesetting of generalization, we follow strong augmentation methods (e.g., mixup) of SVEA [31] and
DrAC [61] to further boost the performance.
It is worth mentioning that since the gradient is stoppedbefore it reaches the encoder, all the data augmentation techniques discussed here do not affect thepre-trained visual representation.
Meanwhile, unlike Rutav et al.
[66] and Simone et al.
[55], wepurely train the agent in a standard RL paradigm without any expert’s demonstration.
[Table]
Experiments
In this section, we investigate the following questions: (1) Can PIE-G improve the agent’s generalization ability?
Specifically, how well does PIE-G deal with jittered color, moving video background,and deformed shapes of robots?
(2) Can PIE-G improve training sample efficiency?
(3) How do thechoice of layers in the encoder and the use of BatchNorm [36] affect the performance?
(4) Can furtherfinetuned RL agents outperform those with frozen visual encoders?
5.1 Setup
We evaluate our method on a wide range of tasks, including DMControl Generalization Benchmark (DMC-GB) [29], DeepMind Manipulation tasks [75], and Drawer World [78].
PIE-G is trainedfor 500k interaction steps with 2 action repeat and evaluated with 100 episodes for every task on thetesting benchmarks.
All the generalization evaluations are in a zero-shot manner.
By default, theencoder uses the ResNet18 architecture [33].
To be more specific, the feature maps of the second layerare flattened and passed through an additional fully connected layer to serve as the representations ofthe observations.
More training details and environment descriptions are in Appendix B.
5.2 Evaluation on Generalization Ability
We compare the generalization ability of PIE-G with state-of-the-art methods and strong baselines:
SAC [27]: a widely used off-policy RL algorithm; DrQ [40]: a SAC-based visual RL algorithm withaugmented observations; DrQ-v2 [82]: the prior state-of-the-art model-free visual RL algorithm interms of sample efficiency; SVEA [31]: the prior state-of-the-art method in terms of generalizationvia reducing the Q-variance through an auxiliary loss; TLDA [87]: another state-of-the-art method ingeneralization by using task-aware data augmentation.
Generalization on color-jittered observations.
The agent’s generalization ability is evaluated on
DMC-GB with randomly jittered color.
For the manipulation tasks, the colors of different objects (e.g.,floors, arms) are modified.
As shown in Table 1, PIE-G obtains better or competitive performance in9 out of 10 instances.
These results suggest that the visual representation from the pre-trained modelis more robust to the color-changing than the one trained by standard RL algorithms.
[Table]
Generalization on unseen and/or moving backgrounds.
We then evaluate PIE-G on the morechallenging settings: video easy and video hard in DMC-GB.
The video hard setting consists of morecomplicated and fast-switching video backgrounds that are drastically different from the trainingenvironments.
Notably, even the reference plane of the ground is removed in this setting.
The comparison results are shown in Table 2.
PIE-G achieves better or comparable performancewith the prior state-of-the-art methods in 9 out of 12 instances.
In particular, PIE-G gains significantimprovement in the video hard setting over all the previous methods with +55% improvement onaverage.
For example, in the Finger Spin , Cup Catch , and Walker Walk tasks, PIE-G outperforms thebest of the other methods by substantial margins 127.0% , 95.0% , and 59.2% respectively.
Attempting to explain the success, we visualize the difference of the normalized feature maps extractedfrom the encoder whose inputs are two Walkers of the same pose but with different backgrounds,as is shown in Figure 2.
Ideally, a well-generalizable encoder would map the observations of the
[Figure] [Figure]
two Walkers to exactly the same embedding, and therefore the difference should be zero.
In practice,as shown in Figure 2, the encoder of PIE-G produces a difference much closer to zero than that of
SVEA.
Numerically, we calculate the average pixel intensity in the difference of normalized featuremaps, and the intensity is decreased by 50.9% with PIE-G than that with SVEA.
Task
Setting
SAC
DrQ-v
SVEA
PIE-G
Drawer-Close
Training 100 % 98 % 70% 99 %
Wood 0% 32% 49% 59 %
Metal 0% 46% 69% 95 %
Blanket 0% 8% 72 % 71 %
Task
Setting
SAC
DrQ-v
SVEA
PIE-G
Drawer-Open
Training 98 % 100 % 75% 97 %
Wood 18% 2% 47% 79 %
Metal 35% 53% 71% 97 %
Blanket 28% 5% 37% 85 %
Furthermore, we conduct experiments on the Drawer Worldbenchmark to test the agent’sgeneralization ability in manipulation tasks with different background textures.
The visualizedobservations are shown in Figure 3a.
Success Rate is adoptedas the evaluation metric for itsgoal-conditioned nature.
Table 3 illustrates that PIE-G canachieve better or comparablegeneralization performance inall the settings with +24% booston average while other approaches may suffer from the CNN’s sensitivity in the face of varioustextures [23].
Generalization on deformed shapes.
To verify agent’s robustness in terms of the deformed shapes,we modify the shapes of the jaco arm and the target objects in the manipulation tasks, as shown in
Figure 4a.
Figure 4b demonstrates that PIE-G also improves the agents’ generalization ability withvarious shapes while other methods could barely generalize to these changes.
We attribute this to thelack of shape changing in previous data augmentation techniques.
Conversely, our pre-trained encoderis learned from a multitude of real-world images with various poses and shapes, thus enhancing itsgeneralization ability on deformed shapes.
[Figure] [Figure]
5.3 Evaluation on Sample Efficiency
We evaluate the sample efficiency of PIE-G on 8 relatively challenging tasks on the DeepMind Control
Suite and the Manipulation task ( Reach Duplo ).
Figure 5 demonstrates that PIE-G achieves betteror comparable sample efficiency and asymptotic performance than DrQ-v2 in 7 out of 8 tasks.
Toeliminate the effect of varied network size, we also include random enc , a baseline that has the samenetwork architecture with frozen random initialized parameters.
In addition, Figure 3b shows that
PIE-G is also superior over the baselines on the Drawer world benchmark.
The results demonstratethat the pre-trained encoder inherits the powerful feature extraction ability trained from ImageNetand acquires better training sample efficiency in control tasks.
5.4 Ablation study
To verify the necessity of the design choices in PIE-G, we conduct a series of ablation studies to takea closer look at the proposed method.
More results are shown in Appendix C.
Task
Layer Layer
Layer Layer
Walker Walk
[Table]
Cheetah Run
[Table]
Choice of layers.
In convolutional neural networks, the laterlayers capture high-level semantic features, while the early layers are responsible for extracting low-level information [48,88, 45].
Figure 6 and Table 4 investigate how much control tasks
[Figure]
Layer Layer
Layer can benefit from the features extracted from different layers.
As shown in Figure 6, the early layerspreserve rich details of edges and corners, while the later layers only provide very abstract information.
Intuitively, for control tasks, a trade-off is required between low-level details and high-level semantics.
Table 4 shows that the Layer 2 gains better performance than the other layers.
[Figure]
Batch normalization.
Batch Normalization (BatchNorm) [36] is a popular technique in computer vision.
However, it is not widely adopted in RL algorithms.
Incontrast to conventional wisdom, BatchNorm is found tobe useful and important in PIE-G. Specifically, we find thatcalculating the mean and variance of the observations during evaluation rather than using the statistics from trainingdata would boost the performance.
Figure 7 demonstratesthat, in the most challenging settings, PIE-G with theuse of BatchNorm can further improve the generalizationperformance.
This is largely because the distribution ofobservations is determined by the agent, violating the assumption of independent and identical distribution (i.i.d.).
This use of BatchNorm also reassures the recommendation from Ioffe et al.
[36] that recomputation ofthe statistical means and variances allows the BatchNorm layer to generalize to new data distributions.
[Figure]
Finetuning the pre-trained model.
We also conduct research to finetunethe encoder’s parameters instead ofkeeping it frozen.
Previous works [35,58] have found that finetuning pretrained models is challenging.
Consistent with these studies, Table 5 suggests that compared with the frozenrepresentations from pre-trained models, the finetuned representations suffer from the out-ofdistribution problem [8] and lead to a significant drop in terms of the generalization ability.
[Table]
other
Additionally, we investigate the efficacy of other visual representations.
MoCo-v2 [9] is a pre-trained modeloptimized via contrastive learning tolearn representations.
We find that
PIE-G with the pre-trained representations of MoCo-v2 can also obtain acomparable performance in terms of both the sample efficiency and generalization ability.
Moreresults are shown in Appendix C.
Conclusion
In this work, we propose PIE-G, a simple yet effective framework that leverages off-the-shelf featuresof ImageNet pre-trained ResNet models for better generalization in visual RL.
Extensive experimentson a variety of tasks in three RL environments confirm the merits of universal visual representations,which endow the agents with improved sample efficiency and better generalization performance.
Inaddition, we show that the choice of layers and the use of BatchNorm are crucial for the performancegain.
Our exploration may inspire more researchers to dig into the great potential of utilizingpre-trained representations in visual RL.
Limitations.
We study generalization in the simulated environments.
However, there might benew challenges in the real-world applications.
In the future, we would like to establish and test onbenchmarks of real-world scenarios.