Action-Sufficient State Representation Learning for Control with Structural Constraints

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

Reviews and author responses:

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

 1. It is challenging to learn "the environment model and ASRs, with random actions." 
    In some cases it might require an exponential number of samples. 

2. Since $I(\tilde{s}^{ASR}; a_t, R_{t+1})=$ ... "we can estimate ASRs by minimizing $H(a_t | \tilde{s}^{ASR}, R_{t+1})$". 
    That is not true in general.

3. "and meanwhile minimizing the dimensionality of ASRs with sparsity constraints:" The first term is supposed to give the minimal representation regardless the sparsity. Why do you add the sparsity constraints? Also, should it be $\tilde{\vec{s}}^{ASR}_{t-1}$ in the conditioning?

4. Please show explicitly the relationship between (3) and "the information bottleneck method (Tishby et al., 1999).". Please show what terms in (3) correspond to $I[X; \tilde{X}]$, to $I[\tilde{X}; Y]$ and to $\frac{\delta I[\tilde{X}; Y]}{\delta I[X; \tilde{X}]}$ in the original IB paper you cite. 

5. In what sense the first term is sufficient? Please provide an exact definition. 

6. The semi-formal use of the information theoretical quantities might be confusing to the readers. "MINIMAL SUFFICIENT"
 has a particular definition and an operational meaning (cf., e.g., "Elements of information theory" by Thomas M. Cover, Joy A. Thomas). I suggest either to use the original definitions or  to formulate your definitions of "MINIMAL SUFFICIENT". The later requires to show the properties of a new definition and its correspondence to the existing ones. 

7. Figure 2. please explain "Deconv". What component in the objective (4) does require "Deconvolution"?

8. typo: "hyperparameter turning."

9 . Eq. 1 shows $a_{t-1}$ directly  affects $s_t$. Figure 2 does not show this connection. please explain. 

10. Please provide a direction comparison between the proposed ASR and methods which are focused on minimal sufficient representation of state. The comparisons in "Comparison with VRL, SLAC, PlaNet, DBC, and Dreamer" are indirect comparison only. It might be an unfair comparison because these methods are not looking for sufficient state representations. There might be other factors which contribute to the advantage of your particular architecture type and training. 

11. Related to '10' above: please evaluate explicitly the quality of the proposed "sufficient state representations" either theoretically or 
     empirically. Specifically, is "Figure 9: Visualization of estimated structural matrices" typical? How does it look for vizdoom? how 
     does it look for other dynamical control systems?  How does it look without the explicit sparsity constraints? Is it possible to show the 
    gap of the lower bound on simple examples where one can estimate the relevant mutual information by  e.g., 
     clustering/counting/Krasnov MI estimator/etc? 

12. I agree "Common strategies for such state representation learning include reconstructing the observation" is not directly relevant for the current work. The relevant works for the comparison are works which do not reconstruct the observation. E.g, "Pacelli, Vincent, and Anirudha Majumdar. "Task-driven estimation and control via information bottlenecks." 2019 International Conference on Robotics and Automation (ICRA). IEEE, 2019." and others.  It is hard to estimate the contribution without a direct comparison to competitive works on minimal sufficient state representations. 


          

 the discussed problem is important and the proposed solution is reasonable.   
my concerns are a) the lack of direct comparisons to the relevant methods for minimal state representations, b) the lack of evaluation of the quality of the solution (See please 11 above), c) the lack of the definitions for the essential quantities used in the paper, e.g., What definition of sufficiency do you use? which leads to 'c)' above..., d) the lack of demonstration/validation of minimality and sufficiency on simple examples, where the essential quantities  (e.g., MI) can be evaluated. 

This paper considered the problem of compressing the state space representation in a manner that is sufficient for control and of smaller dimension than the original state space. This is done by first running the system to collect data and then learning structural constraints that are sufficient for describing the reward random variable. This is done by an adapted Variational Auto Encoder which trades off reconstructability, reward prediction, action sufficiency, and minimality. The resulting low dimensional MDP can be combined with classic model based RL to potentially improve sample efficiency. This paper addresses an important problem in high dimensional RL, namely improving sample complexity by finding a smaller Markov Decision Process that is sufficient for planning. To my knowledge the approach considered is novel and is well motivated by the provided theory.

That said, I must first admit that I found the technical exposition particular hard to follow, particularly from a notional perspective and so I may have missed a few key details. I have four primary concerns:

1. The variational autoencoder objective presented has a number of parameters which the paper does not specify how to select. While I'm happy to consider them to be hyperparameters, the exact selection of these parameters seems important when claiming that the resulting state space compression will me minimal. I am particularly worried due to the lack of unique solution, see for example the Linear Gaussian case. In particular, I imagine one should be concerned with different trade-offs in minimality vs predictiveness in the non-linear case?

2. I am also confused with the interpretation of the graph structure variables D. They are introduced as binary variables, but then the Gaussian section claims that they denote both the connect *and* the strength of the connection? Did I miss a continuous relaxation?

3. Its unclear to me if the experiments achieve the goal of using less data compared to the best baseline. For example, reading the appendix, is seems to be that 10k rollouts of 500 steps were used to learn the ASR! Shouldn't this overhead be included in the learning rate figures?

4. The Dyna algorithm proposed is usually outperformed by the Prioritized sweep algorithm right? Would this significantly change the results of the experiments? (in the paper's favor?) This paper addresses the important problem of learning simpler (lower dimensional) dynamics which are sufficient for RL and optimal control. While I believe this paper has a strong theoretical contribution, I believe in its current state, it is unclear if it paper achieves its goals of minimal state representation (how to tune the hyper parameters) and reducing the sample complexity. These concerns may stem from my misunderstanding, and I would be happy to revise my review if corrected.

The paper proposes a new state representation learning method for RL, called Action-Sufficient State Representations (ASRs) to learn minimal and sufficient state representations for downstream decision-making tasks. Different from other related work, they explicitly characterize structural relationships among variables (i.e., state features, observation, reward and action) in MDP. Such structure and ASRs are learned by constructing a generative environment model (SS-VAE). The proposed algorithm is empirically evaluated in conjunction with the model-free, model-based algorithm in VizDoom and CarRacing.  Strengths:
- To my knowledge, to leverage structural constraint in MDP especially from a state feature-level perspective for sufficient representation learning is novel and interesting.
- Both the structure constraint and ASR are learned from trajectory data in an end-to-end fashion.
- ASRs are evaluated in conjunction with both model-free and model-based RL.
- The structure constraint learned is insightful and is of potential to obtain interpretable representation.

&nbsp;

Weaknesses:
- The methodology is not clear enough, somewhat hard to read.
- The relation (similarity and difference) between SS-VAE and related models (e.g., RSSM used in PlaNet and Dreamer) are not well discussed.
- The choice of experimental environments and several settings are not well explained.

&nbsp;

My main concern is on the methodology. I list some major ones below.

&nbsp;

First, I am not sure whether I well understand ‘the condition of maximizing cumulative reward’ when ASR is defined. Is this to say the ASR is defined for optimal action selection (i.e., the action sequence in Figure 1 should be from the optimal policy)? Or in other words, the desired representation is minimal and sufficient to optimal policy/value function?

If so, there should be a gap between ASR and the representation learned based on the trajectory data obtained by an arbitrary policy. If I do not understand it right, please correct me and explain more about how such reward maximization is reflected during the representation training process.

&nbsp;

To maximizing the mutual information between the action and the given cumulative reward, the second conditional entropy term is considered. Why is the first term neglected?

The authors may provide some more explanation on why the minimization of condition entropy is transformed into a cross-entropy. Besides, how is the cross entropy $H\left(q_{\phi}, p_{\alpha}\right)$ calculated since $ q_{\phi}$ is a conditional state distribution while $ p_{\alpha}$ is a conditional action distribution?

&nbsp;


To achieve minimality of the representation, the lower bound of the conditional mutual information between observed high-dimensional signals and the ASR is minimized. How is the lower bound derived? I suggest the authors to provide a complete process in the appendix. Moreover, does the minimization of the lower bound surrogate make sense?

&nbsp;

‘$\check{D}^{ASR}$ can be directly derived from the estimated structural matrices $D_{a \rightarrow r}$ and $D_{\vec{s}(\cdot,i)}$.’ How is this done exactly?



&nbsp;

In Equation 4, $\log p_{\theta}\left(o_{t+1} \mid \tilde{\vec{s}}_{t}\right)+\log p_{\theta}\left(r_{t+2} \mid \tilde{\vec{s}}_{t}, a_{t+1}\right)$, can the authors explain these two terms of prediction?

&nbsp;

Why is DeepMind Control Suite not considered as experimental environments which are commonly used for high-dim state representation learning?

&nbsp;


I will raise my score if my questions and concerns mentioned above can be well addressed.

&nbsp;


Minors:
- “Each factor in p_γ and p_ϕ is modeled with a mixture of Gaussians (MoGs), to approximate a wide class of continuous distributions.” p should probably be modified to q.
- “The preprocessor architecture is presented in Figure 17, which takes as input the images, actions and rewards, and its output acts as the input to LSTM.” Figure 17 should probably be modified to Figure 12.
- On page 12, the penultimate sixth and seventh references are duplicated.
- On page 18, in the Appendix G section: the paper does not describe Figure 17 in detail.
 I vote for a borderline rejection at present due to my concerns and questions listed above. I will adjust my rating according to the following discussion.

The paper proposes a novel method for learning a compact and sparse generative model of a partially observable environment, such that the sparsity of the model can be analyzed to identify the subset of latent state variables that are relevant to the reward function of the environment. This subset of the state space, called the action sufficient state representation of the task and environment, can then be used to learn a policy from fewer samples than using the entire reconstructed state space or a predictive world model would need. Empirical verification on two difficult decision problems with high-dimensional observation spaces supports the claim that the proposed method is indeed more sample efficient. The general approach of the paper is very appealing from the point of view of learning a latent space representation of the world in a principled and possibly incremental manner. By learning what is essentially a dynamic Bayesian network (DBN), the algorithm acquires a factored representation of the state space, allowing a compact state descriptor. Moreover, by learning the sparsity structure of the generative model, it becomes possible to weed out components of the state descriptor that might be relevant to the overall dynamics of the environment, but are not relevant to the reward function that is being optimized for a particular task in this environment. If a relatively small set of relevant components of the state vector can be identified reliably, it stands to reason that the decision problem can be solved much more efficiently by paying attention only to them, much like humans solve decision problems by ignoring all components of their mental interpretation of the world around them but the very few that appear to be relevant to the decision problem at hand. 

Although this idea by itself is not new, the specific method of learning the sparsity structure of a DBN appears to be new and highly original. Moreover, the learned generative representation in the form of a DBN is advantageous in terms of incremental building of world models, because multiple DBNs can be composed easily, unlike many other predictive models. 

The authors achieve sparsity by including what they call a structural constraint on the DBN expressed as a sum of various regularization terms involving the edges of the DBN, into the objective function to be minimized by means of a so "structured" variational autoencoder. The resulting rather elaborate loss function (Eq. 4) contains as many as 8 different weight/regulation parameters, and one concern is how to set them - maybe the authors can provide some guidelines?

Analysis for the case of linear models shows that the latent model can be identified from the second-order statistics of the observation sequences, which is consistent with how subspace system identification algorithms work by decomposing the covariance matrices of observations. Because linear state space models are a special case of DBNs, it might be instructive to compare the performance of the proposed ASR with the output of such algorithms. For example, what would happen if the outputs of two completely independent linear state space models are concatenated, and the cost/reward of a task depends on the state vector of only one of the models - will the proposed algorithm correctly construct an ASR with reduced size? 

The empirical verification is on two difficult decision problems with very high-dimensional observation spaces. (Although, maybe it would be worth stating the exact number of dimensions?) The learning curves show clearly that the proposed method learns faster and better than several recent methods that can be considered to define the state of the art. 

The paper is very well written, and although the literature on learning latent models for reinforcement learning is vast, I believe the list of included references to that literature is representative and probably sufficient.  

 The paper proposes a novel method for learning generative models of environments with high-dimensional observations, identifying the sparsity structure of these models, and selecting only a subset of the state variables relevant to a particular decision problem for use by RL algorithms when computing the optimal policy. The empirical verification on two challenging high-dimensional problems suggests that the learned representations are indeed instrumental in improving the sample complexity of policy computation.

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

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

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

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

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

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

Decision: Reject