Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning
Abstract
Learning generalizable policies that can adapt to unseen environments remains challenging in visual Reinforcement Learning (RL).
Existing approaches try to acquire a robust representation via diversifying the appearances of in-domain observations for better generalization.
Limited by the specific observations of the environment, these methods ignore the possibility of exploring diverse real-world image datasets.
In this paper, we investigate how a visual RL agent would benefit from the off-the-shelf visual representations.
Surprisingly, we find that the early 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 yet effective framework that can generalize to the unseen visual scenarios in a zero-shot manner.
Extensive experiments are conducted on DMControl Generalization Benchmark, DMControl Manipulation Tasks, and Drawer World to verify the effectiveness of
PIE-G. Empirical evidence suggests PIE-G can significantly outperforms previous state-of-the-art methods in terms of generalization performance.
In particular, PIE-G boasts a 55 % generalization performance gain on average in the challenging video background setting.
1. Introduction
Visual Reinforcement Learning (RL) has achieved significant success in learning complex behaviors directly from image observations (Mnih et al., 2015;
Kalashnikov et al., 2018; Kostrikov et al., 2020).
Despite the progress, RL
First Workshop of Pre-training: Perspectives, Pitfalls, and Paths Forward at ICML 2022 , Baltimore, Maryland, USA, PMLR 162, 2022.
Copyright 2022 by the author(s).
[Figure]
agents are often plagued by the overfitting problem (Song et al., 2019), especially in high-dimensional observation space.
Previous studies show that it is difficult for the visual agents to generalize to unseen scenarios (Cobbe et al., 2019;
Lee et al., 2019), which severely limits their deployment in real-world applications.
In general, visual RL methods rely on their encoders to learn a visual representation to perceive the world.
Recent studies have found that data augmentation (Shorten & Khoshgoftaar, 2019) leads to more generalizable representations so that the agents can adapt to the unseen environments with different visual appearances (Raileanu et al., 2021;
Wang et al., 2020).
However, most of those approaches only augment the observations of the training environments (Kostrikov et al., 2020;
Laskin et al., 2020; Srinivas et al., 2020), which is unable to provide enough diversity for generalization over large domain gaps.
Furthermore, naively applying data augmentation may damage the robustness of learned representations and decrease training sample efficiency (Hansen et al., 2021;
Yuan et al., 2022).
To overcome these drawbacks, what we require is a universal representation that can generalize to a variety of unseen scenarios.
Recent works in representation learning demonstrate promising results in enabling pre-trained models to [Table] provide strong priors for downstream tasks (He et al., 2020;
Devlin et al., 2018).
The pre-trained models 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 features despite the diversity and the variability (Donahue et al., 2014).
Based on the observations, we would like to ask the following question: is it possible to train a visual RL agent that is augmented with pre-trained visual representations so that it can better generalize to novel tasks?
Towards answering the question, the main contribution of this paper is a surprising discovery that the off-the-shelf features of frozen models trained with ImageNet can be used as universal representations for visual RL.
Based on such findings, we present P retrained I mage E ncoder for
G eneralizable visual reinforcement learning (PIE-G), a visual RL framework that allows agents to obtain enhanced generalization ability via integrating the extracted representations from a pre-trained ResNet (He et al., 2016) encoder into RL training.
Straightforward as the framework appears,
PIE-G enjoys thoughtful details and nuanced design choices to acquire representations that are suitable for control and generalizable to novel scenarios.
Specifically, we show that the choice of early layer features and the ever-updating
Batch Normalization (BatchNorm) (Ioffe & Szegedy, 2015) are crucial for the performance gain.
2. Method
In this section, we introduce PIE-G, a simple yet effective framework for visual RL which benefits from the pre-trained encoders on other domains to facilitate the generalization ability.
2.1. Pre-trained Encoder
PIE-G explicitly leverages the pre-trained models as the representation extractor without any modification.
The pretrained encoder projects high-dimensional image observations into low-dimensional embeddings that are later used by RL policies.
Note that PIE-G is as simple as importing a pre-trained ResNet model from the torchvision (Marcel & Rodriguez, 2010) library.
This avoids the design of any auxiliary tasks to acquire useful representations.
For all the training tasks on different benchmarks, the encoder’s parameters are frozen to obtain universal visual representations.
Since the pre-trained model contains the priors from a wide range of real-world images, we hypothesize that the inherited power from a pre-trained model may help to capture and distinguish the main components of different tasks’ observations regardless of the changes of visual appearances or deformed shapes, and will further improve the 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 to the policy networks (Shang et al., 2021; Pari et al., 2021).
Moreover, we enable BatchNorm (Ioffe & Szegedy, 2015) to keep updating the running mean and running standard deviation during the policy training.
The key findings are: 1) early layers of a neural network would provide better representations for visual RL generalization, which resonates with prior works in imitation learning (Parisi et al., 2022); 2) the always updating statistics in BatchNorm helps better adapt to the shift in observation space and thus improve the generalization ability. [Figure]
2.2. Reinforcement Learning Backbone
We implement DrQ-v2 (Yarats et al., 2021) as the base visual reinforcement learning algorithm.
DrQ-v2 is the state-ofthe-art method for visual continuous control tasks, which adopts DDPG (Lillicrap et al., 2015) coupling with clipped
Double Q-learning (Fujimoto et al., 2018) to alleviate the overestimation bias of target Q-value.
We emphasize that PIE-G does not need any other proprioceptive states and sensory information as the inputs besides the representations extracted from original image observations.
In the setting of generalization, we follow stronger augmentation methods (e.g., mixup) of SVEA (Hansen et al., 2021) and DrAC (Raileanu et al., 2021) to further boost the performance.
It is worth mentioning that since the gradient is stopped before it reaches the encoder, all the data augmentation techniques discussed here do not affect the pre-trained visual representation.
Meanwhile, unlike Rutav et al.
(Shah & Kumar, 2021) and Simone et al.
(Parisi et al., 2022), we purely train the agent in a standard RL paradigm without any expert’s demonstration.
3. 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 moving or unseen video backgrounds, and deformed shapes of robots?
(2) How do the choice of layers in the encoder and the use of BatchNorm affect the performance?
3.1. Evaluation on Generalization Ability
Generalization on unseen or moving backgrounds.
We then evaluate PIE-G on the more challenging settings: video easy and video hard in DMC-GB.
The video hard setting consists of more complicated and fast-switching video backgrounds that are drastically different from the training environments.
Notably, even the reference plane of the ground is removed in this setting.
The comparison results are shown in Table 1.
PIE-G achieves better or comparable performance with the prior state-of-the-art methods in 9 out of 12 instances.
In particular, PIE-G gains significant improvement in the video hard setting over all the previous methods with +55% improvement on average.
For example, in the Finger Spin ,
Cup Catch , and Walker Walk tasks, PIE-G outperforms the best 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 extracted from 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 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 feature maps, and the intensity is decreased by 50.9% with
PIE-G than that with SVEA.
[Figure]
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 3.
Figure 4(a) demonstrates
[Table]
that PIE-G also improves the agents’ generalization ability with various shapes while other methods could barely generalize to these changes.
We attribute this to the lack of shape changing in previous data augmentation techniques.
Conversely, our pre-trained encoder is learned from a multitude of real-world images with various poses and shapes, thus enhancing its generalization ability on deformed shapes.
Furthermore, we conduct experiments on the DrawerWorld benchmark to test the agent’s generalization ability in manipulation tasks with different background textures.
Success Rate is adopted as the evaluation metric for its goalconditioned nature.
Table 2 illustrates that PIE-G can achieve better or comparable generalization performance in all the settings with +24% boost on average while other approaches may suffer from the CNN’s sensitivity in the face of various textures (Geirhos et al., 2018).
[Figure]
3.2. Choice of layers
In convolutional neural networks, the later layers capture high-level semantic features, while the early layers are responsible for extracting low-level information (Ma et al., 2015; Zeiler & Fergus, 2014; Lin et al., 2017).
Tableinvestigate how much control tasks can benefit from the features extracted from different layers.
As shown in Figure 5, the early layers preserve rich details of edges and corners, while the later layers only provide very abstract information.
[Figure]
Layer Layer
Intuitively, for control tasks, a trade-off is required between low-level details and high-level semantics.
Table 3 shows that the Layer 2 gains better performance than the other layers.
[Table]
3.3. Batch Normalization
Batch Normalization (BatchNorm) (Ioffe & Szegedy, 2015) is a popular technique in computer vision.
However, it is not widely adopted in RL algorithms.
In contrast to conventional wisdom, BatchNorm is found to be useful and important in PIE-G. Specifically, we find that calculating the mean and variance of the observations during evaluation rather than using the statistics from training data would boost the performance.
Figure 4(b) demonstrates that, in the most challenging settings, PIE-G with the use of BatchNorm can further improve the generalization performance.
This is largely because the distribution of observations 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.
(Ioffe &
Szegedy, 2015) that recomputation of the statistical means and variances allows the BatchNorm layer to generalize to new data distributions.
4. Conclusion
In this work, we propose PIE-G, a simple yet effective framework that leverages off-the-shelf features of ImageNet pre-trained ResNet models for better generalization in visual RL.
Extensive experiments on a variety of tasks in three RL environments confirm the merits of universal visual representations, which endow the agents with better generalization performance.
In addition, we show that the choice of layers and the use of BatchNorm are crucial for the performance gain.