Learning Finite State Representations of Recurrent Policy Networks

1 INTRODUCTION . Deep reinforcement learning ( RL ) and imitation learning ( IL ) have demonstrated impressive performance across a wide range of applications . Unfortunately , the learned policies are difficult to understand and explain , which limits the degree that they can be trusted and used in high-stakes applications . Such explanations are particularly problematic for policies represented as recurrent neural networks ( RNNs ) ( Mnih et al. , 2016 ; Mikolov et al. , 2010 ) , which are increasingly used to achieve state-of-the-art performance ( Mnih et al. , 2015 ; Silver et al. , 2017 ) . This is because RNN policies use internal memory to encode features of the observation history , which are critical to their decision making , but extremely difficult to interpret . In this paper , we take a step towards comprehending and explaining RNN policies by learning more compact memory representations . Explaining RNN memory is challenging due to the typical use of high-dimensional continuous memory vectors that are updated through complex gating networks ( e.g . LSTMs , GRUs ( Hochreiter & Schmidhuber , 1997 ; Chung et al. , 2014 ; Cho et al. , 2014 ) ) . We hypothesize that , in many cases , the continuous memory is capturing and updating one or more discrete concepts . If exposed , such concepts could significantly aid explainability . This motivates attempting to quantize the memory and observation representation used by an RNN to more directly capture those concepts . In this case , understanding the memory use can be approached by manipulating and analyzing the quantized system . Of course , not all RNN policies will have compact quantized representations , but many powerful forms of memory usage can be captured in this way . Our main contribution is to introduce an approach for transforming an RNN policy with continuous memory and continuous observations to a finite-state representation known as a Moore Machine . To accomplish this we introduce the idea of Quantized Bottleneck Network ( QBN ) insertion . QBNs are simply auto-encoders , where the latent representation is quantized . Given a trained RNN , we train QBNs to encode the memory states and observation vectors that are encountered during the RNN operation . We then insert the QBNs into the trained RNN policy in place of the “ wires ” that propagated the memory and observation vectors . The combination of the RNN and QBN results in a ∗Work done at Oregon State University policy represented as a Moore Machine Network ( MMN ) with quantized memory and observations that is nearly equivalent to the original RNN . The MMN can be used directly or fine-tuned to improve on inaccuracies introduced by QBN insertion . While training quantized networks is often considered to be quite challenging , we show that a simple approach works well in the case of QBNs . In particular , we demonstrate that “ straight through ” gradient estimators as in ( Bengio et al. , 2013 ; Courbariaux et al. , 2016 ) are quite effective . We present experiments in synthetic domains designed to exercise different types of memory use as well as benchmark grammar learning problems . Our approach is able to accurately extract the ground-truth MMNs , providing insight into the RNN memory use . We also did experiments on 6 Atari games using RNNs that achieve state-of-the-art performance . We show that in most cases it is possible to extract near-equivalent MMNs and that the MMNs can be surprisingly small . Further , the extracted MMNs give insights into the memory usage that are not obvious based on just observing the RNN policy in action . For example , we identify games where the RNNs do not use memory in a meaningful way , indicating the RNN is implementing purely reactive control . In contrast , in other games , the RNN does not use observations in a meaningful way , which indicates that the RNN is implementing an open-loop controller . 2 RELATED WORK . There have been efforts made in the past to understand the internals of Recurrent Networks ( Karpathy et al. , 2015 ; Arras et al. , 2017 ; Strobelt et al. , 2016 ; Murdoch & Szlam , 2017 ; Jacobsson , 2005 ; Omlin & Giles , 1996 ) . However , to the best of our knowledge there is no prior work on learning finite-memory representations of continuous RNN policies . Our work , however , is related to a large body of work on learning finite-state representations of recurrent neural networks . Below we summarize the branches of that work and the relationship to our own . There has been a significant history of work on extracting Finite State Machines ( FSMs ) from recurrent networks trained to recognize languages ( Zeng et al. , 1993 ; Tiňo et al. , 1998 ; Cechin et al. , 2003 ) . Typical approaches include discretizing the continuous memory space via gridding or clustering followed by minimization . A more recent approach is to use classic query-based learning algorithms to extract FSMs by asking membership and equivalence queries ( Weiss et al. , 2017 ) . However , none of these approaches directly apply to learning policies , which require extending to Moore Machines . In addition , all of these approaches produce an FSM approximation that is separated from the RNN and thus serve as only a proxy of the RNN behavior . Rather , our approach directly inserts discrete elements into the RNN that preserves its behavior , but allows for a finite state characterization . This insertion approach has the advantage of allowing fine-tuning and visualization using standard learning frameworks . The work most similar to ours also focused on learning FSMs ( Zeng et al. , 1993 ) . However , the approach is based on directly learning recurrent networks with finite memory , which are qualitatively similar to the memory representation of our MMNs . That work , however , focused on learning from scratch rather than aiming to describe the behavior of a continuous RNN . Our work extends that approach to learn MMNs and more importantly introduces the method of QBN insertion as a way of learning via guidance from a continuous RNN.This transforms any pre-trained recurrent policy into a finite representation . We note that there has been prior work on learning fully binary networks , where the activation functions and/or weights are binary ( e.g . ( Bengio et al. , 2013 ; Courbariaux et al. , 2016 ; Hinton , 2012 ) ) . The goal of that line of work is typically to learn more time and space efficient networks . Rather , we focus on learning only discrete representations of memory and observations , while allowing the rest of the network to use arbitrary activations and weights . This is due to our alternative goal of supporting interpretability rather than efficiency . 3 RECURRENT POLICY NETWORKS : CONTINUOUS AND QUANTIZED . Recurrent neural networks ( RNNs ) are commonly used in reinforcement learning to represent policies that require or can benefit from internal memory . At each time step , an RNN is given an observation ot ( e.g . image ) and must output an action at to be taken in the environment . During execution an RNN maintains a continuous-valued hidden state ht , which is updated on each transition and influences the action choice . In particular , given the current observation ot and current state ht , an RNN performs the following operations : 1 ) Extract a set of observation features ft from ot , for example , using a CNN when observations are images , 2 ) Outputting an action at = π ( ht ) according to policy π , which is often a linear softmax function of ht , 3 ) transition to a new state ht+1 = δ ( ft , ht ) where δ is the transition function , which is often implemented via different types of gating networks such as LSTMs or GRUs . The continuous and high dimensional nature of ht and ft can make interpreting the role of memory difficult . This motivates our goal of extracting compact quantized representations of ht and ft . Such representations have the potential to allow investigating a finite system that captures the key features of the memory and observations . For this purpose we introduce Moore Machines and their deep network counterparts . Moore Machines . A classical Moore Machine ( MM ) is a standard finite state machine where all states are labeled by output values , which in our case will correspond to actions . In particular , a Moore Machine is described by a finite set of ( hidden ) states Ĥ , an initial hidden state ĥ0 , a finite set of observations Ô , a finite set of actions A , a transition function δ̂ , and a policy π̂ that maps hidden states to actions . The transition function δ̂ : Ĥ × Ô → Ĥ returns the next hidden state ĥt+1 = δ̂ ( ĥt , ôt ) given the current state ĥt and observation ôt . By convention we will use ht and ĥt to denote continuous and discrete states respectively and similarly for other quantities and functions . Moore Machine Networks . A Moore Machine Network ( MMN ) is a Moore Machine where the transition function δ̂ and policy π̂ are represented via deep networks . In addition , since the raw observations given to an MMN are often continuous , or from an effectively unbounded set ( e.g . images ) , an MMN will also provide a mapping ĝ from the continuous observations to a finite discrete observation space Ô . Here ĝ will also be represented via a deep network . In this work , we consider quantized state and observation representations where each ĥ ∈ Ĥ is a discrete vector and each discrete observation in Ô is a discrete vector that describes the raw observation . We will denote the quantization level as k and the dimensions of ĥ and f̂ by Bh and Bf respectively . Based on the above discussion , an MMN can be viewed as a traditional RNN , where : 1 ) The memory is restricted to be composed of k-level activation units , and 2 ) The environmental observations are intermediately transformed to a k-level representation f̂ before being fed to the recurrent module . Given an approach for incorporating quantized units into the backpropagation process , it is straightforward , in concept , to learn MMNs from scratch via standard RNN learning algorithms . However , we have found that learning MMNs from scratch can be quite difficult for non-trivial problems , even when an RNN can be learned with relative ease . For example , we have not been able to train highperforming MMNs from scratch for Atari games . Below we introduce a new approach for learning MMNs that is able to leverage the ability to learn RNNs . 4 LEARNING MOORE MACHINE NETWORKS FROM RNNS . Given a trained RNN , our key idea is to first learn quantized bottleneck networks ( QBNs ) for embedding the continuous observation features and hidden state into a k-level quantized representation . We will then insert the QBNs into the original recurrent net in such a way that its behavior is minimally changed with the option of fine-tuning after insertion . The resulting network can be viewed as consuming quantized features and maintaining quantized state , which is effectively an MMN . Below we describe the steps in further detail , which are illustrated in Figure 1 . 4.1 QUANTIZED BOTTLENECK NETWORKS . A QBN is simply an autoencoder where the latent representation between the encoder and decoder ( i.e . the bottleneck ) is constrained to be composed of k-level activation units . While , traditional autoencoders are generally used for the purpose of dimensionality reduction in continuous space ( Hinton & Salakhutdinov , 2006 ) , QBNs are motivated by the goal of discretizing a continuous space . Conceptually , this can be done by quantizing the activations of units in the encoding layer . We represent a QBN via a continuous multilayer encoder E , which maps inputs x to a latent encoding E ( x ) , and a corresponding multilayer decoder D. To quantize the encoding , the QBN output is given by b ( x ) = D ( quantize ( E ( x ) ) ) In our case , we use 3-level quantization in the form of +1 , 0 and −1 using the quantize function , which assumes the outputs of E ( x ) are in the range [ −1 , 1 ] .1 One choice for the output nodes of E ( x ) would be the tanh activation . However , since the gradient of tanh is close to 1 near 0 , it can be difficult to produce quantization level 0 during learning . Thus , as suggested in Pitis ( 2017 ) , to support 3-valued quantization we use the following activation function , which is flatter in the region around zero input . φ ( x ) = 1.5 tanh ( x ) + 0.5 tanh ( −3x ) Of course introducing the quantize function in the QBN results in b ( x ) being non-differentiable , making it apparently incompatible with backpropagation , since the gradients between the decoder and encoder will almost always be zero . While there are a variety of ways to deal with this issue , we have found that the straight-through estimator , as suggested and used in prior work ( Hinton , 2012 ; Bengio et al. , 2013 ; Courbariaux et al. , 2016 ) is quite effective . In particular , the standard straightthrough estimator of the gradient simply treats the quantize function as the identity function during back-propagation . Overall , the inclusion of the quantize function in the QBN effectively allows us to view the last layer ofE as producing a k-level encoding . We train a QBN as an autoencoder using the standard L2 reconstruction error ‖x− b ( x ) ‖2 for a given input x .

Summary: RNNs are difficult to explain, understand and analyze due to the continuous-valued memory vectors and observations features they use. Thus, this paper attempts to extract finite representation from RNNs so as to better interpret or understand RNNs. They introduce a new technique called Quantized Bottleneck Insertion to extract Moore Machines (MM). The extracted MM can be analyzed to improve the understanding of memory use and general behavior on the policies. The experiments on synthetic datasets and six Atari games validate the effectiveness of the proposal.