A Hierarchical Bayesian Approach to Inverse Reinforcement Learning with Symbolic Reward Machines

1 INTRODUCTION . Reinforcement Learning ( RL ) agents rely on rewards to measure the utility of each interaction with the environment . As the complexity of RL tasks increases , it becomes difficult for the agent to grasp the intricacies of the task solely from goal-driven reward functions – rewarding the agent only at the end of each episode . Reward machine ( RM ) is a formalism proposed by Icarte et al . ( 2020 ) for representing a reward function as a finite state automaton ( FSA ) . However , the design of RMs can quickly become cumbersome as the complexity of the tasks increases . In this paper , we draw inspiration from symbolic finite automaton ( SFA ) and symbolic finite transducer ( SFT ) Antoni & Veanes ( 2017 ) and propose symbolic reward machine ( SRM ) which augment FSA-based RMs by allowing the internal state transitions of the RM to carry predicates and functions on the trajectory . In addition to improving interpretability and conciseness of the reward model , a salient benefit of SRM is that it allows human engineers to express complex task scenarios and reward design patterns . In many existing works on logically designed reward functions , the reward values are manually assigned based on heuristics which can undermine the effectiveness of the resulting reward functions . For instance , if a learning agent is excessively awarded for the completion of a certain stage of a task , the agent may tend to repeatedly complete the same stage to accumulate rewards instead of commencing the next stage , a phenomenon known as reward hacking Amodei et al . ( 2016 ) . We envision that , in a typical design routine of an SRM , a human engineer constructs the SRM to incorporate high-level task information , but leave the low-level details , such as the right amount of reward for a specific event or sub-task , empty or as holes . The SRM formalism also facilitates the specification of symbolic constraints over the holes for capturing important task-specific nuances . A major contribution of this paper is a learning framework for filling out the holes in an SRM , which we call concretization . The input to the framework includes a set of example trajectories demonstrated by a human expert , an SRM and optionally a symbolic constraint that the holes in the SRM must satisfy . We leverage the generative adversarial approaches from Finn et al . ( 2016 ) ; Jeon et al . ( 2018 ) to construct a discriminator with a neural network reward function to distinguish the expert trajectories from the trajectories of an agent policy . However , our framework works in a hierarchical Bayesian manner . Basically , to circumvent the non-differentiability of SRMs , we employ a sampler to sample candidate instantiations of the holes to concretize the SRM . Then we introduce a stochastic reward signal as the latent variable dependent on the output of the concretized SRM and employ a neural-network reward function to perform importance sampling of the stochastic rewards for trajectory discrimination . We summarize our contributions below . • We propose to use SRMs to represent reward functions for RL tasks . • We develop a hierarchical Bayesian inference framework that can concretize an SRM by learning from expert demonstrations . • Our approach enables RL agents to achieve state-of-the-art performance on highly complex environments with only a few demonstrations . In addition , we show that SRMs generalize well across different environment configurations of the same task . 2 RELATED WORK . Inverse Reinforcement Learning . We first note that the IRL formulation proposed in Ng & Russell ( 2000 ) ; Abbeel & Ng ( 2004 ) has an infinite number of solutions . The Max-Entropy IRL from Ziebart et al . ( 2008 ) , Max-Margin IRL from Abbeel & Ng ( 2004 ) ; Ratliff et al . ( 2006 ) and Bayesian IRL from Ramachandran & Amir ( 2007 ) aim at resolving the ambiguity of IRL . However , those approaches restrict the reward function to be linear on the basis of human designed feature functions . Deep learning approaches proposed in Fu et al . ( 2018 ) ; Ho & Ermon ( 2016 ) ; Jeon et al . ( 2018 ) ; Finn et al . ( 2016 ) have substantially improved the scalablity of IRL by drawing a connection between IRL and Generative Adversarial Networks ( GANs ) introduced by Goodfellow et al . ( 2014 ) . Our work , while embracing the data-driven and generative-adversarial ideologies , further extends IRL to cope with symbolically represented human knowledge . Reward Design . There have been substantial efforts on enriching the information in reward functions . Reward shaping proposed by Ng et al . ( 1999 ) adds state-based potentials to the reward in each state . Exploration driven approaches such as Bellemare et al . ( 2016 ) ; Pathak et al . ( 2017 ) ; Alshiekh et al . ( 2017 ) ; Flet-Berliac et al . ( 2021 ) incentivize agents with intrinsic rewards . Compared with these methods , we do not seek to generate reward functions densely ranging over the entire state space but rather design intepretable ones that selectively or even sparsely produce non-zero rewards . Reward machines from Icarte et al . ( 2020 ) directly represent the reward functions as FSAs . The symbolic reward machine in our work is also automata-based but augments RMs in a similar way to SFA for FSA Veanes et al . ( 2012 ) . There have been efforts on learning a so-called perfect RM as termed in Toro Icarte et al . ( 2019 ) from the experience of an RL agent in partially observable environment . However , the RM is still based on FSA and the rewards are still manually assigned . Regarding leveraging human demonstrations , inverse reward design ( IRD ) proposed in HadfieldMenell et al . ( 2017 ) is analogous to IRL but aims at inferring a true reward function from some proxy reward function perceived by a RL agent . Safety-aware apprenticeship learning from Zhou & Li ( 2018 ) pioneers the incorporation of formal verification in IRL . However , those works confine the reward functions to be linear of features as the generic IRL does . Our work does not have such limitations . Interpretable Reinforcement Learning . There have been continual researches on designing interpretable policies Andre & Russell ( 2001 ; 2002 ) ; Verma et al . ( 2018 ) ; Zhu et al . ( 2019 ) ; Yang et al . ( 2021 ) ; Tian et al . ( 2020 ) . This paper concerns the design of interpretable reward functions rather than interpretable policies . Our motivation is that a well-designed reward function is transferable and can be a powerful complement to the vast literature on RL policy learning . 3 BACKGROUND . An RL environment is a tupleM = 〈S , A , P , d0〉 where S is the state space ; A is an action space ; P ( s′|s , a ) is the probability of reaching a state s′ by performing an action a at a state s ; d0 is an initial state distribution . A policy π ( a|s ) determines the probability of an RL agent performing an action a at state s. By successively performing actions for T steps after initializing from a state s ( 0 ) ∼ d0 , a trajectory τ = s ( 0 ) a ( 0 ) s ( 1 ) a ( 1 ) . . . s ( T ) a ( T ) is produced . A state-action based reward function is a mapping f : S × A → R to the real space . With a slight abuse of notations , we denote the total reward along a trajectory τ as f ( τ ) = ∑T t=0 f ( s ( t ) , a ( t ) ) and similarly for the joint probability p ( τ |π ) = ∏T−1 t=0 P ( s ( t+1 ) |s ( t ) , a ( t ) ) π ( a ( t ) |s ( t ) ) of generating a trajectory τ by following π . The objective of an entropy-regularized RL task is to maximize JRL ( π ) = Eτ∼π [ f ( τ ) ] +H ( π ) where τ ∼ π is an abbreviation for τ ∼ p ( τ |π ) andH ( π ) is the expected entropy of π . When the reward function is unknown but a set of expert trajectories τE is sampled with some expert policy πE , GAIL Ho & Ermon ( 2016 ) trains an agent policy πA to match πE by minimizing Jadv in Eq.2 via RL algorithms such as PPO Schulman et al . ( 2017 ) . Adversarially , GAIL optimizes a discriminator D : S × A → [ 0 , 1 ] to accurately identify τE ’ s from τA ∼ πA by maximizing Jadv . From a probabilistic inference perspective , Bayesian GAIL from Jeon et al . ( 2018 ) labels any expert trajectory τE with 1E and 0E to respectively indicate τE as being sampled from an expert demonstration set E and from some agent policy πA . Likewise , the trajectory τA of πA is labeled with 1A and 0A for the same indications . Assuming that the labels 0A , 1E are known a priori , Bayesian GAIL solves the most likely discriminator D by maximizing p ( D|0A , 1E ; πA ; E ) ∝ p ( D ) p ( 0A , 1E |πA , D ; E ) ∝ ∑ τA p ( τA|πA ) p ( 0A|τA ; D ) ∑ τE p ( τE |E ) p ( 1E |τE ; D ) of which the logarithm as in Eq.1 is lower-bounded due to Jensen ’ s inequality by Eq.2 . It is further proposed in Fu et al . ( 2018 ) that by representing D ( s , a ) = exp ( f ( s , a ) ) exp ( f ( s , a ) ) +πA ( a|s ) with a neural network f , when Eq.2 is maximized , it holds that f ≡ log πE and f equals the expert reward function which πE is optimal w.r.t , given that ∀τ.p ( τ |E ) ≈ p ( τ |πE ) . Hence , by representing D in Eq.1 and 2 with f , an objective of solving the most likely expert reward function f is obtained . log ∑ τA p ( τA|πA ) p ( 0A|τA ; D ) ∑ τE p ( τE |E ) p ( 1E |τE ; D ) ( 1 ) ≥ E τE∼E [ log T∏ t=0 D ( s ( t ) E , a ( t ) E ) ] + E τA∼πA [ log T∏ t=0 ( 1−D ( s ( t ) A , a ( t ) A ) ) ] : = Jadv ( D ) ( 2 ) 4 SYMBOLIC REWARD MACHINES ( SRMS ) . We first give a formal definition of SRM and motivate the use of SRMs with a task from a Mini-Grid environment introduced in Chevalier-Boisvert et al . ( 2018 ) . We highlight that the human insights incorporated in the SRM can hardly be realized with conventional goal-driven reward mappings or function approximations . We then formulate the problem of concretizing SRMs . 4.1 DEFINITION . The definition of SRM is inspired from those of SFA and SFT in Veanes et al . ( 2012 ) . To adapt them to the RL setting , we assume a background theory equipped with fixed interpretations on the state and actions in the RL environment as well as a language of functions . Following Pierce ( 2002 ) , a λ-term is a function written in the form of λx.ρ . The type of a λx.ρ is a mapping from the type of input argument x to the type of function body ρ . The free variable set FV ( ρ ) of ρ is the set of symbolically denoted variables appearing in ρ while not appearing in the input of any λ-term inside ρ . Given some concrete input x̂ , the evaluation of λx.ρ is written as [ [ λx.ρ ] ] ( x̂ ) or [ [ ρ [ x̂/x ] ] ] where [ x̂/x ] represents the replacement of x with x̂ in ρ . The denotation stays the same if x and x̂ are not unary . A predicate is a specific set of λ-terms mapping to Boolean type B = { > , ⊥ } where > , ⊥ mean True and False respectively . The set of predicates is closed under Boolean operations ∧ , ∨ , ¬ . Definition 1 . Given an RL environmentM = 〈S , A , P , d0〉 , a symbolic reward machine ( SRM ) is a tuple L = 〈Q , Ψ , R , δ , q0 , Acc〉 where Q is a set of internal states ; Ψ ⊆ ( S ×A ) ∗ → B is a set of predicates on trajectories inM ; R is a set of λ-terms of type ( S ×A ) ∗ → R ; δ is a set of transition rules ( p , ψ , r , q ) , where p , q ∈ Q , ψ ∈ Ψ , r ∈ R ; q0 ∈ Q is an initial state ; Acc ⊆ Q is a set of accepting states ; the free variables set of L is defined as FV ( L ) = ⋃ ρ∈R∪Ψ FV ( ρ ) . We use the notation p ψ//r−−−→ q for a rule ( p , ψ , r , q ) ∈ δ and call ψ its guard . The input to an SRM is a trajectory τ ∈ ( S ×A ) ∗ . A rule p ψ//r−−−→ q is applicable at p iff [ [ ψ ] ] ( τ ) = > , in which case L outputs a reward [ [ r ] ] ( τ ) for the last state-action pair in τ while the state p transitions to q . If no rule is applicable , the state p does not transition and L outputs a fixed constant reward such as 0 , in which case we dub a dummy transition p ∧ ( ¬ψ ) ∃q∈Q . ( p , ψ , r , q ) ∈δ //0 −−−−−−−−−−−−→ p. An SRM L is called deterministic if given an input trajectory τ at any state p ∈ Q , | { ( p , ψ , r , q ) ∈ δ|ψ ( τ ) = > } | ≤ 1 in which case with a little abuse of notations we write the next state as q = δ ( p , τ ) either obtained from the uniquely applicable transition rule ( p , ψ , r , q ) ∈ δ s.t . ψ ( τ ) = > , or from a dummy transition such that q = p. To deploy a deterministic SRM in an RL task is to construct a synchronous product as defined below . Definition 2 . A synchronous product between an M and a deterministic L is a tuple M⊗ L = 〈 ( S ×A ) ∗ × S ×Q , A , Ψ , R , P δ , d0 , q0 , Acc〉 where a product state in ( S ×A ) ∗ × S ×Q is a pair ( τ : : s , q ) where : : means concatenation ; τ : : s means concatenating a trajectory τ ∈ ( S×A ) ∗ with anM state s ∈ S ; q ∈ Q is an L state ; the product transition rule P δ follows Eq.3 where τ : : s : : a is a trajectory resulted from further concatenating τ : : s with an action a ∈ A ; the initial product state is ( s0 , q0 ) where s0 ∼ d0 ; the rest follows the definitions inM and L. ( P δ ) ( ( τ : : s , p ) , a , ( τ : : s : : a : : s′ , q ) ) = { P ( s′|s , a ) if q = δ ( p , τ : : s : : a ) 0 Otherwise ( 3 ) Proposition 1 . Suppose that L is deterministic . Starting from any p ( 0 ) ∈ Q , as the input trajectory extends from τ = s ( 0 ) a ( 0 ) to τ = s ( 0 ) a ( 0 ) . . . s ( T ) a ( T ) , there can only be at most one path σ = p ( 0 ) p ( 1 ) . . . q ( T+1 ) obtained by applying the uniquely applicable or a dummy transition successively to τ from t = 0 to t = T , i.e. , ∀t = 0 , 1 , . . . , T.p ( t+1 ) = δ ( q ( t ) , s ( 0 ) a ( 0 ) . . . s ( t ) a ( t ) ) . Suppose that at time step t there is a transition from ( τ : : s ( t ) , p ( t ) ) to ( τ : : s ( t ) : : a ( t ) : : s ( t+1 ) , p ( t+1 ) ) with ( P δ ) ( ( τ : : s ( t ) , p ( t ) ) , a ( t ) , ( τ : : s ( t ) : : a ( t ) : : s ( t+1 ) , p ( t+1 ) ) ) > 0 . While L witnesses such product state transition , the RL agent only observes the state transition in M. Furthermore , L returns a reward [ [ r ( t ) ] ] ( τ : : s ( t ) : : a ( t ) ) at time step t , where r ( t ) is either 0 or the λ-term associated with the rule p ( t ) ψ ( t ) //r ( t ) −−−−−−→L p ( t+1 ) . We further inductively define [ [ L ] ] ( τ ) as [ [ L ] ] ( τ : : s ( t ) : : a ( t ) ) = [ [ L ] ] ( τ ) : : [ [ r ( t ) ] ] ( τ : : s ( t ) : : a ( t ) ) . For simplicity , we also write τ : : s : : a as τ : : ( s , a ) as if τ is a list of ( s , a ) ’ s , if it does not raise ambiguity in the context . We denote L ( τ ) = ∑T−1 t=0 [ [ L ] ] ( τ [ : t ] ) where τ [ : t ] is the partial trajectory from initialization up until step t. For clarification , we remark that an SRM is an SFT only under stricter conditions , in which case the SRM retains all the properties of SFT , e.g. , composability and decidability . While SFT emphasizes the acceptance of inputs , SRM emphasizes more on computing the rewards for the trajectories .

Summary: The paper presents a method for inverse reinforcement learning for learning the non-Markovian (trajectory-based) rewards for symbolic transitions of a pre-speciefied (up to free variables that need to be learned) automaton. For this purpose, Symbolic Reward Machines (SRM) are introduced, which define symbolic states, transitions and rewards, whereby MDP-trajectories can be unambiguously mapped to SRM-trajectories. The method uses adversarial imitation learning, where the rewards that are computed by the SRM (for a given grounding of the free variables) are used by a discriminator to classify expert- and agent-trajectories. Following the framework of Bayesian GAIL, a distribution over groundings/rewards is learned. However, whereas Bayesian GAIL learns a particle approximation to the posterior, the submission learns a variational Gaussian distribution. The method is evaluted on (rather challenging) gridworlds, where it is able to learn rewards for concepts such as "picking up a key after finding the goal".