Provable Rich Observation Reinforcement Learning with Combinatorial Latent States

1 INTRODUCTION . Most reinforcement learning ( RL ) algorithms scale polynomially with the size of the state space , which is inadequate for many real world applications . Consider for example a simple navigation task in a room with furniture where the set of furniture pieces and their locations change from episode to episode . If we crudely approximate the room as a 10× 10 grid and consider each element in the grid to contain a single bit of information about the presence of furniture , then we end up with a state space of size 2100 , as each element of the grid can be filled independent of others . This is intractable for RL algorithms that depend polynomially on the size of state space . The notion of factorization allows tractable solutions to be developed . For the above example , the room can be considered a state with 100 factors , where the next value of each factor is dependent on just a few other parent factors and the action taken by the agent . Learning in factored Markov Decision Processes ( MDP ) has been studied extensively ( Kearns & Koller , 1999 ; Guestrin et al. , 2003 ; Osband & Van Roy , 2014 ) with tractable solutions scaling linearly in the number of factors and exponentially in the number of parent factors whenever planning can be done efficiently . However , factorization alone is inadequate since the agent may not have access to the underlying factored state space , instead only receiving a rich-observation of the world . In our room example , the agent may have access to an image of the room taken from a megapixel camera instead of the grid representation . Naively , treating each pixel of the image as a factor suggests there are over a million factors and a prohibitively large number of parent factors for each pixel . Counterintuitively , thinking of the observation as the state in this way leads to the conclusion that problems become harder as the camera resolution increases or other sensors are added . It is entirely possible , that these pixels ( or more generally , observation atoms ) are generated by a small number of latent factors with a small number of parent factors . This motivates us to ask : can we achieve PAC RL guarantees that depend polynomially on the number of latent factors and very weakly ( e.g. , logarithmically ) on the size of observation space ? Recent work has addressed this for a rich-observation setting with a non-factored latent state space when certain supervised learning problems are tractable ( Du et al. , 2019 ; Misra et al. , 2020 ; Agarwal et al. , 2020 ) . However , addressing the rich-observation setting with a latent factored state space has remained elusive . Specifically , ignoring the factored structure in the latent space or treating observation atoms as factors yields intractable solutions . ∗Correspondence at : dimisra @ microsoft.com Contributions . We combine two threads of research on rich-observation RL and factored MDP by proposing a new problem setup called Factored Block MDP ( Section 2 ) . In this setup , observations are emitted by latent states that obey the dynamics of a factored MDP . We assume observations to be composed of atoms ( which can be pixels for an image ) that are emitted by the latent factors . A single factor can emit a large number of atoms but no two factors can control the same atom . Following existing rich-observation RL literature , we assume observations are rich enough to decode the current latent state . We introduce an algorithm FactoRL that achieves the desired guarantees for a large class of Factored Block MDPs under certain computational and realizability assumptions ( Section 4 ) . The main challenge that FactoRL handles is to map atoms to the parent factor that emits them . We achieve this by reducing the identification problem to solving a set of independence test problems with distributions satisfying certain properties . We perform independence tests in a domain-agnostic setting using noise-contrastive learning ( Section 3 ) . Once we have mapped atoms to their parent factors , FactoRL then decodes the factors , estimates the model , recovers the latent structure in the transition dynamics , and learns a set of exploration policies . Figure 1 shows the different steps of FactoRL . This provides us with enough tools to visualize the latent dynamics , and plan for any given reward function . Due to the space limit , we defer the discussion of related work to Appendix B . To the best of our knowledge , our work represents the first provable solution to rich-observation RL with a combinatorially large latent state space .. 2 THE FACTORED BLOCK MDP SETTING . There are many possible ways to add rich observations to a factored MDP resulting in inapplicability or intractability . Our goal here is to define a problem setting that is tractable to solve and covers potential real-world problems . We start with the definition of Factored MDP ( Kearns & Koller , 1999 ) , but first review some useful notation that we will be using : Notations : For any n ∈ N , we use [ n ] to denote the set { 1 , 2 , · · · , n } . For any ordered set ( or a vector ) U of size n , and an ordered index set I ⊆ [ n ] and length k , we use the notation U [ I ] to denote the ordered set ( U [ I [ 1 ] ] , U [ I [ 2 ] ] , · · · , U [ I [ k ] ] ) . Definition 1 . A Factored MDP ( S , A , T , R , H ) consists of a d-dimensional discrete state space S ⊆ { 0 , 1 } d , a finite action space A , an unknown transition function T : S × A → ∆ ( S ) , an unknown reward function R : S × A → [ 0 , 1 ] and a time horizon H . Each state s ∈ S consists of d factors with the ith factor denoted as s [ i ] . The transition function satisfies T ( s′ | s , a ) =∏d i=1 Ti ( s ′ [ i ] | s [ pt ( i ) ] , a ) for every s , s′ ∈ S and a ∈ A , where Ti : { 0 , 1 } |pt ( i ) | ×A → ∆ ( { 0 , 1 } ) defines a factored transition distribution and a parent function pt : [ d ] → 2 [ d ] defines the set of parent factors that can influence a factor at the next timestep . We assume a deterministic start state . We also assume , without loss of generality , that each state and observation is reachable at exactly one time step . This can be easily accomplished by concatenating the time step information to state and observations . This allows us to write the state space as S = ( S1 , S2 , · · · , SH ) where Sh is the set of states reachable at time step h. A natural question to ask here is why we assume factored transition . In tabular MDPs , the lower bound for sample complexity scales linearly w.r.t . the size of the state set ( Kakade , 2003 ) . If we do not assume a factorized transition function then we can encode an arbitrary MDP with a state space of size 2d , which would yield a lower bound of Ω ( 2d ) rendering the setting intractable . Instead , we will prove sample complexity guarantees for FactoRL that scales in number of factors as dO ( κ ) where κ : = maxi∈ [ d ] |pt ( i ) | is the size of the largest parent factor set . The dependence of κ in the exponent is unavoidable as we have to find the parent factors from all possible ( d κ ) combinations , as well as learn the model for all possible values of the parent factor . However , for real-world problems we expect κ to be a small constant such as 2 . This yields significant improvement , for example , if κ = 2 and d = 100 then dκ = 100 while 2d ≈ 1030 . Based on the definition of Factored MDP , we define the main problem setup of this paper , called Factored Block MDP , where the agent does not observe the state but instead receives an observation containing enough information to decode the latent state . Definition 2 . A Factored Block MDP consists of an observation space X = X m and a latent state space S ⊆ { 0 , 1 } d. A single observation x ∈ X is made of m atoms with the kth denoted by x [ k ] ∈ X . Observations are generated stochastically given a latent state s ∈ S according to a factored emission function q ( x | s ) = ∏di=1 qi ( x [ ch ( i ) ] | s [ i ] ) where qi : { 0 , 1 } → ∆ ( X |ch ( i ) | ) and ch : [ d ] → 2 [ m ] is a child function satisfying ch ( i ) ∩ch ( j ) = ∅ whenever i 6= j . The emission function satisfies the disjointness property : for every i ∈ [ d ] , we have supp ( qi ( · | 0 ) ) ∩ supp ( qi ( · | 1 ) ) = ∅.1 The dynamics of the latent state space follows a Factored MDP ( S , A , T , R , H ) , with parent function pt and a deterministic start state . The notion of atoms generalizes commonly used abstractions . For example , if the observation is an image then atoms can be individual pixels or superpixels , and if the observation space is a natural language text then atoms can be individual letters or words . We make no assumption about the structure of the atom space X or its size , which can be infinite . An agent is responsible for mapping each observation x ∈ X to individual atoms ( x [ 1 ] , · · · , x [ m ] ) ∈X m. For the two examples above , this mapping is routinely performed in practice . If observation is a text presented to the agent as a string , then it can use off-the-shelf tokenizer to map it to sequence of tokens ( atoms ) . Similar to states , we assume the set of observations reachable at different time steps is disjoint . Additionally , we also allow the parent ( pt ) and child function ( ch ) to change across time steps . We denote these functions at time step h by pth and chh . The disjointness property was introduced in Du et al . ( 2019 ) for Block MDPs—a class of richobservation non-factorized MDPs . This property removes partial observability concerns and enables tractable learning . We expect this property to hold in real world problems whenever sufficient sensor data is available to decode the state from observation . For example , disjointness holds true for the navigation task with an overhead camera in Figure 1 . In this case , the image provides us with enough information to locate all objects in the room , which describes the agent ’ s state .. Disjointness allows us to define a decoder φ ? i : X |ch ( i ) | → { 0 , 1 } for every factor i ∈ [ d ] , such that φ ? i ( x [ ch ( i ) ] ) = s [ i ] if x [ ch ( i ) ] ∈ supp ( qi ( . | s [ i ] ) ) . We define a shorthand φ ? i ( x ) = φ ? i ( x [ ch ( i ) ] ) whenever ch is clear from the context . Lastly , we define the state decoder φ ? : X → { 0 , 1 } d where φ ? ( x ) [ i ] = φ ? i ( x ) . The agent interacts with the environment by taking actions according to a policy π : X → ∆ ( A ) . These interactions consist of episodes { s1 , x1 , a1 , r1 , s2 , x2 , a2 , r2 , · · · , aH , sH } with s1 = ~0 , xh ∼ q ( . | sh ) , rh = R ( xh , ah ) , and sh+1 ∼ T ( . | sh , ah ) . The agent never observes { s1 , · · · , sH } . Technical Assumptions . We make two assumptions that are specific to the FactoRL algorithm . The first is a margin assumption on the transition dynamics that enables us to identify different values of a factor . This assumption was introduced by Du et al . ( 2019 ) , and we adapt it to our setting . Assumption 1 ( Margin Assumption ) . For every h ∈ { 2 , 3 , · · · , H } , i ∈ [ d ] , let ui be the uniform distribution jointly over actions and all possible reachable values of sh−1 [ pt ( i ) ] . Then we assume : ‖Pui ( · , · | sh [ i ] = 1 ) − Pui ( · , · | sh [ i ] = 0 ) ‖TV ≥ σ where Pui ( sh−1 [ pt ( i ) ] , a | sh [ i ] ) is the backward dynamics denoting the probability over parent values and last action given sh [ i ] and roll-in distribution ui , and σ > 0 is the margin . 1The notation supp ( p ) denotes the support of the distribution p. Formally , supp ( p ) = { z | p ( z ) > 0 } . Assumption 1 captures a large set of problems , including all deterministic problems for which the value of σ is 1 . Assumption 1 helps us identify the different values of a factor but it does not help with mapping atoms to the factors from which they are emitted . In order to identify if two atoms come from the same factor , we make the following additional assumption to measure their dependence . Assumption 2 ( Atom Dependency Bound ) . For any h ∈ [ H ] , u , v ∈ [ m ] and u 6= v , if ch−1 ( u ) = ch−1 ( v ) , i.e. , atoms xh [ u ] and xh [ v ] have the same factor . Then under any distribution D ∈ ∆ ( Sh ) we have ‖PD ( xh [ u ] , xh [ v ] ) − PD ( xh [ u ] ) PD ( xh [ v ] ) ‖TV ≥ βmin . Dependence assumption states that atoms emitted from the same factor will be correlated . This is true for many real-world problems . For example , consider a toy grid-based navigation task . Each state factor s [ i ] represents a cell in the grid which can be empty ( s [ i ] = 0 ) or occupied ( s [ i ] = 1 ) . In the latter case , a randomly sampled box from the set { red box , yellow box , black box } , occupies its place . We expect Assumption 2 to hold in this case as pixels emitted from the same factor come from the same object and hence will be correlated . More specifically , if one pixel is red in color , then another pixel from the same cell will also be red as the object occupying the cell is a red box . This assumption does not remove the key challenge in identifying factors . As atoms from different factors can still be dependent due to actions and state distributions from previous time steps . Model Class . We use two regressor classes F and G. The first regressor class F : X ×X → [ 0 , 1 ] takes a pair of atoms and outputs a scalar in [ 0 , 1 ] . To define the second class , we first define a decoder class Φ : X ∗ → { 0 , 1 } . We allow this class to be defined on any set of atoms . This is motivated by empirical research where commonly used neural network models operate on inputs of arbitrary lengths . For example , the LSTM model can operate on a text of arbitrary length ( Sundermeyer et al. , 2012 ) . However , this is without loss of generality as we can define a different model class for different numbers of atom . We also define a model class U : X × A × { 0 , 1 } → [ 0 , 1 ] . Finally , we define the regressor class G : X × A×X ∗ → [ 0 , 1 ] as { ( x , a , x̌ ) 7→ u ( x , a , φ ( x̌ ) ) | u ∈ U , φ ∈ Φ } . We assume F and G are finite classes and derive sample complexity guarantees which scale as log |F| and log |G| . However , since we only use uniform convergence arguments extending the guarantees to other statistical complexity measures such as Rademacher complexity is straightforward . Let Πall : S → A denote the set of all non-stationary policies of this form . We then define the class of policies Π : X → A by { x 7→ ϕ ( φ ? ( x ) ) | ∀ϕ ∈ Πall } , which we use later to define our task . We use Pπ [ E ] to denote probability of an event E under the distribution over episodes induced by policy π. Computational Oracle . We assume access to two regression oracles REG for model classes F and G. Let D1 be a dataset of triplets ( x [ u ] , x [ v ] , y ) where u , v denote two different atoms and y ∈ { 0 , 1 } . Similarly , let D2 be a dataset of quads ( x , a , x′ , y ) where x ∈ X , a ∈ A , x̌ ∈ X ∗ , and y ∈ { 0 , 1 } . Lastly , let ÊD [ · ] denote the empirical mean over dataset D. The two computational oracles compute : REG ( D1 , F ) =arg min f∈F ÊD1 [ ( f ( x [ u ] , x [ v ] ) − y ) 2 ] , REG ( D2 , G ) =arg min g∈GN ÊD2 [ ( g ( x , a , x̌ ) − y ) 2 ] . We also assume access to a ∆pl-optimal planning oracle planner . Let Ŝ = ( Ŝ1 , · · · , Ŝh ) be a learned state space and T̂ = ( T̂1 , · · · , T̂H ) with T̂h : Ŝh−1 ×A → ∆ ( Ŝh ) be the learned dynamics , and R̂ : Ŝ × A → [ 0 , 1 ] be a given reward function . Let ϕ : Ŝ → A be a policy and V ( ϕ ; T̂ , R̂ ) be the policy value . Then for any ∆pl > 0 the output of planner ϕ̂ = planner ( T̂ , R̂ , ∆pl ) satisfies V ( ϕ̂ ; T̂ , R̂ ) ≥ supϕ V ( ϕ ; T̂ , R̂ ) −∆pl , where supremum is taken over policies of type Ŝ → A . Task Definition . We focus on a reward-free setting with the goal of learning a state decoder and estimating the latent dynamics T . Since the state space is exponentially large , we can not visit every state . However , the factorization property allows us to estimate the model by reaching factor values . In fact , we show that controlling the value of at most 2κ factors is sufficient for learning the model . Let C≤k ( U ) denote the space of all sets containing at most k different elements selected from the set U including ∅ . We define the reachability probability ηh ( K , Z ) for a given h ∈ [ H ] , K ⊆ [ d ] , and Z ∈ { 0 , 1 } |K| , and the reachability parameter ηmin as : ηh ( K , Z ) : = sup π∈ΠNS Pπ ( sh [ K ] = Z ) , ηmin : = inf h∈ [ H ] inf s∈Sh inf K∈C≤2κ ( [ d ] ) ηh ( K , s [ K ] ) . Our sample complexity scales polynomially with η−1min . Note that we only require that if sh [ K ] = Z is reachable , then it is reachable with at least ηmin probability , i.e. , either ηh ( K , Z ) = 0 or it is at least ηmin . These requirements are similar to those made by earlier work for non-factored state space ( Du et al. , 2019 ; Misra et al. , 2020 ) . The key difference being that instead of requiring every state to be reachable with ηmin probability , we only require a small set of factor values to be reachable . For reference , if every policy induces a uniform distribution over S = { 0 , 1 } d , then probability of visiting any state is 2−d but the probability of two factors taking certain values is only 0.25 . This gives us a more practical value for ηmin . Besides estimating the dynamics and learning a decoder , we also learn an α-policy cover to enable exploration of different reachable values of factors . We define this below : Definition 3 ( Policy Cover ) . A set of policies Ψ is an α-policy cover of Sh for any α > 0 and h if : ∀s ∈ Sh , K ∈ C≤2κ ( [ d ] ) , sup π∈Ψ Pπ ( sh [ K ] = s [ K ] ) ≥ αηh ( K , s [ K ] ) .

Summary: The paper considers the problem of partitioning the atoms (e.g., pixels of an image) of a reinforcement learning task to latent states (e.g., a grid that determines whether there exists furniture in each cell). The number of states grows exponentially with the number of cells of the grid. So the algorithms that are polynomial in the number of states are not efficient. The paper considers the factored block Markov decision process (MDP) model and adds a few more assumptions. Generally, this model and assumptions guarantee that the cells of the grid partition the atoms (i.e., each atom depends on only one cell), the atoms in a cell are dependent (in the probabilistic sense), the conditional probability of the parent value of the states and the action given the next state is 0 or 1 is separated (i.e., the difference is bounded away from zero), and the regressor classes that are used are realizable. The paper shows that this is enough to give an algorithm that partitioned the atoms in each step with high probability and its time complexity is polynomial in the number of cells and logarithmic in the number of atoms.