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: This paper deals with the problem of learning non-Markov task specifications (or rewards) encoded as a reward machine by inferring the real-value parameters of a reward machine (typically the edge transition rewards in an RM). The authors adopt a Bayesian deep learning approach to simultaneously learn a reward approximator, a distribution over the parameters of the reward machine, and the policy to optimize the most-likely reward instance.   The authors then go on to demonstrate in grid world environments that benefit from non-Markov rewards and policy representations that the inferred RM outperforms standard reward definition using standard RL approaches. And they further demonstrate that the inferred reward machines are beneficial  in larger instances within the same domain.