Dynamics-Aware Comparison of Learned Reward Functions

1 INTRODUCTION . Reward functions provide a general way to specify objectives for reinforcement learning . Tasks such as game playing have simple , natural reward functions that enable learning of effective policies ( Tesauro , 1995 ; Mnih et al. , 2013 ; Silver et al. , 2017 ) ; however , many tasks , such as those involving human interaction , multiple competing objectives , or high-dimensional state spaces , do not have easy to define reward functions ( Knox et al. , 2021 ; Holland et al. , 2013 ; Cabi et al. , 2019 ) . In these cases , reward functions can be learned from demonstrations , explicit preferences , or other forms of reward supervision ( Abbeel & Ng , 2004 ; Christiano et al. , 2017 ; Cabi et al. , 2019 ) . However , evaluating the quality of learned reward functions , which may be too complicated to manually inspect , remains an open challenge . Prior work on comparing reward functions typically involves a computationally expensive and indirect process of ( i ) learning a reward function , ( ii ) learning a policy on that reward function , and ( iii ) evaluating that policy on a known , ground-truth reward function . Gleave et al . ( 2020 ) refer to this as the “ rollout method ” for evaluating methods for learning reward functions . Furthermore , the rollout method can be potentially dangerous in safety critical environments because ( i ) it can not distinguish deficiencies between the reward function and the policy optimization procedure , and ( ii ) rewards are evaluated on the distribution of transitions induced by an optimized policy , which may not capture all the transitions required to accurately quantify reward differences . To address these challenges , Gleave et al . ( 2020 ) introduce Equivalent-Policy Invariant Comparison ( EPIC ) . EPIC quantifies differences in reward functions directly , without learning policies . EPIC uses an offline dataset of transitions , and consists of two stages . First , transitions from this dataset are used to convert reward functions to a canonical form that is invariant to reward transformations that do not affect the optimal policy . Second , the correlation between reward values is computed on transition samples , yielding a pseudometric capturing reward function similarity . However , when computing this canonical form , EPIC evaluates rewards on transitions between all state-state combinations , regardless of whether those state-state transitions are possible or not . While this makes it possible for EPIC to quantify reward distances that are robust to arbitrary changes in transition dynamics , in practice evaluating learned reward functions on physically impossible transitions leads to unreliable reward values , as these are outside the distribution of the transitions on which the reward functions are learned . We show empirically in Section 5 that these out-of-distribution reward evaluations can render EPIC an ineffective reward comparison metric . To address this issue , we propose a new method for quantifying differences in reward functions , Dynamics-Aware Reward Distance ( DARD ) , that transforms reward functions into a form allowing for reward-shaping-invariant comparisons while ensuring that learned reward functions are only evaluated on transitions that are physically realizable . We achieve this by assuming access to a transition model for the environment of interest , and evaluate reward functions on transitions sampled from this transition model . This ensures that reward functions are not evaluated on impossible transitions , which will necessarily be out-of-distribution , and is well motivated in a variety of physical environments in which transition models are unlikely to change significantly between learning and deployment . Experiments suggest that DARD better reflects reward function similarity than does EPIC in comparing learned reward functions in two simulated , physics-based environments : the first involving multi-agent navigation , and the second involving manipulating a robotic arm . Videos available at https : //sites.google.com/view/dard-paper . 2 RELATED WORK . 2.1 LEARNING REWARD FUNCTIONS . A variety of methods exist for learning reward functions , which can be categorized based on the form of supervision used . Rewards can be learned from supervision provided by explicit labels . For example , human-labeled preferences between demonstrations can be used to learn reward functions ( Christiano et al. , 2017 ) . Rewards can also be labeled on a per-timestep basis , allowing for learning via supervised regression ( Cabi et al. , 2019 ) . Alternatively , inverse reinforcement learning ( IRL ) methods typically assume access to expert demonstrations , and attempt to find a reward function under which the expert is uniquely optimal ( Ziebart et al. , 2008 ; Abbeel & Ng , 2004 ) . IRL is underconstrained in that many reward functions can make the expert policy appear optimal . This problem is often addressed by favoring reward functions that induce high-entropy optimal policies ( Ziebart et al. , 2008 ) . Recent work in IRL focuses on learning reward functions with expressive function approximators such as neural networks , particularly in an adversarial formulation ( Fu et al. , 2017 ) . There has also been recent interest in offline IRL methods ( Lee et al. , 2019 ) . In this work , we study how these various reward learning methods can be evaluated efficiently and accurately in practice . 2.2 EVALUATING METHODS OF LEARNING REWARD FUNCTIONS . Reward functions are typically evaluated by first learning an optimal policy under the reward function . The value of the induced policy is found by aggregating the returns from online trajectory rollouts under the known , ground-truth reward function . This requirement of online rollouts from experimental policies is often prohibitive in safety-critical domains such as autonomous driving or decision making in healthcare ( Thapa et al. , 2005 ) . Off-policy evaluation ( OPE ) methods attempt to address these issues by estimating the value of a policy in a purely offline setting by using previously collected data . Traditional OPE methods address policy mismatch by using a form of importance sampling ( Precup et al. , 2000 ) . While this approach can be effective , it is unbiased only if the underlying expert policy is known , and can often be subject to large variance . Direct OPE methods ( Dudik et al . ( 2011 ) ; Voloshin et al . ( 2020 ) ) , apply regression-based techniques in either a model-based ( Paduraru , 2013 ) or model-free setting ( Le et al. , 2019 ) . Contemporary works combine these two methods using Doubly Robust Estimation , which originated in a contextual bandits setting ( Dudik et al. , 2011 ) , and was later extended to reinforcement learning by Jiang & Li ( 2016 ) . More recent works seek to forego models and importance sampling entirely , such as Irpan et al . ( 2019 ) , which recasts OPE as a positiveunlabeled classification problem with the Q-function as the decision function . An overview of all of these methods and empirical analysis of their performance can be found in Voloshin et al . ( 2020 ) . All OPE methods indirectly evaluate the reward by estimating the value of the induced optimal policy . This introduces additional complexity into the evaluation procedure and can obfuscate sources of error as resulting either from aspects of policy learning or due to the underlying reward function . EPIC ( Gleave et al. , 2020 ) seeks to circumvent these shortcomings by comparing reward functions directly . Our work builds on EPIC as covered in Section 4 , but uses a transition model of the environment to significantly improve the accuracy of comparisons involving learned reward functions . 3 PRELIMINARIES . 3.1 MARKOV DECISION PROCESSES . A Markov decision process ( MDP ) is a tuple ( S , A , T , R , γ , d0 ) where S and A are the state and action spaces . The transition model T : S × A × S → [ 0 , 1 ] maps a state and action to a probability distribution over next states , the reward R : S × A × S → R measures the quality of a transition . Finally , γ ∈ [ 0 , 1 ] and d0 are respectively the discount factor and initial state distribution . A trajectory τ consists of a sequence of state-action pairs : τ = { ( s0 , a0 ) , ( s1 , a1 ) , ... } , and the return g of a trajectory is defined as the sum of discounted rewards along that trajectory : g ( τ ) =∑∞ t=0 γ tR ( st , at , st+1 ) . The goal in an MDP is to find a policy π : S × A → [ 0 , 1 ] that maximizes the expected return , E [ g ( τ ) ] , where s0 ∼ d0 , at ∼ π ( at|st ) , and st+1 ∼ T ( st+1|st , at ) . We denote a distribution over a space S as DS , and the set of distributions over a space S as ∆ ( S ) . 3.2 REWARD FUNCTION EQUIVALENCE . Our goal is to define a metric for comparing reward functions without learning policies . Nevertheless , the measure of similarity we are interested in is the extent to which the optimal policies induced by two reward functions are the same . Because different reward functions can induce the same set of optimal policies , we in fact seek a pseudometric as formalized in the following definition . Definition 3.1 ( Pseudometric ( Definition 3.5 from Gleave et al . ( 2020 ) ) ) . Let X be a set and d : X × X → [ 0 , ∞ ] a function . d is a premetric if d ( x , x ) = 0 for all x ∈ X . d is a pseudometric if , furthermore , it is symmetric , d ( x , y ) = d ( y , x ) for all x , y ∈ X ; and satisfies the triangle inequality , d ( x , z ) ≤ d ( x , y ) + d ( y , z ) for all x , y , z ∈ X . d is a metric if , furthermore , for all x , y ∈ X , d ( x , y ) = 0 =⇒ x = y . Which rewards induce the same set of optimal policies ? Ng et al . ( 1999 ) showed that without additional prior knowledge about the MDP only reward functions that are related through a difference in state potentials are equivalent . This “ reward shaping ” refers to an additive transformation F applied to an initial reward function R to compute a shaped reward R′ : R′ ( s , a , s′ ) = R ( s , a , s′ ) + F ( s , a , s′ ) . Reward shaping is typically applied in a reward design setting to produce denser reward functions that make policy optimization easier in practice ( Ng et al. , 2003 ; Schulman et al. , 2015 ) . Ng et al . ( 1999 ) showed that if F is of the form F ( s , a , s′ ) = γΦ ( s′ ) − Φ ( s ) for arbitrary state potential function Φ ( s ) , the set of optimal policies under R and R′ are the same . This provides a definition for reward equivalence : Definition 3.2 ( Reward Equivalence , Definition 3.3 from Gleave et al . ( 2020 ) ) . Bounded rewards RA and RB , are equivalent if and only if ∃λ > 0 , and Φ : S → R bounded s.t . RB ( s , a , s ′ ) = λRA ( s , a , s ′ ) + γΦ ( s′ ) − Φ ( s ) ∀s , s′ ∈ S , a ∈ A . This form of equivalence is only possible in reward functions that depend on ( s , a , s′ ) , which motivates our focus on these types of rewards in this paper . This choice is further motivated by the fact that there are cases where conditioning on the next state s′ can make reward learning simpler , for example in highly stochastic MDPs , or in cases where rewards are a simple function of ( s , a , s′ ) , but not of ( s , a ) . We next describe an existing pseudometric that in theory is capable of identifying these reward equivalence classes .

Summary: This paper proposes a novel method for comparing reward functions without policy optimization, the Dynamics-Aware Reward Distance. DARD improves on the state-of-the-art EPIC distance by using an approximate transition model to evaluate reward functions on transitions close to the training distribution, while EPIC evaluates on arbitrary transitions (which can be infeasible). The authors prove that DARD is also invariant to reward shaping, and experimentally demonstrate on the Bouncing Balls and Reacher environments that DARD is a better predictor of policy return than EPIC and works for coverage distributions induced by either random or expert actions.