Cross-domain Random Pre-training with Prototypes for Reinforcement Learning

End-to-end RL shows sample inefficiency when faced with high-dimensional observations like images. One solution is to employ auxiliary objectives to do SSL, for example, predicting the designed properties of the environments [27, 28]. Following these, CURL [29] introduced a contrastive auxiliary task into end-to-end RL to improve sample efficiency. SPR [30] combined data augmentation with an auxiliary SSL objective. CltrFormer [18] employed transformer [31] in visual-control tasks, using contrastive learning to get transferable representation between different domains on DMControl [15]. Recently, a novel SSL approach: prototypical representation learning [24] based on the Sinkhorn-Knopp algorithm [32] was introduced into RL by Proto-RL [14]. Their success motivated DreamerPro [25], ProtoCAD [26] and our CRPTpro. Different from previous prototypical algorithm, we introduce an extra intrinsic loss between prototypes along with original comparative loss, proposing efficient prototypical learning, a novel and effective self-supervised algorithm. In addition, CRPTpro performs efficient prototypical learning on a static dataset as original SwAV [24] does.

Inspired by the success that SSL can pre-train a strong feature extractor without labels in the field of CV [33, 34] and NLP [35, 36], ATC [37] tried to decouple the representation learning from downstream policy learning and first achieved considerable results. MVP [38] utilized the offline dataset from the internet, pre-training a visual encoder for different motor tasks. SGI [39] employed finetuning on the pre-trained encoder, achieving efficient policy learning on Atari 100k benchmark [40]. Inspired by unsupervised RL  [20, 41, 21], APT [13] pre-trained an extra exploration agent via an unsupervised intrinsic reward for data collection, along with visual representation learning by SimCLR [34]. Subsequently, Proto-RL [14] used prototypical representation learning [24] to simultaneously achieve encoder pre-training and strengthen APT unsupervised reward, setting state-of-the-art downstream policy performance on DMControl. Different from these active pre-training methods [13, 14, 23], CRPTpro gives up unsupervised RL but decouples data collection from encoder pre-training by an off-the-shelf exploration policy, achieving state-of-the-art cross-domain pre-training performance with dramatically improved pre-training efficiency.

The dataset for pre-training should be diverse to cover the observation space as much as possible. As a branch of unsupervised RL, active pre-training methods [13, 14] design an intrinsic reward related to exploration, based on which they train extra agents to explore the observation space along with learning visual encoders. These methods are able to explore and collect far-reaching states in hard-exploration domains and achieve state-of-the-art performance in single-domain visual-control pre-training. However, due to the requirements of extra agent training, these methods suffer a severe pre-training burden that even exceeds that of downstream RL many times. Especially when facing multiple domains, this extra burden is dramatically increased. Meanwhile, these approaches exhibit huge performance drops when transplanted into cross-domain pre-training due to a chicken-and-egg problem: pre-training an effective encoder requires effective exploration agents for qualified data collection, while effective exploration agents rely on an effective encoder. This problem is severely amplified in cross-domain pre-training because multiple exploration policies are required for multiple domains. It’s hard for current active methods to train multiple qualified exploration strategies from scratch on a substandard visual encoder, and it’s also hard to achieve ideal pre-training on unqualified data sampled by unqualified exploration policies.

Due to the above problems, CRPTpro gives up training extra unsupervised exploration agents but decouples data sampling from encoder pre-training, proposing decoupled random collection for cross-domain pre-training. It employs an off-the-shelf random policy across multiple domains for pre-training dataset collection, which dramatically reduces the pre-training burden. Specifically, it chooses a simple uniform distribution to sample actions for environment interaction and data (image observations) collection. Data from different domains is saved into different data buffers, forming the cross-domain pre-training dataset, which is collected once and used permanently. This cross-domain pre-training dataset is then used for SSL.

In terms of the ability to explore far-reaching states, the uniform distribution policy is not as good as active exploration. However, for visual-control pre-training across multiple domains, our decoupled random collection provides better and qualifier data than active exploration due to the following reasons: (i) The above chicken-and-egg conflict makes it hard to train qualified exploration policies for active exploration, while our off-the-shelf uniform distribution policy is not affected. (ii) The difficulty of seeking far-reaching states in common motor control is much lower than that in a ’maze’. The uniform distribution policy may have trouble finding the far-reaching maze end but can drive a cheetah robot to exhibit different behaviors. (iii) A single exploration agent may explore very deeply, but it is hard to explore widely like a random policy. For example, a uniform distribution policy can manipulate a robotic arm to extend in all directions due to its randomness, but a single exploration agent may deeply explore only one direction. Here, a wider exploration is more important because more diverse data can help the encoder better understand motion changes. In conclusion, decoupled random collection can sample a qualified dataset (see visualization in Section IV.G). In addition, due to the changing exploration agent, active pre-training methods train their encoders on a changing dataset, which is unstable. By contrast, our approach decouples data sampling from encoder learning, achieving data collection before pre-training, which enables a steady learning process. Therefore, the proposed decoupled random is the better cross-domain data collection method due to (i) a qualifier pre-training dataset and (ii) a more stable learning process.

In summary, employing decoupled random collection brings the following advantages compared with recent advanced active pre-training methods: (i) CRPTpro has an unparalleled efficiency advantage in the pre-training stage. It employs an off-the-shelf data collection policy and doesn’t need to spend time on extra RL for exploration agents. (ii) CRPTpro achieves better pre-training performance, especially when pre-training efficiency is required. Recent advanced active pre-training suffers from the severe chicken-and-egg problem mentioned above, while CRPTpro enables robust and steady learning on a qualified cross-domain pre-training dataset.

CRPTpro learns a cross-domain visual encoder and several basic vectors called prototypes over the cross-domain pre-training dataset produced by decoupled random collection. As illustrated in Fig. 2, CRPTpro selects data buffers from different domains sequentially and cyclically to achieve the pre-training of one generic encoder. Concretely, in each step of the training, one data buffer is chosen to update the encoder only once, while another one will be chosen in the next step. This prevents the encoder from favoring the latest domain.

Over the chosen buffer, CRPTpro employs efficient prototypical learning, a novel prototypical self-supervised algorithm. Following previous prototypical algorithms [24, 14, 25, 26], it sets several trainable prototypes as the cluster centers in the latent space, comparing observations projected onto prototypes with their cluster assignment targets. In addition to the comparison, a novel intrinsic loss is proposed to facilitate the diffusion of prototypes in the latent space, further improving the pre-training performance, seen in Fig. 3.

The comparative loss is calculated as the following. First of all, M frames $\{x_{t_{i}}\}_{i=1}^{M}$ and their next frames $\{x_{t_{i}+1}\}_{i=1}^{M}$ are sampled randomly from the data buffer. The subscript represents the timing number. The current frames $\{x_{t_{i}}\}_{i=1}^{M}$ are used to predict the cluster assignment targets computed over the next frames $\{x_{t_{i}+1}\}_{i=1}^{M}$ and M trainable vectors $\{c_{1},...,c_{M}\}$ called prototypes. $x_{t_{i}}$ is one of the current frames. It undergoes augmentation by random image shifts, encoding by the generic convolutional encoder $f_{\theta}$, and projection by the Multi-Layer Perceptron (MLP) $g_{\theta}$ in turn to produce a vector $z_{t_{i}}$ in the latent space where prototypes exist. Here, another MLP $v_{\psi}$ is used to predict $z_{t_{i}}$ into $u_{t_{i}}$ in order to avoid collapse to trivial solutions. $v_{\psi}$ doesn’t change the dimension of input, so both $z_{t_{i}}$ and $u_{t_{i}}$ have the same dimension as prototypes. We then take a softmax over the dot products of $u_{t_{i}}$ and all the prototypes $\{c_{j}\}_{j=1}^{M}$:

To obtain the target compared with $p_{t_{i}}$, all next frames $\{x_{t_{i}+1}\}_{i=1}^{M}$ undergo augmentation by random image shifts, encoding by target encoder $f_{\bar{\theta}}$, and projection by target MLP $g_{\bar{\theta}}$ sequentially to produce latent embeddings $\{z_{t_{i}+1}\}_{i=1}^{M}$ analogously. To avoid collapse to trivial solutions, parameters $\bar{\theta}$ of these target networks are updated using the Exponential Moving Average (EMA) [44] of $\theta$:

Now we can apply the Sinkhorn-Knopp algorithm [32] on $l_{2}$-normalized embeddings $\{\hat{z}_{t_{i}+1}\}_{i=1}^{M}$ and prototypes $\{\hat{c}_{j}\}_{j=1}^{M}$. Concretely, the algorithm begins with the square matrix $C$, whose elements are computed by the dot product over each embedding and prototype:

Then it employs iterative doubly-normalization on the matrix $C$ to obtain target matrix $T$, constraining every column and row to have the same sum with as little change of original $C$ as possible. The row normalization and column normalization are used in doubly-normalization. The row normalization is formulated as the following:

where ${\rm sum}(C,1)$ denotes the row addition and $\rm diag(\cdot)$ denotes the diagonalization of a matrix. Similarly, the column normalization is defined as the following:

where ${\rm sum}(C,0)$ denotes the column addition. The doubly-normalization, $\rm dou(\cdot)$, consists of one row normalization and one column normalization:

Three times of doubly-normalization is applied on $C$ to obtain target matrix $T$. The $i$-th row of $T$ is the cluster assignment target $q_{t_{i}+1}$ of the frame $x_{t_{i}}$. Combined with the probabilities $p_{t_{i}}$ computed by Eq.(1), the comparative loss over all frames is calculated as the following:

During the training process, the prototypes serve as cluster centers for the selected samples. They are randomly initialized at the beginning and then gradually expand their coverage to the visited states in the latent space, which is done through the traction of samples.

Intuitively, facilitating the diffusion and coverage of prototypes can (i) accelerate the training process and (ii) make prototypes wider cluster centers, leading to a wider and more distinguishable latent space, i.e., a stronger encoder. To this end, we design a novel intrinsic loss aiming to accelerate the diffusion of the prototypes, which is done by increasing the difference between prototypes. Specifically, the difference between two $l_{2}$-normalized prototypes (unit vectors) can be well measured by their cosine similarity:

where the hats over $c_{j}$ and $c_{k}$ mean $l_{2}$-normalization. Our intrinsic loss is based on this cosine similarity:

where $\rm det(\cdot)$ denotes the detach operation which prevents the gradient from backpropagation. $w$ is a weight hyper-parameter. It is employed to balance different update speeds of different prototypes. In practice, in the case of the dot product as the loss, a remote prototype is updated slower than a near prototype after $l_{2}$-normalization. The denominator defined as the sum of $\hat{c}_{j}^{T}\hat{c}_{k}$ and $w$ increases the gradient of the remote prototypes, further accelerating the diffusion of prototypes.

With the comparative loss $\mathcal{L}_{comp}$ computed by Eq.(7) and the intrinsic loss $\mathcal{L}_{intr}$ computed by Eq.(9), the overall self-supervised loss $\mathcal{L}_{SSL}$ of efficient prototypical learning in practice is:

where $\alpha$ is a coefficient scaling the intrinsic reward. Significantly, CRPTpro updates the parameters $\theta$, $\psi$ and prototypes $\{c_{j}\}_{j=1}^{M}$ during the pre-training. The cross-domain encoder $f_{\theta}$ and prototypes $\{c_{j}\}_{j=1}^{M}$ are then used to conduct efficient downstream RL in different domains.

After pre-training, the encoder can be frozen and used to perform efficient downstream policy learning on challenging visual-control tasks from different domains either seen or unseen, with the help of frozen prototypes. Specifically, the encoder maps the state space $\mathcal{X}$ into the embedding space $\mathcal{Y}$, converting the old MDP $(\mathcal{X},\mathcal{A},\mathcal{P},\mathcal{R},\gamma,d_{0})$ into $(\mathcal{Y},\mathcal{A},\mathcal{P},\mathcal{R},\gamma,d_{0})$. Correspondingly, the transition $(x_{t_{i}},a_{t_{i}},r_{t_{i}},x_{t_{i}+1})$ is converted into $(y_{t_{i}},a_{t_{i}},r_{t_{i}},y_{t_{i}+1})$. Following [14], we augment the extrinsic reward $r_{t_{i}}$ by adding an exploration reward $\hat{r}_{t_{i}}$ based on prototypes to encourage exploration. $\hat{r}_{t_{i}}$ is a particle-based entropy estimation [45] which is positively correlated with the exploration ability of the current strategy:

In the cross-domain setting, we don’t employ finetuning to maintain cross-domain versatility. However, we note that finetuning the cross-domain encoder on a single target domain is able to further improve its performance but reduce its cross-domain versatility, especially in the unseen domains. After finetuning, the encoder focuses more on a certain domain and reduces its versatility, which means CRPTpro-finetuning can be regarded as a single-domain pre-training method. We provide the full pseudo-code of CRPTpro in Algorithm 1.

We compare CRPTpro with 4 cross-domain pre-training baselines over 4 groups defined in Table I. CRPTpro hyper-parameter settings are shown in Table VII. DrQ is set as an end-to-end expert to show the score level of different downstream tasks and calculate an normalized score. Results are summarized in Table III. CRPTpro significantly improves APT(C), ATC and BeCL on 11/12 tasks. Compared with Proto-RL(C), CRPTpro achieves better downstream policy learning on 11/12 tasks and similar performance on the rest task. Numerically, we calculate the mean score and the mean expert-normalized score over all 12 downstream tasks of 4 groups, where our CRPTpro improves upon the next best cross-domain pre-training baseline by 23.8% and 78.5% respectively. Perhaps the most exciting result is that CRPTpro achieves a 1.956 mean expert-normalized score, demonstrating the huge advantages of cross-domain visual pre-training for image-based RL, which is ignored by previous works.

In summary, our CRPTpro significantly beats all baselines, becoming a novel state-of-the-art cross-domain pre-training method. This is attributed to the following reasons: First, our decoupled random collection avoids the severe chicken-and-egg problem between exploration agent training and visual encoder pre-training (detailed description in section III.A), generating a qualified and diverse cross-domain pre-training dataset and enabling stable pre-training process. Second, our efficient prototypical learning helps the encoder learn more effective image embeddings.

In addition to the seen domains, the encoder pre-trained by CRPTpro can generalize well to unseen domains without finetuning. We conduct three sets of representative experiments to show the generalization comparison of all cross-domain methods. (i) We pre-train encoders on Group-A which contains only robot domains. Then we freeze the pre-trained encoders to directly perform 500k steps downstream RL on the rest tasks defined in five unseen and different types of domains: manipulation (Manipulation), robot (Hopper), and balance control (Pendulum, Cartpole, and Finger). (ii) We pre-train encoders on Group-C which contains three different types of domains: manipulation (Manipulation), robot (Hopper) and balance control (Finger), and then freeze the encoder to perform 500k steps RL on the other five unseen tasks. (iii) For encoders pre-trained on Group-C, we further change the background color (by exchanging the RGB value in reverse order) of three different types of control tasks from both the seen domain (Manipulation) and unseen domains (Cheetah and Pendulum), testing the generalization to unseen background colors. CRPTpro hyper-parameters are provided in Table VII.

The results are shown in Table IV. CRPTpro overall exceeds all cross-domain baselines across all three sets of experiments. In some unseen tasks (e.g., Pendulum-Swingup in (i) and Walker-Run in (ii)), CRPTpro even enables better policy learning than the end-to-end expert DrQ. Note that CRPTpro has not trained its encoder in these unseen domains, while DrQ trains its encoder for 1M steps on each task. This indicates that the encoder of CRPTpro can well capture the movement changes that are generic across domains. Moreover, CRPTpro is the only method that doesn’t exhibit performance degradation on Cheetah-Run when facing unseen color, which means it can effectively filter our useless information irrelevant to movements. We mainly attribute these movement understanding advantages to the proposed decoupled random collection because we note that CRPTpro on Group-C (containing three different types of domains) performs much better than CRPTpro on Group-A (containing only robot domains). For example, CRPTpro on Group-C can effectively solve all unseen tasks, including Cartpole-Swingup sparse, while CRPTpro on Group-A can not. Even on Walker-Run which is included in Group-A but not in Group-C, CRPTpro on Group-C performs better than CRPTpro on Group-A (result in Table III). This phenomenon shows that data diversity is the key to pre-training performance, which is also verified in the ablation study (Section IV.F). The more diverse dataset sampled by decoupled random collection enables CRPTpro to better understand generic movement changes than all baselines.

As mentioned in section III.C, the cross-domain encoder obtained by CRPTpro can be finetuned on a target domain for better downstream policy learning, especially in unseen domains. To highlight this ability, we finetune the encoder pre-trained on Group-A by CRPTpro in five different domains (three unseen domains: Pendulum, Manipulation, and Cartpole & two seen domains: Cheetah and Walker) respectively with default hyper-parameters in Table VII. Then, we freeze the finetuned encoder to conduct 500k steps RL on the corresponding downstream task. CRPTpro-finetuning could be regarded as a single-domain pre-training method and we compare it with three non-cross-domain baselines.

The results are shown in Fig. 4, demonstrating that CRPTpro-finetuning is an effective single-domain pre-training method. It is comparable with Proto-RL(S), the best single-domain pre-training method and also one of state-of-the-art image-based RL methods. In addition, CRPTpro-finetuning uses only 58.5k update times on encoder (50k pre-training update times on Group-A, 2k finetuning update times in Pendulum, 2k finetuning update times in Manipulation, 3k finetuning update times in Cartpole, 1k finetuning update times in Cheetah, and 0.5k finetuning update times in Walker) to surpass most non-cross-domain baselines. As a comparison, the total pre-training update times for Proto-RL(S) and APT(S) in 5 domains is 625k. In each unseen domain, CRPTpro only updates its encoder at most 3k times while the baselines update their encoders at least 125k times. This indicates that the cross-domain prior knowledge learned with our novel self-supervised algorithm can greatly promote representation learning even in unseen domains.

In this section, we demonstrate the pre-training efficiency comparison between different pre-training methods when facing different numbers of domains. Three cross-domain pre-training methods, including the current state-of-the-art cross-domain pre-training method, Proto-RL(C), and two single-domain pre-training methods, including the current state-of-the-art single-domain pre-training method, Proto-RL(S), are employed for comparison with our CRPTpro. ATC is not in comparison because it doesn’t provide a task-agnostic pre-training approach in their paper (see baseline details in Section IV.A.c). The outcomes are based on experiments conducted on the NVIDIA Tesla P100.

The results are shown in Table V. CRPTpro becomes a novel state-of-the-art cross-domain pre-training method with greatly improved pre-training efficiency. Concretely, it spends only 54% wall-clock pre-training time to outperform the next best Proto-RL(C) by 78.5% on the mean expert-normalized score. The main contributor is our decoupled random collection, which employs an off-the-shelf exploration policy to avoid the severe training burden of extra RL. Compared with the state-of-the-art single-domain pre-training method (Proto-RL(S)), all the cross-domain pre-training methods exhibit huge efficiency improvements which multiply with the number of domains. This is because cross-domain pre-training methods can pre-train a generic encoder across multiple domains, while Proto-RL(S) cannot. However, the downstream policy performance of all cross-domain baselines is much worse than that of Proto-RL(S), which makes the efficiency comparison between them meaningless. In contrast, CRPTpro can achieve competitive downstream policy learning compared with Proto-RL(S) after few-shot finetuning, with much less pre-training consumption (only 9.3% update times, see Section IV.D).

In downstream RL, CRPTpro augments the extrinsic reward with a k-NN-based unsupervised reward. According to our experimental tests on the NVIDIA Tesla P100, the consumption of this reward is tiny, about 1‰ of the downstream RL consumption. This is due to (i) the small size of the k-NN buffer and (ii) the low dimension of the k-NN particles, which we detail in Section III.C. In addition, the current state-of-the-art Proto-RL(C) and Proto-RL(S) also utilize this reward. Therefore, it doesn’t affect the efficiency advantage of our approach.

CRPTpro employs the proposed efficient prototypical learning, which introduces a novel $\mathcal{L}_{intr}$ to facilitate the diffusion of prototypes and improve pre-training. In this section, we demonstrate its effectiveness and rationality through three different sets of experiments

First, we prove its numerical effectiveness in three different prototypical pre-training settings. Ablation study is also done here. (i) We ablate the $\mathcal{L}_{intr}$ from CRPTpro with default hyper-parameter settings to observe the difference, shown in Fig. 5a. (ii) We add $\mathcal{L}_{intr}$ into Proto-RL(C) with default hyper-parameter settings to show its effectiveness in Fig. 5a. (iii) We add $\mathcal{L}_{intr}$ into Proto-RL(S) of 200k steps task-agnostic pre-training and 500k steps downstream RL. The settings are all shown in Table VII except $\alpha$ set to $1e-3$ for Cheetah-Run and $5e-4$ for Hopper-Hop. The curves are shown in Fig. 5b.

The numerical results verify the effectiveness of $\mathcal{L}_{intr}$ in efficient prototypical learning across all prototypical pre-training settings. Since the Proto-RL(C) w/ $\mathcal{L}_{intr}$ is actually CRPTpro w/o decoupled random collection, Fig. 5a can also serve as the ablation study, verifying the effectiveness of both the decoupled random collection and efficient prototypical learning in CRPTpro.

Second, we show that the coverage and diffusion of prototypes are facilitated when employing efficient prototypical learning. Two metrics are used to evaluate prototypes’ coverage: All-Neighbor Estimation (ANE) proportional to the cosine similarity between all prototype pairs; K-Neighbor Estimation (KNE) proportional to k-nearest neighbor’s cosine similarity of all prototypes. We test the pre-trained models’ prototypes in 3 different prototypical pre-training settings. Results in Table VI show that $\mathcal{L}_{intr}$ makes prototypes finally cover wider in the latent space. In addition, we show the curve of metric difference between Proto-RL(S) and Proto-RL(S) w/ $\mathcal{L}_{intr}$ during pre-training in Fig. 6a. This indicates that the $\mathcal{L}_{intr}$ continues to function throughout the training process of efficient prototypical learning, demonstrating that $\mathcal{L}_{intr}$ makes prototypes cover faster in the self-supervised pre-training. In summary, our efficient prototypical learning indeed improves the diffusion and coverage of prototypes, leading to better encoder pre-training and the final performance.

Finally, we observe that optimizing $\mathcal{L}_{intr}$ is an instinctual process in vanilla prototypical algorithms. We show the $\mathcal{L}_{intr}$ curve in the single-domain active pre-training setting, as shown in Fig. 6b. We find that even in Proto-RL(S), where $\mathcal{L}_{intr}$ is not employed, this loss is also instinctively decreased, which means reducing $\mathcal{L}_{intr}$ is an instinctual process in vanilla prototypical representation learning. This phenomenon motivates us to optimize $\mathcal{L}_{intr}$ in a targeted manner to facilitate the training and accelerate the convergence.

To further analyze the effectiveness of our decoupled random collection and pre-trained encoder, we employ PCA to visualize two pre-training datasets: the cross-domain pre-training dataset produced by decoupled random collection (CRPTpro) and the single-domain exploration dataset produced by Proto-RL(S) [14]. Note that Proto-RL(S) is a state-of-the-art image-based RL method, providing the state-of-the-art unsupervised exploration. We don’t visualize Proto-RL(C) because its performance is much less than CRPTpro and Proto-RL(S). For cross-domain dataset produced by decoupled random collection (CRPTpro), we sample 125 frames from each domain. For Proto-RL(S), we conduct adequate single-domain pre-training and then sample 125 frames with the pre-trained exploration. All the sampled frames are encoded by corresponding CRPTpro encoder and then reduced dimensionality to 4 by PCA (4-dimensional PCA results retain more than half of the original information here).

The results are shown in Fig. 7. In a single domain, the decoupled random collection in CRPTpro (orange) is not as exploratory as the exploration agents in Proto-RL(S) (blue), but not a lot worse. It also explores some observations that Proto-RL(S) cannot reach. Moreover, the cross-domain pre-training dataset produced by CRPTpro (orange, green and red) overall exhibits wider coverage than the single-domain exploration dataset produced by Proto-RL(S) (blue) because of the cross-domain diversity. In addition, we observe that our encoder can ideally handle unseen data sampled by exploration agents and distinguish them clearly, showing the powerful generalization. These explain why CRPTpro can pre-train a generic encoder conducting efficient downstream policy learning in multiple domains, even reaching the same level as the best active pre-training method (also a state-of-the-art image-based RL method), Proto-RL(S), after few-shot finetuning.