A CTION -S UFFICIENT S TATE R EPRESENTATION L EARNING FOR C ONTROL WITH S TRUCTURAL C ONSTRAINTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Perceived signals in real-world scenarios are usually high-dimensional and noisy, and ﬁnding and using their representation that contains essential and sufﬁcient information required by downstream decision-making tasks will help improve computational efﬁciency 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 sufﬁcient information for decision-making, termed Action-Sufﬁcient 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 efﬁciency.
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 difﬁcult 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 trafﬁc lights.
The useful state of trafﬁc 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 efﬁciency 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 efﬁciently (Lesort et al., 2018).
Representative methods along this line include deep Kalman ﬁlters (Krishnan et al., 2015), deep variational Bayes ﬁlters (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 sufﬁcient and necessary for downstream policy learning.
State representations that contain insufﬁcient information may lead to sub-optimal policies, while those with redundant information may require more samples and more complex models for training.
We address this problem by modeling the generative process and selection procedure induced by reward maximization; by considering a generative environment model involving observed states, state-transition dynamics, and rewards, and explicitly characterizing structural relationships among variables in the RL system, we propose a principled approach to learning minimal sufﬁcient 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 sufﬁcient set of state representations, termed Action-Sufﬁcient state Representations (ASRs), for the downstream policy learning by making use of structural constraints and the goal of maximizing cumulative reward in policy learning.
• In light of the characterization, we develop Structured Sequential Variational Auto-Encoder (SSVAE), which explicitly encodes structural relationships among variables, for reliable identiﬁcation of ASRs.
• Accordingly, policy learning can be done separately from representation learning, and the policy function only relies on a set of low-dimensional state representations, which improve both model and sample efﬁciency.
Moreover, the estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space, which effectively reduce possibly risky explorations.
E NVIRONMENT M ODEL WITH S TRUCTURAL C ONSTRAINTS
In order to characterize a set of minimal sufﬁcient state representations for downstream policy learning, we ﬁrst formulate a generative environment model in partially observable Markov decision process (POMDP), and then show how to explicitly embed structural constraints over variables in the
RL system and leverage them.
Suppose we have sequences of observations {(cid:104)
o t , a t , r t (cid:105)}
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 (cid:126)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 , (cid:15) t , and η t are corresponding independent and identically distributed (i.i.d.)
random noises.
The latent states (cid:126)s t form an MDP: given (cid:126)s t − 1 and a t − 1 , (cid:126)s t are independent of states and actions before t − 1 .
Moreover, the action a t − 1 directly inﬂuences latent states (cid:126)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 (cid:126)s t , contaminated by random noise e t .
We also consider noise (cid:15) 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 inﬂuence every dimension of (cid:126)s t , and the reward r t may not be inﬂuenced by every dimension of (cid:126)s t − 1 as well, and furthermore there are structural relationships among different dimensions of (cid:126)s t .
Figure 1 gives an illustrative graphical representation, where s 3 ,t − 1 inﬂuences 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 reﬂect such constraints, we explicitly encode the graph structure over variables, including the structure over different dimensions of (cid:126)s and the structures from a t − 1 to (cid:126)s t , (cid:126)s t − 1 to r t , and (cid:126)s t to o t .
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
Here, a t − 1 does not have an edge to s 3 ,t , and only s 2 ,t − 1 and s 3 ,t − 1 have edges to r t , and moreover, we take into account the structural relationships among different dimensions of latent states (cid:126)s t .
Accordingly, we re-formulate (1) as follows:
[Equation]
for i = 1 , · · · , d , where (cid:126)s t = ( s 1 ,t , · · · , s d,t ) (cid:62) , (cid:12) denotes element-wise product, and D ( · ) are binary matrices indicating the graph structure over variables.
Speciﬁcally, D (cid:126)s (cid:41) o ∈ { 0 , 1 } d × 1 represents the graph structure from d -dimensional (cid:126)s t to o t , D (cid:126)s (cid:41) r ∈ { 0 , 1 } d × 1 the structure from (cid:126)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 (cid:126)s ∈ { 0 , 1 } d × d denotes the graph structure from d -dimensional (cid:126)s t − 1 to d -dimensional (cid:126)s t and
D (cid:126)s ( · ,i ) is its i -th column, and D a (cid:41) (cid:126)s ∈ { 0 , 1 } 1 × d corresponds to the graph structure from a t − 1 to (cid:126)s t with D a (cid:41) (cid:126)s ( · ,i ) representing its i -th column.
For example, D (cid:126)s ( j,i ) =
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 {(cid:104) o t , a t , r t (cid:105)}
Tt =1 , we aim to learn minimal sufﬁcient state representations for the downstream policy learning.
In the following, we ﬁrst 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 sufﬁciency and minimality of the estimated state representations, when only {(cid:104) o t , a t , r t (cid:105)}
Tt =1 , but not the environment model, is given.
Finding minimal sufﬁcient state dimensions with a given environment model.
RL agents learn to choose appropriate actions according to the current state vector (cid:126)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 sufﬁcient to choose optimal actions?
We call such state dimensions Action-Sufﬁcient state Representations (ASRs) , and denote it by (cid:126)s ASR t .
We can leverage the (conditional) independence/dependence relations among the quantities, under the condition of maximizing cumulative reward, which policy learning aims to achieve, to determine which dimensions of the states are the ASRs.
Moreover, such independence/dependence relations can be directly seen from the graphical representation of the environment model, under the Markov condition and faithfulness assumption (Pearl, 2000; Spirtes et al., 1993).
Below, we give the graphical condition that ASRs are expected to satisfy.
Proposition 1.
Given the graphical representation corresponding to the environment model, such as the representation in Figure 1, which is assumed to be Markov and faithful to the measured data, each dimension in (cid:126)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, which connects d-separation 1 to conditional independence/dependence relations.
A proof is given in
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 the above proposition, it is easy to see that for the graph given in Figure 1, we have (cid:126)s ASR t =
( s 2 ,t , s 3 ,t ) (cid:62) .
That is, we only need ( s 2 ,t , s 3 ,t ) (cid:62) , instead of (cid:126)s t , for the downstream policy learning.
Minimal sufﬁcient 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 {(cid:104)
o t , a t , r t (cid:105)}
Tt =1 .
Then how can we learn the ASRs from the raw high-dimensional inputs such as images?
We denote by ˜ (cid:126)s t the estimated whole latent state representations and ˜ (cid:126)s ASR t ⊆
˜ (cid:126)s t the estimated minimal sufﬁcient 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 (˜ (cid:126)s ASR t ; a t | R t +1 ) , where I denotes mutual information.
Since I (˜ (cid:126)s ASR t
[Equation]
) , where H ( · ) denotes the conditional entropy, and the ﬁrst term H ( a t | R t +1 ) does not contain ASRs, we can estimate ASRs by minimizing H ( a t | ˜ (cid:126)s ASR t , R t +1 ) , with
H ( a t | ˜ (cid:126)s ASR t ,
[Equation]
where p α denote the probabilistic predictive model of a t with parameters α , q φ,α is the joint distribution over ˜ (cid:126)s t and a t with q φ,α = q φ p α , and q φ (˜ (cid:126)s t | ˜ (cid:126)s t − 1 , y 1: t , a 1: t − 1 ) is the probabilistic inference model of ˜ (cid:126)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 ˜ (cid:126)s t are in ˜ (cid:126)s ASR t , so
˜ D ASR (cid:12) ˜ (cid:126)s t gives ASRs ˜ (cid:126)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 ˜ (cid:126)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 ( y
[Equation]
; ˜ (cid:126)s ASR t
[Equation]
with p γ being the transition dynamics of ˜ (cid:126)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 ˜ (cid:126)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 ˜ (cid:126)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 the following function:
[Equation]
[Table]
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 (cid:126)s ( · ,i ) .
The constraint in (3) provides a principled way to achieve minimal sufﬁcient 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 sufﬁcient state representations from the view of generative process underlying the RL problem and the selection procedure induced by reward maximization, which enjoys the following advantages.
The structural infor- mation provides an interpretable and intuitive picture of the generating process.
2) Accordingly, it also provides an interpretable and intuitive way to characterize a minimal sufﬁcient set of state representations for policy learning, which removes unrelated information.
There is no information loss when representation learning and policy learning are done separately, which is computationally more efﬁcient. 4)
The generative environment model is ﬁxed, independent of the behavior policy that is performed.
Furthermore, based on the estimated environment model and ASRs, it is ﬂexible to use a wide range of policy learning methods, and one can also perform model-based RL, which effectively reduce possibly risky explorations.
S TRUCTURED S EQUENTIAL VAE FOR THE E STIMATION OF ASR S
In this section, we give estimation procedures for the environment model and ASRs, as well as the identiﬁability guarantee in linear cases.
Identiﬁability in Linear-Gaussian Cases.
Below, we ﬁrst show the identiﬁability guarantee in the linear case, as a special case of Eq.
In the linear case (see the environment model given in Eq. (5) in Appendix D), D (cid:126)s (cid:41) o , D (cid:126)s (cid:41) r , D a (cid:41) r , D (cid:126)s , and D a (cid:41) (cid:126)s are linear coefﬁcients, indicating corresponding graph structures and also the strength.
Denote the covariance matrices of e t and (cid:15) t by Σ e and Σ (cid:15) , respectively.
Further let ¨ D (cid:126)s (cid:41) o :=
( D (cid:62) (cid:126)s (cid:41) o , D (cid:62) (cid:126)s (cid:41) r ) (cid:62) .
The following proposition shows that the environment model in the linear case is identiﬁable up to some orthogonal transformation on certain coefﬁcient matrices from observed data {(cid:104)
o t , a t , r t (cid:105)} Tt =1 .
Proposition 2 (Identiﬁability) .
Suppose the perceived signal o t , the reward r t , and the latent states (cid:126)s t follow a linear environment model.
If assumptions A1 ∼ A4 (given in Appendix C) hold and with the second-order statistics of the observed data {(cid:104) o t , a t , r t (cid:105)} Tt =1 , the noise variances Σ e and Σ (cid:15) , D a (cid:41) r , ¨
D (cid:126)s (cid:41) o D (cid:126)sk D a (cid:41) (cid:126)s (with k ≥ 0 ), and ¨ D (cid:126)s (cid:41) o ¨ D (cid:62) (cid:126)s (cid:41) o are uniquely identiﬁed.
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 (cid:126)s (cid:41) o , D
(cid:126)s (cid:41) r , D a (cid:41) r , D (cid:126)s , D a (cid:41) (cid:126)s , Σ e , Σ (cid:15) ) and that with ( ˜ D (cid:126)s (cid:41) o , ˜ D (cid:126)s (cid:41) r , ˜ D a (cid:41) r , ˜ D
(cid:126)s , ˜ D a (cid:41) (cid:126)s , ˜Σ ˜ e , ˜Σ ˜ (cid:15) ) are observationally equivalent.
Then we have ˜¨ D (cid:126)s (cid:41) o = ¨ D (cid:126)s (cid:41) o U , ˜
D a (cid:41) r
= D a (cid:41) r , ˜ D (cid:126)s =
U (cid:62) D (cid:126)s U , ˜
D a (cid:41) (cid:126)s = D a (cid:41) (cid:126)s U , ˜Σ ˜ e =Σ e , and ˜Σ ˜ (cid:15) =Σ (cid:15) , where U is an orthogonal matrix.
General Nonlinear Cases.
To handle general nonlinear cases with the generative process given in
Eq.
(2), we develop a Structured Sequential VAE (SS-VAE) to learn the model (including structural constraints) and infer latent state representations ˜ (cid:126)s t and ASRs ˜ (cid:126)s ASR t , with the input {(cid:104)
o t , a t , r t (cid:105)}
Tt =1 .
Speciﬁcally, the latent state dimensions are organized with structures, captured by D (cid:126)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 (cid:126)s (cid:41) o , D (cid:126)s (cid:41) r , D a (cid:41) r , D a (cid:41) (cid:126)s ) into SS-VAE.
Moreover, we aim to learn state representations ˜ (cid:126)s t and ASRs ˜ (cid:126)s ASR t that satisfy the following properties: ( i ) ˜
(cid:126)s t should capture sufﬁcient 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 ) ˜
(cid:126)s ASR t are minimal sufﬁcient state representations for the downstream policy learning.
Let y 1: T = { ( o (cid:62) t , r (cid:62) t ) (cid:62) } Tt =1 .
To achieve the above properties, we maximize the following objective function: L ( y 1: T ; ( θ, [Equation] , γ, α, [Equation]
) =
[Equation]
−
[Equation]
=
[Equation] φ (cid:8) [Equation] [Equation]
|
[Equation]
[Table]
[Equation] which contains the reconstruction error at each time instance, the one-step prediction error of observations, the KL divergence to constrain the latent space, and moreover, the MDP restrictions on transition dynamics, the sufﬁciency and minimality guarantee of state representations for policy learning, as well as sparsity constraints on the graph structure.
We denote by p θ the generative model with parameters θ and structural constraints D ( · ) , q φ the inference model with parameters φ , p γ the transition dynamics, and p
α action prediction with ASRs given the cumulative reward R t +1 .
Each factor in p γ and q φ is modeled with a mixture of Gaussians (MoGs), to approximate a wide class of continuous distributions.
Below are the details of each component in the above objective function:
• Reconstruction and prediction components: These two parts are commonly used in sequential VAE. They
aim to minimize the reconstruction error and prediction error of the perceived signal 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 γ (˜ (cid:126)s t | ˜
(cid:126)s t − 1 , a t − 1 ; D (cid:126)s , D a (cid:41) (cid:126)s ) .
In particular, ˜ (cid:126)s t | ˜ (cid:126)s t − 1 is modelled with a mixture of Gaussians: (cid:80) Kk =1 π k N (cid:0) µ k (˜ (cid:126)s t − 1 ,a t − 1 ) , Σ k (˜ (cid:126)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 ˜ (cid:126)s t − 1 and a t − 1 , parameters γ , and structural constraints D (cid:126)s and D a (cid:41) (cid:126)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 ˜ (cid:126)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 (cid:54) =
j , (2) and also to achieve minimality of ASRs.
• Sufﬁciency & minimality constraints: We achieve minimal sufﬁcient state representations for the downstream policy learning by leveraging action prediction given the cumulative reward, the conditional mutual information between y t and ˜ (cid:126)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 identiﬁability.
In particular, we use L 1 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.
Speciﬁcally, the encoder, which is used to learn the inference model q φ (˜ (cid:126)s t | ˜ (cid:126)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 (cid:104)
o t +1 , r t +1 , a t (cid:105) is projected to the encoder and a sample of ˜ (cid:126)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 (cid:126)s (cid:41) o , D (cid:126)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 satisﬁes a Markov process and is embedded with structural constraints D (cid:126)s and D a (cid:41) (cid:126)s , is modeled with an MLP and MDN, marked with red in Figure 2.
The action prediction part (denoted by AP ), which helps sufﬁcient 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.
P OLICY L EARNING WITH ASR S
After estimating the generative environment model, we are ready to learn the optimal policy, where the policy function only depends on low-dimensional ASRs, instead of high-dimensional images.
The entire procedure roughly contains the following three parts: (1) data collection with a random policy, (2) environment model estimation (with details in Section 3), and (3) policy learning with
ASRs.
Notably, the generative environment model is ﬁxed, 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 ﬂexible 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 ˜ (cid:126)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 ﬂexible 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 in Appendix F gives the detailed procedure of model-free policy learning with
ASRs in partially observable environments.
Model-Based Policy Learning.
The downside of model-free RL algorithms is that they are usually data hungry, requiring very large amounts of interactions.
On the contrary, model-based
RL algorithms enjoy much better sample efﬁciency.
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-efﬁciency and mitigate heavy and possibly risky interactions with the environment.
We present the procedure of the classic Dyna algorithm (Sutton, 1990; Sutton & Barto, 2018) with ASRs in Algorithm 2 in Appendix F.
E XPERIMENTS
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 ﬁrst known solution to achieve the score required to solve the task.
Without stated otherwise, all results were averaged across ﬁve random seeds, with standard deviation shown in the shaded area.
5. 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).
Speciﬁcally, we collected a dataset of 10 k random rollouts of the environment, each consisting of 1000 time steps.
The dimensionality of latent states ˜ (cid:126)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 (cid:126)s (cid:41) o , D (cid:126)s (cid:41) r , D (cid:126)s , and D a (cid:41) (cid:126)s , as shown in Figure 9 in Appendix
G. Intuitively, we can see that D (cid:126)s (cid:41) r and D a (cid:41) (cid:126)s have many values close to zero, meaning that the reward is only inﬂuenced by a small number of state dimensions, and not many state dimensions are inﬂuenced by the action.
Furthermore, from D (cid:126)s , we found that there are inﬂuences from ˜ (cid:126)s i,t to ˜ (cid:126)s i,t + (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 signiﬁcantly improve the policy learning performance in terms of both efﬁciency and efﬁcacy.
Comparison Between Model-Free and Model-Based ASRs.
We applied both model-free (DDPG) (Lillicrap et al., 2015) and model-based (Dyna and Prioritized Sweeping) algorithms (Sutton, 1990) to ASRs (with 21-dims).
As shown in Figure 3, interestingly, by taking advantage of the learned generative model, model-based ASRs is superior to model-free ASRs at a faster rate, which
[Figure] [Figure]
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-efﬁciency.
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 sufﬁcient state representations as in Eq. (3)), 2) SLAC
(Lee et al., 2019), 3) PlaNet (Hafner et al., 2018), 4) DBC (Zhang et al., 2021), and 5) Dreamer (Hafner et al., 2019).
For a fair comparison, the latent dimensions of VRL, PlaNet, SLAC, DBC and Dreamer are set to 21 as well, and we require all of them to have the model capacity similar to ours (i.e., similar model architectures).
From Figure 3, we can see that our methods, both model-free and model-based, obviously outperform others.
It is worth noting that the huge performance difference between ASRs and VRL shows that the components for minimal sufﬁcient 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 ﬁtness value is deﬁned as the average cumulative reward of the 16 random rollouts), we show the ﬁtness values of the best performer (max) and the population (mean) at each generation (Figure 4).
We also took the best performing agent at the end of every 25 generations and tested it over 1024 random rollout scenarios to record the average (best avg score).
It is obvious that our method (denoted by ASR-* ) has a more efﬁcient and also efﬁcacy 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.
Speciﬁcally, we followed Zhang et al.
(2021) to incorporate natural video from the Kinetics dataset (Kay et al., 2017) as background in CarRacing.
Similarly, for a fair comparison, we require all of them to have the same latent dimensions and have the similar model capacity.
As shown in Figure 5, we can see that our method outperforms both Dreamer and DBC.
Ablation Study.
We further performed ablation studies on latent dynamics prediction; that is, we compared with the case when the transition dynamics in Eq.
(4) is not explicitly modeled, but is replaced with a standard normal distribution.
Figure 10 in Appendix G shows that by explicitly modelling the transition dynamics (denoted by with LDP ), the cumulative reward has an obvious improvement over the one without modelling the transition dynamics (denoted by without LDP ).
5. 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 deﬁned 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 [Figure] [Figure]
[Figure] POMDPs: DRQN (Hausknecht & Stone, 2015). As shown in Figure 6, DQN on ASRs achieve a much better performance than all other comparisons [Figure]
and in particular, DQN on ASRs outperforms DRQN on observations by around 400 on average in terms of cumulative reward.
Similarly, we applied model-based (Dyna) algorithms (Sutton, 1990) to ASRs (with 21-dims).
As shown in Figure 6, we can draw the same conclusion that by taking advantage of the learned generative model, model-based ASRs is superior to model-free ASRs at a faster rate.
We also applied ASRs to world models, where Figure 7 shows that our method with ASRs (denoted by ASR-* ) achieves a better performance [Figure]
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 and achieves the sufﬁciency 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 sufﬁciency 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 beneﬁts in several machine learning tasks (Zhang* et al., 2020; Huang et al., 2019), and in this paper, we show its advantages in POMDPs.
C ONCLUSIONS AND F UTURE W ORK
In this paper, we develop a principled framework to characterize a minimal set of state representations that sufﬁce for policy learning, by making use of structural constraints and action predictions.
Accordingly, we propose SS-VAE to reliably extract such a set of state representations from raw observations.
The estimated environment model and ASRs allow learning behaviors from imagined outcomes in the compact latent space, which effectively reduce sample complexity and possibly risky interactions with the environment.
The proposed approach shows promising results on complex environments–CarRacing and Vizdoom.
The future work along this direction include investigating identiﬁability conditions in general nonlinear cases and extending the approach to cover heterogeneous environments, where the generating processes may change over time or across domains.
R EPRODUCIBILITY S TATEMENT
Proofs of all our theoretical results are given in Appendix A, B, C, D with disclosure of all assumptions.
More details about the model estimation and policy learning are given in Appendix E and Appendix
F, respectively.
Appendix G provides more experimental details (including a description of the hyperparameters) and results.
All datasets used are publicly available or instructions are provided on how to generate them in Appendix G. A description of the network architectures used is included in Appendix H. Appendix I introduces the platform and license.
Source code for all our experiments will be made available on GitHub after deanonymization, providing also a complete description of our experimental environment, conﬁguration ﬁles and instructions on the reproduction of our experiments.
A PPENDICES FOR “A CTION -S UFFICIENT S TATE R EPRESENTATION L EARNING FOR C ONTROL WITH S TRUCTURAL C ONSTRAINTS " A P ROOF OF P ROPOSITION
We ﬁrst give the deﬁnitions of the Markov condition and the faithfulness assumption, which will be used in the proof.
Deﬁnition 1 (Global Markov Condition (Spirtes et al., 1993; Pearl, 2000)) .
The distribution p over V satisﬁes 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 ) .
Deﬁnition 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 ﬁrst show that if s i,t ∈ (cid:126)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 (cid:126)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 ∈ (cid:126)s ASR t .
Similarly, by contradiction suppose that s i,t is not a dimension in (cid:126)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 ˜ (cid:126)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 ˜ (cid:126)s t − 1 , instead of ˜ (cid:126)s ASR t − 1 , which two give different conditional mutual information.
It can be shown by contradiction.
Suppose I ( y t ; ˜ (cid:126)s ASR t | y 1: t − 1 , a 1: t − 1 , ˜ (cid:126)s t − 1 )
=
I ( y t ; ˜ (cid:126)s ASR t | y 1: t − 1 , a 1: t − 1 , ˜ (cid:126)s ASR t − 1 ) , and denote ˜ (cid:126)s C = ˜ (cid:126)s \ ˜ (cid:126)s ASR .
Then the equivalence implies that ˜ (cid:126)s Ct − 1 is independent of o t (where o t ∈ y t ) given { y 1: t − 1 , a 1: t − 1 , ˜ (cid:126)s ASR t − 1 } .
It is obviously violated for the example given in Figure 1, where ˜ (cid:126)s C = s 1 and ˜ (cid:126)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 ˜ (cid:126)s t − 1 and ˜ (cid:126)s ASR t − 1 give different conditional mutual information.
Therefore, in the above conditional mutual information, we need to condition on the previous states ˜ (cid:126)s t − 1 .
Moreover, the conditional mutual information
I ( y t ; ˜ (cid:126)s ASR t
[Equation]
1 ) can be upper bound by a KL-divergence:
[Equation]
with p γ being the transition dynamics of ˜ (cid:126)s t with parameters γ .
C A SSUMPTIONS OF P ROPOSITION
To show the identiﬁability 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 (cid:62) (cid:126)s (cid:41) o , D
cid:62) (cid:126)s (cid:41) r ) (cid:62) is full column rank and D (cid:126)s is full rank.
A3.
The control signal a t is i.i
d. and the state (cid:126)s t is stationary.
A4.
The process noise has a unit variance,
i
e., var ( η t ) = I .
D 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 identiﬁability proof in the linear-Gaussian case with control signals:
[Equation]
a
[Equation]
− 1 + η t . Let y t +1 = [ o (cid:62) t , r (cid:62) t +1 ] (cid:62) , ¨ D (cid:126)s (cid:41) o = [ D (cid:62) (cid:126)s (cid:41) o , D (cid:62) (cid:126)s (cid:41) r ] (cid:62) , ¨ D a (cid:41) r =
[ (cid:126) 0 (cid:62) , D (cid:62) a (cid:41) r ] (cid:62) , and ¨ e t =
[ e (cid:62) t , (cid:15) (cid:62) t
+1 ] (cid:62) .
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 identiﬁability.
We ﬁrst 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 (cid:126)s (cid:41) [Equation] D a [Equation]
, ¨ D a (cid:41) r , and ¨ D (cid:126)s
[Equation]
D k(cid:126)s D a (cid:41) (cid:126)s for k > 0 .
Next, we consider the auto-covariance function of (cid:126)s .
Deﬁne the auto-covariance function of (cid:126)s at lag k as R (cid:126)s ( k ) =
E [ (cid:126)s t (cid:126)s (cid:62) t + k ] , and similarly for R y ( k ) .
Clearly, R (cid:126)s ( − k ) =
R (cid:126)s ( k ) (cid:62) and R y ( − k ) =
R y ( k ) (cid:62) .
Then we have
[Equation]
Below, we ﬁrst consider the case where d o + d r = d s .
Let ˜ y t
=
¨ D (cid:126)s (cid:41) o (cid:126)s t , so y t = ˜
y t + ¨ D a (cid:41) r a t − 1 + ¨ e t and R ˜ y
( k ) =
¨ D (cid:126)s (cid:41) o R (cid:126)s t ( k ) ¨ D (cid:62) (cid:126)s (cid:41) o .
R ˜
y ( k ) satisﬁes the recursive property:
[Equation]
if k >
[Equation]
where Ω = ¨ D (cid:126)s (cid:41) o D (cid:126)s
[Equation]
D −
[Equation]
(cid:41) o .
Denote S k
= ¨ D (cid:126)s (cid:41) o
D k − 1 (cid:126)s
D a (cid:41) (cid:126)s · Var ( a t ) .
Then we can derive the recursive property for R y ( k ) :
[Equation]
[Equation]
When k = 2 , we have
[Equation]
The above equation can be re-organized as
[Equation]
Because ¨ D a (cid:41) r and S k are identiﬁable, and suppose (cid:0) R y (1) − ¨ D a (cid:41) r · S (cid:62) 1 (cid:1) is invertible, Ω = ¨ D (cid:126)s (cid:41) o D (cid:126)s ¨ D − 1 (cid:126)s (cid:41) o is identiﬁable.
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 ¨
[Equation] Var ( a t − 1 ) D (cid:62) a (cid:41) (cid:126)s + I ) ¨ D (cid:62) (cid:126)s (cid:41) o , and because ¨ D (cid:126)s (cid:41) o D a (cid:41) (cid:126)s is identiﬁable, ¨ D (cid:126)s (cid:41) o ¨ D (cid:62) (cid:126)s (cid:41) o is identiﬁable.
In summary, we have shown the identiﬁability of ¨ D a (cid:41) r , ¨ D (cid:126)s (cid:41) o D a (cid:41) (cid:126)s , ¨ D (cid:126)s (cid:41) o D k(cid:126)s D a (cid:41) (cid:126)s , ¨ D (cid:126)s (cid:41) o ¨ D (cid:62) (cid:126)s (cid:41) o , and Σ e .
Furthermore, ¨ D (cid:126)s (cid:41) o , D (cid:126)s , and D a (cid:41) (cid:126)s are identiﬁed up to some orthogonal transformations.
That is, suppose the model in Eq.
(3) with parameters ( D (cid:126)s (cid:41) o , D
(cid:126)s (cid:41) r , D a (cid:41) r , D (cid:126)s , D a (cid:41) (cid:126)s , Σ e , Σ (cid:15) ) and that with ( ˜ D (cid:126)s (cid:41) o , ˜ D (cid:126)s (cid:41) r , ˜ D a (cid:41) r , ˜ D
(cid:126)s , ˜ D a (cid:41) (cid:126)s , ˜Σ ˜ e , ˜Σ ˜ (cid:15) ) are observationally equivalent, we then have ˜¨ D (cid:126)s (cid:41) o = ¨ D (cid:126)s (cid:41) o U , ˜
D a (cid:41) r
= D a (cid:41) r , ˜ D (cid:126)s =
U (cid:62) D (cid:126)s U , ˜
D a (cid:41) (cid:126)s = D a (cid:41) (cid:126)s U , ˜Σ ˜ e = Σ e , and ˜Σ ˜ (cid:15) = Σ (cid:15) , where U is an orthogonal matrix.
Next, we extend the above results to the case where d o + d r > d s .
Let ¨ D (cid:62) (cid:126)s (cid:41) o
( i, · ) be the i -th row of ¨ D (cid:126)s (cid:41) o .
Recall that ¨ D (cid:126)s (cid:41) o is of full column rank.
Then for any i , one can show that there always exist d s − 1 rows of ¨ D (cid:126)s (cid:41) o ,
such that they, together with ¨ D (cid:62) (cid:126)s (cid:41) o ( i, · ) , form a d s × d s full-rank matrix, denoted by ¯¨ D (cid:126)s (cid:41) o
( i, · ) .
Then from the observed data corresponding to ¯¨ D (cid:126)s (cid:41) o
( i, · ) , ¯¨ D (cid:126)s (cid:41) o
( i, · ) is determined up to orthogonal transformations.
Thus, ¨ D (cid:126)s (cid:41) o is identiﬁed up to orthogonal transformations.
Similarly, D a (cid:41) r , D (cid:126)s , and D a (cid:41) (cid:126)s are identiﬁed up to orthogonal transformations.
Furthermore, Cov ( ¨ D (cid:126)s (cid:41) o (cid:126)s t + D a (cid:41) r a t ) is determined by ¨ D (cid:126)s (cid:41) o , ¨ D a (cid:41) r , D (cid:126)s , and D a (cid:41) (cid:126)s .
Because Cov ( y t ) =
Cov ( ¨ D (cid:126)s (cid:41) o (cid:126)s t + D a (cid:41) r a t ) + Σ ¨ e , Σ ¨ e is identiﬁable.
One may further add sparsity constraints on D (cid:126)s (cid:41) o , D (cid:126)s (cid:41) r , D (cid:126)s , and D a (cid:41) (cid:126)s , to select more sparse structures among the equivalent ones.
For example, one may add sparsity constraints on the columns of D (cid:126)s (cid:41) o .
Note this corresponds to the mask on the elements of (cid:126)s t in Eq.
2; if the full column is 0, then the corresponding dimension of (cid:126)s t is not selected.
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]
where both p θ ( o t | ˜ (cid:126)s t ; D (cid:126)s (cid:41) o ) and p θ ( r t | ˜ (cid:126)s t − 1 , a t − 1 ; D (cid:126)s (cid:41) r , D a (cid:41) r ) are modelled by mixture of Gaussians, with D (cid:126)s (cid:41) o indicating the existence of edges from ˜ (cid:126)s t to o t and D (cid:126)s (cid:41) r indicating the existence of edges from ˜ (cid:126)s t − 1 to r t .
The inference model q φ (˜ (cid:126)s [Equation] 1: T
[Equation]
1 ) is factorized as
[Equation]
where both q φ (˜ (cid:126)s 1 | y 1 , a 0 ) and q φ (˜ (cid:126)s t | ˜ (cid:126)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 ˜ (cid:126)s t | ˜ (cid:126)s t − 1 modelled with mixture of Gaussians.
Thus, the KL divergence can be represented as follows:
[Equation]
In practice [Equation] KL divergence with mixture of Gaussians is hard to implement [Equation]
so instead
[Equation]
we used the following objective function:
[Equation]
(13) where p γ (cid:48) is a standard multivariate Gaussian N ( (cid:126) 0 , I d ) .
F M ORE D ETAILS FOR P OLICY L EARNING WITH ASR S
Algorithm 1 gives the procedure of model-free policy learning with ASRs in partially observable environments.
Speciﬁcally, it starts from model initialization (line 1) and data collection with a random policy (line 2).
Then it updates the environment model and identiﬁes 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 (cid:126)s t , on lines 8 and 12, we infer q φ ( (cid:126)s ASR t +1 | o ≤ t +1 , r ≤ t +1 , a ≤ t ) and sample (cid:126)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 (cid:126)s ASR t +1 and r t +1 , and update the Q function n times.
Speciﬁcally, 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 efﬁciency.
G A DDITIONAL E XPERIMENTS AND D ETAILS G. C AR R ACING E XPERIMENT
CarRacing (with an illustration in Figure 8) is a continuous control task with three continuous actions: steering left/right, acceleration, and brake.
Reward is − 0 . 1 every frame and +1000 /N for every track tile visited, where N is the total number of tiles in track. It is obvious that the CarRacing environment 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.
Receive initial observations o 1 and r 1 .
Infer the posterior q φ ( (cid:126)s ASR| o 1 , r 1 ) and sample (cid:126)s ASR.
9: for t = 1, ..., T do 10:
Select action a t = π ( (cid:126)s ASR t )
+ N t according to the current policy and exploration noise.
Execute action a t and receive reward r t +1 and observation o t +1 .
Infer the posterior q φ ( (cid:126)s ASR t +1 | o ≤ t +1 , r ≤ t +1 , a ≤ t ) and sample (cid:126)s ASR t +1 .
Store transition ( (cid:126)s ASR t ,(cid:126)s t , a t , r t +1 ,(cid:126)s ASR t +1 ,(cid:126)s t +1 , o t +1 ) in B .
Sample a random minibatch of N transitions ( (cid:126)s ASR i , a i , r i +1 ,(cid:126)s ASR i +1 ) from B .
Update network parameters using a speciﬁed RL algorithm (e.g., DQN or DDPG).
Update the model given in (2) with the recorded data from B (according to Section 3).
17: for p = 1, ..., n do 18:
Sample a random minibatch of pairs of ( (cid:126)s t , a t ) from B .
Predict ( (cid:126)s ASR t +1 , r t +1 ) according to the model given in (2).
Update network parameters using a speciﬁed RL algorithm (e.g., DQN or DDPG).
21: end for 22: end for 23: end for is partially observable: by just looking at the current frame, although we can tell the position of the car, we know neither its direction nor velocity that are essential for controlling the car.
For a fair comparison, we followed the same setting as in Ha & Schmidhuber (2018).
Speciﬁcally, 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 ˜ (cid:126)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 (cid:126)s (cid:41) o , D (cid:126)s (cid:41) r , D (cid:126)s , and D a (cid:41) (cid:126)s , as shown in Figure 9.
Intuitively, we can see that D (cid:126)s (cid:41) r and D a (cid:41) (cid:126)s have many values close to zero, meaning that the reward is only inﬂuenced by a small number of state dimensions, and not many state dimensions are inﬂuenced by the action.
Furthermore, from D (cid:126)s , we found that there are inﬂuences from ˜ (cid:126)s i,t to ˜ (cid:126)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 signiﬁcantly improve the policy learning performance in terms of both efﬁciency and efﬁcacy.
[Figure]
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.
Figureshows 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
[Figure]
and (2) it guarantees minimal sufﬁcient state representations for policy learning. Previous approaches usually fail to take into account whether the extracted state representations are sufﬁcient and necessary for downstream policy learning. Moreover, as for the component of recurrent networks, SS-VAE uses LSTM that only contains the stochastic part, while PlaNet and Dreamer use RSSM that contains both deterministic and stochastic components.
G. 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 (Figure 11).
Unlike CarRacing, take cover is a discrete control problem with two actions: move left and move right.
Reward is +1 at each time step while alive, and the cumulative reward is deﬁned 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 ﬁreballs shot by monsters from the other side of the room.
In this task, solving is deﬁned 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 ˜ (cid:126)s t [Figure] 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 ﬁnally reported all the results on the 21-dim ASRs, which achieved the best results in all the experiments.
H 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 Figure 12, which takes as input the images, actions and rewards, and its output acts as the input to LSTM.
We used 256 hidden units in the LSTM and used a ﬁve-component Gaussian mixture in the MDN.
The decoder also consists of three components: a current observation reconstructor (Figure 13), a next observation predictor (Figure 14), and a reward predictor (Figure 15).
The architecture of the transition/dynamics is shown in Figure 16, and its output is also modelled by an MDN with a ﬁve-component Gaussian mixture.
The architecture of the action prediction is given in Figure 17, which is a two-layer MLP taking states and rewards as input and predicted action as output.
In the VizDoom experiment, we used the same image size and the same architectures except that the LSTM has 512 hidden units and the action has one dimension.
It is worth emphasising that we applied weight normalization to all the parameters of the architectures above except for the structural matrices D ( · ) .
In DDPG, both actor network and critic network are modelled by two fully connected layers of size 300 with ReLU and batch normalisation.
Similarly, in DQN (Mnih et al., 2013) on both ASRs and
SSSs, the Q network is also modelled by two fully connected layers of size 300 with ReLU and batch normalisation.
However, in DQN on observations, it is modelled by three convolutional layers (i.e., relu conv 32 × 8 × 8 −→ relu conv 64 × 4 × 4 −→ relu conv 64 × 3 × 3 ) followed by two additional fully connected layers of size 64.
In DRQN (Hausknecht & Stone, 2015) on observations, we used the same architecture as in DQN on observations but padded an extra LSTM layer with 256 hidden units as the ﬁnal layer.
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: Tensorﬂow (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).
[Figure]
ReluConv
[Figure]
ReluConv
[Figure]
ReluConv 128x4x
[Figure]
ReluConv 256x4x
[Figure]
ReluFully
[Figure]
Connected
ReluFully
[Figure]
Connected
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
[Figure]