Visual Pretraining via Contrastive Predictive Model for Pixel-Based Reinforcement Learning

Abstract

In an attempt to overcome the limitations of reward-driven representation learning in vision-based reinforcement learning (RL), an unsupervised learning framework referred to as the visual pretraining via contrastive predictive model (VPCPM) is proposed to learn the representations detached from the policy learning. Our method enables the convolutional encoder to perceive the underlying dynamics through a pair of forward and inverse models under the supervision of the contrastive loss, thus resulting in better representations. In experiments with a diverse set of vision control tasks, by initializing the encoders with VPCPM, the performance of state-of-the-art vision-based RL algorithms is significantly boosted, with 44% and 10% improvement for RAD and DrQ at 100 steps, respectively. In comparison to the prior unsupervised methods, the performance of VPCPM matches or outperforms all the baselines. We further demonstrate that the learned representations successfully generalize to the new tasks that share a similar observation and action space.

1. Introduction

Recent advances in deep reinforcement learning (RL) have allowed agents to perform complex control tasks directly from raw sensory observations. Remarkable successes have been achieved, ranging from learning to play video games from raw pixels [1], solving complex tasks from first-person-view observations [2,3], to autonomously performing robotic tasks [4,5,6,7].

As a standard practice, deep RL methods jointly learn a visual encoder and a policy in an end-to-end manner. In this paradigm, the visual representations are learned under the strong supervision of task-specific rewards. While the simplicity of end-to-end methods is appealing, learning the representations relying on the rewards can have several limits. First, the representations are acquired harder under sparse rewards, thus requiring more data for convergence. Additionally, in practice, the reward function is commonly designed and retested until a suitable function is selected, resulting in the representation learning process being repeated, thus it is inefficient. Furthermore, whenever the agent encounters the new tasks, the representation learning is performed again even if the environment appearance is identical across tasks.

In this work, we pursue an alternative paradigm as depicted in Figure 1, wherein the visual encoder is first pretrained without any reward supervision, detached from the policy learning, and then the learned task-agnostic representations are transferred for learning policy on a specific task, under the reward supervision. This two-stage training enables the representations learned in an environment to be reused for other tasks that use the same environment, such as the cheetah environment, containing moving forward and backward tasks. Furthermore, the policy learning stage also requires fewer samples for training, and is thus more sample efficient.

In the first stage, a natural choice is pretraining the visual encoder on a large dataset such as ImageNet. However, previous works show that naively pretraining in such a dataset does not lead to a significant impact [8,9,10]. This ineffectiveness might stem from the domain discrepancy, and more importantly, the pretraining data is lacking in reflecting natural dynamics relation, i.e., the Markov decision process (MDP) property, which is inherent in RL tasks. As we show later in our experiments, pretraining the visual encoder with even in-domain data without considering dynamics is still underperformed. Early works [11,12,13] commonly learn the compact representation by reconstructing the pixels in the current or subsequent frame, which is very challenging with high-dimensional observations. Recently, ref. [14] introduced the contrastive prediction to bypass the use of reconstruction-based prediction. However, this method is lacking in capturing the environment dynamics, thus resulting in low performance in complex tasks such as Cheetah or Walker. To mitigate this shortcoming, ref. [15] introduced augmented temporal contrast (ATC), which additionally trains a marginal forward dynamic model as an auxiliary task for learning representation. Despite effectiveness, the proposed marginal forward model does not condition actions while learning, thus making ATC focus on temporal coherence between observations rather than modeling environment dynamics. To address these issues, this paper introduces a method referred to as the visual pretraining via contrastive predictive model (VPCPM). Specifically, VPCPM utilizes the reward-free data from the environment to learn the visual encoder by jointly optimizing the forward and inverse dynamic models through the contrastive objective. The visual encoder firstly maps the observation into the latent state. Given a latent state and a latent action, the forward model predicts the next latent state. Meanwhile, the inverse model infers the executed action given the consecutive latent states. Learning the forward and inverse models serves as the constraint to enable the visual encoder to consistently follow the underlying dynamics of the environment.

The proposed method is evaluated on a diverse set of image-based continuous control tasks from DeepMind Control Suite [16]. The experiment results show that VPCPM consistently improves the performance and sample efficiency of the state-of-the-art vision-based RL algorithms. Specifically, at 100k environment step, VPCPM improves the mean returns 44% over reinforcement learning with augmented data (RAD) [17], 10% over data-regularized Q (DrQ) [18]. Compared to prior unsupervised pretraining methods, VPCPM matches or outperforms all the baselines as tested across all environments. Moreover, the investigation of unseen tasks shows that the VPCPM-initialized encoder successfully generalizes to unseen tasks that operate in the same environment.

The rest of the paper is organized as follows: Section 2 discusses related works, while Section 3 describes setup in vision-based RL and the related base algorithm. Section 4 details our proposed method. Section 5 presents extensive experiments of the proposed methods, while the conclusion of our paper is presented in Section 6.

2. Related Work

Vision-based Deep Learning. Since the breakthrough results of convolutional neural network (CNN) [19] in the 2012 ILSVRC image classification task, the deep learning approaches have dominated in almost every recognition and detection task. Many works have already implemented the vision-based deep learning methods to detect the number and location of occupants [20,21,22], perform instance segmentation from RGB images [23,24,25] or point cloud [26,27,28], and reconstruct 3D scenes [29,30,31]. In robotics control, by leveraging the expressive CNN together with the advanced reinforcement learning methods, there are many works that successfully perform the complex manipulation tasks based on visual perception [6,7,32,33]. Despite these successes, there are still challenges lying ahead in achieving versatile, robust, and generalized representations.

Representation Learning for RL. Several methods have been proposed in the literature to learn better representations for vision-based RL. Ref. [2] proposed auxiliary tasks including next-observation and reward-based prediction in improving the performance on DMLab [34]. Representation learning using the reconstruction-based task is also a common method, where a beta variation autoencoder (β-VAE) [35] or a deterministic AE is applied for learning a low-dimensional representation, and subsequently, a policy was learned using this representation. Recent work [36] attempted to jointly learn VAE/AE together with RL objective, while [12] proposed to train the two objectives one after the other in an alternating fashion. Ref. [13] also used the VAE for learning the representation in the goal-conditioned setting. These methods only focus on learning the representation of an observation rather than modeling the environment dynamics. Moreover, the use of reconstruction-based tasks makes algorithms hard to scale in high-dimensional environments.

Leveraging the predictive model, which is the combination of the forward and inverse dynamic models, for learning representation is also commonly used. Ref. [37] constructed a joint inverse-forward model to learn representation, in which the forward model is used as a regularizer for training the inverse model features. Ref. [38] used only the inverse model to learn the representations, and then the learned representations were used for the forward model to compute intrinsic reward. Ref. [39] learned the representation by combining the forward and inverse models using cycle consistency losses. In this approach, the use of the forward model makes the representations more Markovian [40,41], i.e., the next latent state only depends on the current latent state and action. In the meanwhile, the inverse model encourages the encoder to capture the controllable features. However, this also makes the representation ignore the uncontrollable features that cannot be captured by the inverse model, which might be useful for solving tasks. Our method proposes the use of contrastive loss while learning the forward model to prevent this degeneracy.

Several methods leverage the recent advances in self-supervised learning for acquiring the representation [42,43,44,45,46] in order to improve the performance of the base RL algorithms. Ref. [42] applied a variant of the noise-contrastive estimation loss on future steps on top of A2C [47] to learn better representation and improve sample efficiency, although the results were limited. In the Atari [48] domain, ref. [49] introduced a new contrastive loss to learn better representation in a fully unsupervised setting, while in [45], the authors adapted the BYOL objective [50] for learning the visual encoder in parallel with the RL objective, which led to state-of-the-art results on the Atari benchmark. In the vision-based continuous control domain, ref. [14] introduced a general framework that is a combination of contrastive loss with image augmentation for learning the representation in a more sample efficient way. Ref. [46] proposed a contrastive version of the conditional entropy bottleneck objective [51] to learn a compressed representation of the predictive information of the environment dynamics. Furthermore, to learn the robustness representation, ref. [52,53] attempted to embed the bisimulation metric [54] into the latent space to encourage the representations invariant to the distractors and generalized to unseen environments. Recently, refs. [17,18] showed that a modest degree of image augmentation can significantly improve sample efficiency for learning directly from visual observation. Most of the existing methods learn the visual encoder by optimizing the unsupervised/self-supervised auxiliary and the RL objective simultaneously, which can be considered as learning from scratch, in an end-to-end manner. By contrast, our framework considers training in two stages, and we further propose a new self-supervised auxiliary objective in the pretraining stage to provide a meaningful representation for directly using or further fine-tuning in the testing environment.

Unsupervised Pretraining Representation for RL. Pretraining the representations without supervision using an unsupervised or self-supervised framework is a common practice in other fields, such as natural language processing (NLP) [55,56,57] or computer vision [42,43,44,50,58,59]. These studies consider effective ways to learn the visual encoder from massive unlabeled data and reduce sample complexity when learning a new task. In deep RL, early studies attempted to pretrain the visual encoder by using pixel-reconstruction task [11,12] or object detection task [60]; the pretrained representations are then fine-tuned on a specific task. Some recent studies [8,9,10] showed that naively pretraining on ImageNet is not helpful for performing the downstream RL tasks. It hints that the visual encoder of the RL agent would be pretrained on the data closely related to its environments. Along this line of research, ref. [14] attempted to use their proposed objective, originally designated for online training, to learn the visual encoder detached from policy learning. However, due to the lack of encoding of the environment dynamics, the performance was limited in complex tasks such as Cheetah. The study closest to ours is [15], which also attempted to learn the encoder detached from the RL objective. However, this method implicitly estimates the next observation by marginalizing over actions. In contrast, our method relies on the forward and inverse dynamics models with contrastive loss, where the proposed forward model is conditioned on action such that in the latent space, the representation of states and action consistently follow the underlying Markov decision process of the environment.

In addition to pretraining the visual encoder, there is also another line of research that tries to pretrain the policy with self-supervised intrinsic rewards [38,61,62,63,64,65]. In this setting, the agent is first allowed to freely interact with the environment for a long period without access to extrinsic rewards, then it is exposed to task-specific rewards to learn on downstream tasks. The intrinsic reward is commonly formulated to encourage the RL agent to gain new knowledge about the environment [38,61,62], maximize diversity of collected data [9,66], or learn diverse skills [67,68,69,70]. For the vision-based RL tasks that are applicable in this setting, the visual encoder is concurrently trained with the policy during pretraining. Recently, ref. [9] proposed to use a particle-based estimator [71] to estimate entropy for observations, with its representations learned by using contrastive loss from SimCLR [44]. Alternatively, ref. [66] proposed a self-supervised pretraining scheme that allows detaching the representation learning from exploration (i.e., learning from intrinsic rewards) to enable the generalization of representations for unseen tasks. In this method, the representations are learned by a variant of clustering-based contrastive loss SwAV [72]. These works are promising to acquire the general policy as well as the generalized representations. However, they are still required to freely interact with the environment during pretraining, which is potentially unsafe in the real world. In contrast, our framework allows us to learn the encoder entirely from an offline dataset, thus resulting as safer and enabling the reuse of past data.

3. Background

In this section, the framework for vision-based reinforcement learning is presented together with a representative off-policy model-free algorithm, soft actor-critic.

3.1. Reinforcement Learning from Images

The problem of solving control task from high-dimensional observations is formulated as a partially observable Markov decision process (POMDP) [73,74], which can be defined as a tuple $(\mathcal{O},\mathcal{A},p,r,\gamma )$. Here, $\mathcal{O}$ is the high-dimensional observation space, $\mathcal{A}$ is the action space, the transition dynamics $p=Pr\left(o_{t+1}\right|o_{\le t},a_t)$ represent the probability distribution over the next observation $o_{t+1}$ given the history of previous observations $o_{\le t}$ and current action $a_t$, the reward function $r:\mathcal{O}\times \mathcal{A}\to \mathbb{R}$ that maps the current observation and action to a reward $r_t=r(o_t,a_t)$, and $\gamma \in [0,1)$ is a discount factor. Following common practice [1], the POMDP is reformulated as an MDP [73] by stacking consecutive observations into a state $s_t=\{o_t,o_{t-1},o_{t-2},\dots \}$. For simplicity of notation, the transition dynamics and reward function are redefined as $p=Pr\left(s_{t+1}\right|s_t,a_t)$ and $r_t=r(s_t,a_t)$, respectively. The goal of RL is to find a policy $\pi \left(a_t\right|s_t)$ that maximizes the expected return defined as the total of accumulated reward ${\mathbb{E}}_{\pi }\left[{\sum }_{t=0}^T{\gamma }^t{r}_t|a_t\sim \pi (.|s_t),s_{t+1}\sim p(.|s_t,a_t),s_0\sim p_0(.)\right]$, where T is the length of episode and $p_0$ is the probability distribution of initial state.

3.2. Soft Actor-Critic

Soft actor-critic (SAC) [75] is an off-policy actor-critic method based on the maximum entropy RL framework [76], which encourages the exploration and robustness to noise by maximizing a weighted objective of the reward and the policy entropy. To update the parameters, SAC performs the soft policy evaluation and improvement steps. The soft policy evaluation step fits a parametric Q-function $Q(s_t,a_t)$ using transitions from the replay buffer $\mathcal{D}$ by minimizing the soft Bellman residual:

$$J\left(Q\right)={\mathbb{E}}_{(s_t,a_t,r_t,s_{t+1})\sim \mathcal{D}}\left[\left(Q(s_t,a_t)-r_t-\gamma \overline{V}{\left(s_{t+1}\right)}^2\right)\right].$$

(1)

The target value function $\overline{V}$ is approximated via a Monte Carlo estimate of the following expectation:

$$\overline{V}\left(s_t\right)={\mathbb{E}}_{a_t\sim \pi }\left[\overline{Q}(s_t,a_t)-\alpha log\pi \left(a_t\right|s_t)\right],$$

(2)

where $\overline{Q}$ is the target Q-function, with its parameters obtained from an exponentially moving average of the Q-function parameters for stabilizing training. The soft policy improvement step then updates the stochastic policy $\pi$ by minimizing the following objective:

$$J\left(\pi \right)={\mathbb{E}}_{s_t\sim \mathcal{D}}\left[{\mathbb{E}}_{a_t\sim \pi }\left[\alpha log\pi \left(a_t\right|s_t)-Q(s_t,a_t)\right]\right].$$

(3)

In this work, the learnable version of temperature α is used instead of the pre-fixed value, which is optimized by the following objective:

$$J\left(\alpha \right)={\mathbb{E}}_{(s_t,a_t)\sim \mathcal{D}}\left[-\alpha log\pi \left(a_t\right|s_t)-\alpha \overline{\mathcal{H}}\right],$$

(4)

where $\overline{\mathcal{H}}\in \mathbb{R}$ is the target entropy hyperparameter that policy attempts to match, which in practice is usually set equal to -$\left|\mathcal{A}\right|$.

SAC is one of the state-of-the-art RL algorithms for continuous control [75]. It is also widely used as a backbone for solving vision-based control tasks [5,14,17,18,36,46]. In this work, we adopt RAD [17] and DrQ [18], which are built on top of SAC, for policy learning.

4. Method

In this section, the proposed visual pretraining via contrastive predictive model (VPCPM) is described. The proposed method can be used to pretrain the visual encoder, which is then utilized for policy learning by common model-free vision-based RL algorithms.

4.1. Network Architecture

The control policy network takes the input as the state $s_t$ and outputs the action $a_t$. It consists of the visual encoder ${\pi }_e$ parameterized by $\varphi$ and the policy ${\pi }_a$ parameterized by ${\theta }_a$, as depicted in Figure 1. This design enables the encoder to be trained independently with the policy part, i.e., without requiring the RL objective for training. The goal of the proposed method is to learn useful representations from an amount of given data without rewards, such that ${\pi }_a$ can be efficiently trained on top of that to solve RL tasks.

4.2. Visual Pretraining via Contrastive Predictive Model

VPCPM introduces a useful prior for the vision-based RL training procedure by enforcing the representations not only representing the semantic information but also conforming to the basis of dynamics from the environment. During the pretraining stage, for a given environment, it is assumed that there is a pre-collected dataset $\mathcal{D}$ consisting of N transitions without task-specific rewards: $(s_t^{\left(i\right)},a_t^{\left(i\right)},s_{t+1}^{\left(i\right)})$ with the index $i=\{1,\dots ,N\}$. The visual encoder is desired to effectively encode the semantics and consistently follow dynamics only from the primitive elements, i.e., observations and actions.

An overview of the proposed method is shown in Figure 2. The visual encoder ${\pi }_e:\mathcal{O}\to \mathcal{Z}$ learns the mapping from the observation into the latent space. VPCPM alternates between learning the forward dynamic model (forward step) and inverse dynamics model (inverse step) while optimizing the underlying encoder ${\pi }_e$ (Algorithm 1). At the forward step, the forward model F parameterized by $\psi$ takes the inputs as the current latent state and the latent action in predicting the next latent state. To optimize F together with ${\pi }_e$, the InfoNCE loss [42] is employed, which contrasts between the predicted next latent state and the ground truth. Formally, let $f:\mathcal{Z}\times \mathcal{Z}\to \mathbb{R}$ be a similarity metric; the objective of forward model is described as:

$$J_F(\psi ,\varphi )=-\frac{1}{K}\sum _{i}log\frac{expf({\widehat{z}}_{t+1}^{\left(i\right)},{\overline{z}}_{t+1}^{\left(i\right)})}{\frac{1}{K}{\sum }_{j}expf({\widehat{z}}_{t+1}^{\left(i\right)},{\overline{z}}_{t+1}^{\left(j\right)})},$$

(5)

where ${\widehat{z}}_{t+1}^{\left(i\right)}=F({\pi }_e(s_t^{\left(i\right)}),{\overline{a}}_t^{\left(i\right)};\psi )$ is the predicted latent state, and ${\overline{z}}_{t+1}^{\left(i\right)}={\pi }_e(s_{t+1}^{\left(i\right)})$ is the target latent state without encoder parameter update. The expectation is computed over K samples of $(s_t,a_t,s_{t+1})$. Operating in the latent space bypasses the prediction of the forward model in pixel space, which would be extremely challenging given the large uncertainty in pixel prediction. The use of InfoNCE helps to learn the discriminative representation, where the dissimilar states are repelled and the similar states are pulled close. Additionally, this objective also prevents trivial collapsed solutions in which the constant features are obtained for every state.

Figure 2. VPCPM for the encoder ${\pi }_e$: At the forward step, the forward model F takes the inputs as the latent state $z_t$ and the latent action ${\overline{a}}_t$ for predicting next latent state $z_{t+1}$. In this step, both F and ${\pi }_e$ are optimized together. At the inverse step, the inverse model I takes the inputs as two latent states $z_t$ and $z_{t+1}$ for predicting the action $a_t$. In this step, both I and ${\pi }_e$ are optimized together.

Figure 2. VPCPM for the encoder ${\pi }_e$: At the forward step, the forward model F takes the inputs as the latent state $z_t$ and the latent action ${\overline{a}}_t$ for predicting next latent state $z_{t+1}$. In this step, both F and ${\pi }_e$ are optimized together. At the inverse step, the inverse model I takes the inputs as two latent states $z_t$ and $z_{t+1}$ for predicting the action $a_t$. In this step, both I and ${\pi }_e$ are optimized together.

Algorithm 1 Visual pretraining via contrastive predictive model (VPCPM)

1: Input: Dataset $\mathcal{D}={\{(s_t^{\left(i\right)},a_t^{\left(i\right)},s_{t+1}^{\left(i\right)})\}}_{i=1}^N$, the encoder’s parameters ϕ, batch size K 2: Output: The encoder’s parameters ϕ 3: Initialize: The parameters ϕ, ψ, ρ 4: fork = 1 to ∞ do 5:    Sample a batch $B={\{(s_t^{\left(i\right)},a_t^{\left(i\right)},s_{t+1}^{\left(i\right)})\}}_{i=1}^K$ 6:    Compute the latent states:                   $z_t^{\left(i\right)}={\pi }_e(AUG(s_t^{\left(i\right)}))$                    $z_{t+1}^{\left(i\right)}={\pi }_e(AUG(s_{t+1}^{\left(i\right)}))$ 7:    Train forward model: JF(ψ, ϕ)    ▹ Equation (5) 8:    Train inverse model: JI(ρ, ϕ)             ▹ Equation (6) 9: end for

In the inverse step, the inverse model I parameterized by ρ takes the states before and after transition and predicts the action in between. In this work, the inverse model operates in the latent space extracted by the visual encoder ${\pi }_e$. The encoder is jointly optimized with the inverse model by minimizing the following objective:

$$J_I(\rho ,\varphi )=\frac{1}{K}\sum _i\ell (I({\pi }_e(s_t^{\left(i\right)}),{\pi }_e(s_{t+1}^{\left(i\right)});\rho ),a_t^{\left(i\right)}).$$

(6)

For continuous action, ℓ can be defined as the mean squared error or mean absolute error between ground truth and the predicted action. When predicting the action, the inverse model pays attention to the controllable features and the temporal difference between consecutive states in latent space, which also encourages the encoder to capture the discriminative features. Pretraining the representations satisfying both forward and inverse dynamics models strengthens the relations of state and actions in latent space, establishing an initialization point of the encoder for further fine-tuning inside a region of parameter space in which the parameters are henceforth restricted.

5. Experimental Results

5.1. Experiment Setup

The proposed method is evaluated on a diverse set of the image-based control tasks from DeepMind Control suite [16], which was recently considered as a standard for benchmarking the sample efficiency of RL from images [14,17,77,78]. DMControl consists of different robot models (environments), where each model can be associated with a particular MDP representing a specific task. The selected benchmark includes six environments from the PlaNet benchmark [77], as shown in Figure 3, where the action repeat is treated as a hyperparameter (

Table 1. Hyperparameters for training the vision-based RL algorithms: RAD [17] and DrQ [18].

Parameter	RAD [17]	DrQ [18]
Action repeat	2 Finger, spin; Walker, walk
	8 Cartpole, swingup
	4 otherwise
Batch size	512	512
Augmentation	Crop [17]	Crop [18]
Replay buffer size	105	105
Initial steps	1000	1000
Stacked frames	3	3
Discount γ	0.99	0.99
Optimizer	Adam	Adam
Learning rate (πe, πa, Q)	2 × 10−4 cheetah	1× 10−3 (all)
	1 × 10−3 otherwise
Learning rate (α)	1 × 10−4	1 × 10−4
Critic target update frequency	2	2
Critic soft-update rate τ	0.01	0.01
Actor update frequency	2	2
Actor log stddev bounds	[−10, 2]	[−10, 2]
Initial temperature	0.1	0.1

). The settings of visual observation follow [14,17,18,36], which consider a stack of three consecutive 84 × 84 RGB renderings as a state.

For the vision-based RL algorithm, we use two state-of-the-art methods including RAD [17] and DrQ [18], which are based on soft-actor critic (SAC) [75]. The network architecture used is identical to [17,18,36]. Unless stated otherwise, the configurations of the algorithm are as follows: the actor and critic neural network are trained using Adam optimizer [35] and a mini-batch size of 512. For SAC, the initial temperature is 0.1, the soft target update rate τ is 0.01, and the target network and the actor updates are performed every two critic updates similar [17,18,36]. The random cropping [14,17] is used as image augmentation during pretraining. The learning rate for actor, critic, and the parameter α of RAD and DrQ is followed by the setup from each method. In our CPM, the forward model is parameterized by four 50-d hidden layers followed by ReLU activation except the last, and the inverse model is parameterized by three 1024-d hidden layers followed by ReLUs, except the last layer, which uses tanh to normalize the actions. The action is encoded by an MLP consisting of two 50-d hidden layers followed by ReLU except for the last. Input to the forward model MLP is a concatenation of the current latent state and the current encoded action. Input to the inverse model MLP is a concatenation of the current and next latent states. Both the forward and inverse model have separated Adam optimizer [79]. During the pretraining stage, the encoder, inverse, and forward model are trained with a learning rate of 1 × 10⁻⁴ for Walker, 2 × 10⁻⁴ for Cheetah, and 1 × 10⁻³ otherwise, with the batch size of 512.

The performance of the agent is evaluated across five seeds; for every seed, the average returns of 10 episodes are computed every 10k environment steps. The figures plot the mean performance together with ±1 standard deviation shading. The performance is reported over the true environment steps as a common practice [14,17,18,36,78], thus are invariant to the action repeat hyperparameter. Throughout experiments, pretraining data are collected by a random policy. Specifically, for Cheetah and Walker domain, 50k transitions are collected, and 25k for the others. The encoder is pretrained within 50k iterations for Cheetah and Walker, and 25k otherwise, which corresponds to one update step per transition. The full set of parameters is shown in Table 1.

5.2. Effects of Pretrained Representation

In this section, the effectiveness of VPCPM in pretraining the visual encoder for different vision-based algorithms is investigated. Two state-of-the-art algorithms, including RAD [17] and DrQ [18], are evaluated, where the hyperparameters reported in each method are used. The random crop augmentation from each method is applied. Crop image augmentation in RAD is formed by cropping 84 × 84 frames from an input frame of 100 × 100, while in DrQ, the 84 × 84 frames are padded each side by ±4, then cropped back to the 84 × 84 size. Moreover, in the RAD paper, there are some tasks using translation augmentation; we instead use crop augmentation across tasks, thus the results may be varied. The parameters of the pretrained encoder are fine-tuned by the base RL algorithm in a specific task.

Figure 4 and

Table 2. Episodic return at 100k and 500k environment steps with mean and standard deviation for 5 seeds. The scores are bold, indicating the improvement compared to the corresponding base RL algorithm.

	RAD [17]		DrQ [18]
100k Steps	Scratch	+VPCPM	Scratch	+VPCPM
Finger, spin	860 ± 29	974 ± 15	901 ± 104	958 ± 20
Cartpole, swingup	454 ± 155	798 ± 42	759 ± 92	825 ± 40
Reacher, easy	704 ± 212	855 ± 117	601 ± 213	720 ± 65
Cheetah, run	365 ± 31	420 ± 81	344 ± 67	384 ± 45
Walker, walk	493 ± 175	614 ± 128	612 ± 164	651 ± 115
Ball in cup, catch	421 ± 247	918 ± 25	913 ± 53	959 ± 7
500k steps
Finger, spin	982 ± 2	985 ± 15	938 ± 103	988 ± 5
Cartpole, swingup	867 ± 10	870 ± 13	868 ± 10	877 ± 11
Reacher, easy	945 ± 34	965 ± 41	942 ± 71	945 ± 43
Cheetah, run	573 ± 35	670 ± 7	660 ± 96	698 ± 47
Walker, walk	948 ± 11	952 ± 15	921 ± 45	960 ± 10
Ball in cup, catch	962 ± 5	964 ± 13	963 ± 9	968 ± 4

compare these methods with and without pretraining. We provide the result at both 100k and 500k steps as common report for DMControl [14,17,18,46]. At 100k step, VPCPM enhances over RAD ranging from 13 to 118%, with the largest magnitude on Ball in cup-catch. For DrQ, the enhancements are in the range from 6 to 20%, with the largest magnitude on Reacher-easy. From Figure 4, the improvement is shown clearer at the early stage of training on the sparse reward tasks such as Ball in cup-catch and Finger-spin. The reason is that these tasks usually failed to complete the task at the beginning, thus observing less reward signal to learn the visual representation. With our VPCPM-initialized representations, the policy part can be quickly learned, resulting in significantly accelerating learning progress. For the Cartpole-swingup task, the action space is very small (with the dimension of one); thus, with the well-presented representation of states, the task can be quickly solved. The tasks including Reacher-easy, Walker-walk, and Cheetah-run are more challenging because of the exploration problem, thus requiring more samples to complete the tasks even with good representations, resulting in a lower magnitude of the improvement. Overall, our method improves 44% over RAD and 10% over DrQ at 100k steps. At 500k steps, the enhancement is smaller, with 3.4% over RAD and 2.9% over DrQ. This is because the base algorithms are almost converged around 500k; thus, the effect of pretraining is moderate.

5.3. Comparison with Prior Methods

In this section, the comparison of pretraining by using different unsupervised learning methods is conducted. VPCPM is compared against two representation learning approaches: non-model-based and model-based. The non-model-based representation learning approach includes (i) reconstruction loss as in VAE [35], and (ii) the contrastive loss from single observation as in CURL [14]. The model-based approach includes (iii) a simple predictive model (PM), where the forward and inverse model are learned by mean square error loss, (iv) augmented temporal contrast (ATC) [15], where the marginal forward model, i.e., without conditioning on actions, is parameterized by a residual network and learned using contrastive loss, and (v) predictive coding-consistency-curvature (PC3) [80], where the forward model is learned by the weighted sum of three losses: contrastive loss, mean square error loss, and low curvature loss. In PC3, the current latent state-action pair $(z_t,a_t)$ is used as the source of negative samples and used for the contrastive prediction of the next latent state $z_{t+1}$. In contrast, we use the predicted next latent state ${\widehat{z}}_{t+1}$ as the source of negative samples. For the implementation of ATC, we use our implementation with the modification as follows, disable the inverse model, and remove the action input of the forward model. For PC3, we use the author’s provided code (https://github.com/VinAIResearch/PC3-pytorch.git, accessed on 17 March 2022). For a fair comparison, the same amount of samples is used during pretraining. The procedure for evaluation is similar to the previous section, but only the RAD algorithm is considered.

The performance in Figure 5 shows that RAD initialized by VPCPM outperforms across all environments. These improvements suggest the importance of imposing the dynamics to the visual encoder during pretraining, which is lacking in the methods that only focus on semantic information such as reconstruction and contrastive. In comparison to the simple PM, the proposed method benefits from the contrastive objective. Indeed, the inverse model is limited since it cannot capture the changes in the sensory stream beyond the agent’s control, and the use of contrastive is helpful to prevent this degeneracy. Moreover, learning in a contrastive manner represents states more discriminative in the latent space. Compared with ATC, VPCPM shows the importance of the action-condition forward model together with the inverse model in learning the controllable features. PC3 is originally designed to use for model-based planning algorithms such as iLQR, which requires the system to be locally linear. Thus, the features learned by PC3 might not suitable for vision-based RL algorithms in the highly nonlinear system, as in our considered environments. Indeed, the results show that the representations from VPCPM are more useful for vision-based RL algorithms. Overall, the sample efficiency in deep RL should be attained from representations that are discriminative and follow the dynamics.

5.4. Effects of Components during Pretraining

Ablation tests were performed to determine the effects of individual components in VPCPM. The performances of RAD with the encoder pretrained by using the contrastive forward dynamic model (cFDM), inverse dynamic model (IDM), and both of them (VPCPM) are shown in Figure 6. Overall, the base RL agent benefits from pretraining by any type of dynamic model, but the cFDM shows more impacts. Together with the constraint of IDM, the proposed method significantly accelerates the sample efficiency of the base algorithm. Moreover, training the visual encoder together with FDM purely in latent space does not suffer from the collapsed problem, where the encoder outputs a constant across states.

5.5. Generalization over Unseen Tasks

In this section, the generalization of the pretrained encoder for unseen tasks is examined. Specifically, the encoder pretrained from a source task is used for the unseen target tasks. Subsequently, the RAD agent is trained on top of the pretrained encoder until converged. The considered tasks are shown in

Table 3. The encoder pretrained with data from a source task is reused for new target tasks. Almost the encoders are successfully generalized to target tasks, except for Reacher-hard, needing to fine-tune.

Source Task	Target Task
Cartpole-swingup	Cartpole-swingup_sparse
	Cartpole-balance
	Cartpole-balance_sparse
Walker-walk	Walker-stand
Reacher-easy	Reacher-hard

. The target tasks are different on the reward function but share the same observation space, except for Reacher-hard, where the size of the visual indicator is different (see Figure 7). The performance of the base agent is evaluated in both “fine-tuning” and “frozen” settings, where the pretrained representation is frozen or fine-tuned. The results are averaged across five seeds and compared against the agent learning from scratch.

The results are shown in Figure 8. In almost all tasks, the frozen representation is sufficient in learning optimal policy. When fine-tuning from the pretrained initialization, the performance is slightly improved. The major exception is the Reacher-hard task, where the frozen encoder significantly underperforms. However, the fine-tuning encoder shows more sample efficiency than that learning from scratch. The difference in the observation space causes this downgraded performance, i.e., the different size of the target indicator. The enhancement in the performance of the base RL agent shows that VPCPM is successfully learning the abstract features without reward supervision.

5.6. Pretraining with Classification

To show the importance of imposing the dynamics in representation learning, we investigate the case where the visual encoder is trained to capture the semantic only, without knowing about the dynamics. To investigate that, we consider six-way classification which corresponds to six robotic models from DeepMind Control Suite [16], as indicated in Figure 3. The dataset is generated by an expert policy. In each class, the training and test sets contain 50k and 10k samples, respectively. In total, there are 300k samples for training and 60k samples for testing. The visual encoder and the classifier are trained using Adam optimizer [79] with the learning rate 3 × 10⁻⁴, and β = (0.9, 0.999). We use the data augmentation methods used in [43] and the random crop [17]. The pretrained encoder is then frozen and used for policy learning. The results are shown in Figure 9. In Ball in cup-catch and Cartpole-swingup, pretraining by classification slightly improves performance, while other tasks show no gain, or even degraded performance. The results indicate the importance of making the representation encode dynamic information for learning representation offline for RL tasks.

6. Conclusions

In this paper, a new self-supervised representation learning method is proposed to pretrain the visual encoder for the vision-based RL. By leveraging plenty of reward-free data, the proposed method successfully learns the meaningful initial representations that provide sufficient information and consistently follow underlying dynamics from an environment. Experimental results show that the state-of-the-art vision-based RL algorithms benefit from our method, with the gain of 44% over RAD and 10% over DrQ at 100k steps. Additionally, we benchmark several leading self-supervised methods for pretraining visual encoders. The results show that the performance of the policy learned on top of the VPCPM-trained encoder matches or outperforms all others. Furthermore, the independence of task-specific rewards during pretraining allows our learned representations to be reused for different tasks sharing similar observation and action space.

In this paper, we investigate the effectiveness of pretraining the visual encoder, in which the testing and training environments are similar. However, this condition is brittle in practice. Future works should improve the robustness of the pretrained representation such that it is invariant to the visual distractions from the environment such as variations in background, color, and camera pose.