Learning Two-Step Hybrid Policy for Graph-Based Interpretable Reinforcement Learning

1 INTRODUCTION . Recent years have witnessed the rapid developments of deep reinforcement learning across various domains such as mastering board games ( Silver et al. , 2016 ; Schrittwieser et al. , 2020 ) playing video games ( Mnih et al. , 2015 ) , and chip design ( Mirhoseini et al. , 2021 ) , etc . The larger and complicated architecture of deep reinforcement learning models empowered the capabilities of resolving challenging tasks while brought in significant challenges of interpreting the decision making process of those complex policies ( Puiutta & Veith , 2020 ) . This trade-off between performance and interpretability becomes an inevitable issue when DRL is applied to high stakes applications such as health care ( Rudin , 2019 ) . In this work , we focus on graph-based interpretable reinforcement learning ( Zambaldi et al. , 2018 ; Waradpande et al. , 2020 ) as the graph representation is expressive in various domains including drug discovery ( Patel et al. , 2020 ) , visual question answering ( Hildebrandt et al. , 2020 ) , and embodied AI ( Chaplot et al. , 2020 ; Huang et al. , 2019 ) , etc . Another benefit of studying graph-based RL is that the graph structure can provide natural explanations of the decisionmaking process without the necessity of introducing new programms ( Verma et al. , 2018 ) or heavy domain knowledge ( Bastani et al. , 2018 ) for interpretation . Prior works in interpretable reinforcement learning ( Verma et al. , 2018 ; Madumal et al. , 2020 ; Liu et al. , 2018 ) either works on restricted policy class ( Liu et al. , 2018 ) that leads to downgraded performance , or the interpretablility ( Shu et al. , 2017 ; Zambaldi et al. , 2018 ) is limited . Another common issue of interpretable RL is that the provided explanation is generally difficult to comprehend for non-experts ( Du et al. , 2019 ) . To resolve the challenges mentioned above , we propose a novel two-step hybrid decision-making process for general deep reinforcement learning methods with graph input . Our approach is inspired by the observation of human decision-making . When confront complicated tasks that involve expertise from multiple domains , we typically identify which domain of expert we would like to consult first and then search for specific knowledge or solutions to solve the problem . Recognizing the scope of the problem significantly reduces the search space of downstream tasks , which leads to a more simplified problem compared to find the exact solution in all domains directly . As an analogy of this procedure , we disentangle a complicated deep reinforcement learning policy into a classification the problem for problem type selection and rule-miner . The classification establishes a mapping from complex graph input into an action type , which handles high-order logical interactions among node and edge representations with graph neural network . The rule miner conducts explicit one-hop reasoning over the graph and provides user-friendly selective explanations ( Du et al. , 2019 ) by mining several decisive edges . This two-step decision making is essential not only for providing interpretability , but also for generalization and robustness . It is intuitive to see that the simplified classification is easier to achieve better generalization and robustness than the original complicated RL policy . Furthermore , the rule-miner identifying key edges in the graph is much more robust to the noisy perturbations on the irrelevant graph components . In summary , our contributions are three folds : • We formalize an interpretable deep reinforcement learning problem based on graph input . • We propose a two-step decision-making framework that achieves far better performance in terms of generalization and robustness and provides human-friendly interpretations . • Extensive experiments on several text-based games ( Côté et al. , 2018 ) demonstrated that the proposed approach achieves a new state-of-the-art performance for both generalization and robustness . 2 RELATED WORKS . 2.1 INTERPRETABLE RL . Instead of using deep RL as a black box , researchers also worked on making deep RL more interpertable . ( Mott et al. , 2019 ; Zambaldi et al. , 2018 ) introduce some sorts of attention mechanisms into policy networks and explain decision making process by analyzing attention weights . ( Verma et al. , 2018 ) combines program synthesis and PIDs in classic control to solve continuous control problems , while keep the decision making interpretable by explicitly writing the program . Combining symbolic planning and deep RL has similar effects , like ( Lyu et al. , 2019 ) . Tree-based policies ( Bastani et al. , 2018 ) are favorable in interpretable RL , since they are more human-readable and easy to verify . However , many existing works of interpretable RL sacrifices the performance for the interpretability . In contrast , our method achieves even better generalization and robustness while providing interpretability . 2.2 HIERARCHICAL RL . Hierarchical RL focuses on decomposing a long-horizon tasks into several sub-tasks , and applying a policy with hierarchical structure to solve the task . ( Vezhnevets et al. , 2017 ; Frans et al. , 2018 ) try to solve generic long-horizon tasks with a two-level policy , where a high-level policy will generate a latent sub-goal for low-level policies , or directly select a low-level policy , and the low-level policy is responsible for generating the final action . In some robot locomotion tasks , people usually manually define the choices of sub-goals to boost the performance and lower the learning difficulty , like ( Nachum et al. , 2018 ; Levy et al. , 2018 ) . Our two-step hybrid policy is reminiscent of the two-level policy architecture in hierarchical RL . However , the high-level policy in hiereachical RL is mainly used to generate a sub-goal while the action pruner in our hybrid policy is used to reduce the number of action candidates . And the motivations of our method and hierarchical RL are also different : hierarchical RL targets at decomposing a long-horizon tasks into several pieces , but our method focuses on interpertability , generalizability and robustness . 2.3 REFACTORING POLICY . Given a trained neural network-based policies , one may want to refactor the policy into another architecture , so as to improve the generalizibility or interpretablity . And this refactoring processs is usually done by imitation learning . The desired architectures of the new policy can be decision tree ( Bastani et al. , 2018 ) , symbolic policy ( Landajuela et al. , 2021 ) , a mixture of program and neural networks ( Sun et al. , 2018 ) , or graph neural networks ( Mu et al. , 2020 ) . In our method , we also use similar techniques to refactor a reinforcement learning policy into a two-step hybrid policy . 2.4 COMBINE NEURAL NETWORKS WITH RULE-BASED MODELS . Neural networks and rule-based models excel at different aspects , and many works have explored how to bring them together , like ( Chiu et al. , 1997 ; Goodman et al. , 1990 ; Ray & Chakrabarti , 2020 ; Okajima & Sadamasa , 2019 ; Greenspan et al. , 1992 ) . Therefore various ways to combine neural networks with rule-based models . For example , ( Wang , 2019 ) combines them in a ” horizontal ” way , i.e. , some of the data are assigned to rule-based models while the others are handled by neural networks . In contrast , our work combines neural networks and rule-based models in a ” vertical ” way , i.e. , the two models collaborate to make decisions in a two-step manner . 3 A GENERAL FRAMEWORK FOR TWO-STEP HYBRID DECISION MAKING . In this section , we first describe our problem setting including key assumptions and a general framework that formulates the decision-making process in a two-step manner . We consider a discrete time Markov Decision Process ( MDP ) with a discrete state space S and the action spaceA is finite . The environment dynamics is denoted as P = { p ( s′|s , a ) , ∀s , s′ ∈ S , a ∈ A } . Given an action set A and the current state st ∈ S , our goal is to learn a policy ( π ) that select an action at ∈ A to maximize the long-term reward Eπ [ ∑T i=1 r ( si , ai ) ] . Assume we are able to group the actions into several mutual exclusive action types ( Ak ) according to its semantic meanings . More concretely , the k-th action type Ak = { a1k , ... , ank } denotes a subset of actions in original action space Ak ⊆ A . Then we have A1 , A2 , .. , AK ⊆ A , Ai ∩ Aj = ∅ ( i 6= j ) , ∪Ki=1Ai = A . It is worth noting that the number of action type K is usually more than an order of magnitude smaller than original actions ( K |A| ) . Let the policy π = 〈fp , fs〉 represented by a hybrid model that consists of action pruner fp and action selector fs . The action pruner is used to prune all available action candidates to a single action type , i.e. , fp ( si ) = k , where k is the index of chosen action type and Ak ∈ { A1 , A2 , .. , AK } . Then the action selector is used to select a specific action given the action type chosen by the action pruner , i.e. , fs ( si , Ak ) = ai , where ai ∈ Ak and k = fp ( si ) . Intuitively , this design intends to disentangle different phases in decision-making process to two different modules . On one hand , determining the action type typically involves high-order relationships and needs a model with strong expressive power , then neural network is a good candidate in this regard . On the other hand , selecting an action within a specific action type can resort to rule-based approaches , which is essential in providing strong intrepretability , generalizability and robustness . Figure 1 shows the overall pipeline of our two-step hybrid decision making . The agent will receive a state si at each time step i . At time step i , we will first call the action pruner fp ( si ) to select the action type Ai . Then the rule-based action selector fs ( si , Ai ) will take as inputs the current state si and the action type Ai given by action pruner to select the specific action to be executed in this step . 4 TWO-STEP HYBRID POLICY FOR TEXT-BASED GAMES WITH GRAPH INPUT . 4.1 OVERVIEW . In this section , we instantiate our proposed framework in the setting of text-based games ( Côté et al. , 2018 ) with graph input . In text-based games , the agent receives a knowledge graph ( as shown in Figure 3 ) that describes the current state of the environment and the task to be completed , e.g. , the task can be represented by several edges like ( “ potato ” , “ needs ” , “ diced ” ) . Our goal is to learn a policy that maps the input knowledge graph to an action from the provided action set A . Each action aj ∈ A is a short sentence , e.g. , “ take apple from fridge ” . In this setting , we use graph neural networks ( GNNs ) as the action pruner fp , and a rule-based model as the action selector fs . We will elaborate the details of the action pruner and the action selector in Sec 4.3 and Sec 4.4 , respectively . Training the GNN-based action pruner and the rule-based action selector by reinforcement learning is nontrivial since the whole pipeline is not end-to-end differential . Therefore , we propose to learn both models separately from a demonstration dataset , and the demonstration dataset can be obtained by a trained reinforcement learning agent . We will elaborate the details in Sec 4.2.1 . This process is inspired by existing works like ( Bastani et al. , 2018 ; Landajuela et al. , 2021 ; Sun et al. , 2018 ; Mu et al. , 2020 ) , where policies trained by reinforcement learning can be refactorized into other forms .

Summary: This paper introduces a two-step method for RL problems with graph-based inputs. This method trains a teacher policy which it then uses to construct a policy "pruner" which predicts the kind of action needed. It then uses a ideas from retrieval to select the edges needed to complete the action.