Title: Reinforcement Learning Under Latent Dynamics: Toward Statistical and Algorithmic Modularity
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Representation Learning

Latent Dynamics

Function Approximation
Abstract: Real-world applications of reinforcement learning often involve  environments where agents operate on complex, high-dimensional observations, but the underlying (``latent'')  dynamics are comparatively simple. However, beyond restrictive settings
  such as tabular latent dynamics,  the fundamental statistical requirements and algorithmic principles for *reinforcement learning under latent dynamics* are poorly
  understood.

  This paper addresses the question of reinforcement learning under *general latent dynamics* from a
  statistical and algorithmic perspective.  On the statistical side, our main negative
result shows that *most* well-studied settings for reinforcement learning with function approximation become intractable when composed with rich observations; we complement this with a positive result, identifying *latent pushforward coverability* as a
general condition that enables statistical tractability. Algorithmically, we develop provably efficient *observable-to-latent* reductions ---that is, reductions that transform an arbitrary algorithm for the
  latent MDP into an algorithm that can operate on rich observations--- in two settings: one where the agent has access to hindsight
observations of the latent dynamics (Lee et al., 2023) and one
where the agent can estimate *self-predictive* latent models (Schwarzer et al., 2020). Together, our results serve as a
  first step toward a unified statistical and algorithmic theory for
reinforcement learning under latent dynamics.

Reinforcement Learning under Latent Dynamics: Toward Statistical and Algorithmic Modularity

Abstract

Real-world applications of reinforcement learning often involve environments where agents operate on complex, high-dimensional observations, but the underlying ('latent') dynamics are comparatively simple. However, outside of restrictive settings such as small latent spaces, the fundamental statistical requirements and algorithmic principles for reinforcement learning under latent dynamics are poorly understood.

This paper addresses the question of reinforcement learning under general latent dynamics from a statistical and algorithmic perspective. On the statistical side, our main negative result shows that most well-studied settings for reinforcement learning with function approximation become intractable when composed with rich observations; we complement this with a positive result, identifying latent pushforward coverability as a general condition that enables statistical tractability. Algorithmically, we develop provably efficient observable-to-latent reductions-that is, reductions that transform an arbitrary algorithm for the latent MDP into an algorithm that can operate on rich observations-in two settings: one where the agent has access to hindsight observations of the latent dynamics [LADZ23], and one where the agent can estimate self-predictive latent models [SAGHCB20]. Together, our results serve as a first step toward a unified statistical and algorithmic theory for reinforcement learning under latent dynamics.

1 Introduction

Many application domains for reinforcement learning (RL) require the agent to operate on rich, high-dimensional observations of the environment, such as images or text [WSD15; LFDA16; KFPM21; NRKFG22; Bak+22; Bro+22]. However, the environment itself can often be summarized by latent dynamics for a low-dimensional or otherwise simple latent state space. The decoupling of latent dynamics from the complex observation process naturally suggests a modular framework for algorithm design: first learn a representation that decodes the latent state from observations, then apply a reinforcement learning algorithm for the latent dynamics on top of the learned representation. This paper investigates the algorithmic and statistical foundations of this framework. We ask: Can we take existing algorithms and sample complexity guarantees for reinforcement learning in the latent state space and lift them to the observation space in a modular fashion?

There is a growing body of theoretical and empirical work developing algorithms that combine representation learning and reinforcement learning to develop scalable algorithms. On the empirical side, a plethora of representation learning objectives have been deployed to varying degrees of success [PAED17; Tan+17; ZMCGL21; LSA20; YFK21; Lam+24; Guo+22; HPBL23], but we lack

a mathematical framework to systematically compare these objectives and understand when one might be preferred to another. On the theoretical side, all existing approaches suffer from three primary drawbacks: (a) they are tailored to restricted classes of latent dynamics models (tabular MDPs [KAL16; DKJADL19; MHKL20; ZSUWAS22; MFR23], LQR [DR21; Mha+20], or factored MDPs [MLJL21]), limiting generality; (b) the analyses, despite focusing on restrictive settings, are unwieldy, limiting progress in algorithm development; and (c) they are not modular , in the sense that the representation learning procedures are specialized to specific choices of latent reinforcement learning algorithm, limiting ease of use.

1.1 Contributions

We address the aforementioned limitations by introducing a new framework, reinforcement learning under general latent dynamics .

Reinforcement learning under general latent dynamics (Section 2). In our framework, the agent performs control based on high-dimensional observations, but the dynamics of the environment are governed by an unobserved latent state space. Following prior work (particularly the so-called Block MDP formulation [DKJADL19]), we assume that the latent states can be uniquely decoded from observations , but that the true decoder is unknown and must be learned. To aid in the decoding process, we supply the learner with a class of representations that is realizable in the sense that it is powerful enough to represent the true decoder. Our point of departure from prior theoretical works is that we do not assume specific structure (e.g., tabular or linear dynamics) on the Markov decision process (MDP) that governs the latent dynamics. Instead, we make the minimal assumption that the latent dynamics belong to a base MDP class which is statistically tractable , in the sense that when the latent states are directly observed there exists some reinforcement learning algorithm with low sample complexity that is capable of learning a near-optimal policy for every MDP in the class. We take the first steps toward building a unified and modular theory for reinforcement learning in this setting.

Contributions: Statistical modularity (Section 3). A central consideration for reinforcement learning under latent dynamics is that representation learning and exploration must be intertwined: an accurate decoder is required to explore the latent state space, but exploration is required to learn an accurate decoder. To develop provable sample complexity guarantees, one must prevent errors from compounding during this interleaving process, a challenging statistical problem which prior work addresses through strong structural assumptions on the base MDP [KAL16; DKJADL19; MHKL20; ZSUWAS22; MFR23; DR21; Mha+20; MLJL21]. For the general latent-dynamics setting we consider, it is unclear whether similar techniques can be applied, or whether the setting is even statistically tractable, ignoring computational considerations. Thus, our first contribution considers the question of statistical modularity : 2

If a base MDP class is tractable when observed directly, is the corresponding latent-dynamics problem tractable?

Statistical modularity adopts a minimax perspective by assuming that the base MDP lies in a given class, and demands that the sample complexity of the latent-dynamics setting is controlled by a natural bound on the sample complexity of the base MDP class. We show, perhaps surprisingly, that most well-studied reinforcement learning settings involving function approximation [RVR13; JKALS17; SJKAL19; MJTS20; AJSWY20; Li09; DVRZ19; WSY20; ZGS21; Du+21; JLM21; FKQR21] do not admit statistical modularity (Theorem 3.1). In other words, statistical tractability of an MDP class does not extend to the latent-dynamics setting . We complement these negative findings with a positive result, identifying pushforward coverability as a general structural condition on the latent dynamics that enables sample efficiency (Theorem 3.2).

Contributions: Algorithmic modularity (Section 4). Beyond developing a modular understanding of the statistical landscape, we investigate modular algorithm design principles for RL under general latent dynamics. Specifically, we consider the question of observable-to-latent reductions , whereby RL under latent dynamics can be reduced to the simpler problem of RL with latent states directly observed:

Can we generically lift algorithms for a base MDP class to solve the corresponding latent-dynamics problem?

This property, which we refer to as algorithmic modularity , enables modular, greatly simplified algorithm design, allowing one to use an arbitrary base algorithm for the base MDP class to solve the corresponding latent-dynamics problem. Algorithmic modularity is a stronger property than mere statistical modularity, and thus is subject to our statistical lower bound. Accordingly, we consider two settings that sidestep the lower bound through additional feedback and modeling assumptions. Our first algorithmic result considers hindsight observability [LADZ23], where latent states are revealed during training, but not at deployment (Theorem 4.1). Our second considers stronger function approximation conditions that enable the estimation of self-predictive latent models [SAGHCB20] through representation learning (Theorem A.1). Both results are fully modular : they transform any sample-efficient algorithm for the base MDP class into a sample-efficient algorithm for the latent-dynamics setting. Thus, they constitute the first general-purpose algorithms for RL under latent dynamics.

Together, we believe our results can serve as a foundation for further development of practical, general-purpose algorithms for RL under latent dynamics. To this end, we highlight a number of fascinating and challenging open problems for future research (Section 5).

2 Reinforcement Learning under General Latent Dynamics

In this section we formally introduce our framework, reinforcement learning under general latent dynamics.

MDP preliminaries. We consider an episodic finite-horizon online reinforcement learning setting. With H denoting the horizon, a Markov decision process (MDP) M ⋆ = { X , A , { P ⋆ h } H h =0 , { R ⋆ h } H h =1 , H } consists of a state space X , an action space A , a reward distribution R ⋆ h : X × A → ∆([0 , 1]) (with expectation r ⋆ h ( x, a ) ), and a transition kernel P ⋆ h : X × A → ∆( X ) (with the convention that P ⋆ 0 ( · | ∅ ) is the initial state distribution). 3

At the beginning of the episode, the learner selects a randomized, non-stationary policy π = ( π 1 , . . . , π H ) , where π h : X → ∆( A ) ; we let Π rns denote the set of all such policies. The episode evolves through the following process; beginning from x 1 ∼ P ⋆ 0 ( · | ∅ ) , the MDP generates a trajectory ( x 1 , a 1 , r 1 ) , . . . , ( x H , a H , r H ) via a h ∼ π h ( x h ) , r h ∼ R ⋆ h ( x h , a h ) , and x h +1 ∼ P ⋆ h ( · | x h , a h ) . We let P M⋆ ,π denote the law under this process, and let E M⋆ ,π denote the corresponding expectation, and likewise let P M ,π and E M ,π denote the analogous laws and expectations in another MDP M . We assume that ∑ H h =1 r h ∈ [0 , 1] almost surely for any trajectory in M ⋆ .

For a policy π and MDP M , the expected reward for π is given by J M ( π ) := E M ,π [∑ H h =1 r h ] , and the value functions are given by V M ,π h ( x ) := E M ,π [∑ H h ′ = h r h ′ | x h = x ] , and Q M ,π h ( x, a ) := E M ,π [∑ H h ′ = h r h ′ | x h = x, a h = a ] . We let π M = { π M ,h } H h =1 denote an optimal deterministic policy of M , which maximizes V M ,π (over π ) at all states (and in particular, satisfies π M ∈ arg max π ∈ Π rns J M ( π ) ), and write Q M ,⋆ := Q M ,π M . For f : X × A → R , we write π f ( x ) := arg max a f ( x, a ) as well as V f ( x ) = max a f ( x, a ) . For MDP M , horizon h ∈ [ H ] , and g : X → R , we let T M h denote the Bellman (optimality) operator defined via [ T M h g ]( x, a ) = E M [ r h + g ( x h +1 ) | x h = x, a h = a ] , and we overload notation by letting [ T M h f ]( x, a ) = [ T M h V f ]( x, a ) . We also let T M ,π h denote the Bellman evaluation operator defined via [ T M ,π h f ]( x, a ) = E M [ r h + E a ′ ∼ π h +1 ( ·| x h +1 ) [ f ( x h +1 , a ′ )] | x h = x, a h = a ] , for any π ∈ Π rns . We define the occupancy measures for layer h via d M ,π h ( x ) = P M ,π [ x h = x ] and d M ,π h ( x, a ) = P M ,π [ x h = x, a h = a ] .

Online reinforcement learning. In online reinforcement learning, the learning algorithm ALG repeatedly interacts with an unknown MDP M ⋆ by executing a policy and observing the resulting trajectory. After T rounds of interaction, the algorithm outputs a final policy ̂ π , with the goal of minimizing their risk , defined via

Framework: Reinforcement learning under general latent dynamics. In reinforcement learning under general latent dynamics , we consider MDPs M ⋆ where the dynamics are governed by the evolution of an unobserved latent state s h , while the agent observes and acts on observations x h generated from these latent states. Formally, a latent-dynamics MDP consists of two ingredients: a base MDP M lat = {S , A , { P lat ,h } H h =0 , { R lat ,h } H h =1 , H } defined over a latent state space S , and a decodable emission process ψ := { ψ h : S → ∆( X ) } H h =1 , which maps each latent state to a distribution over observations. The former is an arbitrary MDP defined over S , while the latter is defined as follows.

Definition 2.1 (Emission process) . An emission process is any function ψ := { ψ h : S → ∆( X ) } H h =1 , and is said to be decodable if

̸

When ψ = { ψ h } H h =1 is decodable, we let ψ -1 := { ψ -1 h : X → S} H h =1 denote the associated decoder.

With this, we can formally introduce the notion of a latent-dynamics MDP.

Definition 2.2 (Latent-dynamics MDP) . For a base MDP M lat = {S , A , { P lat ,h } H h =0 , { R lat ,h } H h =1 , H } , and a decodable emission process ψ , the latent-dynamics MDP ⟪ M lat , ψ ⟫ := { X , A , { P obs ,h } H h =0 , { R obs ,h } H h =1 , H } is defined as the MDP where the latent dynamics evolve based on the agent's action a h ∈ A via the process s h +1 ∼ P lat ,h ( s h , a h ) and r h ∼ R lat ,h ( s h , a h ) . The latent state is not observed directly, and instead the agent observes x h ∈ X generated by the emission process x h ∼ ψ h +1 ( s h ) . 4

Note that under these dynamics, the decoder ψ -1 associated with ψ ensures that ψ -1 h ( x h ) = s h almost surely for all h ∈ [ H ] . That is, the latent states can be uniquely decoded from the observations. To emphasize the distinction between the latent-dynamics MDP ⟪ M lat , ψ ⟫ (which operates on the observable state space X ) and the MDP M lat (which operates on the latent state space S ), we refer to the latter as a base MDP rather than, for example, a 'latent MDP', and apply a similar convention to other latent objects whenever possible. 5

Departing from prior work, we do not place any inherent restrictions on the base MDP, and in particular do not assume that the latent space is small (i.e., tabular). Rather, we aim to understand-in a unified fashion-what structural assumptions on the base MDP M lat are required to enable learnability under latent dynamics. To this end, it will be useful to considers specific classes (i.e., subsets) of base MDPs M lat and the classes of latent-dynamics MDPs they induce.

Definition 2.3 (Latent-dynamics MDP class) . Given a set of base MDPs M lat and a set of decoders Φ ⊂ {X → S} , we let

denote the class of induced latent-dynamics MDPs.

Stated another way, ⟪ M lat , Φ ⟫ is the set of all latent-dynamics MDPs ⟪ M lat , ψ ⟫ where (i) the base MDP M lat lies in M lat , and (ii), the emission process ψ is decodable, with the corresponding decoder belonging to Φ . The class M lat represents our prior knowledge about the underlying MDP M lat ; concrete classes considered in prior work include tabular MDPs [KAL16; DKJADL19; MHKL20; ZSUWAS22; MFR23], linear dynamical systems [DMRY20; DR21; Mha+20], and factored MDPs [MLJL21]. In particular, the class M lat may itself warrant using function approximation. At the same time, the class Φ represents our prior knowledge or inductive bias about the emission process, enabling representation learning. In what follows, we investigate what conditions on M lat make the induced class ⟪ M lat , Φ ⟫ tractable, both statistically (statistical modularity; Section 3) and via reduction (algorithmic modularity; Section 4).

3 Statistical Modularity: Positive and Negative Results

This section presents our main statistical results. We begin by formally defining the notion of statistical modularity introduced in Section 1, present our main impossibility result (lower bound) and its implications (Section 3.2), then give positive results for the general class of pushforward-coverable MDPs (Section 3.3).

3.1 Statistical modularity: A formal definition

We first define the statistical complexity for a MDP class (or, model class ) M .

Definition 3.1 (Statistical complexity) . We say that an MDP class M can be learned up to ε -optimality using comp ( M , ε, δ ) samples if there exists an algorithm ALG which, for every M ∈ M , attains

with probability at least 1 -δ , after T = comp ( M , ε, δ ) rounds of online interaction in M .

We say that a base MDP class M lat admits statistically modularity if, for any decoder class Φ , the induced latent-dynamics MDP class ⟪ M lat , Φ ⟫ can be learned with statistical complexity that is polynomial in: (i) the statistical complexity for the base class, and (ii) the capacity of the decoder class.

Definition 3.2 (Statistical modularity) . We say the MDP class M lat is statistically modular under complexity comp ( M lat , ε, δ ) if, for every decoder class Φ , we have

We say that M lat admits strong statistical modularity if Eq. (4) holds when comp ( M lat , ε, δ ) is the minimax sample complexity for M lat .

In the sequel, we examine well-studied MDP classes M lat (e.g., those which admit low Bellman rank [JKALS17]) and choose comp ( M lat , ε, δ ) based on natural upper bounds on their optimal sample complexity; in this case we will simply say they are (or are not) statistical modular, leaving the complexity upper bound comp implicit. Following prior work [KAL16; DKJADL19; MHKL20; ZSUWAS22; MFR23; DR21; Mha+20; MLJL21], we use log | Φ | as a proxy for the statistical complexity of supervised learning with the decoder class Φ . 6

The two most notable examples of statistical modularity covered by prior work are: (i) taking M lat as the set of tabular MDPs admits strong statistical modularity [DKJADL19; MHKL20; MFR23], and (ii) taking M lat as the set of linear MDPs admits statistical modularity with complexity poly ( d, H, |A| , ε -1 , log ( δ -1 ) ) [AKKS20; UZS22; MCKJA24; MBFR23]. 7 Interestingly, the latter does not admit strong statistical modularity, because the optimal rate for M lat does not scale with |A| , but the rate for ⟪ M lat , Φ ⟫ necessarily does [LS20; HLSW21]. The results of Mhammedi et al.; Misra et al.; Song et al. [Mha+20; MLJL21; SWFK24] can also be viewed as instances of statistical modularity for other base MDP classes.

3.2 Lower bounds: Impossibility of statistical modularity

Our main result in this section is to show that for most MDP classes M lat considered in the literature on sample-efficient reinforcement learning with function approximation [RVR13; JKALS17; SJKAL19; MJTS20; AJSWY20; Li09; DVRZ19; WSY20; ZGS21; Du+21; JLM21; FKQR21], statistical modularity (under the natural complexity upper bound for the class of interest) is impossible . Our central technical result is the following lower bound, which shows that statistical modularity can be impossible even when the base MDP is known to the learner a-priori . The lower bound is a significant generalization of the result from Song et al. [SWFK24]; we first state the lower bound, then discuss implications.

Theorem 3.1 (Impossibility of statistical modularity) . For every N ≥ 4 , there exists a decoder class Φ with | Φ | = N and a family of base MDPs M lat satisfying (i) |M lat | = 1 , (ii) H ≤ O (log( N )) , (iii) |S| = |X| ≤ N 2 , (iv) |A| = 2 , and such that

1. For all ε, δ > 0 , we have comp ( M lat , ε, δ ) = 0 .

For an absolute constant c > 0 , comp ( ⟪ M lat , Φ ⟫ , c, c ) ≥ Ω( N/ log( N )) .

In other words, even when the base dynamics are fully known, strong statistical modularity (in this case, poly (log | Φ | ) complexity) is impossible; any algorithm will require at least min { √ S, 2 Ω( H ) / H , | Φ | / log | Φ | } episodes to learn a near-optimal policy for a latent-dynamics MDP ⟪ M lat , ψ ⟫ ∈ ⟪ M lat , Φ ⟫ .

Base MDP class

M

lat

Statistical

Modularity?

Tabular

✓

Contextual Bandits

✓

Low-Rank MDP

✓

Known Deterministic MDP (

|M

lat

|

= 1

)

✓

Low State Occupancy (

∀

π

:

S →

∆(

A

))

✓

Model Class + Pushforward Coverability

✓

Linear CB/MDP

✗

⋆

Model Class + Coverability (

∀

π

M

:

M

∈ M

)

✗

Known Stochastic MDP (

|M

lat

|

= 1

)

✗

Bellman Rank (

Q

-type or

V

-type)

✗

Eluder Dimension + Bellman Completeness

✗

Q

⋆

-Irrelevant State Abstraction

✗

Linear Mixture MDP

✗

Linear

Q

⋆

/V

⋆

✗

Low State/State-Action Occupancy (

∀

π

M

:

M

∈ M

)

✗

Bisimulation

?

Low State-Action Occupancy (

∀

π

:

S →

∆(

A

)

)

?

⋆

Model Class + Coverability (

∀

π

:

S →

∆(

A

)

)

?

for (suboptimal) |A| complexity.

See Appendix E.2 for precise descriptions of each setting and our choices for their complexities.

Intuition for lower bound. The intuition behind the lower bound in Theorem 3.1 is as follows: the unobserved latent state space consists of N = | Φ | binary trees (indexed from 1 to N ), each with N leaf nodes. The starting distribution is uniform over the roots of the N trees, and the agent receives a reward of 1 if and only if they navigate to the leaf node that corresponds to the index of their current tree. The observed state space is identical to the latent state space, but the emission process shifts the index of the tree by an amount which is unknown to the agent. Despite the base MDP being known and the decoder class satisfying realizability, the agent requires near-exhaustive search to identify the value of the shift and recover a near-optimal policy.

Ataxonomy of statistical modularity. As a corollary, we prove that many (but not all) well-studied function approximation settings do not admit statistical modularity by embedding them into the lower bound construction of Theorem 3.1 (as well as a variant of the result, Theorem E.1). Our results are summarized in Figure 1. Our impossibility results highlight the following phenomenon: many MDP classes M lat that place structural assumptions via the value functions (e.g., MDPs with linearQ ⋆ /V ⋆ [Du+21] or MDPs with a Bellman complete value function class of bounded eluder dimension [JLM21; WSY20]) become intractable under latent dynamics. Intuitively, this is because it is not possible to take advantage of structure in value functions without learning a good representation, and, simultaneously, these assumptions are too weak by themselves to enable learning such a representation. Meanwhile, MDP classes M lat that place structural assumptions on the transition distribution (e.g., MDPs with low state occupancy complexity [Du+21] or low-rank MDPs [AKKS20]) are sometimes (but not always) tractable under latent dynamics. 8

We point to Appendix E.2 for background on all the settings in Figure 1 and proofs that they are (or are not) statistically modular. We remark that it is fairly straightforward to embed most of the MDP classes of Figure 1 into the construction of Theorem 3.1 since it only uses only a single base MDP M lat , and we expect that many other base MDP classes can similarly be shown to be intractable. However, proving the positive results in Figure 1 requires establishing several new results showing that certain base classes are tractable under latent dynamics; most notably, we next discuss the case of pushforward coverability .

3.3 Upper bounds: Pushforward-coverable MDPs are statistically modular

Our main postive result concerning statistical modularity is to highlight pushforward coverability [XJ21; AFK24; MFR24]-a strengthened version of the coverability parameter introduced in Xie et al. [XFBJK23]-as a general structural parameter that enables sample-efficient reinforcement learning under latent dynamics.

Definition 3.3 (Pushforward coverability) . The pushforward coverability coefficient C push for an MDP M lat with transition kernel P lat is defined by

Concrete examples [AFK24; MFR24] include: (i) tabular MDPs M lat admit C push ( M lat ) ≤ |S| ; and (ii) Low-Rank MDPs M lat (with or without known features) in dimension d admit C push ( M lat ) ≤ d . Further examples include analytically sparse Low-Rank MDPs [GMR24] and Exogenous Block MDPs with weakly correlated noise [MFR24]. Our main result is as follows.

Theorem 3.2 (Pushforward-coverable MDPs are statistically modular) . Let M lat be a base MDP class such that each M lat ∈ M lat has pushforward coverability bounded by C push ( M lat ) ≤ C push . Then, for any decoder class Φ , we have:

Theorem 3.2 shows that, modulo a term that is doubly-logarithmic in |S| , latent pushforward coverability enables statistical modularity. That is, when the base (latent) dynamics satisfy pushforward coverability, there exists an algorithm for the latent-dynamics setting which scales with the statistical complexity of the base MDP class and log | Φ | . We suspect that the additional log log |S| factor is not essential and can be removed with a more sophisticated analysis. We note that the complexity comp chosen above is not the minimax complexity for M lat , since every set of pushforward coverable MDPs is also a set of coverable MDPs with a potentially smaller coverability parameter [AFK24].

Let us provide some intuition for this result. We firstly note that when M ⋆ lat has pushforward coverability parameter C push , it holds that for any emission process ψ ⋆ , the observation-level MDP M ⋆ obs := ⟪ M ⋆ lat , ψ ⋆ ⟫ also satisfies pushforward coverability with the same parameter C push (Lemma D.5). Yet, despite access to realizable base MDP class M lat and decoder class Φ , it is unclear whether the latentdynamics MDP M ⋆ obs satisfies any of the observation-level function approximation conditions required by existing approaches that provide sample complexity guarantees under pushforward coverability. In particular, known algorithms for this setting either require a Bellman-complete value function class [XFBJK23], a class realizing certain density ratios [AFJSX24; AFK24], or a realizable model class [AFK24], and it is highly nontrivial to construct these for the latent-dynamics MDP M ⋆ obs = ⟪ M ⋆ lat , ψ ⋆ ⟫ given only the base MDP class M lat and the decoder class Φ . Intuitively, this is because the former observation-level function approximation classes capture properties of the observationlevel dynamics which cannot be obtained without some knowledge of the emission process.

Our main technical contribution is to establish a new structural property for pushforward-coverable MDPs (Lemma F.1): low-dimensional linear embeddings of their latent models can approximate the Bellman updates for an arbitrary set of test functions (as long as the set is not too large). We use this property to construct low-dimensional linear features that can approximate Bellman backups in observation-space , allowing us to (approximately) satisfy the Bellman completeness assumption required to apply GOLF [JLM21] to the latent-dynamics MDP. A fascinating open question is whether a similar approach can be used to establish that standard (as opposed to pushforward) coverable MDPs are statistically modular, which would encompass all other known positive cases of statistical modularity (cf. Figure 1). We refer interested readers to a more detailed technical overview in Appendix F.1, as well as the full proof in Appendix F.2.

4 Algorithmic Modularity

We now turn our attention to algorithmic modularity . Specifically, we aim for observable-to-latent reductions , whereby-via representation learning-RL under latent dynamics can be efficiently reduced to the simpler problem of RL with latent states directly observed. Since algorithmic modularity is a stronger property than statistical modularity, we sidestep the previous lower bounds in Section 3 through additional feedback and modeling assumptions. Our main result for this section is a new meta-algorithm, O2L, which, under these assumptions (and when equipped with an appropriately designed representation learning oracle), acts as a universal reduction in the sense that, whenever the representation learning oracle has low risk, the reduction transforms any sample-efficient algorithm for any base MDP class into a sample-efficient algorithm for the induced latent-dynamics MDP class.

Algorithm 1 O2L: Observable-to-Latent Reduction

Setup and O2L meta-algorithm. For the results in this section, we denote the (unknown) latentdynamics MDP of interest by M ⋆ obs := ⟪ M ⋆ lat , ψ ⋆ ⟫ , and use ϕ ⋆ := ( ψ ⋆ ) -1 to denote the true decoder. The O2L meta-algorithm (Algorithm 1) learns a near-optimal policy for M ⋆ obs by alternating between performing representation learning and executing a black-box 'base' RL algorithm (designed for the base MDP) on the learned representation; this approach is inspired by empirical methods that blend representation learning and RL in the latent space (e.g., [GKBNB19; SAGHCB20; Ni+24]).

Concretely, the algorithm takes as input a representation learning oracle REPLEARN and a base RL algorithm ALG lat that operates in the latent space. In each epoch t ∈ [ T ] , REPLEARN produces a new representation ̂ ϕ ( t ) : X → S based on data observed so far (potentially using additional side information, which we will elaborate on in the sequel). Then, the reduction invokes ALG lat , using ̂ ϕ ( t ) to simulate access to the true latent states. In particular, ALG lat runs for K episodes, where at each episode k : (i) ALG lat produces a latent policy π lat ( t,k ) : S× [ H ] → ∆( A ) , (ii) the latent policy is transformed into an observation-level policy via composition with ̂ ϕ ( t ) , i.e. π lat ( t,k ) ◦ ̂ ϕ ( t ) , which is then deployed to produce a trajectory { x ( t,k ) h , a ( t,k ) h , r ( t,k ) h } H h =1 , and (iii) the trajectory is compressed through ̂ ϕ ( t ) and used to update ALG lat via { ̂ ϕ ( t ) h ( x ( t,k ) h ) , a ( t,k ) h , r ( t,k ) h } H h =1 (cf. Line 8 of Algorithm 1). 9 After the K rounds conclude, ALG lat produces a final latent policy ̂ π ( t ) lat : S × [ H ] → ∆( A ) . The final policy ̂ π chosen by the O2L algorithm is a uniform mixture of ̂ π ( t ) lat ◦ ̂ ϕ ( t ) over all the epochs.

The central assumption behind O2L is that the base algorithm ALG lat can achieve low-risk in the underlying base MDP M ⋆ lat if given access to the true latent states s h = ϕ ⋆ ( x h ) . Beyond this assumption, we require that the representation learning oracle REPLEARN can learn a sufficiently high-quality representation. In our applications, this will be made possible by assuming access to a realizable decoder class Φ and two distinct assumptions: hindsight observability (Section 4.1) and conditions enabling self-predictive representation learning (Section 4.2). We will show that under these conditions, we can instantiate a representation learning oracle such that O2L inherits the sample complexity guarantee for ALG lat , thereby achieving algorithmic modularity.

4.1 Algorithmic modularity via hindsight observability

Our first algorithmic result bypasses the hardness in Section 3 by considering the setting of hindsight observability , which has garnered recent interest in the context of POMDPs [LADZ23; GCWXWB24; SLS23; LXJZV24]. Here, we assume that at training time (but not during deployment), the algorithm has access to additional feedback in the form of the true latent states, which are revealed at the end of each episode.

Assumption 4.1 (Hindsight Observability [LADZ23]) . The latent states ( ϕ ⋆ 1 ( x 1 ) , . . . , ϕ ⋆ H ( x H )) are revealed to the learner after each episode ( x 1 , a 1 , r 1 , . . . , x H , a H , r H ) concludes.

Weemphasize that in the hindsight observability framework, the learner must still execute observationspace policies π obs : X× [ H ] → ∆( A ) , as the latent states are only revealed at the end of each episode . Under hindsight observability, we can instantiate the representation learning oracle in O2L so that the

reduction achieves low risk for any choice of black-box base algorithm ALG lat . In particular, we make use of online classification oracles, which use the revealed latent states to achieve low classification loss with respect to ϕ ⋆ under adaptively generated data. We first state a guarantee based on generic classification oracles, then instantiate it to give a concrete end-to-end sample complexity bound.

Formally, at each step t , the online classification oracle, denoted via REP class , is given the states and hindsight observations collected so far and produces a deterministic estimate ̂ ϕ ( t ) = REP class ( { x ( i ) h , ϕ ⋆ h ( x ( i ) h ) } i<t,h ≤ H ) for the true decoder ϕ ⋆ . We measure the regret of the oracle via the 0 / 1 loss for classification:

̸

where p ( t ) represents a randomization distribution over the policy π ( t ) . Our reduction succeeds under the assumption that the oracle has low expected regret.

Assumption 4.2. For any (possibly adaptive) sequence π ( t ) , with π ( t ) ∼ p ( t ) , the online classification oracle REP class has expected regret bounded by

where Estclass ( T ) is a known upper bound.

We apply such an oracle within O2L as follows: at the end of each iteration t ∈ [ T ] in O2L, we sample k ∼ [ K ] uniformly, and update the classification oracle with the trajectory ( x ( t,k ) 1 , a ( t,k ) 1 , r ( t,k ) 1 ) , . . . , ( x ( t,k ) H a ( t,k ) H , r ( t,k ) H ) ; see the proof of Theorem 4.1 for details. We let Riskobs ( TK ) denote the risk of the O2L reduction when run for T epochs of K episodes, and we let Risk ⋆ ( K ) := E [ Risk ( K, ALG lat , M ⋆ lat )] denote the expected risk of ALG lat when executed on M ⋆ lat with access to the true latent states s h = ϕ ⋆ ( x h ) for K episodes.

Theorem 4.1 (Risk bound for O2L under hindsight observability) . Let ALG lat be a base algorithm with base risk Risk ⋆ ( K ) , and REP class a representation learning oracle satisfying Assumption 4.2. Then Algorithm 1, with inputs T, K, Φ , REP class , and ALG lat , has expected risk

This result shows that we can achieve sublinear risk under latent dynamics as long as (i) the base algorithm achieves sublinear risk Risk ⋆ ( K ) given access to the true latent states, and (ii) the classification oracle achieves sublinear regret Estclass ( T ) . Notably, the result is fully modular, meaning we require no explicit conditions on the latent dynamics or the base algorithm, and is computationally efficient whenever the base algorithm and classification oracle are efficient.

To make Theorem 4.1 concrete, we next provide a representation learning oracle (EXPWEIGHTS.DR; Algorithm 3 in Appendix G.1) based on a derandomization of the classical exponential weights mechanism, which satisfies Assumption 4.2 with Estclass ≲ H log | Φ | whenever it has access to a class Φ that satisfies decoder realizability.

Lemma 4.1 (Online classification via EXPWEIGHTS.DR) . Under decoder realizability ( ϕ ⋆ ∈ Φ) , EXPWEIGHTS.DR (Algorithm 3) satisfies Assumption 4.2 with 10

Instantiating Theorem 4.1 with the above representation learning oracle, we obtain the following algorithmic modularity result.

Corollary 4.1 (Algorithmic modularity under hindsight observability) . For any base algorithm ALG lat , under decoder realizability ( ϕ ⋆ ∈ Φ) , O2L with inputs T, K, Φ , EXPWEIGHTS.DR , and ALG lat achieves

Consequently, for any ALG lat , setting T ≈ KH log | Φ | / Risk ⋆ ( K ) achieves E [ Riskobs ( TK )] ≲ Risk ⋆ ( K ) with a number of trajectories TK = ˜ O ( K 2 H log | Φ | / Risk ⋆ ( K ) ) .

Beyond achieving algorithmic modularity, this result shows that under hindsight observability, we can achieve strong statistical modularity (modulo possible H factors) for every base MDP class M lat , an important result in its own right. 11 As an example, suppose that Risk ⋆ ( K ) = O ( K -1 / 2 ) , which is satisfied by many standard algorithms of interest [JKALS17; JYWJ20; JLM21; FKQR21]. Then, setting T according to Corollary 4.1 obtains an expected risk bound of ε using O ( H log | Φ | / ε 5 ) trajectories.

Remark 4.1 (Online versus offline oracles) . Theorem 4.1 critically uses that assumption that REP class satisfies an online classification error bound to handle the fact that data is generated adaptively based on the estimators ̂ ϕ (1) , . . . , ̂ ϕ ( T ) it produces, which is by now a relatively standard technique in the design of interactive decision making algorithms [FR20; FKQR21; FR23]. We note that under coverability and other exploration conditions, online oracles for classification can be directly obtained from offline (i.e. supervised) classification oracles [XFBJK23; BRS24; FHQR24].

4.2 Algorithmic modularity via self-predictive estimation

We complement the above results by studying the general online RL setting without hindsight observations. To address this more challenging setting, we design an optimistic self-predictive estimation objective (Eq. (7)), which learns a representation by jointly fitting a decoder together with a latent model. We prove that any representation learning oracle that attains low regret with respect to this objective can be used in O2L to obtain observable-to-latent reductions for any low-risk base algorithm ALG lat (for a formal statement, see Theorem A.1). We provide a (computationally inefficient) estimator (SELFPREDICT.OPT; Algorithm 4 in Appendix H.1) which we show attains low optimistic self-regret under certain statistical conditions (namely, coverability of the base MDP and a function approximation condition enabling us to express the self-prediction target as a latent model, see Lemma A.1 for a formal statement), thereby obtaining an end-to-end reduction for the general online RL setting. For lack of space, these results are deferred to Appendix A.

5 Discussion

Our work initiates the study of statistical and algorithmic modularity for reinforcement learning under general latent dynamics. Our positive and negative results serve as a first step toward a unified theory for reinforcement learning in the presence of high-dimensional observations. To this end, we close with some important future directions and open problems.

Statistical modularity. Can we obtain a unified characterization for the statistical complexity of RL under latent dynamics with a given class of base MDPs M lat ? Our results in Section 3 suggest that this will require new tools that go beyond existing notions of statistical complexity. Toward resolving this problem, concrete questions that are not yet understood include: (i) Is coverability [XFBJK23] (as opposed to pushforward coverability) sufficient for learnability under latent dynamics? (ii) Is the Exogenous Block MDP problem [EMKAL22; MFR24]-a special case of our general frameworkstatistically tractable? Lastly, are there additional types of feedback that are weaker than hindsight observability, yet suffice to bypass the hardness results in Section 3?

Algorithmic modularity. Can we derive a unified representation learning objective that enables algorithmic modularity whenever statistical modularity is possible? Ideally, such an objective would be computationally tractable. Alternatively, can we show that algorithmic modularity fundamentally requires stronger modeling assumptions than statistical modularity? Toward addressing the problems above, a first step might be to understand: (i) What are the minimal statistical assumptions under which we can minimize the self-predictive objective in Section 4.2? (ii) How can we encourage finding good representations via self-prediction beyond the use of optimism over the base (latent) models; and (iii) when can we minimize self-prediction in a computationally efficient fashion?

References

Contents

A Omitted Results from Section 4: Algorithmic Modularity via Self-predictive Estimation

In this section, we remove the assumption of hindsight observability used in Section 4.1 and instantiate O2L in the general online RL setting. Rather than assume access to additional side-information, we adopt a model-based representation learning approach, and augment our ability to perform representation learning by equipping the representation learning algorithm with a set of base MDPs M lat in addition to the decoder class Φ . We will learn a representation by jointly fitting a decoder and the base (latent) dynamics, which is a common approach in practice [GKBNB19; HLBN19; Haf+19; HLNB21; Sch+20; SAGHCB20; Guo+22]. We firstly present in Appendix A.1 a new notion of optimistic self-predictive regret which combines self-predictive representation learning with a form of optimism over a learned latent model. We then show in Appendix A.2 that any representation learning oracle that attains low regret, when used within O2L (Algorithm 1), leads to observable-tolatent reductions that ensure low risk for any base algorithm ALG lat , thereby achieving algorithmic modularity. Lastly, in Appendix A.3, we instantiate this oracle under natural structural and function approximation conditions, yielding end-to-end modularity and sample complexity guarantees.

A.1 Self-predictive estimation

Our self-predictive representation learning oracles learn to fit a representation ϕ such that the induced latent transitions ( ϕ h ( x h ) to ϕ h +1 ( x h +1 ) ) can be accurately modeled by some base (latent) MDP M lat ∈ M lat . To describe the objective, let us first introduce some notation. For a given MDP M over either S (resp. X ), we write M h ( r h , s h +1 | s h , a h ) (resp. M h ( r h , x h +1 | x h , a h ) ) for the joint conditional distribution over rewards and next states. Next, for any ϕ ∈ Φ , we define the pushforward model for M ⋆ obs ,h induced by ϕ via:

The pushforward model for ϕ captures the forward probability of the estimated latent state ϕ ( x ′ ) given a current observation x . To measure distance between models, we will use squared Hellinger distance (e.g, Foster et al. [FKQR21]), defined via D 2 H ( P , Q ) = ∫ ( √ d P d ν -√ d Q d ν ) 2 d ν for a common dominating measure ν . Then, for a base model M lat and a decoder ϕ , the self-predictive error of ( M lat , ϕ ) , at state-action pair x h , a h , is given by

This term captures the ability of M lat ,h ( ϕ h ( x h ) , a h ) to predict the next latent state ϕ h +1 ( x h +1 ) which is obtained by the pushforward model [ ϕ h +1 ♯M ⋆ obs ,h ] ( x h , a h ) . Formally, in our model-based representation learning setup, we consider oracles which, for each iteration t within O2L, take as input the trajectories collected so far and produce an estimate ( ̂ M ( t ) lat , ̂ ϕ ( t ) ) for the decoder and base model. The representation learning oracle's self-predictive regret , for the sequence ( ̂ M ( t ) lat , ̂ ϕ ( t ) ) , is then defined as

where p ( t ) represents a randomization distribution over the policy π ( t ) .

On its own, minimizing this regret may lead to degenerate solutions, a widely observed phenomenon in practice [Tan+23]. For example, in a standard combination lock MDP (e.g., Agarwal et al.; Misra et al. [AJKS22; MHKL20]), a degenerate decoder-model pair that maps all observations to a single latent state will have zero self-predictive loss until we reach the goal, which can take exponentially long. 12 We address this via the notion of optimistic estimation used in Zhang; Foster et al. [Zha22; FGQRS23], which biases the objective towards latent models with high return. This leads to the

following optimistic self-predictive regret , defined for a parameter γ > 0 , via

We assume going forward that REP self ; opt obtains low optimistic self-predictive regret; in Appendix A.3 we provide a maximum-likelihood-type estimator and conditions under which this holds.

Assumption A.1. For a parameter γ > 0 and any (possibly adaptive) sequence π ( t ) , with π ( t ) ∼ p ( t ) , the online representation learning oracle REP self ; opt is proper (i.e. outputs ̂ M ( t ) lat ∈ M lat for all t ∈ [ T ] ) and satisfies

where Estself ; opt ( T, γ ) is a known upper bound.

We note that only the decoder ̂ ϕ ( t ) is used within O2L; the model ̂ M ( t ) lat is only used for analysis (and possibly within the representation learner REP self ; opt ).

A.2 Main result

We now state the main guarantee for O2L with self-predictive representation learning. Recall that Riskobs ( TK ) denotes the risk of the O2L reduction. Compared to the hindsight-observable setting, we require a slightly stronger performance guarantee from the base algorithm ALG lat : our result scales with the worst-case expected risk for ALG lat over all M lat ∈ M lat , defined via Riskbase ( K ) := sup M lat ∈M lat E [ Risk ( K, ALG lat , M lat ))] .

Theorem A.1 (Risk bound for O2L under self-predictive estimation) . Suppose REP self ; opt satisfies Assumption A.1 with parameter γ > 0 . Then Algorithm 1, with inputs T, K, Φ , REP self ; opt , and ALG lat has expected risk

for absolute constants c 1 , c 2 , c 3 > 0 .

Theorem A.1 achieves sublinear risk as long as (i) the latent algorithm achieves sublinear risk Riskbase ( K ) given access to the true states, and (ii) the self-predictive representation learning oracle achieves sublinear regret Estself ; opt ( T, γ ) for an appropriate choice of γ . 13 Intuitively, our result scales with Riskbase ( K ) instead of Risk ⋆ ( K ) due to potential symmetries in the self-predictive objective. For example, there might be a representation-model pair ( ̂ M lat , ̂ ϕ ) that is identical to ( M ⋆ lat , ϕ ⋆ ) up to permutations of the latent state space; these cannot be distinguished by a representation learning oracle that does not observe the latent states directly, and thus the base algorithm may be tasked with solving either of these base MDPs. As with Theorem 4.1, this result achieves algorithmic modularity (since O2L inherits the risk of the base algorithm), and is computationally efficient whenever the base algorithm and self-predictive representation learning oracle are efficient.

Let us provide some intuition behind the proof of Theorem A.1. Recall that, within the inner loop of O2L, the latent algorithm ALG lat interacts with the ̂ ϕ ( t ) -compressed dynamics generated by compressing the observations x h , a h through the current decoder ̂ ϕ ( t ) h (Line 8). The crux of the analysis is the following observation: by the self-predictive representation learning guarantee, these dynamics, despite being possibly non-Markovian and generated from a POMDP (Definition I.1), are well approximated in squared Hellinger distance by the base model ̂ M ( t ) lat estimated by REP self ; opt (cf. Lemma I.2). We can then show that ALG lat , when given data from the ̂ ϕ ( t ) -compressed dynamics, has risk (for solving ̂ M ( t ) lat ) that is proportional to: i) its base risk if it were to observe states from ̂ M ( t ) lat , and ii) the Hellinger distance between ̂ M ( t ) lat and the process induced by its ̂ ϕ ( t ) -compressed

dynamics. The last ingredient is the use of latent optimism in Eq. (7), through which the risk on M ⋆ lat is upper bounded by the risk on ̂ M ( t ) lat .

In the above, showing that ALG lat obtains low risk for ̂ M ( t ) lat (despite given data from a different process) is done by establishing a certain form of corruption robustness (Definition I.2). Indeed, Theorem A.1 is a special case of a more general theorem (Theorem H.1), which provides a bound that adapts to ALG lat 's level of robustness. We obtain Theorem A.1 by showing that any algorithm satisfies the property we require (for a suitably slow rate), but we further show that tighter rates can be achieved by analyzing the specifics of various algorithms of interest (Appendix I.1.4).

A.3 Instantiating the self-predictive estimation oracle

We now present an algorithm, SELFPREDICT.OPT (Algorithm 4 in Appendix H.1), which satisfies Assumption A.1 under additional technical conditions, allowing us to instantiate Theorem A.1 to give end-to-end guarantees. Before stating the main guarantee, we highlight a few technical difficulties regarding obtaining finite-sample guarantees for (online) self-predictive estimation, and use them to motivate our statistical assumptions and algorithm design.

̸

The statistics of (online) self-predictive estimation. The first challenge is a realizability issue: when ϕ = ϕ ⋆ , we may not even be able to represent the objective ϕ♯M ⋆ obs as a latent model using only decoder and latent model realizability. Since we can never guarantee that ϕ = ϕ ⋆ exactly in the presence of statistical errors, we must introduce a modelling assumption which lets us capture the pushforward models ϕ♯M ⋆ obs . To this end, we introduce the mismatch functions, which are defined as follows.

Definition A.1 (Mismatch functions) . For a decodable emission process ψ ⋆ and decoder ϕ ∈ Φ , the mismatch function for ϕ , Γ ϕ = { Γ ϕ,h : S → ∆( S ) } H h =1 , is defined, for every h ∈ [ H ] , as the probability kernel

In the context of self-prediction, we show that the following mismatch completeness assumption suffices to capture the pushforward models ϕ♯M ⋆ obs .

Assumption A.2 (Mismatch completeness) . We have a model class L such that, for each ϕ ∈ Φ , and M lat ∈ M lat , we have Γ ϕ ◦ M lat ∈ L , where

In particular, Lemma D.8 establishes that

Accordingly, we view this assumption as a minimal way to realize the pushforward models ϕ♯M ⋆ obs .

The second challenge is a double-sampling issue, which appears because the decoders in Eq. (7) are coupled at different horizons. We address this with a novel 'debiased' maximum likelihood procedure that subtracts a form of excess risk (cf. Eq. (60)) to recover an unbiased estimator [Jia24]. Our debiased estimator and the mismatch completeness assumption can be viewed as analogous to the techniques and assumptions that are required for squared Bellman error minimization in the context of value function approximation [CJ19; JLM21].

The last issue stems from seeking an online estimation guarantee: the policies chosen by the latent algorithm are a function of the estimated decoders, which precludes the use of randomized estimators (e.g. exponential weights). We bypass this issue by appealing to the structural condition of coverability [XFBJK23], which allows us to restrict our attention to estimators that achieve low offline estimation error (via Lemma C.7). 14

Definition A.2 (State Coverability) . The state coverability coefficient for an MDP M and a policy class Π defined over a state space Z , C cov , st ( M, Π) , is given by

We require coverability in M ⋆ obs over the set of (observation-space) policies played by the O2L reduction (cf. Line 7). Again appealing to the mismatch functions, we can express this as an assumption about the base dynamics M ⋆ lat ; we show (Lemma D.1) that the latter is equivalent to assuming coverability in M ⋆ lat over the set of stochastic policies

where Π lat denotes the set of policies that ALG lat may execute. While this set may appear complicated, it is sufficient to assume coverability over the set of all deterministic non-stationary policies on M ⋆ lat . 15

Guarantee for our self-predictive estimation oracle. With these prerequisites, the main guarantee for our estimator, SELFPREDICT.OPT (Algorithm 4), is as follows.

Lemma A.1 (Optimistic self-predictive estimation via SELFPREDICT.OPT) . Let Π lat denote the set of policies played by ALG lat , and C cov , st = C cov , st ( M ⋆ lat , Γ Φ ◦ Π lat ) be the state coverability parameter on M ⋆ lat over the set of stochastic policies Γ Φ ◦ Π lat (Eq. (9)). Then, for any γ > 0 , under decoder realizability ( ϕ ⋆ ∈ Φ) , base model realizability ( M ⋆ lat ∈ M lat ) , and mismatch function completeness with class L lat (Assumption A.2), the estimator in Algorithm 4 with inputs Φ , M lat , L lat , and γ satisfies Assumption A.1 with 16

Instantiating Theorem A.1 with the above representation learning oracle, we obtain the following algorithmic modularity result.

Corollary A.1 (Algorithmic modularity via SELFPREDICT.OPT) . Under the same conditions as in Lemma A.1, and for any base algorithm ALG lat , O2L with inputs T, K, Φ , SELFPREDICT.OPT , and ALG lat achieves

for absolute constants c 1 , c 2 , c 3 . Consequently, for any ALG lat with base risk Riskbase ( K ) , setting γ and T appropriately gives

with a number of trajectories TK = ˜ O ( K 5 H 3 C cov , st |A| log 2 ( |M lat ||L lat || Φ | ) / ( Riskbase ( K )) 4 ) .

For example, if ALG lat is a base algorithm with Riskbase ( K ) = O ( K -1 / 2 ) , setting γ and T appropriately gives an expected risk of ε with a number of trajectories TK = ˜ O ( H 3 C cov , st |A| (log( |M lat ||L lat || Φ | )) 2 / ε 14 ) . This result shows that statistical modularity can be achieved up to log( |L lat | ) factors for every base MDP class M lat which is subsumed by coverability, including tabular MDPs and low-rank MDPs. 17 Compared to our positive result for the case of pushforward coverability (Section 3.3), this imposes less dynamics assumptions (since coverability is implied by pushforward coverability) but requires more representational assumptions (namely, access to the mismatch-complete class L lat ). We further remark that the mismatch completeness assumption always holds for i) the Block MDP setting, since we can always construct L lat such that log( |L lat | ) = O ( HS 2 ) , and ii) every MDP class M lat whenever we also have a realizable set of emission processes ( ψ ⋆ ∈ Ψ ), since we can construct L lat such that log( |L lat | ) = log( | Φ ||M lat || Ψ | ) . However, the mismatch completeness assumption may be more general than either of these settings.

Our results can be viewed as providing a theoretical justification for self-predictive representation learning, which has been widely used in empirical works [GKBNB19; SAGHCB20]. We consider self-prediction's ability to obtain universal observable-to-latent reductions as a strong indicator that it merits further theoretical study. In particular, many empirical works propose heuristics to alleviate the degeneracy/non-uniqueness issues inherent with self-prediction [GKBNB19; SAGHCB20; HPBL23; Tan+23]. Our methods provide a principled way to address these, and it would be interesting to investigate whether this is also empirically effective. In general, however, it is unclear whether our loss admits a computationally efficient implementation, due to the presence of optimism. Towards this, a fascinating direction for future work is understanding how self-predictive estimation can be used to obtain algorithmic modularity without the addition of optimism over the base (latent) models.

B Additional Discussion of Related Work

In this section, we discuss aspects of related work not already covered in greater detail.

Reinforcement learning under latent dynamics (or, with rich observations). Reinforcement learning under latent dynamics (or, with rich observations) has received extensive investigation in recent years, however most works have been focused on the Block MDP model in which the latent state space is tabular/finite [KAL16; DKJADL19; MHKL20; ZSUWAS22; MFR23] (see also the the closely related framework of Low-Rank MDPs [AKKS20; MCKJA24; ZSUWAS22; UZS22; MBFR23]). Beyond tabular spaces, Dean et al.; Dean et al.; Mhammedi et al. [DMRY20; DR21; Mha+20] consider continuous linear dynamics, Misra et al. [MLJL21] considers factored (but discrete) latent dynamics, Efroni et al.; Efroni et al.; Mhammedi et al. [EMKAL22; EFMKL22; MFR24] consider the Exogenous Block MDP problem in which a tabular latent state space is augmented with a non-controllable ('exogenous') factor, and Song et al. [SWFK24] consider Lipshitz continuous dynamics. To our knowledge, our work is the first to: i) explore reinforcement learning under general latent dynamics, in particular in settings where the latent space itself admits function approximation, and ii) take a more modular approach (cf. the taxonomy of Section 3).

On the algorithmic side, the works of Uehara et al. [UZS22] and Zhang et al. [ZSUWAS22], which consider Low-Rank MDPs and Block MDPs respectively, can be viewed as interleaving representation learning with 'latent' reinforcement learning algorithms that assume access to a good representation, and were an inspiration for this work. However, the algorithmic details and analyses are highly specialized to Block/Low-Rank MDPs, and unlikely to be directly applicable to reinforcement learning under general latent dynamics. Other works with a modular flavor include:

Feng et al. [FWYDY20] solve tabular Block MDPs by combining a black-box latent algorithm with an 'unsupervised learning oracle' for representation learning. This approach only leads to guarantees for tabular Block MDPs, and it is unclear whether the unsupervised learning oracle their approach requires can be constructed in natural settings.

Wu et al. [WYDW21] solve tabular block MDPs by combining a corruption-robust latent algorithm with a representation learning procedure based on clustering. Again, this work is restricted to the tabular setting, and requires a separation condition which may not be satisfied in general.

General complexity measures for reinforcement learning. Another line of research provides general complexity measures that enable sample-efficient reinforcement learning, including Bellman rank [JKALS17; SJKAL19; Du+21; JLM21], eluder dimension [RVR13], coverability [XFBJK23], and the Decision-Estimation Coefficient (DEC) [FKQR21; FGH23; FGQRS23]. Bellman rank and other complexity measures based on average Bellman error [JKALS17; SJKAL19; Du+21; JLM21] are insufficient to characterize learnability under general latent dynamics, as there are classes M lat that are known to be learnable, yet do not have bounded Bellman rank or Bellman-Eluder dimension [EMKAL22; XFBJK23]. Meanwhile, variants of Bellman rank based on squared Bellman error or related notions of error can [XFBJK23; AFJSX24] address this problem for some settings, but satisfying the modeling/realizability assumptions (e.g., Bellman completeness) required by these methods in the latent-dynamics setting is non-trivial. For example, the crux of our sample complexity bounds under latent pushforward coverability in Section 3 (Theorem 3.2) is to prove a rather involved structural result which shows that Bellman completeness can indeed be satisfied under this assumption, but it is unclear whether these techniques can be applied to more general latent dynamics classes. We expect that it is possible to bound the Decision-Estimation Coefficient [FKQR21; FGH23; FGQRS23] for the framework, but deriving efficient algorithms using this framework is non-trivial.

C Technical Tools

Lemma C.1. For any sequence of real-valued random variables ( X t ) t ≤ T adapted to a filtration ( F t ) t ≤ T , it holds that with probability at least 1 -δ ,

Lemma C.2 (Freedman's inequality (e.g., Agarwal et al. [AHKLLS14])) . Let ( X t ) t ≤ T be a realvalued martingale difference sequence adapted to a filtration ( F t ) t ≤ T . If | X t | ≤ R almost surely, then for any η ∈ (0 , 1 /R ) , with probability at least 1 -δ ,

Lemma C.3 (Corollary of Lemma C.2) . Let ( X t ) t ≤ T be a sequence of random variables adapted to a filtration ( F t ) t ≤ T . If 0 ≤ X t ≤ R almost surely, then with probability at least 1 -δ ,

Lemma C.4 (Lemma D.2 of Foster et al. [FHQR24]) . Let ( X 1 , F 1 ) , . . . , ( X n , F n ) be a sequence of measurable spaces, and let X ( i ) = ∏ i t =1 X t and F ( i ) = ⊗ i t =1 F t . For each i , let P ( i ) and Q ( i ) be probability kernels from ( X ( i -1) , F ( i -1) ) to ( X i , F i ) . Let P and Q be the laws of X 1 , . . . , X n under X i ∼ P ( i ) ( · | X 1: i -1 ) and X i ∼ Q ( i ) ( · | X 1: i -1 ) , respectively. Then it holds that

Lemma C.5 (Lemma A.11 of Foster et al. [FKQR21]) . Let P and Q be probability measures on ( X , F ) . For all h : X → R with 0 ≤ h ( X ) ≤ R almost surely under P and Q , we have

LemmaC.6 (Lemma 1 of Jiang et al. [JKALS17]) . For any f : X×A → [0 , 1] , π : S× [ H ] → ∆( A ) , we have

Lemma C.7 (Offline-to-online conversion under coverability [XFBJK23; FHQR24]) . Let M be an MDPover state space Z , Π be a policy set, and C cov = C cov ( M, Π) be the (state-action) coverability coefficient for M and Π (Definition D.3). Let p ( t ) ∈ ∆(Π) be a sequence of distributions over Π , and g ( t ) h : Z × A → [0 , 1] be a sequence of functions. Then we have that

and

D Structural Properties of Coverability and Mismatch Functions

This appendix contains structural results regarding coverability and the mismatch functions. We firstly recall the definition of the mismatch functions.

Definition D.1 (Mismatch functions) . For decodable emission process ψ ⋆ , decoder ϕ ∈ Φ and h ∈ [ H ] , we define the mismatch function for ϕ , Γ ϕ,h : S → ∆( S ) , as the probability kernel

We also recall the definition of state coverability.

Definition D.2 (State Coverability) . The coverability coefficient for an MDP M and a policy class Π defined over a state space Z , C cov , st ( M, Π) , is given by

We also define the related notion of state-action coverability.

Definition D.3 (State-Action Coverability) . The coverability coefficient for an MDP M and a policy class Π defined over a state space Z and action space A , C cov ( M, Π) , is given by

In the remainder of the section, we let Π lat ⊆ {S × [ H ] → ∆( A ) } denote an arbitrary set of latent policies, and

Lemma D.1 (State coverability is invariant to rich observations) . Let M ⋆ obs = ⟪ M ⋆ lat , ψ ⋆ ⟫ . Then, we have

Furthermore, letting { µ lat ,h ∈ ∆( S ) } h ∈ [ H ] denote the distribution which witnesses the right-handside, the left-hand-side is witnessed by the distribution

The lemma follows from the following two observations.

Lemma D.2. Let { Γ ϕ } ϕ ∈ Φ denote the mismatch functions for emission ψ ⋆ , and let M obs = ⟪ M lat , ψ ⋆ ⟫ . Then, for any π lat ∈ Π lat , ϕ ∈ Φ , h ∈ [ H ] , x ∈ X , we have

Proof of Lemma D.2. Below, we write s h = ϕ ⋆ ( x h ) . We proceed by induction, simply writing d obs ,h := d M obs ,π lat ◦ ϕ h and d lat ,h := d M lat , Γ ϕ ◦ π lat h . The base case ( h = 1) is obtained by noting that d lat , 1 ( s ) = P lat , 1 ( s | ∅ ) while d obs , 1 ( x ) = P obs , 1 ( x | ∅ ) = ψ ⋆ 1 ( x | s ) P lat , 1 ( s | ∅ ) . For the general case, via the Bellman flow equations, we have

The result is obtained by noting that

where the second line follows from the definition of the mismatch functions.

Lemma D.3 (Equivalence of state coverability and cumulative state reachability) . Let M be an MDP defined over a state space Z . The following definition is equivalent to Definition D.2:

Proof of Lemma D.3. Straightforward adaptation of the proof of Lemma 3 from Xie et al. [XFBJK23].

Proof of Lemma D.1. Using Lemma D.2 and Lemma D.3, we have

Lastly, we show that state-action coverability is bounded by state coverability times the size of the action set.

Lemma D.4 (State-action coverability bound) . For any MDP M and policy set Π , we have

Proof of Lemma D.4. Let µ s ∈ ∆( Z ) witness C cov , st ( M, Π) . Fix h ∈ [ H ] , which we omit below for cleanliness. Then, we have

Lemma D.5 (Pushforward coverability is invariant to rich observations) . Let C push ( M ) denote the pushforward coverability parameter for an MDP M (Definition 3.3), and M ⋆ obs := ⟪ M ⋆ lat , ψ ⋆ ⟫ . Then, we have

Furthermore, letting { µ lat ,h ∈ ∆( S ) } h ∈ [ H ] denote the distribution which witnesses the right-handside, the left-hand-side is witnessed by the distribution

This follows from an analogous equivalence of pushforward coverability and cumulative conditional reachability .

Lemma D.6 (Equivalence of pushforward coverability and cumulative conditional reachability) . Let M be an MDP defined over a state space Z with transition kernel P . The following definition is equivalent to pushforward coverability (Definition 3.3):

Proof of Lemma D.6. Fix h ∈ [ H ] , whose dependence we omit below. For the first direction, letting µ denote the pushforward coverability distribution, we have:

For the second direction, taking µ ( z ′ ) ∝ max z,a P ( z ′ | z, a ) , we have

Proof of Lemma D.5. This result follows by Lemma D.6 since,

We next show that the mismatch functions can be used to express the observation-level backups for any function of the decoders. For any g : S → R , h ∈ [ H ] , we define the function [Γ ϕ,h ◦ g ] : S → R

We further overload the Bellman operator notation and define, for any g : S → R and M lat = ( r lat , P lat ) ,

Lemma D.7. Let M obs = ⟪ M lat , ψ ⋆ ⟫ , ϕ ⋆ := ( ψ ⋆ ) -1 , ϕ ∈ Φ , and Γ ϕ be the mismatch function for emission ψ ⋆ (Definition D.1). Then, for any f lat : S × A → R , h ∈ [ H ] , and ( x, a ) ∈ X × A , we have

Proof of Lemma D.7. Let f := f lat , h ∈ [ H ] , and ( x, a ) ∈ X × A be given. Then, we have:

where the third line follows from the definition of the mismatch function Γ ϕ .

We next show that the mismatch functions can be used to realize the pushforward dynamics ϕ♯M ⋆ obs , which we recall are defined as:

We also recall the notation [Γ ϕ,h +1 ◦ M lat ] h , defined via:

Lemma D.8 (Pushforward model realizability via mismatch functions) . For all ϕ ∈ Φ , h ∈ [ H ] , we have:

Proof of Lemma D.8. Note that Γ ϕ can alternatively be written as:

We have

as desired.

E Proofs and Additional Results for Section 3.2: Impossibility Results

This section contains additional information and proofs related to our impossibility results regarding statistical modularity (Section 3.2), and is organized as follows:

Appendix E.1 contains the statement for an additional lower bound that is useful for establishing the impossibility results of Figure 1.

Appendix E.2 contains details for each entry of Figure 1.

Appendix E.3 contains for proofs for our main lower bound (Theorem 3.1) and the additional lower bound (Theorem E.1).

E.1 Additional Lower Bound

Theorem E.1 (Alternative lower bound) . For every N ≥ 4 , there exists an emission class Ψ and a decoder class Φ with | Ψ | = | Φ | = N and a family of latent MDPs M lat satisfying (i) |M lat | = 1 , (ii) H = 1 , (iii) |S| = |X| = N , (iv) |A| = N , and such that

For all ε, δ > 0 , we have comp ( M lat , ε, δ ) = 0 .

For an absolute constant c > 0 , comp ( ⟪ M lat , Φ ⟫ , c, c ) ≥ Ω( N/ log( N )) .

Proof of Theorem E.1.

E.2 Details for Figure 1

Below, we provide details on each entry in Figure 1. More precisely, for each latent class M lat , we will give a (brief) description of the MDP class M lat , give our choice of latent complexity comp for M lat , and prove that the class is or is not statistically modular for that choice of latent complexity. We view our choices of latent complexities as natural complexities for the respective classes.

Tabular MDPs ( ✓ ).

Latent class M lat : Tabular MDPs M lat = ( S , A , P lat , R lat , H ) . [AOM17]

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( |S| , |A| , H, ε -1 , log δ -1 ) , which is attainable, for example, via the UCB-VI algorithm of Azar et al. [AOM17]

Statistical modularity ( ✓ ): Known Block MDP algorithms (e.g. MUSIK [MFR23], BRIEE [ZSUWAS22]) have sample complexities of poly ( |S| , |A| , H, ε -1 , log δ -1 , log | Φ | ) .

Contextual Bandits ( ✓ ).

Latent class M lat : Contextual bandits with context space S , action space A , reward function r ⋆ lat : S × A → [0 , 1] and a finite realizable function class satisfying r ⋆ ∈ F lat .

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( |A| , log |F lat | , ε -1 , log δ -1 ) , attainable via, e.g., the SQUARE-CB algorithm [FR20].

Statistical modularity ( ✓ ): We note that F lat ◦ Φ = { [ f ◦ ϕ ] | f ∈ F , ϕ ∈ Φ } is a realizable function class for the observation-level reward function r ⋆ obs , since r ⋆ obs = [ r ⋆ lat ◦ ϕ ⋆ ] ∈ F lat ◦ Φ . Thus, applying the SQUARE-CB algorithm directly on the observations x ( t ) , a ( t ) , r ( t ) will give complexity poly ( |A| log( |F lat || Φ | ) , ε -1 , log δ -1 ) = poly ( |A| , log |F lat | , log | Φ | , ε -1 , log δ -1 ) .

Low-rank MDP ( ✓ ).

Latent class M lat : MDPs M lat = ( S , A , H, P lat , r lat ) such that there exists µ ⋆ lat ,h ∈ R d , θ ⋆ lat ,h ∈ R d , and a known set of features Ξ lat = { ξ lat = { ξ lat ,h : S × A → R d } H h =1 } such that for all h ∈ [ H ] we have r lat ( s h , a h ) = ⟨ ξ ⋆ lat ,h ( s h , a h ) , θ ⋆ lat ,h ⟩ as well as

for some ξ ∈ Ξ lat

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, |A| , H, log | Ξ lat | , ε -1 , log δ -1 ) , which is attainable via the VOX algorithm of Mhammedi et al. [MBFR23].

Statistical modularity ( ✓ ): This is obtained by noting that the observation-level dynamics also satisfy the low-rank property with the same dimension. Formally, letting P obs be the transition kernel for ⟪ M lat , ψ ⋆ ⟫ and ϕ ⋆ = ( ψ ⋆ ) -1 , we have

Thus, the transition kernel P obs is a low-rank MDP with µ obs ,h +1 ( x h +1 ) := ∑ s h +1 µ ⋆ lat ,h +1 ( s h +1 ) ψ ⋆ h +1 ( x h +1 | s h +1 ) and feature class

Lastly, since r obs = [ r lat ◦ ϕ ⋆ ] , the reward function is also linear with the same unknown feature class. Thus we can apply VOX directly on top of the observations, with the feature class Ξ lat ◦ Φ , which will achieve a complexity poly ( d, |A| , H, log | Ξ lat | , log | Φ | , ε -1 , log ( δ -1 ) ) .

Known Deterministic MDP ( |M lat | = 1 ) ( ✓ ).

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs of size 1 with both deterministic rewards and deterministic transitions.

Latent complexity comp : We take comp ( M lat , ε, δ ) = 0 , which is attainable as M lat is known and we can simply deploy its optimal policy.

Statistical modularity ( ✓ ): We note that, due to determinism, the latent optimal policy can be chosen to be open-loop without loss of generality, and thus will always experience the same trajectory ( s ⋆ 1 , a ⋆ 1 , . . . , s ⋆ H , a ⋆ H ) . We can define the observation-level policy which commits to this same sequence of actions, i.e. π obs ,h ( x h ) = a ⋆ h for all x h . This will be an optimal policy for any M obs = ⟪ M lat , ψ ⟫ , and can also be learned in 0 samples.

Low State Occupancy ( ∀ π : S → ∆( A )) ( ✓ ).

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs for which we have a realizable value function class, and such that there exists a feature map ζ lat = { ζ lat ,h : S → R d } H h =1 such that for all π : S → ∆( A ) and for all M lat ∈ M lat , we have

Note that the feature map does not need to be known.

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, |A| , H, log |F lat | , ε -1 , log ( δ -1 ) ) , which is attainable by the BILIN-UCB algorithm of Du et al., since i) MDPs with this property have Bilinear rank bounded by d |A| (see Definition 4.3 and Lemma 4.6 of Du et al. [Du+21]), and ii) one can construct the value function class F lat = { Q M lat ,⋆ | M lat ∈ M lat } , which is realizable and has size log |F lat | = log |M lat | .

Statistical modularity ( ✓ ): We firstly note that one can construct a realizable value function class for the set ⟪ M lat , Φ ⟫ , via the set F obs = { Q M lat ,⋆ ◦ ϕ | M lat ∈ M lat , ϕ ∈ Φ } . This is realizable since, for any M obs := ⟪ M lat , ψ ⟫ , letting ϕ ⋆ = ψ -1 , we have Q M obs ,⋆ = Q M lat ,⋆ ◦ ϕ ⋆ , and that this class has size log |M lat || Φ | . We can then show that the occupancies d M obs ,π f obs , for f obs ∈ F obs , can also be expressed as d -dimensional linear function for an appropriate choice of features, which will imply that the BILIN-UCB algorithm run directly on M obs will attain a complexity of poly ( d, |A| , H, log M lat , log Φ , ε -1 , log ( δ -1 ) ) . To obtain this, we recall the following lemma:

Lemma D.2. Let { Γ ϕ } ϕ ∈ Φ denote the mismatch functions for emission ψ ⋆ , and let M obs = ⟪ M lat , ψ ⋆ ⟫ . Then, for any π lat ∈ Π lat , ϕ ∈ Φ , h ∈ [ H ] , x ∈ X , we have

Thanks to the above lemma, we have

and so d π f ◦ ϕ obs is linear with feature mapping ψ ( x h | ϕ ⋆ ( x h ))[ ζ lat ,h ◦ ϕ ⋆ h ] and parameter θ M lat , Γ ϕ ◦ π f . Recall that the feature map need not be known, so that BILIN-UCB can still be applied despite not knowing ψ and ϕ ⋆ .

Model class + Pushforward Coverability ( ✓ ).

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs that all satisfy pushforward coverability C push ( M lat ) ≤ C push (cf. Eq. (28) for the definition).

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( C push , |A| , H, log |M lat | , ε -1 , log ( δ -1 ) ) , which is attainable by the GOLF algorithm via the results of Xie et al. [XFBJK23] (see also Lemma F.3). We obtain this by noting that i) C cov ≤ C push |A| , where C cov is defined in Definition 2 of Xie et al. [XFBJK23], and ii) a realizable model class can be used to construct a realizable value function class F and a Bellman-complete value function helper class G with sizes log |F| = log |M| and log |G| = O (log |M| ) .

Statistical modularity ( ✓ ): This is obtained via Theorem 3.2.

Linear CB/MDP ( ✗ ⋆ ).

Latent class M lat : MDPs M lat = ( S , A , P lat , R lat , H ) that are linear with respect to a known feature map ξ ⋆ lat : S × A → R d (i.e. such that Eq. (16) holds for ξ ⋆ lat ).

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, H, ε -1 , log ( δ -1 ) ) , which is attainable via the LSVI-UCB algorithm of Jin et al. [JYWJ20]. Note that this guarantee does not depend on the number of actions.

Statistical intractability ( ✗ ): The latent model used in the construction of Theorem E.1 is a set (of size 1) of linear MDPs with d = 1 . In particular, that construction was a contextual bandit so we only have to realize a reward function, and since there is only one latent model so we can trivially embed this with d = 1 via ξ ⋆ lat ( s, a ) = r lat ( s, a ) , where r lat is the reward function of the MDP used in Theorem E.1.

Statistical modularity with additional |A| -dependence: As in the Low-rank MDP case above, ⟪ M lat , ψ ⟫ is low-rank with unknown feature set Φ ′ = { ξ ⋆ lat ◦ ϕ | ϕ ∈ Φ } . Thus, by the same conclusion, a the VOX algorithm will have complexity poly ( d, |A| , H, log | Φ | ) , which is of the desired form if we allow suboptimal dependence on |A| .

Model class + Coverability ( ∀ π M : M ∈ M ) ( ✗ ).

Latent assumption: M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs that all satisfy coverability with respect to the policy class Π M = { π M | M ∈ M} , i.e. we have

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( C cov , H, log |M lat | , ε -1 , log ( δ -1 ) ) , which is attainable by the GOLF algorithm via the results of Xie et al. [XFBJK23] (see also Lemma F.3). We obtain this by noting that a realizable model class can be used to construct a realizable value function class F and a complete value function class G of sizes log |F| = log |M| and log |G| = O (log |M| ) .

Statistical intractability ( ✗ ): The latent models used in the construction of Theorem 3.1 are a set of coverable MDPs - in particular, these are trivially coverable with C cov = 1 since there is a single latent model and we can take µ = d M ⋆ lat ,π M⋆ lat . We remark that it is an interesting open question whether this impossibility result continues to hold if we require coverability with respect to the class Π of all possible latent policies.

Known Stochastic MDP ( |M lat | = 1 ) ( ✗ ).

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs of size 1.

Latent complexity comp : We take comp ( M lat , ε, δ ) = 0 , which is attainable as M lat is known and we can simply deploy its optimal policy.

Statistical intractability ( ✗ ): This is precisely the setting of Theorem 3.1, which shows that at least Ω( N/ log( N )) samples will be needed, where N = | Φ | .

Bellman rank ( Q -type or V -type) ( ✗ )

Latent assumption: M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of latent models such that each M lat ∈ M lat has Q -type Bellman rank d or V -type Bellman rank d [JLM21]. Letting F be a realizable value function class for M lat , in the Q -type case, this means that the | Π F | × |F| matrix

admits a rank d factorization. In the V -type case, the matrix

admits a rankd matrix factorization.

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, H, |A| log |F| , ε -1 , log ( δ -1 ) ) for the V -type Bellman rank case, which is achievable by the OLIVE algorithm of Jiang et al. [JKALS17], and comp ( M lat , ε, δ ) = poly ( d, H, log |F| , ε -1 , log ( δ -1 ) ) for Q -type Bellman rank, which is achievable by the BILIN-UCB algorithm of Du et al. [Du+21].

Statistical intractability ( ✗ ): We note that the construction in Theorem 3.1 has |M lat | = 1 , which trivially has Bellman rank equal to 1 , so Theorem 3.1 precludes statistical modularity with complexity comp .

Eluder dimension + Bellman Completeness ( ✗ )

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs such that there is a function class F lat satisfying

Furthermore, each M lat ∈ M lat has Bellman-Eluder dimension bounded by d (see Definition 8 of [JLM21]).

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, H, log |F| , ε -1 , log ( δ -1 ) ) , which is attainable by the GOLF algorithm of Jin et al. [JLM21].

Statistical intractability ( ✗ ): As in the Bellman rank case, the construction in Theorem 3.1 has |M lat | = 1 , so we can take F lat = { Q M lat ,⋆ | M lat ∈ M lat } which is evidently complete for T M lat , and has Eluder dimension 1 , so Theorem 3.1 precludes statistical modularity with complexity comp .

Q ⋆ -irrelevant State Abstraction ( ✗ )

Latent class M lat : M lat = ( S , A , P lat , R lat , H ) such that there is a known state abstraction function ζ lat : S → Z such that ζ lat ( s ) = ζ lat ( s ′ ) implies that Q M lat ,⋆ ( s, a ) = Q M lat ,⋆ ( s ′ , a ) for all a ∈ A .

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( |Z| , |A| , H, ε -1 , log ( δ -1 ) ) which is attainable by the OLIVE algorithm of Jiang et al. [JKALS17].

Statistical intractability ( ✗ ): We take M lat = { M lat } as the MDP class from the construction of Theorem 3.1. Let Q ⋆ lat := Q M lat ,⋆ . Note that we have Q ⋆ lat ( s, a ) ∈ { 0 , 1 } for all s, a , so we can take a latent abstract state space Z = { (0 , 0) , (0 , 1) , (1 , 0) , (1 , 1) } and a state abstraction function ζ lat such that ζ lat ( s ) = ( i, j ) if Q ⋆ lat ( s, 0) = i and Q ⋆ lat ( s, 1) = j . This satisfies the property of a Q ⋆ -irrelevant abstraction, since ζ lat ( s ) = ζ lat ( s ′ ) = ( i, j ) implies that Q ⋆ lat ( s, 0) = Q ⋆ lat ( s ′ , 0) = i and Q ⋆ lat ( s, 1) = Q ⋆ lat ( s ′ , 1) = j . This has a constant-sized abstract space ( |Z| = 4 ) and |A| = 2 , so Theorem 3.1 precludes statistical modularity with complexity comp .

Linear Mixture MDP ( ✗ ).

Latent class M lat : MDPs M lat = ( S , A , P lat , R lat , H ) such that there is a known feature map ζ lat = { ζ lat ,h : s ′ , s, a ↦→ R d } H h =1 such that

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, H, ε -1 , log ( δ -1 ) ) , which is attainable by the UCRL-VTR + algorithm of Zhou et al. [ZGS21]

Statistical intractability ( ✗ ): We take M lat = { M lat } to be the construction of Theorem 3.1. Here, there is a single latent model, so this is trivially embeddable with ζ lat ,h ( s ′ | s, a ) = P ⋆ lat ,h ( s ′ | s, a ) ∈ R 1 . This has dimension d = 1 , so Theorem 3.1 precludes statistical modularity with complexity comp .

Linear Q ⋆ /V ⋆ ( ✗ ).

Latent class M lat : MDPs M lat = ( S , A , P lat , R lat , H ) such that there are known features maps α lat : S × A → R d and β lat : S → R d such that for all M lat ∈ M lat , there exists unknown parameters θ Q , θ V ∈ R d such that Q M lat ,⋆ ( s, a ) = ⟨ α lat ( s, a ) , θ Q ⟩ and V M lat ,⋆ ( s ) = ⟨ β lat ( s ) , θ V ⟩ .

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, H, ε -1 , log ( δ -1 ) ) , which is attainable by the BILIN-UCB algorithm of Du et al. [Du+21].

Statistical intractability ( ✗ ): We can take M lat to be the latent MDP class from the construction of Theorem 3.1. Since there is a single latent model, this is trivially embeddable with dimension 1 , i.e. we can take ζ lat ( s, a ) = Q ⋆ lat ( s, a ) and β lat ( s ) = V ⋆ lat ( s ) . This has dimension d = 1 , so Theorem 3.1 precludes statistical modularity with complexity comp .

Low State or State-Action Occupancy ( ∀ π M : M ∈ M ) ( ✗ ).

Latent class M lat : In the Low State Occupancy model, M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs such that there exists a feature map ζ V lat = { ζ lat ,h : S → R d } H h =1 such that for all π ∈ { π M lat | M lat ∈ M lat } and for all M lat ∈ M lat , we have

For the State-Action Occupancy model, we have that there exists a feature map ζ Q lat = { ζ lat ,h : S × A → R d } H h =1 such that for all π ∈ { π M lat | M lat ∈ M lat } and for all M lat ∈ M lat , we have

Note that the feature map does not need to be known in either case.

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( d, |A| , H, log |F lat | , ε -1 , log ( δ -1 ) ) for the state occupancy case and comp ( M lat , ε, δ ) = poly ( d, H, log |M lat | , ε -1 , log ( δ -1 ) ) . Both are attainable by the BILIN-UCB algorithm of Du et al., since i) MDPs with this property have Bilinear rank bounded by d |A| and d respectively (see Definition 4.3 and Lemma 4.6 of [Du+21]), and ii) one can construct the value function class F lat = { Q M lat ,⋆ | M lat ∈ M lat } which is realizable and has size log |F lat | = log |M lat | .

Intractability: We can take the construction of Theorem 3.1, which has |M lat | = 1 and thus is trivially embeddable with dimension 1 , i.e. we can take ζ V lat ( s ) = d M lat ,π M lat ( s ) and ζ Q lat ( s, a ) = d M lat ,π M lat ( s, a ) .

Bisimulation (?)

Latent class M lat : MDPs M lat = ( S , A , P lat , R lat , H ) such that there is a known state abstraction function ζ lat : S → Z such that ζ lat ( s ) = ζ lat ( ˜ s ) implies that R lat ( s, a ) = R lat ( ˜ s, a ) for all a ∈ A as well as ∑ s ′ : ζ lat ( s ′ )= z ′ P lat ( s ′ | s, a ) = ∑ s ′ : ζ lat ( s ′ )= z ′ P lat ( s ′ | ˜ s, a ) for all z ′ .

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( |Z| , |A| , H, ε -1 , log ( δ -1 ) ) which is attainable by the OLIVE algorithm of [JKALS17].

Openness ( ? ): A negative result does not follow from existing constructions, since the dynamics from the tree-based construction of Theorem 3.1 are not bisimilar unless |Z| = |S| , which allows for the application of tabular methods. At the same time, a positive result does not follow from existing methods, since it is non-trivial to extend existing Block MDP methods to use the bisimulation state abstraction in a way that only pays for |Z| .

Low State-Action Occupancy ( ∀ π : S → ∆( A ) ) (? ⋆ )

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs such that there exists a feature map ζ Q lat = { ζ lat ,h : S × A → R d } H h =1 such that for all π : S → ∆( A ) and for all M lat ∈ M lat , we have

Note that the feature map does not need to be known.

We take comp ( M lat , ε, δ ) = poly ( d, H, log |M lat | , ε -1 , log ( δ -1 ) ) , which is attainable by the BILIN-UCB algorithm of Du et al., since i) MDPs with this property have Bilinear rank bounded by d (see Definition 4.3 and Lemma 4.6 of [Du+21]), and ii) one can construct a realizable value function class of size log |F| = log |M| .

Openness ( ? ): A negative result does not follow from existing constructions, since the dynamics from the tree-based construction of Theorem 3.1 do not have linear occupancies for all π : S → ∆( A ) unless d = |S| , which allows for the application of tabular methods, and the dynamics from the bandit-based construction Theorem E.1 do not have linear occupancies for all π : S → ∆( A ) unless d = |A| . At the same time, unlike the low state occupancy case, a positive result does not follow as it is unclear if we can express the observation-space occupancies linearly.

Statistical tractability with additional (suboptimal) |A| -dependence ( ✓ ): Note that we can reduce to the Low State Occupancy case ( ✓ ), since

However, this blows up the feature norm bound of the feature map ζ V lat ( s ) by a factor of |A| , which will appear logarithmically in the bound obtained by BILIN-UCB.

Model class + Coverability ( ∀ π : S → ∆( A ) ) (?).

Latent class M lat : M lat = { M lat = ( S , A , P lat , R lat , H ) } is a set of MDPs that all satisfy coverability with respect to all policies π lat : S → ∆( A ) , i.e. we have

Latent complexity comp : We take comp ( M lat , ε, δ ) = poly ( C cov , H, log |M lat | , ε -1 , log ( δ -1 ) ) , which is attainable by the GOLF algorithm via the results of Xie, Foster, Bai, Jiang, and Kakade (see also Lemma F.3). We obtain this by noting that a realizable model class can be used to construct a realizable value function class F and a complete value function class G of sizes log |F| = log |M| and log |G| = O (log |M| ) .

Openness ( ? ): A negative result does not follow from the existing constructions. The tree-based construction of Theorem 3.1 satisfies coverability with C cov = exp(Ω( H )) and the bandit-based construction of Theorem E.1 satisfies coverability with C cov = |A| . In both cases, the lower bounds cannot be used to rule out statistical modularity with the above latent complexity. Similarly, it unclear how to obtain a positive result for the latent-dynamics class ⟪ M lat , Φ ⟫ .

E.3 Proofs for Lower Bounds (Theorems 3.1 and E.1)

E.3.1 Main lower bound (Theorem 3.1)

We will prove the following result.

Theorem 3.1 (Impossibility of statistical modularity) . For every N ≥ 4 , there exists a decoder class Φ with | Φ | = N and a family of base MDPs M lat satisfying (i) |M lat | = 1 , (ii) H ≤ O (log( N )) , (iii) |S| = |X| ≤ N 2 , (iv) |A| = 2 , and such that

For all ε, δ > 0 , we have comp ( M lat , ε, δ ) = 0 .

For an absolute constant c > 0 , comp ( ⟪ M lat , Φ ⟫ , c, c ) ≥ Ω( N/ log( N )) .

Proof. Let N be given and assume without loss of generality that it is a power of 2 . We first construct the class of latent-dynamics MDPs, following Song et al. [SWFK24].

Latent MDP. The construction has a single 'known' latent MDP M lat , so that the only uncertainty in the family of latent-dynamics MDPs we construct arises from the emission processes. We set M lat = { M lat } . Set H = log 2 ( N ) + 1 and A = { 0 , 1 } . We define the state space and latent transition dynamics as follows.

The state space can be partitioned as S = S 1 , . . . , S N .

Each block S i corresponds to a standard depthH binary tree MDP with deterministic dynamics (e.g., Osband et al.; Domingues et al. [OVR16; DMKV21]). There is a single 'root' node at layer h = 1 , which we denote by s i root , and N 'leaf' nodes at layer H , which we denote by { s i,j leaf } j ∈ [ N ] . For each h = 1 , . . . , H -1 , choosing action 0 leads to the left successor of the current state deterministically, and choosing action 1 leads to the right sucessor; this process continues until we reach a leaf node at layer H .

The initial state distribution is P lat , 1 ( ∅ ) = Unif ( s 1 root , . . . , s N root ) .

There are no rewards for layers 1 , . . . , H -1 . For layer H , the reward is

This construction can summarized as follows. At layer 1 , we draw the index of one of N binary trees uniformly at random, and initialize into the root of the tree. From here, we receive a reward of 1 if we successfully navigate to the leaf node whose index agrees with the index of the tree itself, and receive a reward of 0 otherwise.

Note that the total number of latent states in this construction is |S| = N · |S 1 | = N (2 N -1)

Observation space and decoder class. Let us introduce some additional notation. For each block S i , let S i h := { s i,j h } j ∈ [2 h -1 ] denote the states in block i that are reachable at layer h , so that S i 1 = { s i root } and S i H = { s i,j leaf } j ∈ [ N ] . We define X = S so that |X| ≤ 4 N 2 , and consider a class of emission processes corresponding to deterministic maps. Let Σ denote the set of cyclic permutations on N elements, excluding the identity permutation. That is, each σ i ∈ Σ takes the form

For each σ ∈ Σ , we consider the emission process

That is, ψ σ shifts the index of the binary tree containing s ( i,j ) h according to σ . Let Ψ = { ψ σ | σ ∈ Σ } . Consider the decoder class

which has | Φ | = N . We consider the class of rich-observation MDPs given by

It is clear that this class of rich-observation MDPs satisfies the decodability assumption for emissions Ψ .

Sample complexity lower bound. To lower bound the sample complexity, we prove a lower bound on the constrained PAC Decision-Estimation Coefficient (DEC) of [FGH23]. For an arbitrary MDP M (defined over the space X ) and ε ∈ [0 , 2 1 / 2 ] , define 18

where M ( π ) denotes the law over trajectories ( x 1 , a 1 , r 1 ) , . . . , ( x H , a H , r H ) induced by executing the policy π in the MDP M , J M ( π ) denotes the expected reward for policy π under M , and π M denotes the optimal policy for M . We further define

where the supremum ranges over all MDPs defined over X and A . We now appeal to the following technical lemma.

In light of Lemma E.1, it follows from Theorem 2.1 in Foster et al. [FGH23] 19 that any PAC RL algorithm that uses T episodes of interaction for T log( T ) ≤ c · N must have E [ J M ( π M ) -J M ( ̂ π )] ≥ c ′ for a worst-case MDP in M , where c, c ′ > 0 are absolute constants. This implies that any PAC RL which has E [ J M ( π M ) -J M ( ̂ π )] ≤ c ′ must have T log( T ) ≥ c · N and thus T ≥ c · N/ log( N ) .

Proof of Lemma E.1. Define M lat as the latent-space MDP that has identical dynamics to M lat but, has zero reward in every state, and define M := ⟪ M lat , id ⟫ as the rich-observation MDP obtained by composing M lat with the 'identity' emission process id that sets x h = s h . Observe that M and M i , induce identical dynamics in observation space if rewards are ignored: For all policies π ,

It follows that for each i , for all policies π , we have

since the learner receives identical feedback in the MDPs M i and M unless they reach the observation x H = s ( ψ i ( j ) ,j ) leaf for some j (corresponding to latent state s ( j,j ) leaf in M i ), in which case they receiver reward 1 in M i but reward 0 in M . We now claim that for any q ∈ ∆(Π) , there exists a set of at least N/ 2 indices I q ⊂ [ N ] such that

for all i ∈ I q . To see this, note that by Eq. (21), we have

where the second inequality uses that ∑ N i =1 P M ,π [ x H = s ( j,ψ -1 i ( j )) leaf | x 1 = s ( j ) root ] ≤ 1 , as the events in the sum are mutually exclusive (and the event we condition on does not depend on i ). We conclude by Markov's inequality that P i ∼ Unif ([ N ]) [ E π ∼ q [ D 2 H ( M i ( π ) , M ( π ) )] ≥ 4 /N ] ≤ 1 / 2 , giving I q ≥ N/ 2 .

From Eq. (26), we conclude that for all ε 2 ≥ 4 /N ,

To lower bound this quantity, observe that for any index i and any policy π , we have

̸

where the third inequality uses Eq. (19). We conclude that for any distribution p, q ∈ ∆(Π) ,

as long as N ≥ 4 , where the second-to-last inequality uses that for all j , the events { x H = s ( j,ψ -1 i ( j )) leaf | x 1 = s ( j ) root } are disjoint for all i . Since this lower bound holds uniformly for all q, p ∈ ∆(Π) , we conclude that

E.3.2 Proof of alternative lower bound (Theorem E.1)

We will prove the following result.

Theorem E.1 (Alternative lower bound) . For every N ≥ 4 , there exists an emission class Ψ and a decoder class Φ with | Ψ | = | Φ | = N and a family of latent MDPs M lat satisfying (i) |M lat | = 1 , (ii) H = 1 , (iii) |S| = |X| = N , (iv) |A| = N , and such that

For all ε, δ > 0 , we have comp ( M lat , ε, δ ) = 0 .

For an absolute constant c > 0 , comp ( ⟪ M lat , Φ ⟫ , c, c ) ≥ Ω( N/ log( N )) .

Proof of Theorem E.1. We repeat more or less repeat the same proof as Theorem 3.1, but with the appropriate modifications to translate from the contextual tree-based construction in Theorem 3.1 to the contextual bandit-based construction in the theorem statement. Let N be given and assume without loss of generality that it is a power of 2 .

Latent MDP. Our construction has a single 'known' latent MDP M lat ; that is, the only uncertainty in the family of rich-observation MDPs we construct arises from the emission processes. Set M lat = { M lat } . Set H = 1 and A = [ N ] . We define the state space and latent transition dynamics as follows.

The state space can be partitioned as S = S 1 , . . . , S N .

Each block S i corresponds to a single state s i with N actions denoted by a i , i ∈ [ N ] .

The initial state distribution is P lat , 1 ( ∅ ) = Unif ( s 1 , . . . , s N ) .

The reward function is

Informally, this construction can summarized as a contextual bandit (with uniform context distribution), with a reward of 1 if and only if we play the action corresponding to the index of the context drawn.

Note that the total number of latent states in this construction is |S| = N and the number of actions is |A| = N .

Observation space and decoder class. We define X = S so that |X| = |S| , and consider a class of emission processes corresponding to deterministic maps. Let Σ denote the set of cyclic permutations on N elements, excluding the identity permutation. That is, each σ i ∈ Σ takes the form

For each σ ∈ Σ , we consider the emission process

That is, ψ σ shifts the context s i according to σ . Let Ψ = { ψ σ | σ ∈ Σ } . Consider the decoder class

which has | Φ | = N . We consider the class of rich-observation MDPs given by

It is clear that this class of rich-observation MDPs satisfies the decodability assumption for emissions Ψ .

Sample complexity lower bound. To lower bound the sample complexity, we prove a lower bound on the constrained PAC Decision-Estimation Coefficient (DEC) of [FGH23]. For an arbitrary MDP M (defined over the space X ) and ε ∈ [0 , 2 1 / 2 ] , define 20

where M ( π ) denotes the law over observations ( x 1 , a 1 , r 1 ) induced by executing the policy π in the MDP M , J M ( π ) denotes the expected reward for policy π under M , and π M denotes the optimal policy for M . We further define

where the supremum ranges over all MDPs defined over X and A . We now appeal to the following technical lemma.

Lemma E.2. For all ε 2 ≥ 4 /N , we have that sup M dec ε ( M , M ) ≥ 1 2 .

In light of Lemma E.2, it follows from Theorem 2.1 in Foster et al. [FGH23] 21 that any PAC RL algorithm that uses T episodes of interaction for T log( T ) ≤ c · N must have E [ J M ( π M ) -J M ( ̂ π )] ≥ c ′ for a worst-case MDP in M , where c, c ′ > 0 are absolute constants. This implies that any PAC RL which has E [ J M ( π M ) -J M ( ̂ π )] ≤ c ′ must have T log( T ) ≥ c · N and thus T ≥ c · N/ log( N ) .

Proof of Lemma E.2. Define M lat as the latent-space MDP that has identical dynamics to M lat but, has zero reward for every state-action pair, and define M := ⟪ M lat , id ⟫ as the rich-observation MDP obtained by composing M lat with the identity emission process that sets x h = s h . In the rest of the proof, we use the shorthand ψ i := ψ σ i . Observe that M and M i , induce identical dynamics in observation space if rewards are ignored, i.e. for all policies π : X → ∆( A ) ,

It follows that for each i , for all policies π , we have

since the learner receives identical feedback in the MDPs M i and M unless they play the action a 1 = a j given observation x 1 = s ψ i ( j ) (corresponding to latent state s i in M i ), in which case they receiver reward 1 in M i but reward 0 in M . We now claim that for any q ∈ ∆(Π) , there exists a set of at least N/ 2 indices I q ⊂ [ N ] such that

for all i ∈ I q . To see this, note that by Eq. (21), we have

We conclude by Markov's inequality that P i ∼ Unif ([ N ]) [ E π ∼ q [ D 2 H ( M i ( π ) , M ( π ) )] ≥ 4 /N ] ≤ 1 / 2 , giving I q ≥ N/ 2 .

From Eq. (26), we conclude that for all ε 2 ≥ 4 /N ,

To lower bound this quantity, observe that for any index i and any policy π , we have

where the third inequality uses Eq. (25). We conclude that for any distribution p, q ∈ ∆(Π) ,

as long as N ≥ 4 . Since this lower bound holds uniformly for all q, p ∈ ∆(Π) , we conclude that

F Proofs for Section 3.3: Positive Results

This section is dedicated to our upper bound establishing that pushforward-coverable MDPs are statistically modular (Theorem 3.2). We provide a technical overview in Appendix F.1, and provide a full proof in Appendix F.2.

F.1 Technical Overview: Low-dimensional embeddings for pushforward-coverable MDPs.

The idea behind our positive result is to show that under the conditions of Theorem 3.2, it is possible to construct an (approximately) Bellman-complete value function class for the latent-dynamics MDP M ⋆ obs , at which point we can apply the GOLF algorithm of Jin et al. [JLM21]. We achieve this via two technical contributions. The first is the introduction of the mismatch functions Γ ϕ , formally defined as follows.

Definition F.1 (Mismatch functions) . For a decodable emission process ψ ⋆ and decoder ϕ ∈ Φ , the mismatch function for ϕ , Γ ϕ = { Γ ϕ,h : S → ∆( S ) } H h =1 , is defined, for every h ∈ [ H ] , as the probability kernel

The mismatch functions allow us to express functions of the decoders as latent objects, and we revisit them in the context of self-predictive estimation (Appendix A). For the present result, we show (Lemma D.7) that the mismatch functions can capture the observation-level Bellman backups for any function of the decoders. That is, for any x h , a h , letting s h = ( ψ ⋆ ) -1 ( x h ) denote the true latent state, we have that for any f lat : S × A → R and ϕ ∈ Φ :

That is, the Bellman update of f lat ◦ ϕ h +1 in the latent-dynamics MDP M ⋆ obs can be expressed as a Bellman update in the base MDP M ⋆ lat for a different (latent) function Γ ϕ,h +1 ◦ V f lat ( s h +1 ) := ∑ s ′ h +1 Γ ϕ,h +1 ( s ′ h +1 | s h +1 ) max a ′ f lat ( s ′ h +1 , a ′ ) .

However, the mismatch functions Γ ϕ embed some knowledge of the emission process, and (with only decoder and base model realizability) are unknown to the learner. Our second technical contribution bypasses this by establishing a new structural property for pushforward-coverable MDPs (Lemma F.1): there exist low-dimensional linear embeddings of their transition kernels which can approximate Bellman backups for an arbitrary and potentially unknown set of functions, as long as the set is not too large.

Lemma F.1 (Pushforward-coverable MDPs admit low-dimensional embeddings) . Let M be a known MDP with reward function r , transition kernel P , and pushforward coverability parameter C push . Let µ = { µ h } h ∈ [ H ] denote its pushforward coverability distribution (i.e. the minimizer of Definition 3.3) and F ⊆ ( S × [ H ] → [0 , 1]) be an arbitrary class of functions. Suppose that we sample W ∈ {± 1 } d ×S as a matrix of independent Rademacher random variables, and define

and

Then for any ε apx ∈ (0 , 1) , as long as we set

we have that for all f ∈ F and h ∈ [ H ] , with probability at least 1 -δ :

as well as max s,a,h ∥ ψ h ( s, a ) ∥ 2 2 ≤ C push (16 log( |S||A| H ) + 11) and max f,h ∥ w f,h ∥ 2 2 ≤ 16 log( |F| H ) + 11 . We emphasize that the feature map ψ = { ψ h } H h =1 is oblivious to F , in the sense that it can be computed directly from M without any knowledge of F .

We use this property, in conjunction with latent model realizability, to construct linear features that can approximate the right-hand-side of Eq. (27), thus yielding an (approximately) Bellman-complete value function class for the latent-dynamics MDP M ⋆ obs .

F.2 Proofs for Latent Model Class + Pushforward Coverability (Theorem 3.2)

In this section, we establish positive results under latent MDP classes which satisfy pushforward coverability. We assume that every model in M lat satisfies pushforward coverability , defined as follows:

Definition F.2 (Pushforward coverability) . The pushforward coverability coefficient C push for an MDP M with transition kernel P is defined by

The pushforward coverability coefficient for an MDP class M is defined by

Note that for any MDP M we always have

where C cov is the state-action coverability coefficient (Definition D.3). Thus, an MDP with low pushforward coverability is also an MDP with low state-action coverability for all policies (upto a dependence on |A| ).

We will show the show the following result.

Theorem 3.2 (Pushforward-coverable MDPs are statistically modular) . Let M lat be a base MDP class such that each M lat ∈ M lat has pushforward coverability bounded by C push ( M lat ) ≤ C push . Then, for any decoder class Φ , we have:

The proof comes in three parts. We will firstly show that MDP that satisfies pushforward coverability admit low-dimensional feature maps that can approximate Bellman backups (Appendix F.2.1), then establish that a regret bound for the GOLF algorithm [XFBJK23] under misspecification (Appendix F.2.2), and then combine these ingredients (Appendix F.2.3).

F.2.1 A structural result: Pushforward-coverable MDPs are approximately low-rank

Our central technical result for this section is Lemma F.1, which is based on a variant of the JohnsonLindenstrauss lemma and establishes that under pushforward coverability, we can define a linear feature class which satisfies an approximate form of Bellman completeness. We define the clipping operator via

We prove the following lemma.

Lemma F.1 (Pushforward-coverable MDPs admit low-dimensional embeddings) . Let M be a known MDP with reward function r , transition kernel P , and pushforward coverability parameter C push . Let µ = { µ h } h ∈ [ H ] denote its pushforward coverability distribution (i.e. the minimizer of Definition 3.3) and F ⊆ ( S × [ H ] → [0 , 1]) be an arbitrary class of functions. Suppose that we sample W ∈ {± 1 } d ×S as a matrix of independent Rademacher random variables, and define

and

Then for any ε apx ∈ (0 , 1) , as long as we set

we have that for all f ∈ F and h ∈ [ H ] , with probability at least 1 -δ :

as well as max s,a,h ∥ ψ h ( s, a ) ∥ 2 2 ≤ C push (16 log( |S||A| H ) + 11) and max f,h ∥ w f,h ∥ 2 2 ≤ 16 log( |F| H ) + 11 . We emphasize that the feature map ψ = { ψ h } H h =1 is oblivious to F , in the sense that it can be computed directly from M without any knowledge of F .

Proof of Lemma F.1. Fix h ∈ [ H ] , whose dependence we omit for cleanliness. We begin by verifying that, in expectation, ⟨ w f , ψ ( s, a ) ⟩ is equal to T f ( s, a ) . For this, note that

⟨

w

f

, ψ

(

s, a

)

⟩

Consequently, we have

̸

Note that this remaining noise term is zero-mean - we will show in the sequel that it can be made small by picking d appropriately. We next examine the norms of the vectors ψ ( s, a ) and w f . Note that we have

̸

where we have used that

by definition of pushforward coverability. Further note that we have

̸

̸

We will now appeal to the following technical lemma to upper bound Eq. (30), Eq. (31), and Eq. (32) by establishing that the Rademacher noise terms concentrate to their expectations. The proof of the lemma will be given in the sequel.

̸

̸

Lemma F.2. Let u, v ∈ R n , and let W ∈ {± 1 } d × n have independent Rademacher entries. Then with probability at least 1 -δ ,

̸

Furthermore, for any set of vectors V ⊂ R n , we also have

̸

Let ( s, a ) ∈ S ×A and f ∈ F . To bound |⟨ ψ ( s, a ) , w f ⟩ -T f ( s, a ) | (cf. Eq. (30)), we apply the first bound of Lemma F.2 with u = ( P ( s ′ | s, a ) /µ 1 / 2 ( s ′ ) ) s ′ ∈S and v = ( µ 1 / 2 ( s ′ ) f ( s ′ ) ) s ′ ∈S , which gives

where we have again used that ∥ u ∥ 2 2 = ∑ s ′ ∈S P 2 ( s ′ | s,a ) µ ( s ′ ) ≤ C push and also that ∥ v ∥ 2 2 = 1 since ∥ f ∥ ∞ ≤ 1 for all f ∈ F . To bound Eq. (31), we apply the second bound of Lemma F.2 with V = {( P h -1 ( s ′ | s,a ) µ 1 / 2 h ( s ′ ) ) s ′ ∈S } s,a ∈S×A h × [ H ] , which gives

Lastly, to bound Eq. (32), we take V = { ( µ 1 / 2 h ( s ′ ) f h ( s ′ ) ) s ′ ∈S } f ∈F h ∈ [ H ] in Lemma F.2, which establishes that

Note that Eq. (34) establishes that the Bellman backup T f ( s, a ) is well-approximated by ⟨ ψ ( s, a ) , w f ⟩ only at a single state-action pair ( s, a ) . We can obtain an L ∞ -approximation guarantee by taking a union bound over S and A , which would incur a dependence on log |S| in the final sample complexity. Here, we bypass this by instead requiring only an approximation guarantee under the L 2 ( µ ⊗ Unif ( A )) norm. Via (pushforward) coverability, this will ensure that E π [ ( ⟨ w f , ψ ( s, a ) ⟩ - T f ( s, a )) 2 ] is well-controlled for all policies π , which will be sufficient for our downstream sample-complexity analysis of GOLF. However, directly establishing an L 2 ( µ ⊗ Unif ( A )) approximation guarantee is technically challenging since it would require establishing a fourth-order (rather than second-order) equivalent of Eq. (33). The remainder of the proof will obtain an L 2 ( µ ⊗ Unif ( A )) approximation guarantee by instead sampling a dataset of size n from µ ⊗ Unif ( A ) and taking a union bound over that dataset to ensure a uniform bound on all state-action pairs in that dataset. Via an additional concentration bound, this will ensure that the error is well-behaved under the L 2 ( µ ⊗ Unif ( A )) norm.

For each h ∈ [ H ] , sample a dataset D = { ( s ( i ) h , a ( i ) h ) } n i =1 i.i.d. from µ h ⊗ Unif ( A ) . By a union bound over n , F , and H , we have that

∀ i ∈ [ n ] , f ∈ F , h ∈ [ H ] : ∣ ∣ 〈 ψ h ( s ( i ) h , a ( i ) h ) , w f,h 〉 -T h f h +1 ( s ( i ) h , a ( i ) h ) ∣ ∣ ≤ ε ( n |F| Hδ -1 ) , (35) where we recall the definition of ε ( · ) from Eq. (34). Now, let

Note that | X f,h ( s, a ) | ≤ 4 and

since T h f h +1 ( s, a ) ∈ [0 , 2] and the clipping operator is 1-Lipshitz. Note that

where this expectation is only over the sampling of the data point ( s, a ) (and not the Rademacher matrix W ). Let

By boundedness of X f,h ( s, a ) and Hoeffding's inequality, we have that with probability at least 1 -δ :

Taking another union bound over F and H as well as the event in Eq. (35) gives that

2 -1

and ∀ i ∈ [ n ] , f ∈ F , h ∈ [ H ] : X i,f,h ≤ ε ( n |F| Hδ ) , (37)

recalling the definition of ε ( · ) from Eq. (34). Then, re-arranging Eq. (36) gives us that

We now conclude the proof by picking n and d appropriately to ensure that the right-hand-side is bounded by ε apx , which will ensure the desired claim that

For convenience, we introduce absolute constants c and c ′ whose precise values may change from line to line. We pick n = 64log ( 2 |F| Hδ -1 ) /ε 2 apx . Plugging this into (38) gives

Noting that n ≤ 128 |F| Hδ -1 ε 2 apx and plugging this into ε (Eq. (34)) gives

Setting ensures that

by Eq. (40). Combining Eq. (38) and Eq. (41), we get

as desired. It only remains to establish the concentration results of Lemma F.2.

Proof of Lemma F.2. We establish the first claim. Let i ∈ [ d ] be fixed, and consider the random variable

̸

Note that E [ Z i ] = 0 by independence of W i,j and W i,k for every j = k . By Exercise 6.9 of Boucheron et al. [BLM13], we have that

̸

Since Z i are independent, it follows that

Hence, ∑ d i =1 Z i is a sub-Gamma random variable with parameters ν = 16 ∥ u ∥ 2 2 ∥ v ∥ 2 2 d and c = 64 ∥ u ∥ 2 2 ∥ v ∥ 2 2 , and it follows from Equation (2.5) on page 29 of Boucheron et al. [BLM13] that for all ε > 0 ,

Taking a union bound, and using that the random variable is symmetric, we obtain the desired claim. We now establish the second claim. Let V ⊂ R n be a subset of vectors. Let i ∈ [ d ] be fixed, and re-consider the random variable

̸

Again appealing to Exercise 6.9 of Boucheron et al. [BLM13], we have that

̸

where B := max v ∈V ∥ v ∥ 4 2 . Since Z i are independent, it follows that

Hence, ∑ d i =1 Z i is a sub-Gamma random variable with parameters ν = 16max v ∈V ∥ v ∥ 4 2 d and c = 64max v ∈V ∥ v ∥ 4 2 , and it follows from Equation (2.5) on page 29 of Boucheron et al. [BLM13] that for all ε > 0 ,

To conclude, it remains only to show the bound E [ Z i ] ≤ max v ∥ v ∥ 2 2 (16 log |V| +9) . This follows by a standard log-sum-exp approach. Below, we abbreviate ρ j := W i,j . We can observe that for any λ > 0 :

̸

̸

Note that X := ∑ j ρ j v j is subGaussian with parameter ∥ v ∥ 2 2 , since:

Then, it follows (e.g. Lemma 1.12 of Rigollet et al. [RH23]) that X 2 -E [ X 2 ] satisfies a subexponential MGF bound with parameter 16 ∥ v ∥ 2 2 , i.e.

We also note that

Adding and subtracting E [ X 2 ] in Eq. (43) gives

Picking λ = 1 16 max v ∥ v ∥ 2 2 concludes the proof.

F.2.2 GOLF with on-policy misspecification

Consider the version of GOLF [JLM21] in Algorithm 2. We have the following guarantee for the regret of GOLF, which extends Jin et al. [JLM21] to allow for on-policy misspecification.

Lemma F.3. Suppose that Q M ⋆ obs ,⋆ ∈ F and G satisfies ε apx -completeness in the sense that for all h ∈ [ H ] and f ∈ F h +1 , there exists g ∈ G h such that E π ( g -T M ⋆ obs h f ) 2 ≤ ε 2 apx for all π ∈ Π F := { π f : f ∈ F} . Let C cov := C cov ( M ⋆ obs , Π F ) (Definition D.3). Then for an appropriate choice of β , Algorithm 2 ensures that

Algorithm 2 GOLF [JLM21]

input: Function classes F and G , confidence width β > 0 . initialize: F (0) ←F , D (0) ←∅ ∀ h ∈ [ H ] .

h

1: for episode t = 1 , 2 , . . . , T do

2: Select policy π ( t ) ← π f ( t ) , where f ( t ) := arg max f ∈F ( t -1) f ( x 1 , π f, 1 ( x 1 )) .

3: Execute π ( t ) for one episode and obtain trajectory ( x ( t ) 1 , a ( t ) 1 , r ( t ) 1 ) , . . . , ( x ( t ) H , a ( t ) H , r ( t ) H ) .

4: Update dataset: D ( t ) h ←D ( t -1) h ∪ {( x ( t ) h , a ( t ) h , x ( t ) h +1 )} ∀ h ∈ [ H ] .

5: Compute confidence set:

6: end for

7: Output ̂ π = Unif ( π (1: T ) ) .

Proof of Lemma F.3. For each f h +1 ∈ F h +1 , let apx [ f h ] = arg min g h ∈G h sup π ∈ Π E π [ ( g h -T h f h +1 ) 2 ] . Let

and note that by Jensen's inequality we have that for all π , E π [ δ ( t ) h ( · , · ) ] ≤ E π [ ˜ δ ( t ) h ( · , · ) ] + ε apx .

We further adopt the shorthand d ( t ) h ( x, a ) := d π ( t ) h ( x, a ) and ˜ d ( t ) h ( x, a ) := ∑ i<t d ( t ) h ( x, a ) . As a consequence of realizability ( Q ⋆ obs ,h ∈ F h ) and approximate Bellman completeness, standard concentration arguments (proved in the sequel) lead to the following result.

Lemma F.4 (Optimism and small in-sample squared Bellman errors) . With probability at least 1 -δ , by taking β = c log( TH |F||G| /δ ) + Tε apx , we have that for all t ∈ [ T ] ,

The rest of the proof proceeds similarly to the analysis of Section 3.2 in Xie et al. [XFBJK23]. Namely, by optimism (Lemma F.4) and a standard Bellman error decomposition (Lemma C.6) we have

Let us defining the burn-in time

where µ ⋆ h is the coverability distribution for the set of policies Π F (i.e., the distribution µ ⋆ h that achieves the minimum in the coverability definition). Using the same decomposition into 'burn-in phase' and 'stable phase' in Xie et al. [XFBJK23], we have:

Applying a change of measure argument on the second term then gives:

By the same reasoning as in Xie et al. [XFBJK23], we have (A) ≤ O ( √ C cov log( T )) , and by Lemma F.4 we have (B) ≤ O ( √ βT ) . Using that β = log( TH |F| /δ ) + Tε 2 apx gives the desired result. It remains to establish the concentration results of Lemma F.4.

Proof of Lemma F.4. For any function f , define a random variable

Next, note that

By Freedman's inequality (Lemma C.2, Lemma C.3), we have that with probability at least 1 -δ :

Taking a union bound over [ T ] × [ H ] ×F , we have that for all t, h, f , with probability at least 1 -δ :

where ι = log( |F| HT/δ ) . We now show that

which will imply, from Eq. (46), that

as desired. To see Eq. (47), let

and then note that:

where the second-to-last line follows from apx [ T h f ( t ) h +1 ] ∈ G and the last line follows from the definition of the confidence set. It remains to show that ∆ t ≤ O ( Tε 2 apx + ι ) , which we do via a similar concentration argument. Namely, let

and note that, as before,

and

by the same calculation as earlier. Thus, by Freedman's inequality and a union bound, we have that, with probability at least 1 -δ ,

where ι = log( |F| HT/δ ) . Recalling the misspecification assumption, this implies that

for all h, f, t , with high probability. This concludes the result for ( ii ) . For ( i ) , this follows identically to the proof of Lemma 40 in Jin et al. [JLM21], since this only uses the property that Q ⋆ ∈ F .

F.2.3 Sample-efficient latent-dynamics RL under pushforward coverability

We conclude by combining the previous two results to obtain the main result for this section.

Theorem 3.2 (Pushforward-coverable MDPs are statistically modular) . Let M lat be a base MDP class such that each M lat ∈ M lat has pushforward coverability bounded by C push ( M lat ) ≤ C push . Then, for any decoder class Φ , we have:

Proof of Theorem 3.2. Let M ⋆ obs := ⟪ M ⋆ lat , ψ ⋆ ⟫ ∈ ⟪ M lat , Φ ⟫ be the unknown latent-dynamics MDP. Define observation-level value functions

so that Q M ⋆ obs ,⋆ = Q M ⋆ lat ,⋆ ◦ ϕ ⋆ ∈ F via decoder and model realizability, and log |F h | ≤ log |M lat || Φ | . Consider any function class L ⊆ {S → [0 , 1] } and MDP M lat = ( r lat , P lat ) . For a given value ε apx > 0 , setting d according to Lemma F.1 implies that there exists a d -dimensional feature map φ M lat ,h ( s, a ) ∈ R d +1 such that for all ℓ ∈ L and h ∈ [ H ] , there exists w ℓ,h ∈ R d +1 such that

where µ M lat is the pushforward coverability distribution for M lat . Moreover, the map φ h is explicitly computed as a function of M lat by a randomized algorithm with success probability 1 -δ , with no knowledge of the class L required. We consider the class

where Γ ϕ : S → ∆( S ) is the mismatch function for decoder ϕ and emission ψ ⋆ , defined in Definition F.1. Note that L has size log |L| ≤ log |M lat || Φ | , and that we have

by Lemma D.7. By Lemma D.1 we have that µ M ⋆ obs ,h ( x ) = ψ ⋆ h ( x | ϕ ⋆ h ( x )) µ M ⋆ lat ,h ( ϕ ⋆ h ( x )) is the coverability distribution for MDP M ⋆ obs , and

Now, define

Recall the definition of w f (for f : S × A → [0 , 1] ) from Lemma F.1, and note that by the norm bound max ℓ ∈L ∥ w ℓ ∥ 2 2 ≤ 11 + 16 log( |M lat || Φ | H ) given by Lemma F.1, we have ( x, a ) ↦→ ⟨ φ M lat ,h ( ϕ ( x ) , a ) , w ℓ ⟩ ∈ G h for every ℓ ∈ L . Next, note that by the norm bound max s,a ∥ ψ ( s, a ) ∥ 2 2 ≤ C push (11 + 16 log( |S||A| H )) , given by Lemma F.1, we have every g h ∈ G M lat ,h satisfies ∥ g h ∥ ∞ ≤ cC 1 / 2 push log( |M lat || Φ ||S||A| H ) := B for some absolute constant c . Therefore, G M lat ,h has size log |G M lat ,h | ≤ ˜ O ( d · log( B ) + log | Φ | ) = ˜ O ( d log log( |S| ) + log | Φ | ) , where the ˜ O notation ignores logarithmic factors of C push , |A| , log |M lat | , and log | Φ | . 22 Define G h = ∪ M lat ∈M lat G M lat ,h , which has size log |G h | ≤ log | M lat | +( ˜ O ( d log log( |S| ) + log | Φ | )) . Together, these results with Lemma F.1 imply that for all f h +1 ∈ F h +1 , there exists g h ∈ G h such that

This, in turn, implies that for all π obs ∈ Π rns we have

since µ M ⋆ obs ,h ⊗ Unif ( A ) satisfies coverability (Definition D.3) with parameter C cov ( M ⋆ obs , Π rns ) ≤ C push |A| (Eq. (29)).

Then, it follows by Lemma F.3 that if we run Algorithm 2 with the classes F and G we will get

Choosing ε apx = 1 √ T to balance leads to

which gives a risk bound of

Equating this to ε gives a sample complexity of

as desired. Note that we have not made much effort to optimize the rate; in particular, a faster rate is likely possible by using the GOLF.DBR algorithm of Amortila et al. [ACK24], which improves over the GOLF algorithm under the presence of misspecification.

G Proofs and Additional Information for Section 4.1: Hindsight RL

This appendix contains additional information and proofs related to algorithmic modularity under hindsight observations (Section 4.1), and is organized as follows:

Appendix G.1 contains the pseudocode and proofs related to the online representation learning oracle EXPWEIGHTS.DR (Lemma 4.1).

Appendix G.2 contains the proof for our risk bound of the O2L algorithm under hindsight observability (Theorem 4.1).

G.1 Pseudocode and Proofs for EXPWEIGHTS.DR (Lemma 4.1)

Algorithm 3 Derandomized Exponential Weights (EXPWEIGHTS.DR)

̸

The main result for this estimator is the following.

Lemma 4.1 (Online classification via EXPWEIGHTS.DR) . Under decoder realizability ( ϕ ⋆ ∈ Φ) , EXPWEIGHTS.DR (Algorithm 3) satisfies Assumption 4.2 with 23

Proof of Lemma 4.1. For each h ∈ [ H ] , consider the realizable online classification problem where x ( t ) h ∼ d π ( t ) h , for π ( t ) chosen adversarially, and y ( t ) h = ϕ ⋆ h ( x ( t ) h ) . Consider the exponential weights estimator

̸

For every sequence ( x ( t ) h ) T t =1 , these distributions satisfy the deterministic regret bound

̸

by Corollary 2.3 of Cesa-Bianchi et al. [CBL06]. Taking conditional expectations over x ( t ) h ∼ d π ( t ) h and using Lemma C.3 gives that with probability at least 1 -δ :

̸

Taking a union bound over h ∈ [ H ] and summing over h ∈ [ H ] we obtain that with probability at least 1 -δ :

̸

Now, recall that at each time t , we define the improper decoder ¯ ϕ ( t ) h via:

̸

Let ℓ h ( x h , q ( t ) h ) = P ϕ ( t ) h ∼ q ( t ) h ( ϕ ( t ) h ( x h ) = ϕ ⋆ h ( x h )) . Note that ℓ satisfies

̸

̸

̸

By abuse of notation we also denote ℓ h ( x h , ¯ ϕ h ) = I [ ¯ ϕ h ( x ) = ϕ ⋆ ( x ) ] . We will show that

from which we will obtain that with probability at least 1 -δ :

̸

Integrating the high-probability regret bound gives

as desired. Towards establishing Eq. (56), let us fix x and let s max denote the argmax in Eq. (53). There are two cases:

P ϕ ( t ) h ∼ q ( t ) h ( ϕ ( t ) h ( x ) = s max ) ≥ 1 2 :

→ If s max = ϕ ⋆ ( x ) , ℓ ( x, ¯ ϕ ( t ) h ) = 0 so we are done.

̸

→ Otherwise, s max = ϕ ⋆ ( x ) and we have ℓ ( x, ¯ ϕ ( t ) h ) = 1 . However, since ϕ ⋆ ( x ) = s max we have ϕ ( t ) h ( x ) = s max = ⇒ ϕ ( t ) h ( x ) = ϕ ⋆ h ( x ) and so

̸

̸

P ϕ ( t ) h ∼ q ( t ) h ( ϕ ( t ) h ( x ) = s max ) < 1 2 :

→ If s max = ϕ ⋆ h ( x ) , ℓ ( x, ¯ ϕ ( t ) h ) = 0 so we are done.

̸

→ Otherwise, s max = ϕ ⋆ ( x ) and we have ℓ ( x, ¯ ϕ ( t ) h ) = 1 . However, by definition of s max as the mode we also have

so in particular we have

̸

G.2 Proofs for O2L Under Hindsight Observability (Theorem 4.1)

Theorem 4.1 (Risk bound for O2L under hindsight observability) . Let ALG lat be a base algorithm with base risk Risk ⋆ ( K ) , and REP class a representation learning oracle satisfying Assumption 4.2. Then Algorithm 1, with inputs T, K, Φ , REP class , and ALG lat , has expected risk

̸

Proof of Theorem 4.1. Let ( ̂ ϕ ( t ) ) t ∈ [ T ] denote the decoders chosen by REP class , and let ρ ( t ) denote the distribution over decoders induced at time t from the interaction of REP class , ALG lat , and M ⋆ obs . Let π ( t,k ) obs := π ( t,k ) lat ◦ ̂ ϕ ( t ) and p ( t,k ) obs denote the distribution over (observation-space) policies played at epoch t and episode k , induced by the interaction of REP class , ALG lat , and M ⋆ obs . We adopt the notation π ( t,K +1) lat := ̂ π ( t ) lat ∼ p ( t,K +1) lat for the final policy output by ALG lat in epoch t and ( x ( t,K +1) h , a ( t,K +1) h , r ( t,K +1) h ) for the trajectory collected from that (observation-level) policy ̂ π ( t ) lat ◦ ̂ ϕ ( t ) . We firstly note that by assumption, we have the guarantee

̸

which follows by applying Assumption 4.2 to the distributions ¯ p ( t ) obs = 1 ( K +1) ∑ K +1 k =1 p ( t,k ) obs and noting that

̸

̸

Let Risk ( K, ALG lat , ϕ, M ⋆ obs ) = J M ⋆ obs ( π ⋆ M ⋆ obs ) -J M ⋆ obs ( ̂ π lat ◦ ϕ ) be the random variable denoting the risk of the final policy output by ALG lat after K rounds of interaction with M ⋆ obs when given feature ϕ in any epoch t . For any ϕ : X → S , let E ϕ denote the law over trajectories ( x ( k ) h , a ( k ) h , r ( k ) h ) k ∈ [ K +1] ,h ∈ [ H ] and policies ( π ( k ) lat ◦ ϕ ) k ∈ [ K +1] generated after K rounds of interaction when ALG lat is given feature ϕ in any epoch. (Recall that, for all of the above definitions, a new instance of ALG lat is initialized at every epoch, so we do not have to specify which epoch it is , only the current feature ϕ ). Finally, let G t be the 'good' event

Recall that, in any round t , ALG lat only observes the latent ('compressed') trajectories ( ̂ ϕ ( t ) h ( x ( t,k ) h ) , a ( t,k ) h , r ( t,k ) h ) as history for choosing policies. We can therefore conclude that, when ̂ ϕ ( t ) ( x ( t,k ) h ) = ϕ ⋆ ( x ( t,k ) h ) for all k ∈ [ K +1] , h ∈ [ H ] , the distribution over final policies ̂ π ( t ) lat chosen by ALG lat will be identical as if we had chosen ϕ ⋆ as our decoder. In particular, this implies

where the second line simply follows by removing the indicator function, recalling that Risk ⋆ ( K ) = E [ Risk ( K, ALG lat , M ⋆ lat )] , and using that Risk ( K, ALG lat , ϕ ⋆ , M ⋆ obs ) = Risk ( K, ALG lat , M ⋆ lat ) .

Then, we have:

where the first equality applies the tower rule for conditional expectation, the second equality applies linearity of conditional expectations and the upper bound Risk ( K, ALG lat , ̂ ϕ ( t ) , M ⋆ obs ) ≤ 1 , and the third lines applies the upper bound Eq. (58). It remains to bound the last term. Here, note that by a union bound,

̸

where we have used that trajectory k in round t is sampled from policy π ( t,k ) , which is in turn sampled from p ( t,k ) . Summing over t and using the bound in Eq. (57) concludes the proof.

H Proofs for Appendix A: Self-Predictive Estimation

This appendix contains additional information and proofs related to algorithmic modularity under self-predictive estimation (Appendix A), and is organized as follows:

Appendix H.1 contains the pseudocode and proofs related to the online representation learning oracle SELFPREDICT.OPT (Lemma A.1).

Appendix H.2 contains the proof for our risk bound of the O2L algorithm under self-predictive estimation (Theorem A.1).

H.1 Pseudocode and Proofs for SELFPREDICT.OPT (Lemma A.1)

The pseudocode for our self-predictive estimation procedure is given in Algorithm 4.

Algorithm 4 Optimistic Self-Predictive Latent Model Estimation (SELFPREDICT.OPT)

1: input : Decoder set Φ , Latent model class M lat , Mismatch-complete class L lat , Optimism parameter γ

2: Set β := 1 2 √ C cov H log( T ) / T

3: for t = 1 , 2 , · · · , T do

4: Get dataset D ( t ) = { x ( i ) h , a ( i ) h , r ( i ) h , x ( i ) h +1 } i ∈ [ t -1] ,h ∈ [ H ]

5: Compute

6: Return ̂ ϕ ( t ) = { ̂ ϕ ( t ) h } h ∈ [ H ] .

7: end for

Our main result concerning the SELFPREDICT.OPT estimator for online optimistic self-predictive estimation is the following. We recall our notation for the instantaneous self-prediction error

Lemma A.1 (Optimistic self-predictive estimation via SELFPREDICT.OPT) . Let Π lat denote the set of policies played by ALG lat , and C cov , st = C cov , st ( M ⋆ lat , Γ Φ ◦ Π lat ) be the state coverability parameter on M ⋆ lat over the set of stochastic policies Γ Φ ◦ Π lat (Eq. (9)). Then, for any γ > 0 , under decoder realizability ( ϕ ⋆ ∈ Φ) , base model realizability ( M ⋆ lat ∈ M lat ) , and mismatch function completeness with class L lat (Assumption A.2), the estimator in Algorithm 4 with inputs Φ , M lat , L lat , and γ satisfies Assumption A.1 with 24

Proof of Lemma A.1. We will firstly establish that the algorithm obtains low offline estimation error.

Lemma H.1 (SELFPREDICT.OPT attains low offline estimation error) . For any γ > 0 , under decoder realizability ( ϕ ⋆ ∈ Φ ), model realizability ( M ⋆ lat ∈ M lat ), and mismatch function completeness with class L lat (Assumption A.2), the estimator in Algorithm 4 with inputs Φ , M lat , L lat , and γ satisfies

that for all t ∈ [ T ] , with probability at least 1 -δ ,

Given this result, we can appeal to offline-to-online conversions to establish the final result. Let C cov := C cov ( M ⋆ obs , Π lat ◦ Φ) denote the (state-action) coverability coefficient in M ⋆ obs over the set of policies Π lat ◦ Φ . Note that by Lemma D.1 we have C cov , st ( M ⋆ obs , Π lat ◦ Φ) = C cov , st and therefore by Lemma D.4 we have C cov ( M ⋆ obs , Π lat ◦ Φ) ≤ C cov , st |A| . Let η > 0 be a parameter to be chosen later, and β off = O ( log ( |M lat ||L lat || Φ | HTδ -1 )) be the offline estimation error guaranteed by Lemma H.1. We abbreviate α := √ C cov H log( T ) , E p ( t ) [ · ] := E π ( t ) ∼ p ( t ) E π ( t ) [ · ] , and E ˜ p ( t ) := ∑ t -1 i =1 E π ( i ) ∼ p ( i ) E π ( i ) [ · ] . Then, we have:

where in the first inequality we have used Lemma C.7 with g ( t ) h = ∆ h ( ̂ M ( t ) , ̂ ϕ ( t ) ) and in the second inequality we have used the AM-GM inequality with parameter η . Collecting terms, we proceed via:

where in the first inequality we have used Lemma H.1 and the definition of γ in Algorithm 4 (cf. Eq. (59)) and in the second inequality we have chosen η = 1 / √ T to balance the terms and used the bound C cov ≤ C cov , st |A| . We convert to an expected regret bound by picking δ appropriately, which gives the final result. It remains to show Lemma H.1.

Proof of Lemma H.1. Fix an iteration t ∈ [ T ] , and abbreviate ̂ M := ̂ M ( t ) and ̂ ϕ := ̂ ϕ ( t ) . We follow the analysis of maximum likelihood estimation from Geer; Zhang; Agarwal et al. [Gee00; Zha06; AKKS20]. In particular, we quote Lemma 24 of [AKKS20], which in an abstract conditional estimation framework with density class F states the following.

LemmaH.2 (Lemma 24 of Agarwal et al. [AKKS20]) . Let D = { ( x i , y i ) } be a dataset collected with x i ∼ p ( i ) ( x 1: i -1 , y 1: i -1 ) and y i ∼ f ⋆ ( · | x i ) , L ( f, D ) = ∑ n i =1 ℓ ( f, ( x i , y i )) be any loss function that decomposes additively, ̂ f : D →F be an estimator, D ′ be a tangent sequence D ′ = { ( ˜ x i , ˜ y i ) }

sampled independently via ˜ x i ∼ p ( i ) ( x 1: i -1 , y 1: i -1 ) and ˜ y i ∼ f ⋆ ( · | ˜ x i ) . Then, with probability at least 1 -δ , we have

For our purposes, we have that F = M lat ◦ Φ , the data distribution is collected adaptively (for each h ∈ [ H ] ) via π ( i ) ∼ p ( i ) , x ( i ) h , a ( i ) h ∼ d M ⋆ obs ,π ( i ) h , and r ( i ) h , x ( i ) h +1 ∼ M ⋆ obs ( · | x ( i ) h , a ( i ) h ) . For the loss function L , we take

We begin by upper bounding the quantity -L (( ̂ M, ̂ ϕ )( D ) , D ) appearing on the right-hand side of Eq. (62), or equivalently lower bounding L (( ̂ M, ̂ ϕ )( D ) , D ) . Let us abbreviate ̂ V = J ̂ M ( π ̂ M ) and V ⋆ = J M ⋆ lat ( π M ⋆ lat ) . Towards this, note that

where in the second line we have used Lemma D.8 with Assumption A.2 and in the third line we have used the ERM property of ̂ M ◦ ̂ ϕ together with decoder and model realizability. We claim that this implies

by concentration. Indeed, for each h ∈ [ H ] , i ∈ [ t ] , and [ M ′ ◦ ϕ ′ ] ∈ L lat ◦ Φ , let

Applying Lemma C.1, we have that

with probability at least 1 -δ , where we have recalled that data is gathered adaptively according to π ( i ) ∼ p ( i ) .We now quote the following lemma from Zhang; Agarwal et al. [Zha06; AKKS20].

Lemma H.3 (Lemma 25 of Agarwal et al. [AKKS20]) . For any D ∈ ∆( X ) and p, q ∈ [ X → ∆( Y )] , we have

Proof of Lemma H.3. We include the proof for completeness. The result follows via the following steps.

By Lemma H.3, we have that the right-hand-side of Eq. (64) is further lower bounded by

where the last line follows from the non-negativity of squared Hellinger. Taking a union bound over M ′ ◦ ϕ ′ ∈ L lat ◦ Φ and h ∈ [ H ] gives the desired lower bound in Eq. (63).

To conclude the proof, it remains to lower bound the left-hand side in Eq. (62). Here, note that:

where we have used that in the 'tangent sequence' D ′ the current sample ( ˜ r ( i ) h , ˜ x ( i ) h +1 ) is independent of ( r ( i ) h , x ( i ) h +1 ) . To bound this term, we again appeal to Lemma H.3, concluding that

Combining everything, we have:

Taking an additional union bound over t ∈ [ T ] , we have that with probability at least 1 -δ :

for all t ∈ [ T ] , as desired.

Corollary A.1 (Algorithmic modularity via SELFPREDICT.OPT) . Under the same conditions as in Lemma A.1, and for any base algorithm ALG lat , O2L with inputs T, K, Φ , SELFPREDICT.OPT , and ALG lat achieves

for absolute constants c 1 , c 2 , c 3 . Consequently, for any ALG lat with base risk Riskbase ( K ) , setting γ and T appropriately gives

with a number of trajectories TK = ˜ O ( K 5 H 3 C cov , st |A| log 2 ( |M lat ||L lat || Φ | ) / ( Riskbase ( K )) 4 ) .

Proof of Corollary A.1. The first inequality simply follows by plugging the bound of Estself ; opt from Lemma A.1 into Theorem A.1. For the second inequality, let ∆ = c 2 √ HC cov , st |A| log( |M lat ||L lat || Φ | ) . The result follows by setting γ s.t. c 3 γ -1 HK = Riskbase ( K ) i.e. γ = c 3 KH Riskbase ( K ) , and T such that γK ∆ √ T = Riskbase ( K ) i.e. T = K 4 ∆ 2 γ 2 Riskbase ( K ) 2 = K 4 ∆ 2 H 2 ( Riskbase ( K )) 4 . Then the result follows by direct substitution and by noting that K T ≤ 1 since Riskbase ( K ) ≤ 1 .

H.2 Proofs for Main Risk Bound (Theorem A.1)

Our main risk bound (Theorem A.1) follows as a special case of a more general theorem (Theorem H.1), which holds for algorithm that satisfies a property we refer to as CorruptionRobust -ness (Definition I.2). We now state the more general theorem, postponing its proof (and a formal definition of corruption robustness) until Appendix I.

Theorem H.1 (Risk bound for O2L under self-predictive estimation and CorruptionRobustness ) . Assume REP self ; opt satisfies Assumption A.1 with parameter γ > 0 and that M lat is realizable (i.e. M ⋆ lat ∈ M lat ). Furthermore, let ALG lat be CorruptionRobust (Definition I.2) with parameter α . Then, O2L (Algorithm 1) with inputs T, K, Φ , ALG lat , and REP self ; opt has expected risk

for absolute constants c 1 , c 2 , c 3 > 0 .

Our main risk bound (Theorem A.1) follows from the following lemma, which establishes that any ALG lat is CorruptionRobust in the sense of Definition I.2 for a sufficiently large corruption robustness parameter. Below, for any POMDP ˜ M over state-action space S × A , we write ˜ M ( s 1: h , a 1: h ) for the conditional probability over reward r h and s h +1 given s 1: h , a 1: h , i.e. ˜ M h ( s 1: h , a 1: h ) = ˜ M h ( r h , s h +1 = · | s 1: h , a 1: h ) .

Lemma H.4. Let M ⋆ be any reference MDP and ˜ M be any POMDP with the same state and action space. Then for any algorithm ALG lat , we have

where c 1 , c 2 > 0 are absolute constants. In particular, ALG lat is CorruptionRobust (Definition I.2) with α = c 2 √ KH .

Proof of Lemma H.4. Let us abbreviate ALG := ALG lat . For i ∈ [ K ] , let τ ( i ) denote the trajectory ( s ( i ) 1 , a ( i ) 1 , r ( i ) 1 , . . . , s ( i ) H , a ( i ) H , r ( i ) H ) . Let P := P M ⋆ , ALG denote the law of { ( π ( i ) , τ ( i ) ) } i ∈ [ K ] under ALG in the true MDP M ⋆ , and Q := P ˜ M, ALG denote the law of { ( π ( i ) , τ ( i ) ) } i ∈ [ K ] under ALG under the POMDP ˜ M . Let us write M ⋆ ( π ) and ˜ M ( π ) for the laws of trajectory τ sampled from policy π in M ⋆ or ˜ M respectively. Let ̂ π denote the policy output by the algorithm after K rounds of interaction with the environment. By Lemma C.5 we have

By the subadditivity property for squared Hellinger distance (Lemma C.4) applied to the sequence π (1) , τ (1) , . . . , π ( K ) , τ ( K ) , we have

where in the second step we have used that P ( π ( k ) | π (1: k ) , τ (1: k -1) ) = Q ( π ( k ) | π (1: k ) , τ (1: k -1) ) since the histories are equivalent, in the third step we have used that the trajectories are generated by the MDP/PODMP M ⋆ and ˜ M , respectively, in the fourth step we have again applied the subadditivity property of the squared Hellinger distance (Lemma C.4) to the sequence ( s 1 , a 1 , r 1 , . . . , s H , a H , r H ) .

Theorem A.1 (Risk bound for O2L under self-predictive estimation) . Suppose REP self ; opt satisfies Assumption A.1 with parameter γ > 0 . Then Algorithm 1, with inputs T, K, Φ , REP self ; opt , and ALG lat has expected risk

for absolute constants c 1 , c 2 , c 3 > 0 .

Proof of Theorem A.1. This follows from Theorem H.1 as well as Lemma H.4, by taking α = c 2 √ KH and simplifying.

I Additional Results for Appendix A: Self-Predictive Estimation

This section contains a more general result for algorithmic modularity under self-predictive estimation (Theorem H.1), from which our main result is derived as a special case, along with associated background, applications, and proofs. This section is organized as follows.

Appendix I.1 presents: definitions for the ϕ -compressed POMDP and CorruptionRobust algorithms (Appendix I.1.1), statements for properties of the ϕ -compressed dynamics (Appendix I.1.2). The risk bound for O2L under self-predictive estimation and CorruptionRobustness (Theorem H.1) is given in Appendix I.1.3, and a statement that the GOLF algorithm is CorruptionRobust (Appendix I.1.4).

Appendix I.2 presents for the proofs for the properties of the ϕ -compressed POMDPs.

Appendix I.3 presents a proof for the risk bound of O2L under self-predictive estimation and CorruptionRobustness .

Appendix I.4 presents a proof that the GOLF algorithm is CorruptionRobust .

I.1 O2L with Self-predictive Estimation and CorruptionRobust Base Algorithms

I.1.1 Definitions: ϕ -compressed POMDP and CorruptionRobustness

Consider iteration k ∈ [ K ] of epoch t ∈ [ T ] within O2L. Suppose that REPLEARN has chosen decoder ϕ = ϕ ( t ) : X → S . Then, the latent algorithm has observed the data D ( t,k ) = { ϕ ( x ( t,k ) h ) , a ( t,k ) h , r ( t,k ) h , ϕ ( x ( t,k ) h +1 ) } collected from the preceding policies in the epoch: π ( t, 1) lat ◦ ϕ ( t ) , . . . , π ( t,k -1) lat ◦ ϕ ( t ) (Line 8). Due to possible inaccuracies in the decoder ϕ , the dataset D ( t,k ) may not be generated from a Markovian process and must instead be viewed as being generated from a PODMP, formally defined as follows.

Definition I.1 ( ϕ -compressed POMDP) . The ϕ -compressed POMDP ˜ M ⋆ ϕ induced by M ⋆ obs and ϕ is defined by:

Latent state space X

Action space A

Observation state space S

Latent reward functions R ⋆ obs ,h : X × A → [0 , 1]

Latent dynamics P ⋆ obs ,h : X × A → ∆( X )

(Deterministic) observation function O h : X → S defined by O h ( x ) = ϕ h ( x ) ,

Horizon H

Initial latent distribution P ⋆ obs ( x 0 | ∅ )

Note that the latent space for the POMDP is the observation space of the latent-dynamics MDP M ⋆ obs , and vice-versa; we adopt this terminology because-from the perspective of the base algorithm, the observations x h can be viewed as a Markovian (yet partially observed process) that generates the learned states ϕ ( x h ) on which the algorithm acts. We write ˜ P π lat ϕ := P ˜ M ⋆ ϕ ,π lat for the probability distribution over trajectories ( x h , s h , a h , r h ) h ∈ [ H ] in the ϕ -compressed POMDP when playing policy π lat : S× [ H ] → ∆( A ) , where x h ∈ X are the POMDP's latent states, s h ∈ S are the observed states, and a h ∈ A are the actions. We let ˜ E π lat ϕ := E ˜ M ⋆ ϕ ,π lat denote the corresponding expectation. We write ˜ P ϕ,h ( s h +1 | s 1: h , a 1: h ) = ˜ P π lat ϕ ( s h +1 | s 1: h , a 1: h ) and ˜ r ϕ,h ( r h | s 1: h , a 1: h ) = ˜ P π lat ϕ ( r h | s 1: h , a 1: h ) for the conditional distributions of next states and rewards given the first h state-action pairs, which are policy-independent. We also write ˜ M ⋆ ϕ ( r h , s h +1 | s 1: h , a 1: h ) = ˜ r ϕ,h ( r h | s 1: h , a 1: h ) ˜ P ϕ,h ( s h +1 | s 1: h , a 1: h ) for the joint one-step probability. We will abbreviate ˜ M ⋆ ϕ ( s 1: h , a 1: h ) := ˜ M ⋆ ϕ ( r h , s h +1 = · | s 1: h , a 1: h ) .

Note that for any π lat , ˜ P π lat ϕ,h ( s h +1 | s h , a h ) is a well-defined (Markovian, policy-dependent) probability kernel, which is equivalent to

Similarly, ˜ r π lat ϕ,h ( r h | s h , a h ) is a Markovian and policy-dependent reward distribution which is equivalent to

Finally, we let

denote the associated one-step model over joint rewards and transitions.

Our CorruptionRobustness condition asserts that the agent-when observing data from the ϕ ( t ) -compressed dynamics ˜ M ⋆ ϕ -attains a risk bound for M lat which is proportional to its risk when observing data from M lat itself, plus a term that captures the degree of misspecification between ˜ M ⋆ ϕ and M lat .

Definition I.2 ( CorruptionRobust algorithm) . We say that ALG lat is CorruptionRobust with parameters α and Riskbase if there exists a constant c 1 such that, for any ( ϕ, M lat ) ∈ Φ ×M lat , we have

where we recall the definition of the random variable Risk ( K, ALG lat , M lat ) from Eq. (1), the expectation E ˜ M ⋆ ϕ , ALG lat denotes the interaction protocol of ALG lat in the ϕ -compressed dynamics ˜ M ⋆ ϕ , and p ( k ) denotes the randomization distribution over latent policies that ALG lat plays.

I.1.2 Basic properties of the ϕ -compressed dynamics (Definition I.1)

We establish a number of basic properties for the ϕ -compressed POMDP and their relation to the self-prediction guarantee obtained by REP self ; opt . These properties are proved in Appendix I.2. Firstly, we have the following change-of-measure lemma:

Lemma I.1 (Change of measure lemma) . For any ϕ ∈ Φ , f ∈ [ S × A → [0 , 1]] , h ∈ [ H ] , and π lat ∈ [ S × [ H ] → ∆( A )] , we have:

The next lemma states that the kernels of the ϕ -compressed POMDP are well-approximated by the (Markovian) latent model fit by REP self ; opt . We recall the instantaneous self-prediction error

Lemma I.2 (Near-markovianity of the ϕ -compressed dynamics) . For any decoder ϕ , base model M lat , and policy π lat : S × [ H ] → ∆( A ) , we have:

Furthermore, we also have

A corollary is the following lemma establishing errors between expectations under M lat , the model estimated by REP self ; opt , and those under the ϕ -compressed POMDP ˜ M ⋆ ϕ .

Lemma I.3 (Simulation lemma) . For any latent model M lat with Markovian transition kernel { P lat ,h } h ∈ [ H ] , latent policy π lat : S × [ H ] → ∆( A ) , and decoder ϕ ∈ Φ , we have that for all f : S × A → [0 , 1] :

and thus for any sequence of policies π ( t ) lat , latent models M ( t ) lat , and decoders ϕ ( t ) , we have:

I.1.3 Risk bound for O2L under CorruptionRobustness

We state the main risk bound for O2L under self-predictive estimation and the above definition of corruption robustness.

Theorem H.1 (Risk bound for O2L under self-predictive estimation and CorruptionRobustness ) . Assume REP self ; opt satisfies Assumption A.1 with parameter γ > 0 and that M lat is realizable (i.e. M ⋆ lat ∈ M lat ). Furthermore, let ALG lat be CorruptionRobust (Definition I.2) with parameter α . Then, O2L (Algorithm 1) with inputs T, K, Φ , ALG lat , and REP self ; opt has expected risk

for absolute constants c 1 , c 2 , c 3 > 0 .

I.1.4 Examples of CorruptionRobust algorithms

In this section, we establish that the GOLF algorithm satisfies the CorruptionRobust definition (Definition I.2) with a parameter α ≈ K -1 / 2 . This improves upon the rate that would be obtained by invoking the generic guarantee in Lemma H.4. We expect that several other algorithms can be analyzed in a similar way, thereby leading to tight rates in the same fashion. We restate the pseudocode in Algorithm 5 for convenience.

Let M lat = ( r lat , P lat ) be given, and we let Q ⋆ lat := Q M lat ,⋆ , and T lat ,h f ( s, a ) := r lat ,h ( s, a ) + E s ′ ∼ P lat ,h ( s,a ) [ V f ( s ′ )] . We assume that the algorithm has a latent function class F alg which realizes Q ⋆ lat , as well as a helper class G alg which is T lat -complete for F alg .

Assumption I.1 ( T lat -completeness) . We have:

For our analysis of GOLF, it is most natural to quantify the corruption levels in the following way.

Assumption I.2 (Corruption levels of M lat and ˜ M ⋆ ϕ ) . Let ε 2 rep be such that, for any sequence of policies π ( k ) lat played by the algorithm when interacting with the ϕ -compressed POMDP, we have

Algorithm 5 GOLF [JLM21]

input: Function classes F and G , confidence width β > 0 . initialize: F (0) ←F , D (0) h ←∅ ∀ h ∈ [ H ] .

1: for episode t = 1 , 2 , . . . , T do

2: Select policy π ( t ) ← π f ( t ) , where f ( t ) := arg max f ∈F ( t -1) f ( x 1 , π f, 1 ( x 1 )) .

3: Execute π ( t ) for one episode and obtain trajectory ( x ( t ) 1 , a ( t ) 1 , r ( t ) 1 ) , . . . , ( x ( t ) H , a ( t ) H , r ( t ) H ) .

4: Update dataset: D ( t ) h ←D ( t -1) h ∪ {( x ( t ) h , a ( t ) h , x ( t ) h +1 )} ∀ h ∈ [ H ] .

5: Compute confidence set:

6: end for

7: Output ̂ π = Unif ( π (1: T ) ) .

We note that

by the data-processing inequality (cf. Eq. (90) and Eq. (80)) and the inequality ∥ p -q ∥ 2 tv ≤ D 2 H ( p, q ) , and thus a CorruptionRobustness bound in terms of ε rep implies a CorruptionRobustness bound in the sense of Definition I.2.

Theorem I.1 (Latent GOLF is CorruptionRobust ) . Under Assumption I.1 and Assumption I.2, Algorithm 5 with β = c ( log ( |F||G| KHδ -1 ) + ε rep ) , has regret

and consequently is CorruptionRobust (Definition I.2) with parameters

Corollary I.1 (GOLF applied in O2L) . Let us suppose that the appropriate assumptions for the estimator in Algorithm 4 to have regret bounded by Estself ( T, γ ) = O ( √ HC cov T log( C cov |M lat ||L lat || Φ | HT ) ) (Lemma A.1) hold. Then, we can take γ ≈ K -1 / 2 and T ≈ K 4 , and the bound Theorem H.1 gives an expected risk of ε with a number of trajectories TK = poly ( C cov , H, log |M lat | , log | Φ | , log |L lat | ) · 1 / ε 10 , improving over the 1 /ε 14 rate of the universal result (Corollary A.1).

I.2 Proofs for Appendix I.1.2: Properties of ϕ -compressed POMDPs

Lemma I.1 (Change of measure lemma) . For any ϕ ∈ Φ , f ∈ [ S × A → [0 , 1]] , h ∈ [ H ] , and π lat ∈ [ S × [ H ] → ∆( A )] , we have:

Proof of Lemma I.1. Recall that ˜ P π lat ϕ denotes the law of ( x h , s h , a h ) h ∈ [ H ] in the ϕ -compressed POMDP when playing policy π lat . For clarity, and to differentiate a random variable from its realization, in the proofs below we will use upper-case notation such as { S h = s h , A h = a h , X h = x h } to indicate realizations of random variables in the POMDP.

Let ˜ d π lat h ( s, a ) = ˜ P π lat ϕ ( S h = s, A h = a ) be the marginalized occupancy measure for in the ϕ -compressed POMDP ˜ M ⋆ ϕ . We write d π lat ◦ ϕ h := d M ⋆ obs ,π lat ◦ ϕ h . The left-hand side in Eq. (72) is equal to:

Meanwhile, the right-hand side is equal to:

So it only remains to show that, for each s ∈ S and a ∈ A , we have ˜ d π lat h ( s, a ) = ∑ x : ϕ ( x )= s d π lat ◦ ϕ h ( x, a ) . Firstly, note that it is enough to show that ∑ x h : ϕ ( x h )= s h d π lat ◦ ϕ h ( x h ) = ˜ d π lat h ( s h ) , since ˜ d π lat h ( s h , a h ) = ˜ d π lat h ( s h ) π lat ( a h | s h ) and ∑ x h : ϕ ( x h )= s h d π lat ◦ ϕ h ( x h , a h ) = ∑ x h : ϕ ( x h )= s h d π lat ◦ ϕ h ( x h ) π lat ( a h | ϕ ( x h )) = π lat ( a h | s h ) ∑ x h : ϕ ( x h )= s h d π lat ◦ ϕ h ( x h ) . Toward this, we have:

At the same time,

where in the second equality we have used the definition of the observation function s h = O ( x h ) = ϕ ( x h ) .

To conclude, it remains to show that for all h , we have:

We do this by induction. Again, note that it is sufficient to establish d π lat ◦ ϕ h ( x h ) = ˜ P π lat ϕ ( X h = x h ) . The case h = 1 is clear. For the general case, we have:

Lemma I.2 (Near-markovianity of the ϕ -compressed dynamics) . For any decoder ϕ , base model M lat , and policy π lat : S × [ H ] → ∆( A ) , we have:

Furthermore, we also have

Proof of Lemma I.2. We begin with the first event. Note that, for any π lat , the PODMP kernel ˜ M ⋆ ϕ,h ( r h , s h +1 = · | s 1: h , a 1: h ) can be written as:

where we have used ˜ M ( r h , s h +1 = · | s 1: h , a 1: h ) = ˜ P π lat ϕ ( r h , s h +1 = · | s 1: h , a 1: h ) , the law of total probability, and that x h , a h is a sufficient statistic for r h and s h +1 . We further note that

since s h +1 = O h +1 ( x h +1 ) = ϕ h +1 ( x h +1 ) is a deterministic function of x h +1 and r h , x h +1 ∼ M ⋆ obs ,h ( x h , a h ) . Thus, for a fixed h and t , and omitting the h indices on the decoder ϕ for cleanliness, the expectation in equation Eq. (73) becomes:

as desired. Summing over h ∈ [ H ] we obtain the desired bound. The bound Eq. (74) is a consequence of Eq. (73) and the data-processing inequality. Namely, using the definition of ˜ M ⋆,π lat ϕ,h from Eq. (71) and the joint convexity of the squared Hellinger distance we have:

Thus, we have

as desired.

Lemma I.3 (Simulation lemma) . For any latent model M lat with Markovian transition kernel { P lat ,h } h ∈ [ H ] , latent policy π lat : S × [ H ] → ∆( A ) , and decoder ϕ ∈ Φ , we have that for all f : S × A → [0 , 1] :

and thus for any sequence of policies π ( t ) lat , latent models M ( t ) lat , and decoders ϕ ( t ) , we have:

Proof of Lemma I.3. Firstly note that, from Lemma I.1, the left-hand-side of Eq. (75) is equivalent to

For any π lat : S × [ H ] → ∆( A ) , let d π lat lat ,h = d M lat ,π lat h denote the occupancy in M lat , and similarly for any π obs : X × [ H ] → ∆( A ) let d π obs obs ,h ( x h , a h ) = d M ⋆ obs ,π obs h ( x h , a h ) denote the occupancy in M ⋆ obs . We overload notation by letting d π lat ◦ ϕ obs ,h ( s, a ) := ∑ x : ϕ ( x )= s d π ◦ ϕ obs ,h ( x, a ) . We will establish the stronger result that

where the tv norm on the left-hand-side is over S × A . Note that this implies the desired bound on Eq. (81) by Holder's inequality. We prove this by induction over h . For the base case ( h = 0 ), we have:

For the general case, let us further overload notation by letting d π ◦ ϕ obs ,h ( s h ) = ∑ a h d π ◦ ϕ obs ,h ( s h , a h ) and P ⋆ obs ( s h | x h -1 , a h -1 ) = ϕ♯P ⋆ obs ( s h | x h -1 , a h -1 ) = ∑ x h : ϕ ( x h )= s h P ⋆ obs ( x h | x h -1 , a h -1 ) . Let us also abbreviate π := π lat . Firstly note that it is sufficient to establish the result for ∑ s h ∈S ∣ ∣ ∣ d π lat ,h ( s h ) -d π ◦ ϕ obs ,h ( s h ) ∣ ∣ ∣ , since

Below, all summations over s h (resp. x h ) with domain unspecified are over S (resp. X ), and likewise for summations over s h , a h or x h , a h . We have:

From which it follows that, for each h , we have:

I.3 Proofs for Appendix I.1.3: Risk Bound Under CorruptionRobustness (Theorem H.1)

Theorem H.1 (Risk bound for O2L under self-predictive estimation and CorruptionRobustness ) . Assume REP self ; opt satisfies Assumption A.1 with parameter γ > 0 and that M lat is realizable (i.e. M ⋆ lat ∈ M lat ). Furthermore, let ALG lat be CorruptionRobust (Definition I.2) with parameter α . Then, O2L (Algorithm 1) with inputs T, K, Φ , ALG lat , and REP self ; opt has expected risk

Proof of Theorem H.1. Let us write π ( t,K +1) lat = ̂ π ( t ) lat and, for any t, k ∈ [ T ] × [ K +1] , π ( t,k ) obs := π ( t,k ) lat ◦ ϕ ( t ) . Let p ( t,k ) obs denote the distributions of played policies π ( t,k ) obs induced by the interaction of ALG lat and REP self ; opt inside the O2L algorithm. Let us write the online sum of self-prediction errors as

Since the final output policy of O2L satisfies ̂ π lat = Unif ( ̂ π (1) lat , . . . , ̂ π ( T ) lat ) (Line 12), we have

We take the following decomposition on the risk

We will show that E [ ∑ T t =1 A t ] ≲ T Reg base ( K ) + α √ T E [ ε rep ] and that E [ ∑ T t =1 B t ] ≲ √ TH E [ ε rep ] , then return to the first term J M ⋆ lat ( π M ⋆ lat ) -J M ( t ) lat ( π M ( t ) lat ) at the end of the proof.

To bound E [ ∑ T t =1 A t ] , we note that

where the first line follows from the CorruptionRobust definition (Definition I.2), the second line follows from Lemma I.2, the third line follows by Cauchy-Schwartz, and the last line recalls the definition of ε rep from Eq. (83).

For the term ∑ T t =1 B t , for any π lat : S × [ H ] → ∆( A ) we let Q π lat lat ( t ) ,h = T M ( t ) lat h Q π lat lat ( t ) ,h +1 be the Q π lat function of the latent MDP M ( t ) lat . Note that

so it is enough to bound

Fix t and h , whose indexing we omit below for cleanliness. Note that, for any π lat : S× [ H ] → ∆( A ) , we have:

where the final line follows from two applications of the data-processing inequality (since M lat ,h ( r h , s h +1 | ϕ h ( x h ) , a h ) = R lat ,h ( r h | ϕ h ( x h ) , a h ) P lat ,h ( s h +1 | ϕ h ( x h ) , a h ) and ϕ h +1 ♯M ⋆ obs ,h ( r h , s h +1 | x h , a h ) = R ⋆ obs ,h ( r h | x h , a h ) ϕ h +1 ♯P ⋆ obs ,h ( s h +1 | x h , a h ) ) as well as the bound ∥ p -q ∥ 2 tv ≤ D 2 H ( p, q ) . Summing this over t, h and using a standard decomposition for

regret (Lemma C.6) gives:

√

4

THε

rep

.

Returning to the decomposition of Eq. (84) and combining everything gives:

where the second inequality follows by AM-GM applied to the middle term and the third inequality follows from: i) Jensen's inequality, ii) Assumption A.1 applied to the distributions ¯ p ( t ) obs = 1 K ∑ K k =1 p ( t,k ) obs , iii) the bound K +1 ≤ 2 K , and iv) the inequality ( x + y ) 2 ≤ 2( x 2 + y 2 ) .

I.4 Proofs for Appendix I.1.4: Examples of CorruptionRobust Algorithms

Theorem I.1 (Latent GOLF is CorruptionRobust ) . Under Assumption I.1 and Assumption I.2, Algorithm 5 with β = c ( log ( |F||G| KHδ -1 ) + ε rep ) , has regret

and consequently is CorruptionRobust (Definition I.2) with parameters

Proof of Theorem I.1. Recall that the agent is observing data from the ϕ -compressed POMDP ˜ M ⋆ ϕ , and thus the datasets are of the form D ( k ) h = D ( k ) ϕ,h = { ϕ ( x ( i ) h ) , a ( i ) h , r ( i ) h , ϕ ( x ( i ) h +1 ) } k -1 i =1 . For any π lat ∈ Π lat , we define

where ˜ r π lat ϕ,h and ˜ P π lat ϕ,h are the policy-dependent Markov operators defined in Eq. (67) and Eq. (69).

As a consequence, we observe the following misspecification guarantee for T lat .

Lemma I.4 (Misspecification guarantee for T lat ) .

Proof of Lemma I.4. Follows from Assumption I.2 and the definitions of ˜ T π ( k ) ϕ,h and T lat ,h .

We begin with the following lemmas, which will be proved in the sequel.

Lemma I.5 (Optimism) . For the choice of β in Theorem I.1, with probability at least 1 -δ , we have that for all k ∈ [ K ] :

Lemma I.6 (Small in-sample squared Bellman errors) . With probability at least 1 -δ , we have that for all k ∈ [ K ] , h ∈ [ H ] , and f ∈ F ( k ) :

Let us write π ( k ) obs := π ( k ) ◦ ϕ . Let us introduce the shorthand ˜ d ( k ) obs ,h := ∑ k -1 i =1 d π ( k ) obs obs ,h , where d π obs is the occupancy for M ⋆ obs , and also the burn-in time

Let us recall, from the analysis of [XFBJK23], that for any h ∈ [ H ] and f : S × A → [0 , 1] we have

as well as

Note that, by change of measure (Lemma I.1) and the misspecification guarantee (Lemma I.4), the second term is bounded by:

Turning to the first term, we have:

where we have used that, from Lemma I.6, we have:

This gives an upper bound on the regret of

Using that β = O ( log ( |F||G| HK δ ) + ε rep ) and simplifying gives

as desired. It only remains to establish the concentrations results.

Concentration analysis. We establish the concentration results of Lemma I.5 and Lemma I.6. Proof of Lemma I.6. Let

and thus that

Next, note that

By Freedman's inequality (Lemma C.2, Lemma C.3), we have that with probability at least 1 -δ :

Taking a union bound over [ K ] × [ H ] ×F , we have that for all k, h, f , with probability at least 1 -δ :

where ι = log( |F| HK/δ ) . We now show that

which will imply, from Eq. (96), that

as desired. To see Eq. (98), let

and then note that:

where the second-to-last line follows from T lat F ⊆ G and the last line follows from the definition of the confidence set. It remains to show that ∆ k ≤ O ( ε rep + ι ) , which we do via a similar concentration argument. Namely, let

and note that, as before,

and

by the same calculation as earlier. Thus, by Freedman's inequality and a union bound, we have that, with probability at least 1 -δ ,

where ι = log( |F| HK/δ ) . Recalling the misspecification assumption Lemma I.4, this implies that

for all h, f, k , with high probability. Applying this to f = f ( k ) concludes the result.

Proof of Lemma I.5. We use similar arguments to the preceding lemma. Let Q ⋆ lat ,h := Q ⋆ M lat ,h . The aim is to show that, for all h ∈ [ H ] , k ∈ [ K ] , g ∈ G , we have:

from which the conclusion will follow. We show that

and also that

For Eq. (101), note that

and that Var [ W t ( h, g ) | F t,h ] ≤ 16 E [ W t ( h, g ) | F t,h ] . By Freedman, this gives

or in other words

using the non-negativity of Eq. (103). For Eq. (102), note that

and that Var [ V t ( h ) | F t,h ] ≤ 16 ˜ E π ( t ) ϕ [ ( T lat ,h Q ⋆ lat ,h -˜ T π ( t ) ϕ,h Q ⋆ lat ,h ) 2 ] . By Freedman, this gives

or in other words

where we have used Eq. (104).

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: All results in this paper are of a theoretical nature, and the stated contributions in the abstract and introduction are given in a precise, formal language.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: All results in this paper are of a theoretical nature - we precisely state the conditions under which our results hold.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Each theoretical result is stated with all necessary assumptions and is accompanied by complete (and correct) proofs.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification: The paper does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: The paper does not include experiments requiring code.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification: The paper does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: The paper does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification: The paper does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conforms with the NeurIPS Code of Etichs.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This is a primarily theoretical work.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This is a purely theoretical work, and as such poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: The paper does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsouring nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work investigates representation learning for decision-making in episodic MDPs. They consider a problem setting in which the goal is to through interactive decision-making learn a good policy. The task involves learning a good \"decoder\", i.e., function that maps the observations to \"state\" representations. They define a notion of \"statistical modularity\" in this problem which means that there exists an algorithm that can learn the optimal policy (with $\\epsilon$ error) with high probability with a number of episodes that is polynomial in the base MDP and the capacity of the decoder function class. They then prove an impossibility result regarding statistical modularity in this problem in general and prove statistical modularity in MDPs under some condition. They also define a notion of \"algorithmic modularity\" by introducing an algorithm in a hindsight observability setting in which one can use any decoder of interest and any standard episodic MDP algorithm for decision-making; they prove a regret bound in terms of the quality of the quality of the base MDP algorithm and the quality of the decoder.\n\nSTRENGTHS:\n- I thought the problem the authors worked on was interesting and meaningful. Overall the ideas of introducing the concepts of statistical and algorithmic modularity was interesting.\n- The authors provide many theoretical results and the result, especially regarding the algorithmic modularity result in 4.1 seems interesting, intuitive, and rather elegant.\n\nWEAKNESSES:\n- I found the notation and terminology used in this paper to be very dense. (I put specific questions / notes about this in the next section). I think I would prefer that the authors to have fewer theoretical results, but more thorough discussion of the results. (E.g., the self-predictive estimation idea seems interesting, but is literally 1 paragraph in the paper.)\n\nQUESTIONS:\n- Could you clarify if you are using the term \"latent states\" in the sense of partially observed MDPs. Or if you really mean the state is \"latent\" in the sense that the best representation of the state is unknown? If this is the case, it seems like this is more a representation learning problem rather than a \"latent state\" problem. Could you add further discussion of how your formulation relates to POMDPs?\n- It would be helpful to provide an intuitive definition of decoder earlier in the paper, as it is used in the intro without much context.\n\nLIMITATIONS:\n- There is no empirical evaluation of the algorithm.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a theoretical analysis of statistical and algorithmic modularity for RL with latent dynamics. Specifically, it offers conditions and theoretical analysis under which RL with latents is tractable. For statistical modularity, both lower and upper bounds are presented. For algorithmic modularity, observation-to-latent reductions are analyzed under two conditions: hindsight observability and self-predictive estimation. Overall, the theory and proofs are technically solid, addressing a critical problem in RL, especially in scenarios where only high-dimensional pixels are observed. Although I am not an expert in RL theory (my focus is more on algorithms and applications), I would give an acceptance rating for this initial review and will be engaged in the discussion.\n\nSTRENGTHS:\n- [**Motivation and Significance**]: The problem of learning from observation for RL is important, and this paper provides fundamental theory on this topic. The statistical and algorithmic guarantees are critical contributions to the field. The theoretical findings, particularly on algorithmic modularity, have the potential to encourage more empirical work on efficiently identifying self-predictive latent states that facilitate RL policy learning.\n- [**Technical Soundness**]: Although I am not an expert in RL theory, I reviewed the main paper thoroughly and found the theoretical foundations and proofs to be solid.\n- [**Presentation**]: The presentation is clear and accessible, even for readers outside the theory domain.\n\nWEAKNESSES:\nSince I am not an expert on RL theory, I have listed most of my questions in this section. The major question from an empirical point of view is how to leverage some of these theoretical results to enhance RL learning from latent dynamics.\n\nQ1: The authors mentioned that block MDP or factored MDP would be a special case of this general framework. Suppose we narrow down the problem to block MDP or factored MDP, can the statistical or algorithmic modularity be easier to achieve?\n\nQ2: Similar to the previous question, what kind of structure (e.g., symmetry, disentanglement) or distribution assumptions in the latent space could mostly benefit the current theoretical framework?\n\nQUESTIONS:\nI listed my questions in the above section.\n\nLIMITATIONS:\nLimitations and discussions are given in the paper. As this is a theoretical work, I do not think it will pose any negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers theoretical aspects of reinforcement learning in a certain class of MDPs whose observations are governed by a separate, potentially smaller, MDP. They formalize this class of MDPs and denote them latent MDPs. The authors then consider when such MDPs are statistically learnable, beginning with a negative result: they show that in general, even with known latent dynamics, statistical modularity is impossible. They then highlight that statistical modularity in this setting is in some sense distinct from previous works which assume regularity in the value function, and mention that this is because this structure might be useless without a good learnt representation. The authors then go through a laundry list of MDP formalisms in previous work, and provide for most a result on whether or not they are statistically modular. They finally consider algorithmic results, and introduce a 'meta-algorithm' which balances representation and RL learning (where the underlying RL algorithm is arbitrary). Under some additional assumptions, they prove that the additional representation learning adds sublinear risk.\n\nSTRENGTHS:\n- I believe that this is an important step in bridging RL theory and the issues of RL in practice. \n- Balancing representation learning and standard RL learning is an important issue many RL practitioners need to balance. This work paves the way for theoretically-guided answers to those questions. \n- I found it rather interesting how the authors demonstrated that much of the structure used in previous work is not amenable to this setting, and that in those cases statistical modularity is not possible.\n\nWEAKNESSES:\n- Assuming that latent states can be uniquely decoded from the observations is a rather strong observation.\n- There are no experiments in the paper. Of course the contribution of this work is theoretical, but theoretical work can still benefit greatly from some simple experiments which illustrate results in their paper. In particular, doing this allows some readers to better understand the result, and importantly, shows that the results obtained (which are often under unrealistic assumptions) do not break down in practice.\n\nQUESTIONS:\n- Do you believe there are any toy experiments which can be done to illustrate any of your results?\n\nLIMITATIONS:\nThey discuss avenues for future work, and are clear on the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-18"}