STRAPPER: Preference-based Reinforcement Learning via Self-training Augmentation and Peer Regularization

PbRL is one of RL paradigms to learn from human feedback BorjaIbarz2018RewardLF . Several works have successfully utilized feedback from real human to train RL agents DilipArumugam2019DeepRL ; PaulFChristiano2017DeepRL ; BorjaIbarz2018RewardLF ; WBradleyKnox2009InteractivelySA ; KiminLee2021PEBBLEFI ; GarrettWarnell2017DeepTI . PaulFChristiano2017DeepRL scales PbRL to utilize modern deep learning techniques, and BorjaIbarz2018RewardLF improves the efficiency of this method by introducing additional forms of feedback such as demonstrations. Recently, KiminLee2021PEBBLEFI proposes a feedback-efficient RL algorithm by utilizing off-policy learning and pre-training. park2022surf uses pseudo-labeling to utilize unlabeled segments and proposes a novel augmentation called temporal cropping to augment labeled data.

Self-training denotes a learning setting where the supervision on unlabeled data is given by the own prediction of the model trained with the labeled data DavidYarowsky1995UNSUPERVISEDWS ; KamalNigam2000AnalyzingTE ; zhu2005semi ; CharlesJRosenberg2005SemiSupervisedSO . Pseudo-labeling refers to a specific variant, where model predictions are converted to hard labels lee2013pseudo . This is often used along with a confidence-based threshold that retains unlabeled examples only when the classifier is sufficiently confident. Consistency regularization was first proposed by bachman2014learning and later referred to as the “$\Pi$-Model” rasmus2015semi , where the main idea is to force the output of the model remain constant with randomly augmented inputs. After that, FixMatch sohn2020fixmatch first combines consistency regularization and pseudo-labeling in one simple method. On the other hand, NoisyStudent QizheXie2020SelfTrainingWN employs a self-training scheme, a teacher model generates pseudo-labels on unlabeled data, and a larger noisy student model is then trained on both labeled and unlabeled data with consistency regularization. Recently, zhu2021rich proposes an analytical framework to unify consistency regularization with explicit and implicit pseudo-labels.

We consider a standard RL framework where an agent interacts with an environment in discrete time RichardSutton1988ReinforcementLA . Formally, at each timestep $t$, the agent receives a state $\mathbf{s}_{t}$ from the environment and chooses an action $\mathbf{a}_{t}$ based on its policy $\pi_{\phi}$. In traditional RL, the environment also returns a reward $r(\mathbf{s}_{t},\mathbf{a}_{t})$. The return $\mathcal{R}_{t}=\sum^{\infty}_{k=0}\gamma^{k}r_{t+k}$ is the discounted sum of rewards from timestep $t$ with discount factor $\gamma\in(0,1]$. RL then maximizes the expected return $\mathcal{R}_{0}$ with respect to $\pi_{\phi}$.

However, for many complex domains and tasks, it is difficult to construct a suitable reward function. We consider the PbRL framework, where human provide preferences between two behavior segments and then the agent trains under this supervision PaulFChristiano2017DeepRL ; BorjaIbarz2018RewardLF ; KiminLee2021PEBBLEFI ; JanLeike2018ScalableAA ; KiminLee2021BPrefBP . In this work, we follow the classic framework to learn a reward function $\widehat{r}_{\phi}$ from preferences, where the function is trained to be consistent with human feedback AaronWilson2012ABA ; PaulFChristiano2017DeepRL . In this framework, a segment $\sigma$ is a sequence of states and actions $\{\mathbf{s}_{h},\mathbf{a}_{h},\dots,\mathbf{s}_{h+H},\mathbf{a}_{h+H}\}$. Transitions $(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1})$ are stored in a replay buffer $\mathcal{B}$, and we sample segments $\sigma$ from $\mathcal{B}$. Then, we elicit preferences $y$ for segments $\sigma^{i}$ and $\sigma^{j}$. More specifically, $y$ indicates which segment is preferred, i.e., $y\in\{0,1,0.5\}$, where $0$ indicates $\sigma^{i}\succ\sigma^{j}$ (the event that segment $\sigma^{i}$ is preferable to $\sigma^{j}$), $1$ indicates $\sigma^{j}\succ\sigma^{i}$ ($\sigma^{j}$ is preferable to $\sigma^{i}$), and $0.5$ implies an equally preferable case. The judgment is recorded in a dataset $D$ as a triple $(\sigma^{i},\sigma^{j},y)$.

By following the Bradley-Terry model Bradley1952RankAO , we have a preference predictor as follows:

where $\widehat{r}_{\psi}$ is the reward function. Intuitively, we assume that the probability of preferring one segment is exponentially proportional to the sum of an underlying reward function over the segment. While $\widehat{r}_{\psi}$ is not a binary classifier, learning $\widehat{r}_{\psi}$ amounts to binary classification with labels $y$ provided by an annotator. Concretely, the reward function, modeled as a neural network with parameters $\psi$, is updated by minimizing the following loss:

where $y(0)$ and $y(1)$ represent whether the preference label is 0 or 1, respectively. Finally, we use the learned reward function $\widehat{r}_{\psi}$ to train the policy.

Here, we model PbRL into the framework of the semi-supervised learning. The labeled dataset $D_{L}$ is first formed with segment pairs $(\sigma^{i},\sigma^{j})$ sampled from trajectory buffer $\mathcal{B}$ and the preferences $y$ queried from human, e.g., $D_{L}:=\{(\sigma^{i},\sigma^{j},y)\}$. We use $D_{U}$ to denote the dataset formed with segment pairs $(\sigma^{i},\sigma^{j},\cdot)$ without human labels.

In STRAPPER, we first train the parameters of the teacher model $\psi^{k}$ on the labeled dataset $D_{L}$ to minimize the Eq.2, which we abbreviate as

Then, we use the teacher model to generate pseudo preferences $y^{k}$ for unlabeled segments pairs $(\sigma^{i},\sigma^{j})$ which are sampled independently from the buffer $\mathcal{B}$, $y^{k}=\left(P_{\psi^{k}}[\sigma^{i}\succ\sigma^{j}],P_{\psi^{k}}[\sigma^{j}\succ\sigma^{i}]\right),(\sigma^{i},\sigma^{j})\sim D_{U}$. We augment the segment $\sigma$ to $\tilde{\sigma}$ for consistency regularization, so $y^{k}=\left(P_{\psi^{k}}[\tilde{\sigma}^{i}\succ\tilde{\sigma}^{j}],P_{\psi^{k}}[\tilde{\sigma}^{j}\succ\tilde{\sigma}^{i}]\right),(\sigma^{i},\sigma^{j})\sim D_{U}$. Then, the triple $(\sigma^{i},\sigma^{j},y^{k})$ is recorded in the pseudo-labeled dataset $D_{U}^{k}$. After that, the student model is trained with $D_{L}$ and $D_{U}^{k}$:

For simplicity, we define the mixed dataset $D_{L}\cup D_{U}^{k}$ as $\tilde{D}={(\sigma^{i},\sigma^{j},\tilde{y})}$, where $\tilde{y}=y,\forall{(\sigma^{i},\sigma^{j})\sim D_{L}}$, and $\tilde{y}=y^{k},\forall{(\sigma^{i},\sigma^{j})\sim D_{U}}$, to simplify the above expression as:

Overall Loss Function. For semi-supervised problems in computer vision, it has been both empirically and theoretically proved that consistency regularization performs well. But as discussed in Section 4, the inherent phenomenon of similarity trap in PbRL poses difficulties for the direct use of consistency regularization. Due to the presence of the similarity trap, consistency regularization improperly increases the consistency of model predictions between disjoint segment pairs and reduces the confidence in reward learning, since the augmented distribution does not match with the original one. Inspired by YangLiu2020PeerLF , we propose peer regularization to deal with such issue. By adding peer regularization to Eq.5, we obtain the final loss function for training the student model as:

where $(\sigma^{i}_{n_{1}},\sigma^{j}_{n_{1}},\tilde{y}_{n_{1}})$ and $(\sigma^{i}_{n_{2}},\sigma^{j}_{n_{2}},\tilde{y}_{n_{2}})$ are two independently sampled segment pairs and corresponding labels. We call these two samples as peer samples. The second term in Eq.6, called peer regularization, encourages the “inconsistency” of model predictions between random sample paires. Intuitively, since $\tilde{y}_{n_{2}}$ would not provide any valuable information to label $(\sigma^{i}_{n_{1}},\sigma^{j}_{n_{1}})$, we penalize the reward model memorizing uninformative labels. Formally, the peer regularization leads the training to generate more confident predictions:

The above theorem implies that peer regularization leads to either $P_{\psi^{k+1}}[\sigma^{i}\succ\sigma^{j}]\rightarrow 1$ or $P_{\psi^{k+1}}[\sigma^{j}\succ\sigma^{i}]\rightarrow 1$, which indicates a confident prediction. This counteracts the negative impact of consistency regularization. At the same time, confident prediction can help distinguish confusing pairs of similar segments $(\sigma^{1},\widetilde{\sigma}^{1})$, tending to pseudo-label $(0,1)$ or $(1,0)$ rather than $(0.5,0.5)$. This would lead the pseudo-label to be closer to the real label.

Learning with Noisy Label. In addition, we further describe our proposed peer regularization in terms of learning with noisy labels. First, the similarity trap results in a noisy pseudo-labeling process. Because consistency regularization improperly increases the consistency of model predictions between disjoint segment pairs. Here, the noise is introduced by the training process and therefore independent of the data distribution. Second, PbRL is a binary classification problem. Such two problem settings are just consistent with that in YangLiu2020PeerLF . Therefore, we can also straightforward use the method proposed in YangLiu2020PeerLF for solving the problem of training with noisy labels to derive Eq.6 built on peer loss function.

To sum up, we provide the full procedure of STRAPPER in Algorithm 1, which is given along with PEBBLE KiminLee2021PEBBLEFI , an off-policy PbRL algorithm.

Simulate human annotators. Similar to prior work PaulFChristiano2017DeepRL ; KiminLee2021PEBBLEFI , we obtain feedback from simulated human instead of real humans. Following KiminLee2021BPrefBP , we first build simulated annotator to provide (rational and deterministic) preferences $y$ for queries $(\sigma^{i},\sigma^{j})$, oracle behavior for short:

where $r(\cdot,\cdot)$ denotes the oracle reward function.

Hyper-parameters. The implementation of STRAPPER is based on SURF and PEBBLE, and thus inherents the hyperparameter setting from PEBBLE KiminLee2021PEBBLEFI , which is shown in Table 1. An ensemble of three reward models are initialised, while the network structures differ among environments. The reward model for DMC consists of 2 fully-connected linear layers with 1024 neurons in each layer, whereas the one for Meta-World uses three hidden layers of 256 neurons instead. After 10,000 pre-training steps, the reward model are trained via ADAM optimiser. Afterwards, the model gets updated at a fixed frequency as long as the feedback budget is available. We set the update frequency to be 5000 steps for tasks in Meta-world, while the frequencies are 20,000 steps and 30,000 steps for Walker Walk and Quadruped Walk, respectively.

For the first-time learning immeadiately after the pre-training phase, we use uniform sampling to acquire segment pairs with a segment length of 50. To improve the feedback-efficiency for the training procedure, we switch to the disagreement-based sampling scheme for all subsequent feedback sessions. The aim of the disagreement-based method is to find queries with high uncertainty based on the outputs among ensemble of reward models. In each session, the number of labeled queries provided to the reward model is dependent on the training task together with the total feedback budget, and the details are specified in Table 1. On the other hand, as a large portion of unlablled queries are dropped due to the confidence threshold, we sample unlabeled queries as 10 times of labeled queries. We further increase this figure to 100 if the maximum feedback budget is equal or greater than 1000.

Computational resources. Experiments are conducted by using a computational cluster with 22x GeForce RTX 2080 Ti, and 4x NVIDIA Tesla V100 32GB for 7 days.

Benchmark tasks. For simulated human annotators, we take two locomotion tasks from DeepMind Control Suit tassa2018deepmind , Walker-walk and Quadruped-walk, and four robotic manipulation tasks from Meta-world yu2021metaworld , Window Open, Button Press, Drawer Open and Door Open.

We first combine the SOTA PbRL algorithm PEBBLE KiminLee2021PEBBLEFI with the following four alternative SSL methods: 1) Pseudo-labeling (PL) lee2013pseudo : the teacher model generates confident pseudo-labels using a fixed threshold and then updates student model using cross-entropy loss. This is the SSL method employed in SURF park2022surf . 2) vanilla Consistency Regularization (CR) xie2019unsupervised : constrains the output of student model between two augmented inputs using MSEloss without using pseudo-labels. 3) FixMatch(FM) sohn2020fixmatch : generates pseudo-labels with weak augmented inputs and then updates the student model with strong augmented input. 4) self-training with Noisy Student (NS) QizheXie2020SelfTrainingWN : only adds noise to the learning process of the student model. The latter three methods all employ consistency regularization as a component and use random amplitude scaling to augment the input. Random amplitude scaling multiplies a uniform random variable $z$ to the state, i.e., $\hat{s}=s\cdot z$, where $z\sim\text{Unif}[0.995,1.005]$. PEBBLE and SAC using the ground truth reward are also added for benchmark comparison.

Figure 3 shows SAC, PEBBLE and PEBBLE with four alternative SSL methods on locomotion tasks and robotic manipulation tasks. As shown in Fig.3a and Fig.3b, all four SSL methods improve the performance of PEBBLE (orange). And the baseline methods that utilize consistency regularization, vanilla Consistency Regularization (red), self-training with Noisy Student (purple) and FixMatch (brown), outperform or are competitive with Pseudo Labeling (green). But when facing with robotic manipulation tasks as shown in Fig.3c and Fig.3d, these three methods perform poorly. We will show in next subsection that STRAPPER we proposed can alleviate such problem on this type of tasks.

Benchmark Experiments on Robotic Manipulation Tasks. As in Section 6.2, SSL methods using consistency regularization perform well on locomotion tasks (like Walker and Quadruped) but fail on robotic manipulation tasks (like Window Open and Button Press). We attribute this to the fact that similarity trap occurs more frequently in robotic manipulation tasks, affecting the direct use of consistency regularization. Specifically, compared to locomotion tasks, most of robotic manipulation tasks require an agent to interact with objects in the environment. In such scenarios, certain states require precise operation, and subtle differences can bring about serious mistakes. Therefore we mainly focus on comparation between STRAPPER and SOTA baselines on robotic manipulation tasks. PEBBLE+SURF stands for the original implementation of SURF park2022surf , which uses PL, as mentioned in Section 6.2. The temporal data augmentation methods described in park2022surf are employed in both PEBBLE+SURF and STRAPPER. As shown in Fig.4, our method (red) outperforms or is competitive with baselines in robotic manipulation tasks. This empirically demonstrates that the peer regularization we proposed in STRAPPER can counteract the side effect of similarity trap when consistency regularization is used.

More Benchmark Experiments on Locomotion Tasks. As we mention above, the baseline methods that utilizing consistency regularization without the Peer Regularization can improve the performance upon the pseudo labeling. However, when we use Peer Regularization on such locomotion tasks, it made the performance drop, shown as STRAPPER in Fig.5. We believe this is due to the fact that, for locomotion tasks, small differences are not enough to cause fatal errors, so the impact of similarity trap is small. On the contrary, the addition of Peer Regularization causes a certain degree of exploratory deficiency, which in turn leads to performance degradation.

We have shown in the experiments above that the introduction of Peer Regularization in some types of environments can improve the effectiveness of consistency regularization. We will further show that STRAPPER has additional benefits when a more realistic experimental setting is considered. In reality, human annotators are not always perfectly rational, which may provide some stochastic preferences for query. To address this, we generate noisy preferences using a stochastic model as:

where $\gamma\in(0,1]$ is a discount factor, modeling myopic behavior, and $\beta\in(0,\infty]$ , modeling stochastic behavior (expert queries become perfectly rational and deterministic as $\beta\to\infty$). To imitate the accidental error of human expert, we flip the preference with probability of $\epsilon$, denoted as mistake behavior. If both segments (for query) do not contain a desired behavior, human would like to discard the query; we model this noisy behavior as skipping behavior. Further, if both segments have similar returns, human would like to provide a preference with $(0.5,0.5)$; we model this behavior as equally behavior²²2For more behavioral details, we refer the readers to B-Pref benchmark KiminLee2021BPrefBP .. To investigate the influence of noisy preference on PbRL algorithms and study whether STRAPPER is robust to the noisy-labels, we conduct experiments with progressively enhanced noisy labels.

In Figure 6a, we plot the normalized return versus the strength of noise. We observe that the straightforward PEBBLE tends to be brittle and sensitive to the noisy preferences, especially under high-level noise. Similarly, we find that the training across three semi-supervised baselines (CR, PL, NS) can still be unstable despite of incorporating the unlabeled segments. In contrast, STRAPPER yields more robust and stable performance, which consistently outperforms the straightforward PbRL as well as semi-supervised baselines. This result serves as an evidence that unlabeled segments can be useful to eliminate the intrinsic noise in preferences.

Further, we consider both quantity-reduced and noisy preferences in Walker-walk task. In Figure 6b, we illustrate the performance of PEBBLE (top row) and PEBBLE+STRAPPER (bottom row). For PEBBLE, in this specific task, the quantity of preference seems to matter more than the quality. By reducing the amount of preference, the performance drops significantly. We observe that, by incorporating the unlabeled segments, STRAPPER eliminates the effect of the lack of labeled segments and thus leads to better performance, even under a high noise ratio.