Intrinsically Motivated Self-supervised Learning in Reinforcement Learning

A Markov Decision Process (MDP) is defined for the information transmission in Figure 2. $x\in\mathcal{X}$ denotes an optimal representation in the latent space, which is necessary and sufficient for the reinforcement learning task $y$. Another random variable $n\in\mathcal{N}$ denotes nuisance for the task, i.e., $n\perp y$ or $I(y;n)=0$; also, we assume $n\perp x$ and $I(x;n)=0$. Observation $o$ is generated by an implicit function $g(\cdot,\cdot)$, i.e., $o=g(x,n)\in\mathcal{O}$. For instance, $o$ denotes the stacked pixel observation captured from the environment; $x$ is the desired optimal representation for downstream tasks, such as state-based proprioceptive features; $n$ denotes task-irrelevant information like textures, shadows and backgrounds, which may be a culprit for the representation learning, thus seriously depreciating sample efficiency and generalization ability to unseen environments of the algorithm.

We now consider a generic setting for contrastive learning as an auxiliary self-supervised task along side reinforcement learning. The encoding function is denoted by $f:\mathcal{O}\rightarrow\mathcal{Z}$. A main principle for contrastive learning is that features from different views of the same observation should be forced to be close in the feature space, i.e., the ideal encoder $f^{\star}$ would satisfy $\text{dist}(f^{\star}(o_{i}),f^{\star}(o_{j}))\propto\text{dist}(x_{i},x_{j})$. The contrastive loss depends on the SSL architecture, which has the form like Equation (5) in CURL or Equation (6) in SODA as shown in Appendix B. Then we use ${\rho}(z_{i},z_{j})$ to denote the underlying pair-wise contrastive loss as a metric defined on a metric space $\mathcal{H}$, measuring the distance between different latent representations $z_{i}\in\mathcal{Z}$. It satisfies the well-known properties: $\rho(x,y)\geq 0$; $\rho(x,y)=\rho(y,x)$; $\rho(x,y)\leq\rho(x,z)+\rho(y,z)$.

For each observation $o=g(x,n)$, we analyze the pair-wise contrastive loss. Augmentations of $o$ are denoted by $o^{\prime}=t(o)$. Since the augmentation function is designed to perturb nuisance and maintain significant information of the original observation, we assume $\exists\ n^{\prime}$, s.t. $o^{\prime}=t(o)=g(x,n^{\prime})$. The corresponding features are $z=f(o)$ and $z^{\prime}=f(o^{\prime})$. We further introduce an optimal representation function $f^{\star}$. Ideally, it would ignore information related to nuisance $n$ and maintain the information in $x$ related to downstream tasks, i.e., $f^{\star}(g(x,n))=f^{\star}(g(x,n^{\prime})),\forall n,n^{\prime}\in\mathcal{N}$. We may use a fixed $n_{0}$ to denote $z^{\star}=f^{\star}(g(x,n))=f^{\star}(g(x,n_{0})),\forall n\in\mathcal{N}$.

The loss can be decomposed as above, where $h/h^{\star}$ denotes the composite function of $g$ and $f/f^{\star}$, and the second inequality is from the triangle inequality of the metric. Each term on the right hand side can be further bounded by

The first term, interpreted as the projected distance of $x$, is indeed an exploration bonus. It is the distance in feature space between different projection functions $h(.,n_{0})$ and $h^{\star}(.,n_{0})$, which contrastive learning is trying to minimize; therefore, the distance would be smaller if the state has been encountered. From the perspective of exploration in reinforcement learning, this term can be treated as a prediction error, and it would bring extra bonus for novel states that the contrastive part is not trained on. The prediction error is often considered to motivate exploration like RND [8] does; however, it requires additional networks for distillation. In the most desired case, this term should vanish, as the function $h$ is optimized to $h^{\star}$.

The second term, related to nuisance elimination, encourages the agent to visit vulnerable states and introduces randomness into the optimization. States that are more sensitive to distortions will give larger distance between $h(x,n)$ and $h^{\star}(x,n_{0})$, which is encouraged to be visited or revisited. To be clarified, since the nuisance $n$ in the considered MDP is independent of $x$, our method can not directly improve the generalization ability of the baseline; however, it can contribute by building more robust representations for vulnerable states. Robustness in representations makes the learned representations more invariant to distortions, thus improving the generalization ability to untrained environments. While the function $h$ converges to $h^{\star}$, the information in $n$ is getting ignored, and the whole term of nuisance elimination would also vanish.

To summarize, the intrinsic reward, i.e., pair-wise contrastive loss of the auxiliary tasks in reinforcement learning, can be interpreted as: exploration for novel states that improves sample efficiency, and robustness in representations from nuisance elimination that improves generalization ability.

Inequality (4) is tightly bounded, since $h$ is optimized towards $h^{\star}$ and $h^{\star}$ will ignore the information in $n$ ideally. Apart from that, with a decaying parameter $\beta_{t}$ in front of the intrinsic reward, our method is able to converge to the optimal policy asymptotically. Specifically, $\rho({z}^{\prime},z)$ can be a pair-wise loss of Equation (5) in CURL or Equation (6) in SODA, which will serve as an intrinsic reward in IM-SSR. The pair-wise loss can also be designed by ourselves to utilize the information in encoders, which will be further explored in Appendix D.

Empirically, we conduct experiments on DMControl suite. All experiments are evaluated over 5 random seeds, except for those mentioned in the related caption. DMControl-GB provides diverse test environments for policy evaluation, which are important for sim-to-real in robotics. The agent is trained in the unmodified training environment, while it is evaluated in some new environments as well as the unmodified environment of DMControl suite. The new test environments can be summarized as two types: changing the color randomly and replacing the background with natural videos. Specifically, it includes color_easy, color_hard, video_easy, video_hard, as shown in Figure 3 and we select training, color_hard and video_hard as the representative environments.

Primarily, to validate the priority in sample efficiency of IM-SSR, we compare IM-CURL, IM-SODA with CURL and SODA in the training environment. Especially, we elaborately design tasks with sparse rewards to further prove the significant exploration brought by IM-SSR. We conduct all experiments with 500K training frames. We denote the total training frames as T, using 0.2T as early-term, 0.5T as mid-term and 1.0T as final-term to represent the training progress of 20%, 50% and 100%. We report the mean and standard deviation in the tables. Additional results and more details can be found in Appendix A and C.

Besides, to verify that IM-SSR is helpful for robustness and generalization, we compare IM-SODA and SODA in challenging generalization environments. First, in Section IV-C, we intend to evaluate the generalization ability of sim-to-real, where agents are trained in simulated unmodified mujoco environment training, and tested in real-world like environment color_hard and video_hard. Second, we design tasks which are trained on one scene and tested on another to further validate the generalization in transferring. The detailed analysis of generalization tasks and results can be found in Appendix E.

In this subsection, we demonstrate that IM-SSR can improve sample efficiency on the training environment as Figure 3(a). The results shown in Table I can be summarized:

• •

  IM-SSR surpasses its underlying baselines in the early-term in most cases and maintains its priority in the mid-term significantly.

• •

  The improvements in mid-term are much larger than in other terms, since IM-SSR converges much faster than underlying baselines. The baseline may still struggle in learning representations or searching for non-trivial rewards, while IM-SSR has already learned semantic representations and obtained positive and effective reward signals to be trained on.

• •

  In the final-term, we expect IM-SSR and baselines to achieve similar results like in walker_walk. However, we surprisingly find that in several tasks like cartpole_swingup_sparse, IM-SSR achieves higher mean and much lower variance in performance than the underlying baselines, which means IM-SODA is much more stable in these environments.

• •

  There are few cases that IM-SSR provides slight improvements, which mainly attribute to the dense reward. When setting the reward signal to be sparse, the improvements of IM-SSR become significant. Related experiments are implemented in Section IV-D.

We demonstrate that IM-SSR can improve robustness and generalization, and the evaluation is conducted on color_hard and video_hard environments as Figure 3(c) and Figure 3(e).

As shown in Table II, in the early-term and mid-term, IM-SODA outperforms SODA nearly on all tasks, which proves that its generalization ability in unseen challenging evaluation environments is much better. Since video_hard is relatively harder, it makes both algorithms perform poorly in some environments. Still, the peak of IM-SODA is higher than SODA, which also shows better generalization ability of IM-SODA.

Although methods like CURL and SODA achieve comparable performance in major tasks in DMControl, they still perform poorly on tasks with sparse rewards like cartpole_swingup_sparse and pendulum_swingup. As discussed above, simply combining RL with self-supervised learning cannot help solve environments with sparse rewards, but the agent of IM-SSR is encouraged to explore for novel states that have non-trivial rewards potentially.

To verify this hypothesis, we visualize the training process on these two environments with sparse reward signals in Figure 4. Each row in Figure 4 corresponds to an environment, pendulum_swingup is on the top, while cartpole_swingup_sparse is on the bottom. The first column shows the comparison between CURL and IM-CURL on training, while other columns on the right show the comparison between SODA and IM-SODA on training, color_hard and video_hard respectively. IM-SSR improves the sample efficiency significantly in almost all cases. Specifically, the result on the left top panel shows that in IM-CURL pendulum_swingup converges to much higher averaged returns at about 400K steps, while CURL still performs extremely bad at 500K steps.

What makes the improvement to be marginal or significant? As we notice in Table I, sometimes IM-SSR and the corresponding baselines achieve similar performance. The decomposition and interpretation of pair-wise contrastive loss in Equation (4) can be an explanation, especially the exploration term. When the reward signal is dense, the baseline already learns well without exploring for novel states, and there is not much room for IM-SSR to improve. However, when the reward signal is sparse, the significance of exploration emerges and the intrinsic reward does help.

To verify this idea, we take walker_walk as an example and modify the reward to be sparse as either 0 or 1. The new task is called walker_walk_sparse. In walker_walk, IM-SSR and the corresponding baselines have already achieved high performance in early-term. However, as shown in Figure 5 with designed sparse settings, IM-CURL and IM-SODA perform much better than CURL and SODA due to the lack of reward supervision.