Learning Markov State Abstractions for Deep Reinforcement Learning
Authors: Cameron Allen
Neev Parikh
Omer Gottesman
George Konidaris
Venue: NeurIPS 2021 Poster (2021)
Keywords: State Abstraction
Representation Learning
Markov Decision Processes
Reinforcement Learning
Deep RL
Markov Property

Abstract: A fundamental assumption of reinforcement learning in Markov decision processes (MDPs) is that the relevant decision process is, in fact, Markov. However, when MDPs have rich observations, agents typically learn by way of an abstract state representation, and such representations are not guaranteed to preserve the Markov property. We introduce a novel set of conditions and prove that they are sufficient for learning a Markov abstract state representation. We then describe a practical training procedure that combines inverse model estimation and temporal contrastive learning to learn an abstraction that approximately satisfies these conditions. Our novel training objective is compatible with both online and offline training: it does not require a reward signal, but agents can capitalize on reward information when available. We empirically evaluate our approach on a visual gridworld domain and a set of continuous control benchmarks. Our approach learns representations that capture the underlying structure of the domain and lead to improved sample efficiency over state-of-the-art deep reinforcement learning with visual features---often matching or exceeding the performance achieved with hand-designed compact state information.

Learning Markov State Abstractions for Deep Reinforcement Learning
Cameron Allen' Brown University
Neev Parikh Brown University
Omer Gottesman Brown University
George Konidaris Brown University
Abstract
A fundamental assumption of reinforcement learning in Markov decision processes (MDPs) that the relevant decision process is, in fact, Markov. However, when MDPs have rich observations, agents typically learn by way of an abstract state representation, and such representations are not guaranteed to preserve the Markov property. We introduce novel set of conditions and prove that they are sufficient for learning a Markov abstract state representation. We then describe a practical training procedure that combines inverse model estimation and temporal contrastive learning to learn an abstraction that approximately satisfies these conditions. Our novel training objective is compatible with both online and offline training: it does not require reward signal, but agents can capitalize on reward information when available. We empirically evaluate approach on visual gridworld domain and a set of continuous control benchmarks. Our approach learns representations that capture the underlying structure of the domain and lead to improved sample efficiency over state-of-the-art deep reinforcement learning with visual features- often matching or exceeding the performance achieved with hand-designed compact state information.
Introduction
Reinforcement learning (RL) in Markov decision processes with rich observations requires a suitable state representation. Typically, such representations are learned implicitly as a byproduct of doing deep RL. However. in domains where precise and suceinet expert state information is available. agents trained on such expert state features usually outperform agents trained on rich observations. Much recent work (Shelhamer al., 2016; Pathak al., 2017: Ha Schmidhuber, 2018: Gelada et al., 2019; Yarats et al., 2019; Kaiser et al., 2020; Laskin et al., 2020a,b; Zhang et al., 2021) has sought to close this representation gap by incorporating a wide range representation-llarnig objectives that help the agent learn abstract representations with various desirable properties.
Perhaps the most obvious property incentivize in a representation is the Markov property, which holds if and only if the representation contains enough information to accurately characterize the rewards and transition dynamics of the decision process. Markov decision processes (MDPs) have this property by definition, and most reinforcement learning algorithms depend on having Markov state representations. For instance, the ubiquitous objective learning a stationary, state-dependent optimal policy that specifies how the agent should behave is only appropriate for Markov states.
But learned abstract state representations MDPs. This is due to the fact that abstraction necessarily throws away information. Discard too much information, and the resulting representation cannot accurately characterize the environment. Discard too little, and agents will fail to close the representation gap. Abstraction must balance between
not necessarily Markov, even when built top of
Please send correspondence to Cameron Allen <csal@brown. edu>>. Code repository available at http://gthhbbbooommmcmmal///m/oo
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Hidden
Observed State
Abstract Representation
T(a
State
Φ
I
I(al2's 22)
2 F
Pr(z' x))
D
Sensor
Abstraction
Pr((')?
Figure 1: (Left) A x666 6 visual gridworld domain with hidden state and unknown sensor o, where an abstraction function maps each high-dimensional observed state x a lower-dimensional abstract state z (orange circle). (Right) Our Markov abstraction training architecture. A shared encoder & maps ground states x; x' abstract states 2, x', which are inputs to an inverse dynamics model / and contrastive model D that discriminates between real and fake state transitions. The agent's policy depends only on the current abstract state.
ignoring irrelevant information and preserving what is important for decision making. If reward feedback is available, an agent can use determine which state information relevant the task at hand. Alternatively, if the agent can predict raw environment observations from learned abstract states, then all available information is preserved (along with the Markov property), However, these approaches are impractical when rewards are sparse non-existent, observations are sufficiently complex.
We introduce new approach to learning Markov state abstractions. We begin by defining set of theoretical conditions that are sufficient for an abstraction retain the Markov property. We next show that these conditions are approximately satisfied by simultancously training an inverse model to predict the action distribution that explains two consecutive states, and a discriminator determine whether two given states were in fact consecutive. Our combined training objective (architecture shown in Fig. right) supports learning Markov abstract representations without requiring reward information or observation prediction.
Our method is effective for learning Markov state abstractions that are highly beneficial for decision making. We perform evaluations in two settings with rich (visual) observations: gridworld naviga- tion task (Fig. 1, left) and a set of continuous control benchmarks. In the gridworld. we construct an abstract representation offline- Seline-wwtthot access to reward bedback-thal captures the underlying structure of the domain and fully closes the representation gap between visual and expert features. In the control benchmarks, we combine our training objective (online) with traditional RI., where it leads to a significant performance improvement over late-of-theear visual representation learning.
2 Background
Markov decision process M consists of sets of states X and actions A, reward function R : X A x X R. transition dynamics T X x A Pr(X), and discount factor y. For our theoretical results, we assume M is a Block MDP (Du et al., 2019) whose behavior governed by a much smaller (but unobserved) set of states S, and where X is rich observation generated by noisy sensor function a : S Pr(X), as in Figure (left). Block MDPs conveniently avoid potential issues arising from partial observability by assuming that each observation uniquely identifies the unobserved state that generated it. In other words, there exists perfect "inverse sensor" function 2 (z) + s, which means the observations are themselves Markov, as we define below. Note that S, o, and a are all unknown to the agent.
Definition 1 (Markov State Representation). A decision process M A, and its state representation X are Markov and only if) (Ie41 [ar Ft Ft-1) [k) T( T(x=+110, 2e) and R(0) (Ie41, ) = M(x++1, at ft).for all a A, x X, k
The superscript (k)) denotes that the function is being conditioned on A additional steps of history.
The Markov property means that each state X it is sufficient statistic for predicting the next state and expected reward, for any action the agent might select. Technically, each state must also be sufficient for determining the set of actions available to the agent that state, but here we assume, as is common, that every action is available in every state.

---

The behavior of an RL agent is typically determined by a (Markov) policy π X - Pr(A), and each policy induces value function V X + I. which is defined as the expected sum of future discounted rewards starting from given state and following the policy thereafter. The agent's s objective to learn an optimal policy nT that maximizes value every state. Note that the assumption that the optimal policy is stationary and Markov-the that it only depends on state appropriate only the decision process itself is Markov; almost all RL algorithms simply assume this to be true.
2.1 State Abstraction
To support decision making when X is too nigh-dimensional for tractable learning, we turn to state abstraction. Our objective is to find an abstraction function é X - Z mapping each ground state? to an abstract state
((x),, with the hope that learning is tractable using the abstract representation Z (see Figure I, left). Since our goal to support effective abstract decision making, we are mainly concerned with the policy class TT the set of policies with the same behavior for all ground states that share the same abstract state:
Ip {T+ (0(21) 2(x2))
(π(a)211 π T()l A E 02
(1)
An abstraction o X + Z. when applied to an MDP M, induces a new abstract decision process Ma (Z, A, T 2 7), whose dynamics may depend on the current timestep i or the agent's behavior policy and, crucially, which is not necessarily Markov. Consider the following example.
a
4
Un
aa
X
x
X2
X
ag
ZA
ag
4
a1
Ze
Figure 2: An MDP and a non-Markov abstraction.
The figure above depicts four-state, two-action MDP, and an abstraction & where d(x1) 6(x2) zB. It is common to view state abstraction as aggregating or partitioning ground states into abstract states in this way (Li al., 2006). Generally this involves choosing fixed weighting scheme w(xx to express how much each ground state æ contributes to its abstract state where the weights sum to 1 for the set of x in each z. We can then define the abstract transition dynamies To as a w-weighted sum of the ground dynamics (and similarly for reward): Ta(2") a, 2) = En'ex ExE: T(x' a, x)w()). A natural choice for wxx) is to use the ground-state visitation frequencies. For example, if the agent selects actions uniformly random, this leads to w(xz)) w(xaa) = 1 and wxxx) w(22) 0.5.
In this formulation, the abstract decision process is assumed to be Markov by construction, and T and Rg are assumed not to depend the policy or timestep. But this is an oversimplification. The abstract transition dynamics are not Markov; they change depending on how much history is conditioned on: (z4aaa do, 2 2B) 0.5, whereas Pr(zala, By contrast, the ground MDP's dynamics are fully deterministic (and Markov). Clearly, if we define the abstract MDP in this way, it may not match the behavior of the original MDP. Even worse, if the agent's 's policy changes- bounded sample complexity in the ground MDP to make arbitrarily large number of mistakes in the abstract MDP (Abel et al., 2018).
ag 2 2B. a2-1 01, 21-1 =A) 1.
as during learning-thie discrepancy can cause RI. algorithms with
For the abstract decision process to faithfully simulate the ground MDP's dynamics, w(x) must be allowed to vary such that it always reflects the correct ground-state frequencies. Unfortunately, in the simple example above, maintaining accurate weights for w(z) requires keeping track of an unbounded amount of history: if an agent reaches abstract state 2g and repeats action @1 an arbitrarily large number of times (N), then knowing precisely which ground state will end up in (æ) or x22) requires remembering the abstract state it originally came from N + 1 steps prior. Successful modeling of the transition dynamics for subsequent action ag hinges on exactly this distinction between Dn and Z2. The abstraction introduces partial observability, and to compensate, w(x) must be replaced with belief distribution over ground states, conditioned on the entire history; instead
We refer to X as ground states (and M as the ground MDP), to reflect that these quantities are grounded, as opposed to abstract, i.c. they have firm basis in the true environment. Abel et al. (2018) presented three-state chain MDP where they made similar observation.
3

---

of an abstract MDP, we have an abstract POMDP (Bai et al., 2016). This is especially unsatisfying, not just because learning is challenging in POMDPs (Zhang et al., 2012), but because formulating the problem as a Block MDP was supposed to avoid precisely this type of partial observability: an abstraction exists (namely Φ = 1- that would result in a fully observable abstract MDP, if only the agent knew what it was.
We show that it is possible to learn a state abstraction that both reflects the behavior of the underlying ground MDP and preserves the Markov property in the abstract MDP, without ever estimating or maintaining a belief distribution. Our theoretical result leverages inverse dynamies models and contrastive learning to derive conditions under which state abstraction is Markov. We then adapt these conditions into a corresponding training objective for learning such an abstraction directly from the agent's 's experience in the ground MDP.
3 Related Work
3.1 Bisimulation
The idea of learning Markov state abstractions is related the concept of bisimulation (Dean Givan, 1997), the strictest type of state aggregation discussed in Li et al. (2006), where ground states are equivalent if they have exactly the same expected reward and transition dynamics. Preserving the Markov property is prerequisite for bisimulation abstraction, since the abstraction must also preserve the (Markov) ground-state transition dynamics. Bisimulation- abstraction is appealing because, by definition, it leads to high-fidelity representations. But bisimulation is also very restrictive, because it requires & to be Markov for any policy (rather than just those in lla).
Subsequent work on approximate MDP homomorphisms (Ravindran & Barto, 2004) and bisimulation metries (Ferns et al., 2004, 201 1) relaxed these strict assumptions and allowed ground states have varying degrees of "bisimilarity." Castro (2020) introduced a further relaxation, π z-bisimulation, which measures the behavioral similarity of states under policy π. But whereas full bisimulation can be too strong, since it constrains the representation based on policies the agent may never actually select, z-bisimulation can be too weak, since if the policy deviates from π (e.g. during learning), the metric must be updated, and the representation along with it. Our approach can be thought of as useful compromise between these two extremes.
While bisimulation- based approaches have historically been computationally expensive and difficult scale, recent work has started to change that (Castro, 2020; Lehnert & Littman, 2020; Van der Pol et al., 2020; Biza et al., 2021). Two recent algorithms in particular, DeepMDP (Gelada et al., 2019) and Deep Bisimulation for Control (DBC) (Zhang et al., 2021), learn approximate bisimulation abstractions by training the abstraction end-to-end with an abstract transition model and reward function. This is a rather straightforward way to learn Markov abstract state representations since effectively encodes Definition as loss function.
One drawback of bisimulation-based methods is that learning accurate model can be challenging and typically requires restrictive modeling assumptions, such as deterministic, linear, or Gaussian transition dynamics. Bisimulation methods may also struggle if rewards are sparse must be learned without access to rewards. Jointly training an abstraction á with only the transition model (((x),, a) R ((x') can easily lead to a trivial abstraction like (b(:)) - 0 for all since produces both the inputs and outputs for the model. Our approach to learning a Markov abstraction avoids this type of representation collapse without learning a forward model, and is less restrictive than bisimulation, since it is compatible with both reward- free and
the abstraction
settings.
3.2 Ground-State Prediction and Reconstruction
Ground-state (or pixel) prediction (Watter et al., 2015; Song et al., 2016; Kaiser et al., 2020) mitigates representation collapse by forcing the abstract state to be sufficient not just for predicting future abstract states, but also future ground states. Unfortunately, in stochastic domains, this comes with the challenging task of density estimation over the ground state space, and as result, performance about on-par with end-to-end deep RL (Van Hasselt et al., 2019). Moreover, both pixel prediction and the related task of pixel reconstruction (Mattner al., 2012; Finn et al., 2016; Higgins al., 2017: Corneil et al., 2018; Ha & Schmidhuber, 2018; Yarats et al., 2019; Hafner et al., 2020; Lee et al.,

---

2020) are misaligned with the fundamental goal of state ubstraction. These approaches train models to perfectly reproduce the relevant ground state, ergo the abstract state must effectively throw away no information. By contrast, the objective of state abstraction is to throw away much information as possible, while preserving only what necessary for decision making. Provided the abstraction Markov and accurately simulates the ground MDP, we can safely discard the rest of the observation.
3.3 Inverse Dynamics Models
As an alternative to (or in addition to) learning a forward model, sometimes beneficial to learn an inverse model. An inverse dynamics model I [a]xx, æ) predicts the distribution over actions that could have resulted in a transition between a given pair of states. Inverse models have been used for improving generalization from simulation to real- world problems (Christiano et al., 2016), enabling effective robot motion planning (Agrawal et al., 2016), defining intrinsic reward bonuses for exploration (Pathak et al., 2017; Choi al., 2019), and decoupling representation learning from rewards (Zhang et al., 80. But while inverse models often help with representation learning, we show in Sec. 4.2 that they are insufficient for ensuring Markoy abstraction.
3.4 Contrastive Learning
Since the main barrier to effective next-state prediction is learning an accurate forward model, a compelling alternative is contrastive learning (Gutmann & Hyvärinen, 2010), which sidesteps the prediction problem and instead simply aims to decide whether particular state, or sequence of states, came from one distribution or another. Contrastive loss objectives typically aim distinguish either sequential states from non-sequential ones (Shelhamer et al., 2016; Anand et al., 2019; Stooke al., 2020), real states from predicted ones (Van den Oord et al., 2018), or determine whether two augmented views came from the same or different observations (Laskin et al., 2020b). Contrastive methods learn representations that in some cases lead to empirically substantial improvements in learning performance, but none has explicitly addressed the question of whether the resulting state abstractions actually preserve the Markov property. We are the first show that without forward model estimation, pixel prediction/reconstuuctiio,,, or dependence on reward, the specific combination of inverse model estimation and contrastive learning that we introduce in Section is sufficient to learn Markov abstraction.
3.5 Kinematic Inseparability
One contrastive approach which turns out to be closely related to Markov abstraction is Misra et al.'s (2020) HOMER algorithm, and the corresponding notion of kinematic inseparability (KI) abstractions. Two states oi and z'y are defined to be kinematically inseparable if Pr(x, a|æ1] Pr(x,, a (xr)) and T(="la, x1) T(=" x'2) (which the authors call "backwards" and "forwards" KI, respectively). The idea behind KI abstractions is that unless two states can be distinguished from each other by either their backward forward dynamics-
ought be treated the same abstract state. The KI conditions are slightly stronger than the ones we describe in Section 4, although when we convert our conditions into training objective in Section 5, additionally satisfy novel form of the KI conditions, which helps to prevent representation collapse. While our approach works for both continuous and discrete state spaces, HOMER was only designed for discrete abstract states, and requires specifying-in advance-an an upper bound on the number of abstract states (which is impossible for continuous state spaces), well learning a "policy cover" reach each of those abstract states (which remains impractical even under discretization).
(For more detailed discussion about KI and Markov abstractions, see Appendix F.)
3.6 Other Approaches
The Markov property is just one of many potentially desirable properties that a representation might have. Not all Markov representations RL agent end-to-end on (frame-stacked) image inputs ought to be sufficient, and none of the methods in this section would need to do representation learning at all.
equally beneficial for learning; otherwise, simply training an
Smoothness is another desirable property and its benefits in RI. are well known (Pazis Parr, 2013; Pirotta et al., 2015; Asadi et al., 2018). Both DeepMDP (Gelada et al., 2019) and DBC (Zhang
5

---

et al., 2021), which we compare against, utilize Lipschitz smoothness when learning abstract state representations. We find in Section 7 that a simple smoothness objective helps our approach in similar way. A full investigation of other representaiion-leaningg properties (e.g. value preservation (Abel et al., 2016), symbol construction (Konidaris et al., 2018), suitability for planning (Kurutach et al., 2018), information compression (Abel et al., 2019)) is beyond the scope of this paper.
Since our approach does not require any reward information and is agnostic to the underlying RL algorithm, it would naturally complement exploration methods designed for sparse reward problems (Pathak et al., 2017; Burda et al., 2018). Exploration helps to ensure that the experiences used to learn the abstraction cover as much of the ground MDP's state space as possible. In algorithms like HOMER (above) and the more recent Proto-RL (Yarats al., 2021), the exploration and representation learning objectives are intertwined, whereas our approach is, in principle, compatible with any exploration algorithm. Here we focus solely on the problem of learning Markov state abstractions and view exploration as an exciting direction for future work.
4 Markov State Abstractions
Recall that for a state representation to be Markov (whether ground or abstract), it must be a sufficient statistic for predicting the next state and expected reward, for any action the agent selects. The state representation of the ground MDP is Markov by definition, but learned state abstractions typically have no such guarantees. In this section, we introduce conditions that provide the missing guarantees.
Accurate abstract modeling the ground MDP requires replacing the fixed weighting scheme w(x) of Section 2 with a belief distribution, denoted by B [x) H). that measures the probability of each ground state x, conditioned on the entire history of agent experiences. Our objective is to find an abstraction & such that any amount of history can be summarized with single abstract state 2.
When limited to the most recent abstract state 2, B& may be policy-dependee and non-stationary:
B7, (x(=)) 1[((x)= = [) ERE: (T
PT = Tla
dEABEX
for t 2 1, where [-] denotes the indicator function, I the agent's (possibly non-stationary) behavior policy, and Po is arbitrary initial state distribution. Note that P. and B3, may still be non-stationary even if is stationary.
We generalize to a k-step belief distribution (for & 2 1) by conditioning (2) on additional history:
B (T(k) (xr (0-=);
1 1(b(x)) 21) Expexx T( (t-11, Bet Ba(k--) (fr-11 žt-1, {ar-is 2t Zt (3)
Ease, Ex T(ž, (1,, Tx-1) Ba (k-1) (ã--1 2t-1, fac-is 2- ) ((
where 1x((0) = B Any abstraction induces belief distribution, but the latter is only independent of history for a Markov abstraction. We formalize this concept with the following definition.
Definition 2 (Markov State Abstraction). Given an MDP M = (X, A, R,T.)), initial state dis- tribution Po. and policy class 1lc, a state abstraction X Z is Markov if and only for any policy π E llc, & induces a belief distribution B such that for all x E X, E Z, a E A, and k 2 1:
(=(() (x/2e, (ai-s, 1+ = B (x/22).
In other words, Markov abstractions induce belief distributions that only depend on the most recent abstract state. This property allows an agent avoid belief distributions entirely, and base its decisions solely on abstract states. Note that Definition is stricter than Markov state representations
"This section closely follows Hutter (2016), except here we consider helief distributions over ground states, rather than full histories. An advantage of Hutter's work is that it also considers abstractions over histories, though it only provides rough sketch of how learn such abstractions. Extending our learning objective support histories a natural direction future work.
"Tii can happen, for example, when the policy induces either Markov chain that does not have a stationary distribution, whose stationary distribution different from Po.

---

(Def. 1). An abstraction that collapses every ground state to single abstract state still produces Markov state representation, but for non-trivial ground MDPs also induces a history desended belief distribution.
Given these definitions, (see Eqn. (1)) follows:
define the abstract transitions and rewards for the policy class T.
(7.(2'1a,z) 11, Σ
ET(x aa (x))),
(4)
x'Ez'
R(x', a x T x) B (zlz).
(2 a,
Σ
(5)
T5 (2
x'Es'
DE:
Conditioning the belief distribution on additional history yields k-step versions compatible with Definition In the special case where B is stationary and policy-independee (and rewards defined over state-action pairs), we recover the fixed weighting function w(x) of Li et al. (2006).
4.1 Sufficient Conditions for a Markov Abstraction
The strictly necessary conditions for ensuring an abstraction is Markov over its policy class IIg depend on T' and R, which are typically unknown and hard to estimate due to X's dige--imensionality. However, we can still find sufficient conditions without explicitly knowing T and To this, we require that two quantities are equívalent in M and Mg: the inverse dynamics model, and density ratio that we define below. The inverse dynamics model I (ala', (0)a',))) defined in terms of the transition function T'xx' (a, 2) and expected next-state dynamics PT (z"[z) via Bayes' theorem:
I (a|x', x)) == T(x ([ (ala)) where PF (x'lz) = ΣãEA z)) The same true of
their abstract counterparts, 18, [g,((l2', and P PT. (2'122.
Theorem 1. If Φ X Z is state abstraction of MDP M = A,,,,, such that for any policy n in the policy class ll the following conditions hold every timestep t:
1. Inverse Model. The ground and abstract inverse models are equal: 1a,''' I (a)x', for all E A; z,z' E 7; x,,x' X, such that <(x) and Φ(x') z'.
2. Density Ratio. The ground and abstract next-state density ratios are equal, when con- P. (='la) r (x'
PT(z)) for all 2, Z; x' X. such that d(x') = 2'. where (x² = ΣxEx P7 Bx, and PT(2)) =
ditioned the same abstract state:
(p(2))
Evex (T (
Then o is a Markov state abstraction.
Corollary 1.1. if X + 7 is Markov state abstraction of MDP M (X, A, over the policy class Llg. then the abstract decision process Ms = (Z, R 6,11 Ty, 7) also Markov.
We defer all proofs to Appendix D.
Theorem 1 describes pair of conditions under which e is a Markov abstraction. Of course, the conditions themselves do not constitute training objective- we can only use them to confirm an abstraction is Markov. In Section 5, we adapt these conditions into a practical representation learning objective that is differentiable and suitable for learning & using deep neural networks. First, we show why the Inverse Model condition alone insufficient.
4.2 An Inverse Model Counterexample
The example MDP in Figure 2 additionally demonstrates why the Inverse Model condition alone is insufficient to produce a Markov abstraction. Observe that any valid transition between two ground states uniquely identifies the selected action. The same is true for abstract states since the only way to
For the general definitions that support arbitrary policies, see Appendix C.
7

---

reach 2B is via action 01, and the only way to leave is action 4g- Therefore, the abstraction satisfies the Inverse Model condition for any policy. However, as noted in Section 2.1, conditioning on additional history changes the abstract transition probabilities, and thus the Inverse Model condition not sufficient for an abstraction be Markov. In fact, we show in Appendix D.2 that, given the Inverse Model condition, the Density Ratio condition is actually necessary for a Markov abstraction.
5 Training Markov Abstraction
We now present set of training objectives for approximately satisfying the conditions of Theorem 1. Since the theorem applies for the policy class 1I induced by the abstraction, we restrict the policy by defining и mapping from Z Pr(A), rather than from X + Pr(A). In cases where is defined implicitly via the value function, we ensure that the latter is defined over abstract states.
Inverse Models. To ensure the ground and abstract inverse models are equal, we consider a batch of N experiences (Iz, x,),. encode ground states with 6, and jointly train model fa ((z)), ((x)); 0y) to predict distribution over actions, with a, as the label. This can be achieved by minimizing cross-entropy loss, for either discrete continuous action spaces:
N
Line == N Σ log (a a: (lx)) o(2));
Note that because the policy class is restricted to IIg. if the policy is stationary and deterministic, then 1 (a)2'' 2) = (a(a)z) x (a]x))
1.T (a|x', (olx'xz) x)) and the Inverse Model condition satisfied trivially. Thus we expect Line to be most useful for representation learning when the policy has high entropy or is changing rapidly, such as during early training.
Density Ratios. The second condition, namely that P] = (a'2))
P (a') means we can distinguish
P5 (x,
conditional samples from marginal samples equally well for abstract states or ground states. This objective naturally lends itself to type of contrastive loss. We generate batch of N sequential state pairs (Iy, x as samples of Pr(z'|a (z). and a batch N non-sequential state pairs (I,, a() as samples of Pr(x")), ). where the latter pairs can be obtained, for example. by shuffling the z! states in the first batch. We assign positive labels (y, = 1) to sequential pairs and negative labels to non-sequential pairs. This setup, following the derivation of Tiao (2017), allows to write density ratios in terms of class- p(y=1|z,x') Pr(=') T- 1-p(y=1jmaa" and (2') = Pr(=') T=
posterior probabilities: 8(x')= == Pe(æ (w)
Pr(z'[s)
qiv=llz.=")
where p and are just names for specific probability distributions.? 7 We jointly train abstraction and a classifier g(ylo (tyl(((')) (xx): 0g), minimizing the cross-entropy between predictions and labels ys:
2N
1
Ratio ==
- (l(2))
2N
In doing so, we ensure 9 approaches p and simultaneously, which drives (x'))) - 5(x').
Note that this is stronger than the original condition, which only required the ratios to be equal in expectation. This stronger objective, when combined with the inverse loss, actually encodes a novel form of the kinematic inseparability conditions from Section 3.5, which further helps avoid representation collapse. (See Appendix F for more details.)
Smoothness. To improve robustness and encourage our method to learn smooth representations like those discussed in Section 3.6, we optionally add additional term our loss function:
Lsmooth = (ReLU([)o(z') <(x))12 do))
This term penalizes consecutive abstract states for being more than some predefined distance de away from each other. Appendix L describes an additional experiment and provides further justification for why representation smoothness is an important consideration that complements the Markov property.
For completeness, reproduce the derivation Appendix E.

---

Markoy Abstraction Objective. We generate batch of experiences using a mixture of abstract policies π; E Ic S Hg (for example, with a uniform random policy), then train Φ end-to-end while minimizing a weighted combination of the inverse, ratio, and smoothness losses:
Markov == aLine + 3C Ratio + Ty
where a, 3, and 1) are coefficients that compensate for the relative difficulty of minimizing each individual objective for the domain in question.
The Markov objective avoids the problem of representation collapse without requiring reward information or ground state prediction. A trivial abstraction like 6(x) + 0 would not minimize
Markow, because it contains no useful information for predicting actions or distinguishing authentic
transitions from manufactured ones.
Offline Abstraction Learning for Visual Gridworlds
First, we evaluate our approach for learning an abstraction offline for visual gridworld domain (Fig. 1. left). Each discrete (æ, y) position in the 6 x 6 gridworld is mapped to noisy image (see Appendix
G). We emphasize that the agent only sees these images; it does not have access to the ground-truth (x, y) position. The agent gathers a batch of experiences in version of the gridworld with no rewards or terminal states, using uniform random exploration policy over the four directional actions.
These experiences are then used offline to train an abstraction function @arkko by minimizing Markon (with a 3 1, r = 0). We visualize the learned 2-D abstract state space in Figure 3a (top row) and compare against ablations that train with only Line or Ratio as well as against two baselines that we train via pixel prediction and reconstruction, respectively (see Appendix I for more visualizations). We observe that @arkkou and IIne cluster the noisy observations and recover the 6 x 6 grid structure, whereas the others do not generally have an obvious interpretation. We also observed that @atio and Autoonn frequently failed to converge.
Next froze these abstraction functions and used them map images to abstract states while training DQN (Mnih et al., 2015) on the resulting features. We measured the learning performance of each pre-trained abstraction, well as that of end-to-end DQN with pretraining. We plot learning curves in Figure 3b. For reference, we also include learning curves for a uniform random policy and DQN trained ground-truth (xr y) position with abstraction.
Markov abstractions match the performance of ground -ruth position, and beat every other Iearned representation except @Ine. Note that while @arkkov and ÓIne perform similarly in this domain, there is no reason to expect Line to work on its own for other domains, since it lacks the theoretical motivation of our combined Markov loss. When the combined loss is minimized, the Markov conditions are satisfied. But even if the inverse loss goes to zero on its own, the counterexample in Section 4.2 demonstrates that this is insufficient to learn a Markov abstraction.
0
-00
-60
-80
20
60
80
300
Episode
Exper OU
Auseet
Padd Prdd
Weal
sender
(a)
(b)
Figure 3: (a) Visualization of learning progress selected times (left to right) of 2-D state abstraction for the 6 x visual gridworld domain: (top row) CMrkkmi (middle row) Live only; (bottom row) CRatio only. Color denotes ground-trum truth (r. y) position, which is not shown to agent. (b) Mean episode reward for the visual gridworld navigation task. Markov abstractions significantly outperform end-to-end training with visual inputs, and match the performance of the expert (x., v) position features. (300 seeds: 5-point moving average; shaded regions denote 95% confidence intervals.)

---

7 Online Abstraction Learning for Continuous Control
Next, we evaluate our approach in an online setting with a collection of image-based. continuous control tasks from the DeepMind Control Suite (Tassa et al., 2020). Our training objective is agnostic about the underlying RL algorithm, so we use as our baseline the state-of-the-ar technique that combines Soft Actor-Critic (SAC) (Haarnoja et al., 2018) with random data augmentation (RAD) (Laskin et al., 2020a). We initialize replay buffer with experiences from a uniform random policy, as is typical, but before training with RL, we use those same experiences with reward information removed to pretrain a Markov abstraction. We then continue training with the Markov objective alongside traditional RI. (See Appendix H for implementation details).
In Figure 4, we compare against unmodified RAD, as well as contrastive methods CURL (Laskin et al., 2020b) and CPC (Van den Oord et al., 2018), bisimulation methods DeepMDP (Gelada et al., 2019) and DBC (Zhang et al., 2021), and pixclrreconstruction method SAC-AE (Yarats et al., 2019). As reference, we also include ninuvi SAC with expert features. All methods use the same number of environment steps (the experiences used for pretraining are not additional experiences).
Task Cartpolo, Swingup
Task = Ball-in-up,, Catch
Task Cheotah,
1000
800-
Reward
600-
400-
200
0-
20020
40000
s0onn
anion 100000
20000 40000
800n0 80000 100005
1000on 200003
sprioo asaion 500000
lask Finger, Spmn
Task Reacher, Easy
Taask Walker, Walk
10000
Reward
Mattin
100000
200000
300000
103000
XC0000
500000
100000
200000
300000
Sleps
Sleps
Shaps
Figure 4: Mean episode reward vs. environment steps for DeepMind Control. Adding our Markov objective leads to improved learning performance. (10 sceds; 5-point moving average; shaded regions denote 90% confidence intervals; learning curve data is available the linked code repository.)
Relative to RAD, our method learns faster on four domains and slower on one, typically achieving the same final performance (better in one, worse in onc). performs even more favorably relative to the other baselines, of which CURL is most similar to our method, since it combines contrastive learning with data augmentation similar to that of RAD. Our approach even represents a marginal improvement over hypothetical "best of" oracle that always chooses the best performing baseline. These experiments show that even in an online setting, where the agent can leverage reward informa- tion and a Markov ground state when building its abstract state representation, explicitly encouraging Markov abstractions improves learning performance over ofat-oo-the-aa image-based RL.
Conclusion
We have developed a principled approach to learning abstract state representations that provably results in Markov abstract states, and which does not require estimating transition dynamics nor ground-state prediction. We defined what it means for state abstraction to be Markov while ensuring that the abstract MDP accurately reflects the dynamics of the ground MDP, and introduced sufficient conditions for achieving such abstraction. We adapted these conditions into a practical training objective that combines inverse model estimation and temporal contrastive learning. Our approach learns abstract state representations, with and without reward, that capture the structure of the underlying domain and substantially improve learning performance over existing approaches.
'We ran another experiment with no data augmentation. using different state-of ofte-of-the-art continuous control algorithm, RBF-DQN (Asadi et al., 2021), and found similar results there as well (see Appendix K for details).
10

---

Acknowledgments and Disclosure of Funding
Thank you to Ben Abbatematteo, David Abel, Barrett Ames, Séb Arnold, Kavosh Asadi, Akhil Bagaria, Jake Beck, Jules Becker, Alexander Ivanov, Steve James, Michael Littman, Sam Lobel, and our other colleagues at Brown for thoughtful advice and countless helpful discussions, as well as Amy Zhang for generously sharing baseline learning curve data, and the anonymous ICML'21 and NeurlPS'21 reviewers for their time, consideration, and valuable feedback that substantially improved the paper. This research was supported by the ONR under the PERISCOPE MURI Contract N00014- 17-1-2699. by the DARPA Lifelong Learning Machines program under grant FA8750-18-2-0117, and by NSF grants 1955361, 1717569, and CAREER award 1844960.
References
Abel, D., Hershkowitz, E., and Littman, M. Near optimal behavior via approximate state abstraction. In Proceedings of the 33rd International Conference on Machine Learning, volume 48, pp. 2915-2923, 2016.
Abel, D., Arumugam, D.,, Lehnert, L., Littman, L. State abstractions for lifelong reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 10-19, 2018.
Abel, D., Arumugam, D., Asadi, K., Jinnai, Y., Littman, M. L., and Wong, L. L. State abstraction as compression in apprenticeship learning. In Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence, pp.
3 3134-3142, 2019.
Agrawal, P., Nair. A., Abbeel, P., Malik, J., and Levine, S. Learning to poke by poking: Experiential learning of intuitive physics. In Advances Neural Information Processing Systems, volume 29, pp. 5074-5082, 2016.
Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-.., and Hjelm, R. D. Unsupervised state representation learning in Atari, In Advances Neural Information Processing Systems, volume 32, pp. 8769 8782, 2019.
Asadi, K., Misra, D., and Littman, M. Lipschitz continuity in model-based reinforcement learning. In Proceedings of the 35th International Conference on Muchine Learning, volume 80, pp. 264-273, 2018
Asadi, K., Parikh, N., Parr, E., Konidaris, D., and Littman, Deep radial basis value functions for continuous control. In Proceedings of the Thirty- Fifth AAAI Conference on Artificial Intelligence, pp.
6696 6704, 2021.
Bai, A., Srivastava, S., and Russell, S. J. Markovian state and action abstractions for MDPs via hicrarchical MCTS. In Proceedings of the Twenty-Fifth International Joint Conference Artificial Intelligence, pp. 3029-3037, 2016.
Biza, O., Platt, R., Van de Meent, J., and Wong. L. L. S. Learning discrete state abstractions with deep
variational inference. Third Symposium on Advances in Approximate Bayesian Inference, 2021.
Burda, Y., Edwards, H., Storkey, J., and Klimov, Exploration by random network distillation. arXiv, art. 1810.12894. 2018
Castro, Scalable methods for computing state similarity in deterministic Markov decision processes. In Proceedings of the Fhurty-Fourth AAAI Conference on Artificial Intelligence, pp. 10069-10076, 2020.
Choi, J., Guo, Y., Moczulski, M., Oh, J., Wu, N., Norouzi, M., and Lec, H. Contingency- aware exploration in reinforcement learning. In International Conference on Learning Representations, 2019.
Christiano, P., Shah, 7., Mordatch, I., Schneider, J., Blackwell, T., Tobin, J., Ahbeel, P., and Zaremba, W. Transfer from simulation to real world through learning deep inverse dynamics model. arXiv, art. 1610.03518, 2016.
Coreel, D., Gerstner, W., and Brea, Efficient model- based deep reinforcement learning with variational state tabulation. Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 1049-1058, 2018.
Dean, T. and Givan, R. Model minimization Markov decision processes. In Proceedings of the Fourteenth National Conference on Artificial Intelligence, pp. 106-111, 1997.
Du. S. S., Krishnamurthy, A., Jiang. N., Agarwal, A., Dudík, M., and Langford, J. Provably efficient RL with rich observations via latent state decoding. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 1665-1674 2019.

---

Ferns, N., Panangaden, P., and Precup. D. Metrics for finite Markov decision processes. In Proceedings of the 20th Conference on Uncertainty Artificial Intelligence, volume 4, pp. 162-169, 2004.
Ferns, N., Panangaden, P., and Precup. D. Bisimulation metries for continuous Markov decision processes. SIAM Journal on Computing. 40(6):1662-11)) 1714, 2011.
Finn, C., Tan, X. Y., Duan, Y., Darrell, T., Levine, S., and Abbeel, P. Deep spatial autoencoders for visuomotor learning. In IEEE International Conference Robotics and Automation, pp- 512-519, 2016.
Gelada, C., Kumar, S., Buckman, J., Nachum, O., and Bellemare, M. G. DeepMDP: Learning continuous latent space models for representation learning. In Proceedings of the 36th International Conference on Machine Learning, volume 97. pp. 2170-2179, 2019.
Gutmann, M. and Hyvärinen, A. Noise-contrastive estimation: A estimation principle for unnormalized statistical models. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pp. 297-304 2010.
Ha, D. and Schmidhuber, World models. arXiv, art. 1803.10122, 2018.
Haamoja, T., Zhou, A., Abheel, P., and Levine, S. Soft actor-critic: Off-poliey maximum entropy deep reinforcement learning with stochastic actor. In Proceedings of the International Conference on Machine Learning, pp. 1861-1870, 2018.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Van Hasselt, H., Hessel, M., and Aslanides, J. When to use parametric models reinforcement learning? In Advances in Neural Information Processing Systems, volume 32, pp. 14322-14333, 2019.
Higgins, I., Pal, A., Rusu, A., Matthey, I.., Burgess, C., Pritzel, A., Botvinick, M., Blundell, C., and Lerchner, DARLA: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pp. 1480-1490, 2017.
Hutter, M. Extreme state aggregation beyond Markov decision processes. Theoretical Computer Science, 650: 73-91, 2016.
Jonschkowski, R. and Brock, O. Learning state representations with robotic priors. Autonomous Robots, 39(3): 407-428. 2015.
Kaiser, L., Babacizadeh, M., Milos, P., Osiński, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., Mohiuddin, A., Sepassi, R., Tucker, G., and Michalewski, H. Model based reinforcement learning for Atari. In International Conference on Learning Representations, 2020.
Konidaris, G., Kaclbling, L. P., and Lozano-Percz, T. From skills to symbols: Learning symbolic representations for abstract legh planning. Journal of Artificial Intelligence Research, 61:215-200 2018.
Kurutach, T.. Tamar, A., Yang, G., Russell, S., and Abbeel, P. Leamnin plannable representations with causal InfoGAN. In Advances in Neural Information Processing Systems, volume 31, pp. 8747-8758, 2018.
Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A. Reinforcement Icarning with augmented data. Advances Neural Information Processing Systems, volume 33, pp. 19884 19895, 2020a.
Laskin, M., Srinivas, A., and Abbeel, CURL: Contrastive unsupervised representations for reinforcement learning. Proceedings of the 37th International Conference on Machine Learning, volume 119, pp- 5639-5666 1659,565,,22 2020b.
Lee, X., Nagabandi. A., Abbeel, P., and Levine, S. Stochastic latent actor-critic: Deep reinforvement learning with a latent variable model. In Advances Neural Information Processing Systems, volume 33, pp. 741-752. 2020.
Lehnert, L. and Littman, M. L. Successor features combine elements of model free and model-basee reinforce- ment learning. Journal of Machine Learning Research, 21(196): 53, 2020.
Li, L., Walsh, T. J., and Littman, M. L. Towards a unified theory of state abstraction for MDPs. In International Symposium on Artificial Intelligence and Mathematics, 2006.
Mattner, J., Lange, S., and Riedmiller, M. Learn to swing up and balance a real pole based on raw visual input data. In Proceedings of the 19th International Conference on Neural Information Processing, Part V. pp. 126-133. 133. Springer, 2012.
12

---

Misra, D.. Henaff, M.. Krishnamurthy. A., and Langford, J. Kinematic state abstraction and provably efficient fich-observation reinforcement learning. In Proceedingy of the 37th International Conference on Machine Learning, volume 119, pp. 6961-6971, 2020.
Mnih. V.. Kavukcuuglu, K., Silver, D., Rusu, A. A., Veness, J.. Bellemare, M. G., Graves, A.. Riedmiller, M.. Fidjeland, K., Ostrovski, G., Petersen. S., Beattie, C., Sadik, A., Antonoglou. I., King. H., Kumaran, D.. Wierstra, D., Legg. S., and Hassabis, D. Human-level control through deep reinforcement learning. Nature, (18(7540):529-5388 2015.
Van den Oord, A., Li, Y., and Vinyals, O. Representation learning with contrastive predictive coding. arXiv, art. 1807.03748, 2018.
Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. Curiosity-driven exploration by self-supervised prediction. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pp. 2778-2788 2017.
Pazis, J. and Parr, R. PAC optimal exploration continuous space Markov decision processes. In Proceedings of the Twenty-Sevent AAAI Conference on Artificial Intelligence, pp. 774-781, 2013.
Pirotta, M., Restelli, M., and Bascetta, L. Policy gradient in Lipschitz Markov decision processes. Machine
Learning, 100(2-3):255 203, 2015.
Van der Pol, E., Kipf, T., Oliehoek. A., and Welling, M. Plannable approximations to MDP homomorphisms: Equivariance under actions. In Proceedings of the 19th International Conference Audonomous Agents and Multiagent Systems, pp. 1431-1439, 2020.
Ravindran, B. and Barto, A. G. Approximate bomomorphisms: A framework for non non-exact minimization in Markov decision processes. In Proceedings of the Fifth International Conference Knowledge Based Computer Systems. 2004.
Shelhamer, E., Mahmoudich, P. Argus, M., and Darrell, T. Loss is its reward: Self-supervision for reinforcement learning. arXiv, art. 1612.07307, 2016.
Song. Z., Parr, R., Liao, X., and Carin, L. Linear feature encoding for reinforcement learning. In Advances in Neural Information Processing Systems, volume 29, pp. 4224 4232, 2016.
Stooke, A., Lee, K., Abbeel, P., and Laskin, M. Decoupling representation learning from reinforcement learning. arXiv, art. 2009,08319, 2020.
Tassa, Y., Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap. T., and Heess, dm control: Software and tasks for continuous control. arXiv, art. 2006.12983, 2020.
Tiao. L. A simple illustration of density ratio estimation and KL divergence estimation by probabilistic and-kl- -iveerereeeee-
classification. htt::/louistiao.mm wttp://iiwi..eeooeeaaaaaiaiiiaii simple-illustratinn illustration--ff oe-eoostt--dddni-----oeooo ceassit--rato-sstinii estimation-b by- probablistii--clasifccciiee 2017. [Online; accessed (6-March-2020],
Watter, M., Springenberg, J., Boedecker, J., and Riedmiller, M. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems, volume 28, pp. 2746 2754, 2015.
Yarats, D. and Kostrikov, I. Soft actor-critic (SAC) implementation in PoTorch. https:/ggitubb.oomddeenesaaaza pytorch_sac, 2020.
Yarats, D.. Zhang, A., Kostrikov, L., Amos, B., Pineau, J., and Fergus, R. Improving sample efficiency in model-free reinforcement learning from images. arXiv. art. 1910.01741, 2019.
Yarats, D., Fergus, R., Lazaric, A., and Pinto, L. Reinforcement learning with prototypical representations. arXiv, art. 2102. 11271, 2021.
Zhang. A., Satija, H., and Pineau, J. Decoupling dynamics and reward for transfer learning. arXiv, art. 10089, 2018.
Zhang, A., McAllister, R. T., Calandra, R., Gal, Y., and Levine, S. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
Zhang, Z., Chen, X., and Littman, M. Covering number as complexity measure for POMDP planning and learning. In Proceedings of the Twenty-Sixth AAAI Conference on Artificial Intelligence, pp. 1853-1859, 2012.
13


Reviews:
[{"review_id": "_V6gK7mJV4L", "reviewer": "Reviewer_tW7x", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "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.\nOverall the paper is fairly well written, and the contributions are placed well in the context of prior work.\n\nNonetheless, some aspects of the method could deserve further justification or investigation, and I have a few questions for clarifications:\n\nI 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.\nIt’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.\nThis 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!\n\nI’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.\nAnother 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?\n\nA 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?\n\nRegarding 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?).\nThe results are nice, if not astounding, which is unsurprising as the method builds on previously-explored objectives.\n\nAnother 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?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "In the main review I've made suggestions for some areas of potential limitations that could be discussed further.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "1Zz9MwY9MW3", "reviewer": "Reviewer_vzzj", "summary": "# Learning Markov State Abstractions for Deep Reinforcement Learning\n\nThe paper proposes a self-supervised representation learning method for deep RL agents that ensures the learned representation is Markov and non-trivial (avoids collapse).\n\nIn 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## For the response\n\n### Major concerns:\n1. How much of the improvement is due to RAD + L_inv?\n2. Is the improvement to the SoTA significant?\n\nI 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.\n\n### Minor concerns:\nAs 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?\n\n## Strengths\n\nI 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.\n\nThe 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.\n\n## Weaknesses\n\nThe 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.\n\n## Correctness\n\nThe paper seems correct, with some gaps between the theoretical results and the method implementation.\n\n## Clarity\n\nThe 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.\n\n## Relation to prior work\n\nI 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:\n* https://arxiv.org/pdf/1611.05397.pdf\n* https://arxiv.org/pdf/1906.09237.pdf\n* http://proceedings.mlr.press/v119/guo20g/guo20g.pdf\n* https://arxiv.org/pdf/2007.05929.pdf\n\nFor 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. \n\n## Reproducibility\n\nI think the text and the appendix add enough information for reproducibility. \n\n## Additional Feedback\n\nThank you for submitting your paper for review. I would like to share some comments that I hope you will find useful.\n\nThe 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.\n\nI 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.\n\nAnother 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).\n\nFor 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.\n\nI 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.\n\nYou 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$.\n\nWhere does $\\Pi_C$ show up in Def. 2? I guess in $B$?\n\nI could not see the definition of $y_i$ (but I could infer one).\n\nThe plot in Fig 3(b) has very small font.\n\nIn 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:\nThe Markov training objective includes other components (L_smooth).\nWhy is the method SoTA? In Cheetah run it's outperformed by RAD. In Reacher it's outperformed by CURL.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "0Cb2M4KPsuE", "reviewer": "Reviewer_KWGW", "summary": "Post author response update:\nSome 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.\n\nReasons for increasing score:\n* 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.\n* 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).\n\nRemaining concerns:\n* 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.\n* 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.\n\nEnd of update\n\n\nThis 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The 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.\n\nThe comparison to HOMER (L169-179) is crucial, and more space should be devoted in the main body to it. \nThough 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.\nIt 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.\n\nOne 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.\n\nIn 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?\n\nFor 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).\n\nCombined 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).\n\nThe 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).\n\nNone 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.)\n\nMinor Comments:\n\n-The list of work in L22-27 is not helpful toward understanding this work. It should be shortened / moved elsewhere.\n\n-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).\n\n-The Background would benefit from a different structure (currently only 1 subsection, with remainder of comment in the \"parent\").\n\n-L102-104 is out of place in the Background.\n\n-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)).\n\n-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.\n\n-Since Section 5, Density Ratios largely follows existing work, it can be shortened to make room for other content.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, limitations have been sufficiently addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5 + 2.5 (post response)"}, "extra_text": {}}, {"review_id": "-z4d1WCCrWN", "reviewer": "Reviewer_br5G", "summary": "Summary\n-------\n\n\nThis paper analyzes the problem of learning Markov state abstractions in\nreinforcement learning. The authors first prove that the abstraction\nmust satisfy two conditions: an inverse model equality and a density\nratio equality. These two conditions cannot be enforced in the training\nobjective, so an approximate regularizer is proposed instead. The\nproposed algorithm is compared with several baselines in an offline\nlearning regime (visual gridworld) which evaluates the quality of the\nstate abstraction and an online learning regime (continuous control)\nwhich evaluates the performance.\n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Decision\n--------\n\nOverall, the paper provides a well-written analysis of learning Markov\nstate abstractions and is a valuable addition to the RL literature. The\npaper is held-back by the lack of insight in the experimental\ninvestigations. While the experiments on GridWorld show that their\nproposed method does learn a Markov state abstraction, their further\nexperiments do not clearly show the benefit of a Markov state\nabstraction. Taken together, I am rating this paper tentatively at a 5\n(\"Marginally below the acceptance threshold\").\n\n\nStrengths\n---------\n\n-   The theoretical analysis of learning Markov abstractions is\n    interesting and novel. I particularly like that the work highlights\n    that the density ratio condition is often ignored in the previous\n    literature, comments on why this might be the case and later shows\n    it through an ablation study in Section 6. This \"closing the loop\"\n    from theory to empirics is excellent.\n-   Presentation of the work and the writing quality is very high. While\n    some aspects of the technical sections are difficult to follow, I\n    think the paper does a superb job of motivating and explaining the\n    decisions made throughout the paper.\n\nWeaknesses\n----------\n\n\n-   The experiments are not well motivated, and do not meaningfully\n    answer whether learning Markov abstractions is a worthwhile\n    endeavor.\n-   While out of scope for this paper, I think there is a lacking\n    discussion on the trajectory/memory component of learning Markov\n    state representations. Of course, the goal of this paper is to study\n    state abstractions. However, the Markov property is less important\n    for high dimensional observations like images (where, often the\n    images are used as-is). The real interest and challenge lies in\n    learning a Markov state by remembering information from previous\n    observations.\n\nDetailed Comments\n-----------------\n-   Section 4 builds on technical machinery from Section 2 in such a way\n    that I often needed to go back and forth between the two sections. I\n    feel that part of the problem is that the two sections are separated\n    by related work. However, I do not believe that switching section 2\n    and 3 is necessarily the solution. Instead, I think the authors\n    should contextualize some of the technical machinery such as the\n    belief distributions using their working example (the MDP in figure\n    2).\n\n-   Section 7: I think the results here, while impressive, completely\n    undermine the rest of the paper. If Markov states are important,\n    then surely \"expert\" Markov states are the gold standard for state\n    representation. And yet, you show that this is often not the case.\n    Indeed, your method and RAD seem to outperform the expert in many\n    cases. If this is true, then what is the motivation for learning a\n    markov abstraction to begin with? And what is the reason for the\n    proposed algorithm's success?\n\nMinor Comments\n--------------\n\n-   Figure 3b: The experiment is in the offline learning section, yet\n    the x-axis shows performance over a number of episodes. This seems\n    to suggest that the experiment is actually online.\n-   Figure 3b, 4: Y-axis says reward, should be return?\n-   Line 290: How exactly is it mapped to an image? Are the pixels\n    equidistant and what is the resolution for the image?\n\nPost Rebuttal\n--------------\nIn light of the discussion below, I have increased my score from a 5 -> 6.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

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)