A S TUDY OF C AUSAL C ONFUSION IN P REFERENCE -B ASED R EWARD L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Learning policies via preference-based reward learning is an increasingly popular method for customizing agent behavior, but has been shown anecdotally to be prone to spurious correlations and reward hacking behaviors.
While much prior work focuses on causal confusion in reinforcement learning and behavioral cloning, we aim to study it in the context of reward learning.
We demonstrate that preference-based reward learning often fails to produce causal rewards.
In particular, we perform a series of sensitivity and ablation analyses on three benchmark domains where rewards learned from preferences achieve minimal test error but fail to generalize to out-of-distribution states—resulting in poor policy performance when optimized.
We find that the presence of non-causal distractor features, noise in the stated preferences, and partial state observability can all exacerbate causal confusion.
We also identify a set of methods with which to interpret causally confused learned rewards: we observe that optimizing causally confused rewards drives the policy off the reward’s training distribution, resulting in high predicted (learned) rewards but low true rewards.
These findings illuminate the susceptibility of reward learning to causal confusion, especially in highdimensional environments—failure to consider even one of many factors (data coverage, state definition, etc.) can quickly result in unexpected, undesirable behavior.
I NTRODUCTION
Preference-based reward learning (Wirth et al., 2017; Sadigh et al., 2017; Christiano et al., 2017;
Brown et al., 2020b) is a well-studied technique for learning from pairwise preferences or rankings and holds the potential of enabling AI systems to learn specifications for tasks without requiring a human to write down an explicit reward function, and to adapt the AI system’s behavior to individual preferences and needs.
However, recent anecdotal evidence suggests that these methods are prone to learning rewards that pick up on spurious correlations in the data and miss the true underlying causal structure, especially when learning from limited numbers of preferences (Christiano et al., 2017; Ibarz et al., 2018; Javed et al., 2021).
While the effects of reward misspecification have recently been studied in the context of reinforcement learning agents that optimize a proxy reward function (Pan et al., 2022), and the effects of causal confusion have been emphasized in behavioral cloning approaches that directly mimic an expert (De Haan et al., 2019;
Zhang et al., 2020; Swamy et al., 2022), there has not been a systematic study of causal confusion when learning reward functions.
As an example of the type of causal confusion we study in this paper, consider the assistive feeding task in Fig.
2b.
Note that all successful robot executions will move the spoon toward the mouth in the area in front of the patient’s face—few, if any, executions demonstrate behavior behind the patient’s head (the exact likelihood depending on how trajectories for preference queries are generated).
In practice, we find that a preference-based reward learning approach will often pick up on the signed difference between the spoon and the mouth rather than the absolute value of the distance.
The two correlate on the training data, but using the former instead of the latter leads to the robot thinking that moving behind the patient’s head is even better than feeding the person!
In this work, we contribute the first systematic study of causal confusion as it occurs in preferencebased reward learning.
First, we demonstrate the failure of preference-based reward learning to produce causal rewards that lead to desirable behavior on three benchmark robotic tasks, even when given large amounts of preference data—in these settings, the learned reward has high test accuracy but leads to poor policies when optimized.
We then study the effect on causal confusion of several factors related to reward model specification and training: the presence of non-causal distractor features, reward model capacity, noise in the stated preferences, and partial state observability.
For each of these, we perform analysis of what errors the learned reward has, how it compares to the ground truth reward, and the amount of distribution shift it induces during policy optimization.
Some effects are predictable (e.g. removing distractors helps), some are perhaps surprising (e.g. increasing the amount of data when preferences are noisy can hurt when the data has diversity issues to begin with), and they are not always correlated to changes in the validation accuracy.
Another of our contributions is to point to the importance of data coverage and perhaps interactive techniques that iteratively or actively query for feedback to help learners disambiguate causal features from correlates.
But they also caution that there are many other aspects that can go wrong or make the problem much more challenging, from the very way we choose to define state, i.e. what to include and what not to include in the input to the learned model, to how we manage the relationship between noisy data quantity and diversity.
1. R ELATED W ORK
Reward Hacking and gaming behaviors are known to commonly occur in reinforcement learning (Ng et al., 1999; Krakovna et al., 2020) and reward learning (Christiano et al., 2017; Ibarz et al., 2018; He & Dragan, 2021).
However, these behaviors are often only mentioned anecdotally (Krakovna et al., 2020).
Recently, Pan et al.
(2022) proposed to systematically analyze reward misspecification in RL by creating a set of domains where the agent optimizes a proxy reward function and studying when this leads to incorrect behavior.
By contrast, we study cases where the reward function must be learned from human preferences.
Preference Learning is a popular method for training AI systems when a reward function is unavailable (Wirth et al., 2017; Christiano et al., 2017; Sadigh et al., 2017; Ibarz et al., 2018; Biyik &
Sadigh, 2018; Brown et al., 2019; 2020a).
Preferences are often easier to provide than demonstrations or raw numerical values (Sidney, 1957; Kendall, 1948; Wirth et al., 2017) since they do not require expert proficiency or fine-grained feedback.
However, optimizing a reward function that has been trained using preferences can sometimes lead to unintended behaviors.
Annecdotal evidence of this has been documented for Atari games (Christiano et al., 2017; Ibarz et al., 2018; Brown et al., 2020a), simple pointmass navigation tasks (Javed et al., 2021), and language model fine-tuning (Stiennon et al., 2020); however, there has been no systematic study of causal confusion when learning reward functions.
In this paper, we provide the first concrete and systematic study of causal confusion across three robot control benchmarks when learning and optimizing reward functions from preference labels over trajectories.
Causal Confusion in Imitation Learning has previously been studied in the context of behavioral cloning (Pomerleau, 1988; Torabi et al., 2018), where a policy is trained to directly map from states to actions using demonstrated state-action pairs.
Prior work has shown that behavioral cloning approaches can suffer causal confusion due to “causal misidentification,” where giving imitation learning policies more information leads to worse performance (De Haan et al., 2019) due to temporally correlated noise in expert actions (Swamy et al., 2022).
Similarly, we find strong evidence of causal misidentification when expert noise is present.
Zhang et al.
(2020) use causal diagrams to investigates causal confusion for simplified imitation learning tasks with discrete actions and small numbers of states when the features available to the demonstrator are different from those of the imitator.
We also study what happens when the observation space of the imitator changes; however, we focus on imitation via preference learning, consider tasks with continuous states and actions, and evaluate causal confusion when using non-linear reward function approximation.
Goal Misgeneralization has been studied by Langosco et al.
(2022) and Shah et al.
(2022)—here, the RL setting has a known goal and “features of the environment are correlated and predictive of the reward on the training distribution but not OOD.”
This setting is similar to ours in that there is misidentification of the state features that are causal with respect to the reward (e.g. going to the end of a maze rather than collecting the coin that is the true reward).
However, in contrast to our work, during training the reward signal is always present and always the ground-truth reward.
We show that when learning a reward function the learned reward can be misidentified, leading to poor
RL performance.
By contrast, Langosco et al.
(2022) and Shah et al.
(2022) show that even if the learned reward is perfect, the RL optimization can still fail due to spurious correlations during policy optimization.
1. B ACKGROUND : R EWARD L EARNING FROM P REFERENCES
We model the environment as a finite horizon MDP (Puterman, 2014), with state space S , action space A , horizon T , and reward function r : S ×
A → R .
Note that we assume that the reward function is unobserved and must be learned.
The reward function is learned from preferences over trajectories where, using the popular Bradley-Terry model (Bradley & Terry, 1952), the probability a trajectory τ B is preferred over another trajectory τ A is given by
[Equation]
where r ( τ ) =
(cid:80) ( s,a ) ∈
τ r ( s, a ) and where a trajectory τ is defined to be a sequence of state-action pairs: τ = ( s 0 , a 0 , . . .
, s T , a T ) .
To learn a reward function from preferences, we assume access to a set of pair-wise preference labels
P over trajectories τ 1 , . . .
, τ N , where ( i, j ) ∈ P implies that τ i ≺ τ j .
We then optimize a reward function r θ : S ×
A → R parameterized by θ that maximizes the likelihood:
[Equation]
We describe how we implement this in practice in Alg. 1.
C AUSAL R EWARD M ISIDENTIFICATION
1 displays the causal structure of preference-based reward learning.
In preference-based reward learning, pairwise preferences (in the form of a binary label) are given based on an underlying reward function R ( x ) —usually that of the human user or teacher (Fig. 1a).
Within the state space, there are features x it that are causal and influence R ( x ) , as well as other features z it that are nuisance variables and have no bearing on R ( x ) .
Note that z it may very well exhibit correlations with x it , potentially due to their sharing of the same causal parent or there being biases during data collection.
Given the pairs of trajectories (specified by state information ( x, z ) ) and binary preference labels, the goal of preference-based reward learning is to learn a reward function ˆ R ( x, z ) that best matches the stated preferences.
In Figure 1b, there are state features u which are causal with respect to the true reward, but are unobserved by the learning agent.
In both cases, unobserved human bias and noise, denote by η , also affects the preferences labels.
The learned reward ˆ R may be able to achieve low and even near-perfect performance on a heldout test set by making use of 1) nuisance variables that correlate with causal variables or 2) faulty/incomplete correlations between causal variables and R that happen to hold true for the training data distribution.
Optimizing a misidentified reward function leads to distribution shift when performing reinforcement learning on the learned reward values ˆ R ([ x, z ]) , resulting in behaviors with low performance under the true reward function R .
We define this behavior of learning a reward that achieves low test error but results in poor performance (under the true reward function) when the learned reward function is optimized via RL as causal reward misidentification .
E XPERIMENTAL S ETUP
We now introduce the environments and analysis methods that we use in our paper.
Environments for Preference Learning.
We identify a set of robot learning benchmarks that exhibit causal reward misidentification when learning reward functions from preferences.
In
Reacher (Brockman et al., 2016)
(Fig. 2a), the goal is to move an end effector to a desired goal
[Figure]
location.
In Feeding (Erickson et al., 2020) (Fig. 2b), the goal is to feed the human using a spoon carrying pieces of food.
Finally, in Itch Scratching (Erickson et al., 2020) (Fig. 2c), the goal is to repeatedly scratch a desired itch location on the human’s arm.
True Rewards and Preference Data Generation.
Each domain has a predefined “true” reward function r (outlined in Appx. A.5).
This enables us to create synthetic demonstrations and preference labels by adding increasing amounts of noise to an expert policy trained on r (details in Appx.
A.2).
As shown by Brown et al. (2020b), adding this type of disturbance will result in monotonically decreasing performance in expectation and produce a diverse dataset for preference learning.
Gleave et al.
(2020) also uses a similar way of switching between an expert policy and a random policy, which they find produces the best coverage distribution (over a purely random or purely expert policy).
We do compare this method of diverse trajectory generation against that proposed by Brown et al.
and find no significant difference in performance; see Appx.
A.11 for details.
Note that while we use the ground-truth reward function for obtaining preference labels, we assume no access to this reward function during policy learning, but rather seek to learn a policy that obeys a users preferences by first learning a model r θ of the true reward function from preference labels
P , and then run RL on the learned reward function to produce a policy.
We can then evaluate the learned policy on the true reward function r .
To facilitate reproducibility and encourage future research on causal reward learning, we open-source our code and training datasets.
This combination of domains and training data forms the first set of benchmarks for studying causal confusion when learning reward functions.
Evaluating Learned Reward Functions.
In order to establish that a learned reward is causally confused, we first check for low test error and establish that the learned reward performs well on unseen in-distribution test data (thereby ruling out model selection and training failures like not having enough data, not using high enough capacity, not regularizing enough, etc.).
We then show that the learned reward fails in two ways: 1) it leads to a policy that, when optimized, has poor performance with respect to the true reward – we call this policy ”PREF”, and 2) this poor performance is not due to RL failures, i.e. the learned reward actually does prefer its (poorly-performing) optimized policy, ”PREF”, over the optimal policy with respect to the true reward, which we call ”GT”: the learned reward prefers PREF over GT .
Finally, we analyze the learned reward qualitatively via gradient saliency maps and quantitatively via the Equivalent-Policy Invariant Comparison pseudometric (Gleave et al., 2020), then elucidate the effects of the reward error by quantifying the distribution shift induced by policy optimization of the learned reward.
Saliency maps are one of the few methods that allow one to interpret learned reward functions in an isolated, relatively lightweight manner.
Following Michaud et al.
(2020), we use raw gradient saliencies, or ∂R∂ ( s,a ) —the gradient with respect to each element of the input.
We extend upon this by examining gradient saliencies per timestep along with feature spread maps —maps of each input feature’s variation (standard deviation, variance, range) over the course of a trajectory.
[Table]
EPIC , or Equivalent-Policy Invariant Comparison, is a pseudometric proposed by Gleave et al.
(2020) that quantifies the difference between two reward functions on a given coverage distribution and proposes to be predictive of policy performance without the need for policy optimization.
We apply the EPIC pseudometric to compare the distance between various learned rewards and the ground truth reward on the coverage distributions seen during reward learning and those seen during policy training.
Kullback-Leibler divergence : We use a discriminator trained to minimize the cross-entropy loss on states from two different distributions following the approach proposed by Husz´ar (2017) and
Laidlaw & Dragan (2021).
Specifically, we approximate the KL divergence between the distribution of state-action pairs seen during reward learning and those seen during RL on the learned reward.
We use this to measure the amount of distribution shift from the reward learning distribution induced by optimizing (potentially causally confused) learned rewards during policy training.
Details on how this is done are in Appx.
A.6.
E VIDENCE OF C AUSAL C ONFUSION
Before varying different factors that affect the performance of the learned reward, we start with a generous setting where we provide large amounts of data to the learner in each environment, and we make sure to add features to the default input space for each environment such that all necessary information needed to infer the ground truth reward, TRUE, is available.
We produce the preference training data as detailed in the previous section: again, this ensures that there are comparisons that involve optimal behavior relative to slightly suboptimal one, as well as highly suboptimal, in principle ensuring good variety.
For hyperparameter tuning and other training details, see Appx.
A.1.
Table 1 details the results.
First, note that, especially with the largest data setting (L), the learned reward achieves high test accuracy, comparable to the training accuracy (columns 2-4 in the table).
This is a sanity check that the model learned does not overfit, and that the model capacity is high enough and data plentiful enough to get decent performance.
Note that in later sections, we observe that even models that reach over 99% test accuracy sometimes end up failing to produce good polices.
Second, when optimizing the learned reward to obtain the PREF policy, the optimization works in the sense that the learned reward does prefer PREF to the GT policy (obtained by optimizing the true reward).
This shows that the poor performance is not due to an issue in RL exploration.
We see this in column 5 of the table.
Unfortunately, when looking at the actual performance of the PREF policy in columns 6 and 7, it is disastrous: not only does it have poor TRUE reward (compared to the GT policy), but it leads to poor success rates.
Overall, the learned reward incentivizes poor behavior, despite reasonable test data accuracy.
[Figure]
In the Reacher environment, we observe that the policy optimized using the learned reward chooses to simply spin very fast rather than reaching for the target.
We note, however, that the learned reward correctly classifies the leftmost image in Fig.
2a (where the agent just folds its arm) as being worse than the middle image (where the agent successfully reaches the target), and does so with 96 . 6% accuracy on such diverse pairs of comparisons.
One would think that this would apply to the behavior in the rightmost image, where the agent folds its arm as well (and subsequently spins), but the learned reward turns out to strongly prefer the rightmost case.
As will be explored in Sec. 5.1, this is due to the slight positive gradient for the angular velocity feature in the learned reward (Fig. 4a), which stems from a slight correlation between angular velocity and reward.
Fig.
8a shows this correlation; there is a slight bias in the training data toward states with high reward and high velocity.
Similarly, with the
Feeding environment, the learned reward learns to minimize the signed difference between the spoon and the mouth rather than the absolute value of the difference because the majority of trajectories approach the mouth from in front.
Fig.
8b shows that there are far more states that have low reward at negative ‘diff y’ values than states that have low reward at positive ‘diff y’ values.
As a result, the learned reward (Fig. 9b) correctly identifies spilling food (left) as being worse than feeding (middle), but goes further and incentivizes bringing the spoon ‘into’ and even behind the head in an attempt to minimize the signed difference.
In Itch Scratching, the learned reward correctly identifies flailing (left) as being worse than actual scratching (middle).
However, as seen in Fig.
9c, the Itch Scratching agent (spuriously) learns a higher weight on two components of the action (corresponding to two robot arm joints), which results in the agent turning the last section of the arm in a circle while trying to keep the end effector close to the itch target—another type of flailing!
It also falls into the same trap as the Feeding agent in minimizing the signed rather than the absolute difference (as seen in
Fig. 8c, where there is a bias toward negative values of ‘diff x’ rather than an even distribution of states across both positive and negative values).
Fig.
2 summarizes this behavior.
Appx.
A.3 provides plots of the aforementioned features’ spurious correlations with the true reward and gradient saliency maps of the learned rewards.
Next, we analyze why these failures occur.
[Figure]
−−
R e w a r d
Withdistractors
Withoutdistractors
[Figure]
D ATASET
D ISTRACTORS S IZE
W ITH
W ITHOUT
M 53.
6.762 L 46.
[Figure]
F ACTORS THAT L EAD TO C AUSAL C ONFUSION
In the following sections, we examine various factors of the preference-based reward learning problem setup and their effects on causal reward misidentification.
The motivation for exploring each is as follows:
• Distractor Features draws directly from the findings in Sec. 4, where we find that certain non-causal features may be spuriously correlated with the reward.
• Model Capacity is inspired by Pan et al.
(2022)’s findings on the effects of increased agent capabilities on reward hacking.
• Noise in Stated Preferences is inspired by the fact that preference data collected from real humans is often rife with different types of bias and noise.
• Partial Observability of Causal Features is also due to the fact that it may not always be possible to fully-observe all causal features in the real world.
• Complex Causal Features is a topic that may become increasingly relevant in higherdimensional, complex robotic tasks—unlike toy environments, where certain non-causal features can be ignored entirely, features here may be causal in certain contexts and noncausal in others.
As we will see, these factors (which may often be overlooked) can contribute significantly to causal reward misidentification and the eventual success or failure of the reward model.
We find that causal reward misidentification induces a distribution shift that can make the learned reward appear deceptively close to the true reward during reward learning but terribly misspecified during policy optimization.
5. D ISTRACTOR F EATURES
One of the main reasons for learning a wrong (causally confused) reward is the presence of too many features in the input that spuriously correlate with the causal features.
To study this effect, we incrementally remove such ‘distractor’ features and test the impact this has on the learned reward.
For example, for Reacher, we know that joint angles and angular velocities are not causal to the ground truth reward, which is based purely on distance between end effector and target and the norm of the action (a control penalty)
(see Appx.
A.4 for a complete list of causal and distractor features).
As we see in Fig.
3, removing all distractors drastically improves performance, and, indeed, the more such features are left in the input, the worse the performance gets in Reacher, almost linearly.
Granted, removing distractors does also slightly improve the validation and test performance in this case—indicating that there is some signal in the training data to help discern the spuriousness of certain features.
Nonetheless, over the several experiments in this paper, we find that the validation
[Figure]
Fig.
2a depicts the policy of the learned reward with distractors—the Reacher robot learns to simply fold its arm and spin.
Looking at the gradient saliency maps (Fig. 4) illuminates why this is the case.
Firstly, in Fig.
4a, we observe that the learned reward is misaligned with respect to ‘diff y‘, the feature corresponding to the difference between end effector position and target position along the yaxis; specifically, we note that the learned reward actually rewards an increase in ‘diff y‘, rather than a decrease.
Next, we notice that the reward doesn’t penalize action evenly.
It penalizes ‘torque 1‘ (the one responsible for the spinning), but rewards ‘torque 2‘ across nearly all timesteps.
It should instead be largely a negative penalty, as seen in Fig. 11a.
Lastly and perhaps most importantly, we observe in Fig.
4b that ‘vel 1‘ and ‘vel 2‘ have very large feature ranges, corresponding to the large variations in angular velocity achieved by the Reacher robot’s spinning behavior.
Looking back to the gradient for each of these velocity features in Fig.
4a, we observe that the reward has a slight positive gradient with respect to each.
By spinning fast, the Reacher robot is thus able to achieve much higher (learned) reward than performing the proper reaching action.
The KL divergence (Figure 3c) between the distribution of observation-action pairs seen during reward learning and those seen during policy optimization provides further insights—in incentivizing the Reacher robot to spin fast, it leads the RL optimization toward states that were not seen during reward learning (specifically, states where the robot is spinning very fast).
Appx.
A.7 displays results for the other tasks.
Similar to the Reacher task, we find that removing distractor features generally helps performance.
We note here that the reason removing non-causal distractor features does not appear to benefit for Feeding is that the main spurious correlation (discussed in Section 4) involves one of the causal features—namely, the difference between spoon position and target position.
Thus, removing purely non-causal features fails to address this issue in
Feeding.
5. M ODEL C APACITY
In Appx.
A.8, we study the effect of model capacity on the learned reward.
We find that, despite careful tuning of hyperparameters with each model and dataset size, increasing the capacity of the reward model does not necessarily result in an increase in subsequent policy performance.
5. N OISE IN S TATED P REFERENCES
Because user noise is often inevitable, we explore the effect of various types of user noise on causal reward misidentification.
Following Lee et al.
, we modify Eq.
1 to include a user rationality constant β and a myopic discount factor γ ∈ (0 , 1] :
[Equation]
T
a s k s u cc e ss
Oracle
Stochastic
Myopic
Skip
Mistakes 0.978 0.
— test acc.
(a) Feeding, User Noise
M
L 0.
T
a s k s u cc e ss
Oracle
Stochastic (b) Reacher, Data Size
With
Without 0.
T
a s k s u cc e ss
Oracle
Stochastic (c) Feeding, Distractors
[Table]
Using Eq.
3, we explore four types of user noise (varied independently of each other): STOCHASTIC, where the user is rational with β = 1 ; MYOPIC, where earlier rewards are discounted with a γ = 0 . 99 ; SKIP, where the teacher/preference-giver skips a pair of trajectories if both of have rewards below a certain threshold; and MISTAKES, where the preference label is randomly flipped with probability ϵ = 0 .
ORACLE has β = ∞ .
Results are displayed in Fig.
We find that noise in the stated preferences exacerbates causal reward misidentification—test accuracy stays high while policy performance plunges (Fig. 5a).
Notably, two instances of user noise,
Stochastic and Skip, have test accuracies greater than or equal to the Oracle (despite far worse performance).
In addition, increasing the amount of preference data when noise is present actually has a negative effect, as seen in Fig.
5b.
Further, there is a compounding effect with distractor features, where both noise and distractors result in a large loss in performance (Fig. 5c).
Importantly, the poor performance of the learned policies is not well predicted by the validation accuracies, with models achieving 0.995 resulting in poorer reward and success rate than models trained on less data with lower validation accuracy of 0.985.
Additional results are located in Appx.
A.9.
To examine these results further, we compare the learned rewards directly with the ground truth reward (without performing policy optimization) with the EPIC pseudometric in Table 2.
Interestingly, the EPIC distances between the learned rewards and the ground truth reward vary depending on the choice of coverage distribution.
With the distribution of observation-action pairs seen during reward learning (generated by randomly switching between an expert and random policy), the reward trained with labeling errors appears to be closer (in EPIC distance) to the ground truth than the reward trained without labeling errors.
However, these results are flipped when we instead evaluate the pseudometric on the distribution of observation-action pairs seen during policy optimization.
This suggests that the misidentified reward model more closely ‘mimics’ the ground truth reward on the reward learning distribution (hence the lower EPIC distance) but fails to generalize to the reinforcement learning distribution.
Similarly, we again observe (see Tab. 9 in the Appendix) that the KL divergence between reward learning and reinforcement learning state distributions is greater when the learned reward is misspecified.
This indicates that the learned rewards are misidentified: they have low test error but incentivize RL to deviate from the optimal behavior encountered during reward learning.
5. P ARTIAL O BSERVABILITY OF C AUSAL F EATURES
Partial observability over aspects of the state on which the user’s preferences are based is nearly inevitable in the real world.
It is important to note that RL with the ground truth reward is able to learn proper feeding behavior in spite of this lack of critical state information .
However, as displayed in Fig. 6, reward learning is only successful when the model is able to observe all the features that influence the true reward.
Examining the gradient saliency plots (Fig. 15) in the Appendix, we see that without having access to all the features that play into the true reward, the reward model is much more likely to become causally confused—it may over-weight the causal features that are available or pick up on correlations with non-causal variables that are totally spurious.
With Feeding, we observe both cases: the reward model learns stronger weights on one of the joint angles and two components of the action, all of which have no bearing on the true reward.
Simultaneously, the reward also learns a greater weight on the second component of the feature corresponding to the vector distance between the end effector and the mouth, which is causal in the ground truth reward.
(Note that a similar weighting also occurs in the learned reward that has access to all features of the true reward.)
This results in the behavior depicted in Fig.
6—the robot manipulator is able to successfully maneuver the spoon close to the patient’s mouth (by observing the distance feature), but does so without ensuring that the food particles themselves stay on the spoon and end up in the patient’s mouth.
EPIC distances are displayed in Table 11 in the Appendix.
Unlike with previous factors, EPIC distance on the reward learning distribution is predictive of eventual policy performance.
The KL divergences (Table 3) again reveal that the causally confused reward (this time, due to lack of information in the observation) leads to a greater distribution shift during RL training.
5. C OMPLEX C AUSAL F EATURES
The reward with partial observability on causal features induces a greater distribution shift than the reward with full observability
[Table]
To explore how results may differ as we increase the complexity of the task, we evaluate causal reward misidentification on a more complex task of Itch Scratching, which is challenging as it requires a robot to (1) maneuver around the human body in Euclidean 3-space E 3 to reach itch locations that are randomly generated along the 3D surface of the human arm, and (2) to perform a scratching motion, rather than staying still, at the itch target.
We find that even after increasing the amount of training data for reward learning (Fig. 7a), increasing the reward model capacity (Fig. 7b), removing distractor features that are not causally related to the ground truth reward, having perfectly stated preferences, and ensuring full observability over all the features that are involved in the ground truth reward, performing preference-based reward learning still fails to produce a policy that successfully scratches
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
the itch location on the patient.
However, as shown in Table 1, a standard model-free RL algorithm trained using the ground truth reward is able to solve such a task.
The ground truth reward (over which preferences are given) in this task rewards the robot for performing a ‘scratching’ motion.
This scratching motion entails not only making contact with the target using the end-effector (being within a certain radius of the target coordinates), but also requires that the target contact position be greater than a certain δ away from the previous target contact position and that the exerted force be no more than a specified F max .
Although we provide the reward model all the necessary information to infer this scratching motion (including information about the state at the previous timestep; see Appx.
A.4 for details), we find that the reward model is not able to learn the scratching motion.
As seen in Fig.
7c, once we explicitly include an indicator variable, ‘scratched’, denoting whether the robot has successfully performed the aforementioned scratching motion, performance drastically increases.
We suspect that the reward model’s tendency to pick up on spurious correlations that occur consistently over the course of the trajectory involving just a few variables prevents it from learning the true causal relation that involves many variables, each of which are causal only in a particular context.
This is supported by Fig 9c, which shows gradients staying relatively constant for each feature rather than varying with time and context.
As a result, the learned reward without ‘scratched‘ leads to a greater amount of distribution shift (Table 4).
Future work in reward learning would do well to explore how to learn these more complex, multi-feature causal relationships.
C ONCLUSION
Our work provides an analysis into the generalization of learned reward functions.
We identify three tasks where a standard way to generate diverse preference data often results in poor learned reward models.
These models have good validation and test accuracy—in some of our settings even 99.5%, seeming to distinguish the basics of task success vs. failure.
However, optimizing them pushes the policy outside of the training distribution, where the model falsely believes the reward is highest— we demonstrate this via saliency maps, the EPIC distance, and KL divergence and also qualitatively by examining the resulting behavior (e.g., going behind a patient’s head when feeding them).
This out-of-distribution effect results in policies that the learned model assesses to have high reward, whereas in reality they have poor true reward and low success rates.
Digging into the challenges of reward learning, we find that it is easy for the models to pick up on non-causal features, as some issues go away when we eliminate these non-causal features from the input.
Furthermore, noisy preference data aggravates poor generalization, which is not fixed by simply getting more noisy data from the same distribution—neither is test error predictive of performance, which points to causal confusion as the culprit.
And when not all causal features are observable, the learned model will struggle even though there exists a high-performing reward defined only on what is observable.
Based on our results, we have identified several promising directions for aleviating causal reward misidentification.
First, our results show that causally confused rewards result in greater distribution shifts during policy optimization.
As such, future work should investigate methods for penalizing excess exploration beyond the training data distribution.
In Appx.
A.12, we examine the effect of adding a direct penalty on the KL divergence into the cost function during RL.
We find that, unfortunately, the RL agent is again able to hack the reward—this time, the performance and causal reward misidentification are even worse because of the additional degree of freedom in the reward function afforded by the discriminator penalty.
However, we hypothesize that there may be other less obvious ways to successfully incorporate such a penalty.
Next, our results demonstrate that spurious (“nuisance”) features can significantly increase the chance of causal reward misidentification.
Thus, in cases where we can query for human knowledge on exactly which features are spurious or causal, we can use this feedback to learn which features are non-causal or remove non-causal features entirely.
Further, our results indicate that high-level features such as the ‘scratch’ feature are difficult for neural networks to learn, even when all the necessary low-level information is available.
Future work should examine incorporating methods for learning high-level features.
One promising approach would be to explicitly learn the causal feature(s) by querying for feature traces (Bobu et al., 2021).
Finally, we recognize that active and iterative methods for data acquisition is also a widely popular solution—as such, we provide a preliminary exploration in Appx.
A.13.
Using other metrics to compare learned rewards (such as DARD, by Wulfe et al.
(2022)) may also prove fruitful.
Our work cautions that reward learning is brittle—natural choices for acquiring data, deciding on the amount of data, or defining the input space can lead to models that seem very close to the true reward, but lead to spectacular failures when optimized.
And even though active and iterative methods for data acquisition may alleviate these issues, the problem of learning the best reward model when some of the causal features are unobservable remains wide open.
A A PPENDIX A. P REFERENCE L EARNING T RAINING D ETAILS
To learn a reward function from preferences, we assume access to a set of pair-wise preference labels
P over trajectories τ 1 , . . .
, τ N , where ( i, j ) ∈ P implies that τ i ≺ τ j .
We then optimize a reward function r θ : S ×
A → R parameterized by θ that maximizes the likelihood:
[Equation]
This likelihood function is differentiable, allowing us to leverage non-linear function approximation to learn the reward function from trajectory preferences.
In practice, we use the Adam optimizer in
PyTorch to learn the reward function, r θ , and then use PPO Schulman et al.
(2017) or SAC Haarnoja et al.
(2018) for policy optimization given r θ .
We perform preference learning on three dataset sizes (given in terms of unique pairwise comparisons): SMALL (780), MEDIUM (7140), and LARGE (52326).
Our test set is composed ofunique pairwise preferences drawn from a disjoint set of 60 trajectories.
See Appx.
A.2 for dataset generation details.
Hyperparameters—learning rate and weight decay—are tuned coarsely using the MEDIUM dataset size due to runtime limits and cost of computation.
The tuned hyperparameters (best performance on a held-out validation set) for each environment are as follows:
Reacher: weightdecay=0.0001, lr=0.01,
Feeding: weightdecay=0.00001, lr=0.001, Itch Scratching: weightdecay=0.001, lr=0.001.
Using the dataset of preferences (SMALL, MEDIUM, or LARGE) obtained offline, we train a neural network reward function approximator with two hidden layers (128 units and 64 units, respectively) and Leaky ReLU activations, after which we perform 1,000,000 timesteps of reinforcement learning with PPO Schulman et al.
(2017) (for Feeding and Itch Scratching) and
SAC Haarnoja et al.
(2018) (for Reacher) using the learned reward function in place of the ground-truth reward function.
We optimize the reward function approximator using Adam with weight decay and early-stopping on the validation loss (with a patience of 10 epochs).
The full preference-based reward learning procedure we use is detailed in Alg. 1.
The hyperparameters for the PPO and SAC agents are as follows (if not specified, then hyperparameters are set to RLLib’s default values):
PPO:
• training batch size =• number of SGD iterations =
• SGD minibatch size =• lambda = 0.
• dimensions of fcnet hidden layers =
SAC:
• timesteps per iteration =
• learning starts =
• Q model fcnet hiddens =
[100, 100] • policy model fcnet hiddens =
[100, 100] • train batch size =• actor learning rate = 3e- • critic learning rate = 3e- • entropy learning rate = 3e-
Input: Training dataset of pairwise preferences d train ˆ R ← initialize network for epoch ∈ (0 , 100) do for traj i , traj j , label ∈
d train do ˆ r i ← ˆ R ( traj i )
ˆ r j ← ˆ R ( traj j ) loss ← CrossEntropyLoss (ˆ r i , ˆ r j , label ) loss.backward () end for loss val ← calculate cross-entropy loss on validation set if loss val doesn’t decrease for 10 epochs then break end if end for ˆ π ← Run RL (PPO or SAC) using ˆ R as reward for 1M iterations
Output: ˆ R, ˆ π
A.
S YNTHETIC P REFERENCE G ENERATION
To enhance scalability and reproducibility, we automatically generate a large amount of synthetic trajectory preferences.
This was done using an expert RL policy trained using the ground-truth reward function provided with each of environment.
We then generate a large number of diverse trajectories by adding ϵ -greedy noise during policy rollouts, where ϵ is the probability that the policy takes an action uniformly at random from its action space.
Thus, ϵ = 0 corresponds to the fully trained RL policy and ϵ = 1 corresponds to a uniformly random policy.
As noted by Brown et al.
(2020b), adding this type of disturbance will result in monotonically decreasing performance in expectation.
To generate pairwise preferences over trajectories, we select all pairs of trajectories from a set of 40, 120, and 324 total trajectories (for the SMALL, MEDIUM, and LARGE dataset sizes, respectively) generated with ϵ -greedy rollouts for ϵ ∈ { 0 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 } .
We use held-out sets of trajectories for validation and testing.
We then use the ground-truth reward functions provided by each environment to provide ground-truth preference labels.
A. E VIDENCE OF C AUSAL C ONFUSION
Table 5 displays additional results for the Lunar Lander environment (Brockman et al., 2016), where we again notice that despite high performance on a held-out test set, the resulting optimized policy is poor.
Fig.
8 provides plots of select features’ spurious correlations with the true reward.
Fig.shows gradient saliency maps of the learned rewards that elucidate how the learned reward may improperly ‘weight’ a feature.
A. I NPUT F EATURES FOR E NVIRONMENTS
Features that are distractors are labeled with a D .
Features that are causal are labeled with a C .
Reacher:
• D - Cos of angle of first and second arm (2) • D - Sin of angle of first and second arm (2) • D - Coordinates of target (2) • D - Angular velocity of first and second arm (2) • C - Position fingertip - position target (3) • C - Action — Torque applied at first and second hinge (2)
Feeding:
[Figure] [Figure]
[Table]
• D - Spoon pos real (3) • D - Spoon orient real (4) • C - spoon pos real - target pos real (3) • D - Robot joint angles (7) • D - Head pos real (3) • D - Head orient
real
(4) • C - Spoon force on human (1) • C - Action (7) • C - Foods in mouth (1) • C - Foods on floor (1) • C - Foods hit human (1) • C - Sum food mouth velocities (1)
• C - Prev spoon pos real (3) • C - Robot force on human (1)
Itch Scratching:
• C - Tool pos real (3) • D - Tool orient real (4) • C - Tool pos real - target pos real (3) • D - Target pos real (3) • D - Robot joint angles (7) • D - Shoulder pos real (3) • D - Elbow pos real (3) • D - Wrist pos real (3) • C - Tool force (1) • C - Action (7) •
C - Prev tool pos real (3) • C - Robot force on human (1) • C - Prev tool force (1)
A. G ROUND T RUTH R EWARD F UNCTIONS FOR E NVIRONMENTS
We outline the ground truth reward functions for each environment below.
We also refer the reader to the publicly available code repositories for each environment, which describe the rewards in more detail.
Reacher ( https://github.com/openai/gym ):
[Equation]
Feeding ( https
[Equation]
assistive-
[Equation]
[Equation]
action
[Equation]
=
[Equation]
( Foods in mouth, Foods on floor ) preferences score =
g ( Foods hit human, Sum food mouth velocities,
Spoon pos real, Prev spoon pos real, Robot force on human )
[Equation]
preferences score
Itch Scratching ( https://github.com/Healthcare-Robotics/assistive-gym ):
[Equation]
| 2 r action
[Equation]
scratch =
[Equation]
( Tool pos
[Equation]
, Target pos
[Equation] , Prev tool pos real, Tool force, Prev tool force ) preferences score = g ( Spoon pos real, Prev spoon pos real,
Robot force on human, Tool force, Target pos real )
[Equation]
+ preferences score A. E VALUATING L EARNED R EWARD F UNCTIONS
Saliency maps: We produce saliency maps of the learned reward as follows: we forward propagate a single rollout from the learned reward’s policy through the reward network.
Then, with the reward model’s weights fixed, we backpropagate the output from the forward pass through the network and into the input (the policy rollout) to obtain the gradient with respect to each feature in the observation-action pair at each timestep.
Kullback-Leibler divergence:
For each learned reward, we sample 50 trajectories from the reward’s training data and from the resulting policy (taking care to label each trajectory’s origin distribution).
From these 50x2 trajectories, we create the training and validation splits and then flatten each group of trajectories into a dataset of observation-action pairs.
We train a discriminator model (hidden dimensions of 128x128x128) to distinguish between observation-action pairs seen during reward learning and those during RL by minimizing the binary cross-entropy loss.
We tune hyperparameters (learning rate, weight decay) on the validation loss and accuracy.
With this trained discriminator model, we calculate D KL ( p || q ) by taking the discriminator’s negative mean return/logit of all reward learning observation-action pairs, where p is the reward learning distribution and q is the policy optimization distribution.
Similarly, we calculate D KL ( q || p ) by taking the mean return of all RL observation-action pairs.
Since the KL divergence is not symmetric, we report D KL ( p || q ) + D KL ( q || p ) , or the symmetric KL divergence.
For a proof on why a discriminator can be used to approximate the KL divergence, we refer the reader to Appendix B.2 of
Laidlaw & Dragan (2021).
A. D ISTRACTOR F EATURES
Fig.
10 shows the additional results for removing distractors.
Fig.
11 shows the saliency map for the Reacher learned reward without distractors.
Table 6 shows additional KL divergences for the
Feeding and Itch Scratching environments.
We note here that the reason removing non-causal distractor features appears to not benefit for Feeding is that the main spurious correlation (discussed in [Figure] M
L − [Figure] [Figure] Section 4) involves one of the causal features—namely, the difference between spoon position and target position.
Thus, removing purely non-causal features fails to address this issue in Feeding. [Table] [Figure]
T a s k s u
[Figure]
M e a n c u m u l a t i v e r
[Figure] w a r d [Figure] A. M ODEL C APACITY
We study the effect of model capacity on the learned reward.
We find that, despite careful tuning of hyperparameters with each model and dataset size, increasing the capacity of the reward model does not necessarily result in an increase in subsequent policy performance.
In fact, as seen in
Fig.
12, only the Feeding task trained with the LARGE dataset size (52326 preferences) benefits from steadily increasing reward model capacity.
Indeed, although increasing the capacity for the
Reacher task appears to initially increase performance for the large dataset case (Fig. 12a), performance drops back down to below the performance of the smallest model size when the model size is further increased.
Further, increasing model size seems to decrease the performance on the small datasets.
This is not surprising, as benefiting from larger capacity tends to require increasing the amount of data, and validation-set accuracy tends to agree with learned reward performance.
Interestingly though, the EPIC distance paints the picture that performance improves in distribution but suffers out-of-distribution.
To examine these results further, we compare the learned rewards directly with the ground truth reward (without performing policy optimization) with the EPIC pseudometric in Table 7.
Interestingly, the EPIC distances between the learned rewards and the ground truth reward vary widely depending on the choice of coverage distribution.
With the distribution of states seen during reward learning (generated by randomly switching between an expert and random policy), the larger (256x256x256) model appears closer to the ground truth reward than the smaller (128x64) one.
However, these results are flipped when we instead evaluate the pseudometric on the distribution of states seen during policy optimization.
This suggests that increasing the capacity of the reward model allows it to more closely ‘mimic’ the ground truth reward on the reward learning distribution (hence the lower EPIC
[Figure]
distance)
but worsens its ability to generalize to the reinforcement learning distribution.
We see a similar effect with the divergence metric in Tab. 8. .
A. N OISE IN S TATED P REFERENCES
Fig.
14 shows additional results.
Table 9 shows the KL divergences for each model configuration in
Fig.
Table 10 displays complete results of EPIC distances
[Table]
[Table]
M
L − −−
R e w a
[Figure]
With
Without
[Figure]
[Table] [Table] [Figure]
[Table]
A. P ARTIAL O BSERVABILITY OF C AUSAL F EATURES
Fig.
15 displays gradient maps of a reward that partially observes the causal features and a reward that fully observes the causal features.
Table 11 displays the EPIC distances
. A. D ATA G ENERATION M ETHODS
Fig.
16 displays a comparison of our method of generating a diverse dataset of trajectory preferences and the method of using a checkpointed RL policye (proposed by (Brown et al., 2019)).
A. P ENALIZING THE KL D IVERGENCE
Inspired by the results showing that causally confused rewards result in greater distribution shifts during policy optimization, we attempt to address this by adding a KL divergence penalty term to the reward.
Specifically, in addition to using the reward learned from preferences, we incorporate a pretrained discriminator (trained to discriminate between state-action pairs from the reward learning distribution and the RL distribution) to estimate the KL divergence, as detailed earlier in Sectionand Appx.
A.6.
This KL divergence is scaled by a hyperparameter λ and subtracted from the reward learned from preferences.
We show results for λ = 0 , 0 . 01 , 0 .
1 , 1 , 10 in Fig.
Using a discriminator that is pretrained offline (the discriminator isn’t updated during the RL process) to penalize the KL divergence does not appear to help the performance or alleviate causal reward misidentification in any way.
It appears that the RL agent is again able to hack the reward—this time, the performance and causal reward misidentification are even worse because of the additional degree of freedom in the reward function afforded by the discriminator penalty.
[Figure]
M e a n c u m u l a t i v e r e w a r d T-REX (synthetic checkpoints) RL+noise
[Figure]
d
[Figure]
[Figure] [Figure] u cc e ss
[Figure]
[Figure]
A. I TERATIVE P REFERENCE -B ASED R EWARD L EARNING FROM O NLINE D ATA
Following Christiano et al.
(2017), we verify the occurrence of causal reward misidentification when learning from data that is actively queried online (during the learning process).
As before, we use the initial set of offline trajectory and preference data (acquired via Appx.
A.2) to train a reward function, which we then use to optimize a policy using RL.
However, we then sample 10 trajectories from the most recently optimized policy and generate preferences for all possible pairs involving the sampled trajectories, which we concatenate to our existing training dataset.
Using the (now augmented) dataset, we fine-tune the reward function and policy for 10 epochs and 100 iterations, respectively.
We repeat this process of taking rollouts from the most recent policy and fine-tuning the existing learned reward and policy for 50 iterations.
Results are displayed in Fig.
We observe that causal reward misidentification appears to still be present—test error is very low, while resulting policy performance (measured by true reward and success) remains poor.