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: The authors' propose a regularization loss for inverse reinforcement learning (IRL) based on invariant risk minimisation (IRM). They argue that existing IRL methods suffer from overfitting due to the sparsity of expert trajectories, and that this is analogous to the dataset level overfitting tackled by IRM. This approach is validated by comparing against unregularized and L2 regularized maximum entropy IRL on a grid-world, and is used in a scalable adversarial IRL algorithm when dealing with more complex environments. They close by evaluating the reward function learnt by this regularized adversarial IRL in terms of its utility in training a new policy from scratch.