PROVABLE RICH OBSERVATION REINFORCEMENT LEARNING WITH COMBINATORIAL LATENT STATES
Year: 2021

Abstract:
We propose a novel setting for reinforcement learning that combines two common real-world difficulties: presence of observations (such as camera images) and factored states (such as location of objects). In our setting, the agent receives observations generated stochastically from a latent factored state. These observations are rich enough to enable decoding of the latent state and remove partial observability concerns. Since the latent state is combinatorial, the size of state space is exponential in the number of latent factors. We create a learning algorithm FactoRL (Fact-o-Rel) for this setting which uses noise-contrastive learning to identify latent structures in emission processes and discover a factorized state space. We derive polynomial sample complexity guarantees for FactoRL which polynomially depend upon the number factors, and very weakly depend on the size of the observation space. We also provide a guarantee of polynomial time complexity when given access to an efficient planning algorithm.

heading: 1 INTRODUCTION
text: 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.

heading: To the best of our knowledge, our work represents the first provable solution to rich-observation RL with a combinatorially large latent state space.

heading: 2 THE FACTORED BLOCK MDP SETTING
text: 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 cannot 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]).

heading: 3 DISCOVERING EMISSION STRUCTURE WITH CONTRASTIVE LEARNING
text: Directly applying the prior work (Du et al., 2019; Misra et al., 2020) to decode a factored state from observation results in failure, as the learned factored state need not obey the transition factorization. Instead, the key high-level idea of our approach is to first learn the latent emission structure ch, and then use it to decode each factor individually. We next discuss our approach for learning ch.
Reducing Identification of Latent Emission Structure to Independence Tests. Assume we are able to perfectly decode the latent state and estimate the transition model till time step h− 1. Our goal is to infer the latent emission structure chh at time step h, which is equivalent to: given an arbitrary pair of atoms u and v, determine if they are emitted from the same factor or not. This is challenging since we cannot observe or control the latent state factors at time step h.
Let i = ch−1(u) and j = ch−1(v) be the factors that emit x[u] and x[v]. If i = j, then Assumption 2 implies that these atoms are dependent on each other for any roll-in distribution D ∈ ∆(Sh−1 ×A) over previous state and action. However, if i 6= j then deterministically setting the previous action and values of the parent factors pt(i) or pt(j), makes x[u] and x[v] independent. For the example in Figure 1, fixing the value of s[1], s[2] and a would make x[u] and x[u′] independent of each other.
This observation motivates us to reduce this identification problem to performing independence tests with different roll-in distributions D ∈ ∆(Sh−1 ×A). Naively, we can iterate over all subsets K ∈ C≤2κ([d]) where for each K we create a roll-in distribution such that the values of sh−1[K] and the action ah−1 are fixed, and then perform independence test under this distribution. If two atoms are independent then there must exist a K that makes them independent. Otherwise, they should always be dependent by Assumption 2.
However, there are two problems with this approach. Firstly, we do not have access to the latent states but only a decoder at time step h− 1. Further, it may not even be possible to find a policy that can set the values of factors deterministically. We later show that our algorithm FactoRL can learn a decoder that induces a bijection between learned factors and values, and the real factors and values. Therefore, maximizing the probability of ÊK;Z = {φ̂h−1(xh−1)[K] = Z} for a set of learned factors K and their values Z , implicitly maximizes the probability of EK′;Z′ = {sh−1[K′] = Z ′} for corresponding real factors K′ and their values Z ′. Since the event ÊK;Z is observable we can use rejection sampling to increase its probability sufficiently close to 1 which makes the probability of EK′;Z′ close to 1. The second problem is to perform independence tests in a domain agnostic setting. Directly estimating mutual information I(x[u];x[v]) can be challenging. Instead, we propose an oraclized independence test that reduces the problem to binary classification using noise-contrastive learning.
Oraclized Independent Test. Here, we briefly sketch the main idea of our independence test scheme and defer the details to Appendix C. We comment that the high-level idea of our independence testing subroutine is similar to Sen et al. (2017). Suppose we want to test if two random variables Y and Z are independent. Firstly, we construct a dataset in the following way: sample a Bernoulli random variable w ∼ Bern(1/2), and two pairs of independent realizations (y(1), z(1)) and (y(2), z(2)); if w = 1, add (y(1), z(1), w) to the dataset, and (y(1), z(2), w) otherwise. We repeat the sampling procedure n times and obtain a dataset {(yi, zi, wi)}ni=1. Then we can fit a classifier that predicts the value of wi using (yi, zi). If Y and Z are independent, then (yi, zi) will provide no information about wi and thus no classifier can do better than random guess. However, if Y and Z are dependent, then
the Bayes optimal classifier would perform strictly better than random guess. As a result, by looking at the training loss of the learned classifier, we can determine whether Y and Z are dependent or not.
4 FactoRL: REINFORCEMENT LEARNING IN FACTORED BLOCK MDPS
In this section, we present the main algorithm FactoRL (Algorithm 1). It takes as input the model classes F ,G, failure probability δ > 0, and five hyperparameters σ, ηmin, βmin ∈ (0, 1) and d, κ ∈ N.2 We use these hyperparamters to define three sample sizes nind, nabs, nest and rejection sample frequency k. For brevity, we defer the exact values of these constants to Appendix D.7. FactoRL returns a learned decoder φ̂h : X → {0, 1}dh for some dh ∈ [m], an estimated transition model T̂h, learned parent p̂th and child functions ĉhh, and a 1/2-policy cover Ψh of Sh for every time step h ∈ {2, 3, · · · , H}. We use ŝh to denote the learned state at time step h. Formally, ŝh = (φ̂h1(xh), · · · , φ̂hdh(xh)). In the analysis of FactoRL, we show that dh−1 = d, and ĉhh is equivalent to chh up to permutation with high probability. Further, we show that φ̂h and ĉhh together learn a bijection between learned factors and their values and real factors and their values.
FactoRL operates inductively over the time steps (Algorithm 1, line 2-8). In the hth iteration, the algorithm performs four stages of learning: identifying the latent emission structure, decoding the factors, estimating the model, and learning a policy cover. We describe these below.
Algorithm 1 FactoRL(F ,G, δ, σ, ηmin, βmin, d, κ). RL in Factored Block MDPs. 1: Initialize Ψh = ∅ for every h ∈ [H] and φ̂1 = X → {0}. Set global constants nind, nabs, nest, k. 2: for h ∈ {2, 3, · · · , H} do 3: ĉhh = FactorizeEmission(Ψh−1, φ̂h−1,F) // stage 1: discover latent emission structure 4: φ̂h = LearnDecoder(G,Ψh−1, ĉhh) // stage 2: learn a decoder for factors 5: T̂h, p̂th = EstModel(Ψh−1, φ̂h−1, φ̂h) // stage 3: find latent pth and estimate model 6: for I ∈ C≤2κ([d]),Z ∈ {0, 1}|I| do 7: ϕ̂hIZ = planner(T̂ , RhIZ ,∆pl) where RhIZ := 1{ŝh[I] = Z} // stage 4: planning 8: If V (ϕ̂hIZ ; T̂ , RhIZ) ≥ 3ηmin/4 then Ψh ← Ψh ∪ {ϕ̂hIZ ◦ φ̂h}
return {ĉhh, φ̂h, T̂h, p̂th,Ψh}Hh=2
Identifying Latent Emission Process. The FactorizeEmission collects a dataset of observations for every policy in Ψh−1 and action a ∈ A (Algorithm 2, line 1-4). Policies in Ψh−1 are of the type πI;Z where I ∈ C≤2κ([dh−1]) and Z ∈ {0, 1}|I|. We can inductively assume πI;Z to be maximizing the probability of EI;Z = {ŝh−1[I] = Z}. If our decoder is accurate enough, then we hope that maximizing the probability of this event in turn maximizes the probability of fixing the values of a set of real factors. However, it is possible that PπI;Z (ŝh−1[I] = Z) is only O(ηmin). Therefore, as explained earlier, we use rejection sampling to drive the probability of this event close to 1. Formally, we define a procedure RejectSamp(πI;Z , EI;Z , k) which rolls-in at time step h− 1 with πI;Z to observe xh−1 (line 3). If the event EI;Z holds for xh−1 then we return xh−1, otherwise, we repeat the procedure. If we fail to satisfy the event k times then we return the last sample. We use this to define our main sampling procedure xh ∼ DI,Z,a := RejectSamp(πI;Z , EI;Z , k) ◦ a which first samples xh−1 using the rejection sampling procedure and then takes action a to observe xh. We collect a dataset of observation pairs (x(1), x(2)) sampled independently from DI,Z,a.
For every pair of atoms u, v ∈ [m], we calculate if they are independent under the distribution induced by DI,Z,a using IndTest with dataset DI,Z,a (line 5-7). We share the dataset across atoms for sample efficiency. If there exists at least one (I,Z, a) triple such that we evaluate x[u], x[v] to be independent, then we mark these atoms as coming from different factors. Intuitively, such an I would contain parent factors of at least ch−1h (u) or ch −1 h (v). If no such I exists then we mark these atoms as being emitted from the same factor.
2Our analysis can use any non-zero lower bound on ηmin, βmin, σ and an upper bound on d and κ.
Algorithm 2 FactorizeEmission(Ψh−1, φ̂h−1,F). 1: for (πI;Z , a) ∈ Ψh−1 ×A and i ∈ [nind] do 2: Define EI;Z := 1{φ̂h−1(xh−1)[I] = Z} 3: Sample x(1)h , x (2) h ∼ RejectSamp(πI;Z , EI;Z , k) ◦ a // rejection sampling procedure
4: DI;Z;a ← DI;Z;a ∪ {(x(1)h , x (2) h )} // initialize DI;Z;a = ∅
5: for u ∈ {1, 2, · · · ,m− 1} and v ∈ {u+ 1, · · · ,m} do 6: Mark u, v as coming from the same factor, i.e., ĉh −1 h (u) = ĉh −1 h (v) if ∀(I,Z, a) 7: the oraclized independence test finds xh[u], xh[v] as dependent using DI;Z;a and F return ĉhh // label ordering of parents does not matter.
Algorithm 3 LearnDecoder(G,Ψh−1, ĉhh). Child function has type ĉhh : [dh]→ 2[m]
1: for i in [dh], define ω = ĉhh(i),D = ∅ do 2: for nabs times do // collect a dataset of real (y = 1) and imposter (y = 0) transitions 3: Sample (x(1), a(1), x′(1)), (x(2), a(2), x′(2)) ∼ Unf(Ψh−1) ◦ Unf(A) and y ∼ Bern( 12 ) 4: If y = 1 then D ← D ∪ (x(1), a(1), x′(1)[ω], y) else D ← D ∪ (x(1), a(1), x′(2)[ω], y) 5: ûi, φ̂i = REG(D,G) // train the decoder using noise-contrastive learning
return φ̂ : X → {0, 1}dh where for any x ∈ X and i ∈ [dh] we have φ̂(x)[i] = φ̂i(x[ĉhh(i)]).
Decoding Factors. LearnDecoder partitions the set of atoms into groups based on the learned child function ĉhh (Algorithm 3). For the ith group ω, we learn a decoder φ̂hi : X ? → {0, 1} by adapting the prediction problem of Misra et al. (2020) to Factored Block MDP setting. We define a sampling procedure (x, a, x′) ∼ Unf(Ψh−1) ◦ Unf(A) where x is observed after roll-in with a uniformly selected policy in Ψh−1 till time step h− 1, action a is taken uniformly, and x′ ∼ T (· | x, a) (line 3). We collect a dataset D of real and imposter transitions. A single datapoint in D is collected by sampling two independent transitions (x(1), a(1), x′(1)), (x(2), a(2), x′(2)) ∼ Unf(Ψh−1) ◦ Unf(A) and a Bernoulli random variable y ∼ Bern(1/2). If y = 1 then we add the real transition (x(1), a(1), x′(1)[ω], y) to D, otherwise we add the imposter transition (x(1), a(1), x′(2)[ω], y) (line 4). The key difference from Misra et al. (2020) is our use x′[ω] instead of x′ which allows us to decode a specific latent factor. We train a model to predict the probability that a given transition (x, a, x′[ω]) is real by solving a regression task with model class G (line 5). The bottleneck structure of G allows us to recover a decoder φ̂i from the learned model. The algorithm also checks for the special case where a factor takes a single value. If it does, then we return the decoder that always outputs 0, otherwise we stick with φ̂i. For brevity, we defer the details of this special case to Appendix D.2.2. The decoder for the hth timestep is given by composition of decoders for each group.
Algorithm 4 EstModel(Ψh−1, φ̂h−1, φ̂h). 1: Collect dataset D of nest triplets (x, a, x′) ∼ Unf(Ψh−1) ◦ Unf(A) 2: for I,J ∈ C≤κ([dh−1]) satisfying I ∩ J = ∅ do 3: Estimate P̂(ŝh[k] | ŝh−1[I], ŝh−1[J ], a) from D using φ̂, ∀a ∈ A, k ∈ [dh]. 4: For every k define p̂th(k) as solution to following: (where we bind ŝ ′ = ŝh and ŝ = ŝh−1)
argmin I max u,J1,J2,w1,w2,a ∣∣∣P̂(ŝ′[k] | ŝ[I] = u, ŝ[J1] = w1, a)− P̂(ŝ′[k] | ŝ[I] = u, ŝ[J2] = w2, a) ∣∣∣ TV
5: Define T̂h(ŝ′ | ŝ, a) = ∏ k P̂(ŝ[k] | ŝ[p̂th(k)], a) and return T̂h, p̂th
Estimating the Model. EstModel routine first collects a dataset D of nest independent transitions (x, a, x′) ∼ Unf(Ψh−1) ◦ Unf(A) (Algorithm 4, line 1). We iterate over two disjoint sets of factors I,J of size at most κ. We can view I as the control set and J as the variable set. For every
learned factor k ∈ [dh], factor set I,J and action a ∈ A, we estimate the model P̂(ŝh[k] | ŝh−1[I], ŝh−1[J ], a) using count based statistics on dataset D (line 3).
Consider the case where the ĉht = cht for every t ∈ [h] and where we ignore the label permutation for brevity. If I contains the parent factors pt(k), then we expect the value of P̂(ŝ′[k] | ŝ[I], ŝ[J ], a) ≈ Tk(ŝ
′[k] | ŝ[pt(k)], a) to not change significantly on varying either the set J or its values. This motivates us to define the learned parent set as the I which achieves the minimum value of this gap (line 4). When computing the gap, we take max only over those values of ŝ[I] and ŝ[J ] which can be reached jointly using a policy in Ψh−1. This is important since we can only reliably estimate the model for reachable factor values. The learned parent function p̂th need not be identical to pth even upto relabeling. However, finding the exact parent factor is not necessary for learning an accurate model, and may even be impossible. For example, two factors may always take the same value making it impossible to distinguish between them. We use the learned parent function p̂th to define T̂h similar to the structure of T (line 5).
Learning a Policy Cover. We plan in the latent space using the estimated model {T̂t}ht=1, to find a policy cover for time step h. Formally, for every I ∈ C≤2κ([dh]) and Z ∈ {0, 1}|I|, we find a policy ϕ̂hIZ to reach {ŝh[I] = Z} using the planner (Algorithm 1, line 7). This policy acts on the learned state space and is easily lifted to act on observations by composition with the learned decoder. We add every policy that achieves a return of at least O(ηmin) to Ψh (line 8).

heading: 5 THEORETICAL ANALYSIS AND DISCUSSION
text: In this section, we present theoretical guarantees for FactoRL. For technical reasons, we make the following realizability assumption on the function classes F and G. Assumption 3 (Realizability). For any h ∈ [H], i ∈ [d] and distribution ρ ∈ ∆({0, 1}), there exists gihρ ∈ G, such that for all ∀(x, a, x′) ∈ Xh−1 ×A×Xh and x̌ = x′[chh(i)] we have:
gihρ(x, a, x̌) = Ti(φ
? i (x̌) | φ?(x), a)
Ti(φ?i (x̌) | φ?(x), a) + ρ(φ?i (x̌)) .
For any h ∈ [H], u, v ∈ [m] with u 6= v, and any D ∈ ∆(Sh), there exists fuvD ∈ F satisfying:
∀s ∈ supp(D), x ∈ supp(q(· | s)), fuvD(x[u], x[v]) = D(x[u], x[v])
D(x[u], x[v]) +D(x[u])D(x[v]) .
Assumption 3 requires the function classes to be expressive enough to represent optimal solutions for our regression tasks. Realizability assumptions are common in literature and are in practice satisfied by using deep neural networks (Sen et al., 2017; Misra et al., 2020).
Theorem 1 (Main Theorem). For any δ > 0, FactoRL returns a transition function T̂h, a parent function ĉhh, a decoder φ̂h, and a set of policies Ψh for every h ∈ {2, 3, · · · , H}, that with probability at least 1− δ satisfies: (i) ĉhh is equal to chh upto permutation, (ii) Ψh is a 1/2 policy cover of Sh, (iii) For every h ∈ [H], there exists a permutation mapping θh : {0, 1}d → {0, 1}d such that for every s ∈ Sh−1, a ∈ A, s′ ∈ Sh and x′ ∈ Xh we have:
P(φ̂h(x′)= θ−1h (s ′) | s′) ≥ 1−O ( η2min κH ) , ∥∥∥T (s′ | s, a)−T̂h(θ−1h (s′) | θ−1h−1(s), a) ∥∥∥ TV ≤ ηmin 8H ,
and the sample complexity is poly(d16κ, |A|, H, 1ηmin , 1 δ , 1 βmin , 1σ , lnm, ln |F|, ln |G|).
Discussion. The proof and the detailed analysis of Theorem 1 has been deferred to Appendix C-D. Our guarantees show that FactoRL is able to discover the latent emission structure, learn a decoder, estimate the model, and learn a policy cover for every timestep. We set the hyperparmeters in order to learn a 1/2-policy cover, however, they can also be set to achieve a desired accuracy for the decoder or the transition model. This will give a polynomial sample complexity that depends on this desired accuracy. It is straightforward to plan a near-optimal policy for a given reward function in our learned latent space, using the estimated model and the learned decoder. This incurs zero sample cost apart from the samples needed to learn the reward function.
Our results show that we depend polynomially on the number of factors and only logarithmic in the number of atoms. This appeals to real-world problem where d and m can be quite large. We also depend logarithmically on the size of function classes. This allows us to use exponentially large function classes, further, as stated before, our results can also be easily extended to Rademacher complexity. Our algorithm only makes a polynomial number of calls to computational oracles. Hence, if these oracles can be implemented efficiently then our algorithm has a polynomial computational complexity. The squared loss oracles are routinely used in practice, but planning in a fully-observed factored MDP is EXPTIME-complete (see Theorem 2.24 of Mausam (2012)). However, various approximation strategies based on linear programming and dynamic programming have been employed succesfully (Guestrin et al., 2003). These assumptions provide a black-box mechanism to leverage such efforts. Note that all computational oracles incur no additional sample cost and can be simply implemented by enumeration over the search space. Comparison with Block MDP Algorithms. Our work is closely related to algorithms for Block MDPs, which can be viewed as a non-factorized version of our setting. Du et al. (2019) proposed a model-based approach for Block MDPs. They learn a decoder for a given time step by training a classifier to predict the decoded state and action at the last time step. In our case, this results in a classification problem over exponentially many classes which can be practically undesirable. In contrast, Misra et al. (2020) proposed a model-free approach that learns a decoder by training a classifier to distinguish between real and imposter transitions. As optimal policies for factored MDPs do not factorize, therefore, a model-free approach is unlikely to succeed (Sun et al., 2019). Feng et al. (2020) proposed another approach for solving Block MDPs. They assume access to a purely unsupervised learning oracle, that can learn an accurate decoder using a set of observations. This oracle assumption is significantly stronger than those made in Du et al. (2019) and Misra et al. (2020), and reduces the challenge of learning the decoder. Crucially, these three approaches have a sample complexity guarantee which depends polynomially on the size of latent state space. This yields an exponential dependence on d when applied to our setting. It is unclear if these approaches can be extended to give polynomial dependence on d. For general discussion of related work see Appendix B.
Proof of Concept Experiments. We empirically evaluate FactoRL to support our theoretical results, and to provide implementation details. We consider a problem with d factors each emitting 2 atoms. We generate atoms for factor s[i], by first defining a vector zi = [1, 0] if s[i] = 0, and zi = [0, 1] otherwise. We then sample a scalar Gaussian noise gi with 0 mean and 0.1 standard deviation, and add it to both component of zi. Atoms emitted from each factor are concatenated to generate a vector z ∈ R2d. The observation x is generated by applying a fixed time-dependent permutation to z to shuffle atoms from different factors. This ensures that an algorithm cannot figure out the children function by relying on the order in which atoms are presented. We consider an action space A = {a1, a2, · · · , ad} and non-stationary dynamics. For each time step t ∈ [H], we define σt as a fixed permutation of {1, 2, · · · , d}. Dynamics at time step t are given by: Tt(st+1 | st, a) = ∏d i=1 Tti(st+1[i] | st[i], a), where Tti(st+1[i] = st[i] | st[i], a) = 1 for all a 6= aσt(i), and Tti(st+1[i] = 1− st[i] | st[i], aσt(i)) = 1. We evaluate on the setting d = 10 and H = 10. We implement model classes F and G using feed-forward neural networks. Specifically, for G we apply the Gumbel-softmax trick to model the bottleneck following Misra et al. (2020). We train the models using cross-entropy loss instead of squared loss that we use for theoretical analysis.3 For the independence test task, we declare two atoms to be independent, if the best log-loss on the validation set is greater than c. We train the model using Adam optimization and perform model selection using a held-out set. We defer the full model and training details to Appendix F.
For each time step, we collect 20,000 samples and share them across all routines. This gives a sample complexity of 20, 000×H . We repeat the experiment 3 times and found that each time, the model was able to perfectly detect the latent child function, learn a 1/2-policy cover, and estimate the model with error < 0.01. This is in sync with our theoretical findings and demonstrates the empirical use of FactoRL. We will make the code available at: https://github.com/cereb-rl. Conclusion. We introduce Factored Block MDPs that model the real-world difficulties of richobservation environments with exponentially large latent state spaces. We also propose a provable RL algorithm called FactoRL for solving a large class of Factored Block MDPs. We hope the setting and ideas in FactoRL will stimulate both theoretical and empirical work in this important area.
3We can also easily modify our proof to use cross-entropy loss by using generalization bounds for log-loss (see Appendix E in Agarwal et al. (2020))

heading: APPENDIX ORGANIZATION
text: This appendix is organized as follows.
• Appendix A provides a list of notations used in this paper. • Appendix B covers related work • Appendix C describes the independence test algorithm and its sample complexity guarantees • Appendix D provides sample complexity guarantees for FactoRL • Appendix E provides list of supporting results • Appendix F provides details of the experimental setup and optimization

heading: A NOTATIONS
text: We present notations and their definition in Table 1. In general, calligraphic notations represent sets. All logarithms are with respect to base e.

heading: B RELATED WORK
text: There is a rich literature on sample-efficient reinforcement learning in tabular MDPs with a small number of observed states (Brafman & Tennenholtz, 2002; Strehl et al., 2006; Kearns & Singh, 2002; Jaksch et al., 2010; Azar et al., 2017; Jin et al., 2018). While recent state-of-the-art results along this line achieve near-optimal sample complexity, these algorithms do not exploit the latent structure in the environment, and therefore, cannot scale to many practical settings such as rich-observation environments with possibly a large number of factored latent states.
In order to overcome this challenge, one line of research has been focusing on factored MDPs (Kearns & Koller, 1999; Guestrin et al., 2002; 2003; Strehl et al., 2010), which allow a combinatorial number of observed states with a factorized structure. Planning in factored MDPs is EXPTIME-complete (Mausam, 2012) yet often tractable in practice, with factored MDPs statistically learnable with polynomial samples in the number of parameters that encode the factored MDP (Osband & Van Roy, 2014; Li et al., 2011). There has also been several empirical works that either focus on the factored state space setting (e.g., Kim & Mnih (2018); Thomas et al. (2018); Laversanne-Finot et al. (2018); Miladinović et al. (2019)), or the factored action space setting (e.g., He et al. (2016); Sharma et al. (2017)). However, these works do not directly address our problem and do not provide sample complexity guarantees.
Another line of work focuses on exploration in a rich observation environment. Empirical results (Tang et al., 2017; Chen et al., 2017; Bellemare et al., 2016; Pathak et al., 2017) have achieved inspiring performance on several RL benchmarks, while theoretical works (Krishnamurthy et al., 2016; Jiang et al., 2017) show that it is information-theoretically possible to explore these environments. As discussed before, recent works of Du et al. (2019), Misra et al. (2020) and Feng et al. (2020) provide computationally and sample efficient algorithms for Block MDP which is a rich-observation setting with a latent non-factored state space. Nevertheless, this line of results crucially relies on the number of latent states being relatively small.
Finally, we comment that the contrastive learning technique used in this paper has been used by other reinforcement learning algorithms for learning feature representation (e.g., Kim et al. (2019); Nachum et al. (2019); Srinivas et al. (2020)) without theoretical guarantee.
C INDEPENDENCE TESTING USING NOISE CONTRASTIVE ESTIMATION
In this section, we introduce the independence testing algorithm, Algorithm 5 and provide its theoretic guarantees. Algorithm 5 will be used in Algorithm 2 as a subroutine for determining if two atoms are emitted from the same latent factor. We comment that the high-level idea of Algorithm 5 is similar to Sen et al. (2017), which reduces independence testing to regression by adding imposter samples.

heading: C.1 ALGORITHM DESCRIPTION
text: Let D ∈ ∆(Sh−1 × A) be our roll-in distribution that induces a probability distribution PD ∈ ∆(Xh) over observations at time step h. Let u, v ∈ [m] be two different atoms, and PD(x[u], x[v]),PD(xh[u]) and PD(xh[v]) be the joint and marginal distributions over xh[u], xh[v] with respect to roll-in distribution D. The goal of our algorithm to determine if xh[u] and xh[v] are independent under PD.
Algorithm 5 IndTest(F ,D, u, v, β) Oraclized Independency Test. We initialize Dtrain = ∅. 1: Initialize Dtrain = ∅ and sample z1, z2, · · · , zn ∼ Bern( 12 ). 2: for i ∈ [n] do 3: if zi = 1 then 4: Dtrain ← Dtrain ∪ {(x(i,1)[u], x(i,1)[v], 1)}. 5: else 6: Dtrain ← Dtrain ∪ {(x(i,1)[u], x(i,2)[v], 0)}. 7: Compute f̂ := arg minf∈F L(Dtrain, f), where
L(Dtrain, f) := 1
n
∑
(x[u],x[v],z)∈Dtrain {f(x[u], x[v])− z}2 .
8: return Independent if L(Dtrain, f̂) > 0.25− β2/103 else return Dependent.
We solve this task using the IndTest algorithm (Algorithm 5) which takes as input a dataset of observation pairs D = {(x(i,1), x(i,2))}ni=1 where x(i,1), x(i,2) ∼ PD(·, ·), and a scalar β ∈ (0, 1). We use D to create a dataset Dtrain of real and imposter atom pairs (x[u], x[v]). This is done by taking every datapoint in D and sampling a Bernoulli random variable zi ∼ Bern(1/2) (line 1). If zi = 1 then we add the real pair (x(i,1)[u], x(i,1)[v], 1) to Dtrain (line 4), otherwise we add the imposter pair (x(i,1)[u], x(i,2)[v], 0) (line 6). We train a classifier to predict if a given atom pair (x[u], x[v]) is real or an imposter (line 7). The Bayes optimal classifier for this problem is given by:
∀x ∈ supp(PD), f?D(x[u], x[v]) := PD(z = 1 | x[u], x[v]) = PD(x[u], x[v])
PD(x[u], x[v]) + PD(x[u])P(x[v]) .
If x[u] and x[v] are independent then we have PD(x[u], x[v]) = PD(x[u])PD(x[v]) everywhere on the support of PD. This implies f?D(x) = 12 and its training loss will concentrate around 0.25. Intuitively, this can be interpreted as the classifier having no information to tell real samples from imposter samples. However, if x[u], x[v] are dependent and ‖PD(x[u], x[v])− PD(x[u])PD(x[v])‖TV ≥ β then we can show the training loss of f? is less than 0.25 − O(β2) with high probability. The remainder of this section is devoted to a rigorous proof for this argument.

heading: C.2 ANALYSIS OF ALGORITHM 5
text: Before analyzing Algorithm 5, we want to slightly simplify the problem in terms of notations. We introduce two simple notations X and Y which represents the random variables x[u] and x[v], respectively. We will simply use D to denote the joint distribution of X and Y . Define Dtrain to be the distribution of the training data (Xtrain, Ytrain, z) produced in Algorithm 5. It’s easy to verify that
Dtrain(Xtrain, Ytrain, z) = 1
2 [zD(Xtrain, Ytrain) + (1− z)D(Xtrain)D(Ytrain)] . (1)
Suppose the distribution D is specially designed such that at least one of the following hypothesis holds (which can be guaranteed when we invoke Algorithm 5):
H0 : ‖D(X,Y )−D(X)D(Y )‖1 ≥ β
2
v.s. H1 : ‖D(X,Y )−D(X)D(Y )‖1 ≤ β2
103 .
In the remaining part, we will prove that Algorithm 5 can correctly distinguish between H0 and H1 with high probability.

heading: C.2.1 TWO PROPERTIES OF THE BAYES OPTIMAL CLASSIFIER
text: Our first lemma shows that the Bayes optimal classifier for the optimization problem in line 7 is a constant function equal to 1/2 if X and Y are independent.
Lemma 1 (Bayes Optimal Classifier for Independent Case). In Algorithm 5, if X and Y (atoms u and v) are independent under distribution D, then for the optimization problem in line 7 , the Bayes optimal classifier is given by:
∀(Xtrain, Ytrain) f?(Xtrain, Ytrain) = 1
2 .
Proof. From Bayes rule we have:
f?(Xtrain, Ytrain)
=Dtrain(z = 1 | Xtrain, Ytrain)
= Dtrain(Xtrain, Ytrain | z = 1)Dtrain(z = 1)
Dtrain(Xtrain, Ytrain | z = 1)Dtrain(z = 1) +Dtrain(Xtrain, Ytrain | z = 0)Dtrain(z = 0)
= Dtrain(Xtrain, Ytrain | z = 1)
Dtrain(Xtrain, Ytrain | z = 1) +Dtrain(Xtrain, Ytrain | z = 0) ,
where the last identity uses Dtrain(z = 1) = Dtrain(z = 0) = 1/2.
When z = 0, we collect Xtrain and Ytrain from two independent samples. Therefore, we have Dtrain(Xtrain, Ytrain | z = 0) = D(Xtrain)D(Ytrain). When z = 1, using the fact that Xtrain and Ytrain are independent under distribution D, we also have Dtrain(Xtrain, Ytrain | z = 1) = D(Xtrain, Ytrain) = D(Xtrain)D(Ytrain). Consequently, f?(Xtrain, Ytrain) ≡ 1/2.
Our second lemma provides an upper bound for the expected training loss of the Bayes Optimal Classifier. Later we will use this lemma to show the training loss is less than 0.25−O(β2) with high probability when H0 holds.
Lemma 2 (Square Loss of the Bayes Optimal Classifier). In Algorithm 5 line 7, the Bayes optimal classifier has expected square loss
EDtrainL(f?,Dtrain) ≤ 1 4 − ( 1 2 ED [ ∣∣∣∣ 1 2 − D(X,Y ) D(X)D(Y ) +D(X,Y ) ∣∣∣∣ ])2 ,
Proof. Recall the formula of the Bayes optimal classifier in Lemma 1,
f?(X,Y ) = D(X,Y )
D(X)D(Y ) +D(X,Y ) .
Plugging it into the square loss, we obtain
EDtrain [ (f?(X,Y )− y)2 ]
=EDtrain [ f?(X,Y ) (f?(X,Y )− 1)2 + (1− f?(X,Y )) (f?(X,Y ))2 ]
=EDtrain [f?(X,Y ) (1− f?(X,Y ))]
= 1
4 − EDtrain
[( 1
2 − D(X,Y ) D(X)D(Y ) +D(X,Y )
)2 ]
≤1 4 − ( EDtrain [ ∣∣∣∣ 1 2 − D(X,Y ) D(X)D(Y ) +D(X,Y ) ∣∣∣∣ ])2 ≤1 4 − ( 1 2 ED [ ∣∣∣∣ 1 2 − D(X,Y ) D(X)D(Y ) +D(X,Y ) ∣∣∣∣ ])2 .

heading: C.2.2 THREE USEFUL LEMMAS
text: To proceed, we need to take a detour and prove three useful technical lemmas.
Lemma 3. Let µ and ν be two probability measures defined on a countable set X . If ‖µ− ν‖TV ≥ c, then
Ex∼µ ∣∣∣∣
µ(x) µ(x) + ν(x) − 1 2 ∣∣∣∣ ≥ c 4 .

heading: Proof.
text: Ex∼µ [∣∣∣∣
µ(x) µ(x) + ν(x) − 1 2 ∣∣∣∣ ] = Ex∼µ [ 1 2 ∣∣∣∣ µ(x)− ν(x) µ(x) + ν(x) ∣∣∣∣ ]
≥ 1 2 Ex∼µ
[ 1{µ(x) > ν(x)} ∣∣∣∣ µ(x)− ν(x) µ(x) + ν(x) ∣∣∣∣ ]
= 1
2
∑ x∈X µ(x)
[ 1{µ(x) > ν(x)}µ(x)− ν(x)
µ(x) + ν(x)
]
≥ 1 4
∑ x∈X [1{µ(x) > ν(x)} (µ(x)− ν(x))]
≥ c 4 .
Lemma 4. Fix δ ∈ (0, 1). Then with probability at least 1− δ, we have ∣∣∣L(f̂ ,Dtrain)− EDtrainL(f?,Dtrain) ∣∣∣ ≤ 10 √ C(F , δ)
n ,
where Dtrain is the training set consisting of n i.i.d. samples sampled from Dtrain, f̂ is the empirical minimizer of L(f,Dtrain) over F , f? is the population minimizer, and C(F , δ) := ln |F|δ is the complexity measure of function class F .
Proof. By Hoeffding’s inequality and union bound, with probability at least 1− δ, for every f ∈ F , we have
|L(f,Dtrain)− EDtrainL(f,Dtrain)| ≤ 10 √ C(F , δ)
n .
Because f̂ is the empirical optimizer,
L(f̂ ,Dtrain) ≤ L(f?,Dtrain) ≤ EDtrainL(f?,Dtrain) + 10 √ C(F , δ)
n .
Because f? is the population optimizer,
L(f̂ ,Dtrain) ≥ EDtrainL(f̂ ,Dtrain)− √ C(F , δ)
n ≥ EDtrainL(f?,Dtrain)− 10
√ C(F , δ)
n .
Combining the two inequalities above, we finish the proof.
For notational convenience, we introduce the following factor distribution Dfactor defined on the same domain of (X,Y, z):
Dfactor(X,Y, z) = 1
2 D(X)D(Y ).
Lemma 5. Suppose F contains the constant function f ≡ 1/2. Then with probability at least 1− δ, we have ∣∣∣∣L(f̂ ,Dtrain)− 1
4
∣∣∣∣ ≤ 10 √ C(F , δ)
n + 2‖Dtrain −Dfactor‖TV.
Proof. By Lemma 4, with probability at least 1− δ, we have ∣∣∣L(f̂ ,Dtrain)− EDtrainL(f?,Dtrain) ∣∣∣ ≤ 10 √ C(F , δ)
n . (2)
Noticing that L is bounded by 1, we have for every f ∈ F |EDfactorL(f,Dtrain)− EDtrainL(f,Dtrain)| ≤ 2‖Dtrain −Dfactor‖TV, (3)
where EDfactorL(f,Dtrain) defines the expected loss of f over Dtrain where Dtrain consists of samples i.i.d. sampled from Dfactor.
Since y is a symmetric Bernoulli r.v. independent of (X,Y ) under distribution Dfactor, we have
min f∈F
EDfactorL(f,Dtrain) = 1
4 . (4)
Using the inequality |minf L1(f) − minf L2(f)| ≤ maxf |L1(f) − L2(f)| for any functionals L1, L2, along with (4) and (3) we bound the minimum loss under distribution Dtrain as:
∣∣∣∣minf∈F EDtrainL(f,Dtrain)− 1 4 ∣∣∣∣ ≤ 2‖Dtrain −Dfactor‖TV. (5)
Combining (5) with (2) completes the proof.

heading: C.2.3 MAIN THEOREM FOR ALGORITHM 5
text: Finally, we are ready to state and prove the main theorem for Algorithm 5.
Theorem 2. Under the realizability assumption and n ≥ Ω(C(F,δ)β4 ), Algorithm 5 can distinguish
H0 : ‖D(X,Y )−D(X)D(Y )‖1 ≥ β
2
v.s. H1 : ‖D(X,Y )−D(X)D(Y )‖1 ≤ β2
103
correctly with probability at least 1− δ.
Proof. If H1 holds, by Lemma 5, we have the following lower bound for the training loss of the empirical minimizer,
L(f̂ ,D) ≥ 1 4 − 10
√ C(F , δ)
n − β
2
103 . (6)
In contrast, if H0 is true, applying Lemma 4, we obtain
L(f̂ ,D) ≤ EDtrainL(f?,D) + 10 √ C(F , δ)
n .
Invoke Lemma 2 and Lemma 3,
EDtrainL(f?,Dtrain) ≤ 1 4 − ( 1 2 ED [ ∣∣∣∣ 1 2 − D(X,Y ) D(X)D(Y ) +D(X,Y ) ∣∣∣∣ ])2
≤ 1 4 − β
2
256 .
Therefore, under H0, the training loss of the empirical minimizer is upper bounded as below
L(f̂ ,D) ≤ 1 4 − β
2
256 + 10
√ C(F , δ)
n . (7)
Plugging n ≥ O(C(F,δ)β4 ) into (6) and (7), we complete the proof.
D THEORETICAL ANALYSIS OF FactoRL
In this section, we provide a detailed theoretical analysis of FactoRL. The structure of the algorithm is iterative making an inductive case argument appealing. We will, therefore, make an induction hypothesis for each time step that we will guarantee at the end of the time step.
Induction Hypothesis. We make the following induction assumption for FactoRL under Assumption 1-3 and across all time steps. For all t ∈ {2, 3, · · · , H}, at the end of time step t (Algorithm 1, line 8), FactoRL finds a child function ĉht : [d] → 2[m], a decoder φ̂t : X → {0, 1}d, a transition function T̂t : {0, 1}d ×A → {0, 1}d, and a set of policies Ψt satisfying the following:
IH.1 ĉht : [d] → 2m and cht : [d] → 2m are same upto relabeling, i.e, for all u, v ∈ [m] we have ĉh −1 t (u) = ĉh −1 t (v) if and only if ch −1 t (u) = ch −1 t (v). Note that a child function is
invertible by definition. We can ignore this label permutation and assume ĉht = cht for cleaner expressions. This can be done without any effect. We will assume ĉht = cht when stating the next three induction hypothesis.
IH.2 There exists a permutation mapping θt : {0, 1}d → {0, 1}d and % ∈ (0, 12d ) such that for every i ∈ [d] and s ∈ St we have:
P(φ̂t(xt)[i] = θ−1t (s)[i] | s[i]) ≥ 1− %,
P(φ̂t(xt) = θ−1t (s) | s) ≥ 1− d% ≥ 1
2
The two distributions are independent of the roll-in distribution at time step t. The first one holds as φ̂t(xt)[i] only depends upon the value xt[ĉht(i)] = xt[cht(i)] which only depends on s[i]. The second one holds as xt is independent of everything else given st. The form of % will become clear the end of analysis.
IH.3 For every s ∈ St−1, s′ ∈ St and a ∈ A we have: ∥∥∥T̂t(θ−1t (s′) | θ−1t−1(s), a)− T (s′ | s, a) ∥∥∥ TV ≤ 3d(∆est + ∆app),
where ∆est,∆app > 0 denote estimation and approximation errors whose form will become clear at the end of analysis.
IH.4 For every s ∈ St and K ∈ C≤2κ([d]), let Z = s[K] and Ẑ = θ−1t (s)[K], then there exists a policy πKẐ ∈ Ψt such that:
PπKẐ (st[K] = Z) ≥ αηt(K,Z) ≥ αηmin.
Base Case. In the base case (t = 1), we have a deterministic start state. Therefore, we can without loss of generality assume a single factor and define ĉh1[1] = m. As we can also define ch1[1] = [m] without loss of generality, therefore, this trivially satisfies the induction hypothesis 1. We define φ̂1 : X → [0]d (Algorithm 1, line 1). This satisfies induction hypothesis 2 with θ1 being the identity map. The induction hypothesis 3 is vacuous since there is no transition function before time step 1. For the last condition, we have for any K, Z = [0]|K| and Ẑ = [0]|K|. For any policy π we have Pπ(s1[K] = Z) = Pπ(φ̂1(x1)[K] = Ẑ) = 1 ≥ ηmin2 . Note that we never take any action from this policy, therefore, we can simply define Ψ1 = ∅.

heading: D.1 GRAPH STRUCTURE IDENTIFICATION
text: In this section, we analyze the performance of Algorithm 2, given as input Ψh−1, φ̂h−1,F , β and n. We will analyze the performance a fixed pair of atoms u, v ∈ [m] and then apply the full result using union bound. We first state the result for the rejection sampling.
Lemma 6. For policy πI;Ẑ ∈ Ψh−1, event EI;Ẑ = {ŝh−1[I] = Ẑ} and k ∈ N, let D reject I;Ẑ := RejectSamp(πI;Ẑ , EI;Ẑ , k) be the distribution induced by our rejection sampling procedure. Let
Z = θ(Ẑ) denote the real factor values corresponding Ẑ . Then we have:
PDreject I;Ẑ
(sh−1[I] = Z) ≥ 1− %− (
1− ηmin 4
)k . (8)
Proof. From IH.4 we have PπI;Ẑ (sh−1[I] = Z) ≥ ηmin 2 . This implies:
PπI;Ẑ (ŝh−1[I] = Ẑ) ≥ PπI;Ẑ (ŝh−1[I] = Ẑ | sh−1[I] = Z)PπI;Ẑ (sh−1[I] = Z)
≥ (1− d%)ηmin 2 ≥ ηmin 4 , (using IH.2 and IH.4).
Let a = PπI;Ẑ (ŝh−1[I] = Ẑ) be the acceptance probability of event EI;Ẑ . then it is easy to see that the probability of the event occurring under DrejectI;Ẑ is:
PDreject I;Ẑ
( EI;Ẑ ) = a+ (1−a)a+ (1−a)2a+ · · · (1−a)k−1a = 1− (1−a)k ≥ 1− ( 1− ηmin
4
)k .
We express the desired failure probability as shown:
PDreject I;Ẑ (sh−1[I] 6= Z) = PDreject I;Ẑ
( sh−1[I] 6= Z, ŝh−1[I] 6= Ẑ ) +PDreject
I;Ẑ
( sh−1[I] 6= Z, ŝh−1[I] = Ẑ )
(9)
We bound the two terms below:
PDreject I;Ẑ
( sh−1[I] 6= Z, ŝh−1[I] 6= Ẑ ) ≤ PDreject
I;Ẑ
( ŝh−1[I] 6= Ẑ ) ≤ (
1− ηmin 4
)k , (10)
PDreject I;Ẑ
( sh−1[I] 6= Z, ŝh−1[I] = Ẑ ) ≤ PDreject
I;Ẑ
( ŝh−1[I] = Ẑ | sh−1[I] 6= Z, ) ≤ % (11)
Combining Equation 9, Equation 10 and Equation 11 we get:
PDreject I;Ẑ (sh−1[I] = Z) = 1− PDreject I;Ẑ
(sh−1[I] 6= Z) ≥ 1− %− (
1− ηmin 4
)k . (12)
We now analyze the situation for a given pair of atoms. Recall for any distribution D ∈ ∆(Sh−1) and a ∈ A, we denote D ◦ a as the distribution over Sh where s′ ∼ D ◦ a is sampled by sampling s ∼ D and then s′ ∼ T (. | s, a). We want to derive roll-in distributions at time step h, such that atoms coming from the same parent satisfy hypothesis H0 and atoms coming from different parents satisfy hypothesis H1 under this roll-in distribution. This will allow us to use independence test to identify the parent structure in the emission process. Specifically, we consider the roll-in distributions induced by DrejectI;Ẑ ◦ a for some sets I,J and action a. Instantiating the definition of these hypothesis from Appendix C, with these roll-in distributions and setting β = βmin gives us:
H0 : ‖PDreject I;Ẑ ◦a(x[u], x[v])− PDreject I;Ẑ ◦a(x[u])PDreject I;Ẑ ◦a(x[v])‖1 ≥ βmin 2
v.s. H1 : ‖PDreject I;Ẑ ◦a(x[u], x[v])− PDreject I;Ẑ ◦a(x[u])PDreject I;Ẑ ◦a(x[v])‖1 ≤ β2min 100
Lemma 7 (Same Factors). If for two atoms u, v we have ch−1h (u) = ch −1 h (v), i.e., they are from the same factor then the hypothesis H0 is true for D ◦ a for any D ∈ ∆(Sh−1) and a ∈ A. In particular, this is true for DrejectI;Ẑ ◦ a for any choice of sets I,J and action a.
Proof. Follows trivially from Assumption 2.
Lemma 8 (Different Factors). If for two atoms u, v we have ch−1h (u) = i and ch −1 h (v) = j and i 6= j, then if I contains pth(i) ∪ pth(j), then for % ≤ β 2 min 1200 and k ≥ 8ηmin ln ( 30 βmin ) , the hypothesis H1 holds for D reject I;Ẑ ◦ a for any Ẑ such that πI;Ẑ ∈ Ψh−1 and a ∈ A.
Proof. Let D′ ∈ ∆(Sh−1) be a distribution that deterministically sets sh−1[I] = Z . Then it is easy to verify that PD′◦a(xh[u], xh[v]) = PD′◦a(xh[u])PD′◦a(xh[v]) for any a ∈ A and xh ∈ Xh. Then for any Ẑ and action a ∈ A we have using triangle inequality:
∣∣∣∣PDreject I;Ẑ ◦a(xh[u], xh[v])− PDreject I;Ẑ ◦a(xh[u])PDreject I;Ẑ ◦a(xh[v]) ∣∣∣∣ 1
≤ ∣∣∣∣PDreject
I;Ẑ ◦a(xh[u], xh[v])− PD′◦a(xh[u], xh[v]) ∣∣∣∣ 1 +
∣∣∣∣PD′◦a(xh[u])− PDreject I;Ẑ ◦a(xh[u]) ∣∣∣∣ 1 + ∣∣∣∣PD′◦a(xh[v])− PDreject I;Ẑ ◦a(xh[v]) ∣∣∣∣ 1
As xh[u] and xh[v] come from different factors, therefore, we have
P(xh[u], xh[v] | sh−1, a) = P(xh[u] | sh−1[I], a)P(xh[v] | sh−1[I], a).
We use this to bound the three terms in the summation above.∣∣∣∣PDreject I;Ẑ ◦a(xh[u], xh[v])− PD′◦a(xh[u], xh[v]) ∣∣∣∣ 1
= ∑
xh[u],xh[v]
∣∣∣∣∣∣ ∑ sh−1[I] P(xh[u] | sh−1[I], a)P(xh[v] | sh−1[I], a) { PDreject I;Ẑ (sh−1[I])− PD′(sh−1[I]) }∣∣∣∣∣∣
≤ ∑
sh−1[I]
∑
xh[u],xh[v]
P(xh[u] | sh−1[I], a)P(xh[v] | sh−1[I], a) ∣∣∣∣PDreject
I;Ẑ (sh−1[I])− PD′(sh−1[I])
∣∣∣∣
≤ ∑
sh−1[I]
∣∣∣∣PDreject I;Ẑ (sh−1[I])− PD′(sh−1[I]) ∣∣∣∣
= ∣∣∣∣1− PDreject I;Ẑ (sh−1[I] = Z) ∣∣∣∣+ ∑
sh−1[I] 6=Z PDreject I;Ẑ (sh−1[I])
= 2 ( 1− PDreject
I;Ẑ (sh−1[I] = Z)
) ≤ 2%+ 2 ( 1− ηmin
4
)k .
The other two terms are bounded similarly which gives us: ∣∣∣∣PDreject
I;Ẑ ◦a(xh[u], xh[v])− PDreject I;Ẑ ◦a(xh[u])PD′◦a(xh[v]) ∣∣∣∣ 1 ≤ 6%+ 6 ( 1− ηmin 4 )k .
We want this quantity to be less than β 2 min
100 to satisfy hypothesis H1. We distribute the errors equally and use ln(1 + a) ≤ a for all a > −1 to get:
% ≤ β 2 min 1200 , k ≥ 8 ηmin ln
( 30
βmin
) . (13)
Theorem 3 (Learning ĉhh). Fix δind ∈ (0, 1). If % ≤ β 2 min 1200 and k ≥ 8ηmin ln ( 30 βmin ) and nind ≥
O (
1 β4min
ln m 2|A||F|(2ed)2κ+1
δind
) , then learned ĉhh is equivalent to chh upto label permutation with
probability at least 1− δind.
Proof. For any pair of atom u, v, if they are from the same factor then H0 holds from Lemma 7 and IndTest mark them dependent with probability at least 1− δ. This holds for every triplet of I,Z, a and there are at most |A|(2ed)2κ+1, of them. Hence, from union bound we mark u, v correctly as coming from different factors with probability at least 1− |A|(2ed)2κ+1δ. If u and v have different factors then for any I containing the parents of both of them, and any value of Z and a, H1 always holds from Lemma 8 and IndTest marks them as independent. Note that
such an I will exists since the we iterate over all possible sets of size upto 2κ. Hence, with probability at least 1 − |A|2κδ, we find u and v to be independent for every Z and a. Hence, our algorithm correctly will mark them as coming from different factors.
For a given u and v, we correctly predict their output with probability at least 1− |A|(2ed)2κ+1δ. Therefore, using union bound we correctly output right result for each u and v with probability at least 1−|A|m2(2ed)2κ+1δ. From Theorem 2, we require nind ≥ O ( 1 β4min ln |F|δ ) . Binding |A|(2ed)2κ+1δ to δind then gives us the required value of nind to achieve a success probability of at least 1−δind. If we correctly assess the dependence for every pair of atoms correctly, then trivially partitioning them using the dependence equivalence relation gives us ĉhh which is same as chh upto label permutation.

heading: D.2 LEARNING A STATE DECODER
text: We focus on the task of learning an abstraction at time step h using Algorithm 3. We have access to ĉhh which is same as chh upto label permutation. We showed how to do this in Appendix C. We will ignore the label permutation to avoid having to complicate our notations. This would essentially mean that we will recover a backward decoder φ̂h = ( φ̂h1, · · · , φ̂hd ) , where there is a bijection between { φ̂hi
} i and { φ?j } j .
As we learn each decoder {φ̂hi} independently of each other, therefore, we will focus on learning the decoder φ̂hi for a fixed i. The same analysis will hold for other decoder and with application of union bound, we will establish guarantees for all decoders. Further, since we are learning the decoder at a fixed time step h, therefore, we will drop the h from the subscript for brevity. We use additional shorthand described below and visualize some of them in Figure 2.
• s and x denote a state sh−1 and an observation xh−1 at time step h− 1 • s′ and x′ denotes state sh and observation xh at time step h • s′[i] denotes ith factor of state at time step h • š denotes s[pth(i)] which is the set of parent factors of s′[i]. Recall that from the factoriza-
tion assumption, we have T (s′[i] | s, a) = Ti(s′[i] | š, a) for any s, a. • φ̂i denotes φ̂hi decoder for ith factor at time step h • ω denotes pt(i) which is the set of indices of atoms emitted by s′[i]. • x̌ denotes x′[chh(i)] which is the collection of atoms generated by s′[i]. • N = |Ψh−1| is the size of policy cover for previous time step.
Let D = {(x(k), a(k), x̌(k), y(k))}nabsk=1 be a dataset of nabs real transitions (y = 1) and imposter transitions (y = 0) collected in Algorithm 3, line 2-4. We define the expected risk minimizer (ERM) solution as:
ĝi = arg min g∈G
1
nabs
nabs∑
k=1
{ g(x(k), a(k), x̌(k))− y(k) }2 (14)
Recall that by the structure of G, we have ĝi = (ûi, φ̂i) where ŵi ∈ W2 and φ̂i ∈ Φ : X ∗ → {0, 1} is the learned decoder. Our algorithm only cares about the properties of the decoder and we throw away the regressor ûi.
Let D(x, a, x̌) be the marginal distribution over transitions. We get the marginal distribution by marginalizing out the real (y = 1) and imposter transition (y = 0). We also define D(x, a) as the marginal distribution over x, a. We have D(x, a) = µh−1(x) 1|A| as both real and imposter transitions involve sampling x ∼ µh−1 and taking action uniformly. Recall that µh−1 is generated by roll-in with a uniformly selected policy in Ψh−1 till time step h− 1. Let P (x, a, x̌ | y = 1) be the probability of a transition being real and P (x, a, x̌ | y = 0) be the probability of the transition being imposter. We can express these probabilities as:
P (x, a, x̌ | y = 1) = D(x, a)T (x̌ | x, a), P (x, a, x̌ | y = 0) = D(x, a)ρ(x̌), (15) where ρ(x̌) = E(x,a)∼D[T (x̌ | x, a)] is the marginal distribution over x̌. We will overload the notation ρ to also define ρ(x′) = E(x,a)∼D[T (x′ | x, a)]. Lastly, we can express the marginal distribution over transition as:
D(x, a, x̌) = P (x, a, x̌ | y = 1)P (y = 1) + P (x, a, x̌ | y = 0)P (y = 0)
= µh−1(x)
2|A| {T (x̌ | x, a) + ρ(x̌)}
We start by expressing the Bayes optimal classifier for problem in Equation 14.
Lemma 9 (Bayes Optimal Classifier). Bayes optimal classifier g? for problem in Equation 14 is given by:
∀(x, a, x̌) ∈ suppD, g?(x, a, x̌) = Ti(φ ? i (x̌) | φ?(x)[pt(i)], a)
Ti(φ?i (x̌) | φ?(x)[pt(i)], a) + ρ(φ?i (x̌)) (16)
Proof. The Bayes optimal classifier is given by g?(x, a, x̌) = P (y = 1 | x, a, x̌) which can be expressed using Bayes rule as:
P (y = 1 | x, a, x̌) = P (x, a, x̌ | y = 1)P (y = 1) P (x, a, x̌ | y = 1)P (y = 1) + P (x, a, x̌ | y = 0)P (y = 0)
= P (x, a, x̌ | y = 1)
P (x, a, x̌ | y = 1) + P (x, a, x̌ | y = 0) , using p(y) = Bern( 1/2)
= D(x, a)T (x̌ | x, a)
D(x, a)T (x̌ | x, a) +D(x, a)ρ(x̌)
= T (x̌ | x, a)
T (x̌ | x, a) + ρ(x̌)
= qi(x̌ | φ?i (x̌))Ti(φ?i (x̌) | x, a)
qi(x̌ | φ?i (x̌))Ti(φ?i (x̌) | x, a) + qi(x̌ | φ?i (x̌))ρ(φ?i (x̌))
= Ti(φ
? i (x̌) | x, a)
Ti(φ?i (x̌) | x, a) + ρ(φ?i (x̌)) =
Ti(φ ? i (x̌) | φ?(x)[pt(i)], a)
Ti(φ?i (x̌) | φ?(x)[pt(i)], a) + ρ(φ?i (x̌)) .
Theorem 4 (Decoder Regression Guarantees). For any given δabs ∈ (0, 1) and nabs ∈ N we have the following with probability at least 1− δabs:
Ex,a,x̌∼D [ (ĝi(x, a, x̌)− g?(x, a, x̌))2 ] ≤ ∆(nabs, δabs, |G|),
where ∆(nabs, δabs, |G|) := cnabs ln |G| δabs and c is a universal constant.
Proof. This is a standard regression guarantee derived using Bernstein’s inequality with realizability (3). For example, see Proposition 11 in Misra et al. (2020) for proof.
Corollary 5. For any given δabs ∈ (0, 1) and nabs ∈ N we have the following with probability at least 1− δabs:
Ex,a,x̌∼D [|ĝi(x, a, x̌)− g?(x, a, x̌)|] ≤ √ ∆(nabs, δabs, |G|) (17)
Proof. Applying Jensen’s inequality (E[ √ Z] ≤ √ E[Z]) to Theorem 4 gives us:
Ex,a,x̌∼D [|ĝ(x, a, x̌)− g?(x, a, x̌)|] = Ex,a,x̌∼D [√ |ĝ(x, a, x̌)− g?(x, a, x̌)|2 ]
≤ √ Ex,a,x̌∼D [ (ĝ(x, a, x̌)− g?(x, a, x̌))2 ] ≤ √ ∆(nabs, δabs, |G|).
Coupling Distribution We introduce a coupling distribution following Misra et al. (2020).
Dcoup(x, a, x̌1, x̌2) = D(x, a)ρ(x̌1)ρ(x̌2). (18)
We also define the following quantity which will be useful for stating our results:
ξ(x̌1, x̌2, x, a) = T (x̌1 | x, a) ρ(x̌1) − T (x̌2 | x, a) ρ(x̌2) . (19)
Lemma 10. For any fixed δabs ∈ (0, 1) we have the following with probability at least 1− δabs:
Ex,a,x̌1,x̌2∼Dcoup [ 1{φ̂i(x̌1) = φ̂i(x̌2)} |ξ(x̌1, x̌2, x, a)| ] ≤ 8 √ ∆(nabs, δabs, |G|).
Proof. We define a shorthand notation E = 1{φ̂i(x̌1) = φ̂i(x̌2)} for brevity. We also define a different coupled distribution D′coup given below:
D′coup(x, a, x̌1, x̌2) = D(x, a)D(x̌1 | x, a)D(x̌2 | x, a) (20)
where D(x̌ | x, a) = 12 {T (x̌ | x, a) + ρ(x̌)}. It is easy to see that marginal distribution of D′coup over x, a, x̌1 is same as D(x, a, x̌1).
We first use the definition of ξ (Equation 19) and g? (Equation 9) to express their relation:
|g?(x, a, x̌1)− g?(x, a, x̌2)| = ρ(x̌1)ρ(x̌2) |ξ(x̌1, x̌2, x, a)|
(T (x̌1 | x, a) + ρ(x̌1))(T (x̌1 | x, a) + ρ(x̌2))
= ρ(x̌1)ρ(x̌2)
4D(x̌1 | x, a)D(x̌2 |, x, a) |ξ(x̌1, x̌2, x, a)| . (21)
The second line uses the definition ofD(x̌ | x, a). We can view ρ(x̌1)D(x̌1|x,a) and ρ(x̌2) D(x̌2|x,a) as importance weight terms. Multiplying both sides by E and taking expectation with respect to D′coup then gives us:
ED′coup [E |g?(x, a, x̌1)− g?(x, a, x̌2)|] = 1
4 EDcoup [E|ξ(x̌1, x̌2, x, a)|] (22)
We bound the left hand side of Equation 22 as shown below:
ED′coup [E |g?(x, a, x̌1)− g?(x, a, x̌2)|] ≤ ED′coup [E |g?(x, a, x̌1)− ĝi(x, a, x̌1)|] + ED′coup [E |ĝi(x, a, x̌1)− g?(x, a, x̌2)|] = ED′coup [E |g?(x, a, x̌1)− ĝi(x, a, x̌1)|] + ED′coup [E |ĝi(x, a, x̌2)− g?(x, a, x̌2)|] = 2ED′coup [E |g?(x, a, x̌1)− ĝi(x, a, x̌1)|] = 2ED [E |g?(x, a, x̌)− ĝ(x, a, x̌)|] ≤ 2 √ ∆(nabs, δabs,G)
Here the first inequality follows from triangle inequality. The second step is key where we use ĝi(x, a, x̌1) = ĝi(x, a, x̌2) whenever E = 1. This itself follows from the bottleneck structure of G where ĝi(x, a, x̌i) = ŵi(x, a, φ̂i(x̌)). The third step uses the symmetry of x̌1 and x̌2 inD′coup whereas the fourth step uses the fact that marginal distribution of D′coup is same as D. Lastly, final inequality uses E ≤ 1 and the result of Corollary 5. Combining the derived inequality with Equation 22 proves our result.
We define the quantity P(s′[i] = z | D′) := E(s,a)∼D′ [1{s′[i] = z}] for any distribution D′ ∈ ∆(Sh−1 ×A). From the definition of ρ, we have ρ(s′[i] = z) = P(s′[i] = z | D). Intuitively, as we have policy cover at time step h − 1 and we take actions uniformly, therefore, we expect to have good lower bound on P(s′[i] = z | D) for every i ∈ [d] and reachable z ∈ {0, 1}. Note that if z = 0 (z = 1) is not reachable then it means we always have sh[i] = 1 (sh[i] = 0) from our reachability assumption (see Section 2). We formally prove this next which will be useful later. Lemma 11. For any z ∈ {0, 1} such that s′[i] = z is reachable, we have:
ρ(s′[i] = z) = P(s′[i] = z | D) ≥ αηmin N |A|
Proof. Fix z in {0, 1}. As s′[i] = z is reachable, therefore, from the definition of ηmin we have:
ηmin ≤ sup π∈Π Pπ(s′[i] = z) ≤ sup π
∑
š,a
Pπ(š)T (s′[i] = z | š, a)
≤ ∑
š,a
sup π∈Π
Pπ(š)T (s′[i] | š, a) = ∑
š,a
η(š)T (s′[i] | š, a)
We use the derived inequality to bound P(s′[i] = z | D) as shown:
P(s′[i] = z | D) = ∑
š,a
µh−1(š) |A| T (s ′[i] = z | š, a) ≥ α N |A|
∑
š,a
η(š)T (s′[i] = z | š, a) ≥ αηmin N |A| .
The first inequality uses the fact that µh−1 is created by roll-in with a uniformly selected policy in Ψh−1 which is an α policy cover. Recall that N = |Ψh−1|. The second inequality uses the derived result above.
Lemma 12. For any x̌1, x̌2 such that φ?i (x̌1) and φ?i (x̌2) is reachable, we have:
Ex,a∼D [|ξ(x̌1, x̌2, x, a)|] ≥ 1{φ?i (x̌1) 6= φ?i (x̌2)} αηminσ
2N .
Proof. For any x, a, x̌1, x̌2 we can express ξ (Equation 19) as:
|ξ(x̌1, x̌2, x, a)| = ∣∣∣∣ T (φ?i (x̌1) | φ?(x)[pt(i)], a)
ρ(φ?i (x̌1)) − T (φ
? i (x̌2) | φ?(x)[pt(i)], a)
ρ(φ?(x̌2))
∣∣∣∣
where we use the factorization assumption and decodability assumption. Note that we are implicitly assuming φ?i (x̌1) and φ ? i (x̌2) are reachable, for the quantity ξ(x̌1, x̌2, x, a) to be well defined.
We define Di to be the marginal distribution over S[pt(i)] × A. Taking expectation on both side gives us:
Ex,a∼D [|ξ(x̌1, x̌2, x, a)|] = Eš,a∼Di [∣∣∣∣ T (φ?i (x̌1) | š, a) ρ(φ?i (x̌1)) − T (φ ? i (x̌2) | š, a) ρ(φ?i (x̌2)) ∣∣∣∣ ]
= ∑
š,a
|PDi(š, a | φ?i (x̌1))− PDi(š, a | φ?i (x̌2))|
= 2 ‖PDi(., . | φ?i (x̌1))− PDi(., . | φ?i (x̌2))‖TV
The second equality uses the definition of backward dynamics PDi over S[pt(i)] × A and the identity ρ(s′[i]) = P(s′[i] | D). If φ?i (x̌1) = φ?i (x̌2) then the quantity on the right is 0. Otherwise, this quantity is given by 2 ‖PDi(., . | s′[i] = 1)− PDi(., . | s′[i] = 0‖TV. In the later case, both
s′[i] = 1 and s′[i] = 0 configurations are reachable, and without loss of generality we can assume P(s′[i] = 0 | D) ≥ 1/2. Our goal is to bound this term using the margin assumption (Assumption 1). We do so using importance weight as shown below:
2 ‖PDi(., . | s′[i] = 1)− PDi(., . | s′[i] = 0)‖TV = ∑
š,a
∣∣∣∣Pui(š, a | s′[i] = 1) PDi(š, a | s′[i] = 1) Pui(š, a | s′[i] = 1) − Pui(š, a | s′[i] = 0) PDi(š, a | s′[i] = 0) Pui(š, a | s′[i] = 0) ∣∣∣∣
= ∑
š,a
Di(š, a)
ui(š, a) ∣∣∣∣Pui(š, a | s′[i] = 1) P(s′[i] = 1 | ui) P(s′[i] = 1 | Di) − Pui(š, a | s′[i] = 0) P(s′[i] = 0 | ui) P(s′[i] = 0 | Di) ∣∣∣∣
≥ min š,a
Di(š, a)
ui(š, a) P(s′[i] = 0 | Di) P(s′[i] = 0 | ui) ‖Pui(., . | s′[i] = 1)− Pui(., . | s′[i] = 0)‖TV
≥ min š,a
Di(š, a)
ui(š, a) P(s′[i] = 0 | Di) P(s′[i] = 0 | ui) σ
The first step applies importance weight. As ui has support over all reachable configurations š and actions a ∈ A, hence, we can apply importance weight. The second step uses the definition of backward dynamics (PD,Pui). The third step uses Lemma H.1 of Du et al. Du et al. (2019) (see Lemma 24 in Appendix E for statement). Finally, the last step uses Assumption 1. We bound the two multiplicative terms below:
We have D(š, a) = µh−1(š) 1|A| ≥ αηmin N |A| . The first equality uses the fact that actions are taken uniformly and second inequality uses the fact that µh−1 is an α-policy cover. As ui is the uniform distribution over S[pt(i)] × A, therefore, we have ui(š, a) = 12|pt(i)||A| . This gives us D(š,a) ui(š,a)
≥ αηmin N 2 |pt(i)|. We bound the other multiplicative term as shown below:
P(s′[i] = 0 | Di) P(s′[i] = 0 | ui) ≥ P(s′[i] = 0 | Di) ≥ 1 2 .
Combining the lower bounds for the two multiplicative terms and using 2|pt(i)| ≥ 1 we get: 2 ‖PDi(., . | s′[i] = 1)− PDi(., . | s′[i] = 0)‖TV ≥ αηminσ
2N . (23)
Lastly, recall that our desired result is given by 2 ‖PDi(., . | s′[i] = 1)− PDi(., . | s′[i] = 0)‖TV whenever φ?i (x̌1) 6= φ?i (x̌2) and 0 otherwise. Therefore, using the derived lower bound multiplied by 1{φ?i (x̌1) 6= φ?i (x̌2)} gives us the desired result. Corollary 6. We have the following with probability at least 1− δabs:
Ex̌1,x̌2∼ρ [ 1{φ̂i(x̌1) = φ̂i(x̌2)}1{φ?i (x̌1) 6= φ?i (x̌2)} ] ≤ 16N αηminσ √ ∆(nabs, δabs, |G|).
Proof. The proof trivially follows from applying the bound in Lemma 12 to Lemma 10 as shown below:
Ex,a,x̌1,x̌2∼Dcoup [ 1{φ̂i(x̌1) = φ̂i(x̌2)} |ξ(x̌1, x̌2, x, a)| ]
= Ex̌1,x̌2∼ρ [ 1{φ̂i(x̌1) = φ̂i(x̌2)}Ex,a,∼D [|ξ(x̌1, x̌2, x, a)|] ]
≥ αηminσ 2N
Ex̌1,x̌2∼ρ [ [1{φ̂i(x̌1) = φ̂i(x̌2)}1{φ?i (x̌1) 6= φ?i (x̌2)} ]
The inequality here uses Lemma 12. The left hand side is bounded by 8 √
∆(nabs, δabs, |G|) using Lemma 10. Combining the two bounds and rearranging the terms proves the result.
At this point we analyze the two cases separately. In the first case, s′[i] can be set to both 0 and 1. In the second case, s′[i] can only be set to one of the values and we will call s′[i] as degenerate at time step h. We will show how we can detect the second case, at which point we just output a decoder that always outputs 0. We analyze the first case below.
D.2.1 CASE A: WHEN s′[i] CAN TAKE ALL VALUES
Corollary 6 allows us to define a correspondence between learned state (i.e., output of φ̂i) and the actual state (i.e., output of φ?i ). We show this correspondence in the result. Theorem 7 (Correspondence Theorem). For any state factor s′[i], there exists exists û0 ∈ {0, 1} and û1 = 1− û0 with probability at least 1− δabs such that:
P(φ̂i(x̌) = û0 | s′[i] = 0) ≥ 1− %, P(φ̂i(x̌) = û1 | s′[i] = 1) ≥ 1− %,
where % := 16N 2|A|
α2η2minσ
√ ∆(nabs, δabs, |G|) and x̌ ∼ ρ, provided % ∈ ( 0, 12 ) .
Proof. For any u, z ∈ {0, 1} we define the following quantities:
Pz := Ex̌∼ρ[1{φ?i (x̌) = z}], Puz := Ex̌∼ρ[1{φ̂i(x̌) = u}1{φ?i (x̌) = z}]. It is easy to see that these quantities are related by: Pz = Puz + P(1−u)z . We define û0 = arg maxu∈{0,1} Pu0 and û1 = 1− û0. This can be viewed as the learned bit value which is in most correspondence with s[i] = 0. We will derive lower bound on Pû00/P0 and Pû11/P1 which gives us the desired result. We first derive the following lower bound on Pû00:
Pû00 ≥ Pû00 + Pû10 2 ≥ P0 2 , (24)
where we use the fact that max is greater than average. Further, for any u, z ∈ {0, 1} we have:
Ex̌2,x̌2∼ρ [ 1{φ̂i(x̌1) = φ̂i(x̌2)}1{φ?i (x̌1) 6= φ?i (x̌2)} ]
≥ Ex̌1,x̌2∼ρ [ 1{φ̂i(x̌1) = u}1{φ̂i(x̌2) = u}1{φ?i (x̌1) = z}1{φ?i (x̌2) = 1− z} ]
= PuzPu(1−z)
We define a shorthand notation ∆′ := 16Nαηminσ √
∆(nabs, δabs, |G|). Then from Corollary 6 we have proven that PuzPu(1−z) ≤ ∆′ for any u, z ∈ {0, 1}. This allows us to write:
Pû11 = P1 − Pû01 ≥ P1 − ∆′
Pû00 ⇒ Pû11 P1 ≥ 1− ∆
′
P1Pû00 ≥ 1− 2∆
′
P0P1
where the last inequality uses Equation 24. We will derive the same result for Pû00/P0.
Pû00 = P0 − Pû10 ≥ P0 − ∆′
Pû11 ⇒ Pû00 P0 ≥ 1− ∆
′
P0Pû11 ≥ 1− ∆
′
P0P1 − 2∆′ ,
where the last inequality uses derived bound for Pû11/P1. If we assume ∆′ ≤ P0P14 then we get Pû00 P0 ≥ 1− 2∆′P0P1 .
As P0 + P1 = 1, therefore, we get P0P1 = P0 − P 20 = P1 − P 21 . If P0 ≤ 12 then P0 − P 20 ≥ P02 . Otherwise, P0 > 12 which implies P1 ≤ 12 and P1 − P 21 ≥ P12 . This gives us P0P1 ≥ min{P02 , P12 }. Using lower bounds for P0 and P1 from Lemma 11 gives us P0P1 ≥ αηmin2N |A| , and allows us to write:
Pû11 P1 ≥ 1− 4N |A|∆
′
αηmin ,
Pû00 P0 ≥ 1− 4N |A|∆
′
αηmin .
It is easy to verify that % = 4N |A|∆ ′ αηmin . As Pû00P0 = P(φ̂i(x ′) = û0 | s[i] = 0) and Pû11P1 = P(φ̂i(x ′) = û1 | s[i] = 1), therefore, we prove our result. The only requirement we used is that ∆′ ≤ P0P14 which is ensured if % ∈ ( 0, 12 ) .
D.2.2 CASE B: WHEN s′[i] TAKES A SINGLE VALUE
We want to be able to detect this case with high probability so that we can learn a degenerate decoder that only takes value 0. This would trivially give us a correspondence result similar to Theorem 7.
We describe the general form of the LearnDecoder subroutine in Algorithm 6. The key difference from the case we covered earlier is line 6-10. For a given factor i, we first learn the model ĝi containing the decoder φ̂i, as before using noise contrastive learning. We then sample ndeg iid triplets Ddeg = {(xj , aj , x̌j)}ndegj=1 where (xj , aj) ∼ D and x̌j ∼ ρ (line 6-8). Recall x̌ = x[chh(i)]. Next, we compute the width of prediction values over Ddeg as defined below:
max j,k∈[ndeg]
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| (25)
If the width is smaller than a certain value then we determine the factor to be degenerate and output a degenerate decoder φ̂i := 0, otherwise, we stick to the decoder learned by our regressor task. The form of sample size ndeg will become clear at the end of analysis, and we will determine the reason for the choice of threshold for width in line 9. Intuitively, if the latent factor only takes one value then the optimal classifier will always output 1/2 and so our prediction values should be close to one another. However, if the latent factor takes two values then the model prediction should be distinct.
Algorithm 6 LearnDecoder(G,Ψh−1, ĉhh). Child function has type ĉhh : [dh]→ 2[m]
1: for i in [dh], define ω = ĉhh(i),D = ∅,Ddeg = ∅ do 2: for nabs times do // collect a dataset of real (y = 1) and imposter (y = 0) transitions 3: Sample (x(1), a(1), x′(1)), (x(2), a(2), x′(2)) ∼ Unf(Ψh−1) ◦ Unf(A) and y ∼ Bern( 12 ) 4: If y = 1 then D ← D ∪ (x(1), a(1), x′(1)[ω], y) else D ← D ∪ (x(1), a(1), x′(2)[ω], y) 5: ĝi := ûi, φ̂i = REG(D,G) // train the decoder using noise-contrastive learning 6: for ndeg times do // detect degenerate factors 7: Sample (x(1), a(1), x′(1)), (x(2), a(2), x′(2)) ∼ Unf(Ψh−1) ◦ Unf(A) 8: Ddeg ← Ddeg ∪ {(x(1), a(1), x′(2))} 9: if maxj,k∈[ndeg] ∣∣ĝi(xj , aj , x′j [ω])− ĝi(xk, ak, x′k[ω]) ∣∣≤ α 2η2minσ 40|Ψh−1|2|A| then // max over Ddeg
10: φ̂i := 0 // output a decoder that always returns 0 return φ̂ : X → {0, 1}dh where for any x ∈ X and i ∈ [dh] we have φ̂(x)[i] = φ̂i(x[ĉhh(i)]).
For convenience, we define D′(x, a, x̌) = D(x, a)ρ(x̌), and so (xj , aj , x̌j) ∼ D′. For brevity reasons, we do not add additional qualifiers to differentiate xj , aj , x̌j from the dataset of real and imposter transitions, we used in the previous section for the regression task. In this part alone, we will use xj , aj , x̌j to refer to the transitions collected for the purpose of detecting degenerate factors. Lemma 13 (Markov Bound). Let {(xj , aj , x̌j)}ndegj=1 be a dataset of iid transitions sampled from D′. Fix a > 0. Then with probability at least 1− δabs − 2ndeg √ ∆(nabs,δabs,|G|)
a we have: ∀j ∈ [ndeg], |ĝ(xj , aj , x̌j)− g?(xj , aj , x̌j | ≤ a.
Proof. It is straightforward to verify that for any (xj , aj , x̌j) we have D(xj , aj , x̌j) ≥ D′(xj ,aj ,x̌j)/2. Using Corollary 5 we get:
Ex,a,x̌∼D′ [|ĝ(x, a, x̌)− g?(x, a, x̌)|] ≤ 2 √
∆(nabs, δabs, |G|) Let Ej denote the event {|ĝ(xj , aj , x̌j)− g?(xj , aj , x̌j)| ≤ a} and Ej be its negation, then:
P(∩ndegj=1Ej) ≥ 1− ndeg∑
j=1
P(Ej) ≥ 1− 2ndeg
√ ∆(nabs, δabs, |G|)
a ,
where the first inequality uses union bound and the second inequality uses Markov’s inequality. As Corollary 5 holds with probability δabs, our overall failure probability is at most
δabs + 2ndeg √ ∆(nabs,δabs,|G|) a .
Lemma 14. For any reachable parent factor values š, action a ∈ A and reachable s′[i] ∈ {0, 1}, we have D′(š, a, s′[i]) ≥ α
2η2min N2|A|2 .
Proof. We have D′(š, a, s′[i]) = µh−1(š)|A| ρ(s ′[i]) ≥ α 2η2min N2|A|2 , where used the induction hypothesis IH.4 that Ψ is an α-policy cover of Sh−1 and Lemma 11. Lemma 15 (Degenerate Factors). Fix a > 0. If s′[i] only takes a single value then with probability at least 1− δabs − 2ndeg √ ∆(nabs,δabs,|G|) a we have:
max j,k∈[ndeg]
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| ≤ 2a
Proof. When s′[i] takes a single value then g? is the constant function 12 . For any j and k we get the following using Lemma 13 and triangle inequality.
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| ≤ |ĝ(xj , aj , x̌j)− g?(xj , aj , x̌j)|+ |g?(xk, ak, x̌k)− ĝ(xk, ak, x̌k)| ≤ 2a.
Lemma 16 (Non Degenerate Factors). Fix a > 0 and assume ndeg ≥ N 2|A|2
α2η2min , then we have:
max j,k∈[ndeg]
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| ≥ α2η2minσ
16N2|A| − 2a
with probability at least 1− δabs − 2ndeg √ ∆(nabs,δabs,|G|) a − 4 exp ( −α 2η2minndeg 3N2|A|2 ) .
Proof. Equation 23 implies that there exists š, a such that ∣∣∣∣ T (s′[i] = 1 | š, a) ρ(s′[i] = 1) − T (s ′[i] = 0 | š, a) ρ(s′[i] = 0) ∣∣∣∣ ≥ αηminσ 2N
Combining this with Equation 21 we get:
|g?(š, a, s′[i] = 1)− g?(š, a, s′[i] = 0)| ≥ ρ(s ′[i] = 1)ρ(s′[i] = 0)
4D(s′[i] = 1 | š, a)D(s′[i] = 0 | š, a) αηminσ 2N
≥ α 2η2minσ
16N2|A| (26)
where Equation 26 uses ρ(s′[i] = 1)ρ(s′[i] = 0) ≥ αηmin2N |A| , as one of the terms is at least 1/2 and other can be bounded using Lemma 11.
Say we have two examples in our dataset, say {(x1, a1, x̌1), (x2, a2, x̌2)} without loss of generality, such that φ?(x1)[ω] = φ?(x2)[ω] = š, action a1 = a2 = a, φ?i (x̌1) = 1, and φ ? i (x̌2) = 0. Then we have:
max j,k∈[ndeg]
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| ≥ |ĝ(x1, a1, x̌1)− ĝ(x2, a2, x̌2)|
≥ |g?(š, a, 1)− g?(š, a, 0)| − |ĝ(x1, a1, x̌1)− g?(x1, a1, x̌1)| − |ĝ(x2, a2, x̌2)− g?(x2, a2, x̌2)|
≥ α 2η2minσ
16N2|A| − 2a (using Equation 26 and Lemma 13)
We use Lemma 13 which has a failure probability of δabs + 2ndeg √ ∆(nabs,δabs,|G|) a . Further, we also assume that our dataset contains both (š, a, s′[i] = 1) and (š, a, s′[i] = 0). Probability of one of these
events is given by Lemma 14. Therefore, if ndeg ≥ N 2|A|2
α2η2min then from Lemma 25 and union bound,
the probability that at least one of these transitions does not occur is given by 4 exp ( −α
2η2minndeg 3N2|A|2
) .
The total failure probability is given by union bound and computes to:
δabs + 2ndeg
√ ∆(nabs, δabs, |G|)
a + 4 exp
( −α
2η2minndeg 3N2|A|2
) .
If we fix a = α 2η2minσ
80N2|A| then in the two case we have:
(Degenerate Factor) max j,k∈[ndeg]
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| ≤ α2η2minσ
40N2|A|
(Non-Degenerate Factor) max j,k∈[ndeg]
|ĝ(xj , aj , x̌j)− ĝ(xk, ak, x̌k)| ≥ 3α2η2minσ
80N2|A| Theorem 8 (Detecting Degenerate Case). We correctly predict if s′[i] is a degenerate factor or not when using ndeg =
3N2|A|2 α2η2min
log (
4 δabs
) and % ≤ α
2η2minδabs 30N2|A|2 log
−1 (
4 δabs
) , with probability at least
1− 3δabs.
Proof. The result follows by combining Lemma 16 and Lemma 15, and using the value of a described above. These two results hold with probability at least:
1− δabs − 2ndeg
√ ∆(nabs, δabs, |G|)
a − 4 exp
( −α
2η2minndeg 3N2|A|2
)
Setting the hyperparameters to satisfy the following:
ndeg = 3N2|A|2 α2η2min log
( 4
δabs
) , ∆(nabs, δabs, |G|)−1/2 ≥
480N4|A|3 α4η4minδabsσ log
( 4
δabs
) ,
gives a failure probability of at most 3δabs. The later condition can be expressed in terms of % which gives us the desired bounds (see Theorem 7 for definition of %). Lastly, note that setting ndeg this way also satisfies the requirement in Lemma 16. Lastly, note that the resultant bound on % is much stronger than required for Theorem 7. Therefore, we can significantly improve the complexity bounds in the setting where there are no degenerate state factors.

heading: D.2.3 COMBINING CASE A AND CASE B
text: Theorem 8 shows that we can detect degenerate state factors with high probability. If we have a degenerate state factor and we detect it, then correspondence theorem holds trivially. However, if we don’t have degeneracy and we correctly predict it, then we stick our learned decoder and Theorem 7 holds true. These two results allows us to define a bijection between real states and learned states that we explain below.
Bijective Mapping between real and learned states For a given time step h and state bit s[i] = z, we will define ûhiz as the corresponding learned state bit. When h and i will be clear from the context then we will express this as ûz . We will use the notation ŝ to denote a learned state at time step h − 1 and ŝ′ to denote learned state at time step h. Let pt(i) = (i1, · · · , il) and s[pt(i)] = w := (w1, w2, · · ·wl), then we define ŵ = (û(h−1)i1w1 · · · û(h−1)ilwl) as the learned state bits corresponding to w. More generally, for a given set K ∈ 2d, we denote the real state factors as s[K] (or s′[K]) and the corresponding learned state factors as ŝ[K] (or ŝ′[K]). We define a mapping θh : {0, 1}d → {0, 1}d from learned state to real state. We will drop the subscript h when the time step is clear from the context. We denote the domain of θh by Ŝh which is a subset of {0, 1}d. Note that every real state may not be reachable at time step h. E.g., maybe our decoder outputs ŝ′ = (0, 0) but that the corresponding real state is not reachable at time step h. Figure 3 visualizes the mapping.
For a learned state ŝ we have s = θ(ŝ) if s = (z1, · · · , zd) and ŝ = (uh1z1 , · · · , uhdzd). We would also overload our notation to write w = θ(ŵ) for a given ŝ[K] = ŵ where w = θ(ŝ)[K], whenever factor set K is clear from the context. We call s′ (or s) as reachable if s′ ∈ Sh (or s ∈ Sh−1). Similarly, we call ŝ′ (or ŝ) as reachable if θh(ŝ
′) (or θh−1(ŝ)) are reachable. For a given set of factorsK, we call w ∈ {0, 1}|K| as reachable for K if there exists a reachable state with factorsK taking on valuew. Similarly, we define ŵ ∈ {0, 1}|K| as reachable for a given K if θ(ŵ) is reachable for K. We use the mapping to state the correspondence theorem for the whole state.
Corollary 9 (General Case). If ndeg = 3N 2|A|2 α2η2min log ( 4 δabs ) and % ≤ α 2η2minδabs 30N2|A|2 log −1 ( 4 δabs ) holds, then with probability at least 1− 3dδabs, we have: ∀s′ ∈ Sh, P(ŝ′[i] = θ−1(s′)[i] | s′[i]) ≥ 1− %, P(ŝ′ = θ−1(s′) | s′) ≥ 1− d%.
Proof. The first result directly follows from being able to detect if we are in degenerate setting or not and if not, then applying Theorem 7, and if yes then result holds trivially. This holds with probability at least 1− 3δabs from Theorem 8. Applying union bound over all d learned factors gives us success probability across all factors of at least 1− 3dδabs. The second one follows from union bound.
P(ŝ 6= θ−1(s) | s) = P(∃i : ŝ[i] 6= θ−1(s)[i] | s[i]) ≤ d∑
i=1
P(ŝ[i] 6= θ−1(s)[i] | s[i]) ≤ d%.
As we noticed before, our bounds can be significantly improved in the case of no degenerate factors. This prevents application of expensive Markov inequality. Therefore, we also state bounds for the special case below. Corollary 10 (Degenerate Factors Absent). If % ∈ ( 0, 12 ) , then with probability at least 1 − dδabs, we have:
∀s′ ∈ Sh, P(ŝ′[i] = θ−1(s′)[i] | s′[i]) ≥ 1− %, P(ŝ′ = θ−1(s′) | s′) ≥ 1− d%.
Proof. Same as Corollary 9 except we can directly apply Theorem 7 as we don’t need to do any expensive check for a degenerate factor.

heading: D.3 MODEL ESTIMATION
text: Our next goal is to estimate a model T̂h : Ŝh−1 × A → ∆(Ŝh) and latent parent structure p̂th, given roll-in distribution D ∈ ∆(Sh−1), and the learned decoders {φ̂t}t≤h. Recall that our approach estimates the model by count-based estimation. Let D = {(x(k), a(k), x′(k))}nestk=1 be the sample collected for model estimation.
Recall that we estimate p̂th(i) be using a set of learned factors I that we believe is p̂th(i) and varying a disjoint set of learned factors J . If p̂t(i) ⊆ I then we expect the learned model to behave the same
irrespective of how we vary the factors J . However, if p̂t(i) 6⊆ I then there exists a parent factor that on varying will different values for the learned dynamics. Let K = I ∪ J be the set of all factors in the control group (I) and variable group (J ). We will first analyze the case for a fixed i ∈ Ŝh and control group I and variable group J . For a given v̂ ∈ {0, 1}, ŵ ∈ {0, 1}|K| and a ∈ A, we have P̂D(ŝ′[i] = v̂ | ŝ[K] = ŵ, a) denoting the estimate probability derived from our count based estimation (Algorithm 4, line 3). Let PD(ŝ′[i] = v̂ | ŝ[K] = ŵ, a) be the probabilities that are being estimated. It is important to note that we use subscript D for these notations as the learned states ŝ are not Markovian and K may not contain pt(i), therefore, the estimated probabilities P̂ and expected probabilities P will be dependent on the roll-in distribution D.
In order to estimate P̂D(. | ŝh−1[K] = ŵ, a) we want good lower bounds on PD(ŝh−1[K] = ŵ, a) for every a ∈ A and ŵ reachable for K. Our roll-in distribution D only guarantees lower bound on PD(sh−1[K], a). However, we can use IH.2 to bound the desired quantity below. Lemma 17 (Model Estimation Coverage). If % ≤ 12 then for all K ∈ C≤2κ([d]), a ∈ A and ŵ ∈ {0, 1}|K| reachable for K, we have:
PD(ŝh−1[K] = ŵ, a) ≥ αηmin
4κN |A|
Proof. We can express PD(ŝh−1[K] = ŵ, a) = 1|A|PD(ŝh−1[K] = ŵ) as actions are taken uniformly. Let sh−1 = θh−1(ŝh−1) and w = sh−1[K]. We bound PD(ŝh−1[K] = ŵ) as shown:
PD(ŝh−1[K] = ŵ) ≥ PD (ŝh−1[K] = ŵ, sh−1[K] = w) = PD (ŝh−1[K] = ŵ | sh−1[K] = w)PD (sh−1[K] = w) = ∏
k∈K PD (ŝh−1[k] = ŵk | sh−1[k] = wk)PD (sh−1[K] = w)
≥ (1− %)2καηmin N ≥ αηmin 4κN ,
where the third step uses the fact that value of learned state ŝh−1[k] is independent of other decoders given the real state bit sh−1[k]. The fourth step uses IH.2 and the fact that we have good coverage over all sets of state factors of size at most 2κ. Last inequality uses |K| ≤ 2κ and % ≤ 12 .
We now show that our count-based estimator P̂D converges to PD and derive the rate of convergence. Lemma 18 (Model Estimation Error). Fix δest ∈ (0, 1). Then with probability at least 1− δest for every K ∈ C≤2κ([d]), ŵ ∈ {0, 1}|K| reachable for K, and a ∈ A we have the following: ∑
v̂∈{0,1}
∣∣∣P̂D(ŝh[i] = v̂ | ŝh−1[K] = ŵ, a)− PD(ŝh[i] = v̂ | ŝh−1[K] = ŵ, a) ∣∣∣ ≤ 2∆est(nest, δest),
where ∆est(nest, δest) := 12 ( 2κ+5N |A| αηminnest )1/2 ln ( 4e|A|(ed)2κ δest ) .
Proof. We sample nest samples by roll-in at time step h− 1 with distribution D and taking actions uniformly. We first analyze the failure probability for a given K, ŵ, a. Let E(K, ŵ, a) denote the event {ŝh−1[K] = ŵ, ah−1 = a}. If E(K, ŵ, a) occurs in our dataset at least m times for some m ≥ 16 2 ln(1/δ) then from Corollary 15 we have∑
v̂∈{0,1}
∣∣∣P̂D(ŝh[i] = v̂ | ŝh−1[K] = ŵ, a)− PD(ŝh[i] = v̂ | ŝh−1[K] = ŵ, a) ∣∣∣ ≤ ,
with probability at least 1 − δ. Lemma 17 shows that probability of E(K, ŵ, a) is at least αηmin4κN |A| . Therefore, from Lemma 28 if nest ≥ 2 2κ+1mN |A| αηmin ln ( e δ ) then we get at least m samples of event E(K, ŵ, a) with probability at least 1− δ. Therefore, the total failure probability is at most 2δ: δ due to not getting at least m samples and δ due to Corollary 15 on getting m samples. This holds for every triplet (K, ŵ, a) and Lemma 23 shows that there are at most 2(ed)2κ|A| such triplets. Hence, with application of union bound we get the desired result.
Lemma 19 (Model Approximation Error). For any i ∈ [d],K ∈ C≤2κ([d]), s ∈ Sh−1, a ∈ A, s′ ∈ Sh, let ŝ = θ−1h−1(s) and ŝ′ = θ−1h (s′). Then we have:
|PD(ŝ′[i] | ŝ[K], a)− PD(s′[i] | s[K], a)| ≤ ∆app := 5κ%N
αηmin .
Proof. We will first bound |PD(s′[i] | ŝ[K], a)− PD(s′[i] | s[K], a)| and then use correspondence result (Corollary 9) to prove the desired result. We start by expressing our conditional probabilities as ratio of joint probabilities.
PD(s′[i] | ŝ[K], a) = PD(s′[i], ŝ[K], a)
PD(ŝ[K], a) , PD(s′[i] | s[K], a) = PD(s′[i], s[K], a) PD(s[K], a) .
From Lemma 29 we have:
|PD(s′[i] | ŝ[K], a)− PD(s′[i] | s[K], a)| ≤ ε+ ε2
PD(s[K], a) , where (27)
ε1 := |PD(s′[i], ŝ[K], a) − PD(s′[i], s[K], a)| and ε2 := |PD(ŝ[K], a) − PD(s[K], a)|. We bound these two quantities below:
ε1 = ∣∣∣∣∣∣ ∑ sh−1[K] PD(s′[i], ŝ[K], sh−1[K], a)− ∑ ŝh−1[K] PD(s′[i], ŝh−1[K], s[K], a) ∣∣∣∣∣∣
= ∣∣∣∣∣∣ ∑ sh−1[K] 6=s[K] PD(s′[i], ŝ[K], sh−1[K], a)−
∑
ŝh−1[K] 6=ŝ[K] PD(s′[i], ŝh−1[K], s[K], a) ∣∣∣∣∣∣
≤ max   
∑
sh−1[K]6=s[K] PD(s′[i], ŝ[K], sh−1[K], a) ︸ ︷︷ ︸ Term 1
, ∑
ŝh−1[K] 6=ŝ[K] PD(s′[i], ŝh−1[K], s[K], a) ︸ ︷︷ ︸ Term 2
   ,
Where the first inequality uses |a− b| ≤ max{a, b} for a, b > 0. We bound Term 1 below:
Term 1: 1 |A| ∑
sh−1[K]6=s[K] PD(s′[i] | ŝ[K], sh−1[K], a)P(ŝ[K] | sh−1[K])PD(sh−1[K])
≤ %|A| ∑
sh−1[K]6=s[K] PD(s′[i] | ŝ[K], sh−1[K], a)PD(sh−1[K])
≤ %|A| ∑
sh−1[K]6=s[K] PD(sh−1[K]) ≤
%
|A|
The key inequality here is P(ŝ[K] | sh−1[K]) = ∏ k∈K P(ŝ[k] | sh−1[k]) ≤ %, as there exist at least one j ∈ K such that sh−1[j] 6= s[j] and for this j we have P(ŝ[j] | sh−1[j]) ≤ %. We bound Term 2 similarly:
Term 2: 1 |A| ∑
ŝh−1[K] 6=ŝ[K] PD(s′[i] | ŝh−1[K], s[K], a)PD(ŝh−1[K] | s[K])PD(s[K])
≤ 1|A| ∑
ŝh−1[K]6=ŝ[K] PD(ŝh−1[K] | s[K]) =
1
|A| {1− PD(ŝ[K] | s[K])}
≤ 1|A| (1− (1− %) |K|) ≤ 2κ%|A|
where we use PD(ŝ[K] | s[K]) = ∏ k∈K P(ŝ[k] | s[k]) ≥ (1− %)|K| and |K| ≤ 2κ.
This gives us ε1 ≤ 2κ%|A| . The proof for ε2 is similar.
ε2 = ∣∣∣∣∣∣ ∑ sh−1[K] PD(ŝ[K], sh−1[K], a)− ∑ ŝh−1[K] PD(ŝh−1[K], s[K], a) ∣∣∣∣∣∣
= ∣∣∣∣∣∣ ∑ sh−1[K] 6=s[K] PD(ŝ[K], sh−1[K], a)−
∑
ŝh−1[K] 6=ŝ[K] PD(ŝh−1[K], s[K], a) ∣∣∣∣∣∣
max   
∑
sh−1[K] 6=s[K] PD(ŝ[K], sh−1[K], a) ︸ ︷︷ ︸ Term 3
, ∑
ŝh−1[K]6=ŝ[K] PD(ŝh−1[K], s[K], a) ︸ ︷︷ ︸ Term 4
  
We bound Term 3 below similar to Term 1:
Term 3: 1 |A| ∑
sh−1[K] 6=s[K] PD(ŝ[K]|sh−1[K])PD(sh−1[K])
≤ %|A| ∑
sh−1[K]6=s[K] PD(sh−1[K]) ≤
%
|A|
and Term 4 is bounded similar to Term 2 below:
Term 4: 1 |A| ∑
ŝh−1[K] 6=ŝ[K] PD(ŝh−1[K] | s[K])PD(s[K])
≤ 1|A| {1− PD(ŝ[K] | s[K])} ≤ 1|A| { 1− (1− %)|K| } ≤ 2κ%|A|
This gives us ε2 ≤ 2κ%|A| . Plugging bounds for ε1 and ε2 in Equation 27 and using PD(s[K], a) = PD(s[K]) |A| ≥ αηmin N |A| gives us:
|PD(s′[i] | ŝ[K], a)− PD(s′[i] | s[K], a| ≤ 4κ% |A|PD(s[K], a) ≤ 4κ%N αηmin . (28)
We can use correspondence result to derive a lower bound:
PD(ŝ′[i] | ŝ[K], a) ≥ PD(ŝ′[i] | s′[i])PD(s′[i] | ŝ[K], a) ≥ (1− %)PD(s′[i] | ŝ[K], a) ≥ PD(s′[i] | ŝ[K], a)− %
and an upper bound:
PD(ŝ′[i] | ŝ[K], a) = PD(ŝ′[i] | s′[i])PD(s′[i] | ŝ[K], a)+ P(ŝ′[i] | 1− s′[i])PD(1− s′[i] | ŝ[K], a)
≤ PD(s′[i] | ŝ[K], a) + % Combing the lower and upper bounds with Equation 28 gives us:
|PD(ŝ′[i] | ŝ[K], a)− PD(s′[i] | s[K], a)| ≤ 4κ%N αηmin + % ≤ 5κ%N αηmin .
which is the desired result.
We can merge the estimation error and approximation error to generate the total error.
Lemma 20 (K-Model Error). For any i ∈ [d], K ∈ C≤2κ([d]), s ∈ Sh−1, a ∈ A, s′ ∈ Sh, let ŝ = θ−1h−1(s) and ŝ ′ = θ−1h (s ′). Then we have:
∣∣∣P̂D(ŝ′[i] | ŝ[K], a)− PD(s′[i] | s[K], a) ∣∣∣ ≤ ∆est(nest, δest) + ∆app.
with probability at least 1− δest.
Proof. Follows trivially by combining the estimation error (Lemma 18) and approximation error (Lemma 19) with application of triangle inequality.
D.4 DETECTING LATENT PARENT STRUCTURE IN TRANSITION pth
We are now ready to analyze the performance of learned parent function p̂th. Let K1,K2 ∈ C≤2κ([2d]) and ŵ1 ∈ {0, 1}|K1|, ŵ2 ∈ {0, 1}|K2|. We will assume ŵ1 is reachable for K1 and ŵ2 is reachable for K2. For convenience we will define the following quantity Ω to measure total variation distance between distributions P̂D(s′[i] | ·, ·) conditioned on setting ŝ[K1] = ŵ1 and ŝ[K2] = ŵ2, and for a fixed action a ∈ A:
Ω̂ia(K1, ŵ1,K2, ŵ2) := 1
2
∑
v̂∈{0,1}
∣∣∣P̂D(ŝ′[i] = v̂ | ŝ[K1] = ŵ1, a)− P̂D(ŝ′[i] = v̂ | ŝ[K2] = ŵ2, a) ∣∣∣ .
We can compute Ω̂ for every value of i,K1, ŵ1,K2, ŵ2, a in computational time ofO ( (2ed)3κ+3|A| ) . We also define a similar metric for the true distribution for any K1,K2 and v ∈ {0, 1}, w1 ∈ {0, 1}|K1|, w2 ∈ {0, 1}|K2| and a ∈ A:
Ωia(K1, w1,K2, w2) := 1
2
∑
v∈{0,1} |PD(s′[i] = v | s[K1] = w1, a)− PD(s′[i] = v | s[K2] = w2, a)| .
Recall that [I;J ] denotes concatenation of two ordered sets I and J . We use this notation to state our next result. Lemma 21 (Inclusive Case). Fix i ∈ [d] and I ∈ C≤κ([d]). If pt(i) ⊆ I then for all a ∈ A and û ∈ {0, 1}|I| we get:
max J1,J2,ŵ1,ŵ2
Ω̂ia([I;J1], [û; ŵ1], [I;J2], [û; ŵ2]) ≤ 2∆est(nest, δest) + 2∆app,
where max is taken over J1,J2 ∈ C≤κ([d]), ŵ1 ∈ {0, 1}|J1|, ŵ2 ∈ {0, 1}|J2| such that [û; ŵ1] is reachable for [I;J1], [û; ŵ2] is reachable for [I;J2], and I ∩ J1 = I ∩ J2 = ∅.
Proof. We fix J1,J2, û, ŵ1, ŵ2, a and let K1 = [I;J1], K2 = [I;J2], v = θ(v̂), u = θ(û), w1 = θ(ŵ1), and w2 = θ(ŵ2). As pt(i) ⊆ I, therefore, we have: PD(s′[i] = v | s[K1] = [u;w1], a) = Ti(s′[i] = v | s[I] = u, a) = PD(s′[i] = v | s[K2] = [u;w2], a) Using this result along with Lemma 20 and application of triangle inequality we get: ∣∣∣P̂D(ŝ′[i] = v̂ | ŝ[K1] = [û; ŵ1], a)− P̂D(ŝ′[i] = v̂ | ŝ[K2] = [û; ŵ2], a) ∣∣∣ ≤ 2∆est(nest, δest)+2∆app.
Summing over v̂, dividing by 2, and using the definition of Ω̂ proves the result.
The following is a straightforward corollary of Lemma 21.
Corollary 11. Fix i ∈ [d] then there exists an I such that for all a ∈ A and û ∈ {0, 1}|I|: max
J1,J2,ŵ1,ŵ2 Ω̂ia([I;J1], [û; ŵ1], [I;J2], [û; ŵ2]) ≤ 2∆est(nest, δest) + 2∆app,
where max is taken over J1,J2, ŵ1, ŵ2 satisfy the restrictions stated in Lemma 21.
Proof. Take any I such that pt(i) ⊆ I and apply Lemma 21. Note that we are allowed to pick such an I as |pt(i)| ≤ κ by our assumption.
Recall that we define p̂t(i) as the solution of the following problem:
p̂t(i) := argmin I max a,û,J1,J2,ŵ1,ŵ2
Ω̂ia([I;J1], [û; ŵ1], [I;J2], [û; ŵ2]), (29)
where I ∈ C≤κ([d]), a ∈ A, and û,J1,J2, ŵ1, ŵ2 satisfy the restrictions stated in Lemma 21. We are now ready to state our main result for p̂t.
Theorem 12 (Property of p̂t). For any s ∈ Sh−1, a ∈ A, s′ ∈ Sh, let ŝ = θ−1h−1(s) and ŝ′ = θ−1h (s′). Then the learned parent function p̂t satisfies:
∀ ∈ [d], ∣∣∣P̂D(ŝ′[i] | ŝ[p̂t(i)], a)− Ti(s′[i] | s[pt(i)], a) ∣∣∣ ≤ 3∆est(nest, δest) + 3∆app.
Proof. Fix i ∈ [d]. Let J = pt(i)− p̂t(i) and K = p̂t(i) ∪ J . As pt(i) ⊆ K, therefore, we have: PD(s′[i] | s[K], a) = Ti(s′[i] | s[pt(i)], a)
Combining this result with Lemma 20 we get: ∣∣∣P̂D(ŝ′[i] | ŝ[K], a)− Ti(s′[i] | s[pt(i)], a) ∣∣∣ ≤ ∆est(nest, δest) + ∆app.
From the definition of p̂t(i) (Equation 29) and Corollary 11 we have: ∣∣∣P̂D(ŝ′[i] | ŝ[p̂t(i); ∅], a)− P̂D(ŝ′[i] | ŝ[p̂t(i);J ], a) ∣∣∣ ≤ 2∆est(nest, δest) + 2∆app.
Note that we are allowed to use Corollary 11 as ŝ[p̂t(i); ∅] and ŝ[p̂t(i);J ] are both reachable since they are derived from a reachable real state s, |p̂t(i)| ≤ κ, |[p̂t(i);J ]| ≤ |[p̂t(i); pt(i)]| ≤ 2κ, and p̂t(i) ∩ ∅ = ∅ = p̂t(i) ∩ J . Combining the previous two inequalities using triangle inequality completes the proof.

heading: D.5 BOUND TOTAL VARIATION BETWEEN ESTIMATED MODEL AND TRUE MODEL
text: Given the learned transition parent function p̂th we define the transition model as:
T̂hi ( ŝ′[i] | ŝ[p̂th(i)], a ) = P̂D ( ŝ′[i] | ŝ[p̂th(i)], a ) , T̂h (ŝ ′ | ŝ, a) = d∏
i=1
T̂hi ( ŝ′[i] | ŝ[p̂th(i)], a ) .
From Theorem 12 we have for any i ∈ [d], ŝ ∈ Sh−1, a ∈ A, s′ ∈ Sh, and ŝ = θ−1(s), ŝ′ = θ−1(s′): ∣∣∣T̂hi(ŝ′[i] | ŝ[p̂t(i)], a)− T (s′[i] | s[pt(i)], a) ∣∣∣ ≤ 3∆est(nest, δest) + 3∆app.
Transition Closure. A subtle point remains before we prove the model error between T̂h and T . Theorem 12 only states guarantee for those ŝ that are inverse of a reachable state s. However, as stated before, due to decoder error we can reach a state ŝ which does not have a corresponding reachable state, i.e. θ(ŝ) 6∈ Sh (see Figure 3). We cannot get model guarantees for these unreachable states ŝ since we may reach them with arbitrarily small probability. However, we can still derive model error if we can simply define the real transition probabilities in terms of the learned probabilities for these states. This will not cause a problem since the real model will never reach these states. We start by defining the closure of the transition model T ◦ for time step h as:
∀ŝ ∈ Ŝh−1, a ∈ A, ŝ′ ∈ Sh, T ◦h (θ(ŝ′) | θ(ŝ), a) = { T (θ(ŝ′) | θ(ŝ), a), if θ(ŝ) ∈ Sh−1 T̂h(ŝ ′ | ŝ, a), otherwise
We also define the state space domain of T ◦h as S◦h−1 = {θh−1(ŝ) | ∀ŝ ∈ Ŝh−1}. It is easy to see that θh−1 represents a bijection between Ŝh−1 and S◦h−1. We will derive our guarantees with respect to T ◦ which will allow us to define a bijection between the domain of T̂ and T ◦, and use important lemmas from the literature. The next result shows that our use of T ◦ is harmless as it assigns the same probability as T to any event.
Lemma 22 (Closure Result). Let T ◦ be the closure of transition model with respect to some learned transition model. Then for any policy π ∈ Π and any event E which is a function of an episode sampled using π, we have Pπ(E ;T ) = Pπ(E ;T ◦), where Pπ(E ;T ′) denotes the probability of event E when sampling from π and using transition model T ′.
Proof. The proof follows form observing that when using T ◦ we will never reach a state s 6∈ Sh−1 for any h − 1 by definition of Sh−1. From definition of T ◦ this means that both T ◦ and T will generate the same range of episodes sampled from π and will assign the same probabilities to them. As E is a function of an episode, therefore, its probability remains unchanged.
With the definition of closure, we are now ready to state our last result in this section, which bounds the total variation between the estimated model and the transition closure under the bijection map θ.
Theorem 13 (Model Error). For any ŝ ∈ Ŝh−1 and a ∈ A we have: ∑
ŝ′∈Ŝh
∣∣∣T̂h(ŝ′ | ŝ, a)− T ◦h (θ(ŝ′) | θ(ŝ), a) ∣∣∣ ≤ 6d (∆est(nest, δest) + ∆app) .
Proof. If θ(ŝ) 6∈ Sh−1 then by definition T ◦ the bound holds trivially. Therefore, we focus on θ(ŝ) ∈ Sh−1 for which T ◦ = T . We define the quantity for every j ∈ [d]:
Sj = ∑
ŝ′[j]···ŝ′[d]∈{0,1}
∣∣∣∣∣∣ d∏
i=j
T̂hi(ŝ ′[i] | ŝ[p̂t(i)], a)−
d∏
i=j
Ti(θ(ŝ ′)[i] | θ(ŝ)[pt(i)], a) ∣∣∣∣∣∣ (30)
We claim that Sj ≤ 6(d− j + 1)(∆est + ∆app) for every j ∈ [d]. For base case we have:
Sd = ∑
ŝ′[d]∈{0,1}
∣∣∣T̂hd(ŝ′[d] | ŝ[p̂t(d)], a)− T (θ(ŝ′)[d] | θ(ŝ)[pt(d)], a) ∣∣∣ ≤ 6(∆est + ∆app),
from Theorem 12. We will assume the induction hypothesis to be true for Sk for all k > j. We handle the inductive below with triangle inequality:
Sj ≤ ∑
ŝ′[j]···ŝ′[d]∈{0,1}
d∏
i=j+1
T̂hi(ŝ ′[i] | ŝ[p̂t(i)], a)|T̂hj(ŝ′[j] | ŝ[p̂t(j)], a)−
T (θ(ŝ′)[j] | θ(ŝ)[pt(j)], a)|+ ∑
ŝ′[j]···ŝ′[d]∈{0,1} T (θ(ŝ′)[j] | θ(ŝ)[pt(j)], a)|
d∏
i=j+1
T̂hi(ŝ ′[i] | ŝ[p̂t(i)], a)−
d∏
i=j+1
T (θ(ŝ′)[i] | θ(ŝ)[pt(i)], a)|
The first term is equivalent to ∑ ŝ′[j]∈{0,1} |T̂hj(ŝ′[j] | ŝ[p̂t(j)], a) − T (θ(ŝ′)[j] | θ(ŝ)[pt(j)], a)| which is bounded by 6(∆est + ∆app) following base case analysis. The second term is equivalent to Sj+1 which is bounded by 6(d − j)(∆est + ∆app) by induction hypothesis. Combining these two bounds proves the induction hypothesis and the result then follows from bound for S1.

heading: D.6 LEARNING A POLICY COVER
text: In this section we show how we learn the policy cover. We start by defining some notation.
Two MDPs. After time step h, we can define two Markov Decision Processes (MDPs) at this time Mh and M̂h. Mh is the true MDP consists of state space (S◦1 , · · · ,S◦h), action space A, horizon h, a deterministic start state s1 = {0}d, and transition function T ◦t : S◦t−1 × A → ∆(St) for all t ∈ [h]. Recall that the set Sh ⊆ {0, 1}d denote states which are reachable at time step h, and the set
St ⊆ S◦t ⊆ {0, 1}d represents the closure of state space Sh based on the learned state space Ŝt. For any t ∈ [h], s ∈ St and K ∈ C≤2κ([d]) we know supπ∈ΠNS Pπ(st[K] = s[K]) ≥ ηmin.
The second MDP M̂h consists of the learned state space (Ŝ1, · · · , Ŝh), action space A, horizon h, a deterministic start state ŝ1 = {0}d, and transition function T̂t : Ŝt−1 ×A → ∆(Ŝt). For every t ∈ [h], we have θt : Ŝt → S◦t represent a bijection from the learned state space to the closure of the set of reachable states at time step t. The learned decoder φ̂t predict θt(s) given s ∈ St with high probability for all t < h by IH.2 and for t = h due to Corollary 9.
Lastly, the transition model T ◦t and Tt are close in L1 distance for t < h due to IH.3 and for t = h due to Theorem 13.
These results enable us to utilize the analysis of Du et al. (2019) for learning to learn a policy cover.
Let ϕ̂ : Ŝ → A denote a non-stationary deterministic policy that operates on the learned state space. Similarly, ϕ : S◦ → A denote a non-stationary deterministic policy that operates on the real state. We denote ϕ̂ = ϕ ◦ θ if for every ŝ ∈ Ŝ, ϕ̂(ŝ) = ϕ(θ(ŝ)). Similarly, we denote ϕ = ϕ̂ ◦ θ−1 if for every s ∈ S◦, ϕ(s) = ϕ̂(θ−1(s)). Let π : X → A be a non-stationary deterministic policy operating on the observation space. We say π = ϕ̂ ◦ φ̂ if for every x ∈ X we have π(x) = ϕ̂(φ̂(x)). Similarly, we define π = ϕ ◦ φ? if for every x ∈ X we have π(x) = ϕ(φ?(x)). We will use Pπ[E ] to denote probability of an event E when actions are taken according to policy π : X → A. We will use Pϕ[E ] to denote the probability of event E when we operate directly on the real state and take actions using ϕ. Similarly, we define Pϕ̂[E ] to denote the probability of event Ê when we operate on the learned state space. Lastly, let P̂ϕ̂[E ] denote probability of an event E when actions are taken according to policy ϕ̂ operating directly over the latent state and following our estimated transition dynamics T̂ : Ŝ × A → ∆(Ŝ). Recall that our planner will be optimizing with respect to P̂ϕ̂[E ]. Theorem 14 (Planner Guarantee). Fix ∆pl ≥ 0, h ∈ [H]. Let I ∈ C≤2κ([d]) and ŵ ∈ {0, 1}|I|. We define a reward function R : Ŝ → [0, 1] where R(ŝ) := 1{τ(ŝ) = h ∧ ŝ[I] = ŵ}. Let ϕ̂R = planner(T̂ , R, h,∆pl) be the policy learned by the planner. Let π̂ := ϕ̂R ◦ φ̂ then:
Pπ̂ (sh[I] = θ(ŵ)) ≥ η(sh[I] = θ(w))− 2d%H − 12dH∆est − 12dH∆app −∆pl, (31) further, we have:
P̂ϕ̂R({ŝh[I] = ŵ}) ≥ η(sh[I] = θ(w))− 6dH∆est − 6dH∆app −∆pl, (32) and if {sh[I] = θ(ŵ)} is unreachable, then
P̂ϕ̂R({ŝh[I] = ŵ}) ≤ 6dH∆est + 6dH∆app. (33)
Proof. We define two events E := {sh[I] = θ(ŵ)} and Ê := {ŝh[I] = ŵ}. We define a policy ϕR = ϕ̂R ◦ θ−1 where for every s ∈ S we have ϕR(s) = ϕ̂R(θ−1(s)). We also define π̄ : X → A as π̄(x) = ϕR ◦ φ?(x). If for a given x ∈ X and φ?(x) = s we have φ̂(x) = θ−1(s) then π̄(x) = ϕ̂R(φ̂(x)) = π̂(x). Hence, every time our decoder outputs the correct mapped state θ(s), policies π̄ and π̂ take the same action. We use the result of Du et al. (2019) stated in Lemma 30 (setting ε set to d% using Corollary 9) to write:
|Pπ̂(E)− Pπ̄(E)| = |Pπ̂(E)− PϕR(E)| ≤ 2d%H (34)
Let ϕ : S◦ → A be any policy on real state space and let ϕ̂ : Ŝ → A be the induced policy on learned state space given by ϕ̂(ŝ) = ϕ ◦ θ(ŝ) = ϕ(θ(ŝ)) for any ŝ ∈ Ŝ. We showed in Theorem 13 that T̂ and T have small L1 distance under the bijection θ. Therefore, from the perturbation result of Du et al. (2019) stated in Lemma 31 we have:
∑
sh∈S◦h
∣∣∣P̂ϕ̂(θ−1(ŝh))− Pϕ(sh) ∣∣∣ ≤ hε ≤ Hε,
where ε := 6d (∆est(nest, δest) + ∆app) due to Theorem 13. As {sh[I] = θ(ŵ)} ⇔ {ŝh[I] = ŵ}, therefore, we can derive the following bound:
∣∣∣Pϕ(E)− P̂ϕ̂(Ê) ∣∣∣ = ∣∣∣∣∣∣ ∑
sh∈S◦h;sh[I]=θ(ŵ) Pϕ(sh)−
∑
sh∈S◦h;sh[I]=θ(ŵ) P̂ϕ̂(θ−1(sh))
∣∣∣∣∣∣
≤ ∑
sh∈S◦h
∣∣∣Pϕ(sh)− P̂ϕ̂(θ−1(sh)) ∣∣∣ ≤ Hε (35)
Let ϕ? = arg maxPϕ[E ] be the optimal policy to satisfy {sh[I] = θ(w)}. Note that ϕ? is also the latent policy that optimizes the reward function R on the real dynamics. Let ϕ̂? = ϕ? ◦ θ−1 be the induced policy on learned states. We now bound the desired quantity as shown:
Pπ̂(E) ≥ PϕR [E ]− 2d%H (using Equation 34) ≥ P̂ϕ̂R(Ê)− 2d%H −Hε (using Equation 35) ≥ P̂ϕ̂?(Ê)− 2d%H −Hε−∆pl (ϕR is ∆pl-optimal on T̂ ) ≥ Pϕ?(E)− 2d%H − 2Hε−∆pl (using Equation 35) = η(sh[I] = θ(w))− 2d%H − 2Hε−∆pl.
This proves Equation 31 and Equation 32. Note that our calculations above show:
P̂ϕ̂R(Ê) ≤ PϕR [E ] +Hε If {sh[I] = θ(ŵ)} is unreachable then PϕR [E ] = 0. Plugging this in the above equation proves Equation 33 and completes the proof.
D.7 WRAPPING UP THE PROOF FOR FactoRL
We are almost done. All we need to do is to make sure is to set the hyperparameters and verify each induction hypothesis. We first set hyperparameters.
Setting Hyperparameters. Let {sh[I] = θ(ŵ)} be reachable for some I ∈ C≤2κ([d]) and ŵ ∈ {0, 1}|I|. Then applying Theorem 14 and using the definition of ηmin we have: Pπ̂ (sh[I] = θ(w)) ≥ ηmin − 2d%H − 12dH∆est − 12dH∆app −∆pl As we want the right hand side to be at least αηmin we divide the error equally between the three terms. This gives us:
(Planning Error) ∆pl ≤ (1−α)ηmin/4 (36)
(Model Approximation Error) ∆app ≤ (1−α)ηmin/48dH ⇒ % ≤ α(1− α)η2min
240κdHN (37)
(Model Estimation Error) ∆est ≤ (1−α)ηmin/48dH
⇒ nest ≥ 18432 2κd2H2N |A| α(1− α)2η3min
ln2 (
4e|A|(ed)2κ δest
) (38)
(Decoding Error) % ≤ (1−α)ηmin/8dH (39)
The model approximation error places a more stringent requirement on % than the decoding error for planning. However, throughout the proof for FactoRL in this section, we made other requirements on our hyperparameters. For % this is given by min{β2min/1200, 1/2} = β2min/1200 by combining constraints in Lemma 8 and Theorem 7, and an additional constraint for detecting non-degenerate factors stated in Corollary 9. Due to the inefficiency of the non-degenerate factors detection, we state results separately for the two cases:
% ≤ min { β2min 1200 , α(1− α)η2min 240κdHN } (no non-degenerate factor) % ≤ min { β2min 1200 , α(1− α)η2min 240κdHN , α2η2minδabs 30N2|A|2 log −1 ( 4
δabs
)} (general case)
Using the definition of % from Theorem 7, we get a value of nabs for non-degenerate factor (Equation 40) and general case (Equation 41) given below:
nabs ≥ 38402N4|A|2 α4η4minσ 2 ln ( |G| δabs ) max { κ2d2H2N2 α2(1− α)2η4min , 25 β4min } (40) nabs ≥ 38402N4|A|2 α4η4minσ 2 ln ( |G| δabs ) max { κ2d2H2N2 α2(1− α)2η4min , 25 β4min , N4|A|4 α4η4minδ 2 abs ln2 ( 4 δabs )} (41)
Recall that for detecting degenerate factors we collect ndeg samples. Corollary 9 gives value of this hyperparameter as
ndeg = 3N2|A|2 α2η2min log
( 4
δabs
) ,
which also satisfies the condition in Lemma 16. Lastly, Theorem 3 gives number of samples for independence testing nind and rejection sampling frequency k as:
nind ≥ O ( 1
β4min ln m2|A||F|(2ed)2κ+1 δind
) , k ≥ 8
ηmin ln
( 30
βmin
)
Failure probabilities for a single timestep are bounded by δind due to identification of emission structure (Theorem 3), 3dδabs due to decoding (Corollary 9), and δest due to model estimation (Lemma 17). The total failure probability using union bound for a single step is given by δind + 3dδabs + δest, and for the whole algorithm is given by δindH + 3dδabsH + δestH . Binding δindH 7→ δ/3, 3dδabsH 7→ δ/3, δestH 7→ δ/3, gives us total failure probability of δ and the right value of hyperparameters.
Sample complexity of FactoRL is at most kHnind +Hnabs +Hndeg +Hnest episodes which is order of:
poly { d16κ, |A|, H, 1
ηmin ,
1 δ , 1 βmin , 1 σ , lnm, ln |F|, ln |G|)
} ,
where use the fact that N = |Ψh−1| can be at most 2(ed)2κ from Lemma 23. Note that if we did not have to apply the expensive degeneracy detection step, then we would get logarithmic dependence on 1/δabs. Cheaper ways of detecting degeneracy can, therefore, significantly improve the sample complexity.
We have not attempted to optimize the degree and exponent in the sample complexity above.
For our choice of two hyperparameters nest and nabs, we can bound the model error and decoding failure by:
(Model Error) 6d(∆mod + ∆app) ≤ (1− α)ηmin 4H , (Decoding Failure) % ≤ α(1− α)η 2 min 240κdHN .
Verifying Induction Hypothesis. Finally, we verify the different induction hypothesis below.
1. We already verified IH.1 with Theorem 3. We learn a ĉhh that is equivalent to chh upto label permutation.
2. We already verified IH.2 with Corollary 9. Given a real state s ∈ Sh, our decoder outputs the corresponding learned state with high probability. We also derived the form of %.
3. We already verified IH.3 with Theorem 13. We also derived the form of ∆est and ∆app. 4. Lastly, Theorem 14 and our subsequent calculations for hyperparameter show that Ψh is
an α-policy cover of Sh and that the size of Ψh is at most 2(ed)2κ from Lemma 23. Lastly, for all reachable factor values we get the value of learned policy as at least (1+α)ηmin/2 using Equation 32 and our choice of hyperparameter values. Similarly, from Equation 33 we get the value of learned policy for all unreachable factor values as at most (1−α)ηmin/4. This allows us to filter all unreachable factor values. In the main paper, we focus on the value of α = 1/2, which explains why on Algorithm 1, line 8 we only keep those policies with value at least (1+α)ηmin/2 = 3ηmin/4. This verifies IH.4.
This completes the analysis for FactoRL.

heading: E SUPPORTING RESULT
text: Lemma 23 (Assignment Counting Lemma). For a given k, d ∈ N and k ≤ d, the cardinality of the set {(K, u) | K ∈ C≤k([d]), u ∈ {0, 1}|K|} is bounded by 2(ed)k.
Proof. Assume k ≥ 2. The cardinality of this set is given by∑ki=0 ( d i ) 2i which can be bounded as shown below: k∑
i=0
( d
i
) 2i = 1 + 2d+ k∑
i=2
( d
i
) 2i ≤ 1 + 2d+ k∑
i=2
( ed
i
)i 2i ≤ 1 + 2d+ k∑
i=2
(ed) i <
k∑
i=0
(ed) i .
The first inequality here uses the well-known bound for binomial coefficients ( n i ) ≤ ( ed i )i for any n, i ∈ N and i ≤ n. Further bounding the above result using ed− 1 ≥ ed/2 gives us: k∑
i=0
(ed) i ≤ (ed)
k+1
ed− 1 ≤ 2(ed) k.
The proof is completed by checking that inequality holds for k < 2.
Lemma 24 (Lemma H.1 in Du et al. (2019)). Let u, v ∈ Rd+ with ‖u‖1 = ‖v‖1 = 1 and ‖u−v‖1 ≥ ε. Then for any α > 0 we have ‖αu− u‖1 ≥ ε2 . Lemma 25 (Chernoff Bound). Let q be the probability of an event occurring. Then given n iid samples with n ≥ 1q , the probability that the event occurred at least once is at least 1− 2 exp( −qn 3 ).
Proof. Let Xi be a 0-1 indicator denoting if the event occurred or not, and let X = ∑n i=1Xi. We have E[Xi] = q and E[X] = qn. Let t = 1− 1/qn. We will assume that qn > 1 and so t ∈ (0, 1). Then the probability that the event never occurs is bounded by:
P(X < 1) = P(X < (1− t)qn) ≤ exp (−qnt2
3
) < 2 exp { −qn
3
} .
Lemma 26 (Hoeffding’s Inequality). LetX1, X2, · · · , Xn be independent random variables bounded by the interval [0, 1]. Let empirical mean of these random variables be X = 1n ∑n i=1Xn, then for any t > 0 we have: P( ∣∣X − E[X]
∣∣ ≥ t) ≤ 2 exp(−2nt2). Lemma 27 (Theorem 2.1 of Weissman et al. (2003)). Let P be a probability distribution over a discrete set of size a. Let Xn = X1, X2, · · · , Xn be independent identical distributed random variables distributed according to P . Let P̂Xn be the empirical probability distribution estimated from sample set Xn. Then for all > 0:
P(‖P − PXn‖1 ≥ ) ≤ (2a − 2) exp ( −n 2
8
) .
The next result is a direct corollary of Lemma 27. Corollary 15. For any m ≥ 8a 2 ln (1/δ) samples, we have ‖P − PXn‖1 < with probability at least 1− δ. Lemma 28. Let X1, X2, · · · , Xn be 0-1 independent identically distributed random variables with mean µ. Let X = ∑n i=1Xi. Fix m ∈ N and δ ∈ (0, 1). If n ≥ 2mµ ln ( e δ ) then P(X < m) ≤ δ.
Proof. This is a standard Chernoff bound argument. We have E[X] = nµ. Assuming n ≥ m/µ then from multiplicative Chernoff bound we have: P(X < m) = P ( X ≤ { 1− { 1− m
nµ
}} nµ ) ≤ exp ( −nµ
2
{ 1− m
nµ
}2) ≤ exp ( m− nµ
2
)
Setting right hand side equal to δ and solving gives us n ≥ 2µ ( m+ ln( 1δ ) ) which is satisfied
whenever n ≥ 2mµ ln ( e δ ) .
Lemma 29 (Lemma H.3 of Du et al. (2019)). For any a, b, c, d > 0 with a ≤ b and c ≤ d we have: ∣∣∣a b − c d ∣∣∣ ≤ |a− c|+ |b− d| max{b, d}
The next Lemma is borrowed from Du et al. (2019). They state their Lemma for a specific event (E = α(ŝ) in their notation) but this choice of event is not important and their proof holds for any event.
Lemma 30 (Lemma G.5 of Du et al. (2019)). Let S = (S1, · · · ,SH) and Ŝ = (Ŝ1, · · · , ŜH) be the real and learned state space. We assume access to a decoder φ̂ : X → Ŝ and let φ? : X → S be the oracle decoder. Let θh : Ŝh → Sh be a bijection for every h ∈ [H] and θ : Ŝ → S where θ(ŝ) = θh(ŝ) for ŝ ∈ Ŝh. For any h ∈ [H] and sh ∈ Sh, we assume P(ŝh = θ−1h (sh) | sh) ≥ 1− ε, i.e., given the real state sh, our decoder (φ̂) will map it to θ−1h (sh) with probability at least 1− ε.
Let ϕ : S → A be a deterministic policy on the real state space and ϕ̂ : Ŝ → A be the induced policy given by ϕ̂(ŝ) = ϕ(θ(ŝ)) for every ŝ ∈ Ŝ. Let π, π̂ : X → A where π(x) = ϕ(φ?(x)) and π̂(x) = ϕ̂(φ̂(x)) for every x ∈ X . For every random event E we have:
|Pπ(E)− Pπ̂(E)| ≤ 2εH
Lemma 31 (Lemma H.2. of Du et al. (2019)). Let there be two tabular MDPsM and M̂. Let S = (S1, · · · ,SH) be the state space ofM and Ŝ = (Ŝ1, · · · , ŜH) be the state space of M̂. Both MDPs have a A be the action space of both MDPs and horizon of H . For every h ∈ [H], there exists a bijection θh : Ŝh → Sh. Let T : S ×A → ∆(S) and T̂ : Ŝ × A → ∆(Ŝ) be transition dynamics forM and M̂ satisfying:
∀h ∈ [H], a ∈ A, ŝ ∈ Ŝ, ∑
ŝ∼Ŝh
∣∣∣Th(θ(ŝ′) | θ(ŝ), a)− T̂h(ŝ′ | ŝ, a) ∣∣∣ ≤ ε
Let ϕ : S → A be a policy forM. Let ϕ̂ : Ŝ → A be the induced policy on M̂ such that for any ŝ ∈ Ŝ we have ϕ̂(ŝ) = ϕ(θ(ŝ)). Then for any h ∈ [H] we have:
∑
sh∈Sh
∣∣∣P̂ϕ̂(θ−1(ŝh))− Pϕ(sh) ∣∣∣ ≤ hε

heading: F EXPERIMENT DETAILS
text: We provide details for our proof of concept experiment below.
Modeling Details. We model F , used for performing independence test, using a single-layer feedforward network θF with Leaky ReLu non-linearity (Maas et al., 2013) and a softmax output layer. Give a pair of atoms x[u] and x[v], we concatenate these atoms and map it to a probability distribution over {0, 1} by applying θF . We implement the model class G for learning state decoder following suggestion of Misra et al. (2020). Recall that a function in G maps a transition (x, a, x̌) ∈ X ×A×X ? to a value in [0, 1]. We first map x and x̌ to vectors v1 and v2 respectively, using two separate linear layers. We map the action a to its one-hot vector representation 1a. We map the vector v2 to a probability distribution using the Gumbel-softmax trick (Jang et al., 2016), by computing qi ∝ exp(v2[i] + ϑi) for all i ∈ {1, 2}, where ϑi is sampled independently from the Gumbel distribution. We concatenate the vectors v1,1a and q and map it to a probability distribution over {0, 1}, through a single layer feed-forward network θG with Leaky ReLu non-linearity. We recover a decoder φ from the model that maps a set of atoms x̌ to φ(x̌) = arg maxi∈{0,1} qi+1.
Learning Details. We train the two models using cross-entropy loss. Formally, given a dataset Dind = {(xi[u], xi[v], yi)}nindi=1 for performing independence testing, and a dataset Dabs = {(xi, ai, x̌i, yi)}nabsi=1 for learning a decoder, we optimize the model by minimizing the cross-entropy loss as shown below:
f̂ = arg max f∈F
1
nind
nind∑
i=1
ln f(y | xi[u], xi[v]), ĝ = arg max g∈G
1
nabs
nabs∑
i=1
ln g(y | xi, ai, x̌i).
Here we overload our notation to allow the output of models to be distribution over {0, 1} rather than a scalar value in [0, 1] as we assumed before. This is in sync with how we implement these model class and allows us to conveniently perform cross-entropy loss minimization.
Planner Details. We use a simple planner based on approximate dynamic programming. Given model estimate, reward function and a set of visited learned states, we perform dynamic programming to compute optimal Q-values for the set of visited states. We assume the Q-values for non-visited states to be 0. This allows us to compute Q-values in a computationally-efficient manner. Later, if the agent visits an unvisited state, then it simply takes random action.
Hyperparameters. We set the hidden dimension of θF to 10 and that of θG to 56. We set the threshold c on held-out log-loss value, when performing independence test to be 0.65. For reference, if one uses a random uniform classifier then its performance is − ln(0.5) ≈ 0.69. We train both models for 10 epochs using Adam optimization with learning rate of 0.001, and a batch size of 32. We remove 0.2% of the training data and use it as a validation set. We evaluate on the validation set after every epoch, and use the model with the best performance on the validation set. We used PyTorch 1.6 to develop the code and used default initialization scheme for all parameters.

Review summaries:
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.