CCLF: A Contrastive-Curiosity-Driven Learning Framework for Sample-Efficient Reinforcement Learning

1 Introduction

Despite the success of reinforcement learning (RL), extensive data collection and environment interactions are still required to train the agents [Laskin et al., 2020b]. In contrast, human beings are capable of learning new skills quickly and generalizing well with limited practice. Therefore, bridging the gap of sample efficiency and learning capabilities between machine and human learning has become a main challenge in the RL community [Rakelly et al., 2019; Schwarzer et al., 2020; Yarats et al., 2021c; Malik et al., 2021; Sun et al., 2022].

∗Co-corresponding authors

This challenge is particularly vital in learning directly from raw pixels. More recently, data augmentation methods are leveraged to incorporate more invariances, promote data diversity, and thereby enhance representation learning [Laskin et al., 2020a; Laskin et al., 2020b; Yarats et al., 2021b]. Ideally, injecting a larger number of augmented samples should lead to a better model with invariances. Nevertheless, a noticeable trade-off is the computational complexity introduced. What's worse, simply increasing the number of augmented inputs may alter the semantics of samples, which has been empirically shown in our experimental results. Moreover, the samples used for data augmentation are uniformly drawn from the replay buffer which is inefficient as they are not equally important to learn. These assumptions deviate from human-like intelligence, where humans can learn efficiently by curiously focusing on novel knowledge and revisiting old knowledge less frequently. Therefore, replaying the most under-explored experiences and selecting the most informative augmented inputs are the keys to improving sample efficiency and learning capability.

To tackle these challenges, we propose a ContrastiveCuriosity-driven Learning Framework (CCLF) by introducing contrastive curiosity into four important components of RL including experience replay, training input selection, learning regularization, and task exploration without much computational overhead. Inspired by the psychological curiosity that can be externally stimulated, encompassing complexity, novelty, and surprise [Berlyne, 1960; Spielberger and Starr, 2012; Liquin and Lombrozo, 2020], we define the contrastive curiosity based on the surprise conceptualized by the agent's internal belief towards the augmented inputs. The internal belief is modeled by reusing the contrastive loss term in CURL [Laskin et al., 2020a], which can quantitatively measure the curiosity level without introducing any additional network architecture. With the proposed contrastive curiosity, agents can sample more under-explored transitions from the replay buffer, and select the most informative augmented inputs to encode invariances. This process can significantly reduce the amount of data used in RL without sacrificing the invariances. Thereafter, CCLF further utilizes the contrastive curiosity to regularize both Q-function and encoder by concentrating more on the surprising inputs, and intrinsically rewards agents for exploring under-learned observations.

Our contribution can be highlighted as follows. 1) We empirically demonstrate that not all samples nor their augmentations are equally important in RL. Thus, agents should learn curiously from the most important ones in a self-supervised manner. 2) A surprise-aroused type of curiosity, namely, contrastive curiosity, is proposed by reusing the representation learning module without increasing the network complexity. 3) The proposed CCLF is capable of improving the sample efficiency and adapting the learning process directly from raw pixels, where the contrastive curiosity is fully exploited in different RL components in a self-navigated and coherent way. 4) CCLF is model-agnostic and can be applied to modelfree off-policy and on-policy RL algorithms. 5) Compared to other approaches, CCLF obtains state-of-the-art performance on the DeepMind Control (DMC) suite [Tunyasuvunakool et al., 2020], Atari Games [Bellemare et al., 2013], and MiniGrid [Chevalier-Boisvert et al., 2018] benchmarks.

2 RelatedWorks

Data Augmentation in Sample-Efficient RL. Data augmentation has been widely applied in computer vision but is only recently introduced in RL to incorporate invariances for representation learning [Laskin et al., 2020b; Yarats et al., 2021b; Laskin et al., 2020a]. To further improve the sample efficiency, one approach is to automatically apply the most effective augmentation method on any given task, through a multi-armed bandit or meta-learning the hyper-parameters to adapt [Raileanu et al., 2020]. However, the underlying RL algorithm can become non-stationary and it costs more time to converge. Another approach is to regularize the learning process with observations from different training environments [Wang et al., 2020] or different steps [Yarats et al., 2021a]. By injecting greater perturbations from other tasks and steps, the encoded features and learned policies can become more robust to task invariances. Different from these works, the proposed CCLF primarily focuses on perturbations generated in a single task and step, and selects the most underlearned transition tuples and their augmented inputs. As not all samples nor their augmented inputs are equally important, our work exploits the sample importance to adapt the learning process by concentrating more on under-explored samples. Most importantly, the amount of augmentations can be greatly reduced and the sample efficiency is improved without introducing complicated architectures.

Curiosity-Driven RL. In curiosity-driven RL, agents are intrinsically motivated to explore the environment and perform complex control tasks by incorporating curiosity [Aubret et al., 2019; Sun et al., 2022]. In particular, curiosity is mainly used as a sophisticated intrinsic reward based on state novelty [Bellemare et al., 2016], state prediction errors [Pathak et al., 2017], and uncertainty about outcomes [Li et al., 2021] or environment dynamics [Seo et al., 2021]. Meanwhile, it can also be employed to prioritize the experience replay towards under-explored states [Schaul et al., 2015; Zhao and Tresp, 2019]. However, additional networks are required to model curiosity, which can be computationally inefficient and unstable for high-dimensional inputs with continuous controls. Moreover, none of these works have yet attempted to improve the sample efficiency and resolve the instability caused by data augmentation. CCFDM [Nguyen et al., 2021] is a concurrent work that incorporates CURL with action embedding and forward dynamics to formulate an intrinsic reward. Different from CCFDM, our framework does not require any additional architecture but only reuses the contrastive term in CURL that predicts more stably. More importantly, the proposed CCLF seamlessly integrates the curiosity mechanism into experience replay, training input selection, learning regularization, and environment exploration to concentrate more on under-learned samples and improve the sample efficiency stably.

3 Background

In this paper, we consider a Markov Decision Process (MDP) setting with the state st ∈ S, the action at ∈ A, the transition probability P mapping from the current state st and action at to the next state s′t , and the (extrinsic) reward re t ∈ R. More details are provided in Appendix A 1 .

Soft Actor-Critic (SAC). SAC [Haarnoja et al., 2018] is an off-policy model-free algorithm that learns a stochastic policy πψ (actor) with state-action value functions Qφ1 ,Qφ2 (critics), and a temperature α by encouraging exploration through a γ-discounted maximum-policy-entropy term. However, agents are often required to learn directly from highdimensional observations ot ∈ O rather than states st in practice. In this paper, we demonstrate our framework mainly using SAC with raw pixel inputs as the base algorithm.

ContrastiveUnsupervisedRL (CURL). CURL [Laskin et al., 2020a] utilizes data augmentation and contrastive learning to train an image encoder fθ(o) in a self-supervised way, imposing an instance discrimination between similar (+) and dissimilar (-) encoded states. Given a batch of visual observations o, each is augmented twice and encoded into a query q = fθ(oq) and a key k = f¯ θ(ok). The key encoder f¯ θ is a momentum-based moving average of the query encoder fθ to ensure consistency and stability, and fθ is learned by enforcing q to match with k+ while keeping far apart from k-.

Data Regularized Q-Learning (DrQ). Based on SAC settings, DrQ [Yarats et al., 2021b] incorporates optimality invariant image transformations to regularize the Q-function, improving robust learning directly from raw pixels. Let g(o) represent the random image crop augmentation on observations o. It should ideally preserve the Q-values s.t. Q(o, a) = Q(gi(o), a), ∀o ∈ O, a ∈ A, i = 1, 2, 3, · · · . DrQ then applies data augmentation to each transition tuple τt = (ot, at, re t , dt, o′t ) that is uniformly sampled from the replay buffer B, where dt is the done signal. With K augmented next observations gk(o′t ) and M augmented current observations gm(ot), the critic Qφ can be regularized by averaging over M augmented inputs from ot,

where Tt is the soft target value and it can also be regularized by averaging over K augmented inputs from o′t ,

4 The Proposed CCLF

Contrastive-Curiosity-driven Learning Framework (CCLF) extends the model-free RL to further improve sample efficiency when learning directly from the raw pixels. In particular, it fully exploits sample importance for agents to efficiently learn from the most informative data. Firstly, we repurpose the contrastive term in CURL [Laskin et al., 2020a] without additional architectures to quantify the contrastive curiosity (Section 4.1). Subsequently, this contrastive curiosity is coherently integrated into four components to navigate RL with minimum modification: augmented input selection (Section 4.2), experience replay (Section 4.3), Q-function and encoder regularization (Section 4.4), and environment exploration (Section 4.5), as illustrated in Figure 1. Without loss of generosity, we apply CCLF on the state-of-the-art off-policy RL algorithm, SAC [Haarnoja et al., 2018] as summarized in Algorithm 1. Extensions on other base algorithms are carried out in Section 5.2 and B.2 and B.3.

4.1 Contrastive Curiosity

Curiosity can be aroused by an unexpected stimulus that behaves differently from the agent's internal belief. To quantify such surprise-aroused curiosity, we define the agent's curiosity cij by the prediction error of whether any two augmented observations gi(o), gj(o) are from the same observation o,

where IB represents agent's internal belief of whether gi(o), gj(o) are augmented (e.g., randomly cropped) from the same o with similar representations. Since the contrastive loss can be viewed as the log-loss of a softmax-based classifier to match a query q with the key k from the same observation in a batch, it becomes a natural choice for measuring agent's in-

exp(qT Wk+)

ternal belief IB(gi(o), gj(o)) = exp(qT Wk+)+∑B-1 l=1 exp(qTWkl ) , where B is the batch size and q is the query encoder q = fθ(gi(o)). Moreover, we denote the key encoder k = f¯ θ (gj (o)) as k+ if its input is the same as that in the query encoder q; otherwise, we denote it as k-. An immediate merit is that the contrastive curiosity does not require any additional architecture or auxiliary loss because IB is updated directly through representation learning in a self-supervised way.

A higher contrastive curiosity value indicates that the agent does not believe q is similar to k+ or the agent mistakenly matches q with some k- instead, which ultimately results in a surprise in a self-supervised way. It further implies that the sampled transition tuple contains novel information that has yet been learned by the agent, and the encoder fθ is not optimal to extract a meaningful state representation from raw pixels. With the proposed contrastive curiosity in-place, we can integrate different curiosity-driven mechanisms in the proposed CCLF to achieve sample-efficient RL, which is discussed in the following sub-sections.

4.2 Curiosity-Based Augmented Input Selection

Although DrQ [Yarats et al., 2021b] has shown that increasing [K,M], i.e., the amounts of augmented inputs on next and current observations respectively, can potentially improve agent's performances through regularized Q-learning, a crucial trade-off is the introduced higher computational complexity. In addition, more augmented data does not necessarily lead to better performance, as data transformations might alter the semantics and result in the counterproductive performance. To tackle these challenges, we aim to select the most informative inputs for the subsequent learning. Without loss of generality, we assume two inputs are selected fromM augmented current observations ot and similarly two are selected fromK augmented next observations o′t in this paper.

It should be noted that there are various ways to select the most informative augmented inputs, where one straightforward way is to select by least overlap in pixels,

However, a more human-like way is to select the most representative inputs based on the curious level conceptualized by the internal belief rather than simple visual overlaps. Therefore, we propose to select the augmented inputs that cause highest contrastive curiosity as defined in Eq. (3),

In this way, the augmented inputs that are most challenging for matching can be curiously identified since they potentially contain novel knowledge that has yet been learned; meanwhile, this selection mechanism can help to encode more representative state information from the selected inputs, while the agent's internal belief can be jointly updated. As a result, an improved encoder that is robust to different views of observations can be trained with fewer inputs, potentially yielding an improvement in the sample efficiency.

4.3 Curiosity-Based Experience Replaying

In the conventional off-policy RL, agents uniformly sample transitions τ from the replay buffer to learn policies. Although they can eventually perform a complex task by repeatedly practicing in a trial-and-error fashion, we hypothesize that a more sample-efficient and generalizable way is to revisit the transitions that are relatively new or different more frequently. Therefore, we prioritize the experience replay by assigning different prioritization weights w ∈ [0, 1] to all transitions stored in the replay buffer B. In particular, the prioritization weight is initialized to w0 = 1 for any newly added transition tuple. Thereafter, we propose to update the weights of transitions with the overall contrastive curiosity at each training step s,

where β ∈ [0, 1] is a momentum coefficient, and ci∗ j ∗ , c′i∗′ j ∗′ are the contrastive curiosity about o and o′ respectively. The intuition of the momentum update is to maintain a stable update such that the transitions arousing low curiosity will be gradually de-prioritized for learning. Mathematically, the probability of τi to be replayed is p(τi) = wi ∑Nn=1 wn and it becomes small only when τi has been sampled many times. Hence, more recent and surprising transitions arousing high curiosity can be sampled more frequently to learn.

4.4 Curiosity-Based Regularization

Although agents can benefit from learning the selected complex inputs, it imposes challenges for agents as well, which may cause unstable and poor performances. Hence, it is crucial to adapt the learning process by concentrating more on under-learned knowledge. To achieve this, we propose an adaptive regularization for both Q-function and the observation encoder, guided by the contrastive curiosity in order to learn more from the selected inputs arousing high curiosity. In particular, we modify Eq. (2) and Eq. (1) as

It is worth noting that this regularized Q-function is rather general to recover other state-of-the-art works as special cases. When all augmented inputs arouse exactly moderate level of curiosity ci∗ j ∗ = c′i∗′ j ∗′ = 1 2 , the proposed regularization is equivalent to DrQ with [K,M] = [2, 2]. Moreover, when the agent can perfectly match the two augmented inputs

Algorithm 1 An Implementation of CCLF on SAC

Input: MDP τt = (ot, at, re t , dt, o′t ), numbers of augmented inputs [K,M], replay buffer B, training step T , batch size B Parameter: Observation encoder network θ, actor network ψ, critics networks φi, temperature coefficient α, and bilinear product weight W

Output: Optimal policy π∗ψ

Sample a minibatch {(ol, al, re l , dl, o′l )}Bl=1 wl∼ B based on the prioritization weight wl

for each sample τl in the minibatch do

Augment ol and o′l via g(·) to obtain M and K inputs Evaluate the contrastive curiosity cij by Eq. (3) Select gi∗(ol), gj ∗(ol) fromM augmented ol and select gi∗′ (o′l ), gj ∗′ (o′l ) from K augmented o′l by Eq. (5) rl = re l + ri l with ri l from Eq. (9)

Update wl according to Eq. (6)

end for

Update critics Qφi by Eq. (7)

Update the actor πψ and temperature coefficient α

Update encoder fθ and W by Eq. (8)

ot+1 = o′t

end for

with no contrastive curiosity ci∗ j ∗ = c′i∗′ j ∗′ = 0, it is sufficient to update the Q-function with only one input; when the agent fails to encode any similarity and becomes extremely curious ci∗ j ∗ = c′i∗′ j ∗′ = 1, it should focus completely on the novel input instead. Both cases can reproduce the work of RAD [Laskin et al., 2020b]. Most importantly, our proposed Q-function regularization enables the agent to adapt the learning process in a self-supervised way such that it is fully controlled by the conceptualized contrastive curiosity to exploit sample importance and stabilize the learning process.

Similarly, we also regularize the representation learning in a curious manner, inspired by the solution to the supervised class imbalance problem. To deal the with training set containing under-represented classes, a practical approach is to inversely weight the loss of each class according to their sizes. We follow this motivation to incorporate the contrastive curiosity cbi∗ j ∗ about the current observations as the weight for each log-loss class b to update the encoder fθ,

where samples arousing high contrastive curiosity will be considered as under-represented classes and therefore agents need to adaptively pay more attention during the representation learning by optimizing fθ. Meanwhile, agents also jointly re-calibrate a proper internal belief by updating W.

4.5 Curiosity-Based Exploration

Intrinsic rewards can motivate agents to explore actively [Sun et al., 2022], improving the sample efficiency in the conventional RL. While SAC alone can be viewed as the entropy maximization of agent's actions intrinsically, in the proposed CCLF, we explicitly define an intrinsic reward proportional to the average contrastive curiosity about ot and o′t ,

where λ is a temperature coefficient, η is a decay weight, t is the environment step, remax and rimax are respectively the maximum extrinsic and intrinsic rewards over step t.

With the proposed ri t to supplement re t in Eq. (7), agents can be encouraged to explore the surprising states that arouse high contrastive curiosity substantially. In particular, higher ri t rewards agents for exploration when different views of the same observations produce inconsistent representations. Meanwhile, ri t is decayed with respect to the environment step t to ensure the convergence of policies. As the extrinsic reward re differs across different tasks, the normalization is performed to balance re and ri . This formulation is similar to the intrinsic reward in CCFDM [Nguyen et al., 2021], but the proposed CCLF does not require a forward dynamic model or action embedding that increases the model complexity.

5 Experiments and Results

5.1 Experimental Setup

We empirically evaluate the proposed CCLF in terms of sample efficiency and ultimate performance, on 6 continuous control tasks from the DMC suite [Tunyasuvunakool et al., 2020], 26 discrete control tasks from the Atari Games [Bellemare et al., 2013] and 3 navigation tasks with sparse extrinsic rewards from the MiniGrid [Chevalier-Boisvert et al., 2018]. In this section, we mainly present the experimental results in the DMC suite with SAC being the base algorithm while detailed settings and results in the Atari Games and the MiniGrid are included in Appendix B.2, B.3, C.2, and C.3. For a comprehensive evaluation in the DMC suite, we include the following baselines to compare against:

• Pixel-based SAC (SAC-Pixel) [Haarnoja et al., 2018]
• CURL [Laskin et al., 2020a].
• DrQ [Yarats et al., 2021b] with [K,M] = [2, 2] and a modified augmentation method for consistency.
• Hybrids of CURL and DrQ: CURL+ and CURL++, where contrastive representation learning is integrated to DrQ for [K,M] = [2, 2] and [5, 5] respectively.
• Augmented input selection models: 2 out of 5 inputs for each sample are selected by pixel overlap (Select) via Eq. (4) and contrastive curiosity (Select+) via Eq. (5) without the other curiosity-based components.

The detailed setting of hyper-parameters is provided in Appendix B.1. For our proposed CCLF, we initialize it with [K,M] = [5, 5] to generate a sufficiently large amount of augmented inputs. For simplicity, we fix i∗ randomly and only select j ∗ via Eq. (5) for the augmented input selection.

5.2 Results and Discussion

Not all Samples are Equally Important. In CURL+, data augmentation is applied twice for each sampled transition while 5 times in CURL++. Since CURL++ injects 2.5× larger amount of inputs than CURL+, its computational complexity increases dramatically. Table 1 shows that CURL++ performs worse than CURL+ in 4 tasks at 100K steps and slightly outperforms CURL+ in only 2 tasks at 500K steps. In Figure 2 and Appendix Figure 5, the learning curve of CURL++ is clearly below CURL+ at first and gradually approaches to the same level as CURL+. Since more augmented inputs may not guarantee the consistency of semantics, additional training is often required for convergence. Therefore, we can empirically validate the hypothesis that not all augmented inputs are equally important and simply increasing the number of augmentations is instead inefficient. A similar result can be found in DrQ Appendix F [Yarats et al., 2021b].

Main Results on the DMC Suite

The average sample efficiency and asymptotic performance are shown in Table 1 at 100K and 500K environment steps, respectively. Meanwhile, Figure 2 demonstrates the agent's learning capability over 500K steps. Compared SAC-Pixel to the other models in Figure 2, its performance is not improved in all 6 tasks even until 500K steps while the other models can asymptotically perform well. Thus, it is challenging for the conventional SAC to learn directly from raw pixels and a sample-efficient RL method is needed to aid that.

According to Table 1, Select performs better than Select+, on 3 tasks at 100K and 4 tasks at 500K, with more stable learning curves as shown in Figure 2. Indeed, the inputs in Select may contain some invariances to improve sample efficiency and learning capability. However, more under-learned inputs with even richer invariances are present in Select+, and agents cannot adapt the learning process in this model, causing the instability issue in Select+.

To tackle this issue, the proposed CCLF collaboratively adapts the learning process with the selected inputs and contrastive curiosity, so the learning curves become more smooth than others in Figure 2. In particular, CCLF outperforms all baselines in 5 tasks at both 100K and 500K regimes in Table 1. Moreover, it converges much faster than Select+ according to the results at 100K steps. In fact, the proposed CCLF only requires about 50% environment steps to converge to desirable performances on 3 tasks (Ball in Cup, Walker, and Cartpole) as the other baselines. In addition, it even benchmarks on Cheetah-Run and Finger-Spin tasks at 500K steps. Therefore, we can conclude that our proposed CCLF can improve the sample efficiency and learning capabilities of RL agents, with fewer environment interactions and 60% reduced augmented inputs. We also analyze the computational complexity on the Cartpole task by model sizes and training time in Appendix C.1, where CCLF can avoid increasing the training cost dramatically.

Additional Experiments in Atari Games. In addition to continuous control tasks, CCLF can also be incorporated into Rainbow DQN [Hessel et al., 2018] to perform discrete control tasks. As shown in Appendix C.2, the proposed CCLF attains state-of-the-art performances in 8 out of 26 Atari Games at 100K steps. In particular, CCLF is superior to CURL in 11 games and DrQ in 18 games, which favorably indicates the effectiveness of improving sample efficiency.

Further Investigation on MiniGrid. Apart from offpolicy algorithms, we also investigate the compatibility on the on-policy algorithm. More specifically, we extend the proposed CCLF to A2C [Mnih et al., 2016] and RE3 [Seo et al., 2021] to perform navigation tasks with spares rewards in MiniGrid. We first adapt CCLF to the on-policy algorithm by removing the experience replay component. Note that the input from MiniGrid is already a compact and efficient 7×7×3 embedding of partially-observable 7×7 grids, so even slight augmentation will induce highly inconsistent learned features. Thus, we directly duplicate the embedding without random augmentations to obtain contrastive curiosity for regularization and intrinsic reward. Figure 3 shows that CCLF exhibits superior sample efficiency and learning capabilities in all three tasks, even model-agnostic with the state-of-the-art curiosity-driven method RE3. In the Empty16×16 task, our CCLF can reach the optimal level in about 50% and 55% of the training steps of RE3 and A2C, respectively. By comparing the ultimate performance scores, the proposed CCLF obtains 1.63× higher average performance than RE3 in DoorKey-6×6 and 1.3× in DoorKey-8×8.

Effectiveness of the Proposed RL Components. One might wonder if the proposed CCLF benefits mainly from one or several curiosity-based components in practice. Hence, we empirically examine the effectiveness of all possible combinations of the four curiosity-driven components on the Cartpole task from the DMC suite. The results are included in Appendix C.4, where it can be concluded that all four components are necessary and important to attain state-of-the-art performances. Our proposed CCLF can navigate all four RL components together to improve the sample efficiency and resolve instability, which demonstrates effective collaboration.

6 Conclusion

In this paper, we present CCLF, a contrastive-curiosity-driven learning framework for RL with visual observations, which can significantly improve the sample efficiency and learning capabilities of agents. As we empirically find that not all samples nor their augmented inputs are equally important for RL, CCLF encourages agents to learn in a curious way, exploiting sample complexity and importance systemically.

Acknowledgments

This research is supported in part by the National Research Foundation, Prime Minister's Office, Singapore under its NRF Investigatorship Programme (NRFI Award No. NRFNRFI05-2019-0002). Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore. This research is supported in part by the Alibaba-NTU Singapore Joint Research Institute, Nanyang Technological University, and in part by the Singapore Ministry of Health under its National Innovation Challenge on Active and Confident Ageing (NIC Project No. MOH/NIC/COG04/2017) and (NIC Project No. MOH/NIC/HAIG03/2017). H.Qian thanks the support from the Wallenberg-NTU Presidential Postdoctoral Fellowship.