capbtabboxtable[][]

# Towards Principled Representation Learning from Videos for Reinforcement
Learning

Dipendra Misra111DM and AS contributed equally. Akanksha Saran∗ Tengyang Xie
Alex Lamb John Langford
∗Microsoft Research New York †Sony Research
<EMAIL_ADDRESS><EMAIL_ADDRESS>{tengyangxie, lambalex,
<EMAIL_ADDRESS>Work done while at Microsoft Research NYC.

###### Abstract

We study pre-training representations for decision-making using video data,
which is abundantly available for tasks such as game agents and software
testing. Even though significant empirical advances have been made on this
problem, a theoretical understanding remains absent. We initiate the
theoretical investigation into principled approaches for representation
learning and focus on learning the latent state representations of the
underlying MDP using video data. We study two types of settings: one where
there is iid noise in the observation, and a more challenging setting where
there is also the presence of exogenous noise, which is non-iid noise that is
temporally correlated, such as the motion of people or cars in the background.
We study three commonly used approaches: autoencoding, temporal contrastive
learning, and forward modeling. We prove upper bounds for temporal contrastive
learning and forward modeling in the presence of only iid noise. We show that
these approaches can learn the latent state and use it to do efficient
downstream RL with polynomial sample complexity. When exogenous noise is also
present, we establish a lower bound result showing that the sample complexity
of learning from video data can be exponentially worse than learning from
action-labeled trajectory data. This partially explains why reinforcement
learning with video pre-training is hard. We evaluate these representational
learning methods in two visual domains, yielding results that are consistent
with our theoretical findings.

###  1 Introduction

Representations pre-trained on large amounts of offline data have led to
significant advances in machine learning domains such as natural language
processing (Liu et al., 2019; Brown et al., 2020) and multi-modal learning
(Lin et al., 2021; Radford et al., 2021). This has naturally prompted a
similar undertaking in reinforcement learning (RL) with the goal of training a
representation model that can be used in a policy to solve a downstream RL
task. The natural choice of data for RL problems is trajectory data, which
contains the agent’s observation along with actions taken by the agent and the
rewards received by it (Sutton & Barto, 2018). A line of work has proposed
approaches for learning representations with trajectory data in both offline
(Uehara et al., 2021; Islam et al., 2022) and online learning settings (Nachum
et al., 2018; Bharadhwaj et al., 2022). However, unlike text and image data,
which are abundant on the internet or naturally generated by users, trajectory
data is comparatively limited and expensive to collect. In contrast, video
data, which only contains a sequence of observations (without any action or
reward labeling), is often plentiful, especially for domains such as gaming
and software. This motivates a line of work considering learning
representations for RL using video data (Zhao et al., 2022). _But is there a
principled foundation underlying these approaches?_ _Are representations
learned from video data as useful as representations learned from trajectory
data?_ We initiate a theoretical understanding of these approaches to show
when and how these approaches yield representations that can be used to solve
a downstream RL task efficiently.

Figure 1: A flowchart of our video pre-training phase. Left: We assume access
to a large set of videos (or, unlabeled episodes). Center: A representation
learning method is used to train a model $\phi$ which maps an observation to a
vector representation. Right: This representation can be used in a downstream
task to do reinforcement learning or visualize the latent world state.

Consider a representation learning pipeline shown in Figure 1. We are provided
videos, or equivalently a sequence of observations, from agents navigating in
the world. We make no assumption about the behavior of the agent in the video
data. They can be trying to solve one task, many different tasks, or none at
all. This video data is used to learn a model $\phi$ that maps any given
observation to a vector representation. This representation is subsequently
used to perform downstream RL — defining a policy on top of the learned
representation and only training the policy for the downstream task. We can
also use this representation to define a dynamics model or a critique model.
The representation can also help visualize the agent state space or dynamics
for the purpose of debugging.

A suitable representation for performing RL efficiently is aligned with the
underlying dynamics of the world. Ideally, the representation captures the
latent agent state, which contains information about the world relevant to
decision-making while ignoring any noise in the observation. For example, in
Figure 1, ignoring noise such as the motion of geese in the background is
desirable if the task involves walking on the pavement. We distinguish between
two types of noise: (1) temporally independent noise that occurs at each time
step independent of the history, (2) temporally dependent noise, or exogenous
noise, that can evolve temporally but in a manner independent of the agent’s
actions (such as the motion of geese in Fig. 1).

A range of approaches have been developed that provably recover the latent
agent state from observations using trajectory data (Misra et al., 2020;
Efroni et al., 2022) which contains action labels. However, for many domains
there is relatively little trajectory data that exists naturally, making it
expensive to scale these learning approaches. In contrast, video data is more
naturally available but these prior provable approaches do not work with video
data. On the other hand, it is unknown whether approaches that empirically
work with video data provably recover the latent representation and lead to
efficient RL. Motivated by this, we build a theoretical understanding of three
such video-based representation learning approaches: _autoencoder_ which
trains representations by reconstructing observations, _forward modeling_
which predicts future observations, and _temporal contrastive_ learning which
trains a representation to determine if a pair of observations are causally
related or not.

Our first theoretical result shows that in the absence of exogenous noise,
forward modeling and temporal contrastive learning approaches both provably
work. Further, they lead to efficient downstream RL that is strictly more
sample-efficient than solving these tasks without any pre-training. Our second
theoretical result establishes a lower bound showing that in the presence of
exogenous noise, any compact and frozen representation that is pre-trained
using video data cannot be used to do efficient downstream RL. In contrast, if
the trajectory data was available, efficient pre-training would be possible.
This establishes a statistical gap showing that video-based representation
pre-training can be exponentially harder than trajectory-based representation
pre-training.

We empirically test our theoretical results in three visual domains: GridWorld
(a navigation domain), ViZDoom basic (a first-person 3D shooting game), and
ViZDoom Defend The Center (a more challenging first-person 3D shooting game).
We evaluate the aforementioned approaches along with ACRO (Islam et al.,
2022), a representation pre-trained using trajectory data and designed to
filter out exogenous noise. We observe that in accordance with our theory,
both forward modeling and temporal contrastive learning succeed at RL when
there is no exogenous noise. However, in the presence of exogenous noise,
their performance degrades. Specifically, we find that temporal contrastive
learning is especially prone to fail in the presence of exogenous noise, as it
can rely exclusively on such noise to optimally minimize the contrastive loss.
While we find that forward modeling is somewhat robust to exogenous noise,
however, as exogenous noise increases, its performance quickly degrades as
well. While any finite-sample guarantees for the autoencoding method remain an
open question, empirically, we find that the performance of autoencoder-based
representation learning is unpredictable. On the other hand, ACRO continues to
perform well, highlighting a disadvantage of video pre-training. The code for
all experiments is available as part of the Intrepid codebase at
https://github.com/microsoft/Intrepid.

###  2 Preliminaries and Overview

We cover the theoretical settings and the dataset type in this section. We
start by describing the notations in this work.

##### Mathematical Notation.

We use $[N]$ for $N\in\mathbb{N}$ to define the set $\\{1,2,\cdots,N\\}$. We
use calligraphic notations to define sets (e.g., $\mathcal{U}$,
${\mathcal{S}}$). For a given countable set $\mathcal{U}$, we define
$\Delta(\mathcal{U})$ as the space of all distributions over $\mathcal{U}$. We
denote the cardinality of a set $\mathcal{U}$ by $|\mathcal{U}|$. We assume
all sets to be countable. Finally, we denote ${\tt poly}\\{\cdot\\}$ to denote
that something scales polynomially in listed quantities.

##### Block MDPs.

We study episodic reinforcement learning in Block Markov Decision Processes
(Block MDP) (Du et al., 2019). A Block MDP is defined by the tuple
$(\mathcal{X},{\mathcal{S}},\mathcal{A},T,R,q,\mu,H)$ where $\mathcal{X}$ is a
set of observations that can be infinitely large, ${\mathcal{S}}$ is a finite
set of _latent_ states, and $\mathcal{A}$ is a set of finite actions. The
transition dynamics
$T:{\mathcal{S}}\times\mathcal{A}\rightarrow\Delta({\mathcal{S}})$ define
transitions in the latent state space. The reward function
$R:{\mathcal{S}}\times\mathcal{A}\rightarrow[0,1]$ assigns a reward $R(s,a)$
if action $a$ is taken in the latent state $s$. Observations are generated
from states using the emission function
$q:{\mathcal{S}}\rightarrow\Delta(\mathcal{X})$ that generates an observation
$x\sim q(\cdot\mid s)$ when the agent is in latent state $s$. This emission
process contains temporally independent noise but no exogenous noise. Finally,
$\mu\in\Delta({\mathcal{S}})$ is the distribution over the initial latent
state and $H$ is the horizon denoting the number of actions per episode. The
agent interacts with a block MDP environment by repeatedly generating an
episode $(x_{1},a_{1},r_{1},x_{2},a_{2},r_{2},\cdots,x_{H},a_{H},r_{H})$ where
$s_{1}\sim\mu$ and for all $h\in[H]$ we have $x_{h}\sim q(\cdot\mid s_{h})$,
$r_{h}=R(s_{h},a_{h})$, and $s_{h+1}\sim T(\cdot\mid s_{h},a_{h})$, and all
actions $\\{a_{h}\\}_{h=1}^{H}$ are taken by the agent. The agent never
directly observes the latent states $(s_{1},s_{2},\cdots,s_{H})$.

A key assumption in Block MDPs is that two different latent states cannot
generate the same observation. This is called the _disjoint emission property_
and environments with this property are called rich observation environments,
meaning the observation is _sufficiently rich_ to allow decoding of the latent
state from it. Formally, this property allows us to define a decoder
$\phi^{\star}:\mathcal{X}\rightarrow{\mathcal{S}}$ that maps an observation to
the unique state that can generate it. The agent does not have access to
$\phi^{\star}$. If the agent had access to $\phi^{\star}$, one could map each
observation from an infinitely large space to the finite latent state space,
which allows the use of classical finite RL methods (Kearns & Singh, 2002).

##### Exogenous Block MDPs (Ex-Block MDP).

We also consider reinforcement learning in Exogenous Block MDPs (Ex-Block
MDPs) which can be viewed as an extension of Block MDPs to include exogenous
noise (Efroni et al., 2022). An Ex-Block MDP is defined by
$(\mathcal{X},{\mathcal{S}},\Xi,\mathcal{A},T,T_{\xi},R,q,H,\mu,\mu_{\xi})$
where $\mathcal{X},{\mathcal{S}},\mathcal{A},T,R,H$ and $\mu$ have the same
meaning and type as in Block MDPs. The additional quantities include $\Xi$
which is the space of exogenous noise. We use the notation $\xi\in\Xi$ to
denote the exogenous noise. For the setting in Fig. 1, the exogenous noise
variable $\xi$ captures variables such as the position of geese, the position
of leaves on the trees in the background, and lighting conditions. We allow
the space of exogenous noise to be infinitely large. The exogenous noise $\xi$
changes with time according to the transition function
$T_{\xi}:\Xi\rightarrow\Delta(\Xi)$. Note that unlike the agent state
$s\in{\mathcal{S}}$, the exogenous noise $\xi\in\Xi$, evolves independently of
the agent’s action and does not influence the evolution of the agent’s state.
The emission process $q:{\mathcal{S}}\times\Xi\rightarrow\Delta(\mathcal{X})$
in Ex-Block MDP uses both the current agent state and exogenous noise, to
generate the observation at a given time. For example, the image generated by
the agent’s camera contains information based on the agent’s state (e.g.,
agent’s position and orientation), along with exogenous noise (e.g., the
position of geese). Finally, $\mu_{\xi}$ denotes the distribution over the
initial exogenous noise $\xi$.

Similar to the Block MDP, we assume a disjoint emission property, i.e., we can
decode the agent state and exogenous noise from a given observation. We use
$\phi^{\star}:\mathcal{X}\rightarrow{\mathcal{S}}$ to denote the true decoder
for the agent state ($s$), and $\phi^{\star}_{\xi}:\mathcal{X}\rightarrow\Xi$
to denote the true decoder for exogenous noise ($\xi$).

##### Provable RL.

We assume access to a policy class
$\Pi=\\{\pi:\mathcal{X}\rightarrow\mathcal{A}\\}$ where a policy $\pi\in\Pi$
allows the agent to take actions. For a given policy $\pi$, we use
$\mathbb{E}_{\pi}\left[\cdot\right]$ to denote expectation taken over an
episode generated by sampling actions from $\pi$. We define the value of a
policy $V(\pi)=\mathbb{E}_{\pi}\left[\sum_{h=1}^{H}r_{h}\right]$ as the
expected total reward or expected return. Our goal is to learn a near-optimal
policy $\widehat{\pi}$, i.e.,
$\sup_{\pi\in\Pi}V(\pi)-V(\widehat{\pi})\leq\varepsilon$ with probability at
least $1-\delta$ for a given tolerance parameter $\varepsilon>0$ and failure
probability $\delta\in(0,1)$, using number of episodes that scale polynomially
in $\nicefrac{{1}}{{\varepsilon}}$, $\nicefrac{{1}}{{\delta}}$, and other
relevant quantities. We will call such an algorithm as provably efficient.
There exist several provably efficient RL approaches for solving Block MDPs
(Mhammedi et al., 2023; Misra et al., 2020), and Ex-Block MDPs (Efroni et al.,
2022). These approaches typically assume access to a decoder class
$\Phi=\left\\{\phi:\mathcal{X}\rightarrow[N]\right\\}$ and attempt to learn
$\phi^{\star}$ using it. These algorithms don’t use any pre-training and
instead directly interact with the environment and learn a near-optimal policy
by using samples that scale with ${\tt
poly}(S,A,H,\ln|\Phi|,\nicefrac{{1}}{{\varepsilon}},\nicefrac{{1}}{{\delta}})$.
Crucially, the dependence on $\ln|\Phi|$ cannot be removed. The decoder class
$\Phi$ and all other function classes in this work are assumed to have bounded
statistical complexity measures. For simplicity, we will assume that these
function classes are finite and derive guarantees that scale logarithmically
in their size (e.g., $\ln|\Pi|$).222Our theoretical analyses can be extended
to other statistical complexity metrics such as Rademacher complexity.

##### Representation Pre-training using Videos.

As mentioned earlier, existing algorithms for Block MDPs and Ex-Block MDPs
require online interactions that scale with $\ln|\Phi|$. For real-world
problems, the decoder class $\Phi$ can be a complex neural network and these
approaches may, therefore, require lots of expensive online data. Offline RL
approaches offer a substitute for expensive online interactions but require
access to labeled episodes (with actions and rewards) that are not naturally
available for many problems (e.g., Uehara et al., 2021). Our goal instead is
to pre-train the decoder $\phi$ using video data which is more easily
available.

##### Problem Statement.

Here, we state our problem formally. We are given access to two
hyperparameters $\varepsilon>0$ and $\delta\in(0,1)$ and a sufficiently large
dataset of videos. We are also given a decoder class
$\Phi=\\{\phi:\mathcal{X}\rightarrow[N]\\}$ containing decoders that map an
observation to one of the $N$ possible _abstract states_.333We use a discrete
representation to learn the latent state since Block MDP and Ex-Block MDP have
a finite state space, and a discrete representation suffices for this purpose.
During the pre-training phase, we learn a decoder $\phi\in\Phi$ using the
video data. We then freeze $\phi$ and use it to do RL in a downstream task.
Instead of using any particular choice of algorithm for RL, we assume we are
given a provably efficient tabular RL algorithm $\mathscr{A}$. We convert the
observation-based RL problem to a tabular MDP problem by converting
observations $x$ to its abstract state representation $\phi(x)$ using the
frozen learned decoder $\phi$. The algorithm $\mathscr{A}$ uses $\phi(x)$
instead of $x$ and outputs an _abstract policy_
$\varphi:[N]\rightarrow\mathcal{A}$. We want that
$\sup_{\pi\in\Pi}V(\pi)-V(\varphi\circ\phi)\leq\varepsilon$ with probability
at least $1-\delta$, where $\varphi\circ\phi:x\mapsto\varphi(\phi(x))$ is our
learned policy. Additionally, we want the number of online episodes needed in
the downstream RL task to not scale with the size of the decoder class $\Phi$.
This allows us to minimize online episodes while using more naturally
available video data for pre-training.

###  3 Representation Learning for RL using Video Dataset

We assume access to a dataset $\mathcal{D}$ of $n$ videos
$\mathcal{D}=\\{(x^{(i)}_{1},x^{(i)}_{2},\cdots,x^{(i)}_{H})\\}_{i=1}^{n}$
where $x^{(i)}_{j}$ is the $j^{th}$ observation (or frame) of the $i^{th}$
video. We are provided a decoder class
$\Phi=\\{\phi:\mathcal{X}\rightarrow[N]\\}$, and our goal is to learn a
decoder $\phi\in\Phi$ that captures task-relevant information in the
underlying state $\phi^{\star}(x)$ while throwing away as much exogenous noise
as possible. Instead of proposing a new algorithm, we consider the following
three classes of well-known representation learning methods that only assume
access to video data. Our goal is to understand whether these methods provably
learn useful representations.

We define the loss function for these approaches that we use for our
theoretical analysis. In practice, these methods either use continuous
representations or use a Vector Quantized bottleneck to model discrete
representations. We discuss these practical details at the end of this
section.

##### Autoencoder.

This approach first maps a given observation $x$ to an abstract state
$\phi(x)$ using a decoder $\phi\in\Phi$, and then uses it to reconstruct the
observation $x$ with the aid of a reconstruction model $z\in\mathcal{Z}$ where
$\mathcal{Z}=\left\\{z:[N]\rightarrow\mathcal{X}\right\\}$ is our
reconstruction model class. Formally, we optimize the following loss:

$\ell_{\textrm{auto}}(z,\phi)=\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\|z(\phi(x^{(i)}_{h}))-x^{(i)}_{h}\|^{2}_{2}.$
(1)

In practice, autoencoders are typically implemented using a Vector Quantized
bottleneck trained in a Variational AutoEncoder manner, which is called the
VQ-VAE approach (Oord et al., 2017).

##### Forward Modeling.

This approach is similar to the autoencoder approach but instead of
reconstructing the input observation, we reconstruct a future observation
using a model class
$\mathcal{F}=\left\\{f:[N]\times[K]\rightarrow\Delta(\mathcal{X})\right\\}$
where $N$ is the output size of the decoder class $\Phi$ and $K\in\mathbb{N}$
is a hyperparameter representing the forward time steps from the current
observation. We collect a dataset of _multistep transitions_
$\mathcal{D}_{\textrm{for}}=\\{(x^{(i)},k^{(i)},x^{\prime(i)})\\}_{i=1}^{n}$
sampled iid using the video dataset $\mathcal{D}$ where the observation
$x^{(i)}$ is sampled randomly from the $i^{th}$ video, $k^{(i)}\in[K]$, and
$x^{\prime(i)}$ is the frame $k^{(i)}$-steps ahead of $x^{(i)}$ in the
$i^{th}$ video. We distinguish between two types of sampling procedures, one
where $k^{(i)}$ is always a fixed given value $k\in[K]$, and one where
$k^{(i)}\sim\textrm{Unf}\left([K]\right)$. Given the dataset
$\mathcal{D}_{\textrm{for}}$, we optimize the following loss:

$\ell_{\textrm{for}}(f,\phi)=\frac{1}{n}\sum_{i=1}^{n}\ln
f\left(x^{\prime(i)}\mid\phi(x^{(i)}),k^{(i)}\right).$ (2)

##### Temporal Contrastive Learning.

Finally, this approach trains the decoder $\phi$ to learn to separate a pair
of temporally causal observations from a pair of temporally _acausal_
observations. We collect a dataset of
$\mathcal{D}_{\textrm{temp}}=\\{(x^{(i)},k^{(i)},x^{\prime(i)},z^{(i)})\\}_{i=1}^{\lfloor
n/2\rfloor}$ tuples using the multistep transitions dataset
$\mathcal{D}_{\textrm{for}}.$ We use 2 multistep transitions to create a
single datapoint for $\mathcal{D}_{\textrm{temp}}$ to keep the datapoints
independent. To create the $i^{th}$ datapoint for
$\mathcal{D}_{\textrm{temp}}$, we use the multistep transitions
$(x^{(2i)},k^{(2i)},x^{\prime(2i)})$ and
$(x^{(2i+1)},k^{(2i+1)},x^{\prime(2i+1)})$ and sample
$z^{(i)}\sim\textrm{Unf}(\\{0,1\\})$. If $z^{(i)}=1$, then our $i^{th}$
datapoint is a causal observation pair
$(x^{(2i)},k^{(2i)},x^{\prime(2i)},z^{(i)})$, otherwise, it is an acausal
observation pair $(x^{(2i)},k^{(2i)},x^{(2i+1)},z^{(i)})$. Depending on how we
sample $k$, we collect a different multistep transition dataset
$\mathcal{D}_{\textrm{for}}$, and accordingly a different contrastive learning
dataset $\mathcal{D}_{\textrm{temp}}$. Given the dataset
$\mathcal{D}_{\textrm{temp}}$, we optimize the following loss using a
regression model $g$ belonging to a model class
$\mathcal{G}=\\{g:\mathcal{X}\times[K]\times\mathcal{X}\rightarrow[0,1]\\}$:

$\ell_{\textrm{temp}}(g,\phi)=\frac{1}{\lfloor n/2\rfloor}\sum_{i=1}^{\lfloor
n/2\rfloor}\left(z^{(i)}-g(\phi(x^{(i)}),k^{(i)},x^{\prime(i)})\right)^{2}.$
(3)

##### Practical Implementations.

We use the description of the aforementioned methods for theoretical analysis.
However, their practical implementations differ in a few notable ways.
Firstly, we either use a continuous vector representation
$\phi:\mathcal{X}\rightarrow\mathbb{R}^{d}$ for modeling the decoder class
$\Phi$, or apply a Vector Quantized (VQ) bottleneck on top of the vector
representation to model a discrete-representation decoder. The VQ bottleneck
is applied on the final output of $\phi(x)$ to force the representations into
a discrete state space (Oord et al., 2017). The VQ-bottleneck is expressive
yet forces the model to remove excess information (Liu et al., 2021). The
gradients into $\phi$ are produced using the straight-through estimator, while
an extra loss attracts the output of $\phi$ to the discrete embeddings.

We also optimize all loss functions described above using the Adam optimizer
with mini-batches (instead of directly minimizing the loss in Equation 1-3).
Finally, we use square loss instead of log loss for forward modeling and
SimCLR-style batched loss for temporal contrastive learning.

There are several variations of the above methods in prior literature. Our
goal is to not comprehensively cover all such variants but instead describe
the core approach and analyze it theoretically. We will empirically show that
the findings of our theoretical analysis apply to our practical
implementations.

###  4 Is Video Based Representation Learning Provably Correct?

In this section, we establish the theoretical foundation for when video-based
representation learning is useful for downstream RL. We first present upper
bounds for video representation pre-training via forward modeling and temporal
contrastive learning in Block MDPs. We then present a lower-bound construction
that shows that the video-based representation pre-training setting is
exponentially harder than representation pre-training using trajectories where
actions are available. This also explains why learning from video data is
challenging in practice (Zhao et al., 2022). For most of our results in this
section, we will defer the proof to the Appendix and only provide a sketch
here.

####  4.1 Upper Bound in Block MDP Setting

We start by stating our theoretical setting and our main assumptions.

##### Theoretical Setting.

We assume a Block MDP setting and access to a dataset
$\mathcal{D}=\left\\{(x^{(i)}_{1},x^{(i)}_{2},\cdots,x^{(i)}_{H})\right\\}_{i=1}^{n}$
of $n$ independent and identically distributed (iid) videos sampled from data
distribution $D$. We denote the probability of a video as
$D(x_{1},x_{2},\cdots,x_{H})$. We assume that $D$ is generated by a mixture of
Markovian policies $\Pi_{D}$, i.e., the generative procedure for $D$ is to
sample a policy $\pi\in\Pi_{D}$ with some probability and then generate an
entire episode using it. We assume that observations encode time steps. This
can be trivially accomplished by simply concatenating the time step
information to the observation. We also assume that the video data has good
state space coverage and that the data is collected by _noise-free policies._

######  Assumption 1 (Requirements on Data Collection).

There exists an $\eta_{\textrm{min}}>0$ such that if $s$ is a state reachable
at time step $h$ by some policy in $\Pi$, then
$D\left(\phi^{\star}(x_{h})=s\right)\geq\eta_{\textrm{min}}$. Further, we
assume that every data collection policy $\pi\in\Pi_{D}$ is noise-free, i.e.,
$\pi(a\mid x_{h})=\pi(a\mid\phi^{\star}(x_{h}))$ for all $(a,x_{h})$.

Justification for Assumption 1 In practice, we expect this assumption to hold
for tasks such as gaming, or software debugging, where video data is abundant
and, therefore, can be expected to provide good coverage of the underlying
state space. This assumption is far weaker than the assumption in batch RL
which also requires actions and rewards to be labeled, which makes it more
expensive to collect data that has good coverage (Chen & Jiang, 2019).
Further, unlike imitation learning from observations (ILO) (Torabi et al.,
2019), we don’t require that these videos provide demonstrations of the
desired behavior. E.g., video streaming of games is extremely common on the
internet, and one can get many hours of video data this way. However, this
data wouldn’t come with actions (which will be mouse or keyboard strokes) or
reward labeling, and the game levels or tasks that these players are solving
can be all different or even unrelated to the downstream tasks we want to
solve. As such, the video data do not provide demonstrations of any one
desired task. Further, as the video data is typically generated by humans, we
can expect the data collection policies to be noise-free, as these policies
are realized by humans who would not make decisions based on noise. E.g., a
human player is unlikely to turn left due to the background motion of leaves
that is unrelated to the game’s control or objective.

We analyze the temporal contrastive learning and forward modeling approaches
and derive upper bounds for these methods in Block MDPs. While autoencoder-
based approaches sometimes do well in practice, it is an open question whether
finite-sample bounds exist for them and we leave their theoretical analysis to
future work and instead evaluate them empirically. We also do not use a VQ
bottleneck for our theoretical analysis since our decoder class directly
outputs discrete values. We can view the VQ bottleneck as a way to do the
discrete encoding in practice. In addition to the decoder class $\Phi$, we
assume a function class $\mathcal{F}$ to model $f$ for forward modeling and
$\mathcal{G}$ to model $g$ for temporal contrastive learning. We make a
realizability assumption on these function classes.

######  Assumption 2 (Realizability).

There exists $f^{\star}\in\mathcal{F},g^{\star}\in\mathcal{G}$ and
$\phi_{\textrm{for}},\phi_{\textrm{temp}}\in\Phi$ such that
$f^{\star}(X^{\prime}\mid\phi_{\textrm{for}}(x),k)=\mathbb{P}_{\textrm{for}}(X^{\prime}\mid
x,k)$ and
$g^{\star}(z\mid\phi_{\textrm{temp}}(x),k,x^{\prime})=\mathbb{P}_{\textrm{temp}}(z=1\mid
x,k,x^{\prime})$ on the appropriate support, and where
$\mathbb{P}_{\textrm{for}}$ and $\mathbb{P}_{\textrm{temp}}$ are respectively
the Bayes classifier for the forward modeling and temporal contrastive
learning methods.

Justification for Assumption 2. Realizability is a typical assumption made in
theoretical analysis of RL algorithms (Agarwal et al., 2020). Intuitively, the
assumption states that the function classes are expressive enough to represent
the Bayes classifier of their problem. In practice, this is usually not a
concern as we will use expressive deep neural networks to model these function
classes. We will empirically show the feasibility of this assumption in our
experiments.

Our statement of Assumption 2 implicitly makes use of Assumption 1 which
enables us to write the Bayes classifier for these problems as a function of
just the state $\phi^{\star}(x)$ rather then the actual observation $x$. This
is formally proven in our analysis. This allows us to make a realizability
assumption while still assuming a discrete decoder class $\Phi$.

Finally, we assume that our data distribution has the required information to
separate the latent states. We state this assumption formally below and then
show settings where this is true.

######  Assumption 3 (Margin Assumption).

We assume that the margins $\beta_{\textrm{for}}$ and $\beta_{\textrm{temp}}$
defined below:

$\displaystyle\beta_{\textrm{for}}$
$\displaystyle=\inf_{s_{1},s_{2}\in{\mathcal{S}},s_{1}\neq
s_{2}}\mathbb{E}_{k}\left[\left\|\mathbb{P}_{\textrm{for}}(X^{\prime}\mid
s_{1},k)-\mathbb{P}_{\textrm{for}}(X^{\prime}\mid
s_{2},k)\right\|_{\textrm{TV}}\right]$ $\displaystyle\beta_{\textrm{temp}}$
$\displaystyle=\inf_{s_{1},s_{2}\in{\mathcal{S}},s_{1}\neq
s_{2}}\frac{1}{2}\mathbb{E}_{k,s^{\prime}}\left[\left|\mathbb{P}_{\textrm{temp}}(z=1\mid
s_{1},k,s^{\prime})-\mathbb{P}_{\textrm{temp}}(z=1\mid
s_{2},k,s^{\prime})\right|\right],$

are strictly positive, and where in the definition of $\beta_{\textrm{temp}}$,
we sample $s^{\prime}$ from the video data distribution and $k$ is sampled
according to our data collection procedure.

Justification for Assumption 3. This assumption states that we need margins
($\beta_{\textrm{for}}$) for forward modeling and ($\beta_{\textrm{temp}})$
for temporal contrastive learning. A common scenario where these assumptions
are true is when for any pair of different states $s_{1},s_{2}$, there is a
third state $s^{\prime}$ that is reachable from one but not the other. If the
video data distribution $D$ supports all underlying transitions, then this
immediately implies that $\left\|\mathbb{P}_{\textrm{for}}(X^{\prime}\mid
s_{1},k)-\mathbb{P}_{\textrm{for}}(X^{\prime}\mid
s_{2},k)\right\|_{\textrm{TV}}>0$ which implies $\beta_{\textrm{for}}>0$. This
scenario occurs in almost all navigation tasks. Specifically, it occurs in the
three domains we experiment with. While it is less clear, under this
assumption we also have $\beta_{\textrm{temp}}>0$.

We now state our main result for forward modeling under Assumption 1-3.

######  Theorem 1 (Forward Modeling Result).

Fix $\varepsilon>0$ and $\delta\in(0,1)$ and let $\mathscr{A}$ be any provably
efficient RL algorithm for tabular MDPs with sample complexity
$n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)$. If $n$ is ${\tt
poly}\left\\{S,H,\nicefrac{{1}}{{\eta_{\textrm{min}}}},\nicefrac{{1}}{{\beta_{\textrm{for}}}},\nicefrac{{1}}{{\varepsilon}},\ln(\nicefrac{{1}}{{\delta}}),\right.$
$\left.\ln|\mathcal{F}|,\ln|\Phi|\right\\}$ for a suitable polynomial, then
forward modeling learns a decoder
$\widehat{\phi}:\mathcal{X}\rightarrow[|{\mathcal{S}}|]$. Further, running
$\mathscr{A}$ on the tabular MDP with
$n_{\textrm{samp}}(S,A,H,T,\varepsilon/2,\delta/4)$ episodes returns a latent
policy $\widehat{\varphi}$. Then there exists a bijective mapping
$\alpha:{\mathcal{S}}\rightarrow[|{\mathcal{S}}|]$ such that with probability
at least $1-\delta$ we have:

$\forall s\in{\mathcal{S}},\qquad\mathbb{P}_{x\sim q(\cdot\mid
s)}\left(\widehat{\phi}(x)=\alpha(s)\mid\phi^{\star}(x)=s\right)\geq
1-\frac{4S^{3}H^{2}}{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}\sqrt{\frac{1}{n}\ln\left(\frac{|\mathcal{F}|\cdot|\Phi|}{\delta}\right)},$

and the learned observation-based policy
$\widehat{\varphi}\circ\widehat{\phi}:x\mapsto\widehat{\varphi}(\widehat{\phi}(x))$
is $\varepsilon$-optimal, i.e.,

$V(\pi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})\leq\varepsilon.$

Finally, the number of online episodes used in the downstream RL task is given
by $n_{\textrm{samp}}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)$ and
doesn’t scale with the complexity of function classes $\Phi$ and
$\mathcal{F}$.

The result for temporal contrastive is identical to Theorem 1 but instead of
$\beta_{\textrm{for}}$ we have $\beta_{\textrm{temp}}$ and instead of
$\mathcal{F}$ we have $\mathcal{G}$. These upper bounds provide the desired
result which shows that not only can we learn the right representation and
near-optimal policy but also do so without online episodes scaling with
$\ln|\Phi|$.

Theorem 1 shows that in the absence of exogenous noise, we can expect forward
modeling and temporal contrastive learning to learn the right representation
under Assumption 1-3. The upper bounds for these two approaches are nearly
identical except for the difference in margins $\beta_{\textrm{temp}}$ vs
$\beta_{\textrm{for}}$, and the difference in function class. Typically, we
expect the function class $\mathcal{F}$ for forward modeling to have a higher
statistical complexity than temporal contrastive learning $\mathcal{G}$, as
the former is modeling the entire observation compared to the latter which is
solving a binary classification task. This leaves the comparison between their
margins. We prove a result that relates their margin.

######  Theorem 2 (Margin Relation).

For any Block MDP and $K\in\mathbb{N}$ we have
$\frac{\eta_{\textrm{min}}^{2}}{4H^{2}}\beta_{\textrm{for}}^{(k)}\leq\beta_{\textrm{temp}}^{(k)}\leq\beta_{\textrm{for}}^{(k)}$.

The proof of a generalized statement of this theorem can be found in Appendix
B.4. This result shows that the margin for forward modeling
$\beta_{\textrm{for}}$ is at least as large as the margin for temporal
contrastive approach $\beta_{\textrm{temp}}$. In summary, this leads to a
trade-off between forward modeling and temporal contrastive learning
approaches in terms of performance, where the former has a function class with
a higher statistical complexity but has a higher margin, compared to temporal
contrastive approaches.

####  4.2 Learning from Video is Exponentially Harder Than Learning from
Trajectory Data

As discussed before, when online reinforcement learning (RL) is possible, some
algorithms can extract latent states from observation. These algorithms can be
used to extract a decoder $\widehat{\phi}$ that with high probability learns
to map observations from different states to different values and whose output
range $N$ is roughly of the order $|{\mathcal{S}}|$.

This begs the following question: _Is it always possible to learn a compact
state representation that is useful for control from video data?_ From this
following lower bound result, we argue that this is not always the case.

######  Theorem 3 (Lower Bound for Video).

Suppose $|{\mathcal{S}}|,|\mathcal{A}|,H\geq 2$. Then, for any
$\varepsilon\in(0,1)$, any algorithm $\mathscr{A}_{1}$ that outputs a state
decoder $\phi$ with $\phi_{h}:\mathcal{X}\to[L],~{}L\leq
2^{\nicefrac{{1}}{{4\varepsilon}}-1},~{}\forall h\in[H]$ given a video dataset
$\mathcal{D}$ sampled from some MDP and satisfies Assumption 1, and any online
RL algorithm $\mathscr{A}_{2}$ uses that state decoder $\phi$ in its
interaction with such an MDP (i.e., $\mathscr{A}_{2}$ only observes states
through $\phi$) and output a policy $\widehat{\pi}$, there exists an MDP
instance $M$ in a class of MDPs which satisfies Assumption 3 and is PAC
learnable with $\widetilde{O}({\tt
poly}(|{\mathcal{S}}|,|\mathcal{A}|,H,\nicefrac{{1}}{{\varepsilon}}))$
complexity, such that

$\displaystyle V_{M}(\pi_{M}^{\star})-V_{M}(\widehat{\pi})>\varepsilon,$

regardless of the size of the video dataset $\mathcal{D}$ for algorithm
$\mathscr{A}_{1}$ and the number of episodes of interaction for algorithm
$\mathscr{A}_{2}$.

The basic idea behind that hard instance construction is that, without the
action information, it is impossible for the learning agent to distinguish
between endogenous states and exogenous noise. For example, consider an image
consisting of $N\times N$ identical mazes but where the agent controls just
one maze. Other mazes contain other agents which are exogenous for our
purpose. In the absence of actions, we cannot tell which maze is the one we
are controlling and must memorize the configuration of all $N\times N$ mazes
which grow exponentially with $N$. Another implication from that hard instance
is – if the margin condition (Assumption 3) is violated, the exponentially
large state decoder is also required for the regular block MDP without
exogenous noise; a detailed discussion can also be found in Section B.3.

##### Can we get efficient learning under additional assumptions?

Our lower bound suggests that one can in general not learn efficient and
correct representations with just video data. However, it may be possible in
some cases to do so with an additional assumption. We highlight one example
here and defer a proper formal analysis to future work. One such scenario can
be when the best-in-class error:
$\inf_{f,\phi}\mathbb{E}_{x,k}\left[\left\|f(X^{\prime}\mid
x,k)-\mathbb{P}_{\textrm{for}}(X^{\prime}\mid
x,k)\right\|_{\textrm{TV}}\right]$ is small. A domain where this can happen is
when the endogenous state is more predictive of $x^{\prime}$ than any other
$\ln|{\mathcal{S}}|$ bits of information in $x$. E.g., in a navigation domain,
there can be many people in the background, but memorizing all of them can
easily overwhelm the decoder’s model capacity. Focusing solely on modeling the
agent’s state can simplify the task of predicting the future.

Recently some approaches have also considered recovering _latent actions_ from
video data using an encoder-decoder approach (Ye et al., 2022). In general,
the lower bound in Theorem 3 applies to these methods and they do not provably
work in the hard instances with exogenous noise. For example, the latent
actions can capture _exogenous noise_ instead of actions, if the former is
more predictive of changes in the observations. However, in simpler cases such
as 3D games, where the agent’s action is typically most predictive of changes
in observations, or in settings with no exogenous noise, one can expect these
approaches to do well.

###  5 Experimental Results and Discussion

We empirically evaluate the above video-based representation learning methods
on three visual domains: (1) a 2D GridWorld domain, (2) ViZDoom basic
environment, a first-person 3D shooting game, and (3) ViZDoom Defend The
Center environment, which is a more challenging shooting game. Additional
experimental details and results are mentioned in Appendix C. Our main goal is
to validate our theoretical findings from Section 4 and justify our
assumptions.

Specifically, we study the following questions in our experimental analysis:

  * •

How do autoencoding, forward modeling, and temporal contrastive learning
methods perform in the absence of any noise, in the presence of IID noise, and
the presence of exogenous noise?

  * •

How does the performance of forward modeling vary with harder exogenous noise?

  * •

How do these different video representation learning methods compare to
trajectory-based learning methods such as ACRO?

####  5.1 Experimental Details

##### GridWorld.

We consider navigation in a $12\times 12$ Minigrid environment (Chevalier-
Boisvert et al., 2023). The agent is represented as a red triangle and can
take three actions: move forward, turn left, and turn right (Fig. 3). The goal
of the agent is to reach a yellow key. The position of the agent and key
randomizes each episode. The agent only observes an area around itself (as an
agent-centric-view). Horizon $H=12$, and the agent gets a reward of +1.0 for
reaching the goal and -0.01 in other cases.

##### ViZDoom Basic.

We use the basic ViZDoom environment (Wydmuch et al., 2018; Kempka et al.,
2016), which is a first-person shooting game (Fig. 5). The player needs to
kill a monster to win. The episode ends when the monster is killed or after
300 steps. The map of the environment is a rectangle with gray walls, ceiling,
and floor. The player is spawned along the longer wall in the center. A red,
circular monster is spawned randomly somewhere along the opposite wall. The
player can take one of three actions at each time step (left, right, shoot).
One hit is enough to kill the monster. The episode finishes when the monster
is killed or on timeout. The reward scheme is as follows: +101 for shooting
the enemy , -1 per time step, and -5 for missed shots.

##### VizDoom Defend The Center.

We use an additional ViZDoom Defend the Center environment (Wydmuch et al.,
2018; Kempka et al., 2016), which is a first-person shooting game (Fig. 7).
The player needs to kill a variety of monsters to score in the game. The
episode ends when the monster is killed or after 500 steps. Results for this
environment are shown in Fig. 6 and Fig. 7 and further validate our findings
from theory and experiments.

##### Exogenous Noise.

For all domains, the observation is an RGB image. We add exogenous noise to it
by superimposing 10 generated diamonds of a particular size. The color and
position of these diamonds are our exogenous state. At the start of each
episode, we randomly generate these diamonds, after which they move in a
deterministic path. We also test the setting in which there is exogenous noise
in the reward. We compute a score based on just the exogenous noise and add it
to the reward presented to the agent. However, the agent is still evaluated on
the original reward.

##### Model and Learning.

Our decoder class $\Phi$ is a convolutional neural network. We use a
deconvolutional neural network to model $f$ and $h$. We experimented with both
using a vector representation for $\phi$ and also using a VQ-bottleneck to
discretize the embeddings. We use PPO to do downstream RL and keep $\phi$
frozen during the RL training. We also visualize the learned representations
by training a decoder on them and fixing $\phi$ to reconstruct the input
observations. We then look at the generated images to see what information
from the observation is preserved by the representation.

##### ACRO.

We also evaluate the learned representations against ACRO (Islam et al., 2022)
which uses trajectory data. This approach learns representation $\phi$ by
predicting action given a pair of observations $\mathbb{E}\left[\ln
p(a_{h}\mid\phi(x_{h}),x_{h+k})\right]$. ACRO is designed to filter out
exogenous noise as this information is not predictive of the action. Our goal
is to show how much better do representations get if we had trajectory data
instead of video data.

(a) No Noise

(b) Only Observation Noise

(c) Only Reward Noise

(d) Both

Figure 2: RL experiments in the GridWorld environment.

(a) Original

(b) Forward Model

(c) Autoencoder

(d) Contrastive

Figure 3: Decoded image reconstructions from different latent representation
learning methods in the GridWorld environment. We train a decoder on top of
frozen representations trained with the three video pre-training approaches.
Top row: shows an example from the setting where there is no exogenous noise.
Bottom row: shows an example with exogenous noise.

####  5.2 Empirical Results and Discussion

We present our main empirical results in Fig. 2, Fig. 4, and Fig. 6 and
discuss the results below.

(a) No Noise

(b) Only Observation Noise

(c) Only Reward Noise

(d) Both

Figure 4: RL experiments using different latent representations for the
ViZDoom environment.

(a) Original

(b) Forward Modeling

(c) Autoencoder

(d) Contrastive

Figure 5: Decoded image reconstructions from different representation learning
methods in the ViZDoom environment. We train a decoder on top of frozen
representations trained with the three video pre-training approaches. Here, we
show an example where we also have exogenous noise.

##### Forward modeling and temporal contrastive both work when there is no
exogenous noise.

In accordance with Theorem 1, we observe that in the case of both GridWorld
(Fig. 2) and ViZDoom environments (Fig. 4, Fig. 6), these approaches learn a
decoder $\phi$ that lead to success with reinforcement learning in the absence
of any exogenous noise (both in the presence and absence of exogenous reward).
For GridWorld, we find support for this result with VQ bottleneck during
representation learning (Fig. 2(a),(c) and Fig. 6(a), (c)), whereas for
ViZDoom, we find support for this result even without the use of a VQ
bottleneck (Fig. 4(a),(c)). These results are further supported via
qualitative evaluation through image decoding from the learned latent
representations (Fig. 3) which show that these representations can recover
critical elements like walls. We find that autoencoder performs well in
ViZDoom but not in gridworld, which aligns with a lack of any theoretical
understanding of autoencoders.

(a) No Noise

(b) Only Observation Noise

(c) Only Reward Noise

(d) Both

Figure 6: RL experiments using different latent representations for the
ViZDoom Defend the Center environment.

(a)

(b)

(c)

(d)

(e)

(f)

(g)

(h)

(i) Original

(j) Forward Modeling

(k) Autoencoder

(l) Temporal Contrastive

Figure 7: Decoded image reconstructions from different latent representation
learning methods in the ViZDoom Defend the Center environment. We train a
decoder on top of frozen representations trained with the three video pre-
training approaches.

##### Performance with I.I.D. Noise.

We evaluate iid noise in the gridworld domain. We use the diamond-shaped
exogenous noise that we used in Fig. 2, however, at each time step, we
randomly sample the color and position of each diamond, independent of the
agent’s history. Fig. 8(c) shows the result for forward modeling and Fig. 8(d)
shows the same for ACRO. We also ablate the number of noisy diamonds. As
expected, forward modeling and ACRO can learn a good policy while increase in
the number of noisy diamonds (num noise var) only slightly decreases their
performance. Additional results on performance with iid noise in the VizDoom
basic environment are in Appendix C.

(a) Forward Modeling

(b) ACRO

Figure 8: Experiments with iid noise for the Gridworld environment. ‘Num noise
var’ denotes the number of noisy diamonds constituting the exogenous noise.

##### Performance with exogenous noise.

We find that in the presence of exogenous noise for both GridWorld (Fig. 2)
and ViZDoom (Fig. 4, Fig. 6) environments, representations from forward
modeling continue to succeed at RL albeit with a much lower performance in
gridworld, whereas temporal contrastive representations completely fail. One
hypothesis for the stark failure of temporal contrastive learning is that the
agent can tell whether two observations are causal or not, by simply focusing
on the noisy rhombus that moves in a predictive manner. Therefore, the
contrastive learning loss can be reduced by focusing entirely on noise.
Whereas, forward modeling is more robust as it needs to predict future
observations, and the agent’s state is more helpful for doing that than noise.
This shows in the reconstructions (Fig. 3(d), Fig. 5(d), Fig. 7(d)). As
expected, the reconstructions for forward modeling continue to capture state-
relevant information, whereas for temporal contrastive they focus on noise and
miss relevant state information. We also formally prove that there exists an
instance where forward modeling can recover the latent state for low-levels of
exogenous noise, whereas temporal contrastive cannot do so for any level of
exogenous noise. We defer readers to the instance construction and analysis in
Appendix B.5.

(a) ACRO

(b) Forward Modeling

Figure 9: RL performance with varying size for exogenous noise in the
GridWorld environment.

##### Harder Exogenous Noise.

We increase the size of the exogenous noise variables (diamond shapes
overlayed on the image) in the gridworld domain while keeping the number of
exogenous variables fixed at 10 (Fig. 9). We also increase the number of
exogenous noise variables in the gridworld domain, while keeping their sizes
fixed at 4 pixels (Fig. 10). Both results show significant degradation in the
performance of the forward modeling approach. This supports one of our main
theoretical results that exogenous noise poses a challenge for video pre-
training. Additional results for other video-based representation learning
methods for this setting are shown in Fig. 12.

(a) ACRO

(b) Forward Modeling

Figure 10: Gridworld experiments with exogenous noise of size 4 for a varying
number of exogenous noise variables. Several representation learning
algorithms using video data struggle to learn with a larger number of
exogenous noise variables, whereas ACRO which uses trajectory data, still
performs well.

##### Comparison with ACRO.

Finally, we draw a comparison between the performance of video-pretrained
representation and ACRO which uses trajectory data. ACRO achieves the
strongest performance across all tasks (Fig. 2, Fig. 4, Fig. 6). This agrees
with our theoretical finding (Theorem 3) that the sample complexity of video-
based representations can be exponentially worse than trajectory-based
representations in the presence of exogenous noise. Additionally, we also
observe that as we increase the size of the exogenous noise elements in the
observation space (Fig. 9) or the number of exogenous noise variables (Fig.
10), performance for forward-modeling based representation learning degrades
more drastically compared to ACRO.

###  6 Conclusion

Learning models using unlabeled data such as audio, images, and video have had
a transformative impact on deep learning research. An intriguing question is
whether unlabeled data (videos) can have the same value for general
reinforcement learning problems. Our work has sought to address this question
both in theory and in practice. Theoretically, we have shown that certain
representations trained on video can succeed when there is no exogenous noise.
In the presence of exogenous noise, we present a lower bound establishing a
separation between action-labeled pre-training for reinforcement learning and
pre-training from unlabeled video data. We empirically validate our
theoretical results on three visual domains.

##### Acknowledgement.

We thank Sam Devlin, Ching-An Cheng, Andrey Kolobov, and Adith Swaminathan for
useful discussions.

### References

  * Agarwal et al. (2020) Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun.  Flambe: Structural complexity and representation learning of low rank mdps.  _Advances in neural information processing systems_ , 33:20095–20107, 2020.
  * Aubret et al. (2023) Arthur Aubret, Markus R. Ernst, Céline Teulière, and Jochen Triesch.  Time to augment self-supervised visual representation learning.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=o8xdgmwCP8l.
  * Baker et al. (2022) Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune.  Video pretraining (vpt): Learning to act by watching unlabeled online videos.  _Advances in Neural Information Processing Systems_ , 35:24639–24654, 2022.
  * Bharadhwaj et al. (2022) Homanga Bharadhwaj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine.  Information prioritization through empowerment in visual model-based rl.  _arXiv preprint arXiv:2204.08585_ , 2022.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Chen & Jiang (2019) Jinglin Chen and Nan Jiang.  Information-theoretic considerations in batch reinforcement learning.  In _International Conference on Machine Learning_ , pp. 1042–1051. PMLR, 2019.
  * Chevalier-Boisvert et al. (2023) Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry.  Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks.  _CoRR_ , abs/2306.13831, 2023.
  * Du et al. (2019) 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.
  * Efroni et al. (2022) Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Langford.  Provably filtering exogenous distractors using multistep inverse dynamics.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=RQLLzMCefQu.
  * Escontrela et al. (2023) Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Youngwoon Lee, Danijar Hafner, and Pieter Abbeel.  Video prediction models as rewards for reinforcement learning.  _arXiv preprint arXiv:2305.14343_ , 2023.
  * Geer (2000) Sara A Geer.  _Empirical Processes in M-estimation_ , volume 6.  Cambridge university press, 2000.
  * Goo & Niekum (2019) Wonjoon Goo and Scott Niekum.  One-shot learning of multi-step tasks from observation via activity localization in auxiliary video.  In _2019 international conference on robotics and automation (ICRA)_ , pp. 7755–7761. IEEE, 2019.
  * Guo et al. (2022) Zhaohan Guo, Shantanu Thakoor, Miruna Pîslar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al.  Byol-explore: Exploration by bootstrapped prediction.  _Advances in neural information processing systems_ , 35:31855–31870, 2022.
  * Hafner et al. (2019) Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  _arXiv preprint arXiv:1912.01603_ , 2019.  URL https://arxiv.org/pdf/1912.01603.pdf.
  * Hafner et al. (2023) Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.  Mastering diverse domains through world models.  _arXiv preprint arXiv:2301.04104_ , 2023.
  * Islam et al. (2022) Riashat Islam, Manan Tomar, Alex Lamb, Yonathan Efroni, Hongyu Zang, Aniket Didolkar, Dipendra Misra, Xin Li, Harm van Seijen, Remi Tachet des Combes, et al.  Agent-controller representations: Principled offline rl with rich exogenous information.  _arXiv preprint arXiv:2211.00164_ , 2022.
  * Ke et al. (2019) Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra.  Learning dynamics model in reinforcement learning by incorporating the long term future, 2019.
  * Kearns & Singh (2002) Michael Kearns and Satinder Singh.  Near-optimal reinforcement learning in polynomial time.  _Machine learning_ , 49:209–232, 2002.
  * Kempka et al. (2016) Michał Kempka, Marek Wydmuch, Grzegorz Runc, Jakub Toczek, and Wojciech Jaśkowski.  Vizdoom: A doom-based ai research platform for visual reinforcement learning.  In _2016 IEEE conference on computational intelligence and games (CIG)_ , pp. 1–8. IEEE, 2016.
  * Lamb et al. (2022) Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford.  Guaranteed discovery of controllable latent states with multi-step inverse models.  _arXiv preprint arXiv:2207.08229_ , 2022.  URL https://arxiv.org/pdf/2207.08229.pdf.
  * Lin et al. (2021) Xudong Lin, Gedas Bertasius, Jue Wang, Shih-Fu Chang, Devi Parikh, and Lorenzo Torresani.  Vx2text: End-to-end learning of video-based text generation from multimodal inputs.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 7005–7015, 2021.
  * Liu et al. (2021) Dianbo Liu, Alex M Lamb, Kenji Kawaguchi, Anirudh Goyal ALIAS PARTH GOYAL, Chen Sun, Michael C Mozer, and Yoshua Bengio.  Discrete-valued neural communication.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 2109–2121. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper_files/paper/2021/file/10907813b97e249163587e6246612e21-Paper.pdf.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  Roberta: A robustly optimized bert pretraining approach.  _arXiv preprint arXiv:1907.11692_ , 2019.
  * Mhammedi et al. (2023) Zakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin.  Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation rl.  _arXiv preprint arXiv:2304.05889_ , 2023.
  * Micheli et al. (2023) Vincent Micheli, Eloi Alonso, and François Fleuret.  Transformers are sample-efficient world models.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=vhFu1Acb0xb.
  * Misra et al. (2020) Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford.  Kinematic state abstraction and provably efficient rich-observation reinforcement learning.  In _International conference on machine learning_ , pp. 6961–6971. PMLR, 2020.
  * Nachum et al. (2018) Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine.  Near-optimal representation learning for hierarchical reinforcement learning.  In _International Conference on Learning Representations_ , 2018.
  * Oord et al. (2017) Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu.  Neural discrete representation learning.  _arXiv preprint arXiv:1711.00937_ , 2017.
  * Parthasarathy et al. (2022) Nikhil Parthasarathy, SM Eslami, João Carreira, and Olivier J Hénaff.  Self-supervised video pretraining yields strong image representations.  _arXiv preprint arXiv:2210.06433_ , 2022.
  * Radford et al. (2021) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pp. 8748–8763. PMLR, 2021.
  * Sikchi et al. (2022) Harshit Sikchi, Akanksha Saran, Wonjoon Goo, and Scott Niekum.  A ranking game for imitation learning.  _arXiv preprint arXiv:2202.03481_ , 2022.
  * Sobal et al. (2022) Vlad Sobal, Jyothir SV, Siddhartha Jalagam, Nicolas Carion, Kyunghyun Cho, and Yann LeCun.  Joint embedding predictive architectures focus on slow features.  _arXiv preprint arXiv:2211.10831_ , 2022.
  * Srivastava et al. (2015) Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov.  Unsupervised learning of video representations using lstms.  In Francis Bach and David Blei (eds.), _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pp. 843–852, Lille, France, 07–09 Jul 2015\. PMLR.  URL https://proceedings.mlr.press/v37/srivastava15.html.
  * Sutton & Barto (2018) Richard S Sutton and Andrew G Barto.  _Reinforcement learning: An introduction_.  MIT press, 2018.
  * Tang et al. (2023) Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chandak, Remi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, András György, Shantanu Thakoor, Will Dabney, Bilal Piot, Daniele Calandriello, and Michal Valko.  Understanding self-predictive learning for reinforcement learning.  In _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 33632–33656. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/tang23d.html.
  * Torabi et al. (2019) Faraz Torabi, Garrett Warnell, and Peter Stone.  Recent advances in imitation learning from observation.  _arXiv preprint arXiv:1905.13566_ , 2019.  URL https://arxiv.org/pdf/1905.13566.pdf.
  * Uehara et al. (2021) Masatoshi Uehara, Xuezhou Zhang, and Wen Sun.  Representation learning for online and offline rl in low-rank mdps.  _arXiv preprint arXiv:2110.04652_ , 2021.
  * Wang et al. (2022) Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian.  Denoised mdps: Learning world models better than the world itself.  _arXiv preprint arXiv:2206.15477_ , 2022.  URL https://arxiv.org/pdf/2206.15477.pdf.
  * Wydmuch et al. (2018) Marek Wydmuch, Michał Kempka, and Wojciech Jaśkowski.  Vizdoom competitions: Playing doom from pixels.  _IEEE Transactions on Games_ , 11(3):248–259, 2018.
  * Ye et al. (2022) Weirui Ye, Yunsheng Zhang, Pieter Abbeel, and Yang Gao.  Become a proficient player with limited data through watching pure videos.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Zhang et al. (2020) Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine.  Learning invariant representations for reinforcement learning without reconstruction.  _arXiv preprint arXiv:2006.10742_ , 2020.
  * Zhao et al. (2022) Tony Zhao, Siddharth Karamcheti, Thomas Kollar, Chelsea Finn, and Percy Liang.  What makes representation learning from videos hard for control?  _RSS Workshop on Scaling Robot Learning_ , 2022.  URL https://api.semanticscholar.org/CorpusID:252635608.

##  Part I Appendix

###  Appendix A Additional Related Work

##### Representation Learning for Reinforcement Learning

A line of research on recurrent state space models is essentially concerned
with the next-frame approach, although typically with conditioning on actions.
Moreover, to model uncertainty in the observations, a latent variable with a
posterior depending on the current observation (or even a sequence of future
observations) is typically introduced. (Ke et al., 2019) considered learning
such a sequential prediction model which predicts observations and conditions
on actions. They used a latent variable with a posterior depending on future
observations to model uncertainty. These representations were used for model-
predictive control and improved imitation learning. Dreamer (Hafner et al.,
2019, 2023) uses the next-frame objective but also conditions on actions. The
IRIS algorithm (Micheli et al., 2023) uses the next-frame objective but uses
the transformer architecture, again conditioning on actions. The InfoPower
approach (Bharadhwaj et al., 2022) combines a one-step inverse model with a
temporal contrastive objective. Sobal et al. (2022) explored using semi-
supervised objectives for learning representations in RL, yet used action-
labeled data. Wang et al. (2022) used a decoupled recurrent neural network
approach to learn to extract endogenous states, but relied on action-labeled
data to achieve the factorization. Deep Bisimulation for Control (Zhang et
al., 2020) introduced an objective to encourage observations with similar
value functions to map to similar representations.

Self-prediction methods such as BYOL-explore (Guo et al., 2022) proposed
learning reward-free representations for exploration, but depended on open-
loop prediction of future states conditioned on actions . An analysis paper
studied a simplified action-free version of the self-prediction objective
(Tang et al., 2023) and showed results in the absence of using actions,
although this has not been instantiated empirically to our knowledge.

A further line of work from theoretical reinforcemnt learning has examined
provably efficient objectives for discovering representations. Efroni et al.
(2022) explored representation learning in the presence of exogenous noise,
establishing a sample efficient algorithm. However Efroni et al. (2022) and
the closely related work on filtering exogenous noise required actions (Lamb
et al., 2022; Islam et al., 2022). Other theoretical work on learning
representations for RL has required access to action-labeled data (Misra et
al., 2020).

##### Representation Learning from Videos

Self-supervised representation learning from videos has a long history.
Srivastava et al. (2015) used recurrent neural networks with a pixel
prediction objective on future frames. Parthasarathy et al. (2022) explored
temporal contrastive objectives for self-supervised learning from videos. They
also found that the features learned well aligned with human perceptual
priors, despite the model not being explicitly trained to achieve such
alignment. Aubret et al. (2023) applied temporal contrastive learning to
videos of objects being manipulated in a 3D space, showing that this
outperformed standard augmentations used in computer vision.

##### Using Video Data for Reinforcement Learning

The VIPER method (Escontrela et al., 2023) uses a pre-trained autoregressive
generative model over action-free expert videos as a reward signal for
training an imitation learning agent. The Video Pre-training (VPT) algorithm
(Baker et al., 2022) trained an inverse kinematics model on a small dataset of
Minecraft videos and used the model to label a large set of unlabeled
Minecraft videos from the internet. This larger dataset was then used for
imitation learning and reinforcement learning for downstream tasks. Zhao et
al. (2022) explicitly studied the challenges in using videos for
representation learning in RL, identifying five key factors: task mismatch,
camera configuration, visual feature shift, sub-optimal behaviors in the data,
and robot morphology. Goo & Niekum (2019) learn reward functions for multi-
step tasks from videos by leveraging a single video segmented with action
labels (one-shot learning). Sikchi et al. (2022) propose a two-player ranking
game between a policy and a reward function to satisfy pairwise performance
rankings between behaviors. Their proposed method achieves state-of-the-art
sample efficiency and can solve previously unsolvable tasks in the learning
from observation (no actions) setting.

###  Appendix B Proofs of Theoretical Statements

We state our setting and general assumptions before presenting method specific
results. We also include a table of notations in Table 1.

Notation | Description
---|---
$[N]$ | Denotes the set $\\{1,2,\cdots,N\\}$
$\Delta(\mathcal{U})$ | Denotes the set of all distributions over a set $\mathcal{U}$
$\textrm{Unf}(\mathcal{U})$ | Uniform distribution over $\mathcal{U}$
$\textrm{supp}(\mathbb{P})$ | Support of a distribution $\mathbb{P}\in\Delta(\mathcal{U})$, i.e., $\textrm{supp}(\mathbb{P})=\\{x\in\mathcal{U}\mid\mathbb{P}(x)>0\\}$.
$\mathcal{X}$ | Observation space
${\mathcal{S}}$ | Latent endogenous state
$\mathcal{A}$ | Action space
$T:{\mathcal{S}}\rightarrow\mathcal{A}\rightarrow\Delta({\mathcal{S}})$ | Transition dynamics
$R:{\mathcal{S}}\times\mathcal{A}\rightarrow[0,1]$ | Reward function
$\mu$ | Start state distribution
$H$ | Horizon indicating the maximum number of actions per episode
$\phi^{\star}:\mathcal{X}\rightarrow{\mathcal{S}}$ | Endogenous state decoder
Table 1: Description for mathematical notations.

We are given a dataset
$\mathcal{D}=\left\\{(x^{(i)}_{1},x^{(i)}_{2},\cdots,x^{(i)}_{H})\right\\}_{i=1}^{n}$
of $n$ independent and identically distributed (iid) unlabeled episodes. We
will use the word video and unlabeled episodes interchangeably. We assume the
underlying data distribution is $D$. We denote the probability of an unlabeled
episode as $D(x_{1},x_{2},\cdots,x_{H})$. We assume that $D$ is generated by a
mixture of Markovian policies $\Pi_{D}$, i.e., the generative procedure for
$D$ is to sample a policy $\pi\in\Pi_{D}$ with probability $\Theta_{\pi}$ and
then generate an entire episode using it. For this reason, we will denote
$D=\Theta\circ\Pi_{D}$ where $\Theta$ is the mixture distribution. We assume
_no direct knowledge_ about either $\Pi_{D}$ or $\Theta$, other than that the
set of policies in $\Pi_{D}$ are Markovian. We define the _underlying_
distribution over the action-labeled episode as
$D(x_{1},a_{1},x_{2},\cdots,x_{H},a_{H})$, of which the agent _only_ gets to
observe the $(x_{1},x_{2},\cdots,x_{H})$. We will use the notation $D$ to
refer to any distribution that is derived from the above joint distribution.

We assume that observations encode time steps. This can be trivially
accomplished by simply concatenating the time step information to the
observation. This also implies that observations from different time steps are
different. Because of this property, we can assume that the Markovian policies
used to realize $D$ were time homogenous, i.e., they only depend on
observation and not observation and timestep pair (this is because we include
timesteps in the observation). Therefore, for all $h\in[H]$ and
$k\in\mathbb{N}$ we have:

$D(x_{h+k}=x^{\prime}\mid x_{h}=x)=D(x_{k+1}=x^{\prime}\mid x_{1}=x)$ (4)

We denote $D(x_{h})$ to define the marginal distribution over an observation
$x_{h}$, and $D(x_{h},x_{h+k})$ to denote the marginal distribution over a
pair of observations $(x_{h},x_{h+k})$ in the episode. We similarly define
$D(x_{h},a_{h})$ as the distribution over observation action pairs
$(x_{h},a_{h})$.

We assume that the video data has good coverage. This is stated formally
below:

######  Assumption 4 (State Coverage by $D$).

Given our policy class $\Pi$, there exists an $\eta_{\textrm{min}}>0$ such
thatif $\sup_{\pi\in\Pi}\mathbb{P}_{\pi}(s_{h}=s)>0$ for some
$s\in{\mathcal{S}}$, then we assume
$D\left(\phi^{\star}(x_{h})=s\right)\geq\eta_{\textrm{min}}$.

In practice, Assumption 4 can be satisfied since videos are more easily
available than labeled episodes and we can hope that a large diverse
collection of videos can provide reasonable coverage over the underlying state
action space. E.g., for tasks like gaming, one can use hours of streaming data
from many users.

Further, we also assume that the data policy depends only on the endogenous
state. Recall that for an observation $x\in\mathcal{X}$, its endogenous state
is given by $\phi^{\star}(x)\in{\mathcal{S}}$.

######  Assumption 5 (Noise-Free Video Distribution).

For any $h$, $\pi\in\Pi_{D}$, $x_{h}\in\textrm{supp}~{}\mathbb{P}_{\pi}$ and
$a\in\mathcal{A}$, we have

$\pi(a\mid x_{h})=\pi(a\mid\phi^{\star}(x_{h})).$

##### Justification of Noise-Free Policy.

Typically, video data is created by humans. E.g., a human may be playing a
game and the video data is collected by recording the user’s screen. A user is
unlikely to take actions relying on iid or exogenous noise in the observation
process. Therefore, the collected data can be expected to obey the noise-free
assumption.

##### Multi-step transition.

We choose to analyze a multi-step variant of standard temporal contrastive and
forward modeling algorithms that train on a dataset of pairs of observations
$(x,x^{\prime})$ that can be variable time steps apart. As our proof will
show, this gives the algorithms more expressibility and allows them to learn
correct representations for some problems that their single-step variants
(i.e., the observations are adjacent) or fixed time-step variants (i.e., the
observations are fixed time steps apart) cannot solve. We will use the
variable $k$ to denote the time steps by which these observations differ.
Formally, we will call $(x,k,x^{\prime})$ as a multi-step transition where $x$
was observed at some time step $h$, and $x^{\prime}$ was observed at $h+k$.
For the single-step variant of the algorithms, we have $k=1$. For the fixed
multi-step variant, we have $k>1$ but $k$ is fixed. Finally, in the general
multi-step variant, we will assume that $k$ is picked from $\textrm{Unf}([K])$
where $K$ is a fixed upper bound.

##### Extending episode to $H+K$.

When using $k>1$, we may want to collect a multi-step transition
$(x,k,x^{\prime})$ where $x=x_{H}$ to allow learning state representation for
time step $H$. However, at this point, we don’t have time steps left to
observe $x_{H+k}$. We alleviate this by assuming that we can allow an episode
to run till $H+K$ if necessary. In practice, this is not a problem where the
algorithm sets the horizon and not the environment. However, if we cannot go
past $H$, then we can instead assume that all states are reachable by the time
step $H-K$ and so their state representation can be learned when $x$ is
selected at $x_{H-K}$. In our analysis ahead, we make the former setting that
the episodes can be extended to $H+K$, but it can be easily rephrased to work
with the other setting.

For both the forward model and the temporal contrastive approach, we assume
access to a dataset
$\mathcal{D}_{\textrm{for}}=\left((x^{(i)},k^{(i)},x^{\prime(i)})\right)_{i=1}^{n}$
of pairs of observations. We define a few different distributions that can be
used to generate this set. For a given $k\in[K]$, we define a distribution
$D_{k}$ over $k$-step separate observations as:

$D_{k}\left(X=x,X^{\prime}=x^{\prime}\right)=\frac{1}{H}\sum_{h=1}^{H}D(x_{h}=x,x_{h+k}=x^{\prime})$
(5)

We can sample $(x,k,x^{\prime})\sim D_{k}(X,X^{\prime})$ by sampling an
episode $(x_{1},x_{2},\cdots,x_{H})\sim D$, and then sampling a
$h\sim\textrm{Unf}([H])$, and choosing $x=x_{h}$ and $x^{\prime}=x_{h+k}$.

We also define a distribution $D_{\textrm{unf}}$ where we also sample $k$
uniformly over available choices:

$D_{\textrm{unf}}\left(X=x,k,X^{\prime}=x^{\prime}\right)=\frac{1}{K}D_{k}(x_{h}=x,x_{h+k}=x^{\prime})$
(6)

We can sample $(x,k,x^{\prime})\sim D_{\textrm{unf}}(X,X^{\prime})$ by
sampling an episode $(x_{1},x_{2},\cdots,x_{H})\sim D$, and then sampling
$h\in[H]$, and sampling $k\in[K]$, and choosing $(x_{h},x_{h+k})$ as the
selected pair.

We define a useful notation $\rho\in\Delta(\mathcal{X})$ as:

$\rho(X=x)=\frac{1}{H}\sum_{h=1}^{H}D(x_{h}=x).$ (7)

The distribution $\rho(X)$ is a good distribution to sample from as it covers
states across all time steps. Finally, because of Assumption 4, we have the
following:

$\forall s\in{\mathcal{S}},\qquad\rho(s)\geq\frac{\eta_{\textrm{min}}}{H}$ (8)

This is because we assume every state $s\in{\mathcal{S}}$, is visited at some
time step $t$, and so we have $D(s_{t}=s)\geq\eta_{\textrm{min}}$, and
$\rho(s)=\frac{1}{H}\sum_{h=1}^{H}D(s_{h}=s)\geq\frac{1}{H}D(s_{t}=s)\geq\frac{\eta_{\textrm{min}}}{H}$.

It can be easily verified that for both $D_{k}(X,X^{\prime})$ and
$D_{\textrm{unf}}(X,X^{\prime})$, their marginals over $X$ is given by
$\rho(X)$. Both $D_{k}$ and $D_{\textrm{unf}}$ satisfy the noise-free
property. We prove this using the next two Lemma.s

######  Lemma 1 (Property of Noise-Free policy).

Let $\pi$ be a policy such that for any $x\in\mathcal{X}$, we have $\pi(a\mid
x)=\pi(a\mid\phi^{\star}(x))$. Then for any $h\in[H]$ and $k\in[K]$ we have
$\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid x_{h}=x)$ only depend on
$\phi^{\star}(x)$ and this common value is defined by
$\mathbb{P}_{\pi}(x_{h+k}\mid s_{h}=\phi^{\star}(x))$.

###### Proof.

The proof is by induction on $k$. For $k=1$ we have:

$\displaystyle\mathbb{P}_{\pi}(x_{h+1}=x^{\prime}\mid x_{h}=x)$
$\displaystyle=\sum_{a\in\mathcal{A}}T(x^{\prime}\mid x,a)\pi(a\mid
x_{h}=x)=\sum_{a\in\mathcal{A}}T(x^{\prime}\mid\phi^{\star}(x),a)\pi(a\mid
x_{h}=\phi^{\star}(x)),$

and as the right hand side only depends on $\phi^{\star}(x)$, the base case is
proven. For the general case, we have:

$\displaystyle\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid x_{h}=x)$
$\displaystyle=\sum_{\tilde{x}\in\mathcal{X}}\mathbb{P}_{\pi}(x_{h+k}=x^{\prime},x_{h+k-1}=\tilde{x}\mid
x_{h}=x)$
$\displaystyle=\sum_{\tilde{x}\in\mathcal{X}}\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid
x_{h+k-1}=\tilde{x})\mathbb{P}_{\pi}(x_{h+k-1}=\tilde{x}\mid x_{h}=x)$
$\displaystyle=\sum_{\tilde{x}\in\mathcal{X}}\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid
x_{h+k-1}=\tilde{x})\mathbb{P}_{\pi}(x_{h+k-1}=\tilde{x}\mid
x_{h}=\phi^{\star}(x)),$

where the second step uses the fact that $\pi$ is Markovian and the last step
uses the inductive case for $k-1$. ∎

######  Lemma 2 (Distribution over Pairs).

Let $k\in[K]$, $x\in\textrm{supp}~{}\rho(X)$, then the distribution
$D_{k}(X^{\prime}\mid x)$ only depends on $\phi^{\star}(x)$. This allows us to
define $D_{k}(X^{\prime}\mid\phi^{\star}(x))$ as this common value. Similarly,
the distribution $D_{\textrm{unf}}(X^{\prime}\mid x,k)$ depends only on
$\phi^{\star}(x)$ and $k$. We define this common value as
$D_{\textrm{unf}}(X^{\prime}\mid\phi^{\star}(x),k)$.

###### Proof.

For any $k$ we have:

$\displaystyle D_{k}(X=x,X^{\prime}=x^{\prime})$
$\displaystyle=\frac{1}{H}\sum_{h=1}^{H}D(x_{h}=x,x_{h+k}=x^{\prime})$
$\displaystyle=\frac{1}{H}\sum_{h=1}^{H}\sum_{\pi\in\Pi_{D}}\Theta_{\pi}\mathbb{P}_{\pi}(x_{h}=x,x_{h+k}=x^{\prime})$
$\displaystyle=\frac{1}{H}\sum_{h=1}^{H}\sum_{\pi\in\Pi_{D}}\Theta_{\pi}\mathbb{P}_{\pi}(x_{h}=x)\mathbb{P}(x_{h+k}=x^{\prime}\mid
x_{h}=x)$
$\displaystyle=\frac{1}{H}\sum_{h=1}^{H}\sum_{\pi\in\Pi_{D}}\Theta_{\pi}\mathbb{P}_{\pi}(x_{h}=x)\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid
s_{h}=\phi^{\star}(x)),\quad\mbox{(using~{}\hyperref@@ii[lem:exo-free-
policy]{Lemma~{}\ref*{lem:exo-free-policy}})}$
$\displaystyle=\frac{q(x\mid\phi^{\star}(x))}{H}\sum_{h=1}^{H}\sum_{\pi\in\Pi_{D}}\Theta_{\pi}\mathbb{P}_{\pi}(s_{h}=\phi^{\star}(x))\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid
s_{h}=\phi^{\star}(x))$

The marginal $D_{k}(X=x)$ is given by:

$\displaystyle
D_{k}(X=x)=\frac{1}{H}\sum_{h=1}^{H}\sum_{\pi\in\Pi_{D}}\Theta_{\pi}q(x\mid\phi^{\star}(x))\mathbb{P}_{\pi}(s_{h}=\phi^{\star}(x))=\frac{q(x\mid\phi^{\star}(x))}{H}\sum_{h=1}^{H}D_{k}(s_{h}=\phi^{\star}(x)).$

The conditional $D_{k}(X^{\prime}=x^{\prime}\mid X=x)$ is given by:

$\displaystyle D_{k}(X^{\prime}=x^{\prime}\mid X=x)$
$\displaystyle=\frac{D_{k}(X=x,X^{\prime}=x^{\prime})}{D_{k}(x)}$
$\displaystyle=\frac{\sum_{h=1}^{H}\sum_{\pi\in\Pi_{D}}\Theta_{\pi}\mathbb{P}_{\pi}(s_{h}=\phi^{\star}(x))\mathbb{P}_{\pi}(x_{h+k}=x^{\prime}\mid
s_{h}=\phi^{\star}(x))}{\sum_{h=1}^{H}D_{k}(s_{h}=\phi^{\star}(x))}$

Therefore, the conditional $D_{k}(X^{\prime}=x^{\prime}\mid X=x$ only depends
on $\phi^{\star}(x)$, and we define this common value as
$D_{k}(X^{\prime}=x^{\prime}\mid s=\phi^{\star}(x))$.

The proof for $D_{\textrm{unf}}$ is similar. We can use the property of
$D_{k}$ that we have proven to get:

$\displaystyle D_{\textrm{unf}}(X^{\prime}=x^{\prime}\mid X=x,k)$
$\displaystyle=\frac{D_{\textrm{unf}}(X=x,k,X^{\prime}=x^{\prime})}{\sum_{\tilde{x}\in\mathcal{X}}D_{\textrm{unf}}(X=x,k,X^{\prime}=\tilde{x})}$
$\displaystyle=\frac{D_{k}(X=x,X^{\prime}=x^{\prime})}{\sum_{\tilde{x}\in\mathcal{X}}D_{k}(X=x,X^{\prime}=\tilde{x})}$
$\displaystyle=\frac{D_{k}(X^{\prime}=x^{\prime}\mid
X=x)}{\sum_{\tilde{x}\in\mathcal{X}}D_{k}(X^{\prime}=\tilde{x}\mid X=x)}$
$\displaystyle=\frac{D_{k}(X^{\prime}=x^{\prime}\mid
X=\phi^{\star}(x))}{\sum_{\tilde{x}\in\mathcal{X}}D_{k}(X^{\prime}=\tilde{x}\mid
X=\phi^{\star}(x))}.$

Therefore, $D_{\textrm{unf}}(X^{\prime}=x^{\prime}\mid X=x,k)$ only depends on
$\phi^{\star}(x)$. We will define the common values as
$D_{\textrm{unf}}(X^{\prime}=x^{\prime}\mid s=\phi^{\star}(x),k)$. ∎

Lemma 2 allows us to define $D_{k}(x^{\prime}\mid\phi^{\star}(x))$ and
$D_{\textrm{unf}}(x^{\prime}\mid\phi^{\star}(x),k)$, as the distribution only
depends on the latent state.

####  B.1 Upper Bound for the Forward Model Baseline

Let
$\mathcal{D}_{\textrm{for}}=\\{(x^{(i)},k^{(i)},x^{\prime(i)})\\}_{i=1}^{n}$
be a pair of iid multi-step observations. We will collect this dataset in one
of three ways:

  1. 1.

Single step $(k=1)$, in this case we will sample $(x^{(i)},x^{\prime(i)})\sim
D_{k}(X,X^{\prime})$. As explained before, we can get this sample using the
episode data. We save $(x^{(i)},k,x^{\prime(i)})$ as our sample.

  2. 2.

Fixed multi-step. We use a fixed $k>1$, and sample
$(x^{(i)},x^{\prime(i)})\sim D_{k}(X,X^{\prime})$. We save
$(x^{(i)},k,x^{\prime(i)})$ as our sample.

  3. 3.

Variable multi-step. We sample $(x,k,x^{\prime})\sim
D_{\textrm{unf}}(X,k,X^{\prime})$ and use it as our sample.

We will abstract these three choices using a general notion of
$D_{pr}\in\Delta(\mathcal{X}\times[K]\times\mathcal{X})$. In the first two
cases, we assume we have point-mass distribution over $k$ and given this $k$,
we sample from $D_{k}(X,X^{\prime})$. We will assume
$(x^{(i)},k^{(i)},x^{\prime(i)})\sim D_{pr}$. We can create
$\mathcal{D}_{\textrm{for}}$ from the dataset $\mathcal{D}$ of $n$ episodes
sampled from $D$ using the sampling procedures explained earlier. Note that as
marginals over both $D_{k}(X)$ and $D_{\textrm{unf}}(X)$ is $\rho(X)$,
therefore, the marginals over $D_{pr}(X)$ is also $\rho(X)$. Additionally, we
will define $D_{pr}(k)$ as the marginal over $k$ which is either point-mass in
the first two sampling procedures and $\textrm{Unf}([K])$ in the third
procedure.

We assume access to two function classes. The first is a decoder class
$\Phi_{N}:\mathcal{X}\rightarrow[N]$ where $N$ is a given number that
satisfies $N\geq|{\mathcal{S}}|$. The second is a conditional probability
class $\mathcal{F}:[N]\times[K]\rightarrow\Delta(\mathcal{X})$.

######  Assumption 6.

(Realizability of $\Phi$ and $\mathcal{F}$) We assume that there exists
$\phi^{\circ}\in\Phi_{N}$ and $f^{\circ}\in\mathcal{F}$ such that
$f^{\circ}(x^{\prime}\mid\phi^{\circ}(x),k)=D_{pr}(x^{\prime}\mid
x,k)=D_{pr}(x^{\prime}\mid\phi^{\star}(x),k)$ for all $(x,k)\sim
D_{pr}(\cdot,\cdot)$.

This assumption firstly is non-vacuous as $D_{pr}(x^{\prime}\mid
x)=D_{pr}(x^{\prime}\mid\phi^{\star}(x))$, and therefore, we can apply a
bottleneck function $\phi$ and still assume realizability. For example, we can
assume that $\tilde{\phi}$ is the same as $\phi^{\star}$ up to the relabeling
of its output, and $\tilde{f}(x^{\prime}\mid i)=D_{pr}(x^{\prime}\mid s)$.

Let $\widehat{f},\widehat{\phi}$ be the empirical solution to the following
maximum likelihood problem.

$\widehat{f},\widehat{\phi}=\arg\max_{f\in\mathcal{F},\phi\in\Phi_{N}}\frac{1}{n}\sum_{i=1}^{n}\ln
f\left(x^{\prime(i)}\mid\phi(x^{(i)}),k^{(i)}\right)$ (9)

Note that when $k$ is fixed (we sample from $D_{k}$), then information
theoretically there is no advantage of condition on $k$ and it can be dropped
from optimization.

As we are in a realizable setting (Assumption 6), we can use standard maximum
likelihood guarantees to get the following result.

######  Proposition 4 (Generalization Bound).

Fix $\delta\in(0,1)$, then with probability at least $1-\delta$, we have:

$\mathbb{E}_{(x,k)\sim D_{pr}}\left[\left\|D_{pr}(X^{\prime}\mid
x,k)-\widehat{f}(X^{\prime}\mid\widehat{\phi}(x),k)\right\|_{\textrm{TV}}^{2}\right]\leq\Delta^{2}(n;\delta),$

where
$\Delta^{2}(n;\delta)=\frac{2}{n}\ln\left(\frac{|\Phi|\cdot|\mathcal{F}|}{\delta}\right)$.

For proof see Chapter 7 of Geer (2000).

Finally, we assume that the forward modeling objective is expressive to allow
the separation of states. While, this seems like assuming that the objective
works, our goal is to establish a formal notion of the margin so we can verify
it later in different settings to see when it holds.

######  Assumption 7.

(Forward Modeling Margin). We assume there exists a
$\beta_{\textrm{for}}\in(0,1)$ such that:

$\inf_{s_{1},s_{2}\in{\mathcal{S}},s_{1}\neq s_{2}}\mathbb{E}_{k\sim
D_{pr}}\left[\left\|D_{pr}(X^{\prime}\mid s_{1},k)-D_{pr}(X^{\prime}\mid
s_{2},k)\right\|_{\textrm{TV}}\right]\geq\beta_{\textrm{for}}$

Note that this defines two types of margin depending on $D_{pr}$. When $k$ is
a fixed value, the margin is given by:

$\beta_{\textrm{for}}^{(k)}=\inf_{s_{1},s_{2}\in{\mathcal{S}},s_{1}\neq
s_{2}}\left\|D_{pr}(X^{\prime}\mid s_{1},k)-D_{pr}(X^{\prime}\mid
s_{2},k)\right\|_{\textrm{TV}}.$

When we sample $k\sim\textrm{Unf}([K])$ then the margin is given by:

$\beta_{\textrm{for}}^{(u)}=\inf_{s_{1},s_{2}\in{\mathcal{S}},s_{1}\neq
s_{2}}\frac{1}{K}\sum_{k=1}^{K}\left\|D_{pr}(X^{\prime}\mid
s_{1},k)-D_{pr}(X^{\prime}\mid s_{2},k)\right\|_{\textrm{TV}}.$

We will use the abstract notion $\beta_{\textrm{for}}$ for forward margin
which will be equal to $\beta_{\textrm{for}}^{(k)}$ or
$\beta_{\textrm{for}}^{(u)}$ depending on our sampling procedure. It is easy
to see that
$\beta_{\textrm{for}}^{(u)}=\frac{1}{K}\sum_{k=1}^{K}\beta_{\textrm{for}}^{(k)}$.

We are now ready to state our first main result.

######  Proposition 5 (Recovering Endogenous State.).

Fix $\delta\in(0,1)$, then with probability at least $1-\delta$ we learn
$\widehat{\phi}$ that satisfies:

$\mathbb{P}_{x_{1},x_{2}\sim\rho}\left(\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\land\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right)\leq\frac{2\Delta(n,\delta)}{\beta_{\textrm{for}}}.$

###### Proof.

We start with a coupling argument where we sample $x_{1},x_{2}$ independently
from $D_{pr}(X)$ which is the same as $\rho(X)$.

$\displaystyle\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|D_{pr}(X^{\prime}\mid
x_{1},k)-D_{pr}(X^{\prime}\mid x_{2},k)\right\|_{\textrm{TV}}\right]$
$\displaystyle\leq\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x_{1}),k)-D_{pr}(X^{\prime}\mid
x_{1},k)\right\|_{\textrm{TV}}\right]$
$\displaystyle+\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x_{1}),k)-D_{pr}(X^{\prime}\mid
x_{2},k)\right\|_{\textrm{TV}}\right]$

We bound these two terms separately

$\displaystyle\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x_{1}),k)-D_{pr}(X^{\prime}\mid
x_{1},k))\right\|_{\textrm{TV}}\right]$
$\displaystyle\leq\sqrt{\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim
D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\right]}\cdot\sqrt{\mathbb{E}_{x_{1},x_{2}\sim
D_{pr},k\sim
D_{pr}}\left[\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x_{1}),k)-D_{pr}(X^{\prime}\mid
x_{1},k))\right\|_{\textrm{TV}}^{2}\right]}$
$\displaystyle=\sqrt{\mathbb{E}_{x_{1},x_{2}\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\right]}\cdot\sqrt{\mathbb{E}_{(x,k)\sim
D_{pr}}\left[\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x)-D_{pr}(X^{\prime}\mid
x))\right\|_{\textrm{TV}}^{2}\right]}$ $\displaystyle\leq b\cdot\Delta,$

where $b=\sqrt{\mathbb{E}_{x_{1},x_{2}\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\right]}$ and
the second step uses Cauchy-Schwarz inequality. It is straightforward to
verify that $b\in[0,1]$. We bound the second term similarly

$\displaystyle\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x_{1}),k)-D_{pr}(X^{\prime}\mid
x_{2},k)\right\|_{\textrm{TV}}\right]$
$\displaystyle=\mathbb{E}_{x_{1},x_{2}\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|\widehat{f}(X^{\prime}\mid\widehat{\phi}(x_{2}),k)-D_{pr}(X^{\prime}\mid
x_{2},k)\right\|_{\textrm{TV}}\right]$ $\displaystyle\leq b\cdot\Delta,$

where the second step uses the crucial coupling argument that we can replace
$x_{1}$ with $x_{2}$ because of the indicator ${\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}$, and the last
step follows as we reduce it to the first term except we switch the names of
$x_{1}$ and $x_{2}$. Combining the two upper bounds we get:

$\displaystyle\mathbb{E}_{x_{1},x_{2}\sim D_{pr},k\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left\|D_{pr}(X^{\prime}\mid
x_{1},k)-D_{pr}(X^{\prime}\mid x_{2},k)\right\|_{\textrm{TV}}\right]\leq
2b\cdot\Delta$

or, equivalently,

$\displaystyle\mathbb{E}_{x_{1},x_{2}\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\underbrace{\mathbb{E}_{k\sim
D_{pr}}\left[\left\|D_{pr}(X^{\prime}\mid x_{1},k)-D_{pr}(X^{\prime}\mid
x_{2},k)\right\|_{\textrm{TV}}\right]}_{:=\Gamma(x_{1},x_{2})}\right]\leq
2b\cdot\Delta$

Let $\Gamma(x_{1},x_{2})=\mathbb{E}_{k\sim
D_{pr}}\left[\left\|D_{pr}(X^{\prime}\mid x_{1},k)-D_{pr}(X^{\prime}\mid
x_{2},k)\right\|_{\textrm{TV}}\right]$. For any two observations, if
$\phi^{\star}(x_{1})=\phi^{\star}(x_{2})$, then $\left\|D_{pr}(X^{\prime}\mid
x_{1})-D_{pr}(X^{\prime}\mid x_{2})\right\|_{\textrm{TV}}=0$, and therefore,
$\Gamma(x_{1},x_{2})=0$ because of Lemma 2. Otherwise, $\Gamma(x_{1},x_{2})$
is at least $\beta_{\textrm{for}}$, by Assumption 6. Combining these two
observations we get:

$\displaystyle\Gamma(x_{1},x_{2})\geq\beta_{\textrm{for}}{\bf
1}\\{\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\\}$

Combining the previous two inequalities we get:

$\displaystyle\mathbb{E}_{x_{1},x_{2}\sim D_{pr}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\land\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\right\\}\right]\leq\frac{2b\cdot\Delta}{\beta_{\textrm{for}}}$

This directly gives

$\mathbb{P}_{x_{1},x_{2}\sim
D_{pr}}\left(\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\land\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\right)\leq\frac{2b\Delta}{\beta_{\textrm{for}}}\leq\frac{2\Delta}{\beta_{\textrm{for}}}.$

The proof is completed by recalling that marginal $D_{pr}(X)$ is the same as
$\rho(X)$. ∎

Proposition 5 shows that the learned $\widehat{\phi}$ has one-sided error. If
it merges two observations, then with high probability they are not from the
same state. As $N=|{\mathcal{S}}|$, we will show below that the reverse is
also true.

######  Theorem 6.

If $N=|{\mathcal{S}}|$, then there exists a bijection
$\alpha:[N]\rightarrow{\mathcal{S}}$ such that for any $s\in{\mathcal{S}}$ we
have:

$\mathbb{P}_{x\sim q(\cdot\mid
s)}\left(\widehat{\phi}(x)=\alpha(s)\mid\phi^{\star}(x)=s\right)\geq
1-\frac{4N^{3}H^{2}\Delta}{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}},$

provided
$\Delta<\frac{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}{N^{2}H^{2}}$.

###### Proof.

We define a few shorthand below for any $j\in[N]$ and
$\tilde{s}\in{\mathcal{S}}$

$\displaystyle\mathbb{P}(j,\tilde{s})$
$\displaystyle=\mathbb{P}_{x\sim\rho}\left(\widehat{\phi}(x)=j\land\phi^{\star}(x)=\tilde{s}\right)$
$\displaystyle\rho(j)$
$\displaystyle=\mathbb{P}_{x\sim\rho}\left(\widehat{\phi}(x)=j\right)$
$\displaystyle\rho(\tilde{s})$
$\displaystyle=\mathbb{P}_{x\sim\rho}\left(\phi^{\star}(x)=\tilde{s}\right).$

It is easy to verify that $\mathbb{P}(j,\tilde{s})$ is a joint distribution
with $\rho(j)$ and $\rho(\tilde{s})$ as its marginals.

Fix $i\in[N]$ and $s\in{\mathcal{S}}$.

$\displaystyle\mathbb{P}_{x_{1},x_{2}\sim\rho}\left(\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\land\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\right)$
$\displaystyle=\mathbb{P}_{x_{1},x_{2}\sim\rho}\left(\cup_{\tilde{s}\in{\mathcal{S}},j\in[N]}\left\\{\widehat{\phi}(x_{1})=j\land\widehat{\phi}(x_{2})=j\land\phi^{\star}(x_{1})=\tilde{s}\land\phi^{\star}(x_{2})\neq\tilde{s}\right\\}\right)$
$\displaystyle\geq\mathbb{P}_{x_{1},x_{2}\sim\rho}\left(\widehat{\phi}(x_{1})=i\land\widehat{\phi}(x_{2})=i\land\phi^{\star}(x_{1})=s\land\phi^{\star}(x_{2})\neq
s\right)$
$\displaystyle=\mathbb{P}_{x_{1}\sim\rho}\left(\widehat{\phi}(x_{1})=i\land\phi^{\star}(x_{1})=s\right)\mathbb{P}_{x_{2}\sim\rho}\left(\widehat{\phi}(x_{2})=i\land\phi^{\star}(x_{2})\neq
s\right)$
$\displaystyle=\mathbb{P}_{x\sim\rho}\left(\widehat{\phi}(x)=i\land\phi^{\star}(x)=s\right)\left(\sum_{s^{\prime}\in{\mathcal{S}}}\mathbb{P}_{x\sim\rho}\left(\widehat{\phi}(x)=i\land\phi^{\star}(x)=s^{\prime}\right)-\mathbb{P}_{x\sim\rho}(\widehat{\phi}(x)=i\land\phi^{\star}(x)=s)\right)$
$\displaystyle=\mathbb{P}(i,s)\left(\sum_{s^{\prime}\in{\mathcal{S}}}\mathbb{P}(i,s^{\prime})-\mathbb{P}(i,s)\right)$
$\displaystyle=\mathbb{P}(i,s)\left(\rho(i)-\mathbb{P}(i,s)\right).$

Combining this with Proposition 5, we get:

$\forall
i\in[N],s\in{\mathcal{S}},\qquad\mathbb{P}(i,s)\left(\rho(i)-\mathbb{P}(i,s)\right)\leq\Delta^{\prime}:=\frac{2\Delta}{\beta_{\textrm{for}}}$

where we have used a shorthand $\Delta^{\prime}=2\Delta/\beta_{\textrm{for}}$.
We define a mapping $\alpha:{\mathcal{S}}\rightarrow[N]$ where for any
$s\in{\mathcal{S}}$:

$\alpha(s)=\arg\max_{j\in[N]}\mathbb{P}(j,s)$ (10)

We immediately have:

$\mathbb{P}(\alpha(s),s)=\max_{j\in[N]}\mathbb{P}(j,s)\geq\frac{1}{N}\sum_{j=1}^{N}\mathbb{P}(j,s)=\frac{1}{N}\rho(s)\geq\frac{\eta_{\textrm{min}}}{NH},$
(11)

where we use the fact that max is greater than average in the first
inequality, and Equation 8. Further, for every $s\in{\mathcal{S}}$, we have:

$\mathbb{P}(\alpha(s),s)\left(\rho(\alpha(s))-\mathbb{P}(\alpha(s),s)\right)\leq\Delta^{\prime}.$

Plugging the lower bound
$\mathbb{P}(\alpha(s),s)\geq\frac{\eta_{\textrm{min}}}{NH}$, we get:

$\mathbb{P}(\alpha(s),s)\geq\rho(\alpha(s))-\frac{NH\Delta^{\prime}}{\eta_{\textrm{min}}}.$
(12)

We now show that if
$\Delta^{\prime}<\frac{\eta_{\textrm{min}}^{2}}{2N^{2}H^{2}}$, then
$\alpha(s)$ is a bijection. Let $s_{1}$ and $s_{2}$ be such that
$\alpha(s_{1})=\alpha(s_{2})=i$. Then using the above Equation 12 we get
$\mathbb{P}(i,s_{1})\geq\rho(i)-\frac{NH\Delta^{\prime}}{\eta_{\textrm{min}}}$
and
$\mathbb{P}(i,s_{2})\geq\rho(i)-\frac{NH\Delta^{\prime}}{\eta_{\textrm{min}}}$.
We have:

$\rho(i)=\sum_{\tilde{s}\in{\mathcal{S}}}\mathbb{P}(i,\tilde{s})\geq\mathbb{P}(i,s_{1})+\mathbb{P}(i,s_{2})\geq
2\rho(i)-\frac{2NH\Delta^{\prime}}{\eta_{\textrm{min}}}$

This implies $\frac{2N\Delta^{\prime}}{\eta_{\textrm{min}}}\geq\rho(i)$ but as
$\rho(i)=\rho(\alpha(s_{1}))\geq\mathbb{P}(\alpha(s_{1}),s_{1})\geq\frac{\eta_{\textrm{min}}}{NH}$
(Equation 11), we get
$\frac{2NH\Delta^{\prime}}{\eta_{\textrm{min}}}\geq\frac{\eta_{\textrm{min}}}{NH}$
or $\Delta^{\prime}\geq\frac{\eta_{\textrm{min}}^{2}}{2N^{2}H^{2}}$. However,
as we assume that
$\Delta^{\prime}<\frac{\eta_{\textrm{min}}^{2}}{2N^{2}H^{2}}$, therefore, this
is a contradiction. This implies $\alpha(s_{1})\neq\alpha(s_{2})$ for any two
different states $s_{1}$ and $s_{2}$. Since we assume $|N|=|{\mathcal{S}}|$,
this implies $\alpha$ is a bijection.

Fix $s\in{\mathcal{S}}$ and let $i\neq\alpha(s)$. As $\alpha$ is a bijection,
let $\tilde{s}=\alpha^{-1}(i)$, we can show that $\mathbb{P}(i,s)$ is small:

$\mathbb{P}(i,s)\leq\rho(i)-\mathbb{P}(i,\tilde{s})=\rho(\alpha(\tilde{s}))-\mathbb{P}(\alpha(\tilde{s}),\tilde{s})\leq\frac{NH\Delta^{\prime}}{\eta_{\textrm{min}}}$
(13)

where we use $s\neq\tilde{s}$ and Equation 12.

This allows us to show that $\mathbb{P}(\alpha(s)\mid s)$ is high as follows:

$\displaystyle\mathbb{P}(\alpha(s)\mid
s)=\frac{\mathbb{P}(\alpha(s),s)}{\rho(s)}$
$\displaystyle=\frac{\mathbb{P}(\alpha(s),s)}{\mathbb{P}(\alpha(s),s)+\sum_{i=1,i\neq\alpha(s)}^{N}\mathbb{P}(i,s)}$
$\displaystyle\geq\frac{\mathbb{P}(\alpha(s),s)}{\rho(\alpha(s))+\frac{N^{2}H\Delta^{\prime}}{\eta_{\textrm{min}}}},$
$\displaystyle\geq\frac{\rho(\alpha(s))-\frac{NH\Delta^{\prime}}{\eta_{\textrm{min}}}}{\rho(\alpha(s))+\frac{N^{2}H\Delta^{\prime}}{\eta_{\textrm{min}}}}$
$\displaystyle=1-\frac{\left(\frac{N^{2}H\Delta^{\prime}}{\eta_{\textrm{min}}}+\frac{NH\Delta^{\prime}}{\eta_{\textrm{min}}}\right)}{\rho(\alpha(s))+\frac{N^{2}H\Delta^{\prime}}{\eta_{\textrm{min}}}}$
$\displaystyle\geq
1-\frac{2N^{2}H^{2}\Delta^{\prime}}{\eta_{\textrm{min}}\rho(\alpha(s))}$
$\displaystyle\geq
1-\frac{2N^{3}H^{2}\Delta^{\prime}}{\eta_{\textrm{min}}^{2}},$

where the first inequality uses Eq. 13 and
$\rho(\alpha(s))\geq\mathbb{P}(\alpha(s),s)$, second inequality uses Eq. 12,
and the last step uses
$\rho(\alpha(s))\geq\mathbb{P}(\alpha(s),s)\geq\frac{\eta_{\textrm{min}}}{NH}$.

The proof is completed by noting that:

$\mathbb{P}_{x\sim q(\cdot\mid
s)}\left(\widehat{\phi}(x)=\alpha(s)\right)=\mathbb{P}_{x\sim\rho}\left(\widehat{\phi}(x)=\alpha(s)\mid\phi^{\star}(x)=s\right)=\mathbb{P}(\alpha(s)\mid
s).$

∎

Let $\mathscr{A}$ be a PAC RL algorithm for tabular MDPs. We assume that this
algorithm’s sample complexity is given by
$n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)$ where $S$ and $A$ are the size
of the state space and action space of the tabular MDP, $H$ is the horizon,
and $(\varepsilon,\delta)$ are the typical PAC RL hyperparameters denoting
tolerance and failure probability. Formally, the algorithm $\mathscr{A}$
interacts with a tabular MDP $\mathbb{M}$ for
$n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)$ episodes and outputs a policy
$\widehat{\varphi}:{\mathcal{S}}\times[H]\rightarrow\mathcal{A}$ such that
with probability at least $1-\delta$ we have:

$\sup_{\varphi\in\Psi_{\textrm{all}}}V_{\mathbb{M}}(\varphi)-V_{\mathbb{M}}(\widehat{\varphi})\leq\varepsilon,$

where $\Psi_{\textrm{all}}$ is the space of all policies of the type
${\mathcal{S}}\times[H]\rightarrow\mathcal{A}$.

We assume that we are given knowledge of the desired $(\varepsilon,\delta)$
hyperparameters in the downstream RL task during the representation pre-
training phase so we can use the right amount of data.

##### Induced Finite MDP.

The latent MDP inside a block MDP is a tabular MDP with state space
${\mathcal{S}}$, action space $\mathcal{A}$, horizon $H$, transition dynamics
$T$, reward function $R$, and a start state distribution of $\mu$. If we
directly had access to this latent MDP, say via the true decoding function
$\phi^{\star}$, then we can apply the algorithm $\mathscr{A}$ and learn the
optimal latent policy $\varphi^{\star}$ which we can couple with
$\phi^{\star}$ and learn the optimal observation-based policy. Formally, we
write this observation-based policy as
$\varphi\circ\phi^{\star}:\mathcal{X}\times[H]\rightarrow\mathcal{A}$ given by
$\varphi(\phi^{\star}(x),h)$. We dont have access to $\phi^{\star}$, but we
have access to $\widehat{\phi}$ that with high probability for a given $x$
outputs a state which is same as $\phi^{\star}(x)$ up to the learned
$\alpha$-bijection. We, therefore, define the induced MDP $\mathbb{M}$ as the
finite MDP with state space $\widehat{{\mathcal{S}}}$, action space
$\mathcal{A}$, transition function $\widehat{T}$, reward function
$\widehat{R}$ and start state distribution $\widehat{\mu}$. These same as the
latent Block MDP but where the true state $s$ is replaced by $\alpha(s)$. It
is this induced $\mathbb{M}$ that the tabular MDP algorithm $\mathscr{A}$ will
see with high probability.

######  Proposition 7 (PAC RL Bound).

Let $\mathscr{A}$ be a PAC RL algorithm for tabular MDPs and
$n_{\textrm{samp}}$ is its sample complexity. Let
$\widehat{\phi}:\mathcal{X}\rightarrow[N]$ be a decoder pre-trained using
video data and $\alpha:{\mathcal{S}}\rightarrow[N]$ is a bijection such that:

$\forall s\in{\mathcal{S}},\qquad\mathbb{P}_{x\sim q(\cdot\mid
s)}\left(\widehat{\phi}(x)=\alpha(s)\right)\geq 1-\vartheta,$

then let $\widehat{\varphi}$ be the policy returned by $\mathscr{A}$ on the
tabular MDP induced by $\widehat{\phi}(x)$. Then we have with probability at
least $1-\delta-n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta$:

$\sup_{\pi\in\Pi}V(\pi)-V(\varphi\circ\widehat{\phi})\leq\varepsilon+2H^{2}\vartheta$

###### Proof.

The algorithm runs for $n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)$ episodes.
This implies the agent visits $n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H$
many latent states. If the decoder maps every such state $s$ to the correct
permutation $\alpha(s)$, then the tabular MDP algorithm is running as if it
ran on the induced MDP $\mathbb{M}$. The probability of failure is bounded by
$n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta$ as all these failures
are independent given the state. Further, the failure probability of the
tabular MDP algorithm itself is $\delta$. This leads to the total failure
probability of $\delta+n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta$.

Let $\Pi$ be the set of observation-based policies we are competing with and
which includes the optimal observation-based policy $\pi^{\star}$. We can
write $\sup_{\pi\in\Pi}V(\pi)=V_{\mathbb{M}}(\varphi^{\star})$ where we use
the subscript $\mathbb{M}$ to denote that the latent policy is running in the
induced MDP $\mathbb{M}$. Further, for any latent policy $\varphi$ we have
$V(\varphi\circ\alpha\circ\phi^{\star})=V_{\mathbb{M}}(\varphi)$ as the
decoder $\alpha\circ\phi^{\star}:x\mapsto\alpha(\phi^{\star}(x))$ give me
access to the true state of the induced MDP $\mathbb{M}$. Then with
probability at least $1-\delta$, we have:

$\displaystyle
V_{\mathbb{M}}(\varphi^{\star})-V_{\mathbb{M}}(\widehat{\varphi})\leq\varepsilon$

This allows us to bound the sub-optimality of the learned observation-based
policy $\widehat{\varphi}\circ\widehat{\phi}$ as:

$\displaystyle\sup_{\pi\in\Pi}V(\pi)-V(\widehat{\varphi}\circ\widehat{\phi})$
$\displaystyle=V(\varphi^{\star}\circ\alpha\circ\phi^{\star})-V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})+V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})$
$\displaystyle=V_{\mathbb{M}}(\varphi^{\star})-V_{\mathbb{M}}(\widehat{\varphi})+V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})$
$\displaystyle\leq\varepsilon+V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})$

Here we use $\widehat{\varphi}\circ\alpha\circ\phi^{\star}$ to denote an
observation-based policy that takes action as
$\widehat{\varphi}(\alpha(\phi^{\star}(x)),h)$.

We bound
$V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})$
below. Let
$\mathcal{E}_{h}=\\{\widehat{\phi}(x_{h})=\alpha(\phi^{\star}(x_{h}))\\}$ and
$\mathcal{E}=\cap_{h=1}^{H}\mathcal{E}_{h}$ be two events. We have
$\mathbb{P}(\mathcal{E}_{h})\geq 1-\vartheta$. Further, using union bound we
have
$\mathbb{P}(\mathcal{E}^{c})=\mathbb{P}(\cup_{h=1}^{H}\mathcal{E}^{c}_{h})\leq\sum_{h=1}^{H}\mathbb{P}(\mathcal{E}^{c}_{h})\leq
H\vartheta$.

We first prove an upper bound on
$V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})$:

$\displaystyle V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})$
$\displaystyle=\mathbb{E}_{\widehat{\varphi}\circ\alpha\circ\phi^{\star}}\left[\sum_{h=1}^{H}r_{h}\right]$
$\displaystyle=\mathbb{E}_{\widehat{\varphi}\circ\alpha\circ\phi^{\star}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]\mathbb{P}_{\widehat{\varphi}\circ\alpha\circ\phi^{\star}}(\mathcal{E})+\mathbb{E}_{\widehat{\varphi}\circ\alpha\circ\phi^{\star}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}^{c}\right]\mathbb{P}_{\widehat{\varphi}\circ\alpha\circ\phi^{\star}}(\mathcal{E}^{c})$
$\displaystyle\leq\mathbb{E}_{\widehat{\varphi}\circ\alpha\circ\phi^{\star}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]+H^{2}\vartheta$
$\displaystyle=\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]+H^{2}\vartheta$

Here we have used the fact that value of any policy is in $[0,H]$ since the
horizon is $H$ and the rewards are in $[0,1]$.

We next prove a lower bound on $V(\widehat{\varphi}\circ\widehat{\phi})$:

$\displaystyle V(\widehat{\varphi}\circ\widehat{\phi})$
$\displaystyle=\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\right]$
$\displaystyle=\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]\mathbb{P}_{\widehat{\varphi}\circ\widehat{\phi}}(\mathcal{E})+\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}^{c}\right]\mathbb{P}_{\widehat{\varphi}\circ\widehat{\phi}}(\mathcal{E}^{c})$
$\displaystyle\geq\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]\mathbb{P}_{\widehat{\varphi}\circ\widehat{\phi}}(\mathcal{E})$
$\displaystyle\geq\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]-\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]H\vartheta$
$\displaystyle\geq\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]-H^{2}\vartheta$

Combining the two upper bounds we get:

$V(\widehat{\varphi}\circ\alpha\circ\phi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})\leq\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]+H^{2}\vartheta-\mathbb{E}_{\widehat{\varphi}\circ\widehat{\phi}}\left[\sum_{h=1}^{H}r_{h}\mid\mathcal{E}\right]+H^{2}\vartheta\leq
2H^{2}\vartheta$

Therefore, with probability at least $1-\delta-
n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta$, learn a policy
$\widehat{\varphi}\circ\widehat{\phi}$ such that:

$\sup_{\pi\in\Pi}V(\pi)-V(\widehat{\varphi}\circ\widehat{\phi})\leq\varepsilon+2H^{2}\vartheta.$

∎

######  Theorem 8 (Wrapping up the proof.).

Fix $\varepsilon_{\circ}>0$ and $\delta_{\circ}\in(0,1)$ and let $\mathscr{A}$
be any PAC RL algorithm for tabular MDPs with sample complexity
$n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)$. If $n$ satisfies:

$n=\mathcal{O}\left(\left\\{\frac{N^{4}H^{4}}{\eta_{\textrm{min}}^{4}\beta_{\textrm{for}}^{2}}+\frac{N^{6}H^{8}}{\varepsilon^{2}_{\circ}\eta_{\textrm{min}}^{4}\beta_{\textrm{for}}^{2}}+\frac{N^{6}H^{6}n_{\textrm{samp}}^{2}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)}{\delta_{\circ}^{2}\eta_{\textrm{min}}^{4}\beta_{\textrm{for}}^{2}}\right\\}\ln\left(\frac{|\mathcal{F}||\Phi|}{\delta_{\circ}}\right)\right),$

then forward modeling learns a decoder $\widehat{\phi}:\mathcal{X}\rightarrow
N$. Further, running $\mathscr{A}$ on the tabular MDP with induced by
$\widehat{\phi}$ with hyperparameters $\varepsilon=\varepsilon_{\circ}/2$,
$\delta=\delta_{\circ}/4$, returns a latent policy $\widehat{\varphi}$. Then
there exists a bijective mapping
$\alpha:{\mathcal{S}}\rightarrow[|{\mathcal{S}}|]$ such that with probability
at least $1-\delta$ we have:

$\forall s\in{\mathcal{S}},\qquad\mathbb{P}_{x\sim q(\cdot\mid
s)}\left(\widehat{\phi}(x)=\alpha(s)\mid\phi^{\star}(x)=s\right)\geq
1-\frac{4N^{3}H^{2}\Delta}{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}},$

and

$V(\pi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})\leq\varepsilon_{\circ}$

Further, the amount of online interactions in the downstream RL is given by
$n_{\textrm{samp}}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)$ and doesn’t
scale with $\ln|\Phi|$.

###### Proof.

We showed in Theorem 6 that we learn a $\widehat{\phi}$ such that:

$\mathbb{P}_{x\sim q(\cdot\mid
s)}\left(\widehat{\phi}(x)=\alpha(s)\mid\phi^{\star}(x)=s\right)\geq
1-\frac{4N^{3}H^{2}\Delta}{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}},$

provided
$\Delta<\frac{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}{N^{2}H^{2}}$.

Let
$\vartheta=\frac{4N^{3}H^{2}\Delta}{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}$.
Then from Proposition 7 we learn a $\widehat{\varphi}$ such that:

$V(\pi^{\star})-V(\widehat{\varphi}\circ\widehat{\phi})\leq\varepsilon+2H^{2}\vartheta,$

with probability at least $1-\delta-
n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta$. The failure
probability $\delta-n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta$ was
when condition in Theorem 6 holds which holds with $\delta$ probability.
Hence, total failure probability is:

$2\delta+n_{\textrm{samp}}(S,A,H,\varepsilon,\delta)H\vartheta.$

We set $\delta$ both in our representation learning analysis and in PAC RL to
$\delta_{\circ}/4$. We also set $\varepsilon$ in the PAC RL algorithm to
$\varepsilon_{\circ}/2$. This means the PAC RL algorithm runs for
$n_{\textrm{samp}}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)$ episodes.

We enforce
$\vartheta\leq\frac{\delta_{\circ}}{2n_{\textrm{samp}}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)H}$.
Then the total failure probability becomes:

$2\delta_{\circ}/4)+\delta_{\circ}/4+\delta_{\circ}/2\leq\delta_{\circ}$

We also enforce $2H^{2}\vartheta\leq\varepsilon_{\circ}/2$. The sub-optimality
of the PAC RL policy is given by:

$\varepsilon_{\circ}/2+\varepsilon_{\circ}/2\leq\varepsilon_{\circ}$

This gives us our derived PAC RL bound.

We now accumulate all conditions:

$\displaystyle\Delta$
$\displaystyle=\sqrt{\frac{2}{n}\ln\left(\frac{4|\mathcal{F}||\Phi|}{\delta_{\circ}}\right)}$
$\displaystyle\vartheta$
$\displaystyle=\frac{4N^{3}H^{2}\Delta}{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}$
$\displaystyle\Delta$
$\displaystyle<\frac{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}{N^{2}H^{2}}$
$\displaystyle\vartheta$
$\displaystyle\leq\frac{\delta_{\circ}}{2n_{\textrm{samp}}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)H}$
$\displaystyle 2H^{2}\vartheta$ $\displaystyle\leq\varepsilon_{\circ}/2$

This simplifies to

$\displaystyle\Delta$
$\displaystyle\leq\frac{\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}{N^{2}H^{2}}$
$\displaystyle\Delta$
$\displaystyle\leq\frac{\delta_{\circ}\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}{8N^{3}H^{3}n_{\textrm{samp}}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)}$
$\displaystyle\Delta$
$\displaystyle\leq\frac{\varepsilon_{\circ}\eta_{\textrm{min}}^{2}\beta_{\textrm{for}}}{16N^{3}H^{4}}$

Or,

$n=\mathcal{O}\left(\left\\{\frac{N^{4}H^{4}}{\eta_{\textrm{min}}^{4}\beta_{\textrm{for}}^{2}}+\frac{N^{6}H^{8}}{\varepsilon^{2}_{\circ}\eta_{\textrm{min}}^{4}\beta_{\textrm{for}}^{2}}+\frac{N^{6}H^{6}n_{\textrm{samp}}^{2}(S,A,H,\varepsilon_{\circ}/2,\delta_{\circ}/4)}{\delta_{\circ}^{2}\eta_{\textrm{min}}^{4}\beta_{\textrm{for}}^{2}}\right\\}\ln\left(\frac{|\mathcal{F}||\Phi|}{\delta_{\circ}}\right)\right)$

This completes the proof. ∎

####  B.2 Upper Bound for the Temporal Contrastive Approach

We first convert our video dataset $\mathcal{D}$ into a dataset suitable for
contrastive learning. We first split the datasets into $\lfloor n/2\rfloor$
pairs of videos. For each video pair
$\left\\{\left(x^{(2l)}_{1},x^{(2l)}_{2},\cdots,x^{(2k)}_{H}\right),\left(x^{(2l+1)}_{1},x^{(2l+1)}_{2},\cdots,x^{(2l+1)}_{H}\right)\right\\}$,
we create a tuple $(x,x^{\prime},k,z)$ where $z\in\\{0,1\\}$ as follows. As in
forward modeling, we will either use a fixed value of $k$, or sample
$k\in\textrm{Unf}([K])$. We denote this general distribution over $k$ by
$\omega\in\Delta([K])$ which is either point mass, or $\textrm{Unf}([K])$. We
sample $k\sim\omega$ and $z\sim\textrm{Unf}(\\{0,1\\})$ and
$h\in\textrm{Unf}([H])$. We set $x=x^{(2l)}_{h}$. If $z=1$, then we set
$x^{\prime}=x^{(2l)}_{h+k}$, otherwise, we sample
$h^{\prime}\sim\textrm{Unf}(\\{0,1\\})$ and select
$x^{\prime}=x^{(2l)}_{h^{\prime}}$. This way, we collect a dataset
$\mathcal{D}_{\textrm{cont}}$ of $\lfloor n/2\rfloor$ tuples
$(x,k,x^{\prime},z)$. We view a tuple $(x,k,x^{\prime},z)$ as a _real
observation pair_ when $z=1$, and a _fake observation pair_ when $z=0$. Note
that our sampling process leads to all data points being iid.

We define the distribution $D_{\textrm{cont}}(X,k,X^{\prime},Z)$ as the
distribution over $(x,k,x^{\prime},z)$. We can express this distribution as:

$\displaystyle D_{\textrm{cont}}(X=x,k,X^{\prime}=x^{\prime},Z=1)$
$\displaystyle=\frac{\omega(k)}{2H}\sum_{h=1}^{H}D(x=x_{h},x^{\prime}=x_{h+k})$
$\displaystyle=\frac{\omega(k)}{2}\rho(x)D(x_{k+1}=x^{\prime}\mid x_{1}=x)$
$\displaystyle D_{\textrm{cont}}(X=x,X^{\prime}=x^{\prime},Z=0)$
$\displaystyle=\frac{\omega(k)}{2H^{2}}\sum_{h=1}^{H}D(x=x_{h})\sum_{h^{\prime}=1}^{H}D(x^{\prime}=x_{h^{\prime}})$
$\displaystyle=\frac{\omega(k)}{2}\rho(x)\rho(x^{\prime})$

where we use the time homogeneity of $D$ and definition of $\rho$. We will use
a shorthand to denote $D(x_{k+1}=x^{\prime}\mid x_{1}=x)$ as $D(x^{\prime}\mid
x,k)$ in this analysis. It is easy to verify that $D(x^{\prime}\mid
x,k)=D(x^{\prime}\mid\phi^{\star}(x),k)$. The marginal distribution
$D_{\textrm{cont}}(x,k,x^{\prime})$ is given by:

$D_{\textrm{cont}}(x,k,x^{\prime})=\frac{\omega(k)\rho(x)}{2}\left(D(x^{\prime}\mid
x,k)+\rho(x^{\prime})\right)$ (14)

Note that $D_{\textrm{cont}}(X)$ is the same as $\rho(X)$.

We will use $D_{\textrm{cont}}$ for any marginal and conditional distribution
derived from $D_{\textrm{cont}}(X,k,X^{\prime},Z)$. We assume a model class
$\mathcal{G}:\mathcal{X}\times[K]\times\times[N]\rightarrow[0,1]$ that we use
for solving the prediction problem. We will also reuse the decoder class
$\phi:\mathcal{X}\rightarrow[N]$ that we defined earlier, and we will assume
that $N=|{\mathcal{S}}|$. This can be relaxed by doing clustering or working
with a different induced MDP (e.g., see the clustering algorithm in Misra et
al. (2020)). However, this is not the main point of the analysis.

We define the expected risk minimizer of the squared loss problem below:

$\widehat{g},\widehat{\phi}=\arg\min_{g\in\mathcal{G},\phi\in\Phi}\frac{1}{\lfloor
n/2\rfloor}\sum_{i=1}^{\lfloor
n/2\rfloor}\left(g(\phi(x^{(i)}),k^{(i)},x^{\prime(i)})-z^{(i)}\right)^{2}$
(15)

We express the Bayes classifier of this problem below:

######  Lemma 3 (Bayes Classifier).

The Bayes classifier of the problem posed in Equation 15 is given by
$D_{\textrm{cont}}(z=1\mid x,k,x^{\prime})$ which satisfies:

$\displaystyle D_{\textrm{cont}}(z=1\mid x,k,x^{\prime})$
$\displaystyle=\frac{D(\phi^{\star}(x^{\prime})\mid\phi^{\star}(x),k)}{D(\phi^{\star}(x^{\prime})\mid\phi^{\star}(x),k)+\rho(\phi^{\star}(x^{\prime}))}.$

###### Proof.

We can express the Bayes classifier as:

$\displaystyle D_{\textrm{cont}}(z=1\mid x,k,x^{\prime})$
$\displaystyle=\frac{D_{\textrm{cont}}(x,k,x^{\prime},z=1)}{D_{\textrm{cont}}(x,k,x^{\prime},z=1)+D_{\textrm{cont}}(x,k,x^{\prime},z=0)}$
$\displaystyle=\frac{\omega(k)/2\rho(x)D(x^{\prime}\mid
x)}{\omega(k)/2\rho(x)D(x^{\prime}\mid x)+\omega(k)/2\rho(x)\rho(x^{\prime})}$
$\displaystyle=\frac{D(x^{\prime}\mid x,k)}{D(x^{\prime}\mid
x,k)+\rho(x^{\prime})}$
$\displaystyle=\frac{D(x^{\prime}\mid\phi^{\star}(x),k)}{D(x^{\prime}\mid\phi^{\star}(x),k)+\rho(x^{\prime})}$
$\displaystyle=\frac{q(x^{\prime}\mid\phi^{\star}(x))D(\phi^{\star}(x^{\prime})\mid\phi^{\star}(x),k)}{q(x^{\prime}\mid\phi^{\star}(x))D(\phi^{\star}(x^{\prime})\mid\phi^{\star}(x),k)+q(x^{\prime}\mid\phi^{\star}(x))\rho(\phi^{\star}(x^{\prime}))}$
$\displaystyle=\frac{D(\phi^{\star}(x^{\prime})\mid\phi^{\star}(x),k)}{D(\phi^{\star}(x^{\prime})\mid\phi^{\star}(x),k)+\rho(\phi^{\star}(x^{\prime}))}.$

∎

######  Assumption 8 (Realizability).

There exists $g^{\star}\in\mathcal{G}$ and $\phi^{\circ}\in\Phi$ such that for
all $(x,k,x^{\prime})\in\textrm{supp}~{}D_{\textrm{cont}}(X,k,X^{\prime})$, we
have $D_{\textrm{cont}}(z=1\mid
x,k,x^{\prime})=g^{\star}(\phi^{\circ}(x),k,x^{\prime})$.

We will use the shorthand to denote
$g^{\star}(x,k,x^{\prime})=g^{\star}(\phi^{\circ}(x),k,x^{\prime})$.

As before, we start with typical square loss guarantees in the realizable
setting.

######  Theorem 9.

Fix $\delta\in(0,1)$. Under realizability (Assumption 8), the ERM solution of
$\widehat{f},\widehat{\phi}$ in Eq. 15 satisfies:

$\mathbb{E}_{(x,k,x^{\prime})\sim
D_{\textrm{cont}}}\left[\left(\widehat{g}(\widehat{\phi}(x),k,x^{\prime})-g^{\star}(x,k,x^{\prime})\right)^{2}\right]\leq\Delta_{\textrm{cont}}^{2}=\frac{2}{n}\ln\frac{|\mathcal{G}|.|\Phi|}{\delta}$

For proof see Proposition 12 in Misra et al. (2020).

We will prove a coupling result similar to the case for forward modeling.
However, to do this, we need to define a coupling distribution:

$D_{\textrm{coup}}(X_{1}=x_{1},X_{2}=x_{2},k,X^{\prime}=x^{\prime})=\omega(k)D_{\textrm{cont}}(X=x_{1})D_{\textrm{cont}}(X=x_{2})D_{\textrm{cont}}(X^{\prime}=x^{\prime})$

We will derive a useful importance ratio bound.

$\frac{D_{\textrm{coup}}(x_{1},k,x^{\prime})}{D_{\textrm{cont}}(x_{1},k,x^{\prime})}=\frac{2\rho(x_{1})\rho(x^{\prime})}{\rho(x_{1})D(x^{\prime}\mid
x_{1},k)+\rho(x_{1})\rho(x^{\prime})}\leq 2$ (16)

We now prove an analogous result to Proposition 5.

######  Theorem 10 (Coupling for Temporal Contrastive Learning).

With probability at least $1-\delta$ we have:

$\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|g^{\star}(x_{1},k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]<4\Delta_{\textrm{cont}}(n,\delta)$

###### Proof.

We start with triangle inequality:

$\displaystyle\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|g^{\star}(x_{1},k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]$
$\displaystyle\leq\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|g^{\star}(x_{1},k,x^{\prime})-\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})\right|\right]+$
$\displaystyle\qquad\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]$

We bound the first term as:

$\displaystyle\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|g^{\star}(x_{1},k,x^{\prime})-\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})\right|\right]$
$\displaystyle\leq\underbrace{\sqrt{\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\right]}}_{:=b}\cdot\sqrt{\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[\left|g^{\star}(x_{1},k,x^{\prime})-\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})\right|^{2}\right]}$
$\displaystyle=b\sqrt{\mathbb{E}_{(x_{1},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[\left(g^{\star}(x_{1},k,x^{\prime})-\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})\right)^{2}\right]}$
$\displaystyle=b\sqrt{\mathbb{E}_{(x_{1},k,x^{\prime})\sim
D_{\textrm{cont}}}\left[\frac{D_{\textrm{coup}}(x_{1},k,x^{\prime})}{D_{\textrm{cont}}(x_{1},k,x^{\prime})}\left(g^{\star}(x_{1},k,x^{\prime})-\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})\right)^{2}\right]}$
$\displaystyle\leq b\sqrt{2\mathbb{E}_{(x_{1},k,x^{\prime})\sim
D_{\textrm{cont}}}\left[\left(g^{\star}(x_{1},k,x^{\prime})-\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})\right)^{2}\right]}$
$\displaystyle\leq\sqrt{2}b\Delta_{\textrm{cont}},$

where we use Cauchy-Schwartz’s inequality in the first step and Equation 16 in
the second inequality. The second term is bounded as:

$\displaystyle\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]$
$\displaystyle=\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|\widehat{g}(\widehat{\phi}(x_{2}),k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]$
$\displaystyle=\mathbb{E}_{(x_{1},x_{2},k,x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|\widehat{g}(\widehat{\phi}(x_{1}),k,x^{\prime})-g^{\star}(x_{1},k,x^{\prime})\right|\right]$
$\displaystyle\leq\sqrt{2}b\Delta_{\textrm{cont}},$

where we use the coupling argument in the first step and then reduce it to the
first term using symmetric of $(x_{1},x_{2})$ in $D_{\textrm{coup}}$.
Combining the upper bounds of the two terms and using $b\leq 1$ and
$2\sqrt{2}<4$ completes the proof. ∎

######  Assumption 9 (Temporal Contrastive Margin).

We assume that there exists a $\beta_{\textrm{temp}}>0$ such that for any two
different states $s_{1}$ and $s_{2}$:

$\frac{1}{2}\mathbb{E}_{k\sim\omega,s^{\prime}\sim\rho}\left[\left|g^{\star}(s_{1},k,s^{\prime})-g^{\star}(s_{2},k,s^{\prime})\right|\right]\geq\beta_{\textrm{temp}}$

The factor of $\frac{1}{2}$ is chosen for comparison with forward modeling as
will become clear later at the end of the proof. As before, if $k$ is fixed,
the margin is given by

$\beta_{\textrm{temp}}^{(k)}:=\frac{1}{2}\inf_{s_{1}\neq
s_{2};s_{1},s_{2}\in{\mathcal{S}}}\mathbb{E}_{s^{\prime}\sim\rho}\left[\left|g^{\star}(s_{1},k,s^{\prime})-g^{\star}(s_{2},k,s^{\prime})\right|\right]$

and when $k\sim\textrm{Unf}([K])$ the margin is given by

$\beta_{\textrm{temp}}^{(u)}:=\frac{1}{2}\inf_{s_{1}\neq
s_{2};s_{1},s_{2}\in{\mathcal{S}}}\mathbb{E}_{k\sim\textrm{Unf}([K]),s^{\prime}\sim\rho}\left[\left|g^{\star}(s_{1},k,s^{\prime})-g^{\star}(s_{2},k,s^{\prime})\right|\right]$

We directly have
$\beta_{\textrm{temp}}^{(u)}\geq\frac{1}{K}\sum_{k=1}^{K}\beta_{\textrm{temp}}^{(k)}$.

######  Lemma 4.

$\mathbb{P}_{x_{1},x_{2}\sim\rho}\left(\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\land\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\right)\leq\frac{2\Delta_{\textrm{cont}}(n,\delta)}{\beta_{\textrm{temp}}}$

###### Proof.

We start with the left-hand side in Theorem 10.

$\displaystyle\mathbb{E}_{(x_{1},k,x_{2},x^{\prime})\sim
D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\left|g^{\star}(x_{1},k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]$
$\displaystyle=\mathbb{E}_{(x_{1},x_{2})\sim D_{\textrm{coup}}}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\mathbb{E}_{k\sim\omega,x^{\prime}\sim\rho}\left[\left|g^{\star}(x_{1},k,x^{\prime})-g^{\star}(x_{2},k,x^{\prime})\right|\right]\right]$
$\displaystyle=\mathbb{E}_{(x_{1},x_{2})\sim\rho}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\right\\}\mathbb{E}_{k\sim\omega,s^{\prime}\sim\rho}\left[\left|g^{\star}(x_{1},k,s^{\prime})-g^{\star}(x_{2},k,s^{\prime})\right|\right]\right]$
$\displaystyle\geq
2\beta_{\textrm{temp}}\mathbb{E}_{(x_{1},x_{2})\sim\rho}\left[{\bf
1}\left\\{\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\land\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\right\\}\right]$
$\displaystyle=2\beta_{\textrm{temp}}\mathbb{P}_{(x_{1},x_{2})\sim\rho}\left[\widehat{\phi}(x_{1})=\widehat{\phi}(x_{2})\land\phi^{\star}(x_{1})\neq\phi^{\star}(x_{2})\right],$

where we use the definition of $\beta_{\textrm{temp}}$, the fact that marginal
over $D_{\textrm{coup}}(X)$ is $\rho$, and that $g^{\star}(x,k,x^{\prime})$
only depends on $\phi^{\star}(x^{\prime})$ and $\phi^{\star}(x)$ (Lemma 3).
Combining with the inequality proved in Theorem 10, completes the proof. ∎

We have now reduced this analysis to an almost identical one to the forward
analysis case (Proposition 5). We can, therefore, use the same steps and
derive identical bounds. All what changes is that $\beta_{\textrm{for}}$ is
replaced by $\beta_{\textrm{temp}}$ and in $\Delta$ we replace
$\ln|\mathcal{F}|$ with $\ln|\mathcal{G}|$. At this point, we can clarify that
the factor of $\frac{1}{2}$ was chosen in the definition of
$\beta_{\textrm{temp}}$ so that $\beta_{\textrm{for}}$ can be replaced by
$\beta_{\textrm{temp}}$ rather than $\frac{\beta_{\textrm{temp}}}{2}$ which
will make it harder to compare margins, as we will do later.

####  B.3 Proof of Lower Bound for Exogenous Block MDPs

###### Proof of Theorem 3.

We present a hard instance using a family of exogenous block MDPs, with $H=2$,
$\mathcal{A}=\\{1,2\\}$, and a single binary endogenous factor and $d-1$
exogenous binary factors for each level, where each endogenous and exogenous
factor. We first fix an absolute constant $p\in[0,1]$.

Each MDP $M_{i}$ is indexed by $i\in[d]$, and is specified as follows:

  * •

State space: The state is represented by
$x_{h}\coloneqq[s_{h}^{\scriptscriptstyle(1)},s_{h}^{\scriptscriptstyle(2)},\dotsc,s_{h}^{\scriptscriptstyle(d)}]$,
where the superscript denotes different factors. For MDP $M_{i}$, only the
$i$-th factor $s_{h}^{\scriptscriptstyle(i)}$ is an endogenous state for all
$h$, and the other factors are exogenous. Each factor has values of
$\\{0,1\\}$.

  * •

Transition: For the MDP instance $M_{i}$: it has

    1. 1.

For the $i$-th factor (endogenous factor),
$\mathbb{P}(s_{2}^{\scriptscriptstyle(i)}\mid
s_{1}^{\scriptscriptstyle(i)},a)=\mathds{1}[s_{2}^{\scriptscriptstyle(i)}=\mathds{1}(s_{1}^{\scriptscriptstyle(i)}=a)]$.
That is, the endogenous states have deterministic dynamics. If
$s_{1}^{\scriptscriptstyle(i)}=a$, then it transitions to
$s_{2}^{\scriptscriptstyle(i)}=1$, otherwise it transitions to
$s_{2}^{\scriptscriptstyle(i)}=0$.

    2. 2.

For the $j$-th factor with $j\neq i$ (exogenous factor),
$\mathbb{P}(s_{2}^{\scriptscriptstyle(j)}\mid
s_{1}^{\scriptscriptstyle(j)})=(1-p)\mathds{1}(s_{2}^{\scriptscriptstyle(j)}=s_{1}^{\scriptscriptstyle(j)})+p\mathds{1}(s_{2}^{\scriptscriptstyle(j)}\neq
s_{1}^{\scriptscriptstyle(j)})$ for any $s_{2}^{\scriptscriptstyle(j)}$ and
$s_{1}^{\scriptscriptstyle(j)}$. That is, the $j$-th factor has probability of
$1-p$ of transiting to the same state (i.e.,
$s_{1}^{\scriptscriptstyle(j)}=0\to s_{2}^{\scriptscriptstyle(j)}=0$ or
$s_{1}^{\scriptscriptstyle(j)}=1\to s_{2}^{\scriptscriptstyle(j)}=1$), and
probability of $p$ of transiting to the different state (i.e.,
$s_{1}^{\scriptscriptstyle(j)}=0\to s_{2}^{\scriptscriptstyle(j)}=1$ or
$s_{1}^{\scriptscriptstyle(j)}=1\to s_{2}^{\scriptscriptstyle(j)}=0$).

Note that the MDP terminates at $h=2$.

  * •

Initial state distribution and reward: The marginal distribution of
$s_{1}^{\scriptscriptstyle(j)}$ is uniformly distributed at random over
$\\{0,1\\}$ for all $j\in[d]$, and all factors are independent from each
other. For MDP $M_{i}$, the agent only receive reward signal after taking
action at $h=2$, with
$R(s_{2}^{\scriptscriptstyle(i)},a)=s_{2}^{\scriptscriptstyle(i)}$. That is,
it always reward $1$ at $s_{2}^{\scriptscriptstyle(i)}=1$ and reward $0$ at
$s_{2}^{\scriptscriptstyle(i)}=0$ no matter which action it takes.

  * •

Data collection policy for video data: We assume that the data collection
policy always pick action $0$ with probability $p$ and action $1$ with
probability $1-p$ for all states.

Now we use the following two steps to establish the proof.

##### Uninformative video data for learning the state decoder

Since video data only contains state information, from the MDP family
construction above, we can easily verify that all MDP instances in such a
family will have an identical video data distribution, regardless of the
choice of constant $p$. This implies that the video data is uninformative for
the agent to distinguish the MDP instance from the MDP family. Now, we assume
$\mathcal{D}^{\scriptscriptstyle(i)}$ is the video data from the instance
$M^{\scriptscriptstyle(i)}$, and $\phi^{\scriptscriptstyle(i)}$ is the state
decoder learned from an arbitrary algorithm $\mathscr{A}_{1}$ with
$\mathcal{D}^{\scriptscriptstyle(i)}$. Then, for any arbitrary algorithm
$\mathscr{A}_{2}$ that uses the state decoder $\phi^{\scriptscriptstyle(i)}$
in its execution, it is equivalent to such an $\mathscr{A}_{2}$ that uses the