Towards Principled Representation Learning from Videos for Reinforcement Learning
Authors: Dipendra Misra
Akanksha Saran
Tengyang Xie
Alex Lamb
John Langford
Venue: ICLR 2024 spotlight (2024)
Keywords: Reinforcement Learning
Representation Learning
Primary area: reinforcement learning
TL;DR: Theoretical analysis and experiments concerning the value reinforcement learning can gain from pretrained representations of unlabeled video data.

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.

Published as a conference paper at ICLR 2024
TOWARDS PRINCIPLED REPRESENTATION LEARNING FROM VIDEOS FOR REINFORCEMENT LEARNING
Dipendra Misra!* Akanksha Saran
Tengyang Xie!
Alex Lamb'
John Langford'
Microsoff Research, NY 2Sony Research, CA
ABSTRACT
We study pre-training representations for decision-making using video data, which is abundantly available for tasks such as game agents and software test- ing. Even though significant empirical advances have been made on this problem, a theoretical understanding remains absent. We initiate the theoretical investiga- tion 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 exponentially worse than learn- ing from action-labeled trajectory data. This partially explains why reinforcement learning with video pre- vve-trainin is hard. We evaluate these representational learn- ing methods in two visual domains, yielding results that are consistent with our theoretical findings.
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 similar undertaking reinforcement learning (RL) with the goal of training a representation model that can be used in policy to solve 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 al., 2018; Bharadhwaj al., 2022). However, unlike text and image data, which are abundant on the internet or naturally generated by users, trajectory data comparatively limited and expensive to collect. In contrast, video data, which only contains 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 representations learned from trajectory data? We initiate theoretical understanding of these approaches to show when and how these approaches yield representations that can be used to solve downstream RL task efficiently.
Consider a representation learning pipeline shown in Figure 1. We are provided videos, or equiv- alently a sequence of observations. from agents navigating in the world. We make no assumption
* DM and AS contributed equally. Correspondence should be sent to dimi sra@mi crosoft. and akanksha.. www.naaananannnoee com.

---

Published as a conference paper at ICLR 2024
Preteaining Plasee
Vertor reprosntation R
Dowastream RL task
Visualize Lasent Stute Space
Figure I: A flowehart 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 model & which maps an observation to a vector representation. Right: This representation can be used in downstream task to do reinforcement learning or visualize the latent world state.
about the behavior of the agent in the video data. They can be trying to solve one task, many differ- ent tasks, or none at all. This video data is used to learn a model & that maps any given observation to vector representation. This representation is subsequently used to perform downstream RI. 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 critique model. The rep- resentation 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-make making while ignoring any noise in the observation. For ex- ample, 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 manner independent of the agent's actions (such as the motion of geese in Figure 1).
A range of approaches have been developed that provably recover the latent agent state from observa- tions using trajectory data (Misra et al., 2020; Efroni et al., 2022) which contains actions. However, for many domains there 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 ap- proaches that empirically work with video data provably recover the latent representation and lead to efficient RL. Motivated by this, we build a thcoretical understanding of three such video-based representation learning approaches: autoencoder which trains representations by reconstructing ob- servations, forward modeling which predicts future observations, and temporal contrastive learning which trains representation determine if pair of observations are causally related not.
Our first thcoretical result shows that in the absence of exogenous noise, forward modeling and tem- poral 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 sec- ond 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 ef- ficient downstream RL. In contrast, if the trajectory data was available, efficient pre-training would be possible. This establishes 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, that temporal contrastive learning is especially prone to fail in the presence of exogenous noise, as it can rely exclusively on such noise optimally minimize the contrastive loss. While we find that
find
2

---

Published as a conference paper at ICLR 2024
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 open question, empirically, we find that the performance of autoencoder-basee representation learning is unpredictable. On the other hand, ACRO continues to perform well, highlighting disadvantage of video pre-training. The code for all experiments is available as part com httm http://cccooooooreeo//o/aa
of the Intrepid codebase at
PRELIMINARIES AND OVERVIEW
In this section, we provide a formal overview of our learning setup and problem statement.
Mathematical Notation. We use [N] for N E N to define the set {1, ... N.. We assume all sets to be countable. For a given set U, we denote its cardinality by [24] and define (24)) the space of all distributions over 24. We denote the uniform distribution over 2/ by Unf(4). Finally, polyf-} denotes a term that scales polynomially in the listed quantities.
Block MDPs. We study episodic RL in Block Markov Decision Processes (Block MDP) (Du et al., 2019). A Block MDP is defined by the tuple (X, R, μ. H) where X' is a set of observa- tions that can be infinitely large, S is finite set of latent states, and A is set of finite actions. The transition dynamics T S A- + (S)) define transitions in the latent state space. The reward function R: S x A + [0, 1] assigns reward R(s, a) if action is taken in the latent state When the agent visits a state s, it receives an observation x ~ q(- 8) sampled from an emission function q S + A(X). This emission process contains temporally independent noise but no exogenous noise. Finally, μ E (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 block MDP environment by repeat- edly generating an episode (x1, 01, r1," HH an, TH) where 81 ~ μu and for all h E [H] we have TA ~ q(- $4). Th = R(sh, ax), and sh+1 T(- Sh, a,.), and all actions {an f==1 are taken by the agent. The agent never directly observes the latent states (31, (8,,82, ,H)
A key assumption in Block MDPs is that two different latent states cannot generate the same ob- servation. This is called the disjoint emission property and holds in many game and OS settings. Formally. this property allows us define a decoder o² : X' + S that maps an observation to the unique state that can generate it. The agent does not have access to o*. If the agent had access to o", 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 RL in Exogenous Block MDPs (Ex- Block MDPs) that extend Block MDPs to include exogenous noise (Efroni et al., 2022). An Ex- Block MDP is defined by (X, S, A, T, Te, R, q. H, u, (P)) where X, S, A, T, R, H and μ have the same meaning and type in Block MDPs. The additional quantities include E which is the space of exogenous noise and can be infinitely large. We use the notation & E E to denote the exogenous noise. For the setting in Figure 1, the exogenous noise variable & captures variables such as the position of geese, the position of leaves on the trees in the background, and lighting conditions. The exogenous noise { changes with time according the transition function Tg = - A() and is at start sampled from μlg. Note that unlike the agent state E S, the exogenous noise & E E, evolves independently of the agent's 's action and does not influence the evolution of the agent's 's state. The emission process q S x E + A(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). Similar to the Block MDP, we assume there exists unknown decoders o' XX' -> S and oe X - that can map an observation to the current agent state and exogenous respectively.
We assume access to a policy class II {IT X + A} where a policy π II allows the agent to take actions. For a given policy n, we use E, [1] to denote expectation taken over an episode generated by sampling actions from и. We define the value of policy V(x) =
Provable RL.
E E Th as the expected total reward or expected return. Our goal is to learn near-optimal
policy î,, i.c., supren (TT) -+t) < with probability at least114 for a given tolerance parameter e > 0 and failure probability e (0, 1), using number of episodes that scale polynomially in 1/e, 1/s, and other relevant quantities. We will call such algorithm as provably efficient. There exist
3

---

Published as a conference paper at ICLR 2024
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 decoder class + {Φ: X' - [N]] and attempt to learn o* using it. These algorithms don't 't use any pre-training and instead directly interact with the environment and learn a near-optimal policy by using samples that scale with poly(S, A, H, In [ 1/e, 1/6). Crucially, the dependence cannot be removed. The decoder class & 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.. In IT).
In {
Representation Pre-training using Videos. RL algorithms for the above settings require online episodes that scale with In | which is expensive for real-world problems where & is represented by complex neural network. Offline RI. approaches Uehara et al. (2021) offer a substitute for expensive online interactions but require access to labeled episodes (with actions and rewards) that are not naturally available in many settings such as games and software. In contrast, we focus on pre-training the decoder & using video data which is naturally available in these settings.
Problem Statement. We are given two hyperparameters e > 0 and å E (0. 1) and a sufficiently large dataset of videos. We are also given a decoder class 4 {Ø X - [N]] containing decoders that map an observation to one of the .N possible abstract states. During the pre-training phase, we learn decoder ó E 4 using the video data. We then freeze o and use it to do RL in a downstream task. Instead of using any particular choice of algorithm for RI., we assume we are given a provably efficient tabular RL algorithm sd. We convert the observation- RL problem to a tabular MDP problem by converting an observation x to its abstract state representation 6(x) using the frozen learned decoder c. The algorithm af uses <(x) instead of and outputs an abstract policy P [N] + A. We want that su (T) V(p o @) S e with probability at least 1 1-6, 6, where p o Φ x + v(((x)) is our learned policy. We also require the number of online episodes in the downstream RL phase to not scale with the size of the decoder class O. This allows us to minimize expensive online episodes while using naturally available offline video data for pre-training.
3 REPRESENTATION LEARNING FOR RL USING VIDEO DATASET
where (0) is the jth observation (or frame) of the 7th video. We are provided decoder class +6 {o X - [N]), and our goal is to learn decoder & E & that captures task-relevant information in the underlying state o" (z) while throwing away as much exogenous noise as possible. Instead of proposing new algorithm, we analyze the following three classes of well-known video-based representation learning methods. Our goal is to understand whether these methods provably learn useful representations.
We assume access to a dataset D of n videos D = {(x()) 10000 ......................... ."=
Autoencoder. This approach first maps a given observation æ to an abstract state d(x) using a decoder @ E 4, and then uses it to reconstruct the observation æ with the aid of reconstruction model class 2 = {z: [N] + X².. Formally, we optimize the following loss:
12 ΣΣ
Enuo (z, 0) nH
|=(0<=
(
(1)
In practice, autoencoders are typically implemented using a Vector Quantized bottleneck trained in a Variational AutoEncoder manner, which called the VQ-VAE approach (Oord al., 2017).
Forward Modeling, This approach is similar to the autoencoder approach but instead of recon- structing the input observation, we reconstruct future observation using a model class F { x [K] + ((X')} where N is the output size of the decoder class & and K E N is hyper- parameter representing the forward time steps from the current observation. We collect dataset of multistep transitions Dioe = ((x(() 1000 x²(0 _0)0 } sampled iid using the video dataset 'D where the observation (00)) is sampled randomly from the 2th video, k: [K], and z²(() the frame k(0) ahead of (0) in the 7th video. We distinguish between two types of sampling procedures, one where (0 is always fixed given value & [K], and one where A1()) Unf ([K)). Given the dataset Diors
we optimize the following loss:
1 - In (z((0) (2
(, 0) =
(2)
nt
:=1 Our thcoretical analyses be extended to other complexity metrics such Rademacher complexity.
4

---

Published as a conference paper at ICLR 2024
Temporal Contrastive Learning. Finally, this approach trains the decoder o to learn to separate pair of temporally causal observations from a pair of temporally acausal observations. We collect a dataset of Dremp {xx(1) k(O), xd(i), x i=l
[n/2] tuples using the multistep transitions dataset Dior-
We use 2 multistep transitions to create a single datapoint for Premp to keep the datapoints inde- pendent. To create the jth datapoint for Dumm. we use the multistep transitions (x:(2), ,,(20) _,(2)))211111) _2() and (z(2++1), A:(2++1), 11(2++1) and sample z(0) ~ Unf({0, 1}). If 2 = 1, then our 7th datapoint otherwise, it is an acausal observation pair (I(22, k(2x), 22(++111, z0)).. Depending on how we sample k. we collect different dataset Pior- and accordingly a different dataset Dump. Given the dataset Dumpp we optimize the following loss using regression model belonging to model class G = {g: X x [K] x X
is causal observation pair (x(()), ,K(2)) ,z((2),
[0, 1]::
[n/2] Σ
(50) - g(0(x( A,0 ),00, )(())) 2,,00))011))))
exmmp (g, ④)
n
Practical Implementations. We use the aforementioned description of methods for theoretical analysis. However, their practical implementations differ in few notable ways. Most importantly we either use a continuous vector representation Φ X - Rd for modeling Φ, or apply a Vector Quantized (VQ) bottleneck (Oord et 2017) on top of the vector representation to model a discrete- representation decoder. We also optimize the loss using minibatches and use square loss for training forward modeling and SimCLR loss (Chen al., 2020) for contrastive learning. We experimentally show that our theoretical findings extend to these practical implementations.
4 IS VIDEO BASED REPRESENTATION LEARNING PROVABLY CORRECT?
In this section, we present our main theoretical results. We first prove that both forward modeling and temporal contrastive methods succeed when there is no exogenous noise. We then establish a lower bound showing that video-based representation learning is exponentially harder than trajectory- based representation learning. We defer all proofs 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
(xf) .42 x " of n independent and identically distributed (iid) videos sampled from im1
data distribution D. We denote the probability of a video as D(1 T2. ... , We assume that D is generated by mixture of Markovian policies pI, i.e., the generative procedure for D is to sam- ple a policy π E I 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 vidco data has good state space coverage and that the data is collected by noise-free policies.
Data Collection). There exists Mmin > 0 such that if is state reachable at time step h by some policy in II, then D (Φ" (It)) = 8) Z mmin- Further, we assume that every data collection policy π E II noise-free, i.e., xaa Tn. [ d² (xx.) for all (a, Tn).
Assumption (Requirements
Justification for Assumption In practice, we expect this assumption to hold for tasks such as gam- ing, or software debugging, where video data is abundant and, therefore, can be expected to provide good coverage of the underlying state space. This assumption 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 (IL.)) (Torabi et al., 2019), we don't '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 in the data can be different
unrelated to the downstream tasks want to solve. As such, the video data do not provide demonstrations of the desired task. Further, as the video data is typically generated by humans, we can expect the data collection policies be nreee-free, as these policies realized by
5

---

Published as a conference paper at ICLR 2024
humans who would not make decisions based on noise. E.g., human player is unlikely to turn left due to the background motion of leaves that is unrelated to the game's '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 open question whether finite-sample bounds exist for them and we leave their theoretical analysis to future work and instead evaluate them empirically. In addition to the decoder class 4, we assume function class F to model f for forward modeling and to model g for temporal contrastive learning. We make a realizability assumption on these function classes.
Assumption 2 (Realizability). There exists f* E F. (for k) = for x, k) and g (z (remp(2), k, x²) = P temp(z = 1 2, ki x') on the appropriate support, and where Pjor and P. cemp are respectively the Bayes classifier for the forward modeling and temporal contrastive learning methods.
G and Ofors Oremp E & such that (x''
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.
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 truc.
Assumption (Margin Assumption). We assume that the margins Biar and Bromp defined below:
Bgur
inf $1,82ES,
E₁ [X (1, Pjor 182, k82,))
1 En.o' [[Promp(= = 1 81, ,i, 80)
8romp
inf $1.8268.81
Premp( = 82, ki 5')],,
are strictly positive, and where in the definition of Stemp. we sample s' 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 (Biz) for forward modeling and (Suen) for temporal contrastive learning. A common scenario where these assump- tions are true is when for any pair of different states 81, 82, there is a third state s' that is reachable from one but not the other. If the video data distribution D supports all underlying transitions, then this immediately implies that ||P for (X' 81, k) Proe(X') 82, k) |Ty > 0 which implies 3ror > O. 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 Sremp > 0.
We now state our main result for forward modeling under Assumption 1-3.
Theorem 1 (Forward Modeling Result). Fix = > 0 and 6 E (0, 1) and let s be any prov- ably efficient RL algorithm for tabular MDPs with sample complexity Tramp (S, H. E. If n is poly (S, H, 1/ 1/Be.; l/c,
In(1/), In [FI, In [61 for a suitable polynomial, then forward
modeling learns decoder &. X - [[8]]. Further, running o on the tabular MDP with n namp (S, A, H,T, T, e/2, 5/4) episodes returns a latent policy S. Then there exists a bijective map- ping a S [[S]] such that with probability at least 1 - 5 we have:
Pangtola) (ō(z) ass) (z)))= =8) 2 3
In [ 8
+
Vs S.
W 3for
and the learned observation-baser policy @ x p()(x)) is E-optimal, i.e. V(*))
S
Finally, the number of online episodes used in the downstream RL task is given by Ns A, H, e./2,8./4 and doesn't 't scale with the complexity of function classes and F.
The result for temporal contrastive is identical to Theorem but instead of Bior we have Biemp and instead of J we have 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 In |a]. Typically, the function class for forward modeling F is much more complex than G. however, as we show in Appendix B.5, the margin for forward modeling Ptar is larger than for contrastive learning Semp leading to trade-off between these approaches.

---

Published as a conference paper at ICLR 2024
4.2 LEARNING FROM VIDEO IS EXPONENTIALLY HARDER THAN LEARNING FROM TRAJECTORY DATA
When online RL is possible, there exist algorithms Misra et al. (2020); Efroni et al. (2022) that can learn an accurate latent state decoder & with high probability and use it to learn near-optimal policies. These methods train the decoder using online trajectory data. This begs the following question: Is it possible to learn a latent state decoder that is useful for performing RL using offline video data? As the next result shows, this is not always the case.
Theorem 2 (Lower Bound for Video). Suppose |A], H 2 2. Then, for any E (0, 1), any algorithm d1 that outputs state decoder é with dn X [T], I, 21/4-11 wh E [H] given video dataset DD sampled from some MDP and satisfies Assumption I, and any online RL algorithm d2 uses that state decoder in its interaction with such an MDP (i.e., 22 only observes states through 6) and output policy R. there exists an MDP instance M in a class of MDPs which satisfies Assumption 3 and is PAC learnable with ((poly((S), |A), H, 1/5) complexity, such that
VM(( Vis >
regardless of the size of the video dataset DD for algorithm d and the number of episodes of inter- action for algorithm s/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 image consisting of N x 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, cannot tell which maze is the one we are controlling and must memorize the configuration of all N x 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 Appendix B.3. We also discuss settings where we may be able to efficient-learning with just video data with additional assumptions in Appendix B.4.
EXPERIMENTAL RESULTS AND DISCUSSION
We empirically evaluate the above rideo-based representation learning methods on three visual en- vironments: gridworld environment and two VizDoom environments. We defer the results one of the Vizdoom environments along with additional experimental details and results to Appendix C. Our main goal is to validate our theoretical findings by evaluating these methods in the presence and absence of exogenous noise and comparing their performance with a trajectory-bast based method.
5.1 EXPERIMENTAL DETAILS
GridWorld. We consider navigation in a 12 12 Minigrid environment (Chevalier-Boisver) al., 2023). The agent (red triangle) can only observe an area around itself, and the goal is to reach the key quickly (Figure 3). The position of the agent and key randomizes each episode.
ViZDoom Defend the Center This is first-person shooting game (Wydmuch et al., 2018; Kempka et al., 2016), in which the player needs to kill a variety of monsters score (Figure 5). The episode ends when the monster is killed or after 500 steps.
Exogenous Noise. For all domains, the observation is an RGB image. We add exogenous noise to it by superimposing 10 generated diamonds of 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 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 T is a convolutional neural network. We use a deconvo- lutional neural network to model f and h. We experimented with both using a vector representation for o and also using a VQ-bottleneck to discretize the embeddings. We use PPO to do downstream RL. and keep & frozen during the RL training. We also visualize the learned representations by
7

---

Published as a conference paper at ICLR 2024
training decoder them and fixing reconstruct the input observations. We then look 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 & by predicting action given pair of observations E [In p(a,, " Th+ k)]. ACRO is designed filter out exogenous noise this information not predictive the action. Our goal is to test we get much better representations if we have access trajectory data instead of video data.
ACRO
temporal contrastive
forward modeling
VQ-farwark moceling
autoencocer
VC-autoencoder
3.0
1.0
0.8
Return
-02
-0.2
-0.2
Timestegs
Timesteps
Timesteps
Timesteps
(a) No Noise
(b) Only Observation Noise
(c) Only Reward Noise
(d) Both
Figure 2: RI. experiments in the GridWorld environment.
(a) Original
(b) Forward Model (c) Autoencoder
(d) Contrastive
Figure 3: Decoded image reconstructions for different methods in the GridWorld environment. We train a reconstruction model on top frozen Iearned representations trained with given video- based method. Top row: shows example from the setting where there is exogenous noise. Bottom row: shows an example with exogenous noise (colored diamond shapes).
5.2 EMPIRICAL RESULTS AND DISCUSSION
We present our main empirical results in Figure and Figure and discuss the results below.
forward modeling
temporal contrastive
autoencoder
ACRO
Retum
2500
3002
1000
1000
3000
4000
1000
2200
3002
4000
1805
2003
3500
4000
Timesteps
Timesteps
Timesteps
Timesteps
No Noise
Only Observation Noise
(c) Only Reward Noise
Both
Figure 4: RL experiments using different latent representations for the ViZDoom Defend the Center environment.
8

---

Published as a conference paper at ICLR 2024
(a) Original
(b) Forward Modeling
(c) Autoencoder
(d) Temporal Contrastive
Figure 5: Decoded image reconstructions for different methods in ViZDoom Defend the Center.
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 (Figure 2) and ViZDoom Defend the Center (Figure 4), these approaches learn a decoder ø that lead to success with RL in the absence of any exogenous noise. For GridWorld. we find support for this result with VQ bottleneck during representation learning (Figure 2(a)) whereas for ViZDoom Defend the Center, find support for this result even without the use of a VQ bottleneck (Figure 4(a)). These results are further supported via qualitative evaluation through image decoding from the learned latent representations (Figure 3) which show that these representations can recover critical elements like walls. We find that autoencoder performs well in ViZDoom Defend the Center but not in gridworld, which aligns with a lack of any theoretical understanding of autoencoders.
Performance with exogenous noise. We find that in the presence of exogenous noise (Figure 2, Figure 4), representations from forward modeling achieve lower performance specially in grid- world, whereas temporal contrastive representations completely fail. One hypothesis for the stark failure of temporal contrastive learning that the agent can tell whether two observations are causal or not, by simply focusing the noisy diamonds that move in predictive manner. Therefore, the contrastive learning loss can be reduced by focusing entirely on the exogenous noise. Whereas, forward modeling is more robust as it needs to predict future observations, and the agent's 's state is more helpful for doing that than noise. This shows in the reconstructions (Figure (db) Figure 5(b)d)). As expected, the reconstructions for forward modeling continue capture state-relevant information, whereas for temporal contrastive they focus on noise and miss relevant state informa- tion. In Appendix B.6, formally prove that there exists 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.
Comparison with ACRO. Finally, we draw
pixel
pixel
pixel
pixel size:
comparison between the performance of andeo-pretrained representation and ACRO which uses trajectory data. ACRO achieves the strongest performance across all tasks (Figure 2, Figure 4). Additionally, we also observe that as we increase the size of the exogenous noise elements tion space (Figure 6). the performance of forward modeling, the overall best video- based approach, degrades more drastically compared to ACRO. This agrees with our theoretical finding (Theorem that learn- ing representations from video-based data is significantly harder than trajectory-based data when exogenous noise is present.
0.5
0.8
0.6
0.2
0.2
0.0
0.0
the observa-
Timestops
Timestaps
(a) ACRO
(b) Forward Modeling
Figure 6: RL performance with varying size for ex- ogenous noise in the GridWorld environment.
CONCLUSION
Videos are a naturally available source of data for training representations for RI. In this work, study whether existing video-based representation learning methods are provably effective for downstream RL tasks. We provide both upper and lower bounds for these methods in two theoretical settings and provide empirical validation of our findings on visual domains. Using our theoretical tools to develop better ideo-based representation learning methods and extending our analysis to other formal settings are natural future work directions.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS.
We thank Sam Devlin, Ching-An Cheng, Andrey Kolobov, and Adith Swaminathan for useful dis- cussions. This work was done while AS was postdoctoral researcher at Microsoft Research New York.
ETHICS STATEMENT
In our paper, we run experiments on two open-source simulated RL environments. All data was collected in simulation and real-world dataset used in this work.
REPRODUCIBILITY STATEMENT
The code is publicly available https: /tt//ggtthu//b www/microsoot http///irrooootttttttotieee We used publicly available RL environments for our simulated experiments and used pretrained or random- ized policies for data collection as described in Appendix C.
REFERENCES
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural com- plexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095-2011 2020.
Arthur Aubret, Markus R. Ernst, Céline Teulière, and Jochen Triesch. Time augment self- supervised visual representation learning. The Eleventh International Conference on Learning Representations, 2023. URL https: http////peeeeeeeeeeeeeee/e. /
net tatee/ffrrueee cacrruumiio@@eeee
81.
Bowen Baker, 1lge 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.
Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information prioritization through empowerment visual model-based rl. arXiv preprint urXiv:2204.08585 2022.
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 leamners. Advances in neural information processing systems, 33:1877-1901 2020. 2020.
Jinglin Chen and Nan Jiang. information-theoreti considerations in batch reinforcement learning.
In International Conference on Machine Learning, pp. 1042-1055. PMLR, 2019.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning visual representations. International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Maxime Chevalier-Boisse Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306.1388 2023.
Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems,
Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning, pp. 1665-1674. PMLR, 2019.
Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Provably filtering exogenous distractors using multistep inverse dynamics. In International Con- ference on Learning Representations, 2022. URL https: openrevieww net forum? i FF-
10

---

Published as a conference paper at ICLR 2024
Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Young- woon Lee, Danijar Hafner, and Pieter Abbeel. Video prediction models as rewards for reinforce- ment learning. arXiv preprint arXIv:2305.14343, 2023.
Sara A Geer. Empirical Processes in M-estimation, volume 6. Cambridge university press, 2000.
Wonjoon Goo and Scott Nickum. One-shot learning of multi-step tasks from observation via activity localization in auxiliary video. In 2019 international conference robotics and automation (ICRA), pp. 7755-7761. IEEE, 2019.
Zhaohan Guo, Shantanu Thakoor, Miruna Pislar, 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-315 -1870, 2022.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. arXiv preprint arXiv:19/2.01603, 2019. URL https: //a
01 603 .ddf.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.044 041004 2023.
Riashat Islam, Manan Tomar, Alex Lamb, Yonathan Efroni, Hongyu Zang. Aniket Didolkar, Dipen- dra Misra, Xin Li, Harm van Seijen, Remi Tachet des Combes, al. Agent-controller represen- tations: Principled offline with rich exogenous information. arXiv preprint arXiv:2211.00164 2022.
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.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Ma- chine learning, 49:209-233 2002.
Michat Kempka, Marek Wydmuch, Grzegorz Rune, Jakub Toczek, and Wojciech Jaskowski. Viz- doom: 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.
Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed discovery of control- lable latent states with multi-step inverse models. arXiv preprint arXiv:2207.08229 2022. URL https: //trx//avvvvv//////////eee/aaavv org/ pdf////////// 2207. 08229. pdf.
Xudong Lin, Gedas Bertasius, Jue Wang, Shih-Fu Chang, Devi Parikh, and Lorenzo Torresani. Vx2text: End-to-end learning of video- text generation from multimodal inputs. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7005-7015, 2021.
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. / 1692, 2019.
Zakaria Mhammedi, Dylan 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.
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:/ /ate:///ppeneee /pperevieen net /tor/oormm//
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state ab- straction and provably efficient rich-observation reinforcement learning. In International confer- ence on machine learning, pp. 6961-6971. PMLR, 2020.
11

---

Published as a conference paper at ICLR 2024
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near Near-optimal representation learning for hierarchical reinforcement learning. In International Conference on Learning Representa- tions, 2018.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. arXiv preprint 1xXiv:/711.00933: 2017.
Nikhil Parthasarathy, SM Eslami, João Carreira, and Olivier J Hénaff. Self-supervised video pre- training yields strong image representations. arXiv preprint arXiv:2210.06433, 2022.
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-88 8763. PMLR, 2021.
Harshit Sikchi, Akanksha Saran, Wonjoon Goo, and Scott Nickum. A ranking game for imitation learning. arXiv preprint arXivv2202.03483 2022.
Vlad Sobal, Jyothir SV, Siddhartha Jalagam, Nicolas Carion, Kyunghyun Cho, and Yann Le- Cun. Joint embedding predictive architectures focus on slow features. arXiv preprint arXiv:2211.10831, 2022.
Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video rep- resentations using Istms. In Francis Bach and David Blei (eds.), Proceedings of the 32nd Inter- national Conference Machine Learning, volume 37 of Proceedings of Machine Learning Re- search, pp. 843-852, Lille, France, 07-09 Jul 2015. PMLR. uttp:: httpas://ppccceeeeeee /proceedings. mlr. w.l.ppress/ve /v7 srivastaval5 hhmml.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chan- dak, 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 Proceed- ings 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: / //pppoeeeingg.... mlr. press/v200 v222 http://aannnnaa/ce.. html.
Faraz. Torabi, Garrett Warnell, and Peter Stone. Recent advances in imitation learning from obser- vation. arXiv preprint arXiv: /306. 2019. URL /tctps://arxiv.orrt
905.
13566. pdf.
Masatoshi Uchara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. arXiv preprint 10.04655,, 2021.
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: /atp:///axvvvvvivv org/ wttp://wwwwpppp////ee//////////
pdf.
Marek Wydmuch, Michal Kempka, and Wojciech Jaškowski. Vizdoom competitions: Playing doom from pixels. IEEE Transactions on Games, 11(3):248-22) 259, 2018
Weirui Ye, Yunsheng Zhang, Pieter Abbeel, and Yang Gao. Become a proficient player with lim- ited data through watching pure videos. In The Eleventh International Conference on Learning Representations, 2022.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint 12X1:22006.10742, 2020.
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: //apii. semanticon manniccshollar wwtp/oorpgooor....oooo. :ttp//CooSIII::5235500000
12

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
A Additional Related Work
13
B Proofs of Theoretical Statements
14
B.1 Upper Bound for the Forward Model Baseline
17
B.2 Upper Bound
Temporal Contrastive Approach.
25
B.3 Proof of Lower Bound for Exogenous Block MDPs
28
B.4 Can get efficient learning under additional assumptions?
30
B.5 Relation Between Margins...
30
B.6 Why temporal contrastive leamning is more susceptible to exogenous noise than forward modeling
32
C Additional Experimental Details
36
C.I Details Experimental Sctup
36
Results on Additional Domain
C.3 Additional Ablations.
36
A ADDITIONAL RELATED WORK
Representation Learning for Reinforcement Learning A line of research recurrent state space models essentially concerned with the next-frame approach, although typically with condi- tioning on actions. Moreover, model uncertainty in the observations, latent variable with pos- terior depending on the current observation (or even sequence of future observations) is typically introduced. (Ke et al., 2019) considered learning such sequential prediction model which pre- dicts observations and conditions actions. They used latent variable with posterior depending on future observations model uncertainty. These representations were used for model predictve 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 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 temporal contrastive objective. Sobal et al. (2022) explored using semi-supervised objectives for learning representations in RI, yet used action-labeled data. Wang et al. (2022) used a decoupled recurrent neural network approach to learn to extract endogenous states, but relied action-labeled data to achieve the fac- torization. Deep Bisimulation for Control (Zhang et al., 2020) introduced an objective to encourage observations with similar value functions 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- 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.
further line of work from theoretical reinforcemnt learning has examined provably efficient ob- jectives for discovering representations. Efroni et al. (2022) explored representation learning in the presence of exogenous noise, establishing sample efficient algorithm. However Efroni et (2022) and the closely related work on filtering exogenous noise required actions (Lamb et al., 2022; Is- lam 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 vidcos has long history. Srivastava et al. (2015) used recurrent neural networks with pixel prediction ob-
13

---

Published as a conference paper at ICLR 2024
jective future frames. Parthasarathy 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 3D space, showing that this outperformed standard augmentations used in computer vision.
Using Video Data for Reinforcement Learning The VIPER method (Escontrela 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 small dataset of Minecraft videos and used the model 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 using videos for representation learning 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 single video segmented with action labels (one-shot learning). Sikchi et al. (2022) propose a two-player ranking game between policy and 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 the learning from observation (no actions) setting.
Recently some approaches have also considered recovering latent actions from video data using an encoder-decoder approach (Ye al., 2022). In general, the lower bound in Theorem 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 more predictive of changes in the observations. However, in simpler cases such as 3D games, where the agent's '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.
B PROOFS OF THEORETICAL STATEMENTS
We state our setting and general assumptions before presenting method specific results. We also include table of notations in Table 1.
Notation
Description
[N] A(U)) Unf(l() supp(P)
Denotes the set (1,2, N
Denotes the set of all distributions over a set 4
Uniform distribution over 24
Support of distribution P E A(U), i.e., supp(PP) = (a E4 P(x)) > 0}.
x
Observation space
Latent endogenous state
A
Action space
T:S
+ A(S)
Transition dynamics
R:SxA A [0
Reward function
Start state distribution
H
Horizon indicating the maximum number of actions per episode
Endogenous state decoder
Table 1: Description for mathematical notations.
(1f ( (3)
We given dataset D -
of independent and identically distributed
;-1 (iid) unlabeled episodes. We will use the word video and unlabeled episodes interchangeably. We assume the underlying data distribution as D( .H). We assume that D is generated by mixture of Markovian policies II, ie., the generative procedure for D is to sample a policy æ E II, with probability O₂ and then generate entire episode using it. For this reason, will denote D O o Np where O is the
D. We denote the probability an unlabeled episode
14

---

Published as a conference paper at ICLR 2024
mixture distribution. We assume direct knowledge about either II or O other than that the set of policies in II are Markovian. We define the underlying distribution over the action-labeled episode as ( 01, F22.... HI, an), ) of which the agent only gets to observe the (I, ,Zzr). 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 be trivially accomplished by simply con- catenating 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 E [H] and k E N we have:
D(rN+E = 1 2r x) D( = = 1 [x1= =r)
(4)
We denote (Ix,) to define the marginal distribution over observation The and D( Th+*) to denote the marginal distribution over a pair of observations (I), Tb+4) in the episode. We similarly define D( a,) as the distribution over observation action pairs (Ih, ax).
We assume that the video data has good coverage. This is stated formally below:
Assumption 4 (State Coverage by D). Given our policy class IT, there exists an Mmii O such thatif sup+en Pa (5)+ s) > for some s E S, then we assume D (In) = 8) Z Mmin-
In practice, Assumption 4 can be satisfied since videos are more easily available than labeled episodes and we can hope that 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 E X, its endogenous state is given by o (x) E S. Assumption 5 (Noise-Free Video Distribution). For any h, x E II, 24, E supp P. and a E A. we have
<(a un) = Taa d (2n)).
Justification of Noise-Free Policy. Typically, video data is created by humans. E.g., a human may be playing 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 (z, x²) 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 is to denote the time steps by which these observations differ. Formally, we will call (x,, k x') as multi-step transition where & was observed at some time step h, and x' was observed For the fixed multi-step variant, we have k > 1 but & is fixed. Finally, in the general multi-step variant, we will assume that & is picked from Unf([K] where K is a fixed upper bound.
+ k. For the single-step variant of the algorithms, we have 1.
Extending episode to H + K. When using k > 1, we may want to collect a multi-step transition (r, k,a') where x zF to allow learning state representation for time step H. However, this point, we don't 't have time steps left to observe # We alleviate this by assuming that we can allow an episode to run till H + K if necessary. In practice, this 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 is selected at k- In our analysis ahead, make the former setting that the episodes can be extended to H + K, but can be easily rephrased to work with the other setting.
For both the forward model and the temporal contrastive approach, we assume access to dataset Diarr ((x(0) 100 _)())) n [_ of pairs of observations. We define few different distributions that
15

---

Published as a conference paper at ICLR 2024
can be used to generate this set. For given & E separate observations as:
[K], define distribution Dj. over k-step
1 H
-x=x=xX'' = x, = =x)
Σ 1=1
D(r 2h = x²)
(5)
We can sample (z, (,, x') Dx by sampling an episode (21, .HH D, and then sampling a h ~ Unf((H]), and choosing x = Ih and x' 2h+k-
We also define a distribution Dunt where we also sample uniformly over available choices:
1 Dunt (X K, = x') K Dulan z, 2
(6)
We can sample (xr, k. x') ~ Dunt( (X, X!) by sampling an episode (x), ,22,"" , .H) D, and then sampling h E [H). and sampling &s E [K]. and choosing (ah, Th+kk as the selected pair.
We define a useful notation p E (2') as:
H D(m = =x)
p(x = x)
H
h=1
The distribution p(X) is good distribution to sample from it covers states across all time steps. Finally, because of Assumption 4, we have the following:
Mmnn H
Vs E S,
p(s)
(8)
This is because we assume every state s E S, visited at some time step f, and so we have D(s:
s) 2 Mmin and p(s) = # Σd (ss = s) 2 t D(s = s) 2 ln..
It can be easily verified that for both D& (X, X') and Dunt (X, X'),,, their marginals over X is given by P(X). Both D2 and Dunt satisfy the noise-free property. We prove this using the next two Lemma.s
Lemma 1 (Property of Noise-Free policy). Let π a policy such that for any x X, we have (a x) = Ta o² (x))- Then for any h [H] and k E [K] have P, (In+k x' rn x))
only depend on o² (x) and this common value is defined by P. (Z 8h, = (x)).
Proof. The proof is by induction on k. For k 1 we have:
Pa(an+1 = x' x)
Σ
T(a' xr a) In =x)= = Σ T(d T (1)))) = ==)
aEA
aEA
and as the right hand side only depends on 6" (x), the base case is proven. For the general case, we have:
Pa (* 2' I xx) Σ
P ( 1h4k-1 = =
x:)
BEX
Σ
P ( = 1 = 2)P (I = = = x))
xEX
Σ
P. (xntt = | P f In Φ² (x)),
BEX
where the second step uses the fact that π is Markovian and the last step uses the inductive case for k 1.
Lemma (Distribution over Pairs). Let l E K], x E supp pXX), then the distribution D&(X' |x)
only depends on o= (x). This allows to define (x o² (x)) this common value. Similarly,
the distribution D (X' x, k) depends only on o- (x) and We define this common value as
D X' (x),
16

---

Published as a conference paper at ICLR 2024
Proof. For any h we have:
1
D&(X = x, = =x)=
D 20 =
H
h=1
Σ
(r = =2,
Σ
(z = P()
= 2 2/ zx)
h= aello
Σ
= x'1 118h = 63 (z)),
(using Lemma 1)
H
hElnellp
q(> (x))
H
Σ
OxPx (En)) = sn (x)
kEIl
The marginal Dx(X= = x) is given by:
H1
g(x (x))
H
Da(X = x) H Σ
+ Ong(z 10 (x))
Σ h=1
DR(aa = (x))
H
h= nEllo
The conditional Dx =
X x) is given by:
Dx(X'= = L' 1 = =x)= D&(X = xx,X'=x''
D2(x)
2h-1 Enetl O (sh, = (x))) x'
sh ('(x))
Σ Ds(s) = (x))
Therefore, the conditional Dx x X only depends on o (z(x), and we define this common
value D& (X' = x' | (x)).
The proof for Dunt is similar. We can use the property of D3 that we have proven to get:
Duni = X
Dunt(X' =xx x' |x k,
Etex Dunt = x, X'=
D&(X =x, =
Enex Da(X x' =
Dz(X' = X x)
ΣFEx Dz(X' x=
Da(X' = X = x=
EREX Da(X'= = 0 X = (x))
Therefore, Duni( (X' = x' X = x. k) only depends on 6² (z). We will define the common values as Dune( X' x' s 6² k))
Lemma 2 allows depends on the latent state.
define De(a' (x))) and Dani(a' d)(x),)), as the distribution only
B.1 UPPER BOUND FOR THE FORWARD MODEL BASELINE
Let Dior = (x((1) (,0,, (0)) In be pair iid multi-step observations. We will collect this dataset in one of three ways:
1. Single step (k 1), in this case we will sample (2, s((i)) D&(X, X'). As explained before, get this sample using the episode data. We save (x((1) ki (1)) our sample.
17

---

Published as a conference paper at ICLR 2024
2. Fixed ste. We use fixed k >> and sample x0(1)) Dx(X,X),, We save (z(1) k 2.()) as sample.
3. Variable multi-step. We sample (x, k x²) ~ Dunt X, kr X') and use it as our sample.
We will abstract these three choices using general notion of Dpr A(X [K] X). In the first two cases, we assume we have point-mass distribution over and given this k, we sample from Dx(X,X'), We will assume (x(*, k00 x²((0) Dpr. We can create Dior from the dataset D of n episodes sampled from D using the sampling procedures explained earlier. Note that as marginals over both (X) and Dun((X) is p(X), therefore, the marginals over Dxr(X) ) is also p(X). Additionally, we will define Dpr (k) as the marginal over which is either point-mase mast in the first two sampling procedures and Unf([K] in the third procedure.
We assume access to function classes. The first is decoder class &y X given number that satisfies N 2 [S]. The second is a conditional probability class F= [N] x [K] (xx).
[N] where Nis
Assumption 6. (Realizability o and F) We assume that there exists 6f E &x and f E such that (xc (x), k) prr I k) Ppr ( (x),,, for all (x, Dpr (s.))
This assumption firstly non-vacuous as Dpr | xx) Drr x0 (x)).. and therefore, we can apply a bottleneck function & and still assume realizability. For example, we can assume that ø is the same as ø* up to the relabeling of its output, and f(x' ( = (x' s).
Let f. Ø be the empirical solution to the following maximum likelihood problem.
1 Σ In (z)00 (((a!)), (())) i=1
F = arg max
[EF DE TE
Note that when & is fixed (we sample from D&), then information theoretically there is no advantage of condition on k and it can be dropped from optimization.
As in realizable setting (Assumption 6), we can standard maximum likelihood guaran-
tees to get the following result.
Proposition 3 (Generalization Bound). Fix 5 (0, 1), then with probability at least - 6, have:
|DDx(XX |x, f(x') ((x)),)) s (n;
B(a,k))D
where (n; 2 In [01119) n
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 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 Bjor (0, 1) such that:
inf ER 81 St,
[p (X' (81,k) - Drr (X' 82, k)llyv] 2 3oor
Note that this defines two types of margin depending Dpr. When Is is fixed value, the margin is given by:
inf
[X' $1,, k) - Dpr (X' 82, kllly
For
When we sample ls Un]([K]) then the margin is given by:
inf Saat
IID $1, k)) D (X' 82, k)
K.
We will use the abstract notion Boor for forward margin which will be equal to 3 af5) depending on our sampling procedure. It is easy to see that 2(w)
R Σk ((k)
We now ready state first main result.
18

---

Published as a conference paper at ICLR 2024
Proposition 4 (Recovering Endogenous State.). Fix d (0.1), then with probability at least 1 8 we learn & that satisfies:
())
P PRIREr (0((z) # (22) A x(21) = )(12) -
3jor
Proof. We start with a coupling argument where we sample 22 independently from D (X)
which is the same as p(X).
Bay m2 D An
1 (((m)= = 2((2)) p1 (X' x1, Prr (X' x2, 12)l
(0(0)) = )(12) / |(x' (e(z)),k) Dpr(X' I1, k)
Ex Drr kn Drr
TV.
1 10((z))= - 3(12)
|(X' ((x1), - Dpr(X' 12 kb)
+ E2r or
IV.
We bound these two terms separately
[1 {ō(z)) = {(22)} || |(X )(x)), - (X' 21, k3) [3)
E Dwp
Ben Der || (0(1)),)) - Dpr x bb)) TV.
| (0((z)) Dow
((02)))
BeakesDD ||5XX | ((x) - Dxr xx)) |
- -V Vee [ ((0001))
= ((02)}|
< A.
i (0(00))
(12)) }} and the second step uses
where V Ea
inequal-
ity. It is straightforward to verify that & [0, 1]. We bound the second term similarly
Ex Bo pr
[ -((s))
-
} | ((I1),k) - Dpr (X' 12, (2,2)) k)
Iv.
1 (((z)) - } f(X' 0(12), k) Prr (X' x22 k)
En
[T
S.
where the second step uses the crucial coupling argument that we can replace oi with x2 because of the indicator {ō(z)) = ((22) j. and the last step follows we reduce it to the first term except we switch the names of 1ı and x2. Combining the two upper bounds we get:
1 (ō(z1) = 2(a2)) |1 k1, (xr(X') x2, k) S 26
B2n.a Dar
or, equivalently,
Eg Epn, 1 (0(z1)=
((6))=āle))) 1
0(12)) Dpr [| k) Dpr (X' Dx 12, k)lv :=((=1)22)
S 26
Let [(x1, = EN-D,T prr (X' I1, k)- - Dpr | x2, kt [[ For any two observations, 6² (21) = 6² (r2), then (X' (X' mi) - (X' x2) 32)l = O, and therefore, ((21, =
because of Lemma 2. Otherwise, T 222) is at least Stors by Assumption 6. Combining these two observations we get:
[(2,, x2) (0 (zx))
Combining the previous two inequalities we get:
[ ( = 202)) (x1) # (022) <
2b A Phor
E Dpr
19

---

Published as a conference paper at ICLR 2024
This directly gives
P Py (ō(m)) = 0(12) A d= (21) + (x22)
26A
24
Blve
The proof is completed by recalling that marginal (X) is the same as p(X).
Proposition 4 shows that the learned ø has one-sided error. If it merges two observations, then with high probability they are not from the same state. As N |S], we will show below that the reverse is also true.
Theorem 5. If = (S), then there exists a bijection a : [N] - such that for any E S we have:
P (a (ō(z) ((s) 1 (x) = 42
Umin Stor
provided A
Matt N
Proof. We define few shorthand below for any j E [N] and & ES
F(j,8) 8) P. soup
(((x) j A o0 (z)= 3 p(]) Pamp (6(z) =})
p( Pamp (0 (I) 8)
is easy to verify that P(j, 8) joint distribution with p(j) and p(3) as its marginals.
Fix i E [N] and E S.
Psi xrpp (ō(z)) A (1)))
-1)) =j Paaaao USES.JEIN) = #8})
2 Pa PAI (δ(z)) = in =
= ii (z)) = = #s)
Paioo
= in (o(a))=ind)
(11) = PIYMP (0(2)) = i1 in (x2) #s)
(o(z)= - = (x) =)
Panp (Ō(z) = i = =s) -Panp(@aa=iin0tt)
= PPap
WES
sP(i,s) Σ P((, - -P(,s)) -((()))
EEES
= =P( s) (p(1) (i),
Combining this with Proposition 4, get:
24
Vie [N],SES, Pii, P(i, 8)) < A'
Bive
where we have used a shorthand A' 22 /lor. We define mapping a - [N] where for any E S:
a(s) = arg max s)),)) JE(N]
(10)
We immediately have:
N
(()(ss) = max P(J 8)
P(j,s)
1 N p(s) 2
7min NH
(11)
JEIN]
N
where we use the fact that max greater than average the first inequality, and Equation 8. Further, for every E S, we have:
8() (p(a(s)) P(a(ss) AA.
20

---

Published as a conference paper at ICLR 2024
Plugging the lower bound s)) 2 NH we get:
NHN
P(a((s), 8) 2 p(a(s))
(12)
Mmnin
We now show that if A <
Tie 2 then a(s) is bijection. Let 81 and 82 be such that a(s))
(182) = Then using the above Equation 12 we get (P(1, 81) 2 p(t) NHA and P(i, 82) 2
p(t) NHA! We have:
2NH
p(i)
Σ BES
Pii, 8) 2 8(), P(1, 2 2p(1)
M/min
This implies 2NA' p(t) but p(i) p(a(s1)) P(a((1), s1) 2 (Equation 11), get
Tnii
NH
However, as we assume that A' < H therefore, this is a contradiction. This implies a(s1) & a(s2) for any two different states and $2. Since we assume |N] |S]. this implies a bijection.
2NHA'
Jun Yui NH or A' 2
Vinn 2N
Mann
Fix E S and let i # a(s). As a is bijection, let 3 = a (i), we can show that P(i, s) is small:
p()) - P(i, p(a(3)) - P(a((), < NH
P(i,s)
(13)
7mmn
where we use and Equation 12.
This allows us to show that P(((()) 8) is high as follows: P(((s).s)
P(((s),s)
P(a((s 8) p(s)
P(a(s), + > 2i=1,ipa(0) P(),
P(((s),s)
p(a(s)) + N-HA
Unin
p(a(s)) NHA
p(a(s)) + N
N + NHA
Vinia
Yhm
p(a(s))
NA
2N H
> 1
Mmmin (aanP((((8)) (o()))
H
2
Umin
where the first inequality uses Eq. (13) and p(a(s)) 2 P(a(s), s), second inequality uses Eq. (12), and the last step uses p(a(s)) 2 ((()()),)) 2 NH
Nam
The proof is completed by noting that:
P Pang(-1a)
(5(x) a(s) =Pazp (ō(z)= = a(8) ()))) PP(((s)) 1s).
Let of be PAC RL algorithm for tabular MDPs. We assume that this algorithm's sample complexity given by nsamp (S, A, H, e. 5) where S and A are the size of the state space and action space of the tabular MDP, H is the horizon, and (s, 8) are the typical PAC RL hyperparameters denoting tolerance and failure probability. Formally, the algorithm of interacts with a tabular MDP M for Msamp (, A, 8,,80) 5) episodes and outputs policy ; S [H] A such that with probability least 1 - we have:
sup VM)(p) Vxa(T) S E, WEW
21

---

Published as a conference paper at ICLR 2024
where Vau is the space of all policies of the type S [H] + A.
We assume that we are given knowledge of the desired (e, 6) hyperparameters in the downstream RL task during the representation pre- 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 S, action space A, horizon H, transition dynamics T, reward function R, and start state distribution of j.. If we directly had access to this latent MDP, say via the true decoding function Φ-, then we can apply the algorithm a' and learn the optimal latent policy 3o which we can couple with ø² policy as
and learn the optimal observation- based policy. Formally, we write this observation-base i o² X [H] + A given by (z)* We dont have access o², but we have access to & that with high probability for a given x outputs a state which is same as o² (z) up to the learned a-bijection. We, therefore, define the induced MDP M as the finite MDP with state space S, action space A, transition function T. reward function Nt and start state distribution fī. These same as the latent Block MDP but where the true state s is replaced by (s)). It is this induced MI that the tabular MDP algorithm of will see with high probability.
Proposition 6 (PAC RL Bound). Let a be PAC RL algorithm for tabular MDPs and Rssapp is its sample complexity. Let @ X + [M] be decoder pre- pre-trained using video data and a S -> [N] is a bijection such that:
E S, Pang(ols) (ō(x) = a(ss) 2 21
then let G be the policy returned by a on the tabular MDP induced by Z(x). Then we have with probability at least 1 - 8 naamp! (S, A, H, E, H)H::
sup V(x) - V(p o 3 SE+2H') nell
Proof. The algorithm runs for samp H,, episodes. This implies the agent visits Nsamp( (S, A, H,,,6)H many latent states. If the decoder maps every such state 8 to the correct permutation (s), then the tabular MDP algorithm is rumning as if it ran on the induced MDP M. The probability of failure is bounded by naaap A, H. e, HD as all these failures are independent given the state. Further, the failure probability of the tabular MDP algorithm itself is 8. This leads to the total failure probability of + 7samp (S, A, H, e, H)HH..
Let II be the set of observation-based policies we competing with and which includes the optimal observation-ben based policy n. We can write suPrem V(TT) = Vy*((=) where we use the subscript M denote that the latent policy is running in the induced MDP M. Further, for any latent policy y
we have V (fp o a Φ³) = Vyr(pp as the decoder a Φ* x + (x (x)) give me access to the true state of the induced MDP M. Then with probability at least 8, we have:
Vea(T²) Vgr(G) SE
This allows us to bound the sub-optimality of the learned observation-based policy ; o @ as:
sup (T)) - v = - ((or) xen
V-r(f)) - V
< V'+ V(@.
Here we use @o 6² to denote an observation-based policy that takes action as p(a(ő" (x:)) h).
We bound V Ø*) V o 2) below. Let E = (Ō(z.) az)((())) and E = n - En be two events. We have P(E,) 2 G. Further, using union bound we have P(ET) P( S I#L (P < HO.
We first prove an upper bound on V(@oaod))
H Σ
v(@oao0") = Egocos
Fh
22

---

Published as a conference paper at ICLR 2024
H
H
Σ Th
P
Procod (E) Rooaes
Σ 4=1
rh
E
PP Procco (89)
E
Gooros
Th E
H22
E
Goo
Th
+ H29
LA=1
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 lower bound V(@ Ø):
v
Σ
Th
A=
H
H
Σ
rh
3
P Gog
(E) + Go0
Σ
rh E
P
h=1
H
2 E 6
Σ
3
P
Goo
h=1
BBo
Σ rh
Σn rh & HO -111
L1=1
H
2 B
Σ
Th
22.00
Combining the two upper bounds get:
H
v(@oao0") - v(@ 2) <
Poo
Σ l=
rh E + H22 - Bgo
rh E + 122 21120
Therefore, with probability at least nsamp (, (,) 8)H,, learn a policy & 3 such that:
sup
V(x) V 3) S 2H20.
sell
Theorem 7 (Wrapping up the proof.). Fix Eo O and 6, (0,1) and let s be any PAC RL
algorithm for tabular MDPs with sample complexity Tmmmp (S, A, H,,,8) 8). If n satisfies:
(f AA for
H8 2 for
N6 n 2amp (S, A, H. e. 2,
[ 8.
n n= n=0
In
1 02
then forward modeling learns a decoder 3 X N. Further, running e the tabular MDP with induced by " with hyperparameters E./2, 8 = 5,/4, returns latent policy P. Then there exists bijective mapping S + [[S]]] such that with probability at least 1 we huve:
Pa-g(11)) (o(x) = ass) (0(xx)=) =s) 211 4N3H²A
Vs E S,
Wenn 3jor
and
V(**) - V(@ 3) S E0
Further, the amount of online interactions in the downstream RL is given by
nsamp (S, A, H.&. 8,/2, 2,04///5/4// 8o/4) and doesn't I scale with In (6-
23

---

Published as a conference paper at ICLR 2024
Proof. We showed in Theorem 5 that we learn such that:
4N 1122
P (1)))) (ō(z) = (ss) (x)x) = 21 =
M Stor
provided A
VangSoe Stoe N
HNA Then from Proposition we learn a f such that:
Let
V(r) - ( $ &
with probability at least 1 naamp A, E, H. The failure probability
nsamp (S, A, 6,), was when condition in Theorem holds which holds with å probability. Hence, total failure probability is:
28 + shsmp (S, H,,, 8)..
We set å both in our representation learning analysis and in PAC RL to 6./4. We also set e in the PAC RL algorithm to Eo/2. This means the PAC RL algorithm runs for Rsamp (S, A, H, 6,/2, 6,/4) episodes.
We enforce 0 S
6 2n (S..H.H
Then the total failure probability becomes:
26./4) + 6. 6./2 8.
We also enforce 21130 s s./2. The sub-optimality of the PAC RL policy is given by:
En/2 Eo
This gives us our derived PAC RL bound.
We now accumulate all conditions:
2
4|F||
+
In
T2
8o
H22
Train Bror
Tiin Bror
N2 H2
5.
S 2msamp (S, A, H. /
2H 0 S Eo/2
This simplifies to
Tmin Star
A < H2
Miin Sror
< 3N
(S, A, H, 6,
Eo n Stor
A <
Or,
(f
A Mmin B2
8
N6 n
Samp (S, A, H, 2..//////// 80/4) In
[F
0
E 32
82 y Poor
d
This completes the proof.
24

---

Published as a conference paper at ICLR 2024
B.2 UPPER BOUND FOR THE TEMPORAL CONTRASTIVE APPROACH
We first convert our video dataset D into a dataset suitable for contrastive learn- ing. We first split the datasets into |n/2] pairs of videos. For each video pair (2++))
create tuple "t.,,',,,,) where 10 111a as follows. As in forward modeling, we will either use fixed value of k, or sample E Unf( We denote this general distribution over & by w E A([K] which is either point mass, or Unf([[K). We sample k ~ and Unf( 11) and h E Unf([HH).. We set
[1 (21)
.HH (2R)
(z( (2141) 22 ,.......) HH (2t+1)
+ }.
2h (21)
(2))
= 1, then we set x' = tk otherwise, we sample h [0, 13) and select x' = Ji This way, we collect a dataset Deoms of |n/2] tuples (z. k; x', We view a tuple (æ, k. x', 2) as real observation pair when = and a fake observation pair when = O. Note that sampling process leads to all data points being iid.
We define the distribution Demm (X, k, X', 2)) the distribution over (at, (z,) We express this distribution as:
w(k)
Doanl X = Iz X' 11)==
Σ
D(x 2h, =In+k)
2H
h=1
w(k)
p(x)
(xt+1 x 1/1 11 = xx)
2
w(k)
H
H
= X' =''==''--- =='Z= = 2112
Σ
D(x = =xN)
Σ h=1
D(a' = xw
h=1
w(k)
2 p(z))((x))
where we use the time homogeneity of D and definition of We will use shorthand denote D(xx+1 = x' Iı = x) as D(x' x, k) in this analysis. easy to verify that D(x' x, k) D(r' 00 (x),,, The marginal distribution D (I, k,x") is given by:
Doom (2, k x') w(k)p(z)
(D(a' 1z,k) + p(x'))
(14)
2
Note that (X) is the same as p(X).
We will use Deomt for any marginal and conditional distribution derived from Doom (X, X', Z). We assume a model class problem. We will also reuse the decoder class á X' - [N] that defined earlier, and will assume that N [S] This can be relaxed by doing clustering or working with a different induced MDP (e.g., see the clustering algorithm in Misra al. (2020)). However, this is not the main point of the analysis.
X' x [K] x[N] - [0, 1] that we use for solving the prediction
We define the expected risk minimizer of the squared loss problem below:
[n/2]
g. Φ arg
min gE9
Σ
(1(6(1(()) k00) k0)) x/0) -(0))=(0)) =0))
(15)
[n/2]
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 xr, ks x') which satisfies:
Deom
D(O"(2') k*(2),k)
Dommiz = 1, D(O*(x') (T),), +
Proof. We can express the Bayes classifier as:
Deam (x,, k. x 2 1)
Deonr = 1 x, kt, x') = D Deomla (x, k, x', 2 1) + Doamt (I, ki 0)
25

---

Published as a conference paper at ICLR 2024
w
1x)
(k))/2p(z)DDG
Ø) + (k))2222(2)))(
D(x' 12,k)
D(=' xz k) + p(z')
D(E' (x),
D('' O (²) + p(x') g( 00 (z)))
q<'1 (0) ( (z), q(' (0(2))((0<="))
D(0'(r') 2- ( p (x")
Assumption 8 (Realizability). There exists g" G and 6² such that for all (z, ki k,) supp Deom (x, k. X'), we have Deoml = I, k x') g" (x)
We will use the shorthand to denote g* (z, k, x²) = (,), (z))) As before, start with typical square loss guarantees in the realizable setting.
Theorem 8. Fix 5 E (0, 1). Under realizability (Assumption 8), the ERM solution of J. G in Eq. (15) satisfies:
(inkk''''' (O(@la). g" (z,
2
19
A
n
For proof Proposition 12 in Misra al. (2020).
We will prove coupling result similar to the case for forward modeling. However, to do this, we need define coupling distribution:
Dempp (1 3 X2 22, k, X' x (k) x1 D ( X Deom (X' x¹)
We will derive useful importance ratio bound.
D Deoup (I1, k,)²) D comt (I1, ks, x') p(=1)D(x' 21 k) p(z1)p(z() 22
2p(x1))(x')
(16)
We now prove analogous result to Proposition 4.
Theorem (Coupling for Temporal Contrastive Learning). With probability at least 1 - we have:
B( Dangy [1 ((())) = ((ua)) Ig (x",,,,,x') g" (22, k,), < 40com 8)
Proof. We start with triangle inequality:
(ōla1) = 2(22) - (12,,,, 1 (0(t)) = )(0)) 1 (0 - (00(1)) -
Exz Doup
S B(e Deup
((0(z1)= = 3(0z)) (00(01),
E (o
We bound the first term as:
(pnr) 1 fo(e)) = )((z)) ( -(66m))))
E )e 1 (0(011)
2(22)}] v E Eanm
-
(00000(()2)))))))) 17
26

---

Published as a conference paper at ICLR 2024
b/
E (,r) k,," -(0()))))))))
(1m baik,,,ktt 2²)
VV B(ri
(9 -166 k,))))
Doun (11, k))
2B ( -
b/
V2bbccat:
where Sernhh--ccwwrtz's inequality in the first step and Equation 16 in the second inequality. The second term is bounded as:
(0(m1) = x(1)) - (x2, - g
(aa Doap
( = x(02) ((0(12) k,,") (x2, (r2,
=
| [(0(zz) - ) (( (21, x')
2b
where we use the coupling argument in the first step and then reduce the first term using sym- metric of (I1, (2) in Desup- Combining the upper bounds of the two terms and using b S 1 and 2v2 completes the proof.
Assumption 9 (Temporal Contrastive Margin). We assume that there exists a Bremp 0 such that for any two different states 81 and $2.
E Ex 110 (s1,k, 9 (s2, k, s'Ill
Sremp
The factor of 1 is chosen for comparison with forward modeling as will become clear later at the end of the proof. As before, if h: fixed, the margin is given by
Semp
inf #82:81 EIES
Ex'np llg (s,,k, (s2, k, k,
and when Un]([K]) the margin given by
Premp
(u))
=
1
inf
-ali--g--e (82, k. s')l
-
2 51 s1plNais1,
We directly have ehemp 2 R Ek Lemma 4.
Premp.
(δ(z1) = ((x))) o <
2ccome (n, 8) Bromp
P Paiaarop
Proof. We start with the left-hand side in Theorem 9.
i ((((z)) )(12) } 19 (11,k, 9" (12 2')l]
I ( = Ex (1g k, g (x2 x,)
E (e
E 1 (((x)) = =3(x2) } Ex llg (1n,,, - (x2, k s'Il|
2 E (2
1 (f((z)) 2(12) A (x1) * (x2)
= 2rmmp P (un m2)p
((m)) = )(x2) A 00 (21) 0z (v²)
where we the definition of Stemps the fact that marginal over D Dooup (X) is p. and that 9²(x, x²) only depends on 6* (x') and (x) (Lemma 3). Combining with the inequality proved Theorem 9, completes the proof.
27

---

Published as a conference paper at ICLR 2024
We have now reduced this analysis to an almost identical one to the forward analysis case (Proposi- tion 4). We can, therefore, use the same steps and derive identical bounds. All what changes that Stor is replaced by Semp and in A replace In |F] with In [ At this point, can clarify that the factor of 1 was chosen the definition of Premp so that Aiw can be replaced by Premp rather than Ane which will make it harder compare margins, as we will do later.
B.3 PROOF OF LOWER BOUND FOR EXOGENOUS BLOCK MDPS
Proof of Theorem 2. We present hard instance using a family of exogenous block MDPs, with H 2, A = 21, 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 E [0, 1].
Each MDP M, is indexed by i E [d]. and is specified follows:
State space: The state is represented by Ia == [s] sh 60]] where the superscript denotes different factors. For MDP M,, only the i-th factor all h, and the other factors are exogenous. Each factor has values of {0,1}.
( is an endogenous state for
Transition: For the MDP instance Mj: has
1. For the i-th factor (endogenous factor), P(5 1 a) = 1[8 is, the endogenous states have deterministic dynamics. If s] = a, then it transitions to " 1, otherwise it transitions to 8/2 O.
(0) = 1ss = a)}- That
1
2. For the j-th factor with j * (exogenous factor), P(s x (1 p) (
s o{ + pll(s?" # s(²) for any s²" and s/-) That is, the j-th factor has probability of
1 - of transiting to the same state (i.e., 8f") 0 - 22 = Oor 8 - 82 1),
and probability of p of transiting to the different state (i.e., sf = 0 - 82 1 or
8 = + 2z 0).
Note that the MDP terminates at h 2.
Initial state distribution and reward: The marginal distribution of s is uniformly dis- tributed at random over 1) for all j E [d]. and all factors are independent from each other. For MDP M,, the agent only receive reward signal after taking action R(s a) 22" That is, it always reward 1 at s and reward 0 at s 0 no matter which action it takes.
2, with
Data collection policy for video data: We assume that the data collection policy always pick action 0 with probability p and action with probability - 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 uninformative for the agent to distinguish the MDP instance from the MDP family. Now, assume D(9) is the video data from the instance M and g0 is the state decoder learned from an arbitrary algorithm 5 with D Then, for any arbitrary algorithm 2n that uses the state decoder (66) in its execution, it is equivalent to such Ag that uses the state decoder (10)) in its execution, where j can be selected arbitrarily from [d].
State decoder requiring exponential length Without loss of generality, we further restrict the state decoder ó used in the execution of sq for all MDP instance to some On X' > [L], where h E (1 and L < 2d. Then we will argue that there must exists a l E [d], such that
2d_L
Σ (( = (1 = FL
(s +3+))
(17)
d2d
28

---

Published as a conference paper at ICLR 2024
Note that, Eq. (17) means there must be a probability of at least 2 4/224 that 61 will incorrectly group two different s together.
where == [sf") si .......7
) and 11 [81 31 ...... ".
We now prove Eq. (17). Based on the construct above, we know that X = 2 and each state in X' has the same oceupancy for 21 based on the defined initial state distribution (this holds for all instances in the MDP family, as we are now only talking about the initial state 21). Thus, we have
Σ [[]( = =0 (1))) (of' 34 v (sf
31 V
(51 31 2
(18)
Ji Ex
because we defined Φ1 : + [L]. it means that such Φ1 is only able to distinguish the number of different states from X. Then, we obtain
Σ Σ
P ( oz # )()))
jEId Zι,
Σ
((0)(21) = (( (i1))
(01 =
2d -L L
(by Eq. (18))
2d L
Σ
P(1(()))
(8 *
max
SEId
d2d
21,TEE
So this proves Eq. (17).
From Eq. (17), we know that for the MDP instance M 01 will have probability at least 2d 2122/2202" to mistake the endogenous state, which implies that for any policy that represented using the state decoder ø, it must have sub-optimality least 2" 2"-L/2.025 Therefore, it is easy to verify that, for any 0, we can simply pick 1/4s, and obtain
2d- 2s,
sub-optimality
7L S 1/
2 22d
Then, any arbitrary algorithm s442 that uses the state decoder & in its execution, where Φn X + [4] can be chosen arbitrarily for [1,2] and < 2/4e-1 must have sub-optimality larger than s.
Additional characteristics of MDP family and video data Note that, by combining the argu- ments uninformative video data and a state decoder requiring exponential length, we obtain im- possible results. We now discuss the following:
1. The margin condition defined in Assumption 3 regarding the constructed MDPs
2. The PAC learnability of the constructed MDPs
3. The coverage condition of video data.
For the defined margin condition of forward modeling, we have: for the MDP instance M, with constant p, we can bound the forward margin as below (iPr denotes the video distribution)
Fior Xe
(i) 0) Pror X2
s = 1)IlTv
P Pioe (X2
1/ = 0) Piar(XX
$ = =1]
2
X
Σ
Pine (8) = 0 1 = 0)P(X2
1 = 0) + PPr =1 | = 0)P(X2
8/ = =1)
2
X
Pror(s5
0 1 1)P(X2
8/
+ Pear 88
s/0 1)P((22 16 1
29

---

Published as a conference paper at ICLR 2024
1
Σ
(1- 2p) [P((2] $2 0) P(X2
$2 = 1)]
2 X2
(a)
[1 2p
Σ
6/ 0) + P(X2
1 2pl
Σ X2
P(X2 82 = 1)
2
2
X2
= |1 2p/,
where step (a) because 8 part of X2, and then know P(X2 | 0) and P(X₂ s 1) cannot be nonzero simultaneously. So picking p o 0.5 implies positive forward margin.
For the temporal contrastive learning, is easy to verify that [P 1 s 1, X2) Fer (2
1 of 1, X2)1 1 2pl, picking p * 0.5 also implies positive margin for temporal
contrastive learning.
As for the PAC learnability, since the latent dynamics of our constructed MDPs are deterministic, they are provably PAC learnable by Efroni et al. (2022).
As for the coverage property of the video data, it is easy verify
max (=12 =ETT,a1 EX (fo a1) =E (22(
max (x(22) max (1/p. 1/1-p].
Therefore, we can simply pick = 1/s and obtaín the desired MDP and video data properties. This completes the proof.
Addition remark of Theorem In the proof of Theorem 2. if we pick p = 0.5 for that hard instance, the constructed MDP family reduces to block MDP without exogenous noise, but the margin becomes for both forward modeling and temporal contrastive learning. Therefore, it im- forward margin could make the learnability of the
plies that either the exogenous noise problem impossible.
B.4 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, may be possible in some cases to so with an additional assumption. We highlight one example here and defer proper formal analysis to future work. One path to success is when the gold decoder results the best-in-class error. A domain where this can happen is when the endogenous state is more predictive of x' than any other In S bits of information in x. E.g., a navigation domain, there can be many sources of noise the background, but memorizing all of them can easily overwhelm the decoder's model capacity. Instead focusing solely 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-decodee approach (Ye et al., 2022). In general, the lower bound in Theorem 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 3D games, where the agent's 's action is typically most predictive of changes in observations, or in settings with no exogenous noise, one can expect these approaches do well.
B.5 RELATION BETWEEN MARGINS
We defined margins Bar for forward modeling and Been for temporal contrastive learning. The larger the values of these margins, the more easy it is to separate observations from different endogenous states. This can be directly inferred from the sample complexity bounds which scale inversely with these margins. In particular, both Bix and Stemp depend on the way sample the multi-step variable
k. We consider two special cases: one where l E [K] is fixed, we instantiate these margins as (a()
and B(a)) and second where & is uniformly sampled from [K] and we instantiate those margins as ((u) and p(w)
30

---

Published as a conference paper at ICLR 2024
A natural question is how these margins related. The sample complexity bounds of forward modeling and temporal contrastive are almost identical except for the difference in margins (Stor vs Ptemp) and the function classes (F vs 9). If the function classes were of similar complexity, then having larger margin will make it easier to learn the right representation." (()) Pfor 3/00 (k) Premps Premp 0 are related as:
Theorem 10 (Margin Relation). For any Block MDP and K E N. the margins
K for I 8/o0)
1 (K))
K Promp S Slimp
1 ((6)
Umin g(k)
S 3 8//or
4H2 for
Umin glu)
S Blun) S D)
4H for
Proof. We first prove the first two relations. Fix any l E [K] then, (N-Un(((K]
inf
[I (X'' X (1s k') Dpr (X' 82, k') llv]
lor
#1
K
inf
[Xp 81, Dxr(X' 82, llvv
inf
[Xp (1, PXr 82, k2,1i
K a1
g(k)
K
for
Similarly,
Beemp
inf
(y (s1, (0 k, s')
1
inf
Ev-p 11o0 (81, 9" (82, K, s')l],
2K
01
1
> 2K $1
inf Ev'p ||0 (81, g²(22, s')l
1
(k)
K Pump-
We now prove the next two relations. We will prove these bounds for generic distribution E A([A] over k. Recall that w is point-mass over k for o(k) and Unf([K]) for A
We denote our generie margins as 3tor and Sremp for
w. We use shorthand notation
ss, Dpp(olssk)totss for given pair of states s' and integer k E [K]. It is easy to Vim 2H
p(a")
see that Wi(s, s') < 1 as ppr(s' (s' s, (kk))p(s))) E (0, 1]. Further, we have Wa,s, 2 p(e') where we use ps k, p(s') ) (0, 1]. and Equation 8.
s, k, = g² (s, k, k,s') (pc(s')s,8) (pls(e"sss)) using the definition
We have 9" k, = Doom(2 (z of Doom in Lemma 3 and Assumption 8. We can use the shorthand Wi and the definition of 9* to show
1
Bremp
inf 2 21,5,,8 Bx-w,e' (s1, - -0--gg" k.)
K
inf
Iw(h Σ p(d)) (s1, gg (52, k, k,
2 s1 #1 S
k=l
.'ES
This inference has to be made with caveat that since we comparing upper bounds, cannot guarantee this hold.
31

---

Published as a conference paper at ICLR 2024
K Σ
inf 2 $1
(kk Σ
6)) (or( 11,, D Dar 82, k:
k=1
s'ES
(19)
As (s1. S and
s') S 1 we have
K
Stor
inf
Σ (w)
Σ
Wa(51,
Wh(s,, (pr (s') (11,) Dpr(s' s2,
S'ES
inf S $15
Σ (()) Σ
pp k), - Dpr(s' 822 k2,
e'ES
inf
BRn [1] (s' 111,)) - Dpr(d') 1821
#1 #226
Blor-
This gives us Puemp I 3 and Slepp
I p/(0) Finally, we prove the lower bounds. Starting from Equation 19 and using W s') 2 and Wa s') 217 we get the following:
K
Stor 2 $13
inf
(k)
Σ
Wa(si, Wo (pr (s') (11k) - Dor s2, k)),
s'ES
2 R
20in/221
K
2
inf 4H² P2E
Mmin
Σ (k:)
Σ
1pr (8' 181kk - Ps' 18,2))1
k=1
s'ES
M/iin
inf
TE Ekno pr k) Dpr(s' 82, kTvl
4H² a1 S
Mmin
Bror.
4H²
This gives us Pump
(k)
2
Yann (R) and
(1) temp
2
Bain 41f for which completes the proof.
g(v)
411 for
The main finding of the above theorem is that forward modeling has a higher margin than temporal contrastive learning. However, typically the function class used for forward modeling has higher statistical complexity than those for temporal contrastive learning as the latter is solving a simpler binary classification problem than generating an observation.
B.6 WHY TEMPORAL CONTRASTIVE LEARNING IS MORE SUSCEPTIBLE TO EXOGENOUS NOISE THAN FORWARD MODELING
Theorem 2 shows that in the presence of exogenous noise, video-based representation learning approach can be efficient the worst case. However, this result only presents worst-case analysis. In this section, show an instance-dependent analysis. The main finding is that the temporal contrastive approach is very susceptible to even the smallest amount of exogenous noise, while forward modeling is more robust the presence of exogenous noise. However, both approaches fail when there a significant amount of exogenous noise, consistent with Theorem
Problem Instance. We consider a Block MDP with exogenous noise with a state space of (0, 1), action space of A = {0.1) and exogenous noise space of g = {0. 1). We consider H with a uniform distribution over 51 and ε1. i.e., the start state s and the start exogenous noise variable ε are chosen uniformly from {0, 1}. The transition dynamics are deterministic and given as follows: given action a₃ E {0, and state $1 E (0, 1), we deterministicallyy transition to = 1- if $1 d1, otherwise, transitions from & to & 1 & The reward function is given by R(s2, s1) 1(s2 $ )- We use the indicator notation ({ to denote 1 if the condition & is true and 0 otherwise. The
remain in $2
$1. The exogenous noise variable deterministically
32

---

Published as a conference paper at ICLR 2024
observation space is given by X' {0. 1| where (m + 2) is the dimension of observation space.
Given the endogenous state s and exogenous noise ç, the environment generates an observation stochasticaly
[S, " ... 21 W,,".. : where U; Pssamp ( {) and w ~ Paamp (- for all i E [1] and j E (m - 1]. The distribution Pu generates u = 8 with a probability 0.8 and 1 1-8 with probability 0.2. The hyperparameter: I is fixed integer controlling what portion of the observation is generated by the exogenous noise compared to the endogenous state. If 1, we only have a small amount of exogenous noise, while if m 1 have the maximal amount of exogenous noise. The state s and exogenous noise & are both decodable from the observation x. The optimal policy achieves a return of and takes action dj if $1 = 0 and a1 Of $1 As the optimal policy depends on the value of 81, we must learn the latent state to realize the optimal policy.
Learning Setting. We assume decoder class & { , consisting of the true decoder o' and the incorrect decoder ot which maps observation to the exogenous noise Both decoders take an observation and map it to value in {0, 1}. We assume access consisting of tuples (2) collecting iid using a fixed data policy daata This policy takes action a1 0 in 81 0 and action d1 = in $1 1. Let D(#1, 222) be the data distribution induced by Tdatu- We will use D to define other distributions induced by D(11, T2), for example ((22) or (ss2). We also assume access two model classes F {0, 1] - A(X) and g {0, 1 [O, 1]. We assume these model classes are finite and contain certain constructions that we define later.
arbitrarily large dataset D
Overview: As we increase the value of I, the amount of exogenous noise the environment increases. We will prove that irrespective of the value of I, temporal contrastive learning assigns the same loss for both the correct decoder 6" and the incorrect decoder o². contrast, the forward modeling approach able prefer 6" over of when the noise is limited, specifically, when I < m/2. This will establish that temporal contrastive is very susceptible to exogenous noise whereas forward modeling is more robust. However, both approaches provably fail when there is 1 2 m/2.
As we have H 1, we will denote I2, $2,52 E2 by x', s', E' and $1, by z, S.E respectively. Note that unless specified otherwise, s and & are the endogenous state and exogenous noise of the observation x. Similarly, s' and g' are the endogenous state and exogenous noise of x'. We will also use shorthand q(z') to denote the emission probability q(x' oz o (x')) given its endogenous state and exogenous noise. We first state the conditional data distribution D(x' Ix)
D(d') x) = a('')Te (5' <)
Σ BEA
T(s') s, TTnaaa (a s),
=
= -- s}
(20)
where we use (E' <) 18 = { and EEEA T(s' s, T (a s) 1 (s' - s}
which follows from the definition of dataa Note that D(x' z) only depends via s, E. therefore, we can define (ir') xx) = (i'' {)
Let ã be an observation variable with endogenous state é and exogenous noise غ, i.e., 5 = d* (T) and E (i). We use this to derive the marginal data distribution P over x' as follows:
p(x'= = Σ D(
Σ D(x') | (s))
(.EE(11))
g(x")
Σ = 1) -11
(.E1(
((x')
(21)
where in the second step uses the fact that and Me are uniform and Eq. (20). We are now ready to prove our desired result.
Temporal contrastive learning cannot distinguish between good and bad decoder for all (m 1]. We first recall that temporal contrastive learning approach use the given observed data
(a), 22) to compute of real and fake observation tulles. This is collected into dataset (xc, "o,,'',
33

---

Published as a conference paper at ICLR 2024
where: 2 = 1 indicates that (21 x: D2 x') was observed in the dataset, and indicates that (I1 = 2, = z') was not observed, or is an imposter. We sample uniformly in {0.1}. The fake data is constructed by take start by computing the optimal Bayes classifier for the temporal contrastive learning approach using the definition of Bayes classifier in Lemma
21 from one tuple and z' I2 from another observed tuple. We
D(x' (x)
1
coom = 1 x. x2) = D(x' x)) + p(x)) 1 - =
where we use Lemma 3 in the first step and Eqs. (20) and (21) in the second step. Recall that & = denotes whether given observation tuple (xr, is real rather than imposter/false. Note that since we have k = 1, as it is a H = problem, we drop the notation & from all terms.
The marginal distribution over (xr,x') x²) for the temporal contrastive is given by Eq. (14) which in case instantiates to:
Doone (I, x²) ((x)
(D(x' [x)+ + p(z')),
1 5 (q()) (1 = -- s)1(5 =
1/43,
(22)
where we use Eqs. (20) and (21), and D(x) = q(
9(5)/4.
Let G be any classifier head. Given decoder o. define (z, x') + g(0(z), o(x)) as a model for temporal contrastive learning, with an expected contrastive loss of:
foom (g, (*)
B(r,a') D (ana))) g (x)))
Dewmll (z = 12,,2') (1 29 (6' (1))) + (
((w,)")~Daa Dana
1
Σ 1( -1- -1 -
0(8,8()) + g(8, g(8,882
Similarly, the expected temporal contrastive loss of the model g o² with the bad decoder of is given by:
foom(y (g, (5)
(== (dE(x), (x)))))
Demu(-la,at) (ex,e" Deam, Dea (x,,')
1 Σ 115-=11 = 53115' = 1 c3 1) 1[s==1 =1 ()
1
1(s'=1- s11== -1-1
g(E,E')) o((,E))
8.
Note that by interchanging with & and s' with e', we can show (ooni ot) fcom (g. 0*) There- fore, inf loom (g, Φz) = inf (g, o²). This implies that for any value of I, the temporal contrastive loss assigns the same loss to the good decoder ø² and the bad decoder oē. Hence, in practice, temporal contrastive cannot distinguish between the good and bad decoder and may con- verge the latter leading to poor downstream performance. This convergence the bad decoder may happen if it is easier to overfit to noise. For example, in our gridworld example, it is possibly easier for the model overfit the predictable motion of noise than understand the underlying dynamics of the agent. This is observed in Figure 3 where the representation learned via temporal contrastive tends to overfit to the noisy exogenous pixels and perform poorly on downstream RL tasks (Figure 2).
Forward modeling learns the good decoder if / < [m/2]. We likewise analyze the expected forward modeling loss of the good and bad decoder. For any f F. have f(x u) as the generator head that acts on a given decoder's output u E {0, 1) and generates the next observation x'.
If we use the good decoder o², then we cannot predict the exogenous noise & or g' which can be either 0 or with equal probability. This implies that for the I noisy bits 01,"" 12 in x' the best
34

---

Published as a conference paper at ICLR 2024
prediction is that each one has equal probability of taking 1. To see this, fix i [1] and recall that P(2; = 5' 5') = 0.8 and P(2; = value 0 1, therefore, P(e; u) 10 P(u. since we can deterministically predict s', therefore, we can predict the true distribution over w for all j E [m 1j. Let good be this generator head. Formally, we have:
E') = 0.2. As 5 has equal probability of taking
5') م 0.5. However,
m+1
fgood(# (1' (x)( - (1/2)
(1/2)'
II
Psamp (a's 1- Ø- (x)) 1 1 2 (x))
due duc
1+2
due 2m44 =
due to ui
The Bayes distribution is given by:
D(r') x)
=q'x²) (10) (x)) (
(x0) = - (2
m++1
f2' =1 -1-0((x)) (x(+)}
II Pamp (
(z))
II
Psamp (2/) 101--'(0 (x))) = =1- (x))
As we are optimizing the log-loss, we look at the expected KL divergence but between the D( x)) and fgoot (x') o5* (xx)) which gives:
leel Gggooo )²))
D(r' 2))
ED(e' D(a' |z) In
Spood (x' (x)())
Σ D(x') |x)) In 1f 114 {ay== =1 1 (5((0) } II Peamp (1/411 1 2h (x))
Eg
(1/2)
-1+ (1n)) Err
Σ 1
1x)
1 1 = -dp(x)) II Psamp (4/41 41 11 ofz (x:)
x
+1+ 1) [n()))))))))))))))))) + E., Σ Σ
Psamp
Psamp(1/43 (1/+1 1 1 8²" (xx) In Psamp (2/423 1 1- - (2(())
1100111
= 1l)+1)ln(2) IH (Psamp),
where (Psump) denotes the conditional entropy given by -1/2 2. (0.1) Eve Psump (v 8) In Peamp (u 8). As Psamp(u (u u) 0.8 and Psamp(1 (1 u u) = 0.2, we have H (Psamp) = 0.8 (0.88) 0.2 (0.22) x 0.500. Plugging this in, get Exz (Jgood Φ¹) (ln(2))0.55))
0.5l +n(())((2)
(n( + 0.19
Finally, the analysis when we use the of decoder is identical above. In this case, can predict (xx) and correctly predict the Psamp distribution over all the l-noisy bits V1. However, for the W1:m- bits and the x' + 2]. our best bet to predict uniform distribution. We capture this by the generator fond which gives:
I+1
fbaa(x' (x)) =
(1/2)
(1/2)"-1
II
Pammp (x)
05 (x)) 1a = (x))
due
duc to
due to 2 8
due Va.d
The expected KL loss (foad o can be computed almost exactly before and is equal to In ln(2) +
0. 0.193(m - l). We can that for laa (fgond, Φ²) < faalf ot) must have In(2) + 0.193/ (2) 1m1 1), equivalently, < m/2. This completes the analysis.
35

---

Published as a conference paper at ICLR 2024
Hyperparameter
Value
batch size learning rate epochs
128
0.001
400
of exogenous variables
10
exogenous pixel size # of VQ heads
2
VQ codebook size VQ codebook temperature
100
0
VQ codebook dimension
32
VQ bottleneck dimension
1024
Table 2: Hyperparameters used for experiments with the GridWorld and ViZDoom domains.
C ADDITIONAL EXPERIMENTAL DETAILS
C.1 DETAILS OF EXPERIMENTAL. SETUP
All results are reported with mean and standard еттоr computed over seeds. All the code for this work was run A100, V100, P40 GPUs, with compute time of approx. 12 hours for grid world experiments and hours for ViZDoom experiments. Data collection for gridworld was done using a mixture of random walks, optimal trajectories, deviation from optimal trajectories, and walks to randomly chosen goal positions. Data collection for Vizdoom was done via pretrained PPO policies along with random walks for diversity the observation space.
World Details. We consider navigation in 12x 12 Minigrid environment (Chevalier-Boisver et al., 2023). The agent represented as a red triangle and can take three actions: move forward, turn left, and turn right (Figure 3). The agent needs 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 reward of +1 +1.0 0.0 for reaching the goal and -0.01 in other cases.
ViZDoom Defend The Center Details. We test with ViZDoom environment called Defend the Center (Wydmuch et al., 2018; Kempka et al., 2016), which a first-person shooting game (Figure 5). The map is large circle. A player is spawned in the exact center. monsters spawned along the wall. Monsters are killed after single shot. After dying, each monster is respawned after some time. The episode ends when the player dies. The reward scheme is follows: +1 for killing monsterand 1 for death.
Hyperparameters. In Table 2, report the hyperaparameter values used for experiments in this work with the GridWorld and ViZDoom environments.
C.2 RESULTS ON AN ADDITIONAL DOMAIN
ViZDoom Basic. We use an additional basic ViZDoom environment (Wydmuch et al., 2018: Kempka al., 2016), which a first-person shooting game (Figure 8). The player needs to kill monster to win. The map of the environment is rectangle with gray walls, ceiling, and floor. The player is spawned along the longer wall in the center. 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 enough to kill the monster. The episode finishes when the monster is killed timeout. The reward scheme is follows: +101 for shooting the enemy, -1 per time step, and -5 for missed shots. Results for this environment are shown in Figure 7 and Figure 8 and further validate findings from theory and experiments.
C.3 ADDITIONAL ABLATIONS
Harder Exogenous Noise. Figure 6 showed the results when we increase the size of the exogenous noise variables (diamond shapes overlayed the image) in the gridworld domain while keeping the
36

---

Published as a conference paper at ICLR 2024
temporal contrastive
forward modeling
autoencoder
ACRO
50
30
50
50
Return
-300
1806
250
1003
-150
12:00
2000
30.0
0030
-150
3003
2003
3000
.003
1830
1000
limesteps
limesteps
(a) No Noise
(b) Only Observation Noise
(c) Only Reward Noise
(d) Both
Figure 7: RL experiments using different latent representations for the ViZDoom environment.
(m) Original
(n) Forward Modeling
(0) Autoencoder
(p) Temporal Contrastive
Figure 8: Decoded image reconstructions from different latent representation learning methods in ViZDoom environment. train decoder on top frozen representations trained with three video pre-training approaches.
number exogenous variables fixed 10. We also increase the number of exogenous noise vari- ables in the gridworld domain, while keeping their sizes fixed at 4 pixels and present the results in Figure 9. Both results show significant degradation in the performance of video-based represen- tation learning methods whereas ACRO which uses trajectory data continues to perform well. This supports one of our main theoretical results that exogenous noise poses a challenge for video-based representation learning.
37

---

Published as a conference paper at ICLR 2024
num noise var: 10
num noise var: 15
num noise var: 20
num noise var: 30
1.0
1.0
1.0
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.0
0.0
0.0
80000
120000
80000
80000
120000
Timesteps
Timesteps
Timesteps
(a) Autoencoder
(b) A-ttoeencoder
(c) Forward Modeling
1.0
1.0
1.0
0.9
0.8
0.8
0.6
0.6
0.6
94
04
04
42
022
0.2
0.0
0.0
0.0
0.2
B0000
80000
Timesteps
Timesteps
Timesteps
(d) VQ-Forwand Modeling
(c) Temporal Contrastive
(1) ACRO
Figure 9: Gridworld experiments with exogenous noise of size 4 and different the number of ex- ogenous noise variables. Several video-based representation learning methods struggle to learn as the number of exogenous noise variables increases, whereas ACRO which uses trajectory data, still performs well.
noise var: noise 10
num noise 25
0.4
0.4
0.2
0.2
0.0
0.0
-0.2
40000
80000
120000
-6-2
40000
80000
120000
Timestops
Timestops
(a) Forward Modeling
(b) ACRO
Figure 10: Experiments with iid noise for the Gridworld environment. 'Num noise var' denotes the number noisy diamonds constituting the exogenous noise.
LL.D. Noise in Gridworld. We evaluate iid noise in the gridworld domain. We use the diamond- shaped exogenous noise that we used in Figure 2, however, at each time step, we randomly sample the color and position of each diamond, independent of the agent's history. Figure 10(a) shows the result for forward modeling and Figure 10(b) 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 the increase in the number of noisy diamonds (num noise var) only slightly decreases their performance.
L.ID. Noise in the Basic VIZDoom environment. We evaluate the representation learning meth- ods on the basic ViZDoom domain but with independent and identically distributed (iid) noise. We add iid Gaussian noise to each pixel sampled from a mean Gaussian distribution with standard deviation of 0.001. Based on theory, we expect temporal contrastive objectives to be substantially better at filtering out Gaussian iid noise, which is validated experimentally for the basic ViZDoom
38

---

Published as a conference paper at ICLR 2024
forward modeling
temporal contrastive
autoencoder
ACRO
100
Return
-50
-50
-100
-100
1000
2000
3000
1000
-15
1000
2000
3000
4000
Timesteps
Timesteps
(a) L.L., Gaussian Noise
(b) Exogenous Noise
Figure 11: Experiments with (a) Guassian iid noise for the ViZDoom environment and (b) exogenous noise.
(i) Original
(j) Forward Modeling
(k) Autoencoder
()) Temporal Contrastive
Figure 12: 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.
Environment (Figure 11()).. Figure 11(b) refers to the basic ViZDoom result for convenient com- parison.
Additional reconstructions. We show additional image reconstructions Figure 12 for the Grid- World environment and in Figure 13 for the ViZDoom Defend the Center environment. We highlight
39

---

Published as a conference paper at ICLR 2024
that important parts of the observation space recovered successfully by the forward modeling ap- proach under varying levels of exogenous noise, whereas temporal contrastive learning often fails.
0B MODR
15
(c) Original
(f) Forward Modeling
(g) Autoencoder
(h) Temporal Contrastive
Figure 13: Decoded image reconstructions from different latent representation learning methods in the ViZDoom Defend the Center environment. We train decoder on top of frozen representations trained with the three video pre-training approaches.
40


Reviews:
[{"review_id": "AwcpI5ADxM", "reviewer": "Reviewer_sJQE", "summary": "The paper introduces theoretical analysis for pre-trained representation learning using video data and focuses on two settings: where there is iid noise in the observation and where there is also exogenous noise in the observations. \n\nMore specifically the paper investigates three methods for video pre-training - autoencoding, temporal contrastive learning, and forward modeling, and introduces two main theorems. The first theorem provides an upper bound for the setting where there is only iid noise, and the second provides a lower bond when the observations also include exogenous noises. The first theorem leads to the conclusion that learning a representation from videos is provably correct when there is no exogenous noise, while the second means that learning is exponentially hard when there is exogenous noise (in contrast to learning from trajectory data, where the corresponding actions are available). The proofs were provided for temporal contrastive learning, and forward modeling, while evaluation and comparison to learning form trajectory data (ACRO) are provided for all three learning procedures (vector quantized variational autoencoder, temporal contrastive learning, and forward modeling).", "questions": "I would like to ask the following questions:\n\n1. Why are encoder-based approaches harder to analyze?\n\n2. If Assumption 3 (Margin Assumption) holds also for the exogenous noise in addition to the endogenous states, would learning from video data still be exponentially worse than learning from trajectory data? If the answer is yes, it means that learning a representation from videos is provably correct for cases where the margin assumption holds for all the transitions in the data. \n\n3. In addition to the intuition, is it possible to prove the observation that temporal contrastive representation falls short in the presence of exogenous noise, compared to the forward model?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work introduces, for the first time, theoretical analysis and justification for pre-trained representation learning of policies from video data (under certain assumptions). In addition, the paper validates the theoretical analysis in practice, by experimenting on two challenging visual domains (GridWorld and ViZDoom). \n\nThe paper is well organized and clear to read and understand.", "weaknesses": "Although tested empirically, the paper does not provide a theoretical analysis for autoencoder-based approaches. Adding this analysis would make this work more complete. \nIn addition, the observation that temporal contrastive representation fails in the presence of exogenous noise is only empirical, justified with intuition, and lacks a more formal poof. \n\nThe analysis is restricted to training a fixed representation using only video data, without any fine-tuning stage of the learned representation. It would be nice to see an analysis of the common scenario of the fine-tuning stage.\n\nThe evaluation for iid noise with varying strength is missing (evaluation similar to Figure 6 but with iid noise). This evaluation is important for reliably comparing the performance of the setting with iid noise to that with exogenous noise.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QrCkjC5Vmt", "reviewer": "Reviewer_tXAJ", "summary": "This paper studies representation learning from videos in the context of reinforcement learning. In particular, this work focuses on representation learning in the presence of noice, either iid or exogenous. The theoretical results show that while the current methods should be able to work well with iid noise, the agent may need exponentially more samples when exogenous noise is present.\nThe experiments are conducted on GridWorld and VizDoom, and show that existing representation learning methods such as ACRO [1], temporal contrastive learning, VQ-VAE can learn with iid noise. \n\n[1] Agent-Controller Representations: Principled Offline RL with Rich Exogenous Information, Islam et al, https://arxiv.org/abs/2211.00164", "questions": "Can authors explain the connection between the theoretical part and experiments better? What do the results say in relation to the theoretical conclusions?\n\nIn Figure 4, the results seem to show that forward modeling and VAE are actually able to handle the exogenous noise. This is contrary to the theoretical result, is that right?\nOnly in Figure 6 do we see that indeed when noise is strong enough the forward modeling objective fails.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The work presents thorough theoretical analysis of representation learning from videos with iid and exogenous noise and arrives at an interesting conclusion\n- The experimental results shed light on how temporal contrastive method performs compared to models that output images", "weaknesses": "The experiments feel a little bit detached from the theoretical results: the are no experiments with iid noise, and the results with exogenous noise seem to mainly point to the fact that some representation methods are better than others, not that exogenous noise breaks everything. Only in Figure 6 do we see exogenous noise breaking forward modeling, while ACRO still works.\n\nI'm willing to raise my score if the authors present results with iid noise or justify the absence these results.\n\nSuggestions and comments:\nSection 4.2: \"remembering all of them can easily overcome the network’s capacity focusing on the agent’s state can better help the future predictions.\" reads weird.\nPage 2: \"probne\" should be \"prone\".\nAbove Equation 2, given $(x^{(i)}, k^{(i)}$ is missing a parenthesis\nAssumption 1: noisy-free should be noise-free\nJustification for Assumption 3: missing parenthesis in P_for\n\nSuggested additional related work:\n1. INFOrmation Prioritization through EmPOWERment in Visual Model-Based RL, Bharadhwaj et al, https://arxiv.org/pdf/2204.08585.pdf\n2. Joint Embedding Predictive Architectures Focus on Slow Features, Sobal et al, https://arxiv.org/abs/2211.10831\n3. Learning Invariant Representations for Reinforcement Learning without Reconstruction, Zhang et al, https://arxiv.org/abs/2006.10742", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZPPW4PABKB", "reviewer": "Reviewer_TsKD", "summary": "The paper provides theoretical and empirical results for representations learning for decision-making using video data (without explicit knowledge of the actions). Two settings are studied: one where there is iid noise in the observation, and a setting where there is \"exogenous noise, which is non-iid noise that is temporally correlated, such as the motion of people or cars in the background\". Three techniques are compared: autoencoding, temporal contrastive learning, and forward modeling. Theoretical and empirical results are provided.", "questions": "Unclarities in the text:\n- The abstract mentions \"We evaluate these representational learning methods in two visual domains, proving our theoretical findings.\" Empirical evaluation can never prove theoretical results except if it looks at all possible cases for instance. In general, it can only illustrate them.\n- Beginning of Section 3, it is mentioned that \"Our goal is to learn a decoder $\\phi : X \\rightarrow [N]$ that learns information in the underlying endogenous state $\\phi^*(x)$ while throwing away as much irrelevant information as possible.\". I don't understand this sentence. Isn't it an encoder that is learnt? What is an endogenous state and what is $phi^*$?\n- \"ACRO achieves optimal performance across all tasks.\": do the authors mean better than other algorithms instead of optimal? (the optimal is not exactly reached and also basically not known.\n\nTheorems\n- What is $\\alpha$ in Theorem 1?\n- For Theorem 1, unless I'm mistaken, the only discussion that is directly about the theorem mentions \"These upper bound provide the desired result which shows that not only can we learn the right representation and near-optimal policy but also do without the online episodes scaling with ln |Φ|.\" How can that interpretation be made from the theorem?\n- For theorem 2, it also unclear how the interpretations can be deduced from the theorem itself.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Interesting research questions that can have a big scientific impact\n- overall well-written\n- experiments follow good practice", "weaknesses": "- Some parts of the text are not clear/accurate (see the remarks and questions below in the \"questions\" section). There are also typos, e.g. \"(...) temporal contrastive learning is probne to fail (...)\"\n- It is unclear how the key messages that are supposed to come from the theorems are actually deduced (see questions below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mP9gkzD31A", "reviewer": "Reviewer_ZrZp", "summary": "This study offers a theoretical analysis of representation learning from video-based observations without labeled actions. The problem is formulated as a block Markov Decision Process with exogenous noise. Endogenous states are acquired through an encoder trained with one of three methods: a temporal contrastive loss, a single-state reconstruction loss (autoencoder), or a future state reconstruction loss.\n\nThe primary contribution of this paper is the establishment of a theorem that sets an upper bound on representation learning for future state prediction and the contrastive learning approach without noise. It also provides a lower bound in the presence of exogenous noise for these approaches, indicating that agents cannot distinguish between exogenous and endogenous noise. These results are further validated through experiments conducted in both a grid world and a visual environment. In cases without exogenous noise, representation learning proves successful, but it fails in its presence.", "questions": "Could we see some additional experiments in the revision?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper is well-written, and the proofs are clear and concise.\n- The results address a significant and novel problem, namely, representation learning from noisy video-based data, which is of great interest in the current Reinforcement Learning (RL) community.\n- The theoretical results quantitatively address a major challenge in current methods.\n- The empirical results align well with the theoretical findings.\n- The study explores multiple approaches for representation learning.\n- Assumptions are thoroughly explained and justified.\n\nIn conclusion, I believe this work should be accepted, as it offers significant and relevant insights to the action-free RL research community. The paper's strengths and contributions make it a valuable addition to the field.", "weaknesses": "- While the paper is strong in many aspects, it would be beneficial to expand the experimental evaluation to a wider variety of environments to further validate the results.\n- Minor errors, such as unclosed brackets in equations under Assumption 3, should be corrected for clarity and correctness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)