Learning a Fast Mixing Exogenous Block MDP using a Single Trajectory
Authors: Alexander Levine
Peter Stone
Amy Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning
Reinforcement Learning Theory
Controllable Representations
Representation Learning
Exogenous Noise
Controllable Latent State
Unsupervised Reinforcement Learning
Primary area: reinforcement learning
TL;DR: We propose a provably sample-efficient algorithm for learning controllable representations in the Exogenous Block MDP setting, in the case where data is collected in a single trajectory with no state resets.

Abstract: In order to train agents that can quickly adapt to new objectives or reward functions, efficient unsupervised representation learning in sequential decision-making environments can be important. Frameworks such as the Exogenous Block Markov Decision Process (Ex-BMDP) have been proposed to formalize this representation-learning problem (Efroni et al., 2022b). In the Ex-BMDP framework, the agent's high-dimensional observations of the environment have two latent factors: a controllable factor, which evolves deterministically within a small state space according to the agent's actions, and an exogenous factor, which represents time-correlated noise, and can be highly complex. The goal of the representation learning problem is to learn an encoder that maps from observations into the controllable latent space, as well as the dynamics of this space. Efroni et al. (2022b) has shown that this is possible with a sample complexity that depends only on the size of the controllable latent space, and not on the size of the noise factor. However, this prior work has focused on the episodic setting, where the controllable latent state resets to a specific start state after a finite horizon.

By contrast, if the agent can only interact with the environment in a single continuous trajectory, prior works have not established sample-complexity bounds. We propose STEEL, the first provably sample-efficient algorithm for learning the controllable dynamics of an Ex-BMDP from a single trajectory, in the function approximation setting. STEEL has a sample complexity that depends only on the sizes of the controllable latent space and the encoder function class, and (at worst linearly) on the mixing time of the exogenous noise factor. We prove that STEEL is correct and sample-efficient, and demonstrate STEEL on two toy problems. Code is available at: https://github.com/midi-lab/steel.

Published as a conference paper at ICLR 2025
LEARNING A FAST MIXING EXOGENOUS BLOCK MDP USING A SINGLE TRAJECTORY
Alexander Levine*
Peter Stone
Amy Zhang The University Texas Austin
The University of Texas at Austin
The University of Texas at Austin
Sony AI
ABSTRACT
In order to train agents that can quickly adapt new objectives or reward functions, efficient unsupervised representation learning in sequential decision- making environments can be important. Frameworks such as the Exogenous Block Markov Decision Process (Ex-BMDP) have been proposed to formalize this representation-laate learninng problem (Efroni et al., 2022b). In the Ex-BMDP framework, the agent's high-dimensional observations of the environment have two latent factors: a controllable factor, which evolves deterministically within a small state space according to the agent's 's actions, and exogenous factor, which represents time-correlated noise, and can be highly complex. The goal of the rep- resentation learning problem to learn an encoder that maps from observations into the controllable latent space, as well as the dynamics of this space. Efroni et al. (2022b) has shown that this is possible with sample complexity that de- pends only on the size of the controllable latent space, and not on the size of the noise factor. However, this prior work has focused on the episodic setting, where the controllable latent state resets to a specific start state after finite horizon.
By contrast, if the agent can only interact with the environment in single con- tinuous trajectory, prior works have not established sample-complexin bounds. We propose STÉEL, the first provably sample-efficient algorithm for learning the controllable dynamics of an Ex-BMDP from single trajectory. in the function approximation setting. STEEL has sample complexity that depends only on the sizes of the controllable latent space and the encoder function class, and (at worst linearly) on the mixing time of the exogenous noise factor. We prove that STEEL is correct and sample-efficient, and demonstrate STEEL on two toy problems. Code is available at: hhttp:: /githbb. com/ com/midi midi- -abiiiieiiaaeeeeelllllelea
INTRODUCTION
This work considers the unsupervised representation learning problem in sequential control envi- ronments. Suppose an agent (e.g., robot) is able to make observations and take actions in vironment for some period of time, but does not yet have an externally-defined task to accomplish. We want the agent learn model of the environment that may be useful for many downstream tasks: the question is then how to efficiently explore the environment to learn such a model.
Sequential decision-makee making tasks are often modeled as Markov Decision Processes (MDPs). In the unsupervised setting, an MDP consists of set of possible observations X. set of possible actions X P(X).
A. a distribution over initial observations mo P(X)), and a transition function The agent does not have direct access to T. Instead, at each timestep t, the agent observes X and selects action a, The next observation Z then sampled as F441
(x)
In the totally generic MDP setting, the only model learning possible is to directly learn the transition function T. However, the space of possible observations is large, this task becomes intractable. Therefore, prior works have attempted to simplify the problem by assuming that the MDP has some underlying structure, which a learning algorithm can exploit. One such structural assumption is the EX-BMDP (Exogenous Block MDP) framework, introduced by Efroni et al. (002b)). The Ex-
"Correspondence alevine0@@er www...aas....exei.....

---

Published as a conference paper at ICLR 2025
BMDP framework captures situations where the space of observations X' is very large, but the parts of the environment that the agent has control over can be represented by a much smaller latent state.
An EX-BMDP has observation space X, controllable (or endogenous) latent state space S, and an exogenous state space E. In the version of this setting that we consider, the controllable state St S evolves deterministically according is: Sc+1 = T(s,, az). The exogenous state G E & represents time correlated noise: it evolves transition function Te & + P(E). That is: ex+1 ~ (et).. Neither s nor is directly observed. Instead, the observation I4 is sampled as It Q(5), et), where Q E S x 3 + P(X') is the emission function. We make block assumption on Q with respect to S: that is. we assume that for distinct latent states s, E S,, the sets of possible observations that can be sampled from Q(s,) :) and Q(s'..) are disjoint. In other words, there exists a deterministic partial inverse of 2, which is o² : X + S. such that if x ~ Q(s, e), then o² (ar) = 8. Hence, it is always possible in principle to infer s from x:
latent transition function T : x A - S. That
stochastically, independently of actions, according
As in the general MDP setting, the agent only directly observes I2 E X, and chooses actions de in response. However, rather than attempting to learn the full transition dynamics T of the system (which is determined by T. Te. and Q together), the objective of the agent is instead to efficiently model only the latent encoder o² and the latent transition function T. Together, these models allow the agent plan learn in downstream tasks using the encoded representations e (zz). modeling only the parts of the environment that the agent can actually control (the latent state 8 E S) while ignoring the potentially-comple dynamics of time-correlated noise.
Specifically, the aim of efficient representation learning in this setting is to learn 6" and T, using a number of environment steps of exploration that is dependent only on [8] and the size of the function class F that the encoder ø² belongs to, and is not dependent on the size of X' or E. This allows X' and E to be very large or potentially even infinite, but still allows for representation learning to be tractable. Eironi et al. (2022b) proposes an algorithm, PPE, with this property. However. PPE only works in tinitehhorizon setting, where the agent interacts with the environment in episodes of fixed length H. After each episode, the controllable state (almost) always resets to a deterministic start state 80 E S at the beginning of each episode. In this work, we instead consider the single- trajectory, no-reset setting, where the agent interacts with the environment in single episode of unbounded length, with no ability to reset the state. This better models real-world cases, where, for example, expensive human intervention would be required to "reset" the environment that robot trains in: we would rather not require this intervention. This no-reset EX-BMDP setting was previously considered by Lamb et al. (2023) and Levine et al. (2024); however, the algorithms presented in those works do not have sample-complexity guarantees.
By contrast, the algorithm presented in this work is guaranteed to learn d" and T using samples polynomial in |S] and log [FI, with no dependence on [&] and |.2). We only require that the mixing time of the exogenous noise bounded. In other words, the requirement is that mmix. the mixing time of the Markov chain on & induced by Te. is at most some known quantity (mix- Note that we do not require that the endogenous state mixes quickly under any particular policy (although we do require- as do Lamb et al. (2023) and Levine et al. (2024) - tthtt all states in S are eventually reachable from one another). this setting, we derive an algorithm with the following asymptotic sample complexity (where (f())) = (xf(z) log(f(=))))
FFI + |f log 5 + |S|2D -log
NIFI
FF 5
|F 5
O (NNDSSPIAI log
log
(1)
5
where N a predetermined upper-bound on [S), 8 is the failure rate of the algorithm, D is the maximum distance between any two latent states in S (at most [S]), and is the minimum accuracy of the output Iearned encoder o on any latent state class s E S. Note that this expression is at worst polynomial in (S), and linear in fmix and log [F..
Our algorithm proceeds iteratively, at each iteration taking a certain sequence of actions repeatedly in loop. Because the latent state dynamics T are deterministic, this process is guaranteed to (after some transient period) enter a cycle of latent states, of bounded length. Because the latent states in
Unlike most prior works on Ex BMDDs, we do not make a block assumption on E: we allow the same r to be emitted by Q(s, e) and Q(s,e') e'), for distinct e, e'. Technically, then, our Ex Ex-P framework represents a restricted class of Partially-Osservee MDPs (POMDPs), rather than MDPs, because the complete state is not encoded within the observed X.
2

---

Published as a conference paper at ICLR 2025
this cycle are repeatedly re-visited, the algorithm is then able to predictably collect many samples of the same latent state, without the need to "re-set" the environment. Furthermore, because this looping can be continued indefinitely, the algorithm can "wait out" the mixing time of the exogenous dynamics, in order to collect near-i.i.d. samples of each latent state. We call our algorithm Single- Trajectory Exploration for Ex-BMDDPs via Looping, STEEL, the first provably sample-efficient algorithm for learning Ex-BMDPs in a general function- approximation setting from single trajectory, (2) prove the correctness and sample complexity of STEEL, and (3) empirically test STEEL on two toy problems to demonstrate its efficacy.
STEEL In summary, we: (1) introduce
2
RELATED WORKS
2.1
REPRESENTATION LEARNING FOR Ex- EX-BMDP AND ExO-MDPs
The EX-BMDP model was originally introduced by Efroni et al. (2022b), who propose the PPE algo- rithm to learn the endogenous state encoder oi-) and latent transition dynamics T of an EX-BMDP. PPE has explicit sample complexity does not depend explicitly on |E| or [ However, unlike the method proposed in this work, PPE is restricted to the episodic, finite horizon setting with (nearly) deterministic re- sets. After each episode, the endogenous state is (nearly) always reset the same starting latent state 80 E S, and the exogenous state co E 3 is i.i.d. resampled from a fixed starting distribution. Similarly this work, PPE assumes that the latent transition dynamics T are (close to) determinis- tic. Because both and 7 are nearly deterministic, PPE can collect i.i.d. samples of observations r associated with any latent state s E S simply by executing the same sequence of actions starting from So after each reset. By contrast, in our setting, we cannot reset the Ex- EX-BMDP state, so it is more challenging to collect samples of a given latent state 8.
copple-oomlextty guarantees that are polynomial in SS] and log Fl: crucially, the
Other works have considered the Ex- EX-BMDP setting without latent state resets. Lamb et al. (2023) Lnd[Levine et (2024) consider a setting similar to ours, where the agent interacts with the environ- ment in single trajectory. However, these methods do not provide sample-compp complexy guarantees, and instead are only guaranteed to converge to the correct encoder in the limit of infinite samples.
Efroni et al. (022a) considers a related "EXo-MDP" setting, and proposes the ExoRL algorithm. In this setting, while the environment is episodic, the latent state so resets to a starting value sampled randomly from fixed distribution after each episode. Additionally, the latent transition dynamics may be non-deterministic. However, unlike our work, [Efroni et al. (2022a) does not consider the general cunntion-approximaan
setting for state encoders. Instead, the observation x: is explicitly factorized into d factors, and the controllable state s consists of some unknown subset of l: of these factors: the representation learning problem is reduced identifying which ls of the d factors are action-dependent. ExoRL guarantees a sample-complexity polynomial in 2 = =|S] and log(d) 2
Recently, Mhammedi et al. (2024) have proposed algorithm for provably sample-efficient policy optimization in the episodic Ex- BMDP setting with rewards, that can handle nondeterministic latent dynamics. However, the algorithm requires simulator access to the environment: this means that the agent is able to reset the environment to any observation æ E X that has been previously observed. This requirement is considerably stronger than even the requirement of deterministic resets a single latent state found in Efroni et al. (2022b).
In addition our main claim that our proposed algorithm is the first provably sample-efficies algorithm for representation learning in the sragle-traecctrry EX-BMDP setting, another property of our method is that we do not make a "block" assumption on the exogenous state e. For fixed a S, X' may be emitted by multiple distinct exogenous states
in our setting, the same observation e, e' E E. Prior works (Efroni et al., 2022b|a Lamb et al., 2023 Levine et al., 2024) have stated assumptions that require that may be uniquely inferred from x! Removing this restriction allows one to model a greater range of phenomena. For example, suppose an agent can turn on or turn off
2 Because EXoRL allows for nondeterministic starting latent states, one might be able to adapt ExoRL to the
no-reset setting by considering the single episode as chain of "pseudo-episodes" with nondeterministic start state (c.f.. ffff[uuu et al. (024)). However, one would have ensure that mixes sufficiently between episodes, which may be challenging given that s is not directly observed and has unknown, controllable dynamics. Even still, the resulting algorithm would not apply to our general function-approximatio setting. A Athoughh it not immediately clear why this restriction is necessary for the proposed algorithms.
3

---

Published as a conference paper at ICLR 2025
No-Reset Setting
Sample-Complesity
Function
Nondeterministic
Parially-Ossereed Exogenous State
Nondeterministic
Guaranlees
Approximation
Revel State
Latent Transitions
STEEL Efroni 2.26
Lambb 23
X
X
Levine
Elron 225
X
Table Comparison to Prior Works for learning EX-BMDP Latent Dynamics
a "noisy TV": i..., the agent can control whether or not some source of temporally noise is observable. This is allowed in our version of the EX-BMDP formulation, but is not allowed if e must be fully inferable from z. One prior work, Wu et al. (2024), also (implicitly) removes the block restriction the exogenous state e. That work extends EX-BMDP representation learning to the partially-observed state setting, with the assumption that the observation history within some
known window is sufficient infer the latent state s. However, Wu et al. (2024) does not provide any sample-complxxtty guarantees. [Wang et al. (2022) and Kooi et al. (202)) consider similar settings with continuous controllable latent state. However, the proposed methods require explicitly modeling the exogenous noise state and there are no sample complexity guarantees. [Trimponaas & Dietterich] (2023) considers the sample-efficinccy of reward- -ased reinforcement learning Ex- BMDPs assuming known endogenous and exogenous state encoders; however, it does not address the sample complexity of the representation learning problem.
2.2 REPRESENTATION LEARNING FOR BLOCK MDPs AND LOW-RANK MDPs
generalization of the Block MDP framework (Dann et al. 2018 Du et al. 2019). Like the Ex- EX-BMDP setting, the Block MDP setting models environments where the observed state space X' of the overall MDP is much larger than an action- dependent latent state space S. Some works in the Block MDP framework (Mhammedi et al. 2023 Misra et al., 2020) also allow for nondeterministic latent state transitions: that is St+1 ~ ( at)
The EX-BMDP framework can be considered
However, unlike the EX-BMDP setting, there is no exogenous latent state e E E or exogenous dynamics Te: the observation is simply sampled as I₁ ~ Q(s.). In other words, the Block MDP setting does not allow for time correlated noise outside of the modelled latent state s. Therefore, even when stochastic latent-state transition are allowed, any uime-correlated noise must be captured in S, and so impacts the sample complexity (which is typically polynomial in (S)).
The Low-Rank MDP framework can also be considered as an extension the Block MDP framework, but is an orthogonal extension to the Ex-BMDP framework. In Low Rank MDPs, there exist func- tions é : X Rd and μu: : X' - RRd, such that Pr(æ = z'|z; = I, 0g = a) = p(x, a) p(z'). d and logarithmically
The sample complexity depends only polynomially tion classes for the state encoders ó and μc; it should not depend explicitly on |x'] Works under this framework include Agarwal al. (2020) [ehara al. (2022) and Cheng al. (2023). Other works in the Low Rank MDP framework use a reward signal and only explicitly learn part of the represen- tation (the encoder Φ). including Mhammedi et al. (2023) and Jiang al. (2017); see Mhammedi et al. (2023) for recent, thorough comparison of these works. Note that while BMDPs can be formulated as low-rank MDPs with d |S. this does not hold for Ex-BMDPs: the rank of the transition probabilities on X' depends on | as noted by Efroni et al. (2022b).
the size of the func-
3 NOTATION AND ASSUMPTIONS
The EX-BMDP, M, has observation space X, with discrete endogenous states S that have deterministic, controllable dynamics, and possibly continuous exogenous states & with nondeterministic dynamics that do not depend on actions. Concretely, we have that 81+1 = T(s dz), where T is a deterministic function, and e++1 ~ Te(e). Let Dt Q(st, e,), for Ie X. St E S, e, E, with the block assumption on S. That is, a given x E X' can be emitted only by one particular s E S, which we define as d" (xx) st- We assume that M is accessed in one continuous trajectory. The initial endogenous state an arbitrary Sinit S, and the initial exogenous state Cintt T
init where पे P(E).
The exogenous dynamics on È are irreducible and aperiodic, with stationary distribution Tg. There is known upper bound fmix on the mixing time Lmix, where (as defined in
4

---

Published as a conference paper at ICLR 2025
Levin & Peres (2017) and elsewhere) fmix == mmix (1/4), where tmix (c) is defined such that: Ve E &, || P - e'lex e) (e()')] e.
(2)
This assumption bounds how "temporally correlated" the noise in the Ex-BMDP is: it ensures that the exogenous noise state e, at time is relatively unlikely to affect e++mme"
We have known upper bound on the number of endogenous latent states, N (S]. Additionally, assume that all endogenous latent states be reached from one another in at most steps, for some finite D (note that we do not assume that all pairs of states another in exactly D steps). We assume that there is
in S be reached from known upper bound on this diameter: D 2 D. Trivially, if all endogenous latent states reachable from one another then D < N - 1, so if a tighter bound is not available then we can use D == N - 1. (In fact, it is not very important to use tight bound here: D does not appear the asymptotic sample complexity.)
There is encoder hypothesis class F X + {0,1}, with realizablity for one-vs-rest classification of endogenous states. That is,
Vs E S, 3ʃ F Vo X, f(x) (x)===
In other words, for every latent state 8 E S, there is some function f E F such that (xx) = if and only if o (x:)
The algorithm has access to a training oracle for F, which, given two finite multi-sets Do and Di each with elements from X, returns classifier f E F. The only requirement that we have for this oracle is that, if there exist any classifiers F C F, such that, for I E F*, Vo E Do. f (z) = 0 and Vo E Dr, f*(x) = then the oracle will return some member of Ft. Note that an optimizer that minimizes the 0-1 loss on Do U Dı will satisfy this requirement. However, is not strictly necessary to minimize the 0-1 loss in particular.
General notation: Let M(A) be the set of all multisets of the set A. Let 1 represent an undefined value. For lists v, y, let x y represent their concatenation: that is, [a, b) [c, d] = (a, b, c, d). For multisets A and B, let A l B be their union, where their multiplicities are additive. Let % be the modulo operator, so that a%b = a (mod b) and 0 I a%o I
4 ALGORITHM
The STEEL algorithm is presented in full in Appendix A in this section, we give high-level overview of the algorithm, and present bounds on its sample-complexigg We state the sample- complexity and correctness of STEEL in the following theorem, which is proved in Appendix B
Theorem 1. For an Ex-BMDP M (K,,,,,,,, Q.T, xit starting at an arbitrary endoge- nous latent state Siait E S, with [S] S N, where the exogenous Markov chain Te has mix- ing time at most émice and where all states in S are reachable from another in at most D steps; and corresponding encoder function class F such that Equation 3 holds, the algorithm STEEL(M, F, N, D, imirr 5.e) will output a learned endogenous state space S' transition model T', and encoder Ø', such that, with probability at least 1 ö.
S = (S), and under some bijective function o : S - S', it holds that Ys E S.a a E A: o (T(s a)) ) = (o(s), a). and.
(4)
Under the same bijection a,
is
Pr 2~ e
(0'z)= (x))) 21-6,
(5)
Q(s.e,
where ae is the stationary distribution of Te-
Furthermore, the number of steps that STEEL executes on M scales as:
F
NJF
S
S
imix
|
oº (NDISPPAA) - log 8
+ |S||A|imie log
log
log
6
where O³ (J(xx)) (((()))o((()))))))
5

---

Published as a conference paper at ICLR 2025
a D
2. à
3. 1
5. L.
6. R.
12 U D. R. U.
7. LL B. R]
àL D.
9. L,
a- DJ D.
CO-O=O 199
Figure 1: STEEL discovers the latent dynamics S and T by iteratively adding cycles to the learned dynamics graph. this simple example, the mitiaaly-unknoww "true" dynamics consist states arranged in a grid, where the agent can move (U)p, (D)own, (L)eft, (R)ight. STEEL takes 12 iterations to discover the full dynamics; each pane corresponds to an iteration, and shows the still- unknown parts of the dynamics graph in grey, the already-knowe parts of the dynamics graph in black, and the cycle being explored in red. States are represented as circles and transitions as arrows.
We now present high-level overview of STEEL The algorithm proceeds in three phases. In the first phase, the algorithm learns the transition dynamics; in the second phase, it collects additional samples of observations of each latent state in S; in the final phase, the encoder o' is learned.
STEEL Phase 1: Learning latent dynamics. In the first phase, STEEL constructs a learned latent state space S' and learned transition dynamics 7" by iteratively adding cycles the known transition graph. At each iteration, a sequence of actions à is chosen such that, starting anywhere in the known T", taking the actions process of constructing à below.)
guaranteed to traverse a transition not already in T'. (We explain the
STEEL then takes the actions in à repeatedly, collecting a sequence observations ICF. Because the transitions in T are deterministic, this sequence of transitions must eventually (after at most ||| steps) enter a cycle of latent states, of length neye [ā], for some ney: I [S]. Because was chosen to always escape the known transitions in T', this cycle cannot be contained in T', adding the states and transitions of the new cycle to S' and T" is guaranteed to expand the known dynamics graph by at least one edge: this process will discover the full transition dynamics after at most [S|| || iterations. See Figure for an example of how STEEL constructs S' and T" by adding cycles to the dynamics. Throughout this process, STEEL also collects dataset ((s) E M(X) for each newly- discovered latent state s E S', so that each observation in D(s) has latent state s. The observations within each dataset ((s) are collected at least fmix steps apart from one another, so they are near-i.i.d.
Constructing à. To ensure that à always escapes the known transition graph given by T', each iteration, STEEL uses a recurrent procedure to construct à At the beginning of this procedure, is initialized empty, and set of latent states B is initialized with all of the learned states in S At each step of the procedure, B represents the set of known states E S' that can be reached by starting any arbitrary state s' S' and taking the action sequence given by the partially-constructee action list à, following the known dynamics in T". At each step. STEEL chooses a state E B, and plans the shortest path from s any unknown transition in T". The corresponding sequence of actions, à', is then appended to å, and B is updated by replacing each state b E B with the state that results from starting at and taking the sequence of actions à', according to the learned partial transition function T". If this path reaches an unknown transition in 7', then à is removed from B and not replaced. By construction, we know that the state s will be removed, so B will shrink by at least one state at each step of the process. By the end, B is empty, and we are guaranteed that taking action sequence from any state in S' will lead to an unknown transition in I', as desired. Note that at each step, the shortest distance an unknown transition has length at most D + and the procedure continues for at most [S] steps, so |â] < SS (D+1)
Identifying latent states.At At each iteration, identify the distinct latent states in the cycle, STEEL uses a subroutine called CycleFind. CycleFind additionally collects the datasets D(s) for each newly-discovered latent state. CycleFind itself has phases:

---

Published as a conference paper at ICLR 2025
n" eyo = 3; aao
=3: Di contains only the red latent state; Do does not contain this latent state
Do
n ooyo
11
18)
17
30
Lea qyc
(n" qye
180
"aa/lo eye
oye 18U
1817 (n
"Ew//0 oye IAV
cyy = 4: Ooy
3: Do and Dr contain the same latent states
Do
> ayy n' eyc
1.8
[5]
19
[8]
[4]
13]
[4]
n: [8]
"5a000 1807 (0".. I8U
[41 (Y (41)
181 (1 I40
Figure 2: CycleFind determines the period of the cycle in ICF. See Section under "CycleFind Phase 1." We show the sequence ECF sampled from M: specifically, we show every |â] (a obser- vation, where the same actions à are taken between each one. The observations' latent states are color-coded as red, blue, and green: pattern repeats every 3/2] steps, so Nyye = 3. D₃ consists of the first observation in each (al-cll)),yyle,, and Do the other observations taken between executions of à. (Spans of length 2 fmix are skipped to ensure certain subsets of the datasets are near-i.i.d.)
CycleFind Phase 1: Finding the cycle's periodicity. To identify the latent states in the cycle in ICF, CycleFind first determines cycle's 's period, Teey [ā]. To find Reye. CycleFind tests all possible values of neye from N to 1, in decreasing order. To check whether some candidate value, neye. is in fact feyye. CycleFind constructs two datasets, Do and D2 from ØCF = [z]....]. These datasets are constructed so that if n'ye = Peyye then D2 contains observations of only one controllable latent state and Do contains no observations of s. Therefore, if one attempts to train classifier f to perfectly distinguish all observations in Do from those D1, then such classifier unlikely to exist F if n'ye Teye. but is guaranteed to exist if n'ye Neyy. by the realizability assumption. In this way, CycleFind uses the training oracle to determine Yleye.
Specifically, the datasets used are (see Figure 2 for illustration.):
Do = {2 + (25m + |a) ) 1-11 toffs. E {o,..., - 1}. j j (1,...,,,,, Klye - 1})
(6)
D1 {I + (alimatt lal)stnnss i (O,..., 1}}
where mix [ (n'ye |a()] meyo |â]| is fmix rounded up the nearest multiple of n'yo |á]; the number of samples needed for D, is (which depends on log( [F), log(o), and other parameters); and offs. max((N 1)- |ā], fmix) is constant offset ensure that the Ex-BMDD endogenous state has entered the terminal cycle induced by à, and that the exogenous state has mixed.
To see why Dj and D, are perfectly distinguishable only if n'yc neye- consider the sequence of repeated latent states that compose the cycle, [+0 ....; eye |ā] "th state in the eycle, then the resulting sequence, |sge, eve"" syye repeated states Therefore,
-al------1- If we only consider every
S/2p (nn--1)-1aal cannot contain any
Vn, m, = m (mod neye) + o- (3 = 6² aloms))s)
(
and in particular, if m'ye neye. then Then, through modular arithmetic, we can conclude that Dj contains observations only one latent state of the cycle, while Do contains observations of all of the other latent states.
(mod m'ye) + o" (Ijaj)+offs.) = 6² (z m +
Otteewise, due to the deterministic dynamics and repeated application of the same actions à the en- dogenous dynamics would immediately enter even shorter cycle the first time a state is repeated in aj1 all implying shorter period.
[6 lap-- 8 (rugo 1)-

---

Eii does Published as a conference paper at ICLR 2025
Mcanwhile, if rlye neye- then for each observation in D1. there is corresponding observation in Du with the same latent state, that is saally-dentically and independently distributed (i.e, they are collected 2 fiix steps apart). In partieular, consider the (unknown) subset D(")) ( C Du, defined as:
DOd = (a (22mm [a E {O,..., 1}}
where j' =- (-I (nlye [āl]] Weys 1)%eeye + 1. From arithmetic and applying Equation g see that all observations in D and all observations in D1 have the same endogenous latent state. Additionally, all observations in D6 and D1 are collected at least Emix steps apart.
CycleFind Phase 2: Identifying latent states in the cycle. Once neye is known, CycleFind can identify the latent states which re-occur every neyelāl steps in FCF. These latent states are not necessarily distinct from each other, and may also have been discovered already in a previous iteration of CycleFind. Therefore, CycleFind extracts from FCF datasets Di for each position in the cycle: i E {0.... .... /eye - 1}. CycleFind also uses datasets D(s) collected in previous iterations representing the already-discen -issoodedddeeeeeed states in S CycleFind determines whether two datasets (either collected in this call to CycleFind, or collected in previous calls) represent the same latent state by attempting learn classifier f FF that distinguishes them: if they both consist of near-i.i.d. samples of the same latent state, then it is highly unlikely that such a classifier exists.
To ensure the near-i..d. property "well enough." steps within each individual Di; and that, when trying to distinguish two datasets D', D'j which were both collected during this round of CyceFind, there are two (sufficiently large) subsets of DD and D'
only need that samples are separated by fmix
respectively such that all samples in the two subsets are collected at least émix steps apart ensuring this second condition only doubles the number of samples we must collect. Thus, we do not need to "wait" mixx steps between collecting each usable sample from ICP; rather, we collect usable sample for each latent state once for every roughly 2 max(fmmn neye lāl) steps. This is why not appear in the largest term (in (S)) of our asymptotic sample complexity.
If it is determined that some D; represents a newly-discovered latent state, then a new state s' is inserted into S' and (ss'') is initialized as D,, Once all states in [s] ...., S 1a1-1 have been identified, the action sequence à can be used to add them to the learned transition dynamics T".
STEEL Phase 2: Collecting additional samples to train encoder Once we have the complete latent dynamics graph. the determinism of the latent dynamics allows us to use open-loop planning to efficiently re-visit each latent state, in order to collect enough samples learn a highly-aceurate encoder. Note that we can navigate to any arbitrary latent state in 1) steps, so we can visit every latent state in (S] D steps. STEEL collects datasets D(s) for each latent state s where, within each D(s), the samples are collected at least fmnx steps apart: therefore, it can add one sample to each dataset (s)) at worst roughly every max( D, steps.
STEEL Phase 3: Training the encoder. Finally, STEEL trains the encoder. Specifically, for each latent state s E S', it trains binary classifier f. E F distinguish D(s) from Wy'ss (s) D(s'). To ensure that only the correct binary classifier, (o(o-(=))(ar), returns 1, we ensure that each Ja has an accuracy of -e/\S) on each latent state. We guarantee the accuracy of each classifier each latent state separately and apply a union bound: note that because we use a union bound here, we do not need the samples in different datasets D(s), D(s') to be independent, which is why we are able to collect samples more frequently that every fmir steps. Finally, we define o' (æ) == arg max, (x)).
5 SIMULATION EXPERIMENTS
We test the STEEL algorithm on two toy problems: infinite-horizon environment inspired by the "combination lock" environment from Efroni et al. (2022b), and version of the "multi-maze"
IIn some scenarios, one might not need to learn an encoder at all. Note that the latent state s of the agent known at the last environment timestep t of Phase of STEEL. At this point, the full latent dynamics are already known. Thus, if the agent is "deployed" only once, immediately after training such that the latent state does not reset, then one could keep track of 8 in an entirely open-loop manner while planning or learning rewards, encoder. In this case, the sample complexity terms involving e disappear.
without needing
8

---

Published as a conference paper at ICLR 2025
(a) Combination Lock Environment
(b) Multi-Maze Envirennent
Latent Dynamies
Latenz Dynamnies
Observed Tranition Example
Ozeerved Trazstion Example
Figure 3: Visualisations of the simulation experiment environments. For both environments, show the ground-truth latent dynamics T (in the case of the combination lock, we show an arbitrary instance of T, for some [ag. (k.....--11 and example transition in the observed space X.
environment from Lamb et al. 2023). In our combination lock environment,
fo. 1], {O.... K 1}, and there is some sequence of "correct" actions [4r: "6.:...K-- such that T(i. az) i 1. but T(i, 1 a;) = O. In other words, in order to progress through the states, the agent must select the correct next action from the (arbitrary) sequence a otherwise, the latent state reset to O. The observation space = {0, 132, where L >>> K. Some arbitrary subset size K of the components in X' are indicators for each latent state in S: that is, vi S, [O...,,L- (re) 1 + 8t = i. The other L - K components in X' are independent two-state Markov chains with states {0,1}, each with different arbitrary transition probabilities (bounded such that no transition probability for any of the two-state chains less than 0. Because each component is time-correlated, they all must be contained in E, so |8] 2 In the multi-maze environment, the agent learns to navigate a four-room maze (similar to the one in [Sutton et al. (1999)) using actions A {Up. Down, Left, Right}. The latent state space has size 8 68. However, the observation r E X' in fact consists of nine copies of this maze, each containing a different apparent "agent." " Eight of these "agents" move according only present in one of the mazes. Because the eight distractor mazes can be in any configuration and have temporally-persistenn state, we have that |3| = 68 For both environments, we use the hypothesis class F == {(x - (zr)al E {0, dim(X)) - 1}}. In other words, the hypothesis class assumes that for each latent state s, there is some component of the observations such that (z) if and only if (z), = 1. The two environments are visualized in Figure 3
random actions: the true controllable agent is
There are four sources of potential variability in these simulation experiments: (1) the random ele- ments of the environments' dynamics, Te, Q, and Cinits (2) the starting latent state Sini: (3) steps in Algorithm| [ that allow for arbitrary choices (e.g.. the choice of action = [a] in the first invocation of CycleFind); and (4) the parameters of the environment, such the "correct" action sequence (o0....%%-1 in the combination lock. STEEL is designed in such a way that, with high probability (i.e., if the algorithm succeeds), no choice that the algorithm makes in terms of control flow or ac-
Combo, Lock
Combo. Lock
Combo. Lock
(K 20)
(K 30)
(K 40)
Multi-Maze
Fixed Eny. Accuracy
20/20
20/20
20/20
20/20
Fixed Env. Steps
188658210
4286241:0
7914856-10
41003875:10
Variable Env. Accuracy
20/20
20/20
20/20
20/20
Variable Env. Steps
2.00-10
4.78-10
9.59-10"
4.13-10"
-0
+4.36-10
11.1--106
+1.11-108
Table 2: Success rate and number of steps taken for STEEL on both simulation environments. For all experiments, we set 5 = .05. For the combination lock experiments, we set T, = 512, and use the (intentionally loose) upper bounds N = D K + 10 (= [S] + 10) and fmix 40. For the multi-maze environment, how chose the (loose) bounds fiis 2 fmix-
N D 80 (> |S 68), and [ 300. See Appendix D for

---

Published as a conference paper at ICLR 2025
tions will depend on exogenous noise Therefore, if we hold (2-4) constant, we expect the number of environment steps taken to be constant, regardless of the exogenous noise. To verify this, we test both environments for 20 simulations, in both "fixed environment" setting with (2-4) held con- stant, and a "variable environment" setting with (2-4) set randomly. We test the combination lock environment with latent states K E (20,30, 40}. We measure the success rate in exactly learning o" (x( and I' (up to permutation) and the number of steps taken. Results are shown in Table 2
STEEL correctly learned the latent dynamics T and optimal encoder o² in every simulation run; and we verify that the step counts do not depend on exogenous noise. In the combination lock experiments for large K, which are hard exploration problems, the total step counts were many orders of magnitude smaller than either the size of the observation space (ze 10 10154); or the reciprocal- probability of a uniformly-random policy navigating from state 0 to state K ( 10 for K = 40). This shows that STEEL is effective at learning latent dynamics for hard exploration problems under high-dimensional, time-correlated noise. For the multi-maze experiment (which not a hard exploration task), STEEL took few orders of magnitude greater steps than reported in Lamb et al. (202)) or Levine et al. 20024) for the same environment (+= 10 10a steps). However, unlike these prior methods, STEEL is guaranteed to discover the correct encoder with high probability; this requires the use conservative bounds when defining sample counts d, Nsamp, eye. and Rsamp. in Al Algorithms [ and and in making other adversarial assumptions in the design of the algorithm that ensure that is correct and sample-efficient even in pathological cases. Additionally, note that the encoder hypothesis class F used in this experiment has no spatial priors. By contrast, [amb et al. (2023) choose a neural-networt encoder for this environment with strong spatial priors that favor focusing attention on single maze, using sparsely-gates patch encodings (and Levine et al. (2024) use this same network architecture in order to compare to Lamb et al. (2023) this difference in priors over representations may also account for some of the gap apparent sample efficiency.
In Appendix F] we present an additional set of experiments on a family of tabular Ex-BMDPs which are known to particularly challenging to "multistep inverse" methods, such as those proposed by Lamb et al. (2023) and Levine et al. (2024). We find that, for sufficiently large instances of these environments, STEEL can in fact empirically outperform these prior "practical" methods.
6 LIMITATIONS AND CONCLUSION
major limitation of STEEL that may constrain its real-world applicability is its strict determinism assumption on T. In the episodic setting, EEfron et al. (2022b) can get away with allowing rare devi- ations from deterministic latent transitions (no more often on average than once every 4|5| episodes) because the environment resets "erase" these deviations before they can propagate for too long, By contrast, in the single-trajectory setting, the STEEL algorithm is fragile even rare deviations in latent dynamics; ameliorating this issue may require significant changes to the algorithm.
second barrier to practical applicability of STEEL is the need for an optimal training oracle for F. While this is tractable for, e.g., linear models (with the realizability assumption ensuring linear separability), it becomes computationally intractable for anything much more complicated. How- ever, this kind of assumption is common in cample-complese results; and could be worked around in adapting STEEL to practical settings
Two additional limitations to work are the assumption of reachability of all endogenous latent states E S. and the requirement that an upper-boune on the mixing time of the exogenous noise be known priori. However, in Appendix E we argue that these assumptions are in fact necessary, for any algorithm in the single-trajectory. no-resets EX-BMDP setting.
Finally, the core assumption that S is finite and small is of course a major limitation: sample- efficient reinforcement learning in combinatorial and continuous state spaces is a broad area of ongoing and future work. Despite these limitations, STEEL represents what we hope is important contribution to representation learning in scenarios where resetting the environment during training not possible, and observations are impacted by high-dimensional, time-correlated noise.
This property is theoretically important because it ensures that the decision to collect given observation is independent of all previous observations, given the ground truth dynamics and initial latent state. Ffron et al. (2022h) adapts PPE to practical settings in their experimental section.
Similarly to how Efroni
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
A portion of this research has taken place in the Learning Agents Research Group (LARG) at the Artificial Intelligence Laboratory, The University of Texas at Austin. LARG research is supported in part by the National Science Foundation (FA1N-2019844, NRT-2125858) the Office of Naval Research (000014-18-223)) Army Research Office (W91 MENNF-33-2-000- Lockheed Martin, and Good Systems, research grand challenge at the University of Texas at Austin. The views and conclusions contained in this document are those of the authors alone. Peter Stone serves the Executive Director of Sony AI America and receives financial compensation for this work. The terms of this arrangement have been reviewed and approved by the University of Texas at Austin in accordance with its policy on objectivity in research. Alexander Levine is supported by the NSF Institute for Foundations of Machine Learning (FAIN-2019844), Amy Zhang and Alexander Levine are supported by National Science Foundation (2340651) and Army Research Office (W911NF-24-1-015
101IN-17722-01
REFERENCES
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural com- plexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095-20107, 2020.
Yuan Cheng, Ruiquan Huang, Yingbin Liang, and Jing Yang. Improved sample complexity for reward-free reinforcement learning under low-rank MDPs. In The Eleventh International Confer- ence on Learning Representations, 2023. URL
net/ forum?id=
jpsw 7r
Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. On oracle-efficient pac with rich observations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Infor- mation Processing Systems, volume 31. Curran Associates, Inc., 2018.
Simon Du Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning, pp. 1665-1674. PMLR, 2019.
Yonathan Efroni, Dylan J Foster, Dipendra Misra, Akshay Krishnamurthy, and John Langford. Sample-efficient reinforcement learning in the presence of exogenous information. In Confer- ence on Learning Theory, pp. 5062-5127. PMLR, 2022a.
Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alckh Agarwal, and John Langford. Provably filtering exogenous distractors using multistep inverse dynamics. In International Con- ference on Learning Representations, 2022b. URL
net/ forum?
id=RQLLZMCeeQU
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E. Schapire. Contex- tual decision processes with low Bellman rank are PAC-learnabl.. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp. 1704-1713. PMLR, 06-11 Aug 2017. URI https:://procee wwteeeiiiggg.... ngs ml ww.. press hgw.mlr.pr.es........mmmmmp..
vv /tt varress/v70//iann/ 7c.
Jacob Eeuwe Kooi, Mark Hoogendoorn, and Vincent Francois- Interpretable (un)controllable features in MDP's. In Sixteenth European Workshop on Reinforcement Learning, 2023. URL
https::/cpenreviiee wwwprreewwwweeteeeee.. .et fttet//oormm1dd==222
Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Rajiv Didolkar, Dipendra Misra, Dylan J Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed discovery of control-endogenous latent states with multi-step inverse models. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL https: http:///ppneeviii.. net forum? 1d=TNocbXm5M7
11

---

Published as a conference paper at ICLR 2025
D.A. Levin and Y. Peres. Markov Chains and Mixing Times. MBK. American Mathematical So- ciety, 2017. ISBN 9781470429621. URL httpss://book: google..oo..oooooooooooco www.:/es.cooooocaaaoomo ?id= E2C8DWAAQBAJ
Alexander Levine, Peter Stone, and Amy Zhang. Multistep inverse not all you need. Reinforce- ment Learning Journal, 1, 2024.
Zakaria Mhammedi, Dylan Foster, and Alexander Rakhlin. Representation learning with multi- step inverse kinematics: An efficient and optimal approach rich-observation In International Conference Machine Learning, pp. 24659-24700 PMLR, 2023.
Zakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. The power of resets in online rein- forcement learning. In The eiirty-eighth Annual Conference on Neural Information Processing Systems, 2024. httptttttoppeeeeeeeeeee
net. www.eeeeooormmnde1aaacaaea.
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state ab- straction and provably efficient rich-observation reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, vol- ume 119 of Proceedings of Machine Learning Research, pp. 6961-6971. PMLR, 13-18 Jul 2020. URL
.html
Richard Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181-
211, 1999.
George Trimponias and Thomas G Dietterich. Reinforcement learning with exogenous states and rewards. arXiv preprint arXiv:2303.1299 2023.
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline RI in low-rank MDPs. In International Conference on Learning Representations, 2022. URL
/ net wttp:/ooommiddddiSsS---
Tongzhou Wang. Simon Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised MDPs: Learning world models better than the world itself. In Kamalika Chaud- huri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Pro- ceedings of the 39th International Conference on Machine Learning, volume 162 of Proceed- ings of Machine Learning Research, pp. 22591-22612. PMLR, 17-23 Jul 2022. URL https: //proceedings.. . .ml. www.mm.pp.s....i.mmmmmmmmm http:///ww1l//aa22ccoooottc www./2c..g22.m
Kenneth S Williams. The power of 2x matrix. Mathematics Magazine, 65(5):336-336, 1992.
Lili Wu, Ben Evans, Riashat Islam, Raihan Seraj, Yonathan Efroni, and Alex Lamb. Generalizing multi-step inverse models for representation learning to finite-memory pomdps. arXiv preprint arXiv:2404. 14552, 2024.
Wanqiao Xu, Shi Dong, and Benjamin Van Roy. Posterior sampling for continuing environments. Reinforcement Learning Journal, 1, 2024.
12

---

Published as a conference paper at ICLR 2025
A FULL ALGORITHM
The STEEL algorithm is presented here in full Algorithm with major subroutine, CycleFind, split out as Algorithm
Algorithm 1: STEEL
Input:
Access to EX-BMDP M, access to training oracle for function class F, knowledge of upper bounds N, D, fmix. parameters 6,
Initialize learned latent state set S", initially empty;
Initialize table of collected datasets for each latent state: D: - ++M(X);
Initialize learned latent dynamics: T" (S' U (1]) x A (8' ( (When state added S',
we initially set va E A, T(s, a) == L. Also, we set va E A, T(L,a) == 1 as permanent dcfinition.);
// Phase 1: Discover latent dynamics T. Chose arbitrary a A:
,,,",,""a' Soum. +- CycleFind( (', D,,,") /T/ // Special case for first iteration while 3s E ('
S', D, T
;= do Initialize + and initialize action list à + []:
S" while B non-empty do
Chose arbitrary 8
- 5 {a};
Let := minimum-length sequence of actions such that
(T d%),, (4)),))) "., d()))-)) L. (This can be found using
Dijkstra's algorithm.):
à + à. a':
B B
5+1s {s" S' E. ( (s',"))) (1), 02), ",)),..(a))))11 = s")
S", D,T"" Semm + CycleFind (',
// Phase 2: Collect additienal latent samples to train encoder.
Let = 21
while 3s e S | < d
Let C (s ('[] A Scum
Use T" plan cycle of actions starting at Sun. that visits states and then returns to Scur. by greedily applying Dijkstra's 's algorithm repeatedly;
[a] Emixs use T' to extend by repeatedly inserting the shortest-length self-loop of some state visited à into until |a| fmix ;
Exccute all actions in ā once without collecting data;
while Vs C: [D(s)] < do for in āo, ... d do
Take action
M;
Som (Scur, a):
if Spws is being visited the first time in this execution through à then Let Feur. = the observed state M;
D(Scun) + (ss.m.)) W {
/ Phase
Train latent state encoder &'
for E S' do
Let Do := W.rss (x) D(s'); D1 : D(s);
Apply training oracle to distinguish Dg from D,, yielding Sa F return S', and ø' (xx) = arg max, f=(x);
B PROOFS
B.1 STEEL
Here, we explain the STEEL algorithm (Algorithm I. and prove the correctness of Theorem D STEEL proceeds in three phases: in the first phase, we learn tabular representation of the endoge- nous latent states and associated dynamics of the Ex-BMDP. For each S, we also begin to collect dataset Dss), where for each æ E D(s), have that ø² (xr) and additionally where all samples in D USes D(s) were collected from the EX-BMDP M at least fmix time steps apart.
13

---

Published as a conference paper at ICLR 2025
Algorithm 2: CycleFind Subroutine
Input: Action list à; current learned state set S datasets D, and transition dynamics T Also, access training oracle for function class F, knowledge of upper bounds N, fmix,
Ex-BMDP M,
and D, and parameters 8, find Length cycle,
// Phase 1:
Peye āl
/ 16
Let Thamp. ==
MIIININ-
Let Cinit == (20mix 3.N |a| 2) Nhum eye Collect sequence of observation ZCF [11,, Wieeenn from M by taking the actions in repeatedly in loop, total of Cimi actions. (Action à, % [a] taken after observing x;.):
lmix N |á + max((N - -1)-ja | lmix):
Let I. = [4:]]]+ -11))(())))))))....
Initialize Veye 1: // Default value if no other Maye is Teyy for Plye in IN, N-1...,3,21 do
Let [fmx/ (Tay [a)).;:== Plye. k === Sun + (al/////mmaaaa (max((axx))))))) (N- 11C-N-11110 12-1á Amin)
2r- -rgeal+
Let Do == (e Stsl (0. 11, j E (1, Mayo - 1}}:
Let D₁ == ( i (0, k 1}};
Apply training oracle to distinguish Da D,, yielding (ve E Du. f(x) = and E Du, f(x) = 1) then
Neye 7lyoo
break;
// Phase 2: Assemble datasets for observations from cycle, ident ify new latent states, and update S', D, and T".
Let Phamp. = In (1
In
Let Reyo lål- (Maang. 1)-
+ +1)
+ fmix + max((N neye) [å], fmix);
[al-nige
Extend the sequence of ohservation #Cr length least c by taking actions repeatedly in loop on M, for max(0, - Cam) additional steps, so that Fcr [x1....eeh
Let max((N Meye)) lāl, /mix)
ná + (Nsamp.
|ā
Vlyyo
+ |a]
Vlayy
mmix:
Vi E (0..., Roye (ā) 1}. Let:
D'
x.y 1rk E {0, Nisamp.
1}, 3 offset {nu n663
imix
j = ho |â] Taye
+ affset offfet))%(mo)) lál]:
Reye
for i E .O.... New |o| - 1} do
Initialize StateAlreadyFound? + False; for E S' do
Let Do = D(s); Di = D/;
Apply training oracle to distinguish Do from D1, yielding E
not (Yr f (xx = and Vr. Do. f(x) = 1) then
2 4-8;
(Optionally; and only if D(s) was not initialized modified already during this call to CycleFind:) D(s) +- D(s) w D/;;
StateALreadyFoumm +- True:
break;
not centeAtreayFoundee then
Insert new state s' into S'; D(s')
D/:
for E {0. Deye |a 1} do
(" (.s(aa)
+ 8 eye
((-1)%
Scur. maxia (8))
return S', S'D, Senni
14

---

Published as a conference paper at ICLR 2025
Then, in the second phase, the learned dynamics model T" to efficiently collect additional samples for each learned latent state s E S' and add them to D(s), until [D(s) 2 d, where:
3|8 [n(16]5/911110))
d -
(10)
Finally, in the third phase, we use D learn an encoder d', which approximates 6* with high probability when the exogenous state e of the Ex-BMDP is sampled from its stationary distribution
Te.
STEEL relies on the CycleFind subroutine, which is described in detail and proven correct in Section B.I.1 This subroutine is given list of actions à and the previously-leamed states S', datasets D, and dynamics TV. It identifies and collects samples of all latent states in some state cycle which traversed by taking the actions á repeatedly, and also identifies the latent state transitions this eycle.
We restate Theorem Uhere:
Theorem 1. For an Ex-BMDP M (., Tiill) starting at arbitrary endoge- E S, with [S] S N, where the exogenous Markov chain Te has mix- ing time at most bmito and where all states in S are reachable from one another in at most D steps; and corresponding encoder function class F such that Equation 3 holds, the algorithm STEEL(M,,..) N, D, Emixs 6, E) will output a learned endogenous state space S', transition model T", and encoder &', such that, with probability least 1 6,
nous latent state
functioon S = (S), and under some bijective function o S - S', it holds that
Vs E E
T' (o)),, (sss),)), and,
(4)
Under the same bijection
vs E S, Pr (8(x)= (x))
(5)
IZ Q(s.). eriig
where me is the stationary distribution of To
Furthermore, the number of steps that STEEL executes M scales as:
O (NNDISSPAA log 8
F
SSIIA|imi
log
N
|
log
F
S
log
F
8
where ()()()) O (z)()
For the sake of our proof, we will treat the ground-truth properties of the Ex-BMDP, such as the latent dynamics T, as (unknown, arbitrary) fixed quantities, not random variables. Similarly, will treat the initial latent state 8init as an arbitrary but fixed quantity, rather than a random variable. Furthermore, in the proof, we will treat decisions that specified (implicitly
explicitly) as "arbitrary" in Algorithms/ [ [ndd] (such as the choice of the action à = [a] in the first invocation of CycleFind, or the choice of "shortest" paths in cases of ties when Dijkstra's algorithm used) as being made deterministically, such as by pscudorandom process - crucially, we require that these choices are made way that does not depend of the observations of the Ex-BMDP.
2
This leaves the exogenous noise transitions Te. the emission function Q. and the initial exogenous latent state Cinit as the only sources of randomness in the algorithm. We notate the sample space over these three processes together as Ω. Throughout the algorithm, we will ensure that decisions such as control flow choices, choices actions, and choices how to assemble datasets, are made deterministically, independently of Ω with high probability. That is, unless the algorithm fails, these decisions will be fully determined by so. T, and algorithm parameters. While whether or not the algorithm fails will depend (solely) on 2 we will ultimately bound the total probability of failure as less that by union bound, so that at each step in the proof, we can treat the algorithm's choices as statistically independent of Ω.
STEEL begins by repeatedly applying the CycleFind subroutine. CycleFind identifies a cycle in the latent dynamics 7 of the Ex-BMDP, and collects observations of the states in that cycle. In Phase
15

---

Published as a conference paper at ICLR 2025
of STEEL, throughout these application of CycleFind. the algorithm maintains representation of the learned Ex- BM-D "so far", in the form of an incomplete set of learned states S' learned latent dynamics T' (S' U (++) A (S' u {+}), and table of datasets corresponding each latent
state S - M.
We first describe the CycleFind subroutine in detail:
B.1.1 CYCLEFIND SUBROUTINE
Here, we describe CycleFind, and prove its correctness. CycleFind accepts as input list of actions and the EX-BMDP M starting at an arbitrary state zp. CycleFind also takes the representation of the learned EX-BMDP "so far". in the form S', D and T". We assume that, for each E S' all of the previously-ohserved observations in D(s) were collected with gaps of at least fmix steps between them. Also, we assume that As E S', [D(s) 2 Nsamp. where, in terms of the upper bounds N 2 |S| and mix 2 mmix, and total failure probability 6,
= [ -
8 In 4|AA No. (D 1)-15)
Tdsamp.
==
(1I)
16
CycleFind first proceeds to take the actions âg, ..... -...@ia-1 repeatedly in a loop. CycleFind then uses this collected sequence of observations (which may need be extended by cycling through for additional iterations) to learn new states and update S', D, and T".
We will show that sequence of latent states visited by CycleFind eventually periodie; that is, guaranteed to eventually get stuck in cycle of latent states, with a period in the form Meye [ā], for some Peye I N. The goal of CycleFind is to:
1. Identify the period of this cycle. (That is, determine neyy..)
2. Use this period to extract from the sequence of observed states some new multisets of ob- servations D/= M(X)) for E (O.... fleye |å - 1}. which each contain only one unique latent state, corresponding to the position in the cycle. These multisets will only con- tain observations collected at least fmix timesteps apart, so will be close-to-i.i.d. samples. (Depending on Reye. may need to perform additional cycles of data collection this step.)
3. Identify which of these multisets D/ have the same latent states that have been previously identified in S", and which have new latent state, and determine among the new multisets which ones have the same latent states to each other, and which are distinct. This allows us to update S' and 'D with the new samples from D/ while maintaining the property that Vs E S', Va, x' E
= oz (z'), and Vs, s' S', Va E D(s), x' D(s'), + d² (x').) We also update the learned transitions T".
4. Return the updated learned state set S', datasets D, transition function 7", and the current latent state of M, 8eam. E S'.
Specifically, CycleFind has the following property:
Proposition For any action sequence of length at most (D + 1)N, there exists at least one -oo eyee
sequence of ground-truth states in S, [&0 ,81 .... (al-Tis
Taye -1.. for some Nhye N, such that Vi E {0, |a] Tlaye - 1}, Tss (.%%aa = eyen
eycs
Given a sequence of actions à learned partial state set St, transition dynamics T and datasets D which meet the following inductive assumptions:
There exists an injective mapping - S' + such that
Ys E' A, ('ss - 1 T(ssa)=_Voo(((s__1)
(12)
and additionally,
E VI E (s)s (2) -0 -
(13)
Vs S', ID(s)) Tisapp and for each the samples in D(s) were all sampled from M at least bmix steps apart. Additionally, the choice to add any sample D(s)) was made
16

---

Published as a conference paper at ICLR 2025
fully deterministically (as a function of Smnit, the timestep which collected, and algorithm parameters), and independently of the random processes captured by Q.
The choice of action sequence is similarty fully deterministic and independent of Q.
then, with probability at least:
8
(14)
2. (AA
CycleFind will return updated S', D, T', and Sewn. which meet the same inductive assumptions, and for which additionally:
The image of the updated S (new) o (S( is (non-strict) superset of (S"), which additionally includes all unique states in some [s6
cyyc
cycc 1almme--1
The transition matrix T (new)) is a (non-strict) superset of the old transition matrix T" (in the sense that its domain is now 8 ((nww) 2 ((ddy and (o (s.a) + then T (new (5, a) of 1), and T (new) additionally includes the transitions corresponding to the cycle; that is:
vi E (0,..., (a..://l-nuu. Maye - 13, Js, S' :0::"""""00 (s) s; ryyc no Ao
oyee (1+1)% m
T (nne) (s, d:%(aa) =s'
(15)
The final observation sampled by CycleFind from M is such that (Scurr) 0³ dx(x).
Additionally CycleFind will take at most:
max (22mmi- + 3NN |ā| 2) nzamp. N N|aãl, (fmix + +18)
+ [S]]|a]] 1) -Daamp +1)
(16)
+ max(N [ |a| Emixs actions,
where nsamp. is defined Equation 77 and Reamp. defined in Equation 18
Proof. Determining Plye:
CycleFind initially takes Cinit actions, where, in terms of the upper bounds N 2 |S| and Eiix 2 tmix.
Gintt = -22 + 3NN |ā] - 2) Tesamp. eye.
Emix -1a + max 1) (ā], Emix), (17)
where
16 1
Rsamp. eye. :=
In
(18)
4|A1- N (N- 1)-
CycleFind first takes action āo, then à1 then 02, etc, until taking action ûjal-1, at which point repeats the process starting at âg, for total of Cinit steps. The observation after each of these actions is recorded as 2CF := Femmll Let SCF := [s,, Scmll be the (initially unknown) latent states corresponding to these observations; that is, (x)) for each in xCF. (For indexing purposes, zo and will refer to the observation and latent state, respectively, of the EX-BMDP before the first action was taken by CycleFind. However, these will not be used by the algorithm.)
First, we show that SCF must in fact end in a cycle of period Paye |ā], for some Reye S N. Let
max(0, fmix - (N 1) lā)): Sper. consist of every |å] th element in SCF starting an offset Note that the evolution from one that is, Sper. := [Sm: La] Sm+22aä .... Sm+ ((am (idmmmm//lllaaal state to the next in Sper. is deterministic, because it is caused by the same sequence of actions, (am % [ā] (m++1) [a]. -.ll-..,,mmmt (m)+|||))))))) jai]- being taken after each state. That is, if Sm+ijal = 8 and 8m+j--la = s and (m-(i+1)-1a) = s', then $m+(j+1))laal = s'. As consequence, sm+i = 8, and the next occurrence of the latent state in the sequence Sper. is (i+t(+)))l s, then all subsequent states in the sequence Sper. will consist of repetitions the sequence of (i-(i-))-aa) through (m+()+))-).)
17

---

Published as a conference paper at ICLR 2025
Then Sper. must consist of some sequence of 'transient' latent states which occur only once the beginning of the sequence, followed by a repeated cycle. Because these states never re-occur, the transient part lasts length 7um S -
Then, the period of the cycle is Reyes where Tum + Reye I N. Note that the cycle in Sper. contains repeated states. (Otherwise, the span between the first two repetitions of state in the eyclic sequence in Sper. will itself repeat indefinitely, so we can analyse this smaller cycle as the cycle of length Neye.)
There is a corresponding cycle in SCF, of length Neye Jā]. To see this, note that for all 2 Vlm have that Sm+jjá *m++] Furthermore, for all in {0,-..la| 1}. the sequence of actions taken between Sm+ija and +m is the same as the sequence of actions taken between Therefore sm++(ä++
Smti Lāge and Sm++lll++
Thus, for any gen-
eral /' 2 Tim [ (which always can written as i' i j + j) we have that Sm+r' 8m48 |a)- However, this cycle may contain repeated states.
In order avoid the transient part of SCF, and prevent sampling observations with ex- ogenous noise that is correlated to samples taken in previous iterations of CycleFind, we skip the first max((N - 1) |ā], (fmss) transitions in scr. For convenience, we will let 8 8/ m|+x 11) and similarly I; == #+jja|+ ((nx(((N))-(l,,fmm))) Note that the sequence [5o, $1....] is equivalent to Sper. after skipping the first N - ner.. elements of the sequence.
Because the cycle in Sper, contains no repeated states, we have that
Vi,j E N, 8; = 33 + ii (mod Meye)
(19)
In order to find Neyes we test the hypothesis that Peyy Weye for each Wlye E {N,..., until we identify neyye none the tests pass, then we know that Weye = 1. The test for each hypothesis Mlay eye has a zero false-negative rate. Consequently, the loop will always end before Neyyc Rey at each iteration, it must always be the case that Reye I nyyye A failure can only occur if the test that Reye rlye has a false positive, when in fact neye Pleyo
.......... 2). in order,
Each test proceeds as follows:
Let
[fmix/ / (Meyo [a)) and = q N'ys
Let k Cartr-la lā| mamxaxxaa &anss)
2r-1a t
Let Do = (Tr+(2r 3 i ko k- j E (1,.., - 1}}.
Let D1 = ( (0...., k 1}}.
Use the training oracle to try to learn to distinguish Do from Dr, yielding / E F.
If Maye
Veyo then.
Note that, because Neyel /r
ViE N, (2r ) =0 (mod neye),
(20)
but
Wi E E.EN (I, Klgs 1}, r ++ (2r + neye) + = #0 (mod neye).
(21)
Consequently, by Equation 19 all elements of Dj will have the same latent state, and none of the elements of Do have this latent state. By realizability, f will have 100% accuracy on the training set. (This is the "true positive" case of the test.)
Conversely, if Neye < nyyer there is only a small chance that any classifier f will have 100% accuracy the training set.
18

---

Published as a conference paper at ICLR 2025
Define j' the (unknown) value j' := (-r 1) % neye +1. Noting, by assumption, that Reye < Peye we have that j' E (1,.., V'ye - 1}. Then, Vi {O...., 1}. we have that E (2 jtt Dn. while 2(2 Dj. However, also have that:
(2r )e ) + J"
(2r neyo) (mod Meye)
(22)
which by Equation
9 implies that
8y 844(2++)) 5(2r+ng)))
(23)
Now, we can define D C Do as
D8 (3++(2r+ (4 Mtylic
to....,k- 1}}.
(24)
Fix any arbitrary classifier F.
In order for f' to have 100% accuracy on the training set, we must have f'(x) = 1 for all D1, and (x) for all æ D(
-
Note that all observations in D1 w D are collected at least mmix steps apart from one another. (Specifically, ( jitj' is collected (r+ |a| 2 r- |ā| 2 tmis steps after ž (2 )is and (r (r+ Maye -') la] 2 r |ā] 2 tmiz steps before (nrr+)))
-
Because, additionally, DO6)) and Dj are defined independently of Ω by Lemma we have:
-
vt (t'la E D1 4 ) P. - 1/4 Pr( (āx) = 1[(D, w D oe,,""(()) = 8) P. p. 1/4
(25)
where (D1 H D6 )+t refers the samples in Dj s D collected before H, and:
Vs S. p. : Pr(f((x) 1|r C e). TT).
(26)
Then, by Equation 26 the probability that f' returns on all samples D1, and on all samples in D is at most:
Pa + 1/4) 11 - (P - 1/4)).
(27)
By Equation 23 this is:
(P + 1/4)- -1--0110 (pä 1/4)).
(28)
Rearranging gives us:
FPR(f') < II k-
1=0 (-p² mbye + Pa + 5/16).
(29)
Because Vp, -p² + + 5/16 < 9/16, we can upper-bound this as:
9
FPR(f') < 11-- (-p P Pa mb
+ +5/16) <
(30)
16
As uniform convergence bound:
9 k
FPR() S [Fl
(31)
16
Finally, take union bound over all values of n'ye- To do this, must lower bound for all values of nyye First, note that
imix Nla Lmix mleyl |ā|
2 r|a
(32)
19

---

Published as a conference paper at ICLR 2025
Then:
(fiix + N - 1) (2 ssamp. eyc. 1) + N [a]
Reamp, eye. + r -
k
2r [a] Weyo |a]
M [āl (2
Nsamp. eyc.
-1) + N [ā]
Msamp. eyc. + - |ā|
(33)
2r [ā| + N |
NV Rsamp. eyc.
2 nsamp. eye.
So we have that, by union bound over all values of neyc:
9
Thamp crc.
FPR()) s (N - F)
(34)
16
Collecting D':
We now know that sCF eventually enters a cycle length Reye |ā]. where Reye is known. This is the latent-state cycle [s0 ps ,31 ,..., ,..., cyes -1- mentioned in Proposition
Depending on the value of Neye, we might now need to extend øCr (and, respectively sCF) by making additional loops through à, until the length of ICF is at least c, where:
Fiixx
C: = 2 Peyy |ā|-
(Nsamp. 1)
+ +1)
+ fmix + max((N - -eye)) |ā], mmxx). (35)
Peyye
This will entail taking additional max(c Ginit; 0) steps M. Note that in the worst case, this means that CycleFind takes total of at most:
max c) S max (2fmisx + 3NN. |á] 2) Msamp. eyc.
N |ā].
2 (Emix + |8| |a| 1) nssmm.. +1]
(36)
mmx(N N |ál [o| mmixx 0) actions.
We now define how to collect two datasets for each position in the cycle in SCF, D and D. for each E (0.... Meye |å| 1}. Specifically we take:
D:a
xz Ik E {0,
Nsamp. 1}
(37)
Imix
1)
na + (i -no% meye 1al)
j == k.
|a| Neye
[@]
Meye
where we let
nop max((N- Meye) lål, fmix),
(38)
and
DB x 12 {0, Reamp. 1}:
(39)
Emix
1)
n/ + (i (1-n))) 16)% TTeye [a)]
j = k. |āl Meye
Meyo
where
((a) neye
Emit
1) + |ā|| -Tay + fmix-
n/g = no + (Asamp. 1)
(40)
[al -ayyc
Note that, because we know that SCE enters a cycle of length neyo (ā] after at most (N- -Noye))) ni transitions, have that,
Vi,j 2 no: i=j (mod neye: (ā() = 83-
(41)
20

---

Published as a conference paper at ICLR 2025
Di and D are each near-i.i.d
D and D, are nearly-independent
D
D
DA
B
(0
10 150
Byp (0 140)
16 I9
( 140
Figure 4: Illustration of the sampling procedure for datasets Dj. The first goal is to ensure that for each cycle position i, the samples in Di are sampled mix steps apart from each other, and are therefore nearly i..d. We also want to ensure that for any pair of cycle positions j. there is large subset of D' that only contains samples retrieved at least mix steps apart from some large subset of Dj- (This second goal is meant guarantee that if Di and Dj represent the same latent state, then is unlikely that any classifier exists than can separate the two subsets perfectly, which would be strictly necessary to perfectly separate D/ and Dj). However, it is not necessary for all samples in U,D, to be collected bmix steps apart. Therefore, we collect an observation of each cycle position i E {0, Teye |ā }. all together, every (mix steps (rounded up the cycle period). We continue until Namp observations of each position are collected; then wait fiix steps and collect nsamp additional observations of each cycle position. This process ensures that Dj and DDj contain complementary subsets D and Dj. each with least mammp samples, such all samples in DA U Dj are near-i.i.d.
Therefore, because:
Emix
) + rap + (i (i-no))% -1
k: [|all Reye
Meye lāl) =
(mod neye (ā))
(42)
[@] Meye
(and a similar equivalence holds for n(), we have that, for any fixed i, all observations in D.A y DD must share the same latent state. Also, for any fixed i, all observations in Dj and D. are collected least steps apart. Additionally, DA and DØ defined solely in terms of nlee and å, and so the selection of samples to put in these sets only depends on the sequence of latent states s that the Ex-BMDP traverses, and is therefore defined deterministically and independently of of 5 (assuming Peye is correctly determined). We therefore define
D' = D D.
(43)
and note that all elements in this set both share the same latent state and were collected at least i steps apart from one another.
Additionally, for any fixed pair j, all observations in DA e DJ are collected at least steps apart from one another.
Using the samples in sCF, this allows to construct D.A and D. for each i {0, Meye |å| 1}. where |DA| D samp.. See Figure for an illustration of the sampling procedure.
Identifying new latent states from D':
At this point, each set D/ consists of observations of a single latent state 8, but such sets D/ and D' may represent the same latent state, and Dj may contain the same latent state as some previously-collected (s)) for some 3 E S'.
In order identify the newly-discovered latent states add to S', and appropriately update D(-) and 7", we proceed as follows:
For i {0, Meye |ā| 1}:
21

---

Published as a conference paper at ICLR 2025
For each E S', use the training oracle learn classifier f E F, with Do ((s) and D1 := Dj If f can distinguish Du from D1 with 100% training set accuracy, then conclude (with high probability) that ((s) and Dj represent two different latent states. Otherwise, we conclude that D() and Dj both represent the same latent state.
If Dj is identified as representing some aiready-discovered latent state s E S' then discard D. (Or, we can update ((s) by merging the samples in Dj into it; this choice does not affect our analysis however, should avoid doing this if D(s) was either defined, or already updated, during this call of CycleFind: this is because two datasets Dj and Dj from the same iteration of CycleFind may contain observations that were sampled fewer than mix steps apart from each other, which would break the inductive assumption D(s) if they both merged into (s)).) Record this latent state as:
(44)
Otherwise, if D' does not represent the any latent state E S', then Dj (and D(!) represents a newly-discovered state. We update S' by inserting a new state s' into it, and update (ss) by associating s' with D.
S' S'U
(45)
D(d') = :
Finally, we also record this new latent state as
(46)
To analyse the success rate of using the training oracle to determine if a given ((s) and D' represent the same latent state, consider the following:
If D(s) and Dj contain different latent states, then f will be able to distinguish Do from D1, deterministically, with 100% accuracy on the training set (due to our realiz- ability assumption.)
Otherwise, D(s) and Dj both contain samples entirely of the same latent state, s.
Then, either:
The latent state s was identified before the current run of the CycleFind subroutine. Therefore, some subset of samples D/ C Do = D(s) were added to D(s) before the current run of CycleFind, such that D 2 Tesamp. Let D := D₁ = D', and note also that all samples in D W D' were collected at least fiii steps apart from one another. (This is by inductive hypothesis for D(s), by construction for Di, and by the fact that each run of CycleFind starts by "wasting" at least dmis steps.)
* The latent state was identified during the current run of CycleFind, such that D(s) = D' for some j < i. Then let D/ == Dí C Do and D' = D.¹ C D
Note that [76., DD' 2 nhamp, and all observations D/ w D2 were collected at least fmi steps apart from one another.
* In either case, the choice of samples to include in D U D was made determinis- tically and independently of 52 (by construction and/or assumption).
We define ps as in Equation 72 Note that the samples in D2 and Dí were observed at timesteps. Then Lemma is applica- ble, and we have that the probability that an arbitrary I' E F returns on all samples from Dj 2 D'; and also returns 0 on all samples from Do 2 D is at most:
least bmix steps apart, at deterministiiallly--ooss
(p. + 1/4)"smns.
(1- - (p. 1/4))
9
Thump
(47)
(-r, + Pp
5/16) 1516)+
16
As uniform convergence bound, we then have that:
9
Thomp
FPR()) < |F
(48)
16
22

---

Published as a conference paper at ICLR 2025
Note that at all iterations, |S'| I [S], so we train most |S] -Reye classifiers. Therefore, by union bound, the total failure rate is bounded by:
9
Pr(fail) S [S] neye [āl IFl
(49)
16
Note that the states o now represent the latent states associated with the cyclic part of TCF. Because we know the actions in the cycle, we can use this information to update
T", Specifically, Vi {0, .... [å] Reye - 1}, the action taken after s and before eye S ((+1))6(()))m)))
d.% al- We can then update:
eyC (* (0%)) < eye
(50)
(1++))(((()mm))))
Return the updated S',D, and T", well Scurc::
Returning the updated S,,D, and I" is straightforward. Note that the choice to assign or merge given D in to given ((s) depends only on the latent states in the datasets, and is independent Q.
We have then shown that, CycleFind succeeds, then states |s0 t neye |å] |a|- have been added to S' if they were not present already. These states correspond to the states in the cycle eyes
[=0 ...,, S..g& "Page-1ā|-14 and the corresponding transitions have been added to T"; furthermore, the datasets D(s) have been updated appropriately.
To determine the learned latent state the Ex-BMDP M after CycleFind is run, simply note that this is equivalently the state corresponding to the observation which we know belongs to dataset We then know that this observation must have the same latent state as the rest of eye (7oss1s)1111)
D
D' (a)) that is, the observation
The total failure rate for the CycleFind algorithm can be bounded by union bound from the failure rates of Parts 1 and 3 of the algorithm; that is, Equations 34 and
49 That is:
Phamp oy..
9
Pr(fail) S (N- 1) IF]
+ +|S|
Reye
|ā|
|Fl
16
16
S (N - 1) |F]
+ +N3 (D
+1 |F
(51)
A (N - 1]
4 1)|=1
< 2|A] N
All claims of Proposition have therefore been proven.
B.1.2 STEEL PHASE 1
Note that given fixed à, there might be multiple differrntt state cycles that could be discovered by CycleFind. However, only one will actually be discovered, depending on the state that the Ex- BMDP starts in well the not-yet-discoverll parts of the state dynamics. For example, consider an EX-BMDP A = (L, R). S (i, with the following latent dynamics:
L L i
(iiR 2 DL
If set à [L]. then, depending the initial state, CycleFind will either collect samples of and discover its self-loop transition, collect samples of 2 and discover its self-loop transition.
In order to learn the complete latent dynamics of the Ex-BMDD BMDP, we maintain a representation T' of the partial transition graph that has been discovered so far, and iteratively apply CycleFind using, at each step, an action sequence that is guaranteed to produce a cycle that is not entirely contained in the partial graph discovered far.
23

---

Published as a conference paper at ICLR 2025
Note that this is not simple as choosing sequence of actions that leads to unknown state transition from the final latent state the Ex-BMDP reached in the previous iteration of CycleFind. For example, consider the following partially-leamun latent state dynamics (with A {L, R))):
L
i
R
L
I.
R
4
PL
Here, the only unknown transition from an known state is the effect of the 'R' action from 4. Suppose know from the previous iteration of CycleFind that the current latent state of the Ex-BMDP is Scur.. = 3. Naively, it might seem if running CycleFind with à [R, R| would learn some new transition dynamics or states, because it would navigate through the unknown transition. However, this might not be the case fact. In particular, the 'R' -ranninsitoon from 4 might only be visited Iransiently. For example, suppose the full latent dynamics of the EX-BMDP are as follows (with the currently unknown parts shown in gray):
LO
6
L
R
2
R
R
L
L
R
6
3
LL
Then, if we run CycleFind with à [R, Rl,, it will converge cycle between the nodes i and 2:
R
R
L
6
L
222
2
R
R
L
R
5
3
DL
R
R
Note that the states (i and 2) and associated transitions that CycleFind converges were already explored, so we learn no new information from this application of CycleFind.
Instead, at each iteration, design à that cycle of the actions à be entirely contained within the currently-knowe partial transition graph. We show that the length the resulting is at most (D S).
We proceed as follows. Note that in the first iteration, before any latent states are known, we can simply use à [a] for some arbitrary E A. Otherwise, use the following algorithm:
Initialize with all of the previously-leareed latent states (that is, B < S'.)
While B is non-empty:
Remove some latent state s from B.
Use Dijkstra's 's algorithm to compute a shortest path the partial transition graph that starts at and ends at any not-yet-defined transition. (that, is, any transition for which (-. = L). ( Note this must be possible. Otherwise, because all states can reach each other in the full latent dynamies, if there such undiscovered edge in the
24

---

Published as a conference paper at ICLR 2025
same connected component as then would know that have already found the complete latent dynamics.) Also note that the shortest path through such an edge can have length at most D + 1, simply because:
* The length of the shortest path from & to the state with the undefined edge in the full transition graph 7' is at most D.
*
Suppose that some transition on this shortest path is missing in the partial transition graph T', Concretely, let d be the first state along this path such that the transition out of it is missing. Then, we have path from to d of length less than D, and we know that d is itself missing a transition. Then can be used in place of the original state with the undefined edge: it has missing transition, and it is at most
D steps, in T", from 8.
Let à' be the list of actions on the path we have found from s through an undefined edge. Note that taking actions à' from will result in taking an unknown transition,
and that [a'| < D +.
Replace B with the set of states that can result from starting at any state s' B and then taking actions á', according to the learned partial transition graph T". For a given s' E B. if this path leads to an unknown transition, do not insert any state corresponding to s' into the new B
Concatenate á' to the end of à.
Note that every iteration, [B] decreases by at least 1, the algorithm runs for at most [S'] iter- ations, so the final length of is at most (7 + 1)|S"|- Also note that, by construction, taking all actions in à will traverse an unknown transition at some point, starting any latent state that has been learned so far. As consequence, any cycle traversed by taking à repeatedly must involve at least one transition (and possibly some states) that are not yet included in the partial transition graph. Therefore, applying CycleFind using an â constructed in this way is guaranteed to learn at least one new transition. Therefore, fully learn the transition dynamics, we must apply CycleFind at most [A| [S] times.
Also, note that the process of constructing à at each iteration depends only on the partial latent dynamics model T", which in turn depends only on the choices of á in previous invocations of CycleFind, and ultimately these depend only on the starting latent state Sinit and the ground-truth dynamics T. Therefore à is at every iteration independent of Ω2 required by Proposition [
Then, assuming CycleFind succeeds at each invocation, by the end of Phase I, STEEL will have discovered the complete state set S and transition function T, up to permutation.
B.1.3 STEEL PIIAS 2
In the next phase, once we have completely learned 7'' (that is, once there are no state-action pairs E S E A for which (s a) is undefined), we collect additional samples of each latent state, until the total number of samples collected for each least d, where:
d = 3|8 In (n(16])77)1110)
(52)
We can leverage the fact that we now have complete latent transition graph as well as knowledge of the current latent state 8eum. from the last iteration of CycleFind.
To do this, we proceed as follows:
Use 1" to plan a sequence of actions ā such that:
fmix < |a| S max([S)- D. fmix + D), and
Taking the actions in â starting Scur. traverses cycle. That is,
T T'T' (T' (...T" (Seur,, ao), ū1); ā2), (ia),, a[][-1)) Scurr.
(53)
and,
- Taking the actions in a starting at Seam. visits all latent states S E S' \ (scur.) such that [D(s) < d least once.
25

---

Published as a conference paper at ICLR 2025
Note that planning such sequence ā always must be possible. For example, starting at Seum., we can greedily plan a route to the nearest
latent state 8 e S' (Scum:) J such that (D(s) < d and repeat until all such states have been visited, and then navigate back to 9cum. This takes at most |S D steps. If this sequence has length less than fmix. then we can insert a self-loop at any state in the sequence (such as the state with the shortest self-loop) and repeat this self-loop as many times as necessary until |ā 2 fmix- Because all self-loops are of length at most D + 1, this can "overshoot" by at most D, so we have that |a| S max((S) D, mmix D).
Execute the actions in ā on M once without collecting data, in order to ensure that within each set (s)), the newly-collecter observations are collected least fmix steps after obser- vations added in previous phases of STEEL.
Repeatedly take the actions ā on M, collecting the observation of each latent state s the first time in the cycle that it is visited and inserting the observation into D(s), until Vs e S' [D()) 2 d. Note that for a given latent state s, we collect observations of exactly [a] steps apart. Because |ä] 2 fmix this ensures that the observation added to D(s) are collected at least fmix steps apart. Because each D(s) will already contain at least one sample (from CycleFind), this process will take at most d iterations.
Note that if for some state 8 E S', [D()) reaches during some iteration of taking the actions then for the next iteration, we can re-plan shorter ā that does not necessar- ily visit s. However, when we do this, we must execute the newly-planned cycle ā once without collecting data, in order to ensure that all observation added to any particular D(s) are collected at least miix steps apart. This could require at most |S| additional iterations through some a.
This process will ensure that Vs E S [D(+)] 2 d, in at most
(d - 1 + [S)) max(D) + fmix. |8| D) steps.
(54)
Also, note that all samples collected during this phase are sorted into the appropriate dataset ((s) entirely by open-loop planning on TV, so the choice of samples in each D(s) remains independent of Q2, and, in principle, can be a deterministic function of Sint-
B.1.4 STEEL PHASE 3
Finally, for each learned latent state s E S', we train a classifier fs to distinguish Do == (.ES))()) D(s') from D1 == (ss). This set of classifiers allows us to perform one-versus-rest classification identify the latent state of any observation æ. by defining:
d'(x) == arg max fa(x).
(55)
Along with the learned transition dynamics T", this should be sufficient representation of the latent dynamics.
We want to guarantee that when the exogenous state of the Ex-BMDP is at equilibrium (that is, is sampled from its stationary distribution), for any latent state 8, x ~ Q(s, e), then the probability that fx) 1 and, V s' # f(s')= 0 is at least 1 - By union bound, we can do this by ensuring that the accuracy of each classifier fa, on each latent state s' E S. is at least 1 -e/(S), By realizability,, we know that Vs, there exists some classifier f- E F for which (s) = 1 6' (z)
s. Therefore, we need to upper- bound the probability that 3 f E F, for which Ve E D1. =
and Vo Da, (x)) 0, but for which either
Pr aow @(s,e)enns
(f(xz) = 0) 2
(56)
S
Or, for any s' s' #s. s.
Pr (((x)=1)2) T 2~ (=) joyons
(57)
For all 8, all samples in D(s) are collected at least mmix samples apart at timesteps chosen determin- istically and independently of Ω. Therefore, for any single fixed classifier f.we can use Lemma and the fact that Vs. (Ds)) 2 d to bound the false- calse-positive rates in Equations 56 and 57
as:
Pr Va E Ds), f'(x) =1
Pr O(s.)ee
f'(x) = 0 N
6
< 8e 3i8
(58)
26

---

Published as a conference paper at ICLR 2025
and, As' E ES'\ ( {s),
Pr Vo E(s)), -
Pr Oss
((()) 11 [8]
Se डोेर.
(59)
Taking the union bound bound over and all latent states E S' {8} gives a total false positive rate for learning f' as fs as:
FPR(J', ,8) S 8|S|e $े.
(60)
Taking the union bound over all f E gives:
FPR(S.) & F|e मंड.
(61)
Finally, taking the union bound over each classifier f. gives:
FPR & 8|S 131
(62)
B.1.5 BOUNDING THE OVERALL FAILURE RATE AND SAMPLE COMPLEXITY
Here, we bound the overall failure rate of the STEEL algorithm. We do this by separately bounding the failure rate of the first phase of the algorithm (the repeated applications of CycleFind) and the final phase of the algorithm, the learning of classifiers f.. We let each of these failure rates be at most 8/2. Therefore, must have, over the most S [A] iterations of CycleFind. failure rate of at most
2 |S| |A| Pr(CycleFind Fails).
(33)
2
This is satisfied by Proposition (noting that N |S]. The number samples needed for these [S] |A|] iterations of CycleFind, each with |â] < S (D + 1), is (by Equation 36 at most:
[S] ||/1- max ((2fmix + +N. 18) - 1) - 2) Isamp. eyc.
N 181. (D+(D++)
2 (fmix + + (D + 1) Asamp. +1)
(64)
+ -1) + 1) -miix, 0) +
Which is upper-bounded by:
|S| - ||A] max (22fmxx +3N (2N-18) (S) .(D 1) - -samp, eyc.
2 (fmix 1812 -D + 1) 1) -2samp.
(65)
+2 mN - 1) |S| (D 1) fmixx 01)
where Tesamp. is given by Equation 11 and sampp. eye. is given by Equation 18 Meanwhile, the overall failure rate of the second phase most
8
(66)
2
Solving for d Equation 66 gives:
3 (116))))
(67)
Which is indeed satisfied by Equation 52 given that the structure of the latent dynamics were cor- rectly learned using CycleFind in the first phase of the algorithm, that [8'| [8]-
27

---

Published as a conference paper at ICLR 2025
By Equation 54
then know that the number of samples needed for this phase is at most:
(
max + fmix,
- 1+11+00 11+1111
(68)
Combining the number of samples over both phases and simplifying gives us an overall upper-bound of the number of required samples of:
(| |n(16]
max(1 D + mixx |S] D).
+1s1)
[S| I-A max (2lmis + 3N. |8| (D + 1) - -2).
(In (4/A)- N (N -1) 1) (F11//)) /n(16/9) + 1),
(69)
2 (fmix 1812 (D + 1)--).
(In (4|A) No -D (1+1) 11. (F16) //n(16/1))))1// + 1)
+ + max((N 1) |S| (D + 1) Emix:
This gives us a big-O sample complexity of (using that â S N and |S| S N):
(1s N ...D-
D- |AA (log) AA| + log + log J + log (1))))+
[S]- |A| fmix (log |A| log N) + log | + log(1/6))
(70)
[8|2. D -11/e) (log([S)) log | log 1/
[S] fmix (1/e) (log( (1) + log | + log (1/ 1
Using the notation ('(((z)) = ((f(() (f we can write this as:
F + | log
N 8
2 [F
|
fiix
F 8
o (NDISPIAI log 8
+
log
+
log
(71)
Therefore, have shown that, with high probability, STEEL returns (up permutation) the correct latent dynamics for the EX-BMDP, and a high-accuracy latent-state encoder ø', within the sample- complexity bound stated in Theorem This completes the proof.
C USEFUL LEMMATA
Lemma 1. Consider an Ex-BMDP M (, Tiii)
starting an arbitrary latent endogenous state Syni E S. Let represent the sample space of the three sources of randomness in M: that is, Te, 2. and the initial exogenous latent state einit. Assume that all actions on M are taken deterministically and independenly of Q. Let f E for each E S let
{0, 1}
be a fixed arbitrary function,
P. Pr:PP((((>) 1/2 ~ Q(s, e),e e mg).
(72)
where Te is the stationary distribution Te. Consider trajectory sampled from this Ex-BMDP de- noted as Tmjj == 2/, 2', I'..., with endogenous latent states 8moj == sh, , (so that 8 Simi)), and exogenous states ennj e e lg.... Then, for any fixed 1, N. selected independently of ? where la - 2 miiax we have that:
Pa 1/4 Pr(((x()))) 1| = p. 4p.+1/1 1/4,
(73)
where s'Zt, denotes the observations in the trajectory Fmmj up to and including s'y,-
Note that this does necessarily hold h depend Ω
28

---

Published as a conference paper at ICLR 2025
Proof. From the definition of mixing time:
Ve E & Pr(e/ = -16, - -TellTV S 4
(74)
Then,
Ve E, Pr(f(x)= 1|z Q c'), TT)
Pr (f(z) 1/e/, = s = (75)
2m m4,)
Because o/2 depends on s's only through e/ we have:
Pr(f(x)= = 1/x 2( e'),e' mT))
Pr (f(x) 1| ) S
1
(76)
Then by Equation 72
1
Psi
Pr xo 1isiy
-
(77)
which directly implies Equation[] 73 Note that this is does not hold if t1 and to can depend on S. For example, if we define 12 = (min such that f(x() 0 and t1 tmix), then trivially Equation 73 may not apply.
Lemma 2. Consider irreducible, aperiodic Markov chuin To on states & with mixing time Limx and stationary distribution me, and an arbitrary function f {0, 13. Suppose Prewne (f(e) = 1) < 1- Consider fixed sequence of N timesteps t1, .., tN. where Vi, t4-44-1 2 tmis Now, for a trajectory Ep, ... Ety sampled from the Markov chain, starting at an arbitrary Eo. we have that:
Pr ((e)) 1) S 8c -5
(78)
111
Proof. Define è' as:
: Pr (f(e)= 0). owae
(79)
Note that we know that é' 6. Now, fix any i 2 mmix- Let M denote the linear operator on state distributions corresponding taking i steps of the Markov chain: that is, Mt gives the distribution after time steps. Also, let 11 be the linear operator defined as:
11T := ( T(e)de) Te.
(80)
EE
Define the linear operator A' as:
= M' 11.
(81)
By linearity and noting that both M and II are stochastic operators, we have that, for any x
f JEEE (A'm) =
(82)
Also, from the definition of mixing time, we have, for any function a::
1 | 1 |T||
(83)
(To see this, for any & consider the unit vector ह. Then, note that:
|| 2- - M'ellv S 2'
(84)
Then, for any n, we have:
/ J a(e)Aiddelli S / |m(e)]|A')]] <
| m 1 =
2 1111111)
(85)
29

---

Published as a conference paper at ICLR 2025
Because is stationary distribution of M also have that:
ne ( M) = We Te = 0.
(86)
Additionally, let I be the linear operator defined as:
Ix / T(c)f(e)ēd de
(87)
3EE
In other words, (e) :- f(e) T(e). One useful fact about this operator is that, for any function To such that Seee To(e)de have that
I (Tmo)(e)dee 5 mTlll
(88)
(To this, note that have:
[ Jeee (Tro) + /
JeeE (II O,
(89)
and also that, because I'm and (1-18 T'yTo are nonzero for disjoint e's:
|TTelll + I1(I r))))ll ||
(90)
Then, we also have:
V (Tro)(e)de - |
(91)
and
Ufe (T) = 1f Jeee ((T- S [I(T - molll.
(92)
Combining these equations and inequalities yields Equation 88)
Now, consider the operator Mr This operator, when applied to a probability distribution yields the (unnormalized) probability density that results from applying Te e' i times, where e' is sampled from conditioned on fe'') = 1. More precisely, it is the density of ~ (e()), scaled down by the probability that f(e') 1. other words, is given by:
(MTT))(e) p(c: ele (c(c') A f(e') 1Ad ~ T) Pr (f(c') 1).
(93)
e'
Now, consider any vector v. Note that v always can be uniquely decomposed follows:
= ame + bü where [ v(e)de 0 and || = 1 and b 2 0.
(94)
(Specifically,, we must set = Soeg v(e v(e)de and b := ||u ame || and = (v aT.)
Assume that v is such that a 2 0. Now, consider the equation:
d' = M
(95)
If we consider the above decomposition, we have:
nh D' M'T(ame + bv)
(96)
Note that this can be re-written as:
Tee + D' = [aT + bb) + A'T( + bv).
(97)
Note that by Equation 80 the image of II can be written in the form a'm., while, by Equation 82 the image of A' can be written as b'ū where b' and e are constrained as in Equation 97 Then, we have (using Equation 80
a' = ( (T(ane + be)) = / (Tre))(e)de +b / (TT))
(98)
33
EE
and:
b' = =IAAA(aT.. + bo)ll1 all a|A b|
(99)
30

---

Published as a conference paper at ICLR 2025
Now, from Equation 98 have:
d' =a (Tro) +b / (T))
CEE
EEE
aa
(TTe) +
||011
(by Equation 88
EE
(100)
+ b
2 aa1 E')
|
(by Equation 79)
2
b
<a(1 6) + 2
(by definition of ()
And, from Equation 99
we have:
SallA'TT x. + TTjl
<alAA(( -T- ([)-)) b/ A
< <a|AA Felll + al/ m. 11 b||rT
$ al/ - --) Il + 6|
(by Equation 86
2 (1 1) |1 + 2 ||r
(by Equation 83
(101)
a6'
|
(by Equation 79
2
ae'
(by Equation 90
2
ac'
(by definition of E.)
2
We can summarize these results as:
(102)
where there "<" sign applies elementwise. Also, because the elements of this matrix are all non- negative, have:
xc
S
(103)
2
y
Now, let Tt, represent the probability distribution of the Markov chain at timestep 41 and tume be the probability distribution timestep (y /mix- Applying Equation 93 repeatedly gives that:
M r M -x-
- - TTT
this -Pr(f(e,,) = Pr((ee))) = 1[S(ee)) = =)
(104)
Pr(s(ery .) = N- [(ee) =1) Pr(((((x)= = f(e,) 11)
This gives us that:
AN - (e)de Pr( = 1)
(105)
JeCE
where the right-hand side is the probability that we are ultimately trying to bound.
Now, let Ug == Te, for 1 < i < N - 1, let e == M - I....TT, and finally let uy == Ma Mf'TTMM T I I Tw, Let a; and represent the components a and in the decomposition given in Equation 94 of us. Note that:
ax=
(M TM T M IT,, (e)de= Pr(n =).
(106)
JEEE
Also, note that Vj E [1, N]. we have that By 102
M' for some 2 tmix- Therefore, by Equation
1
e'
0;-1
(107)
ba
b
2
31

---

Published as a conference paper at ICLR 2025
(Additionally, each a, represents probability, and so d; 2 0, so Equation 102 is applicable.) Now, due to the relation shown in Equation 103.
we can apply this inequality recursively:
a2
ai-2
b:
M
VA
b
b
VI
e
ai-1
(108)
1
e'
0:-1
e'
M
6 0
03
e
4i-2
W
b.
b:-2
So that we have:
aN bn
1 e'
0g
(109)
bu
The matrix
has eigenvalues
we can use the closed-form solution to the
2
Nth power of an arbitrary 2 x matrix given by Williams 1992) to exactly write this upper-bound on aN:
N
e) 1
(3-6')-1 e72++
ON S
2
2
V +
e)
e2
e (2--1)+
12
2
Ve
bn
+
Where we are also using that Tt, a normalized probability distribution, so ag = 1. Simplifying:
N
- - +
12
3e'
bo
ON S
+
2
62
1
4
2
2
2
e))
12 +
12
1
3E'
bo
2
02
+ 4
4
2
2
This gives us:
N
3
e)
+
12
+
1
3e'
e2
by
ON Σ
+
12
4
2
(110)
N
22
1
3c
62
bo
2
e2
+
2
Note that because Te, is a normalized probability distribution, we have that by is most (with this maximum achieved when nt, has disjoint support from TT.) Numerically one can see that, for E 1).
3
-7
55
(111)
4
2
2
2
1
3
-1.9
<
0
(112)
4
2
2
32

---

Published as a conference paper at ICLR 2025
Additionally, we can bound
1
2.
Then this gives us:
4e²44
(3
c')
12
22
2
ON
+4 4
(113)
2
(8-5)-- 22
(3-e)+
Because, for e' E [0,
and using Equation 106, we can
bound:
N Pr Pr( f(e,)= = 1 8
( - e') e 1
(114)
This bound is clearly somewhat unwieldy; to obtain more manageable bound, it is helpful to consider the asymptotic behavior near e' = 0. Letting ò := Pr(N - (le, 1), we have:
2 c)+ +
Ind S 1 Nh
(115)
2
For é' small, we have 2 < J. so 2 + 1 ~
Then:
E'
In S In 8 + Nin
1
(116)
Then, using the standard approximation In(1 x) 2 x gives us:
1 S 1n8 Ne' 2
(117)
This approximation would give us 8 Se é', it does hold exactly. Despite this, it does suggest form for our final bound. If we try Se A we find that holds that:
N²' However, while this holds approximately for small
c')
Se on the interval 0 s é S 0.44.
(118)
Combining with Equation
114 this implies that
Nr! on the interval 0 05e 5 0.44.
8 < 8e
(119)
For very large e', we can much simpler bound on Pr(
i= ((e,,) 1). Recall that:
N
N
N-1
= Pr
n ((r)
[(())=))
=1)
PP [(exx) =
f(er,) =1 =1)
1=1
i=1
N-2
Pr [(exx =
0 s(t) =1
(120)
i=1
= 1) Pr(f(en. 1)
However, because Vi, t; t--1 2 tmix, we have that:
i-1
Vi 1, Pr f(er) f(())=1 =1 n f(er,) =1 =1 j=1
s Pr (f(e) = 1)+ + envi
1 5
d
(121)
4 4
Combining these equations gives us:
5
N-1
8 <
E'
Pr(f(en) =1)
(122)
N-1
5
5
N
N
4
/
=
1
4
e
$
8
4
33

---

Published as a conference paper at ICLR 2025
Where used that é' Σ 1 in the last step. Finally, it holds that:
c S e on the interval 0.37 - d -
(123)
This then implies that:
5 8c Nv.
on the interval 0.37 S c' < 1.
(124)
Combining with Equation 119 we have that
N Pr
((ssee)))
88 N
Aé [0, 11,
(125)
Because è' 2 have:
N Pr
(f(e)) 1=1
S 88
4/ < 8E 4/
(126)
which was to be proven.
D UPPER-BOO COPER-BOUNDING MIXING TIMES FOR EXAMPLES
Here, we prove that the values of fmix used the simulation experiments are in fact (somewhat loose) upper bounds the true mixing times of T. for these environments. While in practice, the true mixing times would not known a priori, it important for the validity of our examples that the true mix is in fact < /iix-
We use the following following well-known fact:
For distributions A = Ai & A2 ② @ An and B : B, 2 B2 @ @ B,:
| B
Σ 14. B, LT
(127)
First, we deal with combination lock experiment. We can write the transition matrix for any arbitrary two-state Markov chain as
- 60
61
(128)
E0 1
1-
where 0 2 (eo, ) < 1. Note that in our particular example, have 0.1 S (co 0.9.
This matrix has eigenvalues and corresponding to the eigenvalue 1) is
co - 61, and the stationary distribution (the eigenvector
Th
[EE/(Eo + e1), co/(eo +c))²²
(129)
Using the foseed-form formula for the n'th power of two-state Markov Chain given by Williams have:
1992..
('~o c0
1)
]
1-
61 en
co
61
60
6 - -1
co
=1
co
ca 6
-c
(130)
To compute the mixing time, we compute the state distribution W.. n timesteps after each starting state:
[1-60
!
1
E1 + (1 -1-66 "E0
(131)
60
E1
+ 1
60 (1-60 "E0
and,
1 [-
))
00
1
1) [1 03 6 "
(132)
6g
6
1
+ C
09 +
63 6 "61
34

---

Published as a conference paper at ICLR 2025
Note that the TV distance between either of these distributions and the stationary distribution is most
IT - -ooolooll Toollrv _ (11 6o - 61) max(co,
< 1(11 -0 co - =1)l
(133)
(co + +1)
The parameters {Eo, é1) for each two-state Markov chain are chosen uniformly at random, such that 0.1 I (I I 0.9. Therefore, I [(1 co <1) - 0.8. Then, for any individual chain, have:
T. FTTV [(1 - 60 - e1) 0.8
(134)
In the combination lock experiments, there are up to L 512 of these noise Markov chains; the probability distribution over the exogenous noise & is the product distribution over these chains. Then Equation t27 bound the total TV distance for the chain Te to its stationary distribu-
tion Te; that is:
|| T
lotal
&ellTV S 512-
(135)
Then, by the definition of mixing time, we have:
1
mmxx S min n, such that 512 0.8"
(136)
Which gives us:
1og(2048)
mixx
35.
(137)
log(0.8)
So the value that we use in the experiment, fiix 40, a valid upper bound.
For the multi-maee experiment, the exogenous noise state consists of eight identical mazes, with agents moving uniformly at random in each of them. Unlike the "combination lock" example, where the individual components of the exogenous noise are conditioned parameters co, <1, which can vary, in the multi-maze example the individual mazes always represent instances of exactly the same, specific Markov chain. Let the transition matrix of this chain be M, with stationary distribution nM. Then, by Equation 127 for the whole exogenous state Te, we have:
TEll 1
Imi = min n, such that Vsp, [I(Ta))
4
min n, such that vs 80 3...80
(1) (2) (8) M 80
(1)
(2) M sp 2.. M ( (8)
TM TM 8 2 TM S (138)
1
< minn, such that Vsp. -M
TMll
S
= =min n, such that Vso, ||M - TM S
32
Which is to say that tmir for the entire exogenous noise chain is upper-bounded by tmiz (1/32) for the individual maze chain M. Furthermore, while the state space for the entire chain is of size [8] = 688 the individual maze chain M operates on state of size 68. This is small enough that it is tractable exactly compute fmii (1/32) for M using numerical techniques. We performed the computation (source code is provided with the supplementary materials), and found that
1
min such that Vsg. [|M Tall S 32
293.
(139)
Therefore, for the full exogenous noise chain Te, we have that tyii S 293. Then the value that we in the experiment, fmis 300, is valid upper bound.
35

---

Published as a conference paper at ICLR 2025
E DISCUSSION OF ASSUMPTIONS
In this section, we argue that two major assumptions of the STEEL algorithm are necessary sumptions. In other words, if we remove these assumptions, then, we argue, achieving similar guarantees to our algorithm would become effectively impossible for any algorithm accomplish in the single-trajectory, no-resets EX-BMDP setting (at least without making additional assumptions to compensate). These assumptions are (1) reachability: the assumption that all endogenous states s E S are reachable from one another in T; and (2) the availability of a known upper bound on the exogenous state mixing time mix.
E.1 REACHABILITY
Here, we argue that any algorithm for single-trajectory. no-resets representation learning of Ex- BMDPs must necessarily make reachability assumption: that all endogenous latent states must be reachable from one another. At a high level, we argue that, if all latent states are not reachable from one another, then an algorithm must visit states in a particular order in order to explore the entire dynamics with a single trajectory. However, because the latent dynamics are not known priori, it is impossible for an algorithm to guarantee that it visits states in the appropriate order.
More formally, recall that the reachability assumption is that for any pair of latent states $1, $2, there exists both a path from 81 to 82 and a path from $2 to $1. Consider any EX-BMDP M where this condition does not hold, such that |A|] 2 2, and any learning algorithm A. Firstly, if there is any pair of states $1 and $2 where neither can reach the other, then clearly a single trajectory is insufficient to learn the EX-BMDP, because it cannot visit both 81 and 82. Therefore, we restrict the case where, for each pair of states, either both are reachable from each other, or (without loss of generality) 82 reachable from $1 but $1 is not reachable from $2.
Consider every edge (s, s') of the state transition graph defined by the state transition function T. If, for all such edges, s is reachable from s', then the reachability assumption holds on the entire dynamics (because all single "steps" are invertible by some sequence of actions, so for any pair $1, $2, the path from 81 to $2 implies the existence of a path from 82 to $1). Therefore, reachability does not hold, then there exists some edge (s, a, s') such that s is not reachable from s'. Now, consider the first time that the algorithm encounters the state s. Regardless of the details of A, there must exist some action a' such that, on this first encounter, the probability of A taking action a' is at least A/.
Therefore, if a' a, then probability that the algorithm A never revisits s, and so never takes any other action from apart from a', is at least 1/] Then, with substantial probability, the algorithm A never explores the IA 1 other possible transitions from s, and cannot possibly learn the full dynamics of the Ex-BMDP (To be more precise, the algorithm's output will not depend at all on the ground-truth value of T(s, a") for a" E AA\ (a). and so is highly unlikely to return the correct values for these transitions on arbitrary EX-BMDPs.)
Alternatively, if a' a, then consider the alternative EX-BMDP M', which is identical to M in every way (in terms of dynamics, exogenous state, emission function, etc.), except that the effects of actions a and a' on the latent state are swapped. Note that before first encountering the latent state s, the MDPs M and M' will produce identically-distriutee sequences of observations, so the algorithm A will behave identically on them, and have identical internal memory/state. Then, when first encountering s, the algorithm M on A will take action a' with substantial probability. and then transition to s' and be unable to revisit 8.
Therefore, for any EX-BMDP M that violates reachability, any algorithm A is either likely to fail on M, or to fail on moddhul version M. In any case, no such algorithm will be able to succeed with high probability on any general class of Ex-BMDPs that does not require reachability.
E.2 KNOWN UPPER BOUND ON THE MIXING TIME Mix
Here, we argue that the assumption that an upper bound on the mixing time is provided is also necessary. We argue that:
36

---

Published as a conference paper at ICLR 2025
1. Any general, provably sample-efficient algorithm for learning, with high probability, an Ex-BMDD from a single trajectory must necessarily require an upper bound on the mixing time of the exogenous noise to be provided priori (unless some other information about the exogenous noise is provided.)
2. The runtime of any algorithm for learning an Ex-BMDP must, in the worst case, include some term that scales least linearly with the mixing time of the exogenous noise.
At high level, argue that, with single trajectory, it can be impossible distinguish between a static exogenous dynamics (i.e., where only one exogenous state exists) and an extremely slow- mixing exogenous dynamics, in time sublinear in the mixing time. For any proposed algorithm that does not have access to a bound on the mixing time, and any Ex-BMDP with a single exogenous state (that is, a Block MDP), we can construct an extremely slow-mixing Ex-BMDP that behaves identically to the Block MDP with substantial probability during the entire duration of the runtime of the learning algorithm, but which at equílibrium has a quite different distribution of observations. This will make the learned encoder fail on the Ex- EX-BMDP at equilibrium.
More precisely, to show (I): suppose the converse is true: that there exists some algorithm that learns EX-BMDP latent dynamics and state encoders, which is not provided any upper-bound on the mixing time of the exogenous noise of the Ex-BMDP, or any other information about the exogenous noise distribution. We assume that A can learn the correct endogenous dynamics, encoder with accuracy (for every endogenous state, under the stationary exogenous distribution) of least 1 e, with probability at least 1 - 6, for small values of and 5. Then, consider any EX-BMDP Mi and related parameterized family of Ex-BMDPs M2(1), Ma(7) with the following properties:
well as an
M1 has N endogenous latent states $1,-..; SN with some transition function T, single exogenous latent state 61, and an emission function Q(s, 61). (In other words, M1 is any Block MDP).
M2(7) has the same endogenous states and transition probabilities M1, but has two exogenous latent states e1. eg. The state éj transitions to é2, and e2 transitions to 61, each with probability y. Note that the stationary distribution of the exogenous state is uniform over é1 and 62. We also assume that the initial exogenous state distribution is uniform over ej and 62. Regardless of y, the emission function of M2(7) is defined such that Vs,, QM, (si, = Qm, (s1, e1).
Ma(?) is identical to M2(7), except that its emission function is defined as vs, Qm, (8,,e1) == Qas, (st, e1), but Vs; Qm, (8g, (2) == 2M, (8(1+1)% 62). In other words, when the exogenous state is equal to e2, the emission distributions for the endoge- nous states are permuted in M3 compared to M2.
We also assume that the encoder hypothesis class can represent the inverses of QM,, Qa... and QMy However, note that by construction, any fixed encoder o(-) which has accuracy at least 1 - on M2(7) (for every endogenous state, under the stationary exogenous distribution) can only have accuracy most 0.5 + Ma(7), because, when the exogenous state is 62, any time that the encoder returns the correct latent state for M2(7). it will return the incorrect latent state for M2(7)
Now, consider what happens when we run the algorithm A on Mi. The number of environment steps that A takes on M, forms some distribution; let t be the 90' 'h percentile of this distribution, so that with probability 0.9, A stops sampling before step f.
Now, we can set 7 M3(7) do not change before step t. Therefore, by union bound, on M2(7) or M3(7), if the initial exogenous state is el, then with probability least 0.8, the exogenous state will stay constant e
0.9 that, with probability 0.9, the endogenous state of M2(7) or
*There is some nuance involved here: note that for an encoder ø to achieve high accuracy, it only needs approximate (dd under some choice of permutation However, we know that for any 8t, 8144 S. the á which is highly accurate on Mz(7) must map all observations QM, ($((+1)% e2) to two distinct codes, each with marginal probability at least - 2c. Therefore, this e will map at least 0.5 of the observations sampled from QMa (si,e) for e ~ we to one latent state, and at least another 0.5 of the observations sampled from this distribution to a different latent state. Then under any choice of perturbation o, the accuracy of (6-) on M3 can be at most 0.5 + c.
the support of QMz (8., e1) and
37

---

Published as a conference paper at ICLR 2025
up to timestep and the algorithm A, seeing exactly the same distribution of observations if the MDP was M1, will halt before timestep t.
Considering the 50% probability that the exogenous state starts at e1, this gives least 40% chance that A applied to M2(7) or MM((7) never encounter the exogenous state 62. this case, the distribution of encoders output by A should be the same for the two Ex- 3x-BMDPs: let this distribution be G. (That is, 9 is the conditional distribution of encoders output by when applied to M2(+) or Ma(r), given that A never encounters e2.) However, because A fails on M2(7) with probability at most 6, we can conclude that at least (1 - 16) of the encoders from the distribution G represent successes of A on M2(7).
Because A on Ma(n) also draws from 9 with probability at least 0.4, we can conclude that at least 0.4- of the time, on M2( produces an encoder that is highly accurate on M2(r). As discussed above, any encoder that is highly-accurate on M2(r) cannot be highly accurate on M2(7), so A must fail on M3 with substantial probability.
Therefore, we can conclude that for any algorithm A that has no "hint" about the exogenous dynam- ies (such as knowing the mixing time), there exists some ry such that A cannot possibly succeed with high probability on both M2(7) and M2(7).
To show point (2), simply note that in this construction, in order to ensure that e2 observed 7)]
with probability (1 - 8), an algorithm must observe at least the first [log(26)/log)) timesteps. Mcanwhile, the mixing time of the two-state Markov chain exogenous state is given by (1 2y) For fixed d and y approaches O, the ratio between these quantities approaches a constant: therefore, the number of steps required to ensure that 62 is observed with high probability is linear in the mixing time.
F DOUBLE-PRIME LOOP EXPERIMENTS
In this section, we present additional experiments which demonstrate that, on some types on Ex- BMDPs, the STEEL algorithm can empirically outperform the algorithms proposed in Lamb et al. (023) and Levine et al. (2024), which do not have sample-complexity guarantees. We show that there certain structures of Ex- EX-BMDP latent dynamics which are difficult for these prior methods to learn efficiently, but on which STEEL performs well. Specifically, we look at family of "double- prime loop" tabular Ex Ex-BMDPs which are discussed by Levine al. (2024) First, though, we describe the algorithms proposed by Lamb et al.] (2023) and Levine et al. (2024), and motivate why certain dynamics structures such Mhouble-prime loops" present challenge these methods.
F.1 BACKGROUND
F.1.1 AC-STATE AND ACDF ALGORIITMS
Lamb al. 2023 and Levine et al. (2024) both propose algorithms to learn endogenous state encoders in the BMMD framework.
Lamb et al. (023) first proposed the AC-State algorithm, which aims to learn encoder X' N. such that, under some one-to-one permutation a, a()(x)) = o" (xx). To accomplish this task, AC-State optimizes using multistep inverse dynamics objective.
Specifically, in the theoretical treatment, AC-State tries to find the encoder & which optimizes the following objective:
((Ø) = arg min [Range()) ((o)=min, ((dr)
min JENxNXN-PLA kn<I.....K) (s.a)))))>
E E log (g(0(xz), ((E4)) [a.])
(140)
@AC--Stte
In other words, the loss L(Ø) is minimized on tuples (æz, at, D(++) sampled from a trajectory, con- sisting of observation F,, the following action a,, and the observation k steps in the future, Te+R-
The encoder (xx) is trained to retain any information about the observations o(x,) and d(x++k) that useful for predicting o, Specifically, an inverse-dynamics model g is simultaneously trained with
38

---

Published as a conference paper at ICLR 2025
predict d, given o(xz), <(Fe+k), and k. The loss [()) is then taken as the minimum loss for over all such functions
However, note that, based on this loss function alone, the identity encoder 6(x) := achieves the Therefore, in order to filter exogenous information, the final encoder returned by AC-State is the minimal range encoder that minimizes the loss function.
minimum possible value of [( [
Lamb et al. (2023) claims that, if the data D is collected by policy T(z) that only depends on x through (2) (that is, the policy ignores exogenous noise); and the maximum segment length K is at least the endogenous dynamics diameter D; and the endogenous dynamics are deterministic, then the final encoder AtteState returned will correctly encode the endogenous latent state (in the limit of infinite data; and with sufficient function approximation).
When applied in practice, & and g are learned neural networks, and a vector quantization bottleneck is used to restrict the range of & (so the output of é is quantized vector, rather than an integer.)
Levine al. (2024) subsequently show that, in some cases, the AC-State objective in Equation 140 is not sufficient to learn correct latent state encoder of an Ex BM-M even with unbounded amounts of data. Levine et al. (2024) identify specific flaws in the proofs of Lamb al. (2023). and propose a modified loss function to address these flaws. In particular,
The maximum segment length K required to correctly learn the endogenous encoder must in some cases be larger than D; a corrected upper bound on the necessary segment length 2D2 D is given.
2. If the endogenous dynamics are periodic, then the multistep inverse objective be in- sufficient, for any choice of K. Levine et al. (2024) then propose to add an auxiliary loss function to the loss in Equation [40 namely latent forward dynamics loss:
min Lforward bENx.A-P(N) (21,06
E
log (h(((z)), ac) [6(2411]). (141)
This loss term enforces that the transitions in the learned endogenous latent state space deterministic; Levine et al. (2024) prove that enforcing this constraint is sufficient, when combined with the multistep inverse loss, ensure that correct endogenous representa- tions are learned even when the endogenous latent dynamics are periodic. Levine et al.) (2024) name the AC-State algorithm with this modified loss function ACDF.
Note that because D is in general unknown priori, K must be treated as hyperparameter of the algorithm, so point (1) above is primarily of theoretical interest, but may aid in setting this hyperparameter if there is some prior knowledge of D.
So far, we have only described the learning objectives of the two proposed algorithms: now, we address how these algorithms collect the trajectory from which the tuples (Iz, az, .e+kk are sampled. Recall that a condition on the correctness of AC-State (and ACDF) is that the data-collection policy does not depend on the noise in the observation, but only on the endogenous latent state. This condition naturally raises the question how exactly an algorithm intended to discover o can take actions that depend only on o" (z)), without already knowing o² in the first place.
Lamb et al. (2023) performs experiments using two data-collection policies: (I), a uniformly ran- dom policy, which meets the condition simply by ignoring the observation x entirely; and (2) goal-seeking policy, which aims to maximize latent state coverage. For the latter policy, the data collected simultaneously with training and the partially-traine trained encoder òi is used (imper- fectly) filter out exogenous noise and plan in (an approximation of) the endogenous latent space. Lamb al. (2023) explicitly acknowledge that this bootstrapping approach breaky the condition that the action a, depends only on the true endogenous state o' d2(().. which is necessary in their proof of the correctness of their algorithm. Despite this, they observe that the approach seems to work well empirically. Meanwhile, Levine et al. 2024) only use uniformly random actions in their experiments.
For now, we will assume that data is collected under a sefffommly-randow policy. We will return to discussing the antent-state-con coverae maximizing approach proposed by Lamb et al. 2023 later, in Section
13
Assuming X is countable.
39

---

Published as a conference paper at ICLR 2025
0/1
0/1
0/1
0/1
4
0/1
0/1
0/1
0/1
0/1
0/1
0/1
0/1
0/1
0/1
0
0/1
0/1
6
0/1
5
0/1
0/1
12'
11
0/1
0/1
0/
1
Figure Latent dynamics of DoublePrime( 13).
F.1.2 DOUBLE-PRIME LOOPS
To specifically demonstrate that K O(D²) can be necessary learn correct encoder, even when a forward dynamics loss is also being used, Levine et al. (2024) give a concrete example. This example the family of tabular "double-prime loop" EX-BMDPs. For any two primes p, 9. with < let the Ex-BM BMD DoublePrime(p, be defined as follows:
S (0, 10 (q 1),'',,,,,,,,, (g 1y 10, 1 -= {co)
X = S; co) (eo) cu-
. The endogenous latent state transition function T is defined as:
(s + 1)%g
if (1, (1.... 1] or (s 0 and a = 0)
((s 1)%
ifs (1',.., (q 1y) (s O' and 0) (s = 0 and a = 1)
as,,a)=
(142)
(q -1-p+11
(1 - p p+11
(a 0' and a = 1).
In other words. DoublcPrime(p.. q) is noise-free Ex- Ex-BMDP, where all of the dynamics are de- terministic, and the latent state s directly observed as x. The dynamics consist of two connected loops, each of states. The actions are generally ignored and the agent continues to progress through a loop, except for in states 0 and 0', where taking the action transports the agent the other loop, at position p + 1. As an example, the dynamics of (1ublePrime(11, 13) are shown in Figure
There is also related family "ssingle-prime" loop" Ex-BMDPs. Let SinglePrime(p,, g) defined as (borrowing notation from Levine et al. (2024)):
S = {00, 1* 11,,....6 (g A= 1 (co): Te(ea) en-
X {0, 1, ..., (q 1),0', (4-1))
with probability 1/2
Q(s". _o)
with probability 1/2.
The endogenous latent state transition function 7' is defined as:
T(s*,a) a) ((s + 1) if {1" ..... (g 1)*} (s 0* and 0) (s = and 1).
(143)
(1-
In other words, SinglePrime(p., q) has the same observation space X' and action space A as DoublePrime(ps q). but fewer controllable latent states. In SinglcPrime(p.c) q), when the agent is at latent state 8, = i", either the observation I, = i or the observation z, = is emitted, each with probability 0.5. The agent progreses through the loop of latent states regardless of actions, except
40

---

Published as a conference paper at ICLR 2025
0/1
11
12"
0/1
10
9
0/1
1
1
0/1
2
0/1
0/1
0/1
4"
0/1
0/1
5
Figure 6: Latent dynamics of SinglcPrime(11, 13).
in latent state 0", where taking action transitions the latent state to (g - + 1)". (See Figure for an example.) Note that if we ignore the distinction between the observations and i', then DoublePrime(PP (p, 9) and SinglePrime(p, 9) appear to have identical observed dynamics.
Let (DP(x) be the optimal encoder for DoublePrime(pp. q). and
Øsr (x) == i" if = or x
(144)
be the optimal encoder for SinglePrime 4).
Levine al. (024) show that, on the Ex-BMDP DoublePrime(p, q) under uniformly random behavioral policy, in the limit of infinite training data, the encoders Φpp and ésp have exactly the same loss under the loss function Equation 140 with < (q - 1)p. Note that p and are close, then (q 1)p s D²/4. Because Equation 140 prefers the smallest-range encoder among encoders with identical losses, this means that AC-State will incorrectly return éSp (which only has q distinct outputs, rather than 2q.). Futhermore, the forward-dynamies loss suggested by Levine et al. (2024) does not help in this case: it will be zero for both encoders under data generated by DoublePrime(2. 9)
is important to note that SinglePrime(f, q) and DoublePrime(p, q) have truly distinct controllable latent dynamics: in DoublePrime q). given sufficient lead-time, the agent can (eventually) control whether it is in state 1' or 1 at a future time step, while an agent in SinglePrime(p. q) can never control this. Therefore, AC-State is making an error if returns ΦSP rather than ΦDP-
Levine et al. (2224) only considers the case of unlimited data, and only introduces the double-prime loop Ex BMDss to make the point that the hyperparameter K must be set very high in some cases. Here, however. we will also show empirically that, even when using carefully-tune and very large value of K, AC-State and ACDF can take many samples to correctly learn these dynamics.
F.1.3 TABULAR EXPERIMENTS IN LEVINE ET AL.. 2024
their proposed ACDF Levine et al. 2024 empirically compare AC-State Lamb et al. 2023
method, and empirically explore the effect of the hyperparameter K. To do so from a purely statistical perspective (i.e., controlling for differences in optimization), in one set of experiments, Levine et al. 2024) perform tests on small tabular Ex-BMDPs, including SinglePrime 5) and DoublePrime (3, For these experiments, the final learned d'c-State (or (ACDF) found by haustively computing the losses in Equations [40] and [41 on all possible encoders á E X' + (0. | 1} (up relabeling perturbation of the output). For each encoder o, the multistep inverse model (and, for ACDF, the forward model h) is estimated as a tabular function based on
10The experiment in Levine al. (2024) on SinglePrime(3,, 5) is slightly different variant of the MDP from what is presented here: the controllable latent dynamics are the same, but the noise time-correlatdd.
41

---

Published as a conference paper at ICLR 2025
the collected data:
g(s, (k),s',)([aa := # of instances of (#te (2t, &)) such that ((x)) and b((++))= # of instances of (Iz ((,ee) such that ((=))= and ((I+++))=
(145)
However, due to overfitting, this tabular definition of g. if used naively, would always lead to the identity encoder ((x) :- x having the lowest empirical loss, regardless of the true latent encoder. al. (2024) collect two trajectories for each experiment, and uses one to fil y and for each possible encoder Φ and the other to evaluate the loss on each Φ in order to determine the final output encoder OAC-Stae (or dACDF). We will call these sets of tuples the fitting set and optimization set, repectively.
Therefore,
Levine Levine
Because the number of possible encoders grows very quickly with X Levine al. 2024 limit these experiments to very small tabular Ex-BMDPs, with | S 10.
With this background out of the way, we deseribe our experiments:
F.2 EXPERIMENTS
F.2.1 SETUP
Given that Levine et al. (2024) describe how learning the correct latent state encoder for DoublePrime(p,. 4) using a multistep-inverse loss requires taking into account specific long-duration dependencies between states, we hypothesized that learning such encoder using multistep-inverse methods may be considerably less sample-efficient that doing so using STEEL, particularly for large p and q. We therefore adapted the tabular Ex-BMDP experimental setup from the released code of Levine et al. 2024) and tested the algorithms
However, for large and 9, the process of cxhaustively searching all possible encoders (that is, all partitions of 2q elements) quickly becomes intractable. Therefore, we restricted the hypothesis space of encoders to only include the two hypotheses Øsr and ΦDr. Note that this modification makes the learning task strictly "easier." For a fair comparison, we also restricted the hypothesis class for STEEL to the minimum-possible set of "onneverrus-rest" classifiers that would ensure realizability for both DoublePrime (p. q) and SinglePrime(p, 4). In particular, this is 3q classifiers / E F: for 4). includes a hypothesis which distinguishes from all other observations, a
each i E (O..... hypothesis which distinguishes i' from all other observations, and a hypothesis which distinguishes i or i' from all other observations.
We make the following further modifications to the experimental protocol from Levine et al.
2024
When assessing the "minimum-raagge" minimal loss encoder as in Equation 140 Levine ctt al.) (2024) include an empirical "fudge factor": their protocol returns the minimum- range encoder that achieves loss within 0.1% of the true minimum loss over all possible encoders. We found that even without this fudge factor, the multistep-invedd were still heavily biased towards returning ΦeR. when applied either to SinglePrime 9); DoublePrime(p, q) with too-small K or too-few samples. This observation has simple statistical explanation: if, in the "infinite sample" limit, Príaz 8t = i, 8r+1 = j) Pr(a,]st (x 1' 81+1 j'), then fitting g(i, j, k) both samples of (Ir, i, ar, We+ j) and samples of (z, = i', ae, Zt+k = j()) from the "fitting set" will lead to higher-quality multistep inverse model g. and therefore lower loss on the "optimization set", compared to using only the (smaller number of) samples of (xz = i, at, Ti++ = j) alone. We therefore remove the fudge factor entirely, and in fact return ΦGR the case of exact numerical ties.
methods
Because the forward-model loss in Equation 141 is zero for both 6sn and dan for any dataset collected from DoublePrime(p.. 9). we do not use it in our experiments: that is, we regard AC-State and ACDF equivalent in this setting, and refer to AC-State alone from here onward. (Technically, the forward model loss would help identify the correct encoder ΦsR in data generated from SinglePrime q). However, mentioned above, ÓSR is already returned essentially "by default" by AC-State, so this loss term turns out to be unnecessary.)
To better match the spirit of "learning an Ex-BMDD from single trajectory", we collect the "fitting" and "optimization" sets from one single trajectory, one after the other, with each corresponding to half of the trajectory: we regard the total sample complexfity as the
42

---

Published as a conference paper at ICLR 2025
length of the entire trajectory. Within of these two sub- trajeerreecooisss we collect all available tuples (zt, at 2244) for all & < K. There are therefore slightly fewer tuples with & = K than with & = 1: we weight all ruples equally in the loss (rather than weighting all values of l equally: Lamb [qual et al. 2023) is unclear about the "correct" behavior here, and is not theoretically important).
Based on Equation 145 p(8,,s,,)ffe can be zero, which, by Equation 140 can lead to an infinite loss on a sample in the "optimization" set. In order to avoid this, Levine et al.) (2024) set g(s, [a] to an arbitrary floor value of 10-? As a more principled and scale- invariant solution, we use standard Laplace ("add-one") smoothing when fitting
We empirically compare the sample-efficiency of STEEL to AC-State on these problems with (p.q) = (3, 5), (5, 7). (11, 13), (17, 19), (29, 31), and (41,
In our comparison, we run each algorithm under the "best" choice of hyperparameters. For AC- State, the conduct large hyperparameter search to find the optimal K, while for STEEL, simply choose a single set of hyperparameters that will succeed on both DoublePrime(p, q) and SinglePrime(p::) q) based on our prior knowledge of the problems. Specifically:
For AC-State, for each tested value of (p. q). we first collect 20 validation trajectories of DoublePrime((2 q) (under a uniformly random policy), for increasing dataset sizes starting at T 100 steps. At each tested dataset size T, compute the optimal AAC- for each of the 20 trajectories, with each possible value of K from K = 1 to K = 50, 000. We first repeatedly increase the dataset size T by factors of 10. Once we first identify T = 10 such that, for at least some K, AC-State correctly returns d'C State = ΦDP for all 20 trajectories, we then test with T {2- 10 3- 10m--1..... ..... 9-10m- }. At the earliest of these timesteps at which, for some K, a correct encoder is learned for all 20 trajectories, we record the median such K as the "tuned" value Kap of this hyperparameter. (See Figure for the results of the hyperparameter search.) At this point, we perform the actual test: we evaluate the success rate of AC-State on DoublePrime(p, q), using only K Kap- for values of T starting at T 10T-1 and increasing by increments of 10T=11 on 20 new trajectories for each tested dataset size. We stop when this test-time accuracy reaches 20/20 on trajectories from DoublePrime(pp q). at some T Finally, we verify that AC-State can also consistently, correctly learn øsp on data from SinglePrime(p, 4) with K Kapi and a trajectory length of Trax, for 20 trajectories. (This always held in practice IT
Đ 2q (to match the maximum number of states in DoublePrime(y..g) q) or SinglePrime(p, q)): fiii 1 (because neither DoublePrime(p,. nor SinglePrime(p... q) have time-correlated noise); e = 0.49 (because, due the simple emission distributions Q of SinglePrime and DoublcPrime, 51% encoder accuracy on each latent state implics a 100% encoder accuracy); and 8 = 0.05. We test for 20 trials each of SinglePrime(p, 9) and DoublePrime(p, q). All of these tests returned the correct encoders, so the only metric we needed to consider was the number of environment steps taken for each run.
For STEEL, we set
F.2.2 RESULTS
Our top-line results are reported in Table first reported statistic the "Max Steps": the number of environment steps at which all 20 out of 20 tested trajectories for each of DoublePrime(p, 9) and SinglePrime( q) returned the correct encoder. (For STEEL, this is simply the maximum steps taken over all 40 trajectories; for AC-State, this is the Tiux discovered in the test-time search described in the previous section.) We see that, while for small instances (IS) < 14) of DoublePrime(2 q). AC-State is more sample-eflicient than STEEL, STEEL scales efficiently larger instances of the problem.
"We avoided tends to "default" to returning the encoder dsp, which correct for SinglePrime (p. g): and also because the deterministic dynamics of DoublePrime(p, g) allowed us to take some computational "shortcuts" when evaluating the loss function in Equation o4 on data from DoublePrime(p, g) for very large K, which were not possible for SinglePrime(p. 9). This made more exhaustive experimentation on DoublePrime(p, q) significantly more tractable than it would be SinglePrime(p.9) 9).
extensive testing with SinglePrime(p, 9). both because, as discussed above, AC-State
43

---

Published as a conference paper at ICLR 2025
Hyperparameter Search, p 3,q
Hyperparameter Search, p =
1.0
Success Rate DwublePrimeep....om ()
1.0
40
0.8
120
0.8
DrublePrimee
80
0.6
240
0.6
K
300
120
360
0.4
140
420
160
Success
180
540
1.
2-
3.
5
7
8
9 10-
1.
3
4. 5-
6.
7.
9. 10-
101
10
10
10
10
10
10
10'
10
10*
10* 10* 10
10*
10*
10a
10
10t
10*
Trajectory Length T
Trajectory Length
Hyperparameter Search, 11.
13
Hyperparameter Search, 17,
19
200
Success Rate DoublePrimetp.
800
DoublePrime(pp..
400
0.8
1600
0.8
600
2400
800
0,6
3200
0.6
1200
-0.4
4800
0.4
5600
1600
0.2
6400
0.2
7200
0.0
0.0
10
10
10f
10f
Trajectory Length
Trajectory Length
Hyperparameter Search, q
Hyperparameter Search, p q
1.0
Success on DoublePrimelpo ()
Success Rate on DoublePrimellpp.. pp.4)
3000
5000
6000
10000
9000
15000
15000
25000
30000
21000
35000
24000
40000
27000
45000
0.0
1. 2.
6-
7-
9-
10-
1.
3. 4.
5.
6.
9.
10-
10"
102
10"
107
10
10
10
10a
10a
10a
10a
10 10a
Trajectory Length
Trajectory Length
Figure 7: Results from the hyperparameter search for the maximum-step-coum parameter K in Equation 140 The objective is to find which leads to a high success rate in learning the correct latent state encoder for DoublePrime(p.P0 9). for the lowest possible number of environment steps.
44

---

Published as a conference paper at ICLR 2025
(p.g)
(3,5)
(5,7)
(11,13)
(17,19)
(29,31)
(41,43)
Max Steps for STEEL
5.97 10
1.49
10
9.07 10
2.60 10
1.10 10"
2.83 10
Max Steps for SC-State
4-10
4 4-10a
1 1-10
8-10
1-10
5 5-10
Max Steps: (AC-State / STEEL)
R3 0.07
R8 0.3
R/
Rd 3
9
R 18
Median Steps (STEEL, PoublePrime)
2.50 10"
7.48 10"
5.23 10
1.08 10
1.74 10
3.29 10
Median Steps (AC-State, DoublePrime)
3 3-10
3 10
8- 10
7- 7-10
9 10"
5 1-10
Median Steps: (AC-State STEEL)
2 0.1
Rd 0.4
R3 1.5
R 6
R 50
R 150
Table 3: Comparison of the empirical sample complexity of STEEL and AC-State on prime-loop MDPs. See text of Section F22.
Further, we noticed an interesting trend in the data. AC-State tended to transition very abruptly from learning the wrong encoder on all tests of DoublePrime to learning the correct encoder on all tests as T increased (as can be observed in the hyperparameter search in Figure 77) By contrast, the distribution of steps taken for STEEL on DoublePrime((p, q) was highly skewed: most trials took significantly fewer than the maximum observed number of steps. We therefore also compare the median number of steps taken to learn a correct encoder for DoublePrime((2 q). For STEEL, this is computed simply the median length of the 20 DoublePrime( (p, 9) trajectories tested. For AC- State, it is the first timestep T in the test-time search where AC-State with K = Kopt learns a correct encoder for at least 10 of the 20 trajectories. Here, we find that, especially for large (p, STEEL's advantage over AC-State becomes even more pronounced.
F.3 EXPLORATION POLICIES WITH AC-STATE
So far, we have only compared STEEL to AC-State, where AC-State uses uniformly random exploration policy. This may seem unfair: STEEL takes decidedly nonrandom actions, and Lamb al. (2023) proposes an active exploration method for AC-State (albeit, as mentioned in Section F.1. above, method without strict theoretical grounding)
However, there is reason to believe that exploration policies similar to the one proposed by Lamb et al. (2023) would be unlikely to help correctly identify the dynamics of DoublePrime(p.t) q)-
To start with, the stated goal of the the exploration policy in Lamb et al. (2023) is to "achieve high coverage of the control-endogenous state space." A natural question to ask is: how poorly do uniformly random policies perform at achieving this goal? To quantify this, we examined the state coverage for single step trajectory DoublePrime(44 41, 43), and computed the ratio of the state visitation of the most-visited state to the state visitation of the least-visited state. Surprisingly, this was 2. In other words, DoublePrime is not hard exploration problem, and so techniques designed to improve state coverage would seem to be unlikely to provide much of a benefit over a uniformly random policy.
Still, we will now examine the specific exploration technique proposed by Lamb et al. 2023). At a high level, the technique proceeds follows:
Throughout data collection, the agent maintains an approximate version of the encoder o' obtained by optimizing the AC-State loss (Equation 140), and approximate transi- tion function T", obtained through counting. The algorithm also maintains state visitation counts for each of the learned latent states encoded by d'
At the start of each round of exploration (time f), the agent selects goal learned- latent state Sy. with probability inversely proportional to the visitation count. The agent then takes single random action de, and then, starting at time ( + 1, plans shortest-path to 8g. and computes the number of steps k' it will take to get to 8g. Then, the agent proceeds to navigate to sy in a closed-loop manner, using both o' and T". After k' steps, regardless of whether sy is reached, a new goal is set, and the process starts over.
The encoder ø is only trained the tuples (It, Dt, Te+4(41) which begin with a random action and end with a goal state.
We performed experiments of version of this algorithm adapted for our tabular setting, the DoublePrime environments. Here, we explain the modifications we made to this exploration al- gorithm from Lamb al. (2023) for our tests. Examining the original algorithm, immediately
45

---

Published as a conference paper at ICLR 2025
(p.q)
(3.5)
(5.7)
111,13
17.
(29.31)
Max Steps (AC-State, DoublePrime, Uniform Policy)
4-10
4. 10"
1- 10
8 10
1- 10
Max Steps (AC-State, DoublePrime, Exploration)
8 8-10
6. 10
2 10
2 10
10
Median Steps State, DoublePrime, Uniform Policy)
3 10
3. 10
8- 10
7- 10
9- 10
Median Steps (AC-State, DoublePrime, Exploration)
6 10
10
2 10
2- 10
2- 10
Table 4: Comparison of the empirical sample complexity of AC-State on the DoublePrime(p,, g) environments, with and without active exploration. (Note that we did not perform a final comparison p 41, q 43, because the hyperparameter search for the active exploration policy (see Figure [) with a maximum K of 50000 did not rule out a Kapt > 50000, so we could not fairly tune This hyperparameter for the active exploration policy. We were unable perform hyperparameter search with K > 50000 due to technical limitations.)
notice an issue: the encoder is only trained with segments length k'+ 1,, which (assuming d' and 7" are anywhere close to accurate) will only scale linearly with D. However, properly learning the encoder for the DoublePrime environments requires examining longer segments of trajectories than the diameter of the dynamics graph. Therefore, in order adapt this exploration method to the setting we consider, we climinate this feature of the algorithm and train & on the entire collected trajectory, in Equation 140
Furthermore, we give the exploration method what should be a large, unnatural advantage, by plan- ning, choosing goals, and assessing state visitation the true ground truth dynamics T, with the ground-Iruth state encoder Φ.. This represents what should be the "best possible" case for the algo- rithm. It also enables to test in this setting without having to optimize ø' continuously during data collection, and to assess many values of the hyperparameter K (which would impact the intermedi- ate encoder σ') simultaneously on a single collected trajectory.
In our implementation, we reset the state visitation counts between collecting the "fitting" and "opti- mization" portions of the trajectory, order to ensure that they are distributed in the same way. We also take uniformly random actions in cases where both actions produce a path of the same length to the goal. Otherwise, we simply take the shortest available path to the goal state, and, when reached, pick new goal and take one random action.
Hyperparameter search results are shown in Figure X and top-line results are shown in Table We see in Table that AC-State with the active exploration method seems to slightly underperform random exploration on the DoublePrime((,. q) environments, requiring slightly more samples to find the correct encoder, for all values of (p, q) that were tested 12 While it is not immediately obvious why we see this gap, these results support our initial hypothesis that, because the DoublePrime en- vironments well-explored by uniformly random policies, techniques designed to improve latent state coverage are not helpful in improving AC-State's performance on these environments.
F.4 DISCUSSION
important note that the DoublePrime environments are not card-exploration environments, nor do they have any time-correlated noise, nor do they have rich observations. In fact, adding time-correlated noise or rich observations would likely only scale the sample-complexity of STEEL on these environments only modestly (linearly in (miz and log([Fl)).. Rather, STEEL outperforms AC-State on large instances of these environments because the environments' transition dynamics arranged in way that is in a sense "adversarial" to multistep-inverss methods. By contrast, because STEEL has well-understood sample-complexity that applies for any transition graph struc- ture (as long the reachability assumption holds), we do not expect that any such "adversarial" environments to exist for STEEL.
However, note that because we only tested trajectory lengths I at particular values as described in Section
1.21 some of these performance differences may appear exaggerated; for example, the difference between a trajectory length of 1 10 and 2 10 must be interpreted with the understanding that trajectory lengths benween 1 10 and 2 10 were tested.
46

---

Published as a conference paper at ICLR 2025
Hyperparameter Search,
Hyperparameter Search, p 5,
1.0
20
Success Rate on DoublePrimeep.o
60
Success Rate DeublePrime Whleerimelpp...
60
180
0.6
100
300
120
0.4
360
0.4
160
02
480
0.2
0.0
0.0
10 10
10
Trajectory Length T
Trajectory Length T
Hyperparameter Search, p 11, =
13
Hyperparameter Search, p 17, q
19
1.0
1.0
Success Rate on PoublePrimeio DoublePtime Q)
Success Rate DoublePrimetp.ng
1000
0.8
4000
0.8
1500
6000
K 10000
2500
12000
3500
14000
4000
16000
4500
18000
0.0
0.0
1. 2.
6- 7. 8-
10
1.
2. 3-
4.
7-
8-
9. 10
10a
10a
10
10 10f
10
10'
102
102
107
10'
10
107
Trajectory Length
Trajectory Length T
Hyperparameter Search, = 29,
31
Hyperparameter Search, Search,
41,
43
Success Rate on Double DoublePrimeep.. q)
5000
Rate Double DoublePrimeloo
5000
10000
0.8
10000
20000
0.6
20000
0.6
25000
30000
0.4
30000
0.4
35000
35000
40000
0.2
40000
43000
43000
1. 2. 3
5
6-
7
8.
10
1. 2. 3
4
6.
7
6
9
10
10a 10
10"
16a
10
10
10
10"
10" 10"
10a 10" 10"
10a
10a
10"
10"
10a
10
10"
Trajectory Length
Trajectory Length
Figure 8: Results from the hyperparameter search for the maximum-step-comee parameter K for DoublePrime(p. q), using the active exploration policy described in Section
F3
47


Reviews:
[{"review_id": "zlOn7nfPvP", "reviewer": "Reviewer_fyCo", "summary": "The authors propose Single-Trajectory Exploration for Ex-BMDPs via Looping (STEEL), an algorithm to learn the endogenous (controllable) states in an Exogenous Block Markov Decision Process (Ex-BMDP) when the agent is dealing with one continuous infinite trajectory without resetting to some known states. STEEL achieves this by taking actions that result in a predictable cycle of states and iteratively updating the list of known controllable states and their transitions. They show theoretically the sample complexity and correctness of STEEL with simulations on some small environments.", "questions": "- In the Appendix in equations 33 and 35, could you further explain how the sets $\\mathcal{D}_i^\\mathcal{A}$ and  $\\mathcal{D}_i^\\mathcal{B}$ are constructed?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The introduction and related work highlight this work really well. It explains the existing work nicely and shows where the gaps lie and how this work attempts to extend it.\n- The algorithm stands out in terms of the settings it covers compared to existing work. It deals with infinite trajectories, partial observability, and optimization with function approximators all while providing sample complexity guarantees.\n- The algorithm itself is designed very well and has a lot of interesting features which include: forcing a cycle of states through the repetition of actions and detecting the unique states in a cycle using a classifier oracle.\n- The limitations of the algorithm are clearly discussed with useful insights on how to extend this work in the future.", "weaknesses": "- Section 4 can be a bit hard to follow. To quite understand how the algorithm exactly works one has to switch between reading the section text, the pseudocode, and parts of the Appendix. I suggest moving the pseudocode to the appendix and providing further explanation of the algorithm in the main text such that the reader can get a high-level idea of how the Algorithm works from just reading section 4.\n- There are parts of the algorithm that are not very intuitive and might require some further discussion. For example, it is mentioned that the dataset $D_0, D_1$ used in the CycleFind subroutine are generated in a way such that they are disjoint if $n'_{cyc}$ is equal to $n_cyc$. Intuitively, how does the selection process achieve this? \n- In the experiments section, the authors mention that previous work by Lamb et al.(2023) and Levine et al. (2024) don't have theoretical correctness guarantees, which can be why it seems to have better sample efficiency than STEEL. I suggest also including the percentage of runs where these baselines get the correct states and transition probabilities and how often they fail compared to STEEL which is proven to get it right with high probability. This can add additional value to how STEEL outperforms the baselines in terms of correctness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vLIj5vIVVw", "reviewer": "Reviewer_m7A7", "summary": "The authors violated the instructions and reduced the font size substantially for Algorithm 1 and 2. Given they took a whole 10 pages, I decided to recommend desk rejection. If the AC decides differently, please inform me accordingly.", "questions": "NA", "limitations": "", "rating": 1, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "NA", "weaknesses": "NA", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vYqqW7gIe5", "reviewer": "Reviewer_kqTs", "summary": "The paper proposes a representation learning method for Ex-BMDP called STEEL. This method identifies the small latent state space of Ex-BMDP - which encodes the essential controllable part of the MDP - while jointly learning an encoder that maps observations to the latent state. Notably, this approach can be applied without requiring \"reset\" commands, allowing the algorithm to learn from a single trajectory. The key idea is to repeat sequences of actions to detect cycles in the latent state space, which enables the collection of multiple i.i.d. samples to discover the latent space structure. The sample complexity of the algorithm is shown to be polynomial in the size of the latent space, the mixing rate of the Markovian exogenous process, and the complexity of the encoder function class. The algorithm is demonstrated on two problem scenarios.", "questions": "* A discussion of the block assumption on $\\mathcal{Q}$ with respect to $\\mathcal{S}$ would be helpful. In many practical scenarios, the noisy nature of the emission (or observation) function can make distinguishing between two latent states directly from observations challenging, necessitating filtering techniques. It would be beneficial to clarify whether this assumption is not overly restrictive or if it cannot be easily weakened but is widely adopted.\n* Is there a known lower bound for the sample complexity of Ex-BMDP under deterministic latent dynamics? Is STEEL nearly optimal under this assumption, or is there potential for further improvement?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is clearly written, and the analysis of the key result - specifically, the sample complexity of STEEL being polynomial in the latent space size - is supported by solid mathematical arguments. The algorithm's description is intuitive and effectively conveys its core concepts.\n* Furthermore, representation learning from a single episode has been a long-standing interest in the RL community, making this paper's contribution highly relevant to the field.\n* The paper provides a comprehensive literature review, effectively demonstrating the novelty of the work and differentiating it from recent existing works.", "weaknesses": "* The method relies on several assumptions, particularly concerning the latent state space $\\mathcal{S}$. For example, the assumptions of deterministic latent dynamics and the reachability condition of the latent state space are critical for STEEL's CycleFind to function. Addressing these assumptions seems non-trivial, and overcoming them is posed as future work.\n* Although the sample complexity of STEEL is polynomial in the size of the latent state space, the numerical simulations show that a substantial number of samples (millions) are required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KEGHUJEpvg", "reviewer": "Reviewer_1LWB", "summary": "This paper studies a structured class of MDPs called an Ex-BMDP, where the latent factors of the observations decompose into a lower-dimensional controllable factor (which evolves deterministically according to the agent's action) and high-dimensional exogenous factor (which evolves independent of actions). This paper focuses on the single-episodic setting and proposes sample-efficient algorithms for learning controllable dynamics of an Ex-BMDP with sample complexity that depends only on the sizes of the low-dimensional controllable state, the encoder function class, and the mixing time of the exogenous noise factor. The paper also empirically tests the proposed STEEL algorithm on the infinite-horizon variations of the \"combination lock\" and \"multi-maze\" environments.", "questions": "- The paper presumes the availability of an encoder hypothesis class $\\mathcal{F}$, where the true decoder $f(x)$ is included, and the final complexity depends on the size $\\log\\mathcal{F}$. However, it does not seem to specify how to choose this hypothesis class. The simulation section gives an example of $\\mathcal{F}$ that is specific to the examples. Is there any general procedure for selecting $\\mathcal{F}$ with a reasonable size that also guarantees to include the correct decoder? In general settings beyond the specific examples given in the paper, can you provide guidelines or heuristics for selecting an appropriate hypothesis class?\n- Given $\\mathcal{F}$, the paper also assumes access to a training oracle that optimally distinguishes two sets of observations (e.g., similar to minimizing 0-1 loss). What would be an example of such an oracle without prior knowledge of the true classifier? And what is the sample/computation cost of constructing such an oracle?\n- The STEEL algorithm assumes access to an upper bound on the mixing time $t_{mix}$ for the exogenous dynamics. For a general setting with unknown exogenous latent factors and dynamics, how do you get such an upper bound? Can you discuss potential methods or heuristics for estimating or bounding the mixing time in settings where the exogenous dynamics are not fully known?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The class of Ex-BMDP studied in this paper is a general class of structured POMDPs. It captures problems where, despite having high dimensional observation, the majority of the states are exogenous and only a small controllable state matters for learning. It therefore allows more sample-efficient learning by filtering out the exogenous factors and reducing to a smaller MDP depending only on the controllable states. Such setting fits many applications and gives insight to how to best exploit these hidden structures to optimize learning.\n- The main novelty of this paper compared with prior work in Ex-BMDP is that instead of the episodic setting in Efroni et al. (2022) where one gets to reset to starting state, it assumes the agent interacts with the environment in a single episode. This setting is more challenging given it is more difficult to collect samples of a given latent state without the episodic resets. \n- The paper also assumes a more general assumption on the state and emission function, where only the partial inverse with respect to the controllable state exists, but places no such assumption on the exogenous states. This is more general than assumption a block structure in prior works which allows a full inverse from observation to state.", "weaknesses": "- The proposed algorithm is highly dependent on the assumptions that (1) the dynamics of the latent controllable states is deterministic; (2) the mixing time of the exogenous dynamics. Intuitively, assumption (1) leads to to a cycle of latent states of bounded length that is repeatedly visited and allows repeated collection of the same latent state, which on a high-level is similar to \"resetting\" the environment; assumption (2), given the looping behavior, can wait out the mixing time of the exogenous dynamics and collect near i.i.d. samples of each latent state. However, assuming deterministic dynamics and bounded mixing time seems restrictive, and possibly does not capture many practical setting. How sensitive is the algorithm to the violation of both assumptions? Does non-deterministic dynamics of the controllable latent states break the proposed algorithm?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)