Seeking Visual Discomfort: Curiosity-driven Representations for Reinforcement Learning

Recent work on vision-based RL leverage SRL objectives to improve the sample efficiency of policy search methods [9, 41, 22]. There are two main ways to integrate SRL in reinforcement learning. The first one is to simultaneously update both objectives, and the second is to train the two modules in an alternating fashion [9]. The second option could also mean that SRL is only used to train a feature encoder in a pretraining phase preceding the actual RL. In both cases. the quality of the learned encoder and the resulting representations play a central role in the downstream RL tasks. With a finite amount of data, it is not always possible to collect enough samples to learn a representation that is valid across the state subspace relevant to the task at hand. For instance, in environments with sparse rewards, the SRL training rarely encounters observations corresponding to high-reward regions and their surroundings. The resulting representations for such observations might lack the necessary information for the policy to learn any useful behavior.

In general, this lack of coverage is mostly attributed to the exploitative nature of RL algorithms, which leads to the replay buffer containing a lot of redundant and similar observations. Hence, to improve the quality of the feature extraction and learned representations, it is important to encourage collecting data in states outside of the comfort zone of the SRL model. Formally, that would correspond to maximizing the expected SRL error in the replay buffer $D$:

where $p_{o}$ is the distribution of the observations in the replay buffer and $o$ is an observation.

Observations in the replay buffer are part of observation-action trajectories of length $T$ with $p(o_{0},a_{0},...o_{T},a_{T})=p(o_{0})\prod_{t=0}^{T}\pi(a_{t}|o_{t})p(o_{t+1}|o_{t},a_{t})$ (we omit the generative process $p(o_{t}|s_{t})$ for simplicity). Hence, only the initial state distribution and policy are relevant for (4), since the system dynamics are dependent on the environment and cannot be altered. In RL settings, the initial state distribution is dependent on the environment resetting mechanism. Although it’s interesting to study the effect of this mechanism on performance, we leave this for future work.

Instead, in this paper, we learn curious policies that maximize (4). This corresponds to training a policy $\pi_{cure}$ to maximize the objective in (1) with $\gamma=0$ and the weighted SRL error as an intrinsic reward $r_{cure}=\beta\mathcal{L}_{SRL}$. $\beta$ is a hyperparameter that specifies the degree of curiosity. In our experiments we fix $\beta=1$. Furthermore, we allow $\gamma$ to have values different than zero as it didn’t show any negative influence on the training.

There are two main ways to integrate CuRe in a vision-based RL algorithm. Namely, the intrinsic reward could either be added to the task reward $r_{task}$ to train the main policy $\pi_{task}$, or used separately for training a separate curious policy $\pi_{cure}$. Previous methods mostly use the earlier approach to integrate intrinsic rewards (based on a dynamics model) [31].

In this work, we choose the option with two separate policies. By doing so, we ensure that the task policy is purely optimizing the task reward, and additionally obtain a representation-curious agent capable of exploration for similar tasks in the same environment. More importantly, this choice allows our method to be used with both simultaneous and alternating approaches to SRL integration in RL. In addition, our early experiments indicate that a separate curious policy leads to substantially higher reward areas, while the single policy approach could deteriorate the results in comparison to the baselines. Furthermore, adding the rewards together usually introduces extra hyperparameters to weigh the different terms (e.g. in [29], 3 extra hyperparameters are needed). It is important to note that having a separate policy is only possible when using off-policy RL algorithms such as soft-actor-critic (SAC) [12], which is why we use this method in this work.

Our overall approach is illustrated in Figure 1. CuRe is agnostic to the choice of the SRL algorithm. Besides the encoder and the two policies, our architecture includes an SRL model. This model could refer to different modules depending on the SRL approach used. For instance, when using an AE-based method, it would correspond to a decoder. It could also refer to a dynamics model, an identity transformation, as well as any computational block that is used by representation learning methods to constrain the latent space. Furthermore, the updates of both policies affect the encoder parameters $\phi$. The SRL model parameters $\theta$ are only affected by the SRL update. At every step, we either sample actions from the main policy or the curious one. The choice of which policy to use at every step is based on a hyperparameter $p_{c}$ which specifies the percentage of times exploration actions should be sampled. Intuitively, the curious policy is trained to reach states which have high SRL error. By occasionally sampling actions from this policy, the replay buffer ends up containing more problematic and diverse samples which helps to learn a better representation and to avoid overfitting. This interaction between the curious policy and the SRL model/loss results in an interplay similar to the one observed in generative adversarial networks [11], as both modules are mutually beneficial to each other, and are trained in an adversarial setting. This interplay is illustrated in figure 1. The overall approach is summarized in algorithm 1.

To answer these questions, we experimentally evaluate our method on six continuous control tasks from the Deepmind Control Suite [37]. The chosen tasks aim to cover a wide range of common RL challenges, such as contact dynamics and sparse rewards. The tasks we use are reacher_easy, cartpole_swingup, ball_in_cup, finger_spin, finger_turn and reacher_hard. As deep learning models could be energy inefficient [35], we use only subsets of these tasks for minor experiments that are only aimed at validating simple aspects of our method.

The main goal of our experiments is to validate the effectiveness of CuRe on improving the performance of already existing SRL-based approaches to vision-based RL. To do so we use two such algorithms as baselines and compare their performance with and without CuRe. To validate, that the method is agnostic to the choice of SRL algorithms, we experiment with two different methods. Namely we use a combination of SAC with RAEs as in sac_ae [41] and a combination of SAC with contrastive learning based on curl [22]. We chose those two SRL methods since their integration in RL is fairly recent while also being well-established in robotics applications. We refrain from comparing our approach to classical exploration methods, since the two have different goals: classical exploration in RL is concerned with improving the sample diversity for RL while our method is aimed at encouraging the visitation of SRL-problematic states (discomfort zones). Hence comparing methods from these two categories could be misleading. Both baselines and our method are implemented using PyTorch [30]. For simplicity, we use the same hyperparameters for all experiments except for the action repeat value which changes per task, according to [13]. The actor and critic networks for the RL agent and the curious agent are trained using the Adam optimizer [19], using default parameters. We store trajectories of experiences in a standard replay buffer. For implementing SAC, we follow the training procedure detailed in [41]. For the sake of reproducibility, we provide more information about the training procedure, and an overview of the hyperparameters in Appendix A. Our experiments required a training period of over four months on 6 GPUs (NVIDIA RTX2080 and RTX3090).

Visiting High SRL Error Regions. Figure 2 shows the SRL error encountered when sampling actions from three different policies. The first policy generates random actions within the action space of the environment. The second one is trained with sac_ae, and the last one is a CuRe-based curious policy that maximizes the SRL error without a task reward. While random and sac_ae have similar mean errors per step, our method leads to the visitation of states which have on average an SRL error that is around two orders of magnitude higher. This confirms that CuRe fulfills its goal of increasing the probability of visiting high SRL error states. The exact values are shown in Table III in the appendix.

CuRe-based Exploration During RL. To answer ($Q_{2}$), we study the effect of integrating CuRe into two different baselines, namely sac_ae and curl. The integration is based on algorithm 1. Figure 3 shows the task reward for sac_ae with and without CuRe. In all environments, our method exceeds the performance of the baseline. Specifically, for tasks where the baseline doesn’t show any signs of improvement, such as reacher_hard and finger_turn, CuRe leads to exploring high-reward areas, as can be seen when looking at the maximum rewards achieved in those environments. For simpler tasks such as reacher_easy and finger_spin, our method approaches the maximum environment rewards, while sac_ae converges to $80\%$. In addition, CuRe stabilizes the training and reduces the reward variance significantly. This last feature is not given enough attention in RL research. However, in real-world scenarios, when deploying RL agents, there could be cases where only one training run is possible. An algorithm with lower reward variance could guarantee a sufficiently good policy, while it’s hard to say the same when this condition fails. This effect can also be seen for cartpole_swingup and ball_in_cup. We observe that CuRe has a minor effect on the maximum reached reward for these last two environments. This could be attributed to the already good performance of the baseline on these tasks. In fact, in these environments, sac_ae already approaches the performance achieved by SAC trained with the true states [41]. Nonetheless, the additional curious exploration objective accelerates the convergence of all evaluation tasks, thus improving the sample efficiency, which is one key limitation of state-of-the-art model-free algorithms. In general, our experiments show that CuRe becomes more effective when the task complexity increases.

To study the effect of CuRe on curl [22], we run experiments on the four environments where CuRe had the most influence on sac_ae. Figure 4 shows the reward plots for curl with and without CuRe. Similar to our previous results, CuRe has a positive impact on the overall performance, sample efficiency, reward variance and stability of training. This improvement is not as big as the one observed in our sac_ae experiments. However, this difference is understandable, since curl is a more recent algorithm and has previously shown better results on similar deepmind control suite tasks [22]. Despite that, when looking at results on finger_turn (Figures 3 and  4), CuRe applied to sac_ae reaches a higher final episode reward than vanilla curl. Additionally, we observe that sac_ae+cure has a better sample efficiency than curl in the finger_spin environment.

Effect of Pretraining. In addition to our main results, to assess the quality of the learned representation with CuRe, and to answer ($Q_{3}$), we study the effect of two different pretraining procedures on sac_ae. Namely, we look at pretraining the RAE using samples collected either using a random policy (random-pretraining) or using a policy trained with CuRe only, without any task reward (CuRe-pretraining). For both options, we perform the pretraining for $500$ thousand steps. We also compare the performance of those two variants to the case where no pretraining is performed at all (vanilla). The results are shown in Table I. For all six environments, the best results are obtained when using one of the two pretraining mechanisms. In most cases, CuRe-based pretraining leads to better performance than random-pretraining. This become especially apparent for tasks where the vanilla method struggles, such as reacher_hard and finger_turn. However, for the ball_in_cup environment, CuRe-pretraining seems to deteriorate the performance when compared to both vanilla and random-pretraining. This could be attributed to the simplicity of the task, which reduces the need for SRL and SRL-tailored exploration. In general, although CuRe is beneficial for both SRL pretraining (Table I) and RL (figure 3), we observe that it is more effective during task learning than in the pretraining phase.