Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: Summary
-------


This paper analyzes the problem of learning Markov state abstractions in
reinforcement learning. The authors first prove that the abstraction
must satisfy two conditions: an inverse model equality and a density
ratio equality. These two conditions cannot be enforced in the training
objective, so an approximate regularizer is proposed instead. The
proposed algorithm is compared with several baselines in an offline
learning regime (visual gridworld) which evaluates the quality of the
state abstraction and an online learning regime (continuous control)
which evaluates the performance.


main_review: Decision
--------

Overall, the paper provides a well-written analysis of learning Markov
state abstractions and is a valuable addition to the RL literature. The
paper is held-back by the lack of insight in the experimental
investigations. While the experiments on GridWorld show that their
proposed method does learn a Markov state abstraction, their further
experiments do not clearly show the benefit of a Markov state
abstraction. Taken together, I am rating this paper tentatively at a 5
("Marginally below the acceptance threshold").


Strengths
---------

-   The theoretical analysis of learning Markov abstractions is
    interesting and novel. I particularly like that the work highlights
    that the density ratio condition is often ignored in the previous
    literature, comments on why this might be the case and later shows
    it through an ablation study in Section 6. This "closing the loop"
    from theory to empirics is excellent.
-   Presentation of the work and the writing quality is very high. While
    some aspects of the technical sections are difficult to follow, I
    think the paper does a superb job of motivating and explaining the
    decisions made throughout the paper.

Weaknesses
----------


-   The experiments are not well motivated, and do not meaningfully
    answer whether learning Markov abstractions is a worthwhile
    endeavor.
-   While out of scope for this paper, I think there is a lacking
    discussion on the trajectory/memory component of learning Markov
    state representations. Of course, the goal of this paper is to study
    state abstractions. However, the Markov property is less important
    for high dimensional observations like images (where, often the
    images are used as-is). The real interest and challenge lies in
    learning a Markov state by remembering information from previous
    observations.

Detailed Comments
-----------------
-   Section 4 builds on technical machinery from Section 2 in such a way
    that I often needed to go back and forth between the two sections. I
    feel that part of the problem is that the two sections are separated
    by related work. However, I do not believe that switching section 2
    and 3 is necessarily the solution. Instead, I think the authors
    should contextualize some of the technical machinery such as the
    belief distributions using their working example (the MDP in figure
    2).

-   Section 7: I think the results here, while impressive, completely
    undermine the rest of the paper. If Markov states are important,
    then surely "expert" Markov states are the gold standard for state
    representation. And yet, you show that this is often not the case.
    Indeed, your method and RAD seem to outperform the expert in many
    cases. If this is true, then what is the motivation for learning a
    markov abstraction to begin with? And what is the reason for the
    proposed algorithm's success?

Minor Comments
--------------

-   Figure 3b: The experiment is in the offline learning section, yet
    the x-axis shows performance over a number of episodes. This seems
    to suggest that the experiment is actually online.
-   Figure 3b, 4: Y-axis says reward, should be return?
-   Line 290: How exactly is it mapped to an image? Are the pixels
    equidistant and what is the resolution for the image?

Post Rebuttal
--------------
In light of the discussion below, I have increased my score from a 5 -> 6.
limitations_and_societal_impact: The authors adequately addressed the limitations and potential negative societal impact of their work.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper discusses conditions which are sufficient for a state representation to maintain the Markov property, and describes how a pair of practical objectives can encourage these conditions to be satisfied. An empirical investigation shows these auxiliary objectives lead to structured representations and good sample efficiency.
main_review: This is an interesting paper that connects some previously-studied representation-learning approaches with the principle of preserving the Markov property in an abstract MDP.
Overall the paper is fairly well written, and the contributions are placed well in the context of prior work.

Nonetheless, some aspects of the method could deserve further justification or investigation, and I have a few questions for clarifications:

I think further effort should be made to justify the Markov property as an appropriate driving property for representation learning. Certainly many RL algorithms assume a Markov state, and it’s fairly clear why this would be important to preserve when using a representation learning phase followed by learning with that fixed representation.
It’s less clear why this property needs to be preserved for some specific part of the representation (e.g. penultimate activations) if the policy conditions on a Markov state or on a Markovian history, as in normal ‘end-to-end’ deep RL.
This point could be argued more strongly, and any empirical evidence to suggest that deficiencies of other representation learning objectives are specifically due to a lack of Markovianity would be valuable!

I’m also curious in what settings the authors think preserving the Markov property is too strong. One obvious situation is related to the generic deficiency of reward-agnostic representation learning; if there are distractor features and dynamics that are irrelevant to any reward functions of interest they must still be captured in this case. Of course, there are other advantages to being agnostic to reward, but this should probably be discussed.
Another potential point of comparison is, as far as I understand it, between bisimulation (cares about all Pi), the Markov-preserving abstraction (cares about Pi_phi), and *pi-bisimulation* (cares about pi only): is the latter not strong enough a condition? Do the authors think that the coarsest Markovian state abstraction would be in some way optimal?

A clarification: when the authors write that the method can capitalise on reward when available, does this just refer to the fact that the representation learning objectives can be used as an auxiliary alongside normal reward-maximising RL? Or have I missed some other way reward information is incorporated?

Regarding the experiments, can the authors elaborate on the initialisation of the replay buffer and the pretraining phase used? (i.e. is the x-axis in fig 4 fully comparable between all methods?).
The results are nice, if not astounding, which is unsurprising as the method builds on previously-explored objectives.

Another point of feedback: I found the figures difficult to read, with many overlapping thin lines. Maybe the learning curves could be summarised as bar charts showing the area under curve, with full curves in the appendix at larger size?

limitations_and_societal_impact: In the main review I've made suggestions for some areas of potential limitations that could be discussed further.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: Post author response update:
Some of my concerns have been addressed and some remain. I am more positive about this paper than before the author response, so I am increasing my score.

Reasons for increasing score:
* Relationship to HOMER: The additional explanation provided by the authors (and commitment to extend explanation in the paper) reassures me about this work's novelty.
* Pre-training and performance: I am now mostly convinced that the results (with pretraining) can be compared to the baslines. Specifically, I now understand that the proposed method uses the same number of environment samples. [This is not stated in the appendix as the authors claim.] When using pretraining, their method does indeed improve upon results in past work (though final performance is generally comparable, given the 90% confidence intervals).

Remaining concerns:
* Lack of comparison to "Inverse-only": The authors choose not to compare to "Inverse-only" for Fig4 since their approach is more principled, though Inverse-only performs comparably in Fig3. The authors claim "adding our training objective to encourage Markov abstractions improves learning performance over state-of-the-art image-based RL." Since the focus is on performance, then Inverse-only performing well despite not being as principled would undermine the stated claim.
* Limited comparison of abstraction quality: The authors argue that their approach yields better abstractions and show visualizations of learned abstractions. A quantitative comparison would be more appropriate. In other [sub]fields, when one seeks to make the Markov assumption, one can measure how reasonable this assumption would be (e.g., difference between distributions when [not] conditioning on past observations). A similar approach can be taken here.

End of update


This work notes that past work on learning abstract state representations does not guarantee than the learned representation is Markovian. The authors present two conditions that, when met, ensure the representation is Markovian. Then, they introduce a method that approximately meets these two conditions. This method is compared to existing approaches where it generally learns more quickly and attains a similar final reward.
main_review: The motivation for learning a Markovian abstract representation is great, and the authors correctly note that past work does not guarantee that the abstract representation is Markovian while the proposed conditions would guarantee this. However, this method also does not guarantee an abstraction that is Markovian. Since this work is advocating for combining two existing objectives, the potential value is in the effectiveness of the combination. Unfortunately, the comparison to existing methods demonstrates marginal improvements over the RAD baseline (when removing the pre-training, which is provided only to the proposed method) and is missing comparisons.

The comparison to HOMER (L169-179) is crucial, and more space should be devoted in the main body to it. 
Though HOMER allows only for discrete abstract states, one should compare to HOMER on an environment where continuous abstract states are better suited (domains in Figure 4?) or provide a clear reason why this cannot be done.
It is concerning that the authors note "when we convert our conditions to a training objective in Section 5, we actually recover an alternate form of the KI conditions." The proposed work is less novel than other sections (e.g., Introduction) suggest.

One of the central claims is that the proposed method is better at creating a Markovian abstract representation. An example representation is shown, but the degree to which the representation is Markovian is not measured. An introduction of some suitable metric would help demonstrate that the proposed method achieves its stated goals.

In Figure 3b, "Inverse" performs as well as Markov. The authors note this, as well as that it "does not necessarily produce Markov abstractions." However, the presented method ("Markov") does not necessarily produce Markov abstractions either, right? The degree to which the environment is Markovian is not measured. How are alpha and beta chosen to obtain an environment that is maximally Markovian?

For Figure 4, the authors note that they "pretrain a Markov abstraction on those [random] experiences for 100K steps with the reward information removed." Given that two tasks use 100K steps, 3 tasks use 300K steps, and 1 task uses 500K steps (for shown plots), using an additional 100K during pre-training favors the proposed method. For example, in Appendix H, the authors note that (for the Figure 3 experiment) providing other methods with more pre-training leads to other methods performing better and at least one method matching the proposed method's performance. The authors emphasize that the reward is not provided for these steps, but generally gathering experiences (with a specific policy) is difficult rather than obtaining rewards for certain transitions. Though these 100K steps are with a random policy, a similar pre-training step for other methods may yield better results. In Appendix K, the authors include results without pre-training for their method and acknowledge that this, in general, decreases performance (in particular, Ball-in-cup and Walker become comparable to the baseline).

Combined with the marginal improvements in the tasks where the proposed method does outperform RAD, the Figure 4 experiments do not show how the proposed method can "substantially improve learning performance over existing approaches" (Conclusion).

The introduction of a third objective is motivated well, but this change, again, raises the question of how Markovian is the learned representation (presumably, this extra objective leads to larger losses for the other two).

None of the baseline methods are equivalent to Inverse-only, correct? Given that Inverse-only performs well in Figure 3, including it in Figure 4 would show the benefits of the proposed two-objective approach. (If it is equivalent, then this should be made clear, both when discussing Figure 3 and Figure 4.)

Minor Comments:

-The list of work in L22-27 is not helpful toward understanding this work. It should be shortened / moved elsewhere.

-Figure 1 would be more helpful if it was referenced from the main text sooner (and I think 1b was not referenced at all, though it clearly shows the approach).

-The Background would benefit from a different structure (currently only 1 subsection, with remainder of comment in the "parent").

-L102-104 is out of place in the Background.

-L111-113 can be safely cut (if deemed necessary, then can be moved to end of Sec 1 or head of Sec 2 (not end of Sec 2.1)).

-In general, stating what will be shown later (e.g., "In a moment, we will define [2 lines of text] but first we must make...") is unnecessarily repetitive. Explaining everything once, in whatever order is deemed best, allows more space for other content.

-Since Section 5, Density Ratios largely follows existing work, it can be shortened to make room for other content.
limitations_and_societal_impact: Yes, limitations have been sufficiently addressed.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: # Learning Markov State Abstractions for Deep Reinforcement Learning

The paper proposes a self-supervised representation learning method for deep RL agents that ensures the learned representation is Markov and non-trivial (avoids collapse).

In a rewardless gridworld, the learned representation is competitive with an expert representation and with representations learned via inverse models. The proposed method is also shown to  improve on an RAD baseline in control tasks. Both the proposed method and the RAD baseline outperform the expert state representation. The method is SoTA in some of the tasks.
main_review: ## For the response

### Major concerns:
1. How much of the improvement is due to RAD + L_inv?
2. Is the improvement to the SoTA significant?

I will improve my score if I am convinced that the answer to (1) is negative and the answer to (2) is positive and argued clearly in the text. A table with performance at 100k training steps and 500k training steps would also help put the results into perspective related to previous work, and make it easier to assess significance.

### Minor concerns:
As I argue below, the noise-contrastive estimation for the "Density Ratios" implementation can be hard to make work in other settings. Can you clarify where you expect your representation learning technique to be impactful, and where it will require further research?

## Strengths

I find the framing of the problem interesting: The paper acknowledges that if we are learning the state representation, then we are prone to losing properties that are important for solving the underlying MDP.

The proposed method also outperforms current SoTA in the DeepMind Control Suite (to the extent that I could understand from the literature) in some of the tasks at 100k steps. Breakdown: 1 improved final performance, 2 improved sample complexity with same final performance, 2 with roughly same performance throughout training, 1 with worse performance.

## Weaknesses

The empirical evaluation on control tasks makes it difficult to separate the benefits of augmentations and learning an inverse model from the benefits of the proposed method. The proposed method uses four elements: L_inv, L_ratio and L_smooth and augmentations, but are these four elements together responsible for the performance improvement? L_inv and augmentations have been used separately in previous work. Given that L_inv performs as well as the proposed method in the gridworld, it seems plausible that L_inv plus augmentations would improve on augmentations only.

## Correctness

The paper seems correct, with some gaps between the theoretical results and the method implementation.

## Clarity

The paper is clearly written, though I think some editorial changes would improve it further (see feedback). The critical point in clarity I suggest addressing is helping the reader understand how impactful the performance gains are in the DeepMind Control tasks. The gains provided by the extra losses are smaller than the gains obtained by the augmentations.

## Relation to prior work

I found the discussion informative. There's another group of papers that do not do representation learning on control tasks, but might be otherwise of interest:
* https://arxiv.org/pdf/1611.05397.pdf
* https://arxiv.org/pdf/1906.09237.pdf
* http://proceedings.mlr.press/v119/guo20g/guo20g.pdf
* https://arxiv.org/pdf/2007.05929.pdf

For the empirical evaluation, I have seen in other papers (RAD and DrQ, https://arxiv.org/pdf/2004.13649.pdf , for example) provide a table for 100k steps and 500k steps evaluations. It would benefit comparison with previous work and strengthen the SoTA claims if you could provide a similar table. 

## Reproducibility

I think the text and the appendix add enough information for reproducibility. 

## Additional Feedback

Thank you for submitting your paper for review. I would like to share some comments that I hope you will find useful.

The way the "Density Ratios" is implemented, by treating shuffled experience as a sample from the marginal, effectively means you are treating the marginal as stationary. It's good that this seems to work in this case, but in other cases (e.g., where experience is coming from strongly varying environments/episodes or where there's some notion of time in the state) this breaks down. In my opinion sampling from the correct marginal is one of the challenges leveraging this kind of noise-contrastive estimation across different settings. Therefore it's important to share where you expect this particular implementation of "Density Ratios" to hold up, and where further investigation is needed.

I find it interesting that the paper uses tools from POMDPs (namely belief states). My assessment is that the representation learning potentially makes the setting partially observable. Perhaps the key advantage in this setting is that one can sample from the marginal for L_ratio without a lot of work. In my experience that marginal is a lot harder to sample from in POMDPs.

Another pressing question is how much L_smooth contributes to the final result. Appendix K suggests it has a task-dependent positive or negative effect. Are you proposing this smoothness regularizer? Is it common practice to use it in DeepMind control tasks? I am divided about the presentation for this loss. On one hand, it seems to help with the SoTA performance. On the other hand, it does not quite fit with the way the claims in the paper are stated (that Markov representations lead to improved performance).

For the discussion on the representation gap, I would suggest framing that the representation learning does not really aim at learning expert state information, but rather use the richness of the environment to learn a better representation.

I would also suggest rephrasing the claims in lines 188-190 to include more of the paper's contributions, or state that the paragraph refers to the theoretical contributions.

You use the notation $x \in z$ in eq. (2) to define $B$ (see the denominator). I suggest using $x \in \phi^{-1}(z)$ instead of $x \in X : \phi(x) = z$ and of $x \in z$.

Where does $\Pi_C$ show up in Def. 2? I guess in $B$?

I could not see the definition of $y_i$ (but I could infer one).

The plot in Fig 3(b) has very small font.

In the caption for Figure 4, you state "Adding our Markov training objective leads to state-of-the-art learning performance." I suggest moving this to the main text, because it is a core claim in the paper and it needs discussion in order to be precise:
The Markov training objective includes other components (L_smooth).
Why is the method SoTA? In Cheetah run it's outperformed by RAD. In Reacher it's outperformed by CURL.

limitations_and_societal_impact: The authors have adequately addressed the limitations and potential negative societal impact of their work.
review_title: null

Meta-review: All reviewers agree that the proposed approach is an important contribution to the area of representation learning in reinforcement learning. Although some lingering concerns regarding the empirical evaluation of the method still remain, overall the reviewers think that the paper introduces a fairly elegant and potentially useful recipe for enforcing the Markov property. 

We strongly advise the authors to carefully consider the suggestions made in the reviews and also during the discussion phase. For example, the fact that there exist multiple state abstractions that satisfy the Markov property and that other, additional, properties may also play an important role is a point that should be more clear in the paper. Another observation that came up multiple times in the reviews and was a point of contention in the discussion that followed was the lack of comparisons with a version of the proposed approach that only uses the “inverse loss”. The authors should add a comment to the paper explaining the reasons not to have such a comparison. There were also a few concerns regarding the presentation that should be taken into account when preparing a new version of the submission. 

Decision: Accept (Poster)