Provable Rich Observation Reinforcement Learning with Combinatorial Latent States
Authors: Dipendra Misra
Qinghua Liu
Chi Jin
John Langford
Venue: ICLR 2021 Poster (2021)
Keywords: Reinforcement learning theory
Rich observation
Noise-contrastive learning
State abstraction
Factored MDP

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.

Published as a conference paper at ICLR 2021
PROVABLE RICH OBSERVATION REINFORCEMENT LEARNING WITH COMBINATORIAL LATENT STATES
Dipendra Misra' Microsoft Research
Qinghua Liu Princcton University
Chi Jin
John Langford
Princeton University
Microsoft Research
ABSTRACT
We propose novel setting for reinforcement learning that combines two com- mon 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 obser- vations 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 exponential in the number of latent factors. We create a learning algo- rithm 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.
INTRODUCTION
Most reinforcement learning (RL) algorithms scale polynomially with the size of the state space, which is inadequate for many real world applications. Consider for example simple navigation task in a room with furniture where the set furniture pieces and their locations change from episode to episode. If we crudely approximate the room grid to contain a single bit of information about the presence of furniture, then we end up with a state space of size 2000 as each element of the grid can be filled independent of others. This is intractable for RI. algorithms that depend polynomially on the size of state space.
10 x 10 grid and consider each element in the
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 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 generated by small number of latent factors with small
more generally, observation atoms) 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., logarithnically) 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 al. [0119] Misra et al. 2020 Agarwal et al.) 2020). However, addressing the rich-observation setting with 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 mimi i @mi crosoft.comm com

---

Published as a conference paper at ICLR 2021
Room Navigation Task
Identity Emission Structure
Learn State Decoder
Learn Model and Policy Cover
State Foctore
Obeanction Atare.
f M
+19]
Yipu trained to reach (
[1L]
Desoder Motil
=|M]
[0,
+|2
a|3
[[4]
Figure 1: Left: A room navigation tasks as a Factored Block MDP setting showing atoms and factors. Center and Right: Shows the different stages executed by the FactoRL algorithm. We do not show the observation x emitted by s for brevity. In practice a factor would emit many more atoms.
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 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 set of independence test problems with distributions satisfying certain properties. We perform independence tests in 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 I 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 Ao/Appediix B
To the best of our knowledge, our work represents the first provable solution to rich-observation RL with combinatorially large latent stute space.
2 THE FACTORED BLOCK MDP SETTING
There are many possible ways to add rich observations to a factored MDP resulting in inapplicability or intractability. Our goal here is to define 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:
For any n N, we use [n] to denote the set 21, ... ..... For any ordered set (or vector) U of size and an ordered index set I C [nx] and length k. we use the notation U[Z] to denote the ordered set [[[[[]],,[([22) ... [[][(:]]).
Notations:
Definition 1. A Factored MDP (S,A,T R, H) consists of a d-dimensional discrete state space S S {0,1]* a finite action space A, an unknown transition function T : S x A + A(S), an unknown reward function R. : S A + [0, 1] and time horizon H. Each state s E S consists of d factors with the 7th factor denoted as 8(1). The transition function satisfies T(s' 8, 8,a) =
I T((s'(+] s'p(())),) a) for every 8,8' E S and a E A, where T: : {00, 131pt(0) x A + ( 1))
defines factored transition distribution and a parent function pt [d] - 2 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, ,S) where S₁ is the set of states reachable time step h.
2

---

Published as a conference paper at ICLR 2021
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 (003). If we do not assume a factorized transition function then can encode arbitrary MDP with state space of size 24, which would yield a lower bound of 22(2)) rendering the setting intractable. Instead, we will prove sample complexity guarantees for FactoRL that scales in number of factors as do(") where k == [4] [p)(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 (.) combinations, as well as learn the model for all possible values of the parent factor. However, for real-world problems we expect k to be a small constant such as 2. This yields significant improvement, for example, if = 2 and d 100 then d 100 while 2 RJ 10
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 -.m and a latent state space S C {0, 13s. A single observation X is made of m atoms with the kth denoted by xrkk E X. Observations are generated stochastically given a latent state E S according to a
factored emission function q(> s) = I[= [[<][c((i] s(il) where q: :00,13 - A(2 (ch(8)) and ch : [d] + 2imll is child function salisfying ch((i)cch(j) = 0 whenever i + j. The emission function satisfies the disjointness properly: for every i E [d) we have supp(g,(- | 0)) n supp (ge(- 11) = of
The dynamics of the latent state space follows Factored MDP (S, A. 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 2 or its size, which can be infinite. An agent is responsible for mapping each observation x E X' individual atoms (x[[].... æ/m]) E 2n" " 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 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 plg and cha,
The disjointness property was introduced in Du et al. (2019) for Block MDPs-a class of rich- observation 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 decode the state from observation. For example, disjointness holds true for the navigation task with an overhead camera in Figure I| In this case, the image provides us with enough information to locate all objects in the room, which describes the agent's 's state.. Disjointness allows us to define decoder of : (ch(()) + f0, 1) for every factor E [d]. such that of (x(ch(t)) s|i] if x[ch(1)] E supp (9 s()))) We define a shorthand (z) = of (a(ch(i))) whenever ch is clear from the context. Lastly, we define the state decoder o² X + {0, 1)d where ox())[ = (z))
The agent interacts with the environment by taking actions according to policy X - A(A). These interactions consist of episodes {s1 Ø1, 01 F1, $2, D2, a2, r2, ,H SH) } with $1 = 0, In 4(1 8k), Fh- = R(I,, 4x), and 81+1 ~ T(- 8h; on). The agent never observes (81, ... (H
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 our setting.
Assumption 1 (Margin Assumption). For every h E 32.3.... ... H). i E [d]. let u, be the uniform distribution jointly over actions and all possible reachable values of 1h-1 [i)].. Then we assume: 8h (i]) is the back- ward dynamics denoting the probability over parent values and last action given St ii and roll-in
[PPa, 8h i) 1) - (., sn(i] = +)]lr 2 where Pu, (s)))))pt(()
distribution us and a > 0 is the margin.
'The notation supp(p) denotes the support of the distribution p. Formally, (ormally.supp(p) = [p(2) 0).
3

---

Published as a conference paper at ICLR 2021
Assumption 1 captures large set of problems, including all deterministic problems for which the value of a is I. Assumption Shelps 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 (Atom Dependency Bound). For any h [H), u, v E [m] and u + v. if ch (()) ch 2h-- (v), (r), i.e., atoms Ii [u] and xn []] have the same factor. Then under any distribution D E ((S)) we have ||TP p a [u], (.. (v)) - IP (Tr [u))TP D (#% [v]) 2 (mmin
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 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 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.
x 2 - [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 +b 2 - f0. 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' x A x 10. + [0, 1]. Finally, the regressor class 9 : X' x A x 2 + [0. 1| as { (xx, a, .) - u(x, a, x())) u E U, Ø E }). We assume F and 9 are finite classes and derive sample complexity guarantees which scale as log |F] and log [9]. However, since we only use uniform convergence arguments extending the guarantees to other statistical complexity measures such Rademacher complexity is straightforward. Let Ill S + A denote the set of all non-stationary policies of this form. We then define the class of policies II: X - by {ar + o0 (x)) Vi E I, which we use later to define our task. We use [x[8] to denote probability of an event & under the distribution over episodes induced by policy т.
Model Class. We use two regressor classes F and G. The first regressor class F 2
define
Computational Oracle. We assume access to two regression oracles REG for model classes F and G. Let D2 be a dataset of triplets (z(u], x0[l). v) where u, denote two different atoms and E {0. 1). Similarly, let D2 be a dataset of quads (Ir, a, x', y) where x E X, a E A, 3 E 2 and y E (0. 1). Lastly, let [p] denote the empirical mean over dataset D. The two computational oracles compute:
REG F)) =arg min ED, [[((fl)) sfll) w)²} ()G = min ED, (1(2,, a,,,,,,,, x) w)²] 9E9N
FEF
We also assume access to a Aptoptimal planning oracle planner. Let S = (S..... ,S, be a learned state space and 7 (Ti,... ,TH) with T, : S x A + ASS)) be the learned dynamics, and R: S x A + [0, 1] be a given reward function. Let s : S - be policy and V(p; T, R) be the policy value. Then for any ApI > the output of planner و = planner (T, R, Ap) satisfies Ve: R) 2 sup, (I: R) Apt. 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, show that controlling the value of most 2m factors is sufficient for learning the model. Let (()( denote the space of all sets containing at most l different elements selected from the set 21 including G. We define the reachability probability mn (KC, Z) for given h E [H]. K. S [d]. and Z E {0, 10.11... 1)1, and the reachability parameter Mmin as:
Th (K, Z) = sup (s [K] = 2), Ymmi == inf inf inf ma (KC, 8 s(K]). hE(M) SES, ((4))
=Ellss
Our sample complexity scales polynomially with mma Note that we only require that if s [A]] Z is reachable, then it is reachable with at least Mmn probability, i.c., either nh (K,, Z) = 0 or it is at least Mmin- These requirements similar those made by earlier work for non-factored state
4

---

Published as a conference paper at ICLR 2021
2019: Misra et al. 2020). The key difference being that instead of requiring every space (Du et al. state to be reachable with Mmii 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) then probability of visiting any state is 2- but the probability of two factors taking certain values is only 0.25. This gives more practical value for Tjmin-
Besides estimating the dynamics and learning decoder, we also learn an a-policy cover to enable exploration of different reachable values of factors. We define this below:
Definition 3 (Policy Cover). set of policies V is an a-policy cover of Sn for any a > and h if:
Vs E Sn, E (g2n((d),
sup (SA[K]] = s s|K]) 2 am (K, s|K]). FEV
3 DISCOVERING EMISSION STRUCTURE WITH CONTRASTIVE LEARNING
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 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. 1. Our goal is to infer the latent emission structure ch,, at time step h. which is equivalent to: given an arbitrary pair of atoms u and 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 = (u) and j ch (v)) ) be the factors that emit z|u] and z(uj. If j. then Assumption 2 implies that these atoms are dependent on each other for any roll-in distribution D E (()) x A)) over previous state and action. However, if i + j then deterministically setting the previous action and values of the parent factors pi(i) or pl()), makes x sr(u and x[v] independent. For the example in Figure 1 fixing the value of s {11), [2] and would make srfu] and ar(w²) independent of each other.
This observation motivates us to reduce this identification problem to performing independence tests with different roll-in distributions D E A(Sh-1 X A). Naively, we can iterate over all subsets K. E (6<2(([d) where for each K. we create a roll-in distribution such that the values of 8h [K] and the action aę-1 are fixed, and then perform independence test under this distribution. If two atoms are independent then there must exist IC 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 decoder at time step h 1. Further, 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 εx,Z = (1 their values Z, implicitly maximizes the probability of Ex,2 ( KK'] = 2') for corresponding real factors K"' and their values z'. Since the event Êx,Z is observable we can use rejection sampling to increase its probability sufficiently close to which makes the probability of Ex close to 1.
= 2} for a set of learned factors K. and
The second problem is to perform independence tests in a domain agnostic setting. Directly estimating mutual information I(x [xx[u]; æ(u)) 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 [ We comment that the high-level idea of our independence testing subroutine is similar to Seenn al. (2017). Suppose we want 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 (2 2() and (())) if w = 1, add (1) 20), w) to the dataset, and (, 22), w) otherwise. We repeat the sampling procedure n times and obtain a dataset {(y,, 24, w. } Then can fit classifier that predicts the value of w; using (3/. 2.. If Y and are independent, then (y,, 2)) will provide no information about w, and thus no classifier can do better than random guess. However, if Y and Z are dependent, then
5

---

Published as a conference paper at ICLR 2021
the Bayes optimal classifier would perform strictly better than random guess. As 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, 9, failure probability 8 > 0, and five hyperparameters 0, Imin, 9min E (0, 1) and d, K We use these hyperparamters to define three sample sizes Rind: Rabs; fest and rejection sample frequency k. For brevity, we defer the exact values of these constants to Appendix D.7 FactoRL
returns learned decoder Φh X' {0, 1]dd for some d4. [mn], an estimated transition model The learned parent ptg, and child functions cha, and a 1/2-policy cover n of S₁, for every time We
step h E (2.3. ..,,,, H). (or1 (1)), ... And, (zn)). In the analysis of FactoRL, we show that d4-1 = d, and ch₃, is equivalent to cha up to permutation with high probability. Further, we show that Φb, and cha together learn bijection between learned factors and their values and real factors and their values.
8/2 to denote the learned state time step h. Formally, $n,
FactoRL operates inductively over the time steps Algorithm line 8 In the 1 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, o, Tmin Smin, d, k). RL in Factored Block MDPs.
I: Initialize , = 0 for every h E [H] and Φ1 = X + {0}. Set global constants Fejnd, Taass Feests k.
2: for h {2.3. ... .H do
3: cha Factorrieemmsssioo((111111 2n- F) stage 1:
discover latent emission structure // stage 2: learn decoder for factors I/ stage 3: find latent pt, and estimate model
4: Gn LearnDecoder ch,1,
5:
Th ,t, = EstModel 1
6:
for I ((2) 2 E {0, 13 Σ do
7: GAIZ = pllnner RAIZ, Ap)) where RAIZ == [ 2} stage 4: planning
8: If V T, RhIz)= 2 3mmin/4 4 then V + V, U (GAZz @.)
return {che, oh, Th, pl,, Mn
Identifying Latent Emission Process. The FactorizeEmiss ss on collects a dataset of observations for every policy in 1 and action A (Algorithm 2 line 1. Policies in F are of the inductively assume πIZ to be maximizing the probability of &T.Z { [Z] 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 IP Pry,z ($%=1|Z) Z) is only O(pmin). Therefore, as explained earlier, we use rejection sampling to drive the probability of this event close to 1. Formally, we define procedure RejectSamp( TI; εr,z, k) which rolls-in at time step h - with TI;Z to observe - (line 3). the event ετ1æ holds for 2 then we return F4 otherwise, repeat the procedure. If we fail to satisfy the event & times then return the last sample. We use this to define our main sampling procedure Zn ~ DI.Z, := RejectSampp (a 2 εr,z, k) a which first samples F4-1 using the rejection sampling procedure and then takes action to observe Th. We collect dataset of observation pairs (zx(1), (2() sampled independently from Dr.z.a.
type NI.Z where I E (22n((dh-1) and 2 10, 1J We
For every pair of atoms u,v E [m]. we calculate if they are independent under the distribution induced by DI,2, using IndTest with dataset Dz,z. (line 7). We share the dataset across atoms for sample efficiency. If there exists at least one 2I a) triple such that we evaluate x z|u], 2 x[v] to be independent, then mark these atoms as coming from different factors. Intuitively, such an I would contain parent factors of at least ch (u) or ch,, (v). If no such I exists then we mark these atoms being emitted from the same factor.
"Our analysis can use any non-zero lower bound on Mmin Amin and upper bound d and N.

---

Published as a conference paper at ICLR 2021
Algorithm 2 FactorizeEEiissioon - 1h- F).
1: for (TI,2,0) E Wn-1 x A and [ndd do
2: Define εn2 == 1 3}
3: Sample Th2 th
(1) (2) SajjetSamp(TT:2:.. Ez,2, oa (2))
|| rejection sampling procedure I/ initialize DI,2,a 0
D1:2;0 + Dz U((a) (1) th
4:
5: for E (1, , - }} and E {u+ 1,,,, ; mmddo
6: Mark u, v as coming from the same factor, i.e., ch4, (w)= (v) if
7: the oraclized independence test finds on rn[u], vn [v] as dependent using Dz,zın and F
return cha,
// label ordering of parents does not matter.
Algorithm 3 LearnDecoder([, 1,, cha)-
Child function has type chy, [dr] 2
1: for in [dr). define w = cha (t), D do
2: for Raas times do
collect dataset of real (y 1) and imposter (y = transitions
3: Sample (1 a1)) 1))1),, (xx12) al2), ,2((1) Unf n-1) Unf (AA) and y ~ Bern())
4: If 1 then D <- DU (x(1), aa)) ('(D[() y) else D <- 'D u (x(1),, (1), (2) [w].
5:
train the decoder using noise-contrastive learning return Ø : X + {0, 13dhh where for any x E X' and i E [dr] we have [x) = (x(ch,, (i))).
Ui T. = REG
Decoding Factors. LearnDecoder partitions the set of atoms into groups based on the learned child function chin (Algorithm 3). For the ith group we learn a decoder Tu. 2 {0, 1) by adapting the prediction problem of Misra et al. (220)) to Factored Block MDP setting. We define a sampling
procedure (z, x²) Unf x-) Unf (AA) where is observed after roll-in with uniformly selected policy in - till time step h - 1, action a is taken uniformly, and x' ~ T(- ar, a) (ine 3).
We collect dataset DD of real and imposter transitions. A single datapoint in D is collected by
sampling two independent transitions (zx(1), (a (()),, (x) ((()),(()) .x(2) Unf ( o Unf (AA)
and a Bernoulli random variable ~ Bern(1/2). If = then we add the real transition
(x(1), aL)), [w]] y) to D, otherwise we add the imposter transition (x(2), al)), [W]] y) qline 4). The key difference from Misra et al. (2020) is our use x' [W] instead of x' which allows us to decode specific latent factor. We train a model to predict the probability that a given transition (xr a, x' (w]) is real by solving a regression task with model class (line 5). The bottleneck structure of G allows us recover decoder de 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 O, otherwise we stick with d;. For brevity, we defer the details of this special case to Appendix D.2.2 The decoder for the Ah timestep is given by composition of decoders for each group.
Algorithm 4 EstModel (Va
2n). I: Collect dataset D of nea triplets (x, a, x²) Unf( o (nf(A)
2: for I, E (gnh(())))) satisfying I J (1 do
3: Estimate ((( &
a) from D using G. Va E A, [da)-
4: For every & define pl,, (k) as solution to following: (where we bind s' = 84 and 8 = 8n-1)
argmin
max J. a22,a
s[Z] = u. [[] w1, a) ('('(T)) s|[] = u, 8[ u22 a)) TV
5: Define (8' 8) 1h. (((k) 5 (k),, a) and return Th.Pin
(k)),
Estimating the Model. EstModel routine first collects a dataset D of nest independent transitions (x., a, 2") ~ Unf ( o (A4) (Algorithm 4. [ine 1). We iterate over two disjoint sets factors
I, I of size at most K. We can view I as the control set and J as the variable set. For every
7

---

Published as a conference paper at ICLR 2021
learned factor & E [d,], factor set I. J and action E .. we estimate the model [() -/- -1 a) using count based statistics on dataset D
(line 3).
Consider the case where the ch, ch, for every t E [h] and where we ignore the label permutation for brevity. If I contains the parent factors pt(),, then expect the value of P( [k] 8 ( [[J]. a) x [3 [k] (lpt(k)]. a) to not change significantly on varying either the set J or its values. This motivates us 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[Z] and $[J] which can be reached jointly using a policy -n This is important since we can only reliably estimate the model for reachable factor values. The learned parent function pt, need not be identical to pt, 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 distinguish between them. We use the learned parent function pt, define Th similar to the structure of T (line 5).
Learning a Policy Cover. We plan in the latent space using the estimated model H to find a policy cover for time step h. Formally, for every I E (6x))(())) and Z E {0, 13 |z], we find a policy PPIz to reach (&x [Z] 2) ) using the planner (Algorithm 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 return of least (17in)) to v (line 8).
THEORETICAL ANALYSIS AND DISCUSSION
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 E [M], i [d] and distribution p E A{O 13), there exists Gihp E G, such for all V(x, a,,²) E Xx-1 x A x Mr and i = x' d'Ich (i)) we have:
(e (3) 6" (x).
(I (I, a, x) = (+) 6* (z), + p())
For any h E [H), u. D E [m] with u # v, and any D E ((S)), there exists fueD F satisfying:
[(
Vs E Supp(D), supp(al: s));
fuD(<cl)),(()) D(a[u) x D ((slu)
Assumption 3 requires the function classes be expressive enough represent optimal solutions for our regression tasks. Realizability assumptions are common in literature and are in practice satisfied by using deep neural networks (Sen al. 2017 Misra et al., 2020).
Theorem (Main Theorem). For any > O, FactoRL returns a transition function Thn a parent function chy, decoder Pn. and a set of policies Vn for every h. E (2.3.... HH,, that with probability at least 1 satisfies: (i) chy, is equal chy, upto permutation, (ii) V a 1/2 policy cover of Sn, (ii)) For every h E [H), there exists permutation mapping 0, 10, - {0, 1 1 such that for every E Sp-1, E A,s' E ESp, and x' E Xn, we have:
1-0 1-10
| sssa) as, (6)) 0 (S), a)
<
Taii 8H
KH
and the sample complexity is poly |AA,, H1 Timan 8 Bus
In m, In [F In (
Discussion. The proof and the detailed analysis of Theorem has been deferred to Appendix C [ Our guarantees show that FactoRL is able to discover the latent emission structure, learn decoder, estimate the model, and learn a policy cover for every timestep. We set the hyperparmeters to learn 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.
order
8

---

Published as a conference paper at ICLR 2021
Our results show that depend polynomially on the number of factors and only logarithmic in the number of atoms. This appeals to real-world problem where and m can be quite large. We also depend logarithmically 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 computational oracles. Hence, if these oracles can be implemented efficiently then our algorithm has polynomial computational complexity. The squared loss oracles are routinely used in practice, but planning in a fully-observed factored MDP is EXPTIME- -copleee (see Theorem 2.24 of Mausam (2012). However, various ap- proximation 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.
the size of function classes. This allows us to use exponentially large
Comparison with Block MDP Algorithms. Our work is closely related to algorithms for Block MDPs, which can be viewed as foctorccdd 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.) ) Feng et al. (2020) proposed another approach for solving Block MDPs. They assumne access to 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 al. (2019) and Misra et al. (002)) 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 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 dl factors each emitting 2 atoms. We generate atoms for factor s(i]. by first defining a vector 2; = [1,0] if s = 0, and 2 = =00, 1] otherwise. We then sample a scalar Gaussian noise 9. with mean and 0.1 standard deviation, and add it to both component of zj. Atoms emitted from each factor concatenated to generate a vector 2 E Id."d. The observation x is generated by applying a fixed time-dependent permutation 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,a,,... aa) and non-stationary dynamics. For each time step E [H], we define o₁ as fixed permutation of {1, 2, ... ,d. Dynamics at time step / are given by: (s 8t, a) II == (es(s141 [i] s a), where (s.+ [i] = seli] seli], a) 1 for all & (ox()). and To (5 [ = 1 - sc[i] [r] (o) = 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 9 we apply the Gumbel-soffnax 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 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 AAppendix F]
For each time step, we collect 20,000 samples and share them across all routines. This gives a sample complexity of 20, 000 x 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: / //ggggttuuu gighhub............... com/ coom/cereh cwreberreebbrreeee rl.
Conclusion. We introduce Factored Block MDPs that model the real-world difficulties of rich- observation environments with exponentially large latent state spaces. We also propose 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.
We can also easily modify our proof to use cross-entropy loss by using generalization bounds for log-loss (sce Appendix E. in Agarwal et al. 2020

---

Published as a conference paper at ICLR 2021
REFERENCES
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in Neural Information Processing Systems, 2020.
Mohammad Gheshlaghi Azar, lan Osband, and Rémi Munos. Minimax regret bounds for reinforce- ment learning. In International Conference on Machine Learning. 2017.
Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. In Advances in Neural Information Processing Systems, 2016.
Ronen I. Brafman and Moshe Tennenholtz. R-MAX- A general polynomial time algorithm for near-optimal reinforcement learning. The Journal of Machine Learning Research, 2002.
Richard Y Chen, Szymon Sidor, Pieter Abbeel, and John Schulman. UCB exploration via Q-
Ensembles. VwXiv:/706.01:: 2017.
Simon S Du, Akshay Krishnamurthy, Nan Jiang. Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient RI. with rich observations via latent state decoding. In International Conference on Machine Learning, 2019.
Fei Feng, Ruosong Wang, Wotao Yin, Simon S Du, and Lin Yang. Provably efficient exploration for reinforcement learning using unsupervised learning. Advances in Neural Information Processing Systems, 2020.
Carlos Guestrin, Relu Patrascu, and Dale Schuurmans. Agorithm-ddirectd exploration for model- based reinforcement learning factored mdps. In International Conference on Machine Learning. 2002.
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 2003.
He, Mari Ostendorf, Xiaodong Hc, Jianshu Chen, Jianfeng Gao, Lihong Li, and Li Deng. Deep reinforcement learning with combinatorial action space for predicting popular reddit threads. In Proceedings of the 2016 Conference on Empirical Methods Natural Language Processing, 2016.
Thomas Jaksch, Ronald Ortner, and Peter Auer. Ncar-optimal regret bounds for reinforcement
learning. Journal of Machine Learning Research, 2010.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzaioo with Gumbel-Softmax. In International Conference on Learning Representations, 2016.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contex- tual decision processes with low Bellman rank PAC-leamable. In International Conference on Machine Learning, 2017.
Chi Jin, Zeyuan Allen-Zhu Zhu, Sebastien Bubeck, and Michael I Jordan. Is q-learning provably efficient?
In Advances Neural Information Processing Systems, 2018
Sham Kakade. On the Sample Complexity of Reinforcement Learning. PhD thesis, Gatsby Computa- tional Neuroscience Unit, University College London, 2003.
Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In Interna- fional Joint Conference on Artificial Intelligence, 1999.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning polynomial time. Machine learning, 2002.
Hyoungseok Kim, Jackyeom Kim, Yeonwoo Jeong, Sergey Levine, and Hyun Oh Song. Emi: Exploration with mutual information. In International Conference on Machine Learning, 2019.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, arXiv preprint arXiv: 008880009999 2018.
10

---

Published as a conference paper at ICLR 2021
Akshay Krishnamurthy, Alekh Agarwal, and John Langford. PAC reinforcement learning with rich observations. In Advances in Neural Information Processing Systems, 2016.
Adrien Laversanne-Finot, Alexandre Péré, and Pierre-Yves Oudeyer. Curiosity driven exploration of learned disentangled goal spaces. arXiv preprint arXiv:180.001555 2018.
Lihong Li, Michael L Littman, Thomas J Walsh, and Alexander L Strehl. Knows what it knows: framework for self-awan learning. Machine learning, 2011.
Andrew L. Maas, Awni Y. Hannun, and Andrew Y. Ng. Rectifier nonlinearities improve neural network acoustic models. In ICML Workshop Deep Learning for Audio, Speech and Language Processing, 2013.
Andrey Kolobov Mausam. Planning with markov decision processes: an ai perspective. Morgan & Claypool Publishers, 2012.
Dorde Miladinovié. Muhammad Waleed Gondal, Bernhard Schölkopf, Joachim M Buhmann, and Stefan Bauer. Disentangled state space representations. arXiv preprint arXiv:1906.03255 2019.
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstrac- tion and provably efficient rich-observation reinforcement learning. In International Conference on Machine Learning. 2020.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hierarchical reinforcement learning. In International Conference on Learning Representations, 2019.
lan Osband and Benjamin Van Roy. Near-optimal reinforcement learning in factored mdps. In Advances in Neural Information Processing Systems, 2014.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity -riven exploration by self-supervised prediction. In International Conference on Machine Learning, 2017.
Rajat Sen, Ananda Theertha Suresh, Karthikeyan Shanmugam, Alexandros Dimakis, and Sanjay Shakkottai. Model-powered conditional independence test. In Advances in Neural Information Processing Systems, 2017.
Sahil Sharma, Aravind Suresh, Rahul Ramesh, and Balaraman Ravindran. Learning to factor policies and action-value functions: Factored action space representations for deep reinforcement learning. arXiv preprint arXiv: 1 700.0775577 2017.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint S.ti::2004.0113333 2020.
Alexander L. Strehl, Lihong Li, Eric Wiewiora, John Langford, and Michael L. Littman. PAC
model-free reinforcement learning. In International Conference on Machine Learning, 2006.
Alexander L Strehl, Carlos Diuk, and Michael L Lituman. Efficient structure learning in factored-state mdps. In Proceedings of the Twenty- AAAI Conference on Artificial Intelligence, 2010.
Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference Learning Theory, 2019.
Martin Sundermeyer, Ralf Schlüter, and Hermann Ney. Lstm neural networks for language modeling. In Thirteenth annual conference of the international speech communication association, 2012.
Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAl Xi Chen, Yan Duan, John Schulman, Filip DuTurck, and Pieter Abbeel. #Exploration: A study of count-based exploration for deep reinforcement learning. In Advances in Neural Information Processing Systems, 2017.
Valentin Thomas, Emmanuel Bengio, William Fedus, Jules Pondard, Philippe Beaudoin, Hugo Larochelle, Joelle Pineau, Doina Precup, and Yoshua Bengio. Disentangling the independently controllable factors of variation by interacting with the world. arXiv preprint arXiv: 1rXiv:1822.088 094844 2018.
11

---

Published as a conference paper at ICLR 2021
Tsachy Weissman, Erik Ordentlich, Gadiel Seroussi, Sergio Verdu, and Marcelo J Weinberger. Inequalities for the deviation of the empirical distribution. Hewlen- Packkrd Labs, Tech. Rep, 2003.
12

---

Published as a conference paper at ICLR 2021
APPENDIX ORGANIZATION
This appendix is organized as follows.
Appendix A
provides a list of notations used in this paper.
Appendix
covers related work
Appendix C
describes the independence test algorithm and its sample complexity guaranteess
Appendix provides sample complexity guarantees for FactoRL
Appendix provides list of supporting results
Appendix provides details of the experimental setup and optimization
A
NOTATIONS
We present notations and their definition in Table In general, calligraphic notations represent sets. All logarithms are with respect base e.
B RELATED WORK
There is rich literature sample-efficient reinforcement learning in tabular MDPs with small number of observed states (Brafman & Tennenholtz 2002; Strehl et al. 2006 Kearns Singh 2002 Jaksch et al.) 2010 Azar et al. [2017]Uin al. 2018). While recent state- oftte-of-the-eeee 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 rich-observation environments with possibly large number of factored latent states.
In order overcome this challenge, line of research has been focusing factored MDPs (Kearns & Koller 1999] Guestrin et al. 2002, [2003] Sttrhl et al.) 2010). which allow a combinatorial number of observed states with factorized structure. Planning in factored MDPs is EXPTIME-Complete MMusam 2012) yet often tractable in practice, with factored MDPs statistically learnable with polynomial samples in the number parameters that encode the factored MDP (Osband & Van Roy 2014 et al., 2011). There has also been several empirical works that either focus on the factored state space setting (e.g.,/Kin Kim & Mnih (018): Thomas al.) (2018): Laversanne- et al. (2018); Miladinović et al. (2019)), or the factored action space setting (e.g.. He et al. (201 20; /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 rich observation environment, Empirical results (Tang et al. 2017, Chen et al. 2017 Bellemare et al. 2016 Pathak et al. (017) have achieved inspiring performance on several RL benchmarks, while theoretical works (Krishnamurthy et al., 2016 Jiang
Chen
al. 2017) show that it is information-theorrttialll possible to explore these environments. As discussed hefore, recent works Du et al.] (2019), Misra et (2020) and Feng et al. (220)) provide computationally and sample efficient algorithms for Block MDP which is rich-observation setting with 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.. (.. 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 will be used in Algorithm 2 as subroutine for determining if two atoms emitted from the same latent factor. We comment that the high-level idea of Algorithm ( similar to Sen et al. (017). which reduces independence testing regression by adding imposter samples.
13

---

Published as a conference paper at ICLR 2021
Notation
Description
Space of natural numbers {1,
Space of integers {.... - 22 -1, 1, 0,
Zzo
equivalent to N U {0}. Given n E N, this notation denotes the set (1.2...., NN..
Space of positive integers 2.
[V]
(()(1)
Given ordered set U and k E Zza, this denotes the set of all ordered subsets of L with at most elements including the empty set.
w[Z]
For a given n-dimensional vector u and I = (i1 i2. ik) E 2 u[I] denotes the k-dimensional vector (u[i1). ulf2). u(ix))
[
Denotes concatenation of two ordered sets I and J Set of all possible distributions over the set U. Set of atoms
A (2))
2
x
Set of all observations
Xx.
Set of all observations reachable at time step
m
Number of atoms in the observation
Observation consisting of m atoms denoted by (((m)) Set latent states
Sn
Set of latent states reachable at time step / Number of latent factors
Latent state represented by vector in 13 Decoder function which maps an observation E X' to latent state E S. Decoder function which maps x E X to its 8th latent state value. Formally, Va E X' and [d], of (x) oz (x)
6"
6,
A
Action space
T:8 A A
Transition function on the latent dynamics Horizon of the problem denoting the maximum number of actions an agent takes single episode.
H
phy [d] 2(d]
Parent function for time step h. We drop h when is clear.
chy [d] 2(m]
Child function for time step h. By definition, for any i, E [d] and A
we have ch,, (i) n (c,()) = 0. We drop h when clear.
9: A(X)
Emission function that generates x g(: s) given
S.
T;
ith product term in the transition function. Formally, for s, e S and
E A we have T(s' tt product term in emission function. Formally for X with d (x:)
a) I[== (s' [d] s|pt(i)), a).
we have q== s) = I 9: (x]ch(i)] s(i))
F,9
Regressor classes. We will denote individual member of the class E and y E G Decoder class.
&: X + (0.1)
Table List of notations and their definition.
14

---

Published as a conference paper at ICLR 2021
C.1 ALGORITHM DESCRIPTION
Let D E A(S x A) be our roll-in distribution that induces a probability distribution PD ((X)) over observations at time step h. Let u, v [m/] be two different atoms, and Pp(x[u), afvi), ( [u]) and (D(xx(e)) be the joint and marginal distributions over vn [u], 2. [u] with respect to roll-in distribution D. The goal of our algorithm determine I' [u] and 2n.uu] are independent under PD.
Algorithm 5 IndTest(F, D, u, 3) Oraclized Independency Test. We initialize Dirain 0.
1: Initialize Dtrain 0 and sample 21, #2. in Bern (
2: for E [n] do
3:
if 2 =1 then
4:
Dirain + Derain ((x(h:1((u) su),,((.11110 (b). 1)}-
5: else
6:
Derain + Derain U (1l) [t]. 0)).
7:
Compute f= arg fEF L(D f), where
L(Pusiin: f) =
1
Σ
(f([[u) a[u]) =}²
Ti (()
8: return Independent if L f) > 0.25 - B²/103 else return Dependent.
We solve this task using the IndTest algorithm (Algorithm 5) which takes as input dataset of observation pairs D = (x((11) ) where o(51), [,(22] ~ PD(.,.), and scalar 8 E (0, 1). We use D to create dataset Perain of real and imposter atom pairs (z|u), =/o)). This is done by taking every datapoint in 'D and sampling Bernoulli random variable 2; Bern (1/2 (iiee 1}. If
= 1 then we add the real pair (c(().))u) (0) 1) to Derain (line 4p, otherwise we add the imposter pair (x(..) [u], [v), 0) (line 6). We train a classifier to predict if a given atom pair (ar(u), =(l)) real or an imposter (line 7p. The Bayes optimal classifier for this problem is given by:
P (p(x(u),<<v))
Vor E supp(Ppp), (b(x()),<(v)) == (p
((lu).z(v))
arfu and x xrw) are independent then we have (p PD (offu,,,x0u + PP( the support of Pp. This implies ( itively, this can interpreted as the classifier having no information to tell real samples from imposter 2
p ( (>fc)) everywhere on and its training loss will concentrate around 0.25. Intu-
samples. However, if 2/u), æ2vv are dependent and | (=[u) xxee) P
then we can show the training loss of f* is less than 0.25 O(8²) with high probability. The remainder of this section is devoted to a rigorous proof for this argument.
C.2 ANALYSIS ALGORITIMM
Before analyzing Algorithm 5 we want to slightly simplify the problem in terms of notations. We introduce two simple notations X and which represents the random variables æ|u] and a[o], respectively. We will simply use D to denote the joint distribution of X and Y. Define Drrain to be the distribution of the training data (Xtrains Yurain 2) produced in Algorithm It's easy to verify that
Derain (X Xtrain Yurain: 2) == JaD(X Ytrain) + (1 - 2) ((rrri))) DY
(XV) (((rrain))
(1)
Suppose the distribution D is specially designed such that at least of the following hypothesis holds (which can be guaranteed when we invoke Algorithm 5)
0
Ho |D(X, Y) - D(X) 2
g²
H1 : (D D(X)) - In the remaining part, we will prove that Algorithm 5 correctly distinguish between Ho and H with high probability.
103*
15

---

Published as a conference paper at ICLR 2021
C.2.1 Two PROPERTIES OF THE BAYES OPTIMAL CLASSIFIER
Our first lemma shows that the Bayes optimal classifier for the optimization problem line is constant function equal to 1/2 if X and Y are independent.
Lemma (Bayes Optimal Classifier for Independent Case). In AIloogiithmm 5 if X and Y (atoms and v) are independent under distribution D, then for the optimization problem in line 7 the Bayes optimal classifier is given by:
V(Xtraim: Yurain) (X Yurain)
2
Proof. From Bayes rule we have:
("(traanns Yurai))
=Duain(z = Xuran Yerain)
Durain (Xtrain: Yurain 2 == Drriin 1)
DDeaain Perain(Xrrain (Xtrains Yurain 1)Darain 1) Dirain (Xtrain: Ytrain 0) (=
Dirain (X Xtrain: Vtrain 1) Derain (Xtrain Yusin = 1) + Drrain (Xtraiu: Ytrain 0)¹
where the last identity uses Dtrain 1) Dirain (=- 0) = 1/2.
When z O, we collect Xtrain and Ytrain from two independent samples. Therefore, we have 0) D(Xural))(Yuuati
Derain (Xtrain Yirain
When = 1, using the fact that Xtrain and Yurain are independent under distribution D, we also have Drrain (Xtrain Xtrain 1) D(Xtrain: Ytrain) DDX
Consequently, f*(XXraan Yurain) = 1/2.
Our second lemma provides an upper bound for the expected training loss the Bayes Optimal Classifier. Later will use this lemma show the training loss is less than 0.25 O(82) with high probability when Hn holds.
Lemma (Square Loss of the Bayes Optimal Classifier). In Algorithm line 7 the Bayes optimal classifier has expected square loss
1
((X)D(Y)+D(X) D(X,Y) 1))²
D(X,Y)
ED L(f", D Derain)
Ep
Proof. Recall the formula of the Bayes optimal classifier in Lemma
D(X,Y) Y) ('X +
Plugging it into the square loss, we obtain
Epaaio
((*(XXY)) y) -)
= -Eperatn
- -1--0)--------0
-Epani
1 E
(
D(X,Y)
)
D(X) D(X,Y) )
4
<-
1
EDrnin
1/2 D(X)D(Y) +D(X,Y)
D(X.Y)
1)
D(X,Y)
2
EE
4
2
D(X)D(Y) + D(X,Y)
16

---

Published as a conference paper at ICLR 2021
C.2.2 THREE USEFUL LEMMAS
To proceed, we need to take a detour and prove three useful technical lemmas.
Lemma 3. Let μl and two probability measures defined on a countable X. If /1 v/|TV 2 c. then
p(aa)
1
Bzzp
2
u(xx+ +((x)
2
4
Proof.
p(x)
u(x) - v(a)
Exzyy
2 Ewwny
(xx + (2x)
2 u1xx) + u(x)
(u(x) - v(x) m(xx) + v(x)
2 Bznp
l(u(z) > v(x)}
2
Σ m(x) 1(((z) x(x)) (4xx) (10)) - v(x))
2
u(x) v(x)
TEX
[ (m(x) (x)) v(x)
aEX
Lemma 4. Fix 1). Then with probability at least
we have
C(F,8) L(f, Drvain) - - L(f", Diman)) $ 10 n
where Duran is the training set consisting of i.i.d. samples sampled from Dirain, f is the empirical minimizer of L(f, Dresin) over F, f is the population minimizer, and C(F, == In FF is the complexity measure of function class F.
Proof. By Hoeffding's 's inequality and union bound, with probability at least - 6, for every ʃ E F, we have
[L[, Derain) - Ep L(S, Derain)) S 101 C(F,0)
Because f is the empirical optimizer,
O(F,0)
L(J, Dirain) S L(J", Derain) S E Dreain L(5" Derain) + +10
n
Because f* is the population optimizer,
C(F,5)
C(F,5)
L(f, Dirain) EDimi L(f. Durain)
ED L(f", Durain) -10
n
Combining the two inequalities above, we finish the proof.
For notational convenience, we introduce the following factor distribution Diactor defined on the same domain of (X, Y))):
Dinctor (X,Y,z) 1 - D(X)D(Y)
Lemma 5. Suppose F contains the constant functtinn f = 1/2. Then with probability at least 1 - 8, we have
1 Mf. Durain)
C(F,6) 6)
s $10
+ 2 Durain D D I|Tv-
4
17

---

Published as a conference paper at ICLR 2021
Proof. By Lemma 4. with probability at least 6, we have
- 110 C(F,6)
T.J, Dtrain) D.p L Darain)
(2)
Noticing that L is bounded by 1, we have for every f E
E Druotor L(S, Parain) ED [JJ, Ptrain)] 2 Derain Diactor IV,, where Ey L(f, Drrain) defines the expected loss of f over Durain where Draain consists of samples i.i.d. sampled from Diactor-
(3)
Since y is symmetric Bernoulli n.v. independent of (X, under distribution Diactors we have
min E Drastor L(f. Derain) 1
feF
4
Using the inequality miin Li( miny 22(l) & max, |L - L2(f)l for any functionals Li, L2, along with and (3) bound the minimum loss under distribution Draain as:
min D L(f, Dtrain) FEF
D Diccoo ITV.
(5)
4
Combining (5 with completes the proof.
C.2.3 MAIN THEOREM FOR ALGORITHME
5
Finally, we are ready to state and prove the main theorem for Algorithm Theorem 2. Under the realizability assumption and 2 Q
C(F.A) 3
Algorithm
can distinguish
3
:o: Ho:1ID(X.Y) DXX)D(()) 2
2
g
U.S. H1 (1:1DDXX,Y) D(X)D(Y)I S
103
correctly with probability least 1 8.
Proof. If H1 holds, by Lemma 5 we have the following lower bound for the training loss of the empirical minimizer,
1
C(F,8)
g
L(fD) D) 4 10
(6)
103"
In contrast, if Ho true, applying Lemma 4. obtain
C(F,0)
L(J.D) < E E D) + 10
Invoke Lemma and Lemma 3
D(X,Y) D(X)D(Y) + D(X,Y) Y) 11)²
E. D M( D Derain) <
ED
4
B2
<
256
Therefore, under Ho, the training loss of the empirical minimizer is upper bounded below
32
+ 10 C(F,8)
L(f,D) S 4
256
Plugging 2 0 OCCCCCC Br
C(F
into
and (7)
complete the proof.
18

---

Published as a conference paper at ICLR 2021
D THEORETICAL ANALYSIS OF FactoRL
In this section, we provide detailed theoretical analysis FactoRL. The structure of the algorithm iterative making an inductive case argument appealing. We will, therefore, make 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 Assump- tion 13 and across all time steps. For all (2,3, H). at the end of time step AAlgo
rithm line 8). FactoRL finds a child function chę : [d] + 2(m), decoder Φ : X - {0. 13dd f
transition function Te {0, 11d x A - f0, 1] and a set of policies V, satisfying the following:
IH.1 che : [d] + 2m and ch₃ [d] + 2 are same upto relabeling, i., for all u, v E [m]] we have ch, (w) = ch, (v) if and only if ' (u) = ch, ((v). Note that child function is invertible by definition. We can ignore this label permutation and assume ch, che for cleaner expressions. This be done without any effect. We will assume chę- ch, when stating the next three induction hypothesis.
IH.2 There exists permutation mapping 0. {0, 1] - {0, 13d and e E (0, 2 such that for every E [d] and E St we have:
((O((1)]1) = 0r (s) s(1) 2 0.
(().(1)) (s) () 2 1 de
The two distributions are independent of the roll-in distribution at time step t. The first one holds as de (are) only depends upon the value x; (ch;, (f) xz (t)] which only depends s(i). The second one holds as is independent of everything else given 8t. The form of @ will become clear the end of analysis.
IH.3 For every s S.-1. St and a A we have:
T, (s). T(s' s,,0) Iyy 3d Aapp),
where Aest Aap denote estimation and approximation errors whose form will become clear at the end of analysis.
IH.4 For every Sz and K. E (622((d) let Z =(C] and 2 (s){(, then there exists policy xx E V such that:
Pake [sr = Z) 2 am, (KC, 2) 2 a7mmin
Base Case. In the base case (4=1), 1), we have a deterministic start state. Therefore, we can without loss of generality assume a single factor and define chi, [1] m. As we can also define ch, [1] [m] without loss of generality. therefore, this trivially satisfies the induction hypothesis We define Φ1 : X + [ojá Algorithm 1 line 1.. This satisfies induction hypothesis 2 with 0, being the identity map. The induction hypothesis 3 is vacuous since there is transition function before time step For the last condition, we have for any K., = [o]|k] and 2 = [o] For any policy π we have
Fa (s) [K]] = Z) = Pa (o (XC) 2) = 1 > Tmmin policy, therefore, we can simply define Vi
Note that we never take any action from this
2
D.1 GRAPH STRUCTURE IDENTIFICATION
In this section, we analyze the performance of Algorithm given as input Wn-1,ōn-1,33 B and
We will analyze the performance fixed pair atoms u, v E [m]] and then apply the full result using union bound. We first state the result for the rejection sampling.
Lemma 6. For policy πTZ² E Vn-1, event E = ( 2) and k N, let D Pejjctt I:
Re eejetStamppr..... E 22,k be the distribution induced by rejection sampling procedure. Let
19

---

Published as a conference paper at ICLR 2021
Z ((2 denote the real factor values corresponding Z. Then we have:
Z) 1 D (1
Mmin
P
D
(8h
2:2
Proof. From IH.4 we have
P xz,2 (8h-112 = 2) 2 Mm
This implies:
2z.,2 (
2)
P xz (6/- [ 2 86-1 = 2)I xz,2 (8h-1 = 2)
(1 de)mmmii
Mmén
(using
IH.2 and IH.4
Let a Fx (8)-1 2) be the acceptance probability of event ε. then it is easy see that
the probability of the event occurring under D Deject
is:
P Deiia (52 (fg,2= a (1- (1-
a)*
2 1-
Tman
&
4
We express the desired failure probability as shown:
P Dane (5h-1 [I] * Z)
IP
Drsisn 1
2, Sn
2
+P
Daisa 222
shhiiI # Z, 5 [Z]
2
We bound the two terms below:
TP Dripar
8/1
2,
8h
[I]
+
2
< PP
D 2
8h [I] + 2 s 5(1
7mnn
(10)
2
4
IP
8h [
+ 2, 84 I 2
8n [ 2 8h-1 [T]
E,
se
(11)
I
Combining Equation 9
Equation 10 and Equation 11
we get:
P [Z] = 2) = 1 PP
preea 7
Driher
(5)-1
+ 2)
2
(1 0
Mmin
(12)
4
We now analyze the situation for given pair of atoms. Recall for any distribution D (() and a E A. we denote D as the distribution over S₁, where s' ~ D a is sampled by sampling 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 Ho and atoms coming from different parents satisfy hypothesis H, under this roll-in distribution. This will allow us use independence test to identify the parent structure in the emission process. Specifically, we consider the roll-in distributions induced by Droject for some sets I, J and action a. Instantiating the definition of these hypothesis z from Appendixx C. with these roll-in distributions and setting 3 Smin gives us:
Amin
Hu prscet (x/4), x[e)) - pyeat oa (z/u)) paie (x(v))]12 >
2
.22
U.S. H1 : P Drss (zfu), z(o)) -
(z) (r(u])
Dries oa (xfo)))12 -
7.2
100
have eh 1 (u) 1 i.e., they are from the same factor then the hypothesis Ho is true for D o a for any E A(Sm-1) and a E A. In particular, this is true for D Devve o for any choice of sets I, I and action a.
Lemma (Same Factors), for two atoms u,
Proof. Follows trivially from Assumption
Lemma 8 (Different Factors). for two atoms u, v we have ch
ch (u) i and ch (1) = and
i + j. then I contains ptg, (i) U (t),)), then for Q S
82
and k: 2 8 In 30
the hypothesis
1200
Mmmn
Bren
H1 holds for Dreject a a for any 2 such that Tz E 1n- and a E A
20

---

Published as a conference paper at ICLR 2021
Proof. Let D' A(Sh-1) be distribution that deterministically sets 8h-1 2. Then it is easy to verify that oa (In, [u] 24 [,)e) = Pp
(an(e)) for any a E A and 2n E Zn.
Then for any E and action a E A we have using triangle inequality:
Drea 2 oa (In[ In [u])
Dain 2:2 00
(oh (2)(2))
Dr 00 (2
< D on (In[ (u])
D'oa (era (wn//// (n (nl)))
+
P (x - Dusi oa (Ir/ + Prroa(tN(l))
D ca (2
As mn[u] and zn [u] come from different factors, therefore, we have [u]. 24h [u]
PP
8h-1, a) P( [u] 86-1 (PPx 8/ (2),
We use this to bound the three terms in the summation above.
IP Drsin
(En[u],zn(V) - P Parea (2n (zn(u),
P(Zn[u] 9 9h-1 91-11 sh-1[Z],a) a) {
Σ
Dosu (sh-1 - P (s [T]
an/u).ana [v]
5h-
Σ
Σ
P( 9h-1 [Z). 1 (Z),a) (sn [Z] IP (s [T]
[2] 2a(u).2a11
D (sh (8h-1 (T] - pp (8 (8h-1 (Z)
Drsea (81-12 [I] = 2) +
Σ
Droea (8h-1 (Z])
8h-1
[
< 20 (1
Mmin k
2
pops (8 2) 2,2
The other two terms are bounded similarly which gives us:
(o s 6e 6 6 Mmin k
p 2 oa
[x) 2n [u])
pm 2.2 oa
2 We want this quantity be less than and use In(1 + a) & for all > -1 to get:
100 satisfy hypothesis H1. We distribute the errors equally
S 1200'
8 Jmmin
30 Brin
@ 5
(13)
Theorem 3 (Learning chi). Fix 8mn (0. If
1200
62 and
Venin
In
30 Bran
and
Pind
0
In F( then learned clip is equivalent cha upto label permatation with And
probability at least
8ind-
Proof. For any pair of atom u, v. if they are from the same factor then Ho holds from Lemmaa and IndTest mark them dependent with probability least 1 - 5. This holds for every triplet of I, Z, and there are at most (2ed)))xx*+ of them. Hence, from union bound we mark u, correctly as 5.
coming from different factors with probability at least
and have different factors then for any containing the parents of both of them, and any value of 2 and a, H, always holds from Lemma 8 and IndTest marks them as independent. Note that
21

---

Published as a conference paper at ICLR 2021
Time stop
stop
State Factors
Observation
C
Figure 2: Scheme showing the important variables for the decoding step.
such an will exists since the iterate over all possible sets of size upto 2x. Hence, with probability at least |A|2" 8, we find and to be independent for every 2 and a. Hence, our algorithm correctly will mark them coming from different factors.
For given u and we correctly predict their output with probability at least 1 - (2ed) (2 2x &. Therefore, using union bound we correctly output right result for each and with probability at least
1 (2ed) 2x+18. From Theorem 2 we require Wind
In [P Binding (_Al(2ed) 2x+18
to dind then gives us the required value of Ain 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 chy, which is same as ch,, upto label permutation. L.
D.2 LEARNING A STATE DECODER
We focus the task of learning an abstraction time step i using Algorithm We have access cha, which is same as ch,, upto label permutation. We showed how 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 (ΦrI,-.. Tha).. ( -
where there a bijection
between {ons), }. and {
As we learn each decoder (0ms) independently of each other, therefore, we will focus on learning the decoder The for a fixed i. The same analysis will hold for other decoder and with application of union bound, will establish guarantees for all decoders. Further, since 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
and x denote state $ and an observation 2h-1 at time step h 1
and x' denotes state 8y and observation Ty at time step
8 [i] denotes 1th factor of state at time step
8 denotes s(pta which is the set of parent factors of s'ij. Recall that from the factoriza- tion assumption, we have T(s s, a) (s [i] a) for any s, a.
de denotes Th. decoder for ith factor at time step h
w denotes pi) which the set of indices of atoms emitted by [i].
t denotes z' [chn, (i)] which is the collection of atoms generated by s'()..
N , is the size of policy cover for previous time step.
Let D = {x (k) (k), (M), (()) ( Paan be a dataset of naas real transitions (y 1) and imposter k=
transitions (y 0) collected in Algorithm 3 line 2
We define the expected risk minimizer (ERM)
solution as:
Tlais
Σ (ola (k))) (()))))))) -y y())
gs = arg min DEG Tlabs
(14)
22

---

Published as a conference paper at ICLR 2021
Recall that by the structure of G, have 9: (ûz, ,,) where úù, W2 and Φ; E + D {0. 1] the learned decoder. Our algorithm only cares about the properties of the decoder and we throw away the regressor ü,
Let D(x o, t) 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(a,, a) as the marginal distribution over 2; a. We have ((x, a) = Ph-1 (2) [ as both real and imposter transitions involve sampling ~ 14-1 and taking action uniformly. Recall that 14-1 is generated by roll-in with a uniformly selected policy in A- till time step h 1. Let P(x, a, y = 1) be the probability of a transition being real and P(x., o, y = 0) be the probability of the transition being imposter. We can express these probabilities as:
P(x, u, # y = 1) D(x, a) x, a, where p( = S(z.a)>([T(
(15) x, a)] the marginal distribution over æ. We will overload the notation p to also define p(x') = B( x, v.a)]. Lastly, we can express the marginal
P(x,a a = 0) Dax, a)p()),
distribution over transition as:
D(2, a, 8) = a yy I)P(w 1) Pla,a, a y- -)
(x) (T(E x, + *()}
2|A
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:
(T) 0 (z) (x) (T) (z) a) (z)
V(x, a) suppD,
(16)
Proof. The Bayes optimal classifier is given by g" (z, Ply xx x, which can be expressed using Bayes rule as:
P(x, a. 1 =)
Ply = 1|,,,,8) = P(r, a. y 1)P() (1 a P) = 0)
P(r,a, y=
using p(y) = Bern
= a(x,a, 3 y + P(z,0,
D(r, T( x.a)
D(r, a) x,a) + D(x,a)p(t))
T( (x,a)
T(x J., a) + p(t)
(1 (T)) a)
qi(E (+) x, + q. (T(3)))((((i)
()) x,
(T))(() of (x)(pt()).
((()(()) x, a) (())())) ((( of (z) (x) a) + (1)
Theorem 4 (Decoder Regression Guarantees). For any given Sabs E (0, 1) and Nabs N we have the following with probability at least 8absi
Exaa,nD
(0., a, (z, a. a)}}} < (lnabo Sabss 19)),
where (DDater Babes (91)
Øl and e is a universal constant.
Proof. This a standard regression guarantee derived using Bernstein's inequality with realizability 3 For example, Proposition 11 in Misra et al. (0200 for proof.
23

---

Published as a conference paper at ICLR 2021
Corollary 5. For any given dabs E (0, 1) and Raas N we have the following with probability at least babs:
Ex (1ā a, - (x, a, x)ll 2 V D Sabus (
(17)
Proof. Applying Jensen's inequality (L[V S [ to Theorem 4 gives us:
E a [10(2, a) 9 (z, E||l = L a) 90 9x,,
V Buaann (0(o,a, -gr gr ax, 8))
( Sabss (9).
Coupling Distribution
We introduce coupling distribution following Misra et al. (2020).
Decop (2, a, 11, .2) (2x, ())((1))((2).
(18)
We also define the following quantity which will be useful for stating results:
g( $2, x, a) = T( xx,a)
T(12 | x,a)
(19)
p(ā1))
p(i2))
Lemma 10. For any fixed Babs E (0, 1) we have the following with probability at least Sabsi
E2 B2r Domp (())(()) 5x (i2)) (E( $2, a)l
S 88 (Tabss Sabs: [G]).
Proof. We define shorthand notation & 1(0((z1) (122) for brevity. We also define
different coupled distribution D'cup given below:
D' (x, a, b2) = D(x, a) D(1 x2 D(r2 xr
(20)
where D( x. over x, 31 is same as D(I, a, 21).
(T 2, + p(t)}. It is easy to see that marginal distribution of Dooup
first use the definition (Equation 19 and g" [qquation 9 to express their relation:
p((1)p((2) [6 T2, a)l
Ig" az) (x. 0,322 (T x, + p( (T [1 x, a) p()2))
p(t1) ((ā2)
[ #2,
(21)
4D
xr, a) ( a))
The second line uses the definition D(Ā x, a). We can view p(81) and p(2g)
D(áa(æ,a) as importance
(2,))
weight terms. Multiplying both sides by & and taking expectation with respect to D Doup then gives us:
E E [E Ig (Ix, a, x1) g" (a, a, 22
4 ED..p
[81((11, a)ll
(22)
We bound the left hand side of Equation 22 as shown below:
EDico E ig" (x, a, T1) - g" (x, a, ]])
< EDjo E 1g (x, a, 01) (=(2, a, i1] + E IE (0: a, Æ1) - -----g-- g0 a, d2))
-Ep EE (19g (x, a, .1) (z(z, (,), Er, (E [911 a, .2) 9" (z, a #2)]
=2 Ig (s, 81) 01 a,), b1 2E [ Ig (x, a.),
2 1-2A((tate A(Tabos Satss 9)
24

---

Published as a conference paper at ICLR 2021
Here the first inequality follows from triangle inequality. The second step is key where we Ge(2, a, 21) = 9e (2, a, .22 whenever & = This itself follows from the bottleneck structure of g where (x, a; = (z, (;))). The third step uses the symmetry of 21 and Ø2 in Dooup whereas the fourth step uses the fact that marginal distribution of Damup is same as D. Lastly, final inequality uses & I and the result of Corollary 5 Combining the derived inequality with Equation 22 proves
our result.
We define the quantity P(s [1] - D`) B(o,a)= [1 = =}] for any distribution D' E AS x A). From the definition of p. we have p(s [i] = 2) = P(s' [i] = 2 D). Intuitively, as we have policy cover at time step h and we take actions uniformly, therefore,
expect to have
good lower bound on P( = = D) for every E [d] and reachable 2 E (0. 1). Note that if = = 0 (z 1) is not reachable then it means we always have sh [i] = (5%, [i] 0) from our reachability assumption (see (see|Section 2). We formally prove this next which will be useful later.
Lemma 11. For any 2 E {0, 1} such that s' s'i = z is reachable, we have:
p(s' [ 2) P(s
D) 2
a7hmin
N|
Proof. Fix 2 in (0. 13. As s' [i] is reachable, therefore, from the definition of Ammin we have:
I sup 1r, (s =) S sup Σ (5)T((5 = 5,a)
Tmin
*ell
8,a
Σ Sx 8,a) Σ ( 1,
TEN
We use the derived inequality to bound P(5'[i]] = = D) as shown:
P(s' = D) Σ
1h-1 T (s' = 5, a)
NJA Σ
(8))((s11) =3,
ammmn NIA
The first inequality uses the fact that 14-11 is created by roll-in with a uniformly selected policy in Vn-1 which is an a policy cover. Recall that N = // The second inequality uses the derived result above.
Lemma 12. For any 1i, d2 such that o (i1) and of (fz) is reachable, we have:
E Ez [IE(ă1, 82, x, a) 2 1 (#1) + () of]mnn
2N
Proof. For any 2, a, 11, 32 we can express {
Equation 19 as:
KK 02, a)l (z)) o² (x)[pt(i), a) (()) or (x)p(ti)) p()) (X))
((x)(12))
where we use the factorization assumption and decodability assumption. Note that are implicitly assuming (ă1) and of (ã2) are reachable, for the quantity {(i,, 22, x, a) to be well defined.
We define D, to be the marginal distribution over S Sip((i)) x A. Taking expectation on both side gives us:
(())((1) 3,a) (())(22) 3)
By [|5(0), 82, x. a)ll Bs
(())
(iz)
Σ PD. (+())) (8, (
5.a
2 Pp, (,,- (b))) - Fp. 1 (
The second equality uses the definition of backward dynamics Pp, over S(pt(i)) A and the
a identity p(s'[i)) = P( D). If (2,) (Ø))= = or (#z) then the quantity on the right is 0. Otherwise,
this quantity given by ( s Pp, s'( = Oll In the later case, both
25

---

Published as a conference paper at ICLR 2021
s'[ = and s'[ 0 configurations are reachable, and without loss of generality assume
P(s [1] = 0 D) 2 1/2. Our goal is to bound this term using the margin assumption (Assumption 1).
We do using importance weight as shown below:
2 2 (T- [ 1) (.. 8
0)
s'[i] (5,
PD. (5, a s' = 1)
(5, s'( 0)
Σ
Pu, (8,
Pu, (8, [5]0] 0)
(5, s'( =)
P., (5, a
s' 0)
(1 a)
Pu. a s'[] = P(s'(i] = (u)
P(s'[i] = D:) (8, | = 0) P(s =0 m4)
4,(3,a a)
P(s = 0 D,)
8,0
D.(5, a) P(s'[] = 0 Dj)
min
||,, G. | = 1) Pu, 6. | 0>]lv
8,0
(=(8,a) P( 0 ui)
2 min D:(3,a) [(s"[] = 0 D:)
P(s ()
0 =0)))
The first step applies importance weight. As u; has support over all reachable configurations š and actions a A, hence, we can apply importance weight. The second step uses the definition of backward dynamics (Fn,Pu. P.,). 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 Assumppion We bound the
two multiplicative terms below:
We have D($, a) = ( N orimin
NI. The first equality uses the fact that actions are taken uniformly and second inequality uses the fact that Pa-1 is an a-policy cover. As u, is the uniform distribution over Sipt(i)] x A, therefore, we have u, a) 2| This gives D(8,a)
(,(8,a) 2
o11man 2)p(())) We bound the other multiplicative term shown below:
[ 0 D.) (P = uz) 2 P(s' 0 D.) 2 2'
Combining the lower bounds for the two multiplicative terms and using 2(p((()
2 1 we get:
2 DI (.. | = 1) - P.. (.. s' = 0) 2 ammin
(23)
2N
Lastly, recall that our desired result is given by 2 LIPP, s 1) Pp. ( 0) TV whenever of (i,) * oi (ā²) and 0 otherwise. Therefore, using the derived lower bound multiplied by 1 (()(()) + oj (i2)} gives us the desired result.
Corollary 6. We have the following with probability at least 1 - Sabss
16NV
Exp (10((t)) = (()))) + (i2))
v
A (Taber 8abss [G]..
Of]min
The proof trivially follows from applying the bound in Lemma 12 to Lemma 10 as shown below:
Proof.
[10)((1)) (.())) E(61, 82,
B,, Damp
- - (2, a)]]
0Mpain
Ez = - + (12)))
2N
The inequality here uses Lemma 12 The left hand side is bounded by 8 A (naass Sabs (9]) us- ing Lemma 10 Combining the two bounds and rearranging the terms proves the result.
11
At this point we analyze the two cases separately. In the first case, s' can be set to both 0 and 1. In the second case, s' s'[i] can only be set one of the values and we will call s' [i] 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 O. We analyze the first case below.
26

---

Published as a conference paper at ICLR 2021
D.2.1 CASE A: WHEN s'[i] CAN TAKE ALL VALUES
Corollary Glllows define correspondence between learned state (i.e., output of @)) and the actual state (i.e., output of ot). We show this correspondence in the result.
Theorem (Correspondence Theorem). For any state factor s {'1,, there exists exists ùn E {0, 1] and úg with probability at least dabs such that:
(+) 0o s 0)
N 1 2,
(FO = û1 s' = 1) 2 1 2,
where e :== 16N LAI 2
A (Taas: dabsi (T) and ē p. provided (0, 2)
Proof. For any u, {0,1] we define the following quantities:
[1) =}}, Puz =
2]}.
It is easy that these quantities
related by: P2 P.: P1--)s- We define úo
arg (0,1) Pvo and ú = 1 - ug. This can be viewed as the learned bit value which is in most
correspondence with s[i] = O. We will derive lower bound on P500/ro and Par1/13 which gives us the desired result. We first derive the following lower bound Pago:
Pago + Paio
Po 2
Pago
(24)
2
where we use the fact that max is greater than average. Further, for any u, {0, 13 we have:
Exg - ( o
2 E2, -\)()))))(((())) x) -1000) 2]}
PusePa
We define a shorthand notation A' ==
anmin a v A (Tato: Sabs: (9). Then from Corollary 6 we have
16.V
proven that Puz Pu u(1-2) S A for any u, 2 E {0, 1}. This allows us to write:
A'
Pa Par1
A
2
Pe = Pao1 P
21
Pago
Pi
PiPa
PoP
where the last inequality uses
Equation 24
We will derive the same result for Pso0/
/o-
A'
Pago
A'
A
Poço = Po Po.0 2 Po
2
Payi
Po
PoPaji
PoP -
where the last inequality uses derived bound for Pa/2/P1. If we assume A < PoP: then we get
Page Pg
>
24' P.P.
As Po +
therefore, get PoP Po - Pg = -
then P. Pg Po
2
Otherwise, Po 1 which implies P1 and P1 PP Z P. This gives us PoP 2 min( 2 A.
Using lower bounds for Po and P1 from Lemma 11 gives us PoP 2
allmin and allows us to write: 2N
P P
4NJAA/
Poo
AN
2 1
2 1
Of7mmn
Po
am7mmn
It is easy to verify that e ANJAAA asmmin
As
Paga
P(0)(x') â s = 0) and Payl ((O,(x))= P7
Po
â s[i] = 1), therefore, prove our result. The only requirement used is that A' < P.P. which is ensured if 2).
27

---

Published as a conference paper at ICLR 2021
D.2.2 CASE B: WHEN s'[t] TAKES A SINGLE VALUE
We want be able detect this case with high probability so that we can learn degenerate decoder that only takes value O. 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 For given factor 2, we first learn the model 9: containing the decoder ø., as before using noise contrastive learning. We then sample Rdde iid triplets Daeg {Zz, aj, ر where 0,) ~ D and Ø; ~ p (line 6). Recall I (t)c Next, we compute the width of prediction values over Daeg as defined below:
max a5. g 0k. #x)] j.kE|maal
(25)
If the width is smaller than a certain value then we determine the factor to be degenerate and output degenerate decoder Φ: == O. otherwise, we stick to the decoder learned by our regressor task. The form of sample size Vdeg 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((, 14 chn).
Child function has type ch,, [dw] + 2 [mm)
1: for in [da], define (t), D 0, Dáeg @ do 1/ collect dataset of real (y = 1) and imposter (y = 0) transitions
2: for Rabs times do Sample (x(1), a1)) ,0),,,, (x(?), ,a(?), x(1)) ~ Unf (1n-1) o Unf (.A) and y ~ Bern(
3:
4:
If then D < D U (x(1, (1)) x²(1)) [W), y) else + D u (a(1), a (2) [w]
5:
d. = -- 3. = R)
// train the decoder using noise-contrastive learning I/ detect degenerate factors Sample (x( (1) x'(1)),, (x(2), ad²), x(11) ~ Unf((/n-1) o (nf
6:
for Ndleg times do
7:
8:
Ddeg + Daey ((((1) _ x(1))
a iniing
9: if max [m 9 (2), ay, [w] 9 (z ax. o [w])] S 11 then// max over Daeg // output decoder that always returns return X' 10 where for any x E X' and E [da] we have o(x)[i] = Φi (a(chn (i))).
22 ==
For convenience, define (a(s,, a, 8) D(r, p))( and so ay; D'. For brevity reasons, we do not add additional qualifiers to differentiate 2y, az, it 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 Zj, aj. Ry to refer to the transitions collected for the purpose of detecting degenerate factors. Lemma 13 (Markov Bound). Let {1I,, (g.;)))))))) T
Padee be a dataset of iid transitions sampled from D'.
Fix a > O. Then with probability at least 1 - 8ahe 21h / Alnaa,
we have:
Vj
[ndeg]- l0(z,, k.) g" (I, aj Fz $ a.
Proof. It is straightforward to verify that for any (Izz, az, (3) we have D(a #z) 2 D' 02,)/1
Using Corollary 5 we get:
Ex [l0(,, a, ă) g" (z. a, 1] < 2/ A (natss Sahss (9] Let E, denote the event (1ā(1,, ,g, - g (1.z, Ojj 3,)l S a} and E, be its negation, then:
2T1deg v A(rias Oahs 9
P(N 2
Σ P(E)) 1=1
a
where the first inequality uses union bound and the second inequality uses Markov's inequal- ity. As Corollary 5 holds with probability dabs, our overall failure probability is at most
Jabs 2nss / A(m 62.
28

---

Published as a conference paper at ICLR 2021
Lemma 14. For any reachable parent factor values s, action A and reachable s'[i] {0, 1),
we have (s(s, a, s'(i)) Z
Vinii
Proof. We have D' a, s'()) =
(()
p(s
o Shmin where used the induction hypothe- N
sis IH.4 that v is an a-policy cover of - and Lemma 11
Lemma 15 (Degenerate Factors), Fix a > If s'(]] only takes a single value then with probability 2n v A(Man, Sat.,199) 191)
at least - dabs
have:
max (9( (2 ag. ag, #x)l 2a
K (Rdeg)
Proof. When s'[ takes single value then g" is the constant function y. For any j and we get the following using Lemma 13 triangle inequality.
10(r,, ,, Ø5) - ( ak Tx)l S (0(z,, aj, B3)) g (zz, aj #;)l
Ig" (Iz, ag. Tt) 0( : az. .x)l I 2a.
L²12 then we have:
Lemma 16 (Non Degenerate Factors). Fix a > and assume Fldeg 2
a²
a² a
Mmmn
az g(Tk, Ck. 2x)1
2a
16N
2n V Aat
a Minan 3N"AAA
with probability at least
dabs
4exp
Proof. Equation 23 implies that there exists 8, a such that
[s' = 8, a)
T(s = 0 88
O7mmin o
2
(( 1)
p = 0)
2N
Combining this with Equation 21 we get:
p 1)p(s'(1 =0)
a, = 1) g a111= 2 4D [ = 1 8, a (s'(l) = 3,a)
2.V
2 Tiiin o
(26)
16N
where Equation 26 uses p(s' 1))(s'(i] 0) 2 AN
atmin 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 {(11, 41 F1); (Iz, 0g F2)] without loss of generality,
such that [xx1))<w] = (x2)[w] 3, action a₁ d2 a, (5) = 1, and of (x2)) 0. Then have:
max Uj. - g(ak, ak &x)) 2 (0(11, a1 x1) - 5( a2 x2)1
J.KE (mas))
2 lg (5, (8 a, 0)
10(11,01, 21) 9" (*), 01 #1)
10(12, a2,) g" (æ2, 02 Ø2)1
anminin 2a (using Equation 26 and Lemma [3)
16N2[A
We use Lemma 13 which has a failure probability of 8 2race / M Further, we also assume That our dataset contains both (8, a, s' [i] 1) and (8, a, s'[i] 0). Probability of one of these
29

---

Published as a conference paper at ICLR 2021
events given by Lemma 14 Therefore, if ndeg
then from Lemma and union bound,
a a
the probability that at least one these transitions does not occur is given by cxp a
A
The total failure probability is given by union bound and computes to:
2ddg v (Mabs Sabs, 9
a a Rddeg
dabs
4xp
3N |A
2 o 80N- [A]
we fix a
then the two case we have:
0 a.mmmsn Umin 0
(Degenerate Factor)
max (5. [( ag. #x)1 (4)1 S J.KE (aagg)
40 |A
30 mminn o
(Non- -lon-Degeneertee Factor)
max m
10(zz, 05) g( ak, (e,,
80N?
Theorem (Detecting Degenerate Case), We correctly predict if s'i] is a degenerate factor or not a a2 Aat log
3. LA
when using Tddeg
log
and e S 6at
with probability at least
o -
30N A
1 -5
Proof. The result follows by combining Lemma and Lemma 15 and using the value of a described above. These two results hold with probability least:
20dg V A (Nabas Jabss 191)
a2 Ddeg
abs
4 exx
a
3N22 |A|2
Setting the hyperparameters to satisfy the following:
A (Taatss Jabss 1/1 480N"1
Rdeg
a= a2
Log
Jabs
a2 Sabs
log
Jabs
gives failure probability of at most 36 The later condition can be expressed in terms of e which gives us the desired bounds (see Theorem 7 for definition of @). Lastly, note that setting Vlde this way also satisfies the requirement in Lemma 16 Lastly, note that the resultant bound 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.
D.2.3 COMBINING CASE A AND CASE B
Theorem 8 shows that can detect degenerate state factors with high probability. If we have 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 learned decoder and [Theorem 7 holds true. These two results allows us to define bijection between real states and Iearned states that we explain below.
Bijective Mapping between real and learned states For a given time step hi and state bit s[i] we will define úhis as the corresponding learned state bit. When and will be clear from the context then we will express this as ú,. We will use the notation & to denote learned state at time step - and s' to denote learned state at time step h. Let pf(1) = (i...... ..) and s[pt(i)]
:= (W1, u.......i) then define v (u( wy ... --- ((hh state bits corresponding w. More generally, for given K. 2d, we denote the real state factors as s(KC] (or s' [AC]) and the corresponding learned state factors as S[C] (or (((C)).
the learned
We define mapping 0n 1) + {0, from learned state to real state. We will drop the subseript à when the time step is clear from the context. We denote the domain of 6, by S, which is subset of f0, 1 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.
30

---

Published as a conference paper at ICLR 2021
10, 1y
{0.1)*
Oi
Sh
SSh
unreachable states
Figure 3: Bijection between learned state space S₁ and the state space at time step h. The bijection maps every state in S₁, to unique state in S₁ However S, also contains learned states that do not correspond to reachable state at this time step. This happens due error in the decoder.
For learned state we have s = (()) if (21, za) and (Wbies: -.. Uhdd,)) ). We would also overload our notation to write w = (w)) for given & S = á where w = [K]] whenever factor K. clear from the context.
We call s' (or s) as reachable if s' E S₁ (or E SB-1). Similarly, we call 8' (or 3) as reachable if Og (S') (or On (8)) are reachable. For a given set of factors K,, we call w E {0, | as reachable for K. if there exists reachable state with factors K. taking on value w. Similarly, we define v E {0, 13 as reachable for given K. if 0 ( is reachable for K.
We use the mapping to state the correspondence theorem for the whole state.
Corollary (General Case). If Tideg =
3N log ammmme
1
and S
a²n² nenn Sate log 30N-1A
holds,
Outa
then with probability least 1 3d8 we have:
Vs E Sh
P( 0 (()
de.
Proof. The first result directly follows from being able to detect if in degenerate setting and if not, then applying Theorem and if yes then result holds trivially. This holds with probability least 1 38 from Theorem 8| Applying union bound over all learned factors gives success probability across all factors of at least 1 3d5 The second one follows from union bound.
P(8 + 0-) PP::: SM + (+) 1 S
P( + 0 (s) a(1) S de.
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, 2). then with probability at least 1 dd we have:
Vs' E Sh,
> 1 e. P(6' = 0-) > dg.
Proof. Same as Corollary 9 except we can directly apply Theorem as we don't need to do any expensive check for degenerate factor.
D.3 MODEL. ESTIMATION
Our next goal is to estimate model Tí Sn-1 x A + ((Sa) and latent parent structure pÊa, given roll-in distribution D E A(Sh-1), and the learned decoders { hish Recall that our approach estimates the model by count-based estimation. Let 'D (k) (k)) x((k)) te be the sample collected for model estimation.
Recall that estimate pi,, (1) be using set of learned factors I that believe is pt, (i) and varying disjoint set of learned factors J. If pi) C I then expect the learned model behave the same
31

---

Published as a conference paper at ICLR 2021
irrespective of how vary the factors J. However, if pt()) g I then there exists parent factor that on varying will different values for the learned dynamics. Let = I U 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 Sr and control group I and variable group J. For a given ê E 1), ④ E {0. 111k and a E A, we have (5 ē & 8 = w. a) denoting the estimate probability derived from our count based estimation Algorithm 4 line 3). Let D
S[C] rù a) be the probabilities that are being estimated. is important to note that we use subscript D for these notations the learned states are not Markovian and K. may not contain pt)),, therefore, the estimated probabilities P and expected probabilities TP will be dependent on the roll-in distribution D.
In order to estimate PD(. | w a) we want good lower bounds on PD( wi a) for every E A and v reachable for K. Our roll-in distribution D only guarantees lower bound on P PD(Ss-1 [C], a). However, we can use IH.2 to bound the desired quantity below.
Lemma 17 (Model Estimation Coverage). If S 2 then for all K E (g2n((d)), E A and e E {0, 13 [ reachable for K, we have:
Pn(šp- |1] w. 0) N 4
07mmin
Proof. We can express Pp = uù a) = [Ail PD 6/ [K]] = w as actions are taken uniformly.
Let 5h-1 = On- (6 and w -h-1 [K]. We bound PD ($ [K] v0) shown:
P PD [K] = w) > PD (S [C] úZ - [K] w)
Pp (6)- [K] = á 8h-1 [K] = w) Pn (s [K]] w)
II PD ( [k) wx 8h-1 [k] wx) PD (5h-1 =w)
ke
2 (1 e) 2x 077mmnn
a77min
N
where the third step uses the fact that value of learned state - given the real state bit [k]] The fourth step uses [I.22 and the fact that we have good coverage over all sets of state factors of size at most 2x. Last inequality uses [K] < 2x and 2.
[k] is independent of other decoders
We now show that our count-based estimator PD converges to PD and derive the rate of convergence. Lemma 18 (Model Estimation Error). Fix Sesr E 10, Then with probability at least 1 Sest for every X. ([d). á {0, |K| reachable for K., and A we have the following:
Σ
P P = 8 [K]] = w. a) - (P = 6n-1(K] = w.a) S 22 (Near Sea);
$E(0.1)
where Aew (Rest, 8est) =
NIA! 1/2
In
AelLA (ed) Ben
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
Proof. We sample nea samples by roll-in at time step
event {8)-1 [X] w. an-1== a}. E(K, w a) occurs in our dataset at least m times for some 18
2 In(1/0 then from Corollary 15 we have
Σ ée (0.1)
[K]] = u a) - PD = 0 8n S = w. a) Sc
with probability at least ō. Lemma shows that probability of E(C, w;, a) is at least Comia
T<NAA
( In (5) then we get at least m samples of event auijmin
Therefore, from Lemma 28 if Nest 2
EXC, w. a) with probability at least 1 - 5. Therefore, the total failure probability is at most 26: 5 due to not getting at least samples and 5 due to Corollary 15 on getting m samples. This holds for every triplet (AC, w. a) and Lemma 23 shows that there are at most 2(ed)2 [ such triplets. Hence, with application of union bound get the desired result.
32

---

Published as a conference paper at ICLR 2021
Lemma 19 (Model Approximation Error). For any i E [d] K & (52x(d]), E Sh-1, E A, Sp, let 0 (s) and s' =0-'(s') Then we have:
5kpN
[D ((K), (D(s' a[K],a)]= App
and then use correspondence result (Croollary 9) to prove the desired result. We start by expressing our conditional probabilities as ratio joint probabilities.
Proof. We will first bound [D S|K] a) PD( s(K].a) a)]
[ (p(((K]]a)
PD(s (p(s"[],
Pp [
s[]],a)
PD(s |[C],
Pp(s([],a)
From Lemma 29 have:
8 E S2 (p))(K]),))
[> [K]],a) - Pn(6'[i [[[],
where
(27)
1 == (P.(s')i.888 - Pp(s'
and E2 [PD($[KC), - PD(ss We bound
these two quantities below:
Σ
8n - Σ
E1
Bh-1 [C]
86 86-1 [K]]
Σ
Σ
Sh- 5h-11 [K]
max
Σ
((], ()),
Σ
PD S s(C),a))
[
[
Term
Term
Where the first inequality uses We bound Term below:
la - bl S maxfa, b) for a. b O.
Term
Σ
PD(
SI 5h (K], 8 (K)(())PPp9s
Σ
(s 5n 110
A
[]
0
Σ
PD [C]] <
A
[[
The key inequality here is [P((
8h [A]) = ILEE P[([k]
sh-1(k) < as there exist at least
one j E such that 1n * s(y] and for this j we have (P]] sh (5) S g.
We bound Term 2 similarly:
Term 2:
Σ
PD(s'[]] 9h-1 () ((XT))P(s((K)
A
&n [K]]
Σ
PP [K] s(K]) TA (P s(])
A An [
260
< [Ai (1 - (1 - (01) -
TA
where use (G [|C] IIEE P(((k]
s(k) > (1 e)
and KC| 2x.
33

---

Published as a conference paper at ICLR 2021
This gives E1 S
2ne [A]
The proof for Eg is similar.
52
9h
E PD],A_
Fh
An $n-1(R]
Σ
PD 100
84-1 [
max
E
9h
Σ
PD PD 18))
8h-1 [
8n [[|]]
Temm
We bound Term below similar to Term
Term 3:
E
A
[
Σ
PD((h--(KT) - [A
[K]
and Term 4 is bounded similar to Term 2 below:
Term 4:
Σ
PD( (((())
[AI
[
1
{1 PD(((K]] s(K])}
[A
A (1- -) <
2xg
[A
This gives us 82 S 2ng [A
Plugging bounds for E1 and 82 in Equation 27 and using (s((s(]).a) =
0lman gives us:
N
4kg
4kgN
(D(5 ((KT),)) - PD( [K
(28)
| ( K.).a)
O7pmin
We can use correspondence result to derive lower bound:
(D a) 2 PD [1] s (s'( s(K],a)
p a]C] 2 PD(s'[ 3[K], e
and upper bound:
PD a) P (s'( s' [8] a)+
P[ 3 (1) s'[il s)
< PPos' 8) +) +
Combing the lower and upper bounds with Equation 28 gives us:
4x afmin
5kpN
(s S[K] a) (D(s [)) <
0
077mn
which is the desired result.
We merge the estimation error and approximation generate the total error.
34

---

Published as a conference paper at ICLR 2021
Lemma 20 (K-Model Error). For any i E [d], K (s2n((d), s Sh-1. A s' E Sh. let 3 0 1 and 8' (s'). Then have:
P (l].a)-PP(ff s(K], S Dear(Mest dea) + Dapp
with probability at least 1 den-
Proof. Follows trivially by combining the estimation error (Lemma 18) and approximation error (emma 19) with application of triangle inequality.
D
D.4 DETECTING LATENT PARENT STRUCTURE IN TRANSITION pla
We are now ready to analyze the performance of learned parent function pla- Let K1, Eg2n (2d) and úi E {0. 1) 1h {0, 1) We will assume tú is reachable for 1r and
i2 is reachable for K2. For convenience we will define the following quantity $2 measure total
variation distance between distributions (s' :).) conditioned on setting s[K]] = i1 and
[[K2] t22, and for a fixed action E A::
La (K1, (K1, ú1, K2, (22 =
Σ (E(1)))
PD(S'T = | = 01, a) P |||]2 = u2 a)
We can compute 52 for every value of i, K1, ú0, K2, d22 a in computational time of Shet Al). We also define a similar metric for the true distribution for any KI, K2 and v E {0, 13, wi {0. 1)(ckal, W2 E {0, 13ikal and a E A:
((2ed)3r+
Pea (KC, W1, K2, w22 = 2
Σ (E(U,1)
PD(S/] = ||[[[]]]]] = w1, a) - PD(s'[ = s(122] = W2, a)l
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 [d] and I E (cn((d]) pt) S I then for all e A and {0, 1} म we get:
Dig (II: JT), (i; w1). [I; Jel. [ā; #2)) S 24 (nest: Sest) + 2Dapp
max Je.Frn01.
where max is taken over J1, E (ld). in (0, 1315,, u2 E (0, 13 such that [ā; d1) is
reachable for [I: Ji). (û; (d2) is reachable for [I; [2), and I Ji = II J2 01.
Proof. We fix Jr. J2. à, 11 t2z a and let K2 = [ı], N2 = [2: v = -(()) u = A(ū). W1 (w))) and we (w2)) As pi) C I, therefore, we have:
[s' = s[K,]] = [u; w1). a) = ( [ = s s| = u, a) = (s'l = s|K2] = [u: w2). a)
Using this result along with Lemma 20 and application of triangle inequality we get:
IP (S 6 [C,] [û; 1 a) P (8 $ [ [ w2) a) s 24ca (Mests Sest) 2
Summing over 0, dividing by 2, and using the definition of Ô proves the result.
The following is a straightforward corollary of Lemma 21
Corollary 11. Fix E [d] then there exists an I such that for all a e A and E {0, II:
Dra (I:; [ās w1). [I: Jz], [û; (2)) & (essteest Sea) + 22
max Jr.J2.01,
where max is taken over Jr, J2, 1 uh satisfy the restrictions stated in Lemma 27
Proof. Take any such that pt(1) I and apply Lemma 21 Note that we are allowed to pick such an I as [p]( & is by our assumption.
35

---

Published as a conference paper at ICLR 2021
Recall that we define (i) as the solution of the following problem:
(i) argmin 2, dig
max
Qua (Z; Jıl, [u; wi1, [2], [ū; 222))
(29)
where I E B5, ([d)), a A, and à, J1, J2, 11, Wy satisfy the restrictions stated in Lemma 21
We ready state main result for pt.
Theorem 12 (Property of pl). For any 8 1h-1, a A. s E Sh let 8 0 and 8' = =0 (6').
Then the learned parent function pt satisfies:
a [d).
P (lpt(i)), a) Ti(s'[i) s(pt(i)).a) I 33 8es) 3.
Proof. Fix i E [d] Let J pt(i) - pt) and K. = pi(() U.J As pt(i) S K, therefore, we have:
P PD s a) (s s(pt(t)).a)
Combining this result with Lemma 20 we get:
[ 8[[], a) [+] sip( a) S Dea(nest Sest) + upp
From the definition of pi()) Equation 29) and Corollary 11 we have:
[ ()) 0), ( | (1); J., a) S 22e (Meas Sea) + 20p
Note that we are allowed to use Corollary 11 as (i) 0] and (i) J] are both reachable since they are derived from reachable real state s, [pt( S K. (i) J] < p((( S 2x, and (i) 0 0 = 0 = (i) n J. Combining the previous two inequalities using triangle inequality completes the proof.
D.5 BOUND TOTAL VARIATION BETWEEN ESTIMATED MODEL AND TRUE MODEL
Given the learned transition parent function pt,, define the transition model as:
Thi (s' (1), = PD (5'l) (o).))))))) T 8, II The (3'T1) (1),
From Theorem 12 we have for any E [d], 8 Sp-1, E A, Sp and s = (8), 0-'
[s ((pt T(' slpt(i)). s (nsst dea) + 3Dapp
Transition Closure. subtle point remains before we prove the model error between Th and T. Theorem 12 only states guarantee for those that are inverse of reachable state However, as stated before, due to decoder error we can reach a state á which does not have a corresponding reachable state, i.e. 0(s) S (see Figure 3). We cannot get model guarantees for these unreachable states 8 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 To for time step as:
(B(s')) 0))) = T(0(s') 10(3),),, Tn((' 8,a),
if 0(3) E Sh-1 otherwise
V' E S A. E Sh.
We also define the state space domain of TV as S = (0.-1 (8) V8 S.-1). ). easy to see that 8 represents bijection between S and -
We will derive our guarantees with respect to T which will allow us to define a bijection between the domain of T and 7 and use important lemmas from the literature. The next result shows that use of 7 is harmless it assigns the same probability T to any event.
36

---

Published as a conference paper at ICLR 2021
Lemma 22 (Closure Result). Let T be the closure transition model with respect some learned transition model. Then for any policy π E and any event & which is a function of an episode sampled using n, we have P (E: (T; where Fr (E; denotes the probability of event when sampling from n and using transition model T
Proof. The proof follows form observing that when using T we will never reach a state 8 & Sh-1 for any h 1 by definition of - 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 & is 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 3 Sh-1 and A we have:
Σ 8'ES
(s' (,,0) (0(3)))) (()) a) I 6d (Den(Messs Sess) aapp)
Proof. If ((3) S then by definition To the bound holds trivially. Therefore, we focus on (()) E S for which 7 = T. We define the quantity for every j E [d]:
S =
Σ
The (s slpt(i)). II ((((5)) ((3)pt(i)). II
(30)
8'1 (1 (0.1)
We claim that Sj S 6(d- + 1) + app) for every j E [d]. For base case we have:
Sa
Σ
(3)))t())) (8' 1 sl)( (((()))))) [()
ad s 6(Aeu + Aupp),
from Theorem 12
We will assume the induction hypothesis to be true for Sa for all lo > j. We handle the inductive below with triangle inequality:
d II (5'l ()),) on 18()))
S S
Σ
3 3(d imj+1
T(0(3) (3))
d II (8 ())) (=++1
Σ
d II T(0(3) (())))t(())))))))))))))))))))))) a)) 1=11
The first term is equivalent to E) [Thy (B'Li)
(TT()). | T(O (() a)l
which is bounded by 6( 6(Aest + App)) following base case analysis. The second term is equivalent to Sj+1 which is bounded by 6(d - ( + App)) by induction hypothesis. Combining these two bounds proves the induction hypothesis and the result then follows from bound for ST.
D.6 LEARNING A POLICY COVER
In this section show how 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 Mr and M₁. M₁, is the true MDP consists of state space (S), ... ,S)), action space A, horizon {0]d, and transition function Tp S x A A(S,) for all
h. deterministic start state
t [h]. Recall that the set S₁ C (0, 1 denote states which are reachable at time step h, and the set
37

---

Published as a conference paper at ICLR 2021
S, S Si S {0, 1] represents the closure of state space Sn based on the learned state space St. For any E [h], s E St and K E ([d) we know suprellss (x [K] = =s[]] 2 7min.
The second MDP Mi consists of the learned state space (S) deterministic start state $1 = {0]a, and transition function T, S A.S,)...
(G,,... ..... Sn), action space A, horizon h,
For every 1 E [M.. we have 0₂ S. -> S represent a bijection from the learned state space to the
closure of the set of reachable states at time step t. The learned decoder de predict (s) given E S. with high probability for < by [H.2 for = h due to Corolllay
I Lastly, the transition model T and T, are close in L,1 distance for I. < h due to IH.3 and for / = h due to Theorem
These results enable us to utilize the analysis of Du et al. (2019) for learning to learn policy cover.
Let L : S - AA denote a non-stationary deterministic policy that operates on the learned state space. Similarly, yp S + A denote a non-stationary deterministic policy that operates on the real state. We denote 0 9 o 0 if for every& & E S, (()) = (0(())) Similarly, we denote
& 0-1 if for every s E S (s) = *(0 (s)). Let : X + be a non-stationary deterministic policy operating on the observation space. We say π = 9 o ø if for every x E X we have x (xx) = (x)))) Similarly, y o- if for every E X' have (2x) ( (x:).
define
We will use P2 [E] to denote probability of an event & when actions are taken according to policy T: X' + A. We will use Pg]E]] denote the probability of event & when we operate directly on the real state and take actions using yp Similarly, we define PP [E) to denote the probability of event E when operate the learned state space. Lastly, let Po[E]] denote probability of an event & when actions are taken according to policy & operating directly over the latent state and following our estimated transition dynamics T : S x A A(S). Recall that our planner will be optimizing with respect to Pp]EE.
Theorem 14 (Planner Guarantee). Fix Apt 2 0, [H]. Let I E ([2)((d) and i E {0, 1
We define reward function I S - [0, I| where R(3) 1(T($) = h A 8 w). Let RR = planner (T, R. h, Ap)) be the policy learned by the planner. Let == &R Φ then:
Ps (sn [ = ((w)) 2 n( = ((w)) 2dgH 12dH Dear 12dH Dupp Aph
(31)
further, we have:
= (18 (p w)) 2 [| - 6dH Des Adp Apts
(32)
and if [sn[Z] {(w)} is unreachable, then
(enn((n(T) = w)) S 6dH Aest + AGdHAapp
(33)
Proof. We define two events E = [s = {(w)) and E = [T = w]. We define a policy
P = ΦR 0-1 where for every s E S we have (s) = (
We also define : X A
as (x)) = o (x). If for a given x E X' and o (x) = s we have Ō(z) = 1 then [(x) = (z))(z)) = *xx). 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 de using Corollary 9) to write:
(E)) - (x) [P((E) Pn (E) < 2dgH
(34)
Let s S + A be any policy on real state space and let و S + A be the induced policy on learned state space given by (($) = &p 0(5) = (0(()) for any 8 E S. We showed in Theorem 13 that T and T have small L, distance under the bijection 0. Therefore, from the perturbation result of Du et al. al.. (2019) stated in Lemma 3 we have:
Σ
(sx)) -P (sN) S he S He,
ES
38

---

Published as a conference paper at ICLR 2021
where 6d Aesi(neae Ses)) Appp due therefore, we can derive the following bound:
Theorem 13
As {sh[I] 0(()) + [8 20)
(P(()) -P(()
(p((s))
E
Pp(00'(sh))
e e 18h [2] (1())
(P(( P=((0 < He
(35)
Let y0" arg max P [&] be the optimal policy to satisfy {sx[I] O(w)). Note that yp is also the latent policy that optimizes the reward function It on the real dynamics. Let 6 = Dp o 0-11 be the
induced policy learned states. We now bound the desired quantity shown:
Pa (E) 2 P. Pn 2dgH
(using Equation 34)
Po (E) 2dgH HE
(using Equation 35)
Po. (E) 2dpH HE Ap
(Gp is p on T')
Pp (E) 2dgH 2H Ap
(using Equation 35)
(8h I = ((w)) 2deH 2H Ap--
This proves Equation 31 and Equation 32 Note that our calculations above show:
Pox ( < IP p + He
If [sh 7 { is unreachable then PP [8] 0. Plugging this in the above equation proves Equa- tion 33 and completes the proof.
D.7 WRAPPING UP THE PROOF FOR FactoRL
We almost done. All need to do is make sure is the hyperparameters and verify each induction hypothesis. We first set hyperparameters.
Setting Hyperparameters. Let fsh [
((w)) be reachable for some I E ((d) and i
(0, 121, Then applying Theorem 14 and using the definition of 1/min have:
Ps [sx[Z] ((w)) 2 Mmin 2dgHf 12dH Aest 12dH Aupp
Ap
As we want the right hand side to be at least Ofpann we divide the error equally between the three terms. This gives us:
(Planning Error)
Apt S (1 (1-a)Mmin/4
(36)
a(1 a)Mmin
(Model Approximation Error)
A upp < (1 (1-a) Middm 2
(37)
240ndHN
(Model Estimation Error)
Aest S (1-a)Mmon/4Saf)
In? [4e]Al(ed)2
+ Best 2 a(1 -2)
(38)
dest
(Decoding Error)
S (1-0)Mmun/SdHH
(39)
The model approximation еπrо places more stringent requirement on than the decoding error for planning. However, throughout proof for FactoRL in this section, we made other requirements on our hyperparameters. For this is given by min min( (2m/1200, 1/2] 22m/1200 by combining constraints Lemma and Theorem 7 and an additional constraint for detecting non-degenerate factors stated in Corollary Due to the inefficiency of the non-degenerate factors detection, we state results separately for the two cases:
a(1 - a a)m
e S min
min
(no non-degenerate factor)
1200'
240mdHNN
182
a(1 - a))
2 2 8ahs
4
0 S min
log
(general case)
1200'
240n:dHN
30.22 Al2
Sabs
39

---

Published as a conference paper at ICLR 2021
Using the definition of o from Theorem 7
get value of nabs for non-degenerate factor (Equa
tion 40) and general case (Equation 41) given below:
3840 38402N LAP a Umin o² IAP nabs
[9] 8abs
12
25
Taass
In
max
(1 a 2 Umin
B
(40)
|9| 8abs
1sa 222
25
|4
4
ammmm
In
max
a² 2
P
a Tmin
82
In
6abs
(41)
abs
Recall that for detecting degenerate factors we collect Medy samples. Corollary 9 givee value of this hyperparameter as
3N2 log
Ylddg
0221 Frrii
8abs
which also satisfies the condition in Lemma 16 Lastly, Theorem 3 gives number of samples for independence testing Tind and rejection sampling frequency A: as:
1 m
In m² 1 (2ed]
30
Fiind O
h:
In
ônd
Mrii
Biii
Failure probabilities for single timestep are bounded by Snnd due to identification of emission structure (Theorem 3). 3d0abs due to decoding (Crolllary 9). and dest due to model estimation (Lemma 17p. The total failure probability using union bound for single step is given by dnd add + Sea. and for the whole algorithm is given by 8ma + 3.. + Eest H. Binding And H + 6/3, 3d5 H - 8/3, 0esH H + 6/a, gives total failure probability of 5 and the right value of hyperparameters.
Sample complexity of FactoRL most kH wind + H + Hrieg + Hres episodes which is order of:
poly 16m |A]. H.
1
8' Bmin mnm, In Inl.
1
Tmin
where use the fact that N | can be at most 2(ed)2" from Lemma 23 Note that if we did not have apply the expensive degeneracy detection step, then would get logarithmic dependence on 1/8a.- 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. bound the model error and decoding
For our choice of two hyperparameters Test and Plass failure by:
6d mood Appp) I (1 ammnn
a(1 - ammmmi
(Model Error)
(Decoding Failure) @
4H
240xdHN
Verifying Induction Hypothesis.
Finally, we verify the different induction hypothesis below.
1. We already verified IH.1 with Theorem 3
We learn a ch, that is equivalent to chi, upto
label permutation.
2. We already verified with Corollary 9 Given a real state E Sn, our decoder outputs the corresponding learned state with high probability. We also derived the form of p.
3. We already verified [H.3 with Theorem 13 We also derived the form of A and Aapp-
4. Lastly, Theorem 14 and our subsequent calculations for hyperparameter show that N is a-policy cover of S₁, and that the size of Wr, most 2(ed) 2x from Lemma 23 Lastly, for all reachable factor values we get the value of learned policy as at least (1+a))mma using [Equation 32 and our choice of hyperparameter values. Similarly, from /oomEEqation 33 we get the value of learned policy for all unreachable factor values as at most (1 (laa))miin/ 4. This allows to filter all unreachable factor values. the main paper, we focus on the value of a 1/2, which explains why on Algorithm line 8 we only keep those policies with value
at least (1+a))maan2 = 31min/4. This verifies [H.4
This completes the analysis for FactoRL.
40

---

Published as a conference paper at ICLR 2021
E SUPPORTING RESULT
Lemma 23 (Assignment Counting Lemma). For a given k. d and k S d, the cardinality of the set {(K,, u) IC E (gx(ld)), a E {0, 1) Kij bounded by 2
Proof. Assume k 2 2. The cardinality of this set is given by Ei (9) which can be bounded as shown below:
(1) 1 51+2d
ed
Σ i=0
2' = 1+ 2d Σ 2=
2 2 2d
(ed)
(ed)'.
2=2
==0
The first inequality here uses the well-known bound for binomial coefficients (?) 1 (d)) for any
n, E N and < n. Further bounding the above result using ed ed/2 gives us:
(ed) Σ (ed) < ed - & 2(ed)*.
1=0
The proof is completed by checking that inequality holds for k < 2.
Lemma 24 (Lemma H.1 in Du et al. (2019)). Leru, u, RQ. with ||ul| ||l| and |1--o|1 Then for any > O we have I 2ul1 2'
Lemma 25 (Chernoff Bound). Let q be the probability of an event occurring. Then given n iid samples with 2 1 the probability that the event occurred at least once is at least 1 - 2 exxp
Proof. Let X; be 0-1 indicator denoting if the event occurred or not, and let Σi= X, We have E[X.] = q and E[X] = qn. Let t 1 -1- 1/p/ We will assume that qn > 1 and so / E (0,
Then the probability that the event never occurs bounded by:
f)gm) S exp ( (-qnl²
< 2 exx { qn
{-
PP( < 1) P(X <
3
Lemma 26 (Hoeffding's Inequality). Let X1, X2.... Xn, be independent random variables bounded by the interval [0, 1]. Let empirical mean of these random variables be = - Xns then for any > have:
(X--EXX E 2 t) < exxp(-2ndt
Lemma 27 (Theorem 2.1 of Weissman al. (203)). Let P be probability distribution over discrete set of size a. Let X" -X1, X1,X2, ... Xn be independent identical distributed random variables distributed according to P. Let Pxn be the empirical probability distribution estimated from sample set X". Then for all >>>0 0:
P(IP Px 2 20) (2 - cxpp (-
ne²
The next result a direct corollary of Lemma 27
Corollary 15. For any m 2 22 In (1/6) samples, we have ||P Px= ll1 < e with probability at least 1-
Lemma 28. Let X1, X2,... X, be 0-1 independent identically distributed random variables with mean μ. Let X = Din X,. Fix m N and 8 E (0, 1). If
2m In (9) then (PX m) S 8.
p
Proof. This a standard Chernoff bound argument. We have E[X] = np. Assuming n 2 m/u then from multiplicative Chernoff bound we have:
P(X m) P x < {
}
-nyu
2 ))
nn
S
exp
S exp
ni
2
nf
Setting right hand side equal to and solving gives us n whenever 2m In (g).
>
2 R
(m + in())) which is satisfied
41

---

Published as a conference paper at ICLR 2021
Lemma 29 (Lemma H.3 of Du al. (2019) ) For any a, b, c, > 0 with and S d we have:
a c + b-66 d)
b
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 (Lemma G.5 of Du et al. (2019).
(Si,... ,St) and S = =Si,,... ,.,) be
the real and learned state space. We assume access to a decoder & X + S and let o* : X + S
be the oracle decoder. Let 0n Sr + Sp. be bijection for every [H] and 0 S + S where
(()) (5) for å Sn. For any [H] and Sh Sh assume P( 0% (sh) ) 9h) 2 1 i.e., given the real state Sh, our decoder (0) will map it to 07²² (sn) with probability at least E.
Lef yp : - A be deterministic policy the real state spuce and S - A be the induced policy given by ((() =(((()) for every S. Let n, it X A where πTx) (p((=())))) (I)) and
(x)) = ((((x)) for every x E X. For every random event & we have:
[x(E))
P+(E)| 2
Lemma (Lemma H.2. of Du et al. 2019)). Let there be two tabular MDPs M and M. Let S= (S ... ,SH) be the state space of M and S = (S,.... ,SH) be the state space of M. Both MDPs have A be the action space of both MDPs and horizon of H. For every h E [H], there exists a bijection 0g: Sp. - Sn. Let T S x A - A(S) and f: S A A(3) be transition dynamics for M and M satisfying:
vh a
(T((((5)) ((3). a) -a 8, 8. se S
Σ
Let up: + Ab a policy for M. Let ; S + A be the induced policy on M such that for any & E S we have ((8) =()())). Then for any h E [H] we have:
Σ
P
P.p (sw) S he
F EXPERIMENT DETAILS
We provide details for our proof of concept experiment below.
Modeling Details. We model F, used for performing independence test, using a single-layer feed- forward network 0y with Leaky ReLu non-line non-linearity (Maas et al. 2013) and a softmax output layer. Give a pair of atoms arfu] and xfv]. we concatenate these atoms and map a probability distribution over {0, 1} by applying 0y.
We implement the model class G for learning state decoder following suggestion Misra et (2020). Recall that a function in G maps transition (ær, a, 2) E X x A x 2 to value in [0, 1j. We first map æ and T to vectors e1 and V2 respectively, using two separate linear layers. We map the action to its one-hot vector representation 1, We map the vector V2 to probability distribution using the Gumbel-softmax trick (Jang et al. 2016), by computing 9i x exp(v2 [i] + v.) for all E (1, 2). where O; is sampled independently from the Gumbel distribution. We concatenate the vectors U1 1. and and map to a probability distribution over {0, 1), through a single layer feed- forward network G with Leaky ReLu non- -onearira We recover decoder from the model that maps set of atoms
T to (()) = arg (0.1) 1-+1-
42

---

Published as a conference paper at ICLR 2021
Learning Details. We train the two models using cross-entropy loss. Formally, given dataset Dind [xr,[u] sylp], y:) Fine for performing independence testing, and a dataset Dabe = {1o,, dg, 2i, 31) B:) J== for learning decoder, optimize the model by minimizing the cross-entropy
loss as shown below:
Nish
f = arg max JEF 7ind
In f(y 2.[u], 2:(e)).
= arg max gEO 7labs
Σ
In gly Dyy ag, R;,
Here we overload our notation to allow the output of models to be distribution over {0, 1 rather than scalar value in [0, as we assumed before. This is in syne with how we implement these model class and allows us to conveniently perform cross-entropy loss minimization.
Planner Details. We use simple planner based on approximate dynamic programming. Given model estimate, reward function and 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 O. This allows us to compute Q-values in computationally-eenn agent visits unvisited state, then it simply takes random action.
manner. Later, if the
Hyperparameters. We set the hidden dimension of 0z 10 and that of 0g 56. We the threshold e on held-out log-loss value, when performing independence test to be 0.65. For reference, one uses random uniform classifier then its performance is ln(0.)) 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.
0.69. We train both
43


Reviews:
[{"review_id": "d2aIxpwh3eO", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe model and assumptions are very natural in many applications, and the theoretical guarantee is very nice. Moreover, the paper is very well-written and clear. However, I would like to see some experiments regarding this model. Even an experiment on a synthesized dataset with the grid and furniture example would be very interesting.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Y1Z2QynJ8zh", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\nSummary: \nThis paper studies reinforcement learning under the setting of factored block MDP, where the observations are generated from latent factors. The paper presents a framework with separate parts for learning the policy, the emission structure, and the model, respectively. The performance of the proposed framework is theoretically analyzed. \n\n##########################################################################\nI find the paper theoretically sound but lean towards rejection at this point. \n##########################################################################\nMajor comments:\nIt is appreciated that the authors provide the complete theoretical analysis of the proposed framework, and the analysis is reasonable and sound to me.\n\nMy main concern is that the paper combines several settings together without much technical novelty. The settings are somehow artificial in the sense that a lot of assumptions are made only for obtaining theoretical results, and in addition, no experimental study was provided.\n\nDetailed comments:\n\n“Our goal here is to define a problem setting that …” This sounds to me that the paper aims to find a setting that can be analyzed with proofs but not the settings that are natural and have real applications. In addition, it would be good to inform the reader which settings are inapplicability or intractability. \n\nThe discussion on the rationale of using factored transition is not necessary (as it is well-known as the basic part of graphical models).\n\nThe paper adopts many assumptions, but not all of them are well-justified. E.g., \n-\t“For our purposes, we assume a deterministic start state and assume without loss generality that each state is reachable at exactly one time step” It is not clear to me what the \"purposes\" are. What is the consequence if the start state is not deterministic? Why do we assume that each state is reachable at exactly one time step?\n-\tIn addition to making the problem trackable, do we have the disjointness property and Assumption 1& 2 in real applications (such as the image example in the paper)? \n-\tThe paper assumes two computational oracles and one planning oracle. Please provide the accessibility of these oracles so one can evaluate the feasibility of the proposed method. \n\n\n\nThe presentation is often flowery without being clear. E.g.,\n-\t“There are many possible ways to add rich observations to a factored MDP resulting in inapplicability or intractability.” Please clarify with respect to what the inapplicability or intractability is taken (e.g., learnability or computability). \n-\tPlease be consistent. For example, “the observations are emitted by latent states” and “the observations are emitted by latent factors”; “decode the latent state” and “decode the factors”. They are making the paper hard to follow.\n-\tIn Def 2, ch() is defined to be [d]->2^[m], making x[chi(i)] not well-defined. \n-\t“An agent is responsible for mapping each observation x ∈ X to individual atom …” What is the difficulty in having such a mapping? From the paper, the atoms are simply given by the observations.\n\nThere are almost 30 pages proof in the appendix. Are all the proofs new? If not, please state clearly that which proofs can be found in the existing works and which can be obtained by adapting the proofs in other works.  Otherwise, it is hard to evaluate the contribution of this paper.  \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "VCieEIyCHcZ", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a new result for provably efficient exploration in rich observation RL with assuming that the underlying structure is a factored block MDP. Under this setting along with some other dynamics assumptions, e.g., reachability and identifiability, the algorithm achieves a sample complexity independent of the number of latent states |S| but the number of factors d (which can be much smaller than |S|). Although the algorithm uses a similar approach as in Du et al. 2019, the adaptation is not that straightforward. I appreciate the novelty of this result in the direction of combining factored MDP with rich observation RL.\n\nMy major concerns are:\n1. Since this work takes a similar approach as in Du et al. 2019, i.e., learn decoding functions and transition model level-by-level, it also requires additional dynamics assumptions e.g., reachability and identifiability which are described by hyperparameters such as \\eta_{min} and \\beta_{min}. And the sample complexity can still be pretty large if \\eta_{min} and \\beta_{min} are very small even with a small d. \n2. In one recent work by Feng et al. (https://arxiv.org/abs/2003.06898), a general framework is proposed to attack the rich observation RL problem under the block MDP setting without any dynamics assumption. It might be good to include a mention of this work and some discussions about whether the approach in Feng et al. can be extended to the factored block MDP setting (if not, what is the main challenge).\n3. For the hyperparameters such as \\eta_{min}, \\beta_{min}, \\kappa, and d, how should a practitioner select these values without any prior knowledge? Can this be done in an adaptive manner as we run the algorithm? Or must we choose very conservative values (e.g., very small \\eta, \\beta, and big \\kappa) initially to guarantee the learning accuracy?\n\nMinors:\n- Some notations are a little bit confusing for first-time readers. More clarification can be added, e.g., 1) A parent function is a set-valued map from an integer in [d] to a subset of [d]; 2) s[pt(i)] := {s[j] | j \\in pt[i]}. A simple example would be even better for readers to quickly understand.\n- I appreciate this work as a theoretical result. It would be more promising if some preliminary experiments can be conducted to compare with prior works, e.g. Du et al. 2019 and Feng et al. 2020, to further demonstrate the benefit of a factored MDP setting in practice.\n\nI would like to change my score if the authors can address my concerns.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "MZQjn9ISODb", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies reinforcement learning in spaces with a large number of states by modeling the states using a factored / latent representation. This problem has been studied in the non-factored setting by Du et al (2019), and this paper extends to factored settings. This is an interesting problem.\n\nOverall, I found the paper difficult to follow since it is not presented well. It would help to distill the novelty and main ideas in the algorithms early. Section 2 on the setting includes too many technical details that can be deferred to later sections. \n\nThis paper also does not verify the feasibility of their algorithms through experiments. Compare this to the cited papers which include experiments to show the validity and feasibility of their algorithms.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper considers the reinforcement learning in rich observation setting. Concretely, the authors provide a provable sample efficient algorithm for the rich-observation factored MDP. As the majority of the reviewers commented, although the techniques used in the proof share some similarities to the existing work, the analysis for the whole algorithm is still challenging. As a theoretical oriented paper, I think this paper should have a position in ICLR. 

The major concern of the paper is the necessity of the assumptions (R2). The validation and justification of the Assumptions should be stated clearly in main text, even they are adapted from the prior work. 
Decision: Accept (Poster)