Learning a Fast Mixing Exogenous Block MDP using a Single Trajectory

Reviewer 1LWB,

We thank the reviewer for a providing a substantial and constructive review.

The review mentions two substantial limitations: the assumptions of determinism and a known upper bound on the mixing time.

Determinism: We agree that this is a substantial limitation: we address this concern thoroughly in our Response #1 to Reviewer kqTs.

Known Upper bound on the mixing time:

The review mentions that “ assuming [...] bounded mixing time seems restrictive, and possibly does not capture many practical settings,” and asks how such an upper bound may be obtained by an algorithm. Here, we argue that the assumption that an upper bound on the mixing time is provided is in fact absolutely necessary. We argue that:

1. Any general, provably sample-efficient algorithm for learning, with high probability, an Ex-BMDP from a single trajectory must necessarily require an upper bound on the mixing time of the exogenous noise to be provided a 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, scale linearly with the mixing time of the exogenous noise.

At a high level, we argue that, with a single trajectory, it can be impossible to 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 equilibrium has a quite different distribution of observations. This will make the learned encoder fail on the Ex-BMDP at equilibrium.

More precisely, to show (1): suppose the converse is true: that there exists some algorithm $A$ 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, as well as an encoder with accuracy (for every endogenous state, under the stationary exogenous distribution) of at least $1-\epsilon$, with probability at least $1-\delta$, for small values of $\epsilon$ and $\delta$.

Then, consider any Ex-BMDP $M_1$ and related parameterized family of Ex-BMDPs $M_2(\gamma), M_3(\gamma)$ with the following properties:

• $M_1$ has N endogenous latent states $s_1,...,s_N$ with some transition function $T$, a single exogenous latent state $e_1$, and an emission function $Q(s,e_1)$. (In other words, $M_1$ is any Block MDP).

• $M_2(\gamma)$ has the same endogenous states and transition probabilities as $M_1$, but has two exogenous latent states $e_1, e_2$. The state $e_1$ transitions to $e_2$, and $e_2$ transitions to $e_1$, each with probability $\gamma$. Note that the stationary distribution of the exogenous state is uniform over $e_1$ and $e_2$. We also assume that the initial exogenous state distribution is uniform over $e_1$ and $e_2$. Regardless of $\gamma$, the emission function of $M_2(\gamma)$ is defined such that $\forall s_i, Q_{M_2}(s_i,e_1) = Q_{M_1}(s_i,e_1)$.

• $M_3(\gamma)$ is identical to $M_3(\gamma)$, except that its emission function is defined as $\forall s_i, Q_{M_3}(s_i,e_1) := Q_{M_1}(s_i,e_1)$, but $\forall s_i, Q_{M_3}(s_i,e_2) := Q_{M_2}(s_{(i+1) \text{ mod } N},e_2)$. In other words, when the exogenous state is equal to $e_2$, the emission distributions for the endogenous states are permuted in $M_3$ compared to $M_2$.

We also assume that the encoder hypothesis class can represent the inverses of $Q_{M_1}$, $Q_{M_2}$, and $Q_{M_3}$. However, note that by construction, any fixed encoder $\phi(\cdot)$ which has accuracy at least $1- \epsilon$ on $M_2(\gamma)$ (for every endogenous state, under the stationary exogenous distribution) can only have accuracy at most $0.5 + \epsilon$ on $M_3(\gamma)$, because, when the exogenous state is $e_2$, any time that the encoder returns the correct latent state for $M_2(\gamma)$, it will return the incorrect latent state for $M_3(\gamma)$.

Now, consider what happens when we run the algorithm $A$ on $M_1$. The number of environment steps that $A$ takes on $M_1$ forms some distribution; let $t$ be the 90’th percentile of this distribution, so that with probability 0.9, $A$ stops sampling before step $t$.

(continued)

Reviewer kqTs,

Thanks for your constructive and thoughtful review. We will respond to most of the points that you raised now, and will follow up with additional responses later in the rebuttal period:

Restrictive assumptions: determinism and latent-state reachability:

Determinism: We agree that the determinism assumption on the endogenous dynamics is restrictive. However, recall that even in the episodic setting with resets to a single latent state, PPE (Efroni et al, 2022: the current state-of-the-art algorithm for efficiently learning Ex-BMDP latent state encoders) requires near-determinism, with only very rare deviations from deterministic behavior (one non-deterministic latent transition every $4|\mathcal{S}|$ episodes). In other words, tightening this assumption to full determinism does not represent a major restriction beyond what is found in previous work for an “easier” setting.

Furthermore, we were recently made aware of a new paper (Mhammedi et al., to appear in NeurIPS 2024) which proposes an algorithm for sample-efficient learning in the Ex-BMDP setting with nondeterministic latent dynamics. However, the algorithm proposed in that work requires simulator access to the environment, meaning that the agent can reset the environment to any previously-observed observation $x \in \mathcal{X}$. This suggests a general connection between the ability to handle nondeterminstic dynamics and the ability to reset the environment, leading to a “spectrum” of algorithms:

• Mhammedi et al, 2024: full nondeterminism / reset to arbitrary states
• Efroni et al, 2022: near determinism (only very rare departures from determinism allowed) / reset to same state every episode
• This work: full determinism / no resets, ever

We believe that our work therefore contributes to the understanding of the Ex-BMDP setting in general, by filling out one endpoint of this spectrum. (We have added a reference to Mhammedi et al, 2024 to our paper; see the revision we have added.)

(continued)

Reviewer m7A7,

We apologize for any inconvenience this formatting issue may have caused. However, use of smaller-than-standard fonts in figures, tables, and algorithms seems to be fairly common in ICLR submissions, and font size in figures, etc, was not explicitly mentioned in the formatting instructions.

Regardless, in response to Reviewer fyCo's suggestion, we have moved the algorithms to an appendix (and increased their font size), and incorporated more detail into our high-level description of the algorithms in the main text. This should resolve the issue.

In any case, it seems that this paper was not desk-rejected. If you are willing to provide a substantial review of our paper, we would greatly appreciate the additional feedback. (If not, we ask that you delete your review, so that it is not averaged in to our score in case the AC uses a strict cutoff.)

Reviewer fyCo,

Thank you for your constructive review. We will respond to parts of your review now, and we aim to follow up with additional empirical results as they become available. Note that we have already submitted a revision incorporating the suggested changes to the paper.

• "Section 4 can be a bit hard to follow.... 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.": We have now done this: the pseudocode is now in the Appendix, and we have substantially expanded the text of Section 4 to describe the algorithm in a more self-contained manner.

• "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?": We have added an in-depth explanation of how $D_0$ and $D_1$ are constructed to Section 4 of the paper.

• "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." For the particular "multi-maze" environment in the paper, (for which we mention that Lamb et al. (2023) and Levine et al. (2024)'s algorithms have better empirical sample efficiency), Levine et al. (2024) also show that both algorithms learn the dynamics consistently, each succeeding on 20/20 random seeds. We noted in the paper that this strong performance may be partially due to the particular choice of neural network architecture chosen by Lamb et al. (2023), which has a spatial prior that favors focusing on a single maze (as also referenced by Levine et al; in Footnote 8.)

However, even if we accept that STEEL underperforms Lamb et al. (2023) and Levine et al. (2024)'s algorithms on this particular "multi-maze" environment, the benefit of STEEL is that it is guaranteed to perform within a known sample-complexity bound, on any environment that meets its assumptions. To underscore this point, we are currently running experiments involving a specific family of tabular Ex-BMDPs that were introduced by Levine et al. (2024), which have controllable latent dynamics that are particularly difficult for either of Lamb et al. (2023) or Levine et al. (2024)'s algorithms to learn efficiently. Preliminarily, we have been able to show that, in apples-to-apples comparisons, STEEL empirically outperforms these prior methods on environments with this particular structure of controllable latent dynamics. We hope to have results ready to include in the paper by the end of the rebuttal period.

• "In the Appendix in equations 33 and 35, could you further explain how the sets $\mathcal{D}^A_i$ and $\mathcal{D}^B_i$ are constructed?" We have added an additional figure (Figure 4) with an extensive caption, that explains how these datasets are constructed.