Learning Invariant Reward Functions through Trajectory Interventions

1 INTRODUCTION . Recent advancements in reinforcement learning methods applied to real world scenarios have revealed a major limitation , namely how to design a problem specific reward function . This design challenge typically requires a cumbersome and error-prone process of handcrafting a heuristic function to account for all the intricacies of the task at hand . Eliciting the correct behaviour via the optimization of a reward function is of paramount importance both for training artificial agents , such as robots , as well as human agents , for instance , in the case of teaching skills in a simulated environment , but modelling rewards for behavior learning remains an open challenge in a large number of domains . Inverse reinforcement learning ( IRL ) solves the problem of inferring the reward function of a Markov Decision Process ( MDP ) based on a dataset of temporal behaviors . Such trajectories obtained from an agent are assumed to demonstrate near-optimal performance in the respective MDP . The problem to recover the reward from the statistics of the expert trajectories is ill-posed , as there typically exist many reward functions which satisfy the optimization constraints ( Ng et al. , 1999 ) . Maximum entropy inverse reinforcement learning resolves this ambiguity in a unique way . Typically , the amount of expert demonstrations used to infer the reward is limited , which poses the danger of overfitting to the noise in the expert demonstrations . In particular , high-dimensional , parameterized reward models without appropriate regularization tend to absorb spurious correlations present in the trajectory data which do not generalize outside of the training setup . While optimizing the empirical cumulative reward obtained by such models , such an adaptation causes the agent to learn behaviours which frequently fail at test time . This type of observational overfitting has been shown in the context of reinforcement learning ( Song et al. , 2019 ) . An additional difficulty arises when there is a significant discrepancy between expert demonstrations originating from different experts . Adversarial imitation methods fail to recover accurate behaviour from diverse experts as shown in ( Li et al. , 2017 ) , since such learning results in ambiguities when inferring reward functions . We require to learn reward functions that elicit succinct behaviours of agents . Such rewards should avoid adopting behavioural features of experts that might be irrelevant to the task or even detrimental as a solution strategy . Invariant risk minimization ( IRM ) ( Arjovsky et al. , 2019 ) exploits the concept of invariant causal prediction ( Peters et al. , 2015 ; Heinze-Deml et al. , 2017 ) . This recently introduced approach studies the generalization problem of classification models through the lens of causality . The IRM method postulates that the conditional distribution of the class label must be stable across datasets , in order to avoid absorption of spurious correlations into the model predictions . This adaptation principle can be utilized in the context of reward function learning , where we aim to elicit behavioural policies without exploiting spurious reward features , i.e . features which cause nontransferable behaviours . To achieve this goal , we make the assumption that variations in expert demonstrations are a product of causal interventions on the data generating process of trajectories . Our contributions are as follows : • We formulate the assumption that the variations between experts performing optimally on the same task can be seen as interventional settings of the underlying trajectory distribution . • We propose a regularization principle based on invariant risk minimization . This modelling choice allows us to learn reward functions that are invariant to spurious correlations present in the expert data . • We demonstrate the efficacy of this approach in both tractable , finite state-spaces , which we refer to as the exact setting as well as large continuous state-spaces , which we denote as approximate setting . The rest of the paper is structured as follows : we provide an overview of related work in Section 2 and introduce the problem setting from the perspective of IRL as well as causal considerations in Section 3 . We then present our model for the exact and approximate settings in Section 4 and demonstrate the experimental results in Section 5 . We conclude with a brief discussion and outlook in Section 6 . 2 RELATED WORK . IRL and imitation learning Maximum entropy inverse reinforcement learning has been introduced in ( Ziebart et al. , 2008 ) . An extension to functions parameterized by deep neural networks has been presented in ( Wulfmeier et al. , 2015 ) . A formulation which considers large continuous state spaces where an approximation of the partition functions is required has initially been studied in ( Finn et al. , 2016 ) and expanded to an imitation setting in the work by ( Ho & Ermon , 2016 ) . A disentangled version of the adversarial imitation approach which allows to recover rewards has been presented in ( Fu et al. , 2017 ) . The authors of ( Zolna et al. , 2019 ) propose a model which also tackles the issue of spurious correlations being absorbed from expert data . However , their focus is on visual features in a solely imitation learning setting as opposed to our approach , which recovers reward functions that perform favorably in a transfer setting . Invariant representation learning Invariant risk minimization is a method proposed in ( Arjovsky et al. , 2019 ) which builds upon the invariant causal prediction ( Peters et al. , 2015 ; Heinze-Deml et al. , 2017 ) principle and proposes a way to tackle the issue of spurious correlations impeding out-of-distribution generalization of classification algorithms . Various other formulations of the IRM principle have subsequently been proposed . IRM games ( Ahuja et al. , 2020 ) reformulates the IRM objective from a game theoretic perspective and achieves lower variance compared to the original IRM method . Invariant rationalization proposes an adversarial method for recovering invariant rationales in the context of natural language processing ( Chang et al. , 2020 ) by employing an auxiliary model which is trained to be agnostic to interventional settings . Invariance and causality in RL The concept of invariance has been used in a number of works in the reinforcement learning domain . Invariant causal prediction has been utilized in ( Zhang et al. , 2020 ) to learn model invariant state abstractions in a multiple MDP setting with a shared latent space . Invariant policy optimization ( Sonar et al. , 2021 ) uses the IRM games ( Ahuja et al. , 2020 ) formulation to learn policies invariant to certain domain variations . The authors of ( de Haan et al. , 2019 ) tackle the problem of causal confusion in imitation learning by making use of causal structure of demonstrations . ( Angelov et al. , 2020 ) address the problem of user specification in robotic learning through the causal lens by augmenting the images with specific symbols . To the best of our knowledge , our algorithm is the first proposed method to use invariant causal prediction in the context of inverse reinforcement learning . 3 PROBLEM SETTING . MDP We consider environments modelled by a Markov decision processM = ( S , A , T , P0 , R ) , where S is the state space , A is the action space , T is the family of transition distributions on S indexed by S ×A with p ( s′|s , a ) describing the probability of transitioning to state s′ when taking action a in state s , P0 is the initial state distribution , and R : S × A → R is the reward function . A policy π is a conditional probability distribution of actions a ∈ A given states s ∈ S with state descriptors or features ϕ ( s ) . Inverse reinforcement learning aims to estimate a suitable reward function rψ parameterized by weights ψ based on a dataset of expert trajectories DE = { τi } i≤K where τi = ( s ( i ) 1 : T , a ( i ) 1 : T ) is a sequence of states and actions of expert i of length T . To achieve this goal , the feature expectation statistics under the state occupancy measure of the expert , Eτ∼DE [ ϕ ( τ ) ] , are matched with the statistics of the student state occupancy measure Eτ∼p ( τ |ψ ) [ ϕ ( τ ) ] . ϕ ( τ ) = ∑ s∈τ ϕ ( s ) denotes the sum of the state features ϕ ( s ) . The trajectory distribution p ( τ |ψ ) is induced by the policy πrψ trained on the reward function estimate rψ . Given the ill-posed nature of the problem , a maximum entropy formulation of the IRL problem ( Ziebart et al. , 2008 ) is typically chosen in order to find a unique solution . The model describes the fact that the expert trajectories are sampled from the Gibbs distribution p ( τ |ψ , ϕ ) = 1Zϕ , ψ exp ( rψ ( τ ) ) = 1 Zϕ , ψ exp ( ψTϕ ( τ ) ) which corresponds to the solution of the entropy maximization problem under feature matching and simplex constraints . max ψ H ( p ( τ |ψ ) ) s.t . { Eτ∼p ( τ |ψ ) [ ϕ ( τ ) ] = Eτ∼DE [ ϕ ( τ ) ] ; ∫ T p ( τ |ψ ) dτ = 1 ; p ( τ |ψ ) > 0 } ( 1 ) where Eτ∼DE [ ϕ ( τ ) ] is the sample average of the expert trajectories . In the feature matching case , the reward rψ ( τ ) is typically assumed to be linear in the state features rψ = ψTϕ ( τ ) . We further consider the scenario where the state features ϕ ( τ ) are learned using the DeepMaxEnt model ( Wulfmeier et al. , 2015 ) . In large state spaces , the partition function Z presents an estimation challenge and an importance sampling scheme is typically employed . Adversarial IRL formulation It has been shown ( Finn et al. , 2016 ) that the trajectory distribution for the importance sampling scheme can be obtained by solving an adversarial minimax game implemented by a generative adversarial network ( Goodfellow et al. , 2014 ) architecture . Adversarial inverse RL ( AIRL ) methods yield a reward function by learning to distinguish between the transitions sampled from the dataset of expert trajectories τ ∼ DE and transitions continually sampled from the improving policy τ ∼ π ( τ ) , which maximizes an expected cumulative reward objective based on the discriminator output . We adopt the discriminator structure from ( Fu et al. , 2017 ) which separates into the state-action dependent reward function gξ and the state dependent shaping term hφ . Dξ , φ , θ ( s , a , s ′ ) = exp ( gξ ( s , a ) + γhφ ( s ′ ) − hφ ( s ) ) exp ( gξ ( s , a ) + γhφ ( s′ ) − hφ ( s ) ) + πθ ( a|s ) ( 2 ) We also consider a state-only formulation of AIRL where the reward function gξ ( s ) only depends on the state s. The optimization objective of the discriminator is the binary cross-entropy loss on the discriminator output : LBCE ( ξ , φ , θ ) = E ( s , a ) ∼πθ log ( Dξ , φ , θ ( s , a ) ) − E ( s , a ) ∼πE log ( 1−Dξ , φ , θ ( s , a ) ) ( 3 ) The generator corresponds to the policy optimizing the expected cumulative reward given by the discriminator output and is trained using the proximal policy optimization ( PPO ) ( Schulman et al. , 2017 ) method . As with many neural network approximators , the discriminator model absorbs spurious correlations and this learning effect poses a serious problem ( Arjovsky et al. , 2019 ) . These correlations coincide with the binary label information encoding the optimality of the expert . In this work , we avoid this effect by applying the invariant risk minimization principle to the discriminator part of our model . Invariant risk minimization ( Arjovsky et al. , 2019 ) extends the invariant causal prediction ( Peters et al. , 2015 ; Heinze-Deml et al. , 2017 ) principle to nonlinear settings where the underlying SCM might not be explicitly given . The goal of IRM is to learn correlations which are invariant across training settings and ignore spurious correlations . The set of training settings Etr contains datasets De : = { ( x ( i ) e , y ( i ) e ) } |Etr|i=1 sampled from interventional distributions P̃ ( Xe , Ye ) of the SCM of the underlying data generating process . The IRM method effectively learns an invariant data representation ϕ for the predictor w ◦ ϕ ( where w is a linear classifier ) which enables a stable conditional distribution P ( Y e|Xe ) by optimizing the objective stated in the definition below : Definition 1 The data representation ϕ : X → H mapping the input to a hypothesis space , elicits an invariant predictor w ◦ ϕ across the set of interventional settings E if there is a classifier w : H → Y simultaneously optimal for all settings : min ϕ : X→H , w : H→Y ∑ e∈Etr Le ( w ◦ ϕ ) s. t. : w ∈ argmin w̄ : H→Y Le ( w̄ ◦ ϕ ) ∀e ∈ Etr ( 4 ) The aim of this bi-level optimization problem is to minimize the prediction loss on the union of the training sets constrained on selecting the predictor which minimizes the losses on the individual training set e , Le . The tractable approximation of Eq . 4 uses the gradient norm penalty D ( w = 1.0 , ϕ , e ) = ||∇w|w=1.0Le ( w ◦ ϕ ) ||2 which quantifies the violation of the normal equations to measure the optimality of a fixed classifier ( w = 1.0 ) at each setting e. This leads to the following unconstrained formulation of the problem : min ϕ : X→Y ∑ e∈Etr Le ( ϕ ) + λ||∇w|w=1.0Le ( w ◦ ϕ ) ||2 ( 5 ) In the following section , we will show how to incorporate this optimization principle into the IRL setting .

Summary: This work draws inspiration from the invariant risk minimization (IRM) framework to improve the generalization capability of inverse reinforcement learning (IRL) methods. The paper's contributions are deriving a new IRL objective from IRM. This new algorithm includes a penalty regularization term to the standard objective which helps the agent from overfitting to the peculiarities of the expert demonstrations.