Learning a Fast Mixing Exogenous Block MDP using a Single Trajectory

TL;DR

The paper addresses unsupervised representation learning for Exogenous Block MDPs in a no-reset, single-trajectory setting and introduces STEEL, the first provably sample-efficient algorithm for learning the endogenous latent dynamics from a single trajectory. STEEL combines three phases—learning latent dynamics via CycleFind, collecting additional near-iid samples, and training a one-vs-rest encoder—to recover the endogenous state space, the latent transition function, and a high-accuracy encoder with a formal sample-complexity bound that depends on $|\mathcal{S}|$, the encoder class size $|\mathcal{F}|$, and the exogenous mixing time bound $\hat{t}_{\text{mix}}$. The paper proves correctness and complexity guarantees, and validates STEEL on toy problems, illustrating robust latent recovery under high-dimensional, time-correlated noise. Acknowledging limitations, the approach relies on deterministic latent dynamics and known mixing-time bounds, and it assumes reachability and a tractable encoder oracle, but it provides a principled route to representation learning in no-reset, real-world-like settings. Overall, STEEL advances sample-efficient latent dynamics learning in Ex-BMDPs and offers a foundation for more robust unsupervised representations in long-horizon, noisy environments.

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.

Figures (8)

• Figure 1: STEEL discovers the latent dynamics $\mathcal{S}$ and $T$ by iteratively adding cycles to the learned dynamics graph. In this simple example, the initially-unknown "true" dynamics consist of 6 states arranged in a grid, where the agent can move (U)p, (D)own, (L)eft, or (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-known parts of the dynamics graph in black, and the cycle being explored in red. States are represented as circles and transitions as arrows.
• Figure 2: CycleFind determines the period of the cycle in $x_{CF}$. See Section \ref{['sec:alg']} under "CycleFind Phase 1." We show the sequence $x_{CF}$ sampled from $M$: specifically, we show every $|\hat{a}|$'th observation, where the same actions $\hat{a}$ are taken between each one. The observations' latent states are color-coded as red, blue, and green: a pattern repeats every $3|\hat{a}|$ steps, so $n_{cyc} = 3$. $D_1$ consists of the first observation in each $(n_{cyc}'|\hat{a}|)$-cycle, and $D_0$ the other observations taken between executions of $\hat{a}$. (Spans of length $\geq \hat{t}_{\text{mix}}$ are skipped to ensure certain subsets of the datasets are near-i.i.d.)
• Figure 3: Visualisations of the simulation experiment environments. For both environments, we show the ground-truth latent dynamics $T$ (in the case of the combination lock, we show an arbitrary instance of $T$, for some $[a^*_0, ... a^*_{K-1}]$), and an example transition in the observed space $\mathcal{X}$.
• Figure 4: Illustration of the sampling procedure for datasets $\mathcal{D}'_i$. The first goal is to ensure that for each cycle position $i$, the samples in $\mathcal{D}'_i$ are sampled $t_{\text{mix}}$ steps apart from each other, and are therefore nearly i.i.d. We also want to ensure that for any pair of cycle positions $i,j$, there is a large subset of $\mathcal{D}'_i$ that only contains samples retrieved at least $t_{\text{mix}}$ steps apart from some large subset of $\mathcal{D}'_j$. (This second goal is meant to guarantee that if $\mathcal{D}'_i$ and $\mathcal{D}'_j$ represent the same latent state, then it is unlikely that any classifier exists than can separate the two subsets perfectly, which would be strictly necessary to perfectly separate $\mathcal{D}'_i$ and $\mathcal{D}'_j$). However, it is not necessary for all samples in $\cup_i \mathcal{D}_i$ to be collected $t_{\text{mix}}$ steps apart. Therefore, we collect an observation of each cycle position $i \in \{0, n_{\text{cyc}} \cdot |\hat{a}|\}$, all together, every $\hat{t}_{\text{mix}}$ steps (rounded up to the cycle period). We continue until $n_\text{samp}$ observations of each position are collected; then wait $\hat{t}_{\text{mix}}$ steps and collect $n_\text{samp}$ additional observations of each cycle position. This process ensures that $\mathcal{D}'_i$ and $\mathcal{D}'_j$ contain complementary subsets $\mathcal{D}^{A}_i$ and $\mathcal{D}^{B}_j$, each with at least $n_\text{samp}$ samples, such all samples in $\mathcal{D}^{A}_i \cup \mathcal{D}^{B}_j$ are near-i.i.d.
• Figure 5: Latent dynamics of $\text{DoublePrime}(11,13)$.

Theorems & Definitions (8)

• Theorem 1
• Theorem 1
• Proposition 1
• proof
• Lemma 1
• proof
• Lemma 2
• proof