Action-Sufficient State Representation Learning for Control with Structural Constraints

1 INTRODUCTION . State-of-the-art reinforcement learning ( RL ) algorithms leveraging deep neural networks are usually data hungry and lack interpretability . For example , to attain expert-level performance on tasks such as chess or Atari games , deep RL systems usually require many orders of magnitude more training data than human experts ( Tsividis et al. , 2017 ) . One of the reasons is that our perceived signals in real-world scenarios , e.g. , images , are usually high-dimensional and may contain much irrelevant information for decision-making of the task at hand . This makes it difficult and expensive for an agent to directly learn optimal policies from raw observational data . Fortunately , the underlying states that directly guide decision-making could be much lower-dimensional ( Schölkopf , 2019 ; Bengio , 2019 ) . One example is that when crossing the street , our decision on when to cross relies on the traffic lights . The useful state of traffic lights ( e.g. , its color ) can be represented by a single binary variable , while the perceived image is high-dimensional . It is essential to extract and exploit such lower-dimensional states to improve the efficiency and interpretability of the decision-making process . Recently , representation learning algorithms have been designed to learn abstract features from high-dimensional and noisy observations . Exploiting the abstract representations , instead of the raw data , has been shown to perform subsequent decision-making more efficiently ( Lesort et al. , 2018 ) . Representative methods along this line include deep Kalman filters ( Krishnan et al. , 2015 ) , deep variational Bayes filters ( Karl et al. , 2016 ) , world models ( Ha & Schmidhuber , 2018 ) , PlaNet ( Hafner et al. , 2018 ) , DeepMDP ( Gelada et al. , 2019 ) , stochastic latent actor-critic ( Lee et al. , 2019 ) , SimPLe ( Kaiser et al. , 2019 ) , Bisimulation-based methods ( Zhang et al. , 2021 ) , Dreamer ( Hafner et al. , 2019 ; 2020 ) , and others ( Srinivas et al. , 2020 ; Shu et al. , 2020 ) . Moreover , if we can properly model and estimate the underlying transition dynamics , then we can perform model-based RL or planning , which can effectively reduce interactions with the environment ( Ha & Schmidhuber , 2018 ; Hafner et al. , 2018 ; 2019 ; 2020 ) . Despite the effectiveness of the above approaches to learning abstract features , current approaches usually fail to take into account whether the extracted state representations are sufficient and necessary for downstream policy learning . State representations that contain insufficient information may lead to sub-optimal policies , while those with redundant information may require more samples and more complex models for training . We address this problem by modeling the generative process and selection procedure induced by reward maximization ; by considering a generative environment model involving observed states , state-transition dynamics , and rewards , and explicitly characterizing structural relationships among variables in the RL system , we propose a principled approach to learning minimal sufficient state representations . We show that only the state dimensions that have direct or indirect edges to the reward variable are essential and should be considered for decision making . Furthermore , they can be learned by maximizing their ability to predict the action , given that the cumulative reward is included in the prediction model , while at the same time achieving their minimality w.r.t . the mutual information with observations as well as their dimensionality . The contributions of this paper are summarized as follows : • We construct a generative environment model , which includes the observation function , transition dynamics , and reward function , and explicitly characterizes structural relationships among variables in the RL system . • We characterize a minimal sufficient set of state representations , termed Action-Sufficient state Representations ( ASRs ) , for the downstream policy learning by making use of structural constraints and the goal of maximizing cumulative reward in policy learning . • In light of the characterization , we develop Structured Sequential Variational Auto-Encoder ( SSVAE ) , which explicitly encodes structural relationships among variables , for reliable identification of ASRs . • Accordingly , policy learning can be done separately from representation learning , and the policy function only relies on a set of low-dimensional state representations , which improve both model and sample efficiency . Moreover , the estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space , which effectively reduce possibly risky explorations . 2 ENVIRONMENT MODEL WITH STRUCTURAL CONSTRAINTS . In order to characterize a set of minimal sufficient state representations for downstream policy learning , we first formulate a generative environment model in partially observable Markov decision process ( POMDP ) , and then show how to explicitly embed structural constraints over variables in the RL system and leverage them . Suppose we have sequences of observations { 〈ot , at , rt〉 } Tt=1 , where ot ∈ O denotes perceived signals at time t , such as high-dimensional images , withO being the observation space , at ∈ A is the performed action with A being the action space , and rt ∈ R represents the reward variable withR being the reward space . We denote the underlying states , which are latent , by ~st ∈ S , with S being the state space . We describe the generating process of the environment model as follows : { ot = f ( ~st , et ) , rt = g ( ~st−1 , at−1 , t ) , ~st = h ( ~st−1 , at−1 , ηt ) , ( 1 ) where f , g , and h represent the observation function , reward function , and transition dynamics , respectively , and et , t , and ηt are corresponding independent and identically distributed ( i.i.d . ) random noises . The latent states ~st form an MDP : given ~st−1 and at−1 , ~st are independent of states and actions before t − 1 . Moreover , the action at−1 directly influences latent states ~st , instead of perceived signals ot , and the reward is determined by the latent states ( and the action ) as well . The perceived signals ot are generated from the underlying states ~st , contaminated by random noise et . We also consider noise t in the reward function to capture unobserved factors that may affect the reward , as well as measurement noise . It is commonplace that the action variable at−1 may not influence every dimension of ~st , and the reward rt may not be influenced by every dimension of ~st−1 as well , and furthermore there are structural relationships among different dimensions of ~st . Figure 1 gives an illustrative graphical representation , where s3 , t−1 influences s2 , t , at−1 does not have an edge to s3 , t , and among the states , only s2 , t−1 and s3 , t−1 have edges to rt . We use Rt = ∑∞ τ=t γ τ−trτ to denote the discounted cumulative reward starting from time t , where γ ∈ [ 0 , 1 ] is the discounted factor that determines how much immediate rewards are favored over more distant rewards . To reflect such constraints , we explicitly encode the graph structure over variables , including the structure over different dimensions of ~s and the structures from at−1 to ~st , ~st−1 to rt , and ~st to ot . ot−1 ot ot+1 s1 , t−1 s1 , t s1 , t+1 s2 , t−1 s2 , t s2 , t+1 s3 , t−1 s3 , t s3 , t+1 at−1 at rt rt+1 Rt Figure 1 : A graphical illustration of the generative environment model . Grey nodes denote observed variables and white nodes represent unobserved variables . Here , at−1 does not have an edge to s3 , t , and only s2 , t−1 and s3 , t−1 have edges to rt , and moreover , we take into account the structural relationships among different dimensions of latent states ~st . Accordingly , we re-formulate ( 1 ) as follows : { ot = f ( D~s ) o ~st , et ) , rt = g ( D~s ) r ~st−1 , Da ) r at−1 , t ) , si , t = hi ( D~s ( · , i ) ~st−1 , Da ) ~s ( · , i ) at−1 , ηi , t ) , ( 2 ) for i = 1 , · · · , d , where ~st = ( s1 , t , · · · , sd , t ) > , denotes element-wise product , and D ( · ) are binary matrices indicating the graph structure over variables . Specifically , D~s ) o ∈ { 0 , 1 } d×1 represents the graph structure from d-dimensional ~st to ot , D~s ) r ∈ { 0 , 1 } d×1 the structure from ~st−1 to the reward variable rt , Da ) r ∈ { 0 , 1 } the structure from the action variable at−1 to the reward variable rt , D~s ∈ { 0 , 1 } d×d denotes the graph structure from d-dimensional ~st−1 to d-dimensional ~st and D~s ( · , i ) is its i-th column , and Da ) ~s ∈ { 0 , 1 } 1×d corresponds to the graph structure from at−1 to ~st with Da ) ~s ( · , i ) representing its i-th column . For example , D~s ( j , i ) = 0 means that there is no edge from sj , t−1 to si , t . Here , we assume that the environment model , as well as the structural constraints , is invariant across time instance t . 2.1 MINIMAL SUFFICIENT STATE REPRESENTATIONS . Given observational sequences { 〈ot , at , rt〉 } Tt=1 , we aim to learn minimal sufficient state representations for the downstream policy learning . In the following , we first characterize the state dimensions that are indispensable for policy learning , when the environment model , including structural relationships , is given . Then we provide criteria to achieve sufficiency and minimality of the estimated state representations , when only { 〈ot , at , rt〉 } Tt=1 , but not the environment model , is given . Finding minimal sufficient state dimensions with a given environment model . RL agents learn to choose appropriate actions according to the current state vector ~st to maximize the cumulative reward , in which some dimensions may be redundant for policy learning . Then how can we identify a minimal subset of state dimensions that are sufficient to choose optimal actions ? We call such state dimensions Action-Sufficient state Representations ( ASRs ) , and denote it by ~sASRt . We can leverage the ( conditional ) independence/dependence relations among the quantities , under the condition of maximizing cumulative reward , which policy learning aims to achieve , to determine which dimensions of the states are the ASRs . Moreover , such independence/dependence relations can be directly seen from the graphical representation of the environment model , under the Markov condition and faithfulness assumption ( Pearl , 2000 ; Spirtes et al. , 1993 ) . Below , we give the graphical condition that ASRs are expected to satisfy . Proposition 1 . Given the graphical representation corresponding to the environment model , such as the representation in Figure 1 , which is assumed to be Markov and faithful to the measured data , each dimension in ~s ASRt has a direct or indirect edge to rt+τ , for τ > 0 . Proposition 1 can be shown based on the global Markov condition and the faithfulness assumption , which connects d-separation1 to conditional independence/dependence relations . A proof is given in Appendix . If si , t has a ( direct or indirect ) edge to rt+τ , the action variable at is dependent on si , t when the cumulative reward is maximized , which can be thought of as a data selection procedure depending on the cumulative reward , and thus si , t is needed for action determination . On the other hand , if si , t does not have an edge to rt+τ , then the action variable at is ( conditionally ) independent 1A path p is said to be d-separated by a set of nodes Z if and only if ( 1 ) p contains a chain i→ m→ j or a fork i ← m → j such that the middle node m is in Z , or ( 2 ) p contains a collider i → m ← j such that the middle node m is not in Z and such that no descendant of m is in Z. on si , t , and thus si , t is not needed for action determination . According to the above proposition , it is easy to see that for the graph given in Figure 1 , we have ~sASRt = ( s2 , t , s3 , t ) > . That is , we only need ( s2 , t , s3 , t ) > , instead of ~st , for the downstream policy learning . Minimal sufficient state representation learning from observed sequences . In practice , we usually do not have access to the latent states or the environment model , but instead only the observed sequences { 〈ot , at , rt〉 } Tt=1 . Then how can we learn the ASRs from the raw high-dimensional inputs such as images ? We denote by ~̃st the estimated whole latent state representations and ~̃sASRt ⊆ ~̃st the estimated minimal sufficient state representations for policy learning . We collected the data , that are used to learn the environment model and ASRs , with random actions . As discussed above , action and ASRs are dependent given the cumulative reward—this is a type of dependence relationship induced by selection on the effect ( reward ) . We can then learn the ASRs by maximizing their dependence with the action given the cumulative reward , I ( ~̃sASRt ; at |Rt+1 ) , where I denotes mutual information . Since I ( ~̃sASRt ; at |Rt+1 ) = H ( at |Rt+1 ) −H ( at | ~̃sASRt , Rt+1 ) , where H ( · ) denotes the conditional entropy , and the first term H ( at |Rt+1 ) does not contain ASRs , we can estimate ASRs by minimizing H ( at | ~̃sASRt , Rt+1 ) , with H ( at | ~̃sASRt , Rt+1 ) =− T∑ t=1 Eqφ , α { log pα ( at|~̃sASRt , Rt+1 ) } = − T∑ t=1 Eqφ , α { log pα ( at|D̃ASR ~̃st , Rt+1 ) } , where pα denote the probabilistic predictive model of at with parameters α , qφ , α is the joint distribution over ~̃st and at with qφ , α = qφpα , and qφ ( ~̃st|~̃st−1 , y1 : t , a1 : t−1 ) is the probabilistic inference model of ~̃st with parameters φ and yt= ( oTt , r T t ) , and D̃ ASR ∈ { 0 , 1 } d̃×1 is a binary vector indicating which dimensions of ~̃st are in ~̃sASRt , so D̃ ASR ~̃st gives ASRs ~̃sASRt . Moreover , we achieve minimality of the representation by minimizing conditional mutual information between observed high-dimensional signals yt and the ASR ~̃sASRt at time t given data at previous time instances , and meanwhile minimizing the dimensionality of ASRs with sparsity constraints : λ1 ∑T t=2 I ( yt ; ~̃s ASR t |y1 : t−1 , a1 : t−1 , ~̃st−1 ) + λ2‖D̃ ASR‖1 , where the conditional mutual information can be upper bound by a KL-divergence : I ( yt ; ~̃s ASR t |y1 : t−1 , a1 : t−1 , ~̃st−1 ) ≤ ∫ qφ ( ~̃s ASR t |~̃st−1 , y1 : t , a1 : t−1 ) log qφ ( ~̃s ASR t |~̃st−1 , y1 : t , a1 : t−1 ) pγ ( ~̃s ASRt |~̃st−1 , at−1 ; D~s , Da ) ~s ) d~̃sASRt = KL ( qφ ( ~̃s ASR t |~̃st−1 , y1 : t , a1 : t−1 ) ‖pγ ( ~̃sASRt |~̃st−1 , at−1 ; D~s , Da ) ~s ) ) , with pγ being the transition dynamics of ~̃st with parameters γ . Furthermore , Proposition 1 shows that given the ( estimated ) environment model , only those state dimensions that have a direct or indirect edge to the reward variable are the ASRs . In our learning procedure , we also take into account the relationship between the learned states ~̃st and the reward , and leverage such structural constraints for learning the ASRs . Denote by ĎASR ∈ { 0 , 1 } d̃×1 a binary vector indicating whether the corresponding state dimension in ~̃st has a direct or indirect edge to the reward variable . Consequently , we enforce the similarity between ĎASR and D̃ASR by adding an L1 norm on ĎASR − D̃ASR . Therefore , the ASRs can be learned by maximizing the following function : Lmin & suff = λ3 ∑T t=1 Eqφ , α { log pα ( at|D̃ASR ~̃st , Rt+1 ) } ︸ ︷︷ ︸ Sufficiency −λ4‖ĎASR − D̃ASR‖1 −λ1 ∑T t=2 KL ( qφ ( D̃ ASR ~̃st|~̃st−1 , y1 : t , a1 : t−1 ) ‖pγ ( D̃ASR ~̃st|~̃st−1 , at−1 ; D~s , Da ) ~s ) ) − λ2‖D̃ASR‖1︸ ︷︷ ︸ Minimality , ( 3 ) where λ ’ s are regularization terms , and note that ĎASR can be directly derived from the estimated structural matrices Da ) r and D~s ( · , i ) . The constraint in ( 3 ) provides a principled way to achieve minimal sufficient state representations , related to the information bottleneck method ( Tishby et al. , 1999 ) . Notice that it is just part of the objective function to maximize , and it will be involved in the complete objective function in ( 4 ) to learn the whole environment model . Remarks . By explicitly involving structural constraints , we achieve minimal sufficient state representations from the view of generative process underlying the RL problem and the selection procedure induced by reward maximization , which enjoys the following advantages . 1 ) The structural infor- mation provides an interpretable and intuitive picture of the generating process . 2 ) Accordingly , it also provides an interpretable and intuitive way to characterize a minimal sufficient set of state representations for policy learning , which removes unrelated information . 3 ) There is no information loss when representation learning and policy learning are done separately , which is computationally more efficient . 4 ) The generative environment model is fixed , independent of the behavior policy that is performed . Furthermore , based on the estimated environment model and ASRs , it is flexible to use a wide range of policy learning methods , and one can also perform model-based RL , which effectively reduce possibly risky explorations .

Summary: The paper proposes a novel method for learning a compact and sparse generative model of a partially observable environment, such that the sparsity of the model can be analyzed to identify the subset of latent state variables that are relevant to the reward function of the environment. This subset of the state space, called the action sufficient state representation of the task and environment, can then be used to learn a policy from fewer samples than using the entire reconstructed state space or a predictive world model would need. Empirical verification on two difficult decision problems with high-dimensional observation spaces supports the claim that the proposed method is indeed more sample efficient.