P ROVABLE R ICH O BSERVATION R EINFORCEMENT L EARNING WITH C OMBINATORIAL L ATENT S TATES
A BSTRACT
We propose a novel setting for reinforcement learning that combines two common real-world difﬁculties: 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 efﬁcient planning algorithm.
I NTRODUCTION
Most reinforcement learning (RL) algorithms scale polynomially with the size of the state space, which is inadequate for many real world applications.
Consider for example a simple navigation task in a room with furniture where the set of furniture pieces and their locations change from episode to episode.
If we crudely approximate the room as a 10 × 10 grid and consider each element in the grid to contain a single bit of information about the presence of furniture, then we end up with a state space of size 2 100 , as each element of the grid can be ﬁlled independent of others.
This is intractable for RL algorithms that depend polynomially on the size of state space.
The notion of factorization allows tractable solutions to be developed.
For the above example, the room can be considered a state with 100 factors, where the next value of each factor is dependent on just a few other parent factors and the action taken by the agent.
Learning in factored Markov
Decision Processes (MDP) has been studied extensively (Kearns & Koller, 1999; Guestrin et al., 2003; Osband & Van Roy, 2014) with tractable solutions scaling linearly in the number of factors and exponentially in the number of parent factors whenever planning can be done efﬁciently.
However, factorization alone is inadequate since the agent may not have access to the underlying factored state space, instead only receiving a rich-observation of the world.
In our room example, the agent may have access to an image of the room taken from a megapixel camera instead of the grid representation.
Naively, treating each pixel of the image as a factor suggests there are over a million factors and a prohibitively large number of parent factors for each pixel.
Counterintuitively, thinking of the observation as the state in this way leads to the conclusion that problems become harder as the camera resolution increases or other sensors are added.
It is entirely possible, that these pixels (or more generally, observation atoms ) are generated by a small number of latent factors with a small number of parent factors.
This motivates us to ask: can we achieve PAC RL guarantees that depend polynomially on the number of latent factors and very weakly (e.g., logarithmically) on the size of observation space?
Recent work has addressed this for a rich-observation setting with a non-factored latent state space when certain supervised learning problems are tractable (Du et al., 2019; Misra et al., 2020; Agarwal et al., 2020).
However, addressing the rich-observation setting with a latent factored state space has remained elusive.
Speciﬁcally, ignoring the factored structure in the latent space or treating observation atoms as factors yields intractable solutions.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
+UKsmvIaZw6ENloBkOqNdSoI+jkSgXkaPXK791+TNKICk04VspzUKL9DEvNCKeTUjdVNMFkhAfUM1TgiCo/mx07gQdG6cMwlqaEhjP1+0SGI6XGUWA6I6yH6rc3Ff/yvFSHdT9jIkk1FWS+KEw51DGcfg77TFKi+dgQTCQzt0IyxBITbfIpmRC+PoX/k3bVdo7t6tVJpXmex1EEe2AfHAEHnIEmuAQt4AICGHgAT+DZEtaj9WK9zlsLVj6zC37AevsEx6SOrA==</latexit>
Contributions.
We combine two threads of research on rich-observation RL and factored MDP by proposing a new problem setup called Factored Block MDP (Section 2).
In this setup, observations are emitted by latent states that obey the dynamics of a factored MDP.
We assume observations to be composed of atoms (which can be pixels for an image) that are emitted by the latent factors.
A single factor can emit a large number of atoms but no two factors can control the same atom.
Following existing rich-observation RL literature, we assume observations are rich enough to decode the current latent state.
We introduce an algorithm FactoRL that achieves the desired guarantees for a large class of Factored Block MDPs under certain computational and realizability assumptions (Section 4).
The main challenge that FactoRL handles is to map atoms to the parent factor that emits them.
We achieve this by reducing the identiﬁcation problem to solving a set of independence test problems with distributions satisfying certain properties.
We perform independence tests in a domain-agnostic setting using noise-contrastive learning (Section 3).
Once we have mapped atoms to their parent factors, FactoRL then decodes the factors, estimates the model, recovers the latent structure in the transition dynamics, and learns a set of exploration policies.
Figure 1 shows the different steps of
FactoRL .
This provides us with enough tools to visualize the latent dynamics, and plan for any given reward function.
Due to the space limit, we defer the discussion of related work to Appendix B.
To the best of our knowledge, our work represents the ﬁrst provable solution to rich-observation RL with a combinatorially large latent state space.
T HE F ACTORED B LOCK MDP S ETTING
There are many possible ways to add rich observations to a factored MDP resulting in inapplicability or intractability.
Our goal here is to deﬁne a problem setting that is tractable to solve and covers potential real-world problems.
We start with the deﬁnition of Factored MDP (Kearns & Koller, 1999), but ﬁrst review some useful notation that we will be using:
Notations: For any n ∈ N , we use [ n ] to denote the set { 1 , 2 , · · · , n } .
For any ordered set (or a vector)
U of size n , and an ordered index set I ⊆
[ n ] and length k , we use the notation
U
[ I ] to denote the ordered set (
U [ I [1]] , U [ I [2]] , · · · , U [ I [ k ]]) .
Deﬁnition 1.
A Factored MDP ( S , A , T, R, H ) consists of a d -dimensional discrete state space
S ⊆ { 0 , 1 } d , a ﬁnite action space A , an unknown transition function T : S ×
A → ∆( S ) , an unknown reward function R :
S × A → [0 , 1] and a time horizon H .
Each state s ∈ S consists of d factors with the i th factor denoted as s [ i ] .
The transition function satisﬁes T ( s (cid:48)
[Equation]
(cid:81) di =1 T i ( s (cid:48)
[ i ] | s [ pt ( i )] , a ) for every s, s (cid:48) ∈ S and a ∈ A , where T i : { 0 , 1 } | pt ( i ) | × A → ∆( { 0 , 1 } ) deﬁnes a factored transition distribution and a parent function pt :
[ d ] → 2 [ d ] deﬁnes the set of parent factors that can inﬂuence 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 =
( S 1 , S 2 , · · · , S H ) where S h is the set of states reachable at time step h .
A natural question to ask here is why we assume factored transition.
In tabular MDPs, the lower bound for sample complexity scales linearly w.r.t.
the size of the state set (Kakade, 2003).
If we do not assume a factorized transition function then we can encode an arbitrary MDP with a state space of size 2 d , which would yield a lower bound of Ω(2 d ) rendering the setting intractable.
Instead, we will prove sample complexity guarantees for FactoRL that scales in number of factors as d O ( κ ) where κ := max i ∈
[ d ] | pt ( i ) | is the size of the largest parent factor set.
The dependence of κ in the exponent is unavoidable as we have to ﬁnd the parent factors from all possible (cid:0) dκ (cid:1) combinations, as well as learn the model for all possible values of the parent factor.
However, for real-world problems we expect κ to be a small constant such as 2 .
This yields signiﬁcant improvement, for example, if κ = 2 and d = 100 then d κ = 100 while 2 d ≈ 10 30 .
Based on the deﬁnition of Factored MDP, we deﬁne 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.
Deﬁnition 2.
A Factored Block MDP consists of an observation space X =
X m and a latent state space S ⊆ { 0 , 1 } d .
A single observation x ∈ X is made of m atoms with the k th denoted by x [ k ] ∈ X .
Observations are generated stochastically given a latent state s ∈ S according to a factored emission function q ( x | s ) =
(cid:81) di
=1 q i ( x [ ch ( i )] | s [ i ]) where q i : { 0 , 1 } → ∆( X | ch ( i ) | ) and ch :
[ d ] → 2 [ m ] is a child function satisfying ch ( i ) ∩ ch ( j ) = ∅ whenever i (cid:54) =
j .
The emission function satisﬁes the disjointness property: for every i ∈ [ d ] , we have supp ( q i ( · | 0))
∩ supp ( q i ( · | 1))
= ∅ .
The dynamics of the latent state space follows a Factored MDP (
S , A , T, R, H ) , with parent function pt and a deterministic start state.
The notion of atoms generalizes commonly used abstractions.
For example, if the observation is an image then atoms can be individual pixels or superpixels, and if the observation space is a natural language text then atoms can be individual letters or words.
We make no assumption about the structure of the atom space X or its size, which can be inﬁnite.
An agent is responsible for mapping each observation x ∈ X to individual atoms ( x [1] , · · · , x [ m ]) ∈
X m .
For the two examples above, this mapping is routinely performed in practice.
If observation is a text presented to the agent as a string, then it can use off-the-shelf tokenizer to map it to sequence of tokens (atoms).
Similar to states, we assume the set of observations reachable at different time steps is disjoint.
Additionally, we also allow the parent ( pt ) and child function ( ch ) to change across time steps.
We denote these functions at time step h by pt h and ch h .
The disjointness property was introduced in Du et al.
for Block MDPs—a class of richobservation non-factorized MDPs.
This property removes partial observability concerns and enables tractable learning.
We expect this property to hold in real world problems whenever sufﬁcient sensor data is available to decode the state from observation.
For example, disjointness holds true for the navigation task with an overhead camera in Figure 1.
In this case, the image provides us with enough information to locate all objects in the room, which describes the agent’s state..
Disjointness allows us to deﬁne a decoder φ (cid:63)i : X
| ch
( i ) | → { 0 , 1 } for every factor i ∈
[ d ] , such that φ (cid:63)i ( x [ ch ( i )]) =
s [ i ] if x [ ch ( i )]
∈ supp ( q i ( .
| s [ i ])) .
We deﬁne a shorthand φ (cid:63)i ( x ) =
φ (cid:63)i ( x [ ch ( i )]) whenever ch is clear from the context.
Lastly, we deﬁne the state decoder φ (cid:63) :
X → { 0 , 1 } d where φ (cid:63) ( x )[ i ] =
φ (cid:63)i ( x ) .
The agent interacts with the environment by taking actions according to a policy π : X → ∆( A ) .
These interactions consist of episodes { s 1 , x 1 , a 1 , r 1 , s 2 , x 2 , a 2 , r 2 , · · · , a H , s H } with s 1 = (cid:126) 0 , x h ∼ q ( .
| s h ) , r h =
R ( x h , a h ) , and s h +1 ∼ T ( .
| s h , a h ) .
The agent never observes { s 1 , · · · , s H } .
Technical Assumptions.
We make two assumptions that are speciﬁc to the FactoRL algorithm.
The ﬁrst 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.
and we adapt it to our setting.
Assumption 1 (Margin Assumption) .
For every h ∈ { 2 , 3 , · · · , H } , i ∈ [ d ] , let u i be the uniform distribution jointly over actions and all possible reachable values of s h − 1 [ pt ( i )] .
Then we assume: (cid:107) P u i ( · , · | s h
[ i ] = 1) − P u
i ( · , · | s h
[ i ] = 0) (cid:107) TV ≥ σ where P u i ( s h − 1 [ pt ( i )] , a | s h
[ i ]) is the back- ward dynamics denoting the probability over parent values and last action given s h
[ i ] and roll-in distribution u i , and σ > 0 is the margin.
Assumption 1 captures a large set of problems, including all deterministic problems for which the value of σ is 1 .
Assumption 1 helps us identify the different values of a factor but it does not help with mapping atoms to the factors from which they are emitted.
In order to identify if two atoms come from the same factor, we make the following additional assumption to measure their dependence.
Assumption 2 (Atom Dependency Bound) .
For any h ∈ [ H ] , u, v ∈ [ m ] and u (cid:54) = v , if ch − 1 ( u ) = ch − 1 ( v ) , i.e., atoms x h
[ u ] and x h [ v ] have the same factor.
Then under any distribution D ∈ ∆( S h ) we have (cid:107) P D ( x h [ u ] , x h [ v ]) − P D ( x h
[ u ]) P D ( x h [ v ]) (cid:107) TV ≥ β min .
Dependence assumption states that atoms emitted from the same factor will be correlated.
This is true for many real-world problems.
For example, consider a toy grid-based navigation task.
Each state factor s [ i ] represents a cell in the grid which can be empty ( s [ i ] = 0) or occupied ( s [ i ] = 1) .
In the latter case, a randomly sampled box from the set { red box , yellow box , black box } , occupies its place.
We expect Assumption 2 to hold in this case as pixels emitted from the same factor come from the same object and hence will be correlated.
More speciﬁcally, if one pixel is red in color, then another pixel from the same cell will also be red as the object occupying the cell is a red box.
This assumption does not remove the key challenge in identifying factors.
As atoms from different factors can still be dependent due to actions and state distributions from previous time steps.
Model Class.
We use two regressor classes F and G .
The ﬁrst regressor class F : X × X → [0 , 1] takes a pair of atoms and outputs a scalar in [0 , 1] .
To deﬁne the second class, we ﬁrst deﬁne a decoder class Φ : X ∗ → { 0 , 1 } .
We allow this class to be deﬁned 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 deﬁne a different model class for different numbers of atom.
We also deﬁne a model class U :
X ×
A × { 0 , 1 } → [0 , 1] .
Finally, we deﬁne the regressor class
G : X × A × X ∗ → [0 , 1] as { ( x, a, ˇ x ) (cid:55)→ u ( x, a, φ (ˇ x )) | u ∈ U , φ ∈ Φ } .
We assume
F and G are ﬁnite classes and derive sample complexity guarantees which scale as log |F| and log |G| .
However, since we only use uniform convergence arguments extending the guarantees to other statistical complexity measures such as Rademacher complexity is straightforward.
Let Π all : S
→ A denote the set of all non-stationary policies of this form.
We then deﬁne the class of policies Π :
X → A by { x (cid:55)→ ϕ
( φ (cid:63) ( x ))
| ∀ ϕ ∈ Π all } , which we use later to deﬁne our task.
We use P π [ E ] to denote probability of an event E under the distribution over episodes induced by policy π .
Computational Oracle.
We assume access to two regression oracles REG for model classes F and G .
Let D 1 be a dataset of triplets ( x [ u ] , x [ v ] , y ) where u, v denote two different atoms and y ∈ { 0 , 1 } .
Similarly, let
D 2 be a dataset of quads ( x, a, x (cid:48) , y ) where x ∈ X , a ∈ A , ˇ x ∈ X ∗ , and y ∈ { 0 , 1 } .
Lastly, let (cid:98)
E D [ · ] denote the empirical mean over dataset D .
The two computational oracles compute:
REG ( D 1 , F )=
arg min f ∈F (cid:98)
E D(cid:104)
( f ( x [ u ] , x [ v ]) − y )
(cid:105) , REG ( D 2 , G )=
arg min g ∈G N (cid:98)
E D(cid:104)
( g ( x, a, ˇ x ) − y )(cid:105) .
We also assume access to a ∆ pl -optimal planning oracle planner .
Let (cid:98)
S =
( (cid:98)
S 1 , · · · , (cid:98)
S h ) be a learned state space and (cid:98)
T =
( (cid:98)
T 1 , · · · , (cid:98)
T H ) with (cid:98)
T h : (cid:98)
S h − 1 ×
A → ∆( (cid:98)
S h ) be the learned dynamics, and (cid:98)
R : (cid:98)
S × A → [0 , 1] be a given reward function.
Let ϕ : (cid:98)
S
→ A be a policy and V ( ϕ ; (cid:98)
T, (cid:98)
R ) be the policy value.
Then for any ∆ pl > 0 the output of planner ˆ ϕ =
planner
( (cid:98)
T, (cid:98)
R, ∆ pl ) satisﬁes
V ( ˆ ϕ ; (cid:98)
T, (cid:98)
R ) ≥ sup ϕ V ( ϕ ; (cid:98)
T, (cid:98)
R ) − ∆ pl , where supremum is taken over policies of type (cid:98)
S → A .
Task Deﬁnition.
We focus on a reward-free setting with the goal of learning a state decoder and estimating the latent dynamics T .
Since the state space is exponentially large, we cannot visit every state.
However, the factorization property allows us to estimate the model by reaching factor values.
In fact, we show that controlling the value of at most 2 κ factors is sufﬁcient for learning the model.
Let C ≤ k ( U ) denote the space of all sets containing at most k different elements selected from the set
U including ∅ .
We deﬁne the reachability probability η h ( K , Z ) for a given h ∈ [ H ] , K ⊆
[ d ] , and
Z ∈ { 0 , 1 } |K| , and the reachability parameter η min as:
[Equation]
Our sample complexity scales polynomially with η − 1 min .
Note that we only require that if s h
[ K ] =
Z is reachable, then it is reachable with at least η min probability, i.e., either η h ( K , Z ) =
or it is at least η min .
These requirements are similar to those made by earlier work for non-factored state space (Du et al., 2019; Misra et al., 2020).
The key difference being that instead of requiring every state to be reachable with η min probability, we only require a small set of factor values to be reachable.
For reference, if every policy induces a uniform distribution over S = { 0 , 1 } d , then probability of visiting any state is 2 − d
but the probability of two factors taking certain values is only 0 .
This gives us a more practical value for η min .
Besides estimating the dynamics and learning a decoder, we also learn an α -policy cover to enable exploration of different reachable values of factors.
We deﬁne this below:
Deﬁnition 3 (Policy Cover) .
A set of policies Ψ is an α -policy cover of S h for any α > 0 and h if:
[Equation]
D ISCOVERING E MISSION S TRUCTURE WITH C ONTRASTIVE L EARNING
Directly applying the prior work (Du et al., 2019; Misra et al., 2020) to decode a factored state from observation results in failure, as the learned factored state need not obey the transition factorization.
Instead, the key high-level idea of our approach is to ﬁrst learn the latent emission structure ch , and then use it to decode each factor individually.
We next discuss our approach for learning ch .
Reducing Identiﬁcation of Latent Emission Structure to Independence Tests.
Assume we are able to perfectly decode the latent state and estimate the transition model till time step h − 1 .
Our goal is to infer the latent emission structure ch h at time step h , which is equivalent to: given an arbitrary pair of atoms u and v , determine if they are emitted from the same factor or not.
This is challenging since we cannot observe or control the latent state factors at time step h .
Let i =
ch − 1 ( u ) and j = ch − 1 ( v ) be the factors that emit x [ u ] and x [ v ] .
If i = j , then Assumptionimplies that these atoms are dependent on each other for any roll-in distribution D ∈ ∆( S h − 1 × A ) over previous state and action.
However, if i (cid:54) =
j then deterministically setting the previous action and values of the parent factors pt ( i ) or pt ( j ) , makes x [ u ] and x [ v ] independent.
For the example in Figure 1, ﬁxing the value of s [1] , s [2] and a would make x [ u ] and x [ u (cid:48) ] independent of each other.
This observation motivates us to reduce this identiﬁcation problem to performing independence tests with different roll-in distributions D ∈ ∆( S h − 1 × A ) .
Naively, we can iterate over all subsets
K ∈ C ≤ 2 κ ([ d ]) where for each K we create a roll-in distribution such that the values of s h − 1 [ K ] and the action a h − 1 are ﬁxed, and then perform independence test under this distribution.
If two atoms are independent then there must exist a K that makes them independent.
Otherwise, they should always be dependent by Assumption 2.
However, there are two problems with this approach.
Firstly, we do not have access to the latent states but only a decoder at time step h − 1 .
Further, it may not even be possible to ﬁnd a policy that can set the values of factors deterministically.
We later show that our algorithm FactoRL can learn a decoder that induces a bijection between learned factors and values, and the real factors and values.
Therefore, maximizing the probability of ˆ E K ; Z = { (cid:98) φ h − 1 ( x h − 1 )[ K ] = Z} for a set of learned factors K and their values
Z , implicitly maximizes the probability of E K (cid:48) ; Z (cid:48)
= { s h − 1 [ K (cid:48) ]
=
Z (cid:48) } for corresponding real factors K (cid:48) and their values Z (cid:48) .
Since the event ˆ E K ; Z is observable we can use rejection sampling to increase its probability sufﬁciently close to 1 which makes the probability of
E K (cid:48) ; Z (cid:48) close to 1.
The second problem is to perform independence tests in a domain agnostic setting.
Directly estimating mutual information
I ( x [ u ]; x [ v ]) can be challenging.
Instead, we propose an oraclized independence test that reduces the problem to binary classiﬁcation using noise-contrastive learning.
Oraclized Independent Test.
Here, we brieﬂy sketch the main idea of our independence test scheme and defer the details to Appendix C. We comment that the high-level idea of our independence testing subroutine is similar to Sen et al.
Suppose we want to test if two random variables Y and Z are independent.
Firstly, we construct a dataset in the following way: sample a Bernoulli random variable w ∼ Bern ( 1 / 2 ) , and two pairs of independent realizations ( y (1) , z (1) ) and ( y (2) , z (2) ) ; if w = 1 , add ( y (1) , z (1) , w ) to the dataset, and ( y (1) , z (2) , w ) otherwise.
We repeat the sampling procedure n times and obtain a dataset { ( y i , z i , w i ) } ni =1 .
Then we can ﬁt a classiﬁer that predicts the value of w i using ( y i , z i ) .
If Y and Z are independent, then ( y i , z i ) will provide no information about w i and thus no classiﬁer can do better than random guess.
However, if Y and Z are dependent, then the Bayes optimal classiﬁer would perform strictly better than random guess.
As a result, by looking at the training loss of the learned classiﬁer, we can determine whether Y and Z are dependent or not.
FactoRL : R EINFORCEMENT L EARNING IN F ACTORED B LOCK MDP S
In this section, we present the main algorithm FactoRL (Algorithm 1).
It takes as input the model classes F , G , failure probability δ > 0 , and ﬁve hyperparameters σ, η min , β min ∈ (0 , 1) and d, κ ∈ N .
We use these hyperparamters to deﬁne three sample sizes n ind , n abs , n est and rejection sample frequency k .
For brevity, we defer the exact values of these constants to Appendix D.7.
FactoRL returns a learned decoder (cid:98) φ h : X → { 0 , 1 } d h for some d h ∈ [ m ] , an estimated transition model (cid:98)
T h , learned parent (cid:98) pt h and child functions (cid:99) ch h , and a 1 / 2 -policy cover Ψ h of S h for every time step h ∈ { 2 , 3 , · · · , H } .
We use ˆ s h to denote the learned state at time step h .
Formally, ˆ s h
= ( (cid:98) φ h 1 ( x h ) , · · · , (cid:98) φ hd h
( x h )) .
In the analysis of FactoRL , we show that d h − 1 = d , and (cid:99) ch h is equivalent to ch h up to permutation with high probability.
Further, we show that (cid:98) φ h and (cid:99) ch h together learn a bijection between learned factors and their values and real factors and their values.
FactoRL operates inductively over the time steps (Algorithm 1, line 2-8).
In the h th iteration, the algorithm performs four stages of learning: identifying the latent emission structure, decoding the factors, estimating the model, and learning a policy cover.
We describe these below.
Algorithm 1 FactoRL ( F , G , δ, σ, η min , β min , d, κ ) .
RL in Factored Block MDPs.
Identifying Latent Emission Process.
The FactorizeEmission collects a dataset of observations for every policy in Ψ h − 1 and action a ∈ A (Algorithm 2, line 1-4).
Policies in Ψ h − 1 are of the type π
I ; Z where I ∈ C ≤ 2 κ ([ d h − 1 ]) and Z ∈ { 0 , 1 } |I| .
We can inductively assume π I ; Z to be maximizing the probability of
E I ; Z = { ˆ s h − 1 [ I ] =
Z} .
If our decoder is accurate enough, then we hope that maximizing the probability of this event in turn maximizes the probability of ﬁxing the values of a set of real factors.
However, it is possible that P π I ; Z (ˆ s h − 1 [ I ]
=
Z ) is only O ( η min ) .
Therefore, as explained earlier, we use rejection sampling to drive the probability of this event close to 1.
Formally, we deﬁne a procedure RejectSamp ( π I ; Z , E I ; Z , k ) which rolls-in at time step h − 1 with π
I ; Z to observe x h − 1 (line 3).
If the event E I ; Z holds for x h − 1 then we return x h − 1 , otherwise, we repeat the procedure.
If we fail to satisfy the event k times then we return the last sample.
We use this to deﬁne our main sampling procedure x h ∼ D
I , Z ,a := RejectSamp ( π I ; Z , E I ; Z , k ) ◦ a which ﬁrst samples x h − 1 using the rejection sampling procedure and then takes action a to observe x h .
We collect a dataset of observation pairs ( x (1) , x (2) ) sampled independently from D I , Z ,a .
For every pair of atoms u, v ∈ [ m ] , we calculate if they are independent under the distribution induced by D I , Z ,a using IndTest with dataset D
I , Z ,a (line 5-7).
We share the dataset across atoms for sample efﬁciency.
If there exists at least one (
I , Z , a ) triple such that we evaluate x [ u ] , x [ v ] to be independent, then we mark these atoms as coming from different factors.
Intuitively, such an I would contain parent factors of at least ch − 1 h ( u ) or ch − 1 h ( v ) .
If no such I exists then we mark these atoms as being emitted from the same factor.
Algorithm
(cid:98) φ h − 1 , F )
1: for ( π I ; Z , a ) ∈ Ψ h − 1 × A
and i ∈ [ n ind ] do 2:
Deﬁne E
I ; Z := 1 { (cid:98) φ h − 1 ( x h − 1 )
[ I ] =
Z} 3:
Sample x (1) h , x (2) h ∼ RejectSamp ( π I ; Z , E I ; Z , k ) ◦ a // rejection sampling procedure 4:
D I ; Z ; a ← D I ; Z ; a ∪ { ( x (1) h , x (2) h ) } // initialize D I ; Z ; a = ∅ 5: for u ∈ { 1 , 2 , · · · , m − 1 } and v ∈ { u + 1 , · · · , m } do 6:
Mark u, v as coming from the same factor, i.e., (cid:99) ch − 1 h ( u ) = (cid:99) ch − 1 h ( v ) if ∀ ( I , Z , a ) 7: the oraclized independence test ﬁnds x h [ u ] , x h [ v ] as dependent using D I ; Z ; a and F return (cid:99) ch h // label ordering of parents does not matter.
Algorithm
(cid:99) ch h ) .
Child function has type (cid:99) ch h :
[ d h ] → 2 [ m ] 1: for i in [ d h ] , deﬁne ω
=
(cid:99) ch h
( i ) , D = ∅ do 2: for n abs times do // collect a dataset of real ( y = 1) and imposter ( y = 0) transitions 3:
Sample ( x (1) , a (1) , x (cid:48) (1) ) , ( x (2) , a (2) , x (cid:48) (2) ) ∼ Unf (Ψ h − 1 ) ◦
Unf ( A ) and y ∼ Bern ( 12 ) 4:
If y = 1 then D ← D ∪ ( x (1) , a (1) , x (cid:48) (1) [ ω ] , y ) else D ← D ∪ ( x (1) , a (1) , x (cid:48) (2) [ ω ] , y ) 5: (cid:98) u i , (cid:98) φ i =
REG ( D , G ) // train the decoder using noise-contrastive learning return (cid:98) φ : X → { 0 , 1 } d h
where for any x ∈ X and i ∈
[ d h ] we have (cid:98) φ ( x )[ i ] = (cid:98) φ i ( x [ (cid:99) ch
h
( i )]) .
Decoding Factors.
LearnDecoder partitions the set of atoms into groups based on the learned child function (cid:99) ch h
(Algorithm 3).
For the i th group ω , we learn a decoder (cid:98) φ hi : X (cid:63) → { 0 , 1 } by adapting the prediction problem of Misra et al.
(2020) to Factored Block MDP setting.
We deﬁne a sampling procedure ( x, a, x (cid:48) )
∼ Unf (Ψ h − 1 ) ◦
Unf ( A ) where x is observed after roll-in with a uniformly selected policy in Ψ h − 1 till time step h − 1 , action a is taken uniformly, and x (cid:48) ∼ T ( · | x, a ) (line 3).
We collect a dataset
D of real and imposter transitions.
A single datapoint in D is collected by sampling two independent transitions ( x (1) , a (1) , x (cid:48) (1) ) , ( x (2) , a (2) , x (cid:48) (2) ) ∼ Unf (Ψ h − 1 ) ◦
Unf ( A ) and a Bernoulli random variable y ∼ Bern ( 1 / 2 ) .
If y = 1 then we add the real transition ( x (1) , a (1) , x (cid:48) (1) [ ω ] , y ) to D , otherwise we add the imposter transition ( x (1) , a (1) , x (cid:48) (2) [ ω ] , y ) (line 4).
The key difference from Misra et al.
(2020) is our use x (cid:48)
[ ω ] instead of x (cid:48) which allows us to decode a speciﬁc latent factor.
We train a model to predict the probability that a given transition ( x, a, x (cid:48) [ ω ]) is real by solving a regression task with model class G (line 5).
The bottleneck structure of G allows us to recover a decoder (cid:98) φ i from the learned model.
The algorithm also checks for the special case where a factor takes a single value.
If it does, then we return the decoder that always outputs 0 , otherwise we stick with (cid:98) φ i .
For brevity, we defer the details of this special case to Appendix D.2.2.
The decoder for the h th timestep is given by composition of decoders for each group.
Algorithm 4 EstModel (Ψ h − 1 , (cid:98) φ h − 1 , (cid:98) φ h ) .
(cid:98) solution following: (
(cid:48)
s h
ˆ s =
ˆ s h − 1 ) argmin
[Equation]
(cid:12)(cid:12)(cid:98)
P (ˆ s (cid:48)
[ k ] | ˆ s [ I ] =
u, ˆ s [ J 1 ] =
w 1 , a ) − (cid:98)
P (ˆ s (cid:48)
[ k ] | ˆ s [ I ] =
u, ˆ s [ J 2 ]
= w 2 , a ) (cid:12) (cid:12)(cid:12)
TV 5: Deﬁne (cid:98)
T h (ˆ s (cid:48) | ˆ s, a ) = (cid:81) k
(cid:98)
P (ˆ s [ k ] | ˆ s [ (cid:98) pt h ( k )] , a ) and return (cid:98)
T h , (cid:98) pt h
Estimating the Model.
EstModel routine ﬁrst collects a dataset D of n est independent transitions ( x, a, x (cid:48) )
∼ Unf (Ψ h − 1 ) ◦
Unf ( A ) (Algorithm 4, line 1).
We iterate over two disjoint sets of factors
I , J of size at most κ .
We can view I as the control set and J as the variable set.
For every learned factor k ∈ [ d h ] , factor set I , J and action a ∈ A , we estimate the model (cid:98)
P (ˆ s h [ k ] | ˆ s h − 1 [ I ] , ˆ s h − 1 [ J ] , a ) using count based statistics on dataset D (line 3).
Consider the case where the (cid:99) ch t = ch t for every t ∈
[ h ] and where we ignore the label permutation for brevity.
If I contains the parent factors pt ( k ) , then we expect the value of (cid:98)
P (ˆ s (cid:48)
[ k ] | ˆ s [ I ] , ˆ s [ J ] , a ) ≈ T k
(ˆ s (cid:48)
[ k ] | ˆ s [ pt ( k )] , a ) to not change signiﬁcantly on varying either the set J or its values.
This motivates us to deﬁne the learned parent set as the I which achieves the minimum value of this gap (line 4).
When computing the gap, we take max only over those values of ˆ s [ I ] and ˆ s [ J ] which can be reached jointly using a policy in Ψ h − 1 .
This is important since we can only reliably estimate the model for reachable factor values.
The learned parent function (cid:98) pt h need not be identical to pt h even upto relabeling.
However, ﬁnding the exact parent factor is not necessary for learning an accurate model, and may even be impossible.
For example, two factors may always take the same value making it impossible to distinguish between them.
We use the learned parent function (cid:98) pt h to deﬁne (cid:98)
T h similar to the structure of T (line 5).
Learning a Policy Cover.
We plan in the latent space using the estimated model { (cid:98)
T t } ht =1 , to ﬁnd a policy cover for time step h .
Formally, for every I ∈ C ≤ 2 κ ([ d h ]) and Z ∈ { 0 , 1 } |I| , we ﬁnd a policy ˆ ϕ h IZ to reach { ˆ s h
[ I ] =
Z} using the planner (Algorithm 1, line 7).
This policy acts on the learned state space and is easily lifted to act on observations by composition with the learned decoder.
We add every policy that achieves a return of at least O ( η min ) to Ψ h (line 8).
T HEORETICAL A NALYSIS AND D ISCUSSION
In this section, we present theoretical guarantees for FactoRL .
For technical reasons, we make the following realizability assumption on the function classes F and G .
Assumption 3 (Realizability) .
For any h ∈ [ H ] , i ∈
[ d ] and distribution ρ ∈ ∆( { 0 , 1 } ) , there exists g ihρ ∈ G , such that for all ∀ ( x, a, x (cid:48) )
∈
X h − 1 ×
A × X h and ˇ x = x (cid:48)
[ ch h ( i )]
we have:
[Equation]
For any h ∈ [ H ] , u, v ∈ [ m ] with u (cid:54) = v , and any D ∈ ∆( S h ) , there exists f uvD ∈ F satisfying: ∀ s ∈ supp ( D ) , x ∈ supp ( q ( · | s )) ,
[Equation]
Assumption 3 requires the function classes to be expressive enough to represent optimal solutions for our regression tasks.
Realizability assumptions are common in literature and are in practice satisﬁed by using deep neural networks (Sen et al., 2017; Misra et al., 2020).
Theorem 1 (Main Theorem) .
For any δ > 0 , FactoRL returns a transition function
[Equation]
T h , a parent function (cid:99) ch h , a decoder (cid:98) φ h , and a set of policies Ψ h for every h ∈ { 2 , 3 , · · · , H } , that with probability at least 1 − δ satisﬁes: (i) (cid:99) ch
h is equal to ch h upto permutation, (ii) Ψ h is a 1 / 2 policy cover of S h , (iii)
For every h ∈ [ H ] , there exists a permutation mapping θ h : { 0 , 1 } d → { 0 , 1 } d such that for every s ∈ S h − 1 , a ∈ A , s (cid:
) ∈ S h and x (cid:48)
∈
X h we have:
[Equation]
and the sample complexity is poly ( d 16 κ , |A| , H, 1 η min , 1 δ ,
[Equation]
, 1 σ , ln m, ln |F| , ln |G| ) .
Discussion.
The proof and the detailed analysis of Theorem 1 has been deferred to Appendix C-D.
Our guarantees show that FactoRL is able to discover the latent emission structure, learn a decoder, estimate the model, and learn a policy cover for every timestep.
We set the hyperparmeters in order to learn a 1 / 2 -policy cover, however, they can also be set to achieve a desired accuracy for the decoder or the transition model.
This will give a polynomial sample complexity that depends on this desired accuracy.
It is straightforward to plan a near-optimal policy for a given reward function in our learned latent space, using the estimated model and the learned decoder.
This incurs zero sample cost apart from the samples needed to learn the reward function.
Our results show that we depend polynomially on the number of factors and only logarithmic in the number of atoms.
This appeals to real-world problem where d and m can be quite large.
We also depend logarithmically on the size of function classes.
This allows us to use exponentially large function classes, further, as stated before, our results can also be easily extended to Rademacher complexity.
Our algorithm only makes a polynomial number of calls to computational oracles.
Hence, if these oracles can be implemented efﬁciently then our algorithm has a polynomial computational complexity.
The squared loss oracles are routinely used in practice, but planning in a fully-observed factored MDP is EXPTIME-complete (see Theorem 2.24 of Mausam (2012)).
However, various approximation strategies based on linear programming and dynamic programming have been employed succesfully (Guestrin et al., 2003).
These assumptions provide a black-box mechanism to leverage such efforts.
Note that all computational oracles incur no additional sample cost and can be simply implemented by enumeration over the search space.
Comparison with Block MDP Algorithms.
Our work is closely related to algorithms for Block
MDPs, which can be viewed as a non-factorized version of our setting.
Du et al.
(2019) proposed a model-based approach for Block MDPs.
They learn a decoder for a given time step by training a classiﬁer to predict the decoded state and action at the last time step.
In our case, this results in a classiﬁcation 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 classiﬁer to distinguish between real and imposter transitions.
As optimal policies for factored MDPs do not factorize, therefore, a model-free approach is unlikely to succeed (Sun et al., 2019).
Feng et al.
(2020) proposed another approach for solving Block MDPs.
They assume access to a purely unsupervised learning oracle, that can learn an accurate decoder using a set of observations.
This oracle assumption is signiﬁcantly stronger than those made in Du et al.
(2019) and Misra et al.
(2020), and reduces the challenge of learning the decoder.
Crucially, these three approaches have a sample complexity guarantee which depends polynomially on the size of latent state space.
This yields an exponential dependence on d when applied to our setting.
It is unclear if these approaches can be extended to give polynomial dependence on d .
For general discussion of related work see Appendix B.
Proof of Concept Experiments.
We empirically evaluate FactoRL to support our theoretical results, and to provide implementation details.
We consider a problem with d factors each emitting 2 atoms.
We generate atoms for factor s
[ i ] , by ﬁrst deﬁning a vector z i =
[1 , 0] if s [ i ] = 0 , and z i = [0 , 1] otherwise.
We then sample a scalar Gaussian noise g i with 0 mean and 0.1 standard deviation, and add it to both component of z i .
Atoms emitted from each factor are concatenated to generate a vector z ∈ R 2 d .
The observation x is generated by applying a ﬁxed time-dependent permutation to z to shufﬂe atoms from different factors.
This ensures that an algorithm cannot ﬁgure out the children function by relying on the order in which atoms are presented.
We consider an action space
A = { a 1 , a 2 , · · · , a d } and non-stationary dynamics.
For each time step t ∈ [ H ] , we deﬁne σ t as a ﬁxed permutation of { 1 , 2 , · · · , d } .
Dynamics at time step t are given by: T t ( s t +1 | s t , a ) = (cid:81) di =1 T ti ( s t +
[ i ] | s t [ i ] , a ) , where T ti ( s t +
[ i ] = s t
[ i ] | s t [ i ] , a ) = 1 for all a (cid:54)
= a σ t ( i ) , and
T ti ( s t +
[ i ] = 1 − s t [ i ] | s t
[ i ] , a σ t ( i ) ) = 1 .
We evaluate on the setting d = 10 and H = 10 .
We implement model classes F and G using feed-forward neural networks.
Speciﬁcally, for G we apply the Gumbel-softmax trick to model the bottleneck following Misra et al.
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 Appendix F.
For each time step, we collect 20,000 samples and share them across all routines.
This gives a sample complexity of 20 ,
× 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 .
This is in sync with our theoretical ﬁndings and demonstrates the empirical use of
FactoRL .
We will make the code available at: https://github.com/cereb-rl .
Conclusion.
We introduce Factored Block MDPs that model the real-world difﬁculties of richobservation environments with exponentially large latent state spaces.
We also propose a provable
RL algorithm called FactoRL for solving a large class of Factored Block MDPs.
We hope the setting and ideas in FactoRL will stimulate both theoretical and empirical work in this important area.
A PPENDIX O RGANIZATION
This appendix is organized as follows. • Appendix A provides a list of notations used in this paper
• Appendix B covers related work • Appendix C describes the independence test algorithm and its sample complexity guarantees • Appendix D provides sample complexity guarantees for FactoRL • Appendix E provides list of supporting results • Appendix F provides details of the experimental setup and optimization
A N OTATIONS
We present notations and their deﬁnition in Table 1.
In general, calligraphic notations represent sets.
All logarithms are with respect to base e .
B R ELATED W ORK
There is a rich literature on sample-efﬁcient reinforcement learning in tabular MDPs with a small number of observed states (Brafman & Tennenholtz, 2002; Strehl et al., 2006; Kearns & Singh, 2002;
Jaksch et al., 2010; Azar et al., 2017; Jin et al., 2018).
While recent state-of-the-art results along this line achieve near-optimal sample complexity, these algorithms do not exploit the latent structure in the environment, and therefore, cannot scale to many practical settings such as rich-observation environments with possibly a large number of factored latent states.
In order to overcome this challenge, one line of research has been focusing on factored MDPs (Kearns & Koller, 1999; Guestrin et al., 2002; 2003; Strehl et al., 2010), which allow a combinatorial number of observed states with a factorized structure.
Planning in factored MDPs is EXPTIME-complete (Mausam, 2012) yet often tractable in practice, with factored MDPs statistically learnable with polynomial samples in the number of parameters that encode the factored MDP (Osband & Van Roy, 2014; Li et al., 2011).
There has also been several empirical works that either focus on the factored state space setting (e.g., Kim & Mnih (2018); Thomas et al.
(2018); Laversanne-Finot et al.
Miladinovi´c et al.
(2019)), or the factored action space setting (e.g.,
He et al.
; Sharma et al.
However, these works do not directly address our problem and do not provide sample complexity guarantees.
Another line of work focuses on exploration in a rich observation environment.
Empirical results (Tang et al., 2017; Chen et al., 2017; Bellemare et al., 2016; Pathak et al., 2017) have achieved inspiring performance on several RL benchmarks, while theoretical works (Krishnamurthy et al., 2016; Jiang et al., 2017) show that it is information-theoretically possible to explore these environments.
As discussed before, recent works of Du et al.
(2019), Misra et al.
(2020) and Feng et al.
(2020) provide computationally and sample efﬁcient algorithms for Block MDP which is a rich-observation setting with a latent non-factored state space.
Nevertheless, this line of results crucially relies on the number of latent states being relatively small.
Finally, we comment that the contrastive learning technique used in this paper has been used by other reinforcement learning algorithms for learning feature representation (e.g., Kim et al.
Nachum et al.
Srinivas et al.
without theoretical guarantee.
C I NDEPENDENCE TESTING USING N OISE C ONTRASTIVE E STIMATION
In this section, we introduce the independence testing algorithm, Algorithm 5 and provide its theoretic guarantees.
Algorithm 5 will be used in Algorithm 2 as a subroutine for determining if two atoms are emitted from the same latent factor.
We comment that the high-level idea of Algorithm 5 is similar to
Sen et al.
(2017), which reduces independence testing to regression by adding imposter samples.
[Table]
Given an ordered set U and k ∈ Z ≥ 0 , this denotes the set of all ordered subsets of
U with at most k elements including the empty set.
[Table]
For a given n -dimensional vector u and I
[Table]
( i 1 ,
[Table]
k
[Table]
[Table]
u
[Table]
I
[Table]
concatenation of two ordered sets I and J
[Table]
Set of all possible distributions over the set U .
[Table] Set of atoms [Table] observations
[Table]
Set of all observations reachable at time step h
[Table]
Number of atoms in the observation
[Table]
Observation consisting of m atoms denoted by ( x [1] , · · · , x [ m ]) .
[Table]
of latent states
[Table]
latent states reachable at time step h
[Table]
Number of latent factors
[Table]
Latent state represented by a vector in { 0 , 1 } d .
[Table]
Decoder function which maps an observation x ∈ X to latent state s ∈ S .
[Table]
Decoder function which maps x ∈ X to its i th latent state value.
Formally, ∀ x ∈ X and i ∈ [ d ] , φ (cid:63)i ( x ) =
φ (cid:63) ( x )
[ i ] .
A Action space
[Table]
Transition function on the latent dynamics H
Horizon of the problem denoting the maximum number of actions an agent takes in a single episode.
[Table]
[ d ] → 2 [ d ]
Parent function for time step h .
We drop h when it is clear.
ch h :
[ d ] → 2 [ m ]
Child function for time step h .
By deﬁnition, for any i, j ∈ [ d ] and i (cid:54) =
j we have ch h ( i ) ∩ ch h ( j ) = ∅ .
We drop h when clear.
[Equation]
Emission function that generates x ∼ q ( · | s ) given s ∈ S .
T i i th product term in the transition function.
Formally, for s, s (cid:48) ∈ S and a ∈ A we have [Equation]
[Table]
i th product term in emission function.
Formally for x ∈ X with φ (cid:63) ( x )
=
s we have
[Equation]
F , G
Regressor classes.
We will denote individual member of the class as
[Equation]
and
[Equation]
Φ :
[Equation]
Decoder
[Equation]
C. A LGORITHM DESCRIPTION
Let D ∈ ∆( S h − 1 × A ) be our roll-in distribution that induces a probability distribution P
D ∈ ∆( X h ) over observations at time step h .
Let u, v ∈ [ m ] be two different atoms, and P
D ( x [ u ] , x [ v ]) , P D ( x h [ u ]) and P D ( x h [ v ]) be the joint and marginal distributions over x h [ u ] , x h [ v ] with respect to roll-in distribution D .
The goal of our algorithm to determine if x h [ u ] and x h [ v ] are independent under P D .
We solve this task using the IndTest algorithm (Algorithm 5) which takes as input a dataset of observation pairs D
= { ( x ( i, 1) , x ( i, 2) ) }
ni =1 where x ( i, 1) , x ( i, 2) ∼ P D ( · , · ) , and a scalar β ∈ (0 , 1) .
We use D to create a dataset D train of real and imposter atom pairs ( x [ u ] , x [ v ]) .
This is done by taking every datapoint in D and sampling a Bernoulli random variable
z i ∼ Bern ( 1 / 2 ) (line 1).
If z i =
then we add the real pair ( x ( i, 1) [ u ] , x
( i, 1) [ v ] , 1) to D train (line 4), otherwise we add the imposter pair ( x ( i, 1) [ u ] , x ( i, 2) [ v ] , 0) (line 6).
We train a classiﬁer to predict if a given atom pair ( x [ u ] , x [ v ]) is real or an imposter (line 7).
The Bayes optimal classiﬁer for this problem is given by:
[Equation]
If x [ u ] and x [ v ] are independent then we have P D ( x [ u ] , x [ v ]) =
P D ( x [ u ]) P D ( x [ v ]) everywhere on the support of P D .
This implies f (cid:63)D ( x ) = 12 and its training loss will concentrate around 0 .
Intuitively, this can be interpreted as the classiﬁer having no information to tell real samples from imposter samples.
However, if x [ u ] , x [ v ] are dependent and (cid:107) P D ( x [ u ] , x [ v ]) − P D ( x [ u ]) P D ( x [ v ]) (cid:107) TV ≥
β then we can show the training loss of f (cid:63) is less than 0 .
25 − O ( β 2 ) with high probability.
The remainder of this section is devoted to a rigorous proof for this argument.
C. A NALYSIS OF A LGORITHM
Before analyzing Algorithm 5, we want to slightly simplify the problem in terms of notations.
We introduce two simple notations X and Y which represents the random variables x [ u ] and x [ v ] , respectively.
We will simply use D to denote the joint distribution of X and Y .
Deﬁne D train to be the distribution of the training data ( X train , Y train , z ) produced in Algorithm 5.
It’s easy to verify that
[Equation]
Suppose the distribution D is specially designed such that at least one of the following hypothesis holds (which can be guaranteed when we invoke Algorithm 5):
[Equation]
In the remaining part, we will prove that Algorithm 5 can correctly distinguish between H 0 and Hwith high probability.
C.2. T WO PROPERTIES OF THE B AYES OPTIMAL CLASSIFIER
Our ﬁrst lemma shows that the Bayes optimal classiﬁer for the optimization problem in line 7 is a constant function equal to 1 / 2 if X and Y are independent.
Lemma 1 (Bayes Optimal Classiﬁer for Independent Case) .
In Algorithm 5, if X and Y (atoms u and v ) are independent under distribution D , then for the optimization problem in line 7 , the Bayes optimal classiﬁer is given by:
[Equation]
Proof.
From Bayes rule we have:
[Equation]
, where the last identity uses D train ( z = 1)
=
D train ( z = 0)
= 1 / 2 .
When z = 0 , we collect X train and Y train from two independent samples.
Therefore, we have
D train ( X train , Y train | z = 0)
=
D ( X train )
D ( Y train ) .
When z = 1 , using the fact that X train and Y train are independent under distribution D , we also have
D train ( X train , Y train | z = 1) =
D ( X train , Y train )
=
D ( X train )
D ( Y train ) .
Consequently, f (cid:63) ( X train , Y train ) ≡ 1 / 2 .
Our second lemma provides an upper bound for the expected training loss of the Bayes Optimal
Classiﬁer.
Later we will use this lemma to show the training loss is less than 0 .
25 − O ( β 2 ) with high probability when H 0 holds.
Lemma 2 (Square Loss of the Bayes Optimal Classiﬁer) .
In Algorithm 5 line 7, the Bayes optimal classiﬁer has expected square loss
[Equation]
Proof.
Recall the formula of the Bayes optimal classiﬁer in Lemma
[Equation]
Plugging it into the square loss, we obtain
[Equation]
(cid:21)(cid:19) ≤ 4 − (cid:18) 2 E D (cid:20) (cid:12)(cid:12)(cid:12)(cid:12) 2 − D ( X, Y ) D ( X ) D
( Y )
+ D ( X, Y ) (cid:12)(cid:12)(cid:12)(cid:12) (cid:21)(cid:19) .
C.2. T HREE USEFUL LEMMAS
To proceed, we need to take a detour and prove three useful technical lemmas.
Lemma 3.
Let µ and ν be two probability measures deﬁned on a countable set X .
If (cid:107) µ − ν (cid:107) TV ≥ c , then
E x ∼ µ (cid:12)(cid:12)(cid:12)(cid:12) µ ( x ) µ ( x ) + ν ( x ) −  (cid:12)(cid:12)(cid:12)(cid:12) ≥ c 4 .
Proof.
[Equation]
Lemma 4.
Fix δ ∈ (0 , 1) .
Then with probability at least 1 − δ , we have
[Equation]
where D train is the training set consisting of n i.i.d.
samples sampled from D train , ˆ f is the empirical minimizer of L ( f, D train ) over F , f (cid:63) is the population minimizer, and C ( F , δ ) :
= ln |F| δ is the complexity measure of function class
F .
Proof.
By Hoeffding’s inequality and union bound, with probability at least 1 − δ , for every f ∈ F , we have
[Equation]
Because ˆ f is the empirical optimizer,
[Equation]
Because f (cid:63) is the population optimizer,
[Equation]
Combining the two inequalities above, we ﬁnish the proof.
For notational convenience, we introduce the following factor distribution D factor deﬁned on the same domain of ( X, Y, z ) : [Equation] factor
[Equation]
Lemma 5.
Suppose F contains the constant function f ≡ 1 / 2 .
Then with probability at least 1 − δ , we have (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
Proof.
By Lemma 4, with probability at least 1 − δ , we have
[Equation]
Noticing that L is bounded by 1 , we have for every f ∈ F
[Equation]
where E D factor L ( f, D train ) deﬁnes the expected loss of f over D train where D train consists of samples i.i.d.
sampled from D factor .
Since y is a symmetric Bernoulli r.v. independent of ( X, Y ) under distribution D factor , we have
[Equation]
Using the inequality | min f L 1 ( f ) − min f L 2 ( f ) | ≤ max f | L 1 ( f ) − L 2 ( f ) | for any functionals L 1 , L 2 , along with (4) and (3) we bound the minimum loss under distribution D train as:
[Equation]
Combining (5) with (2) completes the proof.
C.2. M AIN THEOREM FOR A LGORITHM
Finally, we are ready to state and prove the main theorem for Algorithm 5.
Theorem 2.
Under the realizability assumption and n ≥ Ω( C ( F ,δ ) β 4 ) , Algorithm 5 can distinguish
[Equation]
correctly with probability at least 1 − δ .
Proof.
If H 1 holds, by Lemma 5, we have the following lower bound for the training loss of the empirical minimizer,
[Equation]
In contrast, if H 0 is true, applying Lemma 4, we obtain
[Equation]
Invoke Lemma 2 and Lemma 3,
[Equation]
Therefore, under H 0 , the training loss of the empirical minimizer is upper bounded as below
[Equation]
Plugging n ≥
O
( C ( F ,δ ) β 4 ) into (6) and (7), we complete the proof.
D T HEORETICAL A NALYSIS OF FactoRL
In this section, we provide a detailed theoretical analysis of FactoRL .
The structure of the algorithm is iterative making an inductive case argument appealing.
We will, therefore, make an induction hypothesis for each time step that we will guarantee at the end of the time step.
Induction Hypothesis.
We make the following induction assumption for FactoRL under Assumption 1-3 and across all time steps.
For all t ∈ { 2 , 3 , · · · , H } , at the end of time step t (Algorithm 1, line 8), FactoRL ﬁnds a child function
[Equation]
transition function
[Equation]
T [Equation] 0 , 1 } d × A → { 0 , 1 } d , and a set of policies Ψ t satisfying the following:
IH.
[Equation]
ch t : [ d ] → 2 m and ch t : [ d ] → 2 m are same upto relabeling, i.e, for all u, v ∈ [ m ] we have (cid:99) [Equation] ch
[Equation]
1 t ( v ) if and only if ch − 1 t ( u ) = ch − 1 t ( v ) .
Note that a child function is invertible by deﬁnition.
We can ignore this label permutation and assume (cid:99)
ch t = ch t for cleaner expressions.
This can be done without any effect.
We will assume (cid:99) ch t = ch t when stating the next three induction hypothesis.
IH.
There exists a permutation mapping θ t : { 0 , 1 } d → { 0 , 1 } d and (cid:37) ∈ (0 , 12 d ) such that for every i ∈ [ d ] and s ∈ S t we have:
[Equation]
The two distributions are independent of the roll-in distribution at time step t .
The ﬁrst one holds as [Equation]
φ t ( x t )
[ i ] only depends upon the value x t [
[Equation]
ch t [Equation] i )]
[Equation]
x t [ ch t ( i )] which only depends on s [ i ] .
The second one holds as x t is independent of everything else given s t .
The form of (cid:37) will become clear the end of analysis.
IH.
For every s ∈ S t − 1 , s (cid:48)
∈ S t and a ∈ A we have:
[Equation]
where ∆ est , ∆ app > 0 denote estimation and approximation errors whose form will become clear at the end of analysis.
IH.
For every s ∈ S t and K ∈ C ≤ 2 κ ([ d ]) , let Z
= s [ K ] and
[Equation]
Z = θ − 1 t ( s )[ K ] , then there exists a policy π
K (cid:98) Z ∈ Ψ t such that:
[Equation]
Base Case.
In the base case ( t = 1) , we have a deterministic start state.
Therefore, we can without loss of generality assume a single factor and deﬁne (cid:99) ch 1 [1] = m .
As we can also deﬁne ch 1 [1] =
[ m ] without loss of generality, therefore, this trivially satisﬁes the induction hypothesis 1.
We deﬁne (cid:98) φ 1 : X → [0] d
(Algorithm 1, line 1).
This satisﬁes induction hypothesis 2 with θ 1 being the identity map.
The induction hypothesis 3 is vacuous since there is no transition function before time step 1.
For the last condition, we have for any K , Z = [0] |K| and (cid:98)
Z = [0] |K| .
For any policy π we have
P π ( s 1 [ K ] = Z )
=
P π ( (cid:98) φ 1 ( x 1 )
[ K ] = (cid:98)
Z )
= 1 ≥ η min 2 .
Note that we never take any action from this policy, therefore, we can simply deﬁne Ψ 1 = ∅ .
D. G RAPH STRUCTURE IDENTIFICATION
In this section, we analyze the performance of Algorithm 2, given as input Ψ h − 1 , (cid:98) [Equation] [Equation] and n .
We will analyze the performance a ﬁxed pair of atoms u, v ∈ [ m ] and then apply the full result using union bound.
We ﬁrst state the result for the rejection sampling.
Lemma 6.
For policy π I ; (cid:98) Z ∈ Ψ h − 1 , event E I ;
(cid:98) Z
= { ˆ s h − 1 [ I ] = (cid:98)
Z} and k ∈ N , let D reject I ;
(cid:98) Z :=
RejectSamp ( π I ; (cid:98) Z , E I ; (cid:98) Z , k ) be the distribution induced by our rejection sampling procedure.
Let
Z = [Equation]
Z ) denote the real factor values corresponding [Equation]
Z .
Then we have:
[Equation]
Proof.
From IH.4 we have P
π I ;
[Equation]
s h
[Equation]
1 [ I ] = Z ) ≥ η min 2 .
This implies:
[Equation]
(using IH.2 and IH.4).
Let a = P
π I ; (cid:98) Z (ˆ s h − 1 [ I ] = (cid:98)
Z ) be the acceptance probability of event E I ; (cid:98) Z .
then it is easy to see that the probability of the event occurring under D reject
I ; (cid:98) Z is:
P D reject I ;
[Equation]
Z
[Equation]
We express the desired failure probability as shown:
[Equation]
We bound the two terms below:
[Equation]
Combining Equation
[Equation]
Equation 10 and Equation 11 we get:
[Equation]
We now analyze the situation for a given pair of atoms.
Recall for any distribution D ∈ ∆( S h − 1 ) and a ∈ A , we denote D ◦ a as the distribution over S h where s (cid:48) ∼ D ◦ a is sampled by sampling s ∼ D and then s (cid:48) ∼ T ( .
| s, a ) .
We want to derive roll-in distributions at time step h , such that atoms coming from the same parent satisfy hypothesis H 0 and atoms coming from different parents satisfy hypothesis H 1 under this roll-in distribution.
This will allow us to use independence test to identify the parent structure in the emission process.
Speciﬁcally, we consider the roll-in distributions induced by D reject
I ; (cid:98)
Z ◦ a for some sets I , J and action a .
Instantiating the deﬁnition of these hypothesis from Appendix C, with these roll-in distributions and setting β = β min gives us:
[Equation]
Lemma 7 (Same Factors) .
If for two atoms u, v we have ch − 1 h ( u ) = ch − 1 h ( v ) , i.e., they are from the same factor then the hypothesis H 0 is true for D ◦ a for any D ∈ ∆( S h − 1 ) and a ∈ A .
In particular, this is true for D reject
I ; (cid:98)
Z ◦ a for any choice of sets I , J and action a .
Proof.
Follows trivially from Assumption 2.
Lemma 8 (Different Factors) .
If for two atoms u, v we have ch − 1 h ( u ) =
i and ch − 1 h ( v ) =
j and i (cid:54)
=
j , then if I contains pt h ( i ) ∪ pt h ( j ) , then for (cid:37) ≤ β 2 min 1200 and k ≥ 8 η min ln (cid:16) 30 β min (cid:17) , the hypothesis
H 1 holds for D reject
I ; (cid:98)
Z ◦ a for any (cid:98)
Z such that π I ; (cid:98) Z ∈ Ψ
h − 1 and a ∈ A .
Proof.
Let
D (cid:48) ∈
∆( S h − 1 ) be a distribution that deterministically sets s h − 1 [ I ] =
Z .
Then it is easy to verify that P
[Equation]
for any
[Equation]
and
[Equation]
Then for any (cid:98)
Z and action a ∈ A we have using triangle inequality:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
As x h [ u ] and x h [ v ] come from different factors, therefore, we have
[Equation]
We use this to bound the three terms in the summation above.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation]
[Table]
[Equation]
The other two terms are bounded similarly which gives us:
[Equation]
We want this quantity to be less than β 2 min 100 to satisfy hypothesis H 1 .
We distribute the errors equally and use ln(1 + a ) ≤ a for all a > − 1 to get:
[Equation]
Theorem 3 (Learning
[Equation]
ch h ) .
Fix δ ind ∈ (0 , 1) .
If (cid:37) ≤ β 2 min 1200 and k ≥ 8 η min ln (cid:16)
[Equation]
β min (cid:17) and n ind ≥
O (cid:16)
[Equation]
2 |A||F| (2 ed )
[Equation]
+1 δ ind (cid:17) , then learned (cid:99) ch h is equivalent to ch h upto label permutation with probability at least 1 − δ ind .
Proof.
For any pair of atom u, v , if they are from the same factor then H 0 holds from Lemma 7 and
IndTest mark them dependent with probability at least 1 − δ .
This holds for every triplet of I , Z , a and there are at most |A| (2 ed )
2 κ +1 , of them.
Hence, from union bound we mark u, v correctly as coming from different factors with probability at least 1 − |A| (2 ed )
2 κ +1 δ .
If u and v have different factors then for any I containing the parents of both of them, and any value of Z and a , H 1 always holds from Lemma 8 and IndTest marks them as independent.
Note that
[Table]
[Figure]
[Table]
=">AAAB6nicdVDLSgNBEOyNrxhfUY9eBoPgxWU2Gk1uQS8eI5oHJEuYncwmQ2YfzMwKYcknePGgiFe/yJt/4yRZQUULGoqqbrq7vFhwpTH+sHJLyyura/n1wsbm1vZOcXevpaJEUtakkYhkxyOKCR6ypuZasE4sGQk8wdre+Grmt++ZVDwK7/QkZm5AhiH3OSXaSLejE6dfLGG7gp3auYOwjecwpFzBtSpGTqaUIEOjX3zvDSKaBCzUVBClug6OtZsSqTkVbFroJYrFhI7JkHUNDUnAlJvOT52iI6MMkB9JU6FGc/X7REoCpSaBZzoDokfqtzcT//K6ifarbsrDONEspItFfiKQjtDsbzTgklEtJoYQKrm5FdERkYRqk07BhPD1KfqftMq2c2qXb85K9cssjjwcwCEcgwMXUIdraEATKAzhAZ7g2RLWo/VivS5ac1Y2sw8/YL19AusPjZA=</latexit> h <latexit sha1_base64="PsSJ9d0Il77Fr7nM2GFC4w9rrHY=">AAAB6HicdVDLSgNBEOyNrxhfUY9eBoPgaZmNRpNb0IvHBMwDkiXMTmaTMbMPZmaFsOQLvHhQxKuf5M2/cZKsoKIFDUVVN91dXiy40hh/WLmV1bX1jfxmYWt7Z3evuH/QVlEiKWvRSESy6xHFBA9ZS3MtWDeWjASeYB1vcj33O/dMKh6Ft3oaMzcgo5D7nBJtpOZ4UCxhu4Kd2oWDsI0XMKRcwbUqRk6mlCBDY1B87w8jmgQs1FQQpXoOjrWbEqk5FWxW6CeKxYROyIj1DA1JwJSbLg6doROjDJEfSVOhRgv1+0RKAqWmgWc6A6LH6rc3F//yeon2q27KwzjRLKTLRX4ikI7Q/Gs05JJRLaaGECq5uRXRMZGEapNNwYTw9Sn6n7TLtnNml5vnpfpVFkcejuAYTsGBS6jDDTSgBRQYPMATPFt31qP1Yr0uW3NWNnMIP2C9fQIR0I0e</latexit
[Figure]
such an I will exists since the we iterate over all possible sets of size upto 2 κ .
Hence, with probability at least 1 − |A| 2 κ δ , we ﬁnd u and v to be independent for every Z and a .
Hence, our algorithm correctly will mark them as coming from different factors.
For a given u and v , we correctly predict their output with probability at least 1 − |A| (2 ed ) 2 κ +1 δ .
Therefore, using union bound we correctly output right result for each u and v with probability at least 1 −|A| m 2 (2 ed ) 2 κ +1 δ .
From Theorem 2, we require n ind ≥ O
[Equation]
ln
[Equation]
Binding |A|
[Equation]
2 ed )
[Equation] [Equation]
to δ ind then gives us the required value of n ind 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 (cid:99) ch h which is same as ch h upto label permutation.
D. L EARNING A S TATE D ECODER
We focus on the task of learning an abstraction at time step h using Algorithm 3.
We have access to ˆ ch h which is same as ch h upto label permutation.
We showed how to do this in Appendix C. We will ignore the label permutation to avoid having to complicate our notations.
This would essentially mean that we will recover a backward decoder
[Equation]
, where there is a bijection between
[Equation]
As we learn each decoder {
[Equation]
φ hi } independently of each other, therefore, we will focus on learning the decoder
[Equation]
φ hi for a ﬁxed i .
The same analysis will hold for other decoder and with application of union bound, we will establish guarantees for all decoders.
Further, since we are learning the decoder at a ﬁxed time step h , therefore, we will drop the h from the subscript for brevity.
We use additional shorthand described below and visualize some of them in Figure 2.
• s and x denote a state s h − 1 and an observation x h − 1 at time step h −• s (cid:48) and x (cid:48) denotes state s h and observation x h at time step h • s (cid:48)
[ i ] denotes i th factor of state at time step h •
ˇ s denotes s [ pt h ( i )] which is the set of parent factors of s (cid:48)
[ i ] .
Recall that from the factorization assumption, we have T ( s (cid:48)
[ i ] | s, a ) =
T i ( s (cid:48)
[ i ] | ˇ s, a ) for any s, a . • (cid:98) φ i denotes (cid:98) φ hi decoder for i th factor at time step h • ω denotes pt ( i ) which is the set of indices of atoms emitted by s (cid:48)
[ i ] .
• ˇ x denotes x (cid:48)
[ ch h ( i )] which is the collection of atoms generated by s (cid:48)
[ i ] .
• N = | Ψ
h − 1 | is the size of policy cover for previous time step.
Let D =
( x ( k ) , a ( k ) , ˇ x ( k ) , y ( k ) )
} n abs k =1 be a dataset of n abs real transitions ( y = 1) and imposter transitions ( y = 0) collected in Algorithm 3, line 2-4.
We deﬁne the expected risk minimizer (ERM) solution as:
[Equation]
Recall that by the structure of G , we have ˆ g i
=
(ˆ u i , (cid:98) φ i )
where
ˆ w i ∈ W 2 and (cid:98) φ i ∈ Φ : X ∗ → { 0 , 1 } is the learned decoder.
Our algorithm only cares about the properties of the decoder and we throw away the regressor ˆ u i .
Let D ( x, a, ˇ x ) be the marginal distribution over transitions.
We get the marginal distribution by marginalizing out the real ( y = 1) and imposter transition ( y = 0) .
We also deﬁne D ( x, a ) as the marginal distribution over x, a .
We have D ( x, a ) =
µ h − 1 ( x ) 1 |A| as both real and imposter transitions involve sampling x ∼ µ h − 1 and taking action uniformly.
Recall that µ h − 1 is generated by roll-in with a uniformly selected policy in Ψ h − 1 till time step h − 1 .
Let P ( x, a, ˇ x | y = 1) be the probability of a transition being real and P ( x, a, ˇ x | y = 0) be the probability of the transition being imposter.
We can express these probabilities as:
[Equation]
where ρ (ˇ x ) =
E ( x,a ) ∼ D
[ T (ˇ x | x, a )] is the marginal distribution over ˇ x .
We will overload the notation ρ to also deﬁne ρ ( x (cid:48) )
=
E ( x,a ) ∼ D
[ T ( x (cid:48)
| x, a )] .
Lastly, we can express the marginal distribution over transition as:
[Equation]
We start by expressing the Bayes optimal classiﬁer for problem in Equation
[Equation]
Lemma 9 (Bayes Optimal Classiﬁer) .
Bayes optimal classiﬁer g (cid:63) for problem in Equation 14 is given by:
[Equation]
Proof.
The Bayes optimal classiﬁer is given by g (cid:63) ( x, a, ˇ x ) =
P ( y = 1 | x, a, ˇ x ) which can be expressed using Bayes rule as:
[Equation]
using p ( y )
=
Bern ( 1 / 2 )
[Equation]
Theorem 4 (Decoder Regression Guarantees) .
For any given δ abs ∈ (0 , 1) and n abs ∈ N we have the following with probability at least 1 − δ abs :
[Equation]
where ∆( n abs , δ abs , |G| ) :
= cn abs ln |G| δ abs and c is a universal constant.
Proof.
This is a standard regression guarantee derived using Bernstein’s inequality with realizability (3).
For example, see Proposition 11 in Misra et al.
(2020) for proof.
Corollary 5.
For any given δ abs ∈ (0 , 1) and n abs ∈ N we have the following with probability at least 1 − δ abs :
[Equation]
Proof.
Applying Jensen’s inequality ( E [ √ Z ] ≤ (cid:112)
E [ Z ]) to Theorem 4 gives us:
[Equation]
Coupling Distribution
We introduce a coupling distribution following Misra et al.
[Equation]
We also deﬁne the following quantity which will be useful for stating our results:
[Equation]
Lemma 10.
For any ﬁxed δ abs ∈ (0 , 1) we have the following with probability at least 1 − δ abs :
[Equation]
Proof.
We deﬁne a shorthand notation E = 1 { (cid:98) φ i (ˇ x 1 )
= (cid:98) φ i (ˇ x 2 ) } for brevity.
We also deﬁne a different coupled distribution D (cid:48) coup given below:
[Equation]
where D (ˇ x | x, a ) = 12 { T (ˇ x | x, a ) + ρ (ˇ x ) } .
It is easy to see that marginal distribution of D (cid:48) coup over x, a, ˇ x 1 is same as D ( x, a, ˇ x 1 ) .
We ﬁrst use the deﬁnition of ξ (Equation 19) and g (cid:63) (Equation 9) to express their relation:
[Equation] [Equation]
The second line uses the deﬁnition of D (ˇ x | x, a ) .
We can view ρ (ˇ x 1 ) D (ˇ x 1 | x,a ) and ρ (ˇ
[Equation]
D (ˇ
[Equation]
2 | x,a ) as importance weight terms.
Multiplying both sides by E and taking expectation with respect to D (cid:48) coup then gives us:
[Equation] [Equation]
We bound the left hand side of Equation 22 as shown below:
[Equation]
Here the ﬁrst inequality follows from triangle inequality.
The second step is key where we use ˆ g i ( x, a, ˇ x 1 )
=
ˆ g i ( x, a, ˇ x 2 ) whenever E = 1 .
This itself follows from the bottleneck structure of G
where ˆ g i ( x, a, ˇ x i )
=
ˆ w i ( x, a, (cid:98) φ i (ˇ x )) .
The third step uses the symmetry of ˇ x 1 and ˇ x 2 in D (cid:48) coup whereas the fourth step uses the fact that marginal distribution of D (cid:48) coup is same as D .
Lastly, ﬁnal inequality uses E ≤ 1 and the result of Corollary 5.
Combining the derived inequality with Equation 22 proves our result.
We deﬁne the quantity P ( s (cid:48)
[ i ] =
z | D (cid:48) ) :
=
E ( s,a ) ∼ D (cid:48)
[ 1 { s (cid:48)
[ i ] =
z } ] for any distribution D (cid:48)
∈ ∆( S h − 1 × A ) .
From the deﬁnition of ρ , we have ρ ( s (cid:48)
[ i ] =
z ) =
P ( s (cid:48)
[ i ] =
z | D ) .
Intuitively, as we have policy cover at time step h − 1 and we take actions uniformly, therefore, we expect to have good lower bound on P ( s (cid:48)
[ i ] =
z | D ) for every i ∈ [ d ] and reachable z ∈ { 0 , 1 } .
Note that if z = 0 ( z = 1 ) is not reachable then it means we always have s h
[ i ] = 1 ( s h
[ i ] = 0 ) from our reachability assumption (see Section 2).
We formally prove this next which will be useful later.
Lemma 11.
For any z ∈ { 0 , 1 } such that s (cid:48)
[ i ] =
z is reachable, we have:
[Equation]
Proof.
Fix z in { 0 , 1 } .
As s (cid:48)
[ i ] =
z is reachable, therefore, from the deﬁnition of η min we have:
[Equation]
We use the derived inequality to bound P ( s (cid:48)
[ i ] =
z | D ) as shown:
[Equation]
The ﬁrst inequality uses the fact that µ h − 1 is created by roll-in with a uniformly selected policy in Ψ h − 1 which is an α policy cover.
Recall that N = | Ψ
h − 1 | .
The second inequality uses the derived result above.
Lemma 12.
For any ˇ x 1 , ˇ x 2 such that φ (cid:63)i (ˇ x 1 ) and φ (cid:63)i (ˇ x 2 ) is reachable, we have:
[Equation]
Proof.
For any x, a, ˇ x 1 , ˇ x 2 we can express ξ (Equation 19) as:
[Equation]
where we use the factorization assumption and decodability assumption.
Note that we are implicitly assuming φ (cid:63)i (ˇ x 1 ) and φ (cid:63)i (ˇ x 2 ) are reachable, for the quantity ξ (ˇ x 1 , ˇ x 2 , x, a ) to be well deﬁned.
We deﬁne D i to be the marginal distribution over S [ pt ( i )]
× A .
Taking expectation on both side gives us:
[Equation] [Equation]
φ
[Equation]
TV
The second equality uses the deﬁnition of backward dynamics P D i over S [ pt ( i )]
× A and the identity ρ ( s (cid:48)
[ i ])
=
P ( s (cid:48)
[ i ] | D ) .
If φ (cid:63)i (ˇ x 1 ) =
φ (cid:63)i (ˇ x 2 ) then the quantity on the right is 0.
Otherwise, this quantity is given by 2 (cid:107) P D
i ( ., .
| s (cid:48)
[ i ] = 1) − P D i ( ., .
| s (cid:48)
[ i ] =
(cid:107) TV .
In the later case, both s
[Equation]
[ i ] = 1 and s (cid:48)
[ i ] =
conﬁgurations are reachable, and without loss of generality we can assume P ( s (cid:48)
[ i ] =
| D ) ≥ 1 / 2 .
Our goal is to bound this term using the margin assumption (Assumption 1).
We do so using importance weight as shown below:
[Equation]
The ﬁrst step applies importance weight.
As u i has support over all reachable conﬁgurations ˇ s and actions a ∈ A , hence, we can apply importance weight.
The second step uses the deﬁnition of backward dynamics ( P D , P u i ) .
The third step uses Lemma H.1 of Du et al.
Du et al.
(see Lemma 24 in Appendix E for statement).
Finally, the last step uses Assumption 1.
We bound the two multiplicative terms below:
We have D (ˇ s, a ) =
µ h − 1 (ˇ s ) 1 |A| ≥ αη min N |A| .
The ﬁrst equality uses the fact that actions are taken uniformly and second inequality uses the fact that µ h − 1 is an α -policy cover.
As u i is the uniform distribution over S [ pt ( i )]
× A , therefore, we have u
i (ˇ s, a ) = 1 2 | pt ( i ) | |A| .
This gives us D (ˇ s,a ) u i (ˇ s,a ) ≥ αη min N 2 | pt ( i ) | .
We bound the other multiplicative term as shown below:
[Equation]
Combining the lower bounds for the two multiplicative terms and using 2 | pt ( i ) | ≥ 1 we get:
[Equation]
Lastly, recall that our desired result is given by 2 (cid:107) P D
i
[Equation]
|
[Equation]
= 1) − P D
[Equation] [Equation]
s
[Equation]
[ i ]
[Equation] [Equation]
TV whenever φ (cid:63)i (ˇ x 1 ) (cid:54) =
φ (cid:63)i (ˇ x 2 ) and 0 otherwise.
Therefore, using the derived lower bound multiplied by 1 { φ (cid:63)i (ˇ x 1 ) (cid:54) =
φ (cid:63)i (ˇ x 2 ) } gives us the desired result.
Corollary 6.
We have the following with probability at least 1 − δ abs :
[Equation]
Proof.
The proof trivially follows from applying the bound in Lemma 12 to Lemma 10 as shown below:
[Equation]
The inequality here uses Lemma 12.
The left hand side is bounded by
[Equation]
∆
[Equation]
δ abs
[Equation]
|G| )
us-
ing Lemma 10.
Combining the two bounds and rearranging the terms proves the result.
At this point we analyze the two cases separately.
In the ﬁrst case, s (cid:48)
[ i ] can be set to both 0 and 1 .
In the second case, s (cid:48)
[ i ] can only be set to one of the values and we will call s (cid:48)
[ i ] as degenerate at time step h .
We will show how we can detect the second case, at which point we just output a decoder that always outputs 0.
We analyze the ﬁrst case below.
D.2. C ASE A: WHEN s (cid:48) [ i ] CAN TAKE ALL VALUES
Corollary 6 allows us to deﬁne a correspondence between learned state (i.e., output of (cid:98) φ i ) and the actual state (i.e., output of φ (cid:63)i ).
We show this correspondence in the result.
Theorem 7 (Correspondence Theorem) .
For any state factor s (cid:48)
[ i ] , there exists exists ˆ u 0 ∈ { 0 , 1 } and ˆ u 1 = 1 − ˆ u 0 with probability at least 1 − δ abs such that:
[Equation]
where (cid:37) :
=
16 N 2 |A| α 2 η 2 min σ (cid:112) ∆( n abs , δ abs , |G| ) and ˇ x ∼ ρ , provided (cid:37) ∈ (cid:0) 0 ,(cid:1) .
Proof.
For any u, z ∈ { 0 , 1 } we deﬁne the following quantities:
[Equation]
It is easy to see that these quantities are related by: P z
= P uz + P (1 − u )
z .
We deﬁne ˆ u 0 = arg max u ∈{ 0 , 1 } P u 0 and ˆ u 1 = 1 − ˆ u 0 .
This can be viewed as the learned bit value which is in most correspondence with s [ i ] = 0 .
We will derive lower bound on P ˆ u 00 / P 0 and P ˆ u 11 / P 1 which gives us the desired result.
We ﬁrst derive the following lower bound on P ˆ u 0 0 :
[Equation]
where we use the fact that max is greater than average.
Further, for any u, z ∈ { 0 , 1 } we have:
[Equation]
We deﬁne a shorthand notation ∆ (cid:48) :
=
[Equation]
σ (cid:112) ∆( n abs , δ abs , |G| ) .
Then from Corollary 6 we have proven that P uz P u (1 − z ) ≤ ∆ (cid:48) for any u, z ∈ { 0 , 1 } .
This allows us to write:
[Equation]
where the last inequality uses Equation 24.
We will derive the same result for P ˆ u 00 / P 0 .
[Equation]
where the last inequality uses derived bound for P ˆ u 11 / P 1 .
If we assume ∆ (cid:48) ≤ P 0 P 1then we get
P
[Equation]
u
[Equation]
P
[Equation]
As P 0 + P 1 = 1 , therefore, we get P 0 P 1 = P 0 − P 20 = P 1 − P 21 .
If P 0 ≤ 12 then P 0 − P 20 ≥ P 0 2 .
Otherwise, P 0 > 12 which implies P 1 ≤ 12 and P 1 − P 21 ≥ P 1 2 .
This gives us P 0 P 1 ≥ min { P 0 2 , P 1 2 } .
Using lower bounds for P 0 and P 1 from Lemma 11 gives us P 0 P 1 ≥ αη min 2 N |A| , and allows us to write:
[Equation]
It is easy to verify that (cid:37) = 4 N |A| ∆ (cid:48) αη min .
As P ˆ u 00 P 0 =
P ( (cid:98) φ i ( x (cid:48) )
=
ˆ u 0 | s [ i ] = 0) and P ˆ u 11 P 1 = P ( (cid:98) φ i ( x (cid:48) )
=
ˆ u 1 | s [ i ] = 1) , therefore, we prove our result.
The only requirement we used is that ∆ (cid:48) ≤ P 0 P 1which is ensured if (cid:37) ∈ (cid:0) 0 , 1(cid:1) .
D.2. C ASE B: WHEN s (cid:48) [ i ] TAKES A SINGLE VALUE
We want to be able to detect this case with high probability so that we can learn a degenerate decoder that only takes value 0.
This would trivially give us a correspondence result similar to Theorem 7.
We describe the general form of the LearnDecoder subroutine in Algorithm 6.
The key difference from the case we covered earlier is line 6-10.
For a given factor i , we ﬁrst learn the model
ˆ g i containing the decoder (cid:98) φ i , as before using noise contrastive learning.
We then sample n deg iid triplets
D deg = { ( x j , a j , ˇ x j ) } n deg j =1 where ( x j , a j ) ∼ D and ˇ x j ∼ ρ (line 6-8).
Recall ˇ x = x
[ ch h ( i )] .
Next, we compute the width of prediction values over
D deg as deﬁned below:
[Equation]
If the width is smaller than a certain value then we determine the factor to be degenerate and output a degenerate decoder (cid:98) φ i := 0 , otherwise, we stick to the decoder learned by our regressor task.
The form of sample size n deg 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 classiﬁer will always output 1 / 2 and so our prediction values should be close to one another.
However, if the latent factor takes two values then the model prediction should be distinct.
Algorithm 6 LearnDecoder ( G , Ψ h − 1 , (cid:99) ch h ) .
Child function has type (cid:99) ch h :
[ d h ] → 2 [ m ] 1: for i in [ d h ] , deﬁne ω
=
(cid:99) ch h
( i ) , D = ∅ , D deg = ∅ do 2: for n abs times do // collect a dataset of real ( y = 1) and imposter ( y = 0) transitions 3:
Sample ( x (1) , a (1) , x (cid:48) (1) ) , ( x (2) , a (2) , x (cid:48) (2) ) ∼ Unf (Ψ h − 1 ) ◦
Unf ( A ) and y ∼ Bern ( 12 ) 4:
If y = 1 then D ← D ∪ ( x (1) , a (1) , x (cid:48) (1) [ ω ] , y ) else D ← D ∪ ( x (1) , a (1) , x (cid:48) (2) [ ω ] , y ) 5:
ˆ g i := (cid:98) u
i , (cid:98) φ i = REG ( D , G ) // train the decoder using noise-contrastive learning 6: for n deg times do // detect degenerate factors 7:
Sample ( x (1) , a (1) , x (cid:48) (1) ) , ( x (2) , a (2) , x (cid:48) (2) ) ∼ Unf (Ψ h − 1 ) ◦
Unf ( A ) 8:
D deg ← D deg ∪ { ( x (1) , a (1) , x (cid:48) (2) ) } 9: if max j,k ∈ [ n deg ] (cid:12)(cid:12)
ˆ g i ( x j , a j , x (cid:48)
j [ ω ]) − ˆ g i ( x k , a k , x (cid:48) k
[ ω ]) (cid:12)(cid:12) ≤ α 2 η 2 min σ| Ψ
h − 1 | 2 |A| then // max over D deg 10: (cid:98) φ i :=// output a decoder that always returnsreturn (cid:98) φ : X → { 0 , 1 } d h where for any x ∈ X and i ∈
[ d h ] we have
[Equation] [Equation]
For convenience, we deﬁne D (cid:48)
( x, a, ˇ x )
=
D ( x, a ) ρ (ˇ x ) , and so ( x j , a j , ˇ x j ) ∼ D (cid:48) .
For brevity reasons, we do not add additional qualiﬁers to differentiate x j , a j , ˇ x j from the dataset of real and imposter transitions, we used in the previous section for the regression task.
In this part alone, we will use x j , a j , ˇ x j to refer to the transitions collected for the purpose of detecting degenerate factors.
Lemma 13 (Markov Bound) .
Let { ( x j , a j , ˇ x j ) } n deg j =1 be a dataset of iid transitions sampled from D (cid:48) .
Fix a > 0 .
Then with probability at least 1 − δ abs −
[Equation]
a we have:
[Equation]
Proof.
It is straightforward to verify that for any ( x j , a j , ˇ x j ) we have D ( x j , a j , ˇ x j
[Equation]
D
[Equation]
j
[Equation]
ˇ
[Equation]
Using Corollary 5 we get:
[Equation]
Let E j denote the event {| ˆ g ( x j , a j , ˇ x j ) − g (cid:63) ( x j , a j , ˇ x j ) | ≤ a } and E j be its negation, then:
[Equation]
where the ﬁrst inequality uses union bound and the second inequality uses Markov’s inequality.
As Corollary 5 holds with probability δ abs , our overall failure probability is at most
[Equation]
Lemma 14.
For any reachable parent factor values ˇ s , action a ∈ A and reachable s (cid:48)
[ i ] ∈ { 0 , 1 } , we have D (cid:48)
(ˇ s, a, s (cid:48)
[ i ]) ≥ α 2 η 2 min N 2 |A| 2 .
Proof.
We have D (cid:48)
(ˇ s, a, s (cid:48)
[ i ])
=
µ h − 1 (ˇ s ) |A| ρ ( s (cid:48)
[ i ]) ≥ α 2 η 2 min
N 2 |A| 2 , where used the induction hypothesis IH.4 that Ψ is an α -policy cover of S h − 1 and Lemma 11.
Lemma 15 (Degenerate Factors) .
Fix a > 0 .
If s (cid:48)
[ i ] only takes a single value then with probability at least
[Equation]
we have:
[Equation]
Proof.
When s (cid:48)
[ i ] takes a single value then g (cid:63) is the constant function 12 .
For any j and k we get the following using Lemma 13 and triangle inequality.
[Equation]
Lemma 16 (Non Degenerate Factors) .
Fix a > 0 and assume n deg ≥ N 2 |A| 2 α 2 η 2 min , then we have:
[Equation]
with probability at least 1 − δ
[Equation]
−
[Equation]
Proof.
Equation 23 implies that there exists ˇ s, a such that
[Equation]
Combining this with Equation 21 we get:
[Equation]
where Equation 26 uses ρ ( s (cid:48)
[ i ] = 1) ρ ( s (cid:48)
[ i ] = 0) ≥ αη min 2 N |A| , as one of the terms is at least 1 / 2 and other can be bounded using Lemma 11.
Say we have two examples in our dataset, say { ( x 1 , a 1 , ˇ x 1 ) , ( x 2 , a 2 , ˇ x 2 ) } without loss of generality, such that φ (cid:63) ( x 1 )[ ω ] = φ (cid:63) ( x 2 )
[ ω ] =
ˇ s , action a 1 = a 2 = a , φ (cid:63)i (ˇ x 1 ) = 1 , and φ (cid:63)i (ˇ x 2 ) = 0 .
Then we have:
[Equation]
We use Lemma 13 which has a failure probability of δ abs + 2 n deg √ ∆( n abs ,δ abs , |G| ) a .
Further, we also assume that our dataset contains both (ˇ s, a, s (cid:48)
[ i ] = 1) and (ˇ s, a, s (cid:48)
[ i ] = 0) .
Probability of one of these events is given by Lemma 14.
Therefore, if n deg ≥ N 2 |A|
[Equation] [Equation]
2 min then from Lemma 25 and union bound, the probability that at least one of these transitions does not occur is given by 4 exp
[Equation]
The total failure probability is given by union bound and computes to:
[Equation]
If we ﬁx a = α
[Equation] [Equation]
σ
[Equation]
N|A| then in the two case we have:
[Equation]
Theorem 8 (Detecting Degenerate Case) .
We correctly predict if s (cid:48)
[ i ] is a degenerate factor or not when using n deg =
[Equation]
|
[Equation]
and (cid:37) ≤ α
[Equation]
, with probability at least 1 − 3 δ abs .
Proof.
The result follows by combining Lemma 16 and Lemma 15, and using the value of a described above.
These two results hold with probability at least:
[Equation]
Setting the hyperparameters to satisfy the following:
[Equation]
gives a failure probability of at most 3 δ abs .
The later condition can be expressed in terms of (cid:37) which gives us the desired bounds (see Theorem 7 for deﬁnition of (cid:37) ).
Lastly, note that setting n deg this way also satisﬁes the requirement in Lemma 16.
Lastly, note that the resultant bound on (cid:37) is much stronger than required for Theorem 7.
Therefore, we can signiﬁcantly improve the complexity bounds in the setting where there are no degenerate state factors.
D.2. C OMBINING C ASE A AND C ASE B
Theorem 8 shows that we can detect degenerate state factors with high probability.
If we have a degenerate state factor and we detect it, then correspondence theorem holds trivially.
However, if we don’t have degeneracy and we correctly predict it, then we stick our learned decoder and Theoremholds true.
These two results allows us to deﬁne a bijection between real states and learned states that we explain below.
Bijective Mapping between real and learned states For a given time step h and state bit s [ i ]
=
z , we will deﬁne ˆ u hiz as the corresponding learned state bit.
When h and i will be clear from the context then we will express this as ˆ u z .
We will use the notation ˆ s to denote a learned state at time step h − 1 and ˆ s (cid:48) to denote learned state at time step h .
Let pt ( i ) =
( i 1 , · · · , i l ) and s [ pt ( i )]
= w :=
( w 1 , w 2 , · · · w l ) , then we deﬁne ˆ w =
(ˆ u ( h − 1)
i 1 w 1 · · · ˆ u ( h − 1)
i l w l ) as the learned state bits corresponding to w .
More generally, for a given set K ∈ 2 d , we denote the real state factors as s [ K ] (or s (cid:48)
[ K ] ) and the corresponding learned state factors as ˆ s [ K ] (or ˆ s (cid:48)
[ K ] ).
We deﬁne a mapping θ h : { 0 , 1 } d → { 0 , 1 } d from learned state to real state.
We will drop the subscript h when the time step is clear from the context.
We denote the domain of θ h by (cid:98)
S h which is a subset of { 0 , 1 } d .
Note that every real state may not be reachable at time step h .
E.g., maybe our decoder outputs ˆ s (cid:48) =
(0 , 0) but that the corresponding real state is not reachable at time step h .
Figure 3 visualizes the mapping.
[Table]
+sk4FfdVIRxAjxki0V+IjFEeJYA7gvFGciJIZQpYW7FbEgVZWByKpgQvj7F/5Pbsu2c2OXr01L9Iosjjw7QITpGDjpHdXSFGqiJGLpDD+gJPVtj69F6sV4XrTkrm9lHP2C9fQLUTZLM</latexit>
(cid:98)
This
For a learned state ˆ s we have s = θ (ˆ s ) if s = ( z 1 , · · · , z d ) and ˆ s =
( u h 1 z 1 , · · · , u hdz d ) .
We would also overload our notation to write w = θ ( ˆ w ) for a given ˆ s [
K ]
= ˆ w where w = θ (ˆ s )
[ K ] , whenever factor set
K is clear from the context.
We call s (cid:48) (or s ) as reachable if s (cid:48)
∈
S h (or s ∈ S h − 1 ).
Similarly, we call ˆ s (cid:48) (or ˆ s ) as reachable if θ h
(ˆ s (cid:48) )
(or θ h − 1 (ˆ s ) ) are reachable.
For a given set of factors K , we call w ∈ { 0 , 1 } |K| as reachable for
K if there exists a reachable state with factors K taking on value w .
Similarly, we deﬁne ˆ w ∈ { 0 , 1 } |K| as reachable for a given
K if θ ( ˆ w ) is reachable for K .
We use the mapping to state the correspondence theorem for the whole state.
Corollary 9 (General Case) .
If n deg = 3 N 2 |A| 2 α
[Equation]
2 min log (cid:16)
[Equation]
(cid:17) and (cid:37) ≤ α 2 η 2 min δ abs 30 N|A| 2 log −(cid:16) 4 δ
[Equation]
(cid:17) holds, then with probability at least 1 − 3 dδ abs , we have:
[Equation]
Proof.
The ﬁrst result directly follows from being able to detect if we are in degenerate setting or not and if not, then applying Theorem 7, and if yes then result holds trivially.
This holds with probability at least 1 − 3 δ abs from Theorem 8.
Applying union bound over all d learned factors gives us success probability across all factors of at least 1 − 3 dδ abs .
The second one follows from union bound.
[Equation]
As we noticed before, our bounds can be signiﬁcantly 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 (cid:37) ∈
[Equation]
0 ,(cid:1) , then with probability at least 1 − dδ abs , we have:
[Equation]
Proof.
Same as Corollary 9 except we can directly apply Theorem 7 as we don’t need to do any expensive check for a degenerate factor.
D. M ODEL E STIMATION
Our next goal is to estimate a model
[Equation] [Equation]
∆
[Equation]
S
[Equation]
) and latent parent structure (cid:98) pt h , given roll-in distribution D ∈ ∆( S h − 1 ) , and the learned decoders { (cid:98) φ t } t ≤ h .
Recall that our approach estimates the model by count-based estimation.
Let D =
( x ( k ) , a ( k ) , x (cid:48) ( k ) ) } n est k =1 be the sample collected for model estimation.
Recall that we estimate (cid:98)
pt h
( i ) be using a set of learned factors I that we believe is (cid:98) pt h ( i ) and varying a disjoint set of learned factors J .
If (cid:98) pt ( i ) ⊆
I then we expect the learned model to behave the same irrespective of how we vary the factors J .
However, if (cid:98) pt ( i ) (cid:54)⊆ I then there exists a parent factor that on varying will different values for the learned dynamics.
Let
K =
I ∪ J be the set of all factors in the control group (
I ) and variable group ( J ) .
We will ﬁrst analyze the case for a ﬁxed i ∈ (cid:98)
S h and control group I and variable group J .
For a given ˆ v ∈ { 0 , 1 } , ˆ w ∈ { 0 , 1 } |K| and a ∈ A , we have (cid:98)
P D (ˆ s (cid:48)
[ i ]
= ˆ v | ˆ s [ K ]
= ˆ w, a ) denoting the estimate probability derived from our count based estimation (Algorithm 4, line 3).
Let P D
(ˆ s (cid:48)
[ i ]
= ˆ v | ˆ s [ K ]
= ˆ w, a ) be the probabilities that are being estimated.
It is important to note that we use subscript D for these notations as the learned states ˆ s are not Markovian and K may not contain pt ( i ) , therefore, the estimated probabilities (cid:98)
P and expected probabilities P will be dependent on the roll-in distribution D .
In order to estimate (cid:98)
P D ( .
| ˆ s h − 1 [ K ] = ˆ w, a ) we want good lower bounds on P D (ˆ s h − 1 [ K ] =
ˆ w, a ) for every a ∈ A and ˆ w reachable for K .
Our roll-in distribution D only guarantees lower bound on P
D ( s h − 1 [ K ] , a ) .
However, we can use IH.2 to bound the desired quantity below.
Lemma 17 (Model Estimation Coverage) .
If (cid:37) ≤ 12 then for all K ∈ C ≤ 2 κ ([ d ]) , a ∈ A and ˆ w ∈ { 0 , 1 } |K| reachable for K , we have:
[Equation]
A|
Proof.
We can express P D (ˆ s h − 1 [ K ] = ˆ w, a ) = 1 |A| P D (ˆ s h − 1 [ K ] = ˆ w ) as actions are taken uniformly.
Let s h − 1 =
θ h − 1 (ˆ s h − 1 ) and w =
s
h − 1 [ K ] .
We bound P D (ˆ s h − 1 [ K ] = ˆ w ) as shown:
[Equation]
where the third step uses the fact that value of learned state
ˆ s h − 1 [ k ] is independent of other decoders given the real state bit s h − 1 [ k ] .
The fourth step uses IH.2 and the fact that we have good coverage over all sets of state factors of size at most 2 κ .
Last inequality uses |K| ≤ 2 κ and (cid:37) ≤ 12 .
We now show that our count-based estimator (cid:98)
P D converges to P D and derive the rate of convergence.
Lemma 18 (Model Estimation Error) .
Fix δ est ∈ (0 , 1) .
Then with probability at least 1 − δ est for every K ∈ C ≤ 2 κ ([ d ]) , ˆ w ∈ { 0 , 1 } |K| reachable for K , and a ∈ A we have the following:
[Equation]
ˆ w, a ) − P D
(ˆ s
[Equation]
[ i ] =
[Equation]
|
[Equation]
s
[Equation] [Equation]
] = ˆ w, a ) (cid:12)(cid:12)(cid:12) ≤ 2∆ est ( n est , δ est ) , where ∆ est ( n est , δ est ) :=(cid:16) 2 κ +5 N |A| αη min n est (cid:17) 1 / 2 ln (cid:16) 4 e |A| ( ed ) 2 κ δ est (cid:17) .
Proof.
We sample n est samples by roll-in at time step h − 1 with distribution D and taking actions uniformly.
We ﬁrst analyze the failure probability for a given K , ˆ w, a .
Let E ( K , ˆ w, a ) denote the event { ˆ s h − 1 [ K ] = ˆ w, a h − 1 = a } .
If E ( K , ˆ w, a ) occurs in our dataset at least m times for some m ≥ 16 (cid:15) 2 ln( 1 / δ ) then from Corollary 15 we have
[Equation]
with probability at least 1 − δ .
Lemma 17 shows that probability of E ( K , ˆ w, a ) is at least
αη min 4 κ N |A| .
Therefore, from Lemma 28 if n est ≥ 2 2 κ +1 mN |A| αη min ln (cid:0) eδ (cid:1) then we get at least m samples of event
E ( K , ˆ w, a ) with probability at least 1 − δ .
Therefore, the total failure probability is at most 2 δ : δ due to not getting at least m samples and δ due to Corollary 15 on getting m samples.
This holds for every triplet ( K , ˆ w, a ) and Lemma 23 shows that there are at most 2( ed ) 2 κ |A| such triplets.
Hence, with application of union bound we get the desired result.
Lemma 19 (Model Approximation Error) .
For any i ∈ [ d ] , K ∈ C ≤ 2 κ ([ d ]) , s ∈ S h − 1 , a ∈ A , s (cid:48) ∈
S h , let ˆ s = θ − 1 h − 1 ( s ) and ˆ s (cid:48)
= θ − 1 h ( s (cid:48) ) .
Then we have:
[Equation]
Proof.
We will ﬁrst bound | P D ( s (cid:48)
[ i ] | ˆ s [ K ] , a ) − P D ( s (cid:48)
[ i ] | s [ K ] , a ) | and then use correspondence result (Corollary 9) to prove the desired result.
We start by expressing our conditional probabilities as ratio of joint probabilities.
[Equation]
From Lemma 29 we have:
[Equation]
, where
[Equation]
ε
[Equation]
|
[Equation]
( s [ K ] , a ) | .
We bound these two quantities below:
[Equation]
  ,
Where the ﬁrst inequality uses | a − b | ≤ max { a, b } for a, b > 0 .
We bound Term 1 below:
Term 1:
[Equation]
The key inequality here is P (ˆ s [ K ] | s h − 1 [ K ]) =
(cid:81)
[Equation]
[ k ] |
[Equation]
[ k ]) ≤ (cid:37) , as there exist at least one j ∈ K such that s h − 1 [ j ] (cid:54) =
s [ j ] and for this j we have P (ˆ s [ j ] | s h − 1 [ j ]) ≤ (cid:37) .
We bound Term 2 similarly:
Term 2:
[Equation]
A| where we use P D (ˆ s [ K ] | s [ K ]) = (cid:81) k ∈K P (ˆ s [ k ] | s [ k ]) ≥ (1 − (cid:37) ) |K| and |K| ≤ 2 κ .
This gives us ε 1 ≤ 2 κ(cid:37) |A| .
The proof for ε 2 is similar.
[Equation]
 We bound Term 3 below similar to Term 1:
[Equation]
and Term 4 is bounded similar to Term 2 below:
[Equation] [Equation]
This gives us ε 2 ≤ 2 κ(cid:37) |A| .
Plugging bounds for ε 1 and ε 2 in Equation 27 and using P D ( s [ K ] , a ) =
P
D ( s [ K ]) |
[Equation]
gives us:
[Equation]
We can use correspondence result to derive a lower bound:
[Equation]
and an upper bound:
[Equation]
Combing the lower and upper bounds with Equation 28 gives us:
[Equation]
which is the desired result.
We can merge the estimation error and approximation error to generate the total error.
Lemma 20 ( K -Model Error) .
For any i ∈ [ d ] , K ∈ C ≤ 2 κ ([ d ]) , s ∈ S h − 1 , a ∈ A , s (cid:48)
∈ S h , let ˆ s = θ − 1 h − 1 ( s ) and ˆ s (cid:48)
= θ − 1 h ( s (cid:48) ) .
Then we have:
[Equation]
with probability at least 1 − δ est .
Proof.
Follows trivially by combining the estimation error (Lemma 18) and approximation error (Lemma 19) with application of triangle inequality.
D. D ETECTING L ATENT P ARENT S TRUCTURE IN T RANSITION pt h
We are now ready to analyze the performance of learned parent function (cid:98) pt h .
Let K 1 , K 2 ∈ C ≤ 2 κ ([2 d ]) and ˆ w 1 ∈ { 0 , 1 } |K 1 | , ˆ w 2 ∈ { 0 , 1 } |K 2 | .
We will assume ˆ w 1 is reachable for K 1 and ˆ w 2 is reachable for K 2 .
For convenience we will deﬁne the following quantity Ω to measure total variation distance between distributions (cid:98)
P D ( s (cid:48)
[ i ] | · , · ) conditioned on setting ˆ s [ K 1 ]
= ˆ w 1 and ˆ s [
K 2 ]
=
ˆ w 2 , and for a ﬁxed action a ∈ A :
[Equation]
We can compute
[Equation]
Ω for every value of i, K 1 , ˆ w 1 , K 2 , ˆ w 2 , a in computational time of O
[Equation]
2 ed
[Equation]
3 κ +3 |A| (cid:1) .
We also deﬁne a similar metric for the true distribution for any K 1 , K 2 and v ∈ { 0 , 1 } , w 1 ∈ { 0 , 1 } |K 1 | , w 2 ∈ { 0 , 1 } |K 2 | and a ∈ A : Ω ia
[Equation]
[Equation] [Equation]
[Equation]
Recall that [ I ; J ] denotes concatenation of two ordered sets I and J .
We use this notation to state our next result.
Lemma 21 (Inclusive Case) .
Fix i ∈
[ d ] and I ∈ C ≤ κ ([ d ]) .
If pt ( i ) ⊆
I then for all a ∈ A and ˆ u ∈ { 0 , 1 } |I| we get:
[Equation]
where max is taken over J 1 , J 2 ∈ C ≤ κ ([ d ]) , ˆ w 1 ∈ { 0 , 1 } |J 1 | , ˆ w 2 ∈ { 0 , 1 } |J 2 | such that [ˆ u ; ˆ w 1 ] is reachable for [
I ; J 1 ] , [ˆ u ; ˆ w 2 ] is reachable for [ I ; J 2 ] , and I ∩ J 1 = I ∩ J 2 = ∅ .
Proof.
We ﬁx J 1 , J 2 , ˆ u, ˆ w 1 , ˆ w 2 , a and let K 1 =
[ I ; J 1 ] , K 2 =
[ I ; J 2 ] , v = θ (ˆ v ) , u = θ (ˆ u ) , w 1 = θ ( ˆ w 1 ) , and w 2 = θ ( ˆ w 2 ) .
As pt ( i ) ⊆
I , therefore, we have:
[Equation]
Using this result along with Lemma 20 and application of triangle inequality we get:
[Equation]
Summing over ˆ v , dividing by 2, and using the deﬁnition of (cid:98) Ω proves the result.
The following is a straightforward corollary of Lemma 21.
Corollary 11.
Fix i ∈
[ d ] then there exists an I such that for all a ∈ A and ˆ u ∈ { 0 , 1 } |I| :
[Equation]
where max is taken over J 1 , J 2 , ˆ w 1 , ˆ w 2 satisfy the restrictions stated in Lemma 21.
Proof.
Take any I such that pt ( i ) ⊆
I and apply Lemma 21.
Note that we are allowed to pick such an I as | pt ( i ) | ≤ κ by our assumption.
Recall that we deﬁne (cid:98) pt ( i ) as the solution of the following problem:
[Equation]
where I ∈ C ≤ κ ([ d ]) , a ∈ A , and ˆ u, J 1 , J 2 , ˆ w 1 , ˆ w 2 satisfy the restrictions stated in Lemma 21.
We are now ready to state our main result for (cid:98) pt .
Theorem 12 (Property of (cid:98) pt ) .
For any s ∈ S h − 1 , a ∈ A , s (cid:48) ∈ S h , let ˆ s = θ − 1 h − 1 ( s ) and ˆ s (cid:48)
= θ − 1 h ( s (cid:48) ) .
Then the learned parent function
[Equation]
pt satisﬁes:
[Equation]
Proof.
Fix i ∈
[ d ] .
Let J = pt ( i ) − (cid:98) pt
( i ) and K = (cid:98) pt ( i )
∪ J .
As pt ( i ) ⊆
K , therefore, we have:
[Equation]
Combining this result with Lemma 20 we get:
[Equation]
From the deﬁnition of (cid:98) pt ( i ) (Equation 29) and Corollary 11 we have:
[Equation]
Note that we are allowed to use Corollary 11 as ˆ s [
[Equation]
pt ( i ); ∅ ] and ˆ s [ (cid:98) pt ( i ); J ] are both reachable since they are derived from a reachable real state s , | (cid:98)
[Equation]
| [
[Equation]
i ); J
[Equation] [Equation]
pt ( i ); pt ( i )] | ≤ 2 κ , and (cid:98) pt ( i ) ∩ ∅ = ∅ = (cid:98) pt ( i )
∩ J .
Combining the previous two inequalities using triangle inequality completes the proof.
D. B OUND T OTAL V ARIATION BETWEEN E STIMATED M ODEL AND T RUE M ODEL
Given the learned transition parent function (cid:98) pt h
we deﬁne the transition model as:
[Equation]
[ (cid:98) pt h ( i )] , a (cid:1) = (cid:98)
P D (cid:0) ˆ s (cid:48)
[ i ] | ˆ s [ (cid:98) pt h ( i )] , a (cid:1) , (cid:98)
T h (ˆ s (cid:48) | ˆ s, a ) =
d (cid:89) i
= (cid:98)
T hi (cid:0) ˆ s (cid:48)
[ i ] | ˆ s [ (cid:98) pt h ( i )] , a (cid:1) .
From Theorem 12 we have for any i ∈
[ d ] , ˆ s ∈ S h − 1 , a ∈ A , s (cid:48)
∈ S h , and ˆ s = θ − 1 ( s ) , ˆ s (cid:48)
= θ − 1 ( s (cid:48) ) : (cid:12) (cid:12)(cid:12) (cid:98)
T hi (ˆ s (cid:48)
[ i ] | ˆ s [ (cid:98) pt ( i )] , a ) − T ( s (cid:48)
[ i ] | s [ pt ( i )] , a ) (cid:12) (cid:12)(cid:12) ≤ 3∆ est ( n est , δ est ) + 3∆ app .
Transition Closure.
A subtle point remains before we prove the model error between (cid:98)
T h and
T .
Theorem 12 only states guarantee for those ˆ s that are inverse of a reachable state s .
However, as stated before, due to decoder error we can reach a state ˆ s which does not have a corresponding reachable state, i.e. θ (ˆ s ) (cid:54)∈ S h (see Figure 3).
We cannot get model guarantees for these unreachable states ˆ s since we may reach them with arbitrarily small probability.
However, we can still derive model error if we can simply deﬁne 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 deﬁning the closure of the transition model T ◦ for time step h as:
[Equation]
otherwise
We also deﬁne the state space domain of T ◦ h as S ◦ h − 1 = { θ h − 1 (ˆ s )
| ∀ ˆ s ∈
[Equation]
S
[Equation]
− 1 } .
It is easy to see that θ h − 1 represents a bijection between
[Equation]
We will derive our guarantees with respect to T ◦ which will allow us to deﬁne a bijection between the domain of (cid:98)
T and T ◦ , and use important lemmas from the literature.
The next result shows that our use of T ◦ is harmless as it assigns the same probability as T to any event.
Lemma 22 (Closure Result) .
Let T ◦ be the closure of transition model with respect to some learned transition model.
Then for any policy π ∈ Π and any event E which is a function of an episode sampled using π , we have P π ( E ; T ) =
P π ( E ; T ◦ ) , where P π ( E ; T (cid:48) ) denotes the probability of event
E when sampling from π and using transition model T (cid:48) .
Proof.
The proof follows form observing that when using T ◦ we will never reach a state s (cid:54)∈ S h − 1 for any h − 1 by deﬁnition of S h − 1 .
From deﬁnition of T ◦ this means that both T ◦ and T will generate the same range of episodes sampled from π and will assign the same probabilities to them.
As E is a function of an episode, therefore, its probability remains unchanged.
With the deﬁnition of closure, we are now ready to state our last result in this section, which bounds the total variation between the estimated model and the transition closure under the bijection map θ .
Theorem 13 (Model Error) .
For any ˆ s ∈
[Equation]
S h − 1 and a ∈ A we have:
[Equation]
Proof.
If θ (ˆ s ) (cid:54)∈ S h − 1 then by deﬁnition T ◦ the bound holds trivially.
Therefore, we focus on θ (ˆ s ) ∈ S h − 1 for which T ◦ = T .
We deﬁne the quantity for every j ∈ [ d ] :
[Equation]
We claim that S j ≤ 6( d − j + 1)(∆ est + ∆ app ) for every j ∈ [ d ] .
For base case we have:
[Equation]
from Theorem
[Equation]
We will assume the induction hypothesis to be true for S k for all k > j .
We handle the inductive below with triangle inequality:
[Equation]
(ˆ s (cid:48) )
[ j ] | θ (ˆ s )
[ pt ( j )] , a ) | + (cid:88) ˆ s
[Equation]
[ j ] ··· ˆ s
[Equation]
[ d ] ∈{ 0 , 1 }
T ( θ (ˆ s (cid:48) )
[ j ] | θ (ˆ s )
[ pt ( j )] , a ) | d (cid:89) i =
j + (cid:98)
T hi (ˆ s (cid:48)
[ i ] | ˆ s [ (cid:98) pt ( i )] , a ) − d (cid:89) i = j
+
T ( θ (ˆ s (cid:48) )
[ i ] | θ (ˆ s )[ pt ( i )] , a ) |
The ﬁrst term is equivalent to (cid:80)
[Equation]
s )
[Equation]
pt ( j )
[Equation]
a ) | which is bounded by 6(∆ est + ∆ app ) following base case analysis.
The second term is equivalent to
S j +1 which is bounded by 6( d − j )(∆ est + ∆ app ) by induction hypothesis.
Combining these two bounds proves the induction hypothesis and the result then follows from bound for S 1 .
D. L EARNING A P OLICY C OVER
In this section we show how we learn the policy cover.
We start by deﬁning some notation.
Two MDPs.
After time step h , we can deﬁne two Markov Decision Processes (MDPs) at this time
M h and (cid:99)
M h .
M h is the true MDP consists of state space ( S ◦ 1 , · · · , S ◦ h ) , action space A , horizon h , a deterministic start state s 1 = { 0 } d , and transition function T ◦ t : S ◦ t − 1 ×
A → ∆( S t ) for all t ∈
[ h ] .
Recall that the set S h ⊆ { 0 , 1 } d denote states which are reachable at time step h , and the set
S t ⊆
S ◦ t ⊆ { 0 , 1 } d represents the closure of state space S h based on the learned state space (cid:98)
S t .
For any t ∈ [ h ] , s ∈ S t and K ∈ C ≤ 2 κ ([ d ]) we know sup π ∈ Π NS P π
( s t [ K ] = s [ K ]) ≥ η min .
The second MDP (cid:99)
M h consists of the learned state space ( (cid:98)
S 1 , · · · , (cid:98)
S h ) , action space A , horizon h , a deterministic start state
ˆ s 1 = { 0 } d , and transition function
[Equation]
For every t ∈ [ h ] , we have θ t : (cid:98)
S t → S ◦ t represent a bijection from the learned state space to the closure of the set of reachable states at time step t .
The learned decoder (cid:98) φ t predict θ t ( s ) given s ∈ S t with high probability for all t < h by IH.2 and for t
= h due to Corollary 9.
Lastly, the transition model T ◦ t and T t are close in L 1 distance for t < h due to IH.3 and for t = h due to Theorem 13.
These results enable us to utilize the analysis of Du et al.
for learning to learn a policy cover.
Let ˆ ϕ : (cid:98)
S
→ A denote a non-stationary deterministic policy that operates on the learned state space.
Similarly, ϕ :
S ◦
→ A denote a non-stationary deterministic policy that operates on the real state.
We denote ˆ ϕ =
ϕ ◦ θ if for every ˆ s ∈ (cid:98)
S , ˆ ϕ (ˆ s ) =
ϕ ( θ (ˆ s )) .
Similarly, we denote ϕ
=
ˆ ϕ ◦ θ − 1 if for every s ∈ S ◦ , ϕ ( s )
=
ˆ ϕ ( θ − 1 ( s )) .
Let π : X → A be a non-stationary deterministic policy operating on the observation space.
We say π
= ˆ ϕ ◦ (cid:98) φ if for every x ∈ X we have π ( x )
=
ˆ ϕ ( (cid:98) φ ( x )) .
Similarly, we deﬁne π = ϕ ◦ φ (cid:63) if for every x ∈ X we have π ( x ) =
ϕ
( φ (cid:63) ( x )) .
We will use P π [ E ] to denote probability of an event E when actions are taken according to policy π : X → A .
We will use P ϕ
[ E ] to denote the probability of event E when we operate directly on the real state and take actions using ϕ .
Similarly, we deﬁne P ˆ ϕ [
E ] to denote the probability of event ˆ E when we operate on the learned state space.
Lastly, let
[Equation]
P ˆ ϕ
[ E ] denote probability of an event E when actions are taken according to policy ˆ ϕ operating directly over the latent state and following our estimated transition dynamics
[Equation]
Recall that our planner will be optimizing with respect to (cid:98)
P ˆ ϕ
[ E ] .
Theorem 14 (Planner Guarantee) .
Fix ∆ pl ≥ 0 , h ∈
[ H ] .
Let I ∈ C ≤ 2 κ ([ d ]) and ˆ w ∈ { 0 , 1 } |I| .
We deﬁne a reward function R :
[Equation]
S → [0 , 1] where R (ˆ s ) :
= 1 { τ (ˆ s ) =
h
∧ ˆ s [ I ]
= ˆ w } .
Let ˆ ϕ R = planner
( (cid:98)
T, R, h, ∆ pl ) be the policy learned by the planner.
Let ˆ π :=
ˆ ϕ R ◦ (cid:98) φ then:
[Equation]
(31) further, we have:
[Equation]
and if { s h [ I ]
= θ ( ˆ w ) } is unreachable, then
[Equation]
Proof.
We deﬁne two events E :=
{ s h
[ I ] = θ ( ˆ w ) } and ˆ E :=
{ ˆ s h
[ I ] = ˆ w } .
We deﬁne a policy ϕ
R =
ˆ ϕ R ◦ θ − 1 where for every s ∈ S we have ϕ R ( s ) =
ˆ ϕ R ( θ − 1 ( s )) .
We also deﬁne ¯ π : X
→ A as ¯ π ( x ) =
ϕ R ◦ φ (cid:63) ( x ) .
If for a given x ∈ X and φ (cid:63) ( x ) =
s we have (cid:98) φ ( x ) =
θ − 1 ( s ) then ¯ π ( x ) =
ˆ ϕ R ( (cid:98) φ ( x ))
=
ˆ π ( x ) .
Hence, every time our decoder outputs the correct mapped state θ ( s ) , policies ¯ π and ˆ π take the same action.
We use the result of Du et al.
(2019) stated in Lemma(setting ε set to d(cid:37) using Corollary 9) to write:
[Equation]
Let ϕ : S ◦
→ A be any policy on real state space and let ˆ ϕ : (cid:98)
S
→ A be the induced policy on learned state space given by ˆ ϕ (ˆ s ) =
ϕ ◦ θ (ˆ s ) =
ϕ
( θ (ˆ s )) for any ˆ s ∈ (cid:98)
S .
We showed in Theoremthat (cid:98)
T and T have small L 1 distance under the bijection θ .
Therefore, from the perturbation result of Du et al.
(2019) stated in Lemma 31 we have:
[Equation]
where ε := 6 d (∆ est ( n est , δ est ) + ∆ app ) due to Theorem 13.
As { s h [ I ] = θ ( ˆ w ) } ⇔ { ˆ s h
[ I ] = ˆ w } , therefore, we can derive the following bound:
[Equation]
Let ϕ (cid:63) = arg max P
ϕ
[ E ] be the optimal policy to satisfy { s h
[ I ] = θ ( w ) } .
Note that ϕ (cid:63) is also the latent policy that optimizes the reward function R on the real dynamics.
Let ˆ ϕ (cid:63) =
ϕ (cid:63) ◦ θ − 1 be the induced policy on learned states.
We now bound the desired quantity as shown:
[Equation]
(using Equation 35)
[Equation]
This proves Equation 31 and Equation 32.
Note that our calculations above show:
[Equation]
If { s h
[ I ] = θ ( ˆ w ) } is unreachable then
P ϕ R
[ E ] = 0 .
Plugging this in the above equation proves Equa- tion 33 and completes the proof.
D. W RAPPING UP THE PROOF FOR FactoRL
We are almost done.
All we need to do is to make sure is to set the hyperparameters and verify each induction hypothesis.
We ﬁrst set hyperparameters.
Setting Hyperparameters.
Let { s h
[ I ] = θ ( ˆ w ) } be reachable for some I ∈ C ≤ 2 κ ([ d ]) and ˆ w ∈ { 0 , 1 } |I| .
Then applying Theorem 14 and using the deﬁnition of η min we have:
[Equation]
As we want the right hand side to be at least αη min we divide the error equally between the three terms.
This gives us:
[Equation] [Equation]
The model approximation error places a more stringent requirement on (cid:37) than the decoding error for planning.
However, throughout the proof for FactoRL in this section, we made other requirements on our hyperparameters.
For (cid:37) this is given by min { β 2 min / 1200 , 1 / 2 } = β 2 min / 1200 by combining constraints in Lemma 8 and Theorem
[Equation]
and an additional constraint for detecting non-degenerate factors stated in Corollary 9.
Due to the inefﬁciency of the non-degenerate factors detection
[Equation]
we state results separately for the two cases:
[Equation]
Using the deﬁnition of (cid:37) from Theorem 7, we get a value of n abs for non-degenerate factor (Equation 40) and general case (Equation 41) given below:
[Equation]
Recall that for detecting degenerate factors we collect n deg samples.
Corollary 9 gives value of this hyperparameter as
[Equation]
which also satisﬁes the condition in Lemma 16.
Lastly, Theorem 3 gives number of samples for independence testing n ind and rejection sampling frequency k as:
[Equation]
Failure probabilities for a single timestep are bounded by δ ind due to identiﬁcation of emission structure (Theorem 3), 3 dδ abs due to decoding (Corollary 9), and δ est due to model estimation (Lemma 17).
The total failure probability using union bound for a single step is given by δ ind + 3 dδ abs + δ est , and for the whole algorithm is given by δ ind H + 3 dδ abs H + δ est H .
Binding δ ind H (cid:55)→ δ / 3 , 3 dδ abs H
(cid:55)→ δ / 3 , δ est H (cid:55)→ δ / 3 , gives us total failure probability of δ and the right value of hyperparameters.
Sample complexity of FactoRL is at most kHn ind
+ Hn abs
+
Hn deg
+
Hn est episodes which is order of: poly (cid:26)
[Equation]
where use the fact that N = | Ψ h − 1 | can be at most 2( ed ) 2 κ from Lemma 23.
Note that if we did not have to apply the expensive degeneracy detection step, then we would get logarithmic dependence on 1 / δ abs .
Cheaper ways of detecting degeneracy can, therefore, signiﬁcantly improve the sample complexity.
We have not attempted to optimize the degree and exponent in the sample complexity above.
For our choice of two hyperparameters n est and n abs , we can bound the model error and decoding failure by: (
[Equation]
Verifying Induction Hypothesis.
Finally, we verify the different induction hypothesis below.
We already veriﬁed IH.1 with Theorem 3.
We learn a
[Equation]
ch h that is equivalent to ch h upto label permutation.
We already veriﬁed IH.2 with Corollary 9.
Given a real state s ∈ S h , our decoder outputs the corresponding learned state with high probability.
We also derived the form of (cid:37) .
We already veriﬁed IH.3 with Theorem 13.
We also
derived the form of ∆ est and ∆ app .
Lastly, Theorem 14 and our subsequent calculations for hyperparameter show that Ψ h is an α -policy cover of S h and that the size of Ψ h is at most 2( ed ) 2 κ from Lemma 23.
Lastly, for all reachable factor values we get the value of learned policy as at least (1+ α ) η min /using Equation 32 and our choice of hyperparameter values.
Similarly, from Equation 33 we get the value of learned policy for all unreachable factor values as at most (1 − α ) η min / 4 .
This allows us to ﬁlter all unreachable factor values.
In the main paper, we focus on the value of α = 1 / 2 , which explains why on Algorithm 1, line 8 we only keep those policies with value at least (1+ α ) η min / 2 = 3 η min / 4 .
This veriﬁes IH.4.
This completes the analysis for FactoRL .
E S UPPORTING R ESULT
Lemma 23 (Assignment Counting Lemma) .
For a given k, d ∈ N and k ≤ d , the cardinality of the set { ( K , u ) | K ∈ C ≤ k ([ d ]) , u ∈ { 0 , 1 } |K| } is bounded by 2( ed ) k .
Proof.
Assume k ≥ 2 .
The cardinality of this set is given by (cid:80) ki =
[Equation]
(cid:0) di (cid:1) 2 i which can be bounded as shown below:
[Equation]
The ﬁrst inequality here uses the well-known bound for binomial coefﬁcients (cid:0) ni
[Equation]
≤
[Equation]
i for any n, i ∈ N and i ≤ n .
Further bounding the above result using ed − 1 ≥ ed / 2 gives us:
[Equation]
The proof is completed by checking that inequality holds for k < 2 .
Lemma 24 (Lemma H.1 in Du et al.
Let u, v ∈ R d + with (cid:107) u (cid:107)
= (cid:107) v (cid:107) 1 = 1 and (cid:107)
u − v (cid:107) 1 ≥ ε .
Then for any α > 0 we have
[Equation]
Lemma 25 (Chernoff Bound) .
Let q be the probability of an event occurring.
Then given n iid samples with n ≥ 1 q , the probability that the event occurred at least once is at least 1 − 2 exp( − qn 3 ) .
Proof.
Let X i be a 0-1 indicator denoting if the event occurred or not, and let X = (cid:80) ni
=1 X
i .
We have E [ X i ] =
q and E
[ X ] = qn .
Let t = 1 − 1 / qn .
We will assume that qn > 1 and so t ∈ (0 , 1) .
Then the probability that the event never occurs is bounded by:
[Equation]
Lemma 26 (Hoeffding’s Inequality) .
Let X 1 , X 2 , · · · , X n be independent random variables bounded by the interval [0 , 1] .
Let empirical mean of these random variables be X =n
[Equation]
ni
[Equation]
1 X n , then for any t > 0 we have:
[Equation]
Lemma 27 (Theorem 2.1 of Weissman et al.
Let P be a probability distribution over a discrete set of size a .
Let X n = X 1 , X 2 , · · · , X n be independent identical distributed random variables distributed according to P .
Let ˆ P X n be the empirical probability distribution estimated from sample set X n .
Then for all (cid:15) > 0 :
[Equation]
The next result is a direct corollary of Lemma 27.
Corollary 15.
For any m ≥ 8 a(cid:15) 2 ln ( 1 / δ ) samples, we have (cid:107) P − P X n (cid:107) 1 < (cid:15) with probability at least 1 − δ .
Lemma 28.
Let X 1 , X 2 , · · · , X n be 0 - 1 independent identically distributed random variables with mean µ .
Let X = (cid:80) ni
=1 X i .
Fix m ∈ N and δ ∈ (0 , 1) .
If n ≥ 2 mµ ln (cid:0) eδ (cid:1) then P ( X < m ) ≤ δ .
Proof.
This is a standard Chernoff bound argument.
We have E [ X ] = nµ .
Assuming n ≥ m / µ then from multiplicative Chernoff bound we have:
[Equation]
Setting right hand side equal to δ and solving gives us n ≥ 2 µ (cid:0) m
+ ln( 1 δ ) (cid:1) which is satisﬁed whenever n ≥ 2 m µ ln (cid:0) e δ (cid:1) .
Lemma 29 (Lemma H.3 of Du et al.
For any a, b, c, d > 0 with a ≤ b and c ≤ d we have:
[Equation]
The next Lemma is borrowed from Du et al.
They state their Lemma for a speciﬁc event ( E = α
(ˆ s ) in their notation) but this choice of event is not important and their proof holds for any event.
Lemma 30 (Lemma G.5 of Du et al.
Let S =
( S 1 , · · · , S H ) and (cid:98)
S =
( (cid:98)
S 1 , · · · , (cid:98)
S H ) be the real and learned state space.
We assume access to a decoder (cid:98) φ : X → (cid:98)
S and let φ (cid:63) :
X → S be the oracle decoder.
Let θ h : (cid:98)
S h
→ S h be a bijection for every h ∈
[ H ] and θ : (cid:98)
S → S where θ (ˆ s )
=
θ h (ˆ s ) for ˆ s ∈ (cid:98)
S h .
For any h ∈ [ H ] and s h ∈ S h , we assume P (ˆ s h
= θ − 1 h ( s h ) | s h ) ≥ 1 − ε , i.e., given the real state s h , our decoder ( (cid:98) φ ) will map it to θ − 1 h ( s h ) with probability at least 1 − ε .
Let ϕ : S
→ A be a deterministic policy on the real state space and ˆ ϕ : (cid:98)
S
→ A be the induced policy given by ˆ ϕ (ˆ s )
=
ϕ ( θ (ˆ s )) for every ˆ s ∈ (cid:98)
S .
Let π, ˆ π : X → A where π ( x ) =
ϕ ( φ (cid:63) ( x )) and ˆ π ( x ) =
ˆ ϕ (
[Equation]
φ ( x )) for every x ∈ X .
For every random event E we have:
[Equation]
Lemma 31 (Lemma H.2.
of Du et al.
Let there be two tabular MDPs M and (cid:99)
M .
Let
S =
( S 1 , · · · , S H ) be the state space of M and (cid:98)
S =
( (cid:98)
S 1 , · · ·
[Equation]
(cid:98)
S H ) be the state space of (cid:99)
M .
Both MDPs have a
A be the action space of both MDPs and horizon of H .
For every h ∈ [ H ] , there exists a bijection θ h : (cid:98)
S h
→ S h .
Let T : S ×
A → ∆( S ) and (cid:98)
T : (cid:98)
S × A → ∆( (cid:98)
S ) be transition dynamics for M and (cid:99)
M satisfying:
[Equation]
Let ϕ : S
→ A be a policy for M .
Let (cid:98) ϕ
(cid:98)
S
→ A be the induced policy on (cid:99)
M such that for any ˆ s ∈ (cid:98)
S we have
[Equation]
s )) .
Then for any h ∈ [ H ] we have:
[Equation]
F E XPERIMENT D ETAILS
We provide details for our proof of concept experiment below.
Modeling Details.
We model F , used for performing independence test, using a single-layer feedforward network θ F with Leaky ReLu non-linearity (Maas et al., 2013) and a softmax output layer.
Give a pair of atoms x [ u ] and x [ v ] , we concatenate these atoms and map it to a probability distribution over { 0 , 1 } by applying θ F .
We implement the model class G for learning state decoder following suggestion of Misra et al.
Recall that a function in G maps a transition ( x, a, ˇ x ) ∈
X
×
A × X (cid:63) to a value in [0 , 1] .
We ﬁrst map x and ˇ x to vectors v 1 and v 2 respectively, using two separate linear layers.
We map the action a to its one-hot vector representation 1 a .
We map the vector v 2 to a probability distribution using the Gumbel-softmax trick (Jang et al., 2016), by computing q i ∝ exp( v
[ i ] + ϑ i ) for all i ∈ { 1 , 2 } , where ϑ i is sampled independently from the Gumbel distribution.
We concatenate the vectors v 1 , 1 a and q and map it to a probability distribution over { 0 , 1 } , through a single layer feed-forward network θ G with Leaky ReLu non-linearity.
We recover a decoder φ from the model that maps a set of atoms ˇ x to φ (ˇ x ) =
arg max i ∈{ 0 , 1 } q i +1 .
Learning Details.
We train the two models using cross-entropy loss.
Formally, given a dataset
D ind
= { ( x i [ u ] , x i [ v ] , y i ) } n ind
i
=1 for performing independence testing, and a dataset D abs
= { ( x i , a i , ˇ x i , y i ) } n abs i =1 for learning a decoder, we optimize the model by minimizing the cross-entropy loss as shown below:
[Equation]
Here we overload our notation to allow the output of models to be distribution over { 0 , 1 } rather than a scalar value in [0 , 1] as we assumed before.
This is in sync with how we implement these model class and allows us to conveniently perform cross-entropy loss minimization.
Planner Details.
We use a simple planner based on approximate dynamic programming.
Given model estimate, reward function and a set of visited learned states, we perform dynamic programming to compute optimal Q-values for the set of visited states.
We assume the Q-values for non-visited states to be 0.
This allows us to compute Q-values in a computationally-efﬁcient manner.
Later, if the agent visits an unvisited state, then it simply takes random action.
Hyperparameters.
We set the hidden dimension of θ F to 10 and that of θ G to 56.
We set the threshold c on held-out log-loss value, when performing independence test to be 0 .
For reference, if one uses a random uniform classiﬁer then its performance is − ln(0 . 5) ≈ 0 .
We train both models for 10 epochs using Adam optimization with learning rate of 0.001, and a batch size of 32.
We remove 0.2% of the training data and use it as a validation set.
We evaluate on the validation set after every epoch, and use the model with the best performance on the validation set.
We used
PyTorch 1.6 to develop the code and used default initialization scheme for all parameters.