A CTION -S UFFICIENT S TATE R EPRESENTATION L EARNING FOR C ONTROL WITH S TRUCTURAL C ONSTRAINTS
A BSTRACT
Perceived signals in real-world scenarios are usually high-dimensional and noisy, and finding and using their representation that contains essential and sufficient information required by downstream decision-making tasks will help improve computational efficiency and generalization ability in the tasks.
In this paper, we focus on partially observable environments and propose to learn a minimal set of state representations that capture sufficient information for decision-making, termed Action-Sufficient state Representations (ASRs).
We build a generative environment model for the structural relationships among variables in the system and present a principled way to characterize ASRs based on structural constraints and the goal of maximizing cumulative reward in policy learning.
We then develop a structured sequential Variational Auto-Encoder to estimate the environment model and extract ASRs.
Our empirical results on CarRacing and VizDoom demonstrate a clear advantage of learning and using ASRs for policy learning.
Moreover, the estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space to improve sample efficiency.
I NTRODUCTION
State-of-the-art reinforcement learning (RL) algorithms leveraging deep neural networks are usually data hungry and lack interpretability.
For example, to attain expert-level performance on tasks such as chess or Atari games, deep RL systems usually require many orders of magnitude more training data than human experts (Tsividis et al., 2017).
One of the reasons is that our perceived signals in real-world scenarios, e.g., images, are usually high-dimensional and may contain much irrelevant information for decision-making of the task at hand.
This makes it difficult and expensive for an agent to directly learn optimal policies from raw observational data.
Fortunately, the underlying states that directly guide decision-making could be much lower-dimensional (Schölkopf, 2019; Bengio, 2019).
One example is that when crossing the street, our decision on when to cross relies on the traffic lights.
The useful state of traffic lights (e.g., its color) can be represented by a single binary variable, while the perceived image is high-dimensional.
It is essential to extract and exploit such lower-dimensional states to improve the efficiency and interpretability of the decision-making process.
Recently, representation learning algorithms have been designed to learn abstract features from high-dimensional and noisy observations.
Exploiting the abstract representations, instead of the raw data, has been shown to perform subsequent decision-making more efficiently (Lesort et al., 2018).
Representative methods along this line include deep Kalman filters (Krishnan et al., 2015), deep variational Bayes filters (Karl et al., 2016), world models (Ha & Schmidhuber, 2018), PlaNet (Hafner et al., 2018), DeepMDP (Gelada et al., 2019), stochastic latent actor-critic (Lee et al., 2019),
SimPLe (Kaiser et al., 2019), Bisimulation-based methods (Zhang et al., 2021),
Dreamer (Hafner et al., 2019; 2020), and others (Srinivas et al., 2020; Shu et al., 2020).
Moreover, if we can properly model and estimate the underlying transition dynamics, then we can perform model-based RL or planning, which can effectively reduce interactions with the environment (Ha & Schmidhuber, 2018;
Hafner et al., 2018; 2019; 2020).
Despite the effectiveness of the above approaches to learning abstract features, current approaches usually fail to take into account whether the extracted state representations are sufficient and necessary for downstream policy learning.
State representations that contain insufficient information may lead to sub-optimal policies, while those with redundant information may require more samples and more complex models for training.
We address this problem by modeling the generative process and selection procedure induced by reward maximization; by considering a generative environment model involving observed states, state-transition dynamics, and rewards, and explicitly characterizing
Published at the ICLR 2022 workshop on Objects, Structure and Causality structural relationships among variables in the RL system, we propose a principled approach to learning minimal sufficient state representations.
We show that only the state dimensions that have direct or indirect edges to the reward variable are essential and should be considered for decision making.
Furthermore, they can be learned by maximizing their ability to predict the action, given that the cumulative reward is included in the prediction model, while at the same time achieving their minimality w.r.t.
the mutual information with observations as well as their dimensionality.
The contributions of this paper are summarized as follows:
• We construct a generative environment model, which includes the observation function, transition dynamics, and reward function, and explicitly characterizes structural relationships among variables in the RL system.
• We characterize a minimal sufficient set of state representations, termed Action-Sufficient state Representations (ASRs), for the downstream policy learning by making use of structural constraints and the goal of maximizing cumulative reward in policy learning.
• We develop Structured Sequential Variational Auto-Encoder (SS-VAE), which explicitly encodes structural relationships among variables, for reliable identification of ASRs.
• Policy learning can be done separately from representation learning, and the policy function only relies on a set of low-dimensional state representations, which improve both model and sample efficiency.
Moreover, the estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space, which effectively reduce risky explorations.
E NVIRONMENT M ODEL WITH S TRUCTURAL C ONSTRAINTS
To characterize a set of minimal sufficient state representations for downstream policy learning, we first formulate a generative environment model in partially observable Markov decision process (POMDP), and then show how to explicitly embed structural constraints over variables in the RL system and leverage them.
o t −o t o t + s 1 ,t −s 1 ,t s 1 ,t + s 2 ,t −s 2 ,t s 2 ,t + s 3 ,t −s 3 ,t s 3 ,t + a t −a t r t r t +
R t
Suppose we have sequences of observations {⟨ o t , a t , r t ⟩} Tt =1 , where o t ∈ O denotes perceived signals at time t , such as high-dimensional images, with
O being the observation space, a t ∈ A is the performed action with A being the action space, and r t ∈ R represents the reward variable with R being the reward space.
We denote the underlying states, which are latent, by ⃗s t ∈ S , with S being the state space.
We describe the generating process of the environment model as follows:
[Equation]
where f , g , and h represent the observation function, reward function, and transition dynamics, respectively, and e t , ϵ t , and η t are corresponding independent and identi- cally distributed (i.i.d.)
random noises.
The latent states ⃗s t form an MDP: given ⃗s t − 1 and a t − 1 , ⃗s t are indepen- dent of states and actions before t − 1 .
Moreover, the action a t − 1 directly influences latent states ⃗s t , instead of perceived signals o t , and the reward is determined by the latent states (and the action) as well.
The perceived signals o t are generated from the underlying states ⃗s t , contaminated by random noise e t .
We also consider noise ϵ t in the reward function to capture unobserved factors that may affect the reward, as well as measurement noise.
It is commonplace that the action variable a t − 1 may not influence every dimension of ⃗s t , and the reward r t may not be influenced by every dimension of ⃗s t − 1 as well, and furthermore there are structural relationships among different dimensions of ⃗s t .
Fig.
1 gives an illustrative graphical representation, where s 3 ,t − 1 influences s 2 ,t , a t − 1 does not have an edge to s 3 ,t , and among the states, only s 2 ,t − 1 and s 3 ,t − 1 have edges to r t .
We use R t = (cid:80) ∞ τ = t γ τ − t r τ to denote the discounted cumulative reward starting from time t , where γ ∈ [0 , 1] is the discounted factor that determines how much immediate rewards are favored over more distant rewards.
To reflect such constraints, we explicitly encode the graph structure over variables, including the structure over different dimensions of ⃗s and the structures from a t − 1 to ⃗s t , ⃗s t − 1 to r t , and ⃗s t to o t .
Accordingly, we re-formulate Eq.
1 as follows:
[Equation]
) for i = 1 , · · · , d , where ⃗s t = ( s 1 ,t , · · · , s d,t ) ⊤ , ⊙ denotes element-wise product, and D ( · ) are binary matrices indicating the graph structure over variables.
Specifically, D ⃗s (cid:41) o ∈ { 0 , 1 } d × 1 represents the graph structure from d -dimensional ⃗s t to o t , D ⃗s (cid:41) r ∈ { 0 , 1 } d × 1 the structure from ⃗s t − 1 to the reward variable r t , D a (cid:41) r ∈ { 0 , 1 } the structure from the action variable a t − 1 to the reward variable r t , D ⃗s ∈ { 0 , 1 } d × d denotes the graph structure from d -dimensional ⃗s t − 1 to d -dimensional ⃗s t and
D ⃗s ( · ,i ) is its i -th column, and D a (cid:41) ⃗s ∈ { 0 , 1 } 1 × d corresponds to the graph structure from a t − 1 to ⃗s t with D a (cid:41) ⃗s ( · ,i ) representing its i -th column.
For example, D ⃗s ( j,i ) = 0 means that there is no edge from s j,t − 1 to s i,t .
Here, we assume that the environment model, as well as the structural constraints, is invariant across time instance t .
2. M INIMAL S UFFICIENT S TATE R EPRESENTATIONS
Given observational sequences {⟨ o t , a t , r t ⟩} Tt =1 , 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 relationships, is given.
Then we provide criteria to achieve sufficiency and minimality of the estimated state representations, when only {⟨ o t , a t , r t ⟩} T t =1 , but not the environment model, is given.
Finding minimal sufficient state dimensions with a given environment model.
RL agents learn to choose appropriate actions according to the current state vector ⃗s t to maximize the cumulative reward, in which some dimensions may be redundant for policy learning.
Then how can we identify a minimal subset of state dimensions that are sufficient to choose optimal actions?
We call such state dimensions Action-Sufficient state Representations (ASRs) , and denote it by ⃗s ASR t .
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 Fig. 1, which is assumed to be Markov and faithful to the measured data, each dimension in ⃗s ASR t has a direct or indirect edge to r t + τ , for τ > 0 .
Proposition 1 can be shown based on the global Markov condition and the faithfulness assumption (Pearl, 2000; Spirtes et al., 1993).
A proof is given in the appendix.
If s i,t has a (direct or indirect) edge to r t
+ τ , the action variable a t is dependent on s i,t when the cumulative reward is maximized, which can be thought of as a data selection procedure depending on the cumulative reward, and thus s i,t is needed for action determination.
On the other hand, if s i,t does not have an edge to r t + τ , then the action variable a t is (conditionally) independent on s i,t , and thus s i,t is not needed for action determination.
According to Proposition 1, for the graph in Fig.
1, we have ⃗s ASR t
= ( s 2 ,t , s 3 ,t ) ⊤ .
That is, we only need ( s 2 ,t , s 3 ,t ) ⊤ , instead of ⃗s t , for the downstream policy learning.
Minimal sufficient state representation learning from observed sequences.
In practice, we usually do not have access to the latent states or the environment model, but instead only the observed sequences {⟨ o t , a t , r t ⟩} Tt =1 .
Then how can we learn the ASRs from the raw high-dimensional inputs such as images?
We denote by ˜ ⃗s t the estimated whole latent state representations and ˜ ⃗s ASR t ⊆
˜ ⃗s t the estimated minimal sufficient state representations for policy learning.
We collected the data, that are used to learn the environment model and ASRs, with random actions.
As discussed above, action and ASRs are dependent given the cumulative reward—this is a type of dependence relationship induced by selection on the effect (reward).
We can then learn the ASRs by maximizing their dependence with the action given the cumulative reward, I (˜ ⃗s ASR t ; a t | R t +1 ) , where I denotes mutual information.
Since I (˜ ⃗s ASR t
[Equation]
) , where H ( · ) denotes the conditional entropy, and the first term H ( a t | R t +1 ) does not contain ASRs, we can estimate ASRs by minimizing H ( a t | ˜ ⃗s ASR t , R t +1 ) , with
[Equation]
where p α denote the probabilistic predictive model of a t with parameters α , q ϕ,α is the joint distribution over ˜ ⃗s t and a t with q ϕ,α = q ϕ p α , and q ϕ (˜ ⃗s t | ˜ ⃗s t − 1 , y 1: t , a 1: t − 1 ) is the probabilistic inference model of ˜ ⃗s t with parameters ϕ
and y t =( o
Tt , r Tt ) , and ˜ D ASR ∈ { 0 , 1 } ˜ d × 1 is a binary vector indicating which dimensions of ˜ ⃗s t are in ˜ ⃗s ASR t , so
˜ D ASR ⊙ ˜ ⃗s t gives ASRs ˜ ⃗s ASR t .
Moreover, we achieve minimality of the representation by minimizing conditional mutual information between observed high-dimensional signals y t and the ASR ˜ ⃗s ASR t at time t given data at previous time instances, and meanwhile minimizing the dimensionality of ASRs with sparsity constraints:
[Equation]
where the conditional mutual information can be upper bound by a KL-divergence:
I
[Equation]
with p γ being the transition dynamics of ˜ ⃗s t with parameters γ .
Furthermore, Proposition 1 shows that given the (estimated) environment model, only those state dimensions that have a direct or indirect edge to the reward variable are the ASRs.
In our learning procedure, we also take into account the relationship between the learned states ˜ ⃗s t and the reward, and leverage such structural constraints for learning the ASRs.
Denote by ˇ D ASR ∈ { 0 , 1 } ˜ d × 1 a binary vector indicating whether the corresponding state dimension in ˜ ⃗s t has a direct or indirect edge to the reward variable.
Consequently, we enforce the similarity between ˇ D ASR and ˜ D ASR by adding an L 1 norm on ˇ D ASR − ˜ D ASR .
Therefore, the ASRs can be learned by maximizing
L min&suff = [Equation] 3 (cid:88) [Equation]
t =1 E q ϕ,α (cid:8) log
[Equation]
|
[Equation]
ASR
[Equation] [Equation] [Equation] (cid:124) (cid:123)(cid:122) (cid:125) Sufficiency
[Equation]
[Equation] 1 (cid:80) [Equation] =
[Equation]
2 ∥ ˜ D ASR
[Equation] [Equation]
Minimality
[Equation]
(3) where λ ’s are regularization terms, and note that ˇ D ASR can be directly derived from the estimated structural matrices D a (cid:41) r and D ⃗s ( · ,i ) .
The constraint in Eq.
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 Eq.
4 to learn the whole environment model.
S TRUCTURED S EQUENTIAL VAE FOR THE E STIMATION OF ASR S
We focus on estimation procedures for the environment model and ASRs in the nonlinear case.
It is worth noting that in the linear cases, we can further give the identifiability result of the model, which is left in Appendix C. To handle general nonlinear cases with the generative process given in
Eq. 2, we develop a Structured Sequential VAE (SS-VAE) to learn the model (including structural constraints) and infer latent state representations ˜ ⃗s t and ASRs ˜ ⃗s ASR t , with the input {⟨ o t , a t , r t ⟩} Tt =1 .
Specifically, the latent state dimensions are organized with structures, captured by D ⃗s , to achieve conditional independence.
The structural relationships over perceived signals, latent states, the action variable, and the reward variable are also embedded as free parameters (i.e., D ⃗s (cid:41) o , D ⃗s (cid:41) r , D a (cid:41) r , D a (cid:41) ⃗s ) into SS-VAE.
Moreover, we aim to learn state representations ˜ ⃗s t and ASRs ˜ ⃗s ASR t that satisfy the following properties: ( i ) ˜ ⃗s t should capture sufficient information of observations o t , r t , and a t , that is, it should be enough to enable reconstruction.
( ii )
The state representations should allow for accurate predictions of the next state and also the next observation.
( iii )
The transition dynamics should follow an MDP.
( iv ) ˜ ⃗s ASR t are minimal sufficient state representations for the downstream policy learning.
Let y 1: T = { ( o ⊤ t , r ⊤ t ) ⊤ } Tt =1 .
To achieve these properties, we maximize the following objective:
[Equation] [Equation]
We denote by p θ the generative model with parameters θ and structural constraints D ( · ) , q ϕ the inference model with parameters ϕ , p γ the transition dynamics
[Equation]
and p α action prediction with ASRs given the cumulative reward R t +1 .
Each factor in p γ and q ϕ is modeled with a mixture of Gaussians
[Figure]
(MoGs), to approximate a wide class of continuous distributions.
The details of each component in the above objective and their corresponding neural network architecture are given in Appendix D.
P OLICY L EARNING WITH ASR S
After estimating the generative environment model
[Figure]
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)
[Figure]
and (3) policy learning with
ASRs.
Notably, the generative environment model is fixed, regardless of the behavior policy that is used to generate the data, and after learning the environment model, as well as the inference model for ASRs, our framework is flexible for both model-free and model-based policy learning.
The details of the algorithms for policy learning is left in Appendix F.
E XPERIMENTS
To evaluate the proposed approach, we follow Ha & Schmidhuber (2018) to conduct experiments on both CarRacing (Klimov, 2016) and VizDoom (Kempka et al., 2016).
Without stated otherwise, all results were averaged across five random seeds, with standard deviation.
Due to space limit, we only highlight some key results and refer readers to Appendix G for more details.
We conduct a series of experiments on CarRacing to empirically show that the low-dimensional ASRs significantly improve the policy learning performance in terms of both efficiency and efficacy.
Analysis of ASRs (Fig. 7 in Appendix H).
To demonstrate the structures over observed frames, latent states, actions, and rewards, we visualized the learned D ⃗s (cid:41) o , D ⃗s (cid:41) r , D ⃗s , and D a (cid:41) ⃗s .
Compared to the original 32-dim latent states, ASRs have only 21 dimensions.
(2) Comparison Between Model-Free and Model-Based ASRs (Fig. 2a).
We applied both model-free (DDPG) (Lillicrap et al., 2015) and model-based (Dyna and Prioritized Sweeping) algorithms (Sutton, 1990) to ASRs (with 21-dims).
(3) Comparison with VRL, SLAC, PlaNet, DBC, and Dreamer (Fig. 2a).
We also compared the proposed framework of policy learning with ASRs (with 21-dims) with a) the same learning strategy but with vanilla representation learning (VRL, implemented without the components for minimal sufficient state representations as in Eq. 3), b) SLAC
(Lee et al., 2019), c) PlaNet (Hafner et al., 2018), d) DBC (Zhang et al., 2021), and e) Dreamer
(Hafner et al., 2019).
For a fair comparison, the latent dimensions of VRL, PlaNet, SLAC, DBC and Dreamer are set to 21.
(4) Comparison with Dreamer and DBC with Background Distraction (Fig. 2b).
We further compared ASRs (with 21-dims) with
Dreamer and DBC when there are natural video distractors (Zhang et al., 2021).
(5) Ablation Study (Fig. 8 in Appendix H).
We further performed ablation studies on latent dynamics prediction.
We also applied the proposed method to VizDoom take cover scenario (Kempka et al., 2016).
Considering that in the take over scenario the action space is discrete, we applied the widely used
DQN (Mnih et al., 2013) on ASRs for policy learning.
In addition to the comparisons with VRL (as in CarRacing) and DQN on raw observations, we further compared with another common approach to POMDPs: DRQN (Hausknecht & Stone, 2015).
As shown in Fig.
2c, DQN on ASRs achieve a much better performance than all other comparisons.
C ONCLUSION AND R ELATED W ORK
In this paper, we develop a principled framework to characterize a minimal set of state representations that suffice for policy learning, by making use of structural constraints and action predictions.
Accordingly, we propose SS-VAE to reliably extract such a set of state representations from raw observations.
The estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space, which effectively reduce sample complexity and possibly risky interactions with the environment.
We discuss the related work in Appendix J.
A P ROOF OF P ROPOSITION
We first give the definitions of the Markov condition and the faithfulness assumption, which will be used in the proof.
Definition 1 (Global Markov Condition (Spirtes et al., 1993; Pearl, 2000)) .
The distribution p over V satisfies the global Markov property on graph G if for any partition ( A, B, C ) such that B d-separates
A from C , p ( A, C | B ) =
p
A | B ) p ( C | B ) .
Definition 2 (Faithfulness Assumption (Spirtes et al., 1993; Pearl, 2000)) .
There are no independencies between variables that are not entailed by the Markov Condition.
Below, we give the proof of Proposition 1.
Proof.
We first show that if s i,t ∈ ⃗s ASR t , then it has a direct or indirect edge to r t
+ τ .
We prove it by contradiction.
Suppose that s i,t does not have a direct or indirect edge to r t
+ τ .
According to the Markov assumption, s i,t is independent of a t given R .
Hence, s i,t is not necessary for decision making, and thus s i,t is not a dimension in ⃗s ASR t , which contradicts to the assumption.
Since we have a contradiction, it must be that s i,t has a direct or indirect edge to r t
+ τ .
We next show that if s i,t has a direct or indirect edge to r t + τ , then s i,t ∈ ⃗s ASR t .
Similarly, by contradiction suppose that s i,t is not a dimension in ⃗s ASR t .
It means that s i,t is independent on a t given R and some other variables.
Then according to the faithfulness assumption, s i,t does not have a direct or indirect edge to r t
+ τ , which contradicts to the assumption.
B M INIMALITY OF THE R EPRESENTATION
In this section, we give the detailed derivation of the minimality of the state representation given in
Section 2.1.
We achieve minimality of the representation by minimizing conditional mutual information between observed high-dimensional signals y t , where y t =
{ o Tt , r Tt } , and the ASR ˜ ⃗s ASR t at time t given data at previous time instances, and meanwhile minimizing the dimensionality of ASRs with sparsity constraints:
[Equation]
Note that in the above conditional mutual information, we need to conditional on the previous states ˜ ⃗s t − 1 , instead of ˜ ⃗s ASR t − 1 , which two give different conditional mutual information.
It can be shown by contradiction.
Suppose I ( y t ; ˜ ⃗s ASR t
[Equation]
, and denote ˜ ⃗s C = ˜ ⃗s \ ˜ ⃗s ASR .
Then the equivalence implies that ˜ ⃗s Ct − 1 is independent of o t (where o t ∈ y t ) given { y 1: t − 1 , a 1: t − 1 , ˜ ⃗s ASR t − 1 } .
It is obviously violated for the example given in Fig.
1, where ˜ ⃗s C = sand ˜ ⃗s ASR =
{ s 2 , s 3 } , and s 1 ,t − 1 is dependent on o t given { y 1: t − 1 , a 1: t − 1 , s 2 ,t − 1 , s 3 ,t − 1 } .
Hence, conditioning on ˜ ⃗s t − 1 and ˜ ⃗s ASR t − 1 give different conditional mutual information.
Therefore, in the above conditional mutual information, we need to condition on the previous states ˜ ⃗s t − 1 .
Moreover, the conditional mutual information I ( y t ; ˜ ⃗s ASR t | y 1: t − 1 , a 1: t − 1 , ˜ ⃗s t − 1 ) can be upper bound by a KL-divergence:
[Equation]
with p γ being the transition dynamics of ˜ ⃗s t with parameters γ .
C I DENTIFIABILITY IN L INEAR -G AUSSIAN C ASES
Below, we first show the identifiability guarantee in the linear case, as a special case of Eq. 2.
In the linear case (see the environment model given in Eq.
5 in Appendix D), D ⃗s (cid:41) o , D ⃗s (cid:41) r , D a (cid:41) r , D ⃗s , and D a (cid:41) ⃗s are linear coefficients, indicating corresponding graph structures and also the strength.
Denote the covariance matrices of e t and ϵ t by Σ e and Σ ϵ , respectively.
Further let ¨ D ⃗s (cid:41) o :=
( D ⊤ ⃗s (cid:41) o , D ⊤ ⃗s (cid:41) r ) ⊤ .
The following proposition shows that the environment model in the linear case is identifiable up to some orthogonal transformation on certain coefficient matrices from observed data {⟨ o t , a t , r t ⟩}
Tt =1 .
Proposition 2 (Identifiability) .
Suppose the perceived signal o t , the reward r t , and the latent states ⃗s t follow a linear environment model.
If assumptions A1 ∼ A4 (given in Appendix C.1) hold and with the second-order statistics of the observed data {⟨ o t , a t , r t ⟩} Tt =1 , the noise variances Σ e and Σ ϵ ,
D a (cid:41) r , ¨ D ⃗s (cid:41) o D ⃗sk D a (cid:41) ⃗s (with k ≥ 0 ), and ¨ D ⃗s (cid:41) o ¨ D ⊤ ⃗s (cid:41) o are uniquely identified.
This proposition shows that in the linear case, with the second-order statistics of the observed data, we can identify the parameters up to orthogonal transformations.
In particular, suppose the linear environment model with parameters ( D ⃗s (cid:41) o , D ⃗s (cid:41) r , D a (cid:41) r , D ⃗s , D a (cid:41) ⃗s , Σ e , Σ ϵ ) and that with ( ˜ D ⃗s (cid:41) o , ˜ D ⃗s (cid:41) r , ˜ D
a (cid:41) r , ˜ D ⃗s ,
˜ D a (cid:41) ⃗s , ˜Σ ˜ e , ˜Σ ˜ ϵ ) are observationally equivalent.
Then we have ˜¨
D ⃗s (cid:41) o = ¨ D ⃗s (cid:41) o U , ˜
D a (cid:41) r
= D a (cid:41) r , ˜ D ⃗s =
U ⊤ D ⃗s U , ˜ D a (cid:41) ⃗s =
D a (cid:41) ⃗s U , ˜Σ ˜ e =Σ e , and ˜Σ ˜ ϵ
=Σ ϵ , where U is an orthogonal matrix.
C. A SSUMPTIONS OF P ROPOSITION
To show the identifiability of the model in the linear case, we make the following assumptions:
A1.
d o + d r ≥ d s , where | o t | =
d o , | r t | = d r , and | s t | =
d s .
A2.
( D ⊤ ⃗s (cid:41) o , D ⊤ ⃗s (cid:41) r ) ⊤ is full column rank and D ⃗s is full rank.
A3.
The control signal a t is i.i.d.
and the state ⃗s t is stationary.
A4.
The process noise has a unit variance,
i e., var ( η t )
=
I .
C. P ROOF OF P ROPOSITION
Proof.
The proof of the linear case without control signals has been shown in Zhang & Hyvärinen (2011).
Below, we give the identifiability proof in the linear-Gaussian case with control signals:
[Equation]
Let y t +1 =
[ o ⊤ t , r ⊤ t +1 ] ⊤ , ¨ D ⃗s (cid:41) o =
[ D ⊤ ⃗s (cid:41) o , D ⊤ ⃗s (cid:41) r ] ⊤ , ¨ D a (cid:41) r =
[ ⃗ 0 ⊤ , D ⊤ a (cid:41) r ] ⊤ , and ¨ e t =
[ e ⊤ t , ϵ ⊤ t +1 ] ⊤ .
Then the above equation can be represented as:
[Equation]
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 y t + k and a t :
[Equation]
Thus, from the cross-covariance between y t + k and a t , we can identify ¨ D ⃗s (cid:41)
o D a (cid:41) ⃗s , ¨ D a (cid:41) r , and ¨
D ⃗s (cid:41) o D k⃗s D a (cid:41) ⃗s for k > 0 .
Next, we consider the auto-covariance function of ⃗s .
Define the auto-covariance function of ⃗s at lag k as R ⃗s ( k ) =
E [ ⃗s t ⃗s ⊤ t + k ] , and similarly for R y ( k ) .
Clearly, R ⃗s ( − k )
=
R ⃗s ( k ) ⊤ and R y ( − k ) =
R y ( k ) ⊤ .
Then we have
[Equation]
Below, we first consider the case where d o + d r = d s .
Let ˜ y t
=
¨ D ⃗s
[Equation]
⃗s t , so y t = ˜
y [Equation] +
[Equation]
D a
[Equation]
r a t − 1 + ¨ e t and R ˜ y
( k
[Equation]
=
[Equation]
k ) satisfies the recursive property:
[Equation]
where Ω = ¨ D ⃗s (cid:41) o D ⃗s ¨ D − 1 ⃗s (cid:41) o .
Denote S k
=
¨ D ⃗s (cid:41) o
D k − 1 ⃗s
D a (cid:41) ⃗s · Var ( a t ) .
Then we can derive the recursive property for R y ( k ) :
[Equation]
k
[Equation]
[Equation]
if k = 0 .
When k =
[Equation]
we have
[Equation]
The above equation can be re-organized as
[Equation]
Because ¨ D a (cid:41) r and S k are identifiable, and suppose (cid:0) R y (1) − ¨ D a (cid:41) r · S ⊤ 1 (cid:1) is invertible, Ω = ¨ D ⃗s (cid:41) o
D ⃗s ¨ D − 1 ⃗s (cid:41) o is identifiable.
We further consider R y (0) and R y (1) and write down the following form:
[Equation]
From the above two equations we can then identify Σ e and ¨ D ⃗s (cid:41) o
( D a (cid:41) ⃗s Var ( a t − 1 )
D ⊤ a (cid:41) ⃗s + I )
¨ D ⊤ ⃗s (cid:41) o , and because ¨ D ⃗s (cid:41) o D a (cid:41) ⃗s is identifiable, ¨ D ⃗s (cid:41) o
¨ D ⊤ ⃗s (cid:41) o is identifiable.
In summary, we have shown the identifiability of ¨ D a (cid:41)
[Equation]
[Equation]
D k⃗s D a (cid:41) ⃗s , ¨ D ⃗s (cid:41) o ¨ D ⊤ ⃗s (cid:41) o , and Σ e .
Furthermore, ¨ D ⃗s (cid:41) o , D ⃗s , and D a (cid:41) ⃗s are identified up to some orthogonal transformations.
That is, suppose the model in Eq.
(3) with parameters ( D ⃗s (cid:41) o , D ⃗s (cid:41) r , D a (cid:41) r , D ⃗s , D a (cid:41) ⃗s , Σ e , Σ ϵ ) and that with ( ˜ D ⃗s (cid:41) o , ˜ D ⃗s (cid:41) r , ˜ D a (cid:41) r , ˜ D ⃗s , ˜ D
a (cid:41) ⃗s , ˜Σ ˜ e , ˜Σ ˜ ϵ ) are observationally equivalent, we then have ˜¨ D ⃗s (cid:41) o = ¨ D ⃗s (cid:41)
o U , ˜
D a (cid:41) r
= D a (cid:41) r , ˜ D ⃗s =
U ⊤ D ⃗s U , ˜ D a (cid:41) ⃗s =
D a (cid:41) ⃗s U , ˜ Σ
˜ e = Σ e , and ˜ Σ
˜ ϵ
= Σ ϵ , where U is an orthogonal matrix.
Next, we extend the above results to the case where d o + d r > d s .
Let ¨ D ⊤ ⃗s (cid:41) o
( i, · ) be the i -th row of ¨ D ⃗s (cid:41) o .
Recall that ¨ D ⃗s (cid:41) o is of full column rank.
Then for any i , one can show that there always exist d s − 1 rows on
Structure and of ¨ D ⃗s (cid:41) o , such that they, together with ¨ D ⊤ ⃗s (cid:41) o
( i, · ) , form a d s × d s full-rank matrix, denoted by ¯¨ D ⃗s (cid:41) o
( i, · ) .
Then from the observed data corresponding to ¯¨ D ⃗s (cid:41) o
( i, · ) , ¯¨ D ⃗s (cid:41) o
( i, · ) is determined up to orthogonal transformations.
Thus, ¨ D ⃗s (cid:41) o is identified up to orthogonal transformations.
Similarly, D a (cid:41) r , D ⃗s , and D a (cid:41) ⃗s are identified up to orthogonal transformations.
Furthermore, Cov ( ¨ D ⃗s (cid:41) o ⃗s t
+ D a (cid:41) r a t ) is determined by ¨ D ⃗s (cid:41) o , ¨ D a (cid:41) r , D ⃗s , and D a (cid:41) ⃗s .
Because Cov ( y t ) =
Cov ( ¨ D ⃗s (cid:41) o ⃗s t
+ D a (cid:41) r a t ) + Σ ¨ e , Σ ¨ e is identifiable.
One may further add sparsity constraints on D ⃗s (cid:41) o , D ⃗s (cid:41) r , D ⃗s , and D a (cid:41) ⃗s , to select more sparse structures among the equivalent ones.
For example, one may add sparsity constraints on the columns of D ⃗s (cid:41) o .
Note this corresponds to the mask on the elements of ⃗s t in Eq.
2; if the full column is 0, then the corresponding dimension of ⃗s t is not selected.
D M ORE EXPLANATIONS ON THE O BJECTIVE F UNCTION FOR G ENERAL N ONLINEAR M ODELS
Below are the details of each component in the above objective function:
• Reconstruction and prediction components: These two parts are commonly used in sequential VAE.
They aim to minimize the reconstruction error and prediction error of the perceived signal o t and the reward r t .
• Transition component: To achieve the property that state representations satisfy an MDP, we explicitly model the transition dynamics: log p γ (˜ ⃗s t | ˜ ⃗s t − 1 , a t − 1 ; D ⃗s , D a (cid:41) ⃗s ) .
In particular, ˜ ⃗s t | ˜ ⃗s t − 1 is modelled with a mixture of Gaussians: (cid:80)
Kk =1 π k N (cid:0) µ k
(˜ ⃗s t − 1 ,a t − 1 ) , Σ k (˜ ⃗s t − 1 ,a t − 1 ) (cid:1) , where K is the number of mixtures, µ k ( · ) and Σ k ( · ) are given by multi-layer perceptrons (MLP) with inputs ˜ ⃗s t − 1 and a t − 1 , parameters γ , and structural constraints D ⃗s and D a (cid:41) ⃗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 ˜ ⃗s t 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 1: T ) to achieve conditional disentanglement between q ϕ (˜ s i,t |· ) and q ϕ (˜ s j,t |· ) for i ̸
= 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 t and ˜ ⃗s ASR t , and structural constraints.
For details, please refer to Section 2.1.
• Sparsity constraints:
According to the edge-minimality property (Zhang & Spirtes, 2011), we additionally put sparsity constraints on structural matrices to achieve better identifiability.
In particular, we use L 1 norm of the structural matrices as regularizers in the objective function to achieve sparsity of the solution.
Fig.
3 gives the diagram of the neural network architecture in model training.
We use SS-VAE to learn the environment model and ASRs.
Specifically, the encoder, which is used to learn the inference model q ϕ (˜ ⃗s t | ˜ ⃗s t − 1 , y 1: t , a 1: t − 1 ) , includes a Long Short-Term Memory (LSTM (Hochreiter & Schmidhuber, 1997)) to encode the sequential information with output h t and a Mixture Density
Network (MDN (Bishop, 1994)) to output the parameters of MoGs.
At each time instance, the input ⟨ o t +1 , r t +1 , a t ⟩ is projected to the encoder and a sample of ˜ ⃗s t +1 is inferred from q ϕ as output.
The generated sample further acts as an input to the decoder, together with a t +1 and structural matrices
D ⃗s (cid:41) o , D ⃗s (cid:41) r , and D a (cid:41) r .
Then the decoder outputs ˆ o t +1 and ˆ r t +2 .
Moreover, the state dynamics which satisfies a Markov process and is embedded with structural constraints D ⃗s and D a (cid:41) ⃗s , is modeled with an MLP and MDN, marked with red in Fig.
The action prediction part (denoted by AP ), which helps sufficient state representations, uses MLP and is marked with blue.
During training, we approximate the expectation in L by sampling and then jointly learn all parameters by maximizing L using stochastic gradient descent.
E M ORE E STIMATION D ETAILS FOR G ENERAL N ONLINEAR M ODELS
The generative model p θ can be further factorized as follows:
[Equation]
[Equation] where both p θ ( o t | ˜ ⃗s t ; D ⃗s (cid:41) o ) and p θ ( r t | ˜ ⃗s t − 1 , a t − 1 ; D ⃗s (cid:41) r , D a (cid:41) r ) are modelled by mixture of Gaussians, with D ⃗s (cid:41) o indicating the existence of edges from ˜ ⃗s t to o t and D ⃗s (cid:41) r indicating the existence of edges from ˜ ⃗s
t − 1 to r t .
The inference model q ϕ (˜ ⃗s 1: T | y 1: T , a 1: T − 1 ) is factorized as
[Equation]
where both q ϕ (˜ ⃗s 1 | y 1 , a 0 ) and q ϕ (˜ ⃗s t | ˜ ⃗s t − 1 , y 1: t , a 1: t − 1 ) are modelled with mixture of Gaussians.
The transition dynamics p γ is factorized as
[Equation]
with ˜ ⃗s t | ˜ ⃗s t − 1 modelled with mixture of Gaussians.
Thus, the KL divergence can be represented as follows:
[Equation]
In practice, KL divergence with mixture of Gaussians is hard to implement, so instead, we used the following objective function:
[Equation]
13) where p γ ′ is a standard multivariate Gaussian N ( ⃗ 0 , I d ) .
F M ORE D ETAILS FOR P OLICY L EARNING WITH ASR S
For model-free policy learning, we make use of the learned environment model to infer ASRs ˜ ⃗s ASR t from past observed sequences { o ≤ t , r ≤ t , a ≤ t − 1 } and then predict the action with the estimated low-dimensional ASRs.
Our method is flexible to use a wide range of model-free methods; for example, one may use deep Q-learning for discrete actions (Mnih et al., 2015) and deep deterministic policy gradient (DDPG) for continuous actions (Lillicrap et al., 2015).
Algorithm 1 gives the detailed procedure of model-free policy learning with ASRs in partially observable environments.
Model-Based Policy Learning.
The downside of model-free RL algorithms is that they are usually data hungry, requiring very large amounts of interactions.
On the contrary, model-based
RL algorithms enjoy much better sample efficiency.
Hence, we make use of the learned generative environment model, including the transition dynamics, observation function, and reward function, for model-based policy optimization.
Based on the generative environment model, one can learn behaviors from imagined outcomes to increase sample-efficiency and mitigate heavy and possibly risky interactions with the environment.
We present the procedure of the classic Dyna algorithm (Sutton, 1990; Sutton & Barto, 2018) with ASRs in Algorithm 2.
F. A LGORITHMS FOR P OLICY L EARNING
Algorithm 1 gives the procedure of model-free policy learning with ASRs in partially observable environments.
Specifically, it starts from model initialization (line 1) and data collection with a random policy (line 2).
Then it updates the environment model and identifies the set of ASRs with the collected data (line 3), after which, the main procedure of policy optimization follows.
In particular, because we do not directly observe the states ⃗s t , on lines 8 and 12, we infer q ϕ ( ⃗s ASR t +1 | o ≤ t +1 , r ≤ t +1 , a ≤ t ) and sample ⃗s ASR t +1 from the posterior.
The sampled ASRs are then stored in the buffer (line 13).
Furthermore, we randomly sample a minibatch of N transitions to optimize the policy (lines 14 and 15).
One may perform various RL algorithms on the ASRs, such as deep deterministic policy gradient (DDPG (Lillicrap et al., 2015)) or Q-learning (Mnih et al., 2015).
Algorithm 2 presents the procedure of the classic Dyna algorithm with ASRs.
Lines 17-22 make use of the learned environment model to predict the next step, including ⃗s ASR t +1 and r t +1 , and update the Q function n times.
Specifically, in our implementation, the hyper-parameter n is 20.
Based on the learned model, the agent learns behaviors from imagined outcomes in the compact latent space, which helps to increase sample efficiency.
G M AIN E XPERIMENTAL R ESULTS
To evaluate the proposed approach, we conducted experiments on both CarRacing (Klimov, 2016) and VizDoom (Kempka et al., 2016) environments, following the setup in the world model (Ha &
Schmidhuber, 2018) for a fair comparison.
It is known that CarRacing is very challenging—the recent world model (Ha & Schmidhuber, 2018) is the first known solution to achieve the score required to do 6: Receive initial observations do 10: Execute posterior Store do 18:
21: end for 22: end for 23: end for solve the task.
Without stated otherwise, all results were averaged across five random seeds, with standard deviation shown in the shaded area.
G. C AR R ACING E XPERIMENT
CarRacing is a continuous control task with three continuous actions: steering left/right, acceleration, and brake.
Reward is − 0 .
1 every frame and +1000 /N
for every track tile visited, where N is the total number of tiles in track.
It is obvious that the CarRacing environment is partially observable: by just looking at the current frame, although we can tell the position of the car, we know neither its direction nor velocity that are essential for controlling the car.
For a fair comparison, we followed the same setting as in Ha & Schmidhuber (2018).
Specifically, we collected a dataset of 10 k random rollouts of the environment, each consisting of 1000 time steps.
The dimensionality of latent states ˜ ⃗s 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, we visualized the learned D ⃗s (cid:41) o , D ⃗s (cid:41) r , D ⃗s , and D a (cid:41) ⃗s , as shown in Fig.
7 in Appendix H. Intuitively, we can see that D ⃗s (cid:41) r and D a (cid:41) ⃗s have many values close to zero, meaning that the reward is only influenced by a small number of state dimensions, and not many state dimensions are influenced by the action.
Furthermore, from D ⃗s , we found that there are influences from ˜ ⃗s
i,t to ˜ ⃗s i,t +1 (diagonal values) for most state dimensions, which is reasonable because we want to learn an MDP over the underlying states, while the connections across states (off-diagonal values) are much sparser.
Compared to the original 32-dim latent states, ASRs have only 21 dimensions.
Below, we empirically showed that the low-dimensional ASRs significantly improve the policy learning performance in terms of both efficiency and efficacy.
Comparison Between Model-Free and Model-Based ASRs.
We applied both model-free (DDPG) (Lillicrap et al., 2015) and model-based (Dyna and Prioritized Sweeping) algorithms (Sutton, 1990) to ASRs (with 21-dims).
As shown in Fig.
2a, interestingly, by taking advantage of the learned generative model, model-based ASRs is superior to model-free ASRs at a faster rate, which demonstrates the effectiveness of the learned model.
It also shows that with the estimated environment model and ASRs, we can learn behaviors from imagined outcomes to improve sample-efficiency.
Comparison with VRL, SLAC, PlaNet, DBC, and Dreamer.
We also compared the proposed framework of policy learning with ASRs (with 21-dims) with 1) the same learning strategy but with vanilla representation learning (VRL, implemented without the components for minimal sufficient state representations as in Eq. 3), 2) SLAC
(Lee et al., 2019), 3)
PlaNet (Hafner et al., 2018), 4)
[Figure]
DBC (Zhang et al., 2021), and 5) Dreamer
(Hafner et al., 2019).
For a fair comparison, the latent dimensions of VRL, PlaNet, SLAC, DBC and Dreamer are set to 21 as well, and we require all of them to have the model capacity similar to ours (i.e., similar model architectures).
From Fig.
2a, we can see that our methods, both model-free and model-based, obviously outperform others.
It is worth noting that the huge performance difference between ASRs and VRL shows that the components for minimal sufficient state representations play a pivotal role in our objective.
Comparison with World Models.
In light of the fact that world models (Ha & Schmidhuber, 2018) achieved good performance in CarRacing, we further compared our method (with 21-dim ASRs) with the world model.
For a fair comparison, following Ha & Schmidhuber (2018), we also used the Covariance-Matrix Adaptation Evolution Strategy (CMA-ES) (Hansen, 2016) with a population of 64 agents to optimize the parameters of the controller.
In addition, following the same setting as in Ha & Schmidhuber (2018) (where the agent’s fitness value is defined as the average cumulative reward of the 16 random rollouts), we show the fitness values of the best performer (max) and the population (mean) at each generation (Fig. 4).
We also took the best performing agent at the end of every 25 generations and tested it over 1024 random rollout scenarios to record the average (best avg score).
It is obvious that our method (denoted by ASR-* ) has a more efficient and also efficacy training process.
The best average score of ASRs is 65 higher than that of world models.
Comparison with Dreamer and DBC with Background Distraction.
We further compared ASRs (with 21-dims) with Dreamer and DBC when there are natural video distractors in CarRacing; we chose Dreamer and DBC, because their performance are relatively better than other comparisons when there are no distractors.
Specifically, we followed Zhang et al.
(2021) to incorporate natural video from the Kinetics dataset (Kay et al., 2017) as background in CarRacing.
Similarly, for a fair comparison, we require all of them to have the same latent dimensions and have the similar model capacity.
As shown in Fig.
2b, we can see that our method outperforms both Dreamer and DBC.
Ablation Study.
We further performed ablation studies on latent dynamics prediction; that is, we compared with the case when the transition dynamics in Eq.
4 is not explicitly modeled, but is replaced with a standard normal distribution.
Fig.
8 in Appendix H shows that by explicitly modelling the transition dynamics (denoted by with LDP ), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP ).
G. V IZ D OOM E XPERIMENT
We also applied the proposed method to VizDoom take cover scenario (Kempka et al., 2016), which is a discrete control problem with two actions: move left and move right.
Reward is +1 at each time step while alive, and the cumulative reward is defined to be the number of time steps the agent manages to stay alive during an episode.
Considering that in the take over scenario the action space is discrete, we applied the widely used
DQN (Mnih et al., 2013) on ASRs for policy learning.
In addition to the comparisons with VRL (as in CarRacing) and DQN on raw observations, we further compared with another common approach to
POMDPs: DRQN (Hausknecht & Stone, 2015).
As shown in Fig.
2c, DQN on ASRs achieve a much better performance than all other comparisons, and in particular, DQN on ASRs outperforms DRQN on observations by around 400 on average in terms of cumulative reward.
Similarly, we applied model-based (Dyna) algorithms (Sutton, 1990) to ASRs (with 21-dims).
As shown in Fig.
2c, we can
[Figure]
draw the same conclusion that by taking advantage of the learned generative model, model-based
ASRs is superior to model-free ASRs at a faster rate.
We also applied ASRs to world models, where
Fig.
5 shows that our method with ASRs (denoted by ASR-* ) achieves a better performance.
H A DDITIONAL E XPERIMENTS AND D ETAILS H. C AR R ACING E XPERIMENT
CarRacing (with an illustration in Fig. 6) is a continuous control task with three continuous actions: steering left/right, acceleration, and brake.
Reward is − 0 .
1 every frame and +1000 /N
for every track tile visited, where N is the total number of tiles in track.
It is obvious that the CarRacing environment is partially observable: by just looking at the current frame, although we can tell the position of the car, we know neither its direction nor velocity that are essential for controlling the car.
For a fair comparison, we followed the same setting as in Ha & Schmidhuber (2018).
Specifically, we collected a dataset of 10 k random rollouts of the environment, each consisting of 1000 time steps, for model estimation.
The dimensionality of latent states ˜ ⃗s t was set to ˜ d = 32 , and regularization parameters was set to λ 1 = 1 , λ 2 = 1 , λ 3 = 1 , λ 4 = 1 , λ 5 = 1 , λ 6 = 6 , λ 7 = 10 , λ 8 = 0 .
1 , which are determined by hyperparameter turning.
Analysis of ASRs.
To demonstrate the structures over observed frames, latent states, actions, and rewards, we visualized the learned D ⃗s (cid:41) o , D ⃗s (cid:41) r , D ⃗s , and D a (cid:41) ⃗s , as shown in Fig.
Intuitively, we can see that D ⃗s (cid:41) r and D a (cid:41) ⃗s have many values close to zero, meaning that the reward is only influenced by a small number of state dimensions, and not many state dimensions are influenced by the action.
Furthermore, from D ⃗s , we found that there are influences from ˜ ⃗s
i,t to ˜ ⃗s i,t +1 (diagonal values) for most state dimensions, which is reasonable because we want to learn an MDP over the underlying states, while the connections across states (off-diagonal values) are much sparser.
Compared to the original 32-dim latent states, ASRs have only 21 dimensions.
Below, we empirically showed that the low-dimensional ASRs significantly improve the policy learning performance in terms of both efficiency and efficacy.
[Figure]
Ablation Study.
We further performed ablation studies on latent dynamics prediction; that is, we compared with the case when the transition dynamics in Fig.
4 are not explicitly involved.
Fig.shows that by explicitly modelling the transition dynamics (denoted by with LDP ), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP ).
Difference between our SS-VAE and Planet, Dreamer.
Both our method and Planet (Hafner et al., 2018) and Dreamer (Hafner et al., 2019) are world model-based methods.
The differences are mainly in two aspects: (1) our method explicitly considers the structural relationships among variables in the
RL system, and (2) it guarantees minimal sufficient state representations for policy learning.
Previous approaches usually fail to take into account whether the extracted state representations are sufficient and necessary for downstream policy learning.
Moreover, as for the component of recurrent networks,
SS-VAE uses LSTM that only contains the stochastic part, while PlaNet and Dreamer use RSSM that contains both deterministic and stochastic components. [Figure] l [Figure]
t
i
[Figure]
R
[Figure]
H. V IZ D OOM E XPERIMENT
We also applied the proposed method to VizDoom (Kempka et al., 2016).
VizDoom provides many scenarios and we chose the take cover scenario (Fig. 9).
Unlike CarRacing, take cover is a discrete control problem with two actions: move left and move right.
Reward is +1 at each time step while alive, and the cumulative reward is defined to be the number of time steps the agent manages to stay alive during a episode.
Therefore, in order to survive as long as possible, the agent has to learn how to avoid fireballs shot by monsters from the other side of the room.
In this task, solving is defined as attaining the average survival time of greater than 750 time steps over 100 consecutive episodes, each running for a maximum of 2100 time steps.
Following the same setting as in Ha & Schmidhuber (2018), we collected a dataset of 10k random rollouts of the environment, each consisting of 500 time steps.
The dimensionality of latent state ˜ ⃗s t is set to ˜ d
= 32 .
We also set λ 1 = 1 , λ 2 = 1 , λ 3 = 1 , λ 4 = 1 , λ 5 = 1 , λ 6 = 6 , λ 7 = 10 , λ 8 = 0 .
By tuning thresholds, we finally reported all the results on the 21-dim ASRs, which achieved the best results in all the experiments.
I D ETAILED M ODEL A RCHITECTURES
In the car racing experiment, the original screen images were resized to 64 × 64 × 3 pixels.
The encoder consists of three components: a preprocessor, an LSTM, and an MDN.
The preprocessor architecture is presented in Fig.
10, which takes as input the images, actions and rewards, and its output acts as the input to LSTM.
We used 256 hidden units in the LSTM and used a five-component
Gaussian mixture in the MDN.
The decoder also consists of three components: a current observation reconstructor (Fig. 11), a next observation predictor (Fig. 12), and a reward predictor (Fig. 13).
The architecture of the transition/dynamics is shown in Fig.
14, and its output is also modelled by an
MDN with a five-component Gaussian mixture.
The architecture of the action prediction is given in
Fig.
15, which is a two-layer MLP taking states and rewards as input and predicted action as output.
In the VizDoom experiment, we used the same image size and the same architectures except that the LSTM has 512 hidden units and the action has one dimension.
It is worth emphasising that we applied weight normalization to all the parameters of the architectures above except for the structural matrices D ( · ) .
In DDPG, both actor network and critic network are modelled by two fully connected layers of size 300 with ReLU and batch normalisation.
Similarly, in DQN (Mnih et al., 2013) on both ASRs and
SSSs, the Q network is also modelled by two fully connected layers of size 300 with ReLU and batch normalisation.
However, in DQN on observations, it is modelled by three convolutional layers (i.e., relu conv 32 × 8 × 8 −→ relu conv 64 × 4 × 4 −→ relu conv 64 × 3 × 3 ) followed by two additional fully connected layers of size 64.
In DRQN (Hausknecht & Stone, 2015) on observations, we used the same architecture as in DQN on observations but padded an extra LSTM layer with 256 hidden units as the final layer.
[Figure]
J R ELATED W ORK
In the past few years, a number of approaches have been proposed to learn low-dimensional Markovian representations, which capture the variation in the environment generated by the agent’s actions, without direct supervision (Lesort et al., 2018; Krishnan et al., 2015; Karl et al., 2016; Ha &
Schmidhuber, 2018; Watter et al., 2015; Zhang et al., 2018; Kulkarni et al., 2016; Mahadevan &
Maggioni, 2007; Gelada et al., 2019; Gregor et al., 2018; Ghosh et al., 2019; Zhang et al., 2021).
Common strategies for such state representation learning include reconstructing the observation, learning a forward model, or learning an inverse model.
Furthermore, prior knowledge, such as temporal continuity (Wiskott & Sejnowski, 2002), can be added to constrain the state space.
Recently, much attention has been paid to world models, which try to learn an abstract representation of both spatial and temporal aspects of the high-dimensional input sequences (Watter et al., 2015;
Ebert et al., 2017; Ha & Schmidhuber, 2018; Hafner et al., 2018; Zhang et al., 2019b; Gelada et al., 2019; Kaiser et al., 2019; Hafner et al., 2019; 2020).
Based on the learned world model, agents can perform model-based RL or planning.
Our proposed method is also in the class of world models, which models the generative environment model, and additionally, encodes structural constraints
[Figure]
Relu
[Figure]
ReluDeconv
[Figure]
ReluDeconv
[Figure]
ReluDeconv
[Figure]
sigmoid Deconv
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
and achieves the sufficiency and minimality of the estimated state representations from the view of generative and selection process.
In contrast, Shu et al.
(2020) makes use of contrastive loss, as an alternative of reconstruction loss; however, it only focuses on the transition dynamics and also fails to ensure the sufficiency and minimality.
Another line of approaches of state representation learning is based on predictive state representations (PSRs) (Littman & Sutton, 2002; Singh et al., 2004).
A recent approach generalizes PSRs to nonlinear predictive models, by exploiting the coarsest partition of histories into classes that are maximally predictive of the future (Zhang et al., 2019a).
Moreover, bisimulation-based methods have also attracted much attention (Castro, 2020; Zhang et al., 2021).
On the other hand, our work is also related to Bayesian network learning and causal discovery (Spirtes et al., 1993; Pearl, 2000; Huang* et al., 2020).
For example, Strehl et al.
(2007) considers factorized-state MDP with structures being modeled with dynamic Bayesian network or decision trees.
Incorporating such structure information has shown benefits in several machine learning tasks (Zhang* et al., 2020; Huang et al., 2019), and in this paper, we show its advantages in POMDPs.
K P LATFORM AND L ICENSE
We run all the experiments on the servers with 4 NVidia V100 GPUs.
We used the Car Racing in
OpenAI gym and VizDoom environments and we have cited the creators.
In our code, we have used the following libraries: Tensorflow (Apache License 2.0), OpenAI Gym (MIT License), VizDoom (MIT License), OpenCV (Apache 2 License), Numpy (BSD 3-Clause "New" or "Revised" License) and NVIDIA-DALI libraries (Apache 2 License).