Action-Sufficient State Representation Learning for Control with Structural Constraints
Authors: Biwei Huang
Chaochao Lu
Liu Leqi
José Miguel Hernández-Lobato
Clark Glymour
Bernhard Schölkopf
Kun Zhang
Venue: ICLR 2022 Submitted (2022)
Keywords: Representation learning in RL
Minimal sufficient state representations
Graphical model
World model

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 \textit{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.

Under review as a conference paper at ICLR 2022
ACTION-SUFFICIENT STATE REPRESENTATION LEARN- ING FOR CONTROL WITH STRUCTURAL CONSTRAINTS
Anonymous authors Paper under double-blind review
ABSTRACT
Perceived signals in real-world scenarios 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. this paper, we focus on partially observable environments and propose to learn minimal set of state representations that capture sufficient information for decision- making. termed Action-Sufficien 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.
INTRODUCTION
State-of-the-ar reinforcement learning (RL) algorithms leveraging deep neural networks 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 that our perceived signals in real-world scenarios, e.g., images, are usually high-dimensional and may contain much irrelevant information for decision-making maainn 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 single binary variable, while the perceived image high-dimensional. 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 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 al., 2019), stochastic latent actor-critic (Lee et al., 2019), SimPLe (Kaiser et al., 2019), 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, 201 8; Hafner et al., 2018; 2019; 2020).
coanuulation-based methods (Zhang et al., 2021), Dreamer (Hafner
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 generative environment

---

Under review as a conference paper at ICLR 2022
model involving observed states, state-transition dynamics, and rewards, and explicitly characterizing structural relationships among variables in the RL system, we propose 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 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 minimal sufficient set of state representations, termed Action-Sufficcien 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- Aito-EEcoder (SS- VAE), which explicitly encodes structural relationships among variables, for reliable identification of ASRs.
Accordingly, policy learning 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, first formulate generative environment model in partially observable Markov decision process (POMDP), and then show how to explicitly embed structural constraints over variables the RL system and leverage them.
Suppose we have sequences of observations {(o,, de, re) T= where o, O denotes perceived signals at time f, such as high- digheniimn images, with O being the observation space, a₂ E A is the performed action with A being the action space, and r. E R represents the reward variable with R being the reward space. We denote the underlying states, which are latent, by & E S, with S being the state space. We describe the generating process of the environment model follows:
o, f Et),
rt g (3-11 a1 11)
(1)
8: h 1/))
where f. g, and h represent the observation function, reward function, and transition dynamics, respectively, and et, G, and 7), are corresponding independent and identically distributed (i.i.d)) random noises. The latent states 5, form an MDP: given $1-1 and 4--1, S; are independent of states and actions before ( 1-1. 1. Moreover, the action dx-1 directly influences latent states 5,, instead of perceived signals oz, and the reward is determined by the latent states (and the action) as well. The perceived signals o, are generated from the underlying states s,, contaminated by random noise e,- We also consider noise e in the reward function to capture unobserved factors that may affect the reward, well as measurement noise.
commonplace that the action variable a-11 may not influence every dimension of 5,, and the reward F, may not be influenced by every dimension of 5 as well, and furthermore there are structural relationships among different dimensions of s, Figure I gives an illustrative graphical representation, where $3, influences 82,1, at-1 does not have an edge to $3,6, and among the states, only 82,-1 and 83,--1 have edges to r., We use R. E "T- r, to denote the discounted cumulative reward starting from time I, where much immediate rewards are favored over more distant rewards.
[0, is the discounted factor that determines how
To reflect such constraints, we explicitly encode the graph structure over variables, including the structure over different dimensions of and the structures from d:-1 5t, 5.1 to rt,, and &, or-
2

---

Under review as a conference paper at ICLR 2022
O++1
Figure I: A graphical illustration of the generative environment model. Grey nodes denote ob- served variables and white nodes represent unobserved variables. Here, 4t-1 does not have an edge to 83,4+ and only $2,r-1 and 83,1-1 have edges to Tt. and moreover, we take into account the struc- tural relationships among differ- ent dimensions of latent states s,.
Ye++
R
Accordingly, we re-formulate (1) as follows:
0 f(D 2o C (5,
Tt = g Dar O 56-1, Dar O a)),
(2)
8i, h4 (D((s((.) C 51-1, Ak-1.Da-Et (. C 16,),
for io d, where & = ($1,e, ,d,r)
denotes element-wise product, and Dey are binary
matrices indicating the graph structure over variables. Specifically, Dzo E {0, 1) axi represents the graph structure from d-dimensional & to ot, Dar E {0, 1),1 the structure from 5 to the reward variable rt. Darr E {0. 1) the structure from the action variable at-1 the reward variable Ta, Dg E {0, 1] jaxd denotes the graph structure from d-dimensional 5 to d-dimensional &, and Dac.s) is its i-th column, and Davs E {0, 1 I]lid corresponds to the graph structure from 4:-1 to 8 with Davs(-i) representing its i-th column. For example, Dzy..) means that there is no edge from sj,1-1 to $4,,, Here, we assume that the environment model, as well as the structural constraints, is invariant across time instance I.
2.1 MINIMAL SUFFICIENT STATE REPRESENTATIONS
Given observational sequences { at, rz)] we aim to learn minimal sufficient state representa- tions 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 relation- ships, is given. Then we provide criteria to achieve sufficiency and minimality of the estimated state representations, when only { de, re) Y but not the environment model, is given.
Finding minimal sufficient state dimensions with given environment model. RL agents learn to choose appropriate actions according to the current state vector & 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-Sufiiient state Representations (ASRs), and denote it by g,AR We can leverage the (conditional) independence/dependeeee 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/depeeddennc 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 I, which is assumed to be Markov and faithful to the measured data, each dimension in SASR has direct or indirect edge Tttr, for >> 0.
Proposition 1 can be shown based on the global Markov condition and the faithfulness assumption, which connects d-separation' I to conditional independence/deeeeddeee relations. proof is given in Appendix. If si, has a (direct or indirect) edge to T++ the action variable a; is dependent on s when the cumulative reward maximized, which can be thought of as a data selection procedure depending on the cumulative reward, and thus 8,, is needed for action determination. On the other hand, if 8 does not have an edge rt+r, then the action variable a, is (conditionally) independent
A path p is said to be d-separated by set of nodes Z if and only if (1) p contains chain -> m or fork + j such that the middle node m is in Z. or (2) p contains collider m +- j such that the middle node m is in Z and such that descendant of in Z.
3

---

Under review as a conference paper at ICLR 2022
on St,e, and thus St, is not needed for action determination. According the above proposition, is easy to see that for the graph given in Figure 1, we have 5yAR ($2,2, $3,c) That is, we only need (82,2,83,1) ", instead of 5;, for the downstream policy learning.
Minimal sufficient state representation learning from observed sequences. In practice, usually do not have access to the latent states or the environment model, but instead only the observed sequences {(o, de,, Tt)) ) Then how can we learn the ASRs from the raw high-dimensional inputs such images? We denote by the estimated whole latent state representations and EASR C & 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, [5 where I denotes mutual information. Since 7 ($ASR, a R&+1) H(a Re+1) H (ag EASR R++1), where (.) denotes the conditional entropy, and the first term H (az Re41) does not contain ASRs, we can estimate ASRs by minimizing H (at 5ASR RRa-)), with
Re-1);
EASR (a+l 8t
R(+1)= --=1=--2 - Byga (log Po (or 8% EASR Re+1]]= - Fagna (log (ax) [ BASR Re+1)}
where Po denote the probabilistic predictive model of az with parameters a, Tao is the joint distri- bution over št and a, with Qáa qoPa. and qo (5 11 01:t-1) is the probabilistic inference model & with parameters and y: = (oī . ) and D ASR E {0, 10 is a binary vector indicating which dimensions of 5, are in JASR, so D ASR O St gives ASRs SASR
Moreover, we achieve minimality of the representation by minimizing conditional mutual information between observed iighddimensional signals y, and the ASR SASR at time t given data at previous time instances, and meanwhile minimizing the dimensionality ASRs with sparsity constraints:
A1
EASR y1a-1, 01 8 + Az ASR Ty;;;
2-
where the conditional mutual information can be upper bound by KL-divergence:
EASR I(y 8t
y 1: 1, 1,,11,0111111 1) ( 5 Vi.t, 01:t-1) 1) log
4i 1 B. 011. d5ASR
p- (6,58 F- Dovn!
KI.
Go [ASR
(8( 1 11.,, -1---1) Ilp, (5
(AASR) 3 a1-1; Dg; Darz)):
with Py being the transition dynamics of with parameters 7.
Furthermore, Proposition shows that given the (estimated) environment model, only those state dimensions that have direct or indirect edge to the reward variable are the ASRs. In our learning procedure, also take into account the relationship between the learned states 5, and the reward, and leverage such structural constraints for learning the ASRs. Denote by DASR E {0, dxx binary vector indicating whether the corresponding state dimension in & has a direct or indirect edge to the reward variable. Consequently, we enforce the similarity between DASR and DASR by adding an La norm on DASR - 5ASR Therefore, the ASRs can be learned by maximizing the following function:
mmin sull
As
BASR Byna log ( ASR R o B., Ret1)} DDaI) ASK -SAAKll
Suftieiency
- T (L
1. and (D ASR 3. Dows)) - >a D
Minimality
(3)
where A's are regularization terms, and note that AASR can be directly derived from the estimated structural matrices Dørr and D((.))- The constraint (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 repre- sentations 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-
4

---

Under review as a conference paper at ICLR 2022
mation provides interpretable and intuitive picture of the generating process. 2) Accordingly, 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 RI, which effectively reduce possibly risky explorations.
3 STRUCTURED SEQUENTIAL VAE FOR THE ESTIMATION OF ASRS
In this section, we give estimation procedures for the environment model and ASRs, as well as the identifiability guarantee in linear cases.
Identifiability in Gausar-G Cases. Below, we first show the identifiability guarantee in the linear case, as special case of Eq. (2). In the linear case (see the environment model given in Eq. (5) in Appendix D), Davas Dār,. Darr. Dz, and Da-s are linear coefficients, indicating corresponding graph structures and also the strength. Denote the covariance matrices of ez and e, by De and E. respectively. Further let Dz == (D s.g D The following proposition shows that the environment model in the linear case identifiable up to some orthogonal transformation certain coefficient matrices from observed data {(o, ar: Y
Proposition 2 (Identifiabilitty) Suppose the perceived signal on, the reward rt, and the latent states 8. follow a linear environment model. If asssumptions A Al~A4 (given in Appendix C) hold and with the second-order statistics of the observed data {(o,, (or,a,,Td) (t,Te) the noise variances De and E., D De-rs Dz aa. (with k 2 0), and Dz.o DT 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 (Droor Dæ.,, Dorr Dz, Da Ee; Mee and that with (Dinor Dznr, Durry Dz, Dans, Le, Ee) are observationally equivalent. Then we have Dr D, Dar-Darrs D==UT D5U, Daos= Do.sU, Ex Σx=E,. and Σ = E,, where U is an orthogonal matrix.
General Nonlinear Cases. To handle general nonlinear cases with the generative process given in Eq. (2), we develop Structured Sequential VAE (AS- to learn the model (including structural constraints) and infer latent state representations and ASRs JASK, with the input {(o,, ae, rx) ) Specifically, the latent state dimensions are organized with structures, captured by Dz. 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æto- Dær, Derr: Da-s) into SS-VAE. Moreover, we aim to learn state representations R and ASRs RAAR that satisfy the following properties: (i) & should capture sufficient information of observations ot, Tt, and a,, that is, should be enough 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) 5AARR are minimal sufficient state representations for the downstream policy learning.
Let y1.7 = ((oT TT To achieve the above properties, we maximize the following objective function:
T-2
C(y (0, d. 7, D())
Σ
Egg loo (o. Ni; Dino)
log pe (re+1 1 asi Denri Daar)
Reconstraction
+ log po (0c+1) (1)) + log po (Tet2 (ret2l3,, a(+1) + As
E Bggn
log Pa (asl &e, Retl: D ASR))
Prediction
Sufliciency
A
KL B yia
(3.)) 1; Ds. Ders)
>al D ASRRIi
Transision
Conditional disentanglomeen Minimalayy
2Aol + Aall lI + D+Il Da-ell + DaIlDAsK ASR Ask (s).
(4) which contains the reconstruction eтror 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
Spansit
5

---

Under review as a conference paper at ICLR 2022
learning, well sparsity constraints the graph structure. We denote by pe the generative model with parameters and structural constraints D(-). do the inference model with parameters Φ, Py the transition dynamics, and Por action prediction with ASRs given the cumulative reward R2+1- Each factor in Pyy and Oe modeled with a mixture of Gaussians (MoGs), approximate a wide class of continuous distributions.
Below 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 and the reward rt-
Transition component: To achieve the property that state representations satisfy an MDP, we explicitly model the transition dynamics: log Py 11: D;, Da-s). In particular, & 5/ is modelled with mixture of Gaussians: Σ _ TNNVN, (3-1) Eg (3--111))))) where K is the number of mixtures, f4x and Ex(+) are given by multi-layer perceptrons (MLP) with inputs -1 and a;-1, parameters ", and structural constraints D3 and De.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 8, 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 log P(y to achieve conditional disentanglement between 4 ( and 4s (5,,1-) for # 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 y; and SAAR and structural constraints, For details, please refer Section 2.1.
Sparsity constraints:
According to the edge-minimality property (Zhang & Spirtes, 2011), we additionally put sparsity constraints 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 -- a1:t-1), includes a Long Short- Memory (LSTM (Hochreiter & Schmidhuber, 1997)) to encode the sequential information with output h4 and Mixture Density Network (MDN (Bishop, 1994)) to output the parameters of MoGs. At each time instance, the input (0t+1, 0z) is projected to the encoder and a sample of 5 is inferred from ds as output. The generated sample further acts an input the decoder, together with 0+41 and structural matrices D2.0 Dz... and Darr. Then the decoder outputs 6 and ft+2. Moreover, the state dynamics which satisfies Markov process and is embedded with structural constraints D3 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 by sampling and then jointly learn all parameters by maximizing using stochastic gradient descent.
POLICY LEARNING WITH ASRS
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 fixed, regardless of the behavior policy that 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 S EASR from past observed sequences Tee. Decerr1 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

---

Under review as a conference paper at ICLR 2022
10 Peaa
h
h
LSTM
LSTM
LSTM
Encoder
MON
MDN
MDN
Dynamics
Figure Diagram of neural net- work architecture to learn state representations. The correspond- ing structural constraints are in- volved in "Deconv" and "MLP", and "AP" represents the action pre- diction part for sufficient state rep- resentation learning.
MILP
MON
MLP
MON
$14-2
Hi+1
Re+3
G2+2
AP
Decoder
Deconv
Deconv
Deconv
Gc41
Ax12
0,, Pr+-1
θ.t1, f
0/42: F243
et al.. 2015) and deep deterministic policy gradient (DDPG) for continuous actions (Lillierap et al., 2015). Algorithm in Appendix 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 the generative environment model, one 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 in Appendix F.
5 EXPERIMENTS
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 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.
5.1 CARRACING EXPERIMENT
CarRacing is continuous control task with three continuous actions: steering left/right, acceleration, and brake. Reward is -0.1 every frame and +1000/N /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 št was set to d 32, determined by hyperparameter tuning.
Analysis of ASRS. To demonstrate the structures over observed frames, latent states, actions, and rewards, visualized the learned Dž Ds.r, Dz. and Da-s. shown in Figure 9 in Appendix G. Intuitively, we can see that Dž-r and Da-i have many values close to zero, meaning that the reward is only influenced by small number of state dimensions, and not many state dimensions are influenced by the action. Furthermore, from D3. we found that there influences from 3 to 51,411 (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- ooff-diagonal values) are much sparser. Compared the original 32-dim latent states, ASRs have only 21 dimensions. Below, we empirically showed that the fow-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) (Lillierap al., 2015) and model-based (Dyna and Prioritized Sweeping) algorithms (Sutton, 1990) to ASRS (with 21 21-dims). As shown in Figure 3, interestingly. by taking advantage of the learned generative model, model-based ASRs is superior model-free ASRs at faster rate, which
7

---

Under review as a conference paper at ICLR 2022
1200
Modul-fi ou ASRs
mcan
Curuultiive Rewwrrd
1000
Model baged ASRS (Dyna)
900
max
Model bases ASAo (PS)
Fitness Valuue
800-
VRL
600
Hultet
SLAC
AS ReRbbea augh angllcoree
po Ma Mp
400-
200
5.00
10:00
1500
2000
Episodes
Generation
Figure 3: Cumulative rewards of model-based ASRs, model-free ASRs, VRL, SLAC, PlaNet. DBC and Dreamer evaluated on CarRacing.
Figure 4: Fitness Value of ASRs compared to world models evaluated on CarRacing, including mean score, max score, and the best average score.
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 improve sample-efficienyy
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., 201 1018), 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
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 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 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 '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 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.
the Covariance-Maii
Comparison with Dreamer and DBC with Background Distraction. We further compared ASRs (with 21-dims) with Dreamer and DBC when there 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 fair comparison, we require all of them to have the same latent dimensions and have the similar model capacity. As shown 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, compared with the case when the transition dynamics in Eq. (4) not explicitly modeled, but is replaced with a standard normal distribution. Figure 10 in Appendix G shows that by explicitly modelling the transition dynamies (denoted by with LDP), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP).
5.2 VIZDOOM EXPERIMENT
We also applied the proposed method to VizDoom take cover scenario (Kempka et al., 2016), which is 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 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 raw observations, we further compared with another common approach to
8

---

Under review as a conference paper at ICLR 2022
1200
Model
Cumulative Rewards
Cumalativv Reward
1000
1200
mean
800
best ave poore ASR ASRimmen
Dreamer
621+:124.5 803::112.5
Fitmas
2000
DBC
Model-free free ASRs
938 11
DON VAL
Model-based ASRs
9541:068.8
BRok Sibseevatiion
200
DON on Obsorvations
Model pesied ASRt
Figure 5: Comparisons with
5.00
1000
1500
100
Dreamer and DBC in CarRacing with natural video distractors, af- ter 2000 training episodes, with SOTA methods evaluated on Viz- (with CMA-FES) compared to world models evaluated on VizDoom.
Episodes
Generation
Figure 6: Comparing ASRs and Figure 7:
Fitness value of ASRs
standard error.
Doom.
POMDPs: DRQN (Hausknecht& & Stone, 2015). As shown in Figure 6, DQN on ASRs achieve a much better performance than all other comparisons, and 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) ASRs (with 21 21-dims). As shown in Figure we can draw the same conclusion that by taking advantage of the learned generative model, model-based ASRs is superior to model-free ASRs a faster rate. We also applied ASRs world models, where Figure shows that our method with ASRS (denoted by ASR-*) achieves a better performance.
6 RELATED WORK
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 '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 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 al., 2017; Ha & Schmidhuber, 2018; Hafner et al., 2018; Zhang et al., 2019 9b; 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-baset methods have also attracted much attention (Castro, 2020; Zhang et al., 2021).
On the other hand, our work is also related 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, show its advantages in POMDPs.
7 CONCLUSIONS AND FUTURE WORK
In this paper, we develop a principled framework to characterize 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 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-Caree
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.

---

Under review as a conference paper at ICLR 2022
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 given in Appendix E and Appendix F, respectively. Appendix
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. description of the network architectures used is included in Appendix H. Appendix introduces the platform and license. Source code for all our experiments will be made available on GitHub after deanonymization, providing also complete description of our experimental environment, configuration files and instructions on the reproduction of our experiments.
REFERENCES
Y. Bengio. The consciousness prior. arXiv preprint vrXiv:1709.08563 2019.
C. M. Bishop. Mixture density networks. In Technical Report NCRG/4288, Aston University, Birmingham, UK. 1994.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. AAAI, 2020.
Ebert, C. Finn, A. X. Lee, and S. Levine. Self-supervised visual planning with temporal skip connections. ArXiv Preprint ArXiv:1710.05268 2017.
C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference Machine Learning (ICML), 2019.
D. Ghosh, A. Gupta, and Levine. Learning actionable representations with goal conditioned policies. ICLR, 2019.
K. Gregor, G. Papamakarios, F. Besse, L. Buesing, and T. Weber. Temporal difference variational auto-encoder. arXiv preprint 1xXxv:1806.03102 2018.
D. Ha and J. Schmidhuber. World models. In Advances in Neural Information Processing Systems, 2018.
D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. arXiv preprint /8/iv: /tt::////.000 1.04551, 2018.
D. Hafner, T. Lillicrap. J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603, 2019.
D. Hafner, T. Lillicrap. M. Norouzi, and J. Ba. Mastering atari with discrete world models. arXiv preprint ar drXiv22010.02193, 2020.
N. Hansen. The CMA evolution strategy: A tutorial. arXiv preprint 4rXiv:1604.00772,200000 2016.
M. Hausknecht and P. Stone. Deep recurrent q-learning for partially observable mdps. In 2015 AAAI Fall Symposiun Series, 2015.
S. Hochreiter and J. Schmidhuber, Long short-term memory. Neural computation, 9(8):1735-1788 1997.
B. Huang, K. Zhang, M. Gong, and C. Glymour. Causal discovery and forecasting in nonstationary environments with state-space models. In International Conference on Machine Learning (ICML), 2019.
B. Huang* K. Zhang* J. Zhang, J. Ramsey, R. Sanchez-Romero, C. Glymour, and B. Schölkopf. Causal discovery from hetenngencous/noottationaar data. JMLR, (89), 2020.
L. Kaiser, M. Babaeizadeh, P. Milos. B. Osinski, H. Campbell, K. Czechowski, ..., and
H. Michalewski. Model-based reinforcement learning for Atari, arXiv preprint vrXiv:1903.00:: 2019.
10

---

Under review as a conference paper at ICLR 2022
M. Karl, M. Soelch, J. Bayer, and P. van der Smagt. Deep variational bayes filters: Unsupervised learning state space models from raw data. arXiv preprint arXiv:1600 06432, 2016.
W. Kay, J. Carreira, K. Simonyan, B. Zhang, C. Hillier, S. Vijaanarasimman, F. Viola, T. Green, T. Back, P. Natsev, et al. The kinetics human action video dataset. arXiv preprint arXiv: (eX.ix:1705.0695000 2017.
M. Kempka, M. Wydmuch, G. Runc, J. Toczek, and W. Jaskowski. ViZDoom: A Doom-based AI research platform for visual reinforcement learning. In IEEE Conference Computational Intelligence and Games, pp. 341 144-34,, Santorini, Greece, Sep 2016. IEEE. URL org/ org/abs/ / 605. 02097 97. The best paper award.
D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
O. Klimov. Carracing-vO. http://sym.openai.com// 2016.
R.G. Krishnan, U. Shalit, and D. Sontag. Deep kalman filters. arXiv preprint arXiv:1511.05111 2015.
D. Kulkarni, Saeedi, Gautam, and S. Gershman. Deep successor reinforcement learning. arXiv preprint (rXiv:/606.023900 2016.
A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. arXiv preprint arXiv:/907.0m/ 00953,, 2019.
T. Lesort, N. Díaz- Rodrigodrigee, J. F. Goudou, and D. Filliat. State representation learning for control: An overview. Neural Networks, 108:379-392, 2018
T. P. Lillierap, J. Hunt, A. Pritzel, N. Heess, T. Erez, Y. Tassa, D. Silver, and D. Wierstra. Continuous control with deep reinforcement learning. arXiv preprint
2015.
M. L. Littman and R. S. Sutton. Predictive representations of state. In In Advances in neural information processing systems, pp. 1555-1561, 2002.
S. Mahadevan and M.. Maggioni. Proto-value functions: laplacian framework for learning representation and control in markov decision processes. Journal of Machine Learning Research (JMLR), 8:2169-2231, 2007.
2007.
V. Mnih, K. Kavukcuoglu, Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing atari with deep reinforcement learning, arXiv preprint arXiv:1112.5602, 2013.
V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Ried- miller, Fidjeland, Ostrovski, Petersen, C. Beattie, A. Sadik, I. Antonoglou, H. King, D. Kumaran, D. Wierstra, S. Legg, and D. Hassabis. Human-level control through deep reinforce- ment learning. Nature, 518(7540):529-533 2015.
Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, Cambridge, 2000.
B. Schölkopf. Causality for machine learning. arXiv preprint arXiv: 1911.0000 2019.
R. Shu, T. Nguyen, Y. Chow, T. Pham, K. Than, M. Ghavamzadeh, S. Ermon, and H. Bui. Predictive coding for locally-linear control. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research. PMLR, 2020.
S. Singh, M. R James, and M. R Rudary. Predictive state representations: A new theory for modeling dynamical systems. In In Proceedings of the 20th conference on Uncertainty artificial intelligence, pp. 512-519, 2004.
P. Spirtes, C. Glymour, and R. Scheines. Causation, Prediction, and Search. Spring- Lectures in Statistics, 1993.
A. Srinivas, M. Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforce- ment learning. ICML. 2020.
11

---

Under review as a conference paper at ICLR 2022
A. L. Strehl, C. Diuk, and M. L. Littman. Efficient structure learning in factored-state mdps. In AAAI, 2007.
R. Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pp. 216-224. Elsevier, 1990.
R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
N. Tishby, F. C. Pereira, and W. Bialek. The information bottleneck method. The 37th annual Allerton Conference on Communication, Control, and Computing, 1999.
A. Tsividis, T. Pouncy, J. L. Xu, B. Tenenbaum, and S. J. Gershman. Human learning in atari. In AAAI Spring Symposium Series, 2017.
M. Watter, J. Springenberg. J. Boedecker, and M. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. NeurIPS, 2015.
L. Wiskott and T. J. Sejnowski. Slow feature analysis: Unsupervised learning of invariances. Neural Computation, 14(4):715-770, 2002.
A. Zhang, Z. C. Lipton, I.. Pineda, K. Arizzadenesheli, A. Anandkumar, L. Itti, J. Pineau, and
T. Furlanello. Learning causal state representations of partially observable environments. arXiv preprint ar arXiv:I906. 10437, 2019a.
A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. ICLR, 2021.
J. Zhang and P. Spirtes. Intervention, determinism, and the causal minimality condition. Synthese, 183(3):335-347, 2011.
K. Zhang and A. Hyvärinen. A general linear non-gaussian state-space model: Identifiability, identification, and applications. In Asian Conference on Machine Learning, pp. 113-128, 2011.
K. Zhang M. Gong*, P. Stojanov, B. Huang, Q. Liu, and C. Glymour. Domain adaptation problem of inference on graphical models. In NeurIPS, 2020.
M. Zhang, Vikram, L. Smith, P. Abbeel, M. J. Johnson, and Levine. Solar: deep structured representations for model-based reinforcement learning. arXiv preprint 18Xi:: 2015 18.
M. Zhang, S. Vikram, L. Smith, P. Abbeel, M. Johnson, and S. Levine. Self-supervised visual planning with temporal skip connections. ICML, 2019b.
12

---

Under review as a conference paper at ICLR 2022
APPENDICES FOR "ACTION-SUFFICIET STATE REPRESENTATION LEARNING FOR CONTROL WITH STRUCTURAL CONSTRAINTS"
A PROOF OF PROPOSITION
We first give the definitions of the Markov condition and the faithfulness assumption, which will be used in the proof.
Definition (Global Markov Condition (Spirtes et al., 1993; Pearl, 2000)), The distribution p over V satisfies the global Markov property on graph G ifr any partition (A, B,C) such that B d-separates A from C,
Definition 2 (Faithfulness Assumption (Spirtes et al., 1993; Pearl, 2000)). There are no independen- cies between variables that are not entailed by the Markov Condition.
Below, we give the proof of Proposition 1.
Proof. We first show that if st, E 5ASK, then it has a direct or indirect edge to Tt+r-
We prove it by contradiction. Suppose that 8i, does not have direct or indirect edge to Te+r- According to the Markov assumption, $i,, is independent a given R. Hence, Si,t is not necessary for decision making, and thus 8₁ is not dimension in 5ASR, which contradicts the assumption. Since we have contradiction, must be that 8i,t has a direct or indirect edge to Ttar-
We next show that if 81, has direct or indirect edge to rt+t. then 81, E E,ASR
Similarly, by contradiction suppose that 8i, not a dimension in 5ASRR means that 81, independent de given R and some other variables. Then according to the faithfulness assumption, 84, does not have direct or indirect edge rt+r, which contradicts to the assumption.
B MINIMALITY OF THE REPRESENTATION
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 rr'), ). and the ASR eASR at time / given data
observed high- high-dimensional signals yı. where y: = =(of", ,r{]
previous time instances, and meanwhile minimizing the dimensionality of ASRs with sparsity constraints:
A 1=2
Tly:; TYy: y1a-1. 3 Aa |SRR
EASR
Note that in the above conditional mutual information, we need to conditional on the previous states 5 instead of EASR which two give different conditional mutual information. can be shown by contradiction. Suppose [yy;; (1 11:1-1 ./- I [yy:; fy 01:t-1, 8 JASK), denote C ASS Then the equivalence implies that 3 is independent of o, (where o; yt)
and
given {y 01::1111 5A5R) It is obviously violated for the example given in Figure 1, where C s and ASR {sg.. sa). and $1 dependent on o, given (y1.t -1.-1 $2, $2,1-1 $t, Hence, conditioning on _ and AASR give different conditional mutual information. Therefore, in the above conditional mutual information, need condition the previous states ø1-1-
13

---

Under review as a conference paper at ICLR 2022
Moreover, the conditional mutual information [ 2ASR by a KL-divergence:
V1--1.Q1--1,8 can be upper bound
I [y:: S Y1:-1. 01:--1,,
5 FASR { lit- 01:--
(yr)) H ASR
01:1-1 (11-1,
Y1--1,
-y 1, 01:00 5/
--p(yylyyy- 01:-1 1o (h0p::/.....--0.1100000000001110
+ p(y+ 15AS
- 1 logg)
3 :
1, 1
S -P(y, 1 EASR
, Y1
log p(y
+ P(y: FASR
-1 log ply $t
EASR Y1s-1 01 3
= I le (5g
ASR
L1., 01:-1) 81-
(5, 111 dzyAS
p-. (SA & 1,a; 1:Dx,Darx)
= KL (3AAS))
RASR
(FASR 11 4111 Dg. D0.s))
with py being the transition dynamics of with parameters 7.
C ASSUMPTIONS OF PROPOSITION 2
To show the identifiability of the model in the linear case, we make the following assumptions:
A1. d, + d da, where |ox| - do, |ra| = d,, and |sel = ds. A2. (D DT..) is full column rank and Dy is full rank.
A3. The control signal a, is i.i.d. and the state & is stationary.
A4. The process noise has unit variance, i.e., var(m): =
D PROOF OF PROPOSITION 2
Proof. The proof of the linear case without control signals has been shown in Zhang & Hyvärinen case with control signals:
(2011). Below, we give the identifiability proof in the
0: D5.0 + Ct,
Tt-1 = D.. 8t + Darr + E241s
(5)
Dz81_ + Da-su + T/c
and ē, "eTleT Then Let Yt+1 T0, r(+11T Da. DDI,0 DT IT Door 10 DT
the above equation can be represented as:
{
15 + Dar + Et,
S = D; Da-5at-1 T)-
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 Yr+k and oz:
{
Cov(ye+k, at) D Daig Var(at), Cov(ye+k, a.) Darr Var(a:),
if k: > 0. if k = 0.
Thus, from the cross-covariance between Y1+k and az, we can identify D Dava, Dan, and D for k 0.
Next, we consider the auto-covariance function of 5. Define the auto-covariance function of at lag & as R2(k) = E[ and similarly for (k). Clearly, Rx(-k) = Rg(K)* and (-k- = (k)T Then have
Ra(k) R2(k -) 1)- DT.
if k 0,
(8)
(k) = [ (1) DJ Var a.s T
if k = 0.
Below, we first consider the case where d, + d, Let D Di Yt = ÿ: + B Darr 1-1 +
and Ry (k) D...R. R,; TT. Ry (k) satisfies the recursive property:
{ Ry (k: = Rk 1)
if > 0,
Ry (k) = R2 (1) Ω T + Buno Da.svVa((at-))DD)) TT 200
(9)
if = 0,
14

---

Under review as a conference paper at ICLR 2022
where ! DDD
Denote Sk Di Dr Da-s" Var(ar). Then we can derive the recursive property for Ry (k):
Ry Ry(k 1) -Q- Du. 12 ST
St. 22
if k > 1,
(y( Ry(k: 1) 1)-Q2 122== Daor Var (ay (ay) Dan
-EA Darrk
if k = 1,
(k(k) (1) -2 +Dwwr Vaunvaar(a D ger
En)
(eto(DasVaa(on) and + 1) D)0300 3000
if
When & 2, we have
(2 R (1) (1) Sa. ST Dan
The above equation can be re-organized as
(Ry - Da.r 5 = (Ry(1) - Dan 8T) -
Because Dar and S₁ are identifiable, and suppose (Ry (1) Dar (T) invertible, Ω DDDDDDD s.o is identifiable.
We further consider Ry and Ry and write down the following form:
Ryy(0 ( a. Da.>Var(aa (iraa-1)D 1)D am (1(
[ (1) Ry (0)
Dan, Var(ar))
I
22 + -Day, Var ( Sa
+ En
-S2
From the above two equations can then identify Ee and (D..sVaa((aa + I)D &o. and because Da.o is identifiable, D identifiable.
In summary, we have shown the identifiability of Dar. Dzo Da.g. D Da-g: DI 2oo Eg. Furthermore, D.... Da. and Da-s are identified up to some orthogonal transformations. That is, suppose the model in Eq. (3) with parameters (Dz.: Dæ.r, Dair, Dz, Dass, Ee, E.) and that with (Dooor Dar, Dan, Dæ, Da-si Ex, Ex) are observationally equivalent, we then have Dæo Dz D =D Davr, D3 U DzU, Dais Da,sU, Ea Ec. and E. Ec, where U is orthogonal matrix.
DT and
Next, we extend the above results to the case where d, +d, > dg. Let D (.0(1.) be the i-th row of D Recall that Dæ is of full column rank. Then for any i, one can show that there always exist d rows of D. such that they, together with DJ Po0(..); form dg x d, full-rank matrix, denoted by Dzao(t.)- Then from the observed data corresponding to D. Dito(t.)) determined up to orthogonal transformations. Thus, Dr.. is identified up to orthogonal transformations. Similarly, Da.r. Da, and Da-s are identified up to orthogonal transformations. Furthermore, Cov(D &e + Darra) is determined by Dr Dar. Dz, and Da-g. Because Cov(y,) Cov(Da.0 + Da.r(i) + Eg. Eg is identifiable.
One may further add sparsity constraints on D De.r. Dz. and D to select more sparse structures among the equivalent ones. For example, one may add sparsity constraints on the columns of z-o- Note this corresponds to the mask on the elements of & in Eq. 2; if the full column is 0, then the corresponding dimension of 5, is not selected.
E
MORE ESTIMATION DETAILS FOR GENERAL NONLINEAR MODELS
The generative model can be further factorized as follows:
log pa m 1:7
BLT Q1:7-1; D 20: Dans Darr)
= logpo(or.1 Si.: Dzvo) + log (r1:T Fi.T a1:r-1; Dzr, Dar)
(10)
Σi lo (o Dzo)) + log pa (re 1 ar- Dzr, Darr)
15

---

Under review as a conference paper at ICLR 2022
where both po(0; $ Dz.o) and pr(rel5111, 01-1: Dær, Da-r) are modelled by mixture of Gaussians, with Dao indicating the existence edges from 3 or and Dær indicating the existence of edges from 5t-1 rt-
The inference model qo (f1:rll 01:T-1) is factorized as
log go (1ir NI.T 01:T-1)
T
= loggs yy, an) + Σ log qs (1) Qna-1),
where both go(šily), ( and q (5. -1 1, Vi.r, 01: modelled with mixture of Gaussians. The transition dynamics Pyy is factorized as
Dz(.)) Du-s(.))) log
Dz(.))) D ().s
(11)
with 3 modelled with mixture of Gaussians.
Thus, the KI. divergence can be represented as follows:
air 112
T
(12)
KL(go(3ily),. ao) lol,, (51) + Σ In
In practice, KL divergence with mixture Gaussians hard implement, so instead, we used the following objective function:
T Σ
KL (L (31ly ao)llp (31)) +
Ey
(1--1))ll (=1
1=2
(13)
(0111111111 4-11 Da..o, (a)
where py is standard multivariate Gaussian (ū, Ia).
F MORE DETAILS FOR POLICY LEARNING WITH ASRs
Algorithm 1 gives the procedure of model-free policy learning with ASRs in partially observable environments. Specifically, starts from model initialization (line 1) and data collection with 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 &, on lines and 12, we infer q, (5 pog++ Fg++1, age) and sample EASR from the posterior. The sampled ASRs are then stored in the buffer (line 13). Furthermore, we randomly sample 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)) Q-learning (Mnih et al., 2015).
Algorithm presents the procedure the classic Dyna algorithm with ASRs. Lines 17-22 make use of the learned environment model to predict the next step. including =ASR and rt+1, and update 20. Based
the function times. Specifically, in our implementation, the hyper-parameter the learned model, the agent learns hehaviors from imagined outcomes in the compact latent space, which helps to increase sample efficiency.
G ADDITIONAL EXPERIMENTS AND DETAILS
G.1 CARRACING EXPERIMENT
CarRacing (with illustration Figure 8) is 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
16

---

end for Under review as a conference paper at ICLR 2022
Algorithm 1 Model-Free Policy Learning with ASRs Partially Observable Environments
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 Section 3).
4: Identify indices of ASRs according to the learned graph structure and the criteria in Prop. 1.
5: for episode I, ..... M do
6: Initialize a random process / for action exploration.
7: Receive initial observations o1 and f1.
8: Infer the posterior qs 5ASR (1, r1) and sample SYAR,
9:
for tort=1, 1, T do
10:
Select action a (T(AAR)) M, according to the current policy and exploration noise. Execute action a₂ and receive reward Ft-1 and observation 0,+1-
11:
Infer the posterior 90 (5 oge+n. agg) and sample gASR
13:
Store transition (3AS. dt: Tt+1, SASR) B.
14:
Sample random minibatch of transitions (Æ,ASR a, Ts+1, ASR) from B.
15:
Update network parameters using specified RL algorithm (e.g., DQN DDPG).
end for
17:
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 01 and F1.
8: Infer the posterior 4s (5ASR (01, r1) and sample SYASR,
9:
T do Select action a₂ = (T(AAAR) + Mi according to the current policy and exploration noise.
10:
11:
Execute action a; and receive reward Ft-1 and observation 01+1-
12:
Infer the posterior qui (5ASB Og(41, age) and sample $,+1
EASR
13:
Store transition (5ASR S2, az, Te+1, 5ASR
5 52+1,0+11 in B.
14:
Sample random minibatch of N transitions (S a, Ts+1, s from B.
Update network parameters using a specified RL algorithm (e.g.. DQN or DDPG). Update the model given in (2) with the recorded data from B (according to Section 3).
16:
for p .,.., do
18:
Sample random minibatch of pairs of (5,, az) from B.
19:
Predict (#ASR, r2+1) according to the model given in (2).
20:
Update network parameters using specified RL algorithm (e.g., DQN or DDPG).
21:
end for
22:
end for
end for
PL
Figure 8: An illustration of Car Racing environment
17

---

Under review as a conference paper at ICLR 2022
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 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 & was set to d 32, and regularization parameters was set to A = 1, X2 = 1, A = 1, A4 = 1, As = 1, A = 6, X7 = 10, Ag- = 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 Di Divr. Dz. and Da.s: as shown in Figure Intuitively, we can see that Dar and Da.s have many values close to zero, meaning that the reward is only influenced by small number of state dimensions, and not many state dimensions are influenced by the action. Furthermore, from Dg. we found that there are influences from $1, $1141 (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 the original 32-dim latent states, ASRs have only 21 dimensions. Below, we empirically showed that the low-dine wow-diensionnl ASRs significantly improve the policy learning performance in terms of both efficiency and efficacy.
D2
Dg
Dz-j
D;
1.00
1.00
1.00
1.0
0.8
0.75
0.75
0.75
10
0.6
12 15
0.6
0.2
0.00
0.00
0.00
30
0.9
15
15
30
10 15
30
Dimonsion
Dimonston
Dimenzion
Figure Visualization of estimated structural matrices Divor Dārr Da-si and Dg in Car Racing.
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 L.DP), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP).
Difference between our VAE and Planet, Dreamer. Both our method and Planet (Hafner 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, 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 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 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 maximum of 2100 time steps.
Following the same setting in Ha & Schmidhuber (2018), we collected dataset of 10k random rollouts of the environment, each consisting of 500 time steps. The dimensionality of latent state ă,
18

---

Under review as a conference paper at ICLR 2022
1200
Cumulative Reward
with LDP
1000
without LDP
800
600
400
200
0
500
1500 Episodes
1000
2000
Figure 11: An illustration of VizDoom lake cover scenario.
Figure 10: Ablation study of latent dynamics prediction (LDP) evaluated on Car Racing with model-free ASR.
is set to d = 32. We also set λ =1 12 1,, - = As 0.1. By tniing thresholds, we finally reported all the results on the 21-dim ASRs, which achieved the best results in all the experiments.
H DETAILED MODEL ARCHITECTURES
In the car racing experiment, the original screen images were resized to 64 x 64 x 3 pixels. The encoder consists of three components: preprocessor, an LSTM, and MDN. The preprocessor architecture presented in Figure 12, which takes as input the images, actions and rewards, and its output acts the input to LSTM. We used 256 hidden units in the LSTM and used five-component Gaussian mixture in the MDN. The decoder also consists of three components: current observation reconstructor (Figure 13), next observation predictor (Figure 14), and reward predictor (Figure 15). The architecture of the transition/dynamies shown in Figure 16, and its output is also modelled by an MDN with five-component Gaussian mixture. The architecture of the action prediction is given in Figure 17, which is 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 (())
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 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 x 8 relu conv 64 x x 4 - relu conv 64 x 3) followed by two additional fully connected layers of size 64. 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.
PLATFORM AND LICENSE
We run all the experiments the servers with 4 NVidia V100 GPUs. We used the Car Racing in OpenAl gym and VizDoom environments and we have cited the creators. our code, we have used the following libraries: Tensorflow (Apache License 2.0), OpenAl Gym (MIT License), VizDoom (MIT License), OpenCV (Apache License), Numpy (BSD 3-Clause "New" or "Revised" License) and NVIDIA-DALI libraries (Apache 2 License).
19

---

Under review as a conference paper at ICLR 2022
Ingut Image x 64 x
Action
reward
Relu Conv 32x654 with Stride 2
Relu Corv 64x4x4 with Stride 2
Stride
Stride
Fully Connected
Fully Connected
Concatenate
Output (input of LSTM)
Figure 12: Network architecture of preprocessor.
3:32
5:32
action:
wlement-wise Multiplication with A
Concaterate
Relu ColleCCceoeee 1024
Relu Comnectd 1624
Relu Deconv 128:5x5 with Stride
Stride
Rolu Deconv 64x5x5 with Stride
Relu Decomy 6465K5 with Stride
Relu Decomy 32x086 with Stride
Relu Deconv 32H5w5 with Stride
spmoid Decorw Strido
sigmoid Deconv 3x6x6 with Stride
Output Image (Prediction)
Output image (Recanstruction)
Figure 14: Network architecture of observation prediction.
Figure 13: Network architecture of observation reconstruction.
action: 3
32
action3
Elemen-siee Maliglkation with
Eament-wia Muliiiiation wihE
wise Muttplication with B
Element wise Multipiiatioon with
Concaterate
Concaterate
Relu Cornected 1024
Fully Connected
Corresee 256
Euly-oomected
CallyComeeted 480
Reward prediction
MON
Figure 15: Network architecture reward.
Network architecture of transi- tion/dynamics.
Figure 16:
20

---

Under review as a conference paper at ICLR 2022
Concatenate
Relu Full-Corneeceed 512
Relu Full--Coeeee 128
Output (Prediction)
Figure 17: Network architecture of action prediction.
D2-+a
D2
D2-f
Dy
1.00
1.00
1.00
0.6
0.75
0.75
0.75
Value
0.6
15
0.2
15
10
15
30
Dimersion
Dimension Index
Dimension
Figure 18: Visualization of estimated structural matrices Dz Davss D Dangi and D₂ in Car Racing, without the explicit sparsity constraints.
Dj-0
Dj
D2-
D;
1.00
1.00
1.00
1.0
0.6
0.75
0.75
0.75
0.6
0.50
0.50
15
25
0.2
10
15
20
10
25
15
30
10
Dimersion
Dimension Index
Dimmnnsion
Figure 19: Visualization of estimated structural matrices D Dānos D D Dani. and D3 in VizDoom.
21


Reviews:
[{"review_id": "q7Rpp0wK3Mp", "reviewer": "Reviewer_vpGY", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper addresses an important problem in high dimensional RL, namely improving sample complexity by finding a smaller Markov Decision Process that is sufficient for planning. To my knowledge the approach considered is novel and is well motivated by the provided theory.\n\nThat said, I must first admit that I found the technical exposition particular hard to follow, particularly from a notional perspective and so I may have missed a few key details. I have four primary concerns:\n\n1. The variational autoencoder objective presented has a number of parameters which the paper does not specify how to select. While I'm happy to consider them to be hyperparameters, the exact selection of these parameters seems important when claiming that the resulting state space compression will me minimal. I am particularly worried due to the lack of unique solution, see for example the Linear Gaussian case. In particular, I imagine one should be concerned with different trade-offs in minimality vs predictiveness in the non-linear case?\n\n2. I am also confused with the interpretation of the graph structure variables D. They are introduced as binary variables, but then the Gaussian section claims that they denote both the connect *and* the strength of the connection? Did I miss a continuous relaxation?\n\n3. Its unclear to me if the experiments achieve the goal of using less data compared to the best baseline. For example, reading the appendix, is seems to be that 10k rollouts of 500 steps were used to learn the ASR! Shouldn't this overhead be included in the learning rate figures?\n\n4. The Dyna algorithm proposed is usually outperformed by the Prioritized sweep algorithm right? Would this significantly change the results of the experiments? (in the paper's favor?)", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper addresses the important problem of learning simpler (lower dimensional) dynamics which are sufficient for RL and optimal control. While I believe this paper has a strong theoretical contribution, I believe in its current state, it is unclear if it paper achieves its goals of minimal state representation (how to tune the hyper parameters) and reducing the sample complexity. These concerns may stem from my misunderstanding, and I would be happy to revise my review if corrected.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "puK2CDL4eib", "reviewer": "Reviewer_83Y3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The general approach of the paper is very appealing from the point of view of learning a latent space representation of the world in a principled and possibly incremental manner. By learning what is essentially a dynamic Bayesian network (DBN), the algorithm acquires a factored representation of the state space, allowing a compact state descriptor. Moreover, by learning the sparsity structure of the generative model, it becomes possible to weed out components of the state descriptor that might be relevant to the overall dynamics of the environment, but are not relevant to the reward function that is being optimized for a particular task in this environment. If a relatively small set of relevant components of the state vector can be identified reliably, it stands to reason that the decision problem can be solved much more efficiently by paying attention only to them, much like humans solve decision problems by ignoring all components of their mental interpretation of the world around them but the very few that appear to be relevant to the decision problem at hand. \n\nAlthough this idea by itself is not new, the specific method of learning the sparsity structure of a DBN appears to be new and highly original. Moreover, the learned generative representation in the form of a DBN is advantageous in terms of incremental building of world models, because multiple DBNs can be composed easily, unlike many other predictive models. \n\nThe authors achieve sparsity by including what they call a structural constraint on the DBN expressed as a sum of various regularization terms involving the edges of the DBN, into the objective function to be minimized by means of a so \"structured\" variational autoencoder. The resulting rather elaborate loss function (Eq. 4) contains as many as 8 different weight/regulation parameters, and one concern is how to set them - maybe the authors can provide some guidelines?\n\nAnalysis for the case of linear models shows that the latent model can be identified from the second-order statistics of the observation sequences, which is consistent with how subspace system identification algorithms work by decomposing the covariance matrices of observations. Because linear state space models are a special case of DBNs, it might be instructive to compare the performance of the proposed ASR with the output of such algorithms. For example, what would happen if the outputs of two completely independent linear state space models are concatenated, and the cost/reward of a task depends on the state vector of only one of the models - will the proposed algorithm correctly construct an ASR with reduced size? \n\nThe empirical verification is on two difficult decision problems with very high-dimensional observation spaces. (Although, maybe it would be worth stating the exact number of dimensions?) The learning curves show clearly that the proposed method learns faster and better than several recent methods that can be considered to define the state of the art. \n\nThe paper is very well written, and although the literature on learning latent models for reinforcement learning is vast, I believe the list of included references to that literature is representative and probably sufficient.  \n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper proposes a novel method for learning generative models of environments with high-dimensional observations, identifying the sparsity structure of these models, and selecting only a subset of the state variables relevant to a particular decision problem for use by RL algorithms when computing the optimal policy. The empirical verification on two challenging high-dimensional problems suggests that the learned representations are indeed instrumental in improving the sample complexity of policy computation.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "jCkYOod6HX-", "reviewer": "Reviewer_2HkD", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- To my knowledge, to leverage structural constraint in MDP especially from a state feature-level perspective for sufficient representation learning is novel and interesting.\n- Both the structure constraint and ASR are learned from trajectory data in an end-to-end fashion.\n- ASRs are evaluated in conjunction with both model-free and model-based RL.\n- The structure constraint learned is insightful and is of potential to obtain interpretable representation.\n\n&nbsp;\n\nWeaknesses:\n- The methodology is not clear enough, somewhat hard to read.\n- The relation (similarity and difference) between SS-VAE and related models (e.g., RSSM used in PlaNet and Dreamer) are not well discussed.\n- The choice of experimental environments and several settings are not well explained.\n\n&nbsp;\n\nMy main concern is on the methodology. I list some major ones below.\n\n&nbsp;\n\nFirst, I am not sure whether I well understand ‘the condition of maximizing cumulative reward’ when ASR is defined. Is this to say the ASR is defined for optimal action selection (i.e., the action sequence in Figure 1 should be from the optimal policy)? Or in other words, the desired representation is minimal and sufficient to optimal policy/value function?\n\nIf so, there should be a gap between ASR and the representation learned based on the trajectory data obtained by an arbitrary policy. If I do not understand it right, please correct me and explain more about how such reward maximization is reflected during the representation training process.\n\n&nbsp;\n\nTo maximizing the mutual information between the action and the given cumulative reward, the second conditional entropy term is considered. Why is the first term neglected?\n\nThe authors may provide some more explanation on why the minimization of condition entropy is transformed into a cross-entropy. Besides, how is the cross entropy $H\\left(q_{\\phi}, p_{\\alpha}\\right)$ calculated since $ q_{\\phi}$ is a conditional state distribution while $ p_{\\alpha}$ is a conditional action distribution?\n\n&nbsp;\n\n\nTo achieve minimality of the representation, the lower bound of the conditional mutual information between observed high-dimensional signals and the ASR is minimized. How is the lower bound derived? I suggest the authors to provide a complete process in the appendix. Moreover, does the minimization of the lower bound surrogate make sense?\n\n&nbsp;\n\n‘$\\check{D}^{ASR}$ can be directly derived from the estimated structural matrices $D_{a \\rightarrow r}$ and $D_{\\vec{s}(\\cdot,i)}$.’ How is this done exactly?\n\n\n\n&nbsp;\n\nIn Equation 4, $\\log p_{\\theta}\\left(o_{t+1} \\mid \\tilde{\\vec{s}}_{t}\\right)+\\log p_{\\theta}\\left(r_{t+2} \\mid \\tilde{\\vec{s}}_{t}, a_{t+1}\\right)$, can the authors explain these two terms of prediction?\n\n&nbsp;\n\nWhy is DeepMind Control Suite not considered as experimental environments which are commonly used for high-dim state representation learning?\n\n&nbsp;\n\n\nI will raise my score if my questions and concerns mentioned above can be well addressed.\n\n&nbsp;\n\n\nMinors:\n- “Each factor in p_γ and p_ϕ is modeled with a mixture of Gaussians (MoGs), to approximate a wide class of continuous distributions.” p should probably be modified to q.\n- “The preprocessor architecture is presented in Figure 17, which takes as input the images, actions and rewards, and its output acts as the input to LSTM.” Figure 17 should probably be modified to Figure 12.\n- On page 12, the penultimate sixth and seventh references are duplicated.\n- On page 18, in the Appendix G section: the paper does not describe Figure 17 in detail.\n", "summary_of_the_paper": "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. ", "summary_of_the_review": "I vote for a borderline rejection at present due to my concerns and questions listed above. I will adjust my rating according to the following discussion.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gXV7KJoBWmL", "reviewer": "Reviewer_hBsj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. It is challenging to learn \"the environment model and ASRs, with random actions.\" \n    In some cases it might require an exponential number of samples. \n\n2. Since $I(\\tilde{s}^{ASR}; a_t, R_{t+1})=$ ... \"we can estimate ASRs by minimizing $H(a_t | \\tilde{s}^{ASR}, R_{t+1})$\". \n    That is not true in general.\n\n3. \"and meanwhile minimizing the dimensionality of ASRs with sparsity constraints:\" The first term is supposed to give the minimal representation regardless the sparsity. Why do you add the sparsity constraints? Also, should it be $\\tilde{\\vec{s}}^{ASR}_{t-1}$ in the conditioning?\n\n4. Please show explicitly the relationship between (3) and \"the information bottleneck method (Tishby et al., 1999).\". Please show what terms in (3) correspond to $I[X; \\tilde{X}]$, to $I[\\tilde{X}; Y]$ and to $\\frac{\\delta I[\\tilde{X}; Y]}{\\delta I[X; \\tilde{X}]}$ in the original IB paper you cite. \n\n5. In what sense the first term is sufficient? Please provide an exact definition. \n\n6. The semi-formal use of the information theoretical quantities might be confusing to the readers. \"MINIMAL SUFFICIENT\"\n has a particular definition and an operational meaning (cf., e.g., \"Elements of information theory\" by Thomas M. Cover, Joy A. Thomas). I suggest either to use the original definitions or  to formulate your definitions of \"MINIMAL SUFFICIENT\". The later requires to show the properties of a new definition and its correspondence to the existing ones. \n\n7. Figure 2. please explain \"Deconv\". What component in the objective (4) does require \"Deconvolution\"?\n\n8. typo: \"hyperparameter turning.\"\n\n9 . Eq. 1 shows $a_{t-1}$ directly  affects $s_t$. Figure 2 does not show this connection. please explain. \n\n10. Please provide a direction comparison between the proposed ASR and methods which are focused on minimal sufficient representation of state. The comparisons in \"Comparison with VRL, SLAC, PlaNet, DBC, and Dreamer\" are indirect comparison only. It might be an unfair comparison because these methods are not looking for sufficient state representations. There might be other factors which contribute to the advantage of your particular architecture type and training. \n\n11. Related to '10' above: please evaluate explicitly the quality of the proposed \"sufficient state representations\" either theoretically or \n     empirically. Specifically, is \"Figure 9: Visualization of estimated structural matrices\" typical? How does it look for vizdoom? how \n     does it look for other dynamical control systems?  How does it look without the explicit sparsity constraints? Is it possible to show the \n    gap of the lower bound on simple examples where one can estimate the relevant mutual information by  e.g., \n     clustering/counting/Krasnov MI estimator/etc? \n\n12. I agree \"Common strategies for such state representation learning include reconstructing the observation\" is not directly relevant for the current work. The relevant works for the comparison are works which do not reconstruct the observation. E.g, \"Pacelli, Vincent, and Anirudha Majumdar. \"Task-driven estimation and control via information bottlenecks.\" 2019 International Conference on Robotics and Automation (ICRA). IEEE, 2019.\" and others.  It is hard to estimate the contribution without a direct comparison to competitive works on minimal sufficient state representations. \n\n\n          \n\n", "summary_of_the_paper": "This paper proposes to find action-sufficient representation which might be helpful for downstream tasks.\nThe idea is that states and/or observations often include irrelevant for actions information which can be removed by \nextracting only reverent information. The method utilizes the concepts of Identifiability, mutual information and sparsity.\nIt is realized by augmenting VAE by Structured Sequential component.\n\n", "summary_of_the_review": "the discussed problem is important and the proposed solution is reasonable.   \nmy concerns are a) the lack of direct comparisons to the relevant methods for minimal state representations, b) the lack of evaluation of the quality of the solution (See please 11 above), c) the lack of the definitions for the essential quantities used in the paper, e.g., What definition of sufficiency do you use? which leads to 'c)' above..., d) the lack of demonstration/validation of minimality and sufficiency on simple examples, where the essential quantities  (e.g., MI) can be evaluated. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a state representation learning technique aiming at extracting only state features that are relevant to solve a given task. It combines several ideas, in particular (1) keeping only features that are relevant to take actions from an information-theoretic point of view, (2) model-based learning, and (3) sparsity-inducing constraints. Experiments on CarRacing and VizDoom show that the proposed method outperforms existing baselines.

Although authors did a thorough job trying to address reviewers' comments during the discussion period, in the end most reviewers remained unconvinced by the submission in its current state, the main remaining concerns being:
* Unclear description of the methodology and informal maths
* A somewhat complex optimization objective that may require tuning many hyper-parameters, and whose entire relevance isn't clearly demonstrated empirically

Overall I agree with these concerns, in particular the general feeling that the theoretical part is difficult to follow, with some apparent typos / mistakes (clearly the original submission had a lot of issues, given the original reviews that required the authors to fix several points). To give a concrete example, while reading the final revision I first ran into potential issues when defining the objective $H(a_t | ...)$ on p.4:
* Although the left hand side is associated to a single timestep $t$, the right hand side is a sum over all values of $t$
* The definition of $q_{\phi}$ seems weird to me, in particular the fact that it takes $o_t$ as input (through $y_{1:t}$) => this seems like a typo (?) and otherwise I don't really understand how this defines a proper definition over states
* As at least one reviewer pointed out, authors start from the mutual information $I$ but drop the entropy term $H(a_t | R_{t+1})$ by claiming it doesn't matter since the goal is to learn the ASR. However, in that objective the distribution over actions $p_{\alpha}$ seems to be learnable (through the $\alpha$ parameters), so if we try to minimize the mutual information $I$ including over $\alpha$ it would have been important to retain the entropy over actions as well.

In terms of the relevance of the results, they look pretty good but:
* The proposed algorithm ends up being somewhat complex, with a lot of terms in the loss (eq. 4), and a lack of empirical validation of what actually matters. I see a single ablation study in the Appendix (Fig. 10), and possibly also the comparison to VRL (but it isn't entirely clear to me what this baseline is implementing as it lacks details). I would have appreciated a more thorough empirical analysis of how each term in eq. 4 matters.
* CarRacing experiments consistently use 21 dimensions "for a fair comparison", but this dimensionality was chosen specifically for and by ASR. As a result, it doesn't really look "fair" to me: a fairer comparison would have either selected the optimal dimensionality for each method, or shown results across a range of different dimensionalities.

I also have some concerns regarding the applicability of the algorithm:
1. Relying on random actions to build a world model only works if random actions allow sufficient enough exploration of the state space. There are many situations where this isn't a realistic assumption (also alluded to by at least one reviewer).
2. Minor: in the setup of eq. 1 the reward $r_t$ doesn't directly depend on $s_t$. I'm not sure to which extent this matters for the proposed algorithm, but if this is a necessary condition for it to work properly, it may cause issues in many stochastic environments.

As a result, I am recommending rejection as I believe the paper is not quite ready for publication. I would encourage the authors to try and simplify the presentation (the paper is very notation-heavy and not an easy read), focusing on showing convincing theoretical and empirical justification for all components of the proposed technique.
Decision: Reject