SURF: Semi-supervised Reward Learning with Data Augmentation for Feedback-efficient Preference-based Reinforcement Learning

1 INTRODUCTION . Reward function plays a crucial role in reinforcement learning ( RL ) to convey complex objectives to agents . For various applications , where we can design an informative reward function , RL with deep neural networks has been used to solve a variety of sequential decision-making problems , including board games ( Silver et al. , 2017 ; 2018 ) , video games ( Mnih et al. , 2015 ; Berner et al. , 2019 ; Vinyals et al. , 2019 ) , autonomous control ( Schulman et al. , 2015 ; Bellemare et al. , 2020 ) , and robotic manipulation ( Kober & Peters , 2011 ; Kober et al. , 2013 ; Kalashnikov et al. , 2018 ; Andrychowicz et al. , 2020 ) . However , there are several issues in reward engineering . First , designing a suitable reward function requires more human effort as the tasks become more complex . For example , defining a reward function for book summarization ( Wu et al. , 2021 ) is non-trivial because it is hard to quantify the quality of summarization in a scale value . Also , it has been observed that RL agents could achieve high returns by discovering undesirable shortcuts if the hand-engineered reward function does fully specify the desired task ( Amodei et al. , 2016 ; Hadfield-Menell et al. , 2017 ; Lee et al. , 2021a ) . Furthermore , there are various domains , where a single ground-truth function does not exist , and thus personalization is required by modeling different reward functions based on the user ’ s preference . Preference-based RL ( Akrour et al. , 2011 ; Christiano et al. , 2017 ; Ibarz et al. , 2018 ; Lee et al. , 2021a ) provides an attractive alternative to avoid reward engineering . Instead of assuming a handengineered reward function , a ( human ) teacher provides preferences between the two agent behaviors , and an agent learns how to show the desired behavior by learning a reward function , which is consistent with the teacher ’ s preferences . Recent progress of preference-based RL has shown that the teacher can guide the agent to perform novel behaviors ( Christiano et al. , 2017 ; Stiennon et al. , 2020 ; Wu et al. , 2021 ) , and mitigate the effects of reward exploitation ( Lee et al. , 2021a ) . However , existing preference-based approaches often suffer from expensive labeling costs , and this makes it hard to apply preference-based RL to various applications . Meanwhile , recent state-of-the-art system in computer vision , the label-efficiency problem has been successfully addressed through semi-supervised learning ( SSL ) approaches ( Berthelot et al. , 2019 ; 2020 ; Sohn et al. , 2020 ; Chen et al. , 2020b ) . By leveraging unlabeled dataset , SSL methods have improved the performance with low cost . Data augmentation also plays a significant role in improving the performance of supervised learning methods ( Cubuk et al. , 2018 ; 2019 ) . By using multiple augmented views of the same data as input , the performance has been improved by learning augmentation-invariant representations . Inspired by the impact of semi-supervised learning and data augmentation , we present SURF : a Semi-sUpervised Reward learning with data augmentation for Feedback-efficient preference-based RL . To be specific , SURF consists of the following key ingredients : ( a ) Pseudo-labeling ( Lee , 2013 ; Sohn et al. , 2020 ) : We leverage unlabeled data by utilizing the artificial labels generated by learned preference predictor , which makes the reward function produce a confident prediction ( see Figure 1a ) . We remark that such a SSL approach is particularly attractive in our setup as an unlimited number of unlabeled data can be obtained with no additional cost , i.e. , from past experiences stored in the buffer . ( b ) Temporal cropping augmentation : We generate slightly shifted or resized behaviors , which are expected to have the same preferences from a teacher , and utilize them for reward learning ( see Figure 1b ) . Our data augmentation technique enhances the feedback-efficiency by enforcing consistencies ( Xie et al. , 2019 ; Berthelot et al. , 2020 ; Sohn et al. , 2020 ) to the reward function . We remark that SURF is not a naïve application of these two techniques , but a novel combination of semi-supervised learning and the proposed data augmentation , which has not been considered or evaluated in the context of the preference-based RL . Our experiments demonstrate that SURF significantly improves the preference-based RL method ( Lee et al. , 2021a ) on complex locomotion and robotic manipulation tasks from DeepMind Control Suite ( Tassa et al. , 2018 ; 2020 ) and Meta-world ( Yu et al. , 2020 ) , in terms of feedbackefficiency . In particular , our framework could make RL agents achieve ∼100 % of success rate on complex robotic manipulation task using only a few hundred preference queries , while its baseline method only achieves ∼50 % of the success rate under the same condition ( see Figure 3 ) . 2 RELATED WORK . Preference-based RL . In the preference-based RL framework , a ( human ) supervisor provides preferences between the two agent behaviors and the agent uses this feedback to perform the task ( Christiano et al. , 2017 ; Ibarz et al. , 2018 ; Leike et al. , 2018 ; Stiennon et al. , 2020 ; Wu et al. , 2021 ; Lee et al. , 2021a ; b ) . Since this approach is only feasible if the feedback is practical for a human to provide , several strategies have been studied in the literature . Ibarz et al . ( 2018 ) initialized the agent ’ s policy with imitation learning from the expert demonstrations , while Lee et al . ( 2021a ) utilized unsupervised pre-training for policy initialization . Several sampling schemes ( Sadigh et al. , 2017 ; Biyik & Sadigh , 2018 ; Biyik et al. , 2020 ) to select informative queries also have been adopted for improving the feedback-efficiency . Our approach differs in that we utilize unlabeled samples for reward learning , and also provide a novel data augmentation technique for the agent behaviors . Data augmentation for RL . In the context of RL , data augmentation has been widely investigated for improving data-efficiency ( Srinivas et al. , 2020 ; Yarats et al. , 2021 ) , or RL generalization ( Cobbe et al. , 2019 ; Lee et al. , 2019 ) . For example , RAD ( Laskin et al. , 2020 ) demonstrated that data augmentation , such as random crop , can improve both data-efficiency and generalization of RL algorithms . While these methods are known to be beneficial to learn policy in the standard RL setup , they have not been tested for learning rewards . To the best of our knowledge , we present the first data augmentation method specially designed for learning reward function . Semi-supervised learning . The goal of semi-supervised learning ( SSL ) is to leveraging unlabeled samples to improve a model ’ s performance when the amount of labeled samples are limited . In an attempt to leverage the information in the unlabeled dataset , a number of techniques have been proposed , e.g. , entropy minimization ( Grandvalet & Bengio , 2004 ; Lee , 2013 ) and consistency regularization ( Sajjadi et al. , 2016 ; Miyato et al. , 2018 ; Xie et al. , 2019 ; Sohn et al. , 2020 ) . Recently , the combination of these two approaches have shown state-of-the-art performance in benchmarks , e.g. , MixMatch ( Berthelot et al. , 2019 ) , and ReMixMatch ( Berthelot et al. , 2020 ) , when used with advanced data augmentation techniques ( Zhang et al. , 2018 ; Cubuk et al. , 2019 ) . Specifically , FixMatch ( Sohn et al. , 2020 ) revisits pseudo-labeling technique and demonstrates that joint usage of pseudo-labels and consistency regularization achieves remarkable performance due to its simplicity . 3 PRELIMINARIES . Reinforcement learning ( RL ) is a framework where an agent interacts with an environment in discrete time ( Sutton & Barto , 2018 ) . At each timestep t , the agent receives a state st from the environment and chooses an action at based on its policy π ( at|st ) . In conventional RL framework , the environment gives a reward r ( st , at ) and the agent transitions to the next state st+1 . The return Rt = ∑∞ k=0 γ kr ( st+k , at+k ) is defined as discounted cumulative sum of the reward with discount factor γ ∈ [ 0 , 1 ) . The goal of the agent is to learn a policy that maximizes the expected return . Preference-based reinforcement learning . In this paper , we consider a preference-based RL framework , which does not assume the existence of hand-engineered reward . Instead , a ( human ) teacher provides preferences between the agent ’ s behaviors and the agent uses this feedback to perform the task ( Christiano et al. , 2017 ; Ibarz et al. , 2018 ; Leike et al. , 2018 ; Stiennon et al. , 2020 ; Lee et al. , 2021a ; b ; Wu et al. , 2021 ) by learning a reward function , which is consistent with the observed preferences . We formulate a reward learning problem as a supervised learning problem ( Wilson et al. , 2012 ; Christiano et al. , 2017 ) . Formally , a segment σ is a sequence of observations and actions { ( sk , ak ) , ... , ( sk+H−1 , ak+H−1 ) } . Given a pair of segments ( σ0 , σ1 ) , a teacher gives a feedback indicating which segment is preferred , i.e. , y ∈ { 0 , 1 , 0.5 } , where 1 indicates σ1 σ0 , 0 indicates σ0 σ1 , and 0.5 implies an equally preferable case . Each feedback is stored in a dataset D as a triple ( σ0 , σ1 , y ) . Then , we model a preference predictor using the reward function r̂ψ following the Bradley-Terry model ( Bradley & Terry , 1952 ) : Pψ [ σ 1 σ0 ] = exp ( ∑ t r̂ψ ( s 1 t , a 1 t ) ) ∑ i∈ { 0,1 } exp ( ∑ t r̂ψ ( s i t , a i t ) ) , ( 1 ) where σi σj denotes the event that segment i is preferable to segment j . The underlying assumption of this model is that the teacher ’ s probability of preferring a segment depends exponentially on the accumulated sum of the reward over the segment . The reward model is trained through supervised learning with teacher ’ s preferences . Specifically , given a dataset of preferences D , the reward function is updated by minimizing the binary cross-entropy loss : LCE = E ( σ0 , σ1 , y ) ∼D [ LReward ] = − E ( σ0 , σ1 , y ) ∼D [ ( 1− y ) logPψ [ σ0 σ1 ] + y logPψ [ σ1 σ0 ] ] . Algorithm 1 SURF Require : Hyperparameters : unlabeled batch ratio µ , threshold parameter τ , and loss weight λ Require : Set of collected labeled data Dl , and unlabeled data Du 1 : for each gradient step do 2 : Sample labeled batch { ( σ0l , σ1l , y ) ( i ) } Bi=1 ∼ Dl 3 : Sample unlabeled batch { ( σ0u , σ1u ) ( j ) } µBj=1 ∼ Du 4 : // DATA AUGMENTATION FOR LABELED DATA 5 : for i in 1 . . . B do 6 : ( σ̂0l , σ̂ 1 l ) ( i ) ← TDA ( ( σ0l , σ1l ) ( i ) ) in Algorithm 2 7 : end for 8 : // PSEUDO-LABELING AND DATA AUGMENTATION FOR UNLABELED DATA 9 : for j in 1 . . . µB do 10 : Predict pseudo-labels ŷ ( ( σ0u , σ1u ) ( j ) ) 11 : ( σ̂0u , σ̂1u ) ( j ) ← TDA ( ( σ0u , σ1u ) ( j ) ) in Algorithm 2 12 : end for 13 : Optimize LSSL ( 3 ) with respect to ψ 14 : end for The reward function r̂ψ is usually optimized only using labels from real human , which are expensive to obtain in practice . Instead , we propose a simple yet effective method based on semi-supervised learning and data augmentation to improve the feedback-efficiency of preference-based learning .

Summary: This work introduces a semi-supervised reward learning approach to reduce the efforts of reward engineering, which contains two key components. The first is to produce artificial labels for these unlabeled samples leveraging the pseudo-labeling and the learned reference predictor. The second component is to crop consecutive subsequences for data augmentation. The proposed approach is tested on Meta-world and DMControl suites, and the results show that it significantly improves the performances.