Year: 2021

Abstract:
Perceived signals in real-world scenarios are usually high-dimensional and noisy, and finding and using their representation that contains essential and sufficient information required by downstream decision-making tasks will help improve computational efficiency and generalization ability in the tasks. In this paper, we focus on partially observable environments and propose to learn a minimal set of state representations that capture sufficient information for decision-making, termed Action-Sufficient state Representations (ASRs). We build a generative environment model for the structural relationships among variables in the system and present a principled way to characterize ASRs based on structural constraints and the goal of maximizing cumulative reward in policy learning. We then develop a structured sequential Variational Auto-Encoder to estimate the environment model and extract ASRs. Our empirical results on CarRacing and VizDoom demonstrate a clear advantage of learning and using ASRs for policy learning. Moreover, the estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space to improve sample efficiency.

heading: 1 INTRODUCTION
text: 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.

heading: 2 ENVIRONMENT MODEL WITH STRUCTURAL CONSTRAINTS
text: 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.

heading: 2.1 MINIMAL SUFFICIENT STATE REPRESENTATIONS
text: 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.

heading: 3 STRUCTURED SEQUENTIAL VAE FOR THE ESTIMATION OF ASRS
text: In this section, we give estimation procedures for the environment model and ASRs, as well as the identifiability guarantee in linear cases.
Identifiability in Linear-Gaussian Cases. Below, we first show the identifiability guarantee in the linear case, as a special case of Eq. (2). In the linear case (see the environment model given in Eq. (5) in Appendix D), D~s)o, D~s)r, Da)r, D~s, and Da)~s are linear coefficients, indicating corresponding graph structures and also the strength. Denote the covariance matrices of et and t by Σe and Σ , respectively. Further let D̈~s)o := (D>~s)o, D > ~s)r) >. The following proposition shows that the environment model in the linear case is identifiable up to some orthogonal transformation on certain coefficient matrices from observed data {〈ot, at, rt〉}Tt=1. Proposition 2 (Identifiability). Suppose the perceived signal ot, the reward rt, and the latent states ~st follow a linear environment model. If assumptions A1∼A4 (given in Appendix C) hold and with the second-order statistics of the observed data {〈ot, at, rt〉}Tt=1, the noise variances Σe and Σ , Da)r, D̈~s)oD~s
kDa)~s (with k ≥ 0), and D̈~s)oD̈>~s)o are uniquely identified. This proposition shows that in the linear case, with the second-order statistics of the observed data, we can identify the parameters up to orthogonal transformations. In particular, suppose the linear environment model with parameters (D~s)o, D~s)r, Da)r, D~s, Da)~s,Σe,Σ ) and that with (D̃~s)o, D̃~s)r, D̃a)r, D̃~s, D̃a)~s, Σ̃ẽ, Σ̃̃) are observationally equivalent. Then we have ˜̈D~s)o= D̈~s)oU , D̃a)r=Da)r, D̃~s=U>D~sU , D̃a)~s=Da)~sU , Σ̃ẽ=Σe, and Σ̃̃=Σ , where U is an orthogonal matrix.
General Nonlinear Cases. To handle general nonlinear cases with the generative process given in Eq. (2), we develop a Structured Sequential VAE (SS-VAE) to learn the model (including structural constraints) and infer latent state representations ~̃st and ASRs ~̃sASRt , with the input {〈ot, at, rt〉}Tt=1. Specifically, the latent state dimensions are organized with structures, captured by D~s, to achieve conditional independence. The structural relationships over perceived signals, latent states, the action variable, and the reward variable are also embedded as free parameters (i.e., D~s)o, D~s)r, Da)r, Da)~s) into SS-VAE. Moreover, we aim to learn state representations ~̃st and ASRs ~̃sASRt that satisfy the following properties: (i) ~̃st should capture sufficient information of observations ot, rt, and at, that is, it should be enough to enable reconstruction. (ii) The state representations should allow for accurate predictions of the next state and also the next observation. (iii) The transition dynamics should follow an MDP. (iv) ~̃sASRt are minimal sufficient state representations for the downstream policy learning.
Let y1:T = {(o>t , r>t )>}Tt=1. To achieve the above properties, we maximize the following objective function:
L(y1:T ; (θ, φ, γ, α,D(·))) = T−2∑ t=1 Eqφ {
log pθ(ot|~̃st;D~s)o) + log pθ(rt+1|~̃st, at;D~s)r, Da)r)︸ ︷︷ ︸ Reconstruction
+ log pθ(ot+1|~̃st) + log pθ(rt+2|~̃st, at+1)︸ ︷︷ ︸ Prediction
} + λ3 T∑ t=1 Eqφ,α { log pα(at|~̃st, Rt+1; D̃ASR) }︸ ︷︷ ︸
Sufficiency − λ1 T∑ t=2 KL ( qφ(~̃st|~̃st−1,y1:t, a1:t−1)‖ pγ(~̃st|~̃st−1, at−1;D~s, Da)~s)︸ ︷︷ ︸
Transition
) − λ2‖D̃ASR‖1
︸ ︷︷ ︸ Conditional disentanglement & Minimality
− ( λ5‖D~s)o‖1 + λ6‖D~s)r‖1 + λ7‖D~s‖1 + λ8‖Da)~s‖1 + λ4‖ĎASR − D̃ASR‖1 )︸ ︷︷ ︸ Sparsity ,
(4) which contains the reconstruction error at each time instance, the one-step prediction error of observations, the KL divergence to constrain the latent space, and moreover, the MDP restrictions on transition dynamics, the sufficiency and minimality guarantee of state representations for policy
learning, as well as sparsity constraints on the graph structure. We denote by pθ the generative model with parameters θ and structural constraints D(·), qφ the inference model with parameters φ, pγ the transition dynamics, and pα action prediction with ASRs given the cumulative reward Rt+1. Each factor in pγ and qφ is modeled with a mixture of Gaussians (MoGs), to approximate a wide class of continuous distributions.
Below are the details of each component in the above objective function:
• Reconstruction and prediction components: These two parts are commonly used in sequential VAE. They aim to minimize the reconstruction error and prediction error of the perceived signal ot and the reward rt.
• Transition component: To achieve the property that state representations satisfy an MDP, we explicitly model the transition dynamics: log pγ(~̃st|~̃st−1, at−1;D~s, Da)~s). In particular, ~̃st|~̃st−1 is modelled with a mixture of Gaussians: ∑K k=1πkN ( µk(~̃st−1,at−1),Σk(~̃st−1,at−1) ) , where K is
the number of mixtures, µk(·) and Σk(·) are given by multi-layer perceptrons (MLP) with inputs ~̃st−1 and at−1, parameters γ, and structural constraints D~s and Da)~s. This explicit constraint on state dynamics is essential for establishing a Markov chain in latent space and for learning a representation for long-term predictions. Note that unlike in traditional VAE (Kingma & Welling, 2013), we do not assume that different dimensions in ~̃st are marginally independent, but model their structural relationships explicitly to achieve conditional independence.
• KL-divergence constraint: The KL divergence is used to constrain the state space with multiple purposes: (1) It is used in the lower bound of logP (y1:T ) to achieve conditional disentanglement between qφ(s̃i,t|·) and qφ(s̃j,t|·) for i 6= j, (2) and also to achieve minimality of ASRs.
• Sufficiency & minimality constraints: We achieve minimal sufficient state representations for the downstream policy learning by leveraging action prediction given the cumulative reward, the conditional mutual information between yt and ~̃sASRt , and structural constraints. For details, please refer to Section 2.1.
• Sparsity constraints: According to the edge-minimality property (Zhang & Spirtes, 2011), we additionally put sparsity constraints on structural matrices to achieve better identifiability. In particular, we use L1 norm of the structural matrices as regularizers in the objective function to achieve sparsity of the solution.
Figure 2 gives the diagram of the neural network architecture in model training. We use SS-VAE to learn the environment model and ASRs. Specifically, the encoder, which is used to learn the inference model qφ(~̃st|~̃st−1,y1:t, a1:t−1), includes a Long Short-Term Memory (LSTM (Hochreiter & Schmidhuber, 1997)) to encode the sequential information with output ht and a Mixture Density Network (MDN (Bishop, 1994)) to output the parameters of MoGs. At each time instance, the input 〈ot+1, rt+1, at〉 is projected to the encoder and a sample of ~̃st+1 is inferred from qφ as output. The generated sample further acts as an input to the decoder, together with at+1 and structural matrices D~s)o, D~s)r, and Da)r. Then the decoder outputs ôt+1 and r̂t+2. Moreover, the state dynamics which satisfies a Markov process and is embedded with structural constraints D~s and Da)~s, is modeled with an MLP and MDN, marked with red in Figure 2. The action prediction part (denoted by AP), which helps sufficient state representations, uses MLP and is marked with blue. During training, we approximate the expectation in L by sampling and then jointly learn all parameters by maximizing L using stochastic gradient descent.

heading: 4 POLICY LEARNING WITH ASRS
text: After estimating the generative environment model, we are ready to learn the optimal policy, where the policy function only depends on low-dimensional ASRs, instead of high-dimensional images. The entire procedure roughly contains the following three parts: (1) data collection with a random policy, (2) environment model estimation (with details in Section 3), and (3) policy learning with ASRs. Notably, the generative environment model is fixed, regardless of the behavior policy that is used to generate the data, and after learning the environment model, as well as the inference model for ASRs, our framework is flexible for both model-free and model-based policy learning.
Model-Free Policy Learning. For model-free policy learning, we make use of the learned environment model to infer ASRs ~̃sASRt from past observed sequences {o≤t, r≤t, a≤t−1} and then predict the action with the estimated low-dimensional ASRs. Our method is flexible to use a wide range of model-free methods; for example, one may use deep Q-learning for discrete actions (Mnih
Figure 2: Diagram of neural network architecture to learn state representations. The corresponding structural constraints are involved in “Deconv” and “MLP”, and “AP" represents the action prediction part for sufficient state representation learning.
et al., 2015) and deep deterministic policy gradient (DDPG) for continuous actions (Lillicrap et al., 2015). Algorithm 1 in Appendix F gives the detailed procedure of model-free policy learning with ASRs in partially observable environments. Model-Based Policy Learning. The downside of model-free RL algorithms is that they are usually data hungry, requiring very large amounts of interactions. On the contrary, model-based RL algorithms enjoy much better sample efficiency. Hence, we make use of the learned generative environment model, including the transition dynamics, observation function, and reward function, for model-based policy optimization. Based on the generative environment model, one can learn behaviors from imagined outcomes to increase sample-efficiency and mitigate heavy and possibly risky interactions with the environment. We present the procedure of the classic Dyna algorithm (Sutton, 1990; Sutton & Barto, 2018) with ASRs in Algorithm 2 in Appendix F.

heading: 5 EXPERIMENTS
text: To evaluate the proposed approach, we conducted experiments on both CarRacing (Klimov, 2016) and VizDoom (Kempka et al., 2016) environments, following the setup in the world model (Ha & Schmidhuber, 2018) for a fair comparison. It is known that CarRacing is very challenging—the recent world model (Ha & Schmidhuber, 2018) is the first known solution to achieve the score required to solve the task. Without stated otherwise, all results were averaged across five random seeds, with standard deviation shown in the shaded area.

heading: 5.1 CARRACING EXPERIMENT
text: CarRacing is a continuous control task with three continuous actions: steering left/right, acceleration, and brake. Reward is −0.1 every frame and +1000/N for every track tile visited, where N is the total number of tiles in track. It is obvious that the CarRacing environment is partially observable: by just looking at the current frame, although we can tell the position of the car, we know neither its direction nor velocity that are essential for controlling the car. For a fair comparison, we followed the same setting as in Ha & Schmidhuber (2018). Specifically, we collected a dataset of 10k random rollouts of the environment, each consisting of 1000 time steps. The dimensionality of latent states ~̃st was set to d̃ = 32, determined by hyperparameter tuning.
Analysis of ASRs. To demonstrate the structures over observed frames, latent states, actions, and rewards, we visualized the learned D~s)o, D~s)r, D~s, and Da)~s, as shown in Figure 9 in Appendix G. Intuitively, we can see that D~s)r and Da)~s have many values close to zero, meaning that the reward is only influenced by a small number of state dimensions, and not many state dimensions are influenced by the action. Furthermore, from D~s, we found that there are influences from ~̃si,t to ~̃si,t+1 (diagonal values) for most state dimensions, which is reasonable because we want to learn an MDP over the underlying states, while the connections across states (off-diagonal values) are much sparser. Compared to the original 32-dim latent states, ASRs have only 21 dimensions. Below, we empirically showed that the low-dimensional ASRs significantly improve the policy learning performance in terms of both efficiency and efficacy.
Comparison Between Model-Free and Model-Based ASRs. We applied both model-free (DDPG) (Lillicrap et al., 2015) and model-based (Dyna and Prioritized Sweeping) algorithms (Sutton, 1990) to ASRs (with 21-dims). As shown in Figure 3, interestingly, by taking advantage of the learned generative model, model-based ASRs is superior to model-free ASRs at a faster rate, which
demonstrates the effectiveness of the learned model. It also shows that with the estimated environment model and ASRs, we can learn behaviors from imagined outcomes to improve sample-efficiency.
Comparison with VRL, SLAC, PlaNet, DBC, and Dreamer. We also compared the proposed framework of policy learning with ASRs (with 21-dims) with 1) the same learning strategy but with vanilla representation learning (VRL, implemented without the components for minimal sufficient state representations as in Eq. (3)), 2) SLAC (Lee et al., 2019), 3) PlaNet (Hafner et al., 2018), 4) DBC (Zhang et al., 2021), and 5) Dreamer (Hafner et al., 2019). For a fair comparison, the latent dimensions of VRL, PlaNet, SLAC, DBC and Dreamer are set to 21 as well, and we require all of them to have the model capacity similar to ours (i.e., similar model architectures). From Figure 3, we can see that our methods, both model-free and model-based, obviously outperform others. It is worth noting that the huge performance difference between ASRs and VRL shows that the components for minimal sufficient state representations play a pivotal role in our objective.
Comparison with World Models. In light of the fact that world models (Ha & Schmidhuber, 2018) achieved good performance in CarRacing, we further compared our method (with 21-dim ASRs) with the world model. For a fair comparison, following Ha & Schmidhuber (2018), we also used the Covariance-Matrix Adaptation Evolution Strategy (CMA-ES) (Hansen, 2016) with a population of 64 agents to optimize the parameters of the controller. In addition, following the same setting as in Ha & Schmidhuber (2018) (where the agent’s fitness value is defined as the average cumulative reward of the 16 random rollouts), we show the fitness values of the best performer (max) and the population (mean) at each generation (Figure 4). We also took the best performing agent at the end of every 25 generations and tested it over 1024 random rollout scenarios to record the average (best avg score). It is obvious that our method (denoted by ASR-*) has a more efficient and also efficacy training process. The best average score of ASRs is 65 higher than that of world models.
Comparison with Dreamer and DBC with Background Distraction. We further compared ASRs (with 21-dims) with Dreamer and DBC when there are natural video distractors in CarRacing; we chose Dreamer and DBC, because their performance are relatively better than other comparisons when there are no distractors. Specifically, we followed Zhang et al. (2021) to incorporate natural video from the Kinetics dataset (Kay et al., 2017) as background in CarRacing. Similarly, for a fair comparison, we require all of them to have the same latent dimensions and have the similar model capacity. As shown in Figure 5, we can see that our method outperforms both Dreamer and DBC.
Ablation Study. We further performed ablation studies on latent dynamics prediction; that is, we compared with the case when the transition dynamics in Eq. (4) is not explicitly modeled, but is replaced with a standard normal distribution. Figure 10 in Appendix G shows that by explicitly modelling the transition dynamics (denoted by with LDP), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP).

heading: 5.2 VIZDOOM EXPERIMENT
text: We also applied the proposed method to VizDoom take cover scenario (Kempka et al., 2016), which is a discrete control problem with two actions: move left and move right. Reward is +1 at each time step while alive, and the cumulative reward is defined to be the number of time steps the agent manages to stay alive during an episode.
Considering that in the take over scenario the action space is discrete, we applied the widely used DQN (Mnih et al., 2013) on ASRs for policy learning. In addition to the comparisons with VRL (as in CarRacing) and DQN on raw observations, we further compared with another common approach to
Model CumulativeRewards
Dreamer 621±124.5 DBC 803±112.5 Model-free ASRs 938±87.2 Model-based ASRs 954±98.6
Figure 5: Comparisons with Dreamer and DBC in CarRacing with natural video distractors, after 2000 training episodes, with standard error.
POMDPs: DRQN (Hausknecht & Stone, 2015). As shown in Figure 6, DQN on ASRs achieve a much better performance than all other comparisons, and in particular, DQN on ASRs outperforms DRQN on observations by around 400 on average in terms of cumulative reward. Similarly, we applied model-based (Dyna) algorithms (Sutton, 1990) to ASRs (with 21-dims). As shown in Figure 6, we can draw the same conclusion that by taking advantage of the learned generative model, model-based ASRs is superior to model-free ASRs at a faster rate. We also applied ASRs to world models, where Figure 7 shows that our method with ASRs (denoted by ASR-*) achieves a better performance.

heading: 6 RELATED WORK
text: In the past few years, a number of approaches have been proposed to learn low-dimensional Markovian representations, which capture the variation in the environment generated by the agent’s actions, without direct supervision (Lesort et al., 2018; Krishnan et al., 2015; Karl et al., 2016; Ha & Schmidhuber, 2018; Watter et al., 2015; Zhang et al., 2018; Kulkarni et al., 2016; Mahadevan & Maggioni, 2007; Gelada et al., 2019; Gregor et al., 2018; Ghosh et al., 2019; Zhang et al., 2021). Common strategies for such state representation learning include reconstructing the observation, learning a forward model, or learning an inverse model. Furthermore, prior knowledge, such as temporal continuity (Wiskott & Sejnowski, 2002), can be added to constrain the state space.
Recently, much attention has been paid to world models, which try to learn an abstract representation of both spatial and temporal aspects of the high-dimensional input sequences (Watter et al., 2015; Ebert et al., 2017; Ha & Schmidhuber, 2018; Hafner et al., 2018; Zhang et al., 2019b; Gelada et al., 2019; Kaiser et al., 2019; Hafner et al., 2019; 2020). Based on the learned world model, agents can perform model-based RL or planning. Our proposed method is also in the class of world models, which models the generative environment model, and additionally, encodes structural constraints and achieves the sufficiency and minimality of the estimated state representations from the view of generative and selection process. In contrast, Shu et al. (2020) makes use of contrastive loss, as an alternative of reconstruction loss; however, it only focuses on the transition dynamics and also fails to ensure the sufficiency and minimality. Another line of approaches of state representation learning is based on predictive state representations (PSRs) (Littman & Sutton, 2002; Singh et al., 2004). A recent approach generalizes PSRs to nonlinear predictive models, by exploiting the coarsest partition of histories into classes that are maximally predictive of the future (Zhang et al., 2019a). Moreover, bisimulation-based methods have also attracted much attention (Castro, 2020; Zhang et al., 2021).
On the other hand, our work is also related to Bayesian network learning and causal discovery (Spirtes et al., 1993; Pearl, 2000; Huang* et al., 2020). For example, Strehl et al. (2007) considers factorized-state MDP with structures being modeled with dynamic Bayesian network or decision trees. Incorporating such structure information has shown benefits in several machine learning tasks (Zhang* et al., 2020; Huang et al., 2019), and in this paper, we show its advantages in POMDPs.

heading: 7 CONCLUSIONS AND FUTURE WORK
text: In this paper, we develop a principled framework to characterize a minimal set of state representations that suffice for policy learning, by making use of structural constraints and action predictions. Accordingly, we propose SS-VAE to reliably extract such a set of state representations from raw observations. The estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space, which effectively reduce sample complexity and possibly risky interactions with the environment. The proposed approach shows promising results on complex environments–CarRacing and Vizdoom. The future work along this direction include investigating identifiability conditions in general nonlinear cases and extending the approach to cover heterogeneous environments, where the generating processes may change over time or across domains.
REPRODUCIBILITY STATEMENT
Proofs of all our theoretical results are given in Appendix A, B, C, D with disclosure of all assumptions. More details about the model estimation and policy learning are given in Appendix E and Appendix F, respectively. Appendix G provides more experimental details (including a description of the hyperparameters) and results. All datasets used are publicly available or instructions are provided on how to generate them in Appendix G. A description of the network architectures used is included in Appendix H. Appendix I introduces the platform and license. Source code for all our experiments will be made available on GitHub after deanonymization, providing also a complete description of our experimental environment, configuration files and instructions on the reproduction of our experiments.

heading: A PROOF OF PROPOSITION 1
text: We first give the definitions of the Markov condition and the faithfulness assumption, which will be used in the proof.
Definition 1 (Global Markov Condition (Spirtes et al., 1993; Pearl, 2000)). The distribution p over V satisfies the global Markov property on graphG if for any partition (A,B,C) such thatB d-separates A from C,
p(A,C|B) = p(A|B)p(C|B).
Definition 2 (Faithfulness Assumption (Spirtes et al., 1993; Pearl, 2000)). There are no independencies between variables that are not entailed by the Markov Condition.
Below, we give the proof of Proposition 1.
Proof. We first show that if si,t ∈ ~sASRt , then it has a direct or indirect edge to rt+τ . We prove it by contradiction. Suppose that si,t does not have a direct or indirect edge to rt+τ . According to the Markov assumption, si,t is independent of at given R. Hence, si,t is not necessary for decision making, and thus si,t is not a dimension in ~sASRt , which contradicts to the assumption. Since we have a contradiction, it must be that si,t has a direct or indirect edge to rt+τ .
We next show that if si,t has a direct or indirect edge to rt+τ , then si,t ∈ ~sASRt .
Similarly, by contradiction suppose that si,t is not a dimension in ~sASRt . It means that si,t is independent on at given R and some other variables. Then according to the faithfulness assumption, si,t does not have a direct or indirect edge to rt+τ , which contradicts to the assumption.

heading: B MINIMALITY OF THE REPRESENTATION
text: In this section, we give the detailed derivation of the minimality of the state representation given in Section 2.1.
We achieve minimality of the representation by minimizing conditional mutual information between observed high-dimensional signals yt, where yt = {oTt , rTt }, 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.
Note that in the above conditional mutual information, we need to conditional on the previous states ~̃st−1, instead of ~̃sASRt−1 , which two give different conditional mutual information. It can be shown by contradiction. Suppose I(yt; ~̃sASRt |y1:t−1, a1:t−1, ~̃st−1) = I(yt; ~̃sASRt |y1:t−1, a1:t−1, ~̃sASRt−1 ), and denote ~̃sC = ~̃s\~̃sASR. Then the equivalence implies that ~̃sCt−1 is independent of ot (where ot ∈ yt) given {y1:t−1, a1:t−1, ~̃sASRt−1 }. It is obviously violated for the example given in Figure 1, where ~̃sC = s1 and ~̃sASR = {s2, s3}, and s1,t−1 is dependent on ot given {y1:t−1, a1:t−1, s2,t−1, s3,t−1}. Hence, conditioning on ~̃st−1 and ~̃sASRt−1 give different conditional mutual information. Therefore, in the above conditional mutual information, we need to condition on the previous states ~̃st−1.
Moreover, the conditional mutual information I(yt; ~̃sASRt |y1:t−1, a1:t−1, ~̃st−1) can be upper bound by a KL-divergence:
I(yt; ~̃s ASR t |y1:t−1, a1:t−1, ~̃st−1) = I(yt; ~̃s ASR t , {y1:t−1, a1:t−1, ~̃st−1})− I(yt; {y1:t−1, a1:t−1, ~̃st−1})
= [ H(yt)−H(yt|~̃sASRt ,y1:t−1, a1:t−1, ~̃st−1) ] − [ H(yt)−H(yt|y1:t−1, a1:t−1, ~̃st−1) ] = H(yt|y1:t−1, a1:t−1, ~̃st−1)−H(yt|~̃sASRt ,y1:t−1, a1:t−1, ~̃st−1) = −p(yt|y1:t−1, a1:t−1, ~̃st−1) log p(yt|y1:t−1, a1:t−1, ~̃st−1)
+ p(yt|~̃sASRt ,y1:t−1, a1:t−1, ~̃st−1) log p(yt|~̃sASRt ,y1:t−1, a1:t−1, ~̃st−1) ≤ −p(yt|~̃sASRt ,y1:t−1, a1:t−1, ~̃st−1) log p(yt|y1:t−1, a1:t−1, ~̃st−1)
+ p(yt|~̃sASRt ,y1:t−1, a1:t−1, ~̃st−1) log p(yt|~̃sASRt ,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 γ.

heading: C ASSUMPTIONS OF PROPOSITION 2
text: To show the identifiability of the model in the linear case, we make the following assumptions:
A1. do + dr ≥ ds, where |ot| = do, |rt| = dr, and |st| = ds. A2. (D>~s)o, D > ~s)r) > is full column rank and D~s is full rank. A3. The control signal at is i.i.d. and the state ~st is stationary. A4. The process noise has a unit variance, i.e., var(ηt) = I .

heading: D PROOF OF PROPOSITION 2
text: Proof. The proof of the linear case without control signals has been shown in Zhang & Hyvärinen (2011). Below, we give the identifiability proof in the linear-Gaussian case with control signals:{
ot = D~s)o~st + et, rt+1 = D~s)r~st +Da)rat + t+1, ~st = D~s~st−1 +Da)~sat−1 + ηt.
(5)
Let yt+1 = [o>t , r > t+1] >, D̈~s)o = [D>~s)o, D > ~s)r] >, D̈a)r = [~0>, D>a)r] >, and ët = [e>t , > t+1] >. Then the above equation can be represented as:{
yt = D̈~s)o~st + D̈a)rat + ët, ~st = D~s~st−1 +Da)~sat−1 + ηt.
(6)
Because the dynamic system is linear and Gaussian, we make use of the second-order statistics of the observed data to show the identifiability. We first consider the cross-covariance between yt+k and at:{
Cov(yt+k, at) = D̈~s)oDk−1~s Da)~s · Var(at), if k > 0. Cov(yt+k, at) = D̈a)r · Var(at), if k = 0.
(7)
Thus, from the cross-covariance between yt+k and at, we can identify D̈~s)oDa)~s, D̈a)r, and D̈~s)oD k ~sDa)~s for k > 0.
Next, we consider the auto-covariance function of ~s. Define the auto-covariance function of ~s at lag k as R~s(k) = E[~st~s>t+k], and similarly for Ry(k). Clearly, R~s(−k) = R~s(k)> and Ry(−k) = Ry(k)
>. Then we have{ R~s(k) = R~s(k − 1) ·D>~s , if k > 0, R~s(k) = R > ~s (1) ·D>~s +Da)~sVar(at−1)D>a)~s + I, if k = 0.
(8)
Below, we first consider the case where do+dr = ds. Let ỹt = D̈~s)o~st, so yt = ỹt+ D̈a)rat−1 + ët and Rỹ(k) = D̈~s)oR~st(k)D̈
> ~s)o. Rỹ(k) satisfies the recursive property:{
Rỹ(k) = Rỹ(k − 1) · Ω>, if k > 0, Rỹ(k) = R > ỹ (1) · Ω> + D̈~s)o(Da)~sVar(at−1)D>a)~s + I)D̈>~s)o, if k = 0,
(9)
where Ω = D̈~s)oD~sD̈−1~s)o.
Denote Sk = D̈~s)oDk−1~s Da)~s · Var(at). Then we can derive the recursive property for Ry(k): Ry(k) = Ry(k − 1) · Ω> − D̈a)rS>k−1Ω> + D̈a)rS>k , if k > 1, Ry(k) = Ry(k − 1) · Ω> − D̈a)rVar>(at)D̈>a)rΩ> − ΣeΩ> + D̈a)rS>k , if k = 1, Ry(k) = R > y (1) · Ω> + ( D̈a)rVar(at)D̈>a)r + Σe
) + D̈~s)o(Da)~sVar(at)D>a)~s + I)D̈ > ~s)o, if k = 0.
When k = 2, we have
Ry(2) = Ry(1) · Ω> − D̈a)rS>1 Ω> + D̈a)rS>2 . The above equation can be re-organized as(
Ry(2)− D̈a)r · S>2 ) = ( Ry(1)− D̈a)r · S>1 ) · Ω>.
Because D̈a)r and Sk are identifiable, and suppose ( Ry(1) − D̈a)r · S>1 ) is invertible, Ω = D̈~s)oD~sD̈ −1 ~s)o is identifiable. We further consider Ry(0) and Ry(1) and write down the following form:[ Ry(0)− D̈~s)o(Da)~sVar(at−1)D>a)~s + I)D̈>~s)o
Ry(1)
]
= [ R>y (1) Ry(0) ] · Ω> + [ D̈a)rVar(at)D̈>a)r −D̈a)rVar>(at)D̈>a)rΩ> + D̈a)rS>1 ] + Σe [ I −Ω> ] .
From the above two equations we can then identify Σe and D̈~s)o(Da)~sVar(at−1)D>a)~s + I)D̈ > ~s)o, and because D̈~s)oDa)~s is identifiable, D̈~s)oD̈>~s)o is identifiable.
In summary, we have shown the identifiability of D̈a)r, D̈~s)oDa)~s, D̈~s)oDk~sDa)~s, D̈~s)oD̈ > ~s)o, and Σe. Furthermore, D̈~s)o, D~s, and Da)~s are identified up to some orthogonal transformations. That is, suppose the model in Eq. (3) with parameters (D~s)o, D~s)r, Da)r, D~s, Da)~s,Σe,Σ ) and that with (D̃~s)o, D̃~s)r, D̃a)r, D̃~s, D̃a)~s, Σ̃ẽ, Σ̃̃) are observationally equivalent, we then have ˜̈D~s)o = D̈~s)oU , D̃a)r = Da)r, D̃~s = U>D~sU , D̃a)~s = Da)~sU , Σ̃ẽ = Σe, and Σ̃̃ = Σ , where U is an orthogonal matrix.
Next, we extend the above results to the case where do+dr > ds. Let D̈>~s)o(i,·) be the i-th row of D̈~s)o. Recall that D̈~s)o is of full column rank. Then for any i, one can show that there always exist ds−1 rows of D̈~s)o, such that they, together with D̈>~s)o(i,·), form a ds × ds full-rank matrix, denoted by ¯̈D~s)o(i,·). Then from the observed data corresponding to ¯̈D~s)o(i,·), ¯̈D~s)o(i,·) is determined up to orthogonal
transformations. Thus, D̈~s)o is identified up to orthogonal transformations. Similarly, Da)r, D~s, and Da)~s are identified up to orthogonal transformations. Furthermore, Cov(D̈~s)o~st + Da)rat) is determined by D̈~s)o, D̈a)r, D~s, and Da)~s. Because Cov(yt) = Cov(D̈~s)o~st +Da)rat) + Σë, Σë is identifiable.
One may further add sparsity constraints onD~s)o,D~s)r,D~s, andDa)~s, to select more sparse structures among the equivalent ones. For example, one may add sparsity constraints on the columns of D~s)o. Note this corresponds to the mask on the elements of ~st in Eq. 2; if the full column is 0, then the corresponding dimension of ~st is not selected.

heading: E MORE ESTIMATION DETAILS FOR GENERAL NONLINEAR MODELS
text: The generative model pθ can be further factorized as follows:
log pθ(y1:T |~̃s1:T , a1:T−1;D~s)o, D~s)r, Da)r) = log pθ(o1:T |~̃s1:T ;D~s)o) + log pθ(r1:T |~̃s1:T , a1:T−1;D~s)r, Da)r) = ∑T t=1 log pθ(ot|~̃st;D~s)o) + log pθ(rt|~̃st−1, at−1;D~s)r, Da)r),
(10)
where both pθ(ot|~̃st;D~s)o) and pθ(rt|~̃st−1, at−1;D~s)r, Da)r) are modelled by mixture of Gaussians, with D~s)o indicating the existence of edges from ~̃st to ot and D~s)r indicating the existence of edges from ~̃st−1 to rt.
The inference model qφ(~̃s1:T |y1:T , a1:T−1) is factorized as
log qφ(~̃s1:T |y1:T , a1:T−1) = log qφ(~̃s1|y1, a0) + T∑ t=2 log qφ(~̃st|~̃st−1,y1:t, a1:t−1),
where both qφ(~̃s1|y1, a0) and qφ(~̃st|~̃st−1,y1:t, a1:t−1) are modelled with mixture of Gaussians. The transition dynamics pγ is factorized as
log pγ(~̃s1:T |a1:T−1;D~s(·,i), Da)~s(·,i)) = T∑ t=1 log pγ(~̃st|~̃st−1, at−1;D~s(·,i), Da)~s(·,i)), (11)
with ~̃st|~̃st−1 modelled with mixture of Gaussians. Thus, the KL divergence can be represented as follows:
KL ( qφ(~̃s1:T |y1:T , a1:T−1)‖pγ(~̃s1:T ) ) = KL ( qφ(~̃s1|y1, a0)‖pγ(~̃s1) ) +
T∑ t=2 Eqφ [ KL ( qφ(~̃st|~̃st−1,y1:t, a1:t−1)‖pγ(~̃st|~̃st−1) )] . (12)
In practice, KL divergence with mixture of Gaussians is hard to implement, so instead, we used the following objective function:
KL ( qφ(~̃s1|y1, a0)‖pγ′(~̃s1) ) + T∑ t=2 Eqφ [ KL ( qφ(~̃st|~̃st−1,y1:t, a1:t−1)‖pγ′(~̃st|~̃st−1) )] +λ
T∑ t=1 log pγ(~̃st|~̃st−1, at−1;D~s(·,i), Da)~s(·,i)) (13)
where pγ′ is a standard multivariate Gaussian N (~0, Id).

heading: F MORE DETAILS FOR POLICY LEARNING WITH ASRS
text: Algorithm 1 gives the procedure of model-free policy learning with ASRs in partially observable environments. Specifically, it starts from model initialization (line 1) and data collection with a random policy (line 2). Then it updates the environment model and identifies the set of ASRs with the collected data (line 3), after which, the main procedure of policy optimization follows. In particular, because we do not directly observe the states ~st, on lines 8 and 12, we infer qφ(~sASRt+1 |o≤t+1, r≤t+1, a≤t) and sample ~sASRt+1 from the posterior. The sampled ASRs are then stored in the buffer (line 13). Furthermore, we randomly sample a minibatch of N transitions to optimize the policy (lines 14 and 15). One may perform various RL algorithms on the ASRs, such as deep deterministic policy gradient (DDPG (Lillicrap et al., 2015)) or Q-learning (Mnih et al., 2015).
Algorithm 2 presents the procedure of the classic Dyna algorithm with ASRs. Lines 17-22 make use of the learned environment model to predict the next step, including ~sASRt+1 and rt+1, and update the Q function n times. Specifically, in our implementation, the hyper-parameter n is 20. Based on the learned model, the agent learns behaviors from imagined outcomes in the compact latent space, which helps to increase sample efficiency.

heading: G ADDITIONAL EXPERIMENTS AND DETAILS
text: G.1 CARRACING EXPERIMENT
CarRacing (with an illustration in Figure 8) is a continuous control task with three continuous actions: steering left/right, acceleration, and brake. Reward is−0.1 every frame and +1000/N for every track tile visited, where N is the total number of tiles in track. It is obvious that the CarRacing environment
Algorithm 1 Model-Free Policy Learning with ASRs in Partially Observable Environments 1: Randomly initialize neural networks and initialize replay buffer B. 2: Apply random control signals and record multiple rollouts. 3: Estimate the model given in (2) with the recorded data (according to Section 3). 4: Identify indices of ASRs according to the learned graph structure and the criteria in Prop. 1. 5: for episode = 1, . . . , M do 6: Initialize a random process N for action exploration. 7: Receive initial observations o1 and r1. 8: Infer the posterior qφ(~sASR1 |o1, r1) and sample ~sASR1 . 9: for t = 1, . . . , T do 10: Select action at = π(~sASRt ) +Nt according to the current policy and exploration noise. 11: Execute action at and receive reward rt+1 and observation ot+1. 12: Infer the posterior qφ(~sASRt+1 |o≤t+1, r≤t+1, a≤t) and sample ~sASRt+1 . 13: Store transition (~sASRt , at, rt+1, ~s ASR t+1 ) in B. 14: Sample a random minibatch of N transitions (~sASRi , ai, ri+1, ~s ASR i+1 ) from B. 15: Update network parameters using a specified RL algorithm (e.g., DQN or DDPG). 16: end for 17: end for
Algorithm 2 Model-Based Policy Learning with ASRs in Partially Observable Environments 1: Randomly initialize neural networks and initialize replay buffer B. 2: Apply random control signals and record multiple rollouts. 3: Estimate the model given in (2) with the recorded data (according to Section 3). 4: Identify indices of ASRs according to the learned graph structure and the criteria in Prop. 1. 5: for episode = 1, . . . , M do 6: Initialize a random process N for action exploration. 7: Receive initial observations o1 and r1. 8: Infer the posterior qφ(~sASR1 |o1, r1) and sample ~sASR1 . 9: for t = 1, . . . , T do 10: Select action at = π(~sASRt ) +Nt according to the current policy and exploration noise. 11: Execute action at and receive reward rt+1 and observation ot+1. 12: Infer the posterior qφ(~sASRt+1 |o≤t+1, r≤t+1, a≤t) and sample ~sASRt+1 . 13: Store transition (~sASRt , ~st, at, rt+1, ~s ASR t+1 , ~st+1, ot+1) in B. 14: Sample a random minibatch of N transitions (~sASRi , ai, ri+1, ~s ASR i+1 ) from B. 15: Update network parameters using a specified RL algorithm (e.g., DQN or DDPG). 16: Update the model given in (2) with the recorded data from B (according to Section 3). 17: for p = 1, . . . , n do 18: Sample a random minibatch of pairs of (~st, at) from B. 19: Predict (~sASRt+1 , rt+1) according to the model given in (2). 20: Update network parameters using a specified RL algorithm (e.g., DQN or DDPG). 21: end for 22: end for 23: end for
is partially observable: by just looking at the current frame, although we can tell the position of the car, we know neither its direction nor velocity that are essential for controlling the car.
For a fair comparison, we followed the same setting as in Ha & Schmidhuber (2018). Specifically, we collected a dataset of 10k random rollouts of the environment, each consisting of 1000 time steps, for model estimation. The dimensionality of latent states ~̃st was set to d̃ = 32, and regularization parameters was set to λ1 = 1, λ2 = 1, λ3 = 1, λ4 = 1, λ5 = 1, λ6 = 6, λ7 = 10, λ8 = 0.1, which are determined by hyperparameter turning.
Analysis of ASRs. To demonstrate the structures over observed frames, latent states, actions, and rewards, we visualized the learned D~s)o, D~s)r, D~s, and Da)~s, as shown in Figure 9. Intuitively, we can see thatD~s)r andDa)~s have many values close to zero, meaning that the reward is only influenced by a small number of state dimensions, and not many state dimensions are influenced by the action. Furthermore, from D~s, we found that there are influences from ~̃si,t to ~̃si,t+1 (diagonal values) for most state dimensions, which is reasonable because we want to learn an MDP over the underlying states, while the connections across states (off-diagonal values) are much sparser. Compared to the original 32-dim latent states, ASRs have only 21 dimensions. Below, we empirically showed that the low-dimensional ASRs significantly improve the policy learning performance in terms of both efficiency and efficacy.
Ablation Study. We further performed ablation studies on latent dynamics prediction; that is, we compared with the case when the transition dynamics in (4) are not explicitly involved. Figure 10 shows that by explicitly modelling the transition dynamics (denoted by with LDP), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP).
Difference between our SS-VAE and Planet, Dreamer. Both our method and Planet (Hafner et al., 2018) and Dreamer (Hafner et al., 2019) are world model-based methods. The differences are mainly in two aspects: (1) our method explicitly considers the structural relationships among variables in the RL system, and (2) it guarantees minimal sufficient state representations for policy learning. Previous approaches usually fail to take into account whether the extracted state representations are sufficient and necessary for downstream policy learning. Moreover, as for the component of recurrent networks, SS-VAE uses LSTM that only contains the stochastic part, while PlaNet and Dreamer use RSSM that contains both deterministic and stochastic components.
G.2 VIZDOOM EXPERIMENT
We also applied the proposed method to VizDoom (Kempka et al., 2016). VizDoom provides many scenarios and we chose the take cover scenario (Figure 11). Unlike CarRacing, take cover is a discrete control problem with two actions: move left and move right. Reward is +1 at each time step while alive, and the cumulative reward is defined to be the number of time steps the agent manages to stay alive during a episode. Therefore, in order to survive as long as possible, the agent has to learn how to avoid fireballs shot by monsters from the other side of the room. In this task, solving is defined as attaining the average survival time of greater than 750 time steps over 100 consecutive episodes, each running for a maximum of 2100 time steps.
Following the same setting as in Ha & Schmidhuber (2018), we collected a dataset of 10k random rollouts of the environment, each consisting of 500 time steps. The dimensionality of latent state ~̃st
is set to d̃ = 32. We also set λ1 = 1, λ2 = 1, λ3 = 1, λ4 = 1, λ5 = 1, λ6 = 6, λ7 = 10, λ8 = 0.1. By tuning thresholds, we finally reported all the results on the 21-dim ASRs, which achieved the best results in all the experiments.

heading: H DETAILED MODEL ARCHITECTURES
text: In the car racing experiment, the original screen images were resized to 64 × 64 × 3 pixels. The encoder consists of three components: a preprocessor, an LSTM, and an MDN. The preprocessor architecture is presented in Figure 12, which takes as input the images, actions and rewards, and its output acts as the input to LSTM. We used 256 hidden units in the LSTM and used a five-component Gaussian mixture in the MDN. The decoder also consists of three components: a current observation reconstructor (Figure 13), a next observation predictor (Figure 14), and a reward predictor (Figure 15). The architecture of the transition/dynamics is shown in Figure 16, and its output is also modelled by an MDN with a five-component Gaussian mixture. The architecture of the action prediction is given in Figure 17, which is a two-layer MLP taking states and rewards as input and predicted action as output. In the VizDoom experiment, we used the same image size and the same architectures except that the LSTM has 512 hidden units and the action has one dimension. It is worth emphasising that we applied weight normalization to all the parameters of the architectures above except for the structural matrices D(·).
In DDPG, both actor network and critic network are modelled by two fully connected layers of size 300 with ReLU and batch normalisation. Similarly, in DQN (Mnih et al., 2013) on both ASRs and SSSs, the Q network is also modelled by two fully connected layers of size 300 with ReLU and batch normalisation. However, in DQN on observations, it is modelled by three convolutional layers (i.e., relu conv 32× 8× 8 −→ relu conv 64× 4× 4 −→ relu conv 64× 3× 3) followed by two additional fully connected layers of size 64. In DRQN (Hausknecht & Stone, 2015) on observations, we used the same architecture as in DQN on observations but padded an extra LSTM layer with 256 hidden units as the final layer.

heading: I PLATFORM AND LICENSE
text: We run all the experiments on the servers with 4 NVidia V100 GPUs. We used the Car Racing in OpenAI gym and VizDoom environments and we have cited the creators. In our code, we have used the following libraries: Tensorflow (Apache License 2.0), OpenAI Gym (MIT License), VizDoom (MIT License), OpenCV (Apache 2 License), Numpy (BSD 3-Clause "New" or "Revised" License) and NVIDIA-DALI libraries (Apache 2 License).

Review summaries:
This paper proposes to find action-sufficient representation which might be helpful for downstream tasks. The idea is that states and/or observations often include irrelevant for actions information which can be removed by  extracting only reverent information. The method utilizes the concepts of Identifiability, mutual information and sparsity. It is realized by augmenting VAE by Structured Sequential component.  

This paper considered the problem of compressing the state space representation in a manner that is sufficient for control and of smaller dimension than the original state space. This is done by first running the system to collect data and then learning structural constraints that are sufficient for describing the reward random variable. This is done by an adapted Variational Auto Encoder which trades off reconstructability, reward prediction, action sufficiency, and minimality. The resulting low dimensional MDP can be combined with classic model based RL to potentially improve sample efficiency.

The paper proposes a new state representation learning method for RL, called Action-Sufficient State Representations (ASRs) to learn minimal and sufficient state representations for downstream decision-making tasks. Different from other related work, they explicitly characterize structural relationships among variables (i.e., state features, observation, reward and action) in MDP. Such structure and ASRs are learned by constructing a generative environment model (SS-VAE). The proposed algorithm is empirically evaluated in conjunction with the model-free, model-based algorithm in VizDoom and CarRacing. 

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.