Learning Markov State Abstractions for Deep Reinforcement Learning

1 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 succinct expert state information is available , agents trained on such expert state features usually outperform agents trained on rich observations . Much recent work ( Shelhamer et al. , 2016 ; Pathak et 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 of representation-learning objectives that help the agent learn abstract representations with various desirable properties . Perhaps the most obvious property to incentivize in a state 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 of 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 are not necessarily Markov , even when built on top of MDPs . This is due to the fact that abstraction necessarily throws away information . Discard too much information , and the resulting representation can not accurately characterize the environment . Discard too little , and agents will fail to close the representation gap . Abstraction must balance between ∗Please send any correspondence to Cameron Allen < csal @ brown.edu > . Code repository available at https : //github.com/camall3n/markov-state-abstractions . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ignoring irrelevant information and preserving what is important for decision making . If reward feedback is available , an agent can use it to determine which state information is relevant to 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 or non-existent , or observations are sufficiently complex . We introduce a new approach to learning Markov state abstractions . We begin by defining a set of theoretical conditions that are sufficient for an abstraction to retain the Markov property . We next show that these conditions are approximately satisfied by simultaneously training an inverse model to predict the action distribution that explains two consecutive states , and a discriminator to determine whether two given states were in fact consecutive . Our combined training objective ( architecture shown in Fig . 1 , 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 : a gridworld navigation task ( Fig . 1 , left ) and a set of continuous control benchmarks . In the gridworld , we construct an abstract representation offline—without access to reward feedback—that 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 RL , where it leads to a significant performance improvement over state-of-the-art visual representation learning . 2 Background . A Markov decision process M consists of sets of states X and actions A , reward function R : X × A × X → R , transition dynamics T : X × A → Pr ( X ) , and discount factor γ . For our theoretical results , we assume M is a Block MDP ( Du et al. , 2019 ) whose behavior is governed by a much smaller ( but unobserved ) set of states S , and where X is a rich observation generated by a noisy sensor function σ : S → Pr ( X ) , as in Figure 1 ( 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 a perfect “ inverse sensor ” function σ−1 ( x ) 7→ s , which means the observations are themselves Markov , as we define below . Note that S , σ , and σ−1 are all unknown to the agent . Definition 1 ( Markov State Representation ) . A decision process M = ( X , A , R , T , γ ) and its state representation X are Markov if and only if T ( k ) ( xt+1| { at−i , xt−i } ki=0 ) = T ( xt+1|at , xt ) and R ( k ) ( xt+1 , { at−i , xt−i } ki=0 ) = R ( xt+1 , at , xt ) , for all a ∈ A , x ∈ X , k ≥ 1 . The superscript ( k ) denotes that the function is being conditioned on k additional steps of history . The Markov property means that each state X it is a 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 in 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 → R , which is defined as the expected sum of future discounted rewards starting from a given state and following the policy π thereafter . The agent ’ s objective is to learn an optimal policy π∗ that maximizes value at every state . Note that the assumption that the optimal policy is stationary and Markov—that it only depends on state—is appropriate only if 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 high-dimensional for tractable learning , we turn to state abstraction . Our objective is to find an abstraction function φ : X → Z mapping each ground state2 x to an abstract state z = φ ( x ) , with the hope that learning is tractable using the abstract representation Z ( see Figure 1 , left ) . Since our goal is to support effective abstract decision making , we are mainly concerned with the policy class Πφ , the set of policies with the same behavior for all ground states that share the same abstract state : Πφ : = { π : ( φ ( x1 ) = φ ( x2 ) ) =⇒ ( π ( a|x1 ) = π ( a|x2 ) ) , ∀ a ∈ A ; x1 , x2 ∈ X } . ( 1 ) An abstraction φ : X → Z , when applied to an MDP M , induces a new abstract decision process Mφ = ( Z , A , T π φ , t , R π φ , t , γ ) , whose dynamics may depend on the current timestep t or the agent ’ s behavior policy π , and , crucially , which is not necessarily Markov . Consider the following example . The figure above depicts a four-state , two-action MDP , and an abstraction φ where φ ( x1 ) = φ ( x2 ) = zB . It is common to view state abstraction as aggregating or partitioning ground states into abstract states in this way ( Li et al. , 2006 ) . Generally this involves choosing a fixed weighting scheme w ( x ) to express how much each ground state x contributes to its abstract state z , where the weights sum to 1 for the set of x in each z . We can then define the abstract transition dynamics Tφ as a w-weighted sum of the ground dynamics ( and similarly for reward ) : Tφ ( z′|a , z ) = ∑ x′∈z′ ∑ x∈z T ( x ′|a , x ) w ( x ) . A natural choice for w ( x ) is to use the ground-state visitation frequencies . For example , if the agent selects actions uniformly at random , this leads to w ( x0 ) = w ( x3 ) = 1 and w ( x1 ) = w ( x2 ) = 0.5 . In this formulation , the abstract decision process is assumed to be Markov by construction , and Tφ and Rφ are assumed not to depend on 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 : Tφ ( zA|at = a0 , zt = zB ) = 0.5 , whereas Pr ( zA|at = a0 , zt = zB , at−1 = a1 , zt−1 = zA ) = 1 . 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.3 Even worse , if the agent ’ s policy changes—such as during learning—this discrepancy can cause RL algorithms with bounded sample complexity in the ground MDP to make an arbitrarily large number of mistakes in the abstract MDP ( Abel et al. , 2018 ) . 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 , even in the simple example above , maintaining accurate weights for w ( x ) requires keeping track of an unbounded amount of history : if an agent reaches abstract state zB and repeats action a1 an arbitrarily large number of times ( N ) , then knowing precisely which ground state it will end up in ( x1 or x2 ) requires remembering the abstract state it originally came from N + 1 steps prior . Successful modeling of the transition dynamics for a subsequent action a0 hinges on exactly this distinction between x1 and x2 . The abstraction introduces partial observability , and to compensate , w ( x ) must be replaced with a belief distribution over ground states , conditioned on the entire history ; instead 2We refer to x ∈ X as ground states ( and M as the ground MDP ) , to reflect that these quantities are grounded , as opposed to abstract , i.e . they have a firm basis in the true environment . 3Abel et al . ( 2018 ) presented a three-state chain MDP where they made a similar observation . 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 dynamics models and contrastive learning to derive conditions under which a state abstraction is Markov . We then adapt these conditions into a corresponding training objective for learning such an abstraction directly from the agent ’ s experience in the ground MDP .

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.