Hindsight PRIORs for Reward Learning from Human Preferences
Authors: Mudit Verma
Katherine Metcalf
Venue: ICLR 2024 poster (2024)
Keywords: preference based reinforcement learning
world models
return redistribution
Primary area: reinforcement learning
TL;DR: Presents a method to address credit assignment problem in preference-based reinforcement learning by guiding rewards to key states according to relative state importance.

Abstract: Preference based Reinforcement Learning (PbRL) removes the need to hand specify a reward function by learning one from preference feedback over policy behaviors. Current approaches to PbRL do not address the credit assignment problem inherent in determining which parts of a behavior most contributed to a preference resulting in data intensive approaches and subpar reward models. We address such limitations by introducing a credit assignment strategy (PRIOR) that uses a forward dynamics world model to approximate state importance within a trajectory and then guides rewards to be proportional to state importance through an auxiliary predicted return redistribution objective. Incorporating state importance into reward learning improves the speed of policy learning, overall policy performance, and reward recovery on both locomotion and manipulation tasks. For example, PRIOR achieves 80% success rate with half the amount of data compared to baselines. The performance gains and our ablations demonstrate the benefits even a simple credit assignment strategy can have on reward learning and that state importance in forward dynamics prediction is a strong proxy for a state's contribution to a preference decision.

Published as a conference paper at ICLR 2024
HINDSIGHT PRIORS FOR REWARD LEARNING FROM HUMAN PREFERENCES
Mudit Verma*
Katherine Metcalf
Arizona State University
Apple Inc. Cupertino, CA, 95014 kmetcal la @@ apple.com
Tempe, AZ. 85281 mudit tve rma @ adu.edu
ABSTRACT
Preference based Reinforcement Learning (PbRL) removes the need to hand spec- ify a reward function by learning a reward from preference feedback over policy behaviors. Current approaches to PbRL. do not address the credit assignment prob- lem inherent in determining which parts of behavior most contributed to a prefer- ence, which result in data intensive approaches and subpar reward functions. We address such limitations by introducing a credit assignment strategy (Hindsight PRIOR) that uses world model to approximate state importance within a trajec- tory and then guides rewards to be proportional to state importance through an auxiliary predicted return redistribution objective. Incorporating state importance into reward learning improves the speed of policy learning, overall policy per- formance, and reward recovery on both locomotion and manipulation tasks. For example, Hindsight PRIOR recovers on average significantly (p < 0.05) more reward on MetaW (20%) and DMC (15%). The performance gains and our ablations demonstrate the benefits even simple credit assignment strategy can have on reward learning and that state importance in forward dynamics prediction is strong proxy for state's contribution to preference decision.
INTRODUCTION
Preference-basee reinforcement learning (PbRL) learns a policy from preference feedback removing the need to hand specify reward function. Compared to other methods that avoid hand-specifying a reward function (e.g. imitation learning, advisable RI, and learning from demonstrations), PbRI. does not require domain expertise nor the ability generate examples of desired behavior. Ad- ditionally, PbRL can be deployed as human-in-the-loww sub-optimal policies, and has shown be highly effective for complex tasks where reward spec- ification is not feasible (e.g. LLM alignment) AArrour et al. (2011): Ibarz et al. (2018); Lee et al. al. (2023); Hejna III & (2di; Lee al. (2023) Korbak et al. (2023); Leike et al. (2018): 8): Ziegler et al. (2019); Ouyang et al. (2022): ; Zhu et al. (2023). However, exist- ing approaches to PbRL, require large amounts human feedback and are not guaranteed to learn well-aligned reward functions. A reward function is "well-aligned" when policy learned from it is optimal under the target reward function. We address the above limitations by incorporating knowl- edge about key states into the reward function objective.
allowing guidance to adapt on-the-fly to
Current approaches to learning a reward function from preference feedback do not impose credit assignment strategy over how the reward function is leamned. The reward function is learned such that preferred trajectories have a higher sum of rewards (returns) and consequentially are more likely to be preferred via a cross-entropy objective Christiano et al. (211)). Without imposing credit assignment strategy to determine the impact of each state on the preference feedback, there are many possible reward functions that assign higher return to the preferred trajectory. To select between possible reward functions large amounts of preference feedback are required. In the absence of enough preference labelled data, reward selection can become arbitrary, leading to misaligned reward functions. Therefore, we hypothesize that: (H1) guiding reward selection according to state importance will improve reward alignment and decrease the amount of preference feedback required
Work during internship Apple

---

Published as a conference paper at ICLR 2024
to learn well-aligned reward function and (H2) state importance can be approximated the states that in hindsight are predictive of a behavior's 's trajectory.
Hindoggt PROCR
Pradiced Radienaaiin
-1
FL
Altertion
Woldd Modeet
G
Ravard Prediciie
Eraupy
price
Ful1]
Cos
Figure 1: Hindsight PRIOR augments the existing PbRL cross-entropy loss by encouraging the magnitude of reward to be proportional to the state's importance. Each reward update preference labelled trajectories are passed
world model T (yellow) and estimated reward for (red), which assign an importance score and a reward (respectively) to each state-action pair. The return Go is then applied to the importance scores, which then as auxiliary targets for reward learning.
To this end, we introduce PRlor On Reward (PRIOR), a PbRL method that guides credit assignment according estimated state importance. State importance is approximated with an attention-based world model. The reward objective is augmented with state importance as an inductive bias to disambiguate between the possible rewards that explain a preference decision. contrast to previous work, our contribution mitigates the credit assignment problem, which decreases the amount of feedback needed while improving policy and reward quality. In particular, compared baselines, Hindsight PRIOR achieves 2 80% success rate with as little as half the amount of feedback on Meta World and recovers on average significantly (p < 0.05) more reward on MetaWorld (20%) and DMC (15%). Additionally, Hindsight PRIOR is more robust in the presence of incorrect feedback.
2 RELATED WORK
Wirth et al. 2017). train RL agents with human preferences on tasks for which reward design is non- trivialiial and can introduce inexplicable and unwanted behavior (Vamplew et al.) 2018 Krakovna et al. 2020). Christiano et al. (2017) extended PbRL Deep RL PEBBLE (Lee al. 2021a) incorporated unsupervised pre-training, reward relabelling, and offline RL to reduce sample complexity. Subsequent works extended PEBBLE by incorporating pseudolabelling into the reward learning process (Park et al. (022), guiding exploration with reward uncertainty (Liang al. 2022 and monitoring Q-function performance on the preference feedback (Liu et al. 2022).
PbRL
Kim et al. (203)) attempts to address the credit assignment problem by assuming that the preference feedback is based on a weighted sum of rewards and use a modified transformer architecture to assign rewards and weights to each state-action pair. However, introducing transformer-bate based reward function increases reward complexity compared to earlier work and Hindsight PRIOR as well tying the reward model to specific architecture. While Hindsight PRIOR also uses transformer architecture, it is independent of the reward architecture. Additionally, Kim et al. (2023) has not been extended to online RL.
Reinforcement Learning, especially model-based RL, leverage learned world models for tasks such planning (Allen & Koomen 1983 Hatner et al., 201 a), data aug- mentation (Gu et al. 2016 Ball et al. 2021). uncertainty estimation (Feinberg et al. [018 Kalweit Boedecker] 2017). and exploration (Ladosz et al. (022). In this work we learn a world model and use it to estimate the importance of state-action pairs While Hindsight PRIOR can use any transformer-base world-model, we use the current state of the art in terms of sample complexity, Transformer-basee World Models (TWM) (Robine et al. (023). To our knowledge, existing work has not incorporated a world model in reward learning from preferences.
Learning World Models:
Many methods exist to estimate the importance of different parts of an input to the model decision- decision-making process. Some popular methods include gradient saliency based approaches (Greydanus et al. 22088 Selvaraju et al. 2017; [Simonyan et al., 2013 Weitkamp al. 2019) and self-attention based methods Ras et al. (2022) Wiegreffe & Pinter (2019): Vashishth et al. (201). Self-attention based methods have been used for video summarization and extraction of key
Feature Importance:
2

---

Published as a conference paper at ICLR 2024
frames (Feng et al. 2020 Bilkhu al. 2019: Apostolidis al. 2021 Liu et al. 2019). Given use of TWM,, we use a self-attention map based method.
Credit Assignment: Credit assignment challenges typically stem from sparse rewards and large state spaces and solutions aim to boost policy learning (Ke et al., 2018 Goyal et al. 2018: Ferret al. 2020) Past works like Gokal al. (2018) have learned backward dynamics model to sample states that could have led to the current state and Ferret et al. (0200) equips a non- non-autoregressive sequence model to reconstruct a reward function and utilizes model attention for credit assignment. Return redistribution is another credit assignment solution that redistribute the ground-truth, non- stationary reward signal in order denisfy the emitted reward signals (Ren et al.] 2021] AArjna- Medina et al., [219 Patil et al. 2020). This is in contrast to PbRL where predicted rewards are dense to begin with. We adapt the idea of return redistribution for PhRI. by redistributing the predicted returns and is discussed in Section 4
3 PREFERENCE-BANE
REINFORCEMENT LEARNING
To learn policy with cefference-basee reinforcement learning (PbRL), the policy T executes an action a, at each time step in environment & based on its observation or of the environment state 8, For each action the environment & transitions to new state 8141 according to transition function T and emits a reward signal é; ( at). The policy Ty is trained to take actions that maximize the expected discounted return Go E. (8e, 0t). The reward (-) is trained to approximate the human's s target reward function (-).).
To learn (-) dataset D of preference triplets (To, T1, Bp) is collected from a teacher (human or synthetic) over the course of policy training. The preference label 1/p indicates which, if any, of the two trajectory segments
Tι with length / has higher (discounted) return G under the target reward function F.) Following Park et al. (2022) and Lee et al. (2021a), feedback is solicited every K steps of policy training for the M maximally informalive trajectories pairs (To, 7)) (e.g. pairs with the largest (-) uncertainty).
Given preference dataset D, (-) is learned such that preferred trajectories have higher predicted returns G.. than dispreferred trajectories. Using the Bradley-Tery Terry model (Bradley Terry 1952). predicted trajectory returns are used to compute the probability that one trajectory is preferred over the other Ps:
exp E, (s).
Po +-|
Σis exp E, f (s)
where To is preferred over Ti The probability estimate Pv is then used to compute and minimize the cross-entropy between the predicted and the true preference labels:
Lce (p.
yp (0) log Pe [To > Til + vo log Polni > roll-
(2)
The reward function Pe is learned over the course of policy गे training by iterating between updating To according the current estimate of Fú and updating Ag on D, which grown by M preference triplets sampled from T&'S 's experience replay buffer B for each fo update. To avoid training T on completely random F at the start of training, To explores the environment populate D with an initial set of trajectories following either a random policy or during intrinsically motivated pre- pre-training period Christiano et al. (2017): Lee et al. (2021a).
HINDSIGHT PRIORS
PbRL relies on learning high-quality reward function that generalizes and quickly adapts in a few- shot manner to unseen portions of the environment, and given its human in the loop nature, reducing the amount of preference feedback is vital. To learn the reward function (s,, o)), trajectory-Ievel feedback is provided and then is distributed each of the trajectory's states-action pairs. Given two trajectories, return per trajectory (d%. G.,), and preference label, many reward functions assign higher return for preferred trajectories, but do not align with the target reward function F (s,, a,) on unseen data. With a large enough dataset, a (-) that aligns with human preferences in all portions of the environment can be learned. However, given set of reward functions i. (.).
3

---

Published as a conference paper at ICLR 2024
each of which conforms to the preference dataset, reward function will be arbitrarily selected in the absence of additional information or constraints. From insufficient preference feedback, the se- lected ( (-) is likely to represent local minimum with respect to previously unseen trajectories, where the assigned returns Rgg are correct, but the distribution of rewards within trajectories are incorrect. Incorrectly assigning rewards at the state-action level, or incorrectly solving the credit as- signment problem, leads to reward functions that do not generalize outside of the preference dataset, resulting in suboptimal policies relative to the target reward function. Thus, address the credit assignment problem and guide reward distribution within a trajectory through an auxiliary objective that provides prior on state-action pair values computed after the trajectory has been observed (in hindsight).
The priors state-action values are identified by answering the following question, "now that I have seen what happened, which state-action pairs best summarize what happened in the given tra- jectory?" We consider the states that summarize trajectory to be those that are most predictive of future state-action pairs. The most predictive states are then used as proxy for the most important states. The use of summarizing state-action pairs is motivated by previous work demonstrating that people have selective attention when evaluating a behavior they attend only to the state-action pairs necessary to provide the evaluation (Desimone & Duncan 1995; Bundesen 1990, Ke et al.) 2018). We therefore assign greater credit to those states that were likely to have been attended to and therefore influenced the preference feedback. As summarizing states are those that are predic- tive of future state-action pairs, we identify them using an attention-based forward dynamics model, where state-action pair importance is proportional to their weight in the attention layers. For ex- ample, in Figure | the important states (highlighted in red) identified from action sequences in Montezuma'ss 's Review are those where the agent lines up to leap from the platform.
4.1 APPROXIMATING STATE IMPORTANCE WITH FORWARD DYNAMICS
An attention-based forward dynamics model (Figure yellow) is used to identify important (sum- marizing) states and address the PbRL credit assignment problem. The states that are key for a forward dynamics model to predict the future are assumed to be similar to those human evaluator would use to predict future states, and thus summarize trajectory. We use the attention layers in an attention-based forward dynamics model to approximate human attention and guide how feedback credit is distributed across a trajectory. In similar vein as Harutyunyan et al. (019)'s 's State Condi- tioned Hindsight Credit Assignment, we consider the importance of a state in a trajectory given that future state was reached.
World models have played a large role in model- reinforcement learning. Given the power that recent work have shown them to convey in reinforeement learning (Manchin et al. 2019 Hafner ctt al., 2019a Hu et al. 2019), we use world modelling techniques to learn an iitentiin-basee for- ward dynamics model. For a world model 7 to identify important states and approximate human attention, it must have two characteristics. First, must model environment dynamics and be able to predict the next future state ŠT given a history of state-action pairs TI:7-1]: [(TITT-I) $T. Sec- ond, it must expose mechanism to compute state-action importance 011:7-1] vector over given
trajectory segment -1:T-1] when performing the next-state prediction: [ $T) afi.T-1) 1]- Transformer based World Models (TWM) Robine et al. (2023) meets both requirements in addition
to being sample efficient (Robine et al. 2023; Micheli et al. 2023).
TWM a Transformer XL based auto-regressive dynamics model Í that predicts the reward Å.. discount factor 9t, and (latent) next state ě.41 given history of state-action pairs (W(T1=)) = sn). In the PbRL paradigm, predicting a transition's reward re is impractical as the reward function f is learned in conjunction with the world model. Therefore, and discount heads, and use the observation and latent state models:
adapt TWM by removing the reward
1. Observation Encoder and Decoder: 2 ~ (xl(o))));; ô; (ā.l()))
2. Aggregation and Latent State Predictor: h = fus (2[ua)) {[Ia)); 2441 ~ (
Consequentially, the loss function for the dynamics model is updated follows, where H is the cross entropy between the predicted and true latent next states:
T
Dyn.
(3)
4

---

Published as a conference paper at ICLR 2024
The Latent State Predictor is a transformer responsible for predicting the forward dynamics given the trajectory history, and is therefore responsible for approximating state-action importance. For description of the latent state predictor and its architecture, specifically the parts that allow to extract state importance, see Appendix C.2
The world model is learned over the course of policy " and reward F&s training. The observation encoder's s and decoder's s weights i are trained during Te 's exploration period to initially populate D. and then frozen for the remainder of T and fy training. The weights of the dynamies model w are trained during T 's exploration phase and then updated every steps of policy training from the same replay buffer B the preference queries are sampled from. Using B removes the need sample additional transitions or trajectories for the purpose of world model learning.
4.2 COMPUTING THE HINDSIGHT PRIORS
The use of transformer-based Latent State Predictor provides approximations of state importance in the form of attention weights (our second requirement in Section 4.1p When updating Ar the attention weights for each trajectory T in the collected preference triplets (To, Tp) D are com- puted by passing т to the Transformer XL model T (Figure | yellow). The transformer uses multi-headed,, multi-layer attention mechanism, where H is the number of attention heads, L the number of layers, and attn;
= (attn',, attn',,) ) E _ the attention weights of the 1-th layer for state-action pair (st. ae) E TT.T. The matrix A denotes the attention distribution in predicting the next state ēr+1 = [T) across all sequence timesteps and attention layers. The hindsight PRIOR (importance) at for given state-action pair (st, az) is estimated as the mean across layers L at timestep f, 04 1/2 Ele attn
4.3 REWARD REDISTRIBUTION AND CONSTRUCTING THE HINDSIGHT PRIOR Loss
To guide reward function learning according to state-action pair importance, the attention maps A from T are incorporated into the reward learning objective as redistribution guidance (Figure or- ange). The attention map does not form a reward target, as state-action importance for predicting future states does not equate absolute value in the target reward function, therefore return redistri- bution (Arjona-Medina et al. 2019), a strategy typically used to address the challenge of delayed returns reinforcement learning, used to align reward assignment with state-action importance.
Return redistribution addresses the challenge of delayed returns by redistributing a trajectory seg- ment's return among its constituent state-action pairs. The return redistribution case in existing work -Mrjona-Medina et al. 2019 Ren et al. 202 I Patil et al. 2020) relied on known and typically stationary, but sparse, rewards. In PbRL, while the learned reward function is dense, the feedback used to learn it occurs at the end of a trajectory and therefore is delayed and sparse. Therefore, to align rewards with estimated state importance, we introduce predicted return Gp redistribution to obtain state-action pair importance conditioned reward targets for a given trajectory т, where G -. (T.).
To obtain the reward targets for each trajectory 7 in a preference triplet (To, T1, Up), the predicted return G is computed (Figure I red), the attention map A(T) F(T) is extracted from the world model (Figure [ yellow), and the mean attention value per state-action pair is taken over layers + atttna,). Reward value targets are then estimated by redistributing the predicted return G... according a to obtain Itarget a Gu. where a is a vector with length |r| and G2 scalar (Figure orange). The state-action pair importance conditioned reward targets Frargez incorporated into reward learning via an auxiliary mean squared error loss between the predicted rewards fo = [f (s1, 01), fx .... (ST aT)] and Tzarget'
= "/_((attinn
Cprior = MSE(ft, Ftarget)-
The PbRL objective LeE (equation[]an is modified to be a linear combination of the proposed hind- sight PRIOR loss Eprior to guide reward learning with both preference feedback and estimated
5

---

Published as a conference paper at ICLR 2024
state-action importance:
1 [2] [D]
p (D)
Σ (Di + x Sprior (D;):
(5)
where À is constant to ensure LCE and Pyrior are on the same scale,
EMPIRICAL EVALUATION
We evaluate the benefits of Hindsight PRIOR on the Deep Mind Control (DMC) Suite locomotion (Tunyasuvunakool et al. 2020) and MetaWorld control (Yu et al. 2020) tasks, compare against baselines ee et al.) 202 Park et al. 2022 Liu et al. 2022 Liang et al. 2022). and ablate over Hindsight PRIOR's s contributions. Following our baselines, tasks with hand-coded rewards are used to assess algorithm performance. The hand-coded rewards serve the target reward functions (used by human in the loop) and are used to assign synthetic preference feedback (trajectories with the higher return preferred). Therefore, PbRL policy performance is measure and compared ac- cording to how well and how quickly the target reward function is maximized. Additionally, SAC (Haarnoja al. 2018) policy is trained on the target reward function provide a reasonable reference point for PbRL performance. Each PbRL method is compared to SAC using mean nor- malized return for DMC et al. (2021b) and mean normalized success rate for MetaWorld. See Appendix [ for the equations. For each comparison against baselines, mean (+standard deviation) policy learning curves and normalized returns are reported over random seeds (see Appendix F). From the learning curves and normalized scores, feedback sample efficiency. environment interac- tion sample efficiency, and reward recovery are compared between Hindsight PRIOR and baselines.
While using synthetic feedback allows to directly compare between the target Fy and learned As reward functions, humans do not always select the trajectory that maximizes the target reward func- tion. Occasionally, humans will mislabel trajectory pair and flip the preference ordering. There- fore, we evaluate Hindsight PRIOR and PEBBLE (the backbone algorithm for Hindsight PRIOR and the baselines) using synthetic feedback labeller that provides incorrect feedback (10% 20%, 40%) of the preference triplets (mistake labeller from (Lee et al.) (0216)).
percentage
To better understand Hindsight PRIOR's performance gains over baselines (Section 5.1 answer the following questions in Section 5.2
(QI) Is it the use of return redistribution strategy versus Hindsight PRIOR's ' specific strategy (guiding return redistribution according state importance) that leads to the performance improvements?
(Q2) Do the performance gains stem from incorporating environment dynamics?
(Q3) What types of states does TWM identify as important?
and to verify that the incorporation of the world model does not negatively impact PbRL capabilities, we answer the following in Section 5.3
(Q4) Does Hindsight PRIOR's ' scale to longer trajectories in the preference triplets?
(Q5) Does combining Hindsight PRIOR with a complementary baseline improve performance?
(Q6) Does Hindsight PRIOR allow for the removal of preference feedback?
Hindsight PRIOR and all baselines extend PEBBLE- as their underlying PbRL algorithm. The policy takes random actions for the first 1k steps of policy training and then trains with an intrinsically- motivated reward (as suggested by Lee et al. (2021a)) for 9k steps. The experimental set up and task configurations are selected following Park et al. (2022) which is the existing state the art method. Algorithm-specific super-parameters match those used by the corresponding paper and hyper-parameters determining feedback schedules and amounts match those used in |Park et al. (see Appendix E
2002)
5.1
COMPARING AGAINST PBRL. BASELINES
Figure and Table compare the performance Hindsight PRIOR to PEBBLE, SURF Park et al. 2022
RUNE I Liang et al. (2022), and MRN Liu et al. 20022) with perfect feedback. The amount

---

Published as a conference paper at ICLR 2024
Window Open /40
Door Open / 4000
Drawer Open / /400
Sucossa Rata
Success Rete
Succesa Rate
Mata- Pista-oridd
Sweep Into / 10000
Button Press 2000
Door Unlock/ 4000
Rota
Success
150
Dtp
Walker Walk / 100
Cheetah Run / 100
Quadruped Walk / 1000
Oracle
Oracle
Oradl Return
DM Control
400
200
203
18A
300 500
EX
20th
14
PRIOR
PEBBLE
SURF
RUNE
SAC
MRN
Figure 2: PbRL and SAC policy learning curves for six MetaWorld (top and middle rows) and three DMC (bottom row) tasks. Each experiment is specified as: task / feedback amount.
Table 1: Mcan (+ variance) normalized success rates (MetaWorld) and normalized returns (DMC)
across tasks.
Suite
PRIOR
PEBBLE
SURF
RUNE
MRN
MetaWorld
0.66 0.002
0.56
0.007
0.56 0.004
0.53 0.009
0.55 + 0.005
DMC
0.591 0.003
0.48 + 0.018
0.55 + 0.019
0.49 = 0.007
0.53 + 0.002
of feedback is held fixed across methods for a given task and is provided every 5k steps of policy training (X-axis), therefore learning curve performance in Figure rrelatvvee to the number of policy steps indicates both reward and policy sample complexity. For example, at policy step 30k for walker-walk, the preference dataset contains 10 preference triplets and 20 at 50k steps. Table reports the mean normalized return and success rate for each algorithm across tasks and shows that Hindsight PRIOR has the best overall performance across tasks. A two-tailed paired t-test with dependent means was performed over the normalized returns and success rates to determine that Hindsight PRIOR's performance gains are statistically significant. (Appendix F for and p-scores. Task specific normalized returns and success rates are reported in Appendix Fb.
For all tasks, Hindsight PRIOR matches or exceeds baseline performance, and for all except guadruped-walk. either converges to a higher performance point (e.g. 100% versus 80% success rate on window-oper) or requires significantly less preference labels achieve the same perfor- mance point (e.g. 100% success rate at ~ 350k policy steps versus ~ 550k for door-open). The results suggest that Hindsight PRIOR's credit assignment strategy improves PbRL beyond guid- ing exploration with reward uncertainty (Liang al. (022),, increasing the amount of preference feedback through pscudo-labelling (Park et al. 2022) and incorporating information about policy performance in reward learning (Liu et al. 2022).
Figure 3 (returns left and success rates center) shows the performance differences for PEBBLE (Lee et al. 2021a) la) and Hindsight PRIOR on wi window-open across different amounts of preference feed- back mistakes. The mistake amounts are percentages of the maximum feedback amount, specifically 0% (perfect labeller), 10%, 20%, and 40%. We compare against PEBBLE, because has compara-
7

---

Published as a conference paper at ICLR 2024
ble performance to the baselines (Figure 2 and Table and is the underlying PbRL algorithm. For all mistake amount conditions, Hindsight PRIOR outperforms PEBBLE. Furthermore, Hindsight PRIOR trained on a dataset with 20% labelling errors beats the performance of PEBBLE with no la- belling errors. The results suggest that the inclusion of a credit assignment strategy, specifically one guided by estimated state importance, makes reward and policy learning more robust to preference feedback labelling errors.
5.2
UNDERSTANDING THE PERFORMANCE GAINS
Window Open/ / 400
Window 0pen/400
Sweep Into 100000
Cracke Return
Sucoma Rata
Success Rate
471
100
208
UKC
2004
PRIOR FEBGLE
PRIOR PEDBLE
PRIOR PEDBLE
PRIOR PEBBLE
Stop
0.4
0.2
0.1
2004
400k
5022
2008
114
Figure 3: PbRL learning curves over different labelling mistake amounts (left & center : purple & pink for PEBBLE and red & magenta for PRIOR), and different methods for return distribution and dynamics-aware rewards (right).
In order to better understand sources of Hindsight PRIOR's performance gains we evaluate the im- portance of the state-importance guided return redistribution strategy by comparing against different redistribution strategies (QI), assess the impact of Hindsight PRIOR making reward learning dy- namics aware by replacing Lpriir with adapted bisimulation objective (Kemertas & Aumentado-
Armstrong 2021) (Q2), and qualitatively assess what the world model TT identifies as important states (Q3). The results show the benefits of the forward dynamics based state- inte-importance redis- tribution strategy, demonstrate that Hindsight PRIOR's 's contributions extend beyond making reward learning dynamics aware, and that A's attention weight identify reasonable state-action as important.
We compare against PEBBLE as it has comparable performance to the baselines (Figure and Table D and is the underlying PbRL algorithm for all baselines.
Redistribution Strategy (Q1): Hindsight PRIOR's ss redistribution strategy is compared against an uninformed return redistribution strategy, using the mean attention weights a serve as the reward targets (RVAR). The uniform strategy corresponds to assigning uniform importance to each state- action pair in a trajectory and each state-action pair is assumed to equally contribute to the preference feedback. The uniform strategy adapts Ren et al. (202T) (RRD) obtain the reward target Rearget by setting a₁ T. Figure (right green) shows that while uniform predicted return redistribution is on par with PÉBBLE (and in some cases better, see Appendix G.I.. Hindsight PRIOR is superior in feedback and environment sample efficiency.
Given Hindsigh PRIOR's performance relative to a uniform redistribution strategy, we amplify Hind- sight PRIOR's attention weights through min-max normalization of the attention map followed by a softmax (NRP). Amplifying the attention map moves it further from the uniform redistribution strategy and potentially improves it. However, Hindsight PRIOR and NPR have comparable per- formance (Figure 6 in Appendix G.3) showing that explicitly discouraging uniform redistribution strategy not necessary.
Dynamics Aware Reward Learning (Q2): While Hindsight PRIOR does not directly use the for- ward dynamics of the world model T, knowledge of transition dynamics influence how the re- ward function is learned. Therefore, we assess the contribution of dynamics-aware reward learning in the absence of return redistribution credit assignment strategy. To incorporate dynamics, bisimulation centtii representation learning objective, which has been used as data-efficient ap- proach for policy learning, is incorporated into reward learning. See Appendix G.2 for details on incorporating the bisimulation auxiliary encoder loss Kemertas & into Hindsight PRIOR.
(202I)
8

---

Published as a conference paper at ICLR 2024
The results show that making reward learning dynamics aware improves policy learning (Figure (right-yellow)) compared to PEBBLE,but not compared to Hindsight PRIOR. Therefore, while in- corporating of environment dynamics into reward learning explains part of Hindsight PRIOR's per- formance gains. it does not explain all of the performance gains highlighting the importance of Hindsight PRIOR's 's credit assignment strategy.
Examining Important States (Q3): Fig. T shows the attention over a trajectory snippet from Montezuma'ss 's Revenge (analysis in App. [t.. In our qualitative experiments with discrete domains of Atari Brockman et al. ( and control based domains of Metaworld Yu et al. (2020) we found a significant overlap between important states for future state prediction and underlying task.
5.3
ASSESSING SCALABILITY AND COMPATIBILITY
Window Open 4000
Window Open /400
Cheetah Run 100
Success Rate
Succisa Rute
Oracle Episode Return
PROCR
200
PDBBLC
PRECH-
M014
700
PEESLL 9920 Cosesstor FROOR
500
17:55
1000
3004
4.00
Figure 4: Learning curves evaluating different trajectory lengths (left), combining Hindsight PRIOR with SURF (center), and removing the influence of preference feedback (right).
Scalability (Q4): Since Hindsight PRIOR subroutines a forward dynamics model to obtain the at- tention map A we evaluate whether it can identify important states in longer trajectorics that provide more context for human evaluators. Figure o (left) and Appendix H show that, following similar trends as PEBBLE, Hindsight PRIOR's ' performance is consistent given a 4x increase in trajectory length (50 versus 200 query length).
Combining with PEBBLE Extensions (Q5): We investigate the benefits of Hindsight PRIOR when used in parallel with another sample-efficient PbRL techniques, like SURF (Park al. 2022). Figure I (center) shows combining Hindsight PRIOR with SURF (Park et al.] (022) improves policy performance relative to PEBBLE and SURF, but provides no real gain relative to Hindsight PRIOR alone.
Removing Preference Feedback (Q6): The results in Figure 4 (right) show the impact of making A very large (green) in Equation 5 resulting in reward function that is learned solely from Eprior- The inability of Hindsight PRIOR learn anything with a very large À verifies that focusing the reward signal around important states is not sufficient for policy learning.
CONCLUSION
We have presented Hindsight PRIOR, a novel technique to guide credit-assigmment during reward learning in PbRL that significantly improves both policy performance and learning speed by incor- porating state importance into reward learning. We use the attention weights of transformer-basee world model estimate state importance and guide predicted return redistribution to be proportional to state importance. The redistributed prediction rewards are then used as an auxiliary target during reward learning. We present results from extensive experiments on complex robot arm manipula- tion and locomotion tasks and compare against state of the art baselines to demonstrate the impact of Hindsight PRIOR and the importance of addressing the credit assignment problem in reward learning.
Limitations Future Work: Hindsight PRIOR greatly improves PbRL and our qualitative assess- ment shows that the selected important states are reasonable. However, it relies on the assumption that states that are important to the world model are also important to an arbitrary human. Different humans might attribute importance different states. Future work will investigate the alignment between the world model's 's important states and those people focus on when providing preference feedback as well investigation the personalization aspects of important state identification.

---

Published as a conference paper at ICLR 2024
REFERENCES
Riad Akrour, Mare Schoenauer, and Michele Sehag. Preference-based policy learning. Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 12-27. Springer, 2011.
James F Allen and Johannes A Koomen. Planning using a temporal world model. In Proceedings of the Eighth international joint conference on Artificial intelligence- Voluue 2. pp. 741-747, 1983.
Evlampios Apostolidis, Georgios Balaouras, Vasileios Mezaris, and Ioannis Patras. Combining global and local attention with positional encoding for video summarization. In 2021 IEEE inter- national symposium on multimedia (ISM), pp. 226-234 IEEE, 2021.
Jose A Arjona-Medinal Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32, 2019.
Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world models fa- cilitate zero-shot dynamics generalization from a single offline environment. In International Conference on Machine Learning, pp.
PMLR, 2021.
Manjot Bilkhu, Siyang Wang, and Tushar Dobhal. Attention all you need for videos: Self-attention based video summarization using universal transformers. arXiv preprint arXiv: /906.02792, 2019.
Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):224-45, 1952.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv: 0016606.01500 540, 2016.
Claus Bundesen. A theory of visual attention. Psychological review, (7)4))523, 1990.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances neural information processing sys- lems, 30, 2017.
Robert Desimone and John Duncan. Neural mechanisms of selective visual attention. Annual review of neuroscience, 18(1):193-22) 1995.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value expansion for efficient model-free reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), 2018.
Xuming Feng, Lei Wang, and Yaping Zhu, Video summarization with self-attention based encoder- decoder framework. In 2020 International Conference Culture-oriented Science & Technology (ICCST), pp. 208-214, IEEE, 2020.
Patrick Fernandes, Aman Madaan, Emmy Liu, António Farinhas, Pedro Henrique Martins, Amanda Bertsch, José GC de Souza, Shuyan Zhou, Tongshuang Wu, Graham Neubig, et al. Bridging the gap: A survey on integrating (human) feedback for natural language generation. arXiv preprint arXiv:2305.00955. 00955, 2023.
Johan Ferret, Raphaël Marinier, Matthieu Geist, and Olivier Pietquin. Self-attentive credit assign- ment for transfer in reinforcement learning. 2020.
Anirudh Goyal, Philemon Brakel, William Fedus, Soumye Singhal, Timothy Lillicrap, Sergey Levine, Hugo Larochelle, and Yoshua Bengio. Recall traces: Backtracking models for efficient reinforcement learning. arXiv preprint /00.v:/804.0003 08/00000779, 2018.
Samuel Greydanus, Anurag Koul, Jonathan Dodge, and Alan Fern. Visualizing and understanding atari agents. In International conference on machine learning, pp. 1792-1801. PMLR, 2018
Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with model-based acceleration. In International conference on machine learning, pp. 2829-2832 PMLR, 2016.
10

---

Published as a conference paper at ICLR 2024
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and appli- cations. arXiv preprint arXiv:1812.05900 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learming, pp. 2555-2565, PMLR, 2019b.
Anna Harutyunyan, Will Dabney, Thomas Mesnard, Mohammad Gheshlaghi Azar, Bilal Piot, Nico- las Heess, Hado van Hasselt, Gregory Wayne, Satinder Singh, Doina Precup, et al. Hindsight credit assignment. Advances in neural information processing systems, 32, 2019.
Donald Joseph Hejna III and Dorsa Sadigh. Few-shot preference learning for In Conference on Robot Learning, pp. 2015-2025. PMLR, 2023.
rl.
Hangkai Hu, Shiji Song, and Gao Huang. Self-attention-basee temporary curiosity in reinforcement learning exploration. IEEE Transactions on Systems, Man, and Cyhernetics: Systems, 51(9): 5773-5784, 2019.
Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Reward learning from human preferences and demonstrations in atari. Advances in neural information processing systems, 31, 2018.
Gabriel Kalweit and Joschka Boedecker. Uncertainty-driven imagination for continuous deep rein- forcement learning. In Conference on Robot Learning, pp. 195-206. PMLR, 2017.
Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. Advances in neural information processing systems, 31, 2018.
Mete Kemertas and Tristan Aumentado-Armsttrom Towards robust bisimulation metric learning. Advances in Neural Information Processing Systems, 34:4764-4777, 2021.
Changyeon Kim, Jongjin Park, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. Pref- erence transformer: Modeling human preferences using transformers for rl arXiv preprint arXiv:2303.00957. 2023.
D. Kingma and J. Ba. ADAM: A method for stochastic optimization. volume 3, 2015.
Tomasz Korbak, Kejian Shi, Angelica Chen, Rasika Vinayak Bhalerao, Christopher Buckley, Jason Phang, Samuel Bowman, and Ethan Perez. Pretraining language models with human prefer- ences. In International Conference on Machine Learning, pp. 17506-1755 PMLR, 2023.
Krakovna, Uesato, Mikulik, al. Specification gaming: The flip side of ai ingenuity
deepmind, 2020.
Pawel Ladosz, Lilian Weng, Minwoo Kim, and Hyondong Oh. Exploration in deep reinforcement learning: A survey. Information Fusion, 85:1-22, 2022.
Kimin Lee, Laura Smith, and Pieter Abbeel. Pebble: Feedback-efficient interactive rein- forcement learning via relabeling experience and unsupervised pre-training. arXiv preprint arXiv:2106.05091, 2021a.
Kimin Lee, Laura Smith, Anca Dragan, and Pieter Abbeel. B-pref: Benchmarking preference-based reinforcement learning. arXiv preprint arXiv:2111.03026, 2021b.
Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to- soxt-to-image models using human feedback. arXiv preprint arXiv:2302. 12192, 2023.

---

Published as a conference paper at ICLR 2024
Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and Shane Legg. Scalable agent alignment via reward modeling: a research direction. arXiv preprint arXiv:1811.088) 2018.
Xinran Liang, Katherine Shu, Kimin Lee, and Pieter Abbeel. Reward uncertainty for exploration in preference-based reinforcement learning, arXiv preprint arXiv:2205 2022.
Runze Liu, Fengshuo Bai, Yali Du, and Yaodong Yang. Meta-reward-net: Implicitly differentiable
reward learning for preference-based reinforcement learning. Advances in Neural Information Processing Systems, 35:22270-22284, 2022.
Yen-Ting Liu, Yu-Jhe Li, Fu-En Yang, Shang-Fu Chen, and Yu-Chiang Frank Wang. Learning hierarchical self-attention for video summarization. In 2019 IEEE international conference on image processing (ICIP), pp. 3377-3381. IEEE, 2019.
Anthony Manchin, Ehsan Abbasnejad, and Anton Van Den Hengel. Reinforcement learning with at- tention that works: A self-supervised approach. In Neural Information Processing: 26th Interna- tional Conference, ICONIP 2019, Sydney, NSW, Australia, December 12-15, 2019, Proceedings, Part V 26, pp. 223-230. Springer, 2019.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL
lew. net. forum?1d=vhPu AAcb0xb
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Jongjin Park. Younggyo Seo, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. Surf: Semi-supervised reward learning with data augmentation for feedback-efficiee preference- reinforcement learning. arXiv preprint arXiv:2203, 10050, 2022.
Vihang P Patil, Markus Hofmarcher,
Dinu, Matthias Dorfer, Patrick M Blies, Johannes Brandstetter, Jose A Arjona- and Sepp Hochreiter. Align-rudder: Learning from few demonstrations by reward redistribution. arXiv preprint urXiv:2009. 14108, 2020.
Gabrielle Ras, Ning Xie, Marcel Van Gerven, and Derek Doran. Explainable deep learning: A field guide for the uninitiated. Journal of Artificial Intelligence Research, 73:329-396 2022.
Zhizhou Ren, Ruihan Guo, Yuan Zhou, and Jian Peng. Learning long-term reward redistribution via randomized return decomposition. arXiv preprint arXiv:2111.13499 2021.
Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-basee world mod- els are happy with 100k interactions. arXiv preprint arXiv:2303.07109, 2023.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based local- ization. In Proceedings of the IEEE international conference on computer vision, pp. 618-68 618-626, 2017.
Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Vi- sualising image classification models and saliency maps. arXiv preprint arXiv:1312.603: 2013.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6: 6:100022, 2020.
Peter Vamplew, Richard Dazeley, Cameron Foale, Sally Firmin, and Jane Mummery. Human- aligned artificial intelligence is multiobjective problem. Ethics and Information Technology, 20(1))27-40, 2018
Shikhar Vashishth. Shyam Upadhyay. Gaurav Singh Tomar, and Manaal Faruqui. Attention inter- pretability across nlp tasks. arXiv preprint arXiv: 1218, 2019.
12

---

Published as a conference paper at ICLR 2024
Laurens Weitkamp. Elise der Pol, and Zeynep Akata. Visual rationalizations in deep rein- forcement learning for atari games. In Artificial Intelligence: 30th Benelux Conference, BNAIC 2018, He-Hertogenbo,,h The Netherlands, November 8-9, 2018, Revised Selected Papers 30, pp- 151-165. Springer, 2019.
Sarah Wiegreffe and Yuval Pinter. arXiv:1908. 04626, 2019.
Attention is not not explanation.
arXiv preprint
Christian Wirth, Riad Akrour, Gerhard Neumann, Johannes Fürnkranz, al. survey of preference-
based reinforcement learning methods. Journal of Machine Learning Research, 18(136):1-46, 2017.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100, PMLR, 2020.
Banghua Zhu, Jiantao Jiao, and Michael I Jordan. Principled reinforcement learning with human feedback from pairwise or k-wwise comparisons. arXiv preprint arXiv:2301.11227, 2023.
Daniel Ziegler, Nisan Stiennon, Jeffrey Wu, Tom Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.088: 08593, 2019.
13

---

Published as a conference paper at ICLR 2024
APPENDIX
A DOMAINS FOR EMPIRICAL EVALUATION
We consider six domains of Metaworld Yu et al. (2020) and three domains of DM Control Tunyasu- vunakool et al. 2020) for our empirical evaluation.
For our experiments we use the internal state representation as the agent state. This includes reward learning, policy learning and world model training. The state features are described in Yu ct al. (2020). Further we follow Park et al. (2022) to utilize the packaged task rewards in MetaWorld for our synthetic oracles.
B REWARD ARCHITECTURE FOR PRIOR
We follow Park al. (2022) reward architecture for Metaworld and DMControl domains.
C WORLD MODEL LEARNING
We generally follow the training paradigm suggested in TWM Robine et al. (2023) to learn our forward dynamics model. Furter, follow the same architecture with minor changes to work with continuous control domains:
1. TWM is proposed for discrete actions image based atari Brockman et al. (2016) domains, therefore modify the observation encode layer to take in 1-D state vector (instead of channel image vector). Additionally, the original TWM makes use frame-stacking which we do not.
2. We change the state predictor to take continuous valued vector based action space (as in our case) instead of discrete I-D actions (as in Atari).
Finally, we modify the world model sequence length and memory length be same as the query length such that we can feed the whole trajectory as an input to the TWM model.
To obtain an attention map from TWM we auto- feed in a trajectory, i.e. first we feed in the first actton tuple trajectory of size 1, then the next state, action tuple and so on. At the final step we take the attention vector from the attention layers in the architecture.
C.1 OBSERVATION MODEL
The observation model in our world model encodes the input state,action tuple. The observation model has two components, an encode and a decoder both of which are multi-layer perceptrons.
The encoder has two hidden layers of size 512 and an output layer of size (32,). The decoder layer has an input size of (32.) followed by two hidden layers of size 512 and the output shape is same as number of state features. As recommended by RRobine al. (2023) we use SiLU activation in the MLP.
C.2 LATENT STATE PREDICTOR
We follow Robine et al. 20233 contruct the world model architecture with minor changes required to adapt continuous control action space and 1-D state vector. That is, vanilla TWM requires categorical action embeddings (because of the discrete action space) but do not. Finally, TWM Robine et al. (003)) can operate in different modes with respect to the prediction heads from the latent state predictor such as, next state prediction, reward prediction, and discount prediction. For Hindsight PRIOR we only require the next-state prediction output head.
Similarly TWM supports different input modes i.e. it can take (state, tuple and (state, action, reward) tuple. We compute state conditioned hindsight priors and therefore only need the (state, action) input. Robine et al. 2023) is inconclusive on the need for rewards in the inputs. Moreover,
14

---

Published as a conference paper at ICLR 2024
since PRIOR is learning the reward model begin with only use (state, action) input. Robine et al. 2023) uses a Transformer XL architecture which we borrow as is.
C.3 TRAINING WORLD MODEL
We share the replay buffer of the agent with the world model. Similar to PbRL paradigm, the world model first gets access to bank of state transitions during PbRL's pre pre-training step. Once the pretraining step complete the world model trained on this data (to get the observation model). After this pretraining step only the dynamics model is trained on incoming data every jth step of PbRL, loop. The world model is trained on its bank of transitions as suggested by |Robine et al. 2023
D POLICY EVALUATION METRICS
Algorithms are evaluated according to their learning curves the course of policy and reward training along with their normalized returns for Deep Mind Control Suite and normalized success rates for Meta World Lee al. (2021b). An algorithm's 's normalized return success rate measures how well the policies trained jointly with a preference-learned reward function recovers the perfor- mance of policy trained on the target reward function. For each episode of policy training, the PbRL. or SAC policy's 's return or success rate is computed and then the PbRL policy is evaluated based how well it is able recover "optimal" performance approximated with SAC trained on the ground truth reward. The normalized returns are computed as:
Σ Ty (S (a:)) To (St, Th (at))
normalized returns +
(6)
where T is the number of policy training steps, F. is the target reward function, T is the policy trained in conjunction the learned reward function, and n is the policy trained on the target reward function. The normalized success rates are computed as:
normalized success rates T Σ
success T. aT (at))²
(at)
success
where success(-) indicates whether action resulted in the policy reaching the goal state.
E
A results in the paper are reported over five random seeds: [12345, 23456, 34567, 45678, 56789).
E. TRAIN SYYPER-PARAMETE
This section specifies the hyper- ppper-parameters (e.g. learning rate, batch size, etc) used for the exper- iments and results. The SAC, and PEBBLE experiments all match those used in Haarnoja et al. (2018) and [ee et al. (2021a) respectively. The SAC hyper-parameters are specified in Table the PEBBLE hyper-parameters are given in Table 3 the hyper-parameters used to train on with Hind- sight PRIOR are in Table 4 and finally the hy peprparraeeters used for training our world model in Table Table only mentions the dyper-parameten that we change in the prescribed Robine et al. configuration.
Table 5
2023
15

---

Published as a conference paper at ICLR 2024
Table 2: Training haper-parameters for SAC (Haarnoja et al. 2018
HYPER-PARAMETER
VALUE
1e-3 (cheetah), 5e-4 (walker),
Learning rate
lc-4 (quadruped), 3e-4 (Meta-World)
Batch size
512 (DMC), 1024 (Meta-World)
Total timesteps
500k, 1M (quadruped, sweep into)
Optimizer
Adam Kingma & Ba 2015
Critic EMA
5c-3
Critic target update freq.
(B1, (32)
(0.9, 0.999)
Initial Temperature
0.1
Discount
0.99
Table 3: PEBBLE hyper-parameters (Lee et al. 2021a).
PAYPER-PARAMETER
VALUE
3e-4 (Metaworld),
Learning rate PEBBLE
Se-4 (Walker, Cheetah),
le-4 (Quadruped)
Optimizer
Adam (Kingma & Ba 2015
Segment length /
50
1000/100, 100/10 (DMC)
Feedhack amount / number queries (M)
10000/50, 4000/20, 2000/25, 400/10 MMta--orldd
20000 (walker, cheetah), 30000 (quadruped),
Steps between queries (K)
5000 (Meta-World)
Table 4: Hindsight PRIOR hyper-parameters.
HYPER-PARAN PARAMETRR
VALUE
1000 (MetaWorld), 5 (DMC)
update frequency j
2000 steps
Table 5: World Model hyper-parameteen in Hindsight PRIOR
HYPER- PARAMETER
VALUE
world model sequence length
50
world model memory length
world model batch size
100
16

---

Published as a conference paper at ICLR 2024
F NORMALIZED RETURNS AND SUCCESS RATES BY TASK
The mean normalized scores for each algorithm on each task are given for MetaWorld in Table and DMC in Table
two-tailed paired t-test with dependent means (significant at p .05) was performed over the normalized returns and success rates to determine that Hindsight PRIOR's performance gains statistically significant over:
1. Meta World: PEBBLE (t -3.92, p 0.006), SURF (t 2.85, p 0.025), RUNE
(f 5..3. p = 0.001), and MRN (t = -4.91, p= 0.002)
2. DMC: PEBBLE (t -.447, p 0.00844), SURF -2.52, p .03541), RUNE
(t = 7.744967, = 0.00006), and MRN (t = -2399223, = 0.0437)
Table 6: Normalized Success Rate for MetaWorld domains
Env/Algorithm PRIOR
PEBBLE
SURF
RUNE
MRN
Window Open
0.55
0.40
0.45
0.35
0.43
Door Open
0.65
0.62
0.55
0.53
0.53
Drawer Open
0.68
0.59
0.59
0.59
0.53
Sweep Into
0.69
0.51
0.54
0.52
0.62
Button Press
0.69
0.65
0.65
0.67
0.65
Door Unlock
0.68
0.62
0.64
0.53
0.54
Table 7: Normalized Returns for DM Control domains
Env/Algorithm
PRIOR
PEBBLE
SURF
RUNE
MRN
Walker Walk
0.60
0.46
0.66
0.47
0.59
Cheetah Run
0.51
0.33
0.36
0.39
0.46
Quadruped Walk
0.65
0.66
0.64
0.60
0.54
G ADAPTED BASELINES FOR PBRL
Sweep Into / 10000
Door Open 4000
Window Open / 400
Figure 5: Learning curves PRIOR, BISIM, RVAR and baseline PEBBLE
G.1 UNINFORMED RETURN REDISTRIBUTION (RVAR)
Inspired from Ren et al. 2021 we consider the Uninformed return redistribution baseline that es- sentially has reward targets the mean reward in the trjaectory, i.e. Ttarget - This essentially reduces the variance of the trajectory rewards (hence the name RVAR). As discussed previously, RVAR is marginally better than PEBBLE and PRIOR is superior to such uninformed redistribu- tion technique. (See Fig. 5)
G.2 BISIMULATION METRIC BISIM
To incorporate the bisimulation metrics into PbRL, the following bisimulation metric loss:
Jrisim (v) (1]=: 22111 Irs ryl -11=2, (x1,,0)) ((2300))22
(8)
17

---

Published as a conference paper at ICLR 2024
adapted from Equation 4 in Kemertas & Aumentado-AAmstronn Armstrong 2021). We use the reward model to provide the predicted rewards rs. y required by bisim loss. Further we use the penultimate layer as the embedding layer to obtain 2;, 2j- Next,
add an additional head from the penultimate layer that predicts the next state embedding to obtain =4, 2 as next state predictors. Finally, we reuse the trajectory buffer (as used by Hindsight PRIOR) obtain the states on which we compute the bisim- larget distance and finally optimize the loss as above. We verify that the BISIM adapted baseline offers only marginal improvements over baseline PEBBLE as given Fig.
G.3 NORMALIZED REDISTRIBUTION PRIOR (NRP)
Readers may question whether the raw attention values extracted from the forward dynamics predic- tion task are the best choice or certain post processing may further improve PRIOR's performance. We construct variant of PRIOR referred to PRIOR-NRP where we perform min-max normal- ization of the attention vector a. That is:
0g
Oymin
à; so s0f fmarr
Ok O'mas
O'yán
The above equation amplifies the attention values thereby preventing the reward targets to become uniform. Fig. @lshows that default PRIOR (without any postprocessing of attention) performs well and NRP like post-processing is not required.
Window Open / 400
Door Open / 4000
Success Rate
Success Rate
PRIOR-NPP
PRIOR
PEBBLE
PEBBLE
PRIOR
PRIOR-NRP
100
100
80
80
60
60
60
40
20
20
Step
Step
100k
200k
300k
400k
500
400k
600k
800k
Figure 6: Learning curves of PRIOR, PRIOR-NRP and baseline PEBBLE
H LONG TRAJECTORIES
Figure shows the performance of PRIOR compared to PEBBLE when the trajectory length is 4x (200). While the Transformer XL architecture is already known for scaling to much longer trajectory lengths Robine et al. (2023) our experiments conclude that this is indeed the case for the challenging continuous control domains.
Window Open 400
Door Open 4000
Success Rate
Success Rate
PRIOR 50
PRIOR: 288
PEBBLE 59
PEBBLE
50
PRIOR 50
PRIOR: 280
PEBBLE 200
100
100
80
80
60
60
Stap
Stop
200k
400k
600k
800k
Figure 7: Learning curves of PRIOR and PEBBLE on query segment lengths of 50 and 200.
18

---

Published as a conference paper at ICLR 2024
1 QUALITATIVE STATE IMPORTANCE ASSESSMENT
Figure 8: Attention analysis on Montezuma's Revenge. The plot shows a situation where the agent (red animated character) attempts to jump from the green platform towards the rope. Each row represents a layer in the Transformer Model and the columns represent the time step where left most column is time t-T and right most column is the present state of the agent. In hindsight, this attention map to obtain the attention over the past states (given as blue cells) and attention over past actions (given
compute
red cells). Note that the map begins with blue cell (as so, dn, a1; :1;a) ...).. We highlight that the agent attends to past states which corresponds to point of launch which can be considered an important summary state for the complete trajectory of jumping from the platform to the rope.
We qualitatively evaluate the states identified by the attention weights a as important for a given trajectory. We evaluate whether the attention map A salience events. We conduct experiments on Atari (Montezuma's Revenge) and Metaworld (Window-Open Door Open) domains, and seek answer whether the world model attends states in the complete history (even for Markovian transitions) and whether that correlates with "critical" events (similar h al. (2023) describes it). From figure X we can see that the forward model does attend to past states and actions. Moreover, upon closer inspection by executing known maneuvers, such as jumping across the platform the rope Montezuma's Revenge, we find that the agent attending to past critical events like the point of launch. We do not expect attention over states in history to be aligned with human's reward function as human may have arbitrary preference unknown to the dynamics model. However, we do expect the attention to be aligned with "critical events" loosely defined as apriori states enough to summarize a trajectory. To further investigate this we force the reward model to predict rewards aligned with the PRIOR reward targets by very high value Aprior in equation 5 and find that the reward model is unable to learn preference-relevant reward model at all. This shows that PRIOR reward targets are not task specific but contain salience information states which can be considered "critical" in the sampled trajectory.
trajectory can capture critical
19


Reviews:
[{"review_id": "otzSPwVbk8", "reviewer": "Reviewer_4UMz", "summary": "This work presents Hindsight PRIOR, a novel technique to guide credit assignment to improve reward inference in Preference-based Reinforcement Learning. The key contribution in this paper is the utilization of attention weights from a transformer-based world model to estimate state importance and the formulation of return redistribution to be proportional to the attention-deduced state importance. The authors present information regarding related work, their approach, and an empirical evaluation in the Deep Mind Control and MetaWorld Control Suites. The results with a synthetic labeler are positive, displaying PRIOR achieves high success across a variety of tasks.", "questions": "Please address the weaknesses above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The proposed method is an improvement to PbRL frameworks. Given the references in the paper and how humans may utilize attention similarly to transformer models, utilizing attention weights to redistribute return may improve preference-based reinforcement learning with  end-users.\n+  This paper is well-written and contains sufficient detail to understand the proposed approach.\n+ The evaluation is extensive, and touches on several important questions beyond simple performance.", "weaknesses": "- It would be beneficial to note exactly how many trajectory labels such a framework requires. This would help detail whether such a framework would be feasible with actual end-users. Further, including actual tests utilizing this framework with human end-users would provide further evidence that PRIOR works well.\n- Along this thread, it seems the simultaneous learning of a highly parameterized world model and reward model is accomplished faster than other works that simply inferring a reward model, as shown by the sample-efficiency in policy learning. Could you comment on why this is the case? I'm unsure if this relates to a paragraph on page 2 referencing the choice of architecture of the reward network.\n- As PRIOR utilizes PEBBLE as its backbone algorithm, this should be touched on in the related work.\n-  In the evaluation, there are several references that are not labeled correctly and lead to ??. As several of these baselines are not referenced or explained previously, it leads to confusion regarding the results.\n- Could you provide justification on why the attention coefficient for the state and action should be equally waited within the \\alpha coefficient?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ungjaff4y8", "reviewer": "Reviewer_gM8u", "summary": "This paper addresses the credit assignment problem in preference-based Reinforcement Learning (PbRL) algorithms. Given sparse feedback, it is challenging to determine where rewards should be assigned in a trajectory, i.e., which states are significant. The proposed solution combines the classical PbRL algorithm PEBBLE with a prior obtained from a world model. This approach assumes that states receiving high attention in the world model are likely to be rewarding, assigning them higher weight when estimating the reward function. The algorithm is evaluated on simulated problems from the DMC suite and MetaWorld control.", "questions": "How do you expect the performance of your algorithm to compare with more sample-efficient algorithms, such as few-shot preference learning [1]?\n\n[1] Hejna, Joey & Sadigh, Dorsa (2023). Few-shot preference learning for human-in-the-loop rl. In Conference on Robot Learning (pp. 2014-2025). PMLR.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The issue of preference-based RL and the credit assignment problem is highly relevant, particularly considering the need for numerous samples to accurately estimate the reward function.\n\nThe idea of utilizing the attention layers of the learned world model to identify rewarding state-action pairs is creative and seemingly novel, offering a straightforward but effective solution.\n\nThe approach outperforms other methods, demonstrating its effectiveness in comparison.\n\nThe paper is well-written and presents its content in an understandable manner.", "weaknesses": "The primary limitation of this work, as acknowledged in the paper, is its reliance on the assumption that states deemed important by the world model are also significant for reward design. While this insight is valuable, the contribution of the paper might be relatively modest, given that the primary novelty lies in a straightforward implementation of this assumption and its evaluation. With that, the quality of the contribution may not fully meet the requirements for acceptance at ICLR.\n\nThe paper would benefit from additional work to clarify the extent to which the learned attention in world models aids in task characterization for interpretability and transferability, as these are key applications of reward learning (building on Q3 in the paper).\n\n__Typos:__\n\nPage 2, \"Learning World Models\": \"us it to\" should be corrected to \"use it to.\"\n\nPage 4, \"a local minima\" (plural) should be corrected to \"a local minimum.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2J7Kk2l2bF", "reviewer": "Reviewer_7k1k", "summary": "This paper presents a new credit assignment strategy for efficiently learning a reward function in preference-based reinforcement learning (PbRL). The preference-based reward model relies on the well-known Bradley-Terry model, and the basic loss consists of the standard cross-entropy loss between the predicted and the true preference labels. The novel contribution of this paper is an additional loss term that redistributes the expected discounted return to each state-action pair in a particular manner. Specifically, a transformer-based forward dynamics model is learned as an auxiliary task, and the expected discounted return is redistributed to be proportional to the attention weights of the state-action pairs. This additional loss term serves as the prior for reward learning, and the authors hypothesize that it leads to sample efficiency and overall policy performance. The empirical evaluation on Deep Mind Control (DMC) and MetaWorld control tasks suggest positive results.", "questions": "1) In the explanation of equation (3), the authors define H as the entropy. Do you mean “cross entropy” there?\n\n\n2) A recent study [1] finds that reporting the mean and variance for performance evaluation of RL policies is insufficient, and suggests reporting confidence intervals or performance profiles as more objective measures. Have you considered them at all instead of the t-test?\n\n\n[1] Agarwal, Rishabh, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Marc Bellemare. \"Deep reinforcement learning at the edge of the statistical precipice.\" Advances in neural information processing systems 34 (2021): 29304-29320.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "A major strength of the paper is that the proposed framework is general in the sense that it can be applied to different baseline PbRL algorithms to improve sample efficiency and performance.\nIn the beginning it was not intuitive to believe that the return should be redistributed according to attention weights of the forward dynamics, but the explanation provided in Appendix I is somewhat convincing that the attention weights are reflective of critical events that summarize the whole trajectory, and thus it is those critical state-action pairs that contribute the most to the success/failure of the trajectory.", "weaknesses": "The authors propose an intriguing approach to boost the performance of PbRL, but some of the empirical results presented in Section 5 reveal some weaknesses. In particular, in the Drawer Open / 4000 task in Figure 2, the success rate of the proposed approach plateaus after 400k steps and is eventually surpassed by some other baselines. This might imply that the hindsight prior eventually hurts the performance of the learned policy as more preference labels become available. Similarly, in Section 5.3 the authors find that a large coefficient on the prior loss leads to a collapse of the learned policy. Those observations indicate that such a prior is assistive of policy learning only up to a certain point (e.g. relatively small preference data or small prior coefficient), and it may eventually hurt the performance if we exceed those bars.\n\nIn terms of the presentation of the paper, there is much room for improvement. First, some citations and references are missing and appearing as “?” or “??”. Second, some performance plots in Section 5 have too many curves of similar colors and widths, making it difficult to extract information from them. Specifically, I suggest that the authors use different line styles for Figure 2 (in particular for SAC since it’s an oracular baseline). The left two figures of Figure 3 are also hard to read as there are 8 line plots on each tiny figure. Third, the explanation of hindsight prior is confusing and needs elaboration/clarification. (a) The transformer uses H attention heads, but H does not appear in the definition of the attention matrix A. (b) Equation (4) uses the notation $R_{target}$ and $\\hat{R}$, but it is unclear how it is related to $\\hat{\\mathbf{r}}_{\\psi}$ and $\\mathbf{r}_{target}$ in the line above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)